Bug when no particles of given species are present

README.md

@@ -125,21 +125,6 @@ nt2e.makeFramesAndMovie(
 )
 ```
 
-#### Plots for debugging
-
-If the simulation also outputs the ghost cells, `nt2py` will interpret the fields differently, and instead of reading the physical coordinates, will build the coordinates based on the number of cells (including ghost cells). In particular, instead of, e.g., `data.fields.x` it will contain `data.fields.i1`. The data will also contain information about the meshblock decomposition. For instance, if you have `Nx` meshblocks in the `x` direction, each having `nx` cells, the coordinates `data.fields.i1` will go from `0` to `nx * NX + 2 * NGHOSTS * Nx`.
-
-You can overplot both the coordinate grid as well as the active zones of the meshblocks using the following:
-
-```python
-ax = plt.gca()
-data.fields.Ex.isel(t=ti).plot(ax=ax)
-data.plotGrid(ax=ax)
-data.plotDomains(ax=ax)
-```
-
-> Keep in mind, that by default `Entity` converts all quantities to tetrad basis (or contravariant in GR) and interpolates to cell centers before outputting (except for the ghost cells). So when doing plots for debugging, make sure to also set `as_is = true` (together with `ghosts = true`) in the `[output.debug]` section of the `toml` input file. This will ensure the fields are being output as is, with no conversion or interpolation. This does not apply to particle moments, which are never stored in the code and are computed only during the output.
-
 ### Dashboard
 
 Support for the dask dashboard is still in beta, but you can access it by first launching the dashboard client:
@@ -200,3 +185,4 @@ There are unit tests included with the code which also require downloading test
 - [ ] Interactive regime (`hvplot`, `bokeh`, `panel`)
 - [x] Ghost cells support
 - [x] Usage examples
+- [ ] Parse the log file with timings

nt2/containers/data.py

@@ -64,7 +64,7 @@ class MoviePlotAccessor(acc_movie.accessor):
     pass
 
 
-class Data(Fields, Particles):  # pyright: ignore[reportUnsafeMultipleInheritance]
+class Data(Fields, Particles):
     """Main class to manage all the data containers.
 
     Inherits from all category-specific containers.
@@ -220,7 +220,7 @@ def remap_prtl_quantities(name: str) -> str:
 
     def makeMovie(
         self,
-        plot: Callable,  # pyright: ignore[reportMissingTypeArgument, reportUnknownParameterType]
+        plot: Callable,
         time: list[float] | None = None,
         num_cpus: int | None = None,
         **movie_kwargs: Any,

nt2/containers/particles.py

@@ -1,8 +1,6 @@
 from typing import Any, Callable, List, Optional, Sequence, Tuple, Literal
 from copy import copy
 
-import dask
-import dask.array as da
 import dask.dataframe as dd
 from dask.delayed import delayed
 import pandas as pd
@@ -12,7 +10,6 @@
 import matplotlib.axes as maxes
 
 from nt2.containers.container import BaseContainer
-from nt2.readers.base import BaseReader
 
 
 IntSelector = int | Sequence[int] | slice | Tuple[int, int]
@@ -353,7 +350,7 @@ def _load_index_partition(st: int, t: float, index_cols: Tuple[str, ...]):
         ddf = dd.from_delayed(delayed_parts, meta=meta)
         return ddf
 
-    def load(self, cols: Sequence[str] = None) -> pd.DataFrame:
+    def load(self, cols: Sequence[str] | None = None) -> pd.DataFrame:
         if cols is None:
             cols = self.columns
 
@@ -387,12 +384,38 @@ def _attach_columns(part: pd.DataFrame) -> pd.DataFrame:
             .drop(columns=["row"])
         )
 
+    def help(self) -> str:
+        ret = "- use .sel(...) to select particles based on criteria:\n"
+        ret += "  t  : time (float)\n"
+        ret += "  st : step (int)\n"
+        ret += "  sp : species (int)\n"
+        ret += "  id : particle id (int)\n\n"
+        ret += "  # example:\n"
+        ret += "  #   .sel(t=slice(10.0, 20.0), sp=[1, 2, 3], id=[42, 22])\n\n"
+        ret += "- use .isel(...) to select particles based on output step:\n"
+        ret += "  t  : timestamp index (int)\n"
+        ret += "  st : step index (int)\n\n"
+        ret += "  # example:\n"
+        ret += "  #   .isel(t=-1)\n"
+        ret += "\n"
+        ret += "- .sel and .isel can be chained together:\n\n"
+        ret += "  # example:\n"
+        ret += "  #   .isel(t=-1).sel(sp=1).sel(id=[55, 66])\n\n"
+        ret += "- use .load(cols=[...]) to load data into a pandas DataFrame (`cols` defaults to all columns)\n\n"
+        ret += "  # example:\n"
+        ret += "  #  .sel(...).load()\n"
+        return ret
+
     def __repr__(self) -> str:
         ret = "ParticleDataset:\n"
+        ret += "================\n"
+        ret += f"Variables:\n  {self.columns}\n\n"
+        ret += "Current selection:\n"
         for k, v in self.selection.items():
-            ret += f"  {k}: {v}\n"
-        ret += f"  columns: {self.columns}\n"
-        ret += f"  total rows: {len(self._ddf_index)}\n"
+            ret += f"  {k:<5} : {v}\n"
+        ret += "\nHelp:\n"
+        ret += "-----\n"
+        ret += f"{self.help()}"
         return ret
 
     def __str__(self) -> str:
@@ -557,22 +580,22 @@ def __read_particles(self) -> ParticleDataset:
         """Helper function to read all particles data."""
         valid_steps = self.nonempty_steps
 
-        quantities = self.reader.ReadCategoryNamesAtTimestep(
-            self.path, "particles", "p", valid_steps[0]
-        )
+        quantities_ = [
+            self.reader.ReadCategoryNamesAtTimestep(self.path, "particles", "p", step)
+            for step in valid_steps
+        ]
+        quantities = sorted(np.unique([q for qtys in quantities_ for q in qtys]))
+
         unique_quantities = sorted(
             list(
                 set(
-                    [
-                        q.split("_")[0]
-                        for q in quantities
-                        if not q.startswith("pIDX") and not q.startswith("pRNK")
-                    ]
+                    str(q).split("_")[0]
+                    for q in quantities
+                    if not q.startswith("pIDX") and not q.startswith("pRNK")
                 )
             )
         )
-
-        all_species = sorted(list(set([int(q.split("_")[1]) for q in quantities])))
+        all_species = sorted(list(set([int(str(q).split("_")[1]) for q in quantities])))
 
         sp_with_idx = sorted(
             [int(q.split("_")[1]) for q in quantities if q.startswith("pIDX")]
@@ -588,11 +611,14 @@ def remap_quantity(name: str) -> str:
             return name
 
         def GetCount(step: int, sp: int) -> np.int64:
-            return np.int64(
-                self.reader.ReadArrayShapeAtTimestep(
-                    self.path, "particles", f"pX1_{sp}", step
-                )[0]
-            )
+            try:
+                return np.int64(
+                    self.reader.ReadArrayShapeAtTimestep(
+                        self.path, "particles", f"pX1_{sp}", step
+                    )[0]
+                )
+            except:
+                return np.int64(0)
 
         def ReadSteps() -> np.ndarray:
             return np.array(self.reader.GetValidSteps(self.path, "particles"))
@@ -664,6 +690,14 @@ def ReadColumn(step: int, colname: str) -> np.ndarray:
             else:
                 read_colname = f"p{colname}"
 
+            def species_has_quantity(sp: int) -> bool:
+                return (
+                    f"{read_colname}_{sp}"
+                    in self.reader.ReadCategoryNamesAtTimestep(
+                        self.path, "particles", "p", step
+                    )
+                )
+
             def get_quantity_for_species(sp: int) -> np.ndarray:
                 if f"{read_colname}_{sp}" in quantities:
                     return self.reader.ReadArrayAtTimestep(
@@ -673,8 +707,16 @@ def get_quantity_for_species(sp: int) -> np.ndarray:
                     return np.zeros(GetCount(step, sp)) * np.nan
 
             return np.concat(
-                [get_quantity_for_species(sp) for sp in sp_with_idx]
-                + [get_quantity_for_species(sp) for sp in sp_without_idx]
+                [
+                    get_quantity_for_species(sp)
+                    for sp in sp_with_idx
+                    if species_has_quantity(sp)
+                ]
+                + [
+                    get_quantity_for_species(sp)
+                    for sp in sp_without_idx
+                    if species_has_quantity(sp)
+                ]
             )
 
         return ParticleDataset(

nt2/plotters/export.py

@@ -3,7 +3,7 @@
 
 def makeFramesAndMovie(
     name: str,
-    plot: Callable,  # pyright: ignore[reportUnknownParameterType,reportMissingTypeArgument]
+    plot: Callable,
     times: list[float],
     data: Any = None,
     **kwargs: Any,
@@ -96,7 +96,7 @@ def makeMovie(**ffmpeg_kwargs: str | int | float) -> bool:
 
 
 def makeFrames(
-    plot: Callable,  # pyright: ignore[reportUnknownParameterType,reportMissingTypeArgument]
+    plot: Callable,
     times: list[float],
     fpath: str,
     data: Any = None,

nt2/plotters/inspect.py

@@ -1,5 +1,3 @@
-# pyright: reportMissingTypeStubs=false
-
 from typing import Any, Callable
 import matplotlib.pyplot as plt
 import matplotlib.figure as mfigure
@@ -15,7 +13,7 @@ def __init__(self, xarray_obj: xr.Dataset):
     def __axes_grid(
         self,
         grouped_fields: dict[str, list[str]],
-        makeplot: Callable,  # pyright: ignore[reportUnknownParameterType,reportMissingTypeArgument]
+        makeplot: Callable,
         nrows: int,
         ncols: int,
         nfields: int,
@@ -53,8 +51,8 @@ def __axes_grid(
     @staticmethod
     def _fixed_axes_grid_with_cbars(
         fields: list[str],
-        makeplot: Callable,  # pyright: ignore[reportUnknownParameterType,reportMissingTypeArgument]
-        makecbar: Callable,  # pyright: ignore[reportUnknownParameterType,reportMissingTypeArgument]
+        makeplot: Callable,
+        makecbar: Callable,
         nrows: int,
         ncols: int,
         nfields: int,

nt2/plotters/movie.py

@@ -6,31 +6,72 @@
 
 
 class accessor:
+    """
+    Movie plotter for xarray DataArray objects.
+
+    Functions
+    ---------
+    plot(name: str, movie_kwargs: dict[str, Any] | None = None, *args: Any, **kwargs: Any) -> bool
+        Plot a movie of the DataArray over the time dimension 't'.
+    """
+
     def __init__(self, xarray_obj: xr.DataArray) -> None:
         self._obj: xr.DataArray = xarray_obj
 
     def plot(
         self,
         name: str,
-        movie_kwargs: dict[str, Any] | None = None,
-        *args: Any,
+        movie_kwargs: dict[str, Any] = {},
+        fig_kwargs: dict[str, Any] = {},
+        aspect_equal: bool = False,
         **kwargs: Any,
     ) -> bool:
-        if movie_kwargs is None:
-            movie_kwargs = {}
+        """
+        Plot a movie of the DataArray over the time dimension 't'.
+
+        Parameters
+        ----------
+        name : str
+            The name of the output movie file.
+        movie_kwargs : dict[str, Any], optional
+            Additional keyword arguments for movie creation (default is {}).
+        fig_kwargs : dict[str, Any], optional
+            Additional keyword arguments for figure creation (default is {}).
+        aspect_equal : bool, optional
+            Whether to set equal aspect ratio for 2D plots (default is False).
+        **kwargs : Any
+            Additional keyword arguments to pass to the plotting function.
+        Returns
+        -------
+        bool
+            True if the movie was created successfully, False otherwise.
+
+        Notes
+        -----
+        kwargs are passed to the xarray plotting function:
+        ```
+        xarray.plot(**kwargs)
+        ```
+        """
         if "t" not in self._obj.dims:
             raise ValueError("The dataset does not have a time dimension.")
 
         import matplotlib.pyplot as plt
 
         def plot_func(ti: int, _: Any) -> None:
             if len(self._obj.isel(t=ti).dims) == 2:
-                x1, x2 = self._obj.isel(t=ti).dims
-                nx1, nx2 = len(self._obj.isel(t=ti)[x1]), len(self._obj.isel(t=ti)[x2])
-                aspect = nx1 / nx2
-                _ = plt.figure(figsize=(6, 4 * aspect))
-            self._obj.isel(t=ti).plot(*args, **kwargs)
-            if len(self._obj.isel(t=ti).dims) == 2:
+                if aspect_equal:
+                    x1, x2 = self._obj.isel(t=ti).dims
+                    nx1, nx2 = len(self._obj.isel(t=ti)[x1]), len(
+                        self._obj.isel(t=ti)[x2]
+                    )
+                    aspect = nx1 / nx2
+                    figsize = fig_kwargs.get("figsize", (6, 4 * aspect))
+                    _ = plt.figure(figsize=figsize, **fig_kwargs)
+                else:
+                    _ = plt.figure(**fig_kwargs)
+            self._obj.isel(t=ti).plot(**kwargs)
+            if aspect_equal and len(self._obj.isel(t=ti).dims) == 2:
                 plt.gca().set_aspect("equal")
             plt.tight_layout()
 

nt2/readers/adios2.py

@@ -1,14 +1,15 @@
-# pyright: reportMissingTypeStubs=false
-
 from typing import Any
 
 import sys
+
 if sys.version_info >= (3, 12):
     from typing import override
 else:
+
     def override(method):
         return method
 
+
 import re
 import os
 import numpy as np

nt2/readers/hdf5.py

@@ -1,14 +1,15 @@
-# pyright: reportMissingTypeStubs=false
-
 from typing import Any
 
 import sys
+
 if sys.version_info >= (3, 12):
     from typing import override
 else:
+
     def override(method):
         return method
 
+
 import re
 import os
 import numpy as np