CODEWITHSUNDEEP
pythonnumpysliding-window
Physics
Physics

Reputation: 15

setting the values of sliding windows of an array in numpy

Suppose I have a 2D array with shape (3, 3), call it a, and an array of zeros with shape (7, 7, 5, 5), call it b. I want to modify b in the following way:

for p in range(5):
    for q in range(5):
        b[p:p + 3, q:q + 3, p, q] = a

Given:

a = np.array([[4, 2, 2],
              [9, 0, 5],
              [9, 9, 4]])
b = np.zeros((7, 7, 5, 5), dtype=int)

b would end up something like:

>>> b[:, :, 0, 0]
array([[4, 2, 2, 0, 0, 0, 0],
       [9, 0, 5, 0, 0, 0, 0],
       [9, 9, 4, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0]])
>>> b[:, :, 0, 1]
array([[0, 4, 2, 2, 0, 0, 0],
       [0, 9, 0, 5, 0, 0, 0],
       [0, 9, 9, 4, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0]])

Upvotes: 1

Views: 856

Answers (1)

Mad Physicist
Mad Physicist

Reputation: 114330

One way to think about this to make a sliding window view of b (6D), slice out the parts you want (3D or 4D), and assign a to them.

However, there is a simpler way to do this altogether. The way a sliding window view works is by creating a dimension that steps along less than the full size of the dimension you are viewing. For example:

>>> x = np.array([1, 2, 3, 4])
array([1, 2, 3, 4])
>>> window = np.lib.stride_tricks.as_strided(
                 x, shape=(x.shape[0] - 2, 3),
                    strides=x.strides * 2)
[[1 2 3]
 [2 3 4]]

I'm deliberately using np.lib.stride_tricks.as_strided rather than np.lib.stride_tricks.sliding_window_view here because it has a certain flexibility that you need.

You can have a stride that is larger than the axis you are viewing, as long as you are careful. Contiguous arrays are more forgiving in this case, but by no means a requirement. An example of this is np.diag. You can implement it something like this:

>>> x = np.arange(12).reshape(3, 4)
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> diag = np.lib.stride_tricks.as_strided(
               x, shape=(min(x.shape),),
                  strides=(sum(x.strides),))
array([ 0,  5, 10])

The trick is to make a view of only the parts of b you care about in a way that makes the assignment easy. Because of broadcasting rules, you will want the last two dimensions of the view to be a.shape, and the strides to be b.strides[:2], since that's where you want to place a.

The first two dimensions of the view will be responsible for making the copies of a. You want 25 copies, so the shape will be (5, 5). The strides are the trickier part. Let's take a look at a 2D case, just because that's easier to visualize, and then attempt to generalize:

>>> a0 = np.array([1, 2])
>>> b0 = np.zeros((4, 3), dtype=int)
>>> b0[0:2, 0] = b0[1:3, 1] = b0[2:4, 2] = a0

The goal is to make a view that strides along the diagonal of b0 in the first axis. So:

>>> np.lib.stride_tricks.as_strided(
        b0, shape=(b0.shape[0] - a0.shape[0] + 1, a0.shape[0]),
            strides=(sum(b0.strides), b0.strides[0]))[:] = a0
>>> b0
array([[1, 0, 0],
       [2, 1, 0],
       [0, 2, 1],
       [0, 0, 2]])

So that's what you do for b, but adding up every second dimension:

a = np.array([[4, 2, 2],
              [9, 0, 5],
              [9, 9, 4]])
b = np.zeros((7, 7, 5, 5), dtype=int)
vshape = (*np.subtract(b.shape[:a.ndim], a.shape) + 1,
          *a.shape)
vstrides = (*np.add(b.strides[:a.ndim], b.strides[a.ndim:]),
            *b.strides[:a.ndim])
np.lib.stride_tricks.as_strided(b, shape=vshape, strides=vstrides)[:] = a

TL;DR

def emplace_window(a, b):
    vshape = (*np.subtract(b.shape[:a.ndim], a.shape) + 1, *a.shape)
    vstrides = (*np.add(b.strides[:a.ndim], b.strides[a.ndim:]), *b.strides[:a.ndim])
    np.lib.stride_tricks.as_strided(b, shape=vshape, strides=vstrides)[:] = a

I've phrased it this way, because now you can apply it to any number of dimensions. The only expectations is that 2 * a.ndim == b.ndim and that b.shape[a.ndim:] == b.shape[:a.ndim] - a.shape + 1.

Upvotes: 1

Related Questions