python - 沿轴插入 numpy.ndarray 的最佳方法

Question

我在 numpy.ndarray 中有 4 维数据，比如温度。数组的形状是(ntime, nheight_in, nlat, nlon)。

我对每个维度都有对应的一维数组，告诉我某个值对应于哪个时间、高度、纬度和经度，对于这个例子，我需要 height_in 给出以米为单位的高度。

现在我需要将它放到不同的高度维度上，height_out，具有不同的长度。

以下内容似乎符合我的要求:

ntime, nheight_in, nlat, nlon = t_in.shape

nheight_out = len(height_out)
t_out = np.empty((ntime, nheight_out, nlat, nlon))

for time in range(ntime):
    for lat in range(nlat):
        for lon in range(nlon):
            t_out[time, :, lat, lon] = np.interp(
                height_out, height_in, t[time, :, lat, lon]
            )

但是有 3 个嵌套循环，并且在 python 和 numpy 之间进行了大量切换，我认为这不是最好的方法。

有什么改进建议吗？谢谢

Answer 1

scipy 的 interp1d可以帮助:

import numpy as np
from scipy.interpolate import interp1d

ntime, nheight_in, nlat, nlon = (10, 20, 30, 40)

heights = np.linspace(0, 1, nheight_in)

t_in = np.random.normal(size=(ntime, nheight_in, nlat, nlon))
f_out = interp1d(heights, t_in, axis=1)

nheight_out = 50
new_heights = np.linspace(0, 1, nheight_out)
t_out = f_out(new_heights)