提高性能python矩阵

Question

我有一个包含地球物理数据的大型矩阵（形状：2e6, 6）。在我搜索矩阵中的值以分配我的变量之前，我有 3 个 for 循环。

我的第一个解决方案是使用np.where。实在是太慢了！我阅读了it would be better 以使用另一个for 循环来提高性能。但是，我想出的代码更慢。

请问有人知道如何提高性能吗？

第一个解决方案 (np.where)

for lat in LATS:
    for lon in LONS:
        for depth in range(1,401,1):

            node_point_line = matrix[np.where( (matrix[:,0]==lat) * (matrix[:,1]==lon) * (matrix[:,2]==depth) )][0]

            var1 = node_point_line[3]
            var2 = node_point_line[4]
            var3 = node_point_line[5]
            ...

第二个解决方案（额外的for 循环）

for lat in LATS:
    for lon in LONS:
        for depth in range(1,401,1):

            matrix_flat = matrix.flatten()
            for i in range( len( matrix_flat )):
                if matrix_flat[i]==lat and matrix_flat[i+1]==lon and matrix_flat[i+2]==depth:
                    var1 = matrix_flat[i+3]
                    var2 = matrix_flat[i+4]
                    var3 = matrix_flat[i+5]
                    ...

同样，这两种解决方案都太慢了。我避免使用 Fortran 或 C++（我知道它更快）。有什么建议吗？

Answer 1

您可以使用矢量化过程压缩起始矩阵，这将使您的循环更快。

DPTH=np.arange(1,401,1)
mask=np.in1d(matrix[:,0],LATS) * np.in1d(matrix[:,1],LONS) * np.in1d(matrix[:,2],DPTH)
matrix_masked=matrix[mask]

然后只需 for 循环 matrix_masked 或 nditer（复杂）或

for i in range(matrix_masked.shape[0]):
    var1 = matrix_masked[i,3]
   . . . 

Answer 2

考虑到数字，仅lexsort 您的数据数组应该是相当有效的。这是在开始删除所有 where 语句时投入的 2 秒。

代码 sn-p 假定 LATS 和 LONS 已排序：

order = np.lexsort([matrix.T][2::-1] # takes < 2 sec
sorted = matrix[order, :]

data_per_lat = np.split(sorted, sorted[:, 0].searchsorted(LATS[1:], axis=0)
for lat, dpla in zip(LATS, data_per_lat):
    data_per_lon = np.split(dpla, dpla[:, 1].searchsorted[LONS[1:], axis=0)
    for lon, dplo in zip(LONS, data_per_lon):
        depth, var1, var2, var3 = dplo.T[2:]
        # the variables are now vectors, containing all data matching lon and lat sorted by depth

Answer 3

您的 LATS 和 LONS 数组有多大？即使它们是一个元素长，那么您的程序也必须执行大约 1*1*400*1e6*3 ~ 1.2e9 操作，这太多了。即使在 C++ 中正确实现，也可能需要一秒钟或更长时间。

我认为你应该优化你的算法。目前还不清楚您要做什么，但我的猜测是您正在尝试在矩阵中找到任何位置，该位置在您的 LATS 列表中具有纬度，在您的 LONS 列表中具有经度并且深度介于 1 和400.交换循环的位置：沿着外循环中的矩阵运行，然后检查所有三个条件是否成立。之后，您就可以用set of latitudes 替换纬度列表，set 查找比列表查找快得多。