Added struphy-plots integration

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

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
 
 ```

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
+<https://struphy-hub.github.io/plasma-plots>.
+
 ## MPI post-processing
 
 `Output` always uses `MPI.COMM_WORLD`; no communicator is passed to its constructor.

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 <https://struphy-hub.github.io/plasma-plots>`_ with
+   ``pip install "struphy[pproc]"``; ``Output`` loads it automatically.
 
 .. code-block:: python
 

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]",

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,13 +232,28 @@ 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 <https://struphy-hub.github.io/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:
         The simulation output folder, ``sim.env.path_out``.
     """
 
     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()))}")

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