過去用ncl作array slice習(xí)慣了, 不同的dimension之間就是直接乘起來, 例如
>>> arr = new ((/3, 4, 5/), float)
>>> print(dimsizes(arr((/0, 1/), (/0, 1/), 0)))
(0) 2
(1) 2
即可得到一個(gè)2x2的數(shù)組
結(jié)果今天在用numpy處理數(shù)據(jù)的時(shí)候發(fā)現(xiàn)
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個(gè)數(shù)
一查, 原來是numpy的slice自帶這個(gè)broadcast... 就好像會(huì)把不同dim的index給zip起來再取數(shù)一樣...和ncl實(shí)屬不一樣
想了幾種解決方法
1. 用np.ix_
例如:
arr[np.ix_([0, 1], [0, 1], [0])] # get (2, 2, 1) ndarray
這里[0]是必須的, 所以shape會(huì)變成(2, 2, 1), 也挺煩人, 賦值的話得給右邊的把維度對(duì)應(yīng), 有時(shí)候加一個(gè)1維啥的...例如:
arr[np.ix_([0, 1], [0, 1], [0])] = ncf.variables[var][index1, index2, None]
ncf是netCDF4.Dataset, var是個(gè)二維數(shù)組, 但是得加個(gè)None才能合法賦值
2. 用xarray
import xarray as xr
xrd = xr.DataArray(arr)
xrd[[0, 1], [0, 1], 0] # get (2, 2) xr.DataArray
這么看xarray還是挺不錯(cuò)的!
參考:
