Vectorized assignment for numpy array with repeated indices (d[i,j,i,j] = s[i,j])

Question

How can I set

d[i,j,i,j] = s[i,j]

using "NumPy" and without for loop?

I've tried the follow:

l1=range(M)
l2=range(N)
d[l1,l2,l1,l2] = s[l1,l2]

Answer 1

If you think about it, that would be same as creating a 2D array of shape (m*n, m*n) and assigning the values from s into the diagonal places. To have the final output as 4D, we just need a reshape at the end. That's basically being implemented below -

m,n = s.shape
d = np.zeros((m*n,m*n),dtype=s.dtype)
d.ravel()[::m*n+1] = s.ravel()
d.shape = (m,n,m,n)

Runtime test

Approaches -

# @MSeifert's solution
def assign_vals_ix(s):    
    d = np.zeros((m, n, m, n), dtype=s.dtype)
    l1 = range(m)
    l2 = range(n)
    d[np.ix_(l1,l2)*2] = s[np.ix_(l1,l2)]
    return d

# Proposed in this post
def assign_vals(s):
    m,n = s.shape
    d = np.zeros((m*n,m*n),dtype=s.dtype)
    d.ravel()[::m*n+1] = s.ravel()
    return d.reshape(m,n,m,n)

# Using a strides based approach
def assign_vals_strides(a):
    m,n = a.shape
    p,q = a.strides

    d = np.zeros((m,n,m,n),dtype=a.dtype)
    out_strides = (q*(n*m*n+n),(m*n+1)*q)
    d_view = np.lib.stride_tricks.as_strided(d, (m,n), out_strides)
    d_view[:] = a
    return d

Timings -

In [285]: m,n = 10,10
     ...: s = np.random.rand(m,n)
     ...: d = np.zeros((m,n,m,n))
     ...: 

In [286]: %timeit assign_vals_ix(s)
10000 loops, best of 3: 21.3 µs per loop

In [287]: %timeit assign_vals_strides(s)
100000 loops, best of 3: 9.37 µs per loop

In [288]: %timeit assign_vals(s)
100000 loops, best of 3: 4.13 µs per loop

In [289]: m,n = 20,20
     ...: s = np.random.rand(m,n)
     ...: d = np.zeros((m,n,m,n))


In [290]: %timeit assign_vals_ix(s)
10000 loops, best of 3: 60.2 µs per loop

In [291]: %timeit assign_vals_strides(s)
10000 loops, best of 3: 41.8 µs per loop

In [292]: %timeit assign_vals(s)
10000 loops, best of 3: 35.5 µs per loop