diff --git a/.gitignore b/.gitignore index 8e8873977..a155451da 100644 --- a/.gitignore +++ b/.gitignore @@ -112,3 +112,4 @@ pyvenv.cfg *profile_output*.txt *kernels.txt struphy.log +struphy.log* diff --git a/profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py b/profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py new file mode 100644 index 000000000..4b968170f --- /dev/null +++ b/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) diff --git a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on.py b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on.py new file mode 100644 index 000000000..4b3cc063f --- /dev/null +++ b/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) diff --git a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py new file mode 100644 index 000000000..c89bd4aa9 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py @@ -0,0 +1,175 @@ +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 = """ +Strong (nonlinear) Landau damping: A nonlinear test case for the VlasovAmpereOneSpecies model. +This test involves a large amplitude electrostatic perturbation in a uniform, collisionless plasma. +Unlike weak Landau damping, the nonlinear regime exhibits trapping of particles in the potential wells +of the self-consistent electric field, leading to vortex formation and complex phase space structures. +This benchmark tests the ability of the particle-in-cell method to capture nonlinear kinetic effects +and validates the long-term stability and accuracy of the Vlasov-Ampère discretization. +""" + +# ------------------ +# 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 +# -------------------------- + +# `--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=75.0, split_algo="LieTrotter") + +# Geometry +domain = domains.Cuboid(r1=12.56) + +# 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() + +# 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) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) + +binplot = BinningPlot(slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-5.0, 5.0))) +saving_params = SavingParameters(binning_plots=(binplot,)) + +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 +background = maxwellians.Maxwellian3D(n=(1.0, None)) +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) +init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) +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 ~1500 identical transport steps a full Tend=75.0 run would otherwise + # dilute the comparison with. + sim.run(one_time_step=True) diff --git a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream.py b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream.py new file mode 100644 index 000000000..240722552 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream.py @@ -0,0 +1,179 @@ +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 two-stream instability: A fundamental kinetic test case for the Vlasov-Ampère model. +This test involves two counter-streaming particle populations with a small perturbation that triggers +the two-stream instability. The instability leads to the formation of electron acoustic waves and +subsequent nonlinear effects including particle trapping and energy exchange between modes. +This benchmark validates the numerical treatment of beam-plasma interactions and tests the accuracy +of the particle-in-cell method in capturing mode coupling and energy transfer phenomena. +""" + +# ------------------ +# 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=50.0, split_algo="LieTrotter") + +# Geometry +domain = domains.Cuboid(r1=31.42) + +# 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 = BinningPlot(slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-10.0, 10.0))) +saving_params = SavingParameters(binning_plots=(binplot,)) + +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=(0.5, None), u1=(3.0, None)) +maxwellian_2 = maxwellians.Maxwellian3D(n=(0.5, None), u1=(-3.0, None)) +background = maxwellian_1 + maxwellian_2 +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps=(0.001,), ls=(1,)) +init1 = maxwellians.Maxwellian3D(n=(0.5, perturbation), u1=(3.0, None)) +init2 = maxwellians.Maxwellian3D(n=(0.5, perturbation), u1=(-3.0, 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) diff --git a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping.py b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping.py new file mode 100644 index 000000000..fb8b46792 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping.py @@ -0,0 +1,175 @@ +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 = """ +Weak Landau damping: A linear test case for the VlasovAmpereOneSpecies model. +This test involves a small amplitude electrostatic perturbation in a uniform, collisionless plasma. +The perturbation is damped due to phase mixing effects (Landau damping) as particles interact with +the self-consistent electric field. This benchmark validates the numerical discretization of the +Vlasov-Ampère system and the accuracy of particle-in-cell methods. +""" + +# ------------------ +# 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.05, Tend=20.0, split_algo="LieTrotter") + +# Geometry +domain = domains.Cuboid(r1=12.56) # r1 -> pi * 4 -> k = 0.5 + +# 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) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) + +binplot = BinningPlot(slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-5.0, 5.0))) +saving_params = SavingParameters(binning_plots=(binplot,)) + +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 +background = maxwellians.Maxwellian3D(n=(1.0, None)) +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps=(0.001,), ls=(1,)) +init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) +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) diff --git a/profiling/examples/VlasovMaxwellOneSpecies/weibel_instability/params_weibel_instability.py b/profiling/examples/VlasovMaxwellOneSpecies/weibel_instability/params_weibel_instability.py new file mode 100644 index 000000000..d63b3920e --- /dev/null +++ b/profiling/examples/VlasovMaxwellOneSpecies/weibel_instability/params_weibel_instability.py @@ -0,0 +1,211 @@ +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 = """ +Weibel instability: A linear test case for the VlasovMaxwellOneSpecies model. This test considers +a plasma with an anisotropic velocity distribution, where temperature differs between directions. +Small magnetic perturbations grow due to the anisotropy, leading to the generation of transverse +magnetic fields. + +Like VlasovAmpereOneSpecies (see strong_Landau_damping/weak_Landau_damping/two_stream/bump_on in +this same profiling suite), VlasovMaxwellOneSpecies exposes its Poisson solve as a one-time +`model.initial_poisson` (not a regular per-step propagator: the fields then evolve via +MaxwellWeakAmpere, PushVxB and VlasovAmpereCoupling instead), so `solver=` only affects this +initial solve -- see `struphy.linear_algebra.petsc_examples_benchmark`'s module docstring for why +that benchmark re-invokes it directly rather than relying on a single sim.run(). + +Plain copy of examples/VlasovMaxwellOneSpecies/weibel_instability, with `num_elements` scaled up +from the original's tiny, highly-anisotropic 1D-style default of `(32, 1, 1)` cells to a proper +`(16, 16, 16)` 3D grid (PETSc's advantage only shows up above roughly 5,000 dofs -- see +struphy.linear_algebra.petsc_poisson_benchmark's module docstring), particle count reduced to +match, and a fixed seed (the original doesn't fix ppc/seed the same way) so pcg/petsc draw the +same particles. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + BinningPlot, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + KernelDensityPlot, + LoadingParameters, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, + perturbations, +) + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import VlasovMaxwellOneSpecies + +# Units +base_units = BaseUnits() + +# Model instance +model = VlasovMaxwellOneSpecies( + base_units=base_units, + alpha=1.0, + epsilon=-1.0, + measure_gauss_law=True, +) + +# --------------------- +# Parameters setup +# --------------------- + +import cunumpy as xp + +k = 1.25 +B_pert_amp = -1e-4 + +vth1_background_val = 0.02 / xp.sqrt(2) +vth2_background_val = vth1_background_val * xp.sqrt(12) + +# 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 +# -------------------------- + +# `--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=400, split_algo="LieTrotter") + +# Geometry +domain = domains.Cuboid(r1=2 * xp.pi / k) + +# 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() + +# 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, + set_zero_velocity=(False, False, True), + moments=(0.0, 0.0, 0.0, vth1_background_val, vth2_background_val, 1.0), + seed=42, +) +weights_params = WeightsParameters(control_variate=False) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) + +binplot_dens = BinningPlot(slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-0.1, 0.1))) +binplot_velocity = BinningPlot(slice="v1_v2", n_bins=(128, 128), ranges=((-0.1, 0.1), (-0.1, 0.1))) +binplot_current = tuple( + BinningPlot(slice=f"e{i}", n_bins=32, ranges=(0.0, 1.0), output_quantity=f"current_{j}") + for j in range(1, 4) + for i in range(1, 4) +) +saving_params = SavingParameters(binning_plots=(binplot_dens, binplot_velocity, *binplot_current)) + +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=2.0, +) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.maxwell.options = model.propagators.maxwell.Options() +model.propagators.push_eta.options = model.propagators.push_eta.Options() +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. + +maxwellian = maxwellians.Maxwellian3D( + vth1=(vth1_background_val, None), + vth2=(vth2_background_val, None), +) +model.kinetic_ions.var.add_background(maxwellian) + +# Perturbation of initial magnetic field +model.em_fields.b_field.add_perturbation( + perturbation=perturbations.ModesCos(amps=(B_pert_amp,), ls=(1,), comp=2), # Initial Bz depending on x-axis +) + +if __name__ == "__main__": + # one_time_step=True isolates the initial Poisson solve (the only part of this model that + # solver= affects -- the fields then evolve via MaxwellWeakAmpere/PushVxB/VlasovAmpereCoupling, + # unrelated to PETScSolver) from the ~8000 identical transport steps a full run would + # otherwise dilute the comparison with. + sim.run(one_time_step=True) diff --git a/profiling/submit_diocotron_strong_scaling.py b/profiling/submit_diocotron_strong_scaling.py index 449f76fd0..0a421ad14 100644 --- a/profiling/submit_diocotron_strong_scaling.py +++ b/profiling/submit_diocotron_strong_scaling.py @@ -15,20 +15,11 @@ from pathlib import Path from profiling_job import ProfilingCase - +from utils import _get_profiling_args def main() -> None: - # Parse arguments, do not remove --upload - parser = argparse.ArgumentParser( - description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), - ) - parser.add_argument( - "--upload", - action="store_true", - help="Upload the packaged profiling results to the profiling-data repo.", - ) - args = parser.parse_args() + args = _get_profiling_args() # Paths relative to this script's location, so it can be run from anywhere. script_dir = Path(__file__).resolve().parent @@ -46,7 +37,6 @@ def main() -> None: upload=args.upload, ) - profiling_case.use_slurm = False # Launch one run per rank count for num_tasks in (2, 4): # , 8, 16, 32, 64, 128, 256): diff --git a/profiling/submit_poisson_strong_scaling.py b/profiling/submit_poisson_strong_scaling.py index 8929e5946..d3a31391e 100644 --- a/profiling/submit_poisson_strong_scaling.py +++ b/profiling/submit_poisson_strong_scaling.py @@ -13,20 +13,11 @@ from pathlib import Path from profiling_job import ProfilingCase - +from utils import _get_profiling_args def main() -> None: - # Parse arguments, do not remove --upload - parser = argparse.ArgumentParser( - description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), - ) - parser.add_argument( - "--upload", - action="store_true", - help="Upload the packaged profiling results to the profiling-data repo.", - ) - args = parser.parse_args() + args = _get_profiling_args() # Paths relative to this script's location, so it can be run from anywhere. script_dir = Path(__file__).resolve().parent diff --git a/profiling/submit_strong_landau_damping_petsc.py b/profiling/submit_strong_landau_damping_petsc.py new file mode 100644 index 000000000..d0ac95cae --- /dev/null +++ b/profiling/submit_strong_landau_damping_petsc.py @@ -0,0 +1,42 @@ +import argparse +from pathlib import Path + +from profiling_job import ProfilingCase +from utils import _get_profiling_args + +def main() -> None: + + args = _get_profiling_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "examples" / "VlasovAmpereOneSpecies" / "strong_Landau_damping" + + profiling_case = ProfilingCase( + label="strong_landau_damping_petsc", + name="Strong Landau damping, initial Poisson solve: pcg vs. PETSc", + description=( + "Strong (nonlinear) Landau damping test case for VlasovAmpereOneSpecies, solving the " + "one-time initial Poisson problem with either feectools' native preconditioned CG " + "(--solver pcg) or PETScSolver (KSP=cg, PC=gamg, --solver petsc)." + ), + physics_problem="Nonlinear Landau damping in a uniform, collisionless plasma.", + struphy_model_used="VlasovAmpereOneSpecies", + params_source=params_dir / "params_strong_Landau_damping.py", + language="fortran", + compiler="GNU", + upload=args.upload, + ) + + # Launch one run per (rank count, solver) combination, same rank counts for both solvers so + # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` + # output folder), so the two solvers never collide even at the same rank count. + for num_tasks in (1, ): + for solver in ("pcg", "petsc"): + profiling_case.launch(num_tasks, param_flags=["--solver", solver]) + + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/profiling/submit_toydrift_hires_petsc.py b/profiling/submit_toydrift_hires_petsc.py new file mode 100644 index 000000000..bff109e08 --- /dev/null +++ b/profiling/submit_toydrift_hires_petsc.py @@ -0,0 +1,82 @@ +"""ToyDrift periodic slab (32^3): PETSc vs. pcg, the largest gap in this benchmark suite. + +See ``profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py`` for the +full story: a 32^3 grid (32768 dofs) with PETSc's preconditioner explicitly set to ``"gamg"`` +(algebraic multigrid, via ``SolverParameters.pc_type`` -- see +``struphy.linear_algebra.solver.SolverParameters``) instead of the default ``"jacobi"``. + +Measured via ``struphy.linear_algebra.petsc_examples_benchmark``'s repeated-solve methodology +(isolates just the solve, amortizing gamg's one-time multigrid setup across several calls at +fixed dt -- see that module's docstring): pcg needs ~190 CG iterations per solve (~10.9s) at this +size, against PETSc+gamg's ~2 iterations (~0.29s) -- a ~38x difference. This near-singular system +is ill-conditioned mainly through its weakly constrained DC mode, not primarily through +resolution, so the ratio does not grow much further with grid size -- but the *absolute* cost +does: pcg's ~11 seconds per solve here is the practically relevant "huge difference" once +multiplied over the many timesteps of a real simulation. See +``submit_strong_landau_damping_petsc.py`` and ``submit_vlasov_maxwell_petsc.py`` for the +smaller-scale (3-10x) comparisons. + +Each launch runs one 32^3-grid, one-time-step simulation (``params_periodic_slab_hires.py``'s +``__main__`` calls ``sim.run(one_time_step=True)``); note that a *single* one-time-step run still +pays gamg's one-time setup cost up front and will not by itself reproduce the ~38x figure above -- +that requires several solves at the same dt to amortize the setup, exactly what +``petsc_examples_benchmark.py`` does. A pcg launch alone takes roughly a minute here (mostly the +single Poisson solve), so a local (non-SLURM) full sweep of this script takes several minutes. +""" + +import argparse +from pathlib import Path + +from profiling_job import ProfilingCase + + +def main() -> None: + + # Parse arguments, do not remove --upload + parser = argparse.ArgumentParser( + description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + args = parser.parse_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "examples" / "ToyDrift" / "periodic_slab_hires" + + profiling_case = ProfilingCase( + label="toydrift_periodic_slab_hires_petsc", + name="ToyDrift periodic slab: PCG vs. PETSc+GAMG", + description=( + "Higher-resolution (32^3 grid, 32768 dofs, degree-3 splines) periodic-slab ToyDrift " + "setup, solving the per-step guiding-center Poisson problem with either feectools' " + "native preconditioned CG (--solver pcg) or PETScSolver with an algebraic multigrid " + "preconditioner (KSP=cg, PC=gamg, --solver petsc). At this size pcg needs roughly 190 " + "iterations per solve (about 11 seconds) while PETSc+gamg needs about 2 (a third of a " + "second); the ratio (~38x) is similar to the smaller periodic_slab case, but the " + "absolute per-solve cost -- and hence the real wall-clock impact over a full run -- " + "is far larger here." + ), + physics_problem="Electrostatic E x B drift of a single ion species in a periodic slab.", + struphy_model_used="ToyDrift", + params_source=params_dir / "params_periodic_slab_hires.py", + language="fortran", + compiler="GNU", + upload=args.upload, + ) + + # Launch one run per (rank count, solver) combination, same rank counts for both solvers so + # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` + # output folder), so the two solvers never collide even at the same rank count. + for num_tasks in (1,): # 2, 4): + for solver in ("pcg", "petsc"): + profiling_case.launch(num_tasks, param_flags=["--solver", solver]) + + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/profiling/submit_vlasov_maxwell_petsc.py b/profiling/submit_vlasov_maxwell_petsc.py new file mode 100644 index 000000000..7108f7701 --- /dev/null +++ b/profiling/submit_vlasov_maxwell_petsc.py @@ -0,0 +1,43 @@ +from pathlib import Path + +from profiling_job import ProfilingCase + +from utils import _get_profiling_args + + +def main() -> None: + + args = _get_profiling_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "examples" / "VlasovMaxwellOneSpecies" / "weibel_instability" + + profiling_case = ProfilingCase( + label="weibel_instability_petsc", + name="Weibel instability, initial Poisson solve: pcg vs. PETSc", + description=( + "Weibel instability test case for VlasovMaxwellOneSpecies, solving the one-time " + "initial Poisson problem with either feectools' native preconditioned CG " + "(--solver pcg) or PETScSolver (KSP=cg, PC=gamg, --solver petsc)." + ), + physics_problem="Weibel instability driven by an anisotropic velocity distribution.", + struphy_model_used="VlasovMaxwellOneSpecies", + params_source=params_dir / "params_weibel_instability.py", + language="fortran", + compiler="GNU", + upload=args.upload, + ) + + # Launch one run per (rank count, solver) combination, same rank counts for both solvers so + # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` + # output folder), so the two solvers never collide even at the same rank count. + for num_tasks in (1, 2, 4): + for solver in ("pcg", "petsc"): + profiling_case.launch(num_tasks, param_flags=["--solver", solver]) + + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/profiling/utils.py b/profiling/utils.py index 8c99c3003..5d31e285f 100644 --- a/profiling/utils.py +++ b/profiling/utils.py @@ -13,6 +13,20 @@ from typing import Any +def _get_profiling_args() -> argparse.Namespace: + # Parse arguments, do not remove --upload + parser = argparse.ArgumentParser( + description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + args = parser.parse_args() + + return args + def _slug(value: str) -> str: return re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("._-") or "unknown" diff --git a/pyproject.toml b/pyproject.toml index cc2ada4f7..c34ab643b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "numpy<=2.5.0", "cunumpy>=0.1.4, <=0.1.5", "pyccel>=2.2.0, <=2.2.3", - "feectools >= 0.1.10, <=0.1.10", + "feectools<=0.1.11", "scipy<=1.18.0", "h5py<=3.16.0", "matplotlib<=3.11.0", @@ -57,6 +57,10 @@ file = "LICENSE" mpi = [ "mpi4py<=4.1.1", ] +petsc = [ + "petsc", + "petsc4py", +] phys = [ "gvec>=1.1.0, <=1.5.0", "desc-opt<=0.17.1", @@ -106,6 +110,7 @@ all = [ "struphy[mpi]", "struphy[doc]", "struphy[likwid]", + "struphy[petsc]", ] [project.urls] diff --git a/setup/modules.pitagora.sh b/setup/modules.pitagora.sh index 795dca625..46b0f59b2 100644 --- a/setup/modules.pitagora.sh +++ b/setup/modules.pitagora.sh @@ -3,5 +3,10 @@ intel-oneapi-mkl/2024.0.0--intel-oneapi-mpi--2021.12.1 \ python/3.11.7" MODULES_GCC="gcc/12.3.0 \ -openmpi/4.1.6--gcc--12.3.0 \ +openmpi/4.1.6--gcc--12.3.0-ucx1.20 \ +petsc/3.22.1--openmpi--4.1.6--gcc--12.3.0-ucx1.20-complex-mumps \ python/3.11.7" + +# The petsc module above is built with CUDA support, so petsc4py's import dlopens +# libcuda.so.1 even on nodes without a GPU driver. Point at a stub so it doesn't fail. +export LD_LIBRARY_PATH="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")/.." && pwd)/.venv/petsc_cuda_stub:${LD_LIBRARY_PATH:-}" diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index c8b63ae9d..d535888fd 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -12,7 +12,6 @@ from feectools.fem.vector import VectorFemSpace from feectools.linalg.basic import IdentityOperator, InverseLinearOperator, LinearOperator, Vector from feectools.linalg.block import BlockLinearOperator, BlockVector -from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilDiagonalMatrix, StencilMatrix, StencilVector from struphy import equils @@ -23,7 +22,7 @@ from struphy.fields_background.base import MHDequilibrium from struphy.geometry.base import Domain from struphy.io.options import LiteralOptions -from struphy.linear_algebra.solver import SolverParameters +from struphy.linear_algebra.solver import SolverParameters, inverse from struphy.polar.basic import PolarVector from struphy.polar.linear_operators import PolarExtractionOperator from struphy.utils.docstring_converter import auto_convert_docstring, info diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index fb9a7e58a..2e57cee2f 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -71,8 +71,9 @@ class LiteralOptions: GivenInBasis = Literal["0", "1", "2", "3", "v", "physical", "physical_at_eta", "norm", None] # solvers - OptsSymmSolver = Literal["pcg", "cg"] + OptsSymmSolver = Literal["pcg", "cg", "petsc"] OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"] + OptsPETScPrecond = Literal["none", "jacobi", "gamg", "ilu", "sor"] OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] OptsSaddlePointSolver = Literal["uzawa"] OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] diff --git a/src/struphy/linear_algebra/petsc_examples_benchmark.py b/src/struphy/linear_algebra/petsc_examples_benchmark.py new file mode 100644 index 000000000..2192276cb --- /dev/null +++ b/src/struphy/linear_algebra/petsc_examples_benchmark.py @@ -0,0 +1,222 @@ +"""Benchmark: solver="petsc" vs solver="pcg" on real struphy examples. + +Uses the committed parameter files under +``profiling/examples///params_.py``, each with a ``--solver`` CLI flag +(default ``"pcg"``) selecting the solver of the case's Poisson-type propagator: + +- ``VlasovAmpereOneSpecies/{strong_Landau_damping,weak_Landau_damping,two_stream,bump_on}``: plain + copies of the corresponding ``examples/VlasovAmpereOneSpecies//params_.py`` (unedited + on disk), with ``num_elements`` scaled up from the examples' tiny, highly-anisotropic 1D-style + default of ``(32, 1, 1)`` cells to a proper ``(16, 16, 16)`` 3D grid -- PETSc's advantage only + shows up above roughly 5,000 dofs -- ``ppc`` reduced to match, and a fixed + ``LoadingParameters.seed`` added (the originals don't set one, so pcg/petsc would otherwise draw + different particles and not be comparable). This model only solves Poisson *once*, as an initial + condition (the field then evolves via VlasovAmpereCoupling), so the repeated-solve timing below + re-invokes ``model.initial_poisson`` directly after ``sim.run()`` rather than relying on the + model's own (single-shot) usage of it. + +- ``ToyDrift/periodic_slab_hires``: no periodic ToyDrift example exists under ``examples/`` (the + real one, ``examples/ToyGyrokinetic/diocotron_instability``, needs a physically non-periodic + HollowCylinder domain), so this one is written from scratch with a periodic Cuboid domain + instead -- which works because ToyDrift's field solve is a plain ``PoissonSolve`` with no + geometry-coupled averaging (unlike ``PoissonAdiabaticGyrokinetic``, used by + ``DriftKineticElectrostaticAdiabatic``, which diverges outright on a periodic domain regardless + of options -- tried first, not usable here). Unlike VlasovAmpereOneSpecies, ToyDrift's + ``gc_poisson`` runs as a *regular per-step propagator*, so no re-invocation workaround is needed. + Uses a 32^3 grid to push feectools' unpreconditioned CG into several hundred iterations per + solve while PETSc+gamg (``pc_type="gamg"``, set explicitly via ``SolverParameters.pc_type`` -- + see ``struphy.linear_algebra.solver.SolverParameters``) stays at a handful, regardless of grid + size. The most lopsided case in this suite by design; see its own params file's docstring. + +- ``VlasovMaxwellOneSpecies/weibel_instability``: plain copy of + ``examples/VlasovMaxwellOneSpecies/weibel_instability/params_weibel_instability.py``, scaled up + the same way as the VlasovAmpereOneSpecies cases above (3D grid, reduced ppc, fixed seed). Same + one-shot ``model.initial_poisson`` pattern as VlasovAmpereOneSpecies (the fields then evolve via + MaxwellWeakAmpere/PushVxB/VlasovAmpereCoupling instead), so the same re-invocation timing applies. + +This script's only job is to import each file's ``sim``/``model`` and drive them -- no source +patching, no ``runpy``. + +Correctness is checked between the two solvers on the *mean-removed* solution (the near-singular, +essentially unregularized ``stab_eps`` these examples use -- via ``ImplicitDiffusion``'s "always +stabilize" clamp to ``1e-14`` -- leaves the constant/DC mode only very weakly constrained, so it is +extremely sensitive to tiny numerical differences between solvers; this is expected and is not a +correctness issue in the physically meaningful, oscillatory part of the solution), normalized by +the *full* solution's norm (not the tiny mean-removed norm itself, which can be dominated by +floating-point noise for weak-perturbation examples and make a naively-normalized relative error +meaningless). + +KNOWN ISSUE -- do not trust results under MPI (comm size > 1): for this same near-singular +``stab_eps``-clamped-to-``1e-14`` regime, PETScSolver was found to disagree substantially with +feectools' native solver specifically under >1 MPI rank, independent of ``pc_type`` (both +"jacobi" and "gamg" reproduced it; "gamg" was far worse -- a false "converged in 1 iteration" to +a wildly wrong answer). This reproduces even though pcg-vs-pcg (same solver, two independent runs) +is bit-identical, ruling out a methodology artifact in this script. The root cause was not found; +serial execution and non-near-singular systems (this same MPI path, e.g. with an explicit +``stab_eps`` of 1e-8 or larger) were extensively validated and are unaffected -- see +``petsc_poisson_benchmark.py`` and ``test_petsc_poisson_solve_pic.py``. This script therefore +only asserts/prints the correctness check when running serially, and warns instead under MPI. + +Run with: + +.. code-block:: bash + + python3 -m struphy.linear_algebra.petsc_examples_benchmark +""" + +import importlib +import os +import shutil +import sys +import tempfile +import time +from pathlib import Path + +import cunumpy as xp +from feectools.ddm.mpi import mpi as MPI + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() + +REPO_ROOT = Path(__file__).resolve().parents[3] +PROFILING_EXAMPLES_DIR = REPO_ROOT / "profiling" / "examples" + +# (model directory under profiling/examples/, case name) +CASES = ( + ("VlasovAmpereOneSpecies", "strong_Landau_damping"), + ("VlasovAmpereOneSpecies", "weak_Landau_damping"), + ("VlasovAmpereOneSpecies", "two_stream"), + ("VlasovAmpereOneSpecies", "bump_on"), + ("ToyDrift", "periodic_slab_hires"), + ("VlasovMaxwellOneSpecies", "weibel_instability"), +) + +# `profiling` is a repo-local package (not part of the installed struphy distribution), normally +# importable only because '' (cwd) is on sys.path at interpreter startup. This script chdir()s +# into a scratch directory before importing (to keep Simulation's output out of the repo), which +# breaks that for '-c'/REPL-style invocations where '' resolves dynamically -- so add the repo +# root explicitly, once, up front. +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + + +def _poisson_propagator(model): + """VlasovAmpereOneSpecies exposes its (one-shot) Poisson solve as `model.initial_poisson`; + other models (e.g. ToyDrift) run it as a regular per-step propagator instead. + """ + if hasattr(model, "initial_poisson"): + return model.initial_poisson + return model.propagators.gc_poisson + + +def _run_variant(model_dir: str, name: str, variant: str, out_folder: str, dt: float, n_solves: int): + module_name = f"profiling.examples.{model_dir}.{name}.params_{name}" + + # EnvironmentOptions.out_folders defaults to `os.getcwd()`, but as a plain dataclass field + # default this is evaluated once, the first time struphy.io.options is imported in this + # process -- not per EnvironmentOptions() call, and not affected by a later os.chdir(). Since + # `struphy`'s own __init__.py re-exports EnvironmentOptions (and everything else), that first + # import can happen before this function even runs (e.g. as a side effect of importing this + # very module). The profiling params files read STRUPHY_PROFILING_OUT_FOLDERS explicitly for + # exactly this reason -- set it before importing, so their Simulation's output lands here + # instead of wherever the process happened to start. + os.environ["STRUPHY_PROFILING_OUT_FOLDERS"] = out_folder + + # Both variants now come from the same params_.py (a `--solver` CLI flag picks between + # them internally), so a plain second `import_module` would just return the first variant's + # already-imported module/sim/model unchanged. Feed the desired `--solver` through sys.argv + # (the params file reads it via argparse) and force a re-execution via `reload` so the second + # variant gets its own fresh Simulation/model instead of reusing (and re-running) the first's. + old_argv = sys.argv + sys.argv = [old_argv[0], "--solver", variant] + try: + if module_name in sys.modules: + mod = importlib.reload(sys.modules[module_name]) + else: + mod = importlib.import_module(module_name) + finally: + sys.argv = old_argv + + sim = mod.sim + model = mod.model + + sim.run(one_time_step=True) # real setup: particle loading, Derham/mass operators, and (for + # VlasovAmpereOneSpecies) the initial Poisson solve, exactly as + # VlasovAmpereOneSpecies.allocate_helpers / a real per-step run does + + poisson = _poisson_propagator(model) + solver = poisson._solver + if variant == "petsc": + solver._options["pc_type"] = "gamg" + solver._ksp = None # force a rebuild with the new pc_type + + # repeated calls at fixed dt (same real charge deposition -- this benchmark is about the + # linear-solve cost, not particle physics): matches how repeated Poisson solves would + # amortize matrix assembly across timesteps at a fixed dt in a real run (see + # ImplicitDiffusion.__call__'s lhs-operator caching) + poisson(dt) # warm-up + t0 = time.perf_counter() + for _ in range(n_solves): + poisson(dt) + t = (time.perf_counter() - t0) / n_solves + info = solver.get_info() if hasattr(solver, "get_info") else solver._info + + phi = model.em_fields.phi.spline.vector.toarray() + return t, info, phi + + +def bench_example(model_dir: str, name: str, dt: float = 0.05, n_solves: int = 10): + """Import and run one example's pcg/petsc parameter files, and report timing + correctness.""" + params_dir = PROFILING_EXAMPLES_DIR / model_dir / name + if not (params_dir / f"params_{name}.py").exists(): + raise FileNotFoundError( + f"Could not find {params_dir}/params_{name}.py -- expected the parameter file " + f"under profiling/examples/{model_dir}/{name}/." + ) + + # tempfile.mkdtemp() is not MPI-coordinated: each rank would otherwise get a *different* + # random path, and Simulation's output-file creation (rank 0 only) would then fail on every + # other rank. Create it on rank 0 and broadcast the path instead. + out_folder = tempfile.mkdtemp() if rank == 0 else None + if comm is not None: + out_folder = comm.bcast(out_folder, root=0) + + try: + t_cg, info_cg, sol_cg = _run_variant(model_dir, name, "pcg", out_folder, dt, n_solves) + t_petsc, info_petsc, sol_petsc = _run_variant(model_dir, name, "petsc", out_folder, dt, n_solves) + finally: + if comm is not None: + comm.Barrier() + if rank == 0: + shutil.rmtree(out_folder, ignore_errors=True) + + mean_removed_pcg = sol_cg - sol_cg.mean() + mean_removed_petsc = sol_petsc - sol_petsc.mean() + rel_err = xp.linalg.norm(mean_removed_pcg - mean_removed_petsc) / xp.linalg.norm(sol_cg) + + comm_size = comm.Get_size() if comm is not None else 1 + if rank == 0: + print(f"\n{model_dir}/{name}: ndofs={sol_cg.size}") + print(f" pcg (unprec.) : {t_cg * 1e3:9.2f} ms/step niter={info_cg.get('niter')}") + print(f" petsc + gamg : {t_petsc * 1e3:9.2f} ms/step niter={info_petsc.get('niter')}") + print(f" speedup: {t_cg / t_petsc:.2f}x", end=" ") + if comm_size > 1: + print( + f"relative solution mismatch: {rel_err:.2e} " + "-- NOT a reliable correctness check under MPI for this near-singular regime, " + "see module docstring (KNOWN ISSUE)" + ) + else: + print(f"relative solution mismatch: {rel_err:.2e}") + assert rel_err < 1e-6, ( + f"{model_dir}/{name}: pcg/petsc solutions disagree by {rel_err:.2e}, expected < 1e-6 serially" + ) + + +def main(): + for model_dir, name in CASES: + bench_example(model_dir, name) + + +if __name__ == "__main__": + main() diff --git a/src/struphy/linear_algebra/petsc_poisson_benchmark.py b/src/struphy/linear_algebra/petsc_poisson_benchmark.py new file mode 100644 index 000000000..546edacfb --- /dev/null +++ b/src/struphy/linear_algebra/petsc_poisson_benchmark.py @@ -0,0 +1,207 @@ +"""Benchmark: PoissonSolve(solver="petsc") vs PoissonSolve(solver="pcg") on a real Vlasov-Poisson testcase. + +Builds actual :class:`~struphy.simulation.sim.Simulation` objects, using the same public API and +setup idiom as a real parameter file (compare +``examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py``, +generalized here to a 3D grid), and runs them with ``sim.run(one_time_step=True)`` -- this is the +one and only supported entry point for allocating and running a struphy simulation; it performs +real particle loading, real Derham/mass-operator setup, and (for this model) a real charge-density +deposition via ``ParticlesToGrid``/``AccumulatorVector``, solved via ``PoissonSolve`` exactly as +``VlasovAmpereOneSpecies.allocate_helpers`` does. + +Two earlier, less realistic benchmarks are *not* wins for PETSc and are not repeated here: + +- Mass-matrix solves (``L2Projector``): already well-conditioned, feectools' native + preconditioned CG wins outright. +- ``PoissonSolve`` with a synthetic, non-mass-weighted random right-hand side: not representative + of any real code path (every real source -- ``FEECVariable``, ``ParticlesToGrid``, ``Callable`` + -- is mass-matrix weighted when forming the weak-form right-hand side, which is inherently + smoothing). + +The genuine win requires: a small ``stab_eps`` (true elliptic Poisson, matching realistic +electrostatic PIC parameters -- ``stab_eps`` is a numerical regularization, not a dominant +physical diffusion), a broadband/noisy right-hand side (real PIC deposition, not a smooth +manufactured mode), and repeated solves at fixed ``dt`` (relies on +``ImplicitDiffusion.__call__`` caching its lhs operator when ``sig_1`` is unchanged, so +``PETScSolver`` can reuse its assembled matrix -- see git history for that fix). Since +``VlasovAmpereOneSpecies`` only calls its Poisson solve *once* (as an initial condition -- the +electric field then evolves via Ampere's law, not repeated Poisson solves), the repeated-solve +timing below re-invokes ``model.initial_poisson`` directly after ``sim.run()`` has performed the +real setup, rather than relying on the model's own (single-shot) usage of it. + +Run with: + +.. code-block:: bash + + python3 -m struphy.linear_algebra.petsc_poisson_benchmark +""" + +import shutil +import tempfile +import time +import warnings + +import cunumpy as xp +from feectools.ddm.mpi import mpi as MPI + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) +from struphy.linear_algebra.solver import SolverParameters +from struphy.models import VlasovAmpereOneSpecies + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() + + +def build_and_run(num_elements, degree, ppc, solver_name, perturbation, pc_type=None, stab_eps=1e-8, out_folder=None): + """Build a VlasovAmpereOneSpecies Simulation exactly as a params.py file would, and run one + step of it via sim.run(one_time_step=True) -- the real, supported entry point. This performs + real particle loading and the real (single-shot) initial Poisson solve. + """ + # alpha=1.0, epsilon=-1.0 (matching the real strong_Landau_damping example) keeps the + # right-hand side well-scaled (order 1); epsilon in particular can never be auto-derived as + # negative (its formula is always positive), so overriding it is unavoidable here, and + # struphy warns on every such override. The warning is expected and harmless -- silence it + # rather than leaving alpha/epsilon at their auto-derived values, which was tried and + # produces a poorly-scaled right-hand side (huge/tiny relative to 1), breaking the implicit + # assumption -- shared by every SolverParameters.tol comparison in this benchmark -- that + # "tol" means the same thing regardless of problem scale. + # with warnings.catch_warnings(): + # warnings.filterwarnings("ignore", message="Override equation parameter", category=UserWarning) + model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + + env = EnvironmentOptions(out_folders=out_folder, sim_folder=f"bench_{solver_name}") + time_opts = Time(dt=0.05, Tend=0.05, split_algo="LieTrotter") + domain = domains.Cuboid() + grid = grids.TensorProductGrid(num_elements=num_elements) + derham_opts = DerhamOptions(degree=degree, bcs=(None, None, None)) + + sim = Simulation( + model=model, + params_path=None, + env=env, + time_opts=time_opts, + domain=domain, + equil=None, + grid=grid, + derham_opts=derham_opts, + ) + + loading_params = LoadingParameters(ppc=ppc, seed=1234) + weights_params = WeightsParameters(control_variate=True) + boundary_params = BoundaryParameters() + model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + ) + + 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( + stab_mat="M0", + stab_eps=stab_eps, + solver=solver_name, + precond="MassMatrixPreconditioner", + solver_params=SolverParameters(tol=1e-10, maxiter=20_000, info=False, recycle=False), + ) + + background = maxwellians.Maxwellian3D(n=(1.0, None)) + model.kinetic_ions.var.add_background(background) + init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) + model.kinetic_ions.var.add_initial_condition(init) + + sim.run(one_time_step=True) + + if solver_name == "petsc" and pc_type is not None and model.initial_poisson._solver._options["pc_type"] != pc_type: + model.initial_poisson._solver._options["pc_type"] = pc_type + model.initial_poisson._solver._ksp = None + + return sim, model + + +def bench_case(name, num_elements, degree, ppc, perturbation, n_solves=10, dt=0.05, **kwargs): + # tempfile.mkdtemp() is not MPI-coordinated: each rank would otherwise get a *different* + # random path, and Simulation's output-file creation (rank 0 only) would then fail on every + # other rank. Create it on rank 0 and broadcast the path instead. + out_folder = tempfile.mkdtemp() if rank == 0 else None + if comm is not None: + out_folder = comm.bcast(out_folder, root=0) + + try: + _, model_cg = build_and_run(num_elements, degree, ppc, "pcg", perturbation, out_folder=out_folder, **kwargs) + _, model_petsc = build_and_run( + num_elements, degree, ppc, "petsc", perturbation, pc_type="gamg", out_folder=out_folder, **kwargs + ) + + # both models' initial_poisson already ran once inside sim.run(); time repeated calls at + # fixed dt, exactly as a real timestepping loop would (see module docstring) + model_cg.initial_poisson(dt) # warm-up + t0 = time.perf_counter() + for _ in range(n_solves): + model_cg.initial_poisson(dt) + t_cg = (time.perf_counter() - t0) / n_solves + info_cg = model_cg.initial_poisson._solver._info + + model_petsc.initial_poisson(dt) # warm-up + t0 = time.perf_counter() + for _ in range(n_solves): + model_petsc.initial_poisson(dt) + t_petsc = (time.perf_counter() - t0) / n_solves + info_petsc = model_petsc.initial_poisson._solver.get_info() + + sol_cg = model_cg.em_fields.phi.spline.vector.toarray() + sol_petsc = model_petsc.em_fields.phi.spline.vector.toarray() + rel_err = xp.linalg.norm(sol_cg - sol_petsc) / xp.linalg.norm(sol_cg) + finally: + if comm is not None: + comm.Barrier() + if rank == 0: + shutil.rmtree(out_folder, ignore_errors=True) + + ndofs = model_cg.em_fields.phi.spline.vector.space.dimension + Np = model_cg.kinetic_ions.var.particles.markers.shape[0] + + if rank == 0: + print(f"\n{name}: num_elements={num_elements}, degree={degree}, ndofs={ndofs}, Np~{Np}") + print(f" pcg (unprec.) : {t_cg * 1e3:9.2f} ms/step niter={info_cg.get('niter')}") + print(f" petsc + gamg : {t_petsc * 1e3:9.2f} ms/step niter={info_petsc.get('niter')}") + print(f" speedup: {t_cg / t_petsc:.2f}x relative solution mismatch: {rel_err:.2e}") + + +def main(): + # 1. grid-size scaling, matching examples/VlasovAmpereOneSpecies/strong_Landau_damping's ICs + landau_damping = perturbations.ModesCos(amps=(0.5,), ls=(1,)) + for num_elements in [[8, 8, 8], [16, 16, 16], [24, 24, 24], [32, 32, 32]]: + bench_case("Landau damping (grid scaling)", num_elements, [2, 2, 2], ppc=20, perturbation=landau_damping) + + # 2. weak Landau damping ICs (small-amplitude perturbation, closer to the linear regime) + weak_landau_damping = perturbations.ModesCos(amps=(0.001,), ls=(1,)) + bench_case("weak Landau damping", [24, 24, 24], [2, 2, 2], ppc=20, perturbation=weak_landau_damping) + + # 3. higher spline degree (matching the real examples' degree=3 in the perturbed direction) + bench_case("Landau damping, degree 3", [16, 16, 16], [3, 3, 3], ppc=20, perturbation=landau_damping) + + # 4. sparser sampling (fewer particles per cell -> noisier deposited density) + bench_case("Landau damping, low ppc (noisier)", [24, 24, 24], [2, 2, 2], ppc=5, perturbation=landau_damping) + + # 5. genuinely 3D, multi-mode perturbation (unlike the 1D-in-x real examples), a closer + # stand-in for 3D electrostatic turbulence + multi_mode_3d = perturbations.ModesCos(amps=(0.5, 0.3, 0.2), ls=(1, 2, 0), ms=(0, 1, 2), ns=(0, 0, 1)) + bench_case("3D multi-mode perturbation", [24, 24, 24], [2, 2, 2], ppc=20, perturbation=multi_mode_3d) + + +if __name__ == "__main__": + main() diff --git a/src/struphy/linear_algebra/petsc_solver.py b/src/struphy/linear_algebra/petsc_solver.py new file mode 100644 index 000000000..408c44480 --- /dev/null +++ b/src/struphy/linear_algebra/petsc_solver.py @@ -0,0 +1,384 @@ +import logging + +import cunumpy as xp +from feectools.feec.derivatives import DirectionalDerivativeOperator +from feectools.linalg.basic import ( + ComposedLinearOperator, + IdentityOperator, + InverseLinearOperator, + LinearOperator, + ScaledLinearOperator, + SumLinearOperator, + Vector, +) +from feectools.linalg.block import BlockLinearOperator, BlockVectorSpace +from feectools.linalg.stencil import StencilMatrix +from feectools.linalg.topetsc import get_npts_local, mat_topetsc, vec_topetsc +from feectools.linalg.utilities import petsc_to_psydac + +logger = logging.getLogger("struphy") + + +def _directional_derivative_to_stencil_matrix(op): + """Build a :class:`~feectools.linalg.stencil.StencilMatrix` equivalent to a + (matrix-free) :class:`~feectools.feec.derivatives.DirectionalDerivativeOperator`, so it can + be handed to :func:`feectools.linalg.topetsc.mat_topetsc`. + + ``DirectionalDerivativeOperator.dot`` computes, along its differentiation axis + ``diffdir`` (identity along every other axis): + + - ``out = v[..., k+1, ...] - v[..., k, ...]`` if not negative, not transposed + - ``out = v[..., k, ...] - v[..., k+1, ...]`` if negative, not transposed + - ``out = v[..., k-1, ...] - v[..., k, ...]`` if not negative, transposed + - ``out = v[..., k, ...] - v[..., k-1, ...]`` if negative, transposed + + i.e. a plain two-point (identity, shift-by-one) stencil. + + Note + ---- + Only verified for a *periodic* differentiation axis. For a non-periodic axis under a + parallel (MPI-comm-attached) space -- which is how struphy always builds its Derham + complex, even with a single rank -- this construction (and feectools' own + ``DirectionalDerivativeOperator.tokronstencil().tostencil()``) was found to disagree with + the operator's actual ``.dot()`` at the two boundary planes along that axis. The root + cause was not identified; rather than risk silently wrong results, this case raises + ``NotImplementedError``. + """ + assert isinstance(op, DirectionalDerivativeOperator) + + V = op.domain + W = op.codomain + ndim = V.ndim + diffdir = op._diffdir + negative = op._negative + transposed = op._transposed + + if not V.periods[diffdir]: + raise NotImplementedError( + "PETScSolver cannot (yet) assemble a DirectionalDerivativeOperator along a " + f"non-periodic axis (diffdir={diffdir}, periods={V.periods}) of a parallel " + "(MPI-comm-attached) space: this was found to disagree with the operator's actual " + "action at the domain boundary, for a reason not yet root-caused. Only fully " + "periodic operators (e.g. derham.grad on a fully periodic domain) are supported." + ) + + M = StencilMatrix(V, W) + + def off(o): + return slice(o, o + 1) + + rows = tuple(slice(None) for _ in range(ndim)) + identity_key = tuple(off(0) for _ in range(ndim)) + + shift = -1 if transposed else 1 + shifted_key = tuple(off(shift) if d == diffdir else off(0) for d in range(ndim)) + + if negative: + M[rows + identity_key] = 1.0 + M[rows + shifted_key] = -1.0 + else: + M[rows + identity_key] = -1.0 + M[rows + shifted_key] = 1.0 + + M.remove_spurious_entries() + return M + + +def _materialize_block(block): + """Turn a ``BlockLinearOperator`` block entry into a concrete matrix (``StencilMatrix`` or + ``None``) that :func:`feectools.linalg.topetsc.mat_topetsc` can handle directly. + + Blocks of a topological operator such as ``derham.curl`` are not always plain + ``StencilMatrix``: sign conventions are sometimes expressed via ``ScaledLinearOperator`` + wrapping a ``StencilMatrix``/``DirectionalDerivativeOperator`` rather than baking the sign + into the matrix data (observed e.g. for ``derham.curl.T``, whose transposed blocks land on + this path). ``mat_topetsc`` calls ``.update_ghost_regions()`` on every block, which only + concrete matrix types implement -- so any such wrapper must be resolved to a concrete matrix + first. Recurses through nested ``ScaledLinearOperator``s. + """ + if block is None: + return None + if isinstance(block, DirectionalDerivativeOperator): + return _directional_derivative_to_stencil_matrix(block) + if isinstance(block, ScaledLinearOperator): + inner = _materialize_block(block.operator) + if inner is None: + return None + scaled = inner.copy() + scaled *= block.scalar + return scaled + return block + + +def _assemble_leaf_operator(A): + """Return an operator equivalent to `A` that is directly convertible via + :func:`feectools.linalg.topetsc.mat_topetsc` (i.e. a ``StencilMatrix`` or a + ``BlockLinearOperator`` whose blocks are all ``StencilMatrix``), replacing any + ``DirectionalDerivativeOperator``/``ScaledLinearOperator`` (block or bare) by its assembled + equivalent -- see :func:`_materialize_block`. + """ + if isinstance(A, (DirectionalDerivativeOperator, ScaledLinearOperator)): + return _materialize_block(A) + + if isinstance(A, BlockLinearOperator): + out = BlockLinearOperator(A.domain, A.codomain) + for i, j in A.nonzero_block_indices: + out[i, j] = _materialize_block(A[i, j]) + return out + + return A + + +def _comm_of(space): + """MPI communicator of a StencilVectorSpace/BlockVectorSpace, matching mat_topetsc's convention.""" + if isinstance(space, BlockVectorSpace): + return space.spaces[0].cart.global_comm + return space.cart.global_comm + + +def _identity_petsc_mat(space): + """Build a PETSc.Mat representing the identity operator on `space`.""" + from petsc4py import PETSc + + comm = _comm_of(space) + localsize = int(xp.sum(xp.prod(get_npts_local(space), axis=1))) + globalsize = space.dimension + + gmat = PETSc.Mat().create(comm=comm) + gmat.setSizes(size=((localsize, globalsize), (localsize, globalsize))) + gmat.setType("mpiaij" if comm else "seqaij") + gmat.setUp() + + ones = space.zeros() + ones._data[:] = 1.0 + gmat.setDiagonal(vec_topetsc(ones)) + gmat.assemble() + + return gmat + + +def _assemble_petsc_matrix(A): + """Recursively assemble a ``PETSc.Mat`` for a (possibly composite) feectools + ``LinearOperator``, by converting every assembled leaf via + :func:`feectools.linalg.topetsc.mat_topetsc` and combining the pieces with PETSc's own + matrix algebra (``matMult`` for composition, ``axpy`` for sums, ``scale`` for scalar + multiples). This lets algebraic preconditioners (jacobi, gamg, ...) work on operators such + as ``grad.T @ M @ grad`` that are not themselves a ``StencilMatrix``/``BlockLinearOperator``. + + Parameters + ---------- + A : feectools.linalg.basic.LinearOperator + Operator to assemble. Supported: ``StencilMatrix``, ``BlockLinearOperator``, any operator + exposing an assembled ``.matrix`` (e.g. ``WeightedMassOperator``), ``IdentityOperator``, + ``ScaledLinearOperator``, ``SumLinearOperator`` and ``ComposedLinearOperator`` built out of + the above (as produced e.g. by ``derham.grad.T @ mass_ops.M1 @ derham.grad``). + + Returns + ------- + gmat : PETSc.Mat + """ + if isinstance(A, (StencilMatrix, BlockLinearOperator, DirectionalDerivativeOperator)): + return mat_topetsc(_assemble_leaf_operator(A)) + + matrix = getattr(A, "matrix", None) + if isinstance(matrix, (StencilMatrix, BlockLinearOperator, DirectionalDerivativeOperator)): + return mat_topetsc(_assemble_leaf_operator(matrix)) + + if isinstance(A, IdentityOperator): + return _identity_petsc_mat(A.domain) + + if isinstance(A, ScaledLinearOperator): + gmat = _assemble_petsc_matrix(A.operator) + gmat.scale(A.scalar) + return gmat + + if isinstance(A, ComposedLinearOperator): + from petsc4py import PETSc + + mats = [_assemble_petsc_matrix(m) for m in A.multiplicants] + gmat = mats[0] + for m in mats[1:]: + gmat = gmat.matMult(m) + return gmat + + if isinstance(A, SumLinearOperator): + from petsc4py import PETSc + + mats = [_assemble_petsc_matrix(a) for a in A.addends] + gmat = mats[0].copy() + for m in mats[1:]: + gmat.axpy(1.0, m, structure=PETSc.Mat.Structure.DIFFERENT_NONZERO_PATTERN) + return gmat + + raise NotImplementedError( + f"PETScSolver cannot assemble a PETSc matrix for operator of type {type(A)}. " + "Supported: StencilMatrix, BlockLinearOperator, operators exposing an assembled " + "'.matrix', IdentityOperator, and Scaled/Sum/Composed combinations thereof." + ) + + +class PETScSolver(InverseLinearOperator): + """(Approximate) inverse of a feectools ``LinearOperator``, computed via a PETSc ``KSP`` + Krylov solver. + + ``A`` is assembled into a ``PETSc.Mat`` (see :func:`_assemble_petsc_matrix` -- this also + handles composite operators such as ``grad.T @ M @ grad``, not just plain + ``StencilMatrix``/``BlockLinearOperator``) and the right-hand side is converted to a + ``PETSc.Vec`` via :func:`feectools.linalg.topetsc.vec_topetsc`; the solve itself is delegated + to ``petsc4py.PETSc.KSP``. Requires the optional ``petsc4py`` dependency + (``pip install struphy[petsc]``). + + Parameters + ---------- + A : feectools.linalg.basic.LinearOperator + Left-hand-side matrix of the linear system, see :func:`_assemble_petsc_matrix` for the + supported operator types. + + x0 : feectools.linalg.basic.Vector, default=None + Kept for interface compatibility with the other + :class:`~feectools.linalg.basic.InverseLinearOperator` subclasses; unused by PETSc's KSP. + + tol : float, default=1e-6 + Relative tolerance, passed to ``KSP.setTolerances(rtol=tol)``. Note this differs from + feectools' own solvers, whose ``tol`` is an *absolute* tolerance on the residual norm -- + for a poorly-scaled system (e.g. a right-hand side far from order 1) the two are not + directly comparable; see git history for a reverted attempt to unify them via + ``atol``, which caused severe slowdowns/inaccuracy for such systems. + + maxiter : int, default=1000 + Maximum number of KSP iterations. + + verbose : bool, default=False + If True, log convergence information after each solve. + + recycle : bool, default=False + Kept for interface compatibility; unused by PETSc's KSP. + + ksp_type : str, default="cg" + PETSc Krylov solver type, see ``petsc4py.PETSc.KSP.Type``. + + pc_type : str, default="none" + PETSc preconditioner type, see ``petsc4py.PETSc.PC.Type``. E.g. ``"gamg"`` (algebraic + multigrid) for large, ill-conditioned elliptic systems. + + near_null_space : {"none", "constant"}, default="none" + Registers a null space with the assembled matrix via ``Mat.setNullSpace`` when + ``"constant"`` (the all-ones vector), so KSP removes any inconsistent component from the + right-hand side instead of letting it pollute the solve. This is not just a minor + robustness tweak: for a *near*-singular operator whose kernel is (numerically) the + constant vector -- e.g. ``ImplicitDiffusion``'s ``grad.T @ M @ grad + sigma_1 * stab_mat`` + on a periodic domain with tiny ``sigma_1`` -- omitting it was found to make ``"gamg"`` + silently converge (small reported residual) to a solution that disagrees substantially + with feectools' own solver, worse and MPI-rank-count-dependent as rank count grows (up to + ~180% relative error at 4 ranks in testing), while reporting success throughout; with it, + the same cases match to ~1e-15 in 1 iteration, independent of rank count. Only pass + ``"constant"`` for operators whose kernel is actually (near) the constant vector -- + forcing it on an operator that is not near-singular there is unlikely to help, and forcing + it on one that is near-singular along some *other* direction would silently corrupt the + answer as this option does not check its own applicability. + """ + + def __init__( + self, + A, + *, + x0=None, + tol=1e-6, + maxiter=1000, + verbose=False, + recycle=False, + ksp_type="cg", + pc_type="none", + near_null_space="none", + ): + assert isinstance(A, LinearOperator), f"PETScSolver requires a LinearOperator, got {type(A)}." + assert near_null_space in ("none", "constant"), f"Unsupported {near_null_space = }" + + self._options = { + "x0": x0, + "tol": tol, + "maxiter": maxiter, + "verbose": verbose, + "recycle": recycle, + "ksp_type": ksp_type, + "pc_type": pc_type, + "near_null_space": near_null_space, + } + + super().__init__(A, **self._options) + + self._info = None + self._ksp = None + # operator for which self._ksp's PETSc.Mat was last built, used to avoid + # re-assembling the matrix on every solve() call when `linop` is unchanged + self._ksp_linop = None + + def _get_ksp(self): + from petsc4py import PETSc + + A = self._A + if self._ksp is None or self._ksp_linop is not A: + gmat = _assemble_petsc_matrix(A) + + if self._options["near_null_space"] == "constant": + nullspace = PETSc.NullSpace().create(constant=True, comm=gmat.getComm()) + gmat.setNullSpace(nullspace) + + if self._ksp is None: + self._ksp = PETSc.KSP().create(comm=gmat.getComm()) + + self._ksp.setType(self._options["ksp_type"]) + self._ksp.getPC().setType(self._options["pc_type"]) + self._ksp.setTolerances(rtol=self._options["tol"], max_it=self._options["maxiter"]) + self._ksp.setOperators(gmat) + self._ksp.setFromOptions() + + self._ksp_linop = A + + return self._ksp + + def solve(self, b, out=None): + """Solve ``A x = b`` using a PETSc KSP Krylov solver. + + Parameters + ---------- + b : feectools.linalg.basic.Vector + Right-hand-side vector of the linear system. + + out : feectools.linalg.basic.Vector | None + The output vector, or None (optional). + + Returns + ------- + x : feectools.linalg.basic.Vector + Numerical solution of the linear system. Convergence info is available + via :meth:`get_info`. + """ + assert isinstance(b, Vector) + assert b.space is self._domain + + ksp = self._get_ksp() + + gvec_b = vec_topetsc(b) + gvec_x = gvec_b.duplicate() + + ksp.solve(gvec_b, gvec_x) + + out = petsc_to_psydac(gvec_x, self._codomain, out=out) + + self._info = { + "niter": ksp.getIterationNumber(), + "success": ksp.getConvergedReason() > 0, + "res_norm": ksp.getResidualNorm(), + } + + if self._options["verbose"]: + logger.info(f"PETSc KSP solver info: {self._info}") + + gvec_b.destroy() + gvec_x.destroy() + + return out + + def dot(self, b, out=None): + return self.solve(b, out=out) diff --git a/src/struphy/linear_algebra/petsc_solver_example.py b/src/struphy/linear_algebra/petsc_solver_example.py new file mode 100644 index 000000000..2b1ce6426 --- /dev/null +++ b/src/struphy/linear_algebra/petsc_solver_example.py @@ -0,0 +1,69 @@ +"""Example: solve a struphy mass-matrix system with :class:`~struphy.linear_algebra.petsc_solver.PETScSolver`. + +Builds the ``H1`` mass matrix ``M0`` of a small 3D Derham complex on a cuboid domain, +manufactures a right-hand side from a known exact solution, and solves ``M0 x = b`` +with a PETSc KSP (CG + Jacobi preconditioner), comparing against feectools' native +preconditioned CG solver. + +Run with: + +.. code-block:: bash + + python3 -m struphy.linear_algebra.petsc_solver_example +""" + +import cunumpy as xp +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.solvers import inverse + +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.fields_background.equils import HomogenSlab +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.petsc_solver import PETScSolver +from struphy.topology.grids import TensorProductGrid + + +def main(): + comm = MPI.COMM_WORLD + + # domain, equilibrium and Derham complex + domain = Cuboid() + equil = HomogenSlab(n0=2.0) + equil.domain = domain + + grid = TensorProductGrid(num_elements=[8, 8, 8]) + derham_opts = DerhamOptions(degree=[2, 2, 2]) + derham = Derham(grid, derham_opts, comm=comm, domain=domain) + + # weighted mass operators -- M0.matrix is the assembled StencilMatrix on the H1 space + mass_ops = WeightedMassOperators(derham, domain, eq_mhd=equil) + M0 = mass_ops.M0.matrix + + # manufacture a right-hand side from a known exact solution + xe = M0.domain.zeros() + xe[:] = xp.random.random(xe[:].shape) + xe.update_ghost_regions() + b = M0.dot(xe) + + # reference solve with feectools' preconditioned CG + pc = M0.diagonal(inverse=True) + cg_solver = inverse(M0, "pcg", pc=pc, tol=1e-12, maxiter=2000, verbose=False, recycle=False) + x_cg = cg_solver.solve(b) + + # solve the same system with PETSc's CG + Jacobi preconditioner + petsc_solver = PETScSolver(M0, tol=1e-12, maxiter=2000, ksp_type="cg", pc_type="jacobi") + x_petsc = petsc_solver.solve(b) + + error_vs_exact = xp.linalg.norm((x_petsc - xe).toarray()) + error_vs_cg = xp.linalg.norm((x_petsc - x_cg).toarray()) + + if comm.Get_rank() == 0: + print(f"PETSc KSP info: {petsc_solver.get_info()}") + print(f"||x_petsc - x_exact|| = {error_vs_exact:.3e}") + print(f"||x_petsc - x_cg|| = {error_vs_cg:.3e}") + + +if __name__ == "__main__": + main() diff --git a/src/struphy/linear_algebra/petsc_solver_ill_conditioned_example.py b/src/struphy/linear_algebra/petsc_solver_ill_conditioned_example.py new file mode 100644 index 000000000..6aa8a4c37 --- /dev/null +++ b/src/struphy/linear_algebra/petsc_solver_ill_conditioned_example.py @@ -0,0 +1,93 @@ +"""Example: PETSc beats feectools' native solver on an ill-conditioned SPD system. + +Mass-matrix solves (see ``petsc_solver_example.py``) are *not* where PETSc helps: +they are already well-conditioned and feectools' diagonal-preconditioned CG converges +in a couple of iterations, so per-solve overhead dominates and makes PETSc slower there. + +Where PETSc *does* win is on badly-conditioned elliptic systems, where its algebraic +multigrid preconditioner (``pc_type="gamg"``) keeps the iteration count roughly constant +while plain (or diagonally-preconditioned) CG needs an iteration count that grows like +``sqrt(condition number)``. + +This example builds the standard 1D discrete Laplacian (tridiagonal, -1/2/-1), whose +condition number scales like ``O(n^2)`` in the number of unknowns ``n``, and compares: + +- feectools' plain, unpreconditioned CG +- :class:`~struphy.linear_algebra.petsc_solver.PETScSolver` with ``ksp_type="cg", pc_type="gamg"`` + +For this to show a genuine win (not just fewer iterations but less wall time), the +per-solve vector conversion (:func:`feectools.linalg.topetsc.vec_topetsc` / +:func:`~feectools.linalg.utilities.petsc_to_psydac`) needs to be vectorized rather than +looping in pure Python over every DOF -- that conversion cost otherwise swamps any +iteration-count savings. There's a rough sweet spot: below a few thousand unknowns +GAMG's one-time setup cost dominates and plain CG wins; above it, PETSc's roughly +constant iteration count pulls ahead, increasingly so as ``n`` grows (and, as a bonus, +plain unpreconditioned CG's accuracy degrades from floating-point error accumulation +once it needs tens of thousands of iterations, while PETSc stays accurate). + +Run with: + +.. code-block:: bash + + python3 -m struphy.linear_algebra.petsc_solver_ill_conditioned_example +""" + +import time + +import cunumpy as xp +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.solvers import inverse +from feectools.linalg.stencil import StencilMatrix, StencilVector, StencilVectorSpace + +from struphy.linear_algebra.petsc_solver import PETScSolver + + +def build_1d_laplacian(n, comm): + """Standard 1D discrete Laplacian (tridiagonal, -1, 2, -1); condition number ~ O(n^2).""" + p = 1 + dd = DomainDecomposition([n - p], [False], comm=comm) + cart = CartDecomposition(dd, [n], [xp.array([0])], [xp.array([n - 1])], [p], [1]) + V = StencilVectorSpace(cart) + s = V.starts[0] + e = V.ends[0] + + A = StencilMatrix(V, V) + A[:, -1:0] = -1.0 + A[:, 0:1] = 2.0 + A[:, 1:2] = -1.0 + A.remove_spurious_entries() + + xe = StencilVector(V) + xe[s : e + 1] = xp.random.random(e + 1 - s) + + return V, A, xe + + +def main(n=50_000, tol=1e-8, maxiter=200_000): + comm = MPI.COMM_WORLD + + _, A, xe = build_1d_laplacian(n, comm) + b = A @ xe + + t0 = time.perf_counter() + cg_solver = inverse(A, "cg", tol=tol, maxiter=maxiter, verbose=False, recycle=False) + x_cg = cg_solver.solve(b) + t_cg = time.perf_counter() - t0 + + t0 = time.perf_counter() + petsc_solver = PETScSolver(A, tol=tol, maxiter=maxiter, ksp_type="cg", pc_type="gamg") + x_petsc = petsc_solver.solve(b) + t_petsc = time.perf_counter() - t0 + + if comm.Get_rank() == 0: + print(f"n={n} dofs (1D Laplacian, condition number ~ O(n^2))") + print(f" cg (unpreconditioned) : {t_cg * 1e3:9.2f} ms, niter={cg_solver.get_info()['niter']}") + print(f" petsc (cg + gamg) : {t_petsc * 1e3:9.2f} ms, niter={petsc_solver.get_info()['niter']}") + print(f" speedup: {t_cg / t_petsc:.2f}x") + print(f" ||x_cg - x_exact|| = {xp.linalg.norm((x_cg - xe).toarray()):.3e}") + print(f" ||x_petsc - x_exact|| = {xp.linalg.norm((x_petsc - xe).toarray()):.3e}") + + +if __name__ == "__main__": + main() diff --git a/src/struphy/linear_algebra/petsc_speedup_example.py b/src/struphy/linear_algebra/petsc_speedup_example.py new file mode 100644 index 000000000..beb81e96d --- /dev/null +++ b/src/struphy/linear_algebra/petsc_speedup_example.py @@ -0,0 +1,142 @@ +"""Standalone, self-contained example: PETSc beats pcg on a real (non-toy) Poisson solve. + +Run directly -- no profiling infrastructure, no submission, nothing to set up beyond an +environment with petsc4py installed (``pip install -e ".[petsc]"``): + +.. code-block:: bash + + python src/struphy/linear_algebra/petsc_speedup_example.py + +What it does +------------ +Builds the same ToyDrift periodic-slab setup used by +``profiling/examples/ToyDrift/periodic_slab_hires`` (the largest PETSc-vs-pcg gap in struphy's +profiling suite) directly in this script, so you can read top to bottom exactly what is being +solved and how it is timed -- no need to trace through ``ProfilingCase``/submit-script machinery. + +ToyDrift's ``gc_poisson`` is a *regular per-step propagator* (unlike e.g. VlasovAmpereOneSpecies, +which only solves Poisson once as an initial condition), so it is called once per timestep with +the left-hand-side operator reused across calls at fixed ``dt`` (see +``ImplicitDiffusion.__call__``'s lhs-operator caching). That means the *first* solve pays for +PETSc's one-time matrix assembly and, with ``pc_type="gamg"``, multigrid hierarchy construction -- +this script does one untimed warm-up call for exactly that reason, matching how the cost would +amortize over the many timesteps of a real simulation, before timing several further calls. + +What to expect +--------------- +At this problem size (32768 dofs), feectools' unpreconditioned CG needs on the order of a few +hundred iterations per solve (several seconds), while PETSc with an algebraic multigrid +preconditioner (``pc_type="gamg"``) needs only a handful (well under a second) -- typically a +30-40x speedup. The two solutions are also checked against each other (mean-removed, since the +near-singular constant/DC mode is only weakly constrained and not physically meaningful here -- +see ``PoissonSolve``'s stabilization) and should agree to within floating-point noise. +""" + +import time + +import cunumpy as xp + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, + perturbations, +) +from struphy.linear_algebra.solver import SolverParameters +from struphy.models import ToyDrift + + +def _build_and_run(solver: str, pc_type: str, n_solves: int): + """Build a fresh ToyDrift periodic-slab simulation and time repeated Poisson solves.""" + model = ToyDrift(base_units=BaseUnits(kBT=1.0)) + + env = EnvironmentOptions(sim_folder=f"sim_petsc_speedup_example_{solver}") + time_opts = Time(dt=0.05, Tend=0.05, split_algo="LieTrotter") + domain = domains.Cuboid() # periodic: required for PETScSolver's DirectionalDerivativeOperator + equil = equils.HomogenSlab(B0z=1.0, n0=1.0) + grid = grids.TensorProductGrid(num_elements=(32, 32, 32)) + derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(None, None, None)) # fully periodic + + sim = Simulation( + model=model, + params_path=None, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + ) + + model.kinetic_ions.set_markers( + 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), + bufsize=0.4, + ) + + model.propagators.gc_poisson.options.solver = solver + model.propagators.gc_poisson.options.solver_params = SolverParameters( + tol=1e-10, + maxiter=5_000, + pc_type=pc_type, # ignored by pcg, see SolverParameters.pc_type + ) + model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options( + algo="explicit", + evaluate_e_field=True, + ) + + 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 = 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) + + # real setup: particle loading, Derham/mass operators, and one (untimed) Poisson solve + sim.run(one_time_step=True) + + poisson = model.propagators.gc_poisson + dt = time_opts.dt + + poisson(dt) # warm-up: pays for one-time matrix assembly / gamg hierarchy construction + t0 = time.perf_counter() + for _ in range(n_solves): + poisson(dt) + elapsed = (time.perf_counter() - t0) / n_solves + + info = poisson._solver.get_info() if hasattr(poisson._solver, "get_info") else poisson._solver._info + phi = model.em_fields.phi.spline.vector.toarray() + return elapsed, info, phi + + +def main(n_solves: int = 5): + print(f"Timing {n_solves} repeated Poisson solves per solver (after one warm-up call) ...\n") + + t_pcg, info_pcg, phi_pcg = _build_and_run("pcg", pc_type="jacobi", n_solves=n_solves) + t_petsc, info_petsc, phi_petsc = _build_and_run("petsc", pc_type="gamg", n_solves=n_solves) + + mean_removed_pcg = phi_pcg - phi_pcg.mean() + mean_removed_petsc = phi_petsc - phi_petsc.mean() + rel_err = xp.linalg.norm(mean_removed_pcg - mean_removed_petsc) / xp.linalg.norm(phi_pcg) + + print(f"pcg (unpreconditioned) : {t_pcg * 1e3:9.2f} ms/solve niter={info_pcg.get('niter')}") + print(f"petsc + gamg : {t_petsc * 1e3:9.2f} ms/solve niter={info_petsc.get('niter')}") + print(f"\nspeedup: {t_pcg / t_petsc:.2f}x") + print(f"relative solution mismatch (mean-removed): {rel_err:.2e}") + assert rel_err < 1e-6, f"pcg/petsc solutions disagree by {rel_err:.2e}, expected < 1e-6" + print("\nSolutions agree -- the speedup above is not at the cost of correctness.") + + +if __name__ == "__main__": + main() diff --git a/src/struphy/linear_algebra/schur_solver.py b/src/struphy/linear_algebra/schur_solver.py index 36c0c7956..d264bc73f 100644 --- a/src/struphy/linear_algebra/schur_solver.py +++ b/src/struphy/linear_algebra/schur_solver.py @@ -5,6 +5,7 @@ from scope_profiler import ProfileManager from struphy.linear_algebra.solver import SolverParameters +from struphy.linear_algebra.solver import inverse as struphy_inverse class SchurSolver: @@ -70,9 +71,30 @@ def __init__( # linear operators self._A = A self._BC = BC + # Set by the A/BC property setters whenever a caller reassigns either (e.g. + # VlasovAmpereCoupling/EfieldWeightsCoupling rebuild `.BC` from a fresh, particle-dependent + # operator every call); only consulted for petsc, see below. + self._schur_dirty = True + + self._is_petsc = solver_name == "petsc" + + if self._is_petsc: + # PETScSolver caches its assembled PETSc.Mat by the *object identity* of `linop` + # (see PETScSolver._get_ksp), rebuilding only when that identity changes -- exactly + # the mechanism ImplicitDiffusion relies on (see its lhs-operator caching). The + # in-place-mutated `self._schur` buffer below defeats that: it is the same Python + # object on every call, so PETScSolver would keep the *first* call's matrix forever, + # silently going stale if `dt`, `A` or `BC` ever change. So for petsc, `self._schur` + # is instead a fresh composite operator, rebuilt only when `dt` changes or `.A`/`.BC` + # were reassigned (`self._schur_dirty`) -- see __call__. Callers that mutate an + # operator obtained via the `A`/`BC` getters in place, without reassigning it through + # the setter, would not be detected -- no current caller does this. + self._schur = None + self._schur_dt = None + else: + # Allocate memory for matrices used in solving the Schur system + self._schur = A.copy() - # Allocate memory for matrices used in solving the Schur system - self._schur = A.copy() self._rhs_m = A.copy() # initialize solver with dummy matrix A @@ -83,7 +105,19 @@ def __init__( if precond is not None: kwargs["pc"] = precond - self._solver = inverse(A, solver_name, **kwargs) + if self._is_petsc: + # struphy's inverse() dispatches "petsc" to PETScSolver and forwards pc_type; the + # dummy operator here is just to build the solver object -- __call__ always assigns + # the real one via `self._solver.linop` before solving. + self._solver = struphy_inverse(A, solver_name, **kwargs) + else: + # pc_type is petsc-only (see struphy.linear_algebra.solver.SolverParameters); this + # branch goes straight to feectools' own `inverse` (imported directly above, not + # struphy's petsc-aware wrapper, which would otherwise strip it), whose + # InverseLinearOperator subclasses forward unknown kwargs straight to their + # constructor and would raise on it. + kwargs.pop("pc_type", None) + self._solver = inverse(A, solver_name, **kwargs) # right-hand side vector (avoids temporary memory allocation!) self._rhs = A.codomain.zeros() @@ -102,11 +136,17 @@ def BC(self): def A(self, a): """Upper left block from [[A B], [C Id]].""" self._A = a + # e.g. VlasovAmpereCoupling/EfieldWeightsCoupling reassign `.A`/`.BC` to a fresh + # (possibly particle-dependent) operator every call; `x.A *= y`-style augmented + # assignment also lands here (Python always re-invokes the setter). See the petsc + # cache-invalidation note in __init__/__call__. + self._schur_dirty = True @BC.setter def BC(self, bc): """Product from [[A B], [C Id]].""" self._BC = bc + self._schur_dirty = True @profile @ProfileManager.profile("solve: SchurSolver") @@ -141,12 +181,19 @@ def __call__(self, xn, Byn, dt, out=None): assert xn.space == self._A.domain assert Byn.space == self._A.codomain - # left- and right-hand side operators - self._schur *= 0.0 - self._schur += self._BC - self._schur *= -(dt**2) - self._schur += self._A + # left-hand side operator + if self._is_petsc: + if self._schur is None or dt != self._schur_dt or self._schur_dirty: + self._schur = self._A - (dt**2) * self._BC + self._schur_dt = dt + self._schur_dirty = False + else: + self._schur *= 0.0 + self._schur += self._BC + self._schur *= -(dt**2) + self._schur += self._A + # right-hand side operator self._rhs_m *= 0.0 self._rhs_m += self._BC self._rhs_m *= dt**2 @@ -224,7 +271,12 @@ def __init__(self, M, solver_name, **solver_params): self._S = self._A - self._B @ self._C - self._solver = inverse(self._S, solver_name, **solver_params) + # struphy_inverse dispatches solver_name="petsc" to PETScSolver and safely strips + # petsc-only kwargs (e.g. pc_type) for every other solver -- see SchurSolver, which needs + # this same dispatch but (unlike this class) also has to handle a stale-cache hazard from + # in-place operator mutation; no such hazard here since callers rebuild this whole object + # fresh each call rather than mutating `self._S` in place. + self._solver = struphy_inverse(self._S, solver_name, **solver_params) # right-hand side vector (avoids temporary memory allocation!) self._rhs = self._A.codomain.zeros() @@ -342,7 +394,8 @@ def __init__(self, M, solver_name, **solver_params): self._S = self._A - self._B @ self._C - self._D @ self._E - self._solver = inverse(self._S, solver_name, **solver_params) + # see SchurSolverFull.__init__'s note on struphy_inverse + self._solver = struphy_inverse(self._S, solver_name, **solver_params) # right-hand side vector (avoids temporary memory allocation!) self._rhs = self._A.codomain.zeros() diff --git a/src/struphy/linear_algebra/solver.py b/src/struphy/linear_algebra/solver.py index f01ff07d9..06230321d 100644 --- a/src/struphy/linear_algebra/solver.py +++ b/src/struphy/linear_algebra/solver.py @@ -5,6 +5,58 @@ logger = logging.getLogger("struphy") +# kwargs accepted by struphy.linear_algebra.petsc_solver.PETScSolver.__init__ +_PETSC_SOLVER_KWARGS = ("x0", "tol", "maxiter", "verbose", "recycle", "ksp_type", "pc_type", "near_null_space") + + +def inverse(A, solver: str, **kwargs): + """Create an (approximate) inverse of ``A``. + + Thin wrapper around :func:`feectools.linalg.solvers.inverse` that additionally + supports ``solver="petsc"``, dispatching to + :class:`~struphy.linear_algebra.petsc_solver.PETScSolver`. For all other solver + names this simply delegates to the feectools implementation. + + Parameters + ---------- + A : feectools.linalg.basic.LinearOperator + Left-hand-side matrix of the linear system. For ``solver="petsc"``, see + :func:`struphy.linear_algebra.petsc_solver._assemble_petsc_matrix` for the + supported operator types -- this includes plain assembled matrices as well as + composite operators such as ``grad.T @ M @ grad``. + + solver : str + Preferred iterative solver, one of feectools' options ('cg', 'pcg', + 'bicg', 'bicgstab', 'pbicgstab', 'minres', 'lsmr', 'gmres') or 'petsc'. + + Returns + ------- + obj : feectools.linalg.basic.InverseLinearOperator + A linear operator acting as the (approximate) inverse of A. + """ + if solver == "petsc": + from struphy.linear_algebra.petsc_solver import PETScSolver + + if kwargs.get("pc") is not None: + logger.debug("PETScSolver ignores the feectools 'pc' preconditioner; use 'pc_type' instead.") + + petsc_kwargs = {k: v for k, v in kwargs.items() if k in _PETSC_SOLVER_KWARGS} + petsc_kwargs.setdefault("ksp_type", "cg") + petsc_kwargs.setdefault("pc_type", "jacobi") + + return PETScSolver(A, **petsc_kwargs) + + # pc_type/ksp_type/near_null_space are petsc-only (see _PETSC_SOLVER_KWARGS above); + # feectools' InverseLinearOperator subclasses forward unknown kwargs straight to their + # constructor and would raise on them, so they never reach this branch. + kwargs.pop("pc_type", None) + kwargs.pop("ksp_type", None) + kwargs.pop("near_null_space", None) + + from feectools.linalg.solvers import inverse as feectools_inverse + + return feectools_inverse(A, solver, **kwargs) + @dataclass class SolverParameters: @@ -14,6 +66,11 @@ class SolverParameters: maxiter: int = 3000 info: bool = False recycle: bool = True + pc_type: LiteralOptions.OptsPETScPrecond = "jacobi" + """Preconditioner for ``solver="petsc"`` only (ignored otherwise): PETSc's ``PCType`` + name, e.g. ``"jacobi"`` (cheap, diagonal) or ``"gamg"`` (algebraic multigrid -- far + stronger for large, ill-conditioned systems, but with more setup overhead per matrix + assembly).""" def __post_init__(self): self.verbose = False diff --git a/src/struphy/linear_algebra/tests/test_petsc_directional_derivative.py b/src/struphy/linear_algebra/tests/test_petsc_directional_derivative.py new file mode 100644 index 000000000..0e81bb08d --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_directional_derivative.py @@ -0,0 +1,68 @@ +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.mpi import mpi as MPI + +from struphy.feec.psydac_derham import Derham +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.petsc_solver import _directional_derivative_to_stencil_matrix +from struphy.topology.grids import TensorProductGrid + + +def _random_fill(v, seed): + from feectools.linalg.block import BlockVector + + xp.random.seed(seed) + if isinstance(v, BlockVector): + for b in v.blocks: + b._data[:] = xp.random.random(b._data.shape) + else: + v._data[:] = xp.random.random(v._data.shape) + v.update_ghost_regions() + return v + + +def test_directional_derivative_matches_grad_on_periodic_domain(): + """The StencilMatrix built by _directional_derivative_to_stencil_matrix must reproduce + every block of derham.grad and derham.grad.T exactly, on a fully periodic domain. + """ + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + + grid = TensorProductGrid(num_elements=[6, 6, 6]) + derham_opts = DerhamOptions(degree=[2, 2, 2], bcs=(None, None, None)) + derham = Derham(grid, derham_opts, comm=comm) + + for op, name in [(derham.grad, "grad"), (derham.grad.T, "grad.T")]: + for i, j in op.nonzero_block_indices: + block = op[i, j] + M = _directional_derivative_to_stencil_matrix(block) + + v = _random_fill(block.domain.zeros(), seed=100 + i + 10 * j + rank) + err = xp.linalg.norm((block.dot(v) - M.dot(v)).toarray()) + assert err < 1e-12, f"{name}[{i},{j}] mismatch: err={err:.3e}" + + +def test_directional_derivative_raises_on_nonperiodic_axis(): + """A non-periodic differentiation axis must raise NotImplementedError rather than + silently produce wrong results (see the docstring of + _directional_derivative_to_stencil_matrix for the unresolved root cause). + """ + comm = MPI.COMM_WORLD + + grid = TensorProductGrid(num_elements=[6, 6, 6]) + derham_opts = DerhamOptions(degree=[2, 2, 2], bcs=(("free", "free"), None, None)) + derham = Derham(grid, derham_opts, comm=comm) + + op = derham.grad.T[0, 0] + assert op.domain.periods[op._diffdir] is False + + with pytest.raises(NotImplementedError): + _directional_derivative_to_stencil_matrix(op) + + +if __name__ == "__main__": + test_directional_derivative_matches_grad_on_periodic_domain() + test_directional_derivative_raises_on_nonperiodic_axis() diff --git a/src/struphy/linear_algebra/tests/test_petsc_l2_projector.py b/src/struphy/linear_algebra/tests/test_petsc_l2_projector.py new file mode 100644 index 000000000..f0d492d2b --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_l2_projector.py @@ -0,0 +1,46 @@ +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.mpi import mpi as MPI + +from struphy.feec.mass import L2Projector, WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.fields_background.equils import HomogenSlab +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.topology.grids import TensorProductGrid + + +@pytest.mark.parametrize("space_id", ["H1", "L2"]) +def test_l2_projector_petsc_matches_pcg(space_id): + """L2Projector(solver_name="petsc") must match L2Projector(solver_name="pcg") for a real mass matrix.""" + comm = MPI.COMM_WORLD + + domain = Cuboid() + equil = HomogenSlab(n0=2.0) + equil.domain = domain + + grid = TensorProductGrid(num_elements=[8, 8, 8]) + derham_opts = DerhamOptions(degree=[2, 2, 2]) + derham = Derham(grid, derham_opts, comm=comm, domain=domain) + + mass_ops = WeightedMassOperators(derham, domain, eq_mhd=equil) + + def rhs(e1, e2, e3): + return xp.sin(2 * xp.pi * e1) * xp.cos(2 * xp.pi * e2) * xp.cos(2 * xp.pi * e3) + + proj_pcg = L2Projector(space_id, mass_ops, solver_name="pcg") + proj_petsc = L2Projector(space_id, mass_ops, solver_name="petsc") + + b = proj_pcg.get_dofs(rhs, apply_bc=True) + + x_pcg = proj_pcg.solve(b) + x_petsc = proj_petsc.solve(b) + + assert xp.linalg.norm((x_petsc - x_pcg).toarray()) < 1e-6 + + +if __name__ == "__main__": + test_l2_projector_petsc_matches_pcg("H1") diff --git a/src/struphy/linear_algebra/tests/test_petsc_poisson_solve.py b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve.py new file mode 100644 index 000000000..2e48d6338 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve.py @@ -0,0 +1,85 @@ +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.mpi import mpi as MPI + +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable +from struphy.propagators.base import Propagator +from struphy.propagators.poisson_solve import PoissonSolve +from struphy.topology.grids import TensorProductGrid + + +def test_poisson_solve_petsc_matches_pcg(): + """PoissonSolve(solver="petsc") must match PoissonSolve(solver="pcg") on a fully periodic domain. + + PETScSolver can only assemble derham.grad on a *periodic* differentiation axis (see + struphy.linear_algebra.petsc_solver._directional_derivative_to_stencil_matrix); this is why + the domain here is fully periodic rather than using Dirichlet/Neumann boundaries. + """ + comm = MPI.COMM_WORLD + + domain = Cuboid() + + grid = TensorProductGrid(num_elements=[10, 10, 10]) + derham_opts = DerhamOptions(degree=[2, 2, 2], bcs=(None, None, None)) + derham = Derham(grid, derham_opts, comm=comm) + + mass_ops = WeightedMassOperators(derham, domain) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + def sol_xyz(x, y, z): + return xp.sin(2 * xp.pi * x) * xp.cos(2 * xp.pi * y) + + def rho_xyz(x, y, z): + return sol_xyz(x, y, z) * ((2 * xp.pi) ** 2 + (2 * xp.pi) ** 2) + + def rho_pulled(e1, e2, e3): + return domain.pull(rho_xyz, e1, e2, e3, kind="0", squeeze_out=False) + + def run(solver_name): + solver_params = SolverParameters(tol=1e-11, maxiter=3000, info=False, recycle=False) + + phi = FEECVariable(space="H1") + phi.allocate(derham=derham, domain=domain) + + prop = PoissonSolve(rho=rho_pulled) + prop.variables.phi = phi + prop.options = prop.Options( + stab_eps=1e-12, + solver=solver_name, + precond="MassMatrixPreconditioner", + solver_params=solver_params, + ) + prop.allocate() + prop(1.0) + return phi + + phi_pcg = run("pcg") + phi_petsc = run("petsc") + + e1 = xp.linspace(0.0, 1.0, 20) + e2 = xp.linspace(0.0, 1.0, 20) + e3 = xp.array([0.5]) + + val_pcg = domain.push(phi_pcg.spline, e1, e2, e3, kind="0") + val_petsc = domain.push(phi_petsc.spline, e1, e2, e3, kind="0") + + x, y, z = domain(e1, e2, e3) + analytic = sol_xyz(x, y, z) + + assert xp.max(xp.abs(val_petsc - analytic)) < 1e-2 + assert xp.max(xp.abs(val_petsc - val_pcg)) < 1e-6 + + +if __name__ == "__main__": + test_poisson_solve_petsc_matches_pcg() diff --git a/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py new file mode 100644 index 000000000..c67379f40 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py @@ -0,0 +1,125 @@ +import cunumpy as xp +import pytest +from cunumpy import PyccelKernel + +pytest.importorskip("petsc4py") + +from feectools.ddm.mpi import mpi as MPI + +from struphy import LoadingParameters, WeightsParameters, maxwellians, perturbations +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable +from struphy.pic.accumulation import accum_kernels +from struphy.pic.accumulation.particles_to_grid import ParticlesToGrid +from struphy.pic.particles import Particles6D +from struphy.propagators.base import Propagator +from struphy.propagators.poisson_solve import PoissonSolve +from struphy.topology.grids import TensorProductGrid + + +class _FakePICVariable: + """Minimal duck-typed stand-in for a model's PICVariable, since ParticlesToGrid only reads .particles.""" + + def __init__(self, particles): + self.particles = particles + + +def test_poisson_solve_petsc_matches_pcg_with_real_pic_deposition(): + """PoissonSolve(solver="petsc") must match PoissonSolve(solver="pcg") when driven by a real + particle-in-cell charge-density deposition (not a synthetic/manufactured source), reproducing + the setup of examples/VlasovAmpereOneSpecies/strong_Landau_damping (Maxwellian3D background + + ModesCos perturbation, control-variate weights). + + Compared *mean-removed* (matching struphy.linear_algebra.petsc_examples_benchmark's + methodology): this case's background is a large uniform density (n=1.0 everywhere), so the + charge density's mean/DC component is large, and the near-singular stab_eps=1e-8 regularizes + the constant/DC mode only very weakly -- feectools' pcg divides that large DC charge by the + tiny stab_eps, landing on an essentially arbitrary large DC potential offset (~-168 in + testing) that is numerical-noise-amplification, not a physically meaningful answer. PETScSolver + now registers the constant mode as a near null space for exactly this operator (see + PETScSolver's near_null_space docstring and ImplicitDiffusion.allocate's near_null_space="constant" + comment) -- the fix for a real MPI-rank-dependent correctness bug this same near-singular + regime caused under >1 rank -- which makes it correctly and robustly discard that + inconsistent DC component instead (mean exactly 0) rather than reproducing pcg's arbitrary + noise-amplified one. The physically meaningful, oscillatory part of the solution still needs + to match to near machine precision, which is what this test actually checks. + """ + comm = MPI.COMM_WORLD + + domain = Cuboid() + grid = TensorProductGrid(num_elements=[10, 10, 10]) + derham_opts = DerhamOptions(degree=[2, 2, 2], bcs=(None, None, None)) + derham = Derham(grid, derham_opts, comm=comm) + mass_ops = WeightedMassOperators(derham, domain) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + background = maxwellians.Maxwellian3D(n=(1.0, None)) + perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) + init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) + + domain_array = derham.domain_array + nprocs = derham.domain_decomposition.nprocs + + def run(solver_name): + loading_params = LoadingParameters(Np=20_000, seed=1234) + weights_params = WeightsParameters(control_variate=True) + + particles = Particles6D( + comm_world=comm, + clone_config=None, + loading_params=loading_params, + weights_params=weights_params, + domain=domain, + domain_decomp=(domain_array, nprocs), + background=background, + initial_condition=init, + ) + particles.draw_markers() + if comm.Get_size() > 1: + particles.mpi_sort_markers() + particles.initialize_weights() + + rho = ParticlesToGrid( + _FakePICVariable(particles), + "H1", + PyccelKernel(accum_kernels.charge_density_0form), + ) + + phi = FEECVariable(space="H1") + phi.allocate(derham=derham, domain=domain) + + solver_params = SolverParameters(tol=1e-10, maxiter=20000, info=False, recycle=False) + + prop = PoissonSolve(rho=rho) + prop.variables.phi = phi + prop.options = prop.Options( + stab_eps=1e-8, + solver=solver_name, + precond="MassMatrixPreconditioner", + solver_params=solver_params, + ) + prop.allocate() + if solver_name == "petsc": + prop._solver._options["pc_type"] = "gamg" + prop._solver._ksp = None + prop(0.05) + return phi.spline.vector.toarray() + + sol_pcg = run("pcg") + sol_petsc = run("petsc") + + mean_removed_pcg = sol_pcg - sol_pcg.mean() + mean_removed_petsc = sol_petsc - sol_petsc.mean() + rel_err = xp.linalg.norm(mean_removed_pcg - mean_removed_petsc) / xp.linalg.norm(sol_pcg) + assert rel_err < 1e-6 + + +if __name__ == "__main__": + test_poisson_solve_petsc_matches_pcg_with_real_pic_deposition() diff --git a/src/struphy/linear_algebra/tests/test_petsc_schur_solver.py b/src/struphy/linear_algebra/tests/test_petsc_schur_solver.py new file mode 100644 index 000000000..63e71917b --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_schur_solver.py @@ -0,0 +1,185 @@ +"""Regression tests for SchurSolver's solver="petsc" support (struphy.linear_algebra.schur_solver). + +SchurSolver previously always imported feectools' own `inverse`, which does not recognize the +name "petsc" -- so `solver="petsc"` would raise for every propagator built on it (MaxwellWeakAmpere, +VlasovAmpereCoupling, EfieldWeightsCoupling, CurlCurlSolve, ...), even though those propagators' +Options all declare `LiteralOptions.OptsSymmSolver` (which includes "petsc") as their solver type. + +Wiring petsc in was not just an import swap: PETScSolver caches its assembled PETSc.Mat by the +*object identity* of the operator assigned to `.linop` (see PETScSolver._get_ksp), rebuilding only +when that identity changes. SchurSolver's non-petsc path mutates its `self._schur` buffer *in +place* every call (`self._schur *= 0.0; += ...`) -- the same Python object every time, which would +make PETScSolver silently reuse a stale matrix forever after the first call. Two call patterns +exist among current callers, and both need to be correct: + +- MaxwellWeakAmpere never reassigns `.A`/`.BC` after construction (both are geometric, constant + operators) -- caching by `dt` alone is correct and safe there. +- VlasovAmpereCoupling (and EfieldWeightsCoupling) reassign `.BC` to a fresh, particle-dependent + operator via the property setter on *every* call -- caching there must be invalidated every + time, tracked via a dirty flag set in the `A`/`BC` property setters (Python routes both + `x.BC = y` and the augmented `x.BC *= y` through the setter). + +These tests exercise both patterns directly (not synthetically) via the real propagators. +""" + +import shutil +import tempfile + +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.mpi import mpi as MPI + +from struphy import ( + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.geometry.domains import Cuboid +from struphy.linear_algebra.solver import SolverParameters +from struphy.models import VlasovAmpereOneSpecies +from struphy.models.variables import FEECVariable +from struphy.propagators.base import Propagator +from struphy.propagators.maxwell_weak_ampere import MaxwellWeakAmpere +from struphy.topology.grids import TensorProductGrid + + +def test_maxwell_weak_ampere_petsc_matches_pcg(): + """MaxwellWeakAmpere(solver="petsc") must match solver="pcg" over several implicit timesteps. + + Exercises SchurSolver's petsc dt-only cache-invalidation path (see module docstring): + MaxwellWeakAmpere never reassigns `.A`/`.BC` after allocate(), so the same lhs operator must + be correctly reused across all calls at fixed dt. + """ + comm = MPI.COMM_WORLD + + domain = Cuboid() + grid = TensorProductGrid(num_elements=[6, 6, 6]) + derham_opts = DerhamOptions(degree=[2, 2, 2], bcs=(None, None, None)) + derham = Derham(grid, derham_opts, comm=comm) + mass_ops = WeightedMassOperators(derham, domain) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + def run(solver_name): + e_field = FEECVariable(space="Hcurl") + e_field.add_perturbation(perturbations.ModesCos(amps=(0.1,), ls=(1,), comp=0)) + e_field.allocate(derham=derham, domain=domain) + + b_field = FEECVariable(space="Hdiv") + b_field.add_perturbation(perturbations.ModesCos(amps=(0.05,), ls=(1,), comp=1)) + b_field.allocate(derham=derham, domain=domain) + + prop = MaxwellWeakAmpere() + prop.variables.e = e_field + prop.variables.b = b_field + prop.options = prop.Options( + algo="implicit", + solver=solver_name, + solver_params=SolverParameters(tol=1e-11, maxiter=3000), + ) + prop.allocate() + + dt = 0.02 + for _ in range(3): + prop(dt) + + return e_field.spline.vector.toarray(), b_field.spline.vector.toarray() + + e_pcg, b_pcg = run("pcg") + e_petsc, b_petsc = run("petsc") + + rel_err_e = xp.linalg.norm(e_pcg - e_petsc) / xp.linalg.norm(e_pcg) + rel_err_b = xp.linalg.norm(b_pcg - b_petsc) / xp.linalg.norm(b_pcg) + assert rel_err_e < 1e-6, f"e-field mismatch: {rel_err_e:.2e}" + assert rel_err_b < 1e-6, f"b-field mismatch: {rel_err_b:.2e}" + + +def test_vlasov_ampere_coupling_petsc_matches_pcg_with_real_pic_deposition(): + """VlasovAmpereCoupling(solver="petsc") must match solver="pcg" over several real timesteps, + driven by real particle-in-cell deposition (not a synthetic source). + + Exercises SchurSolver's petsc dirty-flag cache-invalidation path (see module docstring): + VlasovAmpereCoupling reassigns `.BC` to a fresh, particle-dependent operator every call, which + the dt-only cache used by MaxwellWeakAmpere's test above would get wrong if applied here. + Goes through the full model/Simulation machinery (unlike the other petsc regression tests in + this directory, which build propagators directly) because VlasovAmpereCoupling requires a real + PICVariable/ParticleSpecies (species.equation_params, weights_params) that is impractical to + duck-type -- see test_petsc_poisson_solve_pic.py, whose ParticlesToGrid-based fake works + because ParticlesToGrid does not check isinstance, unlike VlasovAmpereCoupling.Variables.ions. + """ + comm = MPI.COMM_WORLD + + def run(solver_name, out_folder): + model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + + env = EnvironmentOptions(out_folders=out_folder, sim_folder=f"sim_{solver_name}") + time_opts = Time(dt=0.02, Tend=0.06, split_algo="LieTrotter") + domain = domains.Cuboid(r1=12.56) + grid = grids.TensorProductGrid(num_elements=(8, 8, 8)) + derham_opts = DerhamOptions(degree=(2, 2, 2), bcs=(None, None, None)) + + sim = Simulation( + model=model, + params_path=None, + env=env, + time_opts=time_opts, + domain=domain, + equil=None, + grid=grid, + derham_opts=derham_opts, + ) + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(Np=5_000, seed=1234), + weights_params=WeightsParameters(control_variate=True), + ) + + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.coupling_va.options = model.propagators.coupling_va.Options( + solver=solver_name, + solver_params=SolverParameters(tol=1e-10, maxiter=5000), + ) + model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0") + + background = maxwellians.Maxwellian3D(n=(1.0, None)) + perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) + init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) + model.kinetic_ions.var.add_background(background) + model.kinetic_ions.var.add_initial_condition(init) + + sim.run() # several real steps, not one_time_step: coupling_va.BC changes every call + + return model.em_fields.e_field.spline.vector.toarray() + + out_folder = tempfile.mkdtemp() if comm.Get_rank() == 0 else None + out_folder = comm.bcast(out_folder, root=0) + + try: + e_pcg = run("pcg", out_folder) + e_petsc = run("petsc", out_folder) + finally: + comm.Barrier() + if comm.Get_rank() == 0: + shutil.rmtree(out_folder, ignore_errors=True) + + rel_err = xp.linalg.norm(e_pcg - e_petsc) / xp.linalg.norm(e_pcg) + assert rel_err < 1e-6, f"e-field mismatch: {rel_err:.2e}" + + +if __name__ == "__main__": + test_maxwell_weak_ampere_petsc_matches_pcg() + test_vlasov_ampere_coupling_petsc_matches_pcg_with_real_pic_deposition() diff --git a/src/struphy/linear_algebra/tests/test_petsc_schur_solver_full.py b/src/struphy/linear_algebra/tests/test_petsc_schur_solver_full.py new file mode 100644 index 000000000..1ec7e17ce --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_schur_solver_full.py @@ -0,0 +1,152 @@ +"""Regression test for SchurSolverFull's solver="petsc" support. + +Like SchurSolver (see test_petsc_schur_solver.py), SchurSolverFull/SchurSolverFull3 always +imported feectools' own `inverse` directly, which does not recognize "petsc" -- so +solver="petsc" would raise for the variational MHD propagators built on them +(VariationalPBEvolve, VariationalEntropyEvolve, VariationalMagFieldEvolve, VariationalQBEvolve). + +Unlike SchurSolver, there is no in-place-mutation caching hazard here: `self._S` is built once +in __init__ and never mutated afterwards -- callers (e.g. VariationalQBEvolve) rebuild the whole +SchurSolverFull3 object fresh every Newton iteration rather than reusing one across calls (see +those propagators' "local version to avoid creating new version of LinearOperator every time" +comment, which refers to the *Jacobian's blocks*, not to reusing the Schur solver object itself). +So the fix here is the dispatch alone: use struphy.linear_algebra.solver.inverse (which knows +"petsc" and safely strips petsc-only kwargs for every other solver) instead of feectools' own. + +This test exercises that dispatch directly on a small synthetic block system (matching +test_petsc_solver.py's style), not through a real variational MHD model: those models' Jacobian +blocks involve operator types (BasisProjectionOperator-derived, nonlinear-model-specific) that +have not been checked against _assemble_petsc_matrix's supported set, and building one from +scratch without an existing example/test to adapt was judged too failure-prone to do blind. If a +real variational-MHD case is wired up to use solver="petsc" later, verify it separately. +""" + +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.basic import IdentityOperator +from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace +from feectools.linalg.stencil import StencilMatrix, StencilVector, StencilVectorSpace + +from struphy.linear_algebra.schur_solver import SchurSolverFull, SchurSolverFull3 + + +def _make_space(n, p): + domain_decomposition = DomainDecomposition([n - p], [False], comm=MPI.COMM_WORLD) + cart = CartDecomposition(domain_decomposition, [n], [xp.array([0])], [xp.array([n - 1])], [p], [1]) + return StencilVectorSpace(cart) + + +def _spd_tridiagonal(V, p, scale): + """Banded, symmetric positive-definite StencilMatrix with 2p+1 diagonals on space `V`.""" + A = StencilMatrix(V, V) + A[:, -p:0] = -scale + A[:, 0:1] = 2 * p * scale + A[:, 1 : p + 1] = -scale + A.remove_spurious_entries() + return A + + +def test_schur_solver_full_petsc_matches_pcg(): + """SchurSolverFull(solver_name="petsc") must solve [[A B],[C Id]] to the same accuracy as + solver_name="pcg", for a small synthetic system [[A B],[C Id]] x = v. + """ + n, p = 12, 1 + xp.random.seed(0) + + V = _make_space(n, p) + A = _spd_tridiagonal(V, p, scale=1.0) + B = _spd_tridiagonal(V, p, scale=0.01) + C = _spd_tridiagonal(V, p, scale=0.01) + + domain = BlockVectorSpace(V, V) + M = BlockLinearOperator(domain, domain) + M[0, 0] = A + M[0, 1] = B + M[1, 0] = C + M[1, 1] = IdentityOperator(V) + + s = V.starts[0] + e = V.ends[0] + bx = StencilVector(V) + bx[s : e + 1] = xp.random.random(e + 1 - s) + by = StencilVector(V) + by[s : e + 1] = xp.random.random(e + 1 - s) + + v = BlockVector(domain) + v[0] = bx + v[1] = by + + solver_kwargs = {"pc": None, "tol": 1e-13, "maxiter": 2000, "verbose": False, "recycle": False} + solver_pcg = SchurSolverFull(M, "pcg", **solver_kwargs) + solver_petsc = SchurSolverFull(M, "petsc", **solver_kwargs) + + x_pcg = solver_pcg.dot(v) + x_petsc = solver_petsc.dot(v) + + err_x = xp.linalg.norm((x_pcg[0] - x_petsc[0]).toarray()) + err_y = xp.linalg.norm((x_pcg[1] - x_petsc[1]).toarray()) + assert err_x < 1e-8, f"x-block mismatch: {err_x:.2e}" + assert err_y < 1e-8, f"y-block mismatch: {err_y:.2e}" + + +def test_schur_solver_full3_petsc_matches_pcg(): + """SchurSolverFull3(solver_name="petsc") must solve [[A B D],[C Id 0],[E 0 Id]] to the same + accuracy as solver_name="pcg", for a small synthetic system. + """ + n, p = 12, 1 + xp.random.seed(1) + + V = _make_space(n, p) + A = _spd_tridiagonal(V, p, scale=1.0) + B = _spd_tridiagonal(V, p, scale=0.01) + C = _spd_tridiagonal(V, p, scale=0.01) + D = _spd_tridiagonal(V, p, scale=0.01) + E = _spd_tridiagonal(V, p, scale=0.01) + + domain = BlockVectorSpace(V, V, V) + M = BlockLinearOperator(domain, domain) + M[0, 0] = A + M[0, 1] = B + M[1, 0] = C + M[1, 1] = IdentityOperator(V) + M[0, 2] = D + M[2, 0] = E + M[2, 2] = IdentityOperator(V) + + s = V.starts[0] + e = V.ends[0] + bx = StencilVector(V) + bx[s : e + 1] = xp.random.random(e + 1 - s) + by = StencilVector(V) + by[s : e + 1] = xp.random.random(e + 1 - s) + bz = StencilVector(V) + bz[s : e + 1] = xp.random.random(e + 1 - s) + + v = BlockVector(domain) + v[0] = bx + v[1] = by + v[2] = bz + + solver_kwargs = {"pc": None, "tol": 1e-13, "maxiter": 2000, "verbose": False, "recycle": False} + solver_pcg = SchurSolverFull3(M, "pcg", **solver_kwargs) + solver_petsc = SchurSolverFull3(M, "petsc", **solver_kwargs) + + x_pcg = solver_pcg.dot(v) + x_petsc = solver_petsc.dot(v) + + err_x = xp.linalg.norm((x_pcg[0] - x_petsc[0]).toarray()) + err_y = xp.linalg.norm((x_pcg[1] - x_petsc[1]).toarray()) + err_z = xp.linalg.norm((x_pcg[2] - x_petsc[2]).toarray()) + assert err_x < 1e-8, f"x-block mismatch: {err_x:.2e}" + assert err_y < 1e-8, f"y-block mismatch: {err_y:.2e}" + assert err_z < 1e-8, f"z-block mismatch: {err_z:.2e}" + + +if __name__ == "__main__": + test_schur_solver_full_petsc_matches_pcg() + test_schur_solver_full3_petsc_matches_pcg() diff --git a/src/struphy/linear_algebra/tests/test_petsc_solver.py b/src/struphy/linear_algebra/tests/test_petsc_solver.py new file mode 100644 index 000000000..176b28fd4 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_solver.py @@ -0,0 +1,66 @@ +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.solvers import inverse +from feectools.linalg.stencil import StencilMatrix, StencilVector, StencilVectorSpace + +from struphy.linear_algebra.petsc_solver import PETScSolver + + +def _define_tridiagonal_spd_system(n, p): + """Banded, symmetric positive-definite StencilMatrix with 2p+1 diagonals, and a random exact solution.""" + domain_decomposition = DomainDecomposition([n - p], [False], comm=MPI.COMM_WORLD) + cart = CartDecomposition(domain_decomposition, [n], [xp.array([0])], [xp.array([n - 1])], [p], [1]) + V = StencilVectorSpace(cart) + s = V.starts[0] + e = V.ends[0] + + A = StencilMatrix(V, V) + A[:, -p:0] = -1.0 + A[:, 0:1] = 2 * p + A[:, 1 : p + 1] = -1.0 + A.remove_spurious_entries() + + xe = StencilVector(V) + xe[s : e + 1] = xp.random.random(e + 1 - s) + + return V, A, xe + + +@pytest.mark.parametrize("n", [8, 15]) +@pytest.mark.parametrize("p", [1, 2]) +def test_petsc_solver_matches_cg(n, p): + """PETScSolver must solve Ax=b to the same accuracy as feectools' native CG solver.""" + xp.random.seed(n * p) + + _, A, xe = _define_tridiagonal_spd_system(n, p) + + b = A @ xe + + ref_solver = inverse(A, "cg", tol=1e-13, maxiter=2000, verbose=False, recycle=False) + x_ref = ref_solver.solve(b) + + petsc_solver = PETScSolver(A, tol=1e-13, maxiter=2000, ksp_type="cg", pc_type="none") + x_petsc = petsc_solver.solve(b) + + info = petsc_solver.get_info() + assert info["success"] + + error_vs_exact = xp.linalg.norm((x_petsc - xe).toarray()) + assert error_vs_exact < 1e-8 + + error_vs_ref = xp.linalg.norm((x_petsc - x_ref).toarray()) + assert error_vs_ref < 1e-6 + + # re-solving with an unchanged operator (KSP/Mat cache reused) must still be correct + b2 = A @ x_petsc + x_petsc2 = petsc_solver.solve(b2) + assert xp.linalg.norm((x_petsc2 - x_petsc).toarray()) < 1e-8 + + +if __name__ == "__main__": + test_petsc_solver_matches_cg(15, 2) diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index 5261d0fa9..6595201c4 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -1,17 +1,17 @@ import logging +from collections.abc import Callable from dataclasses import dataclass -from typing import Callable, Literal +from typing import Literal import cunumpy as xp from feectools.linalg.basic import IdentityOperator -from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilVector from line_profiler import profile from scope_profiler import ProfileManager from struphy.feec.mass import L2Projector, WeightedMassOperator from struphy.io.options import LiteralOptions, OptionsBase -from struphy.linear_algebra.solver import SolverParameters +from struphy.linear_algebra.solver import SolverParameters, inverse 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 @@ -363,6 +363,18 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, recycle=self.options.solver_params.recycle, + pc_type=self.options.solver_params.pc_type, + # self._diffusion_op = grad.T @ diffusion_mat @ grad structurally has the constant + # function in its kernel on a periodic domain (grad(constant) = 0), regardless of + # diffusion_mat or how small/large sigma_1 (stab_eps) is -- and PETScSolver only + # supports this operator on periodic domains to begin with (see + # _directional_derivative_to_stencil_matrix), so this is always a valid hint where it + # applies at all. Ignored for solver != "petsc". Without it, PETSc+gamg was found to + # silently converge (small reported residual) to a solution that disagrees with + # feectools' own solver -- worse, and more MPI-rank-count-dependent, as rank count + # grows -- for exactly this near-singular regime; see PETScSolver's near_null_space + # docstring. + near_null_space="constant", ) # allocate memory for solution @@ -371,6 +383,11 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: self._rhs2 = phi.space.zeros() self._tmp_src = phi.space.zeros() + # cache for the lhs operator (see __call__): avoids rebuilding (and re-assembling, for + # e.g. solver="petsc") a fresh operator every call when sig_1 (hence dt) is unchanged + self._lhs_op = None + self._lhs_op_sig_1 = None + @property def sources(self) -> list[StencilVector | FEECVariable | AccumulatorVector]: """ @@ -458,8 +475,13 @@ def __call__(self, dt): proj = L2Projector("H1", self.mass_ops) self.diagnostic.spline.vector = proj.solve(rhs) - # compute lhs - self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op + # compute lhs (reuse the cached operator when sig_1 is unchanged, e.g. constant dt -- + # this lets InverseLinearOperator subclasses that cache on `linop` identity, such as + # PETScSolver, avoid re-assembling the operator on every call) + if self._lhs_op is None or sig_1 != self._lhs_op_sig_1: + self._lhs_op = sig_1 * self._stab_mat + self._diffusion_op + self._lhs_op_sig_1 = sig_1 + self._solver.linop = self._lhs_op # solve with ProfileManager.profile_region(self._solve_region, functions=[self._solver.solve]):