def test_expanded_indexer(self): x = np.random.randn(10, 11, 12, 13, 14) y = np.arange(5) I = ReturnItem() # noqa for i in [I[:], I[...], I[0, :, 10], I[..., 10], I[:5, ..., 0], I[..., 0, :], I[y], I[y, y], I[..., y, y], I[..., 0, 1, 2, 3, 4]]: j = indexing.expanded_indexer(i, x.ndim) assert_array_equal(x[i], x[j]) assert_array_equal(self.set_to_zero(x, i), self.set_to_zero(x, j)) with raises_regex(IndexError, 'too many indices'): indexing.expanded_indexer(I[1, 2, 3], 2)
def test_expanded_indexer(self): x = np.random.randn(10, 11, 12, 13, 14) y = np.arange(5) I = ReturnItem() for i in [I[:], I[...], I[0, :, 10], I[..., 10], I[:5, ..., 0], I[..., 0, :], I[y], I[y, y], I[..., y, y], I[..., 0, 1, 2, 3, 4]]: j = indexing.expanded_indexer(i, x.ndim) self.assertArrayEqual(x[i], x[j]) self.assertArrayEqual(self.set_to_zero(x, i), self.set_to_zero(x, j)) with raises_regex(IndexError, 'too many indices'): indexing.expanded_indexer(I[1, 2, 3], 2)
def __getitem__(self, item): """Override __getitem__.""" # take care of ellipsis and other strange indexes item = list(indexing.expanded_indexer(item, len(self.dimensions))) # Slice to change was_scalar = False sl = item[self.ds] if np.isscalar(sl) and not isinstance(sl, slice): sl = slice(sl, sl + 1) was_scalar = True # Ok, get the indexes right start = sl.start or 0 stop = sl.stop or self._ds_shape if stop < 0: stop += self._ds_shape - 1 stop = np.clip(stop + 1, 0, self._ds_shape) itemr = copy.deepcopy(item) if was_scalar: item[self.ds] = start itemr[self.ds] = start + 1 else: item[self.ds] = slice(start, stop - 1) itemr[self.ds] = slice(start + 1, stop) return 0.5 * (self.ncvar[tuple(item)] + self.ncvar[tuple(itemr)])
def __getitem__(self, item): # take care of ellipsis and other strange indexes item = list(indexing.expanded_indexer(item, len(self.dimensions))) # time is always going to be first dim I hope sl = item[0] if np.isscalar(sl) and not isinstance(sl, slice): sl = slice(sl, sl+1) # Ok, get the indexes right start = sl.start or 0 stop = sl.stop or self._nel if stop < 0: stop += self._nel-1 start -= 1 do_nan = False if start < 0: do_nan = True itemr = copy.deepcopy(item) item[0] = slice(start, stop-1) itemr[0] = slice(start+1, stop) # done var = self.nc.variables[self.accvn] if do_nan: item[0] = slice(0, stop-1) out = var[itemr] out[1:, ...] -= var[item] out[0, ...] = np.NaN else: out = var[itemr] out -= var[item] return out * self._factor
def __getitem__(self, item): # take care of ellipsis and other strange indexes item = list(indexing.expanded_indexer(item, len(self.dimensions))) # time is always going to be first dim I hope sl = item[0] if np.isscalar(sl) and not isinstance(sl, slice): sl = slice(sl, sl + 1) # Ok, get the indexes right start = sl.start or 0 stop = sl.stop or self._nel if stop < 0: stop += self._nel - 1 start -= 1 do_nan = False if start < 0: do_nan = True itemr = copy.deepcopy(item) item[0] = slice(start, stop - 1) itemr[0] = slice(start + 1, stop) # done var = self.nc.variables[self.accvn] if do_nan: item[0] = slice(0, stop - 1) out = var[itemr] out[1:, ...] -= var[item] out[0, ...] = np.NaN else: out = var[itemr] out -= var[item] return out * self._factor
def __getitem__(self, item): """Override __getitem__.""" # take care of ellipsis and other strange indexes item = list(indexing.expanded_indexer(item, len(self.dimensions))) # Slice to change was_scalar = False sl = item[self.ds] if np.isscalar(sl) and not isinstance(sl, slice): sl = slice(sl, sl+1) was_scalar = True # Ok, get the indexes right start = sl.start or 0 stop = sl.stop or self._ds_shape if stop < 0: stop += self._ds_shape-1 stop = np.clip(stop+1, 0, self._ds_shape) itemr = copy.deepcopy(item) if was_scalar: item[self.ds] = start itemr[self.ds] = start+1 else: item[self.ds] = slice(start, stop-1) itemr[self.ds] = slice(start+1, stop) return 0.5*(self.ncvar[item] + self.ncvar[itemr])
def __getitem__(self, item): # take care of ellipsis and other strange indexes item = list(indexing.expanded_indexer(item, len(self.dimensions))) # we need the empty dims for _ncl_slp() to work squeezax = [] for i, c in enumerate(item): if np.isscalar(c) and not isinstance(c, slice): item[i] = slice(c, c+1) squeezax.append(i) # add a slice in the 4th dim item.insert(self.ds, slice(0, self._ds_shape+1)) item = tuple(item) # get data vars = self.nc.variables with ScaledVar(vars['TK']) as var: tk = var[item] with ScaledVar(vars['P']) as p, ScaledVar(vars['PB']) as pb: p = p[item] + pb[item] with ScaledVar(vars['QVAPOR']) as var: q = var[item] with ScaledVar(vars['PH']) as ph, ScaledVar(vars['PHB']) as phb: z = (ph[item] + phb[item]) / 9.81 return np.squeeze(_ncl_slp(z, tk, p, q), axis=tuple(squeezax))
def test_expanded_indexer(self) -> None: x = np.random.randn(10, 11, 12, 13, 14) y = np.arange(5) arr = ReturnItem() for i in [ arr[:], arr[...], arr[0, :, 10], arr[..., 10], arr[:5, ..., 0], arr[..., 0, :], arr[y], arr[y, y], arr[..., y, y], arr[..., 0, 1, 2, 3, 4], ]: j = indexing.expanded_indexer(i, x.ndim) assert_array_equal(x[i], x[j]) assert_array_equal(self.set_to_zero(x, i), self.set_to_zero(x, j)) with pytest.raises(IndexError, match=r"too many indices"): indexing.expanded_indexer(arr[1, 2, 3], 2)
def test_expanded_indexer(self): x = np.random.randn(10, 11, 12, 13, 14) y = np.arange(5) I = ReturnItem() for i in [ I[:], I[...], I[0, :, 10], I[..., 10], I[:5, ..., 0], I[..., 0, :], I[y], I[y, y], I[..., y, y], I[..., 0, 1, 2, 3, 4], ]: j = indexing.expanded_indexer(i, x.ndim) self.assertArrayEqual(x[i], x[j]) self.assertArrayEqual(self.set_to_zero(x, i), self.set_to_zero(x, j)) with self.assertRaisesRegexp(IndexError, "too many indices"): indexing.expanded_indexer(I[1, 2, 3], 2)
def __getitem__(self, item): # take care of ellipsis and other strange indexes item = list(indexing.expanded_indexer(item, len(self.dimensions))) # time is always going to be first dim I hope sl = item[0] was_scalar = False if np.isscalar(sl) and not isinstance(sl, slice): was_scalar = True sl = slice(sl, sl+1) # Ok, get the indexes right start = sl.start or 0 stop = sl.stop or self._nel if stop < 0: stop += self._nel-1 start -= 1 do_nan = False if start < 0: do_nan = True itemr = copy.deepcopy(item) item[0] = slice(start, stop-1) itemr[0] = slice(start+1, stop) # done with ScaledVar(self.nc.variables[self.accvn]) as var: if do_nan: item[0] = slice(0, stop-1) out = var[itemr] try: # in case we have a masked array out.unshare_mask() except: pass out[1:, ...] -= var[item] out[0, ...] = np.NaN else: out = var[itemr] out -= var[item] if was_scalar: out = out[0, ...] return out * self._factor
def __getitem__(self, item): # take care of ellipsis and other strange indexes item = list(indexing.expanded_indexer(item, len(self.dimensions))) # we need the empty dims for _ncl_slp() to work squeezax = [] for i, c in enumerate(item): if np.isscalar(c) and not isinstance(c, slice): item[i] = slice(c, c+1) squeezax.append(i) # add a slice in the 4th dim item.insert(self.ds, slice(0, self._ds_shape+1)) item = tuple(item) # get data vars = self.nc.variables tk = vars['TK'][item] p = vars['P'][item] + vars['PB'][item] q = vars['QVAPOR'][item] z = (vars['PH'][item] + vars['PHB'][item]) / 9.81 return np.squeeze(_ncl_slp(z, tk, p, q), axis=squeezax)
def expand(self, key): """Parse key using xarray utils to ensure we have dimension names.""" if not is_dict_like(key): labels = expanded_indexer(key, self.data_array.ndim) key = dict(zip(self.data_array.dims, labels)) return key