Equivalent to np.eye(k=<something>)

[IndianBoy42]

trafficstars

How to make matrices with a constant off-diagonal?

https://numpy.org/devdocs/reference/generated/numpy.eye.html

[jturner314]

Here's one implementation, with an additional elem parameter, in case that's useful:

use ndarray::prelude::*;
use num_traits::Zero;

pub fn from_off_diag_elem<A>(shape: [usize; 2], k: isize, elem: A) -> Array2<A>
where
    A: Clone + Zero,
{
    let mut a = Array2::zeros(shape);
    if k == 0 {
        a.diag_mut().fill(elem);
    } else if k < 0 && ((-k) as usize) < shape[0] {
        a.slice_mut(s![-k.., ..]).diag_mut().fill(elem)
    } else if k > 0 && (k as usize) < shape[1] {
        a.slice_mut(s![.., k..]).diag_mut().fill(elem)
    } else {
        panic!("k={} is out-of-bounds for shape {:?}", k, shape);
    }
    a
}

fn main() {
    assert_eq!(
        from_off_diag_elem([5, 3], 1, 7.),
        array![
            [0., 7., 0.],
            [0., 0., 7.],
            [0., 0., 0.],
            [0., 0., 0.],
            [0., 0., 0.],
        ],
    );
    assert_eq!(
        from_off_diag_elem([5, 3], 0, 7.),
        array![
            [7., 0., 0.],
            [0., 7., 0.],
            [0., 0., 7.],
            [0., 0., 0.],
            [0., 0., 0.],
        ],
    );
    assert_eq!(
        from_off_diag_elem([5, 3], -2, 7.),
        array![
            [0., 0., 0.],
            [0., 0., 0.],
            [7., 0., 0.],
            [0., 7., 0.],
            [0., 0., 7.],
        ],
    );
    assert_eq!(
        from_off_diag_elem([3, 5], 1, 7.),
        array![
            [0., 7., 0., 0., 0.],
            [0., 0., 7., 0., 0.],
            [0., 0., 0., 7., 0.],
        ],
    );
    assert_eq!(
        from_off_diag_elem([3, 5], 0, 7.),
        array![
            [7., 0., 0., 0., 0.],
            [0., 7., 0., 0., 0.],
            [0., 0., 7., 0., 0.],
        ],
    );
    assert_eq!(
        from_off_diag_elem([3, 5], -2, 7.),
        array![
            [0., 0., 0., 0., 0.],
            [0., 0., 0., 0., 0.],
            [7., 0., 0., 0., 0.],
        ],
    );
    // should panic
    from_off_diag_elem([3, 5], -3, 7.);
}

If you want to duplicate NumPy's behavior, then replace the panic!(...) with // no-op.

This implementation doesn't handle the k == isize::MIN case very well. (The -k will overflow.) If you care about that case, then you should handle it more carefully than I did.