1
a=np.zeros((3,3,3))
b=np.arange(3)
c=np.arange(9).reshape(3,3)

I wanna put the elements of the array b or c along the diagonal (or above/below the diagonal) of the 3d matrix (tensor) a with respect to a specific axis.

I tired numpy.diagflat, but it only works for 2d matrix.

For instance, how to make the following matrix?

array([[[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]],

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

       [[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]]])
kinder chen
  • 1,371
  • 5
  • 15
  • 25

2 Answers2

6

For the main diagonals you can use np.einsum. For example:

>> np.einsum('iii->i', a)[...] = b
>>> a
array([[[ 0.,  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.,  2.]]])

Or:

>>> np.einsum('iji->ji', a)[...] = c
>>> a
array([[[ 0.,  0.,  0.],
        [ 3.,  0.,  0.],
        [ 6.,  0.,  0.]],

       [[ 0.,  1.,  0.],
        [ 0.,  4.,  0.],
        [ 0.,  7.,  0.]],

       [[ 0.,  0.,  2.],
        [ 0.,  0.,  5.],
        [ 0.,  0.,  8.]]])

Edit: Broadcasting works normally:

>>> np.einsum('ijj->ij', a)[...] = b
>>> a
array([[[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]],

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

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

Subdiagonals also work but are more tricky as some manual slicing is required. For example:

>>> a=np.zeros((3,3,3))
>>> np.einsum('iij->ij', a[:2,1:])[...] = c[1:]
>>> a
array([[[ 0.,  0.,  0.],
        [ 3.,  4.,  5.],
        [ 0.,  0.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 6.,  7.,  8.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]]])
Paul Panzer
  • 51,835
  • 3
  • 54
  • 99
  • A rising star in the `numpy` tag! – juanpa.arrivillaga Feb 05 '18 at 23:53
  • You seem in a generous mood today. Or are you talking about `einsum`? – Paul Panzer Feb 06 '18 at 00:06
  • Thx, does `np.einsum` assign the elements one by one or do all of them in the same time? – kinder chen Feb 06 '18 at 00:09
  • Nah, it was just a compliment. Maybe if I finally read some tutorial on einsum I'll be able to get some rep in the `numpy` tag ;) – juanpa.arrivillaga Feb 06 '18 at 00:13
  • It does neither. What it does is to create a writable view into the original array. We then use this view to manually assign. My hunch is that this should be pretty fast at least for large arrays. – Paul Panzer Feb 06 '18 at 00:13
  • @kinderchan, look at the `__array_interface__` of these `a` and the `einsums`. `einsum` is creating a `view`, with special striding to access these diagonal elements. Most cleaver, this is! – hpaulj Feb 06 '18 at 00:14
  • @juanpa.arrivillaga Well, thanks then! I just had a look at the rankings page - I wasn't aware. – Paul Panzer Feb 06 '18 at 00:19
  • thx, I re-editted my question and I wanna get the matrix I added, does `np.einsum` still work and how? – kinder chen Feb 06 '18 at 07:41
  • @kinderchan, the key to using `einsum` here is that it creates a `view`. Anything that can be done with a regular `view` (e.g. `x[1:10,:]`), should work with one produced by `einsum`. But you can test this yourself. Once you understand the pieces you can use them in various ways. (edit) `np.einsum('iii->i',A)[...] += 100` works – hpaulj Feb 06 '18 at 08:00
  • @kinderchan I've edited the answer and added your specific example. – Paul Panzer Feb 06 '18 at 10:45
  • @PaulPanzer, can `np.einsum` work on off-diagonal (above/below the diagonal) ? – kinder chen Feb 06 '18 at 17:13
  • @kinderchan not by itself, but you can combine it with slicing to get arbitrary subdiagonals as I show in the last example. – Paul Panzer Feb 06 '18 at 17:17
  • I appreciate your help, and I'm trying in that way. – kinder chen Feb 06 '18 at 17:48
  • @kinderchan Do you mean subdiagonal (for example i==j+1) or triangle (for example i > j)? If you mean triangle you can't use `einsum` or any view-based approach because triangles cannot be expressed as a product of strided intervals. – Paul Panzer Feb 07 '18 at 05:15
  • @PaulPanzer, I mean subdiagonal (`i==j+1` and `j-1`). I already did it by slicing the original matrix to a specific shape successfully. Thx so much for your help! – kinder chen Feb 07 '18 at 20:56
-1

If a sparse matrix will serve your needs, here are some resources:

http://www.janeriksolem.net/sparray-sparse-n-dimensional-arrays-in.html

sparse 3d matrix/array in Python?

ChootsMagoots
  • 670
  • 1
  • 6
  • 19