CODEWITHSUNDEEP
pythonarraysnumpy
samman
samman

Reputation: 613

Modifying the diagonal array of np.diag using a loop

I've been trying to look up how np.diag_indices work, and for examples of them, however the documentation for it is a bit light. I know this creates a diagonal array through your matrix, however I want to change the diagonal array (I was thinking of using a loop to change its dimensions or something along those lines).

I.E. say we have a 3x2 matrix:

[[1 2]
 [3 4]
 [5  6]]

Now if I use np.diag_indices it will form a diagonal array starting at (0,0) and goes through (1,1). 

[1 4]

However, I'd like this diagonal array to then shift one down. So now it starts at (0,1) and goes through (1,2). 

[3 6]

However there are only 2 arguments for np.diag_indices, neither of which from the looks of it enable me to do this. Am I using the wrong tool to try and achieve this? If so, what tools can I use to create a changing diagonal array that goes through my matrix? (I'm looking for something that will also work on larger matrices like a 200x50).

Upvotes: 1

Views: 253

Answers (1)

hpaulj
hpaulj

Reputation: 231385

The code for diag_indices is simple, so simple that I've never used it:

idx = arange(n)
return (idx,) * ndim

In [68]: np.diag_indices(4,2)                                                   
Out[68]: (array([0, 1, 2, 3]), array([0, 1, 2, 3]))

It just returns a tuple of arrays, the arange repeated n times. It's useful for indexing the main diagonal of a square matrix, e.g.

In [69]: arr = np.arange(16).reshape(4,4)                                       
In [70]: arr                                                                    
Out[70]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])
In [71]: arr[np.diag_indices(4,2)]                                              
Out[71]: array([ 0,  5, 10, 15])

The application is straight forward indexing with two arrays that match in shape.

It works on other shapes - if they are big enogh.

np.diag applied to the same array does the same thing:

In [72]: np.diag(arr)                                                           
Out[72]: array([ 0,  5, 10, 15])

but it also allows for offset:

In [73]: np.diag(arr, 1)                                                        
Out[73]: array([ 1,  6, 11])

===

Indexing with diag_indices does allow us to change that diagonal:

In [78]: arr[np.diag_indices(4,2)] += 10                                        
In [79]: arr                                                                    
Out[79]: 
array([[10,  1,  2,  3],
       [ 4, 15,  6,  7],
       [ 8,  9, 20, 11],
       [12, 13, 14, 25]])

====

But we don't have to use diag_indices to generate the desired indexing arrays:

In [80]: arr = np.arange(1,7).reshape(3,2)                                      
In [81]: arr                                                                    
Out[81]: 
array([[1, 2],
       [3, 4],
       [5, 6]])

selecting values from 1st 2 rows, and columns:

In [82]: arr[np.arange(2), np.arange(2)]                                        
Out[82]: array([1, 4])
In [83]: arr[np.arange(2), np.arange(2)] += 10                                  
In [84]: arr                                                                    
Out[84]: 
array([[11,  2],
       [ 3, 14],
       [ 5,  6]])

and for a difference selection of rows:

In [85]: arr[np.arange(1,3), np.arange(2)] += 20                                
In [86]: arr                                                                    
Out[86]: 
array([[11,  2],
       [23, 14],
       [ 5, 26]])

The relevant documentation section on advanced indexing with integer arrays: https://numpy.org/doc/stable/reference/arrays.indexing.html#purely-integer-array-indexing

Upvotes: 3

Related Questions