如何用NumPy对非结构化坐标计算梯度？

问题背景 #

现有形状为(NX, NY, NZ)的3D数据数组A，需计算其在y维度的梯度。当Y为1D坐标向量时，可直接通过NumPy实现：

import numpy as np
dAdy = np.gradient(A, Y, axis=1)

但当Y为非结构化坐标时，每个(x,z)=(Xi,Zi)对应的y列拥有唯一的单调坐标集合（示例构造如下）：

A = np.random.random((10, 10, 10))
X = np.arange(10)
Y = np.sort(np.random.random((10, 10, 10)), axis=1)
Z = np.arange(10)

此时需对每个独立的(x,z)列计算梯度，手动迭代方法速度极慢：

NX, NY, NZ = A.shape[0], A.shape[1], A.shape[2]
dA_dy = np.zeros((NX, NY, NZ))
for i in range(NX):
    for k in range(NZ):
        dA_dy[i, :, k] = np.gradient(A[i,:,k], Y[i,:,k])

尝试用链式规则优化的方法在测试中失效：

g = np.array([1, 5, 6, 10])  # 非结构化坐标
f = g**2                     # 函数值
grad1 = np.gradient(f, g)                # 正确的df/dg
grad2 = np.gradient(f) / np.gradient(g)  # 错误结果

仅当函数为线性时两者结果一致，需明确向量化实现方案及链式规则失效的理论原因。

向量化实现方案 #

方法1：数组展平+批量处理 #

将A和Y的x、z维度展平，把所有独立的y列转化为二维数组的行，再批量计算梯度：

# 将(NX, NY, NZ)形状展平为(NX*NZ, NY)
A_reshaped = A.reshape(-1, NY)
Y_reshaped = Y.reshape(-1, NY)

# 对每一行计算梯度
dA_dy_reshaped = np.array([np.gradient(a_col, y_col) for a_col, y_col in zip(A_reshaped, Y_reshaped)])

# 恢复原始3D形状
dA_dy = dA_dy_reshaped.reshape(NX, NY, NZ)

方法2：NumPy向量化函数 #

利用np.vectorize封装单列梯度计算逻辑，代码更简洁：

from numpy import vectorize

def single_col_grad(a_col, y_col):
    return np.gradient(a_col, y_col)

# 指定输入输出的数组形状签名，确保向量化正确执行
vec_grad = vectorize(single_col_grad, signature='(n),(n)->(n)')
dA_dy = vec_grad(A, Y)

方法3：Numba加速循环 #

若追求极致性能，用Numba编译原始循环代码，可大幅提升计算速度：

from numba import jit

@jit(nopython=True)
def numba_grad(A, Y):
    NX, NY, NZ = A.shape
    dA_dy = np.zeros_like(A)
    for i in range(NX):
        for k in range(NZ):
            dA_dy[i, :, k] = np.gradient(A[i,:,k], Y[i,:,k])
    return dA_dy

dA_dy = numba_grad(A, Y)

链式规则失效的理论原因 #

np.gradient(f, g)与np.gradient(f)/np.gradient(g)的核心差异在于对非均匀坐标的处理逻辑：

1. np.gradient(f, g)的计算逻辑：
   针对非均匀间隔坐标，该函数会采用加权差分公式计算梯度：

   • 端点使用前向/后向差分：(f[1]-f[0])/(g[1]-g[0]) 或 (f[-1]-f[-2])/(g[-1]-g[-2])
   • 中间点综合左右两侧的差分信息，通过坐标间距加权：

     grad[i] = [(f[i+1]-f[i])/(g[i+1]-g[i])*(g[i]-g[i-1]) + (f[i]-f[i-1])/(g[i]-g[i-1])*(g[i+1]-g[i])] / (g[i+1]-g[i-1])
     

   该逻辑保证了非均匀坐标下的数值精度。

2. np.gradient(f)/np.gradient(g)的计算逻辑：
   此方法先按默认均匀步长（步长=1）计算f和g在索引维度的梯度，再做除法。它完全忽略了坐标的非均匀性，等价于假设坐标是等间隔分布的。

只有当函数为线性时，两种方法结果一致——因为线性函数的差分在任何间隔下都是恒定值，加权逻辑不会改变结果；对于非线性函数，加权差分与均匀步长差分的比值会出现偏差，导致结果错误。

内容的提问来源于stack exchange，提问作者user5162426