Using ndarray slices for element-wise operations

help
1

I want to replicate the behaviour of this nested for-loop using ndarray slices:

use ndarray::prelude::*;

// [ ... ]

     // self.hx and self.ex are Array3<f32> of the same shape
     let [xs, ys, zs] = self.hx.shape();

      for x in 0..*xs {
          for y in 1..*ys {
              for z in 0..*zs {
                  self.hx[[x, y, z]] -= self.ez[[x, y, z]] - self.ez[[x, y - 1, z]]
              }
          }
      }

I have landed on

    self.hx.slice_mut(s![.., 1.., ..]) -=
         self.ez.slice(s![.., 1.., ..]) - self.ez.slice(s![.., ..-1, ..]);

but I get the error

error[E0369]: cannot subtract `ArrayBase<ViewRepr<&f32>, Dim<[usize; 3]>>` from `ArrayBase<ViewRepr<&f32>, Dim<[usize; 3]>>`
    --> src/solver/fdtd.rs:81:44
     |
81   |             self.ez.slice(s![.., 1.., ..]) - self.ez.slice(s![.., ..-1, ..]);
     |             ------------------------------ ^ ------------------------------- ArrayBase<ViewRepr<&f32>, Dim<[usize; 3]>>
     |             |
     |             ArrayBase<ViewRepr<&f32>, Dim<[usize; 3]>>
     |

Where am I going wrong? Is there a more idiomatic way to do this?

As a side-question, would the slice approach allocate a temporary matrix? I would like to avoid creating unnecessary copies for performance reasons.

2

I think this is what you're looking for:

let mut slice = self.hx.slice_mut(s![.., 1.., ..]);
slice -= &self.ez.slice(s![.., 1.., ..]);
slice += &self.ez.slice(s![.., ..-1, ..]);

First, you can't have any expression on the left side of assignments, so you have to create the temporary slice. Then, you can't subtract one immutable view from another immutable view because there is nowhere to store the result. And finally, you need to add & because ndarray has only implemented the assignments for references.

3

Thanks for the help!

The reason I didn't go with doing two assignments (slice -= then slice +=) is that I want to avoid iterating through the large matrix twice.

In C++, I could use a library like Eigen that reduces expressions like this using templates, but I'm guessing ndarray does not operate in this way.
https://eigen.tuxfamily.org/dox/TopicLazyEvaluation.html

I need this to be as fast as possible. This means iterating through both matrices once, using vector instructions where possible. Is there a library in the rust ecosystem that would help me with this?

4

That’s not correct. I don't know ndarray, but there are no restrictions on the expression used on the left of an assignment operator other than that it is a place expression (or a few other things, for destructuring assignment). There might be other reasons like borrow checking errors to organize the code differently, but there is no restriction specifically on the complexity of the left side.

5

It's not a place expression. In other cases you can get a place expression out of a value expression by dereferencing, but that's not how ndarray implemented assignments.

Edit: I forgot you can make it a place expression but it's kind of ugly:

*&mut self.hx.slice_mut(s![.., 1.., ..]) -= &self.ez.slice(s![.., 1.., ..]);
*&mut self.hx.slice_mut(s![.., 1.., ..]) += &self.ez.slice(s![.., ..-1, ..]);
6

Is this better?

for ((hx_item, ez1), ez2) in hx
    .slice_mut(s![.., 1.., ..])
    .into_iter()
    .zip(ez.slice(s![.., 1.., ..]))
    .zip(ez.slice(s![.., ..-1, ..]))
{
    *hx_item -= ez1 - ez2
}
1 Like
7

Looks great!

8

Ah, right, I see what is going on (the returned value implements AddAssign but you can't just then start mutating it as an owned value). Sorry for the confusion. I wouldn’t have put it that way but my reply was at least equally confusing.

1 Like
9

ndarray has a macro azip for lock step iteration which reduces the visible zip nesting level

  azip!((h in hx.slice_mut(s![.., 1.., ..]), ez1 in ez.slice(s![.., 1.., ..]), ez2 in ez.slice(s![.., ..-1, ..]))
      *h -= ez1 - ez2
  );
2 Likes