首页 > 解决方案 > 可以重建 2d 矩阵,其中掩码已与 numpy where 一起使用并展平

问题描述

正如问题所说,我有一个 2D 矩阵 (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,)具有相同的维度,是否有任何方法或函数来重建展平矩阵转换为其原始的二维形式。我正在考虑将索引保存在布尔矩阵中并使用它们来重建矩阵,如果有人知道任何 numpy 方法或任何建议会有很大帮助。

问候。

标签: pythonpython-3.xnumpy

解决方案


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

您可以使用np.ravel_multi_index将 2D 索引转换为 1D。然后您可以使用这些 1D 索引将它们转换回 2D 使用np.unravel_index(尽管由于您已经拥有 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]])

推荐阅读