def _valid_grid_backend(): class DummyBackend: name = "dummy_backend" location = Path() def __init__(self, location: Union[str, Path]) -> None: raise NotImplementedError @staticmethod def is_valid_backend(location: Union[str, Path]) -> bool: ... dataset_name = str() dataset_label = str() dataset_unit = str() axes_names = [] axes_labels = [] axes_units = [] axes_min = np.empty(0) axes_max = np.empty(0) iteration = int() time_step = float() time_unit = str() shape = tuple() dtype = np.dtype("i") ndim = int() assert isinstance(DummyBackend, GridBackend) yield DummyBackend if DummyBackend.name in GridArray.get_backends(): GridArray.remove_backend(DummyBackend)
def slice_grid_array( grid: GridArray, constant: Union[str, int], value: float, ) -> GridArray: """ Takes a slice of a `GridArray` at a constant value of a given axis. Arguments: constant: Name or index that defines the axis taken to be constant in the slice. value: Value of the axis at which the slice is taken. Returns: Slice of ``grid``. Examples: Obtain a slice of a two-dimensional array. ```pycon >>> from nata.containers import GridArray >>> from nata.containers import Axis >>> import numpy as np >>> x = np.arange(5) >>> data = np.arange(25).reshape((5, 5)) >>> grid = GridArray.from_array(data, axes=[Axis(x), Axis(x)]) >>> grid.slice(constant=0, value=1).to_numpy() array([5, 6, 7, 8, 9]) # the second column >>> grid.slice(constant=1, value=1).to_numpy() array([ 1, 6, 11, 16, 21]) # the second row ``` """ if grid.ndim < 1: raise ValueError("slice is not available for 0 dimensional GridArrays") # get slice axis slice_axis = get_slice_axis(grid, constant) axis = grid.axes[slice_axis] if value < np.min(axis.to_dask()) or value >= np.max(axis.to_dask()): raise ValueError(f"out of range value for axis '{constant}'") # get index of nearest neighbour slice_idx = (np.abs(axis.to_dask() - value)).argmin(axis=-1) # build data slice data_slice = [slice(None)] * len(grid.axes) data_slice[slice_axis] = slice_idx return GridArray.from_array( grid.to_dask()[tuple(data_slice)], name=grid.name, label=grid.label, unit=grid.unit, axes=[ax for key, ax in enumerate(grid.axes) if ax is not axis], time=grid.time, )
def test_GridArray_get_valid_backend(valid_grid_backend: GridBackend): valid_grid_backend.is_valid_backend = staticmethod(lambda p: p == Path()) # register backend GridArray.add_backend(valid_grid_backend) assert GridArray.get_valid_backend(Path()) is valid_grid_backend assert GridArray.get_valid_backend(Path("some/invalid/path")) is None
def _grid_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): # make checking for 'is_valid_backend' always succeed monkeypatch.setattr(GridArray, "is_valid_backend", lambda _: True) # make 'basedir' basedir = tmp_path / "grid_files" basedir.mkdir(parents=True, exist_ok=True) # create dummy files dummy_data = np.arange(32).reshape((4, 8)) np.savetxt(basedir / "grid.0", dummy_data, delimiter=",") np.savetxt(basedir / "grid.1", dummy_data, delimiter=",") np.savetxt(basedir / "grid.2", dummy_data, delimiter=",") @register_backend(GridArray) class Dummy_GridFile: name: str = "dummy_backend" def __init__(self, location: Union[str, Path]) -> None: self.location = Path(location) self.name = "dummy_backend" self.data = np.loadtxt(location, delimiter=",") self.dataset_name = "dummy_grid" self.dataset_label = "dummy grid label" self.dataset_unit = "dummy unit" self.ndim = dummy_data.ndim self.shape = dummy_data.shape self.dtype = dummy_data.dtype self.axes_min = (0, 1) self.axes_max = (1, 2) self.axes_names = ("dummy_axis0", "dummy_axis1") self.axes_labels = ("dummy label axis0", "dummy label axis1") self.axes_units = ("dummy unit axis0", "dummy unit axis1") self.iteration = 0 self.time_step = 1.0 self.time_unit = "dummy time unit" @staticmethod def is_valid_backend(location: Union[str, Path]) -> bool: location = Path(location) if location.stem == "grid" and location.suffix in (".0", ".1", ".2"): return True else: return False def get_data(self, indexing: Optional[Any] = None) -> ArrayLike: if indexing: return self.data[indexing] else: return self.data yield basedir GridArray.remove_backend(Dummy_GridFile)
def test_GridArray_from_path_check_time_axis(grid_files): # using time axis implicit grid = GridArray.from_path(grid_files / "grid.0") assert grid.time.name == "time" # using time axis explicit grid = GridArray.from_path(grid_files / "grid.0", time_axis="time") assert grid.time.name == "time" # using time axis explicit grid = GridArray.from_path(grid_files / "grid.0", time_axis="iteration") assert grid.time.name == "iteration"
def test_array_fft_invalid_axis(): for invalid_axis in [2, -3]: with pytest.raises(ValueError, match="invalid axis index"): GridArray.from_array( np.arange(100).reshape((10, 10)), axes=[Axis.from_array(np.arange(10))] * 2, ).fft(axes=[invalid_axis]) for invalid_axis in ["axis2", "abc"]: with pytest.raises(ValueError, match="could not be found"): GridArray.from_array( np.arange(100).reshape((10, 10)), axes=[Axis.from_array(np.arange(10))] * 2, ).fft(axes=[invalid_axis])
def transpose_grid_array( grid: GridArray, axes: Optional[list] = None, ) -> GridArray: """Reverses or permutes the axes of a `GridArray`. Parameters ---------- axes: ``list``, optional List of integers and/or strings that identify the permutation of the axes. The i'th axis of the returned `GridArray` will correspond to the axis numbered/labeled axes[i] of the input. If not specified, the order of the axes is reversed. Returns ------ :class:`nata.containers.GridArray`: Transpose of ``grid``. Examples -------- Transpose a three-dimensional array. >>> from nata.containers import GridArray >>> import numpy as np >>> data = np.arange(96).reshape((8, 4, 3)) >>> grid = GridArray.from_array(data) >>> grid.transpose().shape (3, 4, 8) >>> grid.transpose(axes=[0,2,1]).shape (8, 3, 4) """ # get transpose axes tr_axes = get_transpose_axes(grid, axes) if len(set(tr_axes)) is not grid.ndim: raise ValueError("invalid transpose axes") return GridArray.from_array( da.transpose(grid.to_dask(), axes=tr_axes), name=grid.name, label=grid.label, unit=grid.unit, axes=[grid.axes[axis] for axis in tr_axes], time=grid.time, )
def test_array_fft_peak_1d(): x = np.linspace(0, 10 * np.pi, 101) grid = GridArray.from_array(np.sin(x), axes=[Axis.from_array(x)]) fft_grid = grid.fft() assert (fft_grid.to_dask().argmax() == (np.abs(fft_grid.axes[0].to_dask() + 1)).argmin())
def test_array_transpose_data(): data = np.arange(7 * 6 * 5).reshape((7, 6, 5)) grid = GridArray.from_array(data) tr_grid = grid.transpose(axes=[1, 2, 0]) np.testing.assert_array_equal(tr_grid, np.transpose(data, axes=[1, 2, 0]))
def test_GridArray_ufunc_proxy(): grid = GridArray.from_array([1, 2]) # creation of new object np.testing.assert_array_equal(grid + 1, [2, 3]) # in-place operation grid += 1 np.testing.assert_array_equal(grid, [2, 3])
def test_GridArray_repr(): grid = GridArray.from_array(np.arange(12, dtype=np.int32).reshape((4, 3))) expected_repr = ("GridArray<" "shape=(4, 3), " "dtype=int32, " "time=None, " "axes=(Axis(axis0), Axis(axis1))" ">") assert repr(grid) == expected_repr
def test_array_fft_comps_1d(): x = np.linspace(0, 10 * np.pi, 101) grid = GridArray.from_array(np.sin(x), axes=[Axis.from_array(x)]) for comp, fn in zip(["full", "real", "imag", "abs"], [lambda x: x, np.real, np.imag, np.abs]): fft_grid = grid.fft(comp=comp) fft_data = fft.fftshift(fft.fftn(grid.to_dask())) np.testing.assert_array_equal(fft_grid.to_dask(), fn(fft_data))
def test_GridArray_getitem(case): arr = np.array(case["arr"]) indexing = case["indexing"] expected_arr = case["expected_arr"] if "expected_arr" in case else arr[ indexing] expected_axes = case["expected_axes"] if "expected_axes" in case else None grid = GridArray.from_array(arr) subgrid = grid[indexing] expected_grid = GridArray.from_array(expected_arr, axes=expected_axes) for ax, expected_ax in zip(subgrid.axes, expected_grid.axes): assert ax.name == expected_ax.name assert ax.label == expected_ax.label assert ax.unit == expected_ax.unit np.testing.assert_array_equal(ax, expected_ax) np.testing.assert_array_equal(subgrid, expected_grid)
def test_GridArray_add_remove_backend(valid_grid_backend: GridBackend): assert valid_grid_backend.name not in GridArray.get_backends() GridArray.add_backend(valid_grid_backend) assert valid_grid_backend.name in GridArray.get_backends() GridArray.remove_backend(valid_grid_backend) assert valid_grid_backend.name not in GridArray.get_backends()
def test_array_slice_dimensionality(): grid = GridArray.from_array(np.arange(96).reshape((8, 4, 3))) sliced_grid = grid.slice(constant="axis0", value=0) assert sliced_grid.ndim == 2 assert sliced_grid.axes[0].name == grid.axes[1].name assert sliced_grid.axes[1].name == grid.axes[2].name np.testing.assert_array_equal(sliced_grid, np.arange(12).reshape(4, 3))
def test_array_transpose_axes(): data = np.zeros((7, 6, 5)) grid = GridArray.from_array(data) tr_grid = grid.transpose(axes=[1, 2, 0]) assert tr_grid.axes[0] is grid.axes[1] assert tr_grid.axes[1] is grid.axes[2] assert tr_grid.axes[2] is grid.axes[0]
def test_array_fft_shape_1d(): x = np.arange(10) for axes in [ [Axis.from_array(x)], # [Axis([x] * 2)], ]: grid = GridArray.from_array(x, axes=axes) assert grid.fft().shape == grid.shape assert grid.fft().axes[0].shape == grid.axes[0].shape
def test_array_fft_shape_2d(): x = np.arange(10) data = np.arange(100).reshape((10, 10)) for axes in [ [Axis.from_array(x)] * 2, # [Axis([x] * 2)] * 2, ]: grid = GridArray.from_array(data, axes=axes) assert grid.fft().shape == grid.shape assert grid.fft().axes[0].shape == grid.axes[0].shape
def test_GridArray_from_path(grid_files): grid = GridArray.from_path(grid_files / "grid.0") assert grid.name == "dummy_grid" assert grid.label == "dummy grid label" assert grid.unit == "dummy unit" assert grid.axes[0].name == "dummy_axis0" assert grid.axes[1].name == "dummy_axis1" assert grid.axes[0].label == "dummy label axis0" assert grid.axes[1].label == "dummy label axis1" assert grid.axes[0].unit == "dummy unit axis0" np.testing.assert_array_equal(grid, np.arange(32).reshape((4, 8)))
def test_array_slice_selection(): grid = GridArray.from_array(np.arange(100).reshape((10, 10))) with pytest.raises(ValueError, match="out of range value"): grid.slice(constant="axis0", value=-1) with pytest.raises(ValueError, match="could not be found"): grid.slice(constant="axis2", value=0) with pytest.raises(ValueError, match="invalid axis index"): grid.slice(constant=2, value=0) with pytest.raises(ValueError, match="invalid axis index"): grid.slice(constant=-3, value=0)
def test_array_transpose_shape(): data = np.zeros((7, 6, 5)) grid = GridArray.from_array(data) for tr_axes in [ None, [0, 1, 2], [0, 2, 1], [2, 1, 0], [1, 2, 0], ]: tr_grid = grid.transpose(axes=tr_axes) assert tr_grid.shape == np.transpose(data, axes=tr_axes).shape
def test_array_transpose_invalid_axes(): grid = GridArray.from_array(np.arange(7 * 6).reshape((7, 6))) for invalid_axes in [ [0], [1, 1], ["axis0", "axis0"], ]: with pytest.raises(ValueError, match="invalid transpose axes"): grid.transpose(axes=invalid_axes) for invalid_axis in [2, -3]: with pytest.raises(ValueError, match="invalid axis index"): grid.transpose(axes=[invalid_axis])
def test_GridArray_repr_markdown_(): grid = GridArray.from_array(np.arange(12, dtype=np.int32).reshape((4, 3))) expected_markdown = """ | **GridArray** | | | ---: | :--- | | **name** | unnamed | | **label** | unlabeled | | **unit** | '' | | **shape** | (4, 3) | | **dtype** | int32 | | **time** | None | | **axes** | Axis(axis0), Axis(axis1) | """ assert grid._repr_markdown_() == dedent(expected_markdown)
def streak_grid_array(grid: GridDataset, ) -> GridArray: """Converts a `GridDataset` to a `GridArray`. Only `GridDataset` with axes that do not change over time are supported. Returns ------ :class:`nata.containers.GridArray`: Streak of ``grid``. Examples -------- Convert a one-dimensional dataset with time dependence to a two-dimensional array. >>> from nata.containers import GridDataset >>> import numpy as np >>> data = np.arange(5*7).reshape((5, 7)) >>> grid = GridDataset.from_array(data) >>> stk_grid = grid.streak() >>> stk_grid.shape (5, 7) >>> [axis.shape for axis in stk_grid.axes] [(5,), (7,)] """ if grid.ndim < 2: raise ValueError( "streak is not available for 0 dimensional GridDatasets") for axis in grid.axes[1:]: for i, axis_i in enumerate(axis): if np.any(axis_i.to_dask() != axis[0].to_dask()): raise ValueError("invalid axes for streak") return GridArray.from_array( grid.to_dask(), name=grid.name, label=grid.label, unit=grid.unit, axes=[grid.time] + [axis[0] for axis in grid.axes[1:]], )
def test_GridArray_from_array_raise_invalid_axes(): # invalid number of axes with pytest.raises(ValueError, match="mismatches with dimensionality of data"): GridArray.from_array([], axes=[0, 1]) # axes which are not 1D dimensional with pytest.raises(ValueError, match="only 1D axis for GridArray are supported"): GridArray.from_array([0, 1], axes=[[[0, 1]]]) # axis mismatch with shape of data with pytest.raises(ValueError, match="inconsistency between data and axis shape"): GridArray.from_array([0, 1], axes=[[0, 1, 2, 3]])
def test_GridArray_from_array_default(): grid_arr = GridArray.from_array(da.arange(12, dtype=int).reshape((4, 3))) assert grid_arr.shape == (4, 3) assert grid_arr.ndim == 2 assert grid_arr.dtype == int assert grid_arr.axes[0].name == "axis0" assert grid_arr.axes[0].label == "unlabeled" assert grid_arr.axes[0].unit == "" assert grid_arr.axes[0].shape == (4, ) assert grid_arr.axes[1].name == "axis1" assert grid_arr.axes[1].label == "unlabeled" assert grid_arr.axes[1].unit == "" assert grid_arr.axes[1].shape == (3, ) assert grid_arr.time.name == "time" assert grid_arr.time.label == "time" assert grid_arr.time.unit == "" assert grid_arr.name == "unnamed" assert grid_arr.label == "unlabeled" assert grid_arr.unit == ""
def test_GridArray_len(): assert len(GridArray.from_array(np.zeros((3, )))) == 3 assert len(GridArray.from_array(np.zeros((5, 3)))) == 5 with pytest.raises(TypeError, match="unsized object"): len(GridArray.from_array(np.zeros(())))
def test_GridArray_array(): grid = GridArray.from_array([1, 2]) np.testing.assert_array_equal(grid, [1, 2])
def test_GridArray_array_function_proxy(): grid = GridArray.from_array([1, 2]) np.testing.assert_array_equal(np.fft.fft(grid), np.fft.fft([1, 2]))
def test_array_fft_invalid_ndim(): with pytest.raises(ValueError, match="0 dimensional GridArrays"): GridArray.from_array(1).fft()