Mapping a[i+j] array to A[i,j] matrix

Question

I have an array a, and want to create a new matrix A for which A[i,j] = a[i+j]. Example:

import numpy as np

a = np.random.rand(3)
A = np.zeros((2,2));

for i in range(2):
    for j in range(2):
        A[i,j] = a[i+j]

Is there a way of doing this without a for loop? (with numpy)

Answer 1

Using stride_tricks.as_strided

This would be a perfect use case for stride_tricks:

from np.lib.stride_tricks import as_strided

Set the strides as (8, 8) (i.e. (1, 1) in terms of slots). This way we essentially map the resulting array A as i, j -> k = i + j. A more detailed description would be: we map every i, j pair to a natural number k defined by the strides as k = i*s_i + j*s_j, where s_i and s_j are the strides set to 1s, i.e. k = i + j. Thus ending up with the desired result: A[i, j] = a[k] = a[i + j].

>>> a
array([0.53954179, 0.51789927, 0.33982179])

>>> A = as_strided(a, shape=(2,2), strides=(8,8))
array([[0.53954179, 0.51789927],
       [0.51789927, 0.33982179]])

Additional considerations

A more general solution would be to get the shapes as well as the strides from a's metadata.

• The shape of A is given by (len(a)//2+1,)*2.

• As noted by @Daniel F, the memory slot size does not always equal to 8, this indeed depends on the dtype of your array. It would be better to define strides from a's strides instead: a.strides*2.

This comes down to:

>>> A = as_strided(a, shape=(len(a)//2+1,)*2, strides=a.strides*2)

---

Using grid indexing

Alternatively, you can construct a grid of coordinates (you can do so using itertools.product) then copy the appropriate values from a to A:

i, j = np.array(list(product(range(2), range(2)))).T

Then initialize A and copy:

>>> A = np.zeros((2,2))
>>> A[i, j] = a[i + j]

This will however double the memory usage, compared to the as_strided method.