How to reorder 2D matrix axes based on reordering of a vector in numpy-CodePudding

I have a 2D ndarray, shaped (n_x, n_t)

The 2D matrix stores along its rows (so horizontally, for each line) a quantity Q, so a discrete version of a function of t, for fixed x. Across its columns, so for fixed t, (i.e. down a fixed column), the 2D matrix saves that quantity's Q values for different x's, so I have a discrete version of Q(x).

I will be plotting this matrix as a 2D heatmap.

I have a vector shaped (n_t, ) which contains the times for which the 2D matrix stores the quantity Q's values, for any given x values. So for each row, across the columns, the times underlying those values from the 2D matrix are the same.

Visually, for n_x=3 and n_t=3:

[
[Q_{11}, Q_{12}, Q_{13}],
[Q_{21}, Q_{22}, Q_{23}],
[Q_{31}, Q_{32}, Q_{33}]
]

The row [Q11, Q12, Q13] is basically Q(x_1, ts), the row [Q21, Q22, Q23] is basically Q(x_2, ts) and similarly for the last row.

I have ts as: [t1, t2, t3].

Problem:

I reorder ts (for whatever reason) to be [t2, t3, t1].

I want to reorder the matrix as:

[
[Q_{12}, Q_{13}, Q_{11}], 
[Q_{22}, Q_{23}, Q_{21}], 
[Q_{31}, Q_{33}, Q_{32}]
]

How shall I do it? What shall I read about? The reordering of the ts vector comes from: np.fft.fftshift(np.fft.fftreq(ts)).

Thank you!

EDIT from comments:

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

ts = np.array([100, 200, 300])

tss = np.array([200, 300, 100])

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

CodePudding user response：

Use:

import numpy as np

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

ts = np.array([100, 200, 300])
tss = np.array([200, 300, 100])


# number of rows n_x in the original question
n_x = a.shape[0]

# find the original positions
indices = (ts == tss[:, None]).argmax(1)

res = np.take_along_axis(a, np.tile(indices, (n_x, 1)), axis=1)
print(res)

Output

[[2 3 1]
 [5 6 4]
 [8 9 7]]