可以重建 2d 矩阵，其中掩码已与 numpy where 和 flattened 一起使用

Question

正如问题所说，我有一个二维矩阵 (1000, 2000)，我在其中应用 numpy where 函数的条件，以这种方式：

import numpy as np

A = np.random.randn(1000, 2000)
print(A.shape)
(1000, 2000)
mask = np.where((A >=0.1) & (A <= 0.5))

A = A[mask]
print(A.shape)
(303112,)

我得到一个展平矩阵，我将其用作仅支持一维矩阵的 Fortran 程序中的输入，该程序的输出与输入一维矩阵具有相同的维度 (303112,)，是否有任何方法或功能可以将展平矩阵重建为其原始 2D 形式。我正在考虑将索引保存在布尔矩阵中并使用它们来重建矩阵，如果有人知道任何 numpy 方法或任何建议会有很大帮助。

您好。

Answer 1

IIUC 您需要维护掩码的 1D 索引和 2D 索引，以便当您尝试使用 FORTRAN 程序更新这些值时，您可以切换到 1D 进行输入，然后切换回 2D 以更新原始数组。

您可以使用np.ravel_multi_index 将二维索引转换为一维。然后，您可以使用这些 1D 索引使用np.unravel_index 将它们转换回 2D（尽管由于您已经拥有 2D 蒙版，因此无需再次将 1D 转换为 2D。）

import numpy as np

A = np.random.randn(1000, 2000)
mask = np.where((A >=0.1) & (A <= 0.5))

idx_flat = np.ravel_multi_index(mask, (1000,2000)) #FLAT 1D indexes using original mask
idx_2d = np.unravel_index(idx_flat, (1000,2000))   #2D INDEXES using the FLAT 1D indexes


#Comparing the method of using flat indexes and A[mask]
print(np.allclose(A.ravel()[idx_flat],A[mask])) 
### True

#Comparing the freshly created 2d indexes to the original mask
print(np.allclose(idx_2d,mask))
### True

---

这是一个带有 (3,3) 矩阵的端到端代码的虚拟测试用例。

import numpy as np

#Dummy matrix A and mask
A = np.random.randn(3, 3)  #<---- shape (3,3)
mask = np.where(A <= 0.5)
mask[0].shape  #Number of indexes in 2D mask
###Output: (6,)

#########################################################

#Flatten 2D indexes to 1D
idx_flat = np.ravel_multi_index(mask, (3,3)) #<--- shape (3,3)
idx_flat.shape  #Number of indexes in flattened mask
###Output: (6,)

#########################################################

#Feed the 6 length array to fortran function
def fortran_function(x):
    return x**2

flat_array = A.ravel()[idx_flat]
fortran_output = fortran_function(flat_array)

#Number of values in fortran_output
fortran_output.shape
###Output: (6,)

#########################################################

#Create a empty array 
new_arr = np.empty((3,3))  #<---- shape (3,3)
new_arr[:] = np.nan
new_arr[mask] = fortran_output   #Feed the 1D array to the 2D masked empty array

new_arr

array([[5.63399114e-04,            nan, 7.86255167e-01],
       [3.94992857e+00, 4.88932044e-02, 2.45489069e+00],
       [3.51957270e-02,            nan,            nan]])