Update examples with struphy-plots accessors

docs/src/examples/README.md

@@ -98,9 +98,6 @@ filename.
 - Analyze with the `Output` of the run (`sim.output`), e.g.
   `sim.output.scalars["electric_energy"].struphy.analysis.damping_rate(window=(None, 8.0), amplitude=True)`;
   see `weak-landau-damping.py`. Print the measured results; the page shows the figures.
-- Older scripts still save their own output with the helpers of `struphy_plots.gallery`
-  (`save_figure`, `save_extra_figure`, `merge_metadata`, `export_profiling`); `run_example.py`
-  runs those as they are. Write new scripts in the style above.
 
 ## 2. Generate the structural metadata
 

docs/src/examples/acoustic-pulse.py

@@ -134,14 +134,14 @@ def pproc(sim: Simulation, show: bool = False):
     x = rho.eta1.values * length
     density = rho.values
     exact = np.array([exact_density(x, t) for t in times])
-    error = float(np.max(np.abs(density - exact)) / amplitude)
+    error = float(rho.struphy.analysis.error(exact, norm="max", dims=("t", "eta1"))) / amplitude
     print(f"Maximum error of the density, relative to the pulse height: {error:.3f}")
 
     kinetic = np.asarray(output.scalars["en_U"])
     thermo = np.asarray(output.scalars["en_thermo"])
     total = np.asarray(output.scalars["en_tot"])
     scalar_times = np.asarray(output.time)[: len(total)]
-    energy_drift = float(np.max(np.abs(total / total[0] - 1.0)))
+    energy_drift = float(output.scalars["en_tot"].struphy.analysis.relative_error().max())
     print(f"Maximum relative drift of the total energy: {energy_drift:.2e}")
     energy_scale = float(kinetic.max())  # the largest kinetic energy the pulse reaches
 

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

@@ -106,12 +106,10 @@ def pproc(sim: Simulation, show: bool = False):
 
     # The exchange period is half the wave period; measure it from the kinetic-energy maxima.
     kinetic_values = kinetic.values
-    peaks = np.flatnonzero(
-        (kinetic_values[1:-1] > kinetic_values[:-2]) & (kinetic_values[1:-1] >= kinetic_values[2:])
-    ) + 1
-    if peaks.size < 3:
+    peak_times = kinetic.struphy.analysis.envelope().t.values
+    if peak_times.size < 3:
         raise RuntimeError("Too few kinetic-energy maxima to measure the exchange period")
-    measured_period = float(np.mean(np.diff(time[peaks])))
+    measured_period = float(np.mean(np.diff(peak_times)))
     period_error = abs(measured_period / (0.5 * period) - 1.0)
     print(f"Measured exchange period: {measured_period:.4f} (exact: {0.5 * period:.4f}, error {period_error:.2%})")
 

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

@@ -14,8 +14,6 @@
 
 import argparse
 
-import plotly.graph_objects as go
-
 from struphy import (
     BinningPlot,
     BoundaryParameters,
@@ -125,26 +123,7 @@ def pproc(sim: Simulation, show: bool = False):
     growth_rate = field_energy.struphy.analysis.growth_rate(window=(5.0, 25.0), amplitude=True).rate
     print(f"Measured growth rate: {growth_rate:.4f}")
 
-    figure = go.Figure(
-        data=[
-            go.Scatter(
-                x=field_energy.t.values,
-                y=field_energy.values,
-                mode="lines",
-                name="Struphy (PIC)",
-                line={"color": "#168aad", "width": 3},
-            ),
-        ],
-    )
-    figure.update_layout(
-        title="Bump-on-tail instability: electric field energy",
-        xaxis_title="t [a.u.]",
-        yaxis_title="E² / 2 [a.u.]",
-        yaxis={"type": "log"},
-        template="plotly_white",
-        autosize=True,
-        margin={"l": 70, "r": 30, "t": 80, "b": 60},
-    )
+    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)
 

docs/src/examples/coaxial-waveguide.py

@@ -97,36 +97,40 @@ def save(figure, name: str, *, show: bool = False, frame: int | None = None, sti
 
 
 def pproc(sim: Simulation, show: bool = False):
+    from struphy_plots.analysis import evaluate_on
+
     output = sim.output
     output.pproc(physical=True)
 
     # The axial magnetic field on the (r, theta) evaluation grid, and the exact mode at the same points.
     b_z = output.evaluate("em_fields/b_field_xyz").isel(component=2, eta3=0)  # (t, e1, e2)
     times = b_z.t.values
-    radius = np.hypot(b_z.X.values, b_z.Y.values)
-    angle = np.arctan2(b_z.Y.values, b_z.X.values)
 
-    def exact_b_z(time):
+    def exact_b_z(x, y, z, time):
+        radius, angle = np.hypot(x, y), np.arctan2(y, x)
         profile = jv(mode_number, radius) - bessel_ratio * yn(mode_number, radius)
         return profile * np.cos(mode_number * angle - time)
 
-    error = np.stack([b_z.isel(t=index).values - exact_b_z(times[index]) for index in range(len(times))])
-    amplitude = float(np.abs(exact_b_z(0.0)).max())
-    relative_error = np.abs(error).max(axis=(1, 2)) / amplitude
-    print(f"Largest relative error of B_z: {relative_error.max():.2e}")
+    exact = evaluate_on(b_z, exact_b_z)
+    error = b_z.struphy.analysis.error(exact, norm="pointwise").values
+    amplitude = float(np.abs(exact.isel(t=0)).max())
+    relative_error = b_z.struphy.analysis.error(exact, norm="max") / amplitude
+    print(f"Largest relative error of B_z: {float(relative_error.max()):.2e}")
 
     # The mode's frequency, from a sinusoid fitted to B_z at a probe in the middle of the gap. The
     # exact mode is proportional to cos(m theta - t), so its frequency is 1.
-    probe = (b_z.sizes["eta1"] // 2, 0)
-    signal = b_z.values[:, probe[0], probe[1]]
-    exact_signal = exact_b_z(times[:, None, None])[:, probe[0], probe[1]]
-    fit, _ = curve_fit(lambda t, a, w, phase: a * np.cos(w * t + phase), times, signal, p0=[signal.max(), 1.0, 0.0])
+    probe = {"eta1": b_z.sizes["eta1"] // 2, "eta2": 0}
+    signal = b_z.isel(probe).assign_attrs(label="B_z at the probe")
+    exact_signal = exact.isel(probe)
+    fit, _ = curve_fit(
+        lambda t, a, w, phase: a * np.cos(w * t + phase), times, signal.values, p0=[float(signal.max()), 1.0, 0.0]
+    )
     measured_frequency = float(fit[1])
     frequency_error = abs(measured_frequency - 1.0)
     print(f"Measured frequency: {measured_frequency:.5f} (exact: 1, error {frequency_error:.1e})")
 
     energy = output.scalars["total_energy"]
-    energy_drift = float(np.abs(energy.values / energy.values[0] - 1.0).max())
+    energy_drift = float(energy.struphy.analysis.relative_error().max())
     print(f"Largest relative change of the total energy: {energy_drift:.1e}")
     electric = output.scalars["electric_energy"]
     magnetic = output.scalars["magnetic_energy"]
@@ -242,19 +246,11 @@ def frame_traces(index):
     save(figure, "coaxial-waveguide", height=650, frame=still_position, show=show)
 
     # The probe signal against the exact mode.
-    probe_figure = go.Figure()
-    probe_figure.add_scatter(
-        x=times, y=exact_signal, mode="lines", name="exact mode", line={"color": "#d62828", "width": 3, "dash": "dot"}
-    )
-    probe_figure.add_scatter(x=times, y=signal, mode="lines", name="Struphy", line={"color": "#168aad", "width": 3})
-    probe_figure.update_layout(
+    probe_figure = signal.struphy.plot.timeseries(
+        logy=False,
+        reference={"exact mode": exact_signal},
         title=f"Coaxial waveguide: B_z at a probe, measured frequency {measured_frequency:.4f} (exact: 1)",
-        xaxis_title="t [a.u.]",
-        yaxis_title="B_z at the probe [a.u.]",
-        template="plotly_white",
-        autosize=True,
-        legend={"orientation": "h", "x": 0.5, "xanchor": "center", "y": 1.0, "yanchor": "bottom"},
-        margin={"l": 70, "r": 30, "t": 110, "b": 60},
+        backend="plotly",
     )
 
     # The energies, and the relative change of the total energy.

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

@@ -15,6 +15,7 @@
 
 import numpy as np
 import plotly.graph_objects as go
+import xarray as xr
 
 from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, perturbations
 from struphy.models import ColdPlasma
@@ -161,16 +162,14 @@ def profile_traces(index):
     figure.update_yaxes(title_text="energy / initial energy", row=2, col=1)
     save(figure, "cold-plasma-oscillation", width=900, height=850, show=show)
 
-    # The frequency against the density, on the line omega = omega_p.
+    # 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)
-    scan = go.Figure()
-    scan.add_scatter(x=n_line, y=alpha / epsilon * np.sqrt(n_line), mode="lines", name="ω_p ∝ √n₀",
-                     line={"color": "#111", "width": 2, "dash": "dash"})
-    scan.add_scatter(x=list(measured), y=list(measured.values()), mode="markers", name="Struphy",
-                     marker={"color": "#d62828", "size": 12})
-    scan.update_layout(title="Oscillation frequency against density", template="plotly_white",
-                       xaxis_title="density n₀", yaxis_title="angular frequency ω",
-                       margin={"l": 70, "r": 30, "t": 80, "b": 60})
+    frequencies = xr.DataArray(list(measured.values()), dims="n0", coords={"n0": list(measured)}, name="Struphy")
+    scan = frequencies.struphy.plot.against_theory(
+        {"ω_p ∝ √n₀": (n_line, alpha / epsilon * np.sqrt(n_line))}, show_error=False,
+        xlabel="density n₀", ylabel="angular frequency ω", title="Oscillation frequency against density",
+        backend="plotly",
+    )
     save(scan, "cold-plasma-oscillation-frequency-scan", show=show)
 
 

docs/src/examples/cold-plasma-wave-packet.py

@@ -32,24 +32,6 @@
 amplitude = 0.05
 
 
-def branch_frequency(k, wave):
-    """Frequency of the whistler, the upper R wave or the L wave along B0, in units of the cyclotron frequency.
-
-    With c = 1, k^2 = omega^2 - omega_p^2 omega / (omega -+ Omega_c), i.e. the cubic
-    omega^3 -+ Omega_c omega^2 - (omega_p^2 + k^2) omega +- k^2 Omega_c = 0 with the upper sign for the R wave. The whistler is the
-    middle positive root of the R cubic, the upper R wave its largest root, and the L wave the largest root of the L cubic.
-    """
-    if wave == "whistler":
-        return np.sort(np.roots([1.0, -1.0, -(alpha**2 + k**2), k**2]).real)[-2]
-    if wave == "R wave":
-        return np.sort(np.roots([1.0, -1.0, -(alpha**2 + k**2), k**2]).real)[-1]
-    return np.sort(np.roots([1.0, 1.0, -(alpha**2 + k**2), -(k**2)]).real)[-1]
-
-
-def group_velocity(k, wave, step=1e-4):
-    return (branch_frequency(k + step, wave) - branch_frequency(k - step, wave)) / (2 * step)
-
-
 def envelope(z):
     return amplitude * np.exp(-0.5 * ((z - center) / width) ** 2)
 
@@ -120,6 +102,8 @@ def save(figure, name: str, *, show: bool = False, frame: int | None = None, sti
 
 
 def pproc(sim: Simulation, show: bool = False):
+    from struphy_plots.theory.waves import Species, cold_plasma_waves, group_velocity
+
     time_opts = sim.time_opts
 
     output = sim.output
@@ -128,8 +112,14 @@ def pproc(sim: Simulation, show: bool = False):
     times, z, density = packet_energy(output)
     waves = ("whistler", "R wave", "L wave")
     tracked = ("whistler", "L wave")  # the slow packet holds the whistler and upper R branches, of almost equal speed
-    exact_speed = {wave: group_velocity(k0, wave) for wave in waves}
-    exact_frequency = {wave: branch_frequency(k0, wave) for wave in waves}
+    # The branches along B0 (theta = 0) of the cold electron plasma. At k0 they are, by ascending frequency, the whistler,
+    # the plasma oscillation (omega = omega_p), the L wave and the upper R wave.
+    electrons = Species(plasma_frequency=alpha, cyclotron_frequency=-1.0)
+    branch = {"whistler": "branch 1", "L wave": "branch 3", "R wave": "branch 4"}
+    frequencies = cold_plasma_waves(k0, 0.0, electrons)
+    speeds = group_velocity(lambda k: cold_plasma_waves(k, 0.0, electrons), k0, step=1e-4)
+    exact_speed = {wave: float(speeds[branch[wave]].real) for wave in waves}
+    exact_frequency = {wave: float(frequencies[branch[wave]].real) for wave in waves}
 
     def track_peak(side, speed):
         """Position of the energy peak of the packet on one side of the launch point, in a window around its characteristic."""

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

@@ -14,7 +14,6 @@
 import argparse
 
 import numpy as np
-import plotly.graph_objects as go
 
 from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, perturbations
 from struphy.models import ColdPlasma
@@ -144,58 +143,28 @@ def pproc(sim: Simulation, show: bool = False):
     if not all(np.isfinite(error) for error in errors.values()):
         raise RuntimeError("A wave branch could not be read off the spectrum")
 
-    log_power = np.log10(np.clip(power, 1e-12, None))
-    figure = go.Figure(
-        go.Heatmap(
-            x=k, y=omega, z=log_power, zmin=-8, zmax=0, colorscale="Plasma",
-            colorbar={"title": {"text": "log₁₀ P"}},
-            hovertemplate="k=%{x:.3f}<br>ω=%{y:.3f}<br>log₁₀ P=%{z:.2f}<extra></extra>",
-        )
-    )
-    styles = {
-        "R-wave": ("R wave", "#2a9d8f"),
-        "L-wave": ("L wave", "#e9c46a"),
-        "whistler": ("whistler (R, below Ω_c)", "#48cae4"),
-    }
-    for name, (label, color) in styles.items():
-        figure.add_scatter(
-            x=k_fine, y=branch_curves[name], mode="lines", name=label,
-            line={"color": color, "width": 2.5, "dash": "dash"},
-        )
-    for level, label in ((1.0, "Ω_c"), (omega_R, "ω_R cutoff"), (omega_L, "ω_L cutoff")):
-        figure.add_hline(y=level, line={"color": "rgba(255,255,255,0.55)", "width": 1, "dash": "dot"},
-                         annotation_text=label, annotation_position="bottom right",
-                         annotation_font_color="white")
-    figure.update_layout(
+    # The normalized power spectrum over 8 decades, with the analytic branches and the cutoffs.
+    labels = {"R-wave": "R wave", "L-wave": "L wave", "whistler": "whistler (R, below Ω_c)"}
+    figure = spectrum.struphy.plot.dispersion(
+        kmin=0,
+        kmax=k_top,
+        omega_max=omega_top,
+        dynamic_range=8,
+        cmap="plasma",
+        branches={label: (k_fine, branch_curves[name]) for name, label in labels.items()},
+        frequencies={"Ω_c": 1.0, "ω_R cutoff": omega_R, "ω_L cutoff": omega_L},
         title="Cold-plasma waves along B₀: power spectrum of E_x",
-        xaxis_title="k c / Ω_c", yaxis_title="ω / Ω_c", template="plotly_white", autosize=True,
-        legend={"orientation": "h", "y": -0.18}, margin={"l": 75, "r": 40, "t": 80, "b": 110},
+        backend="plotly",
     )
-    figure.update_xaxes(range=[0, k_top])
-    figure.update_yaxes(range=[0, omega_top])
     save(figure, "cold-plasma-waves", height=700, show=show)
 
     # Energy channels: the noise starts purely electric, then shares its energy with the magnetic field
     # and the electron current, while the sum stays constant.
-    channels = {
-        "electric_energy": ("electric", "#168aad"),
-        "magnetic_energy": ("magnetic", "#d62828"),
-        "kinetic_energy": ("electron current", "#f4a261"),
-        "total_energy": ("total", "#264653"),
-    }
-    energies = {name: output.scalars[name] for name in channels}
-    time = energies["total_energy"].t.values
-    total = energies["total_energy"].values
-    relative_drift = float(np.max(np.abs(total / total[0] - 1.0)))
+    energies = [output.scalars[name] for name in ("electric_energy", "magnetic_energy", "kinetic_energy", "total_energy")]
+    relative_drift = float(energies[-1].struphy.analysis.relative_error().max())
     print(f"Maximum relative drift of the total energy: {relative_drift:.2e}")
-    energy_figure = go.Figure()
-    for name, (label, color) in channels.items():
-        energy_figure.add_scatter(x=time, y=energies[name].values, mode="lines", name=label,
-                                  line={"color": color, "width": 3 if name == "total_energy" else 2})
-    energy_figure.update_layout(
-        title="Energy channels of the cold plasma", xaxis_title="t Ω_c", yaxis_title="energy [a.u.]",
-        template="plotly_white", autosize=True, legend={"orientation": "h", "y": -0.2},
-        margin={"l": 75, "r": 30, "t": 80, "b": 100},
+    energy_figure = energies[0].struphy.plot.timeseries(
+        *energies[1:], logy=False, title="Energy channels of the cold plasma", backend="plotly"
     )
     save(energy_figure, "cold-plasma-waves-energy", show=show)
 

docs/src/examples/dam-break.py

@@ -154,8 +154,8 @@ def pproc(sim: Simulation, show: bool = False):
     density_limit = float(density.max())
 
     def frame_traces(index, webgl=True):
-        # In an animation, plotly.js 3.7 stops drawing a heatmap that shares the figure with SVG scatter
-        # frames, so the markers of the interactive figure are drawn with WebGL.
+        # The interactive figure draws its many markers with WebGL, which is faster; the still image
+        # draws them as SVG.
         markers = go.Scattergl if webgl else go.Scatter
         return [
             markers(