diff --git a/README.md b/README.md index f994ee0..289e774 100644 --- a/README.md +++ b/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 diff --git a/nt2/containers/data.py b/nt2/containers/data.py index de2c511..78b0a7a 100644 --- a/nt2/containers/data.py +++ b/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, diff --git a/nt2/containers/particles.py b/nt2/containers/particles.py index cf5dc6f..5a01dfb 100644 --- a/nt2/containers/particles.py +++ b/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( diff --git a/nt2/plotters/export.py b/nt2/plotters/export.py index db0b19f..939639f 100644 --- a/nt2/plotters/export.py +++ b/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, diff --git a/nt2/plotters/inspect.py b/nt2/plotters/inspect.py index d9d9ab7..4effbd3 100644 --- a/nt2/plotters/inspect.py +++ b/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, diff --git a/nt2/plotters/movie.py b/nt2/plotters/movie.py index 3eec8ad..fd64e05 100644 --- a/nt2/plotters/movie.py +++ b/nt2/plotters/movie.py @@ -6,18 +6,53 @@ 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.") @@ -25,12 +60,18 @@ def plot( 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() diff --git a/nt2/readers/adios2.py b/nt2/readers/adios2.py index 14f9aaf..ac352cf 100644 --- a/nt2/readers/adios2.py +++ b/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 diff --git a/nt2/readers/hdf5.py b/nt2/readers/hdf5.py index 8ec7708..435185f 100644 --- a/nt2/readers/hdf5.py +++ b/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