CODEWITHSUNDEEP
pythonnumpynumpy-ndarray
kevin
kevin

Reputation: 95

numpy filling the diagonal 0 for 3D array

Suppose I have a 3D array, how can I fill the diag of the first two dimensions to zero. For example

a = np.random.rand(2,2,3)
for i in range(3):
    np.fill_diagonal(a[:,:,i], 0)

Is there a way to replace the for loop?

Upvotes: 1

Views: 872

Answers (3)

hpaulj
hpaulj

Reputation: 231335

In [6]: a = np.random.randint(1,10,(2,2,3))
   ...: for i in range(3):
   ...:     np.fill_diagonal(a[:,:,i], 0)
In [7]: a
Out[7]: 
array([[[0, 0, 0],
        [7, 4, 4]],

       [[8, 2, 7],
        [0, 0, 0]]])

Indexing a diagonal is easy - just use the same index array for both dimensions. Thus the 0s we just set are:

In [8]: idx=np.arange(2)
In [9]: a[idx,idx,:]
Out[9]: 
array([[0, 0, 0],
       [0, 0, 0]])

and used to set a value:

In [10]: a[idx,idx,:] = 10
In [11]: a
Out[11]: 
array([[[10, 10, 10],
        [ 7,  4,  4]],

       [[ 8,  2,  7],
        [10, 10, 10]]])

Upvotes: 0

kevin
kevin

Reputation: 95

The following is one of the solution

a = np.random.rand(2,2,3)
np.einsum('iij->ij',a)[...] = 0

Upvotes: 2

Tim Roberts
Tim Roberts

Reputation: 54645

The np.diag function returns a 2D diagonal matrix.

a[:,:,0] = np.diag((1,1))

Upvotes: 0

Related Questions