Numpy - 沿轴取平均值，但在每个单元格中具有该轴的不同子集
Numpy - get mean along axis but with different subset of that axis in each cell
我需要数组 (1) 沿时间轴的平均值。
要点：它不会是沿该轴的所有值的平均值，而是从数组 (2) 中给出的索引开始的子集。
我正在使用的数组：
(array1) 3 axes: time, x, y
array([[[ 820, 820, 720, 720],
[ 860, 860, 500, 500],
[ 860, 860, 500, 500],
[ 860, 860, 500, 500]],
[[5980, 5980, 4760, 4760],
[7500, 7500, 7940, 7940],
[7500, 7500, 7940, 7940],
[7500, 7500, 7940, 7940]],
[[ 740, 740, 440, 440],
[1240, 1240, 1140, 1140],
[1240, 1240, 1140, 1140],
[1240, 1240, 1140, 1140]],
[[3200, 3200, 7600, 7600],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400]]])
(array2) 2 axes: x, y
array([[ 1, 2, 1, 1],
[ 1, 0, 3, 3],
[ 4, 0, 2, 2],
[ 4, 0, 1, 2]])
进一步说明示例：
array1 中的值表示位置 x/y 每天的降雨量。
array2 中的值表示需要从哪一天计算位置 x/y.
的平均值
查看第一个单元格，我们将从计算中排除第一天，因为 array2[0,0] = 1。因此我们的结果将是 np.mean(array1[1:, 0, 0 ]) = 3306.67.
我想不通的是如何根据数组 2 为每个单元格指定子集。我知道我可以沿任何轴使用 np.mean，但我如何动态排除值 (从计算中切片数组）？
我找到了使用 xarray 的方法。
它不是很漂亮，但至少我相信它是矢量化的。
首先，将 numpy 数组转换为 xarray DataArray，然后使用 xr.merge:
将两者放入 Dataset
rainfall = xr.DataArray(rainfall, dims=("day", "x", "y"), name="rainfall")
start_idxs = xr.DataArray(start_idxs, dims=("x", "y"), name="start_idxs")
ds = xr.merge((rainfall, start_idxs))
这是 Dataset 的样子：
>>> ds
<xarray.Dataset>
Dimensions: (day: 4, x: 4, y: 4)
Dimensions without coordinates: day, x, y
Data variables:
rainfall (day, x, y) int64 820 820 720 720 860 ... 400 900 900 400 400
start_idxs (x, y) int64 1 2 1 1 1 0 3 3 4 0 2 2 4 0 1 2
然后我们要根据start_idxs变量进行不同的计算，所以我们将groupby这个数据变量：
>>> groups = ds.groupby("start_idxs")
>>> groups
DatasetGroupBy, grouped over 'start_idxs'
5 groups with labels 0, 1, 2, 3, 4.
您看到如预期的那样有 5 个组。现在我们要为每个组应用计算，所以我们将使用 map.
>>> res = groups.map(mean_start_idxs, args=("day",))
>>> res["rainfall"]
<xarray.DataArray 'rainfall' (x: 4, y: 4)>
array([[3306.66666667, 1970. , 4266.66666667, 4266.66666667],
[3213.33333333, 2625. , 400. , 400. ],
[ nan, 2625. , 770. , 770. ],
[ nan, 2625. , 3160. , 770. ]])
Dimensions without coordinates: x, y
这是预期的输出。注意 nan 要求从索引 4 开始计算平均值的值，只有 4 天是不可能的。
但要实现这一点，我们需要定义 mean_start_idxs 函数，这是棘手的部分。
这很棘手，因为从 map 调用的函数内部检索组的“标签”并不简单，但这里有一个解决方案：
def mean_start_idxs(ds, dim):
# Get the start indice
# groups were made from start_idxs, so we can
# take any value of ds["start_idxs"] as a start indice
start = ds["start_idxs"][0].item()
end = ds.sizes[dim]
return ds.isel({dim: slice(start, end)}).mean(dim=dim)
arr1 = np.array(
[[[ 820, 820, 720, 720],
[ 860, 860, 500, 500],
[ 860, 860, 500, 500],
[ 860, 860, 500, 500]],
[[5980, 5980, 4760, 4760],
[7500, 7500, 7940, 7940],
[7500, 7500, 7940, 7940],
[7500, 7500, 7940, 7940]],
[[ 740, 740, 440, 440],
[1240, 1240, 1140, 1140],
[1240, 1240, 1140, 1140],
[1240, 1240, 1140, 1140]],
[[3200, 3200, 7600, 7600],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400]]]
)
arr2 = np.array(
[[ 1, 2, 1, 1],
[ 1, 0, 3, 3],
[ 3, 0, 2, 2],
[ 3, 0, 1, 2]]
)
我们要做的是使用 arr2 中存储的索引对 arr1 的时间轴进行切片，现在 python 只允许使用 : 进行切片我们只能在字面索引时通过，即不使用另一个数组进行索引。所以我们需要轮流做这件事
一种方法是将 arr1 中的所有值更改为 0
现在要找到要忽略的值的索引，我们这样做
no_days = arr1.shape[0]
arr3 = np.arange(no_days)
arr3.shape = [-1,1,1]
arr3
>>> [[[0]],
[[1]],
[[2]],
[[3]]]
filter = arr3 < arr2
filter.shape
>>> (4, 4, 4)
arr3是时间轴索引数组。我们将它与 arr2 进行了比较，现在我们在 filter 中有了要忽略的布尔值索引，我们可以将它们设置为 0
arr1[filter] = 0
arr1
>>> [[[ 0, 0, 0, 0],
[ 0, 860, 0, 0],
[ 0, 860, 0, 0],
[ 0, 860, 0, 0]],
[[5980, 0, 4760, 4760],
[7500, 7500, 0, 0],
[ 0, 7500, 0, 0],
[ 0, 7500, 7940, 0]],
[[ 740, 740, 440, 440],
[1240, 1240, 0, 0],
[ 0, 1240, 1140, 1140],
[ 0, 1240, 1140, 1140]],
[[3200, 3200, 7600, 7600],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400]]]
我们可能会想使用 arr1.mean(axis= 0)，但这样做也会考虑所有 0 影响均值的合法条目，而不是忽略它们
所以我们在时间轴上对 arr1 求和，然后除以切片中没有的元素
arr1.sum(axis= 0) / (no_days - arr2)
>>> [[3306.66666667, 1970. , 4266.66666667, 4266.66666667],
[3213.33333333, 2625. , 400. , 400. ],
[ 900. , 2625. , 770. , 770. ],
[ 900. , 2625. , 3160. , 770. ]]
如果 t < x*y 那么下面的代码会执行得更快
arr1.sum(axis= 0) / (~filter).astype(int).sum(axis= 0)
我需要数组 (1) 沿时间轴的平均值。
要点：它不会是沿该轴的所有值的平均值，而是从数组 (2) 中给出的索引开始的子集。
我正在使用的数组：
(array1) 3 axes: time, x, y
array([[[ 820, 820, 720, 720],
[ 860, 860, 500, 500],
[ 860, 860, 500, 500],
[ 860, 860, 500, 500]],
[[5980, 5980, 4760, 4760],
[7500, 7500, 7940, 7940],
[7500, 7500, 7940, 7940],
[7500, 7500, 7940, 7940]],
[[ 740, 740, 440, 440],
[1240, 1240, 1140, 1140],
[1240, 1240, 1140, 1140],
[1240, 1240, 1140, 1140]],
[[3200, 3200, 7600, 7600],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400]]])
(array2) 2 axes: x, y
array([[ 1, 2, 1, 1],
[ 1, 0, 3, 3],
[ 4, 0, 2, 2],
[ 4, 0, 1, 2]])
进一步说明示例：
array1 中的值表示位置 x/y 每天的降雨量。 array2 中的值表示需要从哪一天计算位置 x/y.
的平均值查看第一个单元格，我们将从计算中排除第一天，因为 array2[0,0] = 1。因此我们的结果将是 np.mean(array1[1:, 0, 0 ]) = 3306.67.
我想不通的是如何根据数组 2 为每个单元格指定子集。我知道我可以沿任何轴使用 np.mean，但我如何动态排除值 (从计算中切片数组）？
我找到了使用 xarray 的方法。 它不是很漂亮，但至少我相信它是矢量化的。
首先，将 numpy 数组转换为 xarray DataArray，然后使用 xr.merge:
Dataset
rainfall = xr.DataArray(rainfall, dims=("day", "x", "y"), name="rainfall")
start_idxs = xr.DataArray(start_idxs, dims=("x", "y"), name="start_idxs")
ds = xr.merge((rainfall, start_idxs))
这是 Dataset 的样子：
>>> ds
<xarray.Dataset>
Dimensions: (day: 4, x: 4, y: 4)
Dimensions without coordinates: day, x, y
Data variables:
rainfall (day, x, y) int64 820 820 720 720 860 ... 400 900 900 400 400
start_idxs (x, y) int64 1 2 1 1 1 0 3 3 4 0 2 2 4 0 1 2
然后我们要根据start_idxs变量进行不同的计算，所以我们将groupby这个数据变量：
>>> groups = ds.groupby("start_idxs")
>>> groups
DatasetGroupBy, grouped over 'start_idxs'
5 groups with labels 0, 1, 2, 3, 4.
您看到如预期的那样有 5 个组。现在我们要为每个组应用计算，所以我们将使用 map.
>>> res = groups.map(mean_start_idxs, args=("day",))
>>> res["rainfall"]
<xarray.DataArray 'rainfall' (x: 4, y: 4)>
array([[3306.66666667, 1970. , 4266.66666667, 4266.66666667],
[3213.33333333, 2625. , 400. , 400. ],
[ nan, 2625. , 770. , 770. ],
[ nan, 2625. , 3160. , 770. ]])
Dimensions without coordinates: x, y
这是预期的输出。注意 nan 要求从索引 4 开始计算平均值的值，只有 4 天是不可能的。
但要实现这一点，我们需要定义 mean_start_idxs 函数，这是棘手的部分。
这很棘手，因为从 map 调用的函数内部检索组的“标签”并不简单，但这里有一个解决方案：
def mean_start_idxs(ds, dim):
# Get the start indice
# groups were made from start_idxs, so we can
# take any value of ds["start_idxs"] as a start indice
start = ds["start_idxs"][0].item()
end = ds.sizes[dim]
return ds.isel({dim: slice(start, end)}).mean(dim=dim)
arr1 = np.array(
[[[ 820, 820, 720, 720],
[ 860, 860, 500, 500],
[ 860, 860, 500, 500],
[ 860, 860, 500, 500]],
[[5980, 5980, 4760, 4760],
[7500, 7500, 7940, 7940],
[7500, 7500, 7940, 7940],
[7500, 7500, 7940, 7940]],
[[ 740, 740, 440, 440],
[1240, 1240, 1140, 1140],
[1240, 1240, 1140, 1140],
[1240, 1240, 1140, 1140]],
[[3200, 3200, 7600, 7600],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400]]]
)
arr2 = np.array(
[[ 1, 2, 1, 1],
[ 1, 0, 3, 3],
[ 3, 0, 2, 2],
[ 3, 0, 1, 2]]
)
我们要做的是使用 arr2 中存储的索引对 arr1 的时间轴进行切片，现在 python 只允许使用 : 进行切片我们只能在字面索引时通过，即不使用另一个数组进行索引。所以我们需要轮流做这件事
一种方法是将 arr1 中的所有值更改为 0
现在要找到要忽略的值的索引，我们这样做
no_days = arr1.shape[0]
arr3 = np.arange(no_days)
arr3.shape = [-1,1,1]
arr3
>>> [[[0]],
[[1]],
[[2]],
[[3]]]
filter = arr3 < arr2
filter.shape
>>> (4, 4, 4)
arr3是时间轴索引数组。我们将它与 arr2 进行了比较，现在我们在 filter 中有了要忽略的布尔值索引，我们可以将它们设置为 0
arr1[filter] = 0
arr1
>>> [[[ 0, 0, 0, 0],
[ 0, 860, 0, 0],
[ 0, 860, 0, 0],
[ 0, 860, 0, 0]],
[[5980, 0, 4760, 4760],
[7500, 7500, 0, 0],
[ 0, 7500, 0, 0],
[ 0, 7500, 7940, 0]],
[[ 740, 740, 440, 440],
[1240, 1240, 0, 0],
[ 0, 1240, 1140, 1140],
[ 0, 1240, 1140, 1140]],
[[3200, 3200, 7600, 7600],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400],
[ 900, 900, 400, 400]]]
我们可能会想使用 arr1.mean(axis= 0)，但这样做也会考虑所有 0 影响均值的合法条目，而不是忽略它们
所以我们在时间轴上对 arr1 求和，然后除以切片中没有的元素
arr1.sum(axis= 0) / (no_days - arr2)
>>> [[3306.66666667, 1970. , 4266.66666667, 4266.66666667],
[3213.33333333, 2625. , 400. , 400. ],
[ 900. , 2625. , 770. , 770. ],
[ 900. , 2625. , 3160. , 770. ]]
如果 t < x*y 那么下面的代码会执行得更快
arr1.sum(axis= 0) / (~filter).astype(int).sum(axis= 0)