What's The Best Way To Create A "3d Identity Matrix" In Numpy?

I don't know if the title makes any sense. Normally an identity matrix is a 2D matrix like In [1]: import numpy as np In [2]: np.identity(2) Out[2]: array([[ 1., 0.], [ 0

Solution 1:

Starting from a 2d identity matrix, here are two options you can make the "3d identity matrix":

import numpy as npi= np.identity(2)

Option 1: stack the 2d identity matrix along the third dimension

np.dstack([i]*3)
#array([[[ 1.,  1.,  1.],
#        [ 0.,  0.,  0.]],

#       [[ 0.,  0.,  0.],
#        [ 1.,  1.,  1.]]])

Option 2: repeat values and then reshape

np.repeat(i, 3, axis=1).reshape((2,2,3))
#array([[[ 1.,  1.,  1.],
#        [ 0.,  0.,  0.]],

#       [[ 0.,  0.,  0.],
#        [ 1.,  1.,  1.]]])

Option 3: Create an array of zeros and assign 1 to positions of diagonal elements (of the 1st and 2nd dimensions) using advanced indexing:

shape = (2,2,3)
identity_3d = np.zeros(shape)
idx = np.arange(shape[0])
identity_3d[idx, idx, :] = 1  


identity_3d
#array([[[ 1.,  1.,  1.],
#        [ 0.,  0.,  0.]],

#       [[ 0.,  0.,  0.],
#        [ 1.,  1.,  1.]]])

---

Timing:

%%timeit
shape = (100,100,300)
i = np.identity(shape[0])
np.repeat(i, shape[2], axis=1).reshape(shape)

# 10 loops, best of 3: 10.1 ms per loop


%%timeit
shape = (100,100,300)
i = np.identity(shape[0])
np.dstack([i] * shape[2])

# 10 loops, best of 3: 47.2 ms per loop

%%timeit
shape = (100,100,300)
identity_3d = np.zeros(shape)
idx = np.arange(shape[0])
identity_3d[idx, idx, :] = 1

# 100 loops, best of 3: 6.31 ms per loop

Solution 2:

One way would be to initialize a 2D identity matrix and then broadcast it to 3D. Thus, with n as the length along the first two axes and r for the last axis, we could do -

np.broadcast_to(np.identity(n)[...,None], (n,n,r))

Sample run to make things clearer -

In [154]: i = np.identity(3); i # Create an identity matrix
Out[154]: 
array([[ 1.,  0.,  0.],
       [ 0.,  1.,  0.],
       [ 0.,  0.,  1.]])

# Extend it to 3D. This helps us broadcast to reqd. shape later on
In [152]: i[...,None]
Out[152]: 
array([[[ 1.],
        [ 0.],
        [ 0.]],

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

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

# Broadcast to (n,n,r) shape for the 3D identity matrix    
In [153]: np.broadcast_to(i[...,None], (3,3,3))
Out[153]: 
array([[[ 1.,  1.,  1.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]],

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

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

Such an approach leads to performance because it simply generates a view into the identity matrix. Thus, in that form the output would be a read-only array. If you need a write-able array that has its own memory space, simply append a .copy() there.

Asserting on the performance, here's a timing test to create a 3D identity matrix of shape :(100, 100, 300) -

In [140]: n,r = 100,300

In [141]: %timeit np.broadcast_to(np.identity(n)[...,None], (n,n,r))
100000 loops, best of 3: 9.29 µs per loop

Solution 3:

Similar to @Psidom, using advanced np.einsum

%%timeit
shape = (100,100,300)
identity_3d = np.zeros(shape)
np.einsum('iij->ij', identity_3d)[:] = 11000 loops, best of 3: 251 µs per loop

%%timeit
shape = (100,100,300)
identity_3d = np.zeros(shape)
idx = np.arange(shape[0])
identity_3d[idx, idx, :] = 11000 loops, best of 3: 320 µs per loop

%timeit np.broadcast_to(np.identity(100)[...,None], (100,100,300)).copy()

100 loops, best of 3: 12.1 ms per loop

I assume you want a copy() of the identity matrix to write to, because anyarray.dot(identity) is otherwise much more easily calculated by np.broadcast_to(anyarray[..., None], a.shape + (300,)). If you really just want a bare identity matrix, then solution of @Divakar is much faster than any others,

%timeit np.broadcast_to(np.identity(100)[...,None], (100,100,300))

The slowest run took 5.16 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 3: 20.8 µs per loop

but broadcast_to(anyarray[..., None], . . . ) is likely even faster than that.