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 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_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_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_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_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_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_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 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_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_change_unit_by_prop():
grid_arr = GridArray.from_array([])
assert grid_arr.unit == ""
grid_arr.unit = "new unit"
assert grid_arr.unit == "new unit"
def test_GridArray_from_array_check_time():
grid_arr = GridArray.from_array([], time=123)
np.testing.assert_array_equal(grid_arr.time, 123)
assert grid_arr.time.name == "time"
assert grid_arr.time.label == "time"
assert grid_arr.time.unit == ""
def test_GridArray_from_array_check_axes():
grid_arr = GridArray.from_array([0, 1], axes=[[0, 1]])
np.testing.assert_array_equal(grid_arr.axes[0], [0, 1])
assert grid_arr.axes[0].name == "axis0"
assert grid_arr.axes[0].label == "unlabeled"
assert grid_arr.axes[0].unit == ""
def test_GridArray_from_array_raise_invalid_name():
with pytest.raises(ValueError,
match="'name' has to be a valid identifier"):
GridArray.from_array([], name="invalid name")
def test_GridArray_from_array_raise_invalid_time():
with pytest.raises(ValueError, match="time axis has to be 0 dimensional"):
GridArray.from_array([], time=[0, 1, 2])
def test_array_fft_invalid_ndim():
with pytest.raises(ValueError, match="0 dimensional GridArrays"):
GridArray.from_array(1).fft()
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_GridArray_change_label_by_prop():
grid_arr = GridArray.from_array([])
assert grid_arr.label == "unlabeled"
grid_arr.label = "new label"
assert grid_arr.label == "new label"
def test_GridArray_array():
grid = GridArray.from_array([1, 2])
np.testing.assert_array_equal(grid, [1, 2])
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(())))