def test_ParticleDatasets_backends_are_registered():
backends = ParticleDataset.get_backends()
assert backends[Osiris_Hdf5_ParticleFile.name] is Osiris_Hdf5_ParticleFile
assert (backends[Osiris_Dev_Hdf5_ParticleFile.name] is
Osiris_Dev_Hdf5_ParticleFile)
assert backends[Osiris_zdf_ParticleFile.name] is Osiris_zdf_ParticleFile
def _TestParticleBackend():
"""Fixture for returning ParticleBackend"""
class TestParticleBackend:
"""Test backend for ParticleDatasets"""
name = "TestParticleBackend"
location = None
dataset_name = "test_dataset_name"
num_particles = 10
quantity_names = ("quant1", "quant2", "quant3")
quantity_labels = ("quant1 label", "quant2 label", "quant3 label")
quantity_units = ("quant1 unit", "quant2 unit", "quant3 unit")
iteration = 42
time_step = 12.3
time_unit = "time unit"
dtype = np.dtype(
[("quant1", np.int), ("quant2", np.float), ("quant3", np.int)]
)
@staticmethod
def is_valid_backend(location: Union[Path, str]) -> bool:
return (
True
if Path(location) == Path("TestParticleBackend_location")
else False
)
def __init__(self, location) -> None:
pass
# makes sure dummy backend is of valid type
assert isinstance(TestParticleBackend, ParticleBackendType)
ParticleDataset.add_backend(TestParticleBackend)
yield TestParticleBackend
# teardown
ParticleDataset.remove_backend(TestParticleBackend)
def test_ParticleDataset_axes_from_Backend(TestParticleBackend):
"""Tests extraction correct type for .axes and its proper extraction"""
ds = ParticleDataset("TestParticleBackend_location")
assert isinstance(ds.axes, Mapping)
for axes_name, type_ in [
("iteration", AxisType),
("time", AxisType),
]:
assert axes_name in ds.axes
assert isinstance(ds.axes.get(axes_name), type_)
def test_Osiris_Dev_Hdf5_ParticleFile_check_is_valid_backend(
os_hdf5_particle_dev_file, ):
"""Check 'Osiris_Dev_Hdf5_ParticleFile' is a valid backend exclusively"""
assert (Osiris_Dev_Hdf5_ParticleFile.is_valid_backend(
os_hdf5_particle_dev_file) is True)
# backend are registered automatically for ParticleDatasets
for (name, backend) in ParticleDataset.get_backends().items():
if name == Osiris_Dev_Hdf5_ParticleFile.name:
continue
assert backend.is_valid_backend(os_hdf5_particle_dev_file) is False
def filter_particle_dataset(dataset: ParticleDataset,
quantities: List[str]) -> ParticleDataset:
"""Filters a :class:`nata.containers.ParticleDataset` according to a\
selection of quantities.
Parameters
----------
quantities: ``list``
List of quantities to be filtered, ordered by the way they should be
sorted in the returned dataset.
Returns
------
:class:`nata.containers.ParticleDataset`:
Filtered dataset with only the quantities selected in
``quantities``.
Examples
--------
The filter plugin is used to get dataset with only a selection, say
``'x1'`` and ``'x2'``, of its quantities.
>>> from nata.containers import ParticleDataset
>>> ds = ParticleDataset("path/to"file")
>>> ds_flt = ds.filter(quantities=["x1","p1"])
"""
if quantities is not None:
quants = {}
for quant in quantities:
quants[quant] = (dataset.quantities[quant]
if quant in dataset.quantities.keys() else None)
else:
raise ValueError("")
return ParticleDataset(
iteration=dataset.axes["iteration"],
time=dataset.axes["time"],
name=dataset.name,
quantities=quants,
)
def test_ParticleDataset_quantities_from_Backend(TestParticleBackend):
"""Checks if quantities propagate from backend to ParticleDataset"""
ds = ParticleDataset("TestParticleBackend_location")
assert isinstance(ds.quantities, Mapping)
expected_names = TestParticleBackend.quantity_names
expected_labels = TestParticleBackend.quantity_labels
expected_units = TestParticleBackend.quantity_units
# check key entries
for key, expected_name in zip(ds.quantities.keys(), expected_names):
assert key == expected_name
# check value entries
for quant, expected_name in zip(ds.quantities.values(), expected_names):
assert quant.name == expected_name
for quant, expected_label in zip(ds.quantities.values(), expected_labels):
assert quant.label == expected_label
for quant, expected_unit in zip(ds.quantities.values(), expected_units):
assert quant.unit == expected_unit
def plot_particle_dataset(
dataset: ParticleDataset,
fig: Optional[Figure] = None,
axes: Optional[Axes] = None,
style: dict = dict(),
interactive: bool = True,
n: int = 0,
):
"""Plots a single/multiple iteration :class:`nata.containers.ParticleDataset`\
using a :class:`nata.plots.types.ScatterPlot`.
Parameters
----------
fig: :class:`nata.plots.Figure`, optional
If provided, the plot is drawn on ``fig``. The plot is drawn on
``axes`` if it is a child axes of ``fig``, otherwise a new axes
is created on ``fig``. If ``fig`` is not provided, a new
:class:`nata.plots.Figure` is created.
axes: :class:`nata.plots.Axes`, optional
If provided, the plot is drawn on ``axes``, which must be an axes
of ``fig``. If ``axes`` is not provided or is provided without a
corresponding ``fig``, a new :class:`nata.plots.Axes` is created in
a new :class:`nata.plots.Figure`.
style: ``dict``, optional
Dictionary that takes a mix of style properties of
:class:`nata.plots.Figure`, :class:`nata.plots.Axes` and any plot
type (see :class:`nata.plots.types.ScatterPlot`).
interactive: ``bool``, optional
Controls wether interactive widgets should be shown with the plot
to allow for temporal navigation. Only applicable if ``dataset``
has multiple iterations.
n: ``int``, optional
Selects the index of the iteration to be shown initially. Only
applicable if ``dataset`` has multiple iterations, .
Returns
------
:class:`nata.plots.Figure` or ``None``:
Figure with plot built based on ``dataset``. Interactive widgets
are shown with the figure if ``dataset`` has multiple iterations,
in which case this method returns ``None``.
Examples
--------
To get a plot with default style properties in a new figure, simply
call the ``.plot()`` method. The first two quantities in the dataset
``quantities`` dictionary will be represented in the horizontal and
vertical plot axes, respectively. If a third quantity is available, it
will be represented in colors.
>>> from nata.containers import ParticleDataset
>>> import numpy as np
>>> arr = np.arange(30).reshape(1,10,3)
>>> ds = ParticleDataset("path/to/file")
>>> fig = ds.plot()
The list of quantities in the dataset can be filtered with the
:meth:`nata.containers.ParticleDataset.filter` method.
>>> fig = ds.filter(quantities=["x1", "p1", "ene"]).plot()
"""
p_plan = PlotPlan(
dataset=dataset,
style=filter_style(dataset.plot_type(), style),
)
a_plan = AxesPlan(axes=axes,
plots=[p_plan],
style=filter_style(Axes, style))
f_plan = FigurePlan(fig=fig,
axes=[a_plan],
style=filter_style(Figure, style))
if len(dataset) > 1:
if inside_notebook():
if interactive:
f_plan.build_interactive(n)
else:
return f_plan[n].build()
else:
# TODO: remove last line from warn
warn(f"Plotting only iteration with index n={str(n)}." +
" Interactive plots of multiple iteration datasets are not" +
" supported outside notebook environments.")
return f_plan[n].build()
else:
return f_plan.build()
def test_ParticleDataset_num_particles_from_Backend(TestParticleBackend):
"""Tests extraction correct type for . and its proper extraction"""
ds = ParticleDataset("TestParticleBackend_location")
assert isinstance(ds.num_particles, AxisType)
def test_ParticleDataset_attr_propagation_from_Backend(
TestParticleBackend, attr, value
):
"""Parameterized check for different props of ParticleDataset"""
ds = ParticleDataset("TestParticleBackend_location")
assert getattr(ds, attr) == value
def test_ParticleDataset_registration(TestParticleBackend):
"""Check if fixture registers backend properly"""
assert TestParticleBackend.name in ParticleDataset.get_backends()
def test_ParticleDataset_time_axis_from_Backend(
TestParticleBackend, attr, value
):
"""Extraction is correct for iteration axis. Check attributes for axis"""
ds = ParticleDataset("TestParticleBackend_location")
assert getattr(ds.axes["time"], attr) == value
def filter_particle_dataset(
dataset: ParticleDataset,
mask: List[bool] = None,
quantities: List[str] = None,
slicing: slice = None,
) -> ParticleDataset:
"""Filters a :class:`nata.containers.ParticleDataset` according to a\
selection of quantities.
Parameters
----------
mask: ``np.ndarray``, optional
Array of booleans indicating the particles to be filtered. Particles
with ``True`` (``False``) mask entries are selected (hidden). The
shape of ``mask`` must match that of each particle quantity.
slicing: ``slice``, optional
Slice of particles to be filtered. Acts only on particle indices
and not on time, as time slicing should be done on the dataset.
When provided together with the ``mask argument``, slicing is done
on the masked array.
quantities: ``list``, optional
List of quantities to be filtered, ordered by the way they should be
sorted in the returned dataset.
Returns
------
:class:`nata.containers.ParticleDataset`:
Filtered dataset with only the quantities selected in
``quantities``.
Examples
--------
The filter plugin is used to get dataset with only a selection, say
``'x1'`` and ``'x2'``, of its quantities.
>>> from nata.containers import ParticleDataset
>>> ds = ParticleDataset("path/to"file")
>>> ds_flt = ds.filter(quantities=["x1","p1"])
"""
dataset_c = deepcopy(dataset)
if quantities is not None:
if not isinstance(quantities, list) or not all(
isinstance(quant, str) for quant in quantities):
raise TypeError(f"'quantities' must be a list of strings.")
quants = {
quant: dataset_c.quantities[quant]
for quant in quantities if quant in dataset_c.quantities.keys()
}
if not bool(quants):
raise ValueError(
f"None of the quantities in 'quantities' were found in the " +
"dataset.")
else:
quants = dataset_c.quantities
if mask is not None:
if not isinstance(mask, np.ndarray) or not mask.dtype == bool:
raise TypeError(f"'mask' must be an 'np.ndarray' of dtype 'bool'.")
quant = next(iter(quants.values()))
if quant.data.shape != mask.shape:
raise ValueError(
f"'mask.shape' ({mask.shape}) does not match the shape of " +
"each quantity ({quant.data.shape}).")
for name, quant in quants.items():
quants[name].data[~mask] = np.ma.masked
if slicing is not None:
if not isinstance(slicing, slice):
raise TypeError(f"'slicing' must be of type 'slice'.")
if len(dataset_c) > 1:
slicing = (slice(None), slicing)
quants = {
quant.name: ParticleQuantity(
data=quant.data[slicing],
name=quant.name,
label=quant.label,
unit=quant.unit,
dtype=quant.data.dtype,
)
for quant in quants.values()
}
return ParticleDataset(
iteration=dataset_c.axes["iteration"],
time=dataset_c.axes["time"],
name=dataset_c.name,
quantities=quants,
)