Add basic building blocks for a PETSc solver

.gitignore

@@ -112,3 +112,4 @@ pyvenv.cfg
 *profile_output*.txt
 *kernels.txt
 struphy.log
+struphy.log*

profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py

@@ -0,0 +1,199 @@
+import os
+
+# -----------------------------
+# Description of the simulation
+# -----------------------------
+
+description = """
+Periodic-slab variant of the ToyDrift model (the model behind
+examples/ToyGyrokinetic/diocotron_instability, which uses a physically non-periodic
+HollowCylinder domain -- radial confinement is inherent to the diocotron instability). This case
+swaps in a periodic Cuboid domain instead: unlike PoissonAdiabaticGyrokinetic (used by
+DriftKineticElectrostaticAdiabatic), ToyDrift's field solve is a plain PoissonSolve with no
+geometry-coupled averaging, so it works correctly on a periodic domain out of the box. Unlike
+VlasovAmpereOneSpecies (which only solves Poisson once, as an initial condition), gc_poisson runs
+as a *regular per-step propagator* here, so a single `sim.run(one_time_step=True)` call already
+times one full, representative solve.
+
+Grid: 32^3 elements, degree-3 splines (32768 dofs), with PETSc's preconditioner set explicitly to
+algebraic multigrid (`pc_type="gamg"`, via `SolverParameters.pc_type` -- see
+struphy.linear_algebra.solver.SolverParameters), instead of the default `"jacobi"`. This
+resolution is deliberately large enough to push feectools' unpreconditioned CG into several
+hundred iterations per solve.
+
+Measured via struphy.linear_algebra.petsc_examples_benchmark's repeated-solve methodology (which
+isolates just the solve, amortizing one-time matrix/preconditioner setup across several calls --
+see that module's docstring): feectools' unpreconditioned CG took ~10.9 s/solve (190 iterations)
+against PETSc+gamg's ~0.29 s/solve (2 iterations) here -- a ~38x difference, the largest gap in
+this benchmark suite. The near-singular Poisson system ToyDrift solves (`stab_eps` clamped to
+~1e-14 by ImplicitDiffusion's "always stabilize" logic, see PETScSolver's docstring) is
+ill-conditioned mainly through its weakly constrained constant/DC mode, not primarily through raw
+grid resolution, so pcg's iteration count does not grow much further with size in this regime --
+while gamg's convergence stays essentially grid-independent regardless. What *does* grow with
+size is the absolute cost: at this larger, more realistic problem size, pcg's ~11 seconds per
+Poisson solve is the practically relevant "huge difference" -- multiplied over the many timesteps
+of a real simulation, it is the difference between a run finishing in minutes and one taking hours.
+"""
+
+# ------------------
+# Import Struphy API
+# ------------------
+
+from struphy import (
+    BaseUnits,
+    BoundaryParameters,
+    DerhamOptions,
+    EnvironmentOptions,
+    LoadingParameters,
+    Simulation,
+    SortingParameters,
+    Time,
+    WeightsParameters,
+    domains,
+    equils,
+    grids,
+    maxwellians,
+    perturbations,
+)
+from struphy.linear_algebra.solver import SolverParameters
+
+# ---------------------
+# Instance of the model
+# ---------------------
+from struphy.models import ToyDrift
+
+# Units
+base_units = BaseUnits(kBT=1.0)
+
+# Model instance
+model = ToyDrift(base_units=base_units)
+
+# List all variables and decide whether to save their data
+model.em_fields.phi.save_data = True
+model.kinetic_ions.var.save_data = False
+
+# --------------------------
+# Instance of the simulation
+# --------------------------
+
+# `--id` distinguishes runs that share a rank count but differ in something else; the
+# profiling driver passes its launch counter (see `ProfilingJob.build_commands`).
+# Unknown flags are ignored so the driver can forward other parameters as well.
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.")
+parser.add_argument(
+    "--solver", type=str, default="pcg", choices=["pcg", "petsc"], help="Solver for the Poisson-type solve."
+)
+args, _ = parser.parse_known_args()
+
+# scope-profiler label: distinguishes solver/rank-count combinations in post-processing
+# (chart legends, `scope-profiler inspect`), see EnvironmentOptions.profiling_label.
+from feectools.ddm.mpi import mpi as MPI
+
+_comm = MPI.COMM_WORLD
+_num_ranks = _comm.Get_size() if _comm is not None else 1
+_profiling_label = f"{args.solver}, {_num_ranks} rank" + ("s" if _num_ranks != 1 else "")
+
+# Environment options
+env = EnvironmentOptions(
+    sim_folder=f"sim_{args.id:02d}",
+    out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()),
+    profiling_activated=True,
+    profiling_trace=True,
+    profiling_label=_profiling_label,
+)
+
+# Time stepping
+time_opts = Time(dt=0.05, Tend=0.05, split_algo="LieTrotter")
+
+# Geometry
+domain = domains.Cuboid()
+
+# Fluid equilibrium: straight B field, homogeneous density (default n0=1.0)
+equil = equils.HomogenSlab(B0z=1.0, n0=1.0)
+
+# Grid -- 32^3: large enough for pcg's iteration count (hence cost) to blow up while PETSc+gamg
+# stays flat.
+grid = grids.TensorProductGrid(num_elements=(32, 32, 32))
+
+# Derham options -- fully periodic (required: PETScSolver cannot (yet) assemble
+# DirectionalDerivativeOperator along a non-periodic axis, see
+# struphy.linear_algebra.petsc_solver._directional_derivative_to_stencil_matrix)
+derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(None, None, None))
+
+# Simulation object
+sim = Simulation(
+    model=model,
+    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=5, seed=42)
+weights_params = WeightsParameters(control_variate=True)
+boundary_params = BoundaryParameters()
+sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True)
+
+model.kinetic_ions.set_markers(
+    loading_params=loading_params,
+    weights_params=weights_params,
+    boundary_params=boundary_params,
+    sorting_params=sorting_params,
+    bufsize=0.4,
+)
+
+# ------------------
+# Propagator options
+# ------------------
+
+model.propagators.gc_poisson.options.solver = args.solver
+# pc_type only affects solver="petsc" (ignored by pcg, see SolverParameters.pc_type); "gamg" is
+# what makes PETSc's advantage large here. Note this only matters for the *steady-state*
+# per-solve cost (see the module docstring): a single sim.run(one_time_step=True) call, as this
+# file's __main__ performs, still pays gamg's one-time multigrid setup cost up front, so it will
+# not by itself reproduce the headline speedup above -- that requires several solves at the same
+# dt to amortize setup, exactly what petsc_examples_benchmark.py's repeated-solve timing does.
+model.propagators.gc_poisson.options.solver_params = SolverParameters(tol=1e-10, maxiter=5_000, pc_type="gamg")
+model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options(
+    algo="explicit",
+    evaluate_e_field=True,
+)
+
+# ------------------
+# Initial conditions
+# ------------------
+# Initial conditions are the sum of the background(s) and the perturbation(s).
+
+# Background for kinetic species
+background = maxwellians.GyroMaxwellian2D(
+    n=(1.0, None),
+    vth_para=(1.0, None),
+    vth_perp=(1.0, None),
+    equil=equil,
+)
+model.kinetic_ions.var.add_background(background)
+
+# Perturbation, matching the Landau-damping style used elsewhere in this benchmark suite
+perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,))
+init = maxwellians.GyroMaxwellian2D(
+    n=(1.0, perturbation),
+    vth_para=(1.0, None),
+    vth_perp=(1.0, None),
+    equil=equil,
+)
+model.kinetic_ions.var.add_initial_condition(init)
+
+if __name__ == "__main__":
+    # one_time_step=True isolates the (still per-step, unlike VlasovAmpereOneSpecies)
+    # gc_poisson solve for a single-step timing snapshot.
+    sim.run(one_time_step=True)

profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on.py

@@ -0,0 +1,184 @@
+import os
+
+# -----------------------------
+# Description of the simulation
+# -----------------------------
+# Please fill in a verbal description of the simulation.
+# It will be printed at the beginning of the simulation and can be used to keep track of the different runs.
+
+description = """
+Nonlinear bump-on-tail instability: A kinetic plasma instability test case for the Vlasov-Ampère model.
+This test features a "bump" (localized excess) in the high-velocity tail of the electron velocity distribution.
+The bump-on-tail configuration is unstable to the generation of Langmuir waves, leading to energy transfer
+from the hot electron population to the growing wave field. This nonlinear process exhibits complex dynamics
+including mode coupling and particle trapping in the wave potential.
+This benchmark validates the particle-in-cell treatment of velocity-space instabilities and wave-particle interactions.
+"""
+
+# ------------------
+# Import Struphy API
+# ------------------
+
+# For particles:
+from struphy import (
+    BaseUnits,
+    BinningPlot,
+    BoundaryParameters,
+    DerhamOptions,
+    EnvironmentOptions,
+    FieldsBackground,
+    KernelDensityPlot,
+    LoadingParameters,
+    SavingParameters,
+    Simulation,
+    SortingParameters,
+    Time,
+    WeightsParameters,
+    domains,
+    equils,
+    grids,
+    maxwellians,
+    perturbations,
+)
+
+# ---------------------
+# Instance of the model
+# ---------------------
+from struphy.models import VlasovAmpereOneSpecies
+
+# Units
+base_units = BaseUnits()
+
+# Model instance
+model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False)
+
+# List all variables and decide whether to save their data
+model.em_fields.e_field.save_data = True
+model.em_fields.phi.save_data = True
+model.kinetic_ions.var.save_data = True
+
+# --------------------------
+# Instance of the simulation
+# --------------------------
+
+# Environment options
+# `--id` distinguishes runs that share a rank count but differ in something else; the
+# profiling driver passes its launch counter (see `ProfilingJob.build_commands`).
+# Unknown flags are ignored so the driver can forward other parameters as well.
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.")
+parser.add_argument(
+    "--solver", type=str, default="pcg", choices=["pcg", "petsc"], help="Solver for the Poisson-type solve."
+)
+args, _ = parser.parse_known_args()
+
+# scope-profiler label: distinguishes solver/rank-count combinations in post-processing
+# (chart legends, `scope-profiler inspect`), see EnvironmentOptions.profiling_label.
+from feectools.ddm.mpi import mpi as MPI
+
+_comm = MPI.COMM_WORLD
+_num_ranks = _comm.Get_size() if _comm is not None else 1
+_profiling_label = f"{args.solver}, {_num_ranks} rank" + ("s" if _num_ranks != 1 else "")
+
+env = EnvironmentOptions(
+    sim_folder=f"sim_{args.id:02d}",
+    out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()),
+    profiling_activated=True,
+    profiling_trace=True,
+    profiling_label=_profiling_label,
+)
+
+# Time stepping
+time_opts = Time(dt=0.1, Tend=60.0, split_algo="LieTrotter")
+
+# Geometry
+domain = domains.Cuboid(r1=62.83)
+
+# Fluid equilibrium (can be used as part of initial conditions)
+equil = None
+
+# Grid
+grid = grids.TensorProductGrid(num_elements=(16, 16, 16))
+
+# Derham options
+derham_opts = DerhamOptions(degree=(3, 1, 1))
+
+# Simulation object
+sim = Simulation(
+    model=model,
+    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=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0))
+weights_params = WeightsParameters(control_variate=True)
+boundary_params = BoundaryParameters()
+sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True)
+
+binplot_1 = BinningPlot(
+    slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-10.0, 10.0))
+)  # for initial velocity distribution
+binplot_2 = BinningPlot(
+    slice="v1", n_bins=128, ranges=(-10.0, 10.0)
+)  # for progression of velocity and space distribution
+saving_params = SavingParameters(binning_plots=(binplot_1, binplot_2))
+
+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()
+if model.with_B0:
+    model.propagators.push_vxb.options = model.propagators.push_vxb.Options()
+model.propagators.coupling_va.options = model.propagators.coupling_va.Options()
+model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver=args.solver)
+
+# ------------------
+# Initial conditions
+# ------------------
+# Initial conditions are the sum of the background(s) and the perturbation(s).
+# If backgrounds or perturbations are not specified, they are assumed to be zero.
+
+# For kinetic species the background is mandatory.
+# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition.
+# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples).
+
+# Background for kinetic species
+maxwellian_1 = maxwellians.Maxwellian3D(n=(9 / 10, None), u1=(3.0, None))
+maxwellian_2 = maxwellians.Maxwellian3D(n=(1 / 10, None), u1=(-4.5, None), vth1=(0.5, None))
+background = maxwellian_1 + maxwellian_2
+model.kinetic_ions.var.add_background(background)
+
+# Perturbations for (some) kinetic species
+perturbation = perturbations.ModesCos(amps=(0.05,), ls=(1,))
+init1 = maxwellians.Maxwellian3D(n=(9 / 10, None), u1=(3.0, None))
+init2 = maxwellians.Maxwellian3D(n=(1 / 10, perturbation), u1=(-4.5, None), vth1=(0.5, None))
+init = init1 + init2
+model.kinetic_ions.var.add_initial_condition(init)
+
+if __name__ == "__main__":
+    # one_time_step=True isolates the initial Poisson solve (the only part of this model that
+    # solver= affects -- the field then evolves via VlasovAmpereCoupling, unrelated to
+    # PETScSolver) from the many identical transport steps a full run would otherwise dilute
+    # the comparison with.
+    sim.run(one_time_step=True)