From b35b1bdf6ba21a7bd631dc5c2e04af8c150e10cc Mon Sep 17 00:00:00 2001 From: Nicolo' Minichelli Date: Fri, 23 Oct 2026 13:40:56 +0200 Subject: [PATCH] Variational MHD: relative Newton criteria, fixed en_thermo, TAE examples - Newton/Picard stopping criteria of the variational propagators (PB, QB, density, magnetic field, entropy, viscosity, resistivity) relative to the state at t^n (err < tol^2 min(|state|^2, 1)): the absolute tolerance was far too loose for the small states of the linear and perturbative models (energy drift ~1e-5 per step in the low-mode TAE). Linear solves and inverse-mass error norms are done for normalized vectors, so their tolerances are relative too. New stagnation check (3 iterations without decrease). Shared helpers NewtonConvergence, dot_normalized and inv_norm2 in feec/variational_utilities.py. VariationalMomentumAdvection is unchanged. - Schur solves of VariationalPBEvolve and VariationalQBEvolve preconditioned with the fixed MassMatrixPreconditioner instead of an inner PCG solve (2x faster for the p model). - en_thermo (and en_thermo_l1) of VariationalPerturbativeMHD and VariationalMHD_with_p: int p dx / (gamma - 1) (was int p^3 / sqrt(g) d eta, wrong in curved geometry). - project_div_free_velocity: Mv-orthogonal projection of an H1vec velocity onto the kernel of the discrete divergence of the variational schemes. - Examples examples/VariationalPerturbativeLinearMHD/with_p|with_q/lowmode_tae: the low-mode TAE of LinearMHD for the linear variational models (first order in time, divergence-free initial velocity). - test_same_as_full_f: threshold 1e-6 (the runs agree up to the solver tolerance). Co-Authored-By: Claude Opus 5.5 --- examples/README.md | 3 +- .../with_p/lowmode_tae/params_lowmode_tae.py | 251 +++++++++++++++++ .../with_p/lowmode_tae/pproc_lowmode_tae.py | 196 +++++++++++++ .../with_q/lowmode_tae/params_lowmode_tae.py | 258 ++++++++++++++++++ .../with_q/lowmode_tae/pproc_lowmode_tae.py | 196 +++++++++++++ src/struphy/feec/variational_utilities.py | 195 +++++++++++++ ...verif_VariationalPerturbativeMHD_with_q.py | 4 +- src/struphy/models/variational_mhd_with_p.py | 13 +- .../models/variational_perturbative_mhd.py | 15 +- .../propagators/variational_density_evolve.py | 32 ++- .../propagators/variational_entropy_evolve.py | 34 ++- .../variational_mag_field_evolve.py | 34 ++- .../propagators/variational_pb_evolve.py | 43 ++- .../propagators/variational_qb_evolve.py | 46 ++-- .../propagators/variational_resistivity.py | 37 +-- .../propagators/variational_viscosity.py | 39 ++- 16 files changed, 1256 insertions(+), 140 deletions(-) create mode 100644 examples/VariationalPerturbativeLinearMHD/with_p/lowmode_tae/params_lowmode_tae.py create mode 100644 examples/VariationalPerturbativeLinearMHD/with_p/lowmode_tae/pproc_lowmode_tae.py create mode 100644 examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/params_lowmode_tae.py create mode 100644 examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/pproc_lowmode_tae.py diff --git a/examples/README.md b/examples/README.md index 9f1f569f2..d04a81ff1 100644 --- a/examples/README.md +++ b/examples/README.md @@ -58,12 +58,13 @@ with `pip install "struphy[pproc]"`. | `DriftKineticElectrostaticAdiabatic` | `cyclone` (Cyclone base case ITG), `itg_cylindre` (ITG in a cylinder), `slab_drift_wave` | | `GuidingCenter` | `screw_pinch_equilibrium`, `tokamak_equilibrium` | | `GyrokineticPoisson` | `gaussian_hollow_cylinder`, `gaussian_hollow_torus`, `gaussian_on_axis_cylinder`, `gaussian_on_axis_torus` | -| `LinearMHD` | `itpa_tae_benchmark` (toroidal Alfvén eigenmode) | +| `LinearMHD` | `itpa_tae_benchmark`, `lowmode_tae` (toroidal Alfvén eigenmodes) | | `LinearMHDDriftkineticCC` | `lowmode_tae` (toroidal Alfvén eigenmode driven by energetic ions) | | `Poisson` | `gaussian_cube`, `gaussian_hollow_cylinder`, `gaussian_hollow_torus`, `gaussian_on_axis_cylinder`, `gaussian_on_axis_torus` | | `ShearAlfven` | `itpa_tae_benchmark`, `lowmode_tae` (toroidal Alfvén eigenmodes) | | `ToyDrift` | `diocotron_instability` | | `VariationalMHD` | `orszag_tang` (Orszag-Tang vortex), `slab_dispersion` (Alfvén, slow and fast waves) | +| `VariationalPerturbativeLinearMHD` | `with_p/lowmode_tae`, `with_q/lowmode_tae` (the `LinearMHD` low-mode TAE with `VariationalPerturbativeMHD` and `VariationalPerturbativeMHD_with_q`, `linear=True`) | | `VariationalPerturbativeMHD_with_q` | `slab_dispersion` (Alfvén, slow and fast waves), `slab_dispersion_linear` (the same with `linear=True`) | | `Vlasov` | `maxwellian_uniform_field` | | `VlasovAmpereOneSpecies` | `bump_on` (bump-on-tail), `strong_Landau_damping`, `two_stream`, `weak_Landau_damping` | diff --git a/examples/VariationalPerturbativeLinearMHD/with_p/lowmode_tae/params_lowmode_tae.py b/examples/VariationalPerturbativeLinearMHD/with_p/lowmode_tae/params_lowmode_tae.py new file mode 100644 index 000000000..a98b04b21 --- /dev/null +++ b/examples/VariationalPerturbativeLinearMHD/with_p/lowmode_tae/params_lowmode_tae.py @@ -0,0 +1,251 @@ +"""Toroidal Alfven eigenmode (TAE) with low mode numbers for the linear variational MHD model in the pressure +(VariationalPerturbativeMHD with linear=True, without viscosity and resistivity): an n = 1 shear Alfven wave in the +gap of the Alfven continuum at the crossing of the m = 1 and m = 2 branches (q = 1.5, r = 0.5) in a tokamak of aspect +ratio 5. The same case as examples/LinearMHD/lowmode_tae, for the comparison of the variational discretization with +the (non-variational) LinearMHD models.""" + +# ----------------------------- +# Description of the simulation +# ----------------------------- +# Please fill in a verbal description of the simulation. +# It will be printed at the beginning of the simulation and can be used to keep track of the different runs. + +name = "Low-mode TAE (VariationalPerturbativeMHD, linear)" +description = """ +Toroidal Alfven eigenmode (TAE) with low mode numbers with the linear variational MHD model VariationalPerturbativeMHD +with linear=True, here without viscosity and resistivity, i.e. ideal linear MHD around a static (zero-flow) +equilibrium: shear Alfven waves together with the coupling to the magnetosonic waves through the pressure and density +perturbations. Its unknowns are the perturbations of the density, of the pressure and of the magnetic field with +respect to the equilibrium (which enters through `equil`), and the velocity (in H1vec). No energetic particles: the +TAE is excited by the initial condition, not driven. Same setup as examples/LinearMHD/lowmode_tae, so that the +frequencies of the variational and of the LinearMHD discretizations can be compared. + +Physics. Shear Alfven waves propagate along the magnetic field with the local frequency + omega_A(r) = |k_par(r)| v_A(r), v_A(r) = B(r) / sqrt(mu0 rho(r)), +where, for a mode with poloidal mode number m and toroidal mode number n in a torus of major radius R with safety +factor q(r), the parallel wave number is + k_par(r) = (n q(r) - m) / (q(r) R). +The frequencies of two counter-propagating modes with the same n and poloidal mode numbers m and m + 1 cross where +k_par,m = -k_par,m+1, i.e. at the radius r_c with + q(r_c) = (2m + 1) / (2n). +Toroidicity couples the two harmonics and opens a gap in the shear Alfven continuum at this crossing, in which the +toroidal Alfven eigenmode (TAE) exists with the frequency + omega_TAE = n / (2m + 1) * v_A(r_c) / R. +In Struphy units, v_A = v' v_unit and R = R' x_unit, so omega_TAE = n / (2m + 1) * v'(r_c) / R' * v_unit / x_unit, +with the velocity unit v_unit (the Alfven velocity of the base units, set by the model) and the length unit x_unit +(base units). + +Setup. Tokamak with circular flux surfaces (AdhocTorus): major radius R0 = 5, minor radius a = 1, toroidal field +B0 = 1, safety factor q(r) = 1.25 + r^2/a^2 (q = 1.5 at r = 0.5), flat number density n = 1 (as in the ITPA +benchmark of Koenies et al. 2018) and on-axis beta 0.0018, gamma = 5/3 (as in LinearMHD). The full torus is computed +(HollowTorus with tor_period = 1 and a hole of radius 0.1): the logical mode number -1 along eta3 is the physical +n = 1. The velocity is initialized with divergence-free m = 1 and m = 2 harmonics, localized around r = 0.5 (given as +2-form components, projected on H1vec) and then made discretely divergence-free in the sense of the scheme +(div Uv u = 0, as in LinearMHD where the commuting projection to H(div) is exactly divergence-free): the H1vec +projection is replaced by its closest (Mv norm) discretely divergence-free field, a change of the order of the +interpolation error (logged as delta_rel; switched off with --no-div-free). The projection is a CG solve without +preconditioner: on the 16 x 16 x 8 grid it stops at its iteration limit with |div u| reduced from 5.7e-5 to 2.6e-9. +The variational propagators run with model="linear" and their default +(Newton/Picard) solver options. The weights of the transport operators are taken at t^n: first order in time, like +LinearMHD (Crank-Nicolson substeps, Lie-Trotter splitting); the step is energy conserving. + +Expected result. With n = 1 and m = 1 the gap is at q(r_c) = 3/2, i.e. r_c = 0.5, where n = 1 and |B| = 1.0 +(flux-surface average), so v_A(r_c) = 1.0 v_unit and + omega_TAE = 1/3 * 1.0/5 = 0.067 v_unit / x_unit, +a period of about 94 time units. The run up to t = 600 covers about 6.4 periods; the pproc script asserts the +measured frequency within 10% of omega_TAE (the formula gives the center of the gap, whose relative width is about +the inverse aspect ratio r_c/R0 = 0.1) and shows the energies, the components of the fields and the radial profiles +of the (m, n) harmonics. The total energy en_tot (with the auxiliary variables p^ and B^, en_thermo and en_mag_2) is +conserved by the scheme. + +References: +- W. W. Heidbrink, "Basic physics of Alfven instabilities in toroidal plasmas", Phys. Plasmas 15, 055501 (2008). +- C. Z. Cheng, L. Chen and M. S. Chance, "High-n ideal and resistive shear Alfven waves in tokamaks", + Ann. Phys. 161, 21 (1985). +- A. Koenies et al., "Benchmark of gyrokinetic, kinetic MHD and gyrofluid codes for the linear calculation of fast + particle driven TAE dynamics", Nucl. Fusion 58, 126027 (2018). +""" + +import argparse +import logging +from pathlib import Path + +# Command line options +# `--out-folder` and `--sim-folder` set the output folder; `--num-elements` the grid. +# `--test-run` selects a small grid and a short run (three saved time steps) for the tests (test_examples.py). +# Unknown flags are ignored so that a driver can forward other parameters as well. +# no -h here when imported by another script, which has its own help +parser = argparse.ArgumentParser(add_help=__name__ == "__main__") +parser.add_argument( + "--out-folder", + default=str(Path(__file__).resolve().parent), + help="Folder in which the output folder is created (default: beside this file).", +) +parser.add_argument("--sim-folder", default="sim_data", help="Name of the output folder.") +parser.add_argument("--num-elements", type=int, nargs=3, default=[16, 16, 8], help="Number of elements per direction.") +parser.add_argument( + "--no-div-free", + action="store_true", + help="Keep the plain H1vec projection of the initial velocity (not exactly divergence-free in the scheme).", +) +parser.add_argument( + "--test-run", + action="store_true", + help="Small resolution (num_elements = [8, 8, 4]) and a short run with three saved time steps, for the tests " + "(test_examples.py).", +) +args, _ = parser.parse_known_args() +if args.test_run: + args.num_elements = [8, 8, 4] + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +# ------------------ +# Import Struphy API +# ------------------ + +import numpy as np + +from struphy import ( + BaseUnits, + DerhamOptions, + EnvironmentOptions, + ProfilingOptions, + Simulation, + Time, + domains, + equils, + grids, + perturbations, +) + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import VariationalPerturbativeMHD + +# adiabatic index (the one of LinearMHD) +gamma = 5.0 / 3.0 + +# Units +base_units = BaseUnits() + +# Model instance: linear and ideal (no viscosity, no resistivity), as LinearMHD +model = VariationalPerturbativeMHD(base_units=base_units, linear=True, with_viscosity=False, with_resistivity=False) + +# List all variables and decide whether to save their data +model.em_fields.b_field.save_data = True +model.mhd.density.save_data = True +model.mhd.velocity.save_data = True +model.mhd.pressure.save_data = True + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions(out_folders=args.out_folder, save_step=4, sim_folder=args.sim_folder) + +# Time stepping +dt = 0.5 +# --test-run: the initial, a middle and a final saved time step +time_opts = Time(dt=dt, Tend=2 * env.save_step * dt if args.test_run else 600.0) + +# Geometry +domain = domains.HollowTorus(a1=0.1, a2=1.0, R0=5.0, sfl=False, pol_period=1, tor_period=1) + +# Fluid equilibrium (enters the linear model through the projected equilibrium, not as a background of the variables) +equil = equils.AdhocTorus( + a=1.0, R0=5.0, B0=1.0, q_kind=0, p_kind=1, q0=1.25, q1=2.25, n1=0.0, n2=0.0, na=1.0, p1=0.95, p2=0.05, beta=0.0018 +) +# Grid +grid = grids.TensorProductGrid(num_elements=tuple(args.num_elements)) + +# Derham options +# Boundary conditions at the walls (flux surfaces, B0.n = 0), as in LinearMHD: b.n = 0 and u.n = 0 strongly. For the +# velocity in H1vec, Dirichlet acts (as for H(div)) only on the first component, u^1 = u.grad(eta1), so the tangential +# velocity is free (ideal slip wall). "free" would drop b.n = 0 and u.n = 0 as well. +derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(("dirichlet", "dirichlet"), None, None)) + +# Profiling options +profiling_opts = ProfilingOptions(use_line_profiler=False) + +# Simulation object +sim = Simulation( + model=model, + name=name, + description=description, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + profiling_opts=profiling_opts, +) + +# ------------------ +# Propagator options +# ------------------ +# linear=True needs the "linear" variants of the variational propagators (checked by the model) + +model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model="linear", gamma=gamma) +model.propagators.variat_pb.options = model.propagators.variat_pb.Options(model="linear", gamma=gamma) + +# ------------------ +# Initial conditions +# ------------------ +# Linear model: the equilibrium is not part of the variables (it enters through `equil`), the perturbations of the +# density, pressure and magnetic field start at zero. + +# Perturbations for (some) FEEC variables +# The m = 1 and m = 2 harmonics of the TAE (n = 1: logical mode number -1 along eta3), Gaussian around eta1_c (the +# radius r_c = 0.5 of the gap), as 2-form components: divergence free for the poloidal amplitude amp / (2 pi m) +# ("d_exp" is the derivative of "exp") +eta1_c = (0.5 - domain.params["a1"]) / (domain.params["a2"] - domain.params["a1"]) +sigma, amp = 0.15, 1e-3 +ms = (1, 2) +perturbation_radial = perturbations.TorusModesSin( + ms=ms, + ns=(-1, -1), + amps=(amp, amp), + pfuns=("exp", "exp"), + pfun_params=([eta1_c, sigma], [eta1_c, sigma]), + comp=0, + given_in_basis="2", +) +perturbation_poloidal = perturbations.TorusModesCos( + ms=ms, + ns=(-1, -1), + amps=tuple(amp / (2 * np.pi * m) for m in ms), + pfuns=("d_exp", "d_exp"), + pfun_params=([eta1_c, sigma], [eta1_c, sigma]), + comp=1, + given_in_basis="2", +) +model.mhd.velocity.add_perturbation(perturbation=perturbation_radial) +model.mhd.velocity.add_perturbation(perturbation=perturbation_poloidal) + +# Exactly divergence-free initial velocity: after the allocation (projection of the initial conditions), the velocity +# is projected onto the kernel of the discrete divergence div Uv of the variational scheme. Not on a restart, whose +# data overwrites the initial conditions. +if not args.no_div_free and not env.restart: + from struphy.feec.variational_utilities import project_div_free_velocity + + model_post_allocate = model.post_allocate + + def post_allocate(): + model_post_allocate() + info = project_div_free_velocity(model.mhd.velocity.spline.vector, sim.derham, sim.mass_ops, sim.basis_ops) + logging.getLogger("struphy").info( + f"Divergence-free initial velocity: |div u| {info['div_before']:.3e} -> {info['div_after']:.3e}, " + f"delta_rel = {info['delta_rel']:.3e} ({info['niter']} CG iterations, success = {info['success']})." + ) + + model.post_allocate = post_allocate + +if __name__ == "__main__": + sim.run(profiling_activated=True) diff --git a/examples/VariationalPerturbativeLinearMHD/with_p/lowmode_tae/pproc_lowmode_tae.py b/examples/VariationalPerturbativeLinearMHD/with_p/lowmode_tae/pproc_lowmode_tae.py new file mode 100644 index 000000000..00ce078f8 --- /dev/null +++ b/examples/VariationalPerturbativeLinearMHD/with_p/lowmode_tae/pproc_lowmode_tae.py @@ -0,0 +1,196 @@ +"""Post-process and plot the low-mode TAE run (VariationalPerturbativeMHD, linear). + +Run as ``python pproc_lowmode_tae.py [path_out] [flags of params_lowmode_tae.py] [--show]``. The standard +plots of the examples (slices in the physical planes and 3D views of the fields) are made at the first, middle and +last saved time, together with the energy budget and the (m, n) harmonics of the radial components. Vector fields +are shown by their components in the basis of their initial condition (``given_in_basis``, else their FEEC form, +see ``Output.field_basis``), as functions of the physical coordinates; the 3D views show the scalar fields and the +first components (radial eta1 in a logical basis). The measured TAE frequency is compared with the analytical one +and saved; a deviation of more than 10% fails with an ``AssertionError`` after the figures are saved (not with +``--test-run``). Figures are written into ``/results``. The energy and harmonics plots use plasma-plots +(``pip install "struphy[pproc]"``). +""" + +import argparse +import os +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt + +from struphy import Output +from struphy.post_processing.plot_3d import plot_field_3d +from struphy.post_processing.slices import plot_physical_slices + +SCRIPT_DIR = Path(__file__).resolve().parent + +# The flags of the run (parsed from the same command line), for the default output folder. +sys.path.insert(0, str(SCRIPT_DIR)) +from params_lowmode_tae import args, parser # noqa: E402 + +DEFAULT_OUTPUT = Path(args.out_folder) / args.sim_folder + +# standard plots: (key, label, species, variable, vector) per field +FIELDS = [ + ("u", "U", "mhd", "velocity", True), + ("b", "B", "em_fields", "b_field", True), + ("n", "n", "mhd", "density", False), + ("p", "p", "mhd", "pressure", False), +] +# (m, n) harmonics of the radial (eta1) components +HARMONICS = [("u_r", "mhd/velocity"), ("b_r", "em_fields/b_field")] +ENERGIES = ("en_U", "en_thermo", "en_mag_1", "en_mag_2") # energy scalars shown next to en_tot + +# the TAE: toroidal mode number n (the logical mode N_LOGICAL along eta3 times tor_period of the domain) in the gap +# between the poloidal mode numbers M and M + 1, see the description of the params file +M_TAE = 1 +N_LOGICAL = -1 +# relative tolerance of the measured TAE frequency: the formula gives the center of the gap, whose relative width +# is about the inverse aspect ratio r_c/R0 = 0.1 at the gap +TAE_RTOL = 0.1 + + +def tae_frequency(run, basis): + """Measured and expected frequency of the TAE. + + Expected: omega_TAE = n / (2m + 1) * v_A(r_c) / R0 at the radius r_c of the gap, q(r_c) = (2m + 1) / (2n), for the + q profile q0 + (q1 - q0) r^2/a^2 of the AdhocTorus; v_A = |B| / sqrt(n) (in units of the Alfven velocity) with the + flux-surface averages of |B| and of the number density. Measured: the dominant frequency of the (M_TAE, n) + harmonic of the radial velocity on the flux surface r_c over the run, projected on its initial phase so that a + standing or a propagating mode gives the same frequency. + """ + equil, params = run.equil, run.domain.params + q0, q1, a, R0 = (equil.params[k] for k in ("q0", "q1", "a", "R0")) + n = abs(N_LOGICAL) * params["tor_period"] + q_c = (2 * M_TAE + 1) / (2 * n) + r_c = a * np.sqrt((q_c - q0) / (q1 - q0)) + eta1 = np.array([(r_c - params["a1"]) / (params["a2"] - params["a1"])]) + eta2, eta3 = np.linspace(0, 1, 64, endpoint=False), np.linspace(0, 1, 8, endpoint=False) + abs_b = np.mean(equil.absB0(eta1, eta2, eta3)) + density = np.mean(equil.n0(eta1, eta2, eta3)) + expected = n / (2 * M_TAE + 1) * abs_b / np.sqrt(density) / R0 + + # the radial component is the first one of a logical basis (the FEEC form for Cartesian initial conditions) + basis = run.field_form("mhd/velocity") if basis == "physical" else basis + u_r = run.evaluate("mhd/velocity", representation=basis, eta1=eta1, eta2=eta2, eta3=eta3, component=0) + modes = u_r.squeeze("eta1", drop=True).plasma.analysis.mode_spectrum() + coefficient = modes.sel(m=M_TAE, n=N_LOGICAL) + initial = complex(coefficient.isel(t=0)) + # projected on the initial phase; without initial amplitude in this harmonic its real part + phase = initial / abs(initial) if abs(initial) > 0 else 1.0 + signal = (coefficient * np.conj(phase)).real + # a frequency needs a time series (not e.g. the single saved step of a one-step test) + if signal.sizes["t"] < 3: + measured = float("nan") + else: + measured = float(signal.plasma.analysis.spectral_peaks(n_peaks=1).omega_refined[0]) + print( + f"TAE (m = {M_TAE}, {M_TAE + 1}, n = {n}) at r_c = {r_c:.3g}: omega = {measured:.4g}, expected {expected:.4g}" + ) + return measured, expected + + +def minor_radius(run): + """The minor radius of the HollowTorus as a function of eta1.""" + a1, a2 = run.domain.params["a1"], run.domain.params["a2"] + return lambda eta1: a1 + (a2 - a1) * eta1 + + +def finish(figs, results_dir, save_figs): + """Save the figures to ``results_dir`` and close them (``save_figs``), or keep them open for ``plt.show()``.""" + if save_figs: + for fig_name, fig in figs.items(): + fig.savefig(os.path.join(results_dir, f"{fig_name}.png")) + plt.close(fig) + + +def main(path_out=DEFAULT_OUTPUT, save_figs=True): + run = Output(path_out) + snapshots = run.standard_snapshots() + # VTK files (for the 3D views) of the standard snapshots only + run.pproc(create_vtk=True, physical=True, vtk_steps=snapshots) + times = [float(run.time[n]) for n in snapshots] + + # The profiling packaging picks these files up from here and uploads them with the run. + results_dir = os.path.join(run.path_out, "results") + os.makedirs(results_dir, exist_ok=True) + + # the basis of the components of each vector field (None for a scalar field) + bases = {f"{s}/{v}": run.field_basis(f"{s}/{v}") if vector else None for _, _, s, v, vector in FIELDS} + + # the frequency of the TAE, compared with the analytical value at the end (after the figures) + measured, expected = tae_frequency(run, bases["mhd/velocity"]) + np.save(os.path.join(results_dir, "omega_tae.npy"), measured) + np.save(os.path.join(results_dir, "omega_tae_expected.npy"), expected) + + finish({"energies": run.plot.energies(parts=ENERGIES).fig}, results_dir, save_figs) + for n, t in zip(snapshots, times): + # one field at a time, saved right away, so that few figures are open at once + for key, label, species, variable, _ in FIELDS: + name = f"{species}/{variable}" + field = (key, f"{label} at t = {t:.3g}", name, n, None, bases[name]) + slices = plot_physical_slices(run, run.domain, run.grid.num_elements, [field], components=True) + # e.g. u1_slices_rz_t1.5 for the first component of U in the R-z plane at t = 1.5 + finish({f"{fig_name}_t{t:.6g}": fig for fig_name, fig in slices.items()}, results_dir, save_figs) + + figs = {} + for key, name in HARMONICS: + # the logical radial component: in the basis of the initial condition unless it is Cartesian + basis = bases[name] if bases[name] != "physical" else run.field_form(name) + data = run.evaluate(name, representation=basis, t=n, component=0) + data.name = key + result = data.plasma.plot.mode_profiles( + scale=(1, run.domain.params["tor_period"]), + x_of=minor_radius(run), + xlabel="minor radius r", + top=4, + ) + figs[f"{key}_harmonics_t{t:.6g}"] = result.fig + finish(figs, results_dir, save_figs) + + # 3D views of the scalar fields and of the first (radial eta1, in a logical basis) components, saved as + # _3d_t