过去用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还是挺不错的!
参考:
