How to perform in-place operations on parts of a vector provided an index array/vector?

[rcarson3]

First off, I just want to say that I'm still very much getting the hang of ndarray. So far, I understand the basics like how to perform matrix multiplication, slicing, and in place mapping functions. However, I've found the documentation to be at times a little hard to understand.

Now on to my actual question, so in Numpy or Fortran I can easily do something like the following:

vec = np.zeros((6,1))

index = [1, 3, 5]

mat = np.ones((3,3))
x = np.ones((3,1))

vec[index] = mat.dot(vec)

It's the being able to assign values to a vector or an array at these indices that I'm trying to get a better hold of using ndarray.

Currently, the only way I can really think of how to do this is the following:

let mut vec = Array::<f32,_>::zeros((3,3));

let mut mat = Array::<f32, _>::zeros((3,3));
mat.fill(1.0);

let mut x = Array::<f32, _>::zeros((3,1));
x.fill(1.0);

ind = vec!(1, 3, 5);

let temp = mat.dot(&x);

for (i, &index) = ind.iter().enumerate(){
      vec[[index, 0]] = temp[[i, 0]];
}

I could make this into function, but I'd like to believe there's a built in way to do this in ndarray that I'm just missing or have over looked that avoids the need for the temporary arrays/vectors creation.

[jturner314]

I'm not quite sure what your exact objectives are, so here are a few solutions depending on what you're trying to achieve. (Note that all these examples use f64 instead of f32; if you want f32, you can specify it when creating the arrays, e.g. Array::<f32>::zeros(6).)

This is the simplest example:

#[macro_use]
extern crate ndarray;

use ndarray::prelude::*;

fn main() {
    let mut vec = Array1::zeros(6);
    let mat = Array2::from_elem((3, 3), 1.);
    let x = Array1::from_elem(3, 1.);
    vec.slice_mut(s![1..;2]).assign(&mat.dot(&x));
}

Some notes:

• This treats vec and x as one-dimensional arrays (which is what I prefer when working with column vectors since with Array1 you get compile-time checking that they're one-dimensional). Dot products with one-dimensional arrays are the same in ndarray as they are in NumPy: 1D dot 1D is a vector dot product, 1D dot 2D is a row vector times a matrix, 2D dot 1D is a matrix times a column vector, and 2D dot 2D is a matrix product.
• The .from_elem() method is a little cleaner than ::zeros() followed by .fill(). A .ones() method was added a few days ago, but a new version of ndarray hasn't been released yet.
• Slicing provides a nice way to select the portion of vec where you want to assign the result. Note that it's limited to indices of the form start, start + step, start + 2*step, start + 3*step, …, not arbitrary "index arrays". AFAIK, the only support for "index arrays" at the moment is the .select() method.
• The s![] macro is used for the slice argument to express "start at index 1, then increment by 2 until reaching the end". (The s![a..b;c] notation is similar to a:b:c in NumPy but has slightly different behavior for negative c.)
• This approach does allocate a temporary array for the result of the dot before assigning it to vec. (By the way, I'm pretty sure NumPy does too in the vec[index] = mat.dot(vec) line.)

If you want a separate index object, you can create and use a Slice instance like this:

#[macro_use]
extern crate ndarray;

use ndarray::prelude::*;
use ndarray::Slice;

fn main() {
    let mut vec = Array1::zeros(6);
    let index = Slice::new(1, None, 2);
    let mat = Array2::from_elem((3, 3), 1.);
    let x = Array1::from_elem(3, 1.);
    vec.slice_mut(s![index]).assign(&mat.dot(&x));
}

If you want to use 2-D arrays for vec and x, then you can do this:

#[macro_use]
extern crate ndarray;

use ndarray::prelude::*;
use ndarray::Slice;

fn main() {
    let mut vec = Array2::zeros((6, 1));
    let index = Slice::new(1, None, 2);
    let mat = Array2::from_elem((3, 3), 1.);
    let x = Array2::from_elem((3, 1), 1.);
    vec.slice_mut(s![index, 0..1]).assign(&mat.dot(&x));
}

If you really want to avoid creating a temporary array for the result of the dot, you could split up the dot product into dot products of rows and columns (which evaluate to scalars), like this:

#[macro_use]
extern crate ndarray;

use ndarray::prelude::*;

fn main() {
    let mut vec = Array1::zeros(6);
    let mat = Array2::from_elem((3, 3), 1.);
    let x = Array1::from_elem(3, 1.);
    azip!(mut vec (vec.slice_mut(s![1..;2])), ref mat (mat.genrows()) in {
        *vec = mat.dot(&x);
    });
}

This iterates over the elements in vec.slice_mut(s![1..;2]) and the rows in mat in lockstep, assigning the result of the dot product of the mat row with x to the corresponding element in vec. Unfortunately, I'm not aware of a cleaner way at the moment if you really want to avoid a temporary allocation of the result of the dot. (I mean, you could iterate over indices or use Iterator::zip, but I think the azip! is cleaner.)

Does this answer your question?

By the way, do you have suggestions to improve the documentation? I've been using ndarray for a while, so I don't really have a good idea of what it's like to approach ndarray as a newcomer.

Edit: Rereading your question, it sounds like you might be specifically asking about NumPy's "index arrays" feature. Other than .select(), I think currently you'd have to use iteration for arbitrary index arrays. It would be possible to add an .assign_indices() method or something to ArrayBase, though.

[rcarson3]

I would say that this definitely answers my question. I was specifically asking about an equivalent to Numpy's index array feature. It looks like I'll need to use iterations for my use case.

Also as far as some of my objectives, I'm looking into implementing a couple of different krylov subspace methods in Rust. One of my use cases with these methods will have block submatrices that reside in your A matrix for Ax=b, and these could be for example elemental stiffness matrices in a finite element calculation. Often these submatrices will not be operating on contiguous sections of your x or b. In Fortran and Python, I am used to being able to use an index array such that these operations are relatively clean from a code point of view. Though, I guess I would need to dig into how gfortran handles this situation to see if it is also creating temporary arrays during compilation. Now that I'm thinking about it I could probably get away with just two temporary variables while iterating over 100k+ of these submatrices.

On the documentation side of things, I would say if its possible maybe having something similar to Numpy's Numpy for MATLAB users could be quite useful. Though, I would say instead of Matlab ndarray for Numpy users would probably hit a wider range of users. I could see this really be useful in helping newcomers getting up and starting to use ndarray relatively quickly. I can see this being a large undertaking, and I would be glad to help with this in some manner when I have some free time during the weekends to work on my Rust projects. Next, I found comments in the examples and test section to be really sparse. For me this is usually one of the first places I go to get an idea of how a code base works when I find them on Github or some other repository site. So, I could see this as a possible area for improvement as far as documentation goes. I might have more ideas on ways to improve it as continue working with ndarray. Though like I said earlier, I'd be glad to help out in some manner or shape when time permits.