diff --git a/doc/sections/subsections/models-all.rst b/doc/sections/subsections/models-all.rst index d779cc5cc..059024323 100644 --- a/doc/sections/subsections/models-all.rst +++ b/doc/sections/subsections/models-all.rst @@ -26,6 +26,11 @@ Full-orbit Vlasov :undoc-members: :show-inheritance: +.. autoclass:: struphy.models.VlasovPoissonElectrostaticAdiabatic + :members: doc_pde, doc_scalar_quantities, doc_normalization, doc_discretization, doc_long_description, doc_examples, doc_use_cases, doc_cannot_be_used_for + :undoc-members: + :show-inheritance: + Gyrokinetic ^^^^^^^^^^^ diff --git a/doc/sections/subsections/propagators-avail.rst b/doc/sections/subsections/propagators-avail.rst index 7389ab748..7fd276e1c 100644 --- a/doc/sections/subsections/propagators-avail.rst +++ b/doc/sections/subsections/propagators-avail.rst @@ -16,6 +16,11 @@ Field solvers :exclude-members: options, allocate :show-inheritance: +.. automodule:: struphy.propagators.adiabatic_poisson_solve + :members: + :exclude-members: options, allocate + :show-inheritance: + .. automodule:: struphy.propagators.curl_curl_solve :members: :exclude-members: options, allocate diff --git a/examples/DriftKineticElectrostaticAdiabatic/slab_drift_wave/params_slab_drift_wave.py b/examples/DriftKineticElectrostaticAdiabatic/slab_drift_wave/params_slab_drift_wave.py new file mode 100644 index 000000000..ebb8fd068 --- /dev/null +++ b/examples/DriftKineticElectrostaticAdiabatic/slab_drift_wave/params_slab_drift_wave.py @@ -0,0 +1,246 @@ +"""Kinetic drift wave in a slab with drift-kinetic ions and adiabatic electrons (DriftKineticElectrostaticAdiabatic); +the same setup as examples/VlasovPoissonElectrostaticAdiabatic/slab_drift_wave, for comparison with full-orbit ions.""" + +# ----------------------------- +# Description of the simulation +# ----------------------------- + +name = "Slab drift wave (DriftKineticElectrostaticAdiabatic)" +description = """ +Drift wave in a slab with a density gradient, with drift-kinetic (guiding-center, 5D) ions and adiabatic electrons +(DriftKineticElectrostaticAdiabatic). The physical setup is the one of the slab_drift_wave example of the full-orbit +model VlasovPoissonElectrostaticAdiabatic, so that the two models can be compared directly. + +Physics. The guiding-center distribution f(X, v_par, mu, t) is advected by the parallel streaming, the E x B drift and +the parallel acceleration by E_par (uniform B0 = e_z, no curvature). The potential follows from the quasi-neutrality +equation with the ion polarization density and adiabatic electrons, + -eps div_perp(n0 grad_perp phi) + n0 / (Z eps T0) phi = n_gc - n0. +As in the full-orbit example, the electrons respond adiabatically to all modes (which_geometry=None, no +flux-surface average): the two examples use the same electron model. In the shearless slab the electrons +could physically not respond to modes with k_z = 0 (no motion along B), which would be the average along the field +lines (which_geometry="cylindrical", average over z); the mode studied here has k_z != 0, where both agree. +A perturbation with wave vector +k = (k_x, k_y, k_z) in a density gradient kappa_n = -d ln n0 / dx propagates in the electron diamagnetic direction (+y). +In the fluid limit + omega = omega_*e / (1 + k_perp^2 rho_s^2), omega_*e = epsilon k_y Z T_e kappa_n, rho_s = epsilon sqrt(Z T_e), +the polarization term being the long-wavelength limit (k_perp rho_i << 1) of the finite-Larmor-radius response of the +full-orbit model. The pproc solves the kinetic dispersion relations of both models, which differ at order +b = k_perp^2 rho_i^2 since the drift-kinetic model has no gyro-average. + +Setup. Cuboid of size Lx x Ly x Lz = 2 x 2 pi/3 x 2 pi/0.05 with B0 = e_z (HomogenSlab), epsilon = 0.1, Z = 1, +T_i = T_e = 1. The density n0 ~ exp(-kappa_n x) has the uniform gradient kappa_n = 1; the guiding-center Maxwellian +n0(X) M(v_par, mu) is an exact equilibrium. The fluid equilibrium (ion density and electron pressure n T_e of the +adiabatic term, n0 / |B0|^2 of the polarization) is overwritten accordingly. The initial density perturbation +1e-2 n0 sin(pi x/Lx) cos(k_y y + k_z z) has k_x = pi/Lx, k_y = 3, k_z = 0.05, i.e. b = 0.115. Markers leaving at x = 0 or +x = Lx are removed and phi vanishes there; y and z are periodic. Delta-f with the control variate, explicit +guiding-center pushers, Strang splitting with dt = 0.25 (no cyclotron motion to resolve, unlike the full-orbit model +with dt = 0.05), grid of 8 x 8 x 4 cubic elements, end time Tend = 25. + +Expected result. The initial (k_y, k_z) mode on the plane x = Lx/2 oscillates with omega close to the root of the +drift-kinetic dispersion relation, omega = 0.287 (full-orbit model: 0.262). The pproc measures the frequency from the +phase of the Fourier coefficient of phi, plots phi in the (x, y) plane at the first, middle and last saved time and, +with --compare , overlays the mode phase of the full-orbit example. + +References: +- W. Horton, "Drift waves and transport", Rev. Mod. Phys. 71, 735 (1999). +- A. J. Brizard and T. S. Hahm, "Foundations of nonlinear gyrokinetic theory", Rev. Mod. Phys. 79, 421 (2007). +""" + +import argparse +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, few markers 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 the pproc 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_1", help="Name of the output folder.") +parser.add_argument("--num-elements", type=int, nargs=3, default=[8, 8, 4], help="Number of elements per direction.") +parser.add_argument("--ppc", type=int, default=100, help="Markers per cell.") +parser.add_argument( + "--test-run", + action="store_true", + help="Small resolution (num_elements = [8, 4, 4]), few markers and three saved time steps, for the tests " + "(test_examples.py).", +) +args, _ = parser.parse_known_args() +if args.test_run: + args.num_elements = [8, 4, 4] + args.ppc = 5 + +import logging + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +import cunumpy as xp + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, +) +from struphy.initial.base import GenericPerturbation +from struphy.models import DriftKineticElectrostaticAdiabatic + +# --------------------- +# Physical parameters (also used by the pproc) +# --------------------- + +epsilon = 0.1 +kappa_n = 1.0 +Te = Ti = 1.0 +Lx, ky, kz = 2.0, 3.0, 0.05 +Ly, Lz = 2 * xp.pi / ky, 2 * xp.pi / kz +kx = xp.pi / Lx +amp = 1e-2 + +# --------------------- +# Instance of the model +# --------------------- + +base_units = BaseUnits(kBT=1.0) +model = DriftKineticElectrostaticAdiabatic(base_units=base_units, epsilon=epsilon) + +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = False + +# -------------------------- +# Instance of the simulation +# -------------------------- + +env = EnvironmentOptions(out_folders=args.out_folder, sim_folder=args.sim_folder, restart=False) + +# Time stepping: no cyclotron motion to resolve +dt = 0.25 +# --test-run: the initial, a middle and a final saved time step +time_opts = Time(dt=dt, Tend=2 * dt if args.test_run else 25.0, split_algo="Strang") + +# Geometry +domain = domains.Cuboid(l1=0.0, r1=Lx, l2=0.0, r2=Ly, l3=0.0, r3=Lz) + +# Fluid equilibrium: uniform B0 = e_z; density and pressure overwritten below +equil = equils.HomogenSlab(B0x=0.0, B0y=0.0, B0z=1.0) + +# Grid +grid = grids.TensorProductGrid(num_elements=tuple(args.num_elements)) + +# Derham options: phi = 0 at the walls in x +derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(("dirichlet", "dirichlet"), None, None)) + +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, +) + +# ------------------- +# Particle parameters +# ------------------- + +# the magnetic moment of the markers is sampled for the field strength B0 (default 2), here |B0| = 1 +loading_params = LoadingParameters(ppc=args.ppc, seed=1234, B0=1.0) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters(bc=("remove", "periodic", "periodic")) +sorting_params = SortingParameters(do_sort=True, boxes_per_dim=tuple(max(n // 2, 1) for n in args.num_elements)) +saving_params = SavingParameters() + +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize=1.0, +) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options(algo="explicit", evaluate_e_field=True) +# no flux-surface average: the electrons respond to all modes, as in the full-orbit example (with "cylindrical", +# the modes with k_z = 0 have no electron response; their ion convective cells, omega ~ 3, made the explicit push +# unstable at dt = 0.25) +model.propagators.gc_poisson.options = model.propagators.gc_poisson.Options(which_geometry=None) +model.propagators.push_gc_para.options = model.propagators.push_gc_para.Options(algo="explicit", evaluate_e_field=True) + +# ------------------ +# Initial conditions +# ------------------ + +# n0 with mean one over the box (the guiding-center Maxwellian n0(X) M is an exact equilibrium in the slab) +C_n = kappa_n * Lx / (2 * xp.sinh(kappa_n * Lx / 2)) + + +def n0_x(x): + return C_n * xp.exp(-kappa_n * (x - Lx / 2)) + + +def _etas(etas): + """Logical coordinates from either (e1, e2, e3) or a single (3, ...) array.""" + if len(etas) == 1: + return etas[0][0], etas[0][1], etas[0][2] + return etas + + +def n_init(*etas): + e1, _, _ = _etas(etas) + return n0_x(Lx * e1) + + +def pert_func(*etas): + e1, e2, e3 = _etas(etas) + x, y, z = Lx * e1, Ly * e2, Lz * e3 + return amp * n0_x(x) * xp.sin(kx * x) * xp.cos(ky * y + kz * z) + + +# the adiabatic term uses n0 / T0 of the fluid equilibrium, with T0 = p / n0 = T_e, the polarization n0 / |B0|^2 +def n0_xyz(x, y, z): + return n0_x(x) + + +def p_xyz(x, y, z): + return n0_x(x) * Te + + +equil.n_xyz = n0_xyz +equil.p_xyz = p_xyz + +# the magnetic moment mu = v_perp^2 / (2 |B0|) needs the field strength of the slab, |B0| = 1 (the default is 2), +# the same as for the loading of the markers +vth = xp.sqrt(Ti) +background = maxwellians.GyroMaxwellian2D(n=(n_init, None), vth_para=(vth, None), vth_perp=(vth, None), B0=1.0) +model.kinetic_ions.var.add_background(background) +init = maxwellians.GyroMaxwellian2D( + n=(n_init, GenericPerturbation(pert_func)), vth_para=(vth, None), vth_perp=(vth, None), B0=1.0 +) +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + sim.run() diff --git a/examples/DriftKineticElectrostaticAdiabatic/slab_drift_wave/pproc_slab_drift_wave.py b/examples/DriftKineticElectrostaticAdiabatic/slab_drift_wave/pproc_slab_drift_wave.py new file mode 100644 index 000000000..44da3046c --- /dev/null +++ b/examples/DriftKineticElectrostaticAdiabatic/slab_drift_wave/pproc_slab_drift_wave.py @@ -0,0 +1,159 @@ +"""Post-process and plot the slab drift wave of the drift-kinetic model. + +Run as ``python pproc_slab_drift_wave.py [sim_folder] [flags of params_slab_drift_wave.py] [--compare FOLDER] [--show]``; +the flags must be those of the run. The figures are written into ``/results`` (``--show`` shows them +instead): the amplitude and phase of the initial (k_y, k_z) mode against time with the measured and predicted +frequencies, the field energy, and phi in the (x, y) plane at the first, middle and last saved time. With +``--compare FOLDER``, the output folder of a run of examples/VlasovPoissonElectrostaticAdiabatic/slab_drift_wave +(same physical setup), the mode of the full-orbit run is overlaid. +""" + +import argparse +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt +from scipy.optimize import fsolve +from scipy.special import ive, wofz + +from struphy import Output + +SCRIPT_DIR = Path(__file__).resolve().parent + +# The setup of the run (parsed from the same command line). +sys.path.insert(0, str(SCRIPT_DIR)) +import params_slab_drift_wave as params # noqa: E402 + +DEFAULT_OUTPUT = Path(params.args.out_folder) / params.args.sim_folder + + +def plasma_Z(zeta): + """Plasma dispersion function.""" + return 1j * np.sqrt(np.pi) * wofz(zeta) + + +def drift_wave_roots(n_harmonics=8): + """Roots of the local kinetic dispersion relations (quasi-neutral, eta_i = 0) of the full-orbit model and of + the drift-kinetic model, for the wave vector of the initial perturbation (as in the pproc of the full-orbit + example).""" + eps, Te, Ti, kap = params.epsilon, params.Te, params.Ti, params.kappa_n + kperp2 = params.kx**2 + params.ky**2 + b = eps**2 * kperp2 * Ti + w_si = -eps * params.ky * Ti * kap + s = np.sqrt(2) * abs(params.kz) * np.sqrt(Ti) + + def D_vlasov(w): + tot = 1.0 + 0j + for n in range(-n_harmonics, n_harmonics + 1): + tot += ive(n, b) * (w - w_si) / s * plasma_Z((w - n / eps) / s) + return 1 / Te + tot / Ti + + def D_dk(w): + return b / Ti + 1 / Te + (1 + (w - w_si) / s * plasma_Z(w / s)) / Ti + + def root(D): + sol = fsolve(lambda x: [D(x[0] + 1j * x[1]).real, D(x[0] + 1j * x[1]).imag], [-w_si * Te / Ti, 0.0]) + return sol[0] + 1j * sol[1] + + return root(D_vlasov), root(D_dk) + + +def mode_coefficient(run, n_points=32): + """Fourier coefficient c(t) of phi on the plane x = Lx/2 for the initial mode cos(k_y y + k_z z), which is one + period in eta2 and eta3; for a wave cos(k_y y + k_z z - omega t), c(t) = A/2 exp(-i omega t).""" + e = np.linspace(0.0, 1.0, n_points, endpoint=False) + phi = run.evaluate("em_fields/phi", eta1=0.5, eta2=e, eta3=e) + data = np.asarray(phi.transpose("t", "eta2", "eta3").values) + phase = np.exp(-2j * np.pi * (e[:, None] + e[None, :])) + return phi.t.values, np.mean(data * phase[None], axis=(1, 2)) + + +def measured_frequency(t, coeff): + """Angular frequency from a linear fit of the unwrapped phase of c(t) = A/2 exp(-i omega t).""" + if len(t) < 3: + return np.nan + return -np.polyfit(t, np.unwrap(np.angle(coeff)), 1)[0] + + +def saver(run, save_figs): + """A function saving figures as ``.png`` into ``/results`` and closing them (``save_figs``).""" + results_dir = Path(run.path_out) / "results" + results_dir.mkdir(exist_ok=True) + + def save(figs): + if save_figs: + for name, fig in figs.items(): + fig.savefig(results_dir / f"{name}.png") + plt.close(fig) + + return save + + +def main(path_out=DEFAULT_OUTPUT, save_figs=True, compare=None): + run = Output(path_out) + snapshots = run.standard_snapshots() + save = saver(run, save_figs) + + # the initial (k_y, k_z) mode on the plane x = Lx/2 + t, coeff = mode_coefficient(run) + omega_num = measured_frequency(t, coeff) + omega_v, omega_dk = drift_wave_roots() + print(f"measured omega = {omega_num:.4f}, drift-kinetic: {omega_dk:.4f}, full orbit: {omega_v:.4f}") + + fig, axs = plt.subplots(2, 1, sharex=True) + axs[0].plot(t, np.abs(coeff), label="drift-kinetic run") + axs[1].plot(t, np.unwrap(np.angle(coeff)), label="drift-kinetic run") + title = rf"$\omega$ drift-kinetic run {omega_num:.3f}" + + # the same mode of a run of the full-orbit example + if compare is not None: + t_v, coeff_v = mode_coefficient(Output(compare)) + omega_num_v = measured_frequency(t_v, coeff_v) + print(f"measured omega of the full-orbit run = {omega_num_v:.4f}") + axs[0].plot(t_v, np.abs(coeff_v), label="full-orbit run") + axs[1].plot(t_v, np.unwrap(np.angle(coeff_v)), label="full-orbit run") + title += rf", full-orbit run {omega_num_v:.3f}" + + axs[1].plot(t, np.angle(coeff[0]) - omega_dk.real * t, ":", label=f"drift-kinetic theory {omega_dk.real:.3f}") + axs[1].plot(t, np.angle(coeff[0]) - omega_v.real * t, "--", label=f"full-orbit theory {omega_v.real:.3f}") + axs[0].set(ylabel="mode amplitude", yscale="log") + axs[0].legend() + axs[1].set(xlabel="time", ylabel="mode phase") + axs[1].legend() + fig.suptitle(title) + save({"mode": fig}) + + energy = run.evaluate("scalars", variables="en_phi")["en_phi"] + fig, ax = plt.subplots() + ax.plot(energy.t, energy) + ax.set(xlabel="time", yscale="log", title="Field energy en_phi") + save({"en_phi": fig}) + + # phi in the (x, y) plane at the first, middle and last saved time + plane = run.evaluate("em_fields/phi", t=list(snapshots), eta3=0.0) + grid = plane.plot(x="eta2", y="eta1", col="t") + save({"phi_xy": grid.fig}) + + if not save_figs: + plt.show() + + +if __name__ == "__main__": + # the flags of the run are parsed by params_slab_drift_wave.py, the rest here + _, other_flags = params.parser.parse_known_args() + parser = argparse.ArgumentParser(description="Plot a saved simulation run.") + parser.add_argument( + "path_out", + nargs="?", + default=DEFAULT_OUTPUT, + help="Simulation output folder (default: the output folder of the run flags, beside this script)", + ) + parser.add_argument( + "--compare", + default=None, + help="Output folder of a run of examples/VlasovPoissonElectrostaticAdiabatic/slab_drift_wave to overlay.", + ) + parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + args = parser.parse_args(other_flags) + main(args.path_out, save_figs=not args.show, compare=args.compare) diff --git a/examples/VlasovPoissonElectrostaticAdiabatic/slab_drift_wave/params_slab_drift_wave.py b/examples/VlasovPoissonElectrostaticAdiabatic/slab_drift_wave/params_slab_drift_wave.py new file mode 100644 index 000000000..0413bc942 --- /dev/null +++ b/examples/VlasovPoissonElectrostaticAdiabatic/slab_drift_wave/params_slab_drift_wave.py @@ -0,0 +1,252 @@ +"""Kinetic drift wave in a slab with full-orbit ions and adiabatic electrons +(VlasovPoissonElectrostaticAdiabatic); the full-orbit counterpart of the drift-kinetic model +DriftKineticElectrostaticAdiabatic.""" + +# ----------------------------- +# Description of the simulation +# ----------------------------- + +name = "Slab drift wave (VlasovPoissonElectrostaticAdiabatic)" +description = """ +Drift wave in a slab with a density gradient, with full-orbit (6D) ions and adiabatic electrons +(VlasovPoissonElectrostaticAdiabatic). + +Physics. The ions follow full Lorentz-force orbits, + dx/dt = v, dv/dt = (1/epsilon) (-grad phi + v x B0), +and the electrons respond adiabatically, n_e = Z n0 (1 + e phi / T_e). In the quasi-neutral limit the potential +follows from + n0 phi / (Z epsilon T0) = int f d^3v - n0, +i.e. the electrons respond adiabatically to all modes (no flux-surface average, as in the drift-kinetic slab +example examples/DriftKineticElectrostaticAdiabatic/slab_drift_wave with which_geometry=None). +A perturbation with wave vector k = (k_x, k_y, k_z) in a density gradient kappa_n = -d ln n0 / dx (B0 = e_z) +propagates in the electron diamagnetic direction (+y) with the drift-wave frequency. In the fluid limit with +finite-Larmor-radius effects, + omega = Gamma_0(b) omega_*e / (1 + (Z T_e / T_i) (1 - Gamma_0(b))), omega_*e = epsilon k_y Z T_e kappa_n, +with b = k_perp^2 rho_i^2, rho_i = epsilon sqrt(T_i) and Gamma_0(b) = I_0(b) exp(-b). Parallel ion dynamics and ion +Landau damping enter through k_z. The pproc solves the full kinetic dispersion relation (sum over cyclotron +harmonics) and the one of the drift-kinetic model, which differ by O(b) since the drift-kinetic model has no +gyro-average. + +Setup. Cuboid of size Lx x Ly x Lz = 2 x 2 pi/3 x 2 pi/0.05 with B0 = e_z (HomogenSlab), epsilon = 0.1, Z = 1, +T_i = T_e = 1. The density n0 ~ exp(-kappa_n x) has the uniform gradient kappa_n = 1. A Vlasov equilibrium must be a +function of the guiding-center position X = x + epsilon v_y; for the exponential profile + f0 = n0(X) M(v) = n0(x) exp(epsilon^2 kappa_n^2 T_i / 2) M(v - u), u = -epsilon kappa_n T_i e_y, +a Maxwellian shifted by the ion diamagnetic velocity (n0(x) M(v) alone is not stationary and would fill the +delta-f weights with noise of relative size epsilon kappa_n). The fluid equilibrium (ion density and electron +pressure n T_e of the adiabatic term) is overwritten accordingly. The initial +density perturbation 1e-2 n0 sin(pi x/Lx) cos(k_y y + k_z z) has k_x = pi/Lx, k_y = 3, k_z = 0.05, i.e. b = 0.115. +Markers are reflected at x = 0 and x = Lx, where phi vanishes; y and z are periodic. The reflection v_x -> -v_x keeps +X = x + epsilon v_y and |v|, so f0 stays an equilibrium. Delta-f with the control variate, quasi-neutral Poisson solve, +Strang splitting with dt = 0.05 (12 steps per cyclotron period), grid of 8 x 8 x 4 cubic elements, end time Tend = 25 +(about one wave period; later the mode sinks into the PIC noise). In the quasi-neutral limit the ion-acoustic branch +omega = k c_s has no Debye cutoff, so the explicit time stepping needs dt c_s k_max < 2 with k_max ~ pi / (cell size); +with 16 elements in x the run is unstable for dt = 0.05. + +Expected result. The initial (k_y, k_z) mode propagates in +y with omega close to the root of the full-orbit +dispersion relation, omega = 0.262 (drift-kinetic model: 0.287; a run gave 0.2619). The pproc measures the frequency +from the phase of the (k_y, k_z) Fourier coefficient of phi on the plane x = Lx/2 and plots phi in the (x, y) plane +at the first, middle and last saved time. + +References: +- W. Horton, "Drift waves and transport", Rev. Mod. Phys. 71, 735 (1999). +- N. A. Krall and A. W. Trivelpiece, "Principles of Plasma Physics", McGraw-Hill (1973), chapter 9. +""" + +import argparse +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, few markers 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 the pproc 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_1", help="Name of the output folder.") +parser.add_argument("--num-elements", type=int, nargs=3, default=[8, 8, 4], help="Number of elements per direction.") +parser.add_argument("--ppc", type=int, default=100, help="Markers per cell.") +parser.add_argument( + "--test-run", + action="store_true", + help="Small resolution (num_elements = [8, 4, 4]), few markers and three saved time steps, for the tests " + "(test_examples.py).", +) +args, _ = parser.parse_known_args() +if args.test_run: + args.num_elements = [8, 4, 4] + args.ppc = 5 + +import logging + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +import cunumpy as xp + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, +) +from struphy.initial.base import GenericPerturbation +from struphy.models import VlasovPoissonElectrostaticAdiabatic + +# --------------------- +# Physical parameters (also used by the pproc) +# --------------------- + +epsilon = 0.1 +kappa_n = 1.0 +Te = Ti = 1.0 +Lx, ky, kz = 2.0, 3.0, 0.05 +Ly, Lz = 2 * xp.pi / ky, 2 * xp.pi / kz +kx = xp.pi / Lx +amp = 1e-2 + +# --------------------- +# Instance of the model +# --------------------- + +base_units = BaseUnits(kBT=1.0) +model = VlasovPoissonElectrostaticAdiabatic(base_units=base_units, epsilon=epsilon) + +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = False + +# -------------------------- +# Instance of the simulation +# -------------------------- + +env = EnvironmentOptions(out_folders=args.out_folder, sim_folder=args.sim_folder, restart=False) + +# Time stepping: the drift (PushEta) and the Poisson solve come first, so both E-kicks of a Strang step +# use the potential of the current marker positions +dt = 0.05 +# --test-run: the initial, a middle and a final saved time step +time_opts = Time(dt=dt, Tend=2 * dt if args.test_run else 25.0, split_algo="Strang") + +# Geometry +domain = domains.Cuboid(l1=0.0, r1=Lx, l2=0.0, r2=Ly, l3=0.0, r3=Lz) + +# Fluid equilibrium: uniform B0 = e_z; density and pressure overwritten below +equil = equils.HomogenSlab(B0x=0.0, B0y=0.0, B0z=1.0) + +# Grid +grid = grids.TensorProductGrid(num_elements=tuple(args.num_elements)) + +# Derham options: phi = 0 at the walls in x +derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(("dirichlet", "dirichlet"), None, None)) + +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, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=args.ppc, seed=1234) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters(bc=("reflect", "periodic", "periodic")) +sorting_params = SortingParameters(do_sort=True, boxes_per_dim=tuple(max(n // 2, 1) for n in args.num_elements)) +saving_params = SavingParameters() + +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize=1.0, +) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_eta.options = model.propagators.push_eta.Options() +model.propagators.poisson.options = model.propagators.poisson.Options(quasi_neutral=True) +model.propagators.push_v_E.options = model.propagators.push_v_E.Options() +model.propagators.push_vxb.options = model.propagators.push_vxb.Options() + +# ------------------ +# Initial conditions +# ------------------ + +# equilibrium f0 = n0(X) M(v) with X = x + epsilon v_y: for the exponential profile a shifted Maxwellian with the +# ion diamagnetic velocity u_y and density n0(x) exp(epsilon^2 kappa_n^2 T_i / 2), of mean one over the box +u_dia = -epsilon * kappa_n * Ti +C_n = kappa_n * Lx / (2 * xp.sinh(kappa_n * Lx / 2)) + + +def n0_x(x): + return C_n * xp.exp(-kappa_n * (x - Lx / 2)) + + +def _etas(etas): + """Logical coordinates from either (e1, e2, e3) or a single (3, ...) array.""" + if len(etas) == 1: + return etas[0][0], etas[0][1], etas[0][2] + return etas + + +def n_init(*etas): + e1, _, _ = _etas(etas) + return n0_x(Lx * e1) + + +def pert_func(*etas): + e1, e2, e3 = _etas(etas) + x, y, z = Lx * e1, Ly * e2, Lz * e3 + return amp * n0_x(x) * xp.sin(kx * x) * xp.cos(ky * y + kz * z) + + +# the adiabatic term of the Poisson solve uses n0 / T0 of the fluid equilibrium, with T0 = p / n0 = T_e +def n0_xyz(x, y, z): + return n0_x(x) + + +def p_xyz(x, y, z): + return n0_x(x) * Te + + +equil.n_xyz = n0_xyz +equil.p_xyz = p_xyz + +vth = xp.sqrt(Ti) +background = maxwellians.Maxwellian3D( + n=(n_init, None), u2=(u_dia, None), vth1=(vth, None), vth2=(vth, None), vth3=(vth, None) +) +model.kinetic_ions.var.add_background(background) +init = maxwellians.Maxwellian3D( + n=(n_init, GenericPerturbation(pert_func)), u2=(u_dia, None), vth1=(vth, None), vth2=(vth, None), vth3=(vth, None) +) +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + sim.run() diff --git a/examples/VlasovPoissonElectrostaticAdiabatic/slab_drift_wave/pproc_slab_drift_wave.py b/examples/VlasovPoissonElectrostaticAdiabatic/slab_drift_wave/pproc_slab_drift_wave.py new file mode 100644 index 000000000..59fcff386 --- /dev/null +++ b/examples/VlasovPoissonElectrostaticAdiabatic/slab_drift_wave/pproc_slab_drift_wave.py @@ -0,0 +1,147 @@ +"""Post-process and plot the slab drift wave. + +Run as ``python pproc_slab_drift_wave.py [sim_folder] [flags of params_slab_drift_wave.py] [--show]``; the flags +must be those of the run. The figures are written into ``/results`` (``--show`` shows them instead): +the amplitude and phase of the initial (k_y, k_z) mode against time with the measured and predicted frequencies, +the field energy, and phi in the (x, y) plane at the first, middle and last saved time. +""" + +import argparse +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt +from scipy.optimize import fsolve +from scipy.special import ive, wofz + +from struphy import Output + +SCRIPT_DIR = Path(__file__).resolve().parent + +# The setup of the run (parsed from the same command line). +sys.path.insert(0, str(SCRIPT_DIR)) +import params_slab_drift_wave as params # noqa: E402 + +DEFAULT_OUTPUT = Path(params.args.out_folder) / params.args.sim_folder + + +def plasma_Z(zeta): + """Plasma dispersion function.""" + return 1j * np.sqrt(np.pi) * wofz(zeta) + + +def drift_wave_roots(n_harmonics=8): + """Roots of the local kinetic dispersion relations (quasi-neutral, eta_i = 0) of the full-orbit model and of + the drift-kinetic model, for the wave vector of the initial perturbation (slab, uniform density gradient): + + full orbit: 1/T_e + (1/T_i) [1 + sum_n Gamma_n(b) (omega - omega_*i) / s Z((omega - n/eps) / s)] = 0 + drift-kinetic: b/T_i + 1/T_e + (1/T_i) [1 + (omega - omega_*i) / s Z(omega / s)] = 0 + + with s = sqrt(2) |k_z| sqrt(T_i), b = eps^2 k_perp^2 T_i, omega_*i = -eps k_y T_i kappa_n and the plasma + dispersion function Z.""" + eps, Te, Ti, kap = params.epsilon, params.Te, params.Ti, params.kappa_n + kperp2 = params.kx**2 + params.ky**2 + b = eps**2 * kperp2 * Ti + w_si = -eps * params.ky * Ti * kap + s = np.sqrt(2) * abs(params.kz) * np.sqrt(Ti) + + def D_vlasov(w): + tot = 1.0 + 0j + for n in range(-n_harmonics, n_harmonics + 1): + tot += ive(n, b) * (w - w_si) / s * plasma_Z((w - n / eps) / s) + return 1 / Te + tot / Ti + + def D_dk(w): + return b / Ti + 1 / Te + (1 + (w - w_si) / s * plasma_Z(w / s)) / Ti + + def root(D): + sol = fsolve(lambda x: [D(x[0] + 1j * x[1]).real, D(x[0] + 1j * x[1]).imag], [-w_si * Te / Ti, 0.0]) + return sol[0] + 1j * sol[1] + + return root(D_vlasov), root(D_dk) + + +def mode_coefficient(run, n_points=32): + """Fourier coefficient c(t) of phi on the plane x = Lx/2 for the initial mode cos(k_y y + k_z z), which is one + period in eta2 and eta3; for a wave cos(k_y y + k_z z - omega t), c(t) = A/2 exp(-i omega t).""" + e = np.linspace(0.0, 1.0, n_points, endpoint=False) + phi = run.evaluate("em_fields/phi", eta1=0.5, eta2=e, eta3=e) + data = np.asarray(phi.transpose("t", "eta2", "eta3").values) + phase = np.exp(-2j * np.pi * (e[:, None] + e[None, :])) + return phi.t.values, np.mean(data * phase[None], axis=(1, 2)) + + +def measured_frequency(t, coeff): + """Angular frequency from a linear fit of the unwrapped phase of c(t) = A/2 exp(-i omega t).""" + if len(t) < 3: + return np.nan + return -np.polyfit(t, np.unwrap(np.angle(coeff)), 1)[0] + + +def saver(run, save_figs): + """A function saving figures as ``.png`` into ``/results`` and closing them (``save_figs``).""" + results_dir = Path(run.path_out) / "results" + results_dir.mkdir(exist_ok=True) + + def save(figs): + if save_figs: + for name, fig in figs.items(): + fig.savefig(results_dir / f"{name}.png") + plt.close(fig) + + return save + + +def main(path_out=DEFAULT_OUTPUT, save_figs=True): + run = Output(path_out) + snapshots = run.standard_snapshots() + save = saver(run, save_figs) + + # the initial (k_y, k_z) mode on the plane x = Lx/2 + t, coeff = mode_coefficient(run) + omega_num = measured_frequency(t, coeff) + omega_v, omega_dk = drift_wave_roots() + print(f"measured omega = {omega_num:.4f}, full orbit: {omega_v:.4f}, drift-kinetic: {omega_dk:.4f}") + + fig, axs = plt.subplots(2, 1, sharex=True) + axs[0].plot(t, np.abs(coeff)) + axs[0].set(ylabel="mode amplitude", yscale="log") + axs[1].plot(t, np.unwrap(np.angle(coeff)), label="measured") + axs[1].plot(t, np.angle(coeff[0]) - omega_v.real * t, "--", label="full orbit") + axs[1].plot(t, np.angle(coeff[0]) - omega_dk.real * t, ":", label="drift-kinetic") + axs[1].set(xlabel="time", ylabel="mode phase") + axs[1].legend() + fig.suptitle( + rf"$\omega$ measured {omega_num:.3f}, full orbit {omega_v.real:.3f}, drift-kinetic {omega_dk.real:.3f}" + ) + save({"mode": fig}) + + energy = run.evaluate("scalars", variables="en_phi")["en_phi"] + fig, ax = plt.subplots() + ax.plot(energy.t, energy) + ax.set(xlabel="time", yscale="log", title="Field energy en_phi") + save({"en_phi": fig}) + + # phi in the (x, y) plane at the first, middle and last saved time + plane = run.evaluate("em_fields/phi", t=list(snapshots), eta3=0.0) + grid = plane.plot(x="eta2", y="eta1", col="t") + save({"phi_xy": grid.fig}) + + if not save_figs: + plt.show() + + +if __name__ == "__main__": + # the flags of the run are parsed by params_slab_drift_wave.py, the rest here + _, other_flags = params.parser.parse_known_args() + parser = argparse.ArgumentParser(description="Plot a saved simulation run.") + parser.add_argument( + "path_out", + nargs="?", + default=DEFAULT_OUTPUT, + help="Simulation output folder (default: the output folder of the run flags, beside this script)", + ) + parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + args = parser.parse_args(other_flags) + main(args.path_out, save_figs=not args.show) diff --git a/src/struphy/api/models/__init__.py b/src/struphy/api/models/__init__.py index 9a120106c..dde4b4891 100644 --- a/src/struphy/api/models/__init__.py +++ b/src/struphy/api/models/__init__.py @@ -48,6 +48,7 @@ from struphy.models.vlasov import Vlasov from struphy.models.vlasov_ampere_one_species import VlasovAmpereOneSpecies from struphy.models.vlasov_maxwell_one_species import VlasovMaxwellOneSpecies +from struphy.models.vlasov_poisson_electrostatic_adiabatic import VlasovPoissonElectrostaticAdiabatic __all__ = [ "ColdPlasma", @@ -90,4 +91,5 @@ "Vlasov", "VlasovAmpereOneSpecies", "VlasovMaxwellOneSpecies", + "VlasovPoissonElectrostaticAdiabatic", ] diff --git a/src/struphy/models/__init__.py b/src/struphy/models/__init__.py index 17fbe98b9..33e976a66 100644 --- a/src/struphy/models/__init__.py +++ b/src/struphy/models/__init__.py @@ -49,6 +49,7 @@ "Vlasov": "struphy.models.vlasov", "VlasovAmpereOneSpecies": "struphy.models.vlasov_ampere_one_species", "VlasovMaxwellOneSpecies": "struphy.models.vlasov_maxwell_one_species", + "VlasovPoissonElectrostaticAdiabatic": "struphy.models.vlasov_poisson_electrostatic_adiabatic", } if TYPE_CHECKING: # static analysis and IDEs see the eager imports @@ -92,6 +93,7 @@ from struphy.models.vlasov import Vlasov from struphy.models.vlasov_ampere_one_species import VlasovAmpereOneSpecies from struphy.models.vlasov_maxwell_one_species import VlasovMaxwellOneSpecies + from struphy.models.vlasov_poisson_electrostatic_adiabatic import VlasovPoissonElectrostaticAdiabatic __all__ = [ "ColdPlasma", @@ -134,6 +136,7 @@ "Vlasov", "VlasovAmpereOneSpecies", "VlasovMaxwellOneSpecies", + "VlasovPoissonElectrostaticAdiabatic", ] diff --git a/src/struphy/models/cold_plasma_vlasov.py b/src/struphy/models/cold_plasma_vlasov.py index 0abf622cc..de0fa2f44 100644 --- a/src/struphy/models/cold_plasma_vlasov.py +++ b/src/struphy/models/cold_plasma_vlasov.py @@ -186,6 +186,7 @@ def __init__( self.initial_poisson = PoissonSolve( rho=particles_to_grid, rho_coeffs=hot_alpha**2 / hot_epsilon, + delta_f_source=True, ) self.initial_poisson.variables.phi = self.em_fields.phi @@ -204,10 +205,9 @@ def post_allocate(self): """ logger.info("\nINITIAL POISSON SOLVE:") - # charge of f - f0 for the Poisson right-hand side (see below for the weights of the time stepping) - particles = self.hot_elec.var.particles - particles.update_weights() - + # The right-hand side is the charge of f - f0, deposited with the delta-f weights also for full-f markers + # (delta_f_source=True, the full-f weights are restored after the deposit): much less noise than the full-f + # weights minus the background n0. self.initial_poisson.allocate() # Solve with dt=1. and compute electric field @@ -221,11 +221,6 @@ def post_allocate(self): self.em_fields.e_field.spline.vector.update_ghost_regions() logger.info("... Done.") - # The Poisson right-hand side is the charge of f - f0 in any case; the time stepping uses the delta-f - # weights only with the control variate (and the scalars must be computed with the same weights at t=0). - if not particles.control_variate: - particles.weights = particles.weights0.copy() - ## default parameters def generate_default_parameter_file(self, path=None, prompt=True): params_path = super().generate_default_parameter_file(path=path, prompt=prompt) diff --git a/src/struphy/models/tests/verification/test_verif_LinearVlasovAmpereOneSpecies.py b/src/struphy/models/tests/verification/test_verif_LinearVlasovAmpereOneSpecies.py index ecdcc3a9f..6fdb0b4ad 100644 --- a/src/struphy/models/tests/verification/test_verif_LinearVlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/verification/test_verif_LinearVlasovAmpereOneSpecies.py @@ -23,6 +23,7 @@ perturbations, ) from struphy.models import LinearVlasovAmpereOneSpecies +from struphy.propagators.base import Propagator logger = logging.getLogger("struphy") @@ -214,6 +215,57 @@ def test_energy_conservation(exit_before_run: bool = False): shutil.rmtree(test_folder) +def test_initial_poisson_no_background(): + r"""The delta-f markers carry f - f0, so the initial Poisson solve must not subtract the background n0. + + Without perturbation the delta-f weights vanish and, with :math:`\phi = 0` at the walls in :math:`x`, so does the + initial potential. Subtracting :math:`n_0` gives :math:`-\Delta \phi = -\alpha^2 n_0 / \varepsilon`, i.e. + :math:`\phi = -x(1 - x)/2` with root mean square 0.09. This runs zero time steps. + """ + model = LinearVlasovAmpereOneSpecies(alpha=1.0, epsilon=1.0, with_B0=False, with_E0=False) + + test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") + out_folders = os.path.join(test_folder, "LinearVlasovAmpereOneSpecies") + env = EnvironmentOptions(out_folders=out_folders, sim_folder="initial_poisson_no_background") + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=100, seed=1234), + weights_params=WeightsParameters(), + boundary_params=BoundaryParameters(bc=("reflect", "periodic", "periodic")), + sorting_params=SortingParameters(boxes_per_dim=(8, 1, 1), do_sort=True), + saving_params=SavingParameters(), + bufsize=0.4, + ) + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.coupling_Eweights.options = model.propagators.coupling_Eweights.Options() + model.initial_poisson.options = model.initial_poisson.Options() + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + # a separate initial condition without perturbation: its density is set to zero for the delta-f markers, + # while the background f0 keeps n0 = 1 + model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, None))) + + sim = Simulation( + model=model, + env=env, + time_opts=Time(dt=0.05, Tend=0.0), + domain=domains.Cuboid(), + grid=grids.TensorProductGrid(num_elements=(16, 1, 1)), + derham_opts=DerhamOptions(degree=(3, 1, 1), bcs=(("dirichlet", "dirichlet"), None, None)), + ) + sim.run() + + # root mean square of phi over the unit cube (collective, on all ranks) + phi = model.em_fields.phi.spline.vector + phi_rms = float(xp.sqrt(Propagator.mass_ops.M0.dot_inner(phi, phi))) + logger.info(f"{phi_rms =}") + assert phi_rms < 1e-6, f"Initial potential without perturbation: {phi_rms =} (a background n0 gives 0.09)." + + MPI.COMM_WORLD.Barrier() + if MPI.COMM_WORLD.Get_rank() == 0: + shutil.rmtree(test_folder) + + if __name__ == "__main__": test_delta_f_initial_condition() + test_initial_poisson_no_background() test_energy_conservation() diff --git a/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py index 037ba04c0..bdb4cf888 100644 --- a/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py @@ -27,6 +27,7 @@ ) from struphy.linear_algebra.solver import SolverParameters from struphy.models import VlasovAmpereOneSpecies +from struphy.propagators.base import Propagator logger = logging.getLogger("struphy") @@ -261,7 +262,119 @@ def test_energy_conservation(control_variate: bool, exit_before_run: bool = Fals shutil.rmtree(env.path_out) +def test_initial_poisson_full_f(): + r"""The initial Poisson solve with full-f markers deposits the charge of :math:`f - f_0` (delta-f weights, no + background source), i.e. subtracts :math:`n_0` once. + + A uniform Maxwellian is neutral, so with :math:`\phi = 0` at the walls in :math:`x` the initial potential is at + noise level. Subtracting :math:`n_0` twice (delta-f weights plus a full-f background source) gives + :math:`-\Delta \phi = -\alpha^2 n_0 / \varepsilon`, i.e. :math:`\phi = -x(1 - x)/2` with root mean square 0.09. + This runs zero time steps. + """ + model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=1.0, with_B0=False) + + test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") + out_folders = os.path.join(test_folder, "VlasovAmpereOneSpecies") + env = EnvironmentOptions(out_folders=out_folders, sim_folder="initial_poisson_full_f") + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=200, seed=1234), + weights_params=WeightsParameters(control_variate=False), + boundary_params=BoundaryParameters(bc=("reflect", "periodic", "periodic")), + sorting_params=SortingParameters(boxes_per_dim=(8, 1, 1), do_sort=True), + saving_params=SavingParameters(), + bufsize=0.4, + ) + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.coupling_va.options = model.propagators.coupling_va.Options() + model.initial_poisson.options = model.initial_poisson.Options() + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + + sim = Simulation( + model=model, + env=env, + time_opts=Time(dt=0.05, Tend=0.0), + domain=domains.Cuboid(), + grid=grids.TensorProductGrid(num_elements=(16, 1, 1)), + derham_opts=DerhamOptions(degree=(3, 1, 1), bcs=(("dirichlet", "dirichlet"), None, None)), + ) + sim.run() + + # root mean square of phi over the unit cube (collective, on all ranks) + phi = model.em_fields.phi.spline.vector + phi_rms = float(xp.sqrt(Propagator.mass_ops.M0.dot_inner(phi, phi))) + logger.info(f"{phi_rms =}") + assert phi_rms < 0.02, ( + f"Initial potential of a neutral full-f plasma: {phi_rms =} (n0 subtracted twice gives 0.09)." + ) + + MPI.COMM_WORLD.Barrier() + if MPI.COMM_WORLD.Get_rank() == 0: + shutil.rmtree(test_folder) + + +def test_initial_poisson_full_f_noise(): + r"""The initial Poisson solve deposits the delta-f weights also for full-f markers. + + For a density perturbation :math:`\delta n = a \cos(k x)` of a uniform Maxwellian, the exact initial potential is + :math:`\phi = (\alpha^2/\varepsilon)\, a \cos(k x) / k^2`. With the delta-f weights the Monte-Carlo noise is that + of :math:`f - f_0` (relative error of about 1% here), with the full-f weights minus the background :math:`n_0` it + is that of :math:`f` (relative error of about 400%). This runs zero time steps. + """ + model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=1.0, with_B0=False) + + test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") + out_folders = os.path.join(test_folder, "VlasovAmpereOneSpecies") + env = EnvironmentOptions(out_folders=out_folders, sim_folder="initial_poisson_full_f_noise") + + r1, amp = 12.56, 1e-2 + k = 2 * xp.pi / r1 + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=100, seed=1234), + weights_params=WeightsParameters(control_variate=False), + boundary_params=BoundaryParameters(), + sorting_params=SortingParameters(boxes_per_dim=(8, 1, 1), do_sort=True), + saving_params=SavingParameters(), + bufsize=0.4, + ) + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.coupling_va.options = model.propagators.coupling_va.Options() + model.initial_poisson.options = model.initial_poisson.Options() + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + perturbation = perturbations.ModesCos(ls=(1,), amps=(amp,)) + model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, perturbation))) + + sim = Simulation( + model=model, + env=env, + time_opts=Time(dt=0.05, Tend=0.0), + domain=domains.Cuboid(r1=r1), + grid=grids.TensorProductGrid(num_elements=(32, 1, 1)), + derham_opts=DerhamOptions(degree=(3, 1, 1)), + ) + sim.run() + + # relative L2 error of phi along x (spline evaluation is collective, on all ranks) + e1 = xp.linspace(0.0, 1.0, 200, endpoint=False) + num = Propagator.domain.push(model.em_fields.phi.spline, e1, 0.5, 0.5, kind="0").squeeze() + exact = amp / k**2 * xp.cos(k * r1 * e1) + error = float(xp.sqrt(xp.mean((num - exact) ** 2) / xp.mean(exact**2))) + logger.info(f"{error =}") + assert error < 0.05, f"Initial potential too noisy: relative error {error} (full-f weights give about 4)." + + # the full-f weights of the time stepping are restored after the delta-f deposit (rank-local check) + particles = model.kinetic_ions.var.particles + assert xp.allclose(particles.weights, particles.weights0), "Full-f weights not restored after the initial solve." + + MPI.COMM_WORLD.Barrier() + if MPI.COMM_WORLD.Get_rank() == 0: + shutil.rmtree(test_folder) + + if __name__ == "__main__": test_weak_Landau(do_plot=True) + test_initial_poisson_full_f_noise() + test_initial_poisson_full_f() test_energy_conservation(control_variate=False) test_energy_conservation(control_variate=True) diff --git a/src/struphy/models/tests/verification/test_verif_VlasovMaxwellOneSpecies.py b/src/struphy/models/tests/verification/test_verif_VlasovMaxwellOneSpecies.py index 77c9a1b28..ed5c55692 100644 --- a/src/struphy/models/tests/verification/test_verif_VlasovMaxwellOneSpecies.py +++ b/src/struphy/models/tests/verification/test_verif_VlasovMaxwellOneSpecies.py @@ -3,6 +3,7 @@ import shutil import cunumpy as xp +import pytest from maybempi import MPI from struphy import ( @@ -25,8 +26,13 @@ logger = logging.getLogger("struphy") -def test_initial_gauss_error(exit_before_run: bool = False): - """Right after the initial Poisson solve the weak Gauss law holds. +@pytest.mark.parametrize( + "control_variate, Tend, rel_tol", + [(True, 0.0, 0.1), (False, 0.0, 0.1), (False, 1.0, 0.5)], +) +def test_gauss_error(control_variate: bool, Tend: float, rel_tol: float, exit_before_run: bool = False): + """The ``gauss_error`` diagnostic: right after the initial Poisson solve the weak Gauss law holds, and for + full-f markers the residual stays at the level of the time stepping. ``post_allocate`` solves :math:`\\mathbb G^\\top \\mathbb M^1 \\mathbb G \\boldsymbol \\phi = \\frac{\\alpha^2}{\\varepsilon} \\boldsymbol \\rho` and sets :math:`\\mathbf e = -\\mathbb G \\boldsymbol \\phi`, so the ``gauss_error`` scalar @@ -34,16 +40,20 @@ def test_initial_gauss_error(exit_before_run: bool = False): against :math:`\\max |\\mathbb G^\\top \\mathbb M^1 \\mathbf e|` (what is left is the Monte-Carlo net charge, which the periodic Poisson problem cannot represent), independently of the number of MPI ranks. - This runs zero time steps. + The cases with ``Tend = 0`` run zero time steps, with and without the control variate. For full-f markers the + initial solve uses the delta-f weights, and the diagnostic deposits the full-f weights minus the cached offset + of :math:`t=0`; the case with ``Tend = 1`` (20 steps) checks that the residual then stays of the order of the + Gauss-law violation of the time stepping (about 0.1 of the field term here), instead of the noise of the + :math:`f_0` deposit (about 3 times the field term without the offset). """ # alpha^2/epsilon != 1 to check the coefficient of the charge density model = VlasovMaxwellOneSpecies(alpha=2.0, epsilon=0.5, measure_gauss_law=True) test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") out_folders = os.path.join(test_folder, "VlasovMaxwellOneSpecies") - env = EnvironmentOptions(out_folders=out_folders, sim_folder="gauss_error") + env = EnvironmentOptions(out_folders=out_folders, sim_folder=f"gauss_error_{control_variate}_{Tend}") - time_opts = Time(dt=0.05, Tend=0.0) + time_opts = Time(dt=0.05, Tend=Tend) domain = domains.Cuboid(r1=12.56) grid = grids.TensorProductGrid(num_elements=(16, 1, 1)) @@ -51,7 +61,7 @@ def test_initial_gauss_error(exit_before_run: bool = False): model.kinetic_ions.set_markers( loading_params=LoadingParameters(ppc=1000, seed=1234), - weights_params=WeightsParameters(control_variate=True), + weights_params=WeightsParameters(control_variate=control_variate), boundary_params=BoundaryParameters(), sorting_params=SortingParameters(boxes_per_dim=(8, 1, 1), do_sort=True), saving_params=SavingParameters(), @@ -93,14 +103,18 @@ def test_initial_gauss_error(exit_before_run: bool = False): if comm.Get_size() > 1: comm.Allreduce(MPI.IN_PLACE, lhs_max, op=MPI.MAX) - assert error < 0.1 * lhs_max[0], ( + assert error < rel_tol * lhs_max[0], ( f"Assertion for the initial Gauss-law residual failed: {error =} vs. max|G^T M1 e| = {lhs_max[0]}." ) logger.info(f"Assertion for the initial Gauss-law residual passed ({error =}, {lhs_max[0] =}).") + # all ranks are done with this case before the shared folder is removed + comm.Barrier() if comm.Get_rank() == 0: shutil.rmtree(test_folder) + comm.Barrier() if __name__ == "__main__": - test_initial_gauss_error() + test_gauss_error(True, 0.0, 0.1) + test_gauss_error(False, 1.0, 0.5) diff --git a/src/struphy/models/tests/verification/test_verif_VlasovPoissonElectrostaticAdiabatic.py b/src/struphy/models/tests/verification/test_verif_VlasovPoissonElectrostaticAdiabatic.py new file mode 100644 index 000000000..3e6fa33a6 --- /dev/null +++ b/src/struphy/models/tests/verification/test_verif_VlasovPoissonElectrostaticAdiabatic.py @@ -0,0 +1,214 @@ +import logging +import os +import shutil + +import cunumpy as xp +import h5py +import pytest +from matplotlib import pyplot as plt +from maybempi import MPI + +from struphy import ( + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, + perturbations, +) +from struphy.models import VlasovPoissonElectrostaticAdiabatic +from struphy.propagators.base import Propagator + +logger = logging.getLogger("struphy") + + +@pytest.mark.parametrize("quasi_neutral", [True, False]) +def test_ion_acoustic_Landau(quasi_neutral: bool, do_plot: bool = False, exit_before_run: bool = False): + r"""Verification test for Landau damping of an ion-acoustic wave along :math:`\mathbf B_0`. + + For :math:`\mathbf k \parallel \mathbf B_0` the gyration decouples and the dispersion relation of the model is + + .. math:: + + \frac{\varepsilon^2 k^2}{\alpha^2 n_0} + \frac{1}{Z T_e} + \frac{1}{T_i}\left[1 + \zeta Z(\zeta)\right] = 0\,, + \qquad \zeta = \frac{\omega}{\sqrt 2 k \sqrt{T_i}}\,. + + With :math:`Z = 1`, :math:`T_i = 1`, :math:`T_e = 4`, :math:`k = 0.5` and in the quasi-neutral limit this is + exactly the dispersion relation of weak (electron) Landau damping at :math:`k = 0.5`, so + :math:`\omega = 1.4157`, :math:`\gamma = -0.1534` (see ``test_verif_VlasovAmpereOneSpecies.py``). + With the vacuum term (:math:`\varepsilon = 1`, :math:`\alpha = 2`) the root is :math:`\omega = 1.3356`, + :math:`\gamma = -0.1936`. + """ + # light-weight model instance (alpha only matters with the vacuum term, epsilon drops out for k parallel to B0) + model = VlasovPoissonElectrostaticAdiabatic(alpha=2.0, epsilon=1.0) + + # environment options + test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") + out_folders = os.path.join(test_folder, "VlasovPoissonElectrostaticAdiabatic") + env = EnvironmentOptions(out_folders=out_folders, sim_folder=f"ion_acoustic_qn_{quasi_neutral}") + + # time stepping + time_opts = Time(dt=0.05, Tend=15) + + # geometry: k = 2 pi / r1 = 0.5 along B0 = B0 e_x + r1 = 12.56 + domain = domains.Cuboid(r1=r1) + + # fluid equilibrium: n0 = 1 and T_e = p / n0 = beta |B0|^2 / 2 = 4 + equil = equils.HomogenSlab(B0x=1.0, B0z=0.0, beta=8.0, n0=1.0) + + # grid + grid = grids.TensorProductGrid(num_elements=(32, 1, 1)) + + # derham options + derham_opts = DerhamOptions(degree=(3, 1, 1)) + + # markers + ppc = 1000 + loading_params = LoadingParameters(ppc=ppc, seed=1234) + weights_params = WeightsParameters(control_variate=True) + boundary_params = BoundaryParameters() + sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) + saving_params = SavingParameters() + + model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize=0.4, + ) + + # propagator options + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.poisson.options = model.propagators.poisson.Options(quasi_neutral=quasi_neutral) + model.propagators.push_v_E.options = model.propagators.push_v_E.Options() + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() + + # background (T_i = 1) and initial condition + background = maxwellians.Maxwellian3D(n=(1.0, None)) + model.kinetic_ions.var.add_background(background) + + perturbation = perturbations.ModesCos(ls=(1,), amps=(1e-3,)) + init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) + model.kinetic_ions.var.add_initial_condition(init) + + # instance of simulation + sim = Simulation( + model=model, + env=env, + time_opts=time_opts, + equil=equil, + domain=domain, + grid=grid, + derham_opts=derham_opts, + ) + + if exit_before_run: + logger.info("Exiting before running simulation.") + return sim + + sim.run() + + # exact damping rate + gamma = -0.1534 if quasi_neutral else -0.1936 + + dt = time_opts.dt + if MPI.COMM_WORLD.Get_rank() == 0: + pa_data = os.path.join(env.path_out, "data") + with h5py.File(os.path.join(pa_data, "data_proc0.hdf5"), "r") as f: + time = f["time"]["value"][()] + E = f["scalar"]["en_phi"][()] + logE = xp.log10(E) + + # maxima of the field energy + dEdt = (xp.roll(logE, -1) - xp.roll(logE, 1))[1:-1] / (2.0 * dt) + zeros = dEdt * xp.roll(dEdt, -1) < 0.0 + maxima_inds = xp.logical_and(zeros, dEdt > 0.0) + maxima = logE[1:-1][maxima_inds] + t_maxima = time[1:-1][maxima_inds] + + if do_plot: + plt.figure(figsize=(12, 8)) + plt.plot(time, logE, label="numerical") + plt.plot(t_maxima[:5], maxima[:5], "or", markersize=10) + plt.plot(time, logE[0] + 2 * gamma * time / xp.log(10), "k--", label="exact rate") + plt.legend() + plt.xlabel("time") + plt.show() + + # linear fit of log10(E) ~ 2 gamma t / ln(10) + linfit = xp.polyfit(t_maxima[:5], maxima[:5], 1) + gamma_num = linfit[0] * xp.log(10) / 2 + + rel_error = xp.abs(gamma_num - gamma) / xp.abs(gamma) + assert rel_error < 0.22, f"Ion-acoustic Landau damping failed: {gamma_num =} vs. {gamma =}." + logger.info(f"Ion-acoustic Landau damping passed ({gamma_num =}, {rel_error =}).") + + shutil.rmtree(test_folder) + + +def test_neutral_full_f_background(): + r"""A uniform full-f plasma of ions with charge number :math:`Z = 2` is neutral, so the potential is at noise level. + + Without the control variate, the Poisson solve subtracts the neutralizing background :math:`n_0` from the deposited + ion density; the charge number is part of the source coefficient, not of the background. A background + :math:`Z n_0` would give :math:`\phi = -Z (Z - 1) \varepsilon T_0`, here :math:`-2`, everywhere. + """ + model = VlasovPoissonElectrostaticAdiabatic(charge_number=2, mass_number=4.0, alpha=1.0, epsilon=1.0) + + test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") + out_folders = os.path.join(test_folder, "VlasovPoissonElectrostaticAdiabatic") + env = EnvironmentOptions(out_folders=out_folders, sim_folder="neutral_full_f") + + # n0 = 1 and T_e = p / n0 = beta |B0|^2 / 2 = 1 + equil = equils.HomogenSlab(B0z=1.0, beta=2.0, n0=1.0) + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=200, seed=1234), + weights_params=WeightsParameters(control_variate=False), + boundary_params=BoundaryParameters(), + sorting_params=SortingParameters(boxes_per_dim=(8, 1, 1), do_sort=True), + saving_params=SavingParameters(), + bufsize=0.4, + ) + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.poisson.options = model.propagators.poisson.Options() + model.propagators.push_v_E.options = model.propagators.push_v_E.Options() + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + + sim = Simulation( + model=model, + env=env, + time_opts=Time(dt=0.05, Tend=0.05), + equil=equil, + domain=domains.Cuboid(), + grid=grids.TensorProductGrid(num_elements=(16, 1, 1)), + derham_opts=DerhamOptions(degree=(3, 1, 1)), + ) + sim.run() + + # root mean square of phi over the unit cube (collective, on all ranks) + phi = model.em_fields.phi.spline.vector + phi_rms = float(xp.sqrt(Propagator.mass_ops.M0.dot_inner(phi, phi))) + logger.info(f"{phi_rms =}") + assert phi_rms < 0.5, f"Full-f background not neutralizing: {phi_rms =} (a background Z n0 gives 2)." + + MPI.COMM_WORLD.Barrier() + if MPI.COMM_WORLD.Get_rank() == 0: + shutil.rmtree(test_folder) + + +if __name__ == "__main__": + test_ion_acoustic_Landau(True, do_plot=True) + test_neutral_full_f_background() diff --git a/src/struphy/models/vlasov_ampere_one_species.py b/src/struphy/models/vlasov_ampere_one_species.py index ca63b26fb..f6453e2b2 100644 --- a/src/struphy/models/vlasov_ampere_one_species.py +++ b/src/struphy/models/vlasov_ampere_one_species.py @@ -148,6 +148,7 @@ def __init__( self.initial_poisson = PoissonSolve( rho=particles_to_grid, rho_coeffs=alpha**2 / epsilon, + delta_f_source=True, ) self.initial_poisson.variables.phi = self.em_fields.phi @@ -166,10 +167,9 @@ def post_allocate(self): """ logger.info("\nINITIAL POISSON SOLVE:") - # charge of f - f0 for the Poisson right-hand side (see below for the weights of the time stepping) - particles = self.kinetic_ions.var.particles - particles.update_weights() - + # The right-hand side is the charge of f - f0, deposited with the delta-f weights also for full-f markers + # (delta_f_source=True, the full-f weights are restored after the deposit): much less noise than the full-f + # weights minus the background n0. self.initial_poisson.allocate() # Solve with dt=1. and compute electric field @@ -183,11 +183,6 @@ def post_allocate(self): self.em_fields.e_field.spline.vector.update_ghost_regions() logger.info("... Done.") - # The Poisson right-hand side is the charge of f - f0 in any case; the time stepping uses the delta-f - # weights only with the control variate (and the scalars must be computed with the same weights at t=0). - if not particles.control_variate: - particles.weights = particles.weights0.copy() - ## default parameters def generate_default_parameter_file(self, path=None, prompt=True): params_path = super().generate_default_parameter_file(path=path, prompt=prompt) diff --git a/src/struphy/models/vlasov_maxwell_one_species.py b/src/struphy/models/vlasov_maxwell_one_species.py index 055e6b445..66ae1185d 100644 --- a/src/struphy/models/vlasov_maxwell_one_species.py +++ b/src/struphy/models/vlasov_maxwell_one_species.py @@ -153,6 +153,7 @@ def __init__( self.initial_poisson = PoissonSolve( rho=particles_to_grid, rho_coeffs=alpha**2 / epsilon, + delta_f_source=True, ) self.initial_poisson.variables.phi = self.em_fields.phi @@ -174,16 +175,14 @@ def post_allocate(self): """ self._tmp = xp.empty(1, dtype=float) - particles = self.kinetic_ions.var.particles - if self.measure_gauss_law: self.op = Propagator.derham.grad.T @ Propagator.mass_ops.M1 logger.info("\nINITIAL POISSON SOLVE:") - # charge of f - f0 for the Poisson right-hand side (see below for the weights of the time stepping) - particles.update_weights() - + # The right-hand side is the charge of f - f0, deposited with the delta-f weights also for full-f markers + # (delta_f_source=True, the full-f weights are restored after the deposit): much less noise than the full-f + # weights minus the background n0. self.initial_poisson.allocate() # keep the AccumulatorVector built by the propagator for the Gauss-law diagnostic below @@ -200,10 +199,21 @@ def post_allocate(self): self.em_fields.e_field.spline.vector.update_ghost_regions() logger.info("... Done.") - # The Poisson right-hand side is the charge of f - f0 in any case; the time stepping uses the delta-f - # weights only with the control variate (and the scalars must be computed with the same weights at t=0). - if not particles.control_variate: + # Full-f markers: the time stepping conserves the charge of the full-f weights, while the initial solve used + # the delta-f weights. The Gauss-law diagnostic therefore deposits the full-f weights and subtracts the + # constant offset alpha^2/eps (rho_full - rho_delta_f) at t=0 (the initial f0 deposit minus n0). + # On restart, post_allocate runs before the markers are restored, so the offset belongs to freshly loaded + # t=0 markers: consistent with the restored run only for reproducible loading (fixed seed, same decomposition). + self._gauss_offset = None + particles = self.kinetic_ions.var.particles + if self.measure_gauss_law and not particles.control_variate: + self.charge_accum() + rho_full = self.charge_accum.vectors[0].copy() + particles.update_weights() + self.charge_accum() particles.weights = particles.weights0.copy() + rho_full -= self.charge_accum.vectors[0] + self._gauss_offset = self.initial_poisson.coeffs[0] * rho_full def calculate_gauss_error(self): r"""Maximum norm of the weak Gauss-law residual @@ -212,21 +222,23 @@ def calculate_gauss_error(self): \mathbb G^\top \mathbb M^1 \mathbf e + \frac{\alpha^2}{\varepsilon} \boldsymbol \rho\,, - where :math:`\boldsymbol \rho` is the charge density of :math:`f - f_0` deposited as in the initial - Poisson solve. Since :math:`\mathbf e = -\mathbb G \boldsymbol \phi` it vanishes at :math:`t=0`, - up to the net charge (the kernel of the periodic Poisson problem).""" - # control variate method + where :math:`\boldsymbol \rho` is the charge density of :math:`f - f_0`: with the control variate deposited + with the delta-f weights, as in the initial Poisson solve; for full-f markers deposited with the full-f + weights (whose charge the time stepping conserves) minus the constant offset + :math:`\boldsymbol \rho_\textnormal{full}(0) - \boldsymbol \rho_{\delta f}(0)` of the initial solve, which used + the delta-f weights. Since :math:`\mathbf e = -\mathbb G \boldsymbol \phi` it vanishes at :math:`t=0`, up to + the net charge (the kernel of the periodic Poisson problem).""" particles = self.kinetic_ions.var.particles - particles.update_weights() + if particles.control_variate: + particles.update_weights() self.charge_accum() rho = self.charge_accum.vectors[0] - # restore the weights of the time stepping (full-f without the control variate) - if not particles.control_variate: - particles.weights = particles.weights0.copy() e = self.em_fields.e_field.spline.vector residual = self.op.dot(e) residual += self.initial_poisson.coeffs[0] * rho + if self._gauss_offset is not None: + residual -= self._gauss_offset # maximum residual over the local MPI rank, then over all ranks of the domain decomposition # (all clones hold the same accumulated charge density) diff --git a/src/struphy/models/vlasov_poisson_electrostatic_adiabatic.py b/src/struphy/models/vlasov_poisson_electrostatic_adiabatic.py new file mode 100644 index 000000000..25ea4aa1a --- /dev/null +++ b/src/struphy/models/vlasov_poisson_electrostatic_adiabatic.py @@ -0,0 +1,351 @@ +import copy +import logging + +from struphy import BaseUnits +from struphy.io.options import LiteralOptions +from struphy.models.base import StruphyModel +from struphy.models.scalars import FunctionScalarFEEC, KineticEnergyPIC, Scalars +from struphy.models.species import ( + FieldSpecies, + ParticleSpecies, +) +from struphy.models.variables import FEECVariable, PICVariable +from struphy.pic.accumulation.kernels.charge_density_0form import charge_density_0form +from struphy.pic.accumulation.particles_to_grid import ParticlesToGrid +from struphy.propagators.adiabatic_poisson_solve import AdiabaticPoissonSolve +from struphy.propagators.base import Propagator +from struphy.propagators.push_eta import PushEta +from struphy.propagators.push_v_in_force_field import PushVinForceField +from struphy.propagators.push_vxb import PushVxB + +logger = logging.getLogger("struphy") + + +class VlasovPoissonElectrostaticAdiabatic(StruphyModel): + r"""Electrostatic full-orbit (6D) model for a single ion species with adiabatic electrons. + + The full-orbit counterpart of :class:`~struphy.models.DriftKineticElectrostaticAdiabatic`: + the ions are not reduced to guiding centers, so the polarization density is not modelled + but comes out of the ion orbits. The electrons enter through the adiabatic (Boltzmann) response + in the Poisson equation, optionally with the same flux-surface average as in the drift-kinetic model. + + Parameters + ---------- + base_units : BaseUnits + Base units for normalization (default: ``BaseUnits(kBT=1.0)``). + charge_number : int + Charge number (in units of the positive elementary charge) of the ion species (default: 1). + mass_number : float + Mass number (in units of the proton mass) of the ion species (default: 1.0). + alpha : float, optional + Ratio of plasma to cyclotron frequency :math:`\hat \omega_p / \hat \Omega_c`. + If ``None``, computed from ``base_units`` and the charge/mass numbers. + epsilon : float, optional + Normalized inverse cyclotron frequency :math:`1 / (\hat{\Omega}_\mathrm{c} \hat{t})`. + If ``None``, computed from ``base_units`` and the charge/mass numbers. + with_B0 : bool + Whether to include the Lorentz force of the background magnetic field (default: True). + """ + + @classmethod + def model_type(cls) -> LiteralOptions.ModelTypes: + return "Kinetic" + + ## species + + class EMFields(FieldSpecies): + def __init__(self): + self.phi = FEECVariable(space="H1") + self.init_variables() + + class KineticIons(ParticleSpecies): + def __init__( + self, + charge_number: int = 1, + mass_number: float = 1.0, + alpha: float = None, + epsilon: float = None, + ): + self.var = PICVariable(space="Particles6D") + self.init_variables( + charge_number=charge_number, + mass_number=mass_number, + alpha=alpha, + epsilon=epsilon, + ) + + ## propagators + + class Propagators: + def __init__( + self, + phi: FEECVariable = None, + rho: ParticlesToGrid = None, + epsilon: float = 1.0, + alpha: float = 1.0, + Z: int = 1, + with_B0: bool = True, + ): + # Poisson right after the drift: both E-kicks of a Strang step see the potential + # of the current marker positions (drift-kick-kick-drift, i.e. leapfrog) + self.push_eta = PushEta() + self.poisson = AdiabaticPoissonSolve(rho=rho, epsilon=epsilon, alpha=alpha, Z=Z) + self.push_v_E = PushVinForceField(potential=phi) + if with_B0: + self.push_vxb = PushVxB() + + ## abstract methods + + def __init__( + self, + base_units: BaseUnits = BaseUnits(kBT=1.0), + charge_number: int = 1, + mass_number: float = 1.0, + alpha: float = None, + epsilon: float = None, + with_B0: bool = True, + ): + self.with_B0 = with_B0 + + # 0. store input parameters + self.params = copy.deepcopy(locals()) + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.kinetic_ions = self.KineticIons( + charge_number, + mass_number, + alpha, + epsilon, + ) + + # 2. derive units (must be done after instantiating species to access charge and mass numbers) + self.setup_equation_params(base_units=base_units) + alpha = self.kinetic_ions.equation_params.alpha + epsilon = self.kinetic_ions.equation_params.epsilon + + # 3. instantiate all propagators + rho = ParticlesToGrid( + self.kinetic_ions.var, + "H1", + charge_density_0form, + ) + self.propagators = self.Propagators( + phi=self.em_fields.phi, + rho=rho, + epsilon=epsilon, + alpha=alpha, + Z=charge_number, + with_B0=with_B0, + ) + + # 4. assign variables to propagators + self.propagators.push_eta.variables.var = self.kinetic_ions.var + self.propagators.poisson.variables.phi = self.em_fields.phi + self.propagators.push_v_E.variables.var = self.kinetic_ions.var + if with_B0: + self.propagators.push_vxb.variables.ions = self.kinetic_ions.var + + # 5. define scalars to be tracked during simulation + en_phi = FunctionScalarFEEC(self._compute_en_phi) + en_particles = KineticEnergyPIC(self.kinetic_ions.var) + self.scalars = Scalars( + en_phi=en_phi, + en_particles=en_particles, + en_tot=en_phi + en_particles, + ) + + @property + def bulk_species(self): + return self.kinetic_ions + + @property + def velocity_scale(self): + return "thermal" + + def post_allocate(self): + """Solve the initial Poisson equation and allocate helper arrays. + + :meta private: + """ + logger.info("\nINITIAL POISSON SOLVE:") + self.propagators.poisson(1.0) + + phi = self.em_fields.phi.spline.vector + self._tmp_phi = phi.space.zeros() + self._grad_phi = Propagator.derham.grad.codomain.zeros() + + def _compute_en_phi(self): + r"""Field energy :math:`\frac{1}{2\alpha^2}\int |\nabla \phi|^2 + \frac{1}{2Z\varepsilon^2}\int \frac{n_0}{T_0}\phi(\phi - \langle\phi\rangle)` + (the vacuum part only if ``quasi_neutral=False``).""" + poisson = self.propagators.poisson + phi = self.em_fields.phi.spline.vector + epsilon = self.kinetic_ions.equation_params.epsilon + alpha = self.kinetic_ions.equation_params.alpha + Z = self.kinetic_ions.var.species.charge_number + + # the adiabatic term with the same operator as in the solve, stab_mat @ (I - A) + en = 0.5 / (Z * epsilon**2) * phi.inner(poisson._stab_mat.dot(phi, out=self._tmp_phi)) + if not poisson.options.quasi_neutral: + e = Propagator.derham.grad.dot(phi, out=self._grad_phi) + en += 0.5 / alpha**2 * Propagator.mass_ops.M1.dot_inner(e, e) + return en + + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "BaseUnits(" in line: + new_file += ["base_units = BaseUnits(kBT=1.0)\n"] + elif "push_vxb.Options" in line: + new_file += ["if model.with_B0:\n"] + new_file += [" " + line] + elif "saving_params = " in line: + new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] + new_file += ["saving_params = SavingParameters(binning_plots=(binplot,))\n\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + + ## abstract methods for documentation + + @classmethod + def doc_pde(cls): + r"""**PDEs solved by model:** + + Vlasov equation: + + .. math:: + + \frac{\partial f}{\partial t} + \mathbf{v} \cdot \nabla f + \frac{1}{\varepsilon} \left( -\nabla \phi + \mathbf{v} \times \mathbf{B}_0 \right) \cdot \frac{\partial f}{\partial \mathbf{v}} = 0 + + Poisson equation with adiabatic electrons: + + .. math:: + + -\frac{\varepsilon}{\alpha^2} \Delta \phi + \frac{1}{Z \varepsilon} \frac{n_0}{T_0} \left( \phi - \langle \phi \rangle \right) = \int_{\mathbb{R}^3} f \, \mathrm{d}^3 \mathbf{v} - n_0 + + where :math:`f(\mathbf{x}, \mathbf{v}, t)` is the ion distribution function, :math:`n_0` the equilibrium + ion density, :math:`T_0` the electron temperature and :math:`Z` the ion charge number. + The vacuum term :math:`-\varepsilon \alpha^{-2} \Delta \phi` is dropped in the quasi-neutral limit + :math:`\alpha \to \infty` (option ``quasi_neutral``). The flux-surface average + :math:`\langle \phi \rangle` (zonal correction, as in + :class:`~struphy.models.DriftKineticElectrostaticAdiabatic`) is optional; by default + :math:`\langle \phi \rangle = 0`. + """ + + @classmethod + def doc_normalization(cls): + r"""The reference speed is the ion thermal speed and the electrostatic + fields are scaled accordingly: + + .. math:: + + \hat v = \sqrt{k_B \hat T / m_i},\qquad \hat E = \hat v \hat B,\qquad \hat\phi = \hat E \hat x, + + with temperatures in units of :math:`k_B \hat T`. The dimensionless parameters are + + .. math:: + + \varepsilon = \frac{1}{\hat\Omega_c\hat t} = \frac{\hat \rho}{\hat x}, \qquad \alpha = \frac{\hat{\omega}_\mathrm{p}}{\hat{\Omega}_\mathrm{c}}, + + such that :math:`\varepsilon / \alpha` is the normalized Debye length.""" + + @classmethod + def doc_scalar_quantities(cls): + r"""**The following scalars are tracked during simulation** (their sum is conserved by the model equations): + + - Field energy: ``en_phi`` + :math:`= \frac{1}{2\alpha^2} \int |\nabla \phi|^2 \, \mathrm d \mathbf x + \frac{1}{2 Z \varepsilon^2} \int \frac{n_0}{T_0} \phi (\phi - \langle \phi \rangle) \, \mathrm d \mathbf x` + (the first term only if ``quasi_neutral=False``) + - Kinetic energy: ``en_particles`` :math:`= \frac{1}{2} \sum_p w_p |\mathbf{v}_p|^2` + (of :math:`f - f_0` with the control variate) + - Total energy: ``en_tot``""" + + @classmethod + def doc_discretization(cls): + """Time integration is performed by the following propagators (in sequence): + + 1. :class:`~struphy.propagators.push_eta.PushEta` + 2. :class:`~struphy.propagators.adiabatic_poisson_solve.AdiabaticPoissonSolve` + 3. :class:`~struphy.propagators.push_v_in_force_field.PushVinForceField` + 4. :class:`~struphy.propagators.push_vxb.PushVxB` (if :attr:`with_B0` is True) + """ + doc = rf"""Time integration is performed by the following propagators (in sequence). +The Poisson solve comes right after the drift, such that with Strang splitting both velocity kicks use the +potential of the current marker positions (drift-kick-kick-drift). Since a Strang step ends with the half drift, the +saved potential and ``en_phi`` belong to the marker positions half a step earlier. + +**1. PushEta:** + +{PushEta.__doc__} + +**2. AdiabaticPoissonSolve:** + +{AdiabaticPoissonSolve.__doc__} + +**3. PushVinForceField:** + +{PushVinForceField.__doc__} + +**4. PushVxB:** + +{PushVxB.__doc__} +""" + return doc + + @classmethod + def doc_long_description(cls): + r"""This model is the full-orbit (6D) counterpart of the electrostatic drift-kinetic model + :class:`~struphy.models.DriftKineticElectrostaticAdiabatic`. The ions follow full Lorentz-force + orbits in the self-consistent electrostatic field and a static background magnetic field. Electrons + are not evolved; they enter through a linearized Boltzmann response in the Poisson equation, optionally + without response to the flux-surface-averaged potential (zonal correction), as in the drift-kinetic model. + The potential is recomputed from the marker positions at each time step. + + In the strong-field limit :math:`\varepsilon \to 0` with drift ordering, the linear dispersion relation + of this model reduces to that of the drift-kinetic model (up to finite-Larmor-radius corrections of order + :math:`k_\perp^2 \rho_i^2`), which makes it a reference for verifying guiding-center models. The time step must + resolve the cyclotron motion, :math:`\Delta t \lesssim \varepsilon`. In the quasi-neutral limit it must also + resolve the ion-acoustic frequency of the smallest cells, :math:`\Delta t\, c_s \pi / h \lesssim 2` + (there is no Debye cutoff).""" + + @classmethod + def doc_examples(cls): + r"""Create and initialize the full-orbit adiabatic-electron model: + + .. code-block:: python + + from struphy.models import VlasovPoissonElectrostaticAdiabatic + + model = VlasovPoissonElectrostaticAdiabatic() + model.em_fields.phi + model.kinetic_ions.var + + # quasi-neutral (default) or with the vacuum term, optionally with the zonal correction + model.propagators.poisson.options = model.propagators.poisson.Options( + quasi_neutral=False, which_geometry="cylindrical" + ) + """ + + @classmethod + def doc_use_cases(cls): + r"""This model is appropriate for: + + - electrostatic ion dynamics with adiabatic electrons, including cyclotron and finite-Larmor-radius effects + - ion-acoustic waves, ion Bernstein waves, drift waves and slab ITG modes in the full-orbit description + - verification of drift-kinetic and gyrokinetic models in the strong-field limit""" + + @classmethod + def doc_cannot_be_used_for(cls): + r"""This model is not suitable for: + + - electromagnetic dynamics with evolving magnetic perturbations + - electron kinetic effects beyond the adiabatic closure + - long simulations in strongly magnetized plasmas, where resolving the cyclotron motion is too expensive + - multi-species kinetic coupling without extending the model""" diff --git a/src/struphy/propagators/__init__.py b/src/struphy/propagators/__init__.py index 5dbdca9e8..a4b34c03e 100644 --- a/src/struphy/propagators/__init__.py +++ b/src/struphy/propagators/__init__.py @@ -4,6 +4,7 @@ # class name -> module defining it, resolved on first access by __getattr__ below _LAZY_IMPORTS = { "AdiabaticPhi": "struphy.propagators.adiabatic_phi", + "AdiabaticPoissonSolve": "struphy.propagators.adiabatic_poisson_solve", "CurlCurlSolve": "struphy.propagators.curl_curl_solve", "CurrentCoupling5DCurlb": "struphy.propagators.current_coupling_5d_curlb", "CurrentCoupling5DDensity": "struphy.propagators.current_coupling_5d_density", @@ -53,6 +54,7 @@ if TYPE_CHECKING: # static analysis and IDEs see the eager imports from struphy.propagators.adiabatic_phi import AdiabaticPhi + from struphy.propagators.adiabatic_poisson_solve import AdiabaticPoissonSolve from struphy.propagators.curl_curl_solve import CurlCurlSolve from struphy.propagators.current_coupling_5d_curlb import CurrentCoupling5DCurlb from struphy.propagators.current_coupling_5d_density import CurrentCoupling5DDensity @@ -103,6 +105,7 @@ __all__ = [ "AdiabaticPhi", + "AdiabaticPoissonSolve", "CurlCurlSolve", "CurrentCoupling5DCurlb", "CurrentCoupling5DDensity", diff --git a/src/struphy/propagators/adiabatic_poisson_solve.py b/src/struphy/propagators/adiabatic_poisson_solve.py new file mode 100644 index 000000000..88f276091 --- /dev/null +++ b/src/struphy/propagators/adiabatic_poisson_solve.py @@ -0,0 +1,205 @@ +from dataclasses import dataclass +from typing import Callable, Literal + +from feectools.linalg.stencil import StencilVector + +from struphy.io.options import LiteralOptions, OptionsBase +from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable +from struphy.pic.accumulation.filter import FilterParameters +from struphy.pic.accumulation.particles_to_grid import ParticlesToGrid +from struphy.propagators.implicit_diffusion import ImplicitDiffusion +from struphy.utils.utils import check_option + + +class AdiabaticPoissonSolve(ImplicitDiffusion): + r""" + Weak discretization of the Poisson equation for full-orbit ions with adiabatic electrons. + + Find :math:`\phi \in H^1` such that + + .. math:: + + \frac{1}{Z\varepsilon^2} \int_\Omega \frac{n_0}{T_0}\, \psi\, (\phi - \langle \phi \rangle)\,\textrm d \mathbf x + + \frac{1}{\alpha^2} \int_\Omega \nabla \psi \cdot \nabla \phi \,\textrm d \mathbf x + = \frac{1}{\varepsilon} \int_\Omega \psi\, (n_i - n_0)\,\textrm d \mathbf x \qquad \forall \ \psi \in H^1\,, + + where :math:`n_i = \int f \,\textrm d^3 \mathbf v` is the ion density, :math:`n_0` and :math:`T_0` are the + equilibrium ion density and electron temperature, :math:`Z` is the ion charge number, + :math:`\varepsilon = 1/(\hat \Omega_c \hat t)` and :math:`\alpha = \hat \omega_p / \hat \Omega_c`. + The first term is the linearized Boltzmann response of the electrons. The optional flux-surface average + :math:`\langle \phi \rangle` (``which_geometry``) is the zonal correction of + :class:`~struphy.propagators.gyrokinetic_poisson_solve.GyrokineticPoissonSolve`: electrons do not respond + to the averaged potential. The second term is the vacuum (Debye) term; it is dropped in the quasi-neutral + limit :math:`\alpha \to \infty` (``quasi_neutral=True``), where the equation becomes a weighted :math:`L^2` + projection. Boundary terms from integration by parts are assumed to vanish. + + In the quasi-neutral limit the ion-acoustic branch :math:`\omega = k c_s` has no Debye cutoff, so an explicit + particle push needs :math:`\Delta t\, c_s k_\textnormal{max} \lesssim 2`, with :math:`k_\textnormal{max} \approx \pi / h` + for the smallest cell size :math:`h`; refining the grid requires a smaller time step. + + The equation is discretized as (multiplied by :math:`\alpha^2` if ``quasi_neutral=False``) + + .. math:: + + \left( \frac{1}{Z\varepsilon^2}\,\mathbb M^0_{\textnormal{ad},T} (\mathbb I - \mathbb A) + + \frac{1}{\alpha^2} \mathbb G^\top \mathbb M^1 \mathbb G \right)\, \boldsymbol\phi + = \frac{1}{\varepsilon} (\Lambda^0, n_i - n_0 )_{L^2}\,, + + where :math:`\mathbb M^0_{\textnormal{ad},T}` is the mass matrix weighted with :math:`n_0/T_0` + (``stab_mat="M0ad_withT"``) and :math:`\mathbb A` is the weighted average projection + (:math:`\mathbb A = 0` if ``which_geometry=None``). + The particle density is deposited with the marker weights; the background :math:`n_0` is subtracted + either by the control variate or, for full-f markers, by an explicit background source + (see :class:`~struphy.propagators.implicit_diffusion.ImplicitDiffusion`). + """ + + def __init__( + self, + rho: FEECVariable | Callable | ParticlesToGrid | list = None, + rho_coeffs: float | list = None, + epsilon: float = 1.0, + alpha: float = 1.0, + Z: int = 1, + diagnostic: FEECVariable | None = None, + ): + """ + Parameters + ---------- + see ImplicitDiffusion.__init__ docstring for other parameters. + + epsilon : float + Inverse normalized cyclotron frequency :math:`1/(\\hat \\Omega_c \\hat t)`. + + alpha : float + Ratio of plasma to cyclotron frequency :math:`\\hat \\omega_p / \\hat \\Omega_c` + (only used if ``quasi_neutral=False``). + + Z : int + Charge number of the ions. + """ + super().__init__(rho=rho, rho_coeffs=rho_coeffs, diagnostic=diagnostic) + self.epsilon = epsilon + self.alpha = alpha + self.Z = Z + + @dataclass(repr=False) + class Options(OptionsBase): + r"""Configuration options for :class:`AdiabaticPoissonSolve`. + + Parameters + ---------- + quasi_neutral : bool, default=True + If ``True``, drop the vacuum term :math:`\alpha^{-2}\Delta \phi` (limit :math:`\alpha \to \infty`); + the system is then a weighted mass-matrix solve. + + which_geometry : {None, "cylindrical", "toroidal"}, default=None + Zonal correction :math:`\phi - \langle \phi \rangle` in the adiabatic electron response: + no average (``None``), average over ``eta3`` (``"cylindrical"``) or over ``eta2`` and ``eta3`` + (``"toroidal"``), as in + :class:`~struphy.propagators.gyrokinetic_poisson_solve.GyrokineticPoissonSolve`. + Requires ``quasi_neutral=False``: without the vacuum term the averaged potential is undetermined. + + stab_mat : {"M0", "M0ad", "M0ad_withT", "Id"}, default="M0ad_withT" + Mass matrix of the adiabatic term (``"M0ad_withT"`` has the weight :math:`n_0/T_0`). + + x0 : StencilVector, default=None + Initial guess for the iterative linear solver. + + solver : LiteralOptions.OptsSymmSolver, default="pcg" + Name of the symmetric iterative solver. + + precond : LiteralOptions.OptsDiffusionPrecond, default=None + Name of the preconditioner, see :class:`ImplicitDiffusion`. ``"MultiGrid"`` and + ``"StiffnessPreconditioner"`` need ``quasi_neutral=False``. + + precond_params : dict, default=None + Keyword arguments passed to the constructor of the preconditioner, see :class:`ImplicitDiffusion`. + + multigrid : MultiGridOptions, default=None + Options of the multigrid preconditioner (if ``precond="MultiGrid"``). + + solver_params : SolverParameters, default=None + Iterative-solver controls. If ``None``, defaults to ``SolverParameters()``. + + filter_params : dict[PICVariable | SPHVariable, FilterParameters], default=None + If not None, specifies a filter to the accumulation of a specific variable. + + Notes + ----- + ``AdiabaticPoissonSolve`` reuses :class:`ImplicitDiffusion` with ``sigma_2 = 0``, ``divide_by_dt = False``, + ``sigma_1 = 1 / (Z * epsilon**2)``, ``sigma_3 = 1 / epsilon`` and ``diffusion_mat=None`` (quasi-neutral), + or ``sigma_1 = alpha**2 / (Z * epsilon**2)``, ``sigma_3 = alpha**2 / epsilon`` and ``diffusion_mat="M1"``. + """ + + # specific literals + OptsStabMat = Literal["M0", "M0ad", "M0ad_withT", "Id"] + OptsGeometry = Literal["cylindrical", "toroidal", None] + # propagator options + quasi_neutral: bool = True + which_geometry: OptsGeometry = None + stab_mat: OptsStabMat = "M0ad_withT" + x0: StencilVector = None + solver: LiteralOptions.OptsSymmSolver = "pcg" + precond: LiteralOptions.OptsDiffusionPrecond = None + precond_params: dict = None + multigrid: MultiGridOptions = None + solver_params: SolverParameters = None + filter_params: dict[PICVariable | SPHVariable, FilterParameters] = None + # the adiabatic term is physical, not a stabilization: the rhs is not projected + enforce_compatibility: bool = False + + def __post_init__(self): + # checks + check_option(self.stab_mat, self.OptsStabMat) + check_option(self.solver, LiteralOptions.OptsSymmSolver) + check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) + check_option(self.which_geometry, self.OptsGeometry) + if self.precond is not None: + assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'." + if self.quasi_neutral: + assert self.which_geometry is None, ( + "The zonal correction (which_geometry) needs quasi_neutral=False: " + "without the vacuum term the averaged potential is undetermined." + ) + assert self.precond not in ("MultiGrid", "StiffnessPreconditioner"), ( + f"precond={self.precond!r} needs quasi_neutral=False." + ) + + # defaults + if self.precond_params is None: + self.precond_params = {} + if self.solver_params is None: + self.solver_params = SolverParameters() + if self.multigrid is None: + self.multigrid = MultiGridOptions() + + @property + def stab_average(self) -> str | None: + """Directions averaged in the adiabatic term, see ``which_geometry``.""" + return {None: None, "cylindrical": "eta3", "toroidal": "eta2 eta3"}[self.which_geometry] + + @property + def diffusion_mat(self) -> str | None: + """Diffusion matrix of :class:`ImplicitDiffusion`: none in the quasi-neutral limit.""" + return None if self.quasi_neutral else "M1" + + def allocate(self): + scale = 1.0 if self.options.quasi_neutral else self.alpha**2 + self.options.sigma_1 = scale / (self.Z * self.epsilon**2) + self.options.sigma_2 = 0.0 + self.options.sigma_3 = scale / self.epsilon + self.options.divide_by_dt = False + super().allocate() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + self._options = new diff --git a/src/struphy/propagators/gyrokinetic_poisson_solve.py b/src/struphy/propagators/gyrokinetic_poisson_solve.py index b5bc5a89e..d7ab0fb98 100644 --- a/src/struphy/propagators/gyrokinetic_poisson_solve.py +++ b/src/struphy/propagators/gyrokinetic_poisson_solve.py @@ -15,49 +15,54 @@ class GyrokineticPoissonSolve(ImplicitDiffusion): r""" - Weak discretization of the Poisson equation, with adabatic response of electrons. + Weak discretization of the Poisson equation, with adiabatic response of electrons. Find :math:`\phi \in H^1` such that .. math:: - \frac{1}{Z\epsilon^2} \int_\Omega \frac{n_0}{T_0} \psi\, \phi\,\textrm d \mathbf x + \int_\Omega \frac{n_0}{|B_0|²} \nabla \psi^\top \,\mathbb D \nabla \phi \,\textrm d \mathbf x = \sum_i \int_\Omega \psi\, \rho_i(\mathbf x)\,\textrm d \mathbf x \qquad \forall \ \psi \in H^1\,, - - where :math:`\mathbb D` can be an anisotropic (gyrokinetic) diffusion operator (or the identity), - :math:`\epsilon \in \mathbb R` is the gyrokinetic ratio defined in units, and Z the charge number of ions. + \frac{1}{Z\epsilon^2} \int_\Omega \frac{n_0}{T_0} \psi\, (\phi - \langle \phi \rangle)\,\textrm d \mathbf x + \int_\Omega \frac{n_0}{|B_0|^2} \nabla \psi^\top \,\mathbb D \nabla \phi \,\textrm d \mathbf x = \frac{1}{\epsilon} \sum_i c_i \int_\Omega \psi\, \rho_i(\mathbf x)\,\textrm d \mathbf x \qquad \forall \ \psi \in H^1\,, + + where :math:`c_i` are the ``rho_coeffs``, :math:`\mathbb D` can be an anisotropic (gyrokinetic) diffusion operator + (or the identity; the weight :math:`n_0/|B_0|^2` is that of ``diffusion_mat="M1gyro"``, ``"M1"`` and ``"M1perp"`` + are unweighted), + :math:`\epsilon \in \mathbb R` is the gyrokinetic ratio defined in units, Z the charge number of ions and + :math:`\langle \phi \rangle` the (weighted) average over :math:`\eta_3` or over :math:`\eta_2, \eta_3` + (``which_geometry``, or zero for ``which_geometry=None``): the electrons do not respond to the flux-surface + averaged potential. + For a full-f particle source (no control variate) the neutralizing background :math:`-n_0` is added to + :math:`\rho_i` automatically (see :class:`~struphy.propagators.implicit_diffusion.ImplicitDiffusion`). Boundary terms from integration by parts are assumed to vanish. The equation is discretized as .. math:: - \left( \frac{1}{Z\epsilon^2}\,\mathbb M^0_ad + \mathbb G^\top \mathbb M^1_\textnormal{gyro} \mathbb G \right)\, \boldsymbol\phi^{n+1} = \sum_i(\Lambda^0, \rho_i )_{L^2}\,, + \left( \frac{1}{Z\epsilon^2}\,\mathbb M^0_{\textnormal{ad},T} (\mathbb I - \mathbb A) + \mathbb G^\top \mathbb M^1_\textnormal{gyro} \mathbb G \right)\, \boldsymbol\phi^{n+1} = \frac{1}{\epsilon} \sum_i c_i (\Lambda^0, \rho_i )_{L^2}\,, - where :math:`\mathbb M^1_\textnormal{gyro}` is the gyrokinetic mass matrix - and :math:`\mathbb S` is a stabilization matrix. + where :math:`\mathbb M^0_{\textnormal{ad},T}` is the mass matrix weighted with :math:`n_0/T_0` (``stab_mat``, + default ``"M0ad_withT"``), :math:`\mathbb A` the weighted average projection and + :math:`\mathbb M^1_\textnormal{gyro}` the gyrokinetic mass matrix (``diffusion_mat``). Parameters ---------- - phi : StencilVector - FE coefficients of the solution as a discrete 0-form. - - stab_eps : float - Stabilization parameter multiplied on stab_mat (default=0.0). - - stab_mat : str - Name of the stabilizing matrix. - rho : FEECVariable or Callable or ParticlesToGrid or list (List of) right-hand side source term(s) of the Poisson problem (optional, can be set with a setter later). See :class:`~struphy.propagators.implicit_diffusion.ImplicitDiffusion` for the accepted entries; particle sources are passed as :class:`~struphy.pic.accumulation.particles_to_grid.ParticlesToGrid`. - x0 : StencilVector - Initial guess for the iterative solver (optional, can be set with a setter later). + rho_coeffs : float or list + Multiplicative coefficient(s) applied to ``rho`` (default 1.0 for all sources). + + epsilon : float + Gyrokinetic parameter. + + Z : int + Charge number of the ions. - solver : dict - Parameters for the iterative solver (see ``__init__`` for details). + diagnostic : FEECVariable, optional + If given, the L2 projection of the right-hand side is stored in it. """ def __init__( @@ -89,45 +94,28 @@ class Options: Parameters ---------- - epsilon : float, default=1.0 - Gyrokinetic parameter - - Z : float, default=1.0 - Charge number of ions - - stab_mat : {"M0", "M0ad", "Id"}, default="Id" - Stabilization matrix multiplied by ``stab_eps``. + stab_mat : {"M0", "M0ad", "Id", "M0ad_withT"}, default="M0ad_withT" + Mass matrix of the adiabatic term, multiplied by ``1 / (Z * epsilon**2)``. - ``"M0"``: standard weighted 0-form mass operator. - - ``"M0ad"``: adiabatic-electron weighted 0-form mass operator. + - ``"M0ad"``: adiabatic-electron weighted 0-form mass operator (weight :math:`n_0`). - ``"Id"``: identity operator. + - ``"M0ad_withT"``: adiabatic-electron weighted 0-form mass operator (weight :math:`n_0/T_0`). - which_geometry: {"cylindrical", "toroidal"}, default="cylindrical" + which_geometry: {"cylindrical", "toroidal", None}, default="cylindrical" Geometry of the problem, determines the meaning of `<\phi>` in the stabilization term: the average over ``eta3`` (cylindrical) or over ``eta2`` and ``eta3`` (toroidal), i.e. - ``stab_average`` of :class:`ImplicitDiffusion` is ``"eta3"`` or ``"eta2 eta3"``. The average + ``stab_average`` of :class:`ImplicitDiffusion` is ``"eta3"`` or ``"eta2 eta3"``. With ``None`` there is + no average, :math:`\langle \phi \rangle = 0`: the electrons respond adiabatically to all modes (e.g. for + a comparison with the full-orbit model in a slab, see + :class:`~struphy.propagators.adiabatic_poisson_solve.AdiabaticPoissonSolve`). The average is weighted with ``stab_mat`` (for equilibrium density and temperature constant on the averaged surfaces, the flux-surface average with the volume element), which keeps the operator symmetric (:class:`~struphy.feec.mass.WeightedAverageProjection`). - diffusion_mat : {"M1", "M1perp", "M1gyro"}, defaults="M1gyro" + diffusion_mat : {"M1", "M1perp", "M1gyro"}, default="M1gyro" Diffusion matrix. - rho : FEECVariable or Callable or ParticlesToGrid or list, default=None - Right-hand side source term(s) of the Poisson problem. - Accepted entries are: - - - ``None``: zero source. - - ``FEECVariable`` in ``H1``. - - ``Callable`` to be projected to ``H1`` via ``L2Projector``. - - :class:`~struphy.pic.accumulation.particles_to_grid.ParticlesToGrid`, describing a - particle-to-grid (charge/current) deposition. - - a ``list`` containing any mix of the entries above. - - rho_coeffs : float or list, default=None - Multiplicative coefficient(s) applied to ``rho``. - If ``None``, coefficients default to ``1.0`` for all sources. - x0 : StencilVector, default=None Initial guess for the iterative linear solver. @@ -168,7 +156,7 @@ class Options: # specific literals OptsStabMat = Literal["M0", "M0ad", "Id", "M0ad_withT"] OptsDiffusionMat = Literal["M1", "M1perp", "M1gyro"] - OptsGeometry = Literal["cylindrical", "toroidal"] + OptsGeometry = Literal["cylindrical", "toroidal", None] # propagator options stab_mat: OptsStabMat = "M0ad_withT" diffusion_mat: OptsDiffusionMat = "M1gyro" @@ -201,9 +189,9 @@ def __post_init__(self): self.multigrid = MultiGridOptions() @property - def stab_average(self) -> str: + def stab_average(self) -> str | None: """Directions averaged in the stabilization, see ``which_geometry``.""" - return {"cylindrical": "eta3", "toroidal": "eta2 eta3"}[self.which_geometry] + return {"cylindrical": "eta3", "toroidal": "eta2 eta3", None: None}[self.which_geometry] def allocate(self): epsilon = self.epsilon diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index f05598256..f6638ecf2 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -17,6 +17,7 @@ from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.pic.accumulation.filter import FilterParameters from struphy.pic.accumulation.particles_to_grid import AccumulatorVector, ParticlesToGrid +from struphy.pic.particles import DeltaFParticles6D from struphy.propagators.base import Propagator from struphy.utils.utils import check_option @@ -76,14 +77,28 @@ class ImplicitDiffusion(Propagator): Multiplicative coefficient(s) for ``rho`` sources. If a scalar is provided, it is applied to a single source. If a sequence is provided, it has one entry per entry of ``rho`` - (the neutralising background of a full-f accumulator gets the coefficient - of its accumulator), or one entry per collected source. + (the neutralising background of a full-f ``"H1"`` accumulator gets the coefficient + of its accumulator), or one entry per collected source (no background source is collected + with ``delta_f_source=True``, the control variate or delta-f markers). If ``None``, all coefficients default to ``1.0``. diagnostic : FEECVariable, default=None If not None, updates at each call to the propagator, takes the value of the right-hand side. Otherwise does not provide diagnostic. + delta_f_source : bool, default=False + Deposit particle sources with the delta-f weights :math:`(f - f_0)/s` also for full-f markers (without + the control variate): the weights are updated before and the full-f weights restored after the deposit, + and no background source is added. The Monte-Carlo noise is then that of :math:`f - f_0` instead of + :math:`f`, which is much smaller for small perturbations; meant for the initial Poisson solve of full-f models. + + Note + ---- + The deposited particle density is :math:`f - f_0` with the control variate (the weights are updated before + each deposit), with delta-f markers (:class:`~struphy.pic.particles.DeltaFParticles6D`) and with + ``delta_f_source=True``. Otherwise (full-f), for ``"H1"`` accumulators the neutralizing background + :math:`-n_0` is added as a separate source with the coefficient of the particle source. + Notes ----- @@ -118,6 +133,7 @@ def __init__( rho: FEECVariable | ParticlesToGrid | Callable | list = None, rho_coeffs: float | list = None, diagnostic: FEECVariable | None = None, + delta_f_source: bool = False, ): if isinstance(rho, list): for r in rho: @@ -132,6 +148,7 @@ def __init__( self.rho = rho self.rho_coeffs = rho_coeffs self.diagnostic = diagnostic + self.delta_f_source = delta_f_source @dataclass(repr=False) class Options(OptionsBase): @@ -162,12 +179,13 @@ class Options(OptionsBase): If ``True``, divide ``sigma_1``, ``sigma_2`` and ``sigma_3`` by the time step passed to ``__call__(dt)`` before assembling the system. - stab_mat : {"M0", "M0ad", "Id"}, default="M0" + stab_mat : {"M0", "M0ad", "M0ad_withT", "Id"}, default="M0" Stabilization operator used in the term weighted by ``sigma_1`` and ``sigma_2``. - ``"M0"``: standard weighted 0-form mass operator. - - ``"M0ad"``: adiabatic-electron weighted 0-form mass operator. + - ``"M0ad"``: adiabatic-electron weighted 0-form mass operator (weight :math:`n_0`). + - ``"M0ad_withT"``: adiabatic-electron weighted 0-form mass operator (weight :math:`n_0/T_0`). - ``"Id"``: identity operator. stab_average : str, default=None @@ -177,12 +195,14 @@ class Options(OptionsBase): average with the volume element for ``"M0"``), such that the operator stays symmetric. With ``precond="MultiGrid"``, the average is re-built on every coarse level. - diffusion_mat : {"M1", "M1perp"} or WeightedMassOperator, default="M1" + diffusion_mat : {"M1", "M1perp", None} or WeightedMassOperator, default="M1" Diffusion metric in the bilinear form ``grad.T @ diffusion_mat @ grad``. You can pass the name of a pre-built operator in ``mass_ops`` or a custom ``WeightedMassOperator`` compatible with the codomain of - ``grad``. + ``grad``. With ``None`` there is no diffusion term and the system is + ``sigma_1 * stab_mat`` only (a weighted L2 projection); then ``precond`` cannot be + ``"MultiGrid"`` or ``"StiffnessPreconditioner"``. x0 : StencilVector, default=None Initial guess for the iterative linear solver. @@ -254,8 +274,8 @@ class Options(OptionsBase): """ # specific literals - OptsStabMat = Literal["M0", "M0ad", "Id"] - OptsDiffusionMat = Literal["M1", "M1perp"] + OptsStabMat = Literal["M0", "M0ad", "M0ad_withT", "Id"] + OptsDiffusionMat = Literal["M1", "M1perp", None] # propagator options sigma_1: float = 1.0 sigma_2: float = 0.0 @@ -281,6 +301,13 @@ def __post_init__(self): check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) if self.precond is not None: assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'." + if self.diffusion_mat is None: + assert self.precond not in ("MultiGrid", "StiffnessPreconditioner"), ( + f"precond={self.precond!r} needs a diffusion term (diffusion_mat=None)." + ) + assert self.stab_average is None, ( + "Without a diffusion term the averaged part of phi is undetermined (stab_average must be None)." + ) preconditioner._parse_directions(self.stab_average) # defaults @@ -343,11 +370,13 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: filter_params=filter_params, ) - if not rhs.particles.control_variate and rhs.space_id == "H1": + if not self._deposits_delta_f(rhs.particles) and rhs.space_id == "H1": l2_proj = L2Projector("H1", Propagator.mass_ops) - f0e = -rho.pic_variable.species.charge_number * rhs.particles.f0 - # neutralising background: get_dofs already returns (n0, Lambda^0)_L2, - # i.e. the weak-form rhs, so it is added as a StencilVector (no extra M0) + # neutralising background -n0: the accumulator deposits the particle density (weights) and the + # charge (number and sign) is in the coefficient of the source, which the background shares + f0e = -1 * rhs.particles.f0 + # get_dofs already returns (n0, Lambda^0)_L2, i.e. the weak-form rhs, + # so it is added as a StencilVector (no extra M0) rho_eh = l2_proj.get_dofs(f0e.n, apply_bc=True) return [rhs, rho_eh] @@ -398,7 +427,9 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: avg = WeightedAverageProjection(self.derham, stab_mat, average_dirs) stab_mat = stab_mat @ (IdentityOperator(stab_mat.domain, stab_mat.codomain) - avg) - if isinstance(self.options.diffusion_mat, str): + if self.options.diffusion_mat is None: + diffusion_mat = None + elif isinstance(self.options.diffusion_mat, str): diffusion_mat = getattr(self.mass_ops, self.options.diffusion_mat) else: diffusion_mat = self.options.diffusion_mat @@ -408,7 +439,10 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: # Set lhs matrices (without dt) self._stab_mat = stab_mat - self._diffusion_op = self.derham.grad.T @ diffusion_mat @ self.derham.grad + if diffusion_mat is None: + self._diffusion_op = None + else: + self._diffusion_op = self.derham.grad.T @ diffusion_mat @ self.derham.grad # preconditioner and solver for Ax=b # the multigrid and stiffness preconditioners depend on sigma_1, they are updated in __call__ @@ -439,7 +473,7 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: # (only "pcg" takes a preconditioner) pc_kwargs = {"pc": pc} if self.options.solver == "pcg" else {} self._solver = inverse( - self._diffusion_op, + stab_mat if self._diffusion_op is None else self._diffusion_op, self.options.solver, **pc_kwargs, x0=self.x0, @@ -465,11 +499,22 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: self._m0_ones = self.mass_ops.M0.dot(self._ones) self._ones_m0_ones = self._ones.inner(self._m0_ones) + @staticmethod + def _is_full_f(particles) -> bool: + """True if the marker weights are full-f weights f / s (no control variate, no delta-f markers).""" + return not particles.control_variate and not isinstance(particles, DeltaFParticles6D) + + def _deposits_delta_f(self, particles) -> bool: + """True if the particle source deposits f - f0 (so that no neutralizing background is needed).""" + return not self._is_full_f(particles) or self.delta_f_source + @property def operator_is_singular(self) -> bool: """True if the diffusion operator ``grad.T @ diffusion_mat @ grad`` is known to be singular with the constant functions (coefficient vector of all ones) in its kernel: no Dirichlet boundary conditions and tensor-product splines. Polar splines are conservatively treated as regular.""" + if self._diffusion_op is None: + return False no_dirichlet = not any(any(bc) for bc in self.derham.dirichlet_bc) return ( no_dirichlet @@ -548,7 +593,9 @@ def __call__(self, dt): v = src.spline.vector self._rhs2 += sig_3 * coeff * self.mass_ops.M0.dot(v, out=self._tmp_src) elif isinstance(src, AccumulatorVector): - if src.particles.control_variate: + # delta-f weights for full-f markers, restored after the deposit (delta_f_source) + restore_full_f = self.delta_f_source and self._is_full_f(src.particles) + if src.particles.control_variate or restore_full_f: src.particles.update_weights() if src.space_id == "H1": @@ -569,6 +616,9 @@ def __call__(self, dt): else: raise ValueError(f"Unsupported source space {src.space_id}.") + if restore_full_f: + src.particles.weights = src.particles.weights0.copy() + self._rhs2 += sig_3 * coeff * vec rhs += self._rhs2 @@ -582,7 +632,10 @@ def __call__(self, dt): self.diagnostic.spline.vector = proj.solve(rhs) # compute lhs - self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op + if self._diffusion_op is None: + self._solver.linop = sig_1 * self._stab_mat + else: + self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op if sig_1 != self._pc_sig_1: if self._mg is not None: self._mg.update(self._solver.linop) diff --git a/src/struphy/propagators/poisson_solve.py b/src/struphy/propagators/poisson_solve.py index 79137400c..c0b688be7 100644 --- a/src/struphy/propagators/poisson_solve.py +++ b/src/struphy/propagators/poisson_solve.py @@ -166,6 +166,7 @@ def __init__( self, rho: FEECVariable | Callable | ParticlesToGrid | list = None, rho_coeffs: float | list = None, + delta_f_source: bool = False, ): """ Parameters @@ -184,8 +185,11 @@ def __init__( rho_coeffs : float or list, default=None Multiplicative coefficient(s) applied to ``rho``. If ``None``, coefficients default to ``1.0`` for all sources. + + delta_f_source : bool, default=False + Deposit particle sources with delta-f weights also for full-f markers, see :class:`ImplicitDiffusion`. """ - super().__init__(rho=rho, rho_coeffs=rho_coeffs) + super().__init__(rho=rho, rho_coeffs=rho_coeffs, delta_f_source=delta_f_source) @property def options(self) -> Options: diff --git a/src/struphy/propagators/tests/test_adiabatic_poisson.py b/src/struphy/propagators/tests/test_adiabatic_poisson.py new file mode 100644 index 000000000..b48053c53 --- /dev/null +++ b/src/struphy/propagators/tests/test_adiabatic_poisson.py @@ -0,0 +1,211 @@ +import logging + +import cunumpy as xp +import pytest +from maybempi import MPI + +from struphy import domains, equils +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable +from struphy.propagators.adiabatic_poisson_solve import AdiabaticPoissonSolve +from struphy.propagators.base import Propagator +from struphy.propagators.gyrokinetic_poisson_solve import GyrokineticPoissonSolve +from struphy.propagators.implicit_diffusion import ImplicitDiffusion +from struphy.topology.grids import TensorProductGrid + +logger = logging.getLogger("struphy") + +comm = MPI.COMM_WORLD + +Lx, Ly, Lz = 2.0, 1.0, 3.0 + +# HomogenSlab with p = beta / 2 = 1 and n0 = 2, i.e. T0 = p / n0 = 0.5 and n0 / T0 = 4 +N0 = 2.0 +N0_OVER_T0 = 4.0 + + +def _setup(num_elements, degree): + """Periodic Cuboid with a HomogenSlab equilibrium; sets the class attributes of Propagator.""" + domain = domains.Cuboid(l1=0.0, r1=Lx, l2=0.0, r2=Ly, l3=0.0, r3=Lz) + eq_mhd = equils.HomogenSlab(B0z=1.0, beta=2.0, n0=N0) + eq_mhd.domain = domain + + grid = TensorProductGrid(num_elements=num_elements) + derham = Derham(grid, DerhamOptions(degree=degree), comm=comm) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = WeightedMassOperators(derham, domain, eq_mhd=eq_mhd) + + phi = FEECVariable(space="H1") + phi.allocate(derham=derham, domain=domain) + return domain, phi + + +def _max_error(domain, phi, sol_xyz): + e1 = xp.linspace(0.0, 1.0, 21) + e3 = xp.linspace(0.0, 1.0, 21) + num = domain.push(phi.spline, e1, 0.5, e3, kind="0") + x, y, z = domain(e1, 0.5, e3) + exact = sol_xyz(x, y, z) + return xp.max(xp.abs(num - exact)) / xp.max(xp.abs(exact)) + + +def _pulled(domain, fun_xyz): + def fun(e1, e2, e3): + return domain.pull(fun_xyz, e1, e2, e3, kind="0", squeeze_out=False) + + return fun + + +SOLVER_PARAMS = SolverParameters(tol=1.0e-13, maxiter=3000) + + +def test_implicit_diffusion_without_diffusion_term(): + """With diffusion_mat=None and stab_mat="M0ad_withT" the solve is the weighted L2 projection + sigma_1 * n0/T0 * phi = rho.""" + domain, phi = _setup([16, 1, 1], [3, 1, 1]) + sigma_1 = 0.7 + + def rho_xyz(x, y, z): + return xp.sin(2 * xp.pi / Lx * x) + 0.5 + + def sol_xyz(x, y, z): + return rho_xyz(x, y, z) / (sigma_1 * N0_OVER_T0) + + solver = ImplicitDiffusion(rho=_pulled(domain, rho_xyz)) + solver.variables.phi = phi + solver.options = solver.Options( + sigma_1=sigma_1, + sigma_3=1.0, + stab_mat="M0ad_withT", + diffusion_mat=None, + precond="MassMatrixPreconditioner", # an M0-based preconditioner is allowed without diffusion + solver_params=SOLVER_PARAMS, + ) + solver.allocate() + assert not solver.operator_is_singular + solver(1.0) + + err = _max_error(domain, phi, sol_xyz) + logger.info(f"{err =}") + assert err < 1e-4 + + +@pytest.mark.parametrize("which_geometry", [None, "cylindrical"]) +def test_adiabatic_poisson_with_vacuum_term(which_geometry): + """Manufactured solution of (1/(Z eps^2)) n0/T0 (phi - ) - (1/alpha^2) Delta phi = rho / eps, + where is the average over eta3 (cylindrical) or zero.""" + domain, phi = _setup([16, 1, 16], [3, 1, 3]) + epsilon, alpha, Z = 0.5, 2.0, 2 + kx, kz = 2 * xp.pi / Lx, 2 * xp.pi / Lz + + def sol_xyz(x, y, z): + return xp.sin(kx * x) * (1.0 + xp.cos(kz * z)) + + def sol_avg_xyz(x, y, z): + return xp.sin(kx * x) if which_geometry == "cylindrical" else 0.0 * x + + def rho_xyz(x, y, z): + laplace = -(kx**2) * xp.sin(kx * x) * (1.0 + xp.cos(kz * z)) - kz**2 * xp.sin(kx * x) * xp.cos(kz * z) + adiabatic = N0_OVER_T0 / (Z * epsilon**2) * (sol_xyz(x, y, z) - sol_avg_xyz(x, y, z)) + return epsilon * (adiabatic - laplace / alpha**2) + + solver = AdiabaticPoissonSolve(rho=_pulled(domain, rho_xyz), epsilon=epsilon, alpha=alpha, Z=Z) + solver.variables.phi = phi + solver.options = solver.Options( + quasi_neutral=False, + which_geometry=which_geometry, + precond=None, + solver_params=SOLVER_PARAMS, + ) + solver.allocate() + solver(1.0) + + err = _max_error(domain, phi, sol_xyz) + logger.info(f"{which_geometry =}, {err =}") + assert err < 1e-3 + + +def test_adiabatic_poisson_quasi_neutral(): + """In the quasi-neutral limit phi = Z eps T0/n0 rho (default options).""" + domain, phi = _setup([16, 1, 1], [3, 1, 1]) + epsilon, Z = 0.1, 1 + + def rho_xyz(x, y, z): + return xp.cos(2 * xp.pi / Lx * x) + + def sol_xyz(x, y, z): + return Z * epsilon / N0_OVER_T0 * rho_xyz(x, y, z) + + solver = AdiabaticPoissonSolve(rho=_pulled(domain, rho_xyz), epsilon=epsilon, alpha=1e-3, Z=Z) + solver.variables.phi = phi + solver.allocate() + solver(1.0) + + err = _max_error(domain, phi, sol_xyz) + logger.info(f"{err =}") + assert err < 1e-4 + + +@pytest.mark.parametrize("which_geometry", [None, "cylindrical"]) +def test_gyrokinetic_poisson_average(which_geometry): + """GyrokineticPoissonSolve with and without the flux-surface average, for a solution that does not depend on + eta3: (1/(Z eps^2)) n0/T0 (phi - ) - div_perp(n0/|B0|^2 grad_perp phi) = rho / eps, with = phi for + the average over eta3 ("cylindrical") and = 0 without average (None).""" + domain, phi = _setup([16, 16, 1], [3, 3, 1]) + epsilon, Z = 0.5, 1 + kx, ky = 2 * xp.pi / Lx, 2 * xp.pi / Ly + # n0 / |B0|^2 of the HomogenSlab of _setup + n0_over_B2 = N0 + + def sol_xyz(x, y, z): + return xp.sin(kx * x) * xp.cos(ky * y) + + def rho_xyz(x, y, z): + adiabatic = N0_OVER_T0 / (Z * epsilon**2) * sol_xyz(x, y, z) if which_geometry is None else 0.0 * x + return epsilon * (adiabatic + n0_over_B2 * (kx**2 + ky**2) * sol_xyz(x, y, z)) + + solver = GyrokineticPoissonSolve(rho=_pulled(domain, rho_xyz), epsilon=epsilon, Z=Z) + solver.variables.phi = phi + solver.options = solver.Options(which_geometry=which_geometry, solver_params=SOLVER_PARAMS) + solver.allocate() + solver(1.0) + + e1 = xp.linspace(0.0, 1.0, 21) + e2 = xp.linspace(0.0, 1.0, 21) + num = domain.push(phi.spline, e1, e2, 0.5, kind="0") + x, y, z = domain(e1, e2, 0.5) + exact = sol_xyz(x, y, z) + err = xp.max(xp.abs(num - exact)) / xp.max(xp.abs(exact)) + logger.info(f"{which_geometry =}, {err =}") + assert err < 1e-3 + + +def test_invalid_options(): + # no diffusion term: the stiffness-based preconditioners are not available + with pytest.raises(AssertionError): + ImplicitDiffusion.Options(diffusion_mat=None, precond="MultiGrid") + with pytest.raises(AssertionError): + ImplicitDiffusion.Options(diffusion_mat=None, precond="StiffnessPreconditioner") + # no diffusion term: the averaged part of phi would be undetermined + with pytest.raises(AssertionError): + ImplicitDiffusion.Options(diffusion_mat=None, stab_average="eta3") + # the zonal correction leaves undetermined without the vacuum term + with pytest.raises(AssertionError): + AdiabaticPoissonSolve.Options(quasi_neutral=True, which_geometry="cylindrical") + with pytest.raises(AssertionError): + AdiabaticPoissonSolve.Options(quasi_neutral=True, precond="StiffnessPreconditioner") + # valid combinations + ImplicitDiffusion.Options(stab_mat="M0ad_withT", diffusion_mat=None) + AdiabaticPoissonSolve.Options(quasi_neutral=False, which_geometry="toroidal", precond="StiffnessPreconditioner") + + +if __name__ == "__main__": + test_implicit_diffusion_without_diffusion_term() + test_adiabatic_poisson_with_vacuum_term("cylindrical") + test_adiabatic_poisson_quasi_neutral() + test_invalid_options()