Expert

as_strided()¶

The as_strided() function can be used to create complex views of an array.

Consider this 2x4 integer array called foo.

foo = np.array([
    [10,20,30,40],
    [50,60,70,80]
])

Recall: arrays are stored in contiguous, fixed-size memory blocks. In this case, foo is stored in memory like this.

# [ 10 | 20 | 30 | 40 | 50 | 60 | 70 | 80 ]
# ------ row 0 ------- ------ row 1 -------

Since foo is comprised of 64-bit integers, each block of memory is 64 bits. Alternatively stated, each block of memory is 8 bytes.

To fetch the element at position (i,j), NumPy could do the following:

Start at the beginning of the memory block
jump across i * 32 bytes of data to get to row i
jump across j * 8 bytes to get to the jth element in row i

This is exactly what the strides attribute of a numpy array describes. For example, foo.strides returns the tuple (32, 8), meaning "axis 0 iterates by 32 bytes, axis 1 iterates by 8 bytes".

foo.strides  # (32, 8)

With np.lib.stride_tricks.as_strided(), you can create a new view of an existing array by modifying its strides without copying or modifying its data.

Example

print(foo)
# [[10 20 30 40]
#  [50 60 70 80]]

bar = np.lib.stride_tricks.as_strided(x=foo, shape=(3,4), strides=(16,8))
print(bar)
# [[10 20 30 40]
#  [30 40 50 60]
#  [50 60 70 80]]

Here we define a 3x4 array that's based on the data in foo, but we tell numpy to jump across 16 bytes to get to the next row and 8 bytes to get to the next column.

For example, if we request the element at index (1,0), numpy starts at the beginning of foo and then jumps across 16 bytes (one row) plus 0 bytes (0 columns), landing at 30.

print(bar[1,0])  # 30

# foo: [10 | 20 | 30 | 40 | 50 | 60 | 70 | 80 ]
#      ---------^

To get to index (1,3), numpy jumps across 16 bytes (one row) plus 24 bytes (one column) = 40 bytes, landing at 60.

print(bar[1,0])  # 60

# foo: [10 | 20 | 30 | 40 | 50 | 60 | 70 | 80 ]
#      ------------------------^

Warning

It's really important to understand that bar is a view of foo. If we modify bar, foo will be modified as well.

Example

bar[1, 0] = 999

print(foo)
# [[ 10  20 999  40]
#  [ 50  60  70  80]]

foo changed even though we modified bar.

Furthermore, if we print(bar) you can see that element (1,0) and element (0,2) changed.

print(bar)
# [[ 10  20 999  40]
#  [999  40  50  60]
#  [ 50  60  70  80]]

That's because these elements in bar point to the same block of memory.

Danger

When using as_strided(), be careful that your strides make sense. Otherwise you may end up pointing to memory used by a different variable. This can have bad side effects, so beware.

sliding_window_view()¶

NumPy has a convenient sliding_window_view() function for making sliding windows.

Relation to as_strided()

Under the hood, sliding_window_view() is just a fancy wrapper for as_strided(). Anything you can do with sliding_window_view(), you could also do with as_strided().

For example, given this array

foo = np.array([0,1,2,3,4])

we can make a sliding window with length 3 like this.

fooview = np.lib.stride_tricks.sliding_window_view(
    x=foo, 
    window_shape=3
)

print(fooview)
# [[0 1 2]    [0,1,2,3,4]
#  [1 2 3]    [0,1,2,3,4]
#  [2 3 4]]   [0,1,2,3,4]

The result is a read-only view of foo. If you try to modify fooview, you'll get an error!

fooview[0,0] = 999 
# ValueError: assignment destination is read-only

We can make it writeable with writeable=True.

fooview = np.lib.stride_tricks.sliding_window_view(
    x=foo, 
    window_shape=3, 
    writeable=True
)
fooview[0,0] = 999

print(fooview)
# [[999   1   2]
#  [  1   2   3]
#  [  2   3   4]]

foo is changed too!

Since fooview is a view of foo, modifying fooview also modifies foo!

print(foo)
# [999   1   2   3   4]

sliding_window_view() works for multidimensional arrays too. Consider this 3x3 array, zoo.

zoo = np.arange(9).reshape(3,3)

print(zoo)
# [[0 1 2]
#  [3 4 5]
#  [6 7 8]]

We can create various sliding windows from zoo by varying the window_shape and axis parameters.

window_shape=2window_shape=(2,2)

With window_shape=2, NumPy searches for subarrays inside zoo like [a, b].

axis=0axis=1axis=(0,1)axis=(1,0)

zooview = np.lib.stride_tricks.sliding_window_view(
    x=zoo, 
    window_shape=2,
    axis=0
)

print(zooview)
# [
#     [[0 3]
#      [1 4]
#      [2 5]]
#
#     [[3 6]
#      [4 7]
#      [5 8]]
# ]

In pseudocode, you could describe this algorithm as

Iterate over the elements of zoo in row-major order.

Step 1    Step 2    Step 3    Step 4
[[0 1 2]  [[0 1 2]  [[0 1 2]  [[0 1 2]   
 [3 4 5]   [3 4 5]   [3 4 5]   [3 4 5]
 [6 7 8]]  [6 7 8]]  [6 7 8]]  [6 7 8]] ...

At each step, move along axis 0 to fill an array with shape (2,).

Step 1    Step 2    Step 3    Step 4
[[0 1 2]  [[0 1 2]  [[0 1 2]  [[0 1 2]   
 [3 4 5]   [3 4 5]   [3 4 5]   [3 4 5]
 [6 7 8]]  [6 7 8]]  [6 7 8]]  [6 7 8]] ...

zooview = np.lib.stride_tricks.sliding_window_view(
    x=zoo, 
    window_shape=2,
    axis=1
)

print(zooview)
# [
#     [[0 1]
#      [1 2]]
#     
#     [[3 4]
#      [4 5]]
#  
#     [[6 7]
#      [7 8]]
# ]

In pseudocode, you could describe this algorithm as

Iterate over the elements of zoo in row-major order.

Step 1    Step 2    Step 3    Step 4
[[0 1 2]  [[0 1 2]  [[0 1 2]  [[0 1 2]   
 [3 4 5]   [3 4 5]   [3 4 5]   [3 4 5]
 [6 7 8]]  [6 7 8]]  [6 7 8]]  [6 7 8]] ...

At each step, move along axis 1 to fill an array with shape (2,).

Step 1    Step 2    Step 3    Step 4
[[0 1 2]  [[0 1 2]  [[0 1 2]  [[0 1 2]   
 [3 4 5]   [3 4 5]   [3 4 5]   [3 4 5]
 [6 7 8]]  [6 7 8]]  [6 7 8]]  [6 7 8]] ...

zooview = np.lib.stride_tricks.sliding_window_view(
    x=zoo, 
    window_shape=2,
    axis=(0,1)
)

# ValueError: Must provide matching length window_shape and axis; got 1 window_shape elements and 2 axes elements.

zooview = np.lib.stride_tricks.sliding_window_view(
    x=zoo, 
    window_shape=2,
    axis=(1,0)
)

# ValueError: Must provide matching length window_shape and axis; got 1 window_shape elements and 2 axes elements.

With window_shape=2, NumPy will search for subarrays inside zoo like

[[a, b]
 [c, d]]

axis=0axis=1axis=(0,1)axis=(1,0)

zooview = np.lib.stride_tricks.sliding_window_view(
    x=zoo, 
    window_shape=(2,2),
    axis=0
)

# ValueError: Must provide matching length window_shape and axis; got 2 window_shape elements and 1 axes elements.

zooview = np.lib.stride_tricks.sliding_window_view(
    x=zoo, 
    window_shape=(2,2),
    axis=1
)

# ValueError: Must provide matching length window_shape and axis; got 2 window_shape elements and 1 axes elements.

zooview = np.lib.stride_tricks.sliding_window_view(
    x=zoo, 
    window_shape=(2,2),
    axis=(0,1)
)

# [
#     [
#         [[0 1]
#          [3 4]]
#     
#         [[1 2]
#          [4 5]]
#     ]
#     [
#         [[3 4]
#          [6 7]]
#   
#         [[4 5]
#          [7 8]]
#     ]
# ]

In pseudocode, you could describe this algorithm as

Iterate over the elements of zoo in row-major order.

Step 1    Step 2    Step 3    Step 4
[[0 1 2]  [[0 1 2]  [[0 1 2]  [[0 1 2]   
 [3 4 5]   [3 4 5]   [3 4 5]   [3 4 5]
 [6 7 8]]  [6 7 8]]  [6 7 8]]  [6 7 8]] ...

At each step, move along axis 0 first, then axis 1 to fill an array with shape (2,2).

Step 1    Step 2    Step 3    Step 4
[[0 1 2]  [[0 1 2]  [[0 1 2]  [[0 1 2]   
 [3 4 5]   [3 4 5]   [3 4 5]   [3 4 5]
 [6 7 8]]  [6 7 8]]  [6 7 8]]  [6 7 8]] ...

zooview = np.lib.stride_tricks.sliding_window_view(
    x=zoo, 
    window_shape=(2,2),
    axis=(1,0)
)

# [
#     [
#         [[0 3]
#          [1 4]]
#   
#         [[1 4]
#          [2 5]]
#     ]
#     [
#         [[3 6]
#          [4 7]]
#   
#         [[4 7]
#          [5 8]]
#     ]
# ]

In pseudocode, you could describe this algorithm as

Iterate over the elements of zoo in row-major order.

Step 1    Step 2    Step 3    Step 4
[[0 1 2]  [[0 1 2]  [[0 1 2]  [[0 1 2]   
 [3 4 5]   [3 4 5]   [3 4 5]   [3 4 5]
 [6 7 8]]  [6 7 8]]  [6 7 8]]  [6 7 8]] ...

At each step, move along axis 1 first, then axis 0 to fill an array with shape (2,2).

Step 1    Step 2    Step 3    Step 4
[[0 1 2]  [[0 1 2]  [[0 1 2]  [[0 1 2]   
 [3 4 5]   [3 4 5]   [3 4 5]   [3 4 5]
 [6 7 8]]  [6 7 8]]  [6 7 8]]  [6 7 8]] ...

einsum()¶

You can use the einsum() function to quickly and efficiently multiply and sum arrays using einstein sums.

Consider these 1-d arrays A and B.

A = np.arange(4)
B = A + 4

print(A)
# [0 1 2 3]

print(B)
# [4 5 6 7]

If we do np.einsum('i,j->i', A, B) we get back the array [ 0, 22, 44, 66].

np.einsum('i,j->i', A, B)
# [ 0, 22, 44, 66]

In pseudocode, you could describe the above algorithm as

initialize the output, Y, as an array of 0s the same size as A
for each i:
  for each j:
    Yi += Ai * Bj

The first parameter of einsum() is subscripts. Assuming we're operating on two arrays A and B, it always has the form

"subscripts for A's axes, subscripts for B's axes -> subscripts for output axes"

In this case,

A has one dimension so we give it the subscript i
B has one dimension so we give it the subscribt j
By reusing subscript i for the output axes, we're saying "the output array has one dimension and it's the same length as A". In other words, element \(A_i\) will always feed into a corresponding \(Y_i\).

Example 1¶

np.einsum('i,j->', A, B)
# 132

In pseudocode,

initialize the output, Y, equal to the scalar 0 (no dimensions)
for each i:
  for each j:
    Y += Ai * Bj

Since there's no subscript to the right of the arrow in the subscript string, we're telling NumPy that our output should have 0 dimensions. In other words, the output should be a scalar. And like before, we iterate over each i in A and each j in B, adding \(A_i * B_j\) to the sum.

Example 2¶

np.einsum('z,z->z', A, B)
# [ 0,  5, 12, 21]

In pseudocode,

intialize the output, Y, as an array of 0s the same size as A (and B)
for each z:
  Yz += Az * Bz

This example is equivalent to doing the element-wise product, A * B.

Example 3¶

np.einsum('s,t->st', A, B)
# [[ 0,  0,  0,  0],
#  [ 4,  5,  6,  7],
#  [ 8, 10, 12, 14],
#  [12, 15, 18, 21]]

In pseudocode,

initialize output, Y, as a 4x4 array of 0s
for each s:
  for each t:
    Yst += As * Bt

Note

The output should be 2-dimensional because it has two subscripts: s and t. Since s iterates over A (length 4) and t iterates over B (length 4), we know the output should have shape (4,4).

Example 4¶

einsum() really starts to shine in two dimensions. Consider these 2x2 arrays C and D

C = np.arange(4).reshape(2,2)
D = C + 4

print(C)
# [[0 1]
#  [2 3]]

print(D)
# [[4 5]
#  [6 7]]

and observe this einstein sum

np.einsum('ij,ji->', C, D)
# 37

Let's breakdown the subscript string.

ij tells us the first array has exactly two dimensions
ji tells us the second array also has exactly two dimensions. Furthermore, the length of its first dimension matches the length of the first array's second dimension since they both use subscript j, and the length of its second dimension matches the length of the first array's first dimension since they both use subscript i.
The bit after the arrow is empty, so we know the output will have 0 dimensions (it'll be a scalar).

In pseudocode,

initialize output, Y, equal to 0
for each i:
 for each j:
   Y += Cij * Dji

Tip

On the surface, this particular einstein sum is equivalent to doing np.sum(C * D.T). However, einsum() only accesses each element once wheras np.sum(C * D.T) accesses each element twice - first when it does the multiplication and second when it does the sum. More importantly though, np.sum(C * D.T) creates a temporary array that takes up memory before gettting summed into a scalar. einsum() avoids this memory consumption, which, if you're dealing with big arrays, can make a significant difference.