CODEWITHSUNDEEP
pythonarraysnumpyindexing
Daniel
Daniel

Reputation: 19547

Indexing offset diagonal of 4 dimensional array

A rather hard question to phrase. I need to index an array such that the following conditions are met for all a and b:

arr[a, b, a, b] = data[b]

This can be accomplished by the following:

import numpy as np

a = 3
b = 2

arr = np.zeros((a, b, a, b))

data = np.array([10, 20])
rngb = np.arange(b)

for i in range(a):
    arr[i, rngb, i, rngb] += data

print arr

The goal here is to remove the python for loop. A trivial answer would be to use tile and repeat on the index and data arrays, but I believe there is an easier solution that I am missing. Any suggestions would be most appreciated.

Edit: Another, but likely more costly way of computing arr:

diag = np.diag(np.ones(a))
arr = np.einsum('ik,jl->ijkl', diag, np.diag(data))

Upvotes: 2

Views: 94

Answers (1)

Oliver W.
Oliver W.

Reputation: 13459

def meth1(a,b):  # old method
    arr = np.zeros((a,b,a,b))
    data = np.arange(b)+1
    rngb = np.arange(b)
    for i in range(a):
        arr[i, rngb, i, rngb] += data
    return arr

def meth2(a,b):  # revised method
    arr = np.zeros((a, b, a, b))
    data = np.arange(b)+1
    arr.ravel()[::a*b+1] = np.tile(data, a)
    return arr

You could test the equality of both with a simple lambda:

m = lambda a,b: np.any((meth2(a,b)-meth1(a,b)))

Upvotes: 2

Related Questions