Create a 3D (partially) diagonal array from a 2D array

Question

I'd like to ask how can I efficiently generate a numpy 3D array from a 2D array with each row filling the diagonal part of the new array? For example, the input 2D array is

array([[1, 2],
       [3, 4],
       [5, 6],
       [7, 8]])

and I want the output to be

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

       [[3, 0],
        [0, 4]],

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

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

Typically, the size of the first dimensional is very large. Thanks in advance.

How are you currently generating, and why don't you consider it efficient enough? — CristiFati
– CristiFati, Commented Jan 24, 2023 at 13:14
I used a for-loop to diag each row, but this is low efficient since normally the input array contains thousands of rows — user8036269
– user8036269, Commented Jan 24, 2023 at 13:21
np.einsum('jk,kl->jkl', arr, np.eye(2, dtype=int)) or arr[...,None] * np.eye(2, dtype=int)[None] — Michael Szczesny
– Michael Szczesny, Commented Jan 24, 2023 at 13:28
@MichaelSzczesny Thanks so much for your reply. It works nicely with the einsum :) — user8036269
– user8036269, Commented Jan 24, 2023 at 13:50

mozway · Accepted Answer · 2023-01-24 13:52:35Z

1

Assuming a the input and using indexing with unravel_index:

x, y = np.unravel_index(np.arange(a.size), a.shape)

out = np.zeros(a.shape+(a.shape[-1],), dtype=a.dtype)

out[x, y, y] = a.flat

Output:

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

       [[3, 0],
        [0, 4]],

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

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

timings:

edited Jan 24, 2023 at 13:52

answered Jan 24, 2023 at 13:34

mozway

267k13 gold badges56 silver badges106 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

CristiFati · Accepted Answer · 2023-01-24 13:23:44Z

0

arr = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
res = np.apply_along_axis(np.diag, 1, arr)

edited Jan 24, 2023 at 13:23

CristiFati

41.6k9 gold badges68 silver badges115 bronze badges

answered Jan 24, 2023 at 13:21

fabi

312 bronze badges

Collectives™ on Stack Overflow

Create a 3D (partially) diagonal array from a 2D array

2 Answers 2

timings:

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

timings:

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related