Roll array along an axis

[mokasin]

trafficstars

In numpy you can roll an array along an axis. Which is an index shift with wrap around, meaning, that given an array

a = np.arange(1, 10).reshape(3, 3)
[[1, 2, 3],
 [4, 5, 6],
 [7, 8, 9]]

rolling along the 2nd axis by 1 results in

np.roll(a, 1, 1)
[[3, 1, 2],
 [6, 4, 5],
 [9, 7, 8]]

I wonder if this can be cheaply implemented in this crate, since it is basically just a remapping of indices? It would be useful to calculate a gradient with periodic boundaries (wrap around). For example, a central difference (modulo prefactors) in y direction can be implemented like

np.roll(a, -1, 1) - np.roll(a, 1, 1)
[[-1,  2, -1],
 [-1,  2, -1],
 [-1,  2, -1]

[bluss]

ndarray only implements custom strides, no other remapping. Do you know a good data structure for implementing that kind of remapping? Numpy does not implement roll as a view-like object.

This method is basically laying dormant waiting to be used to implement roll along an axis https://github.com/bluss/rust-ndarray/blob/510a0433dcc33c5cf843c691ab560fceff29d61b/src/impl_methods.rs#L1372 (Edited later: Superseded by the Inners/InnersMut producers)

It's a simple way to implement it (You split the array into 1D lanes, and do the roll on each one in turn).

[bluss]

Here are two ways to do the roll by 1 along an axis. Either use the lane methods (if those are first made available by making them pub) or use a slice and assign. The latter is especially good if we use an uninitialized array to assign into.

#[bench]
fn roll_lanes_ax0(bench: &mut Bencher)
{
    let a = Array::linspace(0., 127., ROLL * ROLL).into_shape((ROLL, ROLL)).unwrap();
    bench.iter(|| {
        let mut b = a.clone();
        b.lanes_along_mut(Axis(0), |mut row| {
            let mut last = row[0];
            for elt in row.slice_mut(s![1..]) {
                last = replace(elt, last);
            }
            row[0] = last;
        });
        b
    });
}

#[bench]
fn roll_assign_ax0(bench: &mut Bencher)
{
    let a = Array::linspace(0., 127., ROLL * ROLL).into_shape((ROLL, ROLL)).unwrap();
    bench.iter(|| {
        let mut uninit = Vec::with_capacity(a.len());
        unsafe {
            uninit.set_len(a.len());
        }

        //let mut b = Array::zeros(a.dim());
        let mut b = Array::from_vec(uninit).into_shape(a.dim()).unwrap();
        b.slice_mut(s![1.., ..]).assign(&a.slice(s![..-1, ..]));
        b.slice_mut(s![..1, ..]).assign(&a.slice(s![-1.., ..]));
        b
    });
}

#[bench]
fn roll_lanes_ax1(bench: &mut Bencher)
{
    let a = Array::linspace(0., 127., ROLL * ROLL).into_shape((ROLL, ROLL)).unwrap();
    bench.iter(|| {
        let mut b = a.clone();
        b.lanes_along_mut(Axis(1), |mut row| {
            let mut last = row[0];
            for elt in row.slice_mut(s![1..]) {
                last = replace(elt, last);
            }
            row[0] = last;
        });
        b
    });
}

#[bench]
fn roll_assign_ax1(bench: &mut Bencher)
{
    let a = Array::linspace(0., 127., ROLL * ROLL).into_shape((ROLL, ROLL)).unwrap();
    bench.iter(|| {
        let mut uninit = Vec::with_capacity(a.len());
        unsafe {
            uninit.set_len(a.len());
        }

        let mut b = Array::from_vec(uninit).into_shape(a.dim()).unwrap();
        b.slice_mut(s![.., 1..]).assign(&a.slice(s![.., ..-1]));
        b.slice_mut(s![.., ..1]).assign(&a.slice(s![.., -1..]));
        b
    });
}

Using assign is the quicker way, and we also notice how going with the memory order is faster.

test roll_assign_ax0       ... bench:       1,385 ns/iter (+/- 35)
test roll_assign_ax1       ... bench:       1,761 ns/iter (+/- 58)
test roll_lanes_ax0        ... bench:       5,747 ns/iter (+/- 350)
test roll_lanes_ax1        ... bench:       3,820 ns/iter (+/- 161)

[mokasin]

Awesome! You're the best 👍

At the moment, I've implemented it like that

let mut res: Array<f64, Ix4> = Array::zeros((sx, sy));

// bulk
{
    let mut s = res.slice_mut(s![1..-1, ..]);
    s += &field.slice(s![2.., ..]);
    s -= &field.slice(s![..-2, ..]);
}
// borders
{
    let mut s = res.slice_mut(s![..1, ..]);
    s += &field.slice(s![1..2, ..]);
    s -= &field.slice(s![-1.., ..]);
}
{
    let mut s = res.slice_mut(s![-1..,..]);
    s += &field.slice(s![..1, ..]);
    s -= &field.slice(s![-2..-1, ..]);
}

which seems to be rather fast, but is nearly unreadable. I have to compare this to your solution.

[bluss]

Right, we have a bunch of extra noise because not only do we need to cut out a part of an array, the array we modify in place needs to be cut to the same size as well. We should see if there is something that can be improved around that in general.

[bluss]

A “rolled view” as equivalent to a pair of array views doesn't sound so bad.

[bluss]

Note that there is a different way to write this. It's not necessarily better, but we have an array view’s .split_at(axis, index) method to work with too.

[bluss]

There's even the question of numerical stability / accuracy. Ideally you'd want array += x - y right.

[mokasin]

For a zero initialized array, numerically I don't see the difference between

array += x - y

and

array += x
array -= y

is there? But memory wise, x - y will allocate an intermediate array for the result, before assigning, right? Wheres the second approach does not.

By the way, I don't completely understand, why I need to bind the mutable slice before using it. Why is array.slice_mut(...) += bar not working?

[bluss]

You are right, it's the same if the initial state is zero. With “ideally” array += x - y I was suggesting to compute that in a way that doesn't use the intermediate temporary. That will be possible soon since I'm finishing an n-ary zip (:wink: @SuperFluffy), but it doesn't have that nice syntax.

Good question about the += operator. For us it's annoying, it's the rule of the assign operators. The rule avoids code like string.len() += 1 to be accepted too (which makes sense, the integer there is a temporary that nobody holds).

[bluss]

This issue lead to the addition of https://bluss.github.io/rust-ndarray/master/ndarray/struct.ArrayBase.html#method.uninitialized (as can be seen in the example), since the use case warrants easy access to an uninitialized array.

[bluss]

This issue needs to be reformulated to be actionable, if it should be resolved.

• Ndarray does not offer an existing facility to implement this in “view like” fashion
• I could imagine you implement a new kind of object that has multiple “panes” of data and it knows how to apply operations on that. (Roll along one axis = composite of 2 panes)
  • How do we generalize this? Binary operation of 2 pane × 2 pane will need to be carried out over at most 2 × 2 = 4 subareas

[mokasin]

I was thinking about a RolledView, where the index is internally remapped. At least for the index operation this would be easy to do. So basically, something roughly like this

impl Index<usize> for RolledView<S> where S: Data {
    type Output: S;
    fn index(&self, index: usize) -> &Self::Output {
        let shifted_index = (index + self.shift) % self.shape.0;
        self.view[shifted_index]
    }
}

That won't work for more basic internal operations, that do not rely on the Index operator. But maybe this can be generalised?

[bluss]

Yes that's possible. None of the functionality in ndarray uses the Index impl, so it does not pay off with leverage at all unfortunately. Maybe another Dimension trait impl could do this. Some of the representation assumptions are rather deeply embedded. Now that ndarray has come quite far, it's time to clean it up a bit I think.

[johannes-graeter]

Whats the status on this, is there still no possible way to implement roll with ndarray::ArrayView?

[bluss]

@johannes-graeter the data model of ArrayView has not changed, and it doesn't allow making "rolled views"