对所有维度进行 Numpy 迭代，但最后一个维度数未知
Numpy iteration over all dimensions but the last one with unknown number of dimensions
身体背景
我正在研究一个函数，该函数计算最多四维温度场（时间、经度、纬度、压力作为高度测量）中每个垂直剖面的一些指标。我有一个工作函数，它在一个位置获取压力和温度，并 returns 指标（对流层顶信息）。我想用一个函数包装它，将它应用于传递的数据中的每个垂直剖面。
问题的技术描述
我希望我的函数将另一个函数应用于对应于我的 N-dimensional 数组中最后一个维度的每个一维数组，其中 N <= 4。所以我需要一个高效的循环遍历除最后一个维度之外的所有维度事先不知道维数。
为什么我打开一个新问题
我知道有几个问题（例如，iterating over some dimensions of a ndarray, , , Iterating over a numpy matrix with unknown dimension）询问如何遍历特定维度或如何遍历未知维度的数组。据我所知，这两个问题的结合是新的。以 numpy.nditer 为例，我还没有找到如何只排除最后一个维度而不考虑剩余维度的数量。
编辑
我尝试做一个最小的、可重现的例子：
import numpy as np
def outer_function(array, *args):
"""
Array can be 1D, 2D, 3D, or 4D. Regardless the inner_function
should be applied to all 1D arrays spanned by the last axis
"""
# Unpythonic if-else solution
if array.ndim == 1:
return inner_function(array)
elif array.ndim == 2:
return [inner_function(array[i,:]) for i in range(array.shape[0])]
elif array.ndim == 3:
return [[inner_function(array[i,j,:]) for i in range(array.shape[0])] for j in range(array.shape[1])]
elif array.ndim == 4:
return [[[inner_function(array[i,j,k,:]) for i in range(array.shape[0])] for j in range(array.shape[1])] for k in range(array.shape[2])]
else:
return -1
def inner_function(array_1d):
return np.interp(2, np.arange(array_1d.shape[0]), array_1d), np.sum(array_1d)
请假设无法修改实际的 inner_function 以应用于多个维度，而只能应用于一维数组。
编辑结束
如果它有助于我 have/want 的代码结构：
def tropopause_ds(ds):
"""
wraps around tropopause profile calculation. The vertical coordinate has to be the last one.
"""
t = ds.t.values # numpy ndarray
p_profile = ds.plev.values # 1d numpy ndarray
len_t = ds.time.size
len_lon = ds.lon.size
len_lat = ds.lat.size
nlevs = ds.plev.size
ttp = np.empty([len_t, len_lon, len_lat])
ptp = np.empty([len_t, len_lon, len_lat])
ztp = np.empty([len_t, len_lon, len_lat])
dztp = np.empty([len_t, len_lon, len_lat, nlevs])
# Approach 1: use numpy.ndindex - doesn't work in a list comprehension, slow
for idx in np.ndindex(*t.shape[:-1]):
ttp[idx], ptp[idx], ztp[idx], dztp[idx] = tropopause_profile(t[idx], p_profile)
# Approach 2: use nested list comprehensions - doesn't work for different number of dimensions
ttp, ptp, ztp, dztp = [[[tropopause_profile(t[i,j,k,:], p_profile) for k in range(len_lat)]
for j in range(len_lon)] for i in range(len_t)]
return ttp, ptp, ztp, dztp
内部函数结构如下：
def tropopause_profile(t_profile, p_profile):
if tropopause found:
return ttp, ptp, ztp, dztp
return np.nan, np.nan, np.nan, np.nan
我已经尝试了几种选择。计时案例中的测试数据的形状为 (2, 360, 180, 105):
- xarray's apply_ufunc 似乎将整个数组传递给函数。然而，我的内部函数是基于获取一维数组并且很难重新编程以处理 multi-dimensional 数据
- 嵌套列表推导 工作并且似乎相当快但是如果一个维度（例如时间）只有一个值（timed：每个循环 8.53 s ± 11.9 ms（7 次运行的平均值±标准偏差，每次 1 个循环））
- using numpy's nditer 在标准 for 循环中工作，该循环使用列表理解加速。然而，使用这种方法，该函数不是 return 4 个 ndarrays，而是一个包含每个索引的四个 return 值作为列表元素的列表。 （timed with list comprehension: 1min 4s ± 740 ms per loop (mean ± std.dev. of 7 runs, 1 loop each))
解决这个问题的一个丑陋的方法是检查我的数据有多少维，然后对正确数量的列表推导进行 if else 选择，但我希望 python 有更流畅的解决方法这个。如果有帮助，可以轻松更改维度的顺序。我 运行 2 核，10 GB 内存 jupyterhub 服务器上的代码。
我已经多次使用@hpaulj 的重塑方法。这意味着循环可以通过 1d 切片迭代整个数组。
简化了功能和数据，以便测试。
import numpy as np
arr = np.arange( 2*3*3*2*6 ).reshape( 2,3,3,2,6 )
def inner_function(array_1d):
return np.array( [ array_1d.sum(), array_1d.mean() ])
# return np.array( [np.interp(2, np.arange(array_1d.shape[0]), array_1d), np.sum(array_1d) ])
def outer_function( arr, *args ):
res_shape = list( arr.shape )
res_shape[ -1 ] = 2
result = np.zeros( tuple( res_shape ) ) # result has the same shape as arr for n-1 dimensions, then two
# Reshape arr and result to be 2D arrays. These are views into arr and result
work = arr.reshape( -1, arr.shape[-1] )
res = result.reshape( -1, result.shape[-1] )
for ix, w1d in enumerate( work ): # Loop through all 1D
res[ix] = inner_function( w1d )
return result
outer_function( arr )
结果是
array([[[[[ 15. , 2.5],
[ 51. , 8.5]],
[[ 87. , 14.5],
[ 123. , 20.5]],
...
[[1167. , 194.5],
[1203. , 200.5]],
[[1239. , 206.5],
[1275. , 212.5]]]]])
我相信这可以进一步优化，并考虑到应用程序所需的实际功能。
身体背景
我正在研究一个函数，该函数计算最多四维温度场（时间、经度、纬度、压力作为高度测量）中每个垂直剖面的一些指标。我有一个工作函数，它在一个位置获取压力和温度，并 returns 指标（对流层顶信息）。我想用一个函数包装它，将它应用于传递的数据中的每个垂直剖面。
问题的技术描述
我希望我的函数将另一个函数应用于对应于我的 N-dimensional 数组中最后一个维度的每个一维数组，其中 N <= 4。所以我需要一个高效的循环遍历除最后一个维度之外的所有维度事先不知道维数。
为什么我打开一个新问题
我知道有几个问题（例如，iterating over some dimensions of a ndarray,
编辑
我尝试做一个最小的、可重现的例子：
import numpy as np
def outer_function(array, *args):
"""
Array can be 1D, 2D, 3D, or 4D. Regardless the inner_function
should be applied to all 1D arrays spanned by the last axis
"""
# Unpythonic if-else solution
if array.ndim == 1:
return inner_function(array)
elif array.ndim == 2:
return [inner_function(array[i,:]) for i in range(array.shape[0])]
elif array.ndim == 3:
return [[inner_function(array[i,j,:]) for i in range(array.shape[0])] for j in range(array.shape[1])]
elif array.ndim == 4:
return [[[inner_function(array[i,j,k,:]) for i in range(array.shape[0])] for j in range(array.shape[1])] for k in range(array.shape[2])]
else:
return -1
def inner_function(array_1d):
return np.interp(2, np.arange(array_1d.shape[0]), array_1d), np.sum(array_1d)
请假设无法修改实际的 inner_function 以应用于多个维度，而只能应用于一维数组。
编辑结束
如果它有助于我 have/want 的代码结构：
def tropopause_ds(ds):
"""
wraps around tropopause profile calculation. The vertical coordinate has to be the last one.
"""
t = ds.t.values # numpy ndarray
p_profile = ds.plev.values # 1d numpy ndarray
len_t = ds.time.size
len_lon = ds.lon.size
len_lat = ds.lat.size
nlevs = ds.plev.size
ttp = np.empty([len_t, len_lon, len_lat])
ptp = np.empty([len_t, len_lon, len_lat])
ztp = np.empty([len_t, len_lon, len_lat])
dztp = np.empty([len_t, len_lon, len_lat, nlevs])
# Approach 1: use numpy.ndindex - doesn't work in a list comprehension, slow
for idx in np.ndindex(*t.shape[:-1]):
ttp[idx], ptp[idx], ztp[idx], dztp[idx] = tropopause_profile(t[idx], p_profile)
# Approach 2: use nested list comprehensions - doesn't work for different number of dimensions
ttp, ptp, ztp, dztp = [[[tropopause_profile(t[i,j,k,:], p_profile) for k in range(len_lat)]
for j in range(len_lon)] for i in range(len_t)]
return ttp, ptp, ztp, dztp
内部函数结构如下：
def tropopause_profile(t_profile, p_profile):
if tropopause found:
return ttp, ptp, ztp, dztp
return np.nan, np.nan, np.nan, np.nan
我已经尝试了几种选择。计时案例中的测试数据的形状为 (2, 360, 180, 105):
- xarray's apply_ufunc 似乎将整个数组传递给函数。然而，我的内部函数是基于获取一维数组并且很难重新编程以处理 multi-dimensional 数据
- 嵌套列表推导 工作并且似乎相当快但是如果一个维度（例如时间）只有一个值（timed：每个循环 8.53 s ± 11.9 ms（7 次运行的平均值±标准偏差，每次 1 个循环））
- using numpy's nditer 在标准 for 循环中工作，该循环使用列表理解加速。然而，使用这种方法，该函数不是 return 4 个 ndarrays，而是一个包含每个索引的四个 return 值作为列表元素的列表。 （timed with list comprehension: 1min 4s ± 740 ms per loop (mean ± std.dev. of 7 runs, 1 loop each))
解决这个问题的一个丑陋的方法是检查我的数据有多少维，然后对正确数量的列表推导进行 if else 选择，但我希望 python 有更流畅的解决方法这个。如果有帮助，可以轻松更改维度的顺序。我 运行 2 核，10 GB 内存 jupyterhub 服务器上的代码。
我已经多次使用@hpaulj 的重塑方法。这意味着循环可以通过 1d 切片迭代整个数组。
简化了功能和数据，以便测试。
import numpy as np
arr = np.arange( 2*3*3*2*6 ).reshape( 2,3,3,2,6 )
def inner_function(array_1d):
return np.array( [ array_1d.sum(), array_1d.mean() ])
# return np.array( [np.interp(2, np.arange(array_1d.shape[0]), array_1d), np.sum(array_1d) ])
def outer_function( arr, *args ):
res_shape = list( arr.shape )
res_shape[ -1 ] = 2
result = np.zeros( tuple( res_shape ) ) # result has the same shape as arr for n-1 dimensions, then two
# Reshape arr and result to be 2D arrays. These are views into arr and result
work = arr.reshape( -1, arr.shape[-1] )
res = result.reshape( -1, result.shape[-1] )
for ix, w1d in enumerate( work ): # Loop through all 1D
res[ix] = inner_function( w1d )
return result
outer_function( arr )
结果是
array([[[[[ 15. , 2.5],
[ 51. , 8.5]],
[[ 87. , 14.5],
[ 123. , 20.5]],
...
[[1167. , 194.5],
[1203. , 200.5]],
[[1239. , 206.5],
[1275. , 212.5]]]]])
我相信这可以进一步优化，并考虑到应用程序所需的实际功能。