diff --git a/docs/src/examples/README.md b/docs/src/examples/README.md
index 31dc060..2652295 100644
--- a/docs/src/examples/README.md
+++ b/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
diff --git a/docs/src/examples/acoustic-pulse.py b/docs/src/examples/acoustic-pulse.py
index 1d624fa..01e4d02 100644
--- a/docs/src/examples/acoustic-pulse.py
+++ b/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
diff --git a/docs/src/examples/alfven-standing-wave.py b/docs/src/examples/alfven-standing-wave.py
index f08d729..cbafd55 100644
--- a/docs/src/examples/alfven-standing-wave.py
+++ b/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%})")
diff --git a/docs/src/examples/bump-on-tail.py b/docs/src/examples/bump-on-tail.py
index e3f3794..2a45460 100644
--- a/docs/src/examples/bump-on-tail.py
+++ b/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)
diff --git a/docs/src/examples/coaxial-waveguide.py b/docs/src/examples/coaxial-waveguide.py
index 44d6be6..febbdb2 100644
--- a/docs/src/examples/coaxial-waveguide.py
+++ b/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.
diff --git a/docs/src/examples/cold-plasma-oscillation.py b/docs/src/examples/cold-plasma-oscillation.py
index 49452a1..2176fa4 100644
--- a/docs/src/examples/cold-plasma-oscillation.py
+++ b/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)
diff --git a/docs/src/examples/cold-plasma-wave-packet.py b/docs/src/examples/cold-plasma-wave-packet.py
index 68a8aa3..316e09b 100644
--- a/docs/src/examples/cold-plasma-wave-packet.py
+++ b/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."""
diff --git a/docs/src/examples/cold-plasma-waves.py b/docs/src/examples/cold-plasma-waves.py
index 49ab9a7..a7f42d4 100644
--- a/docs/src/examples/cold-plasma-waves.py
+++ b/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}
ω=%{y:.3f}
log₁₀ P=%{z:.2f}",
- )
- )
- 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)
diff --git a/docs/src/examples/dam-break.py b/docs/src/examples/dam-break.py
index 01e74db..15df174 100644
--- a/docs/src/examples/dam-break.py
+++ b/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(
diff --git a/docs/src/examples/damped-alfven-wave.py b/docs/src/examples/damped-alfven-wave.py
index 1efac3b..b45191b 100644
--- a/docs/src/examples/damped-alfven-wave.py
+++ b/docs/src/examples/damped-alfven-wave.py
@@ -94,16 +94,10 @@ def velocity_profile(run):
return u_y.t.values, u_y.eta1.values * length, u_y.values
-def mode_amplitude(x, values):
- """Amplitude of the sin(k x) mode, from the periodic grid points (the last repeats the first)."""
- return 2.0 * np.mean(values[:, :-1] * np.sin(wavenumber * x[:-1]), axis=1)
-
-
-def envelope_peaks(times, values):
- """The times and heights of the local maxima of |values|, which follow the decaying envelope."""
- magnitude = np.abs(values)
- peaks = np.where((magnitude[1:-1] >= magnitude[:-2]) & (magnitude[1:-1] >= magnitude[2:]))[0] + 1
- return times[peaks], magnitude[peaks]
+def mode_amplitude(run):
+ """Amplitude of the sin(k x) mode of u_y over time, from a post-processed run."""
+ u_y = run.evaluate("mhd/velocity_xyz").isel(component=1, eta2=0, eta3=0)
+ return u_y.struphy.analysis.project_mode(dim="eta1", number=1)
def save(figure, name: str, *, show: bool = False, frame: int | None = None, still=None, width=1100, height=650):
@@ -139,11 +133,11 @@ def pproc(sim: Simulation, show: bool = False):
amplitudes, fitted, measured_frequency = {}, {}, {}
for eta, run in runs.items():
- times, x, values = velocity_profile(run)
- amplitudes[eta] = mode_amplitude(x, values)
- peak_times, peak_values = envelope_peaks(times, amplitudes[eta])
- fitted[eta] = float(-np.polyfit(peak_times, np.log(peak_values), 1)[0])
- crossings = np.where(np.diff(np.sign(amplitudes[eta])) != 0)[0]
+ amplitudes[eta] = mode_amplitude(run)
+ # The local maxima of |amplitude| follow the decaying envelope.
+ fitted[eta] = -abs(amplitudes[eta]).struphy.analysis.damping_rate().rate
+ times = amplitudes[eta].t.values
+ crossings = np.where(np.diff(np.sign(amplitudes[eta].values)) != 0)[0]
measured_frequency[eta] = float(np.pi / np.mean(np.diff(times[crossings])))
exact_rate = {eta: eta * wavenumber**2 / 2 for eta in runs}
if not all(np.isfinite(list(fitted.values()) + list(measured_frequency.values()))):
@@ -156,8 +150,7 @@ def pproc(sim: Simulation, show: bool = False):
run = runs[eta_main]
times, x, values = velocity_profile(run)
gamma = exact_rate[eta_main]
- total = np.asarray(run.scalars["en_tot"])
- energy_drift = float(np.max(np.abs(total / total[0] - 1.0)))
+ energy_drift = float(run.scalars["en_tot"].struphy.analysis.relative_error().max())
print(f"Maximum relative drift of the total energy: {energy_drift:.2e}")
def profile_traces(index):
@@ -173,7 +166,7 @@ def profile_traces(index):
subplot_titles=("Transverse velocity u_y along the field", "Amplitude of the mode"))
for trace in profile_traces(0):
figure.add_trace(trace, row=1, col=1)
- figure.add_scatter(x=times, y=amplitudes[eta_main], mode="lines", name="amplitude", showlegend=False,
+ figure.add_scatter(x=times, y=amplitudes[eta_main].values, mode="lines", name="amplitude", showlegend=False,
line={"color": "#d62828", "width": 2}, row=2, col=1)
for sign in (1, -1):
figure.add_scatter(x=times, y=sign * amplitude * np.exp(-gamma * times), mode="lines", showlegend=False,
@@ -196,8 +189,8 @@ def profile_traces(index):
colors = {0.05: "#168aad", 0.1: "#d62828", 0.2: "#f77f00"}
decay = go.Figure()
for eta, amp in amplitudes.items():
- peak_times, peak_values = envelope_peaks(times, amp)
- decay.add_scatter(x=peak_times, y=peak_values, mode="markers", name=f"η = {eta}: peaks of Struphy",
+ peaks = abs(amp).struphy.analysis.envelope()
+ decay.add_scatter(x=peaks.t.values, y=peaks.values, mode="markers", name=f"η = {eta}: peaks of Struphy",
marker={"color": colors[eta], "size": 9})
decay.add_scatter(x=times, y=amplitude * np.exp(-exact_rate[eta] * times), mode="lines",
name=f"η = {eta}: exp(−η k² t / 2)", line={"color": colors[eta], "width": 1.5, "dash": "dash"})
diff --git a/docs/src/examples/diffusion-methods.py b/docs/src/examples/diffusion-methods.py
index f2877fc..06e5c4a 100644
--- a/docs/src/examples/diffusion-methods.py
+++ b/docs/src/examples/diffusion-methods.py
@@ -138,15 +138,14 @@ def pproc(sim: Simulation, show: bool = False):
profile_exact = 1.0 + exact_amplitude[:, None] * np.cos(wavenumber * x)
# The amplitude of the cosine mode relative to the mean. A bin average lowers a mode by sinc(k dx / 2).
- bin_average = np.sinc(wavenumber * (1.0 / bins) / 2 / np.pi)
-
- def mode_amplitude(values):
- return 2 * np.mean(values * np.cos(wavenumber * x), axis=1) / np.mean(values, axis=1) / bin_average
-
- measured = {name: mode_amplitude(values.values) for name, values in density.items()}
- fitted_rate = {name: float(-np.polyfit(times, np.log(np.abs(amp)), 1)[0]) for name, amp in measured.items()}
+ measured = {
+ name: values.struphy.analysis.project_mode(dim="eta1", number=1, kind="cos", bin_correction=True)
+ / values.mean("eta1")
+ for name, values in density.items()
+ }
+ fitted_rate = {name: -abs(amp).struphy.analysis.growth_rate().rate for name, amp in measured.items()}
rms_error = {
- name: float(np.sqrt(np.mean((values.values - profile_exact) ** 2))) for name, values in density.items()
+ name: float(values.struphy.analysis.error(profile_exact, dims=("t", "eta1"))) for name, values in density.items()
}
if not all(np.isfinite(rate) for rate in fitted_rate.values()):
raise RuntimeError("A decay rate could not be fitted")
@@ -174,7 +173,7 @@ def traces(index):
figure.add_scatter(x=times, y=exact_amplitude, mode="lines", name="exact decay", showlegend=False,
line={"color": "#111", "width": 2, "dash": "dash"}, row=2, col=1)
for name, amp in measured.items():
- figure.add_scatter(x=times, y=np.abs(amp), mode="lines", name=name, showlegend=False,
+ figure.add_scatter(x=times, y=np.abs(amp.values), mode="lines", name=name, showlegend=False,
line={"color": colors[name], "width": 2}, row=2, col=1)
figure.frames = [go.Frame(name=f"{times[i]:.3f}", data=traces(i), traces=[0, 1, 2]) for i in picks]
figure.update_layout(
diff --git a/docs/src/examples/diocotron-instability.py b/docs/src/examples/diocotron-instability.py
index e5b3c78..e534a4f 100644
--- a/docs/src/examples/diocotron-instability.py
+++ b/docs/src/examples/diocotron-instability.py
@@ -143,15 +143,12 @@ def pproc(sim: Simulation, show: bool = False):
# seeded m = 4 mode visible as a measurement rather than just a feature in
# the animation.
theta = np.deg2rad(angle_deg)
- ring = frames_data[:, (radius >= r_minus) & (radius <= r_plus), :]
- mean_density = np.mean(ring, axis=(1, 2))
- mode_amplitudes = {
- m: np.abs(np.mean(ring * np.exp(-1j * m * theta)[None, None, :], axis=(1, 2))) / mean_density
- for m in range(1, 9)
- }
+ ring = density.isel(eta1=(radius >= r_minus) & (radius <= r_plus))
+ spectrum = abs(ring.struphy.analysis.mode_spectrum(dims="eta2", names="m").mean("eta1")) / ring.mean(("eta1", "eta2"))
+ mode_amplitudes = {m: spectrum.sel(m=m) for m in range(1, 9)}
growth_window = (times > 5.0) & (times < 15.0)
- growth_fit = np.polyfit(times[growth_window], np.log(mode_amplitudes[mode_number][growth_window] + 1e-12), 1)
- growth_rate = float(growth_fit[0])
+ growth_fit = mode_amplitudes[mode_number].isel(t=growth_window).struphy.analysis.growth_rate()
+ growth_rate = growth_fit.rate
print(f"Measured m = {mode_number} growth rate: {growth_rate:.4f}")
# Keep the high-resolution movie responsive by using evenly spaced frames
@@ -214,15 +211,15 @@ def pproc(sim: Simulation, show: bool = False):
for m, amplitude in mode_amplitudes.items():
mode_figure.add_trace(go.Scatter(
x=times,
- y=amplitude + 1e-12,
+ y=amplitude.values + 1e-12,
mode="lines",
name=f"m = {m}",
line={"width": 3 if m == mode_number else 1.2, "color": "#168aad" if m == mode_number else None},
opacity=1.0 if m == mode_number else 0.55,
))
mode_figure.add_trace(go.Scatter(
- x=times[growth_window],
- y=np.exp(growth_fit[1] + growth_rate * times[growth_window]),
+ x=growth_fit.time,
+ y=growth_fit.fitted,
mode="lines",
name=f"m = {mode_number} exponential fit",
line={"dash": "dash", "color": "#f08a4b", "width": 2},
diff --git a/docs/src/examples/faraday-rotation.py b/docs/src/examples/faraday-rotation.py
index e3a26e4..c147146 100644
--- a/docs/src/examples/faraday-rotation.py
+++ b/docs/src/examples/faraday-rotation.py
@@ -113,7 +113,7 @@ def pproc(sim: Simulation, show: bool = False):
angle = 0.5 * np.unwrap(np.arctan2(u, q))
rotation_rate = float(np.polyfit(z, angle, 1)[0])
angle_error = float(np.max(np.abs(angle + z / 2)))
- energy_drift = float(np.max(np.abs(energy.values / energy.values[0] - 1.0)))
+ energy_drift = float(energy.struphy.analysis.relative_error().max())
print(f"Measured rotation rate: {rotation_rate:.4f} rad per unit length (exact: -0.5)")
if error > 0.03 or angle_error > 0.02 or energy_drift > 1e-6:
raise RuntimeError(f"Faraday check failed: field={error:.3g}, angle={angle_error:.3g}, energy={energy_drift:.3g}")
diff --git a/docs/src/examples/gas-expansion.py b/docs/src/examples/gas-expansion.py
index 39ea476..e174416 100644
--- a/docs/src/examples/gas-expansion.py
+++ b/docs/src/examples/gas-expansion.py
@@ -155,20 +155,14 @@ def density_error(index):
difference = np.abs(density.isel(t=index).values - exact_density)
return np.trapezoid(difference[window], grid[window]) / np.trapezoid(exact_density[window], grid[window])
- def velocity_error(index):
- _, exact_velocity = exact_solution(marker_position[index], times[index])
- return np.sqrt(np.mean((marker_velocity[index] - exact_velocity) ** 2)) / sound_speed
-
- def velocity_median_error(index):
- _, exact_velocity = exact_solution(marker_position[index], times[index])
- return np.median(np.abs(marker_velocity[index] - exact_velocity)) / sound_speed
-
compared = np.flatnonzero(times >= 0.1)
density_errors = np.array([density_error(i) for i in compared])
- velocity_errors = np.array([velocity_error(i) for i in compared])
- velocity_median_final = float(velocity_median_error(compared[-1]))
- mass = np.array([np.trapezoid(density.isel(t=i).values, grid) for i in range(len(times))])
- mass_error = float(np.max(np.abs(mass / (gas_density * release_point) - 1.0)))
+ exact_velocity = np.array([exact_solution(marker_position[i], times[i])[1] for i in compared]) # (t, marker)
+ velocity_errors = orbits.sel(quantity="v1").isel(t=compared).struphy.analysis.error(exact_velocity, dims="marker")
+ velocity_errors = velocity_errors.values / sound_speed
+ velocity_median_final = float(np.median(np.abs(marker_velocity[compared[-1]] - exact_velocity[-1])) / sound_speed)
+ mass = density.integrate("eta1") * box_length
+ mass_error = float(mass.struphy.analysis.relative_error(ref=gas_density * release_point, skip_first=False).max())
print(
f"Density error (relative L1): mean {density_errors.mean():.4f}, final {density_errors[-1]:.4f}; "
f"velocity error (rms, units of c): mean {velocity_errors.mean():.4f}, final {velocity_errors[-1]:.4f}, "
diff --git a/docs/src/examples/grad-b-drift.py b/docs/src/examples/grad-b-drift.py
index 375dede..b2accfa 100644
--- a/docs/src/examples/grad-b-drift.py
+++ b/docs/src/examples/grad-b-drift.py
@@ -105,6 +105,7 @@ def pproc(sim: Simulation, show: bool = False):
time_opts = sim.time_opts
import plotly.graph_objects as go
from plotly.subplots import make_subplots
+ from struphy_plots.theory.orbits import grad_b_drift
output = sim.output
# Post-processing otherwise tries to reconstruct this script-local class
@@ -128,8 +129,9 @@ def pproc(sim: Simulation, show: bool = False):
# approximate guiding-center coordinate, not a separate guiding-center run.
center_y = y - vx / local_b
measured = np.polyfit(times, center_y, 1)[0]
- gradient = -ripple * 2 * np.pi / box # at x_gc = box/2, B = 1
- reference = 0.5 * speed[0]**2 * gradient
+ # At x_gc = box/2: B = e_z and grad B = -ripple * 2 pi / box e_x; the drift is along y.
+ field_at_center, gradient = (0.0, 0.0, 1.0), (-ripple * 2 * np.pi / box, 0.0, 0.0)
+ reference = grad_b_drift(speed[0], field_at_center, gradient)[:, 1]
drift_error = float(np.max(np.abs(measured / reference - 1.0)))
print(f"Measured drift velocities: {np.round(measured, 5).tolist()} (guiding center: {np.round(reference, 5).tolist()})")
if speed_drift > 1e-8 or drift_error > 0.08:
@@ -144,7 +146,8 @@ def pproc(sim: Simulation, show: bool = False):
figure.add_scatter(x=[speed[0, j]**2], y=[measured[j]], mode="markers", showlegend=False,
marker={"color": color, "size": 11}, row=1, col=2)
square_speed = np.linspace(0, 1, 100)
- figure.add_scatter(x=square_speed, y=0.5 * square_speed * gradient, name="Guiding-center prediction",
+ figure.add_scatter(x=square_speed, y=grad_b_drift(np.sqrt(square_speed), field_at_center, gradient)[:, 1],
+ name="Guiding-center prediction",
line={"color": "#222", "dash": "dash"}, row=1, col=2)
figure.update_xaxes(title_text="x", row=1, col=1)
figure.update_yaxes(title_text="y", scaleanchor="x", scaleratio=1, row=1, col=1)
diff --git a/docs/src/examples/hall-mhd-waves.py b/docs/src/examples/hall-mhd-waves.py
index 1f60bdb..8c317dd 100644
--- a/docs/src/examples/hall-mhd-waves.py
+++ b/docs/src/examples/hall-mhd-waves.py
@@ -31,10 +31,11 @@
def hall_branches(k):
- """Parallel Hall-MHD branches: omega = sqrt(v_A^2 k^2 + h^2) +- h with h = v_A^2 k^2 / (2 Omega_i), and sound."""
- h = alfven_speed**2 * k**2 / (2.0 * cyclotron_frequency)
- root = np.sqrt(alfven_speed**2 * k**2 + h**2)
- return {"whistler": root + h, "ion-cyclotron": root - h, "sound": sound_speed * k}
+ """Parallel Hall-MHD branches: omega = v_A k (sqrt(1 + k^2 d_i^2 / 4) +- k d_i / 2) with d_i = v_A / Omega_i, and sound."""
+ from struphy_plots.theory.waves import hall_mhd_parallel
+
+ hall = hall_mhd_parallel(k, alfven_speed=alfven_speed, ion_inertial_length=alfven_speed / cyclotron_frequency)
+ return {"whistler": hall["whistler"].real, "ion-cyclotron": hall["ion cyclotron"].real, "sound": sound_speed * k}
def create_simulation() -> Simulation:
diff --git a/docs/src/examples/hasegawa-wakatani.py b/docs/src/examples/hasegawa-wakatani.py
index 1410b9a..f09c7c4 100644
--- a/docs/src/examples/hasegawa-wakatani.py
+++ b/docs/src/examples/hasegawa-wakatani.py
@@ -218,7 +218,9 @@ def field_traces(index, *, colorbars=False):
# up to roundoff on the sampled grid.
# The sampled periodic grid includes the repeated right/top boundary.
# Remove it before the FFT so it is not counted as a second grid point.
- phi = potential.transpose("t", "eta1", "eta2").values[:, :-1, :-1]
+ from struphy_plots.spectral import drop_periodic_endpoint
+
+ phi = drop_periodic_endpoint(drop_periodic_endpoint(potential, "eta1"), "eta2").transpose("t", "eta1", "eta2").values
nx, ny = phi.shape[1:]
kx = 2 * np.pi * np.fft.fftfreq(nx, d=length / nx)[:, None]
ky = 2 * np.pi * np.fft.fftfreq(ny, d=length / ny)[None, :]
diff --git a/docs/src/examples/hybrid-alfven-ion-coupling.py b/docs/src/examples/hybrid-alfven-ion-coupling.py
index 3a35b69..a540808 100644
--- a/docs/src/examples/hybrid-alfven-ion-coupling.py
+++ b/docs/src/examples/hybrid-alfven-ion-coupling.py
@@ -125,23 +125,24 @@ def pproc(sim: Simulation, show: bool = False):
# LinearMHDVlasovPC tracks each subsystem's energy as a scalar every step:
# en_B/en_U (field/fluid), en_f (kinetic energetic ions), en_tot (total).
+ scalars = output.scalars
time = np.asarray(output.time)
- en_B = np.asarray(output.scalars["en_B"])
- en_U = np.asarray(output.scalars["en_U"])
- en_f = np.asarray(output.scalars["en_f"])
- en_tot = np.asarray(output.scalars["en_tot"])
+ en_B = np.asarray(scalars["en_B"])
+ en_U = np.asarray(scalars["en_U"])
+ en_f_change = np.asarray(scalars["en_f"].struphy.analysis.drift())
+ total_drift = np.asarray(scalars["en_tot"].struphy.analysis.drift() / scalars["en_tot"].isel(t=0))
- relative_drift = float(np.max(np.abs(en_tot - en_tot[0]) / en_tot[0]))
+ relative_drift = float(scalars["en_tot"].struphy.analysis.relative_error().max())
print(f"Max relative drift in total energy (should be ~0): {relative_drift:.2e}")
figure = go.Figure(
data=[
go.Scatter(x=time, y=en_B, mode="lines", name="Field energy (en_B)", line={"color": "#168aad", "width": 2.5}),
go.Scatter(x=time, y=en_U, mode="lines", name="Fluid kinetic energy (en_U)", line={"color": "#f77f00", "width": 2.5}),
- go.Scatter(x=time, y=en_f - en_f[0], mode="lines", name="Energetic-ion energy change (en_f − en_f₀)", line={"color": "#d62828", "width": 2.5}),
+ go.Scatter(x=time, y=en_f_change, mode="lines", name="Energetic-ion energy change (en_f − en_f₀)", line={"color": "#d62828", "width": 2.5}),
go.Scatter(
x=time,
- y=(en_tot - en_tot[0]) / en_tot[0],
+ y=total_drift,
mode="lines",
name="Total energy drift (relative)",
line={"color": "#6a4c93", "width": 2, "dash": "dot"},
diff --git a/docs/src/examples/hybrid-current-coupling.py b/docs/src/examples/hybrid-current-coupling.py
index 231f07c..bd253cf 100644
--- a/docs/src/examples/hybrid-current-coupling.py
+++ b/docs/src/examples/hybrid-current-coupling.py
@@ -150,7 +150,7 @@ def pproc(sim: Simulation, show: bool = False):
wave_energy = float(values["en_U"][0] + values["en_B"][0])
if wave_energy <= 0:
raise RuntimeError("The initial wave energy must be positive")
- total_error = values["en_tot"] - values["en_tot"][0]
+ total_error = energies["en_tot"].struphy.analysis.drift().values
wave_scaled_error = float(np.max(np.abs(total_error)) / wave_energy)
# Require conservation error below 0.1% of the seeded wave energy.
# Plot the measured error explicitly; the run does not conserve to roundoff.
@@ -168,9 +168,8 @@ def pproc(sim: Simulation, show: bool = False):
("en_f", "Ion energy change", "#d62828"),
("en_p", "Pressure energy change", "#6a4c93"),
):
- energy = values[key]
- if key in ("en_f", "en_p"):
- energy = energy - energy[0]
+ # the ion and pressure energies relative to their initial values
+ energy = energies[key].struphy.analysis.drift().values if key in ("en_f", "en_p") else values[key]
figure.add_scatter(
x=times, y=energy / wave_energy, mode="lines", name=label,
line={"color": color, "width": 2.5}, row=1, col=1,
diff --git a/docs/src/examples/incompressible-shear-relaxation.py b/docs/src/examples/incompressible-shear-relaxation.py
index fe110d2..1b987a4 100644
--- a/docs/src/examples/incompressible-shear-relaxation.py
+++ b/docs/src/examples/incompressible-shear-relaxation.py
@@ -140,13 +140,15 @@ def pproc(sim: Simulation, show: bool = False):
raise RuntimeError("Non-finite SPH result: refusing to publish the run")
# The amplitude of each mode, projected on it. A bin average lowers a mode by sinc(k dx / 2).
- shear_bins = np.sinc(np.pi / height * (height / bins) / 2 / np.pi)
- wave_bins = np.sinc(2 * np.pi * (1.0 / bins) / 2 / np.pi)
- shear = 2 * np.mean(across.values * np.sin(np.pi * y / height), axis=1) / shear_bins
- wave = 2 * np.mean(along.values * np.sin(2 * np.pi * x), axis=1) / wave_bins
+ # The shear mode sin(pi y / H) is half a wave along eta2 = y / H.
+ shear = across.struphy.analysis.project_mode(dim="eta2", number=0.5, bin_correction=True)
+ wave = along.struphy.analysis.project_mode(dim="eta1", number=1, bin_correction=True).values
exact_shear = shear_amplitude * np.exp(-decay_rate * times)
- fitted_rate = float(-np.polyfit(times, np.log(np.abs(shear)), 1)[0])
- profile_error = float(np.max(np.abs(across.values[-1] - exact_shear[-1] * np.sin(np.pi * y / height))))
+ fitted_rate = -abs(shear).struphy.analysis.growth_rate().rate
+ profile_error = float(
+ across.isel(t=-1).struphy.analysis.error(exact_shear[-1] * np.sin(np.pi * y / height), norm="max")
+ )
+ shear = shear.values
wave_left = float(np.abs(wave[times >= 0.1]).max() / compressive_amplitude)
print(f"shear decay rate {fitted_rate:.3f} (exact {decay_rate:.3f}); compressive wave after t = 0.1: "
f"{100 * wave_left:.1f}% of its initial amplitude; profile error at the end {profile_error:.4f}")
diff --git a/docs/src/examples/itg-drift-wave.py b/docs/src/examples/itg-drift-wave.py
index 6c8ee9a..1960c91 100644
--- a/docs/src/examples/itg-drift-wave.py
+++ b/docs/src/examples/itg-drift-wave.py
@@ -119,7 +119,8 @@ def mode_amplitudes(phi):
m >= 0 is the poloidal and n the axial mode number. The periodic end points e2 = 1 and e3 = 1 repeat the
first ones and are dropped. A field `A cos(m*theta + 2*pi*n*z/length)` has |phi_mn| = A.
"""
- field = phi.isel(eta2=slice(None, -1), eta3=slice(None, -1)).transpose("t", "eta1", "eta2", "eta3")
+ field = phi.struphy.analysis.drop_periodic_endpoint("eta2").struphy.analysis.drop_periodic_endpoint("eta3")
+ field = field.transpose("t", "eta1", "eta2", "eta3")
n_theta, n_z = field.sizes["eta2"], field.sizes["eta3"]
spectrum = np.fft.fft(np.fft.rfft(field.values, axis=2), axis=3) / (n_theta * n_z)
spectrum[:, :, 1:] *= 2.0 # fold the negative m
@@ -137,18 +138,6 @@ def radial_rms(spectrum):
return np.sqrt((np.abs(spectrum) ** 2).mean("eta1"))
-def exponential_rate(amplitude, window):
- """Growth rate of each column m of `amplitude(t, m)` from a log-linear fit on `window`."""
- selected = amplitude.sel(t=slice(*window))
- times = selected.t.values
- rates = []
- for m in selected.m.values:
- values = selected.sel(m=m).values
- positive = values > 0
- rates.append(np.polyfit(times[positive], np.log(values[positive]), 1)[0] if positive.sum() > 2 else np.nan)
- return np.asarray(rates)
-
-
def log10_or_nan(values):
values = np.asarray(values, dtype=float)
return np.log10(np.where(values > 0, values, np.nan))
@@ -243,26 +232,13 @@ def pproc(sim: Simulation, show: bool = False):
rho = output.fields.diagnostics.rho
rho = rho.isel(component=0) if "component" in rho.dims else rho
times = np.asarray(rho.t)
- perturbation_energy = np.asarray(rho.struphy.analysis.norm(squared=True))
+ perturbation_energy = rho.struphy.analysis.norm(squared=True).assign_attrs(label="‖δn‖²")
growth_window = (times > GROWTH_WINDOW[0]) & (times < GROWTH_WINDOW[1])
- growth_rate = float(np.polyfit(times[growth_window], np.log(perturbation_energy[growth_window]), 1)[0] / 2)
+ growth_rate = float(np.polyfit(times[growth_window], np.log(perturbation_energy.values[growth_window]), 1)[0] / 2)
print(f"Measured growth rate: {growth_rate:.5f}")
- figure = go.Figure(
- data=[
- go.Scatter(x=times, y=perturbation_energy, mode="lines", name="Struphy (drift-kinetic)", line={"color": "#168aad", "width": 3}),
- ],
- )
- figure.update_layout(
- title="ITG drift wave: density perturbation energy",
- xaxis_title="t [a.u.]",
- yaxis_title="‖δn‖² [a.u.]",
- yaxis={"type": "log"},
- template="plotly_white",
- autosize=True,
- margin={"l": 70, "r": 30, "t": 80, "b": 60},
- )
+ figure = perturbation_energy.struphy.plot.timeseries(logy=True, title="ITG drift wave: density perturbation energy", backend="plotly")
save(figure, "itg-drift-wave", show=show)
@@ -296,7 +272,7 @@ def pproc(sim: Simulation, show: bool = False):
# Amplitude of each poloidal mode (radial rms at the seeded axial mode number) and its growth rate.
amplitude = radial_rms(spectrum.sel(n=mode_toroidal)).sel(m=slice(1, MAX_POLOIDAL_MODE))
- rates = exponential_rate(amplitude, GROWTH_WINDOW)
+ rates = [amplitude.sel(m=m).struphy.analysis.growth_rate(window=GROWTH_WINDOW).rate for m in amplitude.m.values]
wavenumber = amplitude.m.values / float(radius.mean())
colors = sample_colorscale("Viridis", np.linspace(0, 1, amplitude.sizes["m"]))
growth_figure = make_subplots(
@@ -402,7 +378,8 @@ def pproc(sim: Simulation, show: bool = False):
# The flux-surface-averaged (m = n = 0) density change: does the profile flatten?
density = output.evaluate("diagnostics/rho", **points)
density = density.isel(component=0, drop=True) if "component" in density.dims else density
- zonal = density.isel(eta2=slice(None, -1), eta3=slice(None, -1)).mean(("eta2", "eta3"))
+ density = density.struphy.analysis.drop_periodic_endpoint("eta2").struphy.analysis.drop_periodic_endpoint("eta3")
+ zonal = density.mean(("eta2", "eta3"))
zonal = zonal - zonal.isel(t=0)
profile_change = zonal.assign_coords(eta1=radius).struphy.plot.slice(
x="eta1",
diff --git a/docs/src/examples/langmuir-wave-dispersion.py b/docs/src/examples/langmuir-wave-dispersion.py
index 6196cc0..bdb00dd 100644
--- a/docs/src/examples/langmuir-wave-dispersion.py
+++ b/docs/src/examples/langmuir-wave-dispersion.py
@@ -19,8 +19,6 @@
import numpy as np
import plotly.graph_objects as go
-from scipy.optimize import fsolve
-from scipy.special import wofz
from struphy import (
BoundaryParameters,
@@ -43,17 +41,6 @@
amplitude = 0.001
-def kinetic_frequency(k, guess=(1.4, -0.15)):
- """The complex Langmuir frequency, as (omega_r, gamma), from the root of the Vlasov dispersion relation."""
-
- def dispersion(values):
- zeta = (values[0] + 1j * values[1]) / (np.sqrt(2) * k)
- residual = 1 + (1 + zeta * 1j * np.sqrt(np.pi) * wofz(zeta)) / k**2
- return [residual.real, residual.imag]
-
- return fsolve(dispersion, guess, xtol=1e-12)
-
-
def create_simulation(k=0.5, folder="langmuir_wave_dispersion") -> Simulation:
"""Vlasov-Ampère in a periodic box of length 2 pi / k with one cosine mode in the density."""
grid = grids.TensorProductGrid(num_elements=(32, 1, 1))
@@ -94,13 +81,12 @@ def create_simulation(k=0.5, folder="langmuir_wave_dispersion") -> Simulation:
def mode_amplitude(run):
- """Time and the amplitude of the sin(k x) mode of the electric field E_x, from a post-processed run."""
+ """The amplitude of the sin(k x) mode of the electric field E_x over time, from a post-processed run."""
cells = run.grid.num_elements[0]
e_x = run.evaluate(
"em_fields/e_field", eta1=np.linspace(0.0, 1.0, cells + 1), eta2=0.0, eta3=0.0, representation="1"
).isel(component=0)
- e1 = e_x.eta1.values
- return e_x.t.values, 2.0 * np.mean(e_x.values[:, :-1] * np.sin(2 * np.pi * e1[:-1]), axis=1)
+ return e_x.struphy.analysis.project_mode(dim="eta1", number=1)
def save(figure, name: str, *, show: bool = False, frame: int | None = None, still=None, width=1100, height=650):
@@ -132,12 +118,15 @@ def pproc(sim: Simulation, show: bool = False):
for run in runs.values():
run.pproc()
+ from struphy_plots.theory.kinetic import bohm_gross, langmuir
+
exact, measured_frequency, measured_rate, amplitudes = {}, {}, {}, {}
- guess = (1.4, -0.15)
for k in sorted(runs):
- exact[k] = kinetic_frequency(k, guess)
- guess = tuple(exact[k])
- times, values = mode_amplitude(runs[k])
+ # The complex Langmuir frequency omega_r + i gamma, the root of the Vlasov dispersion relation.
+ omega = langmuir(k)
+ exact[k] = (omega.real, omega.imag)
+ mode = mode_amplitude(runs[k])
+ times, values = mode.t.values, mode.values
amplitudes[k] = values
# The first time units hold a transient of phase-mixing modes that are not the Langmuir wave, so the fit starts at t = 1
# and stops at t = 12. The measured values depend on this window at the level of a few per cent.
@@ -148,9 +137,7 @@ def pproc(sim: Simulation, show: bool = False):
values[crossings + 1] - values[crossings]
)
measured_frequency[k] = float(np.pi / np.mean(np.diff(roots)))
- magnitude = np.abs(values)
- peaks = np.where(window[1:-1] & (magnitude[1:-1] >= magnitude[:-2]) & (magnitude[1:-1] >= magnitude[2:]))[0] + 1
- measured_rate[k] = float(np.polyfit(times[peaks], np.log(magnitude[peaks]), 1)[0])
+ measured_rate[k] = abs(mode).struphy.analysis.damping_rate(window=(1.0, 12.0)).rate
if not (np.isfinite(list(measured_frequency.values())).all() and np.isfinite(list(measured_rate.values())).all()):
raise RuntimeError("A frequency or a damping rate could not be measured")
for k in sorted(runs):
@@ -160,21 +147,17 @@ def pproc(sim: Simulation, show: bool = False):
from plotly.subplots import make_subplots
k_line = np.linspace(0.25, 0.65, 60)
- lines, guess = [], (1.15, -0.01)
- for kk in k_line:
- guess = tuple(kinetic_frequency(kk, guess))
- lines.append(guess)
- lines = np.array(lines)
+ lines = langmuir(k_line)
ks = sorted(runs)
figure = make_subplots(rows=1, cols=2, horizontal_spacing=0.12,
subplot_titles=("Oscillation frequency", "Damping rate"))
- figure.add_scatter(x=k_line, y=np.sqrt(1 + 3 * k_line**2), mode="lines", name="Bohm–Gross √(1 + 3k²)",
+ figure.add_scatter(x=k_line, y=bohm_gross(k_line).real, mode="lines", name="Bohm–Gross √(1 + 3k²)",
line={"color": "#888", "width": 1.5, "dash": "dot"}, row=1, col=1)
- figure.add_scatter(x=k_line, y=lines[:, 0], mode="lines", name="kinetic dispersion relation",
+ figure.add_scatter(x=k_line, y=lines.real, mode="lines", name="kinetic dispersion relation",
line={"color": "#111", "width": 2, "dash": "dash"}, row=1, col=1)
figure.add_scatter(x=ks, y=[measured_frequency[k] for k in ks], mode="markers", name="Struphy (PIC)",
marker={"color": "#d62828", "size": 11}, row=1, col=1)
- figure.add_scatter(x=k_line, y=lines[:, 1], mode="lines", showlegend=False,
+ figure.add_scatter(x=k_line, y=lines.imag, mode="lines", showlegend=False,
line={"color": "#111", "width": 2, "dash": "dash"}, row=1, col=2)
figure.add_scatter(x=ks, y=[measured_rate[k] for k in ks], mode="markers", showlegend=False,
marker={"color": "#d62828", "size": 11}, row=1, col=2)
@@ -190,7 +173,7 @@ def pproc(sim: Simulation, show: bool = False):
colors = {0.3: "#168aad", 0.4: "#2a9d8f", 0.5: "#f77f00", 0.6: "#d62828"}
signals = go.Figure()
for k in ks:
- times, _ = mode_amplitude(runs[k])
+ times = mode_amplitude(runs[k]).t.values
# The exact envelope exp(gamma t), scaled to the first peak of the wave (after the initial transient).
after = times > 1.0
first_peak = np.argmax(np.abs(amplitudes[k][after][:40]))
diff --git a/docs/src/examples/linear-dissipative-alfven-wave.py b/docs/src/examples/linear-dissipative-alfven-wave.py
index 65c804d..5b4d0cb 100644
--- a/docs/src/examples/linear-dissipative-alfven-wave.py
+++ b/docs/src/examples/linear-dissipative-alfven-wave.py
@@ -82,6 +82,7 @@ def save(figure, name: str, *, show: bool = False, frame: int | None = None, sti
def pproc(sim: Simulation, show: bool = False):
time_opts = sim.time_opts
from plotly.subplots import make_subplots
+ from struphy_plots.theory.waves import dissipative_alfven
output = sim.output
output.pproc(physical=True)
@@ -94,25 +95,30 @@ def pproc(sim: Simulation, show: bool = False):
raise RuntimeError("The linear MHD run produced non-finite diagnostics")
if abs(times[-1] - time_opts.Tend) > time_opts.dt:
raise RuntimeError("The linear MHD run did not reach the requested end time")
- exact_u = amplitude * np.exp(-diffusivity * times[:, None]) * np.cos(times[:, None]) * np.sin(z[None, :])
- exact_b = amplitude * np.exp(-diffusivity * times[:, None]) * np.sin(times[:, None]) * np.cos(z[None, :])
- error = float(max(np.max(np.abs(velocity.values - exact_u)), np.max(np.abs(magnetic.values - exact_b))) / amplitude)
+ # The k = 1 visco-resistive Alfvén wave: omega = 1 - 0.1i for v_A = 1 and nu = eta = 0.1.
+ omega = dissipative_alfven(1.0, resistivity=diffusivity, viscosity=diffusivity)["forward"]
+ frequency, decay = float(omega.real), -float(omega.imag)
+ envelope = amplitude * np.exp(-decay * times[:, None])
+ exact_u = envelope * np.cos(frequency * times[:, None]) * np.sin(z[None, :])
+ exact_b = envelope * np.sin(frequency * times[:, None]) * np.cos(z[None, :])
+ error = max(float(field.struphy.analysis.error(exact, norm="max", dims=("t", "eta3")))
+ for field, exact in ((velocity, exact_u), (magnetic, exact_b))) / amplitude
energy_time = kinetic.t.values
wave_energy = kinetic.values + magnetic_energy.values
- exact_energy = np.exp(-2.0 * diffusivity * energy_time)
+ exact_energy = np.exp(-2.0 * decay * energy_time)
energy_error = float(np.max(np.abs(wave_energy / wave_energy[0] - exact_energy)))
print(f"Maximum relative error against the exact solution: field {error:.2e}, energy {energy_error:.2e}")
if max(error, energy_error) > 0.03:
raise RuntimeError(f"The dissipative Alfvén wave differs from its exact solution: field={error:.3g}, energy={energy_error:.3g}")
- # The periodic evaluation grid repeats its endpoint; exclude it from mode projections.
- u_mode = 2 * np.mean(velocity.values[:, :-1] * np.sin(z[:-1]), axis=1) / amplitude
- b_mode = 2 * np.mean(magnetic.values[:, :-1] * np.cos(z[:-1]), axis=1) / amplitude
+ # sin(z) and cos(z) are mode 1 along eta3; project_mode drops the repeated periodic endpoint.
+ u_mode = velocity.struphy.analysis.project_mode(dim="eta3", number=1, kind="sin").values / amplitude
+ b_mode = magnetic.struphy.analysis.project_mode(dim="eta3", number=1, kind="cos").values / amplitude
figure = make_subplots(rows=2, cols=1, vertical_spacing=0.18,
subplot_titles=("Damped velocity and magnetic modes", "Quadratic wave energy"))
for values, label, color in ((u_mode, "Velocity mode", "#168aad"), (b_mode, "Magnetic mode", "#d62828")):
figure.add_scatter(x=times, y=values, name=label, line={"color": color}, row=1, col=1)
for sign in (-1, 1):
- figure.add_scatter(x=times, y=sign * np.exp(-diffusivity * times), name="Exact envelope",
+ figure.add_scatter(x=times, y=sign * np.exp(-decay * times), name="Exact envelope",
showlegend=sign == 1, line={"color": "#222", "dash": "dot"}, row=1, col=1)
for values, label, color in ((kinetic.values, "Kinetic", "#168aad"),
(magnetic_energy.values, "Magnetic", "#d62828"), (wave_energy, "Wave total", "#222")):
diff --git a/docs/src/examples/maxwell-cavity-resonances.py b/docs/src/examples/maxwell-cavity-resonances.py
index f886d3c..772f2aa 100644
--- a/docs/src/examples/maxwell-cavity-resonances.py
+++ b/docs/src/examples/maxwell-cavity-resonances.py
@@ -18,6 +18,7 @@
import numpy as np
import plotly.graph_objects as go
+import xarray as xr
from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, grids, perturbations
from struphy.models import Maxwell
@@ -123,25 +124,13 @@ def pproc(sim: Simulation, show: bool = False):
print(f"modes {modes}: omega = {found:.4f} (exact {omega_exact:.4f})")
print(f"Maximum relative frequency error: {float(np.max(np.abs(errors))):.3e}")
- figure = go.Figure()
- figure.add_scatter(x=omega, y=power, mode="lines", name="power spectrum of E_z",
- line={"color": "#168aad", "width": 2})
- for index, (omega_exact, modes) in enumerate(exact):
- # Annotation heights on a log axis are log10 of the value.
- label = ", ".join(f"({l}, {m})" for l, m in modes[:2]) # noqa: E741
- figure.add_vline(x=omega_exact, line={"color": "#d62828", "width": 1.5, "dash": "dash"})
- figure.add_annotation(x=omega_exact, y=0.25 if index % 2 == 0 else -0.15, yref="y", text=label, showarrow=False,
- xanchor="left", xshift=3, font={"size": 11, "color": "#d62828"})
- figure.add_scatter(x=[None], y=[None], mode="lines", name="exact ω = c|k|, modes (l, m)",
- line={"color": "#d62828", "width": 1.5, "dash": "dash"})
- figure.update_layout(
- title="Resonances of a rectangular box", template="plotly_white", autosize=True,
- xaxis_title="ω [a.u.]", yaxis_title="power, normalized (log)", yaxis_type="log",
- legend={"x": 0.02, "xanchor": "left", "y": 0.6, "bgcolor": "rgba(255,255,255,0.82)"},
- margin={"l": 75, "r": 30, "t": 80, "b": 60},
+ # The spectrum with the exact resonances as dotted lines, labeled by their modes (l, m).
+ spectrum = xr.DataArray(power, dims="omega", coords={"omega": omega}, name="power")
+ exact_lines = {"exact, (l, m) = " + ", ".join(f"({l}, {m})" for l, m in modes[:2]): omega_exact # noqa: E741
+ for omega_exact, modes in exact}
+ figure = spectrum.struphy.plot.power_spectrum(
+ frequencies=exact_lines, omega_max=14.0, title="Resonances of a rectangular box", backend="plotly",
)
- figure.update_xaxes(range=[0.0, 14.0])
- figure.update_yaxes(range=[-8, 0.3])
save(figure, "maxwell-cavity-resonances", show=show)
error_figure = go.Figure(go.Scatter(
diff --git a/docs/src/examples/maxwell-curved-mesh.py b/docs/src/examples/maxwell-curved-mesh.py
index 7fc4014..d02ee0b 100644
--- a/docs/src/examples/maxwell-curved-mesh.py
+++ b/docs/src/examples/maxwell-curved-mesh.py
@@ -103,17 +103,19 @@ def pproc(sim: Simulation, show: bool = False):
output = sim.output
output.pproc(physical=True)
- e_z = output.evaluate("em_fields/e_field_xyz").isel(component=2, eta3=0) # (t, e1, e2), on the mesh points
+ # (t, e1, e2), on the mesh points
+ e_z = output.evaluate("em_fields/e_field_xyz").isel(component=2, eta3=0).transpose("t", "eta1", "eta2")
times = e_z.t.values
mesh_x, mesh_y = e_z.X.values, e_z.Y.values
- numeric = e_z.transpose("t", "eta1", "eta2").values
+ numeric = e_z.values
exact = np.array([exact_field(mesh_x, mesh_y, t) for t in times])
scale = float(np.abs(exact[0]).max())
- error = np.sqrt(np.mean((numeric - exact) ** 2, axis=(1, 2))) / np.sqrt(np.mean(exact[0] ** 2))
+ # The rms error over the mesh points at each time, relative to the rms of the initial pulse.
+ error = (e_z.struphy.analysis.error(exact) / np.sqrt(np.mean(exact[0] ** 2))).values
if not np.isfinite(error).all():
raise RuntimeError("Non-finite field")
total = output.scalars["total_energy"].values
- energy_drift = float(np.max(np.abs(total / total[0] - 1.0)))
+ energy_drift = float(output.scalars["total_energy"].struphy.analysis.relative_error().max())
print(f"Largest rms error of E_z, relative to the initial rms: {error.max():.3e}")
print(f"Maximum relative drift of the total energy: {energy_drift:.2e}")
diff --git a/docs/src/examples/mhd-slab-waves.py b/docs/src/examples/mhd-slab-waves.py
index 27715cf..d824e3b 100644
--- a/docs/src/examples/mhd-slab-waves.py
+++ b/docs/src/examples/mhd-slab-waves.py
@@ -19,6 +19,7 @@
from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, perturbations
from struphy.models import LinearMHD
+from struphy_plots.theory.waves import magnetosonic_speeds
# The background: B0 = (0, 1, 1), density 0.7 and a plasma beta of 3 (thermal over magnetic pressure).
B0x, B0y, B0z = 0.0, 1.0, 1.0
@@ -26,15 +27,12 @@
B_squared = B0x**2 + B0y**2 + B0z**2
p0 = beta * B_squared / 2.0
-# Ideal-MHD wave speeds along z: the shear Alfvén wave, and the slow and fast magnetosonic waves.
+# Ideal-MHD wave speeds along z, at the angle theta between z and B0: the shear Alfvén wave, and the slow
+# and fast magnetosonic waves.
alfven_speed = np.sqrt(B_squared / n0)
sound_speed = np.sqrt(gamma * p0 / n0)
-delta = 4 * B0z**2 * sound_speed**2 * alfven_speed**2 / ((sound_speed**2 + alfven_speed**2) ** 2 * B_squared)
-exact_speeds = {
- "alfven": alfven_speed * B0z / np.sqrt(B_squared),
- "slow": np.sqrt(0.5 * (sound_speed**2 + alfven_speed**2) * (1.0 - np.sqrt(1.0 - delta))),
- "fast": np.sqrt(0.5 * (sound_speed**2 + alfven_speed**2) * (1.0 + np.sqrt(1.0 - delta))),
-}
+speeds = magnetosonic_speeds(np.arccos(B0z / np.sqrt(B_squared)), alfven_speed=alfven_speed, sound_speed=sound_speed)
+exact_speeds = {"alfven": float(speeds["shear Alfvén"]), "slow": float(speeds["slow"]), "fast": float(speeds["fast"])}
def create_simulation() -> Simulation:
diff --git a/docs/src/examples/ordinary-mode-dispersion.py b/docs/src/examples/ordinary-mode-dispersion.py
index ab3206f..c86f290 100644
--- a/docs/src/examples/ordinary-mode-dispersion.py
+++ b/docs/src/examples/ordinary-mode-dispersion.py
@@ -19,13 +19,14 @@
from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, perturbations
from struphy.models import ColdPlasma
from struphy.linear_algebra.solver import SolverParameters
+from struphy_plots.theory.waves import plasma_light_wave
stem = "ordinary-mode-dispersion"
length = 8.0 * np.pi
mode_numbers = (1, 2, 4, 8)
amplitude = 0.02 # electric-field amplitude per mode
wavenumbers = 2.0 * np.pi * np.array(mode_numbers) / length
-frequencies = np.sqrt(1.0 + wavenumbers**2) # c = omega_p = 1
+frequencies = plasma_light_wave(wavenumbers).real # omega**2 = omega_p**2 + c**2 k**2, c = omega_p = 1
def create_simulation() -> Simulation:
@@ -90,9 +91,10 @@ def pproc(sim: Simulation, show: bool = False):
times, x = field.t.values, field.eta1.values * length
if not np.isfinite(field.values).all() or not np.isfinite(energy.values).all() or not np.isclose(times[-1], time_opts.Tend):
raise RuntimeError("The ordinary-mode run is incomplete or non-finite")
- # Drop the repeated endpoint of the periodic spatial evaluation grid.
- basis = np.cos(wavenumbers[:, None] * x[None, :-1])
- signals = 2.0 * field.values[:, :-1] @ basis.T / (len(x) - 1)
+ # cos(k x) is mode n along eta1; project_mode drops the repeated periodic endpoint.
+ signals = np.column_stack([
+ field.struphy.analysis.project_mode(dim="eta1", number=n, kind="cos").values for n in mode_numbers
+ ])
measured = []
for signal in signals.T:
crossing = np.flatnonzero(np.diff(np.signbit(signal)))
@@ -104,7 +106,7 @@ def pproc(sim: Simulation, show: bool = False):
reference = amplitude * np.cos(times[:, None] * frequencies[None, :])
frequency_error = float(np.max(np.abs(measured / frequencies - 1.0)))
mode_error = float(np.max(np.abs(signals - reference)) / amplitude)
- energy_drift = float(np.max(np.abs(energy.values / energy.values[0] - 1.0)))
+ energy_drift = float(energy.struphy.analysis.relative_error().max())
print(f"Measured frequencies: {np.round(measured, 4).tolist()} (exact: {np.round(frequencies, 4).tolist()})")
if frequency_error > 0.01 or mode_error > 0.1 or energy_drift > 1e-6:
raise RuntimeError(f"Ordinary-mode check failed: frequency={frequency_error:.3g}, field={mode_error:.3g}, energy={energy_drift:.3g}")
@@ -112,7 +114,7 @@ def pproc(sim: Simulation, show: bool = False):
figure = make_subplots(rows=1, cols=2, horizontal_spacing=0.13,
subplot_titles=("Dispersion and cutoff", "Four independently measured modes"))
k = np.linspace(0, 2.2, 200)
- figure.add_scatter(x=k, y=np.sqrt(1 + k**2), name="ω² = 1 + k²", line={"color": "#222"}, row=1, col=1)
+ figure.add_scatter(x=k, y=plasma_light_wave(k).real, name="ω² = 1 + k²", line={"color": "#222"}, row=1, col=1)
figure.add_scatter(x=wavenumbers, y=measured, mode="markers", name="Measured frequencies",
marker={"size": 10, "color": "#d62828"}, row=1, col=1)
figure.add_hrect(y0=0, y1=1, fillcolor="#168aad", opacity=0.08, line_width=0, row=1, col=1)
diff --git a/docs/src/examples/orszag-tang-vortex.py b/docs/src/examples/orszag-tang-vortex.py
index defe31a..f0e1975 100644
--- a/docs/src/examples/orszag-tang-vortex.py
+++ b/docs/src/examples/orszag-tang-vortex.py
@@ -104,9 +104,8 @@ def pproc(sim: Simulation, show: bool = False):
x, y = np.asarray(b.X)[:, 0], np.asarray(b.Y)[0, :]
bx, by = (b.isel(component=i).transpose("t", "eta1", "eta2").values for i in (0, 1))
rho_values = rho.transpose("t", "eta1", "eta2").values
- entropy_values = entropy.transpose("t", "eta1", "eta2").values
gamma = model.propagators.variat_dens.options.gamma
- pressure = (gamma - 1) * rho_values**gamma * np.exp(entropy_values / rho_values)
+ pressure = ((gamma - 1) * rho**gamma * np.exp(entropy / rho)).assign_attrs(label="p")
# Field lines of the in-plane B are contours of the flux function A_z, with Bx = dA/dy and By = -dA/dx.
# A_z is recovered spectrally (the mean field vanishes); a duplicated periodic end point is dropped first.
@@ -120,7 +119,7 @@ def pproc(sim: Simulation, show: bool = False):
flux = np.pad(flux, ((0, 0), (0, len(x) - nx), (0, len(y) - ny)), mode="wrap")
# The classic picture: the density in a jet colour scale with the magnetic field lines on top.
- if not all(np.isfinite(field).all() for field in (rho_values, pressure, flux)):
+ if not all(np.isfinite(field).all() for field in (rho_values, pressure.values, flux)):
raise RuntimeError("Non-finite field diagnostic")
# Fixed levels over all times, so the lines follow the same flux surfaces as the field evolves.
flux_levels = {"start": float(flux.min()), "end": float(flux.max()), "size": float(np.ptp(flux)) / 14}
@@ -150,7 +149,7 @@ def traces(index):
scalars = output.scalars
energy = scalars.en_tot
- drift = np.abs(energy / energy.isel(t=0) - 1)
+ drift = energy.struphy.analysis.relative_error(skip_first=False)
divergence = np.sqrt(np.maximum(scalars.tot_div_B, 0))
diagnostics = make_subplots(rows=2, cols=1, shared_xaxes=True,
subplot_titles=("Energy channels", "Conservation diagnostics"), vertical_spacing=0.18)
@@ -164,10 +163,10 @@ def traces(index):
diagnostics.update_layout(template="plotly_white", margin={"l": 70, "r": 30, "t": 70, "b": 60})
cut_index = int(np.argmin(np.abs(y - np.pi)))
- cut = go.Figure()
- for index, label in ((0, "initial"), (-1, "final")):
- cut.add_scatter(x=x, y=pressure[index, :, cut_index], name=label, mode="lines")
- cut.update_layout(title="Gas pressure along y = π", xaxis_title="x", yaxis_title="p", template="plotly_white")
+ cut = pressure.struphy.plot.profiles(
+ x="eta1", at=[0, -1], x_of=lambda eta1: period * eta1, xlabel="x", title="Gas pressure along y = π",
+ eta2=cut_index, backend="plotly",
+ )
print(f"Maximum relative total-energy drift: {float(drift.max()):.3e}; maximum ‖div B‖: {float(divergence.max()):.3e}")
save(diagnostics, "orszag-tang-vortex-conservation", show=show)
save(cut, "orszag-tang-vortex-pressure-cut", show=show)
diff --git a/docs/src/examples/poisson-convergence.py b/docs/src/examples/poisson-convergence.py
index de73c68..6a2c3be 100644
--- a/docs/src/examples/poisson-convergence.py
+++ b/docs/src/examples/poisson-convergence.py
@@ -64,16 +64,15 @@ def create_simulation(degree=2, cells=8, alpha=distortion, folder="poisson_conve
)
-def potential_error(run, celldivide):
- """Points and the difference between the computed and the exact potential at the last time.
+def computed_potential(run, celldivide):
+ """The computed potential at the last time, with its physical points X, Y.
The potential is evaluated at `celldivide` points per cell and the endpoints, in the plane eta3 = 0.
"""
plane = {
f"eta{i + 1}": np.linspace(0.0, 1.0, cells * celldivide + 1) for i, cells in enumerate(run.grid.num_elements[:2])
}
- phi = run.evaluate("em_fields/phi", **plane, eta3=0.0).isel(t=-1)
- return phi.X.values, phi.Y.values, phi.values, phi.values - exact_potential(phi.X.values, phi.Y.values)
+ return run.evaluate("em_fields/phi", **plane, eta3=0.0).isel(t=-1)
def save(figure, name: str, *, show: bool = False, frame: int | None = None, still=None, width=1100, height=650):
@@ -100,6 +99,7 @@ def save(figure, name: str, *, show: bool = False, frame: int | None = None, sti
def pproc(sim: Simulation, show: bool = False):
from plotly.subplots import make_subplots
from scipy.interpolate import griddata
+ from struphy_plots.analysis import convergence_order
errors = {} # (alpha, degree) -> rms errors against the resolution
for alpha in (distortion, 0.0):
@@ -109,8 +109,8 @@ def pproc(sim: Simulation, show: bool = False):
run = create_simulation(degree, cells, alpha, f"poisson_convergence_p{degree}_n{cells}_a{alpha}").output
run.pproc(physical=True)
# three sample points per cell, to measure the error inside the cells
- _, _, _, difference = potential_error(run, celldivide=3)
- values.append(float(np.sqrt(np.mean(difference**2))))
+ phi = computed_potential(run, celldivide=3)
+ values.append(float(phi.struphy.analysis.error(exact_potential, norm="rms", args=("X", "Y"))))
errors[(alpha, degree)] = np.array(values)
if not all(np.isfinite(v).all() for v in errors.values()):
raise RuntimeError("Non-finite errors")
@@ -118,7 +118,7 @@ def pproc(sim: Simulation, show: bool = False):
n = np.array(resolutions, dtype=float) # the mesh width is h = 1 / n
h = 1.0 / n
# The slope of the last four resolutions, before the error reaches the level of the solver.
- slopes = {key: float(np.polyfit(np.log(h[-4:]), np.log(v[-4:]), 1)[0]) for key, v in errors.items()}
+ slopes = {key: convergence_order(h[-4:], v[-4:]).order for key, v in errors.items()}
for (alpha, degree), slope in slopes.items():
print(f"alpha = {alpha}, degree {degree}: slope {slope:.2f} (expected {degree + 1}), error at n = {resolutions[-1]}: "
f"{errors[(alpha, degree)][-1]:.2e}")
@@ -148,7 +148,9 @@ def pproc(sim: Simulation, show: bool = False):
# The solution and its error on the distorted mesh, at degree 2 and 8 x 12 cells.
run = create_simulation(2, 8, distortion, "poisson_convergence_map").output
run.pproc(physical=True)
- mesh_x, mesh_y, potential, difference = potential_error(run, celldivide=4)
+ phi = computed_potential(run, celldivide=4)
+ mesh_x, mesh_y, potential = phi.X.values, phi.Y.values, phi.values
+ difference = phi.struphy.analysis.error(exact_potential, norm="pointwise", args=("X", "Y")).values
x_plot, y_plot = np.linspace(0, lx, 100), np.linspace(0, ly, 150)
xx, yy = np.meshgrid(x_plot, y_plot)
points = np.column_stack([mesh_x.ravel(), mesh_y.ravel()])
diff --git a/docs/src/examples/poisson-source.py b/docs/src/examples/poisson-source.py
index 5718e99..a93c3ad 100644
--- a/docs/src/examples/poisson-source.py
+++ b/docs/src/examples/poisson-source.py
@@ -13,7 +13,6 @@
import argparse
import numpy as np
-import plotly.graph_objects as go
from struphy import (EnvironmentOptions, Simulation, Time, domains, grids,
perturbations)
from struphy.models import Poisson
@@ -98,109 +97,20 @@ def phi_exact(x, t):
phi_line = phi.isel(eta2=0, eta3=0)
if "component" in phi_line.dims:
phi_line = phi_line.isel(component=0)
- x = np.asarray(phi_line.X)
- times = np.asarray(phi_line.t)
-
+ # The largest pointwise error over the whole run, relative to the amplitude of the exact potential.
phi_scale = AMPLITUDE / k**2
- max_relative_error = 0.0
- frames = []
- for t in times:
- phi_h = np.asarray(phi_line.sel(t=t))
- phi_e = phi_exact(x, t)
- max_relative_error = max(
- max_relative_error, float(np.max(np.abs(phi_h - phi_e))) / phi_scale
- )
- frames.append(
- go.Frame(
- name=f"{t:.2f}",
- data=[
- go.Scatter(x=x, y=phi_h),
- go.Scatter(x=x, y=phi_e),
- ],
- ),
- )
-
+ max_error = phi_line.struphy.analysis.error(phi_exact, norm="max", dims=("t", "eta1"), args=("X", "t"))
+ max_relative_error = float(max_error) / phi_scale
print(f"Max relative error over the run: {max_relative_error:.5f}")
- figure = go.Figure(
- data=[
- go.Scatter(
- x=x,
- y=np.asarray(phi_line.isel(t=0)),
- mode="lines",
- name="Struphy (FEEC)",
- line={"color": "#168aad", "width": 3},
- ),
- go.Scatter(
- x=x,
- y=phi_exact(x, times[0]),
- mode="lines",
- name="Exact",
- line={"color": "#d62828", "width": 2, "dash": "dot"},
- ),
- ],
- frames=frames,
- )
- figure.update_layout(
+ figure = phi_line.struphy.plot.line_animation(
+ x="eta1",
+ reference={"Exact": phi_exact},
+ x_of=lambda eta1: domain.params["l1"] + Lx * eta1,
+ xlabel="x [a.u.]",
+ ylim=(-1.15 * phi_scale, 1.15 * phi_scale),
title="Poisson potential: FEEC solution vs. exact",
- xaxis_title="x [a.u.]",
- yaxis_title="φ [a.u.]",
- template="plotly_white",
- autosize=True,
- yaxis={"range": [-1.15 * phi_scale, 1.15 * phi_scale]},
- legend={
- "x": 0.02,
- "y": 0.98,
- "bgcolor": "rgba(255,255,255,0.82)",
- "bordercolor": "rgba(44,62,80,0.25)",
- "borderwidth": 1,
- },
- margin={"l": 60, "r": 30, "t": 80, "b": 130},
- updatemenus=[
- {
- "type": "buttons",
- "showactive": False,
- "x": 0.0,
- "xanchor": "left",
- "y": -0.32,
- "yanchor": "top",
- "buttons": [
- {
- "label": "Play",
- "method": "animate",
- "args": [
- None,
- {
- "frame": {"duration": 30, "redraw": True},
- "fromcurrent": True,
- },
- ],
- },
- ],
- },
- ],
- sliders=[
- {
- "steps": [
- {
- "args": [
- [frame.name],
- {
- "frame": {"duration": 0, "redraw": True},
- "mode": "immediate",
- },
- ],
- "label": frame.name,
- "method": "animate",
- }
- for frame in frames
- ],
- "x": 0.12,
- "len": 0.88,
- "y": -0.2,
- "currentvalue": {"prefix": "t = "},
- },
- ],
+ backend="plotly",
)
save(figure, "poisson-source", show=show)
diff --git a/docs/src/examples/pressureless-transport.py b/docs/src/examples/pressureless-transport.py
index 47e111c..8b7dab7 100644
--- a/docs/src/examples/pressureless-transport.py
+++ b/docs/src/examples/pressureless-transport.py
@@ -80,6 +80,7 @@ def save(figure, name: str, *, show: bool = False, frame: int | None = None, sti
def pproc(sim: Simulation, show: bool = False):
time_opts = sim.time_opts
from plotly.subplots import make_subplots
+ from struphy_plots.theory.exact import advected
output = sim.output
output.pproc(physical=True)
@@ -87,17 +88,19 @@ def pproc(sim: Simulation, show: bool = False):
velocity = output.evaluate("fluid/velocity_xyz").isel(component=0, eta2=0, eta3=0)
energy = output.scalars["kinetic_energy"]
times, x, density = rho.t.values, rho.eta1.values * length, rho.values
- exact = 1.0 + amplitude * np.cos(x[None, :] - speed * times[:, None])
+ exact = advected(lambda q: 1.0 + amplitude * np.cos(q), x[None, :], times[:, None], speed)
if not all(np.isfinite(a).all() for a in (density, velocity.values, energy.values)):
raise RuntimeError("The transport run produced non-finite diagnostics")
if not np.isclose(times[-1], time_opts.Tend) or np.min(density) <= 0:
raise RuntimeError("The transport run is incomplete or has non-positive density")
- errors = np.max(np.abs(density - exact), axis=1) / amplitude
+ errors = rho.struphy.analysis.error(exact, norm="max", dims="eta1").values / amplitude
velocity_error = float(np.max(np.abs(velocity.values - speed)) / speed)
- energy_drift = float(np.max(np.abs(energy.values / energy.values[0] - 1.0)))
+ energy_error = energy.struphy.analysis.relative_error(skip_first=False)
+ energy_drift = float(energy_error.max())
# Sampled mass uses the periodic evaluation grid, dropping its repeated endpoint.
- mass = length * np.mean(density[:, :-1], axis=1)
- mass_drift = float(np.max(np.abs(mass / mass[0] - 1.0)))
+ mass = length * rho.struphy.analysis.drop_periodic_endpoint("eta1").mean("eta1")
+ mass_error = mass.struphy.analysis.relative_error(skip_first=False)
+ mass_drift = float(mass_error.max())
print(f"Maximum profile error against exact transport: {errors.max():.2e} (relative to A)")
if errors.max() > 0.05 or velocity_error > 0.01 or max(energy_drift, mass_drift) > 1e-3:
raise RuntimeError("Pressureless transport failed its profile or conservation checks")
@@ -110,8 +113,8 @@ def pproc(sim: Simulation, show: bool = False):
figure.add_scatter(x=x[::4], y=exact[i, ::4], mode="markers", name="Exact transport",
showlegend=fraction == 0.0, marker={"color": color, "symbol": "circle-open"}, row=1, col=1)
for t, values, label in ((times, errors, "Profile error / A"),
- (times, np.abs(mass / mass[0] - 1.0), "Relative sampled-mass drift"),
- (energy.t.values, np.abs(energy.values / energy.values[0] - 1.0), "Relative kinetic-energy drift")):
+ (times, mass_error.values, "Relative sampled-mass drift"),
+ (energy.t.values, energy_error.values, "Relative kinetic-energy drift")):
figure.add_scatter(x=t, y=values, name=label, row=2, col=1)
figure.update_xaxes(title_text="x", row=1, col=1)
figure.update_yaxes(title_text="ρ", row=1, col=1)
diff --git a/docs/src/examples/resistive-diffusion.py b/docs/src/examples/resistive-diffusion.py
index 1dff7bb..a46c515 100644
--- a/docs/src/examples/resistive-diffusion.py
+++ b/docs/src/examples/resistive-diffusion.py
@@ -92,9 +92,10 @@ def field_profile(run):
return b_z.t.values, b_z.eta1.values * length, b_z.values
-def mode_amplitude(x, values):
- """Amplitude of the sin(k x) mode, from the periodic grid points (the last repeats the first)."""
- return 2.0 * np.mean(values[:, :-1] * np.sin(wavenumber * x[:-1]), axis=1)
+def mode_amplitude(run):
+ """Amplitude of the sin(k x) mode of B_z over time, from a post-processed run."""
+ b_z = run.evaluate("em_fields/b_field_xyz").isel(component=2, eta2=0, eta3=0)
+ return b_z.struphy.analysis.project_mode(dim="eta1", number=mode_number)
def save(figure, name: str, *, show: bool = False, frame: int | None = None, still=None, width=1100, height=650):
@@ -130,9 +131,9 @@ def pproc(sim: Simulation, show: bool = False):
amplitudes, fitted = {}, {}
for eta, run in runs.items():
- times, x, values = field_profile(run)
- amplitudes[eta] = mode_amplitude(x, values)
- fitted[eta] = float(-np.polyfit(times, np.log(np.abs(amplitudes[eta])), 1)[0])
+ amplitudes[eta] = abs(mode_amplitude(run))
+ amplitudes[eta].attrs = {"label": f"η = {eta}: Struphy"}
+ fitted[eta] = -amplitudes[eta].struphy.analysis.growth_rate().rate
exact_rate = {eta: eta * wavenumber**2 for eta in runs}
if not all(np.isfinite(list(fitted.values()))):
raise RuntimeError("A decay rate could not be fitted")
@@ -146,7 +147,7 @@ def pproc(sim: Simulation, show: bool = False):
thermal = np.asarray(run.scalars["en_thermo"])
total = np.asarray(run.scalars["en_tot"])
scalar_times = np.asarray(run.time)[: len(total)]
- energy_drift = float(np.max(np.abs(total / total[0] - 1.0)))
+ energy_drift = float(run.scalars["en_tot"].struphy.analysis.relative_error().max())
print(f"Maximum relative drift of the total energy: {energy_drift:.2e}")
def profile_traces(index):
@@ -183,18 +184,15 @@ def profile_traces(index):
figure.update_yaxes(title_text="energy change / magnetic energy lost", row=2, col=1)
save(figure, "resistive-diffusion", width=900, height=850, show=show)
- colors = {0.05: "#168aad", 0.1: "#d62828", 0.2: "#f77f00"}
- decay = go.Figure()
- for eta, amp in amplitudes.items():
- decay.add_scatter(x=times, y=np.abs(amp), mode="lines", name=f"η = {eta}: Struphy",
- line={"color": colors[eta], "width": 3})
- decay.add_scatter(x=times, y=amplitude * np.exp(-exact_rate[eta] * times), mode="lines",
- name=f"η = {eta}: exp(−η k² t)", line={"color": "#111", "width": 1.5, "dash": "dash"})
- decay.update_layout(
- title="Decay of the field amplitude", template="plotly_white", autosize=True,
- xaxis_title="t", yaxis_title="amplitude of the sin(kx) mode", yaxis_type="log",
- margin={"l": 75, "r": 30, "t": 80, "b": 60},
+ first, *others = amplitudes.values()
+ decay = first.struphy.plot.timeseries(
+ *others,
+ logy=True,
+ reference={f"η = {eta}: exp(−η k² t)": lambda t, eta=eta: amplitude * np.exp(-exact_rate[eta] * t) for eta in runs},
+ title="Decay of the field amplitude",
+ backend="plotly",
)
+ decay.fig.update_yaxes(title_text="amplitude of the sin(kx) mode")
save(decay, "resistive-diffusion-decay", show=show)
diff --git a/docs/src/examples/shear-alfven-wave.py b/docs/src/examples/shear-alfven-wave.py
index eeec8a0..a469d3b 100644
--- a/docs/src/examples/shear-alfven-wave.py
+++ b/docs/src/examples/shear-alfven-wave.py
@@ -14,7 +14,6 @@
import argparse
import numpy as np
-import plotly.graph_objects as go
from struphy import (
DerhamOptions,
@@ -106,72 +105,28 @@ def pproc(sim: Simulation, show: bool = False):
output = sim.output.pproc(create_vtk=False)
# The (k, omega) power spectrum of u_x along z, and a fit of its branch.
- from struphy_plots.analysis import fit_dispersion_branches, power_spectrum
+ from struphy_plots.analysis import power_spectrum
velocity = output.fields.mhd.velocity
u_x = velocity.isel(component=0, eta1=0, eta2=0)
u_x = u_x.assign_coords(eta3=u_x.eta3 * (domain.params["r3"] - domain.params["l3"])) # physical z
spectrum = power_spectrum(u_x, dim="eta3")
# Peaks count above half the column's peak amplitude, i.e. a quarter of its peak power.
- branch = fit_dispersion_branches(spectrum, n_branches=1, noise_level=0.5**2, order=10)[0]
+ branch = spectrum.struphy.analysis.fit_branches(n_branches=1, noise_level=0.5**2, order=10)[0]
phase_velocity = float(branch.velocity)
print(f"Measured Alfvén speed: {phase_velocity:.5f} (exact: 1.0)")
- # Build an interactive Plotly view of the normalized power spectrum, for omega, k >= 0.
- quadrant = spectrum.sel(omega=spectrum.omega >= 0, k=spectrum.k >= 0)
- omega = quadrant.omega.values
- kvec = quadrant.k.values
- power = quadrant.values / float(quadrant.max())
- log_power = np.log10(np.clip(power, 1e-15, None))
- fit = phase_velocity * kvec
-
- figure = go.Figure(
- go.Heatmap(
- x=kvec,
- y=omega,
- z=log_power,
- zmin=-15,
- zmax=-1,
- colorscale="Plasma",
- colorbar={
- "title": {"text": "log₁₀ P"},
- "tickvals": [-15, -12, -9, -6, -3],
- "ticktext": ["10⁻¹⁵", "10⁻¹²", "10⁻⁹", "10⁻⁶", "10⁻³"],
- },
- hovertemplate="k=%{x:.3f}
ω=%{y:.3f}
log₁₀ P=%{z:.2f}",
- ),
- )
- figure.add_scatter(
- x=kvec,
- y=kvec,
- mode="lines",
- name="Alfvén wave, v_A = 1",
- line={"color": "#168aad", "width": 3, "dash": "dash"},
- )
- figure.add_scatter(
- x=kvec,
- y=fit,
- mode="lines",
- name=f"Struphy fit, v_A = {phase_velocity:.5f}",
- line={"color": "#d62828", "width": 3, "dash": "dot"},
- )
- figure.update_layout(
+ # The normalized power spectrum over 15 decades, for omega, k >= 0, with the exact and the fitted branch.
+ figure = spectrum.struphy.plot.dispersion(
+ kmin=0,
+ branches={"Alfvén wave, v_A = 1": lambda k: k},
+ fits=[branch],
+ dynamic_range=15,
+ cmap="plasma",
+ omega_max=float(spectrum.k.max()),
title="Shear-Alfvén wave dispersion",
- xaxis_title="k [a.u.]",
- yaxis_title="ω [a.u.]",
- template="plotly_white",
- autosize=True,
- legend={
- "x": 0.02,
- "y": 0.98,
- "bgcolor": "rgba(255,255,255,0.82)",
- "bordercolor": "rgba(44,62,80,0.25)",
- "borderwidth": 1,
- },
- margin={"l": 75, "r": 45, "t": 80, "b": 70},
+ backend="plotly",
)
- figure.update_xaxes(range=[0, float(kvec[-1])])
- figure.update_yaxes(range=[0, float(kvec[-1])])
save(figure, "shear-alfven-wave", show=show)
diff --git a/docs/src/examples/sph-velocity-diffusion.py b/docs/src/examples/sph-velocity-diffusion.py
index 16083ef..c89de6c 100644
--- a/docs/src/examples/sph-velocity-diffusion.py
+++ b/docs/src/examples/sph-velocity-diffusion.py
@@ -13,7 +13,6 @@
import argparse
import numpy as np
-import plotly.graph_objects as go
from struphy import (
BinningPlot,
@@ -113,61 +112,28 @@ def pproc(sim: Simulation, show: bool = False):
output.pproc()
velocity = output.evaluate("euler_fluid/e1_current_1/f")
times = np.asarray(velocity.t.values)
- x = np.asarray(velocity.eta1.values) * length
values = np.asarray(velocity.values)
if not np.isfinite(values).all():
raise RuntimeError("Non-finite SPH velocity: refusing to publish the run")
- # A finite bin records the average of a sine wave, rather than its point value.
- bin_average = np.sinc(wavenumber * (length / bins) / (2 * np.pi))
- mode_amplitude = 2 * np.mean(values * np.sin(wavenumber * x), axis=1) / bin_average
+ # The sine amplitude of the mode; a finite bin records the average of a sine wave, rather than its point value.
+ mode_amplitude = velocity.struphy.analysis.project_mode(dim="eta1", number=1, bin_correction=True)
exact_amplitude = initial_amplitude * np.exp(-exact_decay_rate * times)
- measured_decay_rate = float(-np.polyfit(times, np.log(np.abs(mode_amplitude)), 1)[0])
- rms_error = float(np.sqrt(np.mean((mode_amplitude - exact_amplitude) ** 2)))
+ measured_decay_rate = -mode_amplitude.struphy.analysis.growth_rate().rate
+ rms_error = float(mode_amplitude.struphy.analysis.error(exact_amplitude, norm="rms", dims="t"))
print(
f"decay rate {measured_decay_rate:.4f} (exact {exact_decay_rate:.4f}); "
f"amplitude RMS error {rms_error:.4g}"
)
- dense_x = np.linspace(0.0, length, 301)
- picks = np.unique(np.linspace(0, len(times) - 1, min(80, len(times)), dtype=int))
-
- def profile(index):
- return [
- go.Scatter(
- x=dense_x,
- y=exact_amplitude[index] * np.sin(wavenumber * dense_x),
- mode="lines",
- name="exact",
- line={"color": "#111", "width": 2, "dash": "dash"},
- ),
- go.Scatter(
- x=x,
- y=values[index],
- mode="lines+markers",
- name="SPH",
- line={"color": "#168aad", "width": 2},
- marker={"size": 5},
- ),
- ]
-
- figure = go.Figure(data=profile(0))
- figure.frames = [go.Frame(name=f"{times[i]:.3f}", data=profile(i)) for i in picks]
- figure.update_layout(
+ # The binned profile against the exact decaying mode, in at most 80 frames.
+ figure = velocity.assign_attrs(label="SPH velocity u").struphy.plot.line_animation(
+ x_of=lambda eta1: length * eta1,
+ reference={"exact": lambda x, t: initial_amplitude * np.exp(-exact_decay_rate * t) * np.sin(wavenumber * x)},
+ ylim=(-1.1 * initial_amplitude, 1.1 * initial_amplitude),
+ step=-(-len(times) // 80),
title="SPH velocity diffusion: sinusoidal mode against its exact decay",
- template="plotly_white",
- xaxis_title="x",
- yaxis_title="u(x, t)",
- yaxis={"range": [-1.1 * initial_amplitude, 1.1 * initial_amplitude]},
- margin={"l": 65, "r": 30, "t": 80, "b": 130},
- legend={"orientation": "h", "y": 1.1},
- updatemenus=[{"type": "buttons", "showactive": False, "x": 0, "y": -0.28, "buttons": [
- {"label": "Play", "method": "animate", "args": [None, {"frame": {"duration": 55, "redraw": False}, "fromcurrent": True}]}
- ]}],
- sliders=[{"x": 0.12, "len": 0.88, "y": -0.18, "currentvalue": {"prefix": "t = "}, "steps": [
- {"args": [[frame.name], {"frame": {"duration": 0, "redraw": False}, "mode": "immediate"}], "label": frame.name, "method": "animate"}
- for frame in figure.frames
- ]}],
+ backend="plotly",
)
save(figure, "sph-velocity-diffusion", show=show)
diff --git a/docs/src/examples/strong-landau-damping.py b/docs/src/examples/strong-landau-damping.py
index e8284d8..39793f7 100644
--- a/docs/src/examples/strong-landau-damping.py
+++ b/docs/src/examples/strong-landau-damping.py
@@ -16,7 +16,6 @@
import argparse
import numpy as np
-import plotly.graph_objects as go
from struphy import (
BoundaryParameters,
@@ -114,8 +113,6 @@ def save(figure, name: str, *, show: bool = False, frame: int | None = None, sti
def pproc(sim: Simulation, show: bool = False):
output = sim.output
field_energy_array = output.scalars["electric_energy"]
- time = np.asarray(field_energy_array.t)
- field_energy = np.asarray(field_energy_array)
# The bounce period of trapped particles shows up as the spacing between
# local maxima in the field energy, once the initial (linear) damping
@@ -124,19 +121,8 @@ def pproc(sim: Simulation, show: bool = False):
bounce_period = float(np.mean(np.diff(maxima_t))) if len(maxima_t) > 1 else float("nan")
print(f"Estimated trapped-particle bounce period: {bounce_period:.2f}")
- figure = go.Figure(
- data=[
- go.Scatter(x=time, y=field_energy, mode="lines", name="Struphy (PIC)", line={"color": "#168aad", "width": 3}),
- ],
- )
- figure.update_layout(
- title="Strong Landau damping: 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_array.struphy.plot.timeseries(
+ logy=True, title="Strong Landau damping: electric field energy", backend="plotly",
)
save(figure, "strong-landau-damping", show=show)
diff --git a/docs/src/examples/toroidal-shear-alfven.py b/docs/src/examples/toroidal-shear-alfven.py
index 9e03d60..a2b059e 100644
--- a/docs/src/examples/toroidal-shear-alfven.py
+++ b/docs/src/examples/toroidal-shear-alfven.py
@@ -66,23 +66,17 @@ def radial_mode_amplitudes(output, field_xyz, poloidal_modes=(9, 10, 11, 12), se
radial = physical_radial_component(output, field_xyz)
# Convert to the rotating radial basis BEFORE transforming in toroidal angle.
# Cartesian components themselves are not periodic across the sector seam.
- radial = radial.isel({dim: np.flatnonzero(radial[dim].values < 1.0 - 1e-12)
- for dim in ("eta2", "eta3")})
- coefficients = output.analysis.fft(output.analysis.fft(radial, dim="eta2"), dim="eta3")
- selected = coefficients.sel(
- k_eta2=2 * np.pi * np.asarray(poloidal_modes), k_eta3=2 * np.pi * sector_mode,
- method="nearest",
- )
- if (not np.allclose(selected.k_eta2.values / (2 * np.pi), poloidal_modes)
- or not np.isclose(float(selected.k_eta3) / (2 * np.pi), sector_mode)):
+ # mode_spectrum drops the duplicate periodic endpoints and labels integer (m, n_sector).
+ coefficients = radial.struphy.analysis.mode_spectrum(dims=("eta2", "eta3"), names=("m", "n"))
+ if not set(poloidal_modes) <= set(coefficients.m.values) or sector_mode not in coefficients.n.values:
raise ValueError("Angular sampling does not resolve the requested Fourier modes.")
+ selected = coefficients.sel(m=list(poloidal_modes), n=sector_mode)
amplitude = 2 * abs(selected) # Conjugate-pair amplitude of a real spatial harmonic.
if not np.isfinite(amplitude.values).all():
raise RuntimeError("Non-finite radial Fourier amplitude.")
peak = float(amplitude.max())
- normalized = (amplitude / (peak if peak > 0 else 1.0)).rename({"k_eta2": "m"})
+ normalized = amplitude / (peak if peak > 0 else 1.0)
normalized = normalized.assign_coords(
- m=list(poloidal_modes),
radius=params["a1"] + (params["a2"] - params["a1"]) * normalized.eta1,
).rename("normalized_fft_amplitude")
normalized.attrs.update(normalization_amplitude=peak, sector_mode=sector_mode,
@@ -101,11 +95,10 @@ def fixed_theta_amplitudes(output, field_xyz, angles=(0.0, 45.0), sector_mode=-1
# Interpolate the physical radial field only if an angle is off the display
# grid. The default 0 and 45 degree rays lie exactly on that grid.
rays = radial.interp(eta2=np.asarray(angles) / 360.0)
- rays = rays.isel(eta3=np.flatnonzero(rays.eta3.values < 1.0 - 1e-12))
- coefficients = output.analysis.fft(rays, dim="eta3")
- selected = coefficients.sel(k_eta3=2 * np.pi * sector_mode, method="nearest")
- if not np.isclose(float(selected.k_eta3) / (2 * np.pi), sector_mode):
+ coefficients = rays.struphy.analysis.mode_spectrum(dims="eta3", names="n")
+ if sector_mode not in coefficients.n.values:
raise ValueError("Toroidal sampling does not resolve the requested Fourier mode.")
+ selected = coefficients.sel(n=sector_mode)
amplitude = 2 * abs(selected)
if not np.isfinite(amplitude.values).all():
raise RuntimeError("Non-finite fixed-angle Fourier amplitude.")
@@ -410,7 +403,7 @@ def traces(index):
# transforming a velocity magnitude/energy would change its frequencies.
# Omit the duplicated poloidal endpoint from spatial sums.
plane = velocity.copy(data=components).assign_coords(component=list(labels))
- plane = plane.isel(eta2=np.flatnonzero(periodic)).rename("physical_poloidal_velocity")
+ plane = plane.struphy.analysis.drop_periodic_endpoint("eta2").rename("physical_poloidal_velocity")
temporal = output.analysis.time_fft(plane)
band = output.analysis.filter_time(plane, dims=("eta1", "eta2"), pad_bins=0)
positive = temporal.power.isel(omega=slice(1, None))
@@ -484,9 +477,8 @@ def traces(index):
eta3=0.0,
representation="2",
).isel(t=0, component=0)
- logical_initial = logical_initial.isel(eta2=np.flatnonzero(periodic))
- poloidal_fft = output.analysis.fft(logical_initial, dim="eta2")
- mode_numbers = poloidal_fft.k_eta2.values / (2 * np.pi)
+ poloidal_fft = logical_initial.struphy.analysis.mode_spectrum(dims="eta2", names="m")
+ mode_numbers = poloidal_fft.m.values
modal_amplitude = np.sqrt((abs(poloidal_fft) ** 2).mean("eta1")).values
positive_modes = (mode_numbers > 0) & (mode_numbers <= grid.num_elements[1] // 2)
mode_plot = go.Figure(go.Scatter(
diff --git a/docs/src/examples/two-stream-instability.py b/docs/src/examples/two-stream-instability.py
index b0f7f37..59b1a00 100644
--- a/docs/src/examples/two-stream-instability.py
+++ b/docs/src/examples/two-stream-instability.py
@@ -14,7 +14,7 @@
import argparse
-import plotly.graph_objects as go
+import numpy as np
from struphy import (
BinningPlot,
@@ -118,6 +118,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.kinetic import two_stream
+
domain = sim.domain
output = sim.output
@@ -127,28 +129,11 @@ def pproc(sim: Simulation, show: bool = False):
# Fit the exponential growth rate over the clean linear-growth window
# (before trapping saturates it, roughly t in [5, 25] for this setup).
growth_rate = field_energy.struphy.analysis.growth_rate(window=(5.0, 25.0), amplitude=True).rate
- print(f"Measured growth rate: {growth_rate:.4f} (expected: ~0.2845, from the linear dispersion relation)")
-
- 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="Two-stream 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},
- )
+ # The linear kinetic theory of two Maxwellian beams at u = ±3 with v_th = 1, for the box mode k = 2π/L.
+ expected = two_stream(2 * np.pi / domain.params["r1"], beam_speed=3.0, thermal_speed=1.0).imag
+ print(f"Measured growth rate: {growth_rate:.4f} (expected: ~{expected:.4f}, from the linear dispersion relation)")
+
+ figure = field_energy.struphy.plot.timeseries(logy=True, title="Two-stream instability: electric field energy", backend="plotly")
save(figure, "two-stream-instability", show=show)
diff --git a/docs/src/examples/vlasov-tokamak.py b/docs/src/examples/vlasov-tokamak.py
index 4ef3413..347c95c 100644
--- a/docs/src/examples/vlasov-tokamak.py
+++ b/docs/src/examples/vlasov-tokamak.py
@@ -116,8 +116,9 @@ def pproc(sim: Simulation, show: bool = False):
# A magnetic field alone does no work, so each particle's speed should be
# conserved -- a genuine accuracy check on the pusher, not just a demo.
- speed = np.linalg.norm(np.asarray(output.orbits.kinetic_ions.to_dataarray("quantity").transpose("t", "marker", "quantity").sel(quantity=["v1", "v2", "v3"])), axis=2)
- max_relative_speed_drift = float(np.max(np.abs(speed - speed[0]) / speed[0]))
+ velocity = output.orbits.kinetic_ions.to_dataarray("quantity").sel(quantity=["v1", "v2", "v3"])
+ speed = velocity.struphy.analysis.norm(dims=["quantity"])
+ max_relative_speed_drift = float(speed.struphy.analysis.relative_error().max())
print(f"Max relative drift in particle speed (should be ~0): {max_relative_speed_drift:.5f}")
# The plasma boundary (outer flux surface), for visual context around the orbits.
diff --git a/docs/src/examples/vortex-merger.py b/docs/src/examples/vortex-merger.py
index d6df91e..3210bd3 100644
--- a/docs/src/examples/vortex-merger.py
+++ b/docs/src/examples/vortex-merger.py
@@ -16,7 +16,6 @@
import argparse
import numpy as np
-import plotly.graph_objects as go
from struphy import (
BaseUnits,
@@ -164,20 +163,20 @@ def pproc(sim: Simulation, show: bool = False):
movie = mapped.struphy.plot.animation(
x="x", y="y", max_frames=150, vmin=0.0, vmax=float(density.max()), cmap="viridis",
title="Vortex merger: binned charge density", xlabel="x", ylabel="y", colorbar_label="density",
- backend="plotly",
+ equal_aspect=True, backend="plotly",
)
- movie.fig.update_xaxes(range=[-a2, a2], constrain="domain")
- movie.fig.update_yaxes(range=[-a2, a2], scaleanchor="x", scaleratio=1)
save(movie, "vortex-merger", height=750, frame=len(movie.fig.frames) // 2, show=show)
# The first Poisson solve initializes the field energy after t=0. Compare
# subsequent field energies to that first solved state, not to the zero placeholder.
energy = output.scalars["en_phi"].isel(t=slice(1, None))
- drift = (energy / energy.isel(t=0) - 1).values
- energy_figure = go.Figure(go.Scatter(x=energy.t.values, y=drift, mode="lines", name="field energy"))
- energy_figure.update_layout(title="Vortex merger: electrostatic-energy change", template="plotly_white",
- xaxis_title="t", yaxis_title="(W − W₁) / W₁", margin={"l": 80, "r": 30, "t": 80, "b": 60})
- print(f"Maximum relative drift of the electrostatic energy: {np.abs(drift).max():.2e}")
+ change = energy.struphy.analysis.drift() / energy.isel(t=0)
+ change.attrs.update(label="(W − W₁) / W₁", units="")
+ energy_figure = change.struphy.plot.timeseries(
+ logy=False, title="Vortex merger: electrostatic-energy change", backend="plotly"
+ )
+ drift = float(energy.struphy.analysis.relative_error().max())
+ print(f"Maximum relative drift of the electrostatic energy: {drift:.2e}")
save(energy_figure, "vortex-merger-energy", show=show)
diff --git a/docs/src/examples/weak-landau-damping.py b/docs/src/examples/weak-landau-damping.py
index 8a2884b..85d09f5 100644
--- a/docs/src/examples/weak-landau-damping.py
+++ b/docs/src/examples/weak-landau-damping.py
@@ -15,7 +15,6 @@
import argparse
import numpy as np
-import plotly.graph_objects as go
from struphy import (
BoundaryParameters,
@@ -129,38 +128,11 @@ def field_energy_exact(t):
measured_rate = field_energy.struphy.analysis.damping_rate(window=(None, 8.0), amplitude=True).rate
print(f"Measured damping rate: {measured_rate:.4f} (exact: -0.1533)")
- figure = go.Figure(
- data=[
- go.Scatter(
- x=time,
- y=np.asarray(field_energy),
- mode="lines",
- name="Struphy (PIC)",
- line={"color": "#168aad", "width": 3},
- ),
- go.Scatter(
- x=time,
- y=field_energy_exact(time),
- mode="lines",
- name="Exact envelope",
- line={"color": "#d62828", "width": 2, "dash": "dot"},
- ),
- ],
- )
- figure.update_layout(
+ figure = field_energy.struphy.plot.timeseries(
+ logy=True,
+ reference={"Exact envelope": (time, field_energy_exact(time))},
title="Weak Landau damping: electric field energy",
- xaxis_title="t [a.u.]",
- yaxis_title="E² / 2 [a.u.]",
- yaxis={"type": "log"},
- template="plotly_white",
- autosize=True,
- legend={
- "x": 0.98,
- "y": 0.98,
- "xanchor": "right",
- "bgcolor": "rgba(255,255,255,0.82)",
- },
- margin={"l": 70, "r": 30, "t": 80, "b": 60},
+ backend="plotly",
)
save(figure, "weak-landau-damping", show=show)
diff --git a/docs/src/examples/weibel-instability.py b/docs/src/examples/weibel-instability.py
index 044b653..4e37d60 100644
--- a/docs/src/examples/weibel-instability.py
+++ b/docs/src/examples/weibel-instability.py
@@ -17,7 +17,6 @@
import argparse
import numpy as np
-import plotly.graph_objects as go
from struphy import (
BinningPlot,
@@ -134,37 +133,22 @@ def pproc(sim: Simulation, show: bool = False):
# Total magnetic (B3) field energy at each saved time, summed over the grid.
b_field = output.fields.em_fields.b_field
- times = np.asarray(b_field.t)
grid_shape = tuple(b_field.sizes[dim] for dim in ("eta1", "eta2", "eta3"))
cell_volume = float(np.prod([1.0 / max(n - 1, 1) for n in grid_shape]))
- magnetic_energy = np.asarray((b_field.isel(component=2) ** 2).sum(("eta1", "eta2", "eta3"))) * cell_volume / 2
- time = np.asarray(times)
+ magnetic_energy = (b_field.isel(component=2) ** 2).sum(("eta1", "eta2", "eta3")) * cell_volume / 2
+ magnetic_energy.attrs = {"label": "|B₃|² / 2"}
+ time = np.asarray(magnetic_energy.t)
# Fit the growth rate over the clean exponential window (roughly the
# middle third of the run, before saturation).
growth_window = (time > time[-1] / 5) & (time < 2 * time[-1] / 5)
- growth_rate = float(np.polyfit(time[growth_window], np.log(magnetic_energy[growth_window]), 1)[0] / 2)
+ growth_rate = float(np.polyfit(time[growth_window], np.log(magnetic_energy.values[growth_window]), 1)[0] / 2)
print(f"Measured growth rate (in |B3|, from the energy fit): {growth_rate:.5f}")
- figure = go.Figure(
- data=[
- go.Scatter(
- x=time,
- y=magnetic_energy,
- mode="lines",
- name="|B₃|² / 2 (Struphy)",
- line={"color": "#168aad", "width": 3},
- ),
- ],
- )
- figure.update_layout(
+ figure = magnetic_energy.struphy.plot.timeseries(
+ logy=True,
title="Weibel instability: magnetic field energy",
- xaxis_title="t [a.u.]",
- yaxis_title="|B₃|² / 2 [a.u.]",
- yaxis={"type": "log"},
- template="plotly_white",
- autosize=True,
- margin={"l": 70, "r": 30, "t": 80, "b": 60},
+ backend="plotly",
)
save(figure, "weibel-instability", show=show)
@@ -188,21 +172,11 @@ def pproc(sim: Simulation, show: bool = False):
f = output.evaluate("kinetic_ions/v1_v2_density/f") # (t, v1, v2)
moments = f.struphy.analysis.velocity_moments()
anisotropy = moments.variance_v2 / moments.variance_v1
- anisotropy_figure = go.Figure(
- go.Scatter(
- x=anisotropy.t.values,
- y=anisotropy.values,
- mode="lines",
- line={"color": "#168aad", "width": 3},
- ),
- )
- anisotropy_figure.update_layout(
+ anisotropy.attrs = {"label": "⟨(v₂ − u₂)²⟩ / ⟨(v₁ − u₁)²⟩"}
+ anisotropy_figure = anisotropy.struphy.plot.timeseries(
+ logy=False,
title="Weibel instability: temperature anisotropy",
- xaxis_title="t [a.u.]",
- yaxis_title="⟨(v₂ − u₂)²⟩ / ⟨(v₁ − u₁)²⟩",
- template="plotly_white",
- autosize=True,
- margin={"l": 70, "r": 30, "t": 80, "b": 60},
+ backend="plotly",
)
save(space_time, "weibel-instability-space-time", show=show)
diff --git a/docs/src/examples/zeldovich-caustic.py b/docs/src/examples/zeldovich-caustic.py
index 2153206..ba8e13f 100644
--- a/docs/src/examples/zeldovich-caustic.py
+++ b/docs/src/examples/zeldovich-caustic.py
@@ -143,7 +143,7 @@ def pproc(sim: Simulation, show: bool = False):
# The relative L1 error of the density estimate, before the caustic (a smooth density) and after it.
exact = np.array([exact_density(edges, t) for t in times])
- error = np.abs(density.values - exact).sum(axis=1) / exact.sum(axis=1)
+ error = density.struphy.analysis.error(exact, norm="l1", relative=True).values
before = times < 0.8 * caustic_time
after = times > 1.2 * caustic_time
error_before, error_after = float(error[before].mean()), float(error[after].mean())