python - 通过从 3D 数组中采样和分桶来创建热图

Question

我有一些这样存在的实验数据:

x = array([1, 1.12, 1.109, 2.1, 3, 4.104, 3.1, ...])
y = array([-9, -0.1, -9.2, -8.7, -5, -4, -8.75, ...])
z = array([10, 4, 1, 4, 5, 0, 1, ...])

如果方便的话，我们可以假设数据以 3D 数组甚至 pandas DataFrame 的形式存在:

df = pd.DataFrame({'x': x, 'y': y, 'z': z})

解释是，对于每个位置 x[i], y[i]，某个变量的值是 z[i]。这些是不均匀采样，因此会有一些部分“密集采样”(例如 x 中的 1 到 1.2 之间)和其他非常稀疏的部分(例如介于x 中的 2 和 3)。正因为如此，我不能直接将它们放入 pcolormesh 或 contourf 中。

我想做的是在某个固定间隔内均匀地重新采样 x 和 y，然后聚合 z 的值。根据我的需要，可以对 z 求和或取平均值以获得有意义的值，因此这不是问题。我天真的尝试是这样的:

X = np.arange(min(x), max(x), 0.1)  
Y = np.arange(min(y), max(y), 0.1)
x_g, y_g = np.meshgrid(X, Y)
nx, ny = x_g.shape
z_g = np.full(x_g.shape, np.nan)

for ix in range(nx - 1):
    for jx in range(ny - 1):
        x_min = x_g[ix, jx]
        x_max = x_g[ix + 1, jx + 1]
        y_min = y_g[ix, jx]
        y_max = y_g[ix + 1, jx + 1]
        vals = df[(df.x >= x_min) & (df.x < x_max) & 
                  (df.y >= y_min) & (df.y < y_max)].z.values
        if vals.any():
            z_g[ix, jx] = sum(vals)

这有效，我得到了我想要的输出，使用 plt.contourf(x_g, y_g, z_g) 但它很慢!我有大约 20k 个样本，然后我将其子采样为 x 方向的 800 个样本和 y 方向的 500 个样本，这意味着 for 循环的长度为 400k。

有什么方法可以对其进行矢量化/优化吗？如果已经有一些功能可以做到这一点，那就更好了!

(也将其标记为 MATLAB，因为 numpy/MATLAB 之间的语法非常相似，而且我可以访问这两个软件。)

Answer 1

这是一个使用 NumPy broadcasting 的矢量化 Python 解决方案和 矩阵乘法 与 np.dot对于总和减少部分 -

x_mask = ((x >= X[:-1,None]) & (x < X[1:,None]))
y_mask = ((y >= Y[:-1,None]) & (y < Y[1:,None]))

z_g_out = np.dot(y_mask*z[None].astype(np.float32), x_mask.T)

# If needed to fill invalid places with NaNs
z_g_out[y_mask.dot(x_mask.T.astype(np.float32))==0] = np.nan

请注意，我们避免在那里使用 meshgrid。因此，在使用 meshgrid 创建的网格时节省内存将是巨大的，并且有望在此过程中获得性能改进。

基准测试

# Original app
def org_app(x,y,z):    
    X = np.arange(min(x), max(x), 0.1)  
    Y = np.arange(min(y), max(y), 0.1)
    x_g, y_g = np.meshgrid(X, Y)
    nx, ny = x_g.shape
    z_g = np.full(np.asarray(x_g.shape)-1, np.nan)

    for ix in range(nx - 1):
        for jx in range(ny - 1):
            x_min = x_g[ix, jx]
            x_max = x_g[ix + 1, jx + 1]
            y_min = y_g[ix, jx]
            y_max = y_g[ix + 1, jx + 1]
            vals = z[(x >= x_min) & (x < x_max) & 
                      (y >= y_min) & (y < y_max)]
            if vals.any():
                z_g[ix, jx] = sum(vals)
    return z_g

# Proposed app
def app1(x,y,z):
    X = np.arange(min(x), max(x), 0.1)  
    Y = np.arange(min(y), max(y), 0.1)
    x_mask = ((x >= X[:-1,None]) & (x < X[1:,None]))
    y_mask = ((y >= Y[:-1,None]) & (y < Y[1:,None]))

    z_g_out = np.dot(y_mask*z[None].astype(np.float32), x_mask.T)

    # If needed to fill invalid places with NaNs
    z_g_out[y_mask.dot(x_mask.T.astype(np.float32))==0] = np.nan
    return z_g_out

如上所示，为了公平的基准测试，我在原始方法中使用数组值，因为从数据帧中获取值可能会减慢速度。

时间和验证-

In [143]: x = np.array([1, 1.12, 1.109, 2.1, 3, 4.104, 3.1])
     ...: y = np.array([-9, -0.1, -9.2, -8.7, -5, -4, -8.75])
     ...: z = np.array([10, 4, 1, 4, 5, 0, 1])
     ...: 

# Verify outputs
In [150]: np.nansum(np.abs(org_app(x,y,z) - app1(x,y,z)))
Out[150]: 0.0

In [145]: %timeit org_app(x,y,z)
10 loops, best of 3: 19.9 ms per loop

In [146]: %timeit app1(x,y,z)
10000 loops, best of 3: 39.1 µs per loop

In [147]: 19900/39.1  # Speedup figure
Out[147]: 508.95140664961633