diff --git a/README.md b/README.md index 6b97b9bcd..7b79a276c 100755 --- a/README.md +++ b/README.md @@ -50,6 +50,8 @@ This will create `params_Maxwell.py` in your current working directory (cwd). Yo The default output is in `sim_1/` in your cwd. You can change the output path via the class `EnvironmentOptions` in the parameter file. +For plots and diagnostics of the output, install [plasma-plots](https://struphy-hub.github.io/plasma-plots) with `pip install "struphy[pproc]"`; `Output` loads it automatically, adding `out.plot`, `out.analysis` and `.plasma.plot` on every product. + Parallel simulations are run for example with pip install -U mpi4py diff --git a/README.qmd b/README.qmd index 2d86a1249..d41f7a2c3 100644 --- a/README.qmd +++ b/README.qmd @@ -111,6 +111,8 @@ python params_Maxwell.py The default output is in `sim_1/` in your cwd. You can change the output path via the class `EnvironmentOptions` in the parameter file. +For plots and diagnostics of the output, install [plasma-plots](https://struphy-hub.github.io/plasma-plots) with `pip install "struphy[pproc]"`; `Output` loads it automatically, adding `out.plot`, `out.analysis` and `.plasma.plot` on every product. + Parallel simulations are run for example with ``` diff --git a/doc/markdown/output-api.md b/doc/markdown/output-api.md index 6337ef665..ef5924874 100644 --- a/doc/markdown/output-api.md +++ b/doc/markdown/output-api.md @@ -80,8 +80,8 @@ distribution = out.evaluate("kinetic_ions/f") delta_f = out.evaluate("kinetic_ions/f", dataset="e1_v1_density/delta_f") ``` -To make a figure, select the dimensions to show and call xarray's native `.plot()` methods (see -[Plot data](#plot-data)). +To make a figure, select the dimensions to show and call xarray's native `.plot()` methods, or +use the plots and diagnostics of the plasma-plots package (see [Plot data](#plot-data)). ## Analyze and report data @@ -219,6 +219,32 @@ out.evaluate("diagnostics/rho_xyz", t=-1, eta3=0.5, method="nearest").plot(x="et xarray squeezes size-one dimensions before plotting, so an array that is 2-D on a grid with one cell in some direction plots as a line. Select until the array has the dimensions the plot needs. +### Plots and diagnostics with plasma-plots + +For plots and diagnostics made for Struphy output, install the separate package +[plasma-plots](https://struphy-hub.github.io/plasma-plots): + +```bash +pip install plasma-plots # or: pip install "struphy[pproc]" +``` + +When it is installed, every `Output` loads it: no import is needed. It adds `out.plot` and +`out.analysis` for a whole run, and `.plasma.plot`, `.plasma.analysis` and `.plasma.data` on +every product: + +```python +out.plot.energies() # the energy budget of the run +phi = out.evaluate("em_fields/phi") +phi.plasma.plot.slice(coords="physical", plane="XY", t=-1, eta3=0) +phi.plasma.plot.animation(x="eta1", y="eta2", eta3=0) # over time +phi.plasma.analysis.mode_spectrum() # poloidal and toroidal mode numbers +out.kinetic_ions.orbits.plasma.plot.poloidal() # guiding-center orbits +``` + +`import plasma_plots; help(plasma_plots)` gives an overview of what the package does, and `help()` on any method, +e.g. `help(phi.plasma.plot.slice)`, its parameters. Its guides and full reference are at +. + ## MPI post-processing `Output` always uses `MPI.COMM_WORLD`; no communicator is passed to its constructor. diff --git a/doc/sections/quickstart.rst b/doc/sections/quickstart.rst index d9e32d3ae..9b362aa83 100644 --- a/doc/sections/quickstart.rst +++ b/doc/sections/quickstart.rst @@ -76,7 +76,10 @@ For periodic boundary conditions we will stabilize via ``options``. sim = Simulation(model=model, domain=domain, grid=grid) 6. Run the simulation. ``sim.run()`` returns an :class:`~struphy.Output` object, - the entry point for all post-processing. + the entry point for all post-processing. For plots and diagnostics made for Struphy + output (``out.plot``, ``out.analysis`` and ``.plasma.plot`` on every product), install + the separate package `plasma-plots `_ with + ``pip install "struphy[pproc]"``; ``Output`` loads it automatically. .. code-block:: python diff --git a/feectools b/feectools index a536b67a5..280d21e3a 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit a536b67a5401b9329d783e0e0fe0c61fb5132857 +Subproject commit 280d21e3ac749c15b55f364daecf0ece833bc289 diff --git a/pyproject.toml b/pyproject.toml index 71b7e2a64..666a0d1c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,8 +63,12 @@ phys = [ "gvec>=1.1.0, <=1.5.0", "desc-opt<=0.17.1", ] +pproc = [ + "plasma-plots>=0.1.0", +] dev = [ "struphy[mpi]", + "struphy[pproc]", "notebook", "autopep8", "isort", @@ -83,6 +87,7 @@ dev = [ ] doc = [ "struphy[phys]", + "struphy[pproc]", "jupyter", "nbconvert", "ipykernel", @@ -106,6 +111,7 @@ likwid = [ ] all = [ "struphy[phys]", + "struphy[pproc]", "struphy[dev]", "struphy[mpi]", "struphy[doc]", diff --git a/src/struphy/post_processing/output.py b/src/struphy/post_processing/output.py index 7123fc017..2c166f56f 100644 --- a/src/struphy/post_processing/output.py +++ b/src/struphy/post_processing/output.py @@ -59,6 +59,32 @@ def mpi_comm_world(): return MPI.COMM_WORLD +PLOTS_HINT = ( + "plots and diagnostics of Struphy output come from the plasma-plots package: " + "pip install plasma-plots (or struphy[pproc]); see https://struphy-hub.github.io/plasma-plots" +) +_plots = {"loaded": False} + + +def load_plasma_plots() -> bool: + """Load plasma-plots if it is installed, and tell whether it is. + + Importing ``plasma_plots`` registers ``out.plot``, ``out.analysis`` and the ``.plasma`` + accessor on every product. :class:`Output` calls this when it is created, so none of that + needs an explicit ``import plasma_plots``. Struphy does not depend on the package. + """ + if not _plots["loaded"]: + try: + import plasma_plots # noqa: F401 (registers the accessors) + except ImportError: + return False + except Exception as error: # a broken install must not break reading output + logger.warning("plasma-plots is installed but could not be imported: %s", error) + return False + _plots["loaded"] = True + return True + + class ProductMapping(Mapping[str, xr.DataArray]): """A discoverable mapping whose products are loaded on first access.""" @@ -206,6 +232,20 @@ class Output: reconstructed lazily from saved metadata. No simulation object is created or retained. * Every array carries the run in ``attrs["run"]`` (:attr:`label`) and ``attrs["run_name"]``. + **Plots and diagnostics** of the output live in the separate package + `plasma-plots `_ (``pip install plasma-plots``, + or ``pip install "struphy[pproc]"``). When it is installed, creating an ``Output`` loads it, + which adds: + + * ``out.plot`` and ``out.analysis``: whole-run plots and diagnostics, e.g. + ``out.plot.energies()``, ``out.analysis.time_fft("em_fields/phi")``; + * ``.plasma.plot``, ``.plasma.analysis`` and ``.plasma.data`` on every product, e.g. + ``out.evaluate("em_fields/phi").plasma.plot.slice(x="eta1", y="eta2", t=-1)``, or + ``orbits.plasma.plot.poloidal()`` for an orbits Dataset. + + ``import plasma_plots; help(plasma_plots)`` gives an overview of the package, and ``help()`` + on any accessor method (e.g. ``help(phi.plasma.plot.slice)``) its parameters. + Parameters ---------- path_out: @@ -213,6 +253,7 @@ class Output: """ def __init__(self, path_out): + load_plasma_plots() # out.plot, out.analysis and .plasma on every product, if installed self.path_out = Path(path_out).resolve() self._time_units = "normalized" self.comm = mpi_comm_world() @@ -317,6 +358,8 @@ def evaluate( products, a ``"species/variable"`` name selects the first matching binned product, then density/KDE product, then orbits. Pass ``dataset=`` to select a particular discovered product; use ``out.info("species/variable")`` to list them. + The full key of a particle product, ``"species/dataset/variable"`` (as :meth:`keys` + lists it, e.g. ``"kinetic_ions/e1_v1_density/f"``), selects that product directly. Supplying an ``eta`` evaluates a raw FEEC spline field directly on that logical grid. Each eta can be a scalar, a list, a one-dimensional array, or a ``range``; @@ -337,8 +380,16 @@ def evaluate( raise TypeError("physical is no longer supported; use eta1, eta2, eta3 and representation") if "as_numpy" in selectors: raise TypeError("evaluate() always returns xarray; call .to_numpy() on its result when needed") + if name.count("/") == 2: + # the full key of a particle product: "species/dataset/variable" + if dataset is not None: + raise ValueError("dataset= cannot be combined with a 'species/dataset/variable' name") + species, group, variable = name.split("/") + name, dataset = f"{species}/{variable}", f"{group}/{variable}" if name != "scalars" and name.count("/") != 1: - raise ValueError("evaluate() names must use the 'species/variable' form, or be 'scalars'") + raise ValueError( + "evaluate() names must use the 'species/variable' or 'species/dataset/variable' form, or be 'scalars'" + ) eta = (eta1, eta2, eta3) has_eta = any(value is not None for value in eta) @@ -2056,6 +2107,10 @@ def __getattr__(self, name: str) -> ProductNamespace: attribute.func(self) if isinstance(attribute, property): attribute.fget(self) # the property raised AttributeError itself; show its own error + if name in ("plot", "analysis"): + if load_plasma_plots() and name in type(self).__dict__: + return getattr(self, name) + raise AttributeError(f"Output has no {name!r} without plasma-plots; {PLOTS_HINT}") # the raw output names the species, so an unknown name never starts post-processing if name not in self._raw_species(): raise AttributeError(f"{name!r}; available species: {tuple(sorted(self._raw_species()))}") diff --git a/src/struphy/post_processing/tests/test_output.py b/src/struphy/post_processing/tests/test_output.py index 7ea51b9cd..a3ec7df8b 100644 --- a/src/struphy/post_processing/tests/test_output.py +++ b/src/struphy/post_processing/tests/test_output.py @@ -134,6 +134,35 @@ def save_raw_field(run, *names): file["feec/em_fields"].create_dataset(name, data=np.empty(0)) +# Environment prefixes through which MPI launchers tell a process which rank of a job it is. +MPI_LAUNCHER_ENV_PREFIXES = ( + "OMPI_", + "OPAL_", + "PMIX_", + "PMI_", + "PRTE_", + "MV2_", + "HYDRA_", + "I_MPI_", + "MPI_LOCALRANKID", + "ALPS_", + "PALS_", +) + + +@pytest.fixture +def outside_mpi_job(monkeypatch): + """Let child processes start as independent programs, not as ranks of this test's MPI job. + + Importing struphy initializes MPI. A child that inherits the launcher variables of an + ``mpirun`` rank initializes as that same rank, which hangs or corrupts the parent job. + """ + for name in list(os.environ): + if name.startswith(MPI_LAUNCHER_ENV_PREFIXES): + monkeypatch.delenv(name) + monkeypatch.setenv("STRUPHY_MPI", "0") + + @pytest.fixture def run(tmp_path): return Output(write_tree(str(tmp_path))) @@ -347,6 +376,18 @@ def test_evaluate_scalars_and_particle_defaults(run): selected = run.evaluate("kinetic_ions/f", dataset="e1_v1_density/delta_f") assert selected.name == "delta_f" + # the full key selects the same product, with selections as keywords + by_key = run.evaluate("kinetic_ions/e1_v1_density/delta_f") + xr.testing.assert_identical(by_key, selected) + first = float(by_key.eta1[0]) + xr.testing.assert_identical( + run.evaluate("kinetic_ions/e1_v1_density/delta_f", eta1=first), selected.sel(eta1=first) + ) + with pytest.raises(ValueError, match="dataset="): + run.evaluate("kinetic_ions/e1_v1_density/f", dataset="e1_v1_density/f") + with pytest.raises(KeyError): + run.evaluate("kinetic_ions/no_such_density/f") + orbits = run.evaluate("kinetic_ions/orbits") assert isinstance(orbits, xr.Dataset) and "weight" in orbits.data_vars