numpy slice broadcast的坑, 和ncl或者x

过去用ncl作array slice习惯了, 不同的dimension之间就是直接乘起来, 例如

>>> arr = new ((/3, 4, 5/), float)
>>> print(dimsizes(arr((/0, 1/), (/0, 1/), 0)))
(0)     2
(1)     2

即可得到一个2x2的数组

结果今天在用numpy处理数据的时候发现

import numpy as np

arr = np.arange(60).reshape(3, 4, 5)

arr[[0, 1], [0, 1], 0] # array([0, 25])

arr[[0, 1, 2], [0, 1], 0] # IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (3,) (2,)

就只能得到2个数

一查, 原来是numpy的slice自带这个broadcast... 就好像会把不同dim的index给zip起来再取数一样...和ncl实属不一样

想了几种解决方法

1. 用np.ix_

例如:

arr[np.ix_([0, 1], [0, 1], [0])] # get (2, 2, 1) ndarray

这里[0]是必须的, 所以shape会变成(2, 2, 1), 也挺烦人, 赋值的话得给右边的把维度对应, 有时候加一个1维啥的...例如:

arr[np.ix_([0, 1], [0, 1], [0])] = ncf.variables[var][index1, index2, None]

ncf是netCDF4.Dataset, var是个二维数组, 但是得加个None才能合法赋值

2. 用xarray

import xarray as xr

xrd = xr.DataArray(arr)
xrd[[0, 1], [0, 1], 0]  # get (2, 2) xr.DataArray

这么看xarray还是挺不错的!

参考:

• Numpy extract submatrix
• Selecting specific rows and columns from NumPy array