Kronecker delta in Numpy

Question

Is there a convenient way to initialize a Numpy array defined in terms of a Kronecker delta? In an ideal world, there would be a function called, say kron_delta, such that if I set M=kron_delta('ij,kl',5), then M would be a 5x5x5x5 array where M[i,j,k,l]=1 whenever i=j and k=l, and 0 otherwise. My question is whether a function like kron_delta exists.

Followup: If not, is there still a convenient way to initialize the array delta(ijk) (i.e., an nxnxn array which is 1 whenever all indices are equal, and 0 otherwise)?

Answer 1

Alternative solution using numpy's einsum:

n = 5
M = np.einsum('ij,kl->ijkl', np.eye(n,n), np.eye(n,n))

Answer 2

Your particular example can be done like so

import numpy as np
n = 3
i,k = np.ogrid[:n, :n]
res = np.zeros((n,n,n,n), int)
res[i,i,k,k] = 1
res
array([[[[1, 0, 0],
         [0, 1, 0],
         [0, 0, 1]],

        [[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]],

        [[1, 0, 0],
         [0, 1, 0],
         [0, 0, 1]],

        [[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]],

        [[1, 0, 0],
         [0, 1, 0],
         [0, 0, 1]]]])

And the followup:

res3 = np.zeros((3,3,3),int)
i = np.arange(n)
res3[i,i,i] = 1
res3
array([[[1, 0, 0],
        [0, 0, 0],
        [0, 0, 0]],

       [[0, 0, 0],
        [0, 1, 0],
        [0, 0, 0]],

       [[0, 0, 0],
        [0, 0, 0],
        [0, 0, 1]]])