Move save to struphy plots

docs/src/examples/README.md

@@ -86,10 +86,10 @@ filename.
 - `pproc(sim, show=False)` contains only the physics. Draw the figures with struphy-plots and
   `backend="plotly"`, e.g. `e_x.struphy.plot.slice(x="z", y="t", symmetric=True, backend="plotly")`
   or `spectrum.struphy.plot.dispersion(kmin=0, branches=..., fits=..., backend="plotly")` (or with
-  `plotly.graph_objects` for anything struphy-plots does not draw), and save each with the script's
-  own `save(figure, name, show=show)`, copied from `maxwell-wave.py`, which writes `name.html`,
-  `.png` and `.plotly.json` on MPI rank 0. Name the page's main figure `<script-stem>` and the
-  others `<script-stem>-<key>`.
+  `plotly.graph_objects` for anything struphy-plots does not draw), and save each with
+  `struphy_plots.save_figure(figure, name, show=show)`, which writes `name.html`, `.png` and
+  `.plotly.json` on MPI rank 0 (`frame=` or `still=` choose the image of an animation). Name the
+  page's main figure `<script-stem>` and the others `<script-stem>-<key>`.
 - The website's part is `run_example.py`, which `python cli.py run` and CI call: it runs the script
   in `docs/public/examples/`, profiles its simulation, exports the profiling data into
   `<script-stem>.metadata.json` and copies the PNGs to `docs/public/images/examples/`.
@@ -147,7 +147,7 @@ GitHub-hosted runners execute each example with `run_example.py` in a single pro
 For optional local MPI runs: `python cli.py run <example> --mpi 4`, or, from the repository root,
 `mpirun -n 4 python run_example.py <script-stem>`.
 
-- The simulation, `output.pproc(...)` and the analysis run on every rank; the scripts' `save`
+- The simulation, `output.pproc(...)` and the analysis run on every rank; `save_figure`
   writes on rank 0 only (struphy-plots draws nothing on the other ranks), as does `run_example.py`.
 - The grid must split over the ranks: with four ranks, keep at least a few cells per rank and direction
   (Orszag–Tang uses 32 × 32 × 1, i.e. 16 × 16 cells per rank).

docs/src/examples/acoustic-pulse.py

@@ -30,6 +30,7 @@
 )
 from struphy.linear_algebra.solver import NonlinearSolverParameters
 from struphy.models import VariationalCompressibleFluid
+from struphy_plots import save_figure
 
 length = 2 * np.pi
 width = 0.5  # of the Gaussian pulse
@@ -104,25 +105,6 @@ def create_simulation() -> Simulation:
     return sim
 
 
-def save(figure, name: str, *, show: bool = False, frame: int | None = None, still=None, width=1100, height=650):
-    """Save a Plotly figure as ``<name>.html``, ``<name>.png`` and ``<name>.plotly.json``.
-
-    ``figure`` is a plot of struphy-plots drawn with ``backend="plotly"``, or a
-    ``plotly.graph_objects.Figure``. ``show`` shows it first. For an animation, the PNG shows
-    ``frame`` (default: the first), or the figure ``still`` instead. Under MPI only rank 0 writes.
-    """
-    import struphy_plots
-    from struphy_plots.plotting import PlotResult
-
-    if not struphy_plots.is_plotting_rank():
-        return
-    result = figure if isinstance(figure, PlotResult) else PlotResult(figure, None)
-    if show:
-        result.show()
-    result.save(f"{name}.html")
-    image = PlotResult(still, None) if still is not None else result
-    image.save(f"{name}.png", frame=frame, width=width, height=height, scale=2)
-    result.save(f"{name}.plotly.json")
 def pproc(sim: Simulation, show: bool = False):
     from plotly.subplots import make_subplots
 
@@ -176,7 +158,7 @@ def profile_traces(index):
     figure.update_yaxes(title_text="density", range=[1 - 0.3 * amplitude, 1 + 1.1 * amplitude], row=1, col=1)
     figure.update_xaxes(title_text="t", row=2, col=1)
     figure.update_yaxes(title_text="energy change / largest kinetic energy", row=2, col=1)
-    save(figure, "acoustic-pulse", width=900, height=850, show=show)
+    save_figure(figure, "acoustic-pulse", width=900, height=850, show=show)
 
     density_change = rho.copy(data=density - 1.0)
     space_time = density_change.assign_coords(eta1=x).struphy.plot.slice(
@@ -190,7 +172,7 @@ def profile_traces(index):
         colorbar_label="ρ − 1",
         backend="plotly",
     )
-    save(space_time, "acoustic-pulse-space-time", show=show)
+    save_figure(space_time, "acoustic-pulse-space-time", show=show)
 
 
 if __name__ == "__main__":

docs/src/examples/alfven-standing-wave.py

@@ -22,6 +22,7 @@
 
 from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, perturbations
 from struphy.models import ShearAlfven
+from struphy_plots import save_figure
 
 # A uniform background with B0 along z and n0 = 1, so the Alfvén speed is 1.
 B0z, n0, beta = 1.0, 1.0, 0.1
@@ -70,27 +71,6 @@ def create_simulation() -> Simulation:
     return sim
 
 
-def save(figure, name: str, *, show: bool = False, frame: int | None = None, still=None, width=1100, height=650):
-    """Save a Plotly figure as ``<name>.html``, ``<name>.png`` and ``<name>.plotly.json``.
-
-    ``figure`` is a plot of struphy-plots drawn with ``backend="plotly"``, or a
-    ``plotly.graph_objects.Figure``. ``show`` shows it first. For an animation, the PNG shows
-    ``frame`` (default: the first), or the figure ``still`` instead. Under MPI only rank 0 writes.
-    """
-    import struphy_plots
-    from struphy_plots.plotting import PlotResult
-
-    if not struphy_plots.is_plotting_rank():
-        return
-    result = figure if isinstance(figure, PlotResult) else PlotResult(figure, None)
-    if show:
-        result.show()
-    result.save(f"{name}.html")
-    image = PlotResult(still, None) if still is not None else result
-    image.save(f"{name}.png", frame=frame, width=width, height=height, scale=2)
-    result.save(f"{name}.plotly.json")
-
-
 def pproc(sim: Simulation, show: bool = False):
     output = sim.output
     output.pproc(physical=True)
@@ -132,7 +112,7 @@ def pproc(sim: Simulation, show: bool = False):
         margin={"l": 75, "r": 30, "t": 80, "b": 100},
     )
     figure.update_yaxes(range=[-0.05, 1.15])
-    save(figure, "alfven-standing-wave", show=show)
+    save_figure(figure, "alfven-standing-wave", show=show)
 
     # The velocity along z over time: a standing wave keeps its nodes, so the stripes are vertical,
     # unlike the diagonal stripes of the travelling waves in `shear-alfven-wave`.
@@ -151,7 +131,7 @@ def pproc(sim: Simulation, show: bool = False):
         colorbar_label="u₁ (logical component)",
         backend="plotly",
     )
-    save(space_time, "alfven-standing-wave-space-time", show=show)
+    save_figure(space_time, "alfven-standing-wave-space-time", show=show)
 
 
 if __name__ == "__main__":

docs/src/examples/beltrami-sph.py

@@ -41,6 +41,7 @@
 )
 from struphy.models import PressureLessSPH
 from struphy.ode.utils import ButcherTableau
+from struphy_plots import save_figure
 
 box_min = -0.5
 box_max = 0.5
@@ -139,27 +140,6 @@ def create_simulation() -> Simulation:
     return sim
 
 
-def save(figure, name: str, *, show: bool = False, frame: int | None = None, still=None, width=1100, height=650):
-    """Save a Plotly figure as ``<name>.html``, ``<name>.png`` and ``<name>.plotly.json``.
-
-    ``figure`` is a plot of struphy-plots drawn with ``backend="plotly"``, or a
-    ``plotly.graph_objects.Figure``. ``show`` shows it first. For an animation, the PNG shows
-    ``frame`` (default: the first), or the figure ``still`` instead. Under MPI only rank 0 writes.
-    """
-    import struphy_plots
-    from struphy_plots.plotting import PlotResult
-
-    if not struphy_plots.is_plotting_rank():
-        return
-    result = figure if isinstance(figure, PlotResult) else PlotResult(figure, None)
-    if show:
-        result.show()
-    result.save(f"{name}.html")
-    image = PlotResult(still, None) if still is not None else result
-    image.save(f"{name}.png", frame=frame, width=width, height=height, scale=2)
-    result.save(f"{name}.plotly.json")
-
-
 def pproc(sim: Simulation, show: bool = False):
     output = sim.output
     output.pproc()
@@ -359,7 +339,7 @@ def error_trace(index, values, name, color):
     # The frames name only the markers and errors (the contour stays), so build the final still.
     still = go.Figure(data=final_data, layout=figure.layout)
     still.layout.sliders[0].active = len(figure.frames) - 1
-    save(figure, "beltrami-sph", still=still, width=1100, height=680, show=show)
+    save_figure(figure, "beltrami-sph", still=still, width=1100, height=680, show=show)
 
     # The kernel reconstruction shows the simulated mass density independently of the marker view.
     # The exact divergence-free Beltrami transport preserves the initially uniform rho = 1.
@@ -594,10 +574,10 @@ def signed_cell_area(grid_x, grid_y):
         f"maximum marker-cell area deviation from 1 {max_area_deviation:.3e}"
     )
 
-    save(density_figure, "beltrami-sph-density", frame=-1, show=show)
-    save(compression_figure, "beltrami-sph-compression", frame=-1, show=show)
-    save(area_figure, "beltrami-sph-area-deformation", frame=-1, show=show)
-    save(trajectory_figure, "beltrami-sph-trajectories", show=show)
+    save_figure(density_figure, "beltrami-sph-density", frame=-1, show=show)
+    save_figure(compression_figure, "beltrami-sph-compression", frame=-1, show=show)
+    save_figure(area_figure, "beltrami-sph-area-deformation", frame=-1, show=show)
+    save_figure(trajectory_figure, "beltrami-sph-trajectories", show=show)
 
 
 if __name__ == "__main__":

docs/src/examples/bump-on-tail.py

@@ -31,6 +31,7 @@
     perturbations,
 )
 from struphy.models import VlasovAmpereOneSpecies
+from struphy_plots import save_figure
 
 # A 90% bulk Maxwellian plus a 10% "bump" population drifting at u1 = -4.5.
 perturbation_amplitude = 0.05
@@ -91,27 +92,6 @@ def create_simulation() -> Simulation:
     return sim
 
 
-def save(figure, name: str, *, show: bool = False, frame: int | None = None, still=None, width=1100, height=650):
-    """Save a Plotly figure as ``<name>.html``, ``<name>.png`` and ``<name>.plotly.json``.
-
-    ``figure`` is a plot of struphy-plots drawn with ``backend="plotly"``, or a
-    ``plotly.graph_objects.Figure``. ``show`` shows it first. For an animation, the PNG shows
-    ``frame`` (default: the first), or the figure ``still`` instead. Under MPI only rank 0 writes.
-    """
-    import struphy_plots
-    from struphy_plots.plotting import PlotResult
-
-    if not struphy_plots.is_plotting_rank():
-        return
-    result = figure if isinstance(figure, PlotResult) else PlotResult(figure, None)
-    if show:
-        result.show()
-    result.save(f"{name}.html")
-    image = PlotResult(still, None) if still is not None else result
-    image.save(f"{name}.png", frame=frame, width=width, height=height, scale=2)
-    result.save(f"{name}.plotly.json")
-
-
 def pproc(sim: Simulation, show: bool = False):
     domain = sim.domain
 
@@ -125,7 +105,7 @@ def pproc(sim: Simulation, show: bool = False):
 
     figure = field_energy.struphy.plot.timeseries(logy=True, title="Bump-on-tail instability: electric field energy", backend="plotly")
 
-    save(figure, "bump-on-tail", show=show)
+    save_figure(figure, "bump-on-tail", show=show)
 
     # Evaluate the saved products on their grids, then look at them in more than one way.
     output.pproc()
@@ -162,8 +142,8 @@ def pproc(sim: Simulation, show: bool = False):
         backend="plotly",
     )
 
-    save(phase_space, "bump-on-tail-phasespace", frame=len(phase_space.fig.frames) // 2, show=show)
-    save(velocity_time, "bump-on-tail-velocity-time", show=show)
+    save_figure(phase_space, "bump-on-tail-phasespace", frame=len(phase_space.fig.frames) // 2, show=show)
+    save_figure(velocity_time, "bump-on-tail-velocity-time", show=show)
 
 
 if __name__ == "__main__":

docs/src/examples/coaxial-waveguide.py

@@ -23,6 +23,7 @@
 
 from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, perturbations
 from struphy.models import Maxwell
+from struphy_plots import save_figure
 
 # The inner and outer radii are two zeros of the derivative of the radial profile
 # J_m(r) - 0.28 Y_m(r), so that the tangential electric field vanishes on both conducting walls.
@@ -75,27 +76,6 @@ def create_simulation() -> Simulation:
     return sim
 
 
-def save(figure, name: str, *, show: bool = False, frame: int | None = None, still=None, width=1100, height=650):
-    """Save a Plotly figure as ``<name>.html``, ``<name>.png`` and ``<name>.plotly.json``.
-
-    ``figure`` is a plot of struphy-plots drawn with ``backend="plotly"``, or a
-    ``plotly.graph_objects.Figure``. ``show`` shows it first. For an animation, the PNG shows
-    ``frame`` (default: the first), or the figure ``still`` instead. Under MPI only rank 0 writes.
-    """
-    import struphy_plots
-    from struphy_plots.plotting import PlotResult
-
-    if not struphy_plots.is_plotting_rank():
-        return
-    result = figure if isinstance(figure, PlotResult) else PlotResult(figure, None)
-    if show:
-        result.show()
-    result.save(f"{name}.html")
-    image = PlotResult(still, None) if still is not None else result
-    image.save(f"{name}.png", frame=frame, width=width, height=height, scale=2)
-    result.save(f"{name}.plotly.json")
-
-
 def pproc(sim: Simulation, show: bool = False):
     from struphy_plots.analysis import evaluate_on
 
@@ -243,7 +223,7 @@ def frame_traces(index):
 
     # The still image and thumbnail show the mode a little after the start of the run.
     still_position = len(frame_indices) // 4
-    save(figure, "coaxial-waveguide", height=650, frame=still_position, show=show)
+    save_figure(figure, "coaxial-waveguide", height=650, frame=still_position, show=show)
 
     # The probe signal against the exact mode.
     probe_figure = signal.struphy.plot.timeseries(
@@ -291,8 +271,8 @@ def frame_traces(index):
     energy_figure.update_xaxes(title_text="t [a.u.]", row=2, col=1)
     energy_figure.update_layout(template="plotly_white", autosize=True, margin={"l": 80, "r": 30, "t": 80, "b": 60})
 
-    save(probe_figure, "coaxial-waveguide-frequency", show=show)
-    save(energy_figure, "coaxial-waveguide-energy", show=show)
+    save_figure(probe_figure, "coaxial-waveguide-frequency", show=show)
+    save_figure(energy_figure, "coaxial-waveguide-energy", show=show)
 
 
 if __name__ == "__main__":

docs/src/examples/cold-plasma-oscillation.py

@@ -19,6 +19,7 @@
 
 from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, perturbations
 from struphy.models import ColdPlasma
+from struphy_plots import save_figure
 
 # Time in units of the inverse cyclotron frequency and alpha = omega_p / omega_c = 1: at n0 = 1 the plasma
 # frequency is 1 (as in the cold-plasma-waves example).
@@ -74,27 +75,6 @@ def oscillation_frequency(times, values):
     return float(np.pi / np.mean(np.diff(roots)))
 
 
-def save(figure, name: str, *, show: bool = False, frame: int | None = None, still=None, width=1100, height=650):
-    """Save a Plotly figure as ``<name>.html``, ``<name>.png`` and ``<name>.plotly.json``.
-
-    ``figure`` is a plot of struphy-plots drawn with ``backend="plotly"``, or a
-    ``plotly.graph_objects.Figure``. ``show`` shows it first. For an animation, the PNG shows
-    ``frame`` (default: the first), or the figure ``still`` instead. Under MPI only rank 0 writes.
-    """
-    import struphy_plots
-    from struphy_plots.plotting import PlotResult
-
-    if not struphy_plots.is_plotting_rank():
-        return
-    result = figure if isinstance(figure, PlotResult) else PlotResult(figure, None)
-    if show:
-        result.show()
-    result.save(f"{name}.html")
-    image = PlotResult(still, None) if still is not None else result
-    image.save(f"{name}.png", frame=frame, width=width, height=height, scale=2)
-    result.save(f"{name}.plotly.json")
-
-
 def pproc(sim: Simulation, show: bool = False):
     from plotly.subplots import make_subplots
 
@@ -160,7 +140,7 @@ def profile_traces(index):
     figure.update_yaxes(title_text="E_z", range=[-1.2 * amplitude, 1.2 * amplitude], row=1, col=1)
     figure.update_xaxes(title_text="t", row=2, col=1)
     figure.update_yaxes(title_text="energy / initial energy", row=2, col=1)
-    save(figure, "cold-plasma-oscillation", width=900, height=850, show=show)
+    save_figure(figure, "cold-plasma-oscillation", width=900, height=850, show=show)
 
     # The frequency against the density, on the line omega = omega_p (drawn from n0 = 0).
     n_line = np.linspace(0.0, max(densities) * 1.05, 100)
@@ -170,7 +150,7 @@ def profile_traces(index):
         xlabel="density n₀", ylabel="angular frequency ω", title="Oscillation frequency against density",
         backend="plotly",
     )
-    save(scan, "cold-plasma-oscillation-frequency-scan", show=show)
+    save_figure(scan, "cold-plasma-oscillation-frequency-scan", show=show)
 
 
 if __name__ == "__main__":