将 numpy 数组映射到网格答案

【问题标题】：Map numpy array to mesh将 numpy 数组映射到网格
【发布时间】：2018-08-31 17:36:33
【问题描述】：

数据已存储在具有 3 个索引（对应于空间坐标）的 numpy.array 中。要将其传递给某些 python 模块，需要重新组织（和重塑）数据数组以对应于坐标数组，例如应该只有一个索引对应于 3d（空间）向量。我让它工作了，但是速度很慢（而且我的数组大小可能为 100^3 而不是 MWE 中的 10^3）。这是 MWE，如果可能的话，我想就如何（显着）加快速度获得建议：

import numpy as np
import random

# Create regular coordinates
n = 10
minval = -1
maxval = 1
x = y = z = np.linspace(minval, maxval, n+1)

# Create shuffled coordinate array
mesh = np.meshgrid(x, y, z)
coords = list(zip(*(dim.flat for dim in mesh)))
random.shuffle(coords)
coords = np.array(coords)

# Create data
data = np.zeros((n+1, n+1, n+1))
for i in range(n+1):
    for j in range(n+1):
        for k in range(n+1):
            data[i, j, k] = x[i]-2*y[j]+3*z[k]**3

# Reorganize data to correspond to shuffled coordinate array
data2 = np.zeros(coords.shape[0],)
for i in range(n+1):
    for j in range(n+1):
        for k in range(n+1):
            coord = np.array([x[i], y[j], z[k]])
            idx = [np.isclose(coord, xx).all() for xx in coords].index(True)
            data2[idx] = data[i, j, k]

【问题讨论】：

而您只尝试优化# Reorganize data to correspond to shuffled coordinate array 部分？
是的，这就是需要优化的地方。
摆脱所有列表和列表推导并使用 Numba 编译您的代码。这应该会加快大约 2-3 个数量级。示例：stackoverflow.com/a/49011917/4045774

标签： python arrays performance numpy

【解决方案1】：

一种方法是从x,y,z 生成组合网格，然后使用kd-tree based quick nearest-neighbor lookup 追溯索引。因此，我们将有一个矢量化解决方案，就像这样 -

from itertools import product
from scipy.spatial import cKDTree

# Form all combinations    
combs = np.array(list(product(x, y, z)))

# Find closest indices for each pt in coords
closest_idx = cKDTree(combs).query(coords, k=1)[1]

# Index into data array and retrieve the output
data2Out = data.ravel()[closest_idx]

节省内存

我们可以通过分别在三个坐标上使用 kd-tree 计算来优化内存，从而避免组合创建部分并有望提高性能，就像这样 -

I = cKDTree(x[:,None]).query(coords[:,[0]], k=1)[1]
J = cKDTree(y[:,None]).query(coords[:,[1]], k=1)[1]
K = cKDTree(z[:,None]).query(coords[:,[2]], k=1)[1]
data2Out = data[I,J,K]

【讨论】：