Fill the Diagonal of Each Matrix in a 3D numpy Array with a Vector

Question

Having a 3D numpy array, where each 2D slice represents an individual matrix. I'm looking to replace the diagonal elements of every matrix with a specific set of values.

For instance, if I have a 3x3x3 array:

array([[[a1, a2, a3],
        [a4, a5, a6],
        [a7, a8, a9]],

       [[b1, b2, b3],
        [b4, b5, b6],
        [b7, b8, b9]],

       [[c1, c2, c3],
        [c4, c5, c6],
        [c7, c8, c9]]])

I'd like to replace the diagonals [a1, a5, a9], [b1, b5, b9], and [c1, c5, c9] with a new set of values for each matrix. How can I achieve this?

Answer 1

I would use integer indexing:

import numpy as np

# Setup:
arr = np.zeros((3, 4, 6), dtype=int)
vectors = np.random.randint(1, 9, size=(3, 4))

# Should work for arbitrary `arr` with ndim >= 2
n = min(arr.shape[-2:])
idx = np.arange(n)

# Note that `vectors.shape` must broadcast with
# `(*arr.shape[:-2], n)` for this to work:

arr[..., idx, idx] = vectors

Answer 2

It can be done using np.eye masking.

zeros = np.zeros((3,4,6))
vectors = np.random.randint(1 , 9, (3, 4)) # Generating random 3x4 integer array between 1 and 9

mask = np.eye(zeros.shape[1], zeros.shape[2], dtype=bool)
zeros[:, mask] = vectors

Here is the output of print(zeros).

[[[2. 0. 0. 0. 0. 0.]
  [0. 4. 0. 0. 0. 0.]
  [0. 0. 7. 0. 0. 0.]
  [0. 0. 0. 4. 0. 0.]]

 [[7. 0. 0. 0. 0. 0.]
  [0. 4. 0. 0. 0. 0.]
  [0. 0. 3. 0. 0. 0.]
  [0. 0. 0. 6. 0. 0.]]

 [[3. 0. 0. 0. 0. 0.]
  [0. 5. 0. 0. 0. 0.]
  [0. 0. 1. 0. 0. 0.]
  [0. 0. 0. 3. 0. 0.]]]

Answer 3

You can craft the indices for direct indexing:

# example input
a = np.zeros((3,3,3), dtype=int)

N = min(a.shape)
idx1 = np.repeat(range(N), N)
idx2 = np.tile(range(N), N)

a[idx1, idx2, idx2] = np.arange(1, 10)

Output:

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

       [[4, 0, 0],
        [0, 5, 0],
        [0, 0, 6]],

       [[7, 0, 0],
        [0, 8, 0],
        [0, 0, 9]]])

Intermediates:

# idx1
array([0, 0, 0, 1, 1, 1, 2, 2, 2])

# idx2
array([0, 1, 2, 0, 1, 2, 0, 1, 2])

Answer 4

In the future, when numpy implements the functionality described below, this will be a one-liner. Until then, you'll need to use an indexing or masking method shown in one of the other answers.

---

The documentation for numpy.diagonal (as of v1.25) says:

Starting in NumPy 1.9 it returns a read-only view on the original array. Attempting to write to the resulting array will produce an error.

In some future release, it will return a read/write view and writing to the returned array will alter your original array. The returned array will have the same type as the input array.

In some future release of numpy, you should be able to write to the view returned by np.diagonal directly.

np.diagonal(zeros, axis1=-2, axis2=-1) = vectors

Answer 5

import numpy as np

arr = np.array([[[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9]],

                [[10, 11, 12],
                 [13, 14, 15],
                 [16, 17, 18]],

                [[19, 20, 21],
                 [22, 23, 24],
                 [25, 26, 27]]])

new_values = [[-1, -2, -3], [-4, -5, -6], [-7, -8, -9]]

# Correct vectorized diagonal replacement:
# Create a 2D index array
diagonal_indices = np.arange(arr.shape[1])  
print(diagonal_indices)#[0 1 2]
# Use advanced indexing
arr[:, diagonal_indices, diagonal_indices] = new_values  

print(arr)
"""[[[-1  2  3]
  [ 4 -2  6]
  [ 7  8 -3]]

 [[-4 11 12]
  [13 -5 15]
  [16 17 -6]]

 [[-7 20 21]
  [22 -8 24]
  [25 26 -9]]]
"""

Method 2 :

diagonal_indices = np.diag_indices(arr.shape[1])
arr[:, diagonal_indices[0], diagonal_indices[1]] = np.expand_dims(new_values, axis=0)

print(arr)
[[[-1  2  3]
  [ 4 -2  6]
  [ 7  8 -3]]

 [[-4 11 12]
  [13 -5 15]
  [16 17 -6]]

 [[-7 20 21]
  [22 -8 24]
  [25 26 -9]]]