From ceef222d78cef1683a45254dc260e414adfcad8b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Mon, 8 Dec 2025 09:50:14 +0000 Subject: [PATCH 01/41] Fixed spaces for GMRES version of two fluid propagator. --- src/struphy/io/options.py | 62 ++++++++++++++++++- src/struphy/linear_algebra/saddle_point.py | 2 - src/struphy/propagators/propagators_fields.py | 26 +++++--- 3 files changed, 80 insertions(+), 10 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index d57838fcb..2c626df78 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -2,7 +2,67 @@ from dataclasses import dataclass from typing import Literal -from struphy.utils.utils import __dataclass_repr_no_defaults__, all_class_params_are_default, check_option +import cunumpy as xp +from psydac.ddm.mpi import mpi as MPI + +from struphy.physics.physics import ConstantsOfNature + +## Literal options + +# time +SplitAlgos = Literal["LieTrotter", "Strang"] + +# derham +PolarRegularity = Literal[-1, 1] +OptsFEECSpace = Literal["H1", "Hcurl", "Hdiv", "L2", "H1vec"] +OptsVecSpace = Literal["Hcurl", "Hdiv", "H1vec"] + +# fields background +BackgroundTypes = Literal["LogicalConst", "FluidEquilibrium"] + +# perturbations +NoiseDirections = Literal["e1", "e2", "e3", "e1e2", "e1e3", "e2e3", "e1e2e3"] +GivenInBasis = Literal["0", "1", "2", "3", "v", "physical", "physical_at_eta", "norm", None] + +# solvers +OptsSymmSolver = Literal["pcg", "cg"] +OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"] +OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] +OptsSaddlePointSolver = Literal["Uzawa", "GMRES"] # todo +OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] +OptsNonlinearSolver = Literal["Picard", "Newton"] + +# markers +OptsPICSpace = Literal["Particles6D", "DeltaFParticles6D", "Particles5D", "Particles3D"] +OptsMarkerBC = Literal["periodic", "reflect"] +OptsRecontructBC = Literal["periodic", "mirror", "fixed"] +OptsLoading = Literal[ + "pseudo_random", + "sobol_standard", + "sobol_antithetic", + "external", + "restart", + "tesselation", +] +OptsSpatialLoading = Literal["uniform", "disc"] +OptsMPIsort = Literal["each", "last", None] + +# filters +OptsFilter = Literal["fourier_in_tor", "hybrid", "three_point", None] + +# sph +OptsKernel = Literal[ + "trigonometric_1d", + "gaussian_1d", + "linear_1d", + "trigonometric_2d", + "gaussian_2d", + "linear_2d", + "trigonometric_3d", + "gaussian_3d", + "linear_isotropic_3d", + "linear_3d", +] @dataclass diff --git a/src/struphy/linear_algebra/saddle_point.py b/src/struphy/linear_algebra/saddle_point.py index e638b2802..68b60eb15 100644 --- a/src/struphy/linear_algebra/saddle_point.py +++ b/src/struphy/linear_algebra/saddle_point.py @@ -140,9 +140,7 @@ def __init__( self._verbose = solver_params["verbose"] if self._variant == "Inverse_Solver": - self._BT = B.transpose() - # initialize solver with dummy matrix A self._block_domainM = BlockVectorSpace(self._A.domain, self._B.transpose().domain) self._block_codomainM = self._block_domainM self._blocks = [[self._A, self._B.T], [self._B, None]] diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index d6dbb056f..a78af4f5d 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -7767,7 +7767,7 @@ def allocate(self, verbose: bool = False): self.derham.p, self.derham.spl_kind, domain=self.domain, - dirichlet_bc=[[True, True], [False, False], [False, False]], + dirichlet_bc=((True, True), (False, False), (False, False)), ) self._mass_opsv0 = WeightedMassOperators( @@ -8387,7 +8387,7 @@ def allocate(self, verbose: bool = False): pc=None, ) # Allocate memory for call - self._untemp = u.space.zeros() + self._untemp = self.variables.u.spline.vector.space.zeros() elif self._variant == "Uzawa": self._solver_UzawaNumpy = SaddlePointSolver( @@ -8504,16 +8504,27 @@ def __call__(self, dt): self._solver_GMRES.A = _A self._solver_GMRES.B = _B self._solver_GMRES.F = _F - + if self._lifting: ( _sol1, _sol2, info, - ) = self._solver_GMRES(u0.vector, ue0.vector, phinfeec) - un = _sol1[0] + u_prime.vector - uen = _sol1[1] + ue_prime.vector - phin = _sol2 + ) = self._solver_GMRES(u0.vector, ue0.vector, phinfeeccopy.vector) + + un_temp = self.derham.create_spline_function("u", space_id="Hdiv") + un_temp.vector = _sol1[0] + u_prime.vector + + uen_temp = self.derham.create_spline_function("ue", space_id="Hdiv") + uen_temp.vector = _sol1[1] + ue_prime.vector + + phin_temp = self.derham.create_spline_function("phi", space_id="L2") + phin_temp.vector = _sol2 + + un = un_temp.vector + uen = uen_temp.vector + phin = phin_temp.vector + else: ( _sol1, @@ -8524,6 +8535,7 @@ def __call__(self, dt): uen = _sol1[1] phin = _sol2 # write new coeffs into self.feec_vars + max_du, max_due, max_dphi = self.update_feec_variables(u=un, ue=uen, phi=phin) elif self._variant == "Uzawa": From f8e5e4b77621618358cf23675fb206b2bab42f27 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Mon, 8 Dec 2025 09:59:02 +0000 Subject: [PATCH 02/41] Formatting. --- src/struphy/linear_algebra/saddle_point.py | 1 - src/struphy/propagators/propagators_fields.py | 8 ++++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/struphy/linear_algebra/saddle_point.py b/src/struphy/linear_algebra/saddle_point.py index 68b60eb15..ca5633108 100644 --- a/src/struphy/linear_algebra/saddle_point.py +++ b/src/struphy/linear_algebra/saddle_point.py @@ -140,7 +140,6 @@ def __init__( self._verbose = solver_params["verbose"] if self._variant == "Inverse_Solver": - self._block_domainM = BlockVectorSpace(self._A.domain, self._B.transpose().domain) self._block_codomainM = self._block_domainM self._blocks = [[self._A, self._B.T], [self._B, None]] diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index a78af4f5d..4a38c7c01 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -8504,7 +8504,7 @@ def __call__(self, dt): self._solver_GMRES.A = _A self._solver_GMRES.B = _B self._solver_GMRES.F = _F - + if self._lifting: ( _sol1, @@ -8514,13 +8514,13 @@ def __call__(self, dt): un_temp = self.derham.create_spline_function("u", space_id="Hdiv") un_temp.vector = _sol1[0] + u_prime.vector - + uen_temp = self.derham.create_spline_function("ue", space_id="Hdiv") uen_temp.vector = _sol1[1] + ue_prime.vector - + phin_temp = self.derham.create_spline_function("phi", space_id="L2") phin_temp.vector = _sol2 - + un = un_temp.vector uen = uen_temp.vector phin = phin_temp.vector From 3b584fc38aad2851d079b83e81b0b577413e32ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Thu, 5 Mar 2026 19:31:54 +0000 Subject: [PATCH 03/41] 1D periodic, homogeneous and inhomogeneous Dirichlet test cases working --- .gitignore | 11 + 1D_Verification.py | 228 +++++++ .../linear_algebra/rework_saddle_point.py | 567 ++++++++++++++++++ src/struphy/models/rework_model.py | 124 ++++ src/struphy/propagators/rework_propagator.py | 479 +++++++++++++++ 5 files changed, 1409 insertions(+) create mode 100644 1D_Verification.py create mode 100644 src/struphy/linear_algebra/rework_saddle_point.py create mode 100644 src/struphy/models/rework_model.py create mode 100644 src/struphy/propagators/rework_propagator.py diff --git a/.gitignore b/.gitignore index 0fc738535..8d4a130eb 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ + # Distribution / packaging .Python build/ @@ -94,3 +95,13 @@ src/struphy/io/inp/params_* # models list src/struphy/models/models_list src/struphy/models/models_message + +runs/ +bin/ +share/ + +lib64 +pyvenv.cfg + +2D_Verification.py +Restelli_Verification.py \ No newline at end of file diff --git a/1D_Verification.py b/1D_Verification.py new file mode 100644 index 000000000..baa40a32d --- /dev/null +++ b/1D_Verification.py @@ -0,0 +1,228 @@ +from cunumpy import pi, cos, sin, zeros_like, ones_like +from struphy.io.options import EnvironmentOptions, BaseUnits, Time +from struphy.geometry import domains +from struphy.fields_background import equils +from struphy.topology import grids +from struphy.io.options import DerhamOptions +from struphy.initial import perturbations +from struphy import main + +import os +import glob +import cunumpy as xp +import matplotlib.pyplot as plt + +from struphy.models.rework_model import TwoFluidQuasiNeutralToy + +import warnings +# warnings.filterwarnings("error") + + +BC = 'dirichlet_hom' # 'periodic' | 'dirichlet_hom' | 'dirichlet_inhom' + +name = f"runs/sim_1D_{BC}" + +env = EnvironmentOptions(sim_folder=name) +base_units = BaseUnits(kBT=1.0) + +B0 = 1.0 +nu = 10.0 +nu_e = 1.0 +Nel = (32, 1, 1) +p = (2, 1, 1) +epsilon = 1.0 +dt = 1 +Tend = 1 +sigma = 1 + +time_opts = Time(dt=dt, Tend=Tend) +domain = domains.Cuboid() +equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) +grid = grids.TensorProductGrid(Nel=Nel) + +# ---- boundary conditions ---- +if BC == 'periodic': + spl_kind = (True, True, True) + dirichlet_bc = ((False, False), (False, False), (False, False)) + + bcs_u = bcs_ue = { + (0, -1): "periodic", (0, 1): "periodic", + (1, -1): "periodic", (1, 1): "periodic", + (2, -1): "periodic", (2, 1): "periodic", + } + boundary_data_u = boundary_data_ue = None + +elif BC == 'dirichlet_hom': + spl_kind = (False, True, True) + dirichlet_bc = ((True, True), (False, False), (False, False)) + + bcs_u = bcs_ue = { + (0, -1): "dirichlet", (0, 1): "dirichlet", + (1, -1): "periodic", (1, 1): "periodic", + (2, -1): "periodic", (2, 1): "periodic", + } + boundary_data_u = boundary_data_ue = None + +elif BC == 'dirichlet_inhom': + spl_kind = (False, True, True) + dirichlet_bc = ((False, False), (False, False), (False, False)) + + bcs_u = bcs_ue = { + (0, -1): "dirichlet", (0, 1): "dirichlet", + (1, -1): "periodic", (1, 1): "periodic", + (2, -1): "periodic", (2, 1): "periodic", + } + boundary_data_u = { + (0, -1): lambda x, y, z: (zeros_like(x) + 1, zeros_like(x), zeros_like(x)), + (0, 1): lambda x, y, z: (zeros_like(x) + 2, zeros_like(x), zeros_like(x)), + } + boundary_data_ue = { + (0, -1): lambda x, y, z: (zeros_like(x) + 1, zeros_like(x), zeros_like(x)), + (0, 1): lambda x, y, z: (zeros_like(x) + 2, zeros_like(x), zeros_like(x)), + } + +derham_opts = DerhamOptions( + p=p, + spl_kind=spl_kind, + dirichlet_bc=dirichlet_bc, +) + +# ---- manufactured solutions ---- +if BC == 'periodic': + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x) + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + +elif BC == 'dirichlet_hom': + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + +elif BC == 'dirichlet_inhom': + def mms_phi(x, y, z): + return x + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return x + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return x + 1, xp.zeros_like(x), xp.zeros_like(x) + +# ---- source terms ---- +if BC == 'periodic': + def source_function_u(x, y, z): + fx = 2.0 * pi * cos(2 * pi * x) + nu * 4.0 * pi**2 * sin(2 * pi * x) + fy = (sin(2 * pi * x) + 1.0) * B0 / epsilon + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) + fy = -sin(2 * pi * x) * B0 / epsilon + fz = zeros_like(x) + return fx, fy, fz + +elif BC == 'dirichlet_hom': + def source_function_u(x, y, z): + fx = 2.0 * pi * cos(2 * pi * x) + nu * 4.0 * pi**2 * sin(2 * pi * x) + fy = B0 * sin(2 * pi * x) / epsilon + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) + fy = -sin(2 * pi * x) * B0 / epsilon + fz = zeros_like(x) + return fx, fy, fz + +elif BC == 'dirichlet_inhom': + def source_function_u(x, y, z): + fx = ones_like(x) + fy = B0 * (1 + x) / epsilon + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -ones_like(x) - sigma * (1 + x) + fy = -B0 * (1 + x) / epsilon + fz = zeros_like(x) + return fx, fy, fz + +# ---- model ---- +model = TwoFluidQuasiNeutralToy() +model.ions.set_phys_params() +model.electrons.set_phys_params() + +model.propagators.qn_full.options = model.propagators.qn_full.Options( + nu=nu, + nu_e=nu_e, + eps_norm=epsilon, + stab_sigma=sigma, + source_u=source_function_u, + source_ue=source_function_ue, + solver='gmres', + boundary_conditions_u=bcs_u, + boundary_conditions_ue=bcs_ue, + boundary_data_u=boundary_data_u, + boundary_data_ue=boundary_data_ue, +) + +if __name__ == "__main__": + main.run(model, + params_path=__file__, + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=True, + ) + + path = os.path.join(os.getcwd(), name) + main.pproc(path) + simdata = main.load_data(path) + + n1_vals = simdata.grids_log[0] + x = xp.linspace(0, 1, 100) + + os.makedirs(f'{name}/plots', exist_ok=True) + for f in glob.glob(f'{name}/plots/*.png'): + os.remove(f) + + def save_plot(n1_vals, numerical, analytical, ylabel, title, fname, t): + plt.plot(n1_vals, numerical, label='numerical') + plt.plot(x, analytical, '--', label='manufactured') + plt.plot(n1_vals, numerical, 'k.', markersize=4, label='n1 points') + plt.xlabel('n1 (radial)') + plt.ylabel(ylabel) + plt.title(f'{title} at t={t:.3f}') + plt.legend() + plt.grid(True) + plt.savefig(f'{name}/plots/{fname}_{t:.3f}.png', dpi=300) + plt.clf() + + for t in list(simdata.spline_values['ions']['u_log'].keys()): + + u_ions = simdata.spline_values['ions']['u_log'][t] + u_electrons = simdata.spline_values['electrons']['u_log'][t] + phi = simdata.spline_values['em_fields']['phi_log'][t] + + mms_phi_x, _, _ = mms_phi(x, x*0, x*0) + mms_ion_ux, mms_ion_uy, _ = mms_ion_u(x, x*0, x*0) + mms_el_ux, mms_el_uy, _ = mms_electron_u(x, x*0, x*0) + + save_plot(n1_vals, phi[0][:, 0, 0], mms_phi_x, 'φ', 'Electrostatic potential φ', 'plot_potential', t) + save_plot(n1_vals, u_ions[0][:, 0, 0], mms_ion_ux, 'u_x', 'Ion velocity (u_x)', 'plot_ion_ux', t) + save_plot(n1_vals, u_electrons[0][:, 0, 0], mms_el_ux, 'u_x', 'Electron velocity (u_x)', 'plot_electron_ux', t) \ No newline at end of file diff --git a/src/struphy/linear_algebra/rework_saddle_point.py b/src/struphy/linear_algebra/rework_saddle_point.py new file mode 100644 index 000000000..593dc8524 --- /dev/null +++ b/src/struphy/linear_algebra/rework_saddle_point.py @@ -0,0 +1,567 @@ +from typing import Union + +import cunumpy as xp +import scipy as sc +from psydac.linalg.basic import LinearOperator, Vector +from psydac.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace +from psydac.linalg.direct_solvers import SparseSolver +from psydac.linalg.solvers import inverse + +from struphy.linear_algebra.tests.test_saddlepoint_massmatrices import _plot_residual_norms + + +class SaddlePointSolver: + r"""Solves for :math:`(x, y)` in the saddle point problem + + .. math:: + + \left( \matrix{ + A & B^{\top} \cr + B & 0 + } \right) + \left( \matrix{ + x \cr y + } \right) + = + \left( \matrix{ + f \cr 0 + } \right) + + using either the Uzawa iteration :math:`BA^{-1}B^{\top} y = BA^{-1} f` or using on of the solvers given in :mod:`psydac.linalg.solvers`. The prefered solver is GMRES. + The decission which variant to use is given by the type of A. If A is of type list of xp.ndarrays or sc.sparse.csr_matrices, then this class uses the Uzawa algorithm. + If A is of type LinearOperator or BlockLinearOperator, a solver is used for the inverse. + Using the Uzawa algorithm, solution is given by: + + .. math:: + + y = \left[ B A^{-1} B^{\top}\right]^{-1} B A^{-1} f \,, \qquad + x = A^{-1} \left[ f - B^{\top} y \right] \,. + + Parameters + ---------- + A : list, LinearOperator or BlockLinearOperator + Upper left block. + Either the entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. + Alternative: Give whole matrice A as LinearOperator or BlockLinearOperator. + list: Uzawa algorithm is used. + LinearOperator: A solver given in :mod:`psydac.linalg.solvers` is used. Specified by solver_name. + BlockLinearOperator: A solver given in :mod:`psydac.linalg.solvers` is used. Specified by solver_name. + + B : list, LinearOperator or BlockLinearOperator + Lower left block. + Uzwaw Algorithm: All entries of block B are given either as list of xp.ndarray or sc.sparse.csr_matrix. + Solver: Give whole B as LinearOperator or BlocklinearOperator + + F : list + Right hand side of the upper block. + Uzawa: Given as list of xp.ndarray or sc.sparse.csr_matrix. + Solver: Given as LinearOperator or BlockLinearOperator + + Apre : list + The non-inverted preconditioner for entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. Only required for the Uzawa algorithm. + + method_to_solve : str + Method for the inverses. Choose from 'DirectNPInverse', 'ScipySparse', 'InexactNPInverse' ,'SparseSolver'. Only required for the Uzawa algorithm. + + preconditioner : bool + Wheter to use preconditioners given in Apre or not. Only required for the Uzawa algorithm. + + spectralanalysis : bool + Do the spectralanalyis for the matrices in A and if preconditioner given, compare them to the preconditioned matrices. Only possible if A is given as list. + + dimension : str + Which of the predefined manufactured solutions to use ('1D', '2D' or 'Restelli') + + tol : float + Convergence tolerance for the potential residual. + + max_iter : int + Maximum number of iterations allowed. + """ + + def __init__( + self, + A: Union[list, LinearOperator, BlockLinearOperator], + B: Union[list, LinearOperator, BlockLinearOperator], + F: Union[list, Vector, BlockVector], + Apre: list = None, + method_to_solve: str = "DirectNPInverse", + preconditioner: bool = False, + spectralanalysis: bool = False, + dimension: str = "2D", + solver_name: str = "GMRES", + tol: float = 1e-8, + max_iter: int = 1000, + **solver_params, + ): + assert type(A) is type(B) + if isinstance(A, list): + self._variant = "Uzawa" + for i in A: + assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) + for i in B: + assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) + for i in F: + assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) + for i in Apre: + assert ( + isinstance(i, xp.ndarray) + or isinstance(i, sc.sparse.csr_matrix) + or isinstance(i, sc.sparse.csr_array) + ) + assert method_to_solve in ("SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse") + assert A[0].shape[0] == B[0].shape[1] + assert A[0].shape[1] == B[0].shape[1] + assert A[1].shape[0] == B[1].shape[1] + assert A[1].shape[1] == B[1].shape[1] + + self._method_to_solve = ( + method_to_solve # 'SparseSolver', 'ScipySparse', 'InexactNPInverse', 'DirectNPInverse' + ) + self._preconditioner = preconditioner + + elif isinstance(A, LinearOperator) or isinstance(A, BlockLinearOperator): + self._variant = "Inverse_Solver" + assert A.domain == B.domain + assert A.codomain == B.domain + self._solver_name = solver_name + if solver_params["pc"] is None: + solver_params.pop("pc") + + # operators + self._A = A + self._Apre = Apre + self._B = B + self._F = F + self._tol = tol + self._max_iter = max_iter + self._spectralanalysis = spectralanalysis + self._dimension = dimension + self._verbose = solver_params["verbose"] + + if self._variant == "Inverse_Solver": + self._block_domainM = BlockVectorSpace(self._A.domain, self._B.transpose().domain) + self._block_codomainM = self._block_domainM + self._blocks = [[self._A, self._B.T], [self._B, None]] + _Minit = BlockLinearOperator(self._block_domainM, self._block_codomainM, blocks=self._blocks) + self._solverMinv = inverse(_Minit, solver_name, tol=tol, maxiter=max_iter, **solver_params) + + # Solution vectors + self._P = B.codomain.zeros() + self._U = A.codomain.zeros() + self._Utmp = F.copy() * 0 + # Allocate memory for call + self._rhstemp = BlockVector(self._block_domainM, blocks=[A.codomain.zeros(), self._B.codomain.zeros()]) + + elif self._variant == "Uzawa": + if self._method_to_solve in ("InexactNPInverse", "SparseSolver"): + self._preconditioner = False + + self._Anp = self._A[0] + self._Aenp = self._A[1] + self._B1np = self._B[0] + self._B2np = self._B[1] + + # Instanciate inverses + self._setup_inverses() + + # Solution vectors numpy + self._Pnp = xp.zeros(self._B1np.shape[0]) + self._Unp = xp.zeros(self._A[0].shape[1]) + self._Uenp = xp.zeros(self._A[1].shape[1]) + # Allocate memory for matrices used in solving the system + self._rhs0np = self._F[0].copy() + self._rhs1np = self._F[1].copy() + + # List to store residual norms + self._residual_norms = [] + self._stepsize = 0.0 + + @property + def A(self): + """Upper left block.""" + return self._A + + @A.setter + def A(self, a): + if self._variant == "Uzawa": + need_update = True + A0_old, A1_old = self._A + A0_new, A1_new = a + if self._method_to_solve in ("ScipySparse", "SparseSolver"): + same_A0 = (A0_old != A0_new).nnz == 0 + same_A1 = (A1_old != A1_new).nnz == 0 + else: + same_A0 = xp.allclose(A0_old, A0_new, atol=1e-10) + same_A1 = xp.allclose(A1_old, A1_new, atol=1e-10) + if same_A0 and same_A1: + need_update = False + self._A = a + self._Anp = self._A[0] + self._Aenp = self._A[1] + if need_update: + self._setup_inverses() + elif self._variant == "Inverse_Solver": + self._A = a + + @property + def B(self): + """Lower left block.""" + return self._B + + @B.setter + def B(self, b): + self._B = b + + @property + def F(self): + """Right hand side vector.""" + return self._F + + @F.setter + def F(self, f): + self._F = f + + @property + def Apre(self): + """Preconditioner for upper left block A.""" + return self._Apre + + @Apre.setter + def Apre(self, a): + if self._variant == "Uzawa": + need_update = True + A0_old, A1_old = self._Apre + A0_new, A1_new = a + if self._method_to_solve in ("ScipySparse", "SparseSolver"): + same_A0 = (A0_old != A0_new).nnz == 0 + same_A1 = (A1_old != A1_new).nnz == 0 + else: + same_A0 = xp.allclose(A0_old, A0_new, atol=1e-10) + same_A1 = xp.allclose(A1_old, A1_new, atol=1e-10) + if same_A0 and same_A1: + need_update = False + self._Apre = a + if need_update: + self._setup_inverses() + elif self._variant == "Inverse_Solver": + self._Apre = a + + def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): # todo should have options to use other than uzawa. should solve in full generality. A being block diagonal is a special case of uzawa + """ + Solves the saddle-point problem using the Uzawa algorithm. + + Parameters + ---------- + U_init : Vector, xp.ndarray or sc.sparse.csr.csr_matrix, optional + Initial guess for the velocity of the ions. If None, initializes to zero. Types xp.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. + + Ue_init : Vector, xp.ndarray or sc.sparse.csr.csr_matrix, optional + Initial guess for the velocity of the electrons. If None, initializes to zero. Types xp.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. + + P_init : Vector, optional + Initial guess for the potential. If None, initializes to zero. + + Returns + ------- + U : Vector + Solution vector for the velocity. + + P : Vector + Solution vector for the potential. + + info : dict + Convergence information. + """ + + # TODO this contains two different strategies! favágás and actual uzawa + if self._variant == "Inverse_Solver": # todo not in the """""saddle point solver""""""" + self._P1 = P_init if P_init is not None else self._P + self._U1 = U_init if U_init is not None else self._Utmp[0] + self._U2 = Ue_init if Ue_init is not None else self._Utmp[1] + + _blocksM = [[self._A, self._B.T], [self._B, None]] + _M = BlockLinearOperator(self._block_domainM, self._block_codomainM, blocks=_blocksM) + _RHS = BlockVector(self._block_domainM, blocks=[self._F, self._B.codomain.zeros()]) + + self._blockU = BlockVector(self._A.domain, blocks=[self._U1, self._U2]) + self._solblocks = [self._blockU, self._P1] + # comment out the next two lines if working with lifting and GMRES + x0 = BlockVector(self._block_domainM, blocks=self._solblocks) + self._solverMinv._options["x0"] = x0 + + # use setter to update lhs matrix + self._solverMinv.linop = _M + + # Initialize P to zero or given initial guess + self._sol = self._solverMinv.dot(_RHS, out=self._rhstemp) + self._U = self._sol[0] + self._P = self._sol[1] + + return self._U, self._P, self._solverMinv._info + + elif self._variant == "Uzawa": + info = {} + + if self._spectralanalysis: + self._spectralresult = self._spectral_analysis() + else: + self._spectralresult = [] + + # Initialize P to zero or given initial guess + if isinstance(U_init, xp.ndarray) or isinstance(U_init, sc.sparse.csr.csr_matrix): + self._Pnp = P_init if P_init is not None else self._P + self._Unp = U_init if U_init is not None else self._U + self._Uenp = Ue_init if U_init is not None else self._Ue + + else: + self._Pnp = P_init.toarray() if P_init is not None else self._Pnp + self._Unp = U_init.toarray() if U_init is not None else self._Unp + self._Uenp = Ue_init.toarray() if U_init is not None else self._Uenp + + if self._verbose: + print("Uzawa solver:") + print("+---------+---------------------+") + print("+ Iter. # | L2-norm of residual |") + print("+---------+---------------------+") + template = "| {:7d} | {:19.2e} |" + + for iteration in range(self._max_iter): + # Step 1: Compute velocity U by solving A U = -Bᵀ P + F -A Un + self._rhs0np *= 0 + self._rhs0np -= self._B1np.transpose().dot(self._Pnp) + self._rhs0np -= self._Anp.dot(self._Unp) + self._rhs0np += self._F[0] + if not self._preconditioner: + self._Unp += self._Anpinv.dot(self._rhs0np) + elif self._preconditioner: + self._Unp += self._Anpinv.dot(self._A11npinv @ self._rhs0np) + + R1 = self._B1np.dot(self._Unp) + + self._rhs1np *= 0 + self._rhs1np -= self._B2np.transpose().dot(self._Pnp) + self._rhs1np -= self._Aenp.dot(self._Uenp) + self._rhs1np += self._F[1] + if not self._preconditioner: + self._Uenp += self._Aenpinv.dot(self._rhs1np) + elif self._preconditioner: + self._Uenp += self._Aenpinv.dot(self._A22npinv @ self._rhs1np) + + R2 = self._B2np.dot(self._Uenp) + + # Step 2: Compute residual R = BU (divergence of U) + R = R1 + R2 # self._B1np.dot(self._Unp) + self._B2np.dot(self._Uenp) + residual_norm = xp.linalg.norm(R) + residual_normR1 = xp.linalg.norm(R) + self._residual_norms.append(residual_normR1) # Store residual norm + # Check for convergence based on residual norm + if residual_norm < self._tol: + if self._verbose: + print(template.format(iteration + 1, residual_norm)) + print("+---------+---------------------+") + info["success"] = True + info["niter"] = iteration + 1 + if self._verbose: + _plot_residual_norms(self._residual_norms) + return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult + + # Steepest gradient + alpha = (R.dot(R)) / (R.dot(self._Precnp.dot(R))) + # Minimal residual + # alpha = ((self._Precnp.dot(R)).dot(R)) / ((self._Precnp.dot(R)).dot(self._Precnp.dot(R))) + self._Pnp += alpha.real * R.real + + if self._verbose: + print(template.format(iteration + 1, residual_norm)) + + if self._verbose: + print("+---------+---------------------+") + + # Return with info if maximum iterations reached + info["success"] = False + info["niter"] = iteration + 1 + if self._verbose: + _plot_residual_norms(self._residual_norms) + return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult + + def _setup_inverses(self): + A0 = self._A[0] + A1 = self._A[1] + + # === Preconditioner inverses, if used + if self._preconditioner: + A11_pre = self._Apre[0] + A22_pre = self._Apre[1] + + if hasattr(self, "_A11npinv") and self._is_inverse_still_valid(self._A11npinv, A11_pre, "A11 pre"): + pass + else: + self._A11npinv = self._compute_inverse(A11_pre, which="A11 pre") + + if hasattr(self, "_A22npinv") and self._is_inverse_still_valid(self._A22npinv, A22_pre, "A22 pre"): + pass + else: + self._A22npinv = self._compute_inverse(A22_pre, which="A22 pre") + + # === Inverse for A[0] if preconditioned + if hasattr(self, "_Anpinv") and self._is_inverse_still_valid(self._Anpinv, A0, "A[0]", pre=self._A11npinv): + pass + else: + self._Anpinv = self._compute_inverse(self._A11npinv @ A0, which="A[0]") + + # === Inverse for A[1] + if hasattr(self, "_Aenpinv") and self._is_inverse_still_valid( + self._Aenpinv, + A1, + "A[1]", + pre=self._A22npinv, + ): + pass + else: + self._Aenpinv = self._compute_inverse(self._A22npinv @ A1, which="A[1]") + + else: # No preconditioning: + # === Inverse for A[0] + if hasattr(self, "_Anpinv") and self._is_inverse_still_valid(self._Anpinv, A0, "A[0]"): + pass + else: + self._Anpinv = self._compute_inverse(A0, which="A[0]") + + # === Inverse for A[1] + if hasattr(self, "_Aenpinv") and self._is_inverse_still_valid(self._Aenpinv, A1, "A[1]"): + pass + else: + self._Aenpinv = self._compute_inverse(A1, which="A[1]") + + # Precompute Schur complement + self._Precnp = self._B1np @ self._Anpinv @ self._B1np.T + self._B2np @ self._Aenpinv @ self._B2np.T + + def _is_inverse_still_valid(self, inv, mat, name="", pre=None): + # try: + if pre is not None: + test_mat = pre @ mat + else: + test_mat = mat + I_approx = inv @ test_mat + + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + I_exact = xp.eye(test_mat.shape[0]) + if not xp.allclose(I_approx, I_exact, atol=1e-6): + diff = I_approx - I_exact + max_abs = xp.abs(diff).max() + print(f"{name} inverse is NOT valid anymore. Max diff: {max_abs:.2e}") + return False + print(f"{name} inverse is still valid.") + return True + elif self._method_to_solve == "ScipySparse": + I_exact = sc.sparse.identity(I_approx.shape[0], format=I_approx.format) + diff = (I_approx - I_exact).tocoo() + max_abs = xp.abs(diff.data).max() if diff.nnz > 0 else 0.0 + + if max_abs > 1e-6: + print(f"{name} inverse is NOT valid anymore.") + print(f"Max absolute difference: {max_abs:.2e}") + print(f"Number of differing entries: {diff.nnz}") + return False + print(f"{name} inverse is still valid.") + return True + + def _compute_inverse(self, mat, which="matrix"): + print(f"Computing inverse for {which} using method {self._method_to_solve}") + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + return xp.linalg.inv(mat) + elif self._method_to_solve == "ScipySparse": + return sc.sparse.linalg.inv(mat) + elif self._method_to_solve == "SparseSolver": + solver = SparseSolver(mat) + return solver.solve(xp.eye(mat.shape[0])) + else: + raise ValueError(f"Unknown solver method {self._method_to_solve}") + + def _spectral_analysis(self): + # Spectral analysis + # A11 before + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + eigvalsA11_before, eigvecs_before = xp.linalg.eig(self._A[0]) + condA11_before = xp.linalg.cond(self._A[0]) + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + eigvalsA11_before, eigvecs_before = xp.linalg.eig(self._A[0].toarray()) + condA11_before = xp.linalg.cond(self._A[0].toarray()) + maxbeforeA11 = max(eigvalsA11_before) + maxbeforeA11_abs = xp.max(xp.abs(eigvalsA11_before)) + minbeforeA11_abs = xp.min(xp.abs(eigvalsA11_before)) + minbeforeA11 = min(eigvalsA11_before) + specA11_bef = maxbeforeA11 / minbeforeA11 + specA11_bef_abs = maxbeforeA11_abs / minbeforeA11_abs + # print(f'{maxbeforeA11 = }') + # print(f'{maxbeforeA11_abs = }') + # print(f'{minbeforeA11_abs = }') + # print(f'{minbeforeA11 = }') + # print(f'{specA11_bef = }') + print(f"{specA11_bef_abs =}") + + # A22 before + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + eigvalsA22_before, eigvecs_before = xp.linalg.eig(self._A[1]) + condA22_before = xp.linalg.cond(self._A[1]) + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + eigvalsA22_before, eigvecs_before = xp.linalg.eig(self._A[1].toarray()) + condA22_before = xp.linalg.cond(self._A[1].toarray()) + maxbeforeA22 = max(eigvalsA22_before) + maxbeforeA22_abs = xp.max(xp.abs(eigvalsA22_before)) + minbeforeA22_abs = xp.min(xp.abs(eigvalsA22_before)) + minbeforeA22 = min(eigvalsA22_before) + specA22_bef = maxbeforeA22 / minbeforeA22 + specA22_bef_abs = maxbeforeA22_abs / minbeforeA22_abs + # print(f'{maxbeforeA22 = }') + # print(f'{maxbeforeA22_abs = }') + # print(f'{minbeforeA22_abs = }') + # print(f'{minbeforeA22 = }') + # print(f'{specA22_bef = }') + print(f"{specA22_bef_abs =}") + print(f"{condA22_before =}") + + if self._preconditioner: + # A11 after preconditioning with its inverse + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig(self._A11npinv @ self._A[0]) # Implement this + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig((self._A11npinv @ self._A[0]).toarray()) + maxafterA11_prec = max(eigvalsA11_after_prec) + minafterA11_prec = min(eigvalsA11_after_prec) + maxafterA11_abs_prec = xp.max(xp.abs(eigvalsA11_after_prec)) + minafterA11_abs_prec = xp.min(xp.abs(eigvalsA11_after_prec)) + specA11_aft_prec = maxafterA11_prec / minafterA11_prec + specA11_aft_abs_prec = maxafterA11_abs_prec / minafterA11_abs_prec + # print(f'{maxafterA11_prec = }') + # print(f'{maxafterA11_abs_prec = }') + # print(f'{minafterA11_abs_prec = }') + # print(f'{minafterA11_prec = }') + # print(f'{specA11_aft_prec = }') + print(f"{specA11_aft_abs_prec =}") + + # A22 after preconditioning with its inverse + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + eigvalsA22_after_prec, eigvecs_after = xp.linalg.eig(self._A22npinv @ self._A[1]) # Implement this + condA22_after = xp.linalg.cond(self._A22npinv @ self._A[1]) + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + eigvalsA22_after_prec, eigvecs_after = xp.linalg.eig((self._A22npinv @ self._A[1]).toarray()) + condA22_after = xp.linalg.cond((self._A22npinv @ self._A[1]).toarray()) + maxafterA22_prec = max(eigvalsA22_after_prec) + minafterA22_prec = min(eigvalsA22_after_prec) + maxafterA22_abs_prec = xp.max(xp.abs(eigvalsA22_after_prec)) + minafterA22_abs_prec = xp.min(xp.abs(eigvalsA22_after_prec)) + specA22_aft_prec = maxafterA22_prec / minafterA22_prec + specA22_aft_abs_prec = maxafterA22_abs_prec / minafterA22_abs_prec + # print(f'{maxafterA22_prec = }') + # print(f'{maxafterA22_abs_prec = }') + # print(f'{minafterA22_abs_prec = }') + # print(f'{minafterA22_prec = }') + # print(f'{specA22_aft_prec = }') + print(f"{specA22_aft_abs_prec =}") + + return condA22_before, specA22_bef_abs, condA11_before, condA22_after, specA22_aft_abs_prec + + else: + return condA22_before, specA22_bef_abs, condA11_before diff --git a/src/struphy/models/rework_model.py b/src/struphy/models/rework_model.py new file mode 100644 index 000000000..f38629fcc --- /dev/null +++ b/src/struphy/models/rework_model.py @@ -0,0 +1,124 @@ +import cunumpy as xp +from psydac.ddm.mpi import mpi as MPI + +from struphy.feec.projectors import L2Projector +from struphy.feec.variational_utilities import InternalEnergyEvaluator +from struphy.models.base import StruphyModel +from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable +from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers +from struphy.propagators import rework_propagator + + +rank = MPI.COMM_WORLD.Get_rank() + +class TwoFluidQuasiNeutralToy(StruphyModel): + r"""Linearized, quasi-neutral two-fluid model with zero electron inertia. + + :ref:`normalization`: + + .. math:: + + \hat u = \hat v_\textnormal{th}\,,\qquad e\hat \phi = m \hat v_\textnormal{th}^2\,. + + :ref:`Equations `: + + .. math:: + + \frac{\partial \mathbf u}{\partial t} &= - \nabla \phi + \frac{\mathbf u \times \mathbf B_0}{\varepsilon} + \nu \Delta \mathbf u + \mathbf f\,, + \\[2mm] + 0 &= \nabla \phi - \frac{\mathbf u_e \times \mathbf B_0}{\varepsilon} + \nu_e \Delta \mathbf u_e + \mathbf f_e \,, + \\[3mm] + \nabla & \cdot (\mathbf u - \mathbf u_e) = 0\,, + + where :math:`\mathbf B_0` is a static magnetic field and :math:`\mathbf f, \mathbf f_e` are given forcing terms, + and with the normalization parameter + + .. math:: + + \varepsilon = \frac{1}{\hat \Omega_\textnormal{c} \hat t} \,,\qquad \textnormal{with} \,,\qquad \hat \Omega_{\textnormal{c}} = \frac{(Ze) \hat B}{(A m_\textnormal{H})}\,, + + :ref:`propagators` (called in sequence): + + 1. :class:`~struphy.propagators.propagators_fields.TwoFluidQuasiNeutralFull` + + :ref:`Model info `: + + References + ---------- + [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising + in plasma physics, Journal of Computational Physics 2018. + """ + + ## species + + class EMfields(FieldSpecies): + def __init__(self): + self.phi = FEECVariable(space="L2") + self.init_variables() + + class Ions(FluidSpecies): + def __init__(self): + self.u = FEECVariable(space="Hdiv") + self.init_variables() + + class Electrons(FluidSpecies): + def __init__(self): + self.u = FEECVariable(space="Hdiv") + self.init_variables() + + ## propagators + + class Propagators: + def __init__(self): + self.qn_full = rework_propagator.TwoFluidQuasiNeutralFull() + + ## abstract methods + + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMfields() + self.ions = self.Ions() + self.electrons = self.Electrons() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.qn_full.variables.u = self.ions.u + self.propagators.qn_full.variables.ue = self.electrons.u + self.propagators.qn_full.variables.phi = self.em_fields.phi + + # define scalars for update_scalar_quantities + + @property + def bulk_species(self): + return self.ions + + @property + def velocity_scale(self): + return "thermal" + + def allocate_helpers(self): + pass + + def update_scalar_quantities(self): + pass + + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "BaseUnits()" in line: + new_file += ["base_units = BaseUnits(kBT=1.0)\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) diff --git a/src/struphy/propagators/rework_propagator.py b/src/struphy/propagators/rework_propagator.py new file mode 100644 index 000000000..bf6afcadf --- /dev/null +++ b/src/struphy/propagators/rework_propagator.py @@ -0,0 +1,479 @@ + +import copy +from copy import deepcopy +from dataclasses import dataclass +from typing import Callable, Literal, get_args, cast +from warnings import warn + +import cunumpy as xp +import scipy as sc +from line_profiler import profile +from matplotlib import pyplot as plt +from numpy import zeros +from psydac.api.essential_bc import apply_essential_bc_stencil +from psydac.ddm.mpi import mpi as MPI +from psydac.linalg.basic import ComposedLinearOperator, IdentityOperator, ZeroOperator, InverseLinearOperator +from psydac.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace +from psydac.linalg.solvers import inverse +from psydac.linalg.stencil import StencilVector + +import struphy.feec.utilities as util +from struphy.examples.restelli2018 import callables +from struphy.feec import preconditioner +from struphy.feec.basis_projection_ops import ( + BasisProjectionOperator, BasisProjectionOperatorLocal, + BasisProjectionOperators, CoordinateProjector, +) +from struphy.feec.linear_operators import BoundaryOperator +from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators +from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner, MassMatrixPreconditioner +from struphy.feec.projectors import L2Projector +from struphy.feec.psydac_derham import Derham, SplineFunction +from struphy.feec.variational_utilities import ( + BracketOperator, Hdiv0_transport_operator, InternalEnergyEvaluator, + KineticEnergyEvaluator, Pressure_transport_operator, +) +from struphy.fields_background.equils import set_defaults +from struphy.geometry.utilities import TransformedPformComponent +from struphy.initial import perturbations +from struphy.io.options import ( + OptsDirectSolver, OptsGenSolver, OptsMassPrecond, OptsNonlinearSolver, + OptsSaddlePointSolver, OptsSymmSolver, OptsVecSpace, check_option, +) +from struphy.io.setup import descend_options_dict +from struphy.kinetic_background.base import Maxwellian +from struphy.kinetic_background.maxwellians import GyroMaxwellian2D, Maxwellian3D +from struphy.linear_algebra.saddle_point import SaddlePointSolver +from struphy.linear_algebra.schur_solver import SchurSolver, SchurSolverFull +from struphy.linear_algebra.solver import NonlinearSolverParameters, SolverParameters +from struphy.models.species import Species +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable +from struphy.ode.solvers import ODEsolverFEEC +from struphy.ode.utils import ButcherTableau, OptsButcher +from struphy.pic.accumulation import accum_kernels, accum_kernels_gc +from struphy.pic.accumulation.filter import FilterParameters +from struphy.pic.accumulation.particles_to_grid import Accumulator, AccumulatorVector +from struphy.pic.base import Particles +from struphy.pic.particles import Particles5D, Particles6D +from struphy.polar.basic import PolarVector +from struphy.propagators.base import Propagator +from struphy.utils.pyccel import Pyccelkernel + + +class TwoFluidQuasiNeutralFull(Propagator): + r""":ref:`FEEC ` discretization of the following equations: + find :math:`\mathbf u \in H(\textnormal{div})`, :math:`\mathbf u_e \in H(\textnormal{div})` and :math:`\mathbf \phi \in L^2` such that + + .. math:: + + \int_{\Omega} \partial_t \mathbf{u}\cdot \mathbf{v} \, \textrm d\mathbf{x} &= \int_{\Omega} \phi \nabla \! \cdot \! \mathbf{v} \, \textrm d\mathbf{x} + \int_{\Omega} \mathbf{u}\! \times \! \mathbf{B}_0 \cdot \mathbf{v} \, \textrm d\mathbf{x} + \nu \int_{\Omega} \nabla \mathbf{u}\! : \! \nabla \mathbf{v} \, \textrm d\mathbf{x} + \int_{\Omega} f \mathbf{v} \, \textrm d\mathbf{x} \qquad \forall \, \mathbf{v} \in H(\textrm{div}) \,. + \\[2mm] + 0 &= - \int_{\Omega} \phi \nabla \! \cdot \! \mathbf{v_e} \, \textrm d\mathbf{x} - \int_{\Omega} \mathbf{u_e} \! \times \! \mathbf{B}_0 \cdot \mathbf{v_e} \, \textrm d\mathbf{x} + \nu_e \int_{\Omega} \nabla \mathbf{u_e} \!: \! \nabla \mathbf{v_e} \, \textrm d\mathbf{x} + \int_{\Omega} f_e \mathbf{v_e} \, \textrm d\mathbf{x} \qquad \forall \ \mathbf{v_e} \in H(\textrm{div}) \,. + \\[2mm] + 0 &= \int_{\Omega} \psi \nabla \cdot (\mathbf{u}-\mathbf{u_e}) \, \textrm d\mathbf{x} \qquad \forall \, \psi \in L^2 \,. + + :ref:`time_discret`: fully implicit. + """ + + # ========================================================================= + ### State variables (ion velocity u, electron velocity ue, pressure phi) + # ========================================================================= + + class Variables(): + def __init__(self) -> None: + self._u: FEECVariable | None = None + self._ue: FEECVariable | None = None + self._phi: FEECVariable | None = None + + @property + def u(self) -> FEECVariable | None: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._u = new + + @property + def ue(self) -> FEECVariable | None: + return self._ue + + @ue.setter + def ue(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._ue = new + + @property + def phi(self) -> FEECVariable | None: + return self._phi + + @phi.setter + def phi(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._phi = new + + def __init__(self): + self.variables = self.Variables() + + # ========================================================================= + ### Options + # ========================================================================= + + @dataclass + class Options(): + + nu: float | None = None + nu_e: float | None = None + eps_norm: float | None = None + + # boundary conditions per species + # supported kinds: "periodic", "dirichlet" + # future: "neumann", "robin" + boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_data_u: dict[tuple[int, int], Callable] | None = None + boundary_data_ue: dict[tuple[int, int], Callable] | None = None + + source_u: Callable | None = None + source_ue: Callable | None = None + + stab_sigma: float | None = None + + solver: OptsGenSolver = "gmres" + solver_params: SolverParameters | None = None + + def __post_init__(self): + + # --- required parameters --- + assert self.nu is not None, "nu must be specified" + assert self.nu_e is not None, "nu_e must be specified" + assert self.eps_norm is not None, "eps_norm must be specified" + assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" + assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" + + # --- physical parameter sanity checks --- + if self.nu < 0: + raise ValueError(f"nu must be non-negative, got {self.nu}") + if self.nu_e < 0: + raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") + if self.eps_norm <= 0: + raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") + + # --- check all axes are covered --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for d in range(3): + for side in (-1, 1): + assert (d, side) in bcs, \ + f"{name} is missing entry for axis {d} side {side}" + + # --- periodic consistency: periodic must be paired on both sides --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for (d, side), kind in bcs.items(): + if kind == "periodic": + assert bcs.get((d, -side)) == "periodic", \ + f"{name}: axis {d} side {side} is periodic but opposite side is not" + + # --- ions and electrons must agree on which axes are periodic --- + for d in range(3): + u_left = self.boundary_conditions_u.get((d, -1)) + ue_left = self.boundary_conditions_ue.get((d, -1)) + u_right = self.boundary_conditions_u.get((d, 1)) + ue_right = self.boundary_conditions_ue.get((d, 1)) + u_periodic = (u_left == "periodic") + ue_periodic = (ue_left == "periodic") + if u_periodic != ue_periodic: + raise ValueError( + f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " + f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" + ) + + # --- warn for Dirichlet faces with no boundary data --- + for species, bcs, data, label in [ + ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), + ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), + ]: + has_dirichlet = any(v == "dirichlet" for v in bcs.values()) + if has_dirichlet: + if data is None: + warn(f"Dirichlet BCs specified for {species} but no {label} given " + f"— defaulting to homogeneous Dirichlet on all faces.") + else: + for (d, side), kind in bcs.items(): + if kind == "dirichlet" and (d, side) not in data: + warn(f"No {label} given for axis {d} side {side} " + f"— defaulting to homogeneous Dirichlet.") + + # --- warn if no source terms --- + if self.source_u is None: + warn("No source_u specified — defaulting to zero.") + if self.source_ue is None: + warn("No source_ue specified — defaulting to zero.") + + # --- defaults --- + if self.stab_sigma is None: + warn("stab_sigma not specified, defaulting to 0.0") + self.stab_sigma = 0.0 + + check_option(self.solver, OptsGenSolver) + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + assert hasattr(self, "_options"), "Options not set." + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + # ========================================================================= + ### Boundary condition helpers + # ========================================================================= + + def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): + + dirichlet_bc = [] + for d in range(3): + if spl_kind[d]: # periodic spline — no clamping ever + dirichlet_bc.append((False, False)) + else: + left = boundary_conditions.get((d, -1)) == "dirichlet" + right = boundary_conditions.get((d, 1)) == "dirichlet" + dirichlet_bc.append((left, right)) + return tuple(tuple(bc) for bc in dirichlet_bc) + + def _apply_boundary_conditions(self, vec, boundary_conditions): + """Zero out Dirichlet DOFs on the given stencil vector.""" + for (d, side), kind in boundary_conditions.items(): + if kind == "dirichlet": + apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) + # future: neumann and robin require no zeroing here + + # ========================================================================= + ### Allocate + # ========================================================================= + + def allocate(self): + + self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 + self._dt = None + + # ---- constrained (v0) de Rham complex -------------------------------- + + _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) + _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) + _dirichlet_bc = tuple( + (l_u or l_ue, r_u or r_ue) + for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) + ) + + self._derham_v0 = Derham( + self.derham.Nel, self.derham.p, self.derham.spl_kind, + domain=self.domain, dirichlet_bc=_dirichlet_bc, + ) + self._mass_ops_v0 = WeightedMassOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.mass_ops.weights["eq_mhd"], + ) + self._basis_ops_v0 = BasisProjectionOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.basis_ops.weights["eq_mhd"], + ) + + # ---- unconstrained operators (for RHS assembly) ---------------------- + + self._M2 = self.mass_ops.M2 + self._M2B = - self.mass_ops.M2B + self._div = self.derham.div + self._curl = self.derham.curl + self._S21 = self.basis_ops.S21 + + self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div + + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + + self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl + self._A22 = (- self.options.stab_sigma * IdentityOperator(self._A11.domain) + + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) + + # ---- constrained operators (for system matrix) ----------------------- + + self._M2_v0 = self._mass_ops_v0.M2 + self._M3_v0 = self._mass_ops_v0.M3 + self._M2B_v0 = - self._mass_ops_v0.M2B + self._div_v0 = self._derham_v0.div + self._curl_v0 = self._derham_v0.curl + self._S21_v0 = self._basis_ops_v0.S21 + + self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 + + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) + + self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 + self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._A11_v0.domain) + + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) + + # ---- block saddle-point system ---------------------------------------- + + self._block_domain_v0 = BlockVectorSpace(self._A11_v0.domain, self._A22_v0.domain) + self._block_codomain_v0 = self._block_domain_v0 + self._block_codomain_B_v0 = self._M3_v0.codomain + + self._B1_v0 = - self._M3_v0 @ self._div_v0 + self._B2_v0 = self._M3_v0 @ self._div_v0 + + self._B_v0 = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_B_v0, + blocks=[[self._B1_v0, self._B2_v0]] + ) + + self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) + + _A_init = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0, None], [None, self._A22_v0]] + ) + _M_init = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] + ) + self._Minv = cast(InverseLinearOperator, inverse( + _M_init, self.options.solver, + x0=None, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + )) + + # ---- projector ------------------------------------------------------- + + self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + + # ---- solution spline functions (unconstrained) ----------------------- + + self._u = self.derham.create_spline_function("u", space_id="Hdiv") + self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") + self._phi = self.derham.create_spline_function("phi", space_id="L2") + + # ---- BC lifts (unconstrained) ---------------------------------------- + + self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") + self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") + + for u_prime, boundary_data, boundary_conditions in [ + (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), + (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), + ]: + if boundary_data is None: + continue + for (d, side), f_bc in boundary_data.items(): + if boundary_conditions.get((d, side)) == "dirichlet": + bc_pulled = lambda *etas, f=f_bc: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], + lambda *etas: bc_pulled(*etas)[1], + lambda *etas: bc_pulled(*etas)[2]]) + for (d2, side2), kind2 in boundary_conditions.items(): + if kind2 == "dirichlet" and (d2, side2) != (d, side): + apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) + u_prime.vector += _vec + + self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") + self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") + + self._u_prime_v0.vector = self._u_prime.vector + self._ue_prime_v0.vector = self._ue_prime.vector + + # ---- projected source terms (unconstrained) -------------------------- + + self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") + self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") + + for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: + if source is not None: + src_pulled = lambda *etas, f=source: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], + lambda *etas: src_pulled(*etas)[1], + lambda *etas: src_pulled(*etas)[2]]) + + # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- + + self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") + + # ========================================================================= + ### Time step + # ========================================================================= + + def __call__(self, dt): + + # --- copy current state --- + self._u.vector = self.variables.u.spline.vector + self._ue.vector = self.variables.ue.spline.vector + + # --- rebuild system matrix if dt changed --- + if dt != self._dt: + self._dt = dt + _A = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] + ) + _M = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A, self._B_v0.T], [self._B_v0, None]] + ) + self._Minv.linop = _M + + # --- assemble RHS in unconstrained space, then zero boundary DOFs --- + # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' + # electron: F2 = rhs_ue - A22 * ue' + self._rhs_vec_u.vector = (self._rhs_u.vector + + self._M2.dot(self._u.vector) / dt + - self._A11.dot(self._u_prime.vector) + - self._M2.dot(self._u_prime.vector) / dt) + self._rhs_vec_ue.vector = (self._rhs_ue.vector + - self._A22.dot(self._ue_prime.vector)) + + self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) + self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) + + # --- build block RHS and solve --- + _F = BlockVector(self._block_domain_v0, + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) + _RHS = BlockVector(self._block_domain_M, + blocks=[_F, self._block_codomain_B_v0.zeros()]) + + _sol = self._Minv.dot(_RHS) + info = self._Minv.get_info() + + # --- reconstruct full solution: u = u_0 + u' --- + self._u.vector = _sol[0][0] + self._u_prime_v0.vector + self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector + self._phi.vector = _sol[1] + + # --- update FEEC variables --- + max_diffs = self.update_feec_variables( + u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector + ) + + if self.options.solver_params.info and self._rank == 0: + print(f"Status: {info['success']}, Iterations: {info['niter']}") + print(f"Max diffs: {max_diffs}") \ No newline at end of file From d7a346ebda5993e8ff1a907c54b200e6fc7f5004 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Thu, 5 Mar 2026 21:39:53 +0000 Subject: [PATCH 04/41] Rebased onto version 3.0.4 --- .gitignore | 3 - 1D_Verification.py | 228 --------------------------------------------- split_models.py | 78 ---------------- 3 files changed, 309 deletions(-) delete mode 100644 1D_Verification.py delete mode 100644 split_models.py diff --git a/.gitignore b/.gitignore index 8d4a130eb..afefb317e 100644 --- a/.gitignore +++ b/.gitignore @@ -102,6 +102,3 @@ share/ lib64 pyvenv.cfg - -2D_Verification.py -Restelli_Verification.py \ No newline at end of file diff --git a/1D_Verification.py b/1D_Verification.py deleted file mode 100644 index baa40a32d..000000000 --- a/1D_Verification.py +++ /dev/null @@ -1,228 +0,0 @@ -from cunumpy import pi, cos, sin, zeros_like, ones_like -from struphy.io.options import EnvironmentOptions, BaseUnits, Time -from struphy.geometry import domains -from struphy.fields_background import equils -from struphy.topology import grids -from struphy.io.options import DerhamOptions -from struphy.initial import perturbations -from struphy import main - -import os -import glob -import cunumpy as xp -import matplotlib.pyplot as plt - -from struphy.models.rework_model import TwoFluidQuasiNeutralToy - -import warnings -# warnings.filterwarnings("error") - - -BC = 'dirichlet_hom' # 'periodic' | 'dirichlet_hom' | 'dirichlet_inhom' - -name = f"runs/sim_1D_{BC}" - -env = EnvironmentOptions(sim_folder=name) -base_units = BaseUnits(kBT=1.0) - -B0 = 1.0 -nu = 10.0 -nu_e = 1.0 -Nel = (32, 1, 1) -p = (2, 1, 1) -epsilon = 1.0 -dt = 1 -Tend = 1 -sigma = 1 - -time_opts = Time(dt=dt, Tend=Tend) -domain = domains.Cuboid() -equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) -grid = grids.TensorProductGrid(Nel=Nel) - -# ---- boundary conditions ---- -if BC == 'periodic': - spl_kind = (True, True, True) - dirichlet_bc = ((False, False), (False, False), (False, False)) - - bcs_u = bcs_ue = { - (0, -1): "periodic", (0, 1): "periodic", - (1, -1): "periodic", (1, 1): "periodic", - (2, -1): "periodic", (2, 1): "periodic", - } - boundary_data_u = boundary_data_ue = None - -elif BC == 'dirichlet_hom': - spl_kind = (False, True, True) - dirichlet_bc = ((True, True), (False, False), (False, False)) - - bcs_u = bcs_ue = { - (0, -1): "dirichlet", (0, 1): "dirichlet", - (1, -1): "periodic", (1, 1): "periodic", - (2, -1): "periodic", (2, 1): "periodic", - } - boundary_data_u = boundary_data_ue = None - -elif BC == 'dirichlet_inhom': - spl_kind = (False, True, True) - dirichlet_bc = ((False, False), (False, False), (False, False)) - - bcs_u = bcs_ue = { - (0, -1): "dirichlet", (0, 1): "dirichlet", - (1, -1): "periodic", (1, 1): "periodic", - (2, -1): "periodic", (2, 1): "periodic", - } - boundary_data_u = { - (0, -1): lambda x, y, z: (zeros_like(x) + 1, zeros_like(x), zeros_like(x)), - (0, 1): lambda x, y, z: (zeros_like(x) + 2, zeros_like(x), zeros_like(x)), - } - boundary_data_ue = { - (0, -1): lambda x, y, z: (zeros_like(x) + 1, zeros_like(x), zeros_like(x)), - (0, 1): lambda x, y, z: (zeros_like(x) + 2, zeros_like(x), zeros_like(x)), - } - -derham_opts = DerhamOptions( - p=p, - spl_kind=spl_kind, - dirichlet_bc=dirichlet_bc, -) - -# ---- manufactured solutions ---- -if BC == 'periodic': - def mms_phi(x, y, z): - return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) - - def mms_ion_u(x, y, z): - return xp.sin(2 * xp.pi * x) + 1, xp.zeros_like(x), xp.zeros_like(x) - - def mms_electron_u(x, y, z): - return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) - -elif BC == 'dirichlet_hom': - def mms_phi(x, y, z): - return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) - - def mms_ion_u(x, y, z): - return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) - - def mms_electron_u(x, y, z): - return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) - -elif BC == 'dirichlet_inhom': - def mms_phi(x, y, z): - return x + 1, xp.zeros_like(x), xp.zeros_like(x) - - def mms_ion_u(x, y, z): - return x + 1, xp.zeros_like(x), xp.zeros_like(x) - - def mms_electron_u(x, y, z): - return x + 1, xp.zeros_like(x), xp.zeros_like(x) - -# ---- source terms ---- -if BC == 'periodic': - def source_function_u(x, y, z): - fx = 2.0 * pi * cos(2 * pi * x) + nu * 4.0 * pi**2 * sin(2 * pi * x) - fy = (sin(2 * pi * x) + 1.0) * B0 / epsilon - fz = zeros_like(x) - return fx, fy, fz - - def source_function_ue(x, y, z): - fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) - fy = -sin(2 * pi * x) * B0 / epsilon - fz = zeros_like(x) - return fx, fy, fz - -elif BC == 'dirichlet_hom': - def source_function_u(x, y, z): - fx = 2.0 * pi * cos(2 * pi * x) + nu * 4.0 * pi**2 * sin(2 * pi * x) - fy = B0 * sin(2 * pi * x) / epsilon - fz = zeros_like(x) - return fx, fy, fz - - def source_function_ue(x, y, z): - fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) - fy = -sin(2 * pi * x) * B0 / epsilon - fz = zeros_like(x) - return fx, fy, fz - -elif BC == 'dirichlet_inhom': - def source_function_u(x, y, z): - fx = ones_like(x) - fy = B0 * (1 + x) / epsilon - fz = zeros_like(x) - return fx, fy, fz - - def source_function_ue(x, y, z): - fx = -ones_like(x) - sigma * (1 + x) - fy = -B0 * (1 + x) / epsilon - fz = zeros_like(x) - return fx, fy, fz - -# ---- model ---- -model = TwoFluidQuasiNeutralToy() -model.ions.set_phys_params() -model.electrons.set_phys_params() - -model.propagators.qn_full.options = model.propagators.qn_full.Options( - nu=nu, - nu_e=nu_e, - eps_norm=epsilon, - stab_sigma=sigma, - source_u=source_function_u, - source_ue=source_function_ue, - solver='gmres', - boundary_conditions_u=bcs_u, - boundary_conditions_ue=bcs_ue, - boundary_data_u=boundary_data_u, - boundary_data_ue=boundary_data_ue, -) - -if __name__ == "__main__": - main.run(model, - params_path=__file__, - env=env, - base_units=base_units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=True, - ) - - path = os.path.join(os.getcwd(), name) - main.pproc(path) - simdata = main.load_data(path) - - n1_vals = simdata.grids_log[0] - x = xp.linspace(0, 1, 100) - - os.makedirs(f'{name}/plots', exist_ok=True) - for f in glob.glob(f'{name}/plots/*.png'): - os.remove(f) - - def save_plot(n1_vals, numerical, analytical, ylabel, title, fname, t): - plt.plot(n1_vals, numerical, label='numerical') - plt.plot(x, analytical, '--', label='manufactured') - plt.plot(n1_vals, numerical, 'k.', markersize=4, label='n1 points') - plt.xlabel('n1 (radial)') - plt.ylabel(ylabel) - plt.title(f'{title} at t={t:.3f}') - plt.legend() - plt.grid(True) - plt.savefig(f'{name}/plots/{fname}_{t:.3f}.png', dpi=300) - plt.clf() - - for t in list(simdata.spline_values['ions']['u_log'].keys()): - - u_ions = simdata.spline_values['ions']['u_log'][t] - u_electrons = simdata.spline_values['electrons']['u_log'][t] - phi = simdata.spline_values['em_fields']['phi_log'][t] - - mms_phi_x, _, _ = mms_phi(x, x*0, x*0) - mms_ion_ux, mms_ion_uy, _ = mms_ion_u(x, x*0, x*0) - mms_el_ux, mms_el_uy, _ = mms_electron_u(x, x*0, x*0) - - save_plot(n1_vals, phi[0][:, 0, 0], mms_phi_x, 'φ', 'Electrostatic potential φ', 'plot_potential', t) - save_plot(n1_vals, u_ions[0][:, 0, 0], mms_ion_ux, 'u_x', 'Ion velocity (u_x)', 'plot_ion_ux', t) - save_plot(n1_vals, u_electrons[0][:, 0, 0], mms_el_ux, 'u_x', 'Electron velocity (u_x)', 'plot_electron_ux', t) \ No newline at end of file diff --git a/split_models.py b/split_models.py deleted file mode 100644 index 58a03298b..000000000 --- a/split_models.py +++ /dev/null @@ -1,78 +0,0 @@ -import inspect -import os -import re - -import struphy.models.fluid as fluid -import struphy.models.hybrid as hybrid -import struphy.models.kinetic as kinetic -import struphy.models.toy as toy -from struphy.models.base import StruphyModel - - -def camel_to_snake(name): - s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name) - return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower() - - -imports = """ -import cunumpy as xp -from feectools.ddm.mpi import mpi as MPI -from feectools.linalg.block import BlockVector -from feectools.linalg.stencil import StencilVector - -from struphy.feec.projectors import L2Projector -from struphy.feec.variational_utilities import ( - H1vecMassMatrix_density, - InternalEnergyEvaluator, -) -from struphy.kinetic_background.base import KineticBackground -from struphy.kinetic_background.maxwellians import Maxwellian3D -from struphy.models.base import StruphyModel -from struphy.models.species import ( - DiagnosticSpecies, - FieldSpecies, - FluidSpecies, - ParticleSpecies, -) -from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable -from struphy.pic.accumulation import accum_kernels, accum_kernels_gc -from struphy.pic.accumulation.particles_to_grid import AccumulatorVector -from struphy.polar.basic import PolarVector -from struphy.propagators import ( - propagators_coupling, - propagators_fields, - propagators_markers, -) -from struphy.utils.pyccel import Pyccelkernel - -rank = MPI.COMM_WORLD.Get_rank() -""" - -# Output directory -out_dir = "src/struphy/models" -os.makedirs(out_dir, exist_ok=True) - -model_dict = {} - -# Iterate over all modules and discover subclasses of StruphyModel -for model_type in [toy, fluid, hybrid, kinetic]: - for _, cls in model_type.__dict__.items(): - if isinstance(cls, type) and issubclass(cls, StruphyModel) and cls != StruphyModel: - model_name = cls.__name__ - try: - # Get the source code of the class - model_code = inspect.getsource(cls) - model_dict[model_name] = model_code - except Exception as e: - print(f"Could not get source for {model_name}: {e}") - -# Write each model to its own file -for model_name, model_code in model_dict.items(): - file_name = camel_to_snake(model_name) + ".py" - file_path = os.path.join(out_dir, file_name) - with open(file_path, "w") as f: - f.write(imports) - f.write("\n\n") - f.write(model_code) - -print(f"Written {len(model_dict)} model files to {out_dir}") From ce120c619b951cf3734320a60aae8f1d6f32c74d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Thu, 5 Mar 2026 21:40:56 +0000 Subject: [PATCH 05/41] Rebased onto 3.0.4 v2 --- src/struphy/io/options.py | 6 +- ...rk_saddle_point.py => saddle_point_new.py} | 14 +-- .../{rework_model.py => two_fluid_new.py} | 26 +++--- src/struphy/propagators/base.py | 2 +- ...py => propagators_fields_two_fluid_new.py} | 85 +++++-------------- 5 files changed, 50 insertions(+), 83 deletions(-) rename src/struphy/linear_algebra/{rework_saddle_point.py => saddle_point_new.py} (97%) rename src/struphy/models/{rework_model.py => two_fluid_new.py} (85%) rename src/struphy/propagators/{rework_propagator.py => propagators_fields_two_fluid_new.py} (86%) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 2c626df78..91577d1e0 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -2,10 +2,10 @@ from dataclasses import dataclass from typing import Literal -import cunumpy as xp -from psydac.ddm.mpi import mpi as MPI +from struphy.utils.utils import check_option -from struphy.physics.physics import ConstantsOfNature +import cunumpy as xp +from feectools.ddm.mpi import mpi as MPI ## Literal options diff --git a/src/struphy/linear_algebra/rework_saddle_point.py b/src/struphy/linear_algebra/saddle_point_new.py similarity index 97% rename from src/struphy/linear_algebra/rework_saddle_point.py rename to src/struphy/linear_algebra/saddle_point_new.py index 593dc8524..292313f69 100644 --- a/src/struphy/linear_algebra/rework_saddle_point.py +++ b/src/struphy/linear_algebra/saddle_point_new.py @@ -2,10 +2,10 @@ import cunumpy as xp import scipy as sc -from psydac.linalg.basic import LinearOperator, Vector -from psydac.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace -from psydac.linalg.direct_solvers import SparseSolver -from psydac.linalg.solvers import inverse +from feectools.linalg.basic import LinearOperator, Vector +from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace +from feectools.linalg.direct_solvers import SparseSolver +from feectools.linalg.solvers import inverse from struphy.linear_algebra.tests.test_saddlepoint_massmatrices import _plot_residual_norms @@ -27,7 +27,7 @@ class SaddlePointSolver: f \cr 0 } \right) - using either the Uzawa iteration :math:`BA^{-1}B^{\top} y = BA^{-1} f` or using on of the solvers given in :mod:`psydac.linalg.solvers`. The prefered solver is GMRES. + using either the Uzawa iteration :math:`BA^{-1}B^{\top} y = BA^{-1} f` or using on of the solvers given in :mod:`feectools.linalg.solvers`. The prefered solver is GMRES. The decission which variant to use is given by the type of A. If A is of type list of xp.ndarrays or sc.sparse.csr_matrices, then this class uses the Uzawa algorithm. If A is of type LinearOperator or BlockLinearOperator, a solver is used for the inverse. Using the Uzawa algorithm, solution is given by: @@ -44,8 +44,8 @@ class SaddlePointSolver: Either the entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. Alternative: Give whole matrice A as LinearOperator or BlockLinearOperator. list: Uzawa algorithm is used. - LinearOperator: A solver given in :mod:`psydac.linalg.solvers` is used. Specified by solver_name. - BlockLinearOperator: A solver given in :mod:`psydac.linalg.solvers` is used. Specified by solver_name. + LinearOperator: A solver given in :mod:`feectools.linalg.solvers` is used. Specified by solver_name. + BlockLinearOperator: A solver given in :mod:`feectools.linalg.solvers` is used. Specified by solver_name. B : list, LinearOperator or BlockLinearOperator Lower left block. diff --git a/src/struphy/models/rework_model.py b/src/struphy/models/two_fluid_new.py similarity index 85% rename from src/struphy/models/rework_model.py rename to src/struphy/models/two_fluid_new.py index f38629fcc..e5a569eab 100644 --- a/src/struphy/models/rework_model.py +++ b/src/struphy/models/two_fluid_new.py @@ -1,13 +1,15 @@ -import cunumpy as xp -from psydac.ddm.mpi import mpi as MPI +from feectools.ddm.mpi import mpi as MPI -from struphy.feec.projectors import L2Projector -from struphy.feec.variational_utilities import InternalEnergyEvaluator +from struphy.io.options import LiteralOptions from struphy.models.base import StruphyModel -from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies -from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable -from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.propagators import rework_propagator +from struphy.models.species import ( + FieldSpecies, + FluidSpecies, +) +from struphy.models.variables import FEECVariable +from struphy.propagators import ( + propagators_fields_two_fluid_new, +) rank = MPI.COMM_WORLD.Get_rank() @@ -50,6 +52,10 @@ class TwoFluidQuasiNeutralToy(StruphyModel): in plasma physics, Journal of Computational Physics 2018. """ + @classmethod + def model_type(cls) -> LiteralOptions.ModelTypes: + return "Fluid" + ## species class EMfields(FieldSpecies): @@ -71,7 +77,7 @@ def __init__(self): class Propagators: def __init__(self): - self.qn_full = rework_propagator.TwoFluidQuasiNeutralFull() + self.qn_full = propagators_fields_two_fluid_new.TwoFluidQuasiNeutralFull() ## abstract methods @@ -102,7 +108,7 @@ def bulk_species(self): def velocity_scale(self): return "thermal" - def allocate_helpers(self): + def allocate_helpers(self, verbose=False): pass def update_scalar_quantities(self): diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index e7736fab8..62468a54f 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -13,7 +13,7 @@ from struphy.feec.psydac_derham import Derham from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB from struphy.geometry.base import Domain -from struphy.io.options import check_option +from struphy.utils.utils import check_option from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable diff --git a/src/struphy/propagators/rework_propagator.py b/src/struphy/propagators/propagators_fields_two_fluid_new.py similarity index 86% rename from src/struphy/propagators/rework_propagator.py rename to src/struphy/propagators/propagators_fields_two_fluid_new.py index bf6afcadf..c5cf54039 100644 --- a/src/struphy/propagators/rework_propagator.py +++ b/src/struphy/propagators/propagators_fields_two_fluid_new.py @@ -1,63 +1,22 @@ - -import copy -from copy import deepcopy from dataclasses import dataclass -from typing import Callable, Literal, get_args, cast +from typing import Callable, Literal, cast from warnings import warn -import cunumpy as xp -import scipy as sc -from line_profiler import profile -from matplotlib import pyplot as plt -from numpy import zeros -from psydac.api.essential_bc import apply_essential_bc_stencil -from psydac.ddm.mpi import mpi as MPI -from psydac.linalg.basic import ComposedLinearOperator, IdentityOperator, ZeroOperator, InverseLinearOperator -from psydac.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace -from psydac.linalg.solvers import inverse -from psydac.linalg.stencil import StencilVector - -import struphy.feec.utilities as util -from struphy.examples.restelli2018 import callables -from struphy.feec import preconditioner -from struphy.feec.basis_projection_ops import ( - BasisProjectionOperator, BasisProjectionOperatorLocal, - BasisProjectionOperators, CoordinateProjector, -) -from struphy.feec.linear_operators import BoundaryOperator -from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators -from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner, MassMatrixPreconditioner +from feectools.api.essential_bc import apply_essential_bc_stencil +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.basic import IdentityOperator, InverseLinearOperator +from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace +from feectools.linalg.solvers import inverse + +from struphy.feec.basis_projection_ops import BasisProjectionOperators +from struphy.feec.mass import WeightedMassOperators from struphy.feec.projectors import L2Projector -from struphy.feec.psydac_derham import Derham, SplineFunction -from struphy.feec.variational_utilities import ( - BracketOperator, Hdiv0_transport_operator, InternalEnergyEvaluator, - KineticEnergyEvaluator, Pressure_transport_operator, -) -from struphy.fields_background.equils import set_defaults -from struphy.geometry.utilities import TransformedPformComponent -from struphy.initial import perturbations -from struphy.io.options import ( - OptsDirectSolver, OptsGenSolver, OptsMassPrecond, OptsNonlinearSolver, - OptsSaddlePointSolver, OptsSymmSolver, OptsVecSpace, check_option, -) -from struphy.io.setup import descend_options_dict -from struphy.kinetic_background.base import Maxwellian -from struphy.kinetic_background.maxwellians import GyroMaxwellian2D, Maxwellian3D -from struphy.linear_algebra.saddle_point import SaddlePointSolver -from struphy.linear_algebra.schur_solver import SchurSolver, SchurSolverFull -from struphy.linear_algebra.solver import NonlinearSolverParameters, SolverParameters -from struphy.models.species import Species -from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable -from struphy.ode.solvers import ODEsolverFEEC -from struphy.ode.utils import ButcherTableau, OptsButcher -from struphy.pic.accumulation import accum_kernels, accum_kernels_gc -from struphy.pic.accumulation.filter import FilterParameters -from struphy.pic.accumulation.particles_to_grid import Accumulator, AccumulatorVector -from struphy.pic.base import Particles -from struphy.pic.particles import Particles5D, Particles6D -from struphy.polar.basic import PolarVector +from struphy.feec.psydac_derham import Derham +from struphy.io.options import OptsGenSolver +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel +from struphy.utils.utils import check_option class TwoFluidQuasiNeutralFull(Propagator): @@ -264,8 +223,9 @@ def _apply_boundary_conditions(self, vec, boundary_conditions): ### Allocate # ========================================================================= - def allocate(self): + def allocate(self, verbose=False): + self.verbose = verbose self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 self._dt = None @@ -295,6 +255,7 @@ def allocate(self): # ---- unconstrained operators (for RHS assembly) ---------------------- + self._M2 = self.mass_ops.M2 self._M2B = - self.mass_ops.M2B self._div = self.derham.div @@ -303,9 +264,9 @@ def allocate(self): self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - + self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl - self._A22 = (- self.options.stab_sigma * IdentityOperator(self._A11.domain) + self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) # ---- constrained operators (for system matrix) ----------------------- @@ -319,16 +280,16 @@ def allocate(self): self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) - + self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._A11_v0.domain) + self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) # ---- block saddle-point system ---------------------------------------- - self._block_domain_v0 = BlockVectorSpace(self._A11_v0.domain, self._A22_v0.domain) + self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self._M3_v0.codomain + self._block_codomain_B_v0 = self._derham_v0.Vh["3"] self._B1_v0 = - self._M3_v0 @ self._div_v0 self._B2_v0 = self._M3_v0 @ self._div_v0 From f48b865bf9d28c0204b2861a02e8c830cb4d782e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Thu, 12 Mar 2026 22:23:27 +0000 Subject: [PATCH 06/41] Replaced old propagator in propagators_fields and other minor changes. --- .gitignore | 1 - feectools | 2 +- src/struphy/io/options.py | 4 +- .../linear_algebra/saddle_point_new.py | 567 ------ src/struphy/models/two_fluid_new.py | 130 -- src/struphy/propagators/propagators_fields.py | 1673 ++++++++--------- .../propagators_fields_two_fluid_new.py | 440 ----- src/struphy/utils/utils.py | 6 +- struphy-parameter-files | 2 +- 9 files changed, 755 insertions(+), 2070 deletions(-) delete mode 100644 src/struphy/linear_algebra/saddle_point_new.py delete mode 100644 src/struphy/models/two_fluid_new.py delete mode 100644 src/struphy/propagators/propagators_fields_two_fluid_new.py diff --git a/.gitignore b/.gitignore index afefb317e..4bac15ae5 100644 --- a/.gitignore +++ b/.gitignore @@ -96,7 +96,6 @@ src/struphy/io/inp/params_* src/struphy/models/models_list src/struphy/models/models_message -runs/ bin/ share/ diff --git a/feectools b/feectools index 278908bdb..1981de121 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 278908bdb513a402ae4121843f4887467b4a61b2 +Subproject commit 1981de121fe6949b4a0797a3d61c8089c25c0b9f diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 91577d1e0..cbe656e72 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -102,9 +102,9 @@ class LiteralOptions: # solvers OptsSymmSolver = Literal["pcg", "cg"] - OptsGenSolver = Literal["pbicgstab", "bicgstab", "GMRES"] + OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"] OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] - OptsSaddlePointSolver = Literal["Uzawa", "GMRES"] + OptsSaddlePointSolver = Literal["uzawa"] OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] OptsNonlinearSolver = Literal["Picard", "Newton"] OptsButcher = Literal["rk4", "forward_euler", "heun2", "rk2", "heun3", "3/8 rule"] diff --git a/src/struphy/linear_algebra/saddle_point_new.py b/src/struphy/linear_algebra/saddle_point_new.py deleted file mode 100644 index 292313f69..000000000 --- a/src/struphy/linear_algebra/saddle_point_new.py +++ /dev/null @@ -1,567 +0,0 @@ -from typing import Union - -import cunumpy as xp -import scipy as sc -from feectools.linalg.basic import LinearOperator, Vector -from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace -from feectools.linalg.direct_solvers import SparseSolver -from feectools.linalg.solvers import inverse - -from struphy.linear_algebra.tests.test_saddlepoint_massmatrices import _plot_residual_norms - - -class SaddlePointSolver: - r"""Solves for :math:`(x, y)` in the saddle point problem - - .. math:: - - \left( \matrix{ - A & B^{\top} \cr - B & 0 - } \right) - \left( \matrix{ - x \cr y - } \right) - = - \left( \matrix{ - f \cr 0 - } \right) - - using either the Uzawa iteration :math:`BA^{-1}B^{\top} y = BA^{-1} f` or using on of the solvers given in :mod:`feectools.linalg.solvers`. The prefered solver is GMRES. - The decission which variant to use is given by the type of A. If A is of type list of xp.ndarrays or sc.sparse.csr_matrices, then this class uses the Uzawa algorithm. - If A is of type LinearOperator or BlockLinearOperator, a solver is used for the inverse. - Using the Uzawa algorithm, solution is given by: - - .. math:: - - y = \left[ B A^{-1} B^{\top}\right]^{-1} B A^{-1} f \,, \qquad - x = A^{-1} \left[ f - B^{\top} y \right] \,. - - Parameters - ---------- - A : list, LinearOperator or BlockLinearOperator - Upper left block. - Either the entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. - Alternative: Give whole matrice A as LinearOperator or BlockLinearOperator. - list: Uzawa algorithm is used. - LinearOperator: A solver given in :mod:`feectools.linalg.solvers` is used. Specified by solver_name. - BlockLinearOperator: A solver given in :mod:`feectools.linalg.solvers` is used. Specified by solver_name. - - B : list, LinearOperator or BlockLinearOperator - Lower left block. - Uzwaw Algorithm: All entries of block B are given either as list of xp.ndarray or sc.sparse.csr_matrix. - Solver: Give whole B as LinearOperator or BlocklinearOperator - - F : list - Right hand side of the upper block. - Uzawa: Given as list of xp.ndarray or sc.sparse.csr_matrix. - Solver: Given as LinearOperator or BlockLinearOperator - - Apre : list - The non-inverted preconditioner for entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. Only required for the Uzawa algorithm. - - method_to_solve : str - Method for the inverses. Choose from 'DirectNPInverse', 'ScipySparse', 'InexactNPInverse' ,'SparseSolver'. Only required for the Uzawa algorithm. - - preconditioner : bool - Wheter to use preconditioners given in Apre or not. Only required for the Uzawa algorithm. - - spectralanalysis : bool - Do the spectralanalyis for the matrices in A and if preconditioner given, compare them to the preconditioned matrices. Only possible if A is given as list. - - dimension : str - Which of the predefined manufactured solutions to use ('1D', '2D' or 'Restelli') - - tol : float - Convergence tolerance for the potential residual. - - max_iter : int - Maximum number of iterations allowed. - """ - - def __init__( - self, - A: Union[list, LinearOperator, BlockLinearOperator], - B: Union[list, LinearOperator, BlockLinearOperator], - F: Union[list, Vector, BlockVector], - Apre: list = None, - method_to_solve: str = "DirectNPInverse", - preconditioner: bool = False, - spectralanalysis: bool = False, - dimension: str = "2D", - solver_name: str = "GMRES", - tol: float = 1e-8, - max_iter: int = 1000, - **solver_params, - ): - assert type(A) is type(B) - if isinstance(A, list): - self._variant = "Uzawa" - for i in A: - assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) - for i in B: - assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) - for i in F: - assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) - for i in Apre: - assert ( - isinstance(i, xp.ndarray) - or isinstance(i, sc.sparse.csr_matrix) - or isinstance(i, sc.sparse.csr_array) - ) - assert method_to_solve in ("SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse") - assert A[0].shape[0] == B[0].shape[1] - assert A[0].shape[1] == B[0].shape[1] - assert A[1].shape[0] == B[1].shape[1] - assert A[1].shape[1] == B[1].shape[1] - - self._method_to_solve = ( - method_to_solve # 'SparseSolver', 'ScipySparse', 'InexactNPInverse', 'DirectNPInverse' - ) - self._preconditioner = preconditioner - - elif isinstance(A, LinearOperator) or isinstance(A, BlockLinearOperator): - self._variant = "Inverse_Solver" - assert A.domain == B.domain - assert A.codomain == B.domain - self._solver_name = solver_name - if solver_params["pc"] is None: - solver_params.pop("pc") - - # operators - self._A = A - self._Apre = Apre - self._B = B - self._F = F - self._tol = tol - self._max_iter = max_iter - self._spectralanalysis = spectralanalysis - self._dimension = dimension - self._verbose = solver_params["verbose"] - - if self._variant == "Inverse_Solver": - self._block_domainM = BlockVectorSpace(self._A.domain, self._B.transpose().domain) - self._block_codomainM = self._block_domainM - self._blocks = [[self._A, self._B.T], [self._B, None]] - _Minit = BlockLinearOperator(self._block_domainM, self._block_codomainM, blocks=self._blocks) - self._solverMinv = inverse(_Minit, solver_name, tol=tol, maxiter=max_iter, **solver_params) - - # Solution vectors - self._P = B.codomain.zeros() - self._U = A.codomain.zeros() - self._Utmp = F.copy() * 0 - # Allocate memory for call - self._rhstemp = BlockVector(self._block_domainM, blocks=[A.codomain.zeros(), self._B.codomain.zeros()]) - - elif self._variant == "Uzawa": - if self._method_to_solve in ("InexactNPInverse", "SparseSolver"): - self._preconditioner = False - - self._Anp = self._A[0] - self._Aenp = self._A[1] - self._B1np = self._B[0] - self._B2np = self._B[1] - - # Instanciate inverses - self._setup_inverses() - - # Solution vectors numpy - self._Pnp = xp.zeros(self._B1np.shape[0]) - self._Unp = xp.zeros(self._A[0].shape[1]) - self._Uenp = xp.zeros(self._A[1].shape[1]) - # Allocate memory for matrices used in solving the system - self._rhs0np = self._F[0].copy() - self._rhs1np = self._F[1].copy() - - # List to store residual norms - self._residual_norms = [] - self._stepsize = 0.0 - - @property - def A(self): - """Upper left block.""" - return self._A - - @A.setter - def A(self, a): - if self._variant == "Uzawa": - need_update = True - A0_old, A1_old = self._A - A0_new, A1_new = a - if self._method_to_solve in ("ScipySparse", "SparseSolver"): - same_A0 = (A0_old != A0_new).nnz == 0 - same_A1 = (A1_old != A1_new).nnz == 0 - else: - same_A0 = xp.allclose(A0_old, A0_new, atol=1e-10) - same_A1 = xp.allclose(A1_old, A1_new, atol=1e-10) - if same_A0 and same_A1: - need_update = False - self._A = a - self._Anp = self._A[0] - self._Aenp = self._A[1] - if need_update: - self._setup_inverses() - elif self._variant == "Inverse_Solver": - self._A = a - - @property - def B(self): - """Lower left block.""" - return self._B - - @B.setter - def B(self, b): - self._B = b - - @property - def F(self): - """Right hand side vector.""" - return self._F - - @F.setter - def F(self, f): - self._F = f - - @property - def Apre(self): - """Preconditioner for upper left block A.""" - return self._Apre - - @Apre.setter - def Apre(self, a): - if self._variant == "Uzawa": - need_update = True - A0_old, A1_old = self._Apre - A0_new, A1_new = a - if self._method_to_solve in ("ScipySparse", "SparseSolver"): - same_A0 = (A0_old != A0_new).nnz == 0 - same_A1 = (A1_old != A1_new).nnz == 0 - else: - same_A0 = xp.allclose(A0_old, A0_new, atol=1e-10) - same_A1 = xp.allclose(A1_old, A1_new, atol=1e-10) - if same_A0 and same_A1: - need_update = False - self._Apre = a - if need_update: - self._setup_inverses() - elif self._variant == "Inverse_Solver": - self._Apre = a - - def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): # todo should have options to use other than uzawa. should solve in full generality. A being block diagonal is a special case of uzawa - """ - Solves the saddle-point problem using the Uzawa algorithm. - - Parameters - ---------- - U_init : Vector, xp.ndarray or sc.sparse.csr.csr_matrix, optional - Initial guess for the velocity of the ions. If None, initializes to zero. Types xp.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. - - Ue_init : Vector, xp.ndarray or sc.sparse.csr.csr_matrix, optional - Initial guess for the velocity of the electrons. If None, initializes to zero. Types xp.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. - - P_init : Vector, optional - Initial guess for the potential. If None, initializes to zero. - - Returns - ------- - U : Vector - Solution vector for the velocity. - - P : Vector - Solution vector for the potential. - - info : dict - Convergence information. - """ - - # TODO this contains two different strategies! favágás and actual uzawa - if self._variant == "Inverse_Solver": # todo not in the """""saddle point solver""""""" - self._P1 = P_init if P_init is not None else self._P - self._U1 = U_init if U_init is not None else self._Utmp[0] - self._U2 = Ue_init if Ue_init is not None else self._Utmp[1] - - _blocksM = [[self._A, self._B.T], [self._B, None]] - _M = BlockLinearOperator(self._block_domainM, self._block_codomainM, blocks=_blocksM) - _RHS = BlockVector(self._block_domainM, blocks=[self._F, self._B.codomain.zeros()]) - - self._blockU = BlockVector(self._A.domain, blocks=[self._U1, self._U2]) - self._solblocks = [self._blockU, self._P1] - # comment out the next two lines if working with lifting and GMRES - x0 = BlockVector(self._block_domainM, blocks=self._solblocks) - self._solverMinv._options["x0"] = x0 - - # use setter to update lhs matrix - self._solverMinv.linop = _M - - # Initialize P to zero or given initial guess - self._sol = self._solverMinv.dot(_RHS, out=self._rhstemp) - self._U = self._sol[0] - self._P = self._sol[1] - - return self._U, self._P, self._solverMinv._info - - elif self._variant == "Uzawa": - info = {} - - if self._spectralanalysis: - self._spectralresult = self._spectral_analysis() - else: - self._spectralresult = [] - - # Initialize P to zero or given initial guess - if isinstance(U_init, xp.ndarray) or isinstance(U_init, sc.sparse.csr.csr_matrix): - self._Pnp = P_init if P_init is not None else self._P - self._Unp = U_init if U_init is not None else self._U - self._Uenp = Ue_init if U_init is not None else self._Ue - - else: - self._Pnp = P_init.toarray() if P_init is not None else self._Pnp - self._Unp = U_init.toarray() if U_init is not None else self._Unp - self._Uenp = Ue_init.toarray() if U_init is not None else self._Uenp - - if self._verbose: - print("Uzawa solver:") - print("+---------+---------------------+") - print("+ Iter. # | L2-norm of residual |") - print("+---------+---------------------+") - template = "| {:7d} | {:19.2e} |" - - for iteration in range(self._max_iter): - # Step 1: Compute velocity U by solving A U = -Bᵀ P + F -A Un - self._rhs0np *= 0 - self._rhs0np -= self._B1np.transpose().dot(self._Pnp) - self._rhs0np -= self._Anp.dot(self._Unp) - self._rhs0np += self._F[0] - if not self._preconditioner: - self._Unp += self._Anpinv.dot(self._rhs0np) - elif self._preconditioner: - self._Unp += self._Anpinv.dot(self._A11npinv @ self._rhs0np) - - R1 = self._B1np.dot(self._Unp) - - self._rhs1np *= 0 - self._rhs1np -= self._B2np.transpose().dot(self._Pnp) - self._rhs1np -= self._Aenp.dot(self._Uenp) - self._rhs1np += self._F[1] - if not self._preconditioner: - self._Uenp += self._Aenpinv.dot(self._rhs1np) - elif self._preconditioner: - self._Uenp += self._Aenpinv.dot(self._A22npinv @ self._rhs1np) - - R2 = self._B2np.dot(self._Uenp) - - # Step 2: Compute residual R = BU (divergence of U) - R = R1 + R2 # self._B1np.dot(self._Unp) + self._B2np.dot(self._Uenp) - residual_norm = xp.linalg.norm(R) - residual_normR1 = xp.linalg.norm(R) - self._residual_norms.append(residual_normR1) # Store residual norm - # Check for convergence based on residual norm - if residual_norm < self._tol: - if self._verbose: - print(template.format(iteration + 1, residual_norm)) - print("+---------+---------------------+") - info["success"] = True - info["niter"] = iteration + 1 - if self._verbose: - _plot_residual_norms(self._residual_norms) - return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult - - # Steepest gradient - alpha = (R.dot(R)) / (R.dot(self._Precnp.dot(R))) - # Minimal residual - # alpha = ((self._Precnp.dot(R)).dot(R)) / ((self._Precnp.dot(R)).dot(self._Precnp.dot(R))) - self._Pnp += alpha.real * R.real - - if self._verbose: - print(template.format(iteration + 1, residual_norm)) - - if self._verbose: - print("+---------+---------------------+") - - # Return with info if maximum iterations reached - info["success"] = False - info["niter"] = iteration + 1 - if self._verbose: - _plot_residual_norms(self._residual_norms) - return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult - - def _setup_inverses(self): - A0 = self._A[0] - A1 = self._A[1] - - # === Preconditioner inverses, if used - if self._preconditioner: - A11_pre = self._Apre[0] - A22_pre = self._Apre[1] - - if hasattr(self, "_A11npinv") and self._is_inverse_still_valid(self._A11npinv, A11_pre, "A11 pre"): - pass - else: - self._A11npinv = self._compute_inverse(A11_pre, which="A11 pre") - - if hasattr(self, "_A22npinv") and self._is_inverse_still_valid(self._A22npinv, A22_pre, "A22 pre"): - pass - else: - self._A22npinv = self._compute_inverse(A22_pre, which="A22 pre") - - # === Inverse for A[0] if preconditioned - if hasattr(self, "_Anpinv") and self._is_inverse_still_valid(self._Anpinv, A0, "A[0]", pre=self._A11npinv): - pass - else: - self._Anpinv = self._compute_inverse(self._A11npinv @ A0, which="A[0]") - - # === Inverse for A[1] - if hasattr(self, "_Aenpinv") and self._is_inverse_still_valid( - self._Aenpinv, - A1, - "A[1]", - pre=self._A22npinv, - ): - pass - else: - self._Aenpinv = self._compute_inverse(self._A22npinv @ A1, which="A[1]") - - else: # No preconditioning: - # === Inverse for A[0] - if hasattr(self, "_Anpinv") and self._is_inverse_still_valid(self._Anpinv, A0, "A[0]"): - pass - else: - self._Anpinv = self._compute_inverse(A0, which="A[0]") - - # === Inverse for A[1] - if hasattr(self, "_Aenpinv") and self._is_inverse_still_valid(self._Aenpinv, A1, "A[1]"): - pass - else: - self._Aenpinv = self._compute_inverse(A1, which="A[1]") - - # Precompute Schur complement - self._Precnp = self._B1np @ self._Anpinv @ self._B1np.T + self._B2np @ self._Aenpinv @ self._B2np.T - - def _is_inverse_still_valid(self, inv, mat, name="", pre=None): - # try: - if pre is not None: - test_mat = pre @ mat - else: - test_mat = mat - I_approx = inv @ test_mat - - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - I_exact = xp.eye(test_mat.shape[0]) - if not xp.allclose(I_approx, I_exact, atol=1e-6): - diff = I_approx - I_exact - max_abs = xp.abs(diff).max() - print(f"{name} inverse is NOT valid anymore. Max diff: {max_abs:.2e}") - return False - print(f"{name} inverse is still valid.") - return True - elif self._method_to_solve == "ScipySparse": - I_exact = sc.sparse.identity(I_approx.shape[0], format=I_approx.format) - diff = (I_approx - I_exact).tocoo() - max_abs = xp.abs(diff.data).max() if diff.nnz > 0 else 0.0 - - if max_abs > 1e-6: - print(f"{name} inverse is NOT valid anymore.") - print(f"Max absolute difference: {max_abs:.2e}") - print(f"Number of differing entries: {diff.nnz}") - return False - print(f"{name} inverse is still valid.") - return True - - def _compute_inverse(self, mat, which="matrix"): - print(f"Computing inverse for {which} using method {self._method_to_solve}") - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - return xp.linalg.inv(mat) - elif self._method_to_solve == "ScipySparse": - return sc.sparse.linalg.inv(mat) - elif self._method_to_solve == "SparseSolver": - solver = SparseSolver(mat) - return solver.solve(xp.eye(mat.shape[0])) - else: - raise ValueError(f"Unknown solver method {self._method_to_solve}") - - def _spectral_analysis(self): - # Spectral analysis - # A11 before - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA11_before, eigvecs_before = xp.linalg.eig(self._A[0]) - condA11_before = xp.linalg.cond(self._A[0]) - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA11_before, eigvecs_before = xp.linalg.eig(self._A[0].toarray()) - condA11_before = xp.linalg.cond(self._A[0].toarray()) - maxbeforeA11 = max(eigvalsA11_before) - maxbeforeA11_abs = xp.max(xp.abs(eigvalsA11_before)) - minbeforeA11_abs = xp.min(xp.abs(eigvalsA11_before)) - minbeforeA11 = min(eigvalsA11_before) - specA11_bef = maxbeforeA11 / minbeforeA11 - specA11_bef_abs = maxbeforeA11_abs / minbeforeA11_abs - # print(f'{maxbeforeA11 = }') - # print(f'{maxbeforeA11_abs = }') - # print(f'{minbeforeA11_abs = }') - # print(f'{minbeforeA11 = }') - # print(f'{specA11_bef = }') - print(f"{specA11_bef_abs =}") - - # A22 before - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA22_before, eigvecs_before = xp.linalg.eig(self._A[1]) - condA22_before = xp.linalg.cond(self._A[1]) - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA22_before, eigvecs_before = xp.linalg.eig(self._A[1].toarray()) - condA22_before = xp.linalg.cond(self._A[1].toarray()) - maxbeforeA22 = max(eigvalsA22_before) - maxbeforeA22_abs = xp.max(xp.abs(eigvalsA22_before)) - minbeforeA22_abs = xp.min(xp.abs(eigvalsA22_before)) - minbeforeA22 = min(eigvalsA22_before) - specA22_bef = maxbeforeA22 / minbeforeA22 - specA22_bef_abs = maxbeforeA22_abs / minbeforeA22_abs - # print(f'{maxbeforeA22 = }') - # print(f'{maxbeforeA22_abs = }') - # print(f'{minbeforeA22_abs = }') - # print(f'{minbeforeA22 = }') - # print(f'{specA22_bef = }') - print(f"{specA22_bef_abs =}") - print(f"{condA22_before =}") - - if self._preconditioner: - # A11 after preconditioning with its inverse - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig(self._A11npinv @ self._A[0]) # Implement this - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig((self._A11npinv @ self._A[0]).toarray()) - maxafterA11_prec = max(eigvalsA11_after_prec) - minafterA11_prec = min(eigvalsA11_after_prec) - maxafterA11_abs_prec = xp.max(xp.abs(eigvalsA11_after_prec)) - minafterA11_abs_prec = xp.min(xp.abs(eigvalsA11_after_prec)) - specA11_aft_prec = maxafterA11_prec / minafterA11_prec - specA11_aft_abs_prec = maxafterA11_abs_prec / minafterA11_abs_prec - # print(f'{maxafterA11_prec = }') - # print(f'{maxafterA11_abs_prec = }') - # print(f'{minafterA11_abs_prec = }') - # print(f'{minafterA11_prec = }') - # print(f'{specA11_aft_prec = }') - print(f"{specA11_aft_abs_prec =}") - - # A22 after preconditioning with its inverse - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA22_after_prec, eigvecs_after = xp.linalg.eig(self._A22npinv @ self._A[1]) # Implement this - condA22_after = xp.linalg.cond(self._A22npinv @ self._A[1]) - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA22_after_prec, eigvecs_after = xp.linalg.eig((self._A22npinv @ self._A[1]).toarray()) - condA22_after = xp.linalg.cond((self._A22npinv @ self._A[1]).toarray()) - maxafterA22_prec = max(eigvalsA22_after_prec) - minafterA22_prec = min(eigvalsA22_after_prec) - maxafterA22_abs_prec = xp.max(xp.abs(eigvalsA22_after_prec)) - minafterA22_abs_prec = xp.min(xp.abs(eigvalsA22_after_prec)) - specA22_aft_prec = maxafterA22_prec / minafterA22_prec - specA22_aft_abs_prec = maxafterA22_abs_prec / minafterA22_abs_prec - # print(f'{maxafterA22_prec = }') - # print(f'{maxafterA22_abs_prec = }') - # print(f'{minafterA22_abs_prec = }') - # print(f'{minafterA22_prec = }') - # print(f'{specA22_aft_prec = }') - print(f"{specA22_aft_abs_prec =}") - - return condA22_before, specA22_bef_abs, condA11_before, condA22_after, specA22_aft_abs_prec - - else: - return condA22_before, specA22_bef_abs, condA11_before diff --git a/src/struphy/models/two_fluid_new.py b/src/struphy/models/two_fluid_new.py deleted file mode 100644 index e5a569eab..000000000 --- a/src/struphy/models/two_fluid_new.py +++ /dev/null @@ -1,130 +0,0 @@ -from feectools.ddm.mpi import mpi as MPI - -from struphy.io.options import LiteralOptions -from struphy.models.base import StruphyModel -from struphy.models.species import ( - FieldSpecies, - FluidSpecies, -) -from struphy.models.variables import FEECVariable -from struphy.propagators import ( - propagators_fields_two_fluid_new, -) - - -rank = MPI.COMM_WORLD.Get_rank() - -class TwoFluidQuasiNeutralToy(StruphyModel): - r"""Linearized, quasi-neutral two-fluid model with zero electron inertia. - - :ref:`normalization`: - - .. math:: - - \hat u = \hat v_\textnormal{th}\,,\qquad e\hat \phi = m \hat v_\textnormal{th}^2\,. - - :ref:`Equations `: - - .. math:: - - \frac{\partial \mathbf u}{\partial t} &= - \nabla \phi + \frac{\mathbf u \times \mathbf B_0}{\varepsilon} + \nu \Delta \mathbf u + \mathbf f\,, - \\[2mm] - 0 &= \nabla \phi - \frac{\mathbf u_e \times \mathbf B_0}{\varepsilon} + \nu_e \Delta \mathbf u_e + \mathbf f_e \,, - \\[3mm] - \nabla & \cdot (\mathbf u - \mathbf u_e) = 0\,, - - where :math:`\mathbf B_0` is a static magnetic field and :math:`\mathbf f, \mathbf f_e` are given forcing terms, - and with the normalization parameter - - .. math:: - - \varepsilon = \frac{1}{\hat \Omega_\textnormal{c} \hat t} \,,\qquad \textnormal{with} \,,\qquad \hat \Omega_{\textnormal{c}} = \frac{(Ze) \hat B}{(A m_\textnormal{H})}\,, - - :ref:`propagators` (called in sequence): - - 1. :class:`~struphy.propagators.propagators_fields.TwoFluidQuasiNeutralFull` - - :ref:`Model info `: - - References - ---------- - [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising - in plasma physics, Journal of Computational Physics 2018. - """ - - @classmethod - def model_type(cls) -> LiteralOptions.ModelTypes: - return "Fluid" - - ## species - - class EMfields(FieldSpecies): - def __init__(self): - self.phi = FEECVariable(space="L2") - self.init_variables() - - class Ions(FluidSpecies): - def __init__(self): - self.u = FEECVariable(space="Hdiv") - self.init_variables() - - class Electrons(FluidSpecies): - def __init__(self): - self.u = FEECVariable(space="Hdiv") - self.init_variables() - - ## propagators - - class Propagators: - def __init__(self): - self.qn_full = propagators_fields_two_fluid_new.TwoFluidQuasiNeutralFull() - - ## abstract methods - - def __init__(self): - if rank == 0: - print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - - # 1. instantiate all species - self.em_fields = self.EMfields() - self.ions = self.Ions() - self.electrons = self.Electrons() - - # 2. instantiate all propagators - self.propagators = self.Propagators() - - # 3. assign variables to propagators - self.propagators.qn_full.variables.u = self.ions.u - self.propagators.qn_full.variables.ue = self.electrons.u - self.propagators.qn_full.variables.phi = self.em_fields.phi - - # define scalars for update_scalar_quantities - - @property - def bulk_species(self): - return self.ions - - @property - def velocity_scale(self): - return "thermal" - - def allocate_helpers(self, verbose=False): - pass - - def update_scalar_quantities(self): - pass - - ## default parameters - def generate_default_parameter_file(self, path=None, prompt=True): - params_path = super().generate_default_parameter_file(path=path, prompt=prompt) - new_file = [] - with open(params_path, "r") as f: - for line in f: - if "BaseUnits()" in line: - new_file += ["base_units = BaseUnits(kBT=1.0)\n"] - else: - new_file += [line] - - with open(params_path, "w") as f: - for line in new_file: - f.write(line) diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 4a38c7c01..661e909ae 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -4,6 +4,7 @@ from copy import deepcopy from dataclasses import dataclass from typing import Callable, Literal, get_args +from warnings import warn import cunumpy as xp import scipy as sc @@ -7644,14 +7645,18 @@ class TwoFluidQuasiNeutralFull(Propagator): :ref:`time_discret`: fully implicit. """ - class Variables: - def __init__(self): - self._u: FEECVariable = None - self._ue: FEECVariable = None - self._phi: FEECVariable = None + # ========================================================================= + ### State variables (ion velocity u, electron velocity ue, pressure phi) + # ========================================================================= + + class Variables(): + def __init__(self) -> None: + self._u: FEECVariable | None = None + self._ue: FEECVariable | None = None + self._phi: FEECVariable | None = None @property - def u(self) -> FEECVariable: + def u(self) -> FEECVariable | None: return self._u @u.setter @@ -7661,7 +7666,7 @@ def u(self, new): self._u = new @property - def ue(self) -> FEECVariable: + def ue(self) -> FEECVariable | None: return self._ue @ue.setter @@ -7671,7 +7676,7 @@ def ue(self, new): self._ue = new @property - def phi(self) -> FEECVariable: + def phi(self) -> FEECVariable | None: return self._phi @phi.setter @@ -7683,968 +7688,784 @@ def phi(self, new): def __init__(self): self.variables = self.Variables() + # ========================================================================= + ### Options + # ========================================================================= + @dataclass - class Options: - # specific literals - OptsDimension = Literal["1D", "2D", "Restelli", "Tokamak"] - # propagator options - nu: float = 1.0 - nu_e: float = 0.01 - eps_norm: float = 1.0 - solver: LiteralOptions.OptsGenSolver = "GMRES" - solver_params: SolverParameters = None - a: float = 1.0 - R0: float = 1.0 - B0: float = 10.0 - Bp: float = 12.0 - alpha: float = 0.1 - beta: float = 1.0 - stab_sigma: float = 1e-5 - variant: LiteralOptions.OptsSaddlePointSolver = "Uzawa" - method_to_solve: LiteralOptions.OptsDirectSolver = "DirectNPInverse" - preconditioner: bool = False - spectralanalysis: bool = False - lifting: bool = False - dimension: OptsDimension = "2D" - D1_dt: float = 1e-3 + class Options(): + + nu: float | None = None + nu_e: float | None = None + eps_norm: float | None = None + + # boundary conditions per species + # supported kinds: "periodic", "dirichlet" + # future: "neumann", "robin" + boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_data_u: dict[tuple[int, int], Callable] | None = None + boundary_data_ue: dict[tuple[int, int], Callable] | None = None + + source_u: Callable | None = None + source_ue: Callable | None = None + + stab_sigma: float | None = None + + solver: LiteralOptions.OptsGenSolver = "gmres" + solver_params: SolverParameters | None = None def __post_init__(self): - # checks - check_option(self.solver, LiteralOptions.OptsGenSolver) - check_option(self.variant, LiteralOptions.OptsSaddlePointSolver) - check_option(self.method_to_solve, LiteralOptions.OptsDirectSolver) - check_option(self.dimension, self.OptsDimension) - # defaults + # --- required parameters --- + assert self.nu is not None, "nu must be specified" + assert self.nu_e is not None, "nu_e must be specified" + assert self.eps_norm is not None, "eps_norm must be specified" + assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" + assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" + + # --- physical parameter sanity checks --- + if self.nu < 0: + raise ValueError(f"nu must be non-negative, got {self.nu}") + if self.nu_e < 0: + raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") + if self.eps_norm <= 0: + raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") + + # --- check all axes are covered --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for d in range(3): + for side in (-1, 1): + assert (d, side) in bcs, \ + f"{name} is missing entry for axis {d} side {side}" + + # --- periodic consistency: periodic must be paired on both sides --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for (d, side), kind in bcs.items(): + if kind == "periodic": + assert bcs.get((d, -side)) == "periodic", \ + f"{name}: axis {d} side {side} is periodic but opposite side is not" + + # --- ions and electrons must agree on which axes are periodic --- + for d in range(3): + u_left = self.boundary_conditions_u.get((d, -1)) + ue_left = self.boundary_conditions_ue.get((d, -1)) + u_right = self.boundary_conditions_u.get((d, 1)) + ue_right = self.boundary_conditions_ue.get((d, 1)) + u_periodic = (u_left == "periodic") + ue_periodic = (ue_left == "periodic") + if u_periodic != ue_periodic: + raise ValueError( + f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " + f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" + ) + + # --- warn for Dirichlet faces with no boundary data --- + for species, bcs, data, label in [ + ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), + ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), + ]: + has_dirichlet = any(v == "dirichlet" for v in bcs.values()) + if has_dirichlet: + if data is None: + warn(f"Dirichlet BCs specified for {species} but no {label} given " + f"— defaulting to homogeneous Dirichlet on all faces.") + else: + for (d, side), kind in bcs.items(): + if kind == "dirichlet" and (d, side) not in data: + warn(f"No {label} given for axis {d} side {side} " + f"— defaulting to homogeneous Dirichlet.") + + # --- warn if no source terms --- + if self.source_u is None: + warn("No source_u specified — defaulting to zero.") + if self.source_ue is None: + warn("No source_ue specified — defaulting to zero.") + + # --- defaults --- + if self.stab_sigma is None: + warn("stab_sigma not specified, defaulting to 0.0") + self.stab_sigma = 0.0 + + check_option(self.solver, LiteralOptions.OptsGenSolver) if self.solver_params is None: self.solver_params = SolverParameters() @property def options(self) -> Options: - if not hasattr(self, "_options"): - self._options = self.Options() + assert hasattr(self, "_options"), "Options not set." return self._options @options.setter def options(self, new): assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") self._options = new - @profile - def allocate(self, verbose: bool = False): - self._info = self.options.solver_params.info - if self.derham.comm is not None: - self._rank = self.derham.comm.Get_rank() - else: - self._rank = 0 + # ========================================================================= + ### Boundary condition helpers + # ========================================================================= - self._nu = self.options.nu - self._nu_e = self.options.nu_e - self._eps_norm = self.options.eps_norm - self._a = self.options.a - self._R0 = self.options.R0 - self._B0 = self.options.B0 - self._Bp = self.options.Bp - self._alpha = self.options.alpha - self._beta = self.options.beta - self._stab_sigma = self.options.stab_sigma - self._variant = self.options.variant - self._method_to_solve = self.options.method_to_solve - self._preconditioner = self.options.preconditioner - self._dimension = self.options.dimension - self._spectralanalysis = self.options.spectralanalysis - self._lifting = self.options.lifting + def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): - solver_params = self.options.solver_params + dirichlet_bc = [] + for d in range(3): + if spl_kind[d]: # periodic spline — no clamping ever + dirichlet_bc.append((False, False)) + else: + left = boundary_conditions.get((d, -1)) == "dirichlet" + right = boundary_conditions.get((d, 1)) == "dirichlet" + dirichlet_bc.append((left, right)) + return tuple(tuple(bc) for bc in dirichlet_bc) + + def _apply_boundary_conditions(self, vec, boundary_conditions): + """Zero out Dirichlet DOFs on the given stencil vector.""" + for (d, side), kind in boundary_conditions.items(): + if kind == "dirichlet": + apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) + # future: neumann and robin require no zeroing here + + # ========================================================================= + ### Allocate + # ========================================================================= + + def allocate(self, verbose=False): + + self.verbose = verbose + self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 + self._dt = None + + # ---- constrained (v0) de Rham complex -------------------------------- + + _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) + _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) + _dirichlet_bc = tuple( + (l_u or l_ue, r_u or r_ue) + for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) + ) - u = self.variables.u.spline.vector + self._derham_v0 = Derham( + self.derham.Nel, self.derham.p, self.derham.spl_kind, + domain=self.domain, dirichlet_bc=_dirichlet_bc, + ) + self._mass_ops_v0 = WeightedMassOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.mass_ops.weights["eq_mhd"], + ) + self._basis_ops_v0 = BasisProjectionOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.basis_ops.weights["eq_mhd"], + ) - # Lifting for nontrivial boundary conditions - # derham had boundary conditions in eta1 direction, the following is in space Hdiv_0 - if self._lifting: - self.derhamv0 = Derham( - self.derham.Nel, - self.derham.p, - self.derham.spl_kind, - domain=self.domain, - dirichlet_bc=((True, True), (False, False), (False, False)), - ) + # ---- unconstrained operators (for RHS assembly) ---------------------- - self._mass_opsv0 = WeightedMassOperators( - self.derhamv0, - self.domain, - verbose=solver_params.verbose, - eq_mhd=self.mass_ops.weights["eq_mhd"], - ) - self._basis_opsv0 = BasisProjectionOperators( - self.derhamv0, - self.domain, - verbose=solver_params.verbose, - eq_mhd=self.basis_ops.weights["eq_mhd"], - ) - else: - self.derhamnumpy = Derham( - self.derham.Nel, - self.derham.p, - self.derham.spl_kind, - domain=self.domain, - # dirichlet_bc=self.derham.dirichlet_bc, - # nquads = self.derham._nquads, - # nq_pr = self.derham._nq_pr, - # comm = MPI.COMM_SELF, # self.derham._comm, - # polar_ck= self.derham._polar_ck, - # local_projectors=self.derham.with_local_projectors - ) - # get forceterms for according dimension - if self._dimension in ["2D", "1D"]: - ### Manufactured solution ### - _forceterm_logical = lambda e1, e2, e3: 0 * e1 - _funx = getattr(callables, "ManufacturedSolutionForceterm")( - species="Ions", - comp="0", - b0=self._B0, - nu=self._nu, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - ) - _funy = getattr(callables, "ManufacturedSolutionForceterm")( - species="Ions", - comp="1", - b0=self._B0, - nu=self._nu, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - ) - _funelectronsx = getattr(callables, "ManufacturedSolutionForceterm")( - species="Electrons", - comp="0", - b0=self._B0, - nu_e=self._nu_e, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - ) - _funelectronsy = getattr(callables, "ManufacturedSolutionForceterm")( - species="Electrons", - comp="1", - b0=self._B0, - nu_e=self._nu_e, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - ) + self._M2 = self.mass_ops.M2 + self._M2B = - self.mass_ops.M2B + self._div = self.derham.div + self._curl = self.derham.curl + self._S21 = self.basis_ops.S21 - # get callable(s) for specified init type - forceterm_class = [_funx, _funy, _forceterm_logical] - forcetermelectrons_class = [_funelectronsx, _funelectronsy, _forceterm_logical] - - # pullback callable - funx = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - funy = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - fun_electronsx = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - fun_electronsy = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - self._F1 = l2_proj([funx, funy, _forceterm_logical]) - self._F2 = l2_proj([fun_electronsx, fun_electronsy, _forceterm_logical]) - - elif self._dimension == "Restelli": - ### Restelli ### - - _forceterm_logical = lambda e1, e2, e3: 0 * e1 - _fun = getattr(callables, "RestelliForcingTerm")( - B0=self._B0, - nu=self._nu, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - eps=self._eps_norm, - ) - _funelectrons = getattr(callables, "RestelliForcingTerm")( - B0=self._B0, - nu=self._nu_e, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - eps=self._eps_norm, - ) + self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div + + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + + self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl + self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) + + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) - # get callable(s) for specified init type - forceterm_class = [_forceterm_logical, _forceterm_logical, _fun] - forcetermelectrons_class = [_forceterm_logical, _forceterm_logical, _funelectrons] - - # pullback callable - fun_pb_1 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - fun_pb_2 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - fun_pb_3 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=2, - domain=self.domain, - ) - fun_electrons_pb_1 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - fun_electrons_pb_2 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - fun_electrons_pb_3 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=2, - domain=self.domain, - ) - if self._lifting: - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) - else: - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - self._F1 = l2_proj([fun_pb_1, fun_pb_2, fun_pb_3], apply_bc=self._lifting) - self._F2 = l2_proj([fun_electrons_pb_1, fun_electrons_pb_2, fun_electrons_pb_3], apply_bc=self._lifting) - - ### End Restelli ### - - elif self._dimension == "Tokamak": - ### Tokamak geometry curl-free manufactured solution ### - - _forceterm_logical = lambda e1, e2, e3: 0 * e1 - _funx = getattr(callables, "ManufacturedSolutionForceterm")( - species="Ions", - comp="0", - b0=self._B0, - nu=self._nu, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) - _funy = getattr(callables, "ManufacturedSolutionForceterm")( - species="Ions", - comp="1", - b0=self._B0, - nu=self._nu, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) - _funz = getattr(callables, "ManufacturedSolutionForceterm")( - species="Ions", - comp="2", - b0=self._B0, - nu=self._nu, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) - _funelectronsx = getattr(callables, "ManufacturedSolutionForceterm")( - species="Electrons", - comp="0", - b0=self._B0, - nu_e=self._nu_e, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) - _funelectronsy = getattr(callables, "ManufacturedSolutionForceterm")( - species="Electrons", - comp="1", - b0=self._B0, - nu_e=self._nu_e, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) - _funelectronsz = getattr(callables, "ManufacturedSolutionForceterm")( - species="Electrons", - comp="2", - b0=self._B0, - nu_e=self._nu_e, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) + # ---- constrained operators (for system matrix) ----------------------- - # get callable(s) for specified init type - forceterm_class = [_funx, _funy, _funz] - forcetermelectrons_class = [_funelectronsx, _funelectronsy, _funelectronsz] - - # pullback callable - fun_pb_1 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - fun_pb_2 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - fun_pb_3 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=2, - domain=self.domain, - ) - fun_electrons_pb_1 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - fun_electrons_pb_2 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - fun_electrons_pb_3 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=2, - domain=self.domain, - ) - if self._lifting: - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) - else: - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - self._F1 = l2_proj([fun_pb_1, fun_pb_2, fun_pb_3], apply_bc=self._lifting) - self._F2 = l2_proj([fun_electrons_pb_1, fun_electrons_pb_2, fun_electrons_pb_3], apply_bc=self._lifting) - - ### End Tokamak geometry manufactured solution ### - - if self._variant == "GMRES": - if self._lifting: - self._M2 = getattr(self._mass_opsv0, "M2") - self._M3 = getattr(self._mass_opsv0, "M3") - self._M2B = -getattr(self._mass_opsv0, "M2B") - self._div = self.derhamv0.div - self._curl = self.derhamv0.curl - self._S21 = self._basis_opsv0.S21 - else: - self._M2 = getattr(self.mass_ops, "M2") - self._M3 = getattr(self.mass_ops, "M3") - self._M2B = -getattr(self.mass_ops, "M2B") - self._div = self.derham.div - self._curl = self.derham.curl - self._S21 = self.basis_ops.S21 - - # Define block matrix [[A BT], [B 0]] (without time step size dt in the diagonals) - _A11 = ( - self._M2 - - self._M2B / self._eps_norm - + self._nu - * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - ) - _A12 = None - _A21 = _A12 - _A22 = ( - -self._stab_sigma * IdentityOperator(_A11.domain) - + self._M2B / self._eps_norm - + self._nu_e - * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - ) - _B1 = -self._M3 @ self._div - _B2 = self._M3 @ self._div - - if _A12 is not None: - assert _A11.codomain == _A12.codomain - if _A21 is not None: - assert _A22.codomain == _A21.codomain - assert _B1.codomain == _B2.codomain - if _A12 is not None: - assert _A11.domain == _A12.domain == _B1.domain - if _A21 is not None: - assert _A21.domain == _A22.domain == _B2.domain - assert _A22.domain == _B2.domain - assert _A11.domain == _B1.domain - - self._block_domainA = BlockVectorSpace(_A11.domain, _A22.domain) - self._block_codomainA = self._block_domainA - self._block_domainB = self._block_domainA - self._block_codomainB = _B2.codomain - _blocksA = [[_A11, _A12], [_A21, _A22]] - _A = BlockLinearOperator(self._block_domainA, self._block_codomainA, blocks=_blocksA) - _blocksB = [[_B1, _B2]] - _B = BlockLinearOperator(self._block_domainB, self._block_codomainB, blocks=_blocksB) - _F = BlockVector(self._block_domainA, blocks=[self._F1, self._F2]) # missing M2/dt *un-1 - - elif self._variant == "Uzawa": - # Numpy - if self._lifting: - fun = [] - for m in range(3): - fun += [[]] - for n in range(3): - fun[-1] += [ - lambda e1, e2, e3, m=m, n=n: ( - self._basis_opsv0.G(e1, e2, e3)[:, :, :, m, n] / self._basis_opsv0.sqrt_g(e1, e2, e3) - ), - ] - self._S21 = None - if self.derhamv0.with_local_projectors: - self._S21 = BasisProjectionOperatorLocal( - self.derhamv0._Ploc["1"], - self.derhamv0.Vh_fem["2"], - fun, - transposed=False, - ) + self._M2_v0 = self._mass_ops_v0.M2 + self._M3_v0 = self._mass_ops_v0.M3 + self._M2B_v0 = - self._mass_ops_v0.M2B + self._div_v0 = self._derham_v0.div + self._curl_v0 = self._derham_v0.curl + self._S21_v0 = self._basis_ops_v0.S21 - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - Vbc = self._mass_opsv0.M2._V_boundary_op.toarray_struphy() - Wbc = self._mass_opsv0.M2._W_boundary_op.toarray_struphy() - M2_mat = self._mass_opsv0.M2._mat.toarray() - self._M2np = Wbc @ M2_mat @ Vbc.T - Vbc = self._mass_opsv0.M3._V_boundary_op.toarray_struphy() - Wbc = self._mass_opsv0.M3._W_boundary_op.toarray_struphy() - M3_mat = self._mass_opsv0.M3._mat.toarray() - self._M3np = Wbc @ M3_mat @ Vbc.T - if isinstance(self.derhamv0.div, ComposedLinearOperator): - for mult in self.derhamv0.div.multiplicants: - if isinstance(mult, BlockLinearOperator): - if hasattr(self, "_Dnp"): - self._Dnp = self._Dnp @ mult.toarray() - else: - self._Dnp = mult.toarray() - # print(f"{type(mult.toarray())=}") #with_pads = True - elif isinstance(mult, BoundaryOperator): - if hasattr(self, "_Dnp"): - self._Dnp = self._Dnp @ mult.T.toarray_struphy() - else: - self._Dnp = mult.toarray_struphy() - elif isinstance(self.derhamv0.div, BlockLinearOperator): - self._Dnp = self.derhamv0.div.toarray() - if isinstance(self.derhamv0.curl, ComposedLinearOperator): - for mult in self.derhamv0.curl.multiplicants: - if isinstance(mult, BlockLinearOperator): - if hasattr(self, "_Cnp"): - self._Cnp = self._Cnp @ mult.toarray() - else: - self._Cnp = mult.toarray() - elif isinstance(mult, BoundaryOperator): - if hasattr(self, "_Cnp"): - self._Cnp = self._Cnp @ mult.T.toarray_struphy() - else: - self._Cnp = mult.toarray_struphy() - elif isinstance(self.derhamv0.curl, BlockLinearOperator): - self._Dnp = self.derhamv0.curl.toarray() - - if self._S21 is not None: - self._Hodgenp = self._S21.toarray - else: - self._Hodgenp = self._basis_opsv0.S21.toarray_struphy() # self.basis_ops.S21.toarray - Vbc = self._mass_opsv0.M2B._V_boundary_op.toarray_struphy() - Wbc = self._mass_opsv0.M2B._W_boundary_op.toarray_struphy() - M2B_mat = -self._mass_opsv0.M2B._mat.toarray() # - sign because of the definition of M2B - self._M2Bnp = Wbc @ M2B_mat @ Vbc.T - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - Vbc = self._mass_opsv0.M2._V_boundary_op.toarray_struphy(is_sparse=True) - Wbc = self._mass_opsv0.M2._W_boundary_op.toarray_struphy(is_sparse=True) - M2_mat = self._mass_opsv0.M2._mat.tosparse() - self._M2np = Wbc @ M2_mat @ Vbc.T - Vbc = self._mass_opsv0.M3._V_boundary_op.toarray_struphy(is_sparse=True) - Wbc = self._mass_opsv0.M3._W_boundary_op.toarray_struphy(is_sparse=True) - M3_mat = self._mass_opsv0.M3._mat.tosparse() - self._M3np = Wbc @ M3_mat @ Vbc.T - if self._S21 is not None: - self._Hodgenp = self._S21.tosparse - else: - self._Hodgenp = self._basis_opsv0.S21.toarray_struphy(is_sparse=True) - Vbc = self._mass_opsv0.M2B._V_boundary_op.toarray_struphy(is_sparse=True) - Wbc = self._mass_opsv0.M2B._W_boundary_op.toarray_struphy(is_sparse=True) - M2B_mat = self._mass_opsv0.M2B._mat.tosparse() - self._M2Bnp = -Wbc @ M2B_mat @ Vbc.T # - sign because of the definition of M2B - - if isinstance(self.derhamv0.div, ComposedLinearOperator): - for mult in self.derhamv0.div.multiplicants: - if isinstance(mult, BlockLinearOperator): - if hasattr(self, "_Dnp"): - self._Dnp = self._Dnp @ mult.tosparse() - else: - self._Dnp = mult.tosparse() - elif isinstance(mult, BoundaryOperator): - if hasattr(self, "_Dnp"): - self._Dnp = self._Dnp @ mult.toarray_struphy(is_sparse=True) - else: - self._Dnp = mult.toarray_struphy(is_sparse=True) - elif isinstance(self.derhamv0.div, BlockLinearOperator): - self._Dnp = self.derhamv0.div.tosparse() - - if isinstance(self.derhamv0.curl, ComposedLinearOperator): - for mult in self.derhamv0.curl.multiplicants: - if isinstance(mult, BlockLinearOperator): - if hasattr(self, "_Cnp"): - self._Cnp = self._Cnp @ mult.tosparse() - else: - self._Cnp = mult.tosparse() - elif isinstance(mult, BoundaryOperator): - if hasattr(self, "_Cnp"): - self._Cnp = self._Cnp @ mult.toarray_struphy(is_sparse=True) - else: - self._Cnp = mult.toarray_struphy(is_sparse=True) - elif isinstance(self.derhamv0.curl, BlockLinearOperator): - self._Dnp = self.derhamv0.curl.tosparse() - - else: # no lifting, use original Derham - fun = [] - for m in range(3): - fun += [[]] - for n in range(3): - fun[-1] += [ - lambda e1, e2, e3, m=m, n=n: ( - self.basis_ops.G(e1, e2, e3)[:, :, :, m, n] / self.basis_ops.sqrt_g(e1, e2, e3) - ), - ] - self._S21 = None - if self.derham.with_local_projectors: - self._S21 = BasisProjectionOperatorLocal( - self.derham._Ploc["1"], - self.derham.Vh_fem["2"], - fun, - transposed=False, - ) + self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 + + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) + + self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 + self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) + + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - Vbc = self.mass_ops.M2._V_boundary_op.toarray_struphy() - Wbc = self.mass_ops.M2._W_boundary_op.toarray_struphy() - M2_mat = self.mass_ops.M2._mat.toarray() - self._M2np = Wbc @ M2_mat @ Vbc.T - Vbc = self.mass_ops.M3._V_boundary_op.toarray_struphy() - Wbc = self.mass_ops.M3._W_boundary_op.toarray_struphy() - M3_mat = self.mass_ops.M3._mat.toarray() - self._M3np = Wbc @ M3_mat @ Vbc.T - self._Dnp = self.derhamnumpy.div.toarray() - self._Cnp = self.derhamnumpy.curl.toarray() - - if self._S21 is not None: - self._Hodgenp = self._S21.toarray - else: - self._Hodgenp = self.basis_ops.S21.toarray_struphy() - Vbc = self.mass_ops.M2B._V_boundary_op.toarray_struphy() - Wbc = self.mass_ops.M2B._W_boundary_op.toarray_struphy() - M2B_mat = -self.mass_ops.M2B._mat.toarray() - self._M2Bnp = Wbc @ M2B_mat @ Vbc.T - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - self._M2np = self.mass_ops.M2.tosparse - self._M3np = self.mass_ops.M3.tosparse - if self._S21 is not None: - self._Hodgenp = self._S21.tosparse - else: - self._Hodgenp = self.basis_ops.S21.toarray_struphy(is_sparse=True) - self._M2Bnp = -self.mass_ops.M2B.tosparse + # ---- block saddle-point system ---------------------------------------- - self._Dnp = self.derhamnumpy.div.tosparse() - self._Cnp = self.derhamnumpy.curl.tosparse() + self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) + self._block_codomain_v0 = self._block_domain_v0 + self._block_codomain_B_v0 = self._derham_v0.Vh["3"] - self._A11np_notimedependency = ( - self._nu - * ( - self._Dnp.T @ self._M3np @ self._Dnp - + 1.0 * self._Hodgenp.T @ self._Cnp.T @ self._M2np @ self._Cnp @ self._Hodgenp - ) - - 1.0 * self._M2Bnp / self._eps_norm - ) - A11np = self._M2np + self._A11np_notimedependency - - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - A11np += self._stab_sigma * xp.identity(A11np.shape[0]) - self.A22np = ( - self._stab_sigma * xp.identity(A11np.shape[0]) - + self._nu_e - * ( - self._Dnp.T @ self._M3np @ self._Dnp - + self._Hodgenp.T @ self._Cnp.T @ self._M2np @ self._Cnp @ self._Hodgenp - ) - + self._M2Bnp / self._eps_norm - ) - self._A22prenp = ( - xp.identity(self.A22np.shape[0]) * self._stab_sigma - ) # + self._nu_e * (self._Dnp.T @ self._M3np @ self._Dnp) - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - A11np += self._stab_sigma * sc.sparse.eye(A11np.shape[0], format="csr") - self.A22np = ( - self._stab_sigma * sc.sparse.eye(A11np.shape[0], format="csr") - + self._nu_e - * ( - self._Dnp.T @ self._M3np @ self._Dnp - + self._Hodgenp.T @ self._Cnp.T @ self._M2np @ self._Cnp @ self._Hodgenp - ) - + self._M2Bnp / self._eps_norm - ) - self._A22prenp = self._stab_sigma * sc.sparse.eye(self.A22np.shape[0], format="csr") - - B1np = -self._M3np @ self._Dnp - B2np = self._M3np @ self._Dnp - self._B1np = B1np - self._B2np = B2np - self._F1np = self._F1.toarray() - self._F2np = self._F2.toarray() - _Anp = [A11np, self.A22np] - _Bnp = [B1np, B2np] - _Fnp = [self._F1np, self._F2np] - self._A11prenp_notimedependency = self._nu * (self._Dnp.T @ self._M3np @ self._Dnp) - _A11prenp = self._M2np + self._A11prenp_notimedependency - _Anppre = [_A11prenp, self._A22prenp] - - if self._variant == "GMRES": - self._solver_GMRES = SaddlePointSolver( - A=_A, - B=_B, - F=_F, - solver_name=self.options.solver, - tol=self.options.solver_params.tol, - max_iter=self.options.solver_params.maxiter, - verbose=self.options.solver_params.verbose, - pc=None, - ) - # Allocate memory for call - self._untemp = self.variables.u.spline.vector.space.zeros() - - elif self._variant == "Uzawa": - self._solver_UzawaNumpy = SaddlePointSolver( - Apre=_Anppre, - A=_Anp, - B=_Bnp, - F=_Fnp, - method_to_solve=self._method_to_solve, - preconditioner=self._preconditioner, - spectralanalysis=self.options.spectralanalysis, - tol=self.options.solver_params.tol, - max_iter=self.options.solver_params.maxiter, - verbose=self.options.solver_params.verbose, - ) + self._B1_v0 = - self._M3_v0 @ self._div_v0 + self._B2_v0 = self._M3_v0 @ self._div_v0 + + self._B_v0 = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_B_v0, + blocks=[[self._B1_v0, self._B2_v0]] + ) + + self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) + + _A_init = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0, None], [None, self._A22_v0]] + ) + _M_init = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] + ) + + self._Minv = inverse( + _M_init, self.options.solver, + A11=self._A11_v0, + A22=self._A22_v0, + B1=self._B1_v0, + B2=self._B2_v0, + recycle=self.options.solver_params.recycle, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + ) + + # ---- projector ------------------------------------------------------- + + self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + + # ---- solution spline functions (unconstrained) ----------------------- + + self._u = self.derham.create_spline_function("u", space_id="Hdiv") + self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") + self._phi = self.derham.create_spline_function("phi", space_id="L2") + + # ---- BC lifts (unconstrained) ---------------------------------------- + + self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") + self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") + + for u_prime, boundary_data, boundary_conditions in [ + (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), + (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), + ]: + if boundary_data is None: + continue + for (d, side), f_bc in boundary_data.items(): + if boundary_conditions.get((d, side)) == "dirichlet": + bc_pulled = lambda *etas, f=f_bc: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], + lambda *etas: bc_pulled(*etas)[1], + lambda *etas: bc_pulled(*etas)[2]]) + for (d2, side2), kind2 in boundary_conditions.items(): + if kind2 == "dirichlet" and (d2, side2) != (d, side): + apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) + u_prime.vector += _vec + + self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") + self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") + + self._u_prime_v0.vector = self._u_prime.vector + self._ue_prime_v0.vector = self._ue_prime.vector + + # ---- projected source terms (unconstrained) -------------------------- + + self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") + self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") + + for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: + if source is not None: + src_pulled = lambda *etas, f=source: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], + lambda *etas: src_pulled(*etas)[1], + lambda *etas: src_pulled(*etas)[2]]) + + # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- + + self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") + + # ========================================================================= + ### Time step + # ========================================================================= def __call__(self, dt): - # current variables - unfeec = self.variables.u.spline.vector - uenfeec = self.variables.ue.spline.vector - phinfeec = self.variables.phi.spline.vector - - if self._variant == "GMRES": - if self._lifting: - phinfeeccopy = self.derhamv0.create_spline_function("phi", space_id="L2") - phinfeeccopy.vector = phinfeec - # unfeec in space Hdiv, u0 in space Hdiv_0 - unfeeccopy = self.derhamv0.create_spline_function("u", space_id="Hdiv") - u0 = self.derhamv0.create_spline_function("u", space_id="Hdiv") - u_prime = self.derhamv0.create_spline_function("u", space_id="Hdiv") - u0.vector = uenfeec - unfeeccopy.vector = uenfeec - apply_essential_bc_stencil(u0.vector[0], axis=0, ext=-1, order=0) - apply_essential_bc_stencil(u0.vector[0], axis=0, ext=1, order=0) - u_prime.vector = unfeeccopy.vector - u0.vector - - uenfeeccopy = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue0 = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue_prime = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue0.vector = uenfeec - uenfeeccopy.vector = uenfeec - apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=-1, order=0) - apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=1, order=0) - ue_prime.vector = uenfeeccopy.vector - ue0.vector - - _A11 = ( - self._M2 / dt - - self._M2B / self._eps_norm - + self._nu - * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + + # --- copy current state --- + self._u.vector = self.variables.u.spline.vector + self._ue.vector = self.variables.ue.spline.vector + + # --- rebuild system matrix if dt changed --- + if dt != self._dt: + self._dt = dt + _A = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] ) - _A12 = None - _A21 = _A12 - _A22 = ( - self._nu_e - * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - + self._M2B / self._eps_norm - - self._stab_sigma * IdentityOperator(_A11.domain) + + _M = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A, self._B_v0.T], [self._B_v0, None]] ) + self._Minv.linop = _M + class Variables(): + def __init__(self) -> None: + self._u: FEECVariable | None = None + self._ue: FEECVariable | None = None + self._phi: FEECVariable | None = None - if self._lifting: - _A11prime = -self._M2B / self._eps_norm + self._nu * ( - self.derhamv0.div.T @ self._M3 @ self.derhamv0.div - + self._basis_opsv0.S21.T - @ self.derhamv0.curl.T - @ self._M2 - @ self.derhamv0.curl - @ self._basis_opsv0.S21 - ) - _A22prime = ( - self._nu_e - * ( - self.derhamv0.div.T @ self._M3 @ self.derhamv0.div - + self._basis_opsv0.S21.T - @ self.derhamv0.curl.T - @ self._M2 - @ self.derhamv0.curl - @ self._basis_opsv0.S21 - ) - + self._M2B / self._eps_norm - - self._stab_sigma * IdentityOperator(_A11.domain) - ) - _B1 = -self._M3 @ self._div - _B2 = self._M3 @ self._div - - if _A12 is not None: - assert _A11.codomain == _A12.codomain - if _A21 is not None: - assert _A22.codomain == _A21.codomain - assert _B1.codomain == _B2.codomain - if _A12 is not None: - assert _A11.domain == _A12.domain == _B1.domain - if _A21 is not None: - assert _A21.domain == _A22.domain == _B2.domain - assert _A22.domain == _B2.domain - assert _A11.domain == _B1.domain - - _blocksA = [[_A11, _A12], [_A21, _A22]] - _A = BlockLinearOperator(self._block_domainA, self._block_codomainA, blocks=_blocksA) - _blocksB = [[_B1, _B2]] - _B = BlockLinearOperator(self._block_domainB, self._block_codomainB, blocks=_blocksB) - if self._lifting: - _blocksF = [ - self._M2.dot(self._F1) + self._M2.dot(u0.vector) / dt - _A11prime.dot(u_prime.vector), - self._M2.dot(self._F2) - _A22prime.dot(ue_prime.vector), - ] - else: - _blocksF = [ - self._M2.dot(self._F1) + self._M2.dot(unfeec) / dt, - self._M2.dot(self._F2), - ] - _F = BlockVector(self._block_domainA, blocks=_blocksF) + @property + def u(self) -> FEECVariable | None: + return self._u - # Imported solver - self._solver_GMRES.A = _A - self._solver_GMRES.B = _B - self._solver_GMRES.F = _F + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._u = new - if self._lifting: - ( - _sol1, - _sol2, - info, - ) = self._solver_GMRES(u0.vector, ue0.vector, phinfeeccopy.vector) + @property + def ue(self) -> FEECVariable | None: + return self._ue - un_temp = self.derham.create_spline_function("u", space_id="Hdiv") - un_temp.vector = _sol1[0] + u_prime.vector + @ue.setter + def ue(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._ue = new - uen_temp = self.derham.create_spline_function("ue", space_id="Hdiv") - uen_temp.vector = _sol1[1] + ue_prime.vector + @property + def phi(self) -> FEECVariable | None: + return self._phi - phin_temp = self.derham.create_spline_function("phi", space_id="L2") - phin_temp.vector = _sol2 + @phi.setter + def phi(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._phi = new - un = un_temp.vector - uen = uen_temp.vector - phin = phin_temp.vector + def __init__(self): + self.variables = self.Variables() - else: - ( - _sol1, - _sol2, - info, - ) = self._solver_GMRES(unfeec, uenfeec, phinfeec) - un = _sol1[0] - uen = _sol1[1] - phin = _sol2 - # write new coeffs into self.feec_vars - - max_du, max_due, max_dphi = self.update_feec_variables(u=un, ue=uen, phi=phin) - - elif self._variant == "Uzawa": - # Numpy - A11np = self._M2np / dt + self._A11np_notimedependency - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - A11np += self._stab_sigma * xp.identity(A11np.shape[0]) - _A22prenp = self._A22prenp - A22np = self.A22np - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - A11np += self._stab_sigma * sc.sparse.eye(A11np.shape[0], format="csr") - _A22prenp = self._A22prenp - A22np = self.A22np - - # _Anp[1] and _Anppre[1] remain unchanged - _Anp = [A11np, A22np] - if self._preconditioner: - _A11prenp = self._M2np / dt # + self._A11prenp_notimedependency - _Anppre = [_A11prenp, _A22prenp] - - if self._lifting: - # unfeec in space Hdiv, u0 in space Hdiv_0 - unfeeccopy = self.derhamv0.create_spline_function("u", space_id="Hdiv") - u0 = self.derhamv0.create_spline_function("u", space_id="Hdiv") - u_prime = self.derham.create_spline_function("u", space_id="Hdiv") - u0.vector = unfeec - unfeeccopy.vector = unfeec - apply_essential_bc_stencil(u0.vector[0], axis=0, ext=-1, order=0) - apply_essential_bc_stencil(u0.vector[0], axis=0, ext=1, order=0) - u_prime.vector = unfeeccopy.vector - u0.vector - - uenfeeccopy = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue0 = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue_prime = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue0.vector = uenfeec - uenfeeccopy.vector = uenfeec - apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=-1, order=0) - apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=1, order=0) - ue_prime.vector = uenfeeccopy.vector - ue0.vector - - _F1np = ( - self._M2np @ self._F1np - + 1.0 / dt * self._M2np.dot(u0.vector.toarray()) - - self._A11np_notimedependency.dot(u_prime.vector.toarray()) - ) - _F2np = self._M2np @ self._F2np - self.A22np.dot(ue_prime.vector.toarray()) - _Fnp = [_F1np, _F2np] - else: - _F1np = self._M2np @ self._F1np + 1.0 / dt * self._M2np.dot(unfeec.toarray()) - _F2np = self._M2np @ self._F2np - _Fnp = [_F1np, _F2np] - - if self.rank == 0: - if self._preconditioner: - self._solver_UzawaNumpy.Apre = _Anppre - self._solver_UzawaNumpy.A = _Anp - self._solver_UzawaNumpy.F = _Fnp - if self._lifting: - un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy( - u0.vector, - ue0.vector, - phinfeec, - ) + # ========================================================================= + ### Options + # ========================================================================= + + @dataclass + class Options(): + + nu: float | None = None + nu_e: float | None = None + eps_norm: float | None = None + + # boundary conditions per species + # supported kinds: "periodic", "dirichlet" + # future: "neumann", "robin" + boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_data_u: dict[tuple[int, int], Callable] | None = None + boundary_data_ue: dict[tuple[int, int], Callable] | None = None + + source_u: Callable | None = None + source_ue: Callable | None = None - un += u_prime.vector.toarray() - uen += ue_prime.vector.toarray() - else: - un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy( - unfeec, - uenfeec, - phinfeec, + stab_sigma: float | None = None + + solver: LiteralOptions.OptsGenSolver = "gmres" + solver_params: SolverParameters | None = None + + def __post_init__(self): + + # --- required parameters --- + assert self.nu is not None, "nu must be specified" + assert self.nu_e is not None, "nu_e must be specified" + assert self.eps_norm is not None, "eps_norm must be specified" + assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" + assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" + + # --- physical parameter sanity checks --- + if self.nu < 0: + raise ValueError(f"nu must be non-negative, got {self.nu}") + if self.nu_e < 0: + raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") + if self.eps_norm <= 0: + raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") + + # --- check all axes are covered --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for d in range(3): + for side in (-1, 1): + assert (d, side) in bcs, \ + f"{name} is missing entry for axis {d} side {side}" + + # --- periodic consistency: periodic must be paired on both sides --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for (d, side), kind in bcs.items(): + if kind == "periodic": + assert bcs.get((d, -side)) == "periodic", \ + f"{name}: axis {d} side {side} is periodic but opposite side is not" + + # --- ions and electrons must agree on which axes are periodic --- + for d in range(3): + u_left = self.boundary_conditions_u.get((d, -1)) + ue_left = self.boundary_conditions_ue.get((d, -1)) + u_right = self.boundary_conditions_u.get((d, 1)) + ue_right = self.boundary_conditions_ue.get((d, 1)) + u_periodic = (u_left == "periodic") + ue_periodic = (ue_left == "periodic") + if u_periodic != ue_periodic: + raise ValueError( + f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " + f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" ) - dimlist = [[shp - 2 * pi for shp, pi in zip(unfeec[i][:].shape, self.derham.p)] for i in range(3)] - dimphi = [shp - 2 * pi for shp, pi in zip(phinfeec[:].shape, self.derham.p)] - u_temp = BlockVector(self.derham.Vh["2"]) - ue_temp = BlockVector(self.derham.Vh["2"]) - phi_temp = StencilVector(self.derham.Vh["3"]) - test = 0 - for i, bl in enumerate(u_temp.blocks): - s = bl.starts - e = bl.ends - totaldim = dimlist[i][0] * dimlist[i][1] * dimlist[i][2] - test += totaldim - bl[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = un[ - i * totaldim : (i + 1) * totaldim - ].reshape(*dimlist[i]) - - for i, bl in enumerate(ue_temp.blocks): - s = bl.starts - e = bl.ends - totaldim = dimlist[i][0] * dimlist[i][1] * dimlist[i][2] - bl[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = uen[ - i * totaldim : (i + 1) * totaldim - ].reshape(*dimlist[i]) - - s = phi_temp.starts - e = phi_temp.ends - phi_temp[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = phin.reshape(*dimphi) + # --- warn for Dirichlet faces with no boundary data --- + for species, bcs, data, label in [ + ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), + ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), + ]: + has_dirichlet = any(v == "dirichlet" for v in bcs.values()) + if has_dirichlet: + if data is None: + warn(f"Dirichlet BCs specified for {species} but no {label} given " + f"— defaulting to homogeneous Dirichlet on all faces.") + else: + for (d, side), kind in bcs.items(): + if kind == "dirichlet" and (d, side) not in data: + warn(f"No {label} given for axis {d} side {side} " + f"— defaulting to homogeneous Dirichlet.") + + # --- warn if no source terms --- + if self.source_u is None: + warn("No source_u specified — defaulting to zero.") + if self.source_ue is None: + warn("No source_ue specified — defaulting to zero.") + + # --- defaults --- + if self.stab_sigma is None: + warn("stab_sigma not specified, defaulting to 0.0") + self.stab_sigma = 0.0 + + check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver) + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + assert hasattr(self, "_options"), "Options not set." + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + # ========================================================================= + ### Boundary condition helpers + # ========================================================================= + + def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): + + dirichlet_bc = [] + for d in range(3): + if spl_kind[d]: # periodic spline — no clamping ever + dirichlet_bc.append((False, False)) else: - print("TwoFluidQuasiNeutralFull is only running on one MPI.") + left = boundary_conditions.get((d, -1)) == "dirichlet" + right = boundary_conditions.get((d, 1)) == "dirichlet" + dirichlet_bc.append((left, right)) + return tuple(tuple(bc) for bc in dirichlet_bc) + + def _apply_boundary_conditions(self, vec, boundary_conditions): + """Zero out Dirichlet DOFs on the given stencil vector.""" + for (d, side), kind in boundary_conditions.items(): + if kind == "dirichlet": + apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) + # future: neumann and robin require no zeroing here + + # ========================================================================= + ### Allocate + # ========================================================================= + + def allocate(self, verbose=False): + + self.verbose = verbose + self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 + self._dt = None + + # ---- constrained (v0) de Rham complex -------------------------------- + + _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) + _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) + _dirichlet_bc = tuple( + (l_u or l_ue, r_u or r_ue) + for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) + ) - # write new coeffs into self.feec_vars - max_du, max_due, max_dphi = self.update_feec_variables(u=u_temp, ue=ue_temp, phi=phi_temp) + self._derham_v0 = Derham( + self.derham.Nel, self.derham.p, self.derham.spl_kind, + domain=self.domain, dirichlet_bc=_dirichlet_bc, + ) + self._mass_ops_v0 = WeightedMassOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.mass_ops.weights["eq_mhd"], + ) + self._basis_ops_v0 = BasisProjectionOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.basis_ops.weights["eq_mhd"], + ) - if self._info and self._rank == 0: - print("Status for TwoFluidQuasiNeutralFull:", info["success"]) - print("Iterations for TwoFluidQuasiNeutralFull:", info["niter"]) - print("Maxdiff u for TwoFluidQuasiNeutralFull:", max_du) - print("Maxdiff u_e for TwoFluidQuasiNeutralFull:", max_due) - print("Maxdiff phi for TwoFluidQuasiNeutralFull:", max_dphi) - print() + # ---- unconstrained operators (for RHS assembly) ---------------------- + + + self._M2 = self.mass_ops.M2 + self._M2B = - self.mass_ops.M2B + self._div = self.derham.div + self._curl = self.derham.curl + self._S21 = self.basis_ops.S21 + + self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div + + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + + self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl + self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) + + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) + + # ---- constrained operators (for system matrix) ----------------------- + + self._M2_v0 = self._mass_ops_v0.M2 + self._M3_v0 = self._mass_ops_v0.M3 + self._M2B_v0 = - self._mass_ops_v0.M2B + self._div_v0 = self._derham_v0.div + self._curl_v0 = self._derham_v0.curl + self._S21_v0 = self._basis_ops_v0.S21 + + self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 + + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) + + self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 + self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) + + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) + + # ---- block saddle-point system ---------------------------------------- + + self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) + self._block_codomain_v0 = self._block_domain_v0 + self._block_codomain_B_v0 = self._derham_v0.Vh["3"] + + self._B1_v0 = - self._M3_v0 @ self._div_v0 + self._B2_v0 = self._M3_v0 @ self._div_v0 + + self._B_v0 = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_B_v0, + blocks=[[self._B1_v0, self._B2_v0]] + ) + + self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) + + _A_init = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0, None], [None, self._A22_v0]] + ) + _M_init = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] + ) + + if self.options.solver in get_args(LiteralOptions.OptsSaddlePointSolver): + self._Minv = inverse( + _M_init, self.options.solver, + A11=self._A11_v0, + A22=self._A22_v0, + B1=self._B1_v0, + B2=self._B2_v0, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + recycle=self.options.solver_params.recycle, + ) + else: + self._Minv = inverse( + _M_init, self.options.solver, + recycle=self.options.solver_params.recycle, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + ) + + # ---- projector ------------------------------------------------------- + + self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + + # ---- solution spline functions (unconstrained) ----------------------- + + self._u = self.derham.create_spline_function("u", space_id="Hdiv") + self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") + self._phi = self.derham.create_spline_function("phi", space_id="L2") + + # ---- BC lifts (unconstrained) ---------------------------------------- + + self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") + self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") + + for u_prime, boundary_data, boundary_conditions in [ + (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), + (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), + ]: + if boundary_data is None: + continue + for (d, side), f_bc in boundary_data.items(): + if boundary_conditions.get((d, side)) == "dirichlet": + bc_pulled = lambda *etas, f=f_bc: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], + lambda *etas: bc_pulled(*etas)[1], + lambda *etas: bc_pulled(*etas)[2]]) + for (d2, side2), kind2 in boundary_conditions.items(): + if kind2 == "dirichlet" and (d2, side2) != (d, side): + apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) + u_prime.vector += _vec + + self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") + self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") + + self._u_prime_v0.vector = self._u_prime.vector + self._ue_prime_v0.vector = self._ue_prime.vector + + # ---- projected source terms (unconstrained) -------------------------- + + self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") + self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") + + for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: + if source is not None: + src_pulled = lambda *etas, f=source: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], + lambda *etas: src_pulled(*etas)[1], + lambda *etas: src_pulled(*etas)[2]]) + + # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- + + self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") + + # ========================================================================= + ### Time step + # ========================================================================= + + def __call__(self, dt): + + # --- copy current state --- + self._u.vector = self.variables.u.spline.vector + self._ue.vector = self.variables.ue.spline.vector + + # --- rebuild system matrix if dt changed --- TODO change uzawa internals + if dt != self._dt: + self._dt = dt + _A = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] + ) + + _M = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A, self._B_v0.T], [self._B_v0, None]] + ) + self._Minv.linop = _M + + # --- assemble RHS in unconstrained space, then zero boundary DOFs --- + # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' + # electron: F2 = rhs_ue - A22 * ue' + self._rhs_vec_u.vector = (self._rhs_u.vector + + self._M2.dot(self._u.vector) / dt + - self._A11.dot(self._u_prime.vector) + - self._M2.dot(self._u_prime.vector) / dt) + self._rhs_vec_ue.vector = (self._rhs_ue.vector + - self._A22.dot(self._ue_prime.vector)) + + self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) + self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) + + # --- build block RHS and solve --- + _F = BlockVector(self._block_domain_v0, + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) + _RHS = BlockVector(self._block_domain_M, + blocks=[_F, self._block_codomain_B_v0.zeros()]) + + _sol = self._Minv.dot(_RHS) + info = self._Minv.get_info() + + # --- reconstruct full solution: u = u_0 + u' --- + self._u.vector = _sol[0][0] + self._u_prime_v0.vector + self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector + self._phi.vector = _sol[1] + + # --- update FEEC variables --- + max_diffs = self.update_feec_variables( + u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector + ) + + if self.options.solver_params.info and self._rank == 0: + print(f"Status: {info['success']}, Iterations: {info['niter']}") + print(f"Max diffs: {max_diffs}") + # --- assemble RHS in unconstrained space, then zero boundary DOFs --- + # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' + # electron: F2 = rhs_ue - A22 * ue' + self._rhs_vec_u.vector = (self._rhs_u.vector + + self._M2.dot(self._u.vector) / dt + - self._A11.dot(self._u_prime.vector) + - self._M2.dot(self._u_prime.vector) / dt) + self._rhs_vec_ue.vector = (self._rhs_ue.vector + - self._A22.dot(self._ue_prime.vector)) + + self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) + self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) + + # --- build block RHS and solve --- + _F = BlockVector(self._block_domain_v0, + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) + _RHS = BlockVector(self._block_domain_M, + blocks=[_F, self._block_codomain_B_v0.zeros()]) + + _sol = self._Minv.dot(_RHS) + info = self._Minv.get_info() + + # --- reconstruct full solution: u = u_0 + u' --- + self._u.vector = _sol[0][0] + self._u_prime_v0.vector + self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector + self._phi.vector = _sol[1] + + # --- update FEEC variables --- + max_diffs = self.update_feec_variables( + u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector + ) + + if self.options.solver_params.info and self._rank == 0: + print(f"Status: {info['success']}, Iterations: {info['niter']}") + print(f"Max diffs: {max_diffs}") \ No newline at end of file diff --git a/src/struphy/propagators/propagators_fields_two_fluid_new.py b/src/struphy/propagators/propagators_fields_two_fluid_new.py deleted file mode 100644 index c5cf54039..000000000 --- a/src/struphy/propagators/propagators_fields_two_fluid_new.py +++ /dev/null @@ -1,440 +0,0 @@ -from dataclasses import dataclass -from typing import Callable, Literal, cast -from warnings import warn - -from feectools.api.essential_bc import apply_essential_bc_stencil -from feectools.ddm.mpi import mpi as MPI -from feectools.linalg.basic import IdentityOperator, InverseLinearOperator -from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace -from feectools.linalg.solvers import inverse - -from struphy.feec.basis_projection_ops import BasisProjectionOperators -from struphy.feec.mass import WeightedMassOperators -from struphy.feec.projectors import L2Projector -from struphy.feec.psydac_derham import Derham -from struphy.io.options import OptsGenSolver -from struphy.linear_algebra.solver import SolverParameters -from struphy.models.variables import FEECVariable -from struphy.propagators.base import Propagator -from struphy.utils.utils import check_option - - -class TwoFluidQuasiNeutralFull(Propagator): - r""":ref:`FEEC ` discretization of the following equations: - find :math:`\mathbf u \in H(\textnormal{div})`, :math:`\mathbf u_e \in H(\textnormal{div})` and :math:`\mathbf \phi \in L^2` such that - - .. math:: - - \int_{\Omega} \partial_t \mathbf{u}\cdot \mathbf{v} \, \textrm d\mathbf{x} &= \int_{\Omega} \phi \nabla \! \cdot \! \mathbf{v} \, \textrm d\mathbf{x} + \int_{\Omega} \mathbf{u}\! \times \! \mathbf{B}_0 \cdot \mathbf{v} \, \textrm d\mathbf{x} + \nu \int_{\Omega} \nabla \mathbf{u}\! : \! \nabla \mathbf{v} \, \textrm d\mathbf{x} + \int_{\Omega} f \mathbf{v} \, \textrm d\mathbf{x} \qquad \forall \, \mathbf{v} \in H(\textrm{div}) \,. - \\[2mm] - 0 &= - \int_{\Omega} \phi \nabla \! \cdot \! \mathbf{v_e} \, \textrm d\mathbf{x} - \int_{\Omega} \mathbf{u_e} \! \times \! \mathbf{B}_0 \cdot \mathbf{v_e} \, \textrm d\mathbf{x} + \nu_e \int_{\Omega} \nabla \mathbf{u_e} \!: \! \nabla \mathbf{v_e} \, \textrm d\mathbf{x} + \int_{\Omega} f_e \mathbf{v_e} \, \textrm d\mathbf{x} \qquad \forall \ \mathbf{v_e} \in H(\textrm{div}) \,. - \\[2mm] - 0 &= \int_{\Omega} \psi \nabla \cdot (\mathbf{u}-\mathbf{u_e}) \, \textrm d\mathbf{x} \qquad \forall \, \psi \in L^2 \,. - - :ref:`time_discret`: fully implicit. - """ - - # ========================================================================= - ### State variables (ion velocity u, electron velocity ue, pressure phi) - # ========================================================================= - - class Variables(): - def __init__(self) -> None: - self._u: FEECVariable | None = None - self._ue: FEECVariable | None = None - self._phi: FEECVariable | None = None - - @property - def u(self) -> FEECVariable | None: - return self._u - - @u.setter - def u(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hdiv" - self._u = new - - @property - def ue(self) -> FEECVariable | None: - return self._ue - - @ue.setter - def ue(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hdiv" - self._ue = new - - @property - def phi(self) -> FEECVariable | None: - return self._phi - - @phi.setter - def phi(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "L2" - self._phi = new - - def __init__(self): - self.variables = self.Variables() - - # ========================================================================= - ### Options - # ========================================================================= - - @dataclass - class Options(): - - nu: float | None = None - nu_e: float | None = None - eps_norm: float | None = None - - # boundary conditions per species - # supported kinds: "periodic", "dirichlet" - # future: "neumann", "robin" - boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_data_u: dict[tuple[int, int], Callable] | None = None - boundary_data_ue: dict[tuple[int, int], Callable] | None = None - - source_u: Callable | None = None - source_ue: Callable | None = None - - stab_sigma: float | None = None - - solver: OptsGenSolver = "gmres" - solver_params: SolverParameters | None = None - - def __post_init__(self): - - # --- required parameters --- - assert self.nu is not None, "nu must be specified" - assert self.nu_e is not None, "nu_e must be specified" - assert self.eps_norm is not None, "eps_norm must be specified" - assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" - assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" - - # --- physical parameter sanity checks --- - if self.nu < 0: - raise ValueError(f"nu must be non-negative, got {self.nu}") - if self.nu_e < 0: - raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") - if self.eps_norm <= 0: - raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - - # --- check all axes are covered --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for d in range(3): - for side in (-1, 1): - assert (d, side) in bcs, \ - f"{name} is missing entry for axis {d} side {side}" - - # --- periodic consistency: periodic must be paired on both sides --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for (d, side), kind in bcs.items(): - if kind == "periodic": - assert bcs.get((d, -side)) == "periodic", \ - f"{name}: axis {d} side {side} is periodic but opposite side is not" - - # --- ions and electrons must agree on which axes are periodic --- - for d in range(3): - u_left = self.boundary_conditions_u.get((d, -1)) - ue_left = self.boundary_conditions_ue.get((d, -1)) - u_right = self.boundary_conditions_u.get((d, 1)) - ue_right = self.boundary_conditions_ue.get((d, 1)) - u_periodic = (u_left == "periodic") - ue_periodic = (ue_left == "periodic") - if u_periodic != ue_periodic: - raise ValueError( - f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " - f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" - ) - - # --- warn for Dirichlet faces with no boundary data --- - for species, bcs, data, label in [ - ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), - ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), - ]: - has_dirichlet = any(v == "dirichlet" for v in bcs.values()) - if has_dirichlet: - if data is None: - warn(f"Dirichlet BCs specified for {species} but no {label} given " - f"— defaulting to homogeneous Dirichlet on all faces.") - else: - for (d, side), kind in bcs.items(): - if kind == "dirichlet" and (d, side) not in data: - warn(f"No {label} given for axis {d} side {side} " - f"— defaulting to homogeneous Dirichlet.") - - # --- warn if no source terms --- - if self.source_u is None: - warn("No source_u specified — defaulting to zero.") - if self.source_ue is None: - warn("No source_ue specified — defaulting to zero.") - - # --- defaults --- - if self.stab_sigma is None: - warn("stab_sigma not specified, defaulting to 0.0") - self.stab_sigma = 0.0 - - check_option(self.solver, OptsGenSolver) - if self.solver_params is None: - self.solver_params = SolverParameters() - - @property - def options(self) -> Options: - assert hasattr(self, "_options"), "Options not set." - return self._options - - @options.setter - def options(self, new): - assert isinstance(new, self.Options) - if MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nNew options for propagator '{self.__class__.__name__}':") - for k, v in new.__dict__.items(): - print(f" {k}: {v}") - self._options = new - - # ========================================================================= - ### Boundary condition helpers - # ========================================================================= - - def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): - - dirichlet_bc = [] - for d in range(3): - if spl_kind[d]: # periodic spline — no clamping ever - dirichlet_bc.append((False, False)) - else: - left = boundary_conditions.get((d, -1)) == "dirichlet" - right = boundary_conditions.get((d, 1)) == "dirichlet" - dirichlet_bc.append((left, right)) - return tuple(tuple(bc) for bc in dirichlet_bc) - - def _apply_boundary_conditions(self, vec, boundary_conditions): - """Zero out Dirichlet DOFs on the given stencil vector.""" - for (d, side), kind in boundary_conditions.items(): - if kind == "dirichlet": - apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) - # future: neumann and robin require no zeroing here - - # ========================================================================= - ### Allocate - # ========================================================================= - - def allocate(self, verbose=False): - - self.verbose = verbose - self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 - self._dt = None - - # ---- constrained (v0) de Rham complex -------------------------------- - - _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) - _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) - _dirichlet_bc = tuple( - (l_u or l_ue, r_u or r_ue) - for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) - ) - - self._derham_v0 = Derham( - self.derham.Nel, self.derham.p, self.derham.spl_kind, - domain=self.domain, dirichlet_bc=_dirichlet_bc, - ) - self._mass_ops_v0 = WeightedMassOperators( - self._derham_v0, self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.mass_ops.weights["eq_mhd"], - ) - self._basis_ops_v0 = BasisProjectionOperators( - self._derham_v0, self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.basis_ops.weights["eq_mhd"], - ) - - # ---- unconstrained operators (for RHS assembly) ---------------------- - - - self._M2 = self.mass_ops.M2 - self._M2B = - self.mass_ops.M2B - self._div = self.derham.div - self._curl = self.derham.curl - self._S21 = self.basis_ops.S21 - - self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div - + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - - self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl - self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) - + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) - - # ---- constrained operators (for system matrix) ----------------------- - - self._M2_v0 = self._mass_ops_v0.M2 - self._M3_v0 = self._mass_ops_v0.M3 - self._M2B_v0 = - self._mass_ops_v0.M2B - self._div_v0 = self._derham_v0.div - self._curl_v0 = self._derham_v0.curl - self._S21_v0 = self._basis_ops_v0.S21 - - self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 - + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) - - self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) - + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) - - # ---- block saddle-point system ---------------------------------------- - - self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) - self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self._derham_v0.Vh["3"] - - self._B1_v0 = - self._M3_v0 @ self._div_v0 - self._B2_v0 = self._M3_v0 @ self._div_v0 - - self._B_v0 = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_B_v0, - blocks=[[self._B1_v0, self._B2_v0]] - ) - - self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) - - _A_init = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, - blocks=[[self._A11_v0, None], [None, self._A22_v0]] - ) - _M_init = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, - blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] - ) - self._Minv = cast(InverseLinearOperator, inverse( - _M_init, self.options.solver, - x0=None, - tol=self.options.solver_params.tol, - maxiter=self.options.solver_params.maxiter, - verbose=self.options.solver_params.verbose, - )) - - # ---- projector ------------------------------------------------------- - - self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - - # ---- solution spline functions (unconstrained) ----------------------- - - self._u = self.derham.create_spline_function("u", space_id="Hdiv") - self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") - self._phi = self.derham.create_spline_function("phi", space_id="L2") - - # ---- BC lifts (unconstrained) ---------------------------------------- - - self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") - self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") - - for u_prime, boundary_data, boundary_conditions in [ - (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), - (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), - ]: - if boundary_data is None: - continue - for (d, side), f_bc in boundary_data.items(): - if boundary_conditions.get((d, side)) == "dirichlet": - bc_pulled = lambda *etas, f=f_bc: self.domain.pull( - [lambda x,y,z, f=f: f(x,y,z)[0], - lambda x,y,z, f=f: f(x,y,z)[1], - lambda x,y,z, f=f: f(x,y,z)[2]], - *etas, kind="2") - _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], - lambda *etas: bc_pulled(*etas)[1], - lambda *etas: bc_pulled(*etas)[2]]) - for (d2, side2), kind2 in boundary_conditions.items(): - if kind2 == "dirichlet" and (d2, side2) != (d, side): - apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) - u_prime.vector += _vec - - self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") - self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") - - self._u_prime_v0.vector = self._u_prime.vector - self._ue_prime_v0.vector = self._ue_prime.vector - - # ---- projected source terms (unconstrained) -------------------------- - - self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") - self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") - - for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: - if source is not None: - src_pulled = lambda *etas, f=source: self.domain.pull( - [lambda x,y,z, f=f: f(x,y,z)[0], - lambda x,y,z, f=f: f(x,y,z)[1], - lambda x,y,z, f=f: f(x,y,z)[2]], - *etas, kind="2") - rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], - lambda *etas: src_pulled(*etas)[1], - lambda *etas: src_pulled(*etas)[2]]) - - # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- - - self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") - self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") - - # ========================================================================= - ### Time step - # ========================================================================= - - def __call__(self, dt): - - # --- copy current state --- - self._u.vector = self.variables.u.spline.vector - self._ue.vector = self.variables.ue.spline.vector - - # --- rebuild system matrix if dt changed --- - if dt != self._dt: - self._dt = dt - _A = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, - blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] - ) - _M = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, - blocks=[[_A, self._B_v0.T], [self._B_v0, None]] - ) - self._Minv.linop = _M - - # --- assemble RHS in unconstrained space, then zero boundary DOFs --- - # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' - # electron: F2 = rhs_ue - A22 * ue' - self._rhs_vec_u.vector = (self._rhs_u.vector - + self._M2.dot(self._u.vector) / dt - - self._A11.dot(self._u_prime.vector) - - self._M2.dot(self._u_prime.vector) / dt) - self._rhs_vec_ue.vector = (self._rhs_ue.vector - - self._A22.dot(self._ue_prime.vector)) - - self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) - self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) - - # --- build block RHS and solve --- - _F = BlockVector(self._block_domain_v0, - blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) - _RHS = BlockVector(self._block_domain_M, - blocks=[_F, self._block_codomain_B_v0.zeros()]) - - _sol = self._Minv.dot(_RHS) - info = self._Minv.get_info() - - # --- reconstruct full solution: u = u_0 + u' --- - self._u.vector = _sol[0][0] + self._u_prime_v0.vector - self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector - self._phi.vector = _sol[1] - - # --- update FEEC variables --- - max_diffs = self.update_feec_variables( - u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector - ) - - if self.options.solver_params.info and self._rank == 0: - print(f"Status: {info['success']}, Iterations: {info['niter']}") - print(f"Max diffs: {max_diffs}") \ No newline at end of file diff --git a/src/struphy/utils/utils.py b/src/struphy/utils/utils.py index 6b26ab20f..963cce5d8 100644 --- a/src/struphy/utils/utils.py +++ b/src/struphy/utils/utils.py @@ -101,9 +101,11 @@ def kernels_to_txt(kernels: list, output: str): # print(f"kernels written to {output}.") -def check_option(opt, options): +def check_option(opt, *options): """Check if opt is contained in options; if opt is a list, checks for each element.""" - opts = get_args(options) + opts = [] + for o in options: + opts.extend(get_args(o)) if not isinstance(opt, list): opt = [opt] for o in opt: diff --git a/struphy-parameter-files b/struphy-parameter-files index 5781701ed..7f28854ea 160000 --- a/struphy-parameter-files +++ b/struphy-parameter-files @@ -1 +1 @@ -Subproject commit 5781701ed9d36997bdcdb015a310bf7d7c42ba86 +Subproject commit 7f28854eae0daa78dd9e8f351438f51f30ac651c From e677be773b5d4ae2fc9fc9b4b422d2a6ae3dc139 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Mon, 16 Mar 2026 14:08:25 +0000 Subject: [PATCH 07/41] moved v0 de rham complex construction into the Derham class, with everything that entails --- feectools | 2 +- src/struphy/feec/psydac_derham.py | 40 ++ src/struphy/io/options.py | 8 + src/struphy/io/setup.py | 4 + src/struphy/propagators/propagators_fields.py | 570 ++---------------- struphy-parameter-files | 2 +- 6 files changed, 118 insertions(+), 508 deletions(-) diff --git a/feectools b/feectools index 1981de121..d2a48ef19 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 1981de121fe6949b4a0797a3d61c8089c25c0b9f +Subproject commit d2a48ef19d31ae22ba238369cd5a53c679c6f1d8 diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 5ea6f6df0..19dcbf883 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -98,6 +98,7 @@ def __init__( spl_kind: list | tuple, *, dirichlet_bc: list | tuple = None, + lifting: list | tuple = None, nquads: list | tuple = None, nq_pr: list | tuple = None, comm=None, @@ -125,6 +126,41 @@ def __init__( self._dirichlet_bc = dirichlet_bc + # --- lifting: build constrained (v0) sub-complex --- + self._lifting = lifting + if lifting is not None: + assert len(lifting) == 3 + # lifting only makes sense on non-periodic axes + for d in range(3): + if spl_kind[d]: + assert lifting[d] == (False, False), \ + f"Axis {d} is periodic, lifting must be (False, False)" + + # v0 dirichlet_bc = dirichlet_bc OR lifting + if dirichlet_bc is not None: + v0_dirichlet_bc = tuple( + (d_l or l_l, d_r or l_r) + for (d_l, d_r), (l_l, l_r) in zip(dirichlet_bc, lifting) + ) + else: + v0_dirichlet_bc = lifting + + self._derham_v0 = Derham( + Nel, p, spl_kind, + dirichlet_bc=v0_dirichlet_bc, + nquads=nquads, + nq_pr=nq_pr, + comm=comm, + mpi_dims_mask=mpi_dims_mask, + with_projectors=with_projectors, + polar_ck=polar_ck, + local_projectors=self.with_local_projectors, + domain=domain, + ) + else: + self._derham_v0 = None + + # default p: exact integration of degree 2p+1 polynomials if nquads is None: self._nquads = [pi + 1 for pi in p] @@ -542,6 +578,10 @@ def __init__( xp.array(self.Vh["0"].starts), ) + @property + def derham_v0(self): + return self._derham_v0 + @property def Nel(self): """List of number of elements (=cells) in each direction.""" diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index cbe656e72..ce2408179 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -291,6 +291,11 @@ class DerhamOptions: dirichlet_bc : tuple[tuple[bool]] Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). + lifting : tuple[tuple[bool]] + Whether to build a constrained (v0) sub-complex with additional clamping on each face. + Used for inhomogeneous Dirichlet BCs: the v0 complex clamps faces where + lifting is True, and the propagator builds a lift in the unconstrained space. + nquads : tuple[int] Number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials). @@ -307,6 +312,7 @@ class DerhamOptions: p: tuple = (1, 1, 1) spl_kind: tuple = (True, True, True) dirichlet_bc: tuple = ((False, False), (False, False), (False, False)) + lifting: tuple = ((False, False), (False, False), (False, False)) nquads: tuple = None nq_pr: tuple = None polar_ck: LiteralOptions.PolarRegularity = -1 @@ -332,6 +338,7 @@ def to_dict(self) -> dict: "p": self.p, "spl_kind": self.spl_kind, "dirichlet_bc": self.dirichlet_bc, + "lifting": self.lifting, "nquads": self.nquads, "nq_pr": self.nq_pr, "polar_ck": self.polar_ck, @@ -345,6 +352,7 @@ def from_dict(cls, dct) -> "DerhamOptions": p=dct["p"], spl_kind=dct["spl_kind"], dirichlet_bc=dct["dirichlet_bc"], + lifting=dct.get("lifting", ((False, False), (False, False), (False, False))), nquads=dct["nquads"], nq_pr=dct["nq_pr"], polar_ck=dct["polar_ck"], diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index af5e10f3d..b9a157c5d 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -75,11 +75,14 @@ def setup_derham( # local commuting projectors local_projectors = options.local_projectors + lifting = options.lifting + derham = Derham( Nel, p, spl_kind, dirichlet_bc=dirichlet_bc, + lifting=lifting, nquads=nquads, nq_pr=nq_pr, comm=comm, @@ -90,6 +93,7 @@ def setup_derham( local_projectors=local_projectors, ) + if MPI.COMM_WORLD.Get_rank() == 0 and verbose: print("\nDERHAM:") print("number of elements:".ljust(25), Nel) diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 661e909ae..d31eb9479 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -7699,13 +7699,8 @@ class Options(): nu_e: float | None = None eps_norm: float | None = None - # boundary conditions per species - # supported kinds: "periodic", "dirichlet" - # future: "neumann", "robin" - boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_data_u: dict[tuple[int, int], Callable] | None = None - boundary_data_ue: dict[tuple[int, int], Callable] | None = None + boundary_data_u: dict[tuple[int, int], Callable] | None = None + boundary_data_ue: dict[tuple[int, int], Callable] | None = None source_u: Callable | None = None source_ue: Callable | None = None @@ -7718,11 +7713,9 @@ class Options(): def __post_init__(self): # --- required parameters --- - assert self.nu is not None, "nu must be specified" - assert self.nu_e is not None, "nu_e must be specified" - assert self.eps_norm is not None, "eps_norm must be specified" - assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" - assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" + assert self.nu is not None, "nu must be specified" + assert self.nu_e is not None, "nu_e must be specified" + assert self.eps_norm is not None, "eps_norm must be specified" # --- physical parameter sanity checks --- if self.nu < 0: @@ -7732,52 +7725,6 @@ def __post_init__(self): if self.eps_norm <= 0: raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - # --- check all axes are covered --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for d in range(3): - for side in (-1, 1): - assert (d, side) in bcs, \ - f"{name} is missing entry for axis {d} side {side}" - - # --- periodic consistency: periodic must be paired on both sides --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for (d, side), kind in bcs.items(): - if kind == "periodic": - assert bcs.get((d, -side)) == "periodic", \ - f"{name}: axis {d} side {side} is periodic but opposite side is not" - - # --- ions and electrons must agree on which axes are periodic --- - for d in range(3): - u_left = self.boundary_conditions_u.get((d, -1)) - ue_left = self.boundary_conditions_ue.get((d, -1)) - u_right = self.boundary_conditions_u.get((d, 1)) - ue_right = self.boundary_conditions_ue.get((d, 1)) - u_periodic = (u_left == "periodic") - ue_periodic = (ue_left == "periodic") - if u_periodic != ue_periodic: - raise ValueError( - f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " - f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" - ) - - # --- warn for Dirichlet faces with no boundary data --- - for species, bcs, data, label in [ - ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), - ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), - ]: - has_dirichlet = any(v == "dirichlet" for v in bcs.values()) - if has_dirichlet: - if data is None: - warn(f"Dirichlet BCs specified for {species} but no {label} given " - f"— defaulting to homogeneous Dirichlet on all faces.") - else: - for (d, side), kind in bcs.items(): - if kind == "dirichlet" and (d, side) not in data: - warn(f"No {label} given for axis {d} side {side} " - f"— defaulting to homogeneous Dirichlet.") - # --- warn if no source terms --- if self.source_u is None: warn("No source_u specified — defaulting to zero.") @@ -7789,7 +7736,7 @@ def __post_init__(self): warn("stab_sigma not specified, defaulting to 0.0") self.stab_sigma = 0.0 - check_option(self.solver, LiteralOptions.OptsGenSolver) + check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver) if self.solver_params is None: self.solver_params = SolverParameters() @@ -7811,394 +7758,40 @@ def options(self, new): ### Boundary condition helpers # ========================================================================= - def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): - - dirichlet_bc = [] - for d in range(3): - if spl_kind[d]: # periodic spline — no clamping ever - dirichlet_bc.append((False, False)) - else: - left = boundary_conditions.get((d, -1)) == "dirichlet" - right = boundary_conditions.get((d, 1)) == "dirichlet" - dirichlet_bc.append((left, right)) - return tuple(tuple(bc) for bc in dirichlet_bc) - - def _apply_boundary_conditions(self, vec, boundary_conditions): - """Zero out Dirichlet DOFs on the given stencil vector.""" - for (d, side), kind in boundary_conditions.items(): - if kind == "dirichlet": - apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) - # future: neumann and robin require no zeroing here - - # ========================================================================= - ### Allocate - # ========================================================================= - - def allocate(self, verbose=False): - - self.verbose = verbose - self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 - self._dt = None - - # ---- constrained (v0) de Rham complex -------------------------------- - - _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) - _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) - _dirichlet_bc = tuple( - (l_u or l_ue, r_u or r_ue) - for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) - ) - - self._derham_v0 = Derham( - self.derham.Nel, self.derham.p, self.derham.spl_kind, - domain=self.domain, dirichlet_bc=_dirichlet_bc, - ) - self._mass_ops_v0 = WeightedMassOperators( - self._derham_v0, self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.mass_ops.weights["eq_mhd"], - ) - self._basis_ops_v0 = BasisProjectionOperators( - self._derham_v0, self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.basis_ops.weights["eq_mhd"], - ) - - # ---- unconstrained operators (for RHS assembly) ---------------------- - - - self._M2 = self.mass_ops.M2 - self._M2B = - self.mass_ops.M2B - self._div = self.derham.div - self._curl = self.derham.curl - self._S21 = self.basis_ops.S21 - - self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div - + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - - self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl - self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) - + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) - - # ---- constrained operators (for system matrix) ----------------------- - - self._M2_v0 = self._mass_ops_v0.M2 - self._M3_v0 = self._mass_ops_v0.M3 - self._M2B_v0 = - self._mass_ops_v0.M2B - self._div_v0 = self._derham_v0.div - self._curl_v0 = self._derham_v0.curl - self._S21_v0 = self._basis_ops_v0.S21 - - self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 - + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) - - self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) - + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) - - # ---- block saddle-point system ---------------------------------------- - - self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) - self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self._derham_v0.Vh["3"] - - self._B1_v0 = - self._M3_v0 @ self._div_v0 - self._B2_v0 = self._M3_v0 @ self._div_v0 - - self._B_v0 = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_B_v0, - blocks=[[self._B1_v0, self._B2_v0]] - ) - - self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) - - _A_init = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, - blocks=[[self._A11_v0, None], [None, self._A22_v0]] - ) - _M_init = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, - blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] - ) - - self._Minv = inverse( - _M_init, self.options.solver, - A11=self._A11_v0, - A22=self._A22_v0, - B1=self._B1_v0, - B2=self._B2_v0, - recycle=self.options.solver_params.recycle, - tol=self.options.solver_params.tol, - maxiter=self.options.solver_params.maxiter, - verbose=self.options.solver_params.verbose, - ) - - # ---- projector ------------------------------------------------------- - - self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - - # ---- solution spline functions (unconstrained) ----------------------- - - self._u = self.derham.create_spline_function("u", space_id="Hdiv") - self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") - self._phi = self.derham.create_spline_function("phi", space_id="L2") - - # ---- BC lifts (unconstrained) ---------------------------------------- - - self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") - self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") - - for u_prime, boundary_data, boundary_conditions in [ - (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), - (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), - ]: - if boundary_data is None: - continue - for (d, side), f_bc in boundary_data.items(): - if boundary_conditions.get((d, side)) == "dirichlet": - bc_pulled = lambda *etas, f=f_bc: self.domain.pull( - [lambda x,y,z, f=f: f(x,y,z)[0], - lambda x,y,z, f=f: f(x,y,z)[1], - lambda x,y,z, f=f: f(x,y,z)[2]], - *etas, kind="2") - _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], - lambda *etas: bc_pulled(*etas)[1], - lambda *etas: bc_pulled(*etas)[2]]) - for (d2, side2), kind2 in boundary_conditions.items(): - if kind2 == "dirichlet" and (d2, side2) != (d, side): - apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) - u_prime.vector += _vec - - self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") - self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") - - self._u_prime_v0.vector = self._u_prime.vector - self._ue_prime_v0.vector = self._ue_prime.vector - - # ---- projected source terms (unconstrained) -------------------------- - - self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") - self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") - - for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: - if source is not None: - src_pulled = lambda *etas, f=source: self.domain.pull( - [lambda x,y,z, f=f: f(x,y,z)[0], - lambda x,y,z, f=f: f(x,y,z)[1], - lambda x,y,z, f=f: f(x,y,z)[2]], - *etas, kind="2") - rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], - lambda *etas: src_pulled(*etas)[1], - lambda *etas: src_pulled(*etas)[2]]) - - # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- - - self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") - self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") - - # ========================================================================= - ### Time step - # ========================================================================= - - def __call__(self, dt): - - # --- copy current state --- - self._u.vector = self.variables.u.spline.vector - self._ue.vector = self.variables.ue.spline.vector - - # --- rebuild system matrix if dt changed --- - if dt != self._dt: - self._dt = dt - _A = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, - blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] - ) - - _M = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, - blocks=[[_A, self._B_v0.T], [self._B_v0, None]] - ) - self._Minv.linop = _M - class Variables(): - def __init__(self) -> None: - self._u: FEECVariable | None = None - self._ue: FEECVariable | None = None - self._phi: FEECVariable | None = None - - @property - def u(self) -> FEECVariable | None: - return self._u - - @u.setter - def u(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hdiv" - self._u = new - - @property - def ue(self) -> FEECVariable | None: - return self._ue - - @ue.setter - def ue(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hdiv" - self._ue = new - - @property - def phi(self) -> FEECVariable | None: - return self._phi - - @phi.setter - def phi(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "L2" - self._phi = new - - def __init__(self): - self.variables = self.Variables() - - # ========================================================================= - ### Options - # ========================================================================= - - @dataclass - class Options(): - - nu: float | None = None - nu_e: float | None = None - eps_norm: float | None = None - - # boundary conditions per species - # supported kinds: "periodic", "dirichlet" - # future: "neumann", "robin" - boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_data_u: dict[tuple[int, int], Callable] | None = None - boundary_data_ue: dict[tuple[int, int], Callable] | None = None - - source_u: Callable | None = None - source_ue: Callable | None = None - - stab_sigma: float | None = None - - solver: LiteralOptions.OptsGenSolver = "gmres" - solver_params: SolverParameters | None = None - - def __post_init__(self): - - # --- required parameters --- - assert self.nu is not None, "nu must be specified" - assert self.nu_e is not None, "nu_e must be specified" - assert self.eps_norm is not None, "eps_norm must be specified" - assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" - assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" - - # --- physical parameter sanity checks --- - if self.nu < 0: - raise ValueError(f"nu must be non-negative, got {self.nu}") - if self.nu_e < 0: - raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") - if self.eps_norm <= 0: - raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - - # --- check all axes are covered --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for d in range(3): - for side in (-1, 1): - assert (d, side) in bcs, \ - f"{name} is missing entry for axis {d} side {side}" - - # --- periodic consistency: periodic must be paired on both sides --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for (d, side), kind in bcs.items(): - if kind == "periodic": - assert bcs.get((d, -side)) == "periodic", \ - f"{name}: axis {d} side {side} is periodic but opposite side is not" - - # --- ions and electrons must agree on which axes are periodic --- - for d in range(3): - u_left = self.boundary_conditions_u.get((d, -1)) - ue_left = self.boundary_conditions_ue.get((d, -1)) - u_right = self.boundary_conditions_u.get((d, 1)) - ue_right = self.boundary_conditions_ue.get((d, 1)) - u_periodic = (u_left == "periodic") - ue_periodic = (ue_left == "periodic") - if u_periodic != ue_periodic: - raise ValueError( - f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " - f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" - ) - - # --- warn for Dirichlet faces with no boundary data --- - for species, bcs, data, label in [ - ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), - ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), - ]: - has_dirichlet = any(v == "dirichlet" for v in bcs.values()) - if has_dirichlet: - if data is None: - warn(f"Dirichlet BCs specified for {species} but no {label} given " - f"— defaulting to homogeneous Dirichlet on all faces.") - else: - for (d, side), kind in bcs.items(): - if kind == "dirichlet" and (d, side) not in data: - warn(f"No {label} given for axis {d} side {side} " - f"— defaulting to homogeneous Dirichlet.") + def _get_dirichlet_faces(self): + """Infer which faces have Dirichlet BCs by comparing derham and derham_v0. - # --- warn if no source terms --- - if self.source_u is None: - warn("No source_u specified — defaulting to zero.") - if self.source_ue is None: - warn("No source_ue specified — defaulting to zero.") - - # --- defaults --- - if self.stab_sigma is None: - warn("stab_sigma not specified, defaulting to 0.0") - self.stab_sigma = 0.0 - - check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver) - if self.solver_params is None: - self.solver_params = SolverParameters() - - @property - def options(self) -> Options: - assert hasattr(self, "_options"), "Options not set." - return self._options - - @options.setter - def options(self, new): - assert isinstance(new, self.Options) - if MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nNew options for propagator '{self.__class__.__name__}':") - for k, v in new.__dict__.items(): - print(f" {k}: {v}") - self._options = new + A face is Dirichlet if it is unclamped in derham but clamped in derham_v0 + (i.e. lifting is True there). + """ + faces = [] + derham = self.derham + derham_v0 = derham.derham_v0 - # ========================================================================= - ### Boundary condition helpers - # ========================================================================= + if derham_v0 is None: + return faces - def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): + bc = derham.dirichlet_bc + bc_v0 = derham_v0.dirichlet_bc - dirichlet_bc = [] for d in range(3): - if spl_kind[d]: # periodic spline — no clamping ever - dirichlet_bc.append((False, False)) - else: - left = boundary_conditions.get((d, -1)) == "dirichlet" - right = boundary_conditions.get((d, 1)) == "dirichlet" - dirichlet_bc.append((left, right)) - return tuple(tuple(bc) for bc in dirichlet_bc) - - def _apply_boundary_conditions(self, vec, boundary_conditions): - """Zero out Dirichlet DOFs on the given stencil vector.""" - for (d, side), kind in boundary_conditions.items(): - if kind == "dirichlet": - apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) - # future: neumann and robin require no zeroing here + if derham.spl_kind[d]: + continue # periodic axis, no Dirichlet + for s, side in enumerate((-1, 1)): + # clamped in v0 but not in derham => this is a lifted (inhom Dirichlet) face + unclamped = not bc[d][s] + clamped_v0 = bc_v0[d][s] if bc_v0 is not None else False + if unclamped and clamped_v0: + faces.append((d, side)) + # clamped in both => homogeneous Dirichlet, also need to zero DOFs + elif bc[d][s] and clamped_v0: + faces.append((d, side)) + return faces + + def _apply_essential_bc(self, vec): + """Zero out Dirichlet DOFs, inferred from derham vs derham_v0.""" + for (d, side) in self._dirichlet_faces: + apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) # ========================================================================= ### Allocate @@ -8210,19 +7803,13 @@ def allocate(self, verbose=False): self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 self._dt = None - # ---- constrained (v0) de Rham complex -------------------------------- + # ---- v0 de Rham complex (from derham.derham_v0) ---------------------- - _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) - _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) - _dirichlet_bc = tuple( - (l_u or l_ue, r_u or r_ue) - for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) - ) + assert self.derham.derham_v0 is not None, \ + "derham must be constructed with lifting to use this propagator" + + self._derham_v0 = self.derham.derham_v0 - self._derham_v0 = Derham( - self.derham.Nel, self.derham.p, self.derham.spl_kind, - domain=self.domain, dirichlet_bc=_dirichlet_bc, - ) self._mass_ops_v0 = WeightedMassOperators( self._derham_v0, self.domain, verbose=self.options.solver_params.verbose, @@ -8234,8 +7821,11 @@ def allocate(self, verbose=False): eq_mhd=self.basis_ops.weights["eq_mhd"], ) - # ---- unconstrained operators (for RHS assembly) ---------------------- + # ---- Dirichlet faces (inferred from derham vs derham_v0) ------------- + self._dirichlet_faces = self._get_dirichlet_faces() + + # ---- unconstrained operators (for RHS assembly) ---------------------- self._M2 = self.mass_ops.M2 self._M2B = - self.mass_ops.M2B @@ -8298,10 +7888,10 @@ def allocate(self, verbose=False): A22=self._A22_v0, B1=self._B1_v0, B2=self._B2_v0, + recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, - recycle=self.options.solver_params.recycle, ) else: self._Minv = inverse( @@ -8312,6 +7902,7 @@ def allocate(self, verbose=False): verbose=self.options.solver_params.verbose, ) + # ---- projector ------------------------------------------------------- self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) @@ -8327,14 +7918,14 @@ def allocate(self, verbose=False): self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") - for u_prime, boundary_data, boundary_conditions in [ - (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), - (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), + for u_prime, boundary_data in [ + (self._u_prime, self.options.boundary_data_u), + (self._ue_prime, self.options.boundary_data_ue), ]: if boundary_data is None: continue for (d, side), f_bc in boundary_data.items(): - if boundary_conditions.get((d, side)) == "dirichlet": + if (d, side) in self._dirichlet_faces: bc_pulled = lambda *etas, f=f_bc: self.domain.pull( [lambda x,y,z, f=f: f(x,y,z)[0], lambda x,y,z, f=f: f(x,y,z)[1], @@ -8343,8 +7934,8 @@ def allocate(self, verbose=False): _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], lambda *etas: bc_pulled(*etas)[1], lambda *etas: bc_pulled(*etas)[2]]) - for (d2, side2), kind2 in boundary_conditions.items(): - if kind2 == "dirichlet" and (d2, side2) != (d, side): + for (d2, side2) in self._dirichlet_faces: + if (d2, side2) != (d, side): apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) u_prime.vector += _vec @@ -8385,14 +7976,14 @@ def __call__(self, dt): self._u.vector = self.variables.u.spline.vector self._ue.vector = self.variables.ue.spline.vector - # --- rebuild system matrix if dt changed --- TODO change uzawa internals - if dt != self._dt: + # --- rebuild system matrix if dt changed --- + if dt != self._dt: # TODO change uzawa A11 block too self._dt = dt _A = BlockLinearOperator( self._block_domain_v0, self._block_codomain_v0, blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] ) - + _M = BlockLinearOperator( self._block_domain_M, self._block_domain_M, blocks=[[_A, self._B_v0.T], [self._B_v0, None]] @@ -8402,21 +7993,21 @@ def __call__(self, dt): # --- assemble RHS in unconstrained space, then zero boundary DOFs --- # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' # electron: F2 = rhs_ue - A22 * ue' - self._rhs_vec_u.vector = (self._rhs_u.vector - + self._M2.dot(self._u.vector) / dt - - self._A11.dot(self._u_prime.vector) - - self._M2.dot(self._u_prime.vector) / dt) + self._rhs_vec_u.vector = (self._rhs_u.vector # TODO boundary operator + + self._M2.dot(self._u.vector) / dt + - self._A11.dot(self._u_prime.vector) + - self._M2.dot(self._u_prime.vector) / dt) self._rhs_vec_ue.vector = (self._rhs_ue.vector - - self._A22.dot(self._ue_prime.vector)) + - self._A22.dot(self._ue_prime.vector)) - self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) - self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) + self._apply_essential_bc(self._rhs_vec_u.vector) + self._apply_essential_bc(self._rhs_vec_ue.vector) # --- build block RHS and solve --- _F = BlockVector(self._block_domain_v0, - blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) _RHS = BlockVector(self._block_domain_M, - blocks=[_F, self._block_codomain_B_v0.zeros()]) + blocks=[_F, self._block_codomain_B_v0.zeros()]) _sol = self._Minv.dot(_RHS) info = self._Minv.get_info() @@ -8434,38 +8025,5 @@ def __call__(self, dt): if self.options.solver_params.info and self._rank == 0: print(f"Status: {info['success']}, Iterations: {info['niter']}") print(f"Max diffs: {max_diffs}") - # --- assemble RHS in unconstrained space, then zero boundary DOFs --- - # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' - # electron: F2 = rhs_ue - A22 * ue' - self._rhs_vec_u.vector = (self._rhs_u.vector - + self._M2.dot(self._u.vector) / dt - - self._A11.dot(self._u_prime.vector) - - self._M2.dot(self._u_prime.vector) / dt) - self._rhs_vec_ue.vector = (self._rhs_ue.vector - - self._A22.dot(self._ue_prime.vector)) - - self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) - self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) - - # --- build block RHS and solve --- - _F = BlockVector(self._block_domain_v0, - blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) - _RHS = BlockVector(self._block_domain_M, - blocks=[_F, self._block_codomain_B_v0.zeros()]) - - _sol = self._Minv.dot(_RHS) - info = self._Minv.get_info() - - # --- reconstruct full solution: u = u_0 + u' --- - self._u.vector = _sol[0][0] + self._u_prime_v0.vector - self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector - self._phi.vector = _sol[1] - - # --- update FEEC variables --- - max_diffs = self.update_feec_variables( - u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector - ) - - if self.options.solver_params.info and self._rank == 0: print(f"Status: {info['success']}, Iterations: {info['niter']}") print(f"Max diffs: {max_diffs}") \ No newline at end of file diff --git a/struphy-parameter-files b/struphy-parameter-files index 7f28854ea..5143ca521 160000 --- a/struphy-parameter-files +++ b/struphy-parameter-files @@ -1 +1 @@ -Subproject commit 7f28854eae0daa78dd9e8f351438f51f30ac651c +Subproject commit 5143ca52173bb00766c5032d2b9e652ecabd885e From 578a761597c522bf0ebe6c67a6f36b819e306116 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Mon, 8 Dec 2025 09:50:14 +0000 Subject: [PATCH 08/41] Fixed spaces for GMRES version of two fluid propagator. --- src/struphy/io/options.py | 81 +- src/struphy/propagators/propagators_fields.py | 814 +++++++++++++----- 2 files changed, 680 insertions(+), 215 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 6ecfff7c7..d88eae287 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -3,26 +3,67 @@ from dataclasses import dataclass, fields from typing import Any, Callable, Literal -from struphy.utils.utils import ( - __class_with_params_repr_no_defaults__, - __dataclass_repr_no_defaults__, - all_class_params_are_default, - check_option, -) - -logger = logging.getLogger("struphy") - - -class OptionsBase: - def to_dict(self) -> dict: - """Convert dataclass instance to dictionary.""" - return {field.name: getattr(self, field.name) for field in fields(type(self)) if field.init} - - @classmethod - def from_dict(cls, dct) -> "Any": - """Create dataclass instance from dictionary.""" - valid_fields = {field.name for field in fields(cls) if field.init} - return cls(**{key: value for key, value in dct.items() if key in valid_fields}) +import cunumpy as xp +from psydac.ddm.mpi import mpi as MPI + +from struphy.physics.physics import ConstantsOfNature + +## Literal options + +# time +SplitAlgos = Literal["LieTrotter", "Strang"] + +# derham +PolarRegularity = Literal[-1, 1] +OptsFEECSpace = Literal["H1", "Hcurl", "Hdiv", "L2", "H1vec"] +OptsVecSpace = Literal["Hcurl", "Hdiv", "H1vec"] + +# fields background +BackgroundTypes = Literal["LogicalConst", "FluidEquilibrium"] + +# perturbations +NoiseDirections = Literal["e1", "e2", "e3", "e1e2", "e1e3", "e2e3", "e1e2e3"] +GivenInBasis = Literal["0", "1", "2", "3", "v", "physical", "physical_at_eta", "norm", None] + +# solvers +OptsSymmSolver = Literal["pcg", "cg"] +OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"] +OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] +OptsSaddlePointSolver = Literal["Uzawa", "GMRES"] # todo +OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] +OptsNonlinearSolver = Literal["Picard", "Newton"] + +# markers +OptsPICSpace = Literal["Particles6D", "DeltaFParticles6D", "Particles5D", "Particles3D"] +OptsMarkerBC = Literal["periodic", "reflect"] +OptsRecontructBC = Literal["periodic", "mirror", "fixed"] +OptsLoading = Literal[ + "pseudo_random", + "sobol_standard", + "sobol_antithetic", + "external", + "restart", + "tesselation", +] +OptsSpatialLoading = Literal["uniform", "disc"] +OptsMPIsort = Literal["each", "last", None] + +# filters +OptsFilter = Literal["fourier_in_tor", "hybrid", "three_point", None] + +# sph +OptsKernel = Literal[ + "trigonometric_1d", + "gaussian_1d", + "linear_1d", + "trigonometric_2d", + "gaussian_2d", + "linear_2d", + "trigonometric_3d", + "gaussian_3d", + "linear_isotropic_3d", + "linear_3d", +] @dataclass diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 7528f6157..4392d056a 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -7750,82 +7750,393 @@ def options(self, new): logger.info(f" {k}: {v}") self._options = new - # ========================================================================= - ### Boundary condition helpers - # ========================================================================= - - def _get_dirichlet_faces(self): - """Infer which faces have Dirichlet BCs by comparing derham and derham_v0. - - A face is Dirichlet if it is unclamped in derham but clamped in derham_v0 - (i.e. lifting is True there). - """ - faces = [] - derham = self.derham - derham_v0 = derham - - if derham_v0 is None: - return faces - - bc = derham.dirichlet_bc - bc_v0 = derham_v0.dirichlet_bc - - for d in range(3): - if derham.spl_kind[d]: - continue # periodic axis, no Dirichlet - for s, side in enumerate((-1, 1)): - # clamped in v0 but not in derham => this is a lifted (inhom Dirichlet) face - unclamped = not bc[d][s] - clamped_v0 = bc_v0[d][s] if bc_v0 is not None else False - if unclamped and clamped_v0: - faces.append((d, side)) - # clamped in both => homogeneous Dirichlet, also need to zero DOFs - elif bc[d][s] and clamped_v0: - faces.append((d, side)) - return faces - - def _apply_essential_bc(self, vec): - """Zero out Dirichlet DOFs, inferred from derham vs derham_v0.""" - for d, side in self._dirichlet_faces: - apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) + @profile + def allocate(self, verbose: bool = False): + self._info = self.options.solver_params.info + if self.derham.comm is not None: + self._rank = self.derham.comm.Get_rank() + else: + self._rank = 0 - # ========================================================================= - ### Allocate - # ========================================================================= + self._nu = self.options.nu + self._nu_e = self.options.nu_e + self._eps_norm = self.options.eps_norm + self._a = self.options.a + self._R0 = self.options.R0 + self._B0 = self.options.B0 + self._Bp = self.options.Bp + self._alpha = self.options.alpha + self._beta = self.options.beta + self._stab_sigma = self.options.stab_sigma + self._variant = self.options.variant + self._method_to_solve = self.options.method_to_solve + self._preconditioner = self.options.preconditioner + self._dimension = self.options.dimension + self._spectralanalysis = self.options.spectralanalysis + self._lifting = self.options.lifting - def allocate(self, verbose=False): + solver_params = self.options.solver_params - self.verbose = verbose - self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 - self._dt = None + u = self.variables.u.spline.vector - # ---- v0 de Rham complex (from derham.derham_v0) ---------------------- - self._derham_v0 = self.derham + # Lifting for nontrivial boundary conditions + # derham had boundary conditions in eta1 direction, the following is in space Hdiv_0 + if self._lifting: + self.derhamv0 = Derham( + self.derham.Nel, + self.derham.p, + self.derham.spl_kind, + domain=self.domain, + dirichlet_bc=((True, True), (False, False), (False, False)), + ) - self._mass_ops_v0 = WeightedMassOperators( - self._derham_v0, - self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.mass_ops.weights["eq_mhd"], - ) - self._basis_ops_v0 = BasisProjectionOperators( - self._derham_v0, - self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.basis_ops.weights["eq_mhd"], - ) + self._mass_opsv0 = WeightedMassOperators( + self.derhamv0, + self.domain, + verbose=solver_params.verbose, + eq_mhd=self.mass_ops.weights["eq_mhd"], + ) + self._basis_opsv0 = BasisProjectionOperators( + self.derhamv0, + self.domain, + verbose=solver_params.verbose, + eq_mhd=self.basis_ops.weights["eq_mhd"], + ) + else: + self.derhamnumpy = Derham( + self.derham.Nel, + self.derham.p, + self.derham.spl_kind, + domain=self.domain, + # dirichlet_bc=self.derham.dirichlet_bc, + # nquads = self.derham._nquads, + # nq_pr = self.derham._nq_pr, + # comm = MPI.COMM_SELF, # self.derham._comm, + # polar_ck= self.derham._polar_ck, + # local_projectors=self.derham.with_local_projectors + ) - # ---- Dirichlet faces (inferred from derham vs derham_v0) ------------- + # get forceterms for according dimension + if self._dimension in ["2D", "1D"]: + ### Manufactured solution ### + _forceterm_logical = lambda e1, e2, e3: 0 * e1 + _funx = getattr(callables, "ManufacturedSolutionForceterm")( + species="Ions", + comp="0", + b0=self._B0, + nu=self._nu, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=self.options.D1_dt, + ) + _funy = getattr(callables, "ManufacturedSolutionForceterm")( + species="Ions", + comp="1", + b0=self._B0, + nu=self._nu, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=self.options.D1_dt, + ) + _funelectronsx = getattr(callables, "ManufacturedSolutionForceterm")( + species="Electrons", + comp="0", + b0=self._B0, + nu_e=self._nu_e, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=self.options.D1_dt, + ) + _funelectronsy = getattr(callables, "ManufacturedSolutionForceterm")( + species="Electrons", + comp="1", + b0=self._B0, + nu_e=self._nu_e, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=self.options.D1_dt, + ) - self._dirichlet_faces = self._get_dirichlet_faces() + # get callable(s) for specified init type + forceterm_class = [_funx, _funy, _forceterm_logical] + forcetermelectrons_class = [_funelectronsx, _funelectronsy, _forceterm_logical] + + # pullback callable + funx = TransformedPformComponent( + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, + ) + funy = TransformedPformComponent( + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, + ) + fun_electronsx = TransformedPformComponent( + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, + ) + fun_electronsy = TransformedPformComponent( + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, + ) + l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + self._F1 = l2_proj([funx, funy, _forceterm_logical]) + self._F2 = l2_proj([fun_electronsx, fun_electronsy, _forceterm_logical]) + + elif self._dimension == "Restelli": + ### Restelli ### + + _forceterm_logical = lambda e1, e2, e3: 0 * e1 + _fun = getattr(callables, "RestelliForcingTerm")( + B0=self._B0, + nu=self._nu, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + eps=self._eps_norm, + ) + _funelectrons = getattr(callables, "RestelliForcingTerm")( + B0=self._B0, + nu=self._nu_e, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + eps=self._eps_norm, + ) - # ---- unconstrained operators (for RHS assembly) ---------------------- + # get callable(s) for specified init type + forceterm_class = [_forceterm_logical, _forceterm_logical, _fun] + forcetermelectrons_class = [_forceterm_logical, _forceterm_logical, _funelectrons] + + # pullback callable + fun_pb_1 = TransformedPformComponent( + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, + ) + fun_pb_2 = TransformedPformComponent( + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, + ) + fun_pb_3 = TransformedPformComponent( + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=2, + domain=self.domain, + ) + fun_electrons_pb_1 = TransformedPformComponent( + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, + ) + fun_electrons_pb_2 = TransformedPformComponent( + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, + ) + fun_electrons_pb_3 = TransformedPformComponent( + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=2, + domain=self.domain, + ) + if self._lifting: + l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) + else: + l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + self._F1 = l2_proj([fun_pb_1, fun_pb_2, fun_pb_3], apply_bc=self._lifting) + self._F2 = l2_proj([fun_electrons_pb_1, fun_electrons_pb_2, fun_electrons_pb_3], apply_bc=self._lifting) + + ### End Restelli ### + + elif self._dimension == "Tokamak": + ### Tokamak geometry curl-free manufactured solution ### + + _forceterm_logical = lambda e1, e2, e3: 0 * e1 + _funx = getattr(callables, "ManufacturedSolutionForceterm")( + species="Ions", + comp="0", + b0=self._B0, + nu=self._nu, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=self.options.D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + _funy = getattr(callables, "ManufacturedSolutionForceterm")( + species="Ions", + comp="1", + b0=self._B0, + nu=self._nu, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=self.options.D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + _funz = getattr(callables, "ManufacturedSolutionForceterm")( + species="Ions", + comp="2", + b0=self._B0, + nu=self._nu, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=self.options.D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + _funelectronsx = getattr(callables, "ManufacturedSolutionForceterm")( + species="Electrons", + comp="0", + b0=self._B0, + nu_e=self._nu_e, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=self.options.D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + _funelectronsy = getattr(callables, "ManufacturedSolutionForceterm")( + species="Electrons", + comp="1", + b0=self._B0, + nu_e=self._nu_e, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=self.options.D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + _funelectronsz = getattr(callables, "ManufacturedSolutionForceterm")( + species="Electrons", + comp="2", + b0=self._B0, + nu_e=self._nu_e, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=self.options.D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) - self._M2 = self.mass_ops.M2 - self._M2B = -self.mass_ops.M2B - self._div = self.derham.div - self._curl = self.derham.curl - self._S21 = self.basis_ops.S21 + # get callable(s) for specified init type + forceterm_class = [_funx, _funy, _funz] + forcetermelectrons_class = [_funelectronsx, _funelectronsy, _funelectronsz] + + # pullback callable + fun_pb_1 = TransformedPformComponent( + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, + ) + fun_pb_2 = TransformedPformComponent( + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, + ) + fun_pb_3 = TransformedPformComponent( + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=2, + domain=self.domain, + ) + fun_electrons_pb_1 = TransformedPformComponent( + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, + ) + fun_electrons_pb_2 = TransformedPformComponent( + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, + ) + fun_electrons_pb_3 = TransformedPformComponent( + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=2, + domain=self.domain, + ) + if self._lifting: + l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) + else: + l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + self._F1 = l2_proj([fun_pb_1, fun_pb_2, fun_pb_3], apply_bc=self._lifting) + self._F2 = l2_proj([fun_electrons_pb_1, fun_electrons_pb_2, fun_electrons_pb_3], apply_bc=self._lifting) + + ### End Tokamak geometry manufactured solution ### + + if self._variant == "GMRES": + if self._lifting: + self._M2 = getattr(self._mass_opsv0, "M2") + self._M3 = getattr(self._mass_opsv0, "M3") + self._M2B = -getattr(self._mass_opsv0, "M2B") + self._div = self.derhamv0.div + self._curl = self.derhamv0.curl + self._S21 = self._basis_opsv0.S21 + else: + self._M2 = getattr(self.mass_ops, "M2") + self._M3 = getattr(self.mass_ops, "M3") + self._M2B = -getattr(self.mass_ops, "M2B") + self._div = self.derham.div + self._curl = self.derham.curl + self._S21 = self.basis_ops.S21 self._lapl = ( self._div.T @ self.mass_ops.M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21 @@ -7894,147 +8205,260 @@ def allocate(self, verbose=False): maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, ) - else: - self._Minv = inverse( - _M_init, - self.options.solver, - recycle=self.options.solver_params.recycle, + # Allocate memory for call + self._untemp = self.variables.u.spline.vector.space.zeros() + + elif self._variant == "Uzawa": + self._solver_UzawaNumpy = SaddlePointSolver( + Apre=_Anppre, + A=_Anp, + B=_Bnp, + F=_Fnp, + method_to_solve=self._method_to_solve, + preconditioner=self._preconditioner, + spectralanalysis=self.options.spectralanalysis, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, ) - # ---- projector ------------------------------------------------------- - - self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - - # ---- solution spline functions (unconstrained) ----------------------- - - self._u = self.derham.create_spline_function("u", space_id="Hdiv") - self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") - self._phi = self.derham.create_spline_function("phi", space_id="L2") - - # ---- BC lifts (unconstrained) ---------------------------------------- - - self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") - self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") - - for u_prime, boundary_data in [ - (self._u_prime, self.options.boundary_data_u), - (self._ue_prime, self.options.boundary_data_ue), - ]: - if boundary_data is None: - continue - for (d, side), f_bc in boundary_data.items(): - if (d, side) in self._dirichlet_faces: - bc_pulled = lambda *etas, f=f_bc: self.domain.pull( - [ - lambda x, y, z, f=f: f(x, y, z)[0], - lambda x, y, z, f=f: f(x, y, z)[1], - lambda x, y, z, f=f: f(x, y, z)[2], - ], - *etas, - kind="2", - ) - _vec = self._projector( - [ - lambda *etas: bc_pulled(*etas)[0], - lambda *etas: bc_pulled(*etas)[1], - lambda *etas: bc_pulled(*etas)[2], - ] - ) - for d2, side2 in self._dirichlet_faces: - if (d2, side2) != (d, side): - apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) - u_prime.vector += _vec - - self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") - self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") - - self._u_prime_v0.vector = self._u_prime.vector - self._ue_prime_v0.vector = self._ue_prime.vector - - # ---- projected source terms (unconstrained) -------------------------- - - self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") - self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") - - for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: - if source is not None: - src_pulled = lambda *etas, f=source: self.domain.pull( - [ - lambda x, y, z, f=f: f(x, y, z)[0], - lambda x, y, z, f=f: f(x, y, z)[1], - lambda x, y, z, f=f: f(x, y, z)[2], - ], - *etas, - kind="2", - ) - rhs.vector = self._projector.get_dofs( - [ - lambda *etas: src_pulled(*etas)[0], - lambda *etas: src_pulled(*etas)[1], - lambda *etas: src_pulled(*etas)[2], - ] - ) - - # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- - - self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") - self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") - - # ========================================================================= - ### Time step - # ========================================================================= - def __call__(self, dt): - - # --- copy current state --- - self._u.vector = self.variables.u.spline.vector - self._ue.vector = self.variables.ue.spline.vector - - # --- rebuild system matrix if dt changed --- - if dt != self._dt: # TODO change uzawa A11 block too - self._dt = dt - _A = BlockLinearOperator( - self._block_domain_v0, - self._block_codomain_v0, - blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]], + # current variables + unfeec = self.variables.u.spline.vector + uenfeec = self.variables.ue.spline.vector + phinfeec = self.variables.phi.spline.vector + + if self._variant == "GMRES": + if self._lifting: + phinfeeccopy = self.derhamv0.create_spline_function("phi", space_id="L2") + phinfeeccopy.vector = phinfeec + # unfeec in space Hdiv, u0 in space Hdiv_0 + unfeeccopy = self.derhamv0.create_spline_function("u", space_id="Hdiv") + u0 = self.derhamv0.create_spline_function("u", space_id="Hdiv") + u_prime = self.derhamv0.create_spline_function("u", space_id="Hdiv") + u0.vector = uenfeec + unfeeccopy.vector = uenfeec + apply_essential_bc_stencil(u0.vector[0], axis=0, ext=-1, order=0) + apply_essential_bc_stencil(u0.vector[0], axis=0, ext=1, order=0) + u_prime.vector = unfeeccopy.vector - u0.vector + + uenfeeccopy = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue0 = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue_prime = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue0.vector = uenfeec + uenfeeccopy.vector = uenfeec + apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=-1, order=0) + apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=1, order=0) + ue_prime.vector = uenfeeccopy.vector - ue0.vector + + _A11 = ( + self._M2 / dt + - self._M2B / self._eps_norm + + self._nu + * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) ) - - _M = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, blocks=[[_A, self._B_v0.T], [self._B_v0, None]] + _A12 = None + _A21 = _A12 + _A22 = ( + self._nu_e + * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + + self._M2B / self._eps_norm + - self._stab_sigma * IdentityOperator(_A11.domain) ) - self._Minv.linop = _M - - # --- assemble RHS in unconstrained space, then zero boundary DOFs --- - # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' - # electron: F2 = rhs_ue - A22 * ue' - self._rhs_vec_u.vector = ( - self._rhs_u.vector # TODO boundary operator - + self._M2.dot(self._u.vector) / dt - - self._A11.dot(self._u_prime.vector) - - self._M2.dot(self._u_prime.vector) / dt - ) - self._rhs_vec_ue.vector = self._rhs_ue.vector - self._A22.dot(self._ue_prime.vector) - self._apply_essential_bc(self._rhs_vec_u.vector) - self._apply_essential_bc(self._rhs_vec_ue.vector) + if self._lifting: + _A11prime = -self._M2B / self._eps_norm + self._nu * ( + self.derhamv0.div.T @ self._M3 @ self.derhamv0.div + + self._basis_opsv0.S21.T + @ self.derhamv0.curl.T + @ self._M2 + @ self.derhamv0.curl + @ self._basis_opsv0.S21 + ) + _A22prime = ( + self._nu_e + * ( + self.derhamv0.div.T @ self._M3 @ self.derhamv0.div + + self._basis_opsv0.S21.T + @ self.derhamv0.curl.T + @ self._M2 + @ self.derhamv0.curl + @ self._basis_opsv0.S21 + ) + + self._M2B / self._eps_norm + - self._stab_sigma * IdentityOperator(_A11.domain) + ) + _B1 = -self._M3 @ self._div + _B2 = self._M3 @ self._div + + if _A12 is not None: + assert _A11.codomain == _A12.codomain + if _A21 is not None: + assert _A22.codomain == _A21.codomain + assert _B1.codomain == _B2.codomain + if _A12 is not None: + assert _A11.domain == _A12.domain == _B1.domain + if _A21 is not None: + assert _A21.domain == _A22.domain == _B2.domain + assert _A22.domain == _B2.domain + assert _A11.domain == _B1.domain + + _blocksA = [[_A11, _A12], [_A21, _A22]] + _A = BlockLinearOperator(self._block_domainA, self._block_codomainA, blocks=_blocksA) + _blocksB = [[_B1, _B2]] + _B = BlockLinearOperator(self._block_domainB, self._block_codomainB, blocks=_blocksB) + if self._lifting: + _blocksF = [ + self._M2.dot(self._F1) + self._M2.dot(u0.vector) / dt - _A11prime.dot(u_prime.vector), + self._M2.dot(self._F2) - _A22prime.dot(ue_prime.vector), + ] + else: + _blocksF = [ + self._M2.dot(self._F1) + self._M2.dot(unfeec) / dt, + self._M2.dot(self._F2), + ] + _F = BlockVector(self._block_domainA, blocks=_blocksF) + + # Imported solver + self._solver_GMRES.A = _A + self._solver_GMRES.B = _B + self._solver_GMRES.F = _F + + if self._lifting: + ( + _sol1, + _sol2, + info, + ) = self._solver_GMRES(u0.vector, ue0.vector, phinfeeccopy.vector) + + un_temp = self.derham.create_spline_function("u", space_id="Hdiv") + un_temp.vector = _sol1[0] + u_prime.vector + + uen_temp = self.derham.create_spline_function("ue", space_id="Hdiv") + uen_temp.vector = _sol1[1] + ue_prime.vector + + phin_temp = self.derham.create_spline_function("phi", space_id="L2") + phin_temp.vector = _sol2 + + un = un_temp.vector + uen = uen_temp.vector + phin = phin_temp.vector - # --- build block RHS and solve --- - _F = BlockVector(self._block_domain_v0, blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) - _RHS = BlockVector(self._block_domain_M, blocks=[_F, self._block_codomain_B_v0.zeros()]) + else: + ( + _sol1, + _sol2, + info, + ) = self._solver_GMRES(unfeec, uenfeec, phinfeec) + un = _sol1[0] + uen = _sol1[1] + phin = _sol2 + # write new coeffs into self.feec_vars + + max_du, max_due, max_dphi = self.update_feec_variables(u=un, ue=uen, phi=phin) + + elif self._variant == "Uzawa": + # Numpy + A11np = self._M2np / dt + self._A11np_notimedependency + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + A11np += self._stab_sigma * xp.identity(A11np.shape[0]) + _A22prenp = self._A22prenp + A22np = self.A22np + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + A11np += self._stab_sigma * sc.sparse.eye(A11np.shape[0], format="csr") + _A22prenp = self._A22prenp + A22np = self.A22np + + # _Anp[1] and _Anppre[1] remain unchanged + _Anp = [A11np, A22np] + if self._preconditioner: + _A11prenp = self._M2np / dt # + self._A11prenp_notimedependency + _Anppre = [_A11prenp, _A22prenp] + + if self._lifting: + # unfeec in space Hdiv, u0 in space Hdiv_0 + unfeeccopy = self.derhamv0.create_spline_function("u", space_id="Hdiv") + u0 = self.derhamv0.create_spline_function("u", space_id="Hdiv") + u_prime = self.derham.create_spline_function("u", space_id="Hdiv") + u0.vector = unfeec + unfeeccopy.vector = unfeec + apply_essential_bc_stencil(u0.vector[0], axis=0, ext=-1, order=0) + apply_essential_bc_stencil(u0.vector[0], axis=0, ext=1, order=0) + u_prime.vector = unfeeccopy.vector - u0.vector + + uenfeeccopy = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue0 = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue_prime = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue0.vector = uenfeec + uenfeeccopy.vector = uenfeec + apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=-1, order=0) + apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=1, order=0) + ue_prime.vector = uenfeeccopy.vector - ue0.vector + + _F1np = ( + self._M2np @ self._F1np + + 1.0 / dt * self._M2np.dot(u0.vector.toarray()) + - self._A11np_notimedependency.dot(u_prime.vector.toarray()) + ) + _F2np = self._M2np @ self._F2np - self.A22np.dot(ue_prime.vector.toarray()) + _Fnp = [_F1np, _F2np] + else: + _F1np = self._M2np @ self._F1np + 1.0 / dt * self._M2np.dot(unfeec.toarray()) + _F2np = self._M2np @ self._F2np + _Fnp = [_F1np, _F2np] + + if self.rank == 0: + if self._preconditioner: + self._solver_UzawaNumpy.Apre = _Anppre + self._solver_UzawaNumpy.A = _Anp + self._solver_UzawaNumpy.F = _Fnp + if self._lifting: + un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy( + u0.vector, + ue0.vector, + phinfeec, + ) - _sol = self._Minv.dot(_RHS) - info = self._Minv.get_info() + un += u_prime.vector.toarray() + uen += ue_prime.vector.toarray() + else: + un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy( + unfeec, + uenfeec, + phinfeec, + ) - # --- reconstruct full solution: u = u_0 + u' --- - self._u.vector = _sol[0][0] + self._u_prime_v0.vector - self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector - self._phi.vector = _sol[1] + dimlist = [[shp - 2 * pi for shp, pi in zip(unfeec[i][:].shape, self.derham.p)] for i in range(3)] + dimphi = [shp - 2 * pi for shp, pi in zip(phinfeec[:].shape, self.derham.p)] + u_temp = BlockVector(self.derham.Vh["2"]) + ue_temp = BlockVector(self.derham.Vh["2"]) + phi_temp = StencilVector(self.derham.Vh["3"]) + test = 0 + for i, bl in enumerate(u_temp.blocks): + s = bl.starts + e = bl.ends + totaldim = dimlist[i][0] * dimlist[i][1] * dimlist[i][2] + test += totaldim + bl[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = un[ + i * totaldim : (i + 1) * totaldim + ].reshape(*dimlist[i]) + + for i, bl in enumerate(ue_temp.blocks): + s = bl.starts + e = bl.ends + totaldim = dimlist[i][0] * dimlist[i][1] * dimlist[i][2] + bl[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = uen[ + i * totaldim : (i + 1) * totaldim + ].reshape(*dimlist[i]) + + s = phi_temp.starts + e = phi_temp.ends + phi_temp[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = phin.reshape(*dimphi) + else: + print("TwoFluidQuasiNeutralFull is only running on one MPI.") - # --- update FEEC variables --- - max_diffs = self.update_feec_variables(u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector) + # write new coeffs into self.feec_vars + max_du, max_due, max_dphi = self.update_feec_variables(u=u_temp, ue=ue_temp, phi=phi_temp) if self.options.solver_params.info and self._rank == 0: logger.info(f"Status: {info['success']}, Iterations: {info['niter']}") From f43e394cf55e126464f0aab53ee8fe219cbf285f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Mon, 8 Dec 2025 09:59:02 +0000 Subject: [PATCH 09/41] Formatting. --- src/struphy/propagators/propagators_fields.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 4392d056a..084dcdc2f 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -8323,7 +8323,7 @@ def __call__(self, dt): self._solver_GMRES.A = _A self._solver_GMRES.B = _B self._solver_GMRES.F = _F - + if self._lifting: ( _sol1, @@ -8333,13 +8333,13 @@ def __call__(self, dt): un_temp = self.derham.create_spline_function("u", space_id="Hdiv") un_temp.vector = _sol1[0] + u_prime.vector - + uen_temp = self.derham.create_spline_function("ue", space_id="Hdiv") uen_temp.vector = _sol1[1] + ue_prime.vector - + phin_temp = self.derham.create_spline_function("phi", space_id="L2") phin_temp.vector = _sol2 - + un = un_temp.vector uen = uen_temp.vector phin = phin_temp.vector From 64e9663a6d50812bd2c142be8c3df5912a6c3079 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Thu, 5 Mar 2026 19:31:54 +0000 Subject: [PATCH 10/41] 1D periodic, homogeneous and inhomogeneous Dirichlet test cases working --- .gitignore | 4 + 1D_Verification.py | 228 +++++++ .../linear_algebra/rework_saddle_point.py | 567 ++++++++++++++++++ src/struphy/models/rework_model.py | 124 ++++ src/struphy/propagators/rework_propagator.py | 479 +++++++++++++++ 5 files changed, 1402 insertions(+) create mode 100644 1D_Verification.py create mode 100644 src/struphy/linear_algebra/rework_saddle_point.py create mode 100644 src/struphy/models/rework_model.py create mode 100644 src/struphy/propagators/rework_propagator.py diff --git a/.gitignore b/.gitignore index 36b889d17..673842aa8 100644 --- a/.gitignore +++ b/.gitignore @@ -98,8 +98,12 @@ src/struphy/io/inp/params_* src/struphy/models/models_list src/struphy/models/models_message +runs/ bin/ share/ lib64 pyvenv.cfg + +2D_Verification.py +Restelli_Verification.py diff --git a/1D_Verification.py b/1D_Verification.py new file mode 100644 index 000000000..baa40a32d --- /dev/null +++ b/1D_Verification.py @@ -0,0 +1,228 @@ +from cunumpy import pi, cos, sin, zeros_like, ones_like +from struphy.io.options import EnvironmentOptions, BaseUnits, Time +from struphy.geometry import domains +from struphy.fields_background import equils +from struphy.topology import grids +from struphy.io.options import DerhamOptions +from struphy.initial import perturbations +from struphy import main + +import os +import glob +import cunumpy as xp +import matplotlib.pyplot as plt + +from struphy.models.rework_model import TwoFluidQuasiNeutralToy + +import warnings +# warnings.filterwarnings("error") + + +BC = 'dirichlet_hom' # 'periodic' | 'dirichlet_hom' | 'dirichlet_inhom' + +name = f"runs/sim_1D_{BC}" + +env = EnvironmentOptions(sim_folder=name) +base_units = BaseUnits(kBT=1.0) + +B0 = 1.0 +nu = 10.0 +nu_e = 1.0 +Nel = (32, 1, 1) +p = (2, 1, 1) +epsilon = 1.0 +dt = 1 +Tend = 1 +sigma = 1 + +time_opts = Time(dt=dt, Tend=Tend) +domain = domains.Cuboid() +equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) +grid = grids.TensorProductGrid(Nel=Nel) + +# ---- boundary conditions ---- +if BC == 'periodic': + spl_kind = (True, True, True) + dirichlet_bc = ((False, False), (False, False), (False, False)) + + bcs_u = bcs_ue = { + (0, -1): "periodic", (0, 1): "periodic", + (1, -1): "periodic", (1, 1): "periodic", + (2, -1): "periodic", (2, 1): "periodic", + } + boundary_data_u = boundary_data_ue = None + +elif BC == 'dirichlet_hom': + spl_kind = (False, True, True) + dirichlet_bc = ((True, True), (False, False), (False, False)) + + bcs_u = bcs_ue = { + (0, -1): "dirichlet", (0, 1): "dirichlet", + (1, -1): "periodic", (1, 1): "periodic", + (2, -1): "periodic", (2, 1): "periodic", + } + boundary_data_u = boundary_data_ue = None + +elif BC == 'dirichlet_inhom': + spl_kind = (False, True, True) + dirichlet_bc = ((False, False), (False, False), (False, False)) + + bcs_u = bcs_ue = { + (0, -1): "dirichlet", (0, 1): "dirichlet", + (1, -1): "periodic", (1, 1): "periodic", + (2, -1): "periodic", (2, 1): "periodic", + } + boundary_data_u = { + (0, -1): lambda x, y, z: (zeros_like(x) + 1, zeros_like(x), zeros_like(x)), + (0, 1): lambda x, y, z: (zeros_like(x) + 2, zeros_like(x), zeros_like(x)), + } + boundary_data_ue = { + (0, -1): lambda x, y, z: (zeros_like(x) + 1, zeros_like(x), zeros_like(x)), + (0, 1): lambda x, y, z: (zeros_like(x) + 2, zeros_like(x), zeros_like(x)), + } + +derham_opts = DerhamOptions( + p=p, + spl_kind=spl_kind, + dirichlet_bc=dirichlet_bc, +) + +# ---- manufactured solutions ---- +if BC == 'periodic': + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x) + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + +elif BC == 'dirichlet_hom': + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + +elif BC == 'dirichlet_inhom': + def mms_phi(x, y, z): + return x + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return x + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return x + 1, xp.zeros_like(x), xp.zeros_like(x) + +# ---- source terms ---- +if BC == 'periodic': + def source_function_u(x, y, z): + fx = 2.0 * pi * cos(2 * pi * x) + nu * 4.0 * pi**2 * sin(2 * pi * x) + fy = (sin(2 * pi * x) + 1.0) * B0 / epsilon + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) + fy = -sin(2 * pi * x) * B0 / epsilon + fz = zeros_like(x) + return fx, fy, fz + +elif BC == 'dirichlet_hom': + def source_function_u(x, y, z): + fx = 2.0 * pi * cos(2 * pi * x) + nu * 4.0 * pi**2 * sin(2 * pi * x) + fy = B0 * sin(2 * pi * x) / epsilon + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) + fy = -sin(2 * pi * x) * B0 / epsilon + fz = zeros_like(x) + return fx, fy, fz + +elif BC == 'dirichlet_inhom': + def source_function_u(x, y, z): + fx = ones_like(x) + fy = B0 * (1 + x) / epsilon + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -ones_like(x) - sigma * (1 + x) + fy = -B0 * (1 + x) / epsilon + fz = zeros_like(x) + return fx, fy, fz + +# ---- model ---- +model = TwoFluidQuasiNeutralToy() +model.ions.set_phys_params() +model.electrons.set_phys_params() + +model.propagators.qn_full.options = model.propagators.qn_full.Options( + nu=nu, + nu_e=nu_e, + eps_norm=epsilon, + stab_sigma=sigma, + source_u=source_function_u, + source_ue=source_function_ue, + solver='gmres', + boundary_conditions_u=bcs_u, + boundary_conditions_ue=bcs_ue, + boundary_data_u=boundary_data_u, + boundary_data_ue=boundary_data_ue, +) + +if __name__ == "__main__": + main.run(model, + params_path=__file__, + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=True, + ) + + path = os.path.join(os.getcwd(), name) + main.pproc(path) + simdata = main.load_data(path) + + n1_vals = simdata.grids_log[0] + x = xp.linspace(0, 1, 100) + + os.makedirs(f'{name}/plots', exist_ok=True) + for f in glob.glob(f'{name}/plots/*.png'): + os.remove(f) + + def save_plot(n1_vals, numerical, analytical, ylabel, title, fname, t): + plt.plot(n1_vals, numerical, label='numerical') + plt.plot(x, analytical, '--', label='manufactured') + plt.plot(n1_vals, numerical, 'k.', markersize=4, label='n1 points') + plt.xlabel('n1 (radial)') + plt.ylabel(ylabel) + plt.title(f'{title} at t={t:.3f}') + plt.legend() + plt.grid(True) + plt.savefig(f'{name}/plots/{fname}_{t:.3f}.png', dpi=300) + plt.clf() + + for t in list(simdata.spline_values['ions']['u_log'].keys()): + + u_ions = simdata.spline_values['ions']['u_log'][t] + u_electrons = simdata.spline_values['electrons']['u_log'][t] + phi = simdata.spline_values['em_fields']['phi_log'][t] + + mms_phi_x, _, _ = mms_phi(x, x*0, x*0) + mms_ion_ux, mms_ion_uy, _ = mms_ion_u(x, x*0, x*0) + mms_el_ux, mms_el_uy, _ = mms_electron_u(x, x*0, x*0) + + save_plot(n1_vals, phi[0][:, 0, 0], mms_phi_x, 'φ', 'Electrostatic potential φ', 'plot_potential', t) + save_plot(n1_vals, u_ions[0][:, 0, 0], mms_ion_ux, 'u_x', 'Ion velocity (u_x)', 'plot_ion_ux', t) + save_plot(n1_vals, u_electrons[0][:, 0, 0], mms_el_ux, 'u_x', 'Electron velocity (u_x)', 'plot_electron_ux', t) \ No newline at end of file diff --git a/src/struphy/linear_algebra/rework_saddle_point.py b/src/struphy/linear_algebra/rework_saddle_point.py new file mode 100644 index 000000000..593dc8524 --- /dev/null +++ b/src/struphy/linear_algebra/rework_saddle_point.py @@ -0,0 +1,567 @@ +from typing import Union + +import cunumpy as xp +import scipy as sc +from psydac.linalg.basic import LinearOperator, Vector +from psydac.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace +from psydac.linalg.direct_solvers import SparseSolver +from psydac.linalg.solvers import inverse + +from struphy.linear_algebra.tests.test_saddlepoint_massmatrices import _plot_residual_norms + + +class SaddlePointSolver: + r"""Solves for :math:`(x, y)` in the saddle point problem + + .. math:: + + \left( \matrix{ + A & B^{\top} \cr + B & 0 + } \right) + \left( \matrix{ + x \cr y + } \right) + = + \left( \matrix{ + f \cr 0 + } \right) + + using either the Uzawa iteration :math:`BA^{-1}B^{\top} y = BA^{-1} f` or using on of the solvers given in :mod:`psydac.linalg.solvers`. The prefered solver is GMRES. + The decission which variant to use is given by the type of A. If A is of type list of xp.ndarrays or sc.sparse.csr_matrices, then this class uses the Uzawa algorithm. + If A is of type LinearOperator or BlockLinearOperator, a solver is used for the inverse. + Using the Uzawa algorithm, solution is given by: + + .. math:: + + y = \left[ B A^{-1} B^{\top}\right]^{-1} B A^{-1} f \,, \qquad + x = A^{-1} \left[ f - B^{\top} y \right] \,. + + Parameters + ---------- + A : list, LinearOperator or BlockLinearOperator + Upper left block. + Either the entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. + Alternative: Give whole matrice A as LinearOperator or BlockLinearOperator. + list: Uzawa algorithm is used. + LinearOperator: A solver given in :mod:`psydac.linalg.solvers` is used. Specified by solver_name. + BlockLinearOperator: A solver given in :mod:`psydac.linalg.solvers` is used. Specified by solver_name. + + B : list, LinearOperator or BlockLinearOperator + Lower left block. + Uzwaw Algorithm: All entries of block B are given either as list of xp.ndarray or sc.sparse.csr_matrix. + Solver: Give whole B as LinearOperator or BlocklinearOperator + + F : list + Right hand side of the upper block. + Uzawa: Given as list of xp.ndarray or sc.sparse.csr_matrix. + Solver: Given as LinearOperator or BlockLinearOperator + + Apre : list + The non-inverted preconditioner for entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. Only required for the Uzawa algorithm. + + method_to_solve : str + Method for the inverses. Choose from 'DirectNPInverse', 'ScipySparse', 'InexactNPInverse' ,'SparseSolver'. Only required for the Uzawa algorithm. + + preconditioner : bool + Wheter to use preconditioners given in Apre or not. Only required for the Uzawa algorithm. + + spectralanalysis : bool + Do the spectralanalyis for the matrices in A and if preconditioner given, compare them to the preconditioned matrices. Only possible if A is given as list. + + dimension : str + Which of the predefined manufactured solutions to use ('1D', '2D' or 'Restelli') + + tol : float + Convergence tolerance for the potential residual. + + max_iter : int + Maximum number of iterations allowed. + """ + + def __init__( + self, + A: Union[list, LinearOperator, BlockLinearOperator], + B: Union[list, LinearOperator, BlockLinearOperator], + F: Union[list, Vector, BlockVector], + Apre: list = None, + method_to_solve: str = "DirectNPInverse", + preconditioner: bool = False, + spectralanalysis: bool = False, + dimension: str = "2D", + solver_name: str = "GMRES", + tol: float = 1e-8, + max_iter: int = 1000, + **solver_params, + ): + assert type(A) is type(B) + if isinstance(A, list): + self._variant = "Uzawa" + for i in A: + assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) + for i in B: + assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) + for i in F: + assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) + for i in Apre: + assert ( + isinstance(i, xp.ndarray) + or isinstance(i, sc.sparse.csr_matrix) + or isinstance(i, sc.sparse.csr_array) + ) + assert method_to_solve in ("SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse") + assert A[0].shape[0] == B[0].shape[1] + assert A[0].shape[1] == B[0].shape[1] + assert A[1].shape[0] == B[1].shape[1] + assert A[1].shape[1] == B[1].shape[1] + + self._method_to_solve = ( + method_to_solve # 'SparseSolver', 'ScipySparse', 'InexactNPInverse', 'DirectNPInverse' + ) + self._preconditioner = preconditioner + + elif isinstance(A, LinearOperator) or isinstance(A, BlockLinearOperator): + self._variant = "Inverse_Solver" + assert A.domain == B.domain + assert A.codomain == B.domain + self._solver_name = solver_name + if solver_params["pc"] is None: + solver_params.pop("pc") + + # operators + self._A = A + self._Apre = Apre + self._B = B + self._F = F + self._tol = tol + self._max_iter = max_iter + self._spectralanalysis = spectralanalysis + self._dimension = dimension + self._verbose = solver_params["verbose"] + + if self._variant == "Inverse_Solver": + self._block_domainM = BlockVectorSpace(self._A.domain, self._B.transpose().domain) + self._block_codomainM = self._block_domainM + self._blocks = [[self._A, self._B.T], [self._B, None]] + _Minit = BlockLinearOperator(self._block_domainM, self._block_codomainM, blocks=self._blocks) + self._solverMinv = inverse(_Minit, solver_name, tol=tol, maxiter=max_iter, **solver_params) + + # Solution vectors + self._P = B.codomain.zeros() + self._U = A.codomain.zeros() + self._Utmp = F.copy() * 0 + # Allocate memory for call + self._rhstemp = BlockVector(self._block_domainM, blocks=[A.codomain.zeros(), self._B.codomain.zeros()]) + + elif self._variant == "Uzawa": + if self._method_to_solve in ("InexactNPInverse", "SparseSolver"): + self._preconditioner = False + + self._Anp = self._A[0] + self._Aenp = self._A[1] + self._B1np = self._B[0] + self._B2np = self._B[1] + + # Instanciate inverses + self._setup_inverses() + + # Solution vectors numpy + self._Pnp = xp.zeros(self._B1np.shape[0]) + self._Unp = xp.zeros(self._A[0].shape[1]) + self._Uenp = xp.zeros(self._A[1].shape[1]) + # Allocate memory for matrices used in solving the system + self._rhs0np = self._F[0].copy() + self._rhs1np = self._F[1].copy() + + # List to store residual norms + self._residual_norms = [] + self._stepsize = 0.0 + + @property + def A(self): + """Upper left block.""" + return self._A + + @A.setter + def A(self, a): + if self._variant == "Uzawa": + need_update = True + A0_old, A1_old = self._A + A0_new, A1_new = a + if self._method_to_solve in ("ScipySparse", "SparseSolver"): + same_A0 = (A0_old != A0_new).nnz == 0 + same_A1 = (A1_old != A1_new).nnz == 0 + else: + same_A0 = xp.allclose(A0_old, A0_new, atol=1e-10) + same_A1 = xp.allclose(A1_old, A1_new, atol=1e-10) + if same_A0 and same_A1: + need_update = False + self._A = a + self._Anp = self._A[0] + self._Aenp = self._A[1] + if need_update: + self._setup_inverses() + elif self._variant == "Inverse_Solver": + self._A = a + + @property + def B(self): + """Lower left block.""" + return self._B + + @B.setter + def B(self, b): + self._B = b + + @property + def F(self): + """Right hand side vector.""" + return self._F + + @F.setter + def F(self, f): + self._F = f + + @property + def Apre(self): + """Preconditioner for upper left block A.""" + return self._Apre + + @Apre.setter + def Apre(self, a): + if self._variant == "Uzawa": + need_update = True + A0_old, A1_old = self._Apre + A0_new, A1_new = a + if self._method_to_solve in ("ScipySparse", "SparseSolver"): + same_A0 = (A0_old != A0_new).nnz == 0 + same_A1 = (A1_old != A1_new).nnz == 0 + else: + same_A0 = xp.allclose(A0_old, A0_new, atol=1e-10) + same_A1 = xp.allclose(A1_old, A1_new, atol=1e-10) + if same_A0 and same_A1: + need_update = False + self._Apre = a + if need_update: + self._setup_inverses() + elif self._variant == "Inverse_Solver": + self._Apre = a + + def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): # todo should have options to use other than uzawa. should solve in full generality. A being block diagonal is a special case of uzawa + """ + Solves the saddle-point problem using the Uzawa algorithm. + + Parameters + ---------- + U_init : Vector, xp.ndarray or sc.sparse.csr.csr_matrix, optional + Initial guess for the velocity of the ions. If None, initializes to zero. Types xp.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. + + Ue_init : Vector, xp.ndarray or sc.sparse.csr.csr_matrix, optional + Initial guess for the velocity of the electrons. If None, initializes to zero. Types xp.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. + + P_init : Vector, optional + Initial guess for the potential. If None, initializes to zero. + + Returns + ------- + U : Vector + Solution vector for the velocity. + + P : Vector + Solution vector for the potential. + + info : dict + Convergence information. + """ + + # TODO this contains two different strategies! favágás and actual uzawa + if self._variant == "Inverse_Solver": # todo not in the """""saddle point solver""""""" + self._P1 = P_init if P_init is not None else self._P + self._U1 = U_init if U_init is not None else self._Utmp[0] + self._U2 = Ue_init if Ue_init is not None else self._Utmp[1] + + _blocksM = [[self._A, self._B.T], [self._B, None]] + _M = BlockLinearOperator(self._block_domainM, self._block_codomainM, blocks=_blocksM) + _RHS = BlockVector(self._block_domainM, blocks=[self._F, self._B.codomain.zeros()]) + + self._blockU = BlockVector(self._A.domain, blocks=[self._U1, self._U2]) + self._solblocks = [self._blockU, self._P1] + # comment out the next two lines if working with lifting and GMRES + x0 = BlockVector(self._block_domainM, blocks=self._solblocks) + self._solverMinv._options["x0"] = x0 + + # use setter to update lhs matrix + self._solverMinv.linop = _M + + # Initialize P to zero or given initial guess + self._sol = self._solverMinv.dot(_RHS, out=self._rhstemp) + self._U = self._sol[0] + self._P = self._sol[1] + + return self._U, self._P, self._solverMinv._info + + elif self._variant == "Uzawa": + info = {} + + if self._spectralanalysis: + self._spectralresult = self._spectral_analysis() + else: + self._spectralresult = [] + + # Initialize P to zero or given initial guess + if isinstance(U_init, xp.ndarray) or isinstance(U_init, sc.sparse.csr.csr_matrix): + self._Pnp = P_init if P_init is not None else self._P + self._Unp = U_init if U_init is not None else self._U + self._Uenp = Ue_init if U_init is not None else self._Ue + + else: + self._Pnp = P_init.toarray() if P_init is not None else self._Pnp + self._Unp = U_init.toarray() if U_init is not None else self._Unp + self._Uenp = Ue_init.toarray() if U_init is not None else self._Uenp + + if self._verbose: + print("Uzawa solver:") + print("+---------+---------------------+") + print("+ Iter. # | L2-norm of residual |") + print("+---------+---------------------+") + template = "| {:7d} | {:19.2e} |" + + for iteration in range(self._max_iter): + # Step 1: Compute velocity U by solving A U = -Bᵀ P + F -A Un + self._rhs0np *= 0 + self._rhs0np -= self._B1np.transpose().dot(self._Pnp) + self._rhs0np -= self._Anp.dot(self._Unp) + self._rhs0np += self._F[0] + if not self._preconditioner: + self._Unp += self._Anpinv.dot(self._rhs0np) + elif self._preconditioner: + self._Unp += self._Anpinv.dot(self._A11npinv @ self._rhs0np) + + R1 = self._B1np.dot(self._Unp) + + self._rhs1np *= 0 + self._rhs1np -= self._B2np.transpose().dot(self._Pnp) + self._rhs1np -= self._Aenp.dot(self._Uenp) + self._rhs1np += self._F[1] + if not self._preconditioner: + self._Uenp += self._Aenpinv.dot(self._rhs1np) + elif self._preconditioner: + self._Uenp += self._Aenpinv.dot(self._A22npinv @ self._rhs1np) + + R2 = self._B2np.dot(self._Uenp) + + # Step 2: Compute residual R = BU (divergence of U) + R = R1 + R2 # self._B1np.dot(self._Unp) + self._B2np.dot(self._Uenp) + residual_norm = xp.linalg.norm(R) + residual_normR1 = xp.linalg.norm(R) + self._residual_norms.append(residual_normR1) # Store residual norm + # Check for convergence based on residual norm + if residual_norm < self._tol: + if self._verbose: + print(template.format(iteration + 1, residual_norm)) + print("+---------+---------------------+") + info["success"] = True + info["niter"] = iteration + 1 + if self._verbose: + _plot_residual_norms(self._residual_norms) + return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult + + # Steepest gradient + alpha = (R.dot(R)) / (R.dot(self._Precnp.dot(R))) + # Minimal residual + # alpha = ((self._Precnp.dot(R)).dot(R)) / ((self._Precnp.dot(R)).dot(self._Precnp.dot(R))) + self._Pnp += alpha.real * R.real + + if self._verbose: + print(template.format(iteration + 1, residual_norm)) + + if self._verbose: + print("+---------+---------------------+") + + # Return with info if maximum iterations reached + info["success"] = False + info["niter"] = iteration + 1 + if self._verbose: + _plot_residual_norms(self._residual_norms) + return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult + + def _setup_inverses(self): + A0 = self._A[0] + A1 = self._A[1] + + # === Preconditioner inverses, if used + if self._preconditioner: + A11_pre = self._Apre[0] + A22_pre = self._Apre[1] + + if hasattr(self, "_A11npinv") and self._is_inverse_still_valid(self._A11npinv, A11_pre, "A11 pre"): + pass + else: + self._A11npinv = self._compute_inverse(A11_pre, which="A11 pre") + + if hasattr(self, "_A22npinv") and self._is_inverse_still_valid(self._A22npinv, A22_pre, "A22 pre"): + pass + else: + self._A22npinv = self._compute_inverse(A22_pre, which="A22 pre") + + # === Inverse for A[0] if preconditioned + if hasattr(self, "_Anpinv") and self._is_inverse_still_valid(self._Anpinv, A0, "A[0]", pre=self._A11npinv): + pass + else: + self._Anpinv = self._compute_inverse(self._A11npinv @ A0, which="A[0]") + + # === Inverse for A[1] + if hasattr(self, "_Aenpinv") and self._is_inverse_still_valid( + self._Aenpinv, + A1, + "A[1]", + pre=self._A22npinv, + ): + pass + else: + self._Aenpinv = self._compute_inverse(self._A22npinv @ A1, which="A[1]") + + else: # No preconditioning: + # === Inverse for A[0] + if hasattr(self, "_Anpinv") and self._is_inverse_still_valid(self._Anpinv, A0, "A[0]"): + pass + else: + self._Anpinv = self._compute_inverse(A0, which="A[0]") + + # === Inverse for A[1] + if hasattr(self, "_Aenpinv") and self._is_inverse_still_valid(self._Aenpinv, A1, "A[1]"): + pass + else: + self._Aenpinv = self._compute_inverse(A1, which="A[1]") + + # Precompute Schur complement + self._Precnp = self._B1np @ self._Anpinv @ self._B1np.T + self._B2np @ self._Aenpinv @ self._B2np.T + + def _is_inverse_still_valid(self, inv, mat, name="", pre=None): + # try: + if pre is not None: + test_mat = pre @ mat + else: + test_mat = mat + I_approx = inv @ test_mat + + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + I_exact = xp.eye(test_mat.shape[0]) + if not xp.allclose(I_approx, I_exact, atol=1e-6): + diff = I_approx - I_exact + max_abs = xp.abs(diff).max() + print(f"{name} inverse is NOT valid anymore. Max diff: {max_abs:.2e}") + return False + print(f"{name} inverse is still valid.") + return True + elif self._method_to_solve == "ScipySparse": + I_exact = sc.sparse.identity(I_approx.shape[0], format=I_approx.format) + diff = (I_approx - I_exact).tocoo() + max_abs = xp.abs(diff.data).max() if diff.nnz > 0 else 0.0 + + if max_abs > 1e-6: + print(f"{name} inverse is NOT valid anymore.") + print(f"Max absolute difference: {max_abs:.2e}") + print(f"Number of differing entries: {diff.nnz}") + return False + print(f"{name} inverse is still valid.") + return True + + def _compute_inverse(self, mat, which="matrix"): + print(f"Computing inverse for {which} using method {self._method_to_solve}") + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + return xp.linalg.inv(mat) + elif self._method_to_solve == "ScipySparse": + return sc.sparse.linalg.inv(mat) + elif self._method_to_solve == "SparseSolver": + solver = SparseSolver(mat) + return solver.solve(xp.eye(mat.shape[0])) + else: + raise ValueError(f"Unknown solver method {self._method_to_solve}") + + def _spectral_analysis(self): + # Spectral analysis + # A11 before + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + eigvalsA11_before, eigvecs_before = xp.linalg.eig(self._A[0]) + condA11_before = xp.linalg.cond(self._A[0]) + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + eigvalsA11_before, eigvecs_before = xp.linalg.eig(self._A[0].toarray()) + condA11_before = xp.linalg.cond(self._A[0].toarray()) + maxbeforeA11 = max(eigvalsA11_before) + maxbeforeA11_abs = xp.max(xp.abs(eigvalsA11_before)) + minbeforeA11_abs = xp.min(xp.abs(eigvalsA11_before)) + minbeforeA11 = min(eigvalsA11_before) + specA11_bef = maxbeforeA11 / minbeforeA11 + specA11_bef_abs = maxbeforeA11_abs / minbeforeA11_abs + # print(f'{maxbeforeA11 = }') + # print(f'{maxbeforeA11_abs = }') + # print(f'{minbeforeA11_abs = }') + # print(f'{minbeforeA11 = }') + # print(f'{specA11_bef = }') + print(f"{specA11_bef_abs =}") + + # A22 before + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + eigvalsA22_before, eigvecs_before = xp.linalg.eig(self._A[1]) + condA22_before = xp.linalg.cond(self._A[1]) + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + eigvalsA22_before, eigvecs_before = xp.linalg.eig(self._A[1].toarray()) + condA22_before = xp.linalg.cond(self._A[1].toarray()) + maxbeforeA22 = max(eigvalsA22_before) + maxbeforeA22_abs = xp.max(xp.abs(eigvalsA22_before)) + minbeforeA22_abs = xp.min(xp.abs(eigvalsA22_before)) + minbeforeA22 = min(eigvalsA22_before) + specA22_bef = maxbeforeA22 / minbeforeA22 + specA22_bef_abs = maxbeforeA22_abs / minbeforeA22_abs + # print(f'{maxbeforeA22 = }') + # print(f'{maxbeforeA22_abs = }') + # print(f'{minbeforeA22_abs = }') + # print(f'{minbeforeA22 = }') + # print(f'{specA22_bef = }') + print(f"{specA22_bef_abs =}") + print(f"{condA22_before =}") + + if self._preconditioner: + # A11 after preconditioning with its inverse + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig(self._A11npinv @ self._A[0]) # Implement this + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig((self._A11npinv @ self._A[0]).toarray()) + maxafterA11_prec = max(eigvalsA11_after_prec) + minafterA11_prec = min(eigvalsA11_after_prec) + maxafterA11_abs_prec = xp.max(xp.abs(eigvalsA11_after_prec)) + minafterA11_abs_prec = xp.min(xp.abs(eigvalsA11_after_prec)) + specA11_aft_prec = maxafterA11_prec / minafterA11_prec + specA11_aft_abs_prec = maxafterA11_abs_prec / minafterA11_abs_prec + # print(f'{maxafterA11_prec = }') + # print(f'{maxafterA11_abs_prec = }') + # print(f'{minafterA11_abs_prec = }') + # print(f'{minafterA11_prec = }') + # print(f'{specA11_aft_prec = }') + print(f"{specA11_aft_abs_prec =}") + + # A22 after preconditioning with its inverse + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + eigvalsA22_after_prec, eigvecs_after = xp.linalg.eig(self._A22npinv @ self._A[1]) # Implement this + condA22_after = xp.linalg.cond(self._A22npinv @ self._A[1]) + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + eigvalsA22_after_prec, eigvecs_after = xp.linalg.eig((self._A22npinv @ self._A[1]).toarray()) + condA22_after = xp.linalg.cond((self._A22npinv @ self._A[1]).toarray()) + maxafterA22_prec = max(eigvalsA22_after_prec) + minafterA22_prec = min(eigvalsA22_after_prec) + maxafterA22_abs_prec = xp.max(xp.abs(eigvalsA22_after_prec)) + minafterA22_abs_prec = xp.min(xp.abs(eigvalsA22_after_prec)) + specA22_aft_prec = maxafterA22_prec / minafterA22_prec + specA22_aft_abs_prec = maxafterA22_abs_prec / minafterA22_abs_prec + # print(f'{maxafterA22_prec = }') + # print(f'{maxafterA22_abs_prec = }') + # print(f'{minafterA22_abs_prec = }') + # print(f'{minafterA22_prec = }') + # print(f'{specA22_aft_prec = }') + print(f"{specA22_aft_abs_prec =}") + + return condA22_before, specA22_bef_abs, condA11_before, condA22_after, specA22_aft_abs_prec + + else: + return condA22_before, specA22_bef_abs, condA11_before diff --git a/src/struphy/models/rework_model.py b/src/struphy/models/rework_model.py new file mode 100644 index 000000000..f38629fcc --- /dev/null +++ b/src/struphy/models/rework_model.py @@ -0,0 +1,124 @@ +import cunumpy as xp +from psydac.ddm.mpi import mpi as MPI + +from struphy.feec.projectors import L2Projector +from struphy.feec.variational_utilities import InternalEnergyEvaluator +from struphy.models.base import StruphyModel +from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable +from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers +from struphy.propagators import rework_propagator + + +rank = MPI.COMM_WORLD.Get_rank() + +class TwoFluidQuasiNeutralToy(StruphyModel): + r"""Linearized, quasi-neutral two-fluid model with zero electron inertia. + + :ref:`normalization`: + + .. math:: + + \hat u = \hat v_\textnormal{th}\,,\qquad e\hat \phi = m \hat v_\textnormal{th}^2\,. + + :ref:`Equations `: + + .. math:: + + \frac{\partial \mathbf u}{\partial t} &= - \nabla \phi + \frac{\mathbf u \times \mathbf B_0}{\varepsilon} + \nu \Delta \mathbf u + \mathbf f\,, + \\[2mm] + 0 &= \nabla \phi - \frac{\mathbf u_e \times \mathbf B_0}{\varepsilon} + \nu_e \Delta \mathbf u_e + \mathbf f_e \,, + \\[3mm] + \nabla & \cdot (\mathbf u - \mathbf u_e) = 0\,, + + where :math:`\mathbf B_0` is a static magnetic field and :math:`\mathbf f, \mathbf f_e` are given forcing terms, + and with the normalization parameter + + .. math:: + + \varepsilon = \frac{1}{\hat \Omega_\textnormal{c} \hat t} \,,\qquad \textnormal{with} \,,\qquad \hat \Omega_{\textnormal{c}} = \frac{(Ze) \hat B}{(A m_\textnormal{H})}\,, + + :ref:`propagators` (called in sequence): + + 1. :class:`~struphy.propagators.propagators_fields.TwoFluidQuasiNeutralFull` + + :ref:`Model info `: + + References + ---------- + [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising + in plasma physics, Journal of Computational Physics 2018. + """ + + ## species + + class EMfields(FieldSpecies): + def __init__(self): + self.phi = FEECVariable(space="L2") + self.init_variables() + + class Ions(FluidSpecies): + def __init__(self): + self.u = FEECVariable(space="Hdiv") + self.init_variables() + + class Electrons(FluidSpecies): + def __init__(self): + self.u = FEECVariable(space="Hdiv") + self.init_variables() + + ## propagators + + class Propagators: + def __init__(self): + self.qn_full = rework_propagator.TwoFluidQuasiNeutralFull() + + ## abstract methods + + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMfields() + self.ions = self.Ions() + self.electrons = self.Electrons() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.qn_full.variables.u = self.ions.u + self.propagators.qn_full.variables.ue = self.electrons.u + self.propagators.qn_full.variables.phi = self.em_fields.phi + + # define scalars for update_scalar_quantities + + @property + def bulk_species(self): + return self.ions + + @property + def velocity_scale(self): + return "thermal" + + def allocate_helpers(self): + pass + + def update_scalar_quantities(self): + pass + + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "BaseUnits()" in line: + new_file += ["base_units = BaseUnits(kBT=1.0)\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) diff --git a/src/struphy/propagators/rework_propagator.py b/src/struphy/propagators/rework_propagator.py new file mode 100644 index 000000000..bf6afcadf --- /dev/null +++ b/src/struphy/propagators/rework_propagator.py @@ -0,0 +1,479 @@ + +import copy +from copy import deepcopy +from dataclasses import dataclass +from typing import Callable, Literal, get_args, cast +from warnings import warn + +import cunumpy as xp +import scipy as sc +from line_profiler import profile +from matplotlib import pyplot as plt +from numpy import zeros +from psydac.api.essential_bc import apply_essential_bc_stencil +from psydac.ddm.mpi import mpi as MPI +from psydac.linalg.basic import ComposedLinearOperator, IdentityOperator, ZeroOperator, InverseLinearOperator +from psydac.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace +from psydac.linalg.solvers import inverse +from psydac.linalg.stencil import StencilVector + +import struphy.feec.utilities as util +from struphy.examples.restelli2018 import callables +from struphy.feec import preconditioner +from struphy.feec.basis_projection_ops import ( + BasisProjectionOperator, BasisProjectionOperatorLocal, + BasisProjectionOperators, CoordinateProjector, +) +from struphy.feec.linear_operators import BoundaryOperator +from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators +from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner, MassMatrixPreconditioner +from struphy.feec.projectors import L2Projector +from struphy.feec.psydac_derham import Derham, SplineFunction +from struphy.feec.variational_utilities import ( + BracketOperator, Hdiv0_transport_operator, InternalEnergyEvaluator, + KineticEnergyEvaluator, Pressure_transport_operator, +) +from struphy.fields_background.equils import set_defaults +from struphy.geometry.utilities import TransformedPformComponent +from struphy.initial import perturbations +from struphy.io.options import ( + OptsDirectSolver, OptsGenSolver, OptsMassPrecond, OptsNonlinearSolver, + OptsSaddlePointSolver, OptsSymmSolver, OptsVecSpace, check_option, +) +from struphy.io.setup import descend_options_dict +from struphy.kinetic_background.base import Maxwellian +from struphy.kinetic_background.maxwellians import GyroMaxwellian2D, Maxwellian3D +from struphy.linear_algebra.saddle_point import SaddlePointSolver +from struphy.linear_algebra.schur_solver import SchurSolver, SchurSolverFull +from struphy.linear_algebra.solver import NonlinearSolverParameters, SolverParameters +from struphy.models.species import Species +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable +from struphy.ode.solvers import ODEsolverFEEC +from struphy.ode.utils import ButcherTableau, OptsButcher +from struphy.pic.accumulation import accum_kernels, accum_kernels_gc +from struphy.pic.accumulation.filter import FilterParameters +from struphy.pic.accumulation.particles_to_grid import Accumulator, AccumulatorVector +from struphy.pic.base import Particles +from struphy.pic.particles import Particles5D, Particles6D +from struphy.polar.basic import PolarVector +from struphy.propagators.base import Propagator +from struphy.utils.pyccel import Pyccelkernel + + +class TwoFluidQuasiNeutralFull(Propagator): + r""":ref:`FEEC ` discretization of the following equations: + find :math:`\mathbf u \in H(\textnormal{div})`, :math:`\mathbf u_e \in H(\textnormal{div})` and :math:`\mathbf \phi \in L^2` such that + + .. math:: + + \int_{\Omega} \partial_t \mathbf{u}\cdot \mathbf{v} \, \textrm d\mathbf{x} &= \int_{\Omega} \phi \nabla \! \cdot \! \mathbf{v} \, \textrm d\mathbf{x} + \int_{\Omega} \mathbf{u}\! \times \! \mathbf{B}_0 \cdot \mathbf{v} \, \textrm d\mathbf{x} + \nu \int_{\Omega} \nabla \mathbf{u}\! : \! \nabla \mathbf{v} \, \textrm d\mathbf{x} + \int_{\Omega} f \mathbf{v} \, \textrm d\mathbf{x} \qquad \forall \, \mathbf{v} \in H(\textrm{div}) \,. + \\[2mm] + 0 &= - \int_{\Omega} \phi \nabla \! \cdot \! \mathbf{v_e} \, \textrm d\mathbf{x} - \int_{\Omega} \mathbf{u_e} \! \times \! \mathbf{B}_0 \cdot \mathbf{v_e} \, \textrm d\mathbf{x} + \nu_e \int_{\Omega} \nabla \mathbf{u_e} \!: \! \nabla \mathbf{v_e} \, \textrm d\mathbf{x} + \int_{\Omega} f_e \mathbf{v_e} \, \textrm d\mathbf{x} \qquad \forall \ \mathbf{v_e} \in H(\textrm{div}) \,. + \\[2mm] + 0 &= \int_{\Omega} \psi \nabla \cdot (\mathbf{u}-\mathbf{u_e}) \, \textrm d\mathbf{x} \qquad \forall \, \psi \in L^2 \,. + + :ref:`time_discret`: fully implicit. + """ + + # ========================================================================= + ### State variables (ion velocity u, electron velocity ue, pressure phi) + # ========================================================================= + + class Variables(): + def __init__(self) -> None: + self._u: FEECVariable | None = None + self._ue: FEECVariable | None = None + self._phi: FEECVariable | None = None + + @property + def u(self) -> FEECVariable | None: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._u = new + + @property + def ue(self) -> FEECVariable | None: + return self._ue + + @ue.setter + def ue(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._ue = new + + @property + def phi(self) -> FEECVariable | None: + return self._phi + + @phi.setter + def phi(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._phi = new + + def __init__(self): + self.variables = self.Variables() + + # ========================================================================= + ### Options + # ========================================================================= + + @dataclass + class Options(): + + nu: float | None = None + nu_e: float | None = None + eps_norm: float | None = None + + # boundary conditions per species + # supported kinds: "periodic", "dirichlet" + # future: "neumann", "robin" + boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_data_u: dict[tuple[int, int], Callable] | None = None + boundary_data_ue: dict[tuple[int, int], Callable] | None = None + + source_u: Callable | None = None + source_ue: Callable | None = None + + stab_sigma: float | None = None + + solver: OptsGenSolver = "gmres" + solver_params: SolverParameters | None = None + + def __post_init__(self): + + # --- required parameters --- + assert self.nu is not None, "nu must be specified" + assert self.nu_e is not None, "nu_e must be specified" + assert self.eps_norm is not None, "eps_norm must be specified" + assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" + assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" + + # --- physical parameter sanity checks --- + if self.nu < 0: + raise ValueError(f"nu must be non-negative, got {self.nu}") + if self.nu_e < 0: + raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") + if self.eps_norm <= 0: + raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") + + # --- check all axes are covered --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for d in range(3): + for side in (-1, 1): + assert (d, side) in bcs, \ + f"{name} is missing entry for axis {d} side {side}" + + # --- periodic consistency: periodic must be paired on both sides --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for (d, side), kind in bcs.items(): + if kind == "periodic": + assert bcs.get((d, -side)) == "periodic", \ + f"{name}: axis {d} side {side} is periodic but opposite side is not" + + # --- ions and electrons must agree on which axes are periodic --- + for d in range(3): + u_left = self.boundary_conditions_u.get((d, -1)) + ue_left = self.boundary_conditions_ue.get((d, -1)) + u_right = self.boundary_conditions_u.get((d, 1)) + ue_right = self.boundary_conditions_ue.get((d, 1)) + u_periodic = (u_left == "periodic") + ue_periodic = (ue_left == "periodic") + if u_periodic != ue_periodic: + raise ValueError( + f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " + f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" + ) + + # --- warn for Dirichlet faces with no boundary data --- + for species, bcs, data, label in [ + ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), + ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), + ]: + has_dirichlet = any(v == "dirichlet" for v in bcs.values()) + if has_dirichlet: + if data is None: + warn(f"Dirichlet BCs specified for {species} but no {label} given " + f"— defaulting to homogeneous Dirichlet on all faces.") + else: + for (d, side), kind in bcs.items(): + if kind == "dirichlet" and (d, side) not in data: + warn(f"No {label} given for axis {d} side {side} " + f"— defaulting to homogeneous Dirichlet.") + + # --- warn if no source terms --- + if self.source_u is None: + warn("No source_u specified — defaulting to zero.") + if self.source_ue is None: + warn("No source_ue specified — defaulting to zero.") + + # --- defaults --- + if self.stab_sigma is None: + warn("stab_sigma not specified, defaulting to 0.0") + self.stab_sigma = 0.0 + + check_option(self.solver, OptsGenSolver) + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + assert hasattr(self, "_options"), "Options not set." + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + # ========================================================================= + ### Boundary condition helpers + # ========================================================================= + + def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): + + dirichlet_bc = [] + for d in range(3): + if spl_kind[d]: # periodic spline — no clamping ever + dirichlet_bc.append((False, False)) + else: + left = boundary_conditions.get((d, -1)) == "dirichlet" + right = boundary_conditions.get((d, 1)) == "dirichlet" + dirichlet_bc.append((left, right)) + return tuple(tuple(bc) for bc in dirichlet_bc) + + def _apply_boundary_conditions(self, vec, boundary_conditions): + """Zero out Dirichlet DOFs on the given stencil vector.""" + for (d, side), kind in boundary_conditions.items(): + if kind == "dirichlet": + apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) + # future: neumann and robin require no zeroing here + + # ========================================================================= + ### Allocate + # ========================================================================= + + def allocate(self): + + self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 + self._dt = None + + # ---- constrained (v0) de Rham complex -------------------------------- + + _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) + _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) + _dirichlet_bc = tuple( + (l_u or l_ue, r_u or r_ue) + for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) + ) + + self._derham_v0 = Derham( + self.derham.Nel, self.derham.p, self.derham.spl_kind, + domain=self.domain, dirichlet_bc=_dirichlet_bc, + ) + self._mass_ops_v0 = WeightedMassOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.mass_ops.weights["eq_mhd"], + ) + self._basis_ops_v0 = BasisProjectionOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.basis_ops.weights["eq_mhd"], + ) + + # ---- unconstrained operators (for RHS assembly) ---------------------- + + self._M2 = self.mass_ops.M2 + self._M2B = - self.mass_ops.M2B + self._div = self.derham.div + self._curl = self.derham.curl + self._S21 = self.basis_ops.S21 + + self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div + + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + + self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl + self._A22 = (- self.options.stab_sigma * IdentityOperator(self._A11.domain) + + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) + + # ---- constrained operators (for system matrix) ----------------------- + + self._M2_v0 = self._mass_ops_v0.M2 + self._M3_v0 = self._mass_ops_v0.M3 + self._M2B_v0 = - self._mass_ops_v0.M2B + self._div_v0 = self._derham_v0.div + self._curl_v0 = self._derham_v0.curl + self._S21_v0 = self._basis_ops_v0.S21 + + self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 + + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) + + self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 + self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._A11_v0.domain) + + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) + + # ---- block saddle-point system ---------------------------------------- + + self._block_domain_v0 = BlockVectorSpace(self._A11_v0.domain, self._A22_v0.domain) + self._block_codomain_v0 = self._block_domain_v0 + self._block_codomain_B_v0 = self._M3_v0.codomain + + self._B1_v0 = - self._M3_v0 @ self._div_v0 + self._B2_v0 = self._M3_v0 @ self._div_v0 + + self._B_v0 = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_B_v0, + blocks=[[self._B1_v0, self._B2_v0]] + ) + + self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) + + _A_init = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0, None], [None, self._A22_v0]] + ) + _M_init = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] + ) + self._Minv = cast(InverseLinearOperator, inverse( + _M_init, self.options.solver, + x0=None, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + )) + + # ---- projector ------------------------------------------------------- + + self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + + # ---- solution spline functions (unconstrained) ----------------------- + + self._u = self.derham.create_spline_function("u", space_id="Hdiv") + self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") + self._phi = self.derham.create_spline_function("phi", space_id="L2") + + # ---- BC lifts (unconstrained) ---------------------------------------- + + self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") + self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") + + for u_prime, boundary_data, boundary_conditions in [ + (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), + (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), + ]: + if boundary_data is None: + continue + for (d, side), f_bc in boundary_data.items(): + if boundary_conditions.get((d, side)) == "dirichlet": + bc_pulled = lambda *etas, f=f_bc: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], + lambda *etas: bc_pulled(*etas)[1], + lambda *etas: bc_pulled(*etas)[2]]) + for (d2, side2), kind2 in boundary_conditions.items(): + if kind2 == "dirichlet" and (d2, side2) != (d, side): + apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) + u_prime.vector += _vec + + self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") + self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") + + self._u_prime_v0.vector = self._u_prime.vector + self._ue_prime_v0.vector = self._ue_prime.vector + + # ---- projected source terms (unconstrained) -------------------------- + + self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") + self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") + + for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: + if source is not None: + src_pulled = lambda *etas, f=source: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], + lambda *etas: src_pulled(*etas)[1], + lambda *etas: src_pulled(*etas)[2]]) + + # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- + + self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") + + # ========================================================================= + ### Time step + # ========================================================================= + + def __call__(self, dt): + + # --- copy current state --- + self._u.vector = self.variables.u.spline.vector + self._ue.vector = self.variables.ue.spline.vector + + # --- rebuild system matrix if dt changed --- + if dt != self._dt: + self._dt = dt + _A = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] + ) + _M = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A, self._B_v0.T], [self._B_v0, None]] + ) + self._Minv.linop = _M + + # --- assemble RHS in unconstrained space, then zero boundary DOFs --- + # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' + # electron: F2 = rhs_ue - A22 * ue' + self._rhs_vec_u.vector = (self._rhs_u.vector + + self._M2.dot(self._u.vector) / dt + - self._A11.dot(self._u_prime.vector) + - self._M2.dot(self._u_prime.vector) / dt) + self._rhs_vec_ue.vector = (self._rhs_ue.vector + - self._A22.dot(self._ue_prime.vector)) + + self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) + self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) + + # --- build block RHS and solve --- + _F = BlockVector(self._block_domain_v0, + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) + _RHS = BlockVector(self._block_domain_M, + blocks=[_F, self._block_codomain_B_v0.zeros()]) + + _sol = self._Minv.dot(_RHS) + info = self._Minv.get_info() + + # --- reconstruct full solution: u = u_0 + u' --- + self._u.vector = _sol[0][0] + self._u_prime_v0.vector + self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector + self._phi.vector = _sol[1] + + # --- update FEEC variables --- + max_diffs = self.update_feec_variables( + u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector + ) + + if self.options.solver_params.info and self._rank == 0: + print(f"Status: {info['success']}, Iterations: {info['niter']}") + print(f"Max diffs: {max_diffs}") \ No newline at end of file From dc43aaadb2011707556e0d0b779999af371b134b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Thu, 5 Mar 2026 21:39:53 +0000 Subject: [PATCH 11/41] Rebased onto version 3.0.4 --- .gitignore | 3 - 1D_Verification.py | 228 --------------------------------------------- 2 files changed, 231 deletions(-) delete mode 100644 1D_Verification.py diff --git a/.gitignore b/.gitignore index 673842aa8..daf35eef6 100644 --- a/.gitignore +++ b/.gitignore @@ -104,6 +104,3 @@ share/ lib64 pyvenv.cfg - -2D_Verification.py -Restelli_Verification.py diff --git a/1D_Verification.py b/1D_Verification.py deleted file mode 100644 index baa40a32d..000000000 --- a/1D_Verification.py +++ /dev/null @@ -1,228 +0,0 @@ -from cunumpy import pi, cos, sin, zeros_like, ones_like -from struphy.io.options import EnvironmentOptions, BaseUnits, Time -from struphy.geometry import domains -from struphy.fields_background import equils -from struphy.topology import grids -from struphy.io.options import DerhamOptions -from struphy.initial import perturbations -from struphy import main - -import os -import glob -import cunumpy as xp -import matplotlib.pyplot as plt - -from struphy.models.rework_model import TwoFluidQuasiNeutralToy - -import warnings -# warnings.filterwarnings("error") - - -BC = 'dirichlet_hom' # 'periodic' | 'dirichlet_hom' | 'dirichlet_inhom' - -name = f"runs/sim_1D_{BC}" - -env = EnvironmentOptions(sim_folder=name) -base_units = BaseUnits(kBT=1.0) - -B0 = 1.0 -nu = 10.0 -nu_e = 1.0 -Nel = (32, 1, 1) -p = (2, 1, 1) -epsilon = 1.0 -dt = 1 -Tend = 1 -sigma = 1 - -time_opts = Time(dt=dt, Tend=Tend) -domain = domains.Cuboid() -equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) -grid = grids.TensorProductGrid(Nel=Nel) - -# ---- boundary conditions ---- -if BC == 'periodic': - spl_kind = (True, True, True) - dirichlet_bc = ((False, False), (False, False), (False, False)) - - bcs_u = bcs_ue = { - (0, -1): "periodic", (0, 1): "periodic", - (1, -1): "periodic", (1, 1): "periodic", - (2, -1): "periodic", (2, 1): "periodic", - } - boundary_data_u = boundary_data_ue = None - -elif BC == 'dirichlet_hom': - spl_kind = (False, True, True) - dirichlet_bc = ((True, True), (False, False), (False, False)) - - bcs_u = bcs_ue = { - (0, -1): "dirichlet", (0, 1): "dirichlet", - (1, -1): "periodic", (1, 1): "periodic", - (2, -1): "periodic", (2, 1): "periodic", - } - boundary_data_u = boundary_data_ue = None - -elif BC == 'dirichlet_inhom': - spl_kind = (False, True, True) - dirichlet_bc = ((False, False), (False, False), (False, False)) - - bcs_u = bcs_ue = { - (0, -1): "dirichlet", (0, 1): "dirichlet", - (1, -1): "periodic", (1, 1): "periodic", - (2, -1): "periodic", (2, 1): "periodic", - } - boundary_data_u = { - (0, -1): lambda x, y, z: (zeros_like(x) + 1, zeros_like(x), zeros_like(x)), - (0, 1): lambda x, y, z: (zeros_like(x) + 2, zeros_like(x), zeros_like(x)), - } - boundary_data_ue = { - (0, -1): lambda x, y, z: (zeros_like(x) + 1, zeros_like(x), zeros_like(x)), - (0, 1): lambda x, y, z: (zeros_like(x) + 2, zeros_like(x), zeros_like(x)), - } - -derham_opts = DerhamOptions( - p=p, - spl_kind=spl_kind, - dirichlet_bc=dirichlet_bc, -) - -# ---- manufactured solutions ---- -if BC == 'periodic': - def mms_phi(x, y, z): - return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) - - def mms_ion_u(x, y, z): - return xp.sin(2 * xp.pi * x) + 1, xp.zeros_like(x), xp.zeros_like(x) - - def mms_electron_u(x, y, z): - return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) - -elif BC == 'dirichlet_hom': - def mms_phi(x, y, z): - return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) - - def mms_ion_u(x, y, z): - return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) - - def mms_electron_u(x, y, z): - return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) - -elif BC == 'dirichlet_inhom': - def mms_phi(x, y, z): - return x + 1, xp.zeros_like(x), xp.zeros_like(x) - - def mms_ion_u(x, y, z): - return x + 1, xp.zeros_like(x), xp.zeros_like(x) - - def mms_electron_u(x, y, z): - return x + 1, xp.zeros_like(x), xp.zeros_like(x) - -# ---- source terms ---- -if BC == 'periodic': - def source_function_u(x, y, z): - fx = 2.0 * pi * cos(2 * pi * x) + nu * 4.0 * pi**2 * sin(2 * pi * x) - fy = (sin(2 * pi * x) + 1.0) * B0 / epsilon - fz = zeros_like(x) - return fx, fy, fz - - def source_function_ue(x, y, z): - fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) - fy = -sin(2 * pi * x) * B0 / epsilon - fz = zeros_like(x) - return fx, fy, fz - -elif BC == 'dirichlet_hom': - def source_function_u(x, y, z): - fx = 2.0 * pi * cos(2 * pi * x) + nu * 4.0 * pi**2 * sin(2 * pi * x) - fy = B0 * sin(2 * pi * x) / epsilon - fz = zeros_like(x) - return fx, fy, fz - - def source_function_ue(x, y, z): - fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) - fy = -sin(2 * pi * x) * B0 / epsilon - fz = zeros_like(x) - return fx, fy, fz - -elif BC == 'dirichlet_inhom': - def source_function_u(x, y, z): - fx = ones_like(x) - fy = B0 * (1 + x) / epsilon - fz = zeros_like(x) - return fx, fy, fz - - def source_function_ue(x, y, z): - fx = -ones_like(x) - sigma * (1 + x) - fy = -B0 * (1 + x) / epsilon - fz = zeros_like(x) - return fx, fy, fz - -# ---- model ---- -model = TwoFluidQuasiNeutralToy() -model.ions.set_phys_params() -model.electrons.set_phys_params() - -model.propagators.qn_full.options = model.propagators.qn_full.Options( - nu=nu, - nu_e=nu_e, - eps_norm=epsilon, - stab_sigma=sigma, - source_u=source_function_u, - source_ue=source_function_ue, - solver='gmres', - boundary_conditions_u=bcs_u, - boundary_conditions_ue=bcs_ue, - boundary_data_u=boundary_data_u, - boundary_data_ue=boundary_data_ue, -) - -if __name__ == "__main__": - main.run(model, - params_path=__file__, - env=env, - base_units=base_units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=True, - ) - - path = os.path.join(os.getcwd(), name) - main.pproc(path) - simdata = main.load_data(path) - - n1_vals = simdata.grids_log[0] - x = xp.linspace(0, 1, 100) - - os.makedirs(f'{name}/plots', exist_ok=True) - for f in glob.glob(f'{name}/plots/*.png'): - os.remove(f) - - def save_plot(n1_vals, numerical, analytical, ylabel, title, fname, t): - plt.plot(n1_vals, numerical, label='numerical') - plt.plot(x, analytical, '--', label='manufactured') - plt.plot(n1_vals, numerical, 'k.', markersize=4, label='n1 points') - plt.xlabel('n1 (radial)') - plt.ylabel(ylabel) - plt.title(f'{title} at t={t:.3f}') - plt.legend() - plt.grid(True) - plt.savefig(f'{name}/plots/{fname}_{t:.3f}.png', dpi=300) - plt.clf() - - for t in list(simdata.spline_values['ions']['u_log'].keys()): - - u_ions = simdata.spline_values['ions']['u_log'][t] - u_electrons = simdata.spline_values['electrons']['u_log'][t] - phi = simdata.spline_values['em_fields']['phi_log'][t] - - mms_phi_x, _, _ = mms_phi(x, x*0, x*0) - mms_ion_ux, mms_ion_uy, _ = mms_ion_u(x, x*0, x*0) - mms_el_ux, mms_el_uy, _ = mms_electron_u(x, x*0, x*0) - - save_plot(n1_vals, phi[0][:, 0, 0], mms_phi_x, 'φ', 'Electrostatic potential φ', 'plot_potential', t) - save_plot(n1_vals, u_ions[0][:, 0, 0], mms_ion_ux, 'u_x', 'Ion velocity (u_x)', 'plot_ion_ux', t) - save_plot(n1_vals, u_electrons[0][:, 0, 0], mms_el_ux, 'u_x', 'Electron velocity (u_x)', 'plot_electron_ux', t) \ No newline at end of file From 7e3c622ca9644c0b211380bbcb03145bae842e72 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Thu, 5 Mar 2026 21:40:56 +0000 Subject: [PATCH 12/41] Rebased onto 3.0.4 v2 --- src/struphy/io/options.py | 6 +- ...rk_saddle_point.py => saddle_point_new.py} | 14 +-- .../{rework_model.py => two_fluid_new.py} | 26 +++--- src/struphy/propagators/base.py | 1 + ...py => propagators_fields_two_fluid_new.py} | 85 +++++-------------- 5 files changed, 50 insertions(+), 82 deletions(-) rename src/struphy/linear_algebra/{rework_saddle_point.py => saddle_point_new.py} (97%) rename src/struphy/models/{rework_model.py => two_fluid_new.py} (85%) rename src/struphy/propagators/{rework_propagator.py => propagators_fields_two_fluid_new.py} (86%) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index d88eae287..898881b61 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -3,10 +3,10 @@ from dataclasses import dataclass, fields from typing import Any, Callable, Literal -import cunumpy as xp -from psydac.ddm.mpi import mpi as MPI +from struphy.utils.utils import check_option -from struphy.physics.physics import ConstantsOfNature +import cunumpy as xp +from feectools.ddm.mpi import mpi as MPI ## Literal options diff --git a/src/struphy/linear_algebra/rework_saddle_point.py b/src/struphy/linear_algebra/saddle_point_new.py similarity index 97% rename from src/struphy/linear_algebra/rework_saddle_point.py rename to src/struphy/linear_algebra/saddle_point_new.py index 593dc8524..292313f69 100644 --- a/src/struphy/linear_algebra/rework_saddle_point.py +++ b/src/struphy/linear_algebra/saddle_point_new.py @@ -2,10 +2,10 @@ import cunumpy as xp import scipy as sc -from psydac.linalg.basic import LinearOperator, Vector -from psydac.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace -from psydac.linalg.direct_solvers import SparseSolver -from psydac.linalg.solvers import inverse +from feectools.linalg.basic import LinearOperator, Vector +from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace +from feectools.linalg.direct_solvers import SparseSolver +from feectools.linalg.solvers import inverse from struphy.linear_algebra.tests.test_saddlepoint_massmatrices import _plot_residual_norms @@ -27,7 +27,7 @@ class SaddlePointSolver: f \cr 0 } \right) - using either the Uzawa iteration :math:`BA^{-1}B^{\top} y = BA^{-1} f` or using on of the solvers given in :mod:`psydac.linalg.solvers`. The prefered solver is GMRES. + using either the Uzawa iteration :math:`BA^{-1}B^{\top} y = BA^{-1} f` or using on of the solvers given in :mod:`feectools.linalg.solvers`. The prefered solver is GMRES. The decission which variant to use is given by the type of A. If A is of type list of xp.ndarrays or sc.sparse.csr_matrices, then this class uses the Uzawa algorithm. If A is of type LinearOperator or BlockLinearOperator, a solver is used for the inverse. Using the Uzawa algorithm, solution is given by: @@ -44,8 +44,8 @@ class SaddlePointSolver: Either the entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. Alternative: Give whole matrice A as LinearOperator or BlockLinearOperator. list: Uzawa algorithm is used. - LinearOperator: A solver given in :mod:`psydac.linalg.solvers` is used. Specified by solver_name. - BlockLinearOperator: A solver given in :mod:`psydac.linalg.solvers` is used. Specified by solver_name. + LinearOperator: A solver given in :mod:`feectools.linalg.solvers` is used. Specified by solver_name. + BlockLinearOperator: A solver given in :mod:`feectools.linalg.solvers` is used. Specified by solver_name. B : list, LinearOperator or BlockLinearOperator Lower left block. diff --git a/src/struphy/models/rework_model.py b/src/struphy/models/two_fluid_new.py similarity index 85% rename from src/struphy/models/rework_model.py rename to src/struphy/models/two_fluid_new.py index f38629fcc..e5a569eab 100644 --- a/src/struphy/models/rework_model.py +++ b/src/struphy/models/two_fluid_new.py @@ -1,13 +1,15 @@ -import cunumpy as xp -from psydac.ddm.mpi import mpi as MPI +from feectools.ddm.mpi import mpi as MPI -from struphy.feec.projectors import L2Projector -from struphy.feec.variational_utilities import InternalEnergyEvaluator +from struphy.io.options import LiteralOptions from struphy.models.base import StruphyModel -from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies -from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable -from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.propagators import rework_propagator +from struphy.models.species import ( + FieldSpecies, + FluidSpecies, +) +from struphy.models.variables import FEECVariable +from struphy.propagators import ( + propagators_fields_two_fluid_new, +) rank = MPI.COMM_WORLD.Get_rank() @@ -50,6 +52,10 @@ class TwoFluidQuasiNeutralToy(StruphyModel): in plasma physics, Journal of Computational Physics 2018. """ + @classmethod + def model_type(cls) -> LiteralOptions.ModelTypes: + return "Fluid" + ## species class EMfields(FieldSpecies): @@ -71,7 +77,7 @@ def __init__(self): class Propagators: def __init__(self): - self.qn_full = rework_propagator.TwoFluidQuasiNeutralFull() + self.qn_full = propagators_fields_two_fluid_new.TwoFluidQuasiNeutralFull() ## abstract methods @@ -102,7 +108,7 @@ def bulk_species(self): def velocity_scale(self): return "thermal" - def allocate_helpers(self): + def allocate_helpers(self, verbose=False): pass def update_scalar_quantities(self): diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index fcdcefb45..f48dc1ce7 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -14,6 +14,7 @@ from struphy.feec.psydac_derham import Derham from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB from struphy.geometry.base import Domain +from struphy.utils.utils import check_option from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.utils.utils import check_option diff --git a/src/struphy/propagators/rework_propagator.py b/src/struphy/propagators/propagators_fields_two_fluid_new.py similarity index 86% rename from src/struphy/propagators/rework_propagator.py rename to src/struphy/propagators/propagators_fields_two_fluid_new.py index bf6afcadf..c5cf54039 100644 --- a/src/struphy/propagators/rework_propagator.py +++ b/src/struphy/propagators/propagators_fields_two_fluid_new.py @@ -1,63 +1,22 @@ - -import copy -from copy import deepcopy from dataclasses import dataclass -from typing import Callable, Literal, get_args, cast +from typing import Callable, Literal, cast from warnings import warn -import cunumpy as xp -import scipy as sc -from line_profiler import profile -from matplotlib import pyplot as plt -from numpy import zeros -from psydac.api.essential_bc import apply_essential_bc_stencil -from psydac.ddm.mpi import mpi as MPI -from psydac.linalg.basic import ComposedLinearOperator, IdentityOperator, ZeroOperator, InverseLinearOperator -from psydac.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace -from psydac.linalg.solvers import inverse -from psydac.linalg.stencil import StencilVector - -import struphy.feec.utilities as util -from struphy.examples.restelli2018 import callables -from struphy.feec import preconditioner -from struphy.feec.basis_projection_ops import ( - BasisProjectionOperator, BasisProjectionOperatorLocal, - BasisProjectionOperators, CoordinateProjector, -) -from struphy.feec.linear_operators import BoundaryOperator -from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators -from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner, MassMatrixPreconditioner +from feectools.api.essential_bc import apply_essential_bc_stencil +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.basic import IdentityOperator, InverseLinearOperator +from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace +from feectools.linalg.solvers import inverse + +from struphy.feec.basis_projection_ops import BasisProjectionOperators +from struphy.feec.mass import WeightedMassOperators from struphy.feec.projectors import L2Projector -from struphy.feec.psydac_derham import Derham, SplineFunction -from struphy.feec.variational_utilities import ( - BracketOperator, Hdiv0_transport_operator, InternalEnergyEvaluator, - KineticEnergyEvaluator, Pressure_transport_operator, -) -from struphy.fields_background.equils import set_defaults -from struphy.geometry.utilities import TransformedPformComponent -from struphy.initial import perturbations -from struphy.io.options import ( - OptsDirectSolver, OptsGenSolver, OptsMassPrecond, OptsNonlinearSolver, - OptsSaddlePointSolver, OptsSymmSolver, OptsVecSpace, check_option, -) -from struphy.io.setup import descend_options_dict -from struphy.kinetic_background.base import Maxwellian -from struphy.kinetic_background.maxwellians import GyroMaxwellian2D, Maxwellian3D -from struphy.linear_algebra.saddle_point import SaddlePointSolver -from struphy.linear_algebra.schur_solver import SchurSolver, SchurSolverFull -from struphy.linear_algebra.solver import NonlinearSolverParameters, SolverParameters -from struphy.models.species import Species -from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable -from struphy.ode.solvers import ODEsolverFEEC -from struphy.ode.utils import ButcherTableau, OptsButcher -from struphy.pic.accumulation import accum_kernels, accum_kernels_gc -from struphy.pic.accumulation.filter import FilterParameters -from struphy.pic.accumulation.particles_to_grid import Accumulator, AccumulatorVector -from struphy.pic.base import Particles -from struphy.pic.particles import Particles5D, Particles6D -from struphy.polar.basic import PolarVector +from struphy.feec.psydac_derham import Derham +from struphy.io.options import OptsGenSolver +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel +from struphy.utils.utils import check_option class TwoFluidQuasiNeutralFull(Propagator): @@ -264,8 +223,9 @@ def _apply_boundary_conditions(self, vec, boundary_conditions): ### Allocate # ========================================================================= - def allocate(self): + def allocate(self, verbose=False): + self.verbose = verbose self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 self._dt = None @@ -295,6 +255,7 @@ def allocate(self): # ---- unconstrained operators (for RHS assembly) ---------------------- + self._M2 = self.mass_ops.M2 self._M2B = - self.mass_ops.M2B self._div = self.derham.div @@ -303,9 +264,9 @@ def allocate(self): self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - + self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl - self._A22 = (- self.options.stab_sigma * IdentityOperator(self._A11.domain) + self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) # ---- constrained operators (for system matrix) ----------------------- @@ -319,16 +280,16 @@ def allocate(self): self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) - + self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._A11_v0.domain) + self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) # ---- block saddle-point system ---------------------------------------- - self._block_domain_v0 = BlockVectorSpace(self._A11_v0.domain, self._A22_v0.domain) + self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self._M3_v0.codomain + self._block_codomain_B_v0 = self._derham_v0.Vh["3"] self._B1_v0 = - self._M3_v0 @ self._div_v0 self._B2_v0 = self._M3_v0 @ self._div_v0 From aaaa33ee32740222b03ccc0ecfc5db6e32a46a7e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Thu, 12 Mar 2026 22:23:27 +0000 Subject: [PATCH 13/41] Replaced old propagator in propagators_fields and other minor changes. --- .gitignore | 1 - feectools | 2 +- .../linear_algebra/saddle_point_new.py | 567 ------- src/struphy/models/two_fluid_new.py | 130 -- src/struphy/propagators/propagators_fields.py | 1360 +++++++++-------- .../propagators_fields_two_fluid_new.py | 440 ------ struphy-parameter-files | 1 + 7 files changed, 693 insertions(+), 1808 deletions(-) delete mode 100644 src/struphy/linear_algebra/saddle_point_new.py delete mode 100644 src/struphy/models/two_fluid_new.py delete mode 100644 src/struphy/propagators/propagators_fields_two_fluid_new.py create mode 160000 struphy-parameter-files diff --git a/.gitignore b/.gitignore index daf35eef6..36b889d17 100644 --- a/.gitignore +++ b/.gitignore @@ -98,7 +98,6 @@ src/struphy/io/inp/params_* src/struphy/models/models_list src/struphy/models/models_message -runs/ bin/ share/ diff --git a/feectools b/feectools index 8c88dec79..d2a48ef19 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 8c88dec79510b315024d4b7e0ccc28e76ad8c9e7 +Subproject commit d2a48ef19d31ae22ba238369cd5a53c679c6f1d8 diff --git a/src/struphy/linear_algebra/saddle_point_new.py b/src/struphy/linear_algebra/saddle_point_new.py deleted file mode 100644 index 292313f69..000000000 --- a/src/struphy/linear_algebra/saddle_point_new.py +++ /dev/null @@ -1,567 +0,0 @@ -from typing import Union - -import cunumpy as xp -import scipy as sc -from feectools.linalg.basic import LinearOperator, Vector -from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace -from feectools.linalg.direct_solvers import SparseSolver -from feectools.linalg.solvers import inverse - -from struphy.linear_algebra.tests.test_saddlepoint_massmatrices import _plot_residual_norms - - -class SaddlePointSolver: - r"""Solves for :math:`(x, y)` in the saddle point problem - - .. math:: - - \left( \matrix{ - A & B^{\top} \cr - B & 0 - } \right) - \left( \matrix{ - x \cr y - } \right) - = - \left( \matrix{ - f \cr 0 - } \right) - - using either the Uzawa iteration :math:`BA^{-1}B^{\top} y = BA^{-1} f` or using on of the solvers given in :mod:`feectools.linalg.solvers`. The prefered solver is GMRES. - The decission which variant to use is given by the type of A. If A is of type list of xp.ndarrays or sc.sparse.csr_matrices, then this class uses the Uzawa algorithm. - If A is of type LinearOperator or BlockLinearOperator, a solver is used for the inverse. - Using the Uzawa algorithm, solution is given by: - - .. math:: - - y = \left[ B A^{-1} B^{\top}\right]^{-1} B A^{-1} f \,, \qquad - x = A^{-1} \left[ f - B^{\top} y \right] \,. - - Parameters - ---------- - A : list, LinearOperator or BlockLinearOperator - Upper left block. - Either the entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. - Alternative: Give whole matrice A as LinearOperator or BlockLinearOperator. - list: Uzawa algorithm is used. - LinearOperator: A solver given in :mod:`feectools.linalg.solvers` is used. Specified by solver_name. - BlockLinearOperator: A solver given in :mod:`feectools.linalg.solvers` is used. Specified by solver_name. - - B : list, LinearOperator or BlockLinearOperator - Lower left block. - Uzwaw Algorithm: All entries of block B are given either as list of xp.ndarray or sc.sparse.csr_matrix. - Solver: Give whole B as LinearOperator or BlocklinearOperator - - F : list - Right hand side of the upper block. - Uzawa: Given as list of xp.ndarray or sc.sparse.csr_matrix. - Solver: Given as LinearOperator or BlockLinearOperator - - Apre : list - The non-inverted preconditioner for entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. Only required for the Uzawa algorithm. - - method_to_solve : str - Method for the inverses. Choose from 'DirectNPInverse', 'ScipySparse', 'InexactNPInverse' ,'SparseSolver'. Only required for the Uzawa algorithm. - - preconditioner : bool - Wheter to use preconditioners given in Apre or not. Only required for the Uzawa algorithm. - - spectralanalysis : bool - Do the spectralanalyis for the matrices in A and if preconditioner given, compare them to the preconditioned matrices. Only possible if A is given as list. - - dimension : str - Which of the predefined manufactured solutions to use ('1D', '2D' or 'Restelli') - - tol : float - Convergence tolerance for the potential residual. - - max_iter : int - Maximum number of iterations allowed. - """ - - def __init__( - self, - A: Union[list, LinearOperator, BlockLinearOperator], - B: Union[list, LinearOperator, BlockLinearOperator], - F: Union[list, Vector, BlockVector], - Apre: list = None, - method_to_solve: str = "DirectNPInverse", - preconditioner: bool = False, - spectralanalysis: bool = False, - dimension: str = "2D", - solver_name: str = "GMRES", - tol: float = 1e-8, - max_iter: int = 1000, - **solver_params, - ): - assert type(A) is type(B) - if isinstance(A, list): - self._variant = "Uzawa" - for i in A: - assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) - for i in B: - assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) - for i in F: - assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) - for i in Apre: - assert ( - isinstance(i, xp.ndarray) - or isinstance(i, sc.sparse.csr_matrix) - or isinstance(i, sc.sparse.csr_array) - ) - assert method_to_solve in ("SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse") - assert A[0].shape[0] == B[0].shape[1] - assert A[0].shape[1] == B[0].shape[1] - assert A[1].shape[0] == B[1].shape[1] - assert A[1].shape[1] == B[1].shape[1] - - self._method_to_solve = ( - method_to_solve # 'SparseSolver', 'ScipySparse', 'InexactNPInverse', 'DirectNPInverse' - ) - self._preconditioner = preconditioner - - elif isinstance(A, LinearOperator) or isinstance(A, BlockLinearOperator): - self._variant = "Inverse_Solver" - assert A.domain == B.domain - assert A.codomain == B.domain - self._solver_name = solver_name - if solver_params["pc"] is None: - solver_params.pop("pc") - - # operators - self._A = A - self._Apre = Apre - self._B = B - self._F = F - self._tol = tol - self._max_iter = max_iter - self._spectralanalysis = spectralanalysis - self._dimension = dimension - self._verbose = solver_params["verbose"] - - if self._variant == "Inverse_Solver": - self._block_domainM = BlockVectorSpace(self._A.domain, self._B.transpose().domain) - self._block_codomainM = self._block_domainM - self._blocks = [[self._A, self._B.T], [self._B, None]] - _Minit = BlockLinearOperator(self._block_domainM, self._block_codomainM, blocks=self._blocks) - self._solverMinv = inverse(_Minit, solver_name, tol=tol, maxiter=max_iter, **solver_params) - - # Solution vectors - self._P = B.codomain.zeros() - self._U = A.codomain.zeros() - self._Utmp = F.copy() * 0 - # Allocate memory for call - self._rhstemp = BlockVector(self._block_domainM, blocks=[A.codomain.zeros(), self._B.codomain.zeros()]) - - elif self._variant == "Uzawa": - if self._method_to_solve in ("InexactNPInverse", "SparseSolver"): - self._preconditioner = False - - self._Anp = self._A[0] - self._Aenp = self._A[1] - self._B1np = self._B[0] - self._B2np = self._B[1] - - # Instanciate inverses - self._setup_inverses() - - # Solution vectors numpy - self._Pnp = xp.zeros(self._B1np.shape[0]) - self._Unp = xp.zeros(self._A[0].shape[1]) - self._Uenp = xp.zeros(self._A[1].shape[1]) - # Allocate memory for matrices used in solving the system - self._rhs0np = self._F[0].copy() - self._rhs1np = self._F[1].copy() - - # List to store residual norms - self._residual_norms = [] - self._stepsize = 0.0 - - @property - def A(self): - """Upper left block.""" - return self._A - - @A.setter - def A(self, a): - if self._variant == "Uzawa": - need_update = True - A0_old, A1_old = self._A - A0_new, A1_new = a - if self._method_to_solve in ("ScipySparse", "SparseSolver"): - same_A0 = (A0_old != A0_new).nnz == 0 - same_A1 = (A1_old != A1_new).nnz == 0 - else: - same_A0 = xp.allclose(A0_old, A0_new, atol=1e-10) - same_A1 = xp.allclose(A1_old, A1_new, atol=1e-10) - if same_A0 and same_A1: - need_update = False - self._A = a - self._Anp = self._A[0] - self._Aenp = self._A[1] - if need_update: - self._setup_inverses() - elif self._variant == "Inverse_Solver": - self._A = a - - @property - def B(self): - """Lower left block.""" - return self._B - - @B.setter - def B(self, b): - self._B = b - - @property - def F(self): - """Right hand side vector.""" - return self._F - - @F.setter - def F(self, f): - self._F = f - - @property - def Apre(self): - """Preconditioner for upper left block A.""" - return self._Apre - - @Apre.setter - def Apre(self, a): - if self._variant == "Uzawa": - need_update = True - A0_old, A1_old = self._Apre - A0_new, A1_new = a - if self._method_to_solve in ("ScipySparse", "SparseSolver"): - same_A0 = (A0_old != A0_new).nnz == 0 - same_A1 = (A1_old != A1_new).nnz == 0 - else: - same_A0 = xp.allclose(A0_old, A0_new, atol=1e-10) - same_A1 = xp.allclose(A1_old, A1_new, atol=1e-10) - if same_A0 and same_A1: - need_update = False - self._Apre = a - if need_update: - self._setup_inverses() - elif self._variant == "Inverse_Solver": - self._Apre = a - - def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): # todo should have options to use other than uzawa. should solve in full generality. A being block diagonal is a special case of uzawa - """ - Solves the saddle-point problem using the Uzawa algorithm. - - Parameters - ---------- - U_init : Vector, xp.ndarray or sc.sparse.csr.csr_matrix, optional - Initial guess for the velocity of the ions. If None, initializes to zero. Types xp.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. - - Ue_init : Vector, xp.ndarray or sc.sparse.csr.csr_matrix, optional - Initial guess for the velocity of the electrons. If None, initializes to zero. Types xp.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. - - P_init : Vector, optional - Initial guess for the potential. If None, initializes to zero. - - Returns - ------- - U : Vector - Solution vector for the velocity. - - P : Vector - Solution vector for the potential. - - info : dict - Convergence information. - """ - - # TODO this contains two different strategies! favágás and actual uzawa - if self._variant == "Inverse_Solver": # todo not in the """""saddle point solver""""""" - self._P1 = P_init if P_init is not None else self._P - self._U1 = U_init if U_init is not None else self._Utmp[0] - self._U2 = Ue_init if Ue_init is not None else self._Utmp[1] - - _blocksM = [[self._A, self._B.T], [self._B, None]] - _M = BlockLinearOperator(self._block_domainM, self._block_codomainM, blocks=_blocksM) - _RHS = BlockVector(self._block_domainM, blocks=[self._F, self._B.codomain.zeros()]) - - self._blockU = BlockVector(self._A.domain, blocks=[self._U1, self._U2]) - self._solblocks = [self._blockU, self._P1] - # comment out the next two lines if working with lifting and GMRES - x0 = BlockVector(self._block_domainM, blocks=self._solblocks) - self._solverMinv._options["x0"] = x0 - - # use setter to update lhs matrix - self._solverMinv.linop = _M - - # Initialize P to zero or given initial guess - self._sol = self._solverMinv.dot(_RHS, out=self._rhstemp) - self._U = self._sol[0] - self._P = self._sol[1] - - return self._U, self._P, self._solverMinv._info - - elif self._variant == "Uzawa": - info = {} - - if self._spectralanalysis: - self._spectralresult = self._spectral_analysis() - else: - self._spectralresult = [] - - # Initialize P to zero or given initial guess - if isinstance(U_init, xp.ndarray) or isinstance(U_init, sc.sparse.csr.csr_matrix): - self._Pnp = P_init if P_init is not None else self._P - self._Unp = U_init if U_init is not None else self._U - self._Uenp = Ue_init if U_init is not None else self._Ue - - else: - self._Pnp = P_init.toarray() if P_init is not None else self._Pnp - self._Unp = U_init.toarray() if U_init is not None else self._Unp - self._Uenp = Ue_init.toarray() if U_init is not None else self._Uenp - - if self._verbose: - print("Uzawa solver:") - print("+---------+---------------------+") - print("+ Iter. # | L2-norm of residual |") - print("+---------+---------------------+") - template = "| {:7d} | {:19.2e} |" - - for iteration in range(self._max_iter): - # Step 1: Compute velocity U by solving A U = -Bᵀ P + F -A Un - self._rhs0np *= 0 - self._rhs0np -= self._B1np.transpose().dot(self._Pnp) - self._rhs0np -= self._Anp.dot(self._Unp) - self._rhs0np += self._F[0] - if not self._preconditioner: - self._Unp += self._Anpinv.dot(self._rhs0np) - elif self._preconditioner: - self._Unp += self._Anpinv.dot(self._A11npinv @ self._rhs0np) - - R1 = self._B1np.dot(self._Unp) - - self._rhs1np *= 0 - self._rhs1np -= self._B2np.transpose().dot(self._Pnp) - self._rhs1np -= self._Aenp.dot(self._Uenp) - self._rhs1np += self._F[1] - if not self._preconditioner: - self._Uenp += self._Aenpinv.dot(self._rhs1np) - elif self._preconditioner: - self._Uenp += self._Aenpinv.dot(self._A22npinv @ self._rhs1np) - - R2 = self._B2np.dot(self._Uenp) - - # Step 2: Compute residual R = BU (divergence of U) - R = R1 + R2 # self._B1np.dot(self._Unp) + self._B2np.dot(self._Uenp) - residual_norm = xp.linalg.norm(R) - residual_normR1 = xp.linalg.norm(R) - self._residual_norms.append(residual_normR1) # Store residual norm - # Check for convergence based on residual norm - if residual_norm < self._tol: - if self._verbose: - print(template.format(iteration + 1, residual_norm)) - print("+---------+---------------------+") - info["success"] = True - info["niter"] = iteration + 1 - if self._verbose: - _plot_residual_norms(self._residual_norms) - return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult - - # Steepest gradient - alpha = (R.dot(R)) / (R.dot(self._Precnp.dot(R))) - # Minimal residual - # alpha = ((self._Precnp.dot(R)).dot(R)) / ((self._Precnp.dot(R)).dot(self._Precnp.dot(R))) - self._Pnp += alpha.real * R.real - - if self._verbose: - print(template.format(iteration + 1, residual_norm)) - - if self._verbose: - print("+---------+---------------------+") - - # Return with info if maximum iterations reached - info["success"] = False - info["niter"] = iteration + 1 - if self._verbose: - _plot_residual_norms(self._residual_norms) - return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult - - def _setup_inverses(self): - A0 = self._A[0] - A1 = self._A[1] - - # === Preconditioner inverses, if used - if self._preconditioner: - A11_pre = self._Apre[0] - A22_pre = self._Apre[1] - - if hasattr(self, "_A11npinv") and self._is_inverse_still_valid(self._A11npinv, A11_pre, "A11 pre"): - pass - else: - self._A11npinv = self._compute_inverse(A11_pre, which="A11 pre") - - if hasattr(self, "_A22npinv") and self._is_inverse_still_valid(self._A22npinv, A22_pre, "A22 pre"): - pass - else: - self._A22npinv = self._compute_inverse(A22_pre, which="A22 pre") - - # === Inverse for A[0] if preconditioned - if hasattr(self, "_Anpinv") and self._is_inverse_still_valid(self._Anpinv, A0, "A[0]", pre=self._A11npinv): - pass - else: - self._Anpinv = self._compute_inverse(self._A11npinv @ A0, which="A[0]") - - # === Inverse for A[1] - if hasattr(self, "_Aenpinv") and self._is_inverse_still_valid( - self._Aenpinv, - A1, - "A[1]", - pre=self._A22npinv, - ): - pass - else: - self._Aenpinv = self._compute_inverse(self._A22npinv @ A1, which="A[1]") - - else: # No preconditioning: - # === Inverse for A[0] - if hasattr(self, "_Anpinv") and self._is_inverse_still_valid(self._Anpinv, A0, "A[0]"): - pass - else: - self._Anpinv = self._compute_inverse(A0, which="A[0]") - - # === Inverse for A[1] - if hasattr(self, "_Aenpinv") and self._is_inverse_still_valid(self._Aenpinv, A1, "A[1]"): - pass - else: - self._Aenpinv = self._compute_inverse(A1, which="A[1]") - - # Precompute Schur complement - self._Precnp = self._B1np @ self._Anpinv @ self._B1np.T + self._B2np @ self._Aenpinv @ self._B2np.T - - def _is_inverse_still_valid(self, inv, mat, name="", pre=None): - # try: - if pre is not None: - test_mat = pre @ mat - else: - test_mat = mat - I_approx = inv @ test_mat - - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - I_exact = xp.eye(test_mat.shape[0]) - if not xp.allclose(I_approx, I_exact, atol=1e-6): - diff = I_approx - I_exact - max_abs = xp.abs(diff).max() - print(f"{name} inverse is NOT valid anymore. Max diff: {max_abs:.2e}") - return False - print(f"{name} inverse is still valid.") - return True - elif self._method_to_solve == "ScipySparse": - I_exact = sc.sparse.identity(I_approx.shape[0], format=I_approx.format) - diff = (I_approx - I_exact).tocoo() - max_abs = xp.abs(diff.data).max() if diff.nnz > 0 else 0.0 - - if max_abs > 1e-6: - print(f"{name} inverse is NOT valid anymore.") - print(f"Max absolute difference: {max_abs:.2e}") - print(f"Number of differing entries: {diff.nnz}") - return False - print(f"{name} inverse is still valid.") - return True - - def _compute_inverse(self, mat, which="matrix"): - print(f"Computing inverse for {which} using method {self._method_to_solve}") - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - return xp.linalg.inv(mat) - elif self._method_to_solve == "ScipySparse": - return sc.sparse.linalg.inv(mat) - elif self._method_to_solve == "SparseSolver": - solver = SparseSolver(mat) - return solver.solve(xp.eye(mat.shape[0])) - else: - raise ValueError(f"Unknown solver method {self._method_to_solve}") - - def _spectral_analysis(self): - # Spectral analysis - # A11 before - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA11_before, eigvecs_before = xp.linalg.eig(self._A[0]) - condA11_before = xp.linalg.cond(self._A[0]) - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA11_before, eigvecs_before = xp.linalg.eig(self._A[0].toarray()) - condA11_before = xp.linalg.cond(self._A[0].toarray()) - maxbeforeA11 = max(eigvalsA11_before) - maxbeforeA11_abs = xp.max(xp.abs(eigvalsA11_before)) - minbeforeA11_abs = xp.min(xp.abs(eigvalsA11_before)) - minbeforeA11 = min(eigvalsA11_before) - specA11_bef = maxbeforeA11 / minbeforeA11 - specA11_bef_abs = maxbeforeA11_abs / minbeforeA11_abs - # print(f'{maxbeforeA11 = }') - # print(f'{maxbeforeA11_abs = }') - # print(f'{minbeforeA11_abs = }') - # print(f'{minbeforeA11 = }') - # print(f'{specA11_bef = }') - print(f"{specA11_bef_abs =}") - - # A22 before - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA22_before, eigvecs_before = xp.linalg.eig(self._A[1]) - condA22_before = xp.linalg.cond(self._A[1]) - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA22_before, eigvecs_before = xp.linalg.eig(self._A[1].toarray()) - condA22_before = xp.linalg.cond(self._A[1].toarray()) - maxbeforeA22 = max(eigvalsA22_before) - maxbeforeA22_abs = xp.max(xp.abs(eigvalsA22_before)) - minbeforeA22_abs = xp.min(xp.abs(eigvalsA22_before)) - minbeforeA22 = min(eigvalsA22_before) - specA22_bef = maxbeforeA22 / minbeforeA22 - specA22_bef_abs = maxbeforeA22_abs / minbeforeA22_abs - # print(f'{maxbeforeA22 = }') - # print(f'{maxbeforeA22_abs = }') - # print(f'{minbeforeA22_abs = }') - # print(f'{minbeforeA22 = }') - # print(f'{specA22_bef = }') - print(f"{specA22_bef_abs =}") - print(f"{condA22_before =}") - - if self._preconditioner: - # A11 after preconditioning with its inverse - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig(self._A11npinv @ self._A[0]) # Implement this - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig((self._A11npinv @ self._A[0]).toarray()) - maxafterA11_prec = max(eigvalsA11_after_prec) - minafterA11_prec = min(eigvalsA11_after_prec) - maxafterA11_abs_prec = xp.max(xp.abs(eigvalsA11_after_prec)) - minafterA11_abs_prec = xp.min(xp.abs(eigvalsA11_after_prec)) - specA11_aft_prec = maxafterA11_prec / minafterA11_prec - specA11_aft_abs_prec = maxafterA11_abs_prec / minafterA11_abs_prec - # print(f'{maxafterA11_prec = }') - # print(f'{maxafterA11_abs_prec = }') - # print(f'{minafterA11_abs_prec = }') - # print(f'{minafterA11_prec = }') - # print(f'{specA11_aft_prec = }') - print(f"{specA11_aft_abs_prec =}") - - # A22 after preconditioning with its inverse - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA22_after_prec, eigvecs_after = xp.linalg.eig(self._A22npinv @ self._A[1]) # Implement this - condA22_after = xp.linalg.cond(self._A22npinv @ self._A[1]) - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA22_after_prec, eigvecs_after = xp.linalg.eig((self._A22npinv @ self._A[1]).toarray()) - condA22_after = xp.linalg.cond((self._A22npinv @ self._A[1]).toarray()) - maxafterA22_prec = max(eigvalsA22_after_prec) - minafterA22_prec = min(eigvalsA22_after_prec) - maxafterA22_abs_prec = xp.max(xp.abs(eigvalsA22_after_prec)) - minafterA22_abs_prec = xp.min(xp.abs(eigvalsA22_after_prec)) - specA22_aft_prec = maxafterA22_prec / minafterA22_prec - specA22_aft_abs_prec = maxafterA22_abs_prec / minafterA22_abs_prec - # print(f'{maxafterA22_prec = }') - # print(f'{maxafterA22_abs_prec = }') - # print(f'{minafterA22_abs_prec = }') - # print(f'{minafterA22_prec = }') - # print(f'{specA22_aft_prec = }') - print(f"{specA22_aft_abs_prec =}") - - return condA22_before, specA22_bef_abs, condA11_before, condA22_after, specA22_aft_abs_prec - - else: - return condA22_before, specA22_bef_abs, condA11_before diff --git a/src/struphy/models/two_fluid_new.py b/src/struphy/models/two_fluid_new.py deleted file mode 100644 index e5a569eab..000000000 --- a/src/struphy/models/two_fluid_new.py +++ /dev/null @@ -1,130 +0,0 @@ -from feectools.ddm.mpi import mpi as MPI - -from struphy.io.options import LiteralOptions -from struphy.models.base import StruphyModel -from struphy.models.species import ( - FieldSpecies, - FluidSpecies, -) -from struphy.models.variables import FEECVariable -from struphy.propagators import ( - propagators_fields_two_fluid_new, -) - - -rank = MPI.COMM_WORLD.Get_rank() - -class TwoFluidQuasiNeutralToy(StruphyModel): - r"""Linearized, quasi-neutral two-fluid model with zero electron inertia. - - :ref:`normalization`: - - .. math:: - - \hat u = \hat v_\textnormal{th}\,,\qquad e\hat \phi = m \hat v_\textnormal{th}^2\,. - - :ref:`Equations `: - - .. math:: - - \frac{\partial \mathbf u}{\partial t} &= - \nabla \phi + \frac{\mathbf u \times \mathbf B_0}{\varepsilon} + \nu \Delta \mathbf u + \mathbf f\,, - \\[2mm] - 0 &= \nabla \phi - \frac{\mathbf u_e \times \mathbf B_0}{\varepsilon} + \nu_e \Delta \mathbf u_e + \mathbf f_e \,, - \\[3mm] - \nabla & \cdot (\mathbf u - \mathbf u_e) = 0\,, - - where :math:`\mathbf B_0` is a static magnetic field and :math:`\mathbf f, \mathbf f_e` are given forcing terms, - and with the normalization parameter - - .. math:: - - \varepsilon = \frac{1}{\hat \Omega_\textnormal{c} \hat t} \,,\qquad \textnormal{with} \,,\qquad \hat \Omega_{\textnormal{c}} = \frac{(Ze) \hat B}{(A m_\textnormal{H})}\,, - - :ref:`propagators` (called in sequence): - - 1. :class:`~struphy.propagators.propagators_fields.TwoFluidQuasiNeutralFull` - - :ref:`Model info `: - - References - ---------- - [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising - in plasma physics, Journal of Computational Physics 2018. - """ - - @classmethod - def model_type(cls) -> LiteralOptions.ModelTypes: - return "Fluid" - - ## species - - class EMfields(FieldSpecies): - def __init__(self): - self.phi = FEECVariable(space="L2") - self.init_variables() - - class Ions(FluidSpecies): - def __init__(self): - self.u = FEECVariable(space="Hdiv") - self.init_variables() - - class Electrons(FluidSpecies): - def __init__(self): - self.u = FEECVariable(space="Hdiv") - self.init_variables() - - ## propagators - - class Propagators: - def __init__(self): - self.qn_full = propagators_fields_two_fluid_new.TwoFluidQuasiNeutralFull() - - ## abstract methods - - def __init__(self): - if rank == 0: - print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - - # 1. instantiate all species - self.em_fields = self.EMfields() - self.ions = self.Ions() - self.electrons = self.Electrons() - - # 2. instantiate all propagators - self.propagators = self.Propagators() - - # 3. assign variables to propagators - self.propagators.qn_full.variables.u = self.ions.u - self.propagators.qn_full.variables.ue = self.electrons.u - self.propagators.qn_full.variables.phi = self.em_fields.phi - - # define scalars for update_scalar_quantities - - @property - def bulk_species(self): - return self.ions - - @property - def velocity_scale(self): - return "thermal" - - def allocate_helpers(self, verbose=False): - pass - - def update_scalar_quantities(self): - pass - - ## default parameters - def generate_default_parameter_file(self, path=None, prompt=True): - params_path = super().generate_default_parameter_file(path=path, prompt=prompt) - new_file = [] - with open(params_path, "r") as f: - for line in f: - if "BaseUnits()" in line: - new_file += ["base_units = BaseUnits(kBT=1.0)\n"] - else: - new_file += [line] - - with open(params_path, "w") as f: - for line in new_file: - f.write(line) diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 084dcdc2f..66a1b129d 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from typing import Callable, Literal, get_args from warnings import warn +from warnings import warn import cunumpy as xp import scipy as sc @@ -7652,13 +7653,18 @@ class TwoFluidQuasiNeutralFull(Propagator): ### State variables (ion velocity u, electron velocity ue, pressure phi) # ========================================================================= - class Variables: + # ========================================================================= + ### State variables (ion velocity u, electron velocity ue, pressure phi) + # ========================================================================= + + class Variables(): def __init__(self) -> None: self._u: FEECVariable | None = None self._ue: FEECVariable | None = None self._phi: FEECVariable | None = None @property + def u(self) -> FEECVariable | None: def u(self) -> FEECVariable | None: return self._u @@ -7669,6 +7675,7 @@ def u(self, new): self._u = new @property + def ue(self) -> FEECVariable | None: def ue(self) -> FEECVariable | None: return self._ue @@ -7679,6 +7686,7 @@ def ue(self, new): self._ue = new @property + def phi(self) -> FEECVariable | None: def phi(self) -> FEECVariable | None: return self._phi @@ -7695,16 +7703,26 @@ def __init__(self): ### Options # ========================================================================= + # ========================================================================= + ### Options + # ========================================================================= + @dataclass - class Options: - nu: float = 1.0 - nu_e: float = 1.0 - eps_norm: float = 1e-3 + class Options(): + + nu: float | None = None + nu_e: float | None = None + eps_norm: float | None = None - boundary_data_u: dict[tuple[int, int], Callable] | None = None - boundary_data_ue: dict[tuple[int, int], Callable] | None = None + # boundary conditions per species + # supported kinds: "periodic", "dirichlet" + # future: "neumann", "robin" + boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_data_u: dict[tuple[int, int], Callable] | None = None + boundary_data_ue: dict[tuple[int, int], Callable] | None = None - source_u: Callable | None = None + source_u: Callable | None = None source_ue: Callable | None = None stab_sigma: float | None = None @@ -7713,6 +7731,14 @@ class Options: solver_params: SolverParameters | None = None def __post_init__(self): + + # --- required parameters --- + assert self.nu is not None, "nu must be specified" + assert self.nu_e is not None, "nu_e must be specified" + assert self.eps_norm is not None, "eps_norm must be specified" + assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" + assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" + # --- physical parameter sanity checks --- if self.nu < 0: raise ValueError(f"nu must be non-negative, got {self.nu}") @@ -7721,6 +7747,52 @@ def __post_init__(self): if self.eps_norm <= 0: raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") + # --- check all axes are covered --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for d in range(3): + for side in (-1, 1): + assert (d, side) in bcs, \ + f"{name} is missing entry for axis {d} side {side}" + + # --- periodic consistency: periodic must be paired on both sides --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for (d, side), kind in bcs.items(): + if kind == "periodic": + assert bcs.get((d, -side)) == "periodic", \ + f"{name}: axis {d} side {side} is periodic but opposite side is not" + + # --- ions and electrons must agree on which axes are periodic --- + for d in range(3): + u_left = self.boundary_conditions_u.get((d, -1)) + ue_left = self.boundary_conditions_ue.get((d, -1)) + u_right = self.boundary_conditions_u.get((d, 1)) + ue_right = self.boundary_conditions_ue.get((d, 1)) + u_periodic = (u_left == "periodic") + ue_periodic = (ue_left == "periodic") + if u_periodic != ue_periodic: + raise ValueError( + f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " + f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" + ) + + # --- warn for Dirichlet faces with no boundary data --- + for species, bcs, data, label in [ + ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), + ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), + ]: + has_dirichlet = any(v == "dirichlet" for v in bcs.values()) + if has_dirichlet: + if data is None: + warn(f"Dirichlet BCs specified for {species} but no {label} given " + f"— defaulting to homogeneous Dirichlet on all faces.") + else: + for (d, side), kind in bcs.items(): + if kind == "dirichlet" and (d, side) not in data: + warn(f"No {label} given for axis {d} side {side} " + f"— defaulting to homogeneous Dirichlet.") + # --- warn if no source terms --- if self.source_u is None: warn("No source_u specified — defaulting to zero.") @@ -7732,12 +7804,13 @@ def __post_init__(self): warn("stab_sigma not specified, defaulting to 0.0") self.stab_sigma = 0.0 - check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver) + check_option(self.solver, LiteralOptions.OptsGenSolver) if self.solver_params is None: self.solver_params = SolverParameters() @property def options(self) -> Options: + assert hasattr(self, "_options"), "Options not set." assert hasattr(self, "_options"), "Options not set." return self._options @@ -7745,723 +7818,672 @@ def options(self) -> Options: def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: - logger.info(f"\nNew options for propagator '{self.__class__.__name__}':") + print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - logger.info(f" {k}: {v}") + print(f" {k}: {v}") self._options = new - @profile - def allocate(self, verbose: bool = False): - self._info = self.options.solver_params.info - if self.derham.comm is not None: - self._rank = self.derham.comm.Get_rank() - else: - self._rank = 0 + # ========================================================================= + ### Boundary condition helpers + # ========================================================================= - self._nu = self.options.nu - self._nu_e = self.options.nu_e - self._eps_norm = self.options.eps_norm - self._a = self.options.a - self._R0 = self.options.R0 - self._B0 = self.options.B0 - self._Bp = self.options.Bp - self._alpha = self.options.alpha - self._beta = self.options.beta - self._stab_sigma = self.options.stab_sigma - self._variant = self.options.variant - self._method_to_solve = self.options.method_to_solve - self._preconditioner = self.options.preconditioner - self._dimension = self.options.dimension - self._spectralanalysis = self.options.spectralanalysis - self._lifting = self.options.lifting + def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): - solver_params = self.options.solver_params + dirichlet_bc = [] + for d in range(3): + if spl_kind[d]: # periodic spline — no clamping ever + dirichlet_bc.append((False, False)) + else: + left = boundary_conditions.get((d, -1)) == "dirichlet" + right = boundary_conditions.get((d, 1)) == "dirichlet" + dirichlet_bc.append((left, right)) + return tuple(tuple(bc) for bc in dirichlet_bc) + + def _apply_boundary_conditions(self, vec, boundary_conditions): + """Zero out Dirichlet DOFs on the given stencil vector.""" + for (d, side), kind in boundary_conditions.items(): + if kind == "dirichlet": + apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) + # future: neumann and robin require no zeroing here - u = self.variables.u.spline.vector + # ========================================================================= + ### Allocate + # ========================================================================= - # Lifting for nontrivial boundary conditions - # derham had boundary conditions in eta1 direction, the following is in space Hdiv_0 - if self._lifting: - self.derhamv0 = Derham( - self.derham.Nel, - self.derham.p, - self.derham.spl_kind, - domain=self.domain, - dirichlet_bc=((True, True), (False, False), (False, False)), - ) + def allocate(self, verbose=False): - self._mass_opsv0 = WeightedMassOperators( - self.derhamv0, - self.domain, - verbose=solver_params.verbose, - eq_mhd=self.mass_ops.weights["eq_mhd"], - ) - self._basis_opsv0 = BasisProjectionOperators( - self.derhamv0, - self.domain, - verbose=solver_params.verbose, - eq_mhd=self.basis_ops.weights["eq_mhd"], - ) - else: - self.derhamnumpy = Derham( - self.derham.Nel, - self.derham.p, - self.derham.spl_kind, - domain=self.domain, - # dirichlet_bc=self.derham.dirichlet_bc, - # nquads = self.derham._nquads, - # nq_pr = self.derham._nq_pr, - # comm = MPI.COMM_SELF, # self.derham._comm, - # polar_ck= self.derham._polar_ck, - # local_projectors=self.derham.with_local_projectors - ) + self.verbose = verbose + self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 + self._dt = None - # get forceterms for according dimension - if self._dimension in ["2D", "1D"]: - ### Manufactured solution ### - _forceterm_logical = lambda e1, e2, e3: 0 * e1 - _funx = getattr(callables, "ManufacturedSolutionForceterm")( - species="Ions", - comp="0", - b0=self._B0, - nu=self._nu, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - ) - _funy = getattr(callables, "ManufacturedSolutionForceterm")( - species="Ions", - comp="1", - b0=self._B0, - nu=self._nu, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - ) - _funelectronsx = getattr(callables, "ManufacturedSolutionForceterm")( - species="Electrons", - comp="0", - b0=self._B0, - nu_e=self._nu_e, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - ) - _funelectronsy = getattr(callables, "ManufacturedSolutionForceterm")( - species="Electrons", - comp="1", - b0=self._B0, - nu_e=self._nu_e, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - ) + # ---- constrained (v0) de Rham complex -------------------------------- - # get callable(s) for specified init type - forceterm_class = [_funx, _funy, _forceterm_logical] - forcetermelectrons_class = [_funelectronsx, _funelectronsy, _forceterm_logical] - - # pullback callable - funx = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - funy = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - fun_electronsx = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - fun_electronsy = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - self._F1 = l2_proj([funx, funy, _forceterm_logical]) - self._F2 = l2_proj([fun_electronsx, fun_electronsy, _forceterm_logical]) - - elif self._dimension == "Restelli": - ### Restelli ### - - _forceterm_logical = lambda e1, e2, e3: 0 * e1 - _fun = getattr(callables, "RestelliForcingTerm")( - B0=self._B0, - nu=self._nu, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - eps=self._eps_norm, - ) - _funelectrons = getattr(callables, "RestelliForcingTerm")( - B0=self._B0, - nu=self._nu_e, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - eps=self._eps_norm, - ) + _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) + _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) + _dirichlet_bc = tuple( + (l_u or l_ue, r_u or r_ue) + for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) + ) - # get callable(s) for specified init type - forceterm_class = [_forceterm_logical, _forceterm_logical, _fun] - forcetermelectrons_class = [_forceterm_logical, _forceterm_logical, _funelectrons] - - # pullback callable - fun_pb_1 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - fun_pb_2 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - fun_pb_3 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=2, - domain=self.domain, - ) - fun_electrons_pb_1 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - fun_electrons_pb_2 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - fun_electrons_pb_3 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=2, - domain=self.domain, - ) - if self._lifting: - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) - else: - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - self._F1 = l2_proj([fun_pb_1, fun_pb_2, fun_pb_3], apply_bc=self._lifting) - self._F2 = l2_proj([fun_electrons_pb_1, fun_electrons_pb_2, fun_electrons_pb_3], apply_bc=self._lifting) - - ### End Restelli ### - - elif self._dimension == "Tokamak": - ### Tokamak geometry curl-free manufactured solution ### - - _forceterm_logical = lambda e1, e2, e3: 0 * e1 - _funx = getattr(callables, "ManufacturedSolutionForceterm")( - species="Ions", - comp="0", - b0=self._B0, - nu=self._nu, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) - _funy = getattr(callables, "ManufacturedSolutionForceterm")( - species="Ions", - comp="1", - b0=self._B0, - nu=self._nu, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) - _funz = getattr(callables, "ManufacturedSolutionForceterm")( - species="Ions", - comp="2", - b0=self._B0, - nu=self._nu, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) - _funelectronsx = getattr(callables, "ManufacturedSolutionForceterm")( - species="Electrons", - comp="0", - b0=self._B0, - nu_e=self._nu_e, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) - _funelectronsy = getattr(callables, "ManufacturedSolutionForceterm")( - species="Electrons", - comp="1", - b0=self._B0, - nu_e=self._nu_e, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) - _funelectronsz = getattr(callables, "ManufacturedSolutionForceterm")( - species="Electrons", - comp="2", - b0=self._B0, - nu_e=self._nu_e, - dimension=self._dimension, - stab_sigma=self._stab_sigma, - eps=self._eps_norm, - dt=self.options.D1_dt, - a=self._a, - Bp=self._Bp, - alpha=self._alpha, - beta=self._beta, - ) + self._derham_v0 = Derham( + self.derham.Nel, self.derham.p, self.derham.spl_kind, + domain=self.domain, dirichlet_bc=_dirichlet_bc, + ) + self._mass_ops_v0 = WeightedMassOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.mass_ops.weights["eq_mhd"], + ) + self._basis_ops_v0 = BasisProjectionOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.basis_ops.weights["eq_mhd"], + ) - # get callable(s) for specified init type - forceterm_class = [_funx, _funy, _funz] - forcetermelectrons_class = [_funelectronsx, _funelectronsy, _funelectronsz] - - # pullback callable - fun_pb_1 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - fun_pb_2 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, - ) - fun_pb_3 = TransformedPformComponent( - forceterm_class, - given_in_basis="physical", - out_form="2", - comp=2, - domain=self.domain, - ) - fun_electrons_pb_1 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=0, - domain=self.domain, - ) - fun_electrons_pb_2 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=1, - domain=self.domain, + # ---- unconstrained operators (for RHS assembly) ---------------------- + + + self._M2 = self.mass_ops.M2 + self._M2B = - self.mass_ops.M2B + self._div = self.derham.div + self._curl = self.derham.curl + self._S21 = self.basis_ops.S21 + + self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div + + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + + self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl + self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) + + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) + + # ---- constrained operators (for system matrix) ----------------------- + + self._M2_v0 = self._mass_ops_v0.M2 + self._M3_v0 = self._mass_ops_v0.M3 + self._M2B_v0 = - self._mass_ops_v0.M2B + self._div_v0 = self._derham_v0.div + self._curl_v0 = self._derham_v0.curl + self._S21_v0 = self._basis_ops_v0.S21 + + self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 + + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) + + self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 + self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) + + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) + + # ---- block saddle-point system ---------------------------------------- + + self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) + self._block_codomain_v0 = self._block_domain_v0 + self._block_codomain_B_v0 = self._derham_v0.Vh["3"] + + self._B1_v0 = - self._M3_v0 @ self._div_v0 + self._B2_v0 = self._M3_v0 @ self._div_v0 + + self._B_v0 = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_B_v0, + blocks=[[self._B1_v0, self._B2_v0]] + ) + + self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) + + _A_init = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0, None], [None, self._A22_v0]] + ) + _M_init = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] + ) + + self._Minv = inverse( + _M_init, self.options.solver, + A11=self._A11_v0, + A22=self._A22_v0, + B1=self._B1_v0, + B2=self._B2_v0, + recycle=self.options.solver_params.recycle, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + ) + + # ---- projector ------------------------------------------------------- + + self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + + # ---- solution spline functions (unconstrained) ----------------------- + + self._u = self.derham.create_spline_function("u", space_id="Hdiv") + self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") + self._phi = self.derham.create_spline_function("phi", space_id="L2") + + # ---- BC lifts (unconstrained) ---------------------------------------- + + self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") + self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") + + for u_prime, boundary_data, boundary_conditions in [ + (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), + (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), + ]: + if boundary_data is None: + continue + for (d, side), f_bc in boundary_data.items(): + if boundary_conditions.get((d, side)) == "dirichlet": + bc_pulled = lambda *etas, f=f_bc: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], + lambda *etas: bc_pulled(*etas)[1], + lambda *etas: bc_pulled(*etas)[2]]) + for (d2, side2), kind2 in boundary_conditions.items(): + if kind2 == "dirichlet" and (d2, side2) != (d, side): + apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) + u_prime.vector += _vec + + self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") + self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") + + self._u_prime_v0.vector = self._u_prime.vector + self._ue_prime_v0.vector = self._ue_prime.vector + + # ---- projected source terms (unconstrained) -------------------------- + + self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") + self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") + + for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: + if source is not None: + src_pulled = lambda *etas, f=source: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], + lambda *etas: src_pulled(*etas)[1], + lambda *etas: src_pulled(*etas)[2]]) + + # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- + + self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") + + # ========================================================================= + ### Time step + # ========================================================================= + + def __call__(self, dt): + + # --- copy current state --- + self._u.vector = self.variables.u.spline.vector + self._ue.vector = self.variables.ue.spline.vector + + # --- rebuild system matrix if dt changed --- + if dt != self._dt: + self._dt = dt + _A = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] ) - fun_electrons_pb_3 = TransformedPformComponent( - forcetermelectrons_class, - given_in_basis="physical", - out_form="2", - comp=2, - domain=self.domain, + + _M = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A, self._B_v0.T], [self._B_v0, None]] ) - if self._lifting: - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) - else: - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - self._F1 = l2_proj([fun_pb_1, fun_pb_2, fun_pb_3], apply_bc=self._lifting) - self._F2 = l2_proj([fun_electrons_pb_1, fun_electrons_pb_2, fun_electrons_pb_3], apply_bc=self._lifting) - - ### End Tokamak geometry manufactured solution ### - - if self._variant == "GMRES": - if self._lifting: - self._M2 = getattr(self._mass_opsv0, "M2") - self._M3 = getattr(self._mass_opsv0, "M3") - self._M2B = -getattr(self._mass_opsv0, "M2B") - self._div = self.derhamv0.div - self._curl = self.derhamv0.curl - self._S21 = self._basis_opsv0.S21 + self._Minv.linop = _M + class Variables(): + def __init__(self) -> None: + self._u: FEECVariable | None = None + self._ue: FEECVariable | None = None + self._phi: FEECVariable | None = None + + @property + def u(self) -> FEECVariable | None: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._u = new + + @property + def ue(self) -> FEECVariable | None: + return self._ue + + @ue.setter + def ue(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._ue = new + + @property + def phi(self) -> FEECVariable | None: + return self._phi + + @phi.setter + def phi(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._phi = new + + def __init__(self): + self.variables = self.Variables() + + # ========================================================================= + ### Options + # ========================================================================= + + @dataclass + class Options(): + + nu: float | None = None + nu_e: float | None = None + eps_norm: float | None = None + + # boundary conditions per species + # supported kinds: "periodic", "dirichlet" + # future: "neumann", "robin" + boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None + boundary_data_u: dict[tuple[int, int], Callable] | None = None + boundary_data_ue: dict[tuple[int, int], Callable] | None = None + + source_u: Callable | None = None + source_ue: Callable | None = None + + stab_sigma: float | None = None + + solver: LiteralOptions.OptsGenSolver = "gmres" + solver_params: SolverParameters | None = None + + def __post_init__(self): + + # --- required parameters --- + assert self.nu is not None, "nu must be specified" + assert self.nu_e is not None, "nu_e must be specified" + assert self.eps_norm is not None, "eps_norm must be specified" + assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" + assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" + + # --- physical parameter sanity checks --- + if self.nu < 0: + raise ValueError(f"nu must be non-negative, got {self.nu}") + if self.nu_e < 0: + raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") + if self.eps_norm <= 0: + raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") + + # --- check all axes are covered --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for d in range(3): + for side in (-1, 1): + assert (d, side) in bcs, \ + f"{name} is missing entry for axis {d} side {side}" + + # --- periodic consistency: periodic must be paired on both sides --- + for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), + ("boundary_conditions_ue", self.boundary_conditions_ue)]: + for (d, side), kind in bcs.items(): + if kind == "periodic": + assert bcs.get((d, -side)) == "periodic", \ + f"{name}: axis {d} side {side} is periodic but opposite side is not" + + # --- ions and electrons must agree on which axes are periodic --- + for d in range(3): + u_left = self.boundary_conditions_u.get((d, -1)) + ue_left = self.boundary_conditions_ue.get((d, -1)) + u_right = self.boundary_conditions_u.get((d, 1)) + ue_right = self.boundary_conditions_ue.get((d, 1)) + u_periodic = (u_left == "periodic") + ue_periodic = (ue_left == "periodic") + if u_periodic != ue_periodic: + raise ValueError( + f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " + f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" + ) + + # --- warn for Dirichlet faces with no boundary data --- + for species, bcs, data, label in [ + ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), + ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), + ]: + has_dirichlet = any(v == "dirichlet" for v in bcs.values()) + if has_dirichlet: + if data is None: + warn(f"Dirichlet BCs specified for {species} but no {label} given " + f"— defaulting to homogeneous Dirichlet on all faces.") + else: + for (d, side), kind in bcs.items(): + if kind == "dirichlet" and (d, side) not in data: + warn(f"No {label} given for axis {d} side {side} " + f"— defaulting to homogeneous Dirichlet.") + + # --- warn if no source terms --- + if self.source_u is None: + warn("No source_u specified — defaulting to zero.") + if self.source_ue is None: + warn("No source_ue specified — defaulting to zero.") + + # --- defaults --- + if self.stab_sigma is None: + warn("stab_sigma not specified, defaulting to 0.0") + self.stab_sigma = 0.0 + + check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver) + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + assert hasattr(self, "_options"), "Options not set." + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + # ========================================================================= + ### Boundary condition helpers + # ========================================================================= + + def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): + + dirichlet_bc = [] + for d in range(3): + if spl_kind[d]: # periodic spline — no clamping ever + dirichlet_bc.append((False, False)) else: - self._M2 = getattr(self.mass_ops, "M2") - self._M3 = getattr(self.mass_ops, "M3") - self._M2B = -getattr(self.mass_ops, "M2B") - self._div = self.derham.div - self._curl = self.derham.curl - self._S21 = self.basis_ops.S21 - - self._lapl = ( - self._div.T @ self.mass_ops.M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21 + left = boundary_conditions.get((d, -1)) == "dirichlet" + right = boundary_conditions.get((d, 1)) == "dirichlet" + dirichlet_bc.append((left, right)) + return tuple(tuple(bc) for bc in dirichlet_bc) + + def _apply_boundary_conditions(self, vec, boundary_conditions): + """Zero out Dirichlet DOFs on the given stencil vector.""" + for (d, side), kind in boundary_conditions.items(): + if kind == "dirichlet": + apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) + # future: neumann and robin require no zeroing here + + # ========================================================================= + ### Allocate + # ========================================================================= + + def allocate(self, verbose=False): + + self.verbose = verbose + self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 + self._dt = None + + # ---- constrained (v0) de Rham complex -------------------------------- + + _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) + _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) + _dirichlet_bc = tuple( + (l_u or l_ue, r_u or r_ue) + for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) ) - self._A11 = -self._M2B / self.options.eps_norm + self.options.nu * self._lapl - self._A22 = ( - -self.options.stab_sigma * IdentityOperator(self.derham.V2) - + self._M2B / self.options.eps_norm - + self.options.nu_e * self._lapl + self._derham_v0 = Derham( + self.derham.Nel, self.derham.p, self.derham.spl_kind, + domain=self.domain, dirichlet_bc=_dirichlet_bc, + ) + self._mass_ops_v0 = WeightedMassOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.mass_ops.weights["eq_mhd"], ) + self._basis_ops_v0 = BasisProjectionOperators( + self._derham_v0, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.basis_ops.weights["eq_mhd"], + ) + + # ---- unconstrained operators (for RHS assembly) ---------------------- + + + self._M2 = self.mass_ops.M2 + self._M2B = - self.mass_ops.M2B + self._div = self.derham.div + self._curl = self.derham.curl + self._S21 = self.basis_ops.S21 + + self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div + + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + + self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl + self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) + + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) # ---- constrained operators (for system matrix) ----------------------- - self._M2_v0 = self._mass_ops_v0.M2 - self._M3_v0 = self._mass_ops_v0.M3 - self._M2B_v0 = -self._mass_ops_v0.M2B - self._div_v0 = self._derham_v0.div + self._M2_v0 = self._mass_ops_v0.M2 + self._M3_v0 = self._mass_ops_v0.M3 + self._M2B_v0 = - self._mass_ops_v0.M2B + self._div_v0 = self._derham_v0.div self._curl_v0 = self._derham_v0.curl - self._S21_v0 = self._basis_ops_v0.S21 - - self._lapl_v0 = ( - self._div_v0.T @ self._M3_v0 @ self._div_v0 - + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0 - ) + self._S21_v0 = self._basis_ops_v0.S21 - self._A11_v0 = -self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = ( - -self.options.stab_sigma * IdentityOperator(self._derham_v0.V2) - + self._M2B_v0 / self.options.eps_norm - + self.options.nu_e * self._lapl_v0 - ) + self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 + + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) + + self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 + self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) + + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) # ---- block saddle-point system ---------------------------------------- - self._block_domain_v0 = BlockVectorSpace(self._derham_v0.V2, self._derham_v0.V2) - self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self._derham_v0.V3 + self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) + self._block_codomain_v0 = self._block_domain_v0 + self._block_codomain_B_v0 = self._derham_v0.Vh["3"] - self._B1_v0 = -self._M3_v0 @ self._div_v0 - self._B2_v0 = self._M3_v0 @ self._div_v0 + self._B1_v0 = - self._M3_v0 @ self._div_v0 + self._B2_v0 = self._M3_v0 @ self._div_v0 self._B_v0 = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_B_v0, blocks=[[self._B1_v0, self._B2_v0]] + self._block_domain_v0, self._block_codomain_B_v0, + blocks=[[self._B1_v0, self._B2_v0]] ) self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) _A_init = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, blocks=[[self._A11_v0, None], [None, self._A22_v0]] + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0, None], [None, self._A22_v0]] ) _M_init = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] + self._block_domain_M, self._block_domain_M, + blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] ) if self.options.solver in get_args(LiteralOptions.OptsSaddlePointSolver): self._Minv = inverse( - _M_init, - self.options.solver, + _M_init, self.options.solver, A11=self._A11_v0, A22=self._A22_v0, B1=self._B1_v0, B2=self._B2_v0, - recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, + maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, + recycle=self.options.solver_params.recycle, ) - # Allocate memory for call - self._untemp = self.variables.u.spline.vector.space.zeros() - - elif self._variant == "Uzawa": - self._solver_UzawaNumpy = SaddlePointSolver( - Apre=_Anppre, - A=_Anp, - B=_Bnp, - F=_Fnp, - method_to_solve=self._method_to_solve, - preconditioner=self._preconditioner, - spectralanalysis=self.options.spectralanalysis, + else: + self._Minv = inverse( + _M_init, self.options.solver, + recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, + maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, ) - def __call__(self, dt): - # current variables - unfeec = self.variables.u.spline.vector - uenfeec = self.variables.ue.spline.vector - phinfeec = self.variables.phi.spline.vector - - if self._variant == "GMRES": - if self._lifting: - phinfeeccopy = self.derhamv0.create_spline_function("phi", space_id="L2") - phinfeeccopy.vector = phinfeec - # unfeec in space Hdiv, u0 in space Hdiv_0 - unfeeccopy = self.derhamv0.create_spline_function("u", space_id="Hdiv") - u0 = self.derhamv0.create_spline_function("u", space_id="Hdiv") - u_prime = self.derhamv0.create_spline_function("u", space_id="Hdiv") - u0.vector = uenfeec - unfeeccopy.vector = uenfeec - apply_essential_bc_stencil(u0.vector[0], axis=0, ext=-1, order=0) - apply_essential_bc_stencil(u0.vector[0], axis=0, ext=1, order=0) - u_prime.vector = unfeeccopy.vector - u0.vector - - uenfeeccopy = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue0 = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue_prime = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue0.vector = uenfeec - uenfeeccopy.vector = uenfeec - apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=-1, order=0) - apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=1, order=0) - ue_prime.vector = uenfeeccopy.vector - ue0.vector - - _A11 = ( - self._M2 / dt - - self._M2B / self._eps_norm - + self._nu - * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - ) - _A12 = None - _A21 = _A12 - _A22 = ( - self._nu_e - * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - + self._M2B / self._eps_norm - - self._stab_sigma * IdentityOperator(_A11.domain) - ) - - if self._lifting: - _A11prime = -self._M2B / self._eps_norm + self._nu * ( - self.derhamv0.div.T @ self._M3 @ self.derhamv0.div - + self._basis_opsv0.S21.T - @ self.derhamv0.curl.T - @ self._M2 - @ self.derhamv0.curl - @ self._basis_opsv0.S21 - ) - _A22prime = ( - self._nu_e - * ( - self.derhamv0.div.T @ self._M3 @ self.derhamv0.div - + self._basis_opsv0.S21.T - @ self.derhamv0.curl.T - @ self._M2 - @ self.derhamv0.curl - @ self._basis_opsv0.S21 - ) - + self._M2B / self._eps_norm - - self._stab_sigma * IdentityOperator(_A11.domain) - ) - _B1 = -self._M3 @ self._div - _B2 = self._M3 @ self._div - - if _A12 is not None: - assert _A11.codomain == _A12.codomain - if _A21 is not None: - assert _A22.codomain == _A21.codomain - assert _B1.codomain == _B2.codomain - if _A12 is not None: - assert _A11.domain == _A12.domain == _B1.domain - if _A21 is not None: - assert _A21.domain == _A22.domain == _B2.domain - assert _A22.domain == _B2.domain - assert _A11.domain == _B1.domain - - _blocksA = [[_A11, _A12], [_A21, _A22]] - _A = BlockLinearOperator(self._block_domainA, self._block_codomainA, blocks=_blocksA) - _blocksB = [[_B1, _B2]] - _B = BlockLinearOperator(self._block_domainB, self._block_codomainB, blocks=_blocksB) - if self._lifting: - _blocksF = [ - self._M2.dot(self._F1) + self._M2.dot(u0.vector) / dt - _A11prime.dot(u_prime.vector), - self._M2.dot(self._F2) - _A22prime.dot(ue_prime.vector), - ] - else: - _blocksF = [ - self._M2.dot(self._F1) + self._M2.dot(unfeec) / dt, - self._M2.dot(self._F2), - ] - _F = BlockVector(self._block_domainA, blocks=_blocksF) - - # Imported solver - self._solver_GMRES.A = _A - self._solver_GMRES.B = _B - self._solver_GMRES.F = _F - - if self._lifting: - ( - _sol1, - _sol2, - info, - ) = self._solver_GMRES(u0.vector, ue0.vector, phinfeeccopy.vector) - - un_temp = self.derham.create_spline_function("u", space_id="Hdiv") - un_temp.vector = _sol1[0] + u_prime.vector + # ---- projector ------------------------------------------------------- + + self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + + # ---- solution spline functions (unconstrained) ----------------------- + + self._u = self.derham.create_spline_function("u", space_id="Hdiv") + self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") + self._phi = self.derham.create_spline_function("phi", space_id="L2") + + # ---- BC lifts (unconstrained) ---------------------------------------- + + self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") + self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") + + for u_prime, boundary_data, boundary_conditions in [ + (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), + (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), + ]: + if boundary_data is None: + continue + for (d, side), f_bc in boundary_data.items(): + if boundary_conditions.get((d, side)) == "dirichlet": + bc_pulled = lambda *etas, f=f_bc: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], + lambda *etas: bc_pulled(*etas)[1], + lambda *etas: bc_pulled(*etas)[2]]) + for (d2, side2), kind2 in boundary_conditions.items(): + if kind2 == "dirichlet" and (d2, side2) != (d, side): + apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) + u_prime.vector += _vec + + self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") + self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") + + self._u_prime_v0.vector = self._u_prime.vector + self._ue_prime_v0.vector = self._ue_prime.vector + + # ---- projected source terms (unconstrained) -------------------------- + + self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") + self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") + + for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: + if source is not None: + src_pulled = lambda *etas, f=source: self.domain.pull( + [lambda x,y,z, f=f: f(x,y,z)[0], + lambda x,y,z, f=f: f(x,y,z)[1], + lambda x,y,z, f=f: f(x,y,z)[2]], + *etas, kind="2") + rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], + lambda *etas: src_pulled(*etas)[1], + lambda *etas: src_pulled(*etas)[2]]) + + # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- + + self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") - uen_temp = self.derham.create_spline_function("ue", space_id="Hdiv") - uen_temp.vector = _sol1[1] + ue_prime.vector - - phin_temp = self.derham.create_spline_function("phi", space_id="L2") - phin_temp.vector = _sol2 - - un = un_temp.vector - uen = uen_temp.vector - phin = phin_temp.vector + # ========================================================================= + ### Time step + # ========================================================================= - else: - ( - _sol1, - _sol2, - info, - ) = self._solver_GMRES(unfeec, uenfeec, phinfeec) - un = _sol1[0] - uen = _sol1[1] - phin = _sol2 - # write new coeffs into self.feec_vars - - max_du, max_due, max_dphi = self.update_feec_variables(u=un, ue=uen, phi=phin) - - elif self._variant == "Uzawa": - # Numpy - A11np = self._M2np / dt + self._A11np_notimedependency - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - A11np += self._stab_sigma * xp.identity(A11np.shape[0]) - _A22prenp = self._A22prenp - A22np = self.A22np - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - A11np += self._stab_sigma * sc.sparse.eye(A11np.shape[0], format="csr") - _A22prenp = self._A22prenp - A22np = self.A22np - - # _Anp[1] and _Anppre[1] remain unchanged - _Anp = [A11np, A22np] - if self._preconditioner: - _A11prenp = self._M2np / dt # + self._A11prenp_notimedependency - _Anppre = [_A11prenp, _A22prenp] - - if self._lifting: - # unfeec in space Hdiv, u0 in space Hdiv_0 - unfeeccopy = self.derhamv0.create_spline_function("u", space_id="Hdiv") - u0 = self.derhamv0.create_spline_function("u", space_id="Hdiv") - u_prime = self.derham.create_spline_function("u", space_id="Hdiv") - u0.vector = unfeec - unfeeccopy.vector = unfeec - apply_essential_bc_stencil(u0.vector[0], axis=0, ext=-1, order=0) - apply_essential_bc_stencil(u0.vector[0], axis=0, ext=1, order=0) - u_prime.vector = unfeeccopy.vector - u0.vector - - uenfeeccopy = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue0 = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue_prime = self.derhamv0.create_spline_function("ue", space_id="Hdiv") - ue0.vector = uenfeec - uenfeeccopy.vector = uenfeec - apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=-1, order=0) - apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=1, order=0) - ue_prime.vector = uenfeeccopy.vector - ue0.vector - - _F1np = ( - self._M2np @ self._F1np - + 1.0 / dt * self._M2np.dot(u0.vector.toarray()) - - self._A11np_notimedependency.dot(u_prime.vector.toarray()) - ) - _F2np = self._M2np @ self._F2np - self.A22np.dot(ue_prime.vector.toarray()) - _Fnp = [_F1np, _F2np] - else: - _F1np = self._M2np @ self._F1np + 1.0 / dt * self._M2np.dot(unfeec.toarray()) - _F2np = self._M2np @ self._F2np - _Fnp = [_F1np, _F2np] - - if self.rank == 0: - if self._preconditioner: - self._solver_UzawaNumpy.Apre = _Anppre - self._solver_UzawaNumpy.A = _Anp - self._solver_UzawaNumpy.F = _Fnp - if self._lifting: - un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy( - u0.vector, - ue0.vector, - phinfeec, - ) + def __call__(self, dt): - un += u_prime.vector.toarray() - uen += ue_prime.vector.toarray() - else: - un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy( - unfeec, - uenfeec, - phinfeec, - ) + # --- copy current state --- + self._u.vector = self.variables.u.spline.vector + self._ue.vector = self.variables.ue.spline.vector - dimlist = [[shp - 2 * pi for shp, pi in zip(unfeec[i][:].shape, self.derham.p)] for i in range(3)] - dimphi = [shp - 2 * pi for shp, pi in zip(phinfeec[:].shape, self.derham.p)] - u_temp = BlockVector(self.derham.Vh["2"]) - ue_temp = BlockVector(self.derham.Vh["2"]) - phi_temp = StencilVector(self.derham.Vh["3"]) - test = 0 - for i, bl in enumerate(u_temp.blocks): - s = bl.starts - e = bl.ends - totaldim = dimlist[i][0] * dimlist[i][1] * dimlist[i][2] - test += totaldim - bl[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = un[ - i * totaldim : (i + 1) * totaldim - ].reshape(*dimlist[i]) - - for i, bl in enumerate(ue_temp.blocks): - s = bl.starts - e = bl.ends - totaldim = dimlist[i][0] * dimlist[i][1] * dimlist[i][2] - bl[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = uen[ - i * totaldim : (i + 1) * totaldim - ].reshape(*dimlist[i]) - - s = phi_temp.starts - e = phi_temp.ends - phi_temp[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = phin.reshape(*dimphi) - else: - print("TwoFluidQuasiNeutralFull is only running on one MPI.") + # --- rebuild system matrix if dt changed --- TODO change uzawa internals + if dt != self._dt: + self._dt = dt + _A = BlockLinearOperator( + self._block_domain_v0, self._block_codomain_v0, + blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] + ) + + _M = BlockLinearOperator( + self._block_domain_M, self._block_domain_M, + blocks=[[_A, self._B_v0.T], [self._B_v0, None]] + ) + self._Minv.linop = _M + + # --- assemble RHS in unconstrained space, then zero boundary DOFs --- + # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' + # electron: F2 = rhs_ue - A22 * ue' + self._rhs_vec_u.vector = (self._rhs_u.vector + + self._M2.dot(self._u.vector) / dt + - self._A11.dot(self._u_prime.vector) + - self._M2.dot(self._u_prime.vector) / dt) + self._rhs_vec_ue.vector = (self._rhs_ue.vector + - self._A22.dot(self._ue_prime.vector)) + + self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) + self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) + + # --- build block RHS and solve --- + _F = BlockVector(self._block_domain_v0, + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) + _RHS = BlockVector(self._block_domain_M, + blocks=[_F, self._block_codomain_B_v0.zeros()]) + + _sol = self._Minv.dot(_RHS) + info = self._Minv.get_info() + + # --- reconstruct full solution: u = u_0 + u' --- + self._u.vector = _sol[0][0] + self._u_prime_v0.vector + self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector + self._phi.vector = _sol[1] + + # --- update FEEC variables --- + max_diffs = self.update_feec_variables( + u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector + ) - # write new coeffs into self.feec_vars - max_du, max_due, max_dphi = self.update_feec_variables(u=u_temp, ue=ue_temp, phi=phi_temp) + if self.options.solver_params.info and self._rank == 0: + print(f"Status: {info['success']}, Iterations: {info['niter']}") + print(f"Max diffs: {max_diffs}") + # --- assemble RHS in unconstrained space, then zero boundary DOFs --- + # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' + # electron: F2 = rhs_ue - A22 * ue' + self._rhs_vec_u.vector = (self._rhs_u.vector + + self._M2.dot(self._u.vector) / dt + - self._A11.dot(self._u_prime.vector) + - self._M2.dot(self._u_prime.vector) / dt) + self._rhs_vec_ue.vector = (self._rhs_ue.vector + - self._A22.dot(self._ue_prime.vector)) + + self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) + self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) + + # --- build block RHS and solve --- + _F = BlockVector(self._block_domain_v0, + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) + _RHS = BlockVector(self._block_domain_M, + blocks=[_F, self._block_codomain_B_v0.zeros()]) + + _sol = self._Minv.dot(_RHS) + info = self._Minv.get_info() + + # --- reconstruct full solution: u = u_0 + u' --- + self._u.vector = _sol[0][0] + self._u_prime_v0.vector + self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector + self._phi.vector = _sol[1] + + # --- update FEEC variables --- + max_diffs = self.update_feec_variables( + u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector + ) if self.options.solver_params.info and self._rank == 0: - logger.info(f"Status: {info['success']}, Iterations: {info['niter']}") - logger.info(f"Max diffs: {max_diffs}") - logger.info(f"Status: {info['success']}, Iterations: {info['niter']}") - logger.info(f"Max diffs: {max_diffs}") + print(f"Status: {info['success']}, Iterations: {info['niter']}") + print(f"Max diffs: {max_diffs}") \ No newline at end of file diff --git a/src/struphy/propagators/propagators_fields_two_fluid_new.py b/src/struphy/propagators/propagators_fields_two_fluid_new.py deleted file mode 100644 index c5cf54039..000000000 --- a/src/struphy/propagators/propagators_fields_two_fluid_new.py +++ /dev/null @@ -1,440 +0,0 @@ -from dataclasses import dataclass -from typing import Callable, Literal, cast -from warnings import warn - -from feectools.api.essential_bc import apply_essential_bc_stencil -from feectools.ddm.mpi import mpi as MPI -from feectools.linalg.basic import IdentityOperator, InverseLinearOperator -from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace -from feectools.linalg.solvers import inverse - -from struphy.feec.basis_projection_ops import BasisProjectionOperators -from struphy.feec.mass import WeightedMassOperators -from struphy.feec.projectors import L2Projector -from struphy.feec.psydac_derham import Derham -from struphy.io.options import OptsGenSolver -from struphy.linear_algebra.solver import SolverParameters -from struphy.models.variables import FEECVariable -from struphy.propagators.base import Propagator -from struphy.utils.utils import check_option - - -class TwoFluidQuasiNeutralFull(Propagator): - r""":ref:`FEEC ` discretization of the following equations: - find :math:`\mathbf u \in H(\textnormal{div})`, :math:`\mathbf u_e \in H(\textnormal{div})` and :math:`\mathbf \phi \in L^2` such that - - .. math:: - - \int_{\Omega} \partial_t \mathbf{u}\cdot \mathbf{v} \, \textrm d\mathbf{x} &= \int_{\Omega} \phi \nabla \! \cdot \! \mathbf{v} \, \textrm d\mathbf{x} + \int_{\Omega} \mathbf{u}\! \times \! \mathbf{B}_0 \cdot \mathbf{v} \, \textrm d\mathbf{x} + \nu \int_{\Omega} \nabla \mathbf{u}\! : \! \nabla \mathbf{v} \, \textrm d\mathbf{x} + \int_{\Omega} f \mathbf{v} \, \textrm d\mathbf{x} \qquad \forall \, \mathbf{v} \in H(\textrm{div}) \,. - \\[2mm] - 0 &= - \int_{\Omega} \phi \nabla \! \cdot \! \mathbf{v_e} \, \textrm d\mathbf{x} - \int_{\Omega} \mathbf{u_e} \! \times \! \mathbf{B}_0 \cdot \mathbf{v_e} \, \textrm d\mathbf{x} + \nu_e \int_{\Omega} \nabla \mathbf{u_e} \!: \! \nabla \mathbf{v_e} \, \textrm d\mathbf{x} + \int_{\Omega} f_e \mathbf{v_e} \, \textrm d\mathbf{x} \qquad \forall \ \mathbf{v_e} \in H(\textrm{div}) \,. - \\[2mm] - 0 &= \int_{\Omega} \psi \nabla \cdot (\mathbf{u}-\mathbf{u_e}) \, \textrm d\mathbf{x} \qquad \forall \, \psi \in L^2 \,. - - :ref:`time_discret`: fully implicit. - """ - - # ========================================================================= - ### State variables (ion velocity u, electron velocity ue, pressure phi) - # ========================================================================= - - class Variables(): - def __init__(self) -> None: - self._u: FEECVariable | None = None - self._ue: FEECVariable | None = None - self._phi: FEECVariable | None = None - - @property - def u(self) -> FEECVariable | None: - return self._u - - @u.setter - def u(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hdiv" - self._u = new - - @property - def ue(self) -> FEECVariable | None: - return self._ue - - @ue.setter - def ue(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hdiv" - self._ue = new - - @property - def phi(self) -> FEECVariable | None: - return self._phi - - @phi.setter - def phi(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "L2" - self._phi = new - - def __init__(self): - self.variables = self.Variables() - - # ========================================================================= - ### Options - # ========================================================================= - - @dataclass - class Options(): - - nu: float | None = None - nu_e: float | None = None - eps_norm: float | None = None - - # boundary conditions per species - # supported kinds: "periodic", "dirichlet" - # future: "neumann", "robin" - boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_data_u: dict[tuple[int, int], Callable] | None = None - boundary_data_ue: dict[tuple[int, int], Callable] | None = None - - source_u: Callable | None = None - source_ue: Callable | None = None - - stab_sigma: float | None = None - - solver: OptsGenSolver = "gmres" - solver_params: SolverParameters | None = None - - def __post_init__(self): - - # --- required parameters --- - assert self.nu is not None, "nu must be specified" - assert self.nu_e is not None, "nu_e must be specified" - assert self.eps_norm is not None, "eps_norm must be specified" - assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" - assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" - - # --- physical parameter sanity checks --- - if self.nu < 0: - raise ValueError(f"nu must be non-negative, got {self.nu}") - if self.nu_e < 0: - raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") - if self.eps_norm <= 0: - raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - - # --- check all axes are covered --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for d in range(3): - for side in (-1, 1): - assert (d, side) in bcs, \ - f"{name} is missing entry for axis {d} side {side}" - - # --- periodic consistency: periodic must be paired on both sides --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for (d, side), kind in bcs.items(): - if kind == "periodic": - assert bcs.get((d, -side)) == "periodic", \ - f"{name}: axis {d} side {side} is periodic but opposite side is not" - - # --- ions and electrons must agree on which axes are periodic --- - for d in range(3): - u_left = self.boundary_conditions_u.get((d, -1)) - ue_left = self.boundary_conditions_ue.get((d, -1)) - u_right = self.boundary_conditions_u.get((d, 1)) - ue_right = self.boundary_conditions_ue.get((d, 1)) - u_periodic = (u_left == "periodic") - ue_periodic = (ue_left == "periodic") - if u_periodic != ue_periodic: - raise ValueError( - f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " - f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" - ) - - # --- warn for Dirichlet faces with no boundary data --- - for species, bcs, data, label in [ - ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), - ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), - ]: - has_dirichlet = any(v == "dirichlet" for v in bcs.values()) - if has_dirichlet: - if data is None: - warn(f"Dirichlet BCs specified for {species} but no {label} given " - f"— defaulting to homogeneous Dirichlet on all faces.") - else: - for (d, side), kind in bcs.items(): - if kind == "dirichlet" and (d, side) not in data: - warn(f"No {label} given for axis {d} side {side} " - f"— defaulting to homogeneous Dirichlet.") - - # --- warn if no source terms --- - if self.source_u is None: - warn("No source_u specified — defaulting to zero.") - if self.source_ue is None: - warn("No source_ue specified — defaulting to zero.") - - # --- defaults --- - if self.stab_sigma is None: - warn("stab_sigma not specified, defaulting to 0.0") - self.stab_sigma = 0.0 - - check_option(self.solver, OptsGenSolver) - if self.solver_params is None: - self.solver_params = SolverParameters() - - @property - def options(self) -> Options: - assert hasattr(self, "_options"), "Options not set." - return self._options - - @options.setter - def options(self, new): - assert isinstance(new, self.Options) - if MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nNew options for propagator '{self.__class__.__name__}':") - for k, v in new.__dict__.items(): - print(f" {k}: {v}") - self._options = new - - # ========================================================================= - ### Boundary condition helpers - # ========================================================================= - - def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): - - dirichlet_bc = [] - for d in range(3): - if spl_kind[d]: # periodic spline — no clamping ever - dirichlet_bc.append((False, False)) - else: - left = boundary_conditions.get((d, -1)) == "dirichlet" - right = boundary_conditions.get((d, 1)) == "dirichlet" - dirichlet_bc.append((left, right)) - return tuple(tuple(bc) for bc in dirichlet_bc) - - def _apply_boundary_conditions(self, vec, boundary_conditions): - """Zero out Dirichlet DOFs on the given stencil vector.""" - for (d, side), kind in boundary_conditions.items(): - if kind == "dirichlet": - apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) - # future: neumann and robin require no zeroing here - - # ========================================================================= - ### Allocate - # ========================================================================= - - def allocate(self, verbose=False): - - self.verbose = verbose - self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 - self._dt = None - - # ---- constrained (v0) de Rham complex -------------------------------- - - _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) - _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) - _dirichlet_bc = tuple( - (l_u or l_ue, r_u or r_ue) - for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) - ) - - self._derham_v0 = Derham( - self.derham.Nel, self.derham.p, self.derham.spl_kind, - domain=self.domain, dirichlet_bc=_dirichlet_bc, - ) - self._mass_ops_v0 = WeightedMassOperators( - self._derham_v0, self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.mass_ops.weights["eq_mhd"], - ) - self._basis_ops_v0 = BasisProjectionOperators( - self._derham_v0, self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.basis_ops.weights["eq_mhd"], - ) - - # ---- unconstrained operators (for RHS assembly) ---------------------- - - - self._M2 = self.mass_ops.M2 - self._M2B = - self.mass_ops.M2B - self._div = self.derham.div - self._curl = self.derham.curl - self._S21 = self.basis_ops.S21 - - self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div - + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - - self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl - self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) - + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) - - # ---- constrained operators (for system matrix) ----------------------- - - self._M2_v0 = self._mass_ops_v0.M2 - self._M3_v0 = self._mass_ops_v0.M3 - self._M2B_v0 = - self._mass_ops_v0.M2B - self._div_v0 = self._derham_v0.div - self._curl_v0 = self._derham_v0.curl - self._S21_v0 = self._basis_ops_v0.S21 - - self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 - + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) - - self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) - + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) - - # ---- block saddle-point system ---------------------------------------- - - self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) - self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self._derham_v0.Vh["3"] - - self._B1_v0 = - self._M3_v0 @ self._div_v0 - self._B2_v0 = self._M3_v0 @ self._div_v0 - - self._B_v0 = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_B_v0, - blocks=[[self._B1_v0, self._B2_v0]] - ) - - self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) - - _A_init = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, - blocks=[[self._A11_v0, None], [None, self._A22_v0]] - ) - _M_init = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, - blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] - ) - self._Minv = cast(InverseLinearOperator, inverse( - _M_init, self.options.solver, - x0=None, - tol=self.options.solver_params.tol, - maxiter=self.options.solver_params.maxiter, - verbose=self.options.solver_params.verbose, - )) - - # ---- projector ------------------------------------------------------- - - self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - - # ---- solution spline functions (unconstrained) ----------------------- - - self._u = self.derham.create_spline_function("u", space_id="Hdiv") - self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") - self._phi = self.derham.create_spline_function("phi", space_id="L2") - - # ---- BC lifts (unconstrained) ---------------------------------------- - - self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") - self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") - - for u_prime, boundary_data, boundary_conditions in [ - (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), - (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), - ]: - if boundary_data is None: - continue - for (d, side), f_bc in boundary_data.items(): - if boundary_conditions.get((d, side)) == "dirichlet": - bc_pulled = lambda *etas, f=f_bc: self.domain.pull( - [lambda x,y,z, f=f: f(x,y,z)[0], - lambda x,y,z, f=f: f(x,y,z)[1], - lambda x,y,z, f=f: f(x,y,z)[2]], - *etas, kind="2") - _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], - lambda *etas: bc_pulled(*etas)[1], - lambda *etas: bc_pulled(*etas)[2]]) - for (d2, side2), kind2 in boundary_conditions.items(): - if kind2 == "dirichlet" and (d2, side2) != (d, side): - apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) - u_prime.vector += _vec - - self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") - self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") - - self._u_prime_v0.vector = self._u_prime.vector - self._ue_prime_v0.vector = self._ue_prime.vector - - # ---- projected source terms (unconstrained) -------------------------- - - self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") - self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") - - for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: - if source is not None: - src_pulled = lambda *etas, f=source: self.domain.pull( - [lambda x,y,z, f=f: f(x,y,z)[0], - lambda x,y,z, f=f: f(x,y,z)[1], - lambda x,y,z, f=f: f(x,y,z)[2]], - *etas, kind="2") - rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], - lambda *etas: src_pulled(*etas)[1], - lambda *etas: src_pulled(*etas)[2]]) - - # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- - - self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") - self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") - - # ========================================================================= - ### Time step - # ========================================================================= - - def __call__(self, dt): - - # --- copy current state --- - self._u.vector = self.variables.u.spline.vector - self._ue.vector = self.variables.ue.spline.vector - - # --- rebuild system matrix if dt changed --- - if dt != self._dt: - self._dt = dt - _A = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, - blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] - ) - _M = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, - blocks=[[_A, self._B_v0.T], [self._B_v0, None]] - ) - self._Minv.linop = _M - - # --- assemble RHS in unconstrained space, then zero boundary DOFs --- - # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' - # electron: F2 = rhs_ue - A22 * ue' - self._rhs_vec_u.vector = (self._rhs_u.vector - + self._M2.dot(self._u.vector) / dt - - self._A11.dot(self._u_prime.vector) - - self._M2.dot(self._u_prime.vector) / dt) - self._rhs_vec_ue.vector = (self._rhs_ue.vector - - self._A22.dot(self._ue_prime.vector)) - - self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) - self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) - - # --- build block RHS and solve --- - _F = BlockVector(self._block_domain_v0, - blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) - _RHS = BlockVector(self._block_domain_M, - blocks=[_F, self._block_codomain_B_v0.zeros()]) - - _sol = self._Minv.dot(_RHS) - info = self._Minv.get_info() - - # --- reconstruct full solution: u = u_0 + u' --- - self._u.vector = _sol[0][0] + self._u_prime_v0.vector - self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector - self._phi.vector = _sol[1] - - # --- update FEEC variables --- - max_diffs = self.update_feec_variables( - u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector - ) - - if self.options.solver_params.info and self._rank == 0: - print(f"Status: {info['success']}, Iterations: {info['niter']}") - print(f"Max diffs: {max_diffs}") \ No newline at end of file diff --git a/struphy-parameter-files b/struphy-parameter-files new file mode 160000 index 000000000..5143ca521 --- /dev/null +++ b/struphy-parameter-files @@ -0,0 +1 @@ +Subproject commit 5143ca52173bb00766c5032d2b9e652ecabd885e From 3a30a3cc9554c3598bf9b3f7fd8bcb70567d04b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Mon, 16 Mar 2026 14:08:25 +0000 Subject: [PATCH 14/41] moved v0 de rham complex construction into the Derham class, with everything that entails --- src/struphy/feec/psydac_derham.py | 216 ++++--- src/struphy/io/options.py | 52 +- src/struphy/io/setup.py | 89 +++ src/struphy/propagators/propagators_fields.py | 575 ++---------------- 4 files changed, 295 insertions(+), 637 deletions(-) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index adcb1e81b..09048cb7c 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -67,117 +67,115 @@ class DiscreteDerham: Parameters ---------- - V0 : TensorFemSpace - First space of the de Rham sequence : H1 space - V1 : VectorFemSpace - Second space of the de Rham sequence : Hcurl space - V2 : VectorFemSpace - Third space of the de Rham sequence : Hdiv space - V3 : TensorFemSpace - Fourth space of the de Rham sequence : L2 space + Nel : list[int] + Number of elements in each direction. - Notes - ----- - On construction, differential operators are created and attached to the - input spaces as convenience attributes: - - - ``V0.grad`` and ``V0.diff`` - - ``V1.curl`` and ``V1.diff`` - - ``V2.div`` and ``V2.diff`` - """ - - def __init__(self, V0: TensorFemSpace, V1: VectorFemSpace, V2: VectorFemSpace, V3: TensorFemSpace): - spaces = (V0, V1, V2, V3) - assert all(isinstance(space, (TensorFemSpace, VectorFemSpace)) for space in spaces) - - self._V0 = V0 - self._V1 = V1 - self._V2 = V2 - self._V3 = V3 - self._spaces = spaces - self._dim = 3 + p : list[int] + Spline degree in each direction. - D0 = Gradient3D(V0, V1) - D1 = Curl3D(V1, V2) - D2 = Divergence3D(V2, V3) + spl_kind : list[bool] + Kind of spline in each direction (True=periodic, False=clamped). - V0.diff = V0.grad = D0 - V1.diff = V1.curl = D1 - V2.diff = V2.div = D2 - - # -------------------------------------------------------------------------- - @property - def dim(self) -> int: - """Dimension of the physical and logical domains, which are assumed to be the same.""" - return self._dim - - @property - def V0(self) -> TensorFemSpace: - """First space of the de Rham sequence : H1 space""" - return self._V0 + dirichlet_bc : list[list[bool]] + Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). - @property - def V1(self) -> VectorFemSpace: - """Second space of the de Rham sequence : Hcurl space""" - return self._V1 + nq_pr : list[int] + Number of Gauss-Legendre quadrature points in each direction for geometric projectors (default = p+1, leads to exact integration of degree 2p+1 polynomials). - @property - def V2(self) -> VectorFemSpace: - """Third space of the de Rham sequence : Hdiv space""" - return self._V2 + nquads : list[int] + Number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials). - @property - def V3(self) -> TensorFemSpace: - """Fourth space of the de Rham sequence : L2 space""" - return self._V3 + comm : mpi4py.MPI.Intracomm + MPI communicator (within a clone if domain cloning is used, otherwise MPI.COMM_WORLD) - @property - def spaces(self) -> tuple[TensorFemSpace | VectorFemSpace, ...]: - """Spaces of the proper de Rham sequence (excluding Hvec).""" - return self._spaces + mpi_dims_mask: list of bool + True if the dimension is to be used in the domain decomposition (=default for each dimension). + If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed. - @property - def derivatives_as_matrices(self): - """Differential operators of the De Rham sequence as LinearOperator objects.""" - return tuple(V.diff.linop for V in self.spaces[:-1]) + with_projectors : bool + Whether to add global commuting projectors to the diagram. - @property - def derivatives(self): - """Differential operators of the De Rham sequence as `DiffOperator` objects. + polar_ck : int + Smoothness at a polar singularity at eta_1=0 (default -1 : standard tensor product splines, OR 1 : C1 polar splines) - Those are objects with `domain` and `codomain` properties that are `FemSpace`, - they act on `FemField` (they take a `FemField` of their `domain` as input and return - a `FemField` of their `codomain`. - """ - return tuple(V.diff for V in self.spaces[:-1]) + local_projectors : bool + Whether to build the local commuting projectors based on quasi-inter-/histopolation. - # -------------------------------------------------------------------------- - def projectors(self, *, kind="global", nquads=None) -> tuple[GlobalGeometricProjector, ...]: - """Projectors mapping callable functions of the physical coordinates to a - corresponding `FemField` object in the De Rham sequence. + domain : struphy.geometry.base.Domain + Mapping from logical unit cube to physical domain (only needed in case of polar splines polar_ck=1). + """ - Parameters - ---------- - kind : str - Type of the projection : at the moment, only global is accepted and - returns geometric commuting projectors based on interpolation/histopolation - for the De Rham sequence (GlobalGeometricProjector objects). + def __init__( + self, + Nel: list | tuple, + p: list | tuple, + spl_kind: list | tuple, + *, + dirichlet_bc: list | tuple = None, + lifting: list | tuple = None, + nquads: list | tuple = None, + nq_pr: list | tuple = None, + comm=None, + mpi_dims_mask: list = None, + with_projectors: bool = True, + polar_ck: int = -1, + local_projectors: bool = False, + domain: Domain = None, + ): + # number of elements, spline degrees and kind of splines in each direction (periodic vs. clamped) + assert len(Nel) == 3 + assert len(p) == 3 + assert len(spl_kind) == 3 + + self._Nel = Nel + self._p = p + self._spl_kind = spl_kind + self._with_local_projectors = local_projectors - nquads : list(int) | tuple(int) - Number of quadrature points along each direction, to be used in Gauss - quadrature rule for computing the (approximated) degrees of freedom. + # boundary conditions at eta=0 and eta=1 in each direction (None for periodic, 'd' for homogeneous Dirichlet) + if dirichlet_bc is not None: + assert len(dirichlet_bc) == 3 + # make sure that boundary conditions are compatible with spline space + assert xp.all([bc == (False, False) for i, bc in enumerate(dirichlet_bc) if spl_kind[i]]) + + self._dirichlet_bc = dirichlet_bc + + # --- lifting: build constrained (v0) sub-complex --- + self._lifting = lifting + if lifting is not None: + assert len(lifting) == 3 + # lifting only makes sense on non-periodic axes + for d in range(3): + if spl_kind[d]: + assert lifting[d] == (False, False), \ + f"Axis {d} is periodic, lifting must be (False, False)" + + # v0 dirichlet_bc = dirichlet_bc OR lifting + if dirichlet_bc is not None: + v0_dirichlet_bc = tuple( + (d_l or l_l, d_r or l_r) + for (d_l, d_r), (l_l, l_r) in zip(dirichlet_bc, lifting) + ) + else: + v0_dirichlet_bc = lifting - Returns - ------- - P0, ..., Pn : callables - Projectors that can be called on any callable function that maps - from the physical space to R (scalar case) or R^d (vector case) and - returns a FemField belonging to the i-th space of the De Rham sequence - """ + self._derham_v0 = Derham( + Nel, p, spl_kind, + dirichlet_bc=v0_dirichlet_bc, + nquads=nquads, + nq_pr=nq_pr, + comm=comm, + mpi_dims_mask=mpi_dims_mask, + with_projectors=with_projectors, + polar_ck=polar_ck, + local_projectors=self.with_local_projectors, + domain=domain, + ) + else: + self._derham_v0 = None - if not (kind == "global"): - raise NotImplementedError("only global projectors are available") + # default p: exact integration of degree 2p+1 polynomials if nquads is None: nquads = [degree + 1 for degree in self.V0.degree] elif isinstance(nquads, int): @@ -837,27 +835,17 @@ def __init__( # collect arguments for kernels self._args_derham = DerhamArguments( - xp.array(self.degree), - self.V0fem.knots[0], - self.V0fem.knots[1], - self.V0fem.knots[2], - xp.array(self.V0.starts), + xp.array(self.p), + self.Vh_fem["0"].knots[0], + self.Vh_fem["0"].knots[1], + self.Vh_fem["0"].knots[2], + xp.array(self.Vh["0"].starts), ) - if MPI.COMM_WORLD.Get_rank() == 0 and verbose: - logger.info("\nDERHAM:") - logger.info(f"{'number of elements:'.ljust(25)} {num_elements}") - logger.info(f"{'spline degrees:'.ljust(25)} {degree}") - logger.info(f"{'boundary conditions:'.ljust(25)} {bcs}") - logger.info(f"{'GL quad pts (L2):'.ljust(25)} {nquads}") - logger.info(f"{'GL quad pts (hist):'.ljust(25)} {nquads_proj}") - logger.info(f"{'MPI proc. per dir.:'.ljust(25)} {self.domain_decomposition.nprocs}") - logger.info(f"{'use polar splines:'.ljust(25)} {self.polar_splines}") - logger.info(f"{'domain on process 0:'.ljust(25)} {self.domain_array[0]}") - - # ----------------------------- - # Input arguments as properties - # ----------------------------- + @property + def derham_v0(self): + return self._derham_v0 + @property def grid(self) -> TensorProductGrid: """The FEEC grid.""" diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 898881b61..b9f0b782a 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -255,9 +255,13 @@ class DerhamOptions(OptionsBase): Use ``None`` in a direction for periodic boundaries, or a tuple ``(left, right)`` with entries in ``{"free", "dirichlet"}`` for non-periodic boundaries. - nquads : tuple[int, int, int] | None - Number of Gauss-Legendre quadrature points per direction for cell - integrals. If ``None``, backend defaults are used. + lifting : tuple[tuple[bool]] + Whether to build a constrained (v0) sub-complex with additional clamping on each face. + Used for inhomogeneous Dirichlet BCs: the v0 complex clamps faces where + lifting is True, and the propagator builds a lift in the unconstrained space. + + nquads : tuple[int] + Number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials). nquads_proj : tuple[int, int, int] | None Number of Gauss-Legendre quadrature points per direction for geometric @@ -274,15 +278,13 @@ class DerhamOptions(OptionsBase): quasi-inter-/histopolation. """ - degree: tuple[int, int, int] = (1, 1, 1) - bcs: tuple[ - None | tuple[NonTrivialBC, NonTrivialBC], - None | tuple[NonTrivialBC, NonTrivialBC], - None | tuple[NonTrivialBC, NonTrivialBC], - ] = (None, None, None) - nquads: tuple[int, int, int] | None = None - nquads_proj: tuple[int, int, int] | None = None - polar_splines: bool = False + p: tuple = (1, 1, 1) + spl_kind: tuple = (True, True, True) + dirichlet_bc: tuple = ((False, False), (False, False), (False, False)) + lifting: tuple = ((False, False), (False, False), (False, False)) + nquads: tuple = None + nq_pr: tuple = None + polar_ck: LiteralOptions.PolarRegularity = -1 local_projectors: bool = False def __post_init__(self): @@ -306,6 +308,32 @@ def __repr_no_defaults__(self): def is_default(self): return all_class_params_are_default(self) + def to_dict(self) -> dict: + dct = { + "p": self.p, + "spl_kind": self.spl_kind, + "dirichlet_bc": self.dirichlet_bc, + "lifting": self.lifting, + "nquads": self.nquads, + "nq_pr": self.nq_pr, + "polar_ck": self.polar_ck, + "local_projectors": self.local_projectors, + } + return dct + + @classmethod + def from_dict(cls, dct) -> "DerhamOptions": + return cls( + p=dct["p"], + spl_kind=dct["spl_kind"], + dirichlet_bc=dct["dirichlet_bc"], + lifting=dct.get("lifting", ((False, False), (False, False), (False, False))), + nquads=dct["nquads"], + nq_pr=dct["nq_pr"], + polar_ck=dct["polar_ck"], + local_projectors=dct["local_projectors"], + ) + @dataclass class FieldsBackground(OptionsBase): diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index 19d3ed5b9..cb75f4219 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -31,6 +31,95 @@ def import_parameters_py(params_path: str, name: str = "parameters") -> ModuleTy return params_in +def setup_derham( + grid: TensorProductGrid, + options: DerhamOptions, + comm: MPI.Intracomm = None, + domain: Domain = None, + verbose=False, +): + """ + Creates the 3d derham sequence for given grid parameters. + + Parameters + ---------- + grid : TensorProductGrid + The FEEC grid. + + comm: Intracomm + MPI communicator (sub_comm if clones are used). + + domain : Domain, optional + The Struphy domain object for evaluating the mapping F : [0, 1]^3 --> R^3 and the corresponding metric coefficients. + + verbose : bool + Show info on screen. + + Returns + ------- + derham : struphy.feec.psydac_derham.Derham + Discrete de Rham sequence on the logical unit cube. + """ + + from struphy.feec.psydac_derham import Derham + + # number of grid cells + Nel = grid.Nel + # mpi + mpi_dims_mask = grid.mpi_dims_mask + + # spline degrees + p = options.p + # spline types (clamped vs. periodic) + spl_kind = options.spl_kind + # boundary conditions (Homogeneous Dirichlet or None) + dirichlet_bc = options.dirichlet_bc + # Number of quadrature points per histopolation cell + nq_pr = options.nq_pr + # Number of quadrature points per grid cell for L^2 + nquads = options.nquads + # C^k smoothness at eta_1=0 for polar domains + polar_ck = options.polar_ck + # local commuting projectors + local_projectors = options.local_projectors + + lifting = options.lifting + + derham = Derham( + Nel, + p, + spl_kind, + dirichlet_bc=dirichlet_bc, + lifting=lifting, + nquads=nquads, + nq_pr=nq_pr, + comm=comm, + mpi_dims_mask=mpi_dims_mask, + with_projectors=True, + polar_ck=polar_ck, + domain=domain, + local_projectors=local_projectors, + ) + + + if MPI.COMM_WORLD.Get_rank() == 0 and verbose: + print("\nDERHAM:") + print("number of elements:".ljust(25), Nel) + print("spline degrees:".ljust(25), p) + print("periodic bcs:".ljust(25), spl_kind) + print("hom. Dirichlet bc:".ljust(25), dirichlet_bc) + print("GL quad pts (L2):".ljust(25), nquads) + print("GL quad pts (hist):".ljust(25), nq_pr) + print( + "MPI proc. per dir.:".ljust(25), + derham.domain_decomposition.nprocs, + ) + print("use polar splines:".ljust(25), derham.polar_ck == 1) + print("domain on process 0:".ljust(25), derham.domain_array[0]) + + return derham + + def descend_options_dict( d: dict, out: list | dict, diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 66a1b129d..23c94842c 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -7703,10 +7703,6 @@ def __init__(self): ### Options # ========================================================================= - # ========================================================================= - ### Options - # ========================================================================= - @dataclass class Options(): @@ -7714,13 +7710,8 @@ class Options(): nu_e: float | None = None eps_norm: float | None = None - # boundary conditions per species - # supported kinds: "periodic", "dirichlet" - # future: "neumann", "robin" - boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_data_u: dict[tuple[int, int], Callable] | None = None - boundary_data_ue: dict[tuple[int, int], Callable] | None = None + boundary_data_u: dict[tuple[int, int], Callable] | None = None + boundary_data_ue: dict[tuple[int, int], Callable] | None = None source_u: Callable | None = None source_ue: Callable | None = None @@ -7733,11 +7724,9 @@ class Options(): def __post_init__(self): # --- required parameters --- - assert self.nu is not None, "nu must be specified" - assert self.nu_e is not None, "nu_e must be specified" - assert self.eps_norm is not None, "eps_norm must be specified" - assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" - assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" + assert self.nu is not None, "nu must be specified" + assert self.nu_e is not None, "nu_e must be specified" + assert self.eps_norm is not None, "eps_norm must be specified" # --- physical parameter sanity checks --- if self.nu < 0: @@ -7747,52 +7736,6 @@ def __post_init__(self): if self.eps_norm <= 0: raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - # --- check all axes are covered --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for d in range(3): - for side in (-1, 1): - assert (d, side) in bcs, \ - f"{name} is missing entry for axis {d} side {side}" - - # --- periodic consistency: periodic must be paired on both sides --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for (d, side), kind in bcs.items(): - if kind == "periodic": - assert bcs.get((d, -side)) == "periodic", \ - f"{name}: axis {d} side {side} is periodic but opposite side is not" - - # --- ions and electrons must agree on which axes are periodic --- - for d in range(3): - u_left = self.boundary_conditions_u.get((d, -1)) - ue_left = self.boundary_conditions_ue.get((d, -1)) - u_right = self.boundary_conditions_u.get((d, 1)) - ue_right = self.boundary_conditions_ue.get((d, 1)) - u_periodic = (u_left == "periodic") - ue_periodic = (ue_left == "periodic") - if u_periodic != ue_periodic: - raise ValueError( - f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " - f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" - ) - - # --- warn for Dirichlet faces with no boundary data --- - for species, bcs, data, label in [ - ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), - ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), - ]: - has_dirichlet = any(v == "dirichlet" for v in bcs.values()) - if has_dirichlet: - if data is None: - warn(f"Dirichlet BCs specified for {species} but no {label} given " - f"— defaulting to homogeneous Dirichlet on all faces.") - else: - for (d, side), kind in bcs.items(): - if kind == "dirichlet" and (d, side) not in data: - warn(f"No {label} given for axis {d} side {side} " - f"— defaulting to homogeneous Dirichlet.") - # --- warn if no source terms --- if self.source_u is None: warn("No source_u specified — defaulting to zero.") @@ -7804,13 +7747,12 @@ def __post_init__(self): warn("stab_sigma not specified, defaulting to 0.0") self.stab_sigma = 0.0 - check_option(self.solver, LiteralOptions.OptsGenSolver) + check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver) if self.solver_params is None: self.solver_params = SolverParameters() @property def options(self) -> Options: - assert hasattr(self, "_options"), "Options not set." assert hasattr(self, "_options"), "Options not set." return self._options @@ -7827,394 +7769,40 @@ def options(self, new): ### Boundary condition helpers # ========================================================================= - def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): - - dirichlet_bc = [] - for d in range(3): - if spl_kind[d]: # periodic spline — no clamping ever - dirichlet_bc.append((False, False)) - else: - left = boundary_conditions.get((d, -1)) == "dirichlet" - right = boundary_conditions.get((d, 1)) == "dirichlet" - dirichlet_bc.append((left, right)) - return tuple(tuple(bc) for bc in dirichlet_bc) - - def _apply_boundary_conditions(self, vec, boundary_conditions): - """Zero out Dirichlet DOFs on the given stencil vector.""" - for (d, side), kind in boundary_conditions.items(): - if kind == "dirichlet": - apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) - # future: neumann and robin require no zeroing here - - # ========================================================================= - ### Allocate - # ========================================================================= - - def allocate(self, verbose=False): - - self.verbose = verbose - self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 - self._dt = None - - # ---- constrained (v0) de Rham complex -------------------------------- - - _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) - _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) - _dirichlet_bc = tuple( - (l_u or l_ue, r_u or r_ue) - for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) - ) - - self._derham_v0 = Derham( - self.derham.Nel, self.derham.p, self.derham.spl_kind, - domain=self.domain, dirichlet_bc=_dirichlet_bc, - ) - self._mass_ops_v0 = WeightedMassOperators( - self._derham_v0, self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.mass_ops.weights["eq_mhd"], - ) - self._basis_ops_v0 = BasisProjectionOperators( - self._derham_v0, self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.basis_ops.weights["eq_mhd"], - ) - - # ---- unconstrained operators (for RHS assembly) ---------------------- - - - self._M2 = self.mass_ops.M2 - self._M2B = - self.mass_ops.M2B - self._div = self.derham.div - self._curl = self.derham.curl - self._S21 = self.basis_ops.S21 - - self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div - + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - - self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl - self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) - + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) - - # ---- constrained operators (for system matrix) ----------------------- - - self._M2_v0 = self._mass_ops_v0.M2 - self._M3_v0 = self._mass_ops_v0.M3 - self._M2B_v0 = - self._mass_ops_v0.M2B - self._div_v0 = self._derham_v0.div - self._curl_v0 = self._derham_v0.curl - self._S21_v0 = self._basis_ops_v0.S21 - - self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 - + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) - - self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) - + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) - - # ---- block saddle-point system ---------------------------------------- + def _get_dirichlet_faces(self): + """Infer which faces have Dirichlet BCs by comparing derham and derham_v0. - self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) - self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self._derham_v0.Vh["3"] - - self._B1_v0 = - self._M3_v0 @ self._div_v0 - self._B2_v0 = self._M3_v0 @ self._div_v0 - - self._B_v0 = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_B_v0, - blocks=[[self._B1_v0, self._B2_v0]] - ) - - self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) - - _A_init = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, - blocks=[[self._A11_v0, None], [None, self._A22_v0]] - ) - _M_init = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, - blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] - ) - - self._Minv = inverse( - _M_init, self.options.solver, - A11=self._A11_v0, - A22=self._A22_v0, - B1=self._B1_v0, - B2=self._B2_v0, - recycle=self.options.solver_params.recycle, - tol=self.options.solver_params.tol, - maxiter=self.options.solver_params.maxiter, - verbose=self.options.solver_params.verbose, - ) - - # ---- projector ------------------------------------------------------- - - self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - - # ---- solution spline functions (unconstrained) ----------------------- - - self._u = self.derham.create_spline_function("u", space_id="Hdiv") - self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") - self._phi = self.derham.create_spline_function("phi", space_id="L2") - - # ---- BC lifts (unconstrained) ---------------------------------------- - - self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") - self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") - - for u_prime, boundary_data, boundary_conditions in [ - (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), - (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), - ]: - if boundary_data is None: - continue - for (d, side), f_bc in boundary_data.items(): - if boundary_conditions.get((d, side)) == "dirichlet": - bc_pulled = lambda *etas, f=f_bc: self.domain.pull( - [lambda x,y,z, f=f: f(x,y,z)[0], - lambda x,y,z, f=f: f(x,y,z)[1], - lambda x,y,z, f=f: f(x,y,z)[2]], - *etas, kind="2") - _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], - lambda *etas: bc_pulled(*etas)[1], - lambda *etas: bc_pulled(*etas)[2]]) - for (d2, side2), kind2 in boundary_conditions.items(): - if kind2 == "dirichlet" and (d2, side2) != (d, side): - apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) - u_prime.vector += _vec - - self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") - self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") - - self._u_prime_v0.vector = self._u_prime.vector - self._ue_prime_v0.vector = self._ue_prime.vector - - # ---- projected source terms (unconstrained) -------------------------- - - self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") - self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") - - for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: - if source is not None: - src_pulled = lambda *etas, f=source: self.domain.pull( - [lambda x,y,z, f=f: f(x,y,z)[0], - lambda x,y,z, f=f: f(x,y,z)[1], - lambda x,y,z, f=f: f(x,y,z)[2]], - *etas, kind="2") - rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], - lambda *etas: src_pulled(*etas)[1], - lambda *etas: src_pulled(*etas)[2]]) - - # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- - - self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") - self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") - - # ========================================================================= - ### Time step - # ========================================================================= - - def __call__(self, dt): - - # --- copy current state --- - self._u.vector = self.variables.u.spline.vector - self._ue.vector = self.variables.ue.spline.vector - - # --- rebuild system matrix if dt changed --- - if dt != self._dt: - self._dt = dt - _A = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, - blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] - ) - - _M = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, - blocks=[[_A, self._B_v0.T], [self._B_v0, None]] - ) - self._Minv.linop = _M - class Variables(): - def __init__(self) -> None: - self._u: FEECVariable | None = None - self._ue: FEECVariable | None = None - self._phi: FEECVariable | None = None - - @property - def u(self) -> FEECVariable | None: - return self._u - - @u.setter - def u(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hdiv" - self._u = new - - @property - def ue(self) -> FEECVariable | None: - return self._ue - - @ue.setter - def ue(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hdiv" - self._ue = new - - @property - def phi(self) -> FEECVariable | None: - return self._phi - - @phi.setter - def phi(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "L2" - self._phi = new - - def __init__(self): - self.variables = self.Variables() - - # ========================================================================= - ### Options - # ========================================================================= - - @dataclass - class Options(): - - nu: float | None = None - nu_e: float | None = None - eps_norm: float | None = None - - # boundary conditions per species - # supported kinds: "periodic", "dirichlet" - # future: "neumann", "robin" - boundary_conditions_u: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_conditions_ue: dict[tuple[int, int], Literal["periodic", "dirichlet"]] | None = None - boundary_data_u: dict[tuple[int, int], Callable] | None = None - boundary_data_ue: dict[tuple[int, int], Callable] | None = None - - source_u: Callable | None = None - source_ue: Callable | None = None - - stab_sigma: float | None = None - - solver: LiteralOptions.OptsGenSolver = "gmres" - solver_params: SolverParameters | None = None - - def __post_init__(self): - - # --- required parameters --- - assert self.nu is not None, "nu must be specified" - assert self.nu_e is not None, "nu_e must be specified" - assert self.eps_norm is not None, "eps_norm must be specified" - assert self.boundary_conditions_u is not None, "boundary_conditions_u must be specified" - assert self.boundary_conditions_ue is not None, "boundary_conditions_ue must be specified" - - # --- physical parameter sanity checks --- - if self.nu < 0: - raise ValueError(f"nu must be non-negative, got {self.nu}") - if self.nu_e < 0: - raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") - if self.eps_norm <= 0: - raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - - # --- check all axes are covered --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for d in range(3): - for side in (-1, 1): - assert (d, side) in bcs, \ - f"{name} is missing entry for axis {d} side {side}" - - # --- periodic consistency: periodic must be paired on both sides --- - for name, bcs in [("boundary_conditions_u", self.boundary_conditions_u), - ("boundary_conditions_ue", self.boundary_conditions_ue)]: - for (d, side), kind in bcs.items(): - if kind == "periodic": - assert bcs.get((d, -side)) == "periodic", \ - f"{name}: axis {d} side {side} is periodic but opposite side is not" - - # --- ions and electrons must agree on which axes are periodic --- - for d in range(3): - u_left = self.boundary_conditions_u.get((d, -1)) - ue_left = self.boundary_conditions_ue.get((d, -1)) - u_right = self.boundary_conditions_u.get((d, 1)) - ue_right = self.boundary_conditions_ue.get((d, 1)) - u_periodic = (u_left == "periodic") - ue_periodic = (ue_left == "periodic") - if u_periodic != ue_periodic: - raise ValueError( - f"Axis {d}: ions and electrons must both be periodic or both non-periodic, " - f"got u={u_left}/{u_right}, ue={ue_left}/{ue_right}" - ) - - # --- warn for Dirichlet faces with no boundary data --- - for species, bcs, data, label in [ - ("u", self.boundary_conditions_u, self.boundary_data_u, "boundary_data_u"), - ("ue", self.boundary_conditions_ue, self.boundary_data_ue, "boundary_data_ue"), - ]: - has_dirichlet = any(v == "dirichlet" for v in bcs.values()) - if has_dirichlet: - if data is None: - warn(f"Dirichlet BCs specified for {species} but no {label} given " - f"— defaulting to homogeneous Dirichlet on all faces.") - else: - for (d, side), kind in bcs.items(): - if kind == "dirichlet" and (d, side) not in data: - warn(f"No {label} given for axis {d} side {side} " - f"— defaulting to homogeneous Dirichlet.") - - # --- warn if no source terms --- - if self.source_u is None: - warn("No source_u specified — defaulting to zero.") - if self.source_ue is None: - warn("No source_ue specified — defaulting to zero.") - - # --- defaults --- - if self.stab_sigma is None: - warn("stab_sigma not specified, defaulting to 0.0") - self.stab_sigma = 0.0 - - check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver) - if self.solver_params is None: - self.solver_params = SolverParameters() + A face is Dirichlet if it is unclamped in derham but clamped in derham_v0 + (i.e. lifting is True there). + """ + faces = [] + derham = self.derham + derham_v0 = derham.derham_v0 - @property - def options(self) -> Options: - assert hasattr(self, "_options"), "Options not set." - return self._options + if derham_v0 is None: + return faces - @options.setter - def options(self, new): - assert isinstance(new, self.Options) - if MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nNew options for propagator '{self.__class__.__name__}':") - for k, v in new.__dict__.items(): - print(f" {k}: {v}") - self._options = new + bc = derham.dirichlet_bc + bc_v0 = derham_v0.dirichlet_bc - # ========================================================================= - ### Boundary condition helpers - # ========================================================================= - - def _bc_to_dirichlet_flags(self, boundary_conditions, spl_kind): - - dirichlet_bc = [] for d in range(3): - if spl_kind[d]: # periodic spline — no clamping ever - dirichlet_bc.append((False, False)) - else: - left = boundary_conditions.get((d, -1)) == "dirichlet" - right = boundary_conditions.get((d, 1)) == "dirichlet" - dirichlet_bc.append((left, right)) - return tuple(tuple(bc) for bc in dirichlet_bc) - - def _apply_boundary_conditions(self, vec, boundary_conditions): - """Zero out Dirichlet DOFs on the given stencil vector.""" - for (d, side), kind in boundary_conditions.items(): - if kind == "dirichlet": - apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) - # future: neumann and robin require no zeroing here + if derham.spl_kind[d]: + continue # periodic axis, no Dirichlet + for s, side in enumerate((-1, 1)): + # clamped in v0 but not in derham => this is a lifted (inhom Dirichlet) face + unclamped = not bc[d][s] + clamped_v0 = bc_v0[d][s] if bc_v0 is not None else False + if unclamped and clamped_v0: + faces.append((d, side)) + # clamped in both => homogeneous Dirichlet, also need to zero DOFs + elif bc[d][s] and clamped_v0: + faces.append((d, side)) + return faces + + def _apply_essential_bc(self, vec): + """Zero out Dirichlet DOFs, inferred from derham vs derham_v0.""" + for (d, side) in self._dirichlet_faces: + apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) # ========================================================================= ### Allocate @@ -8226,19 +7814,13 @@ def allocate(self, verbose=False): self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 self._dt = None - # ---- constrained (v0) de Rham complex -------------------------------- + # ---- v0 de Rham complex (from derham.derham_v0) ---------------------- - _dirichlet_u = self._bc_to_dirichlet_flags(self.options.boundary_conditions_u, self.derham.spl_kind) - _dirichlet_ue = self._bc_to_dirichlet_flags(self.options.boundary_conditions_ue, self.derham.spl_kind) - _dirichlet_bc = tuple( - (l_u or l_ue, r_u or r_ue) - for (l_u, r_u), (l_ue, r_ue) in zip(_dirichlet_u, _dirichlet_ue) - ) + assert self.derham.derham_v0 is not None, \ + "derham must be constructed with lifting to use this propagator" + + self._derham_v0 = self.derham.derham_v0 - self._derham_v0 = Derham( - self.derham.Nel, self.derham.p, self.derham.spl_kind, - domain=self.domain, dirichlet_bc=_dirichlet_bc, - ) self._mass_ops_v0 = WeightedMassOperators( self._derham_v0, self.domain, verbose=self.options.solver_params.verbose, @@ -8250,8 +7832,11 @@ def allocate(self, verbose=False): eq_mhd=self.basis_ops.weights["eq_mhd"], ) - # ---- unconstrained operators (for RHS assembly) ---------------------- + # ---- Dirichlet faces (inferred from derham vs derham_v0) ------------- + self._dirichlet_faces = self._get_dirichlet_faces() + + # ---- unconstrained operators (for RHS assembly) ---------------------- self._M2 = self.mass_ops.M2 self._M2B = - self.mass_ops.M2B @@ -8314,11 +7899,11 @@ def allocate(self, verbose=False): A22=self._A22_v0, B1=self._B1_v0, B2=self._B2_v0, + recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, - recycle=self.options.solver_params.recycle, ) else: self._Minv = inverse( @@ -8330,6 +7915,7 @@ def allocate(self, verbose=False): verbose=self.options.solver_params.verbose, ) + # ---- projector ------------------------------------------------------- self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) @@ -8345,14 +7931,14 @@ def allocate(self, verbose=False): self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") - for u_prime, boundary_data, boundary_conditions in [ - (self._u_prime, self.options.boundary_data_u, self.options.boundary_conditions_u), - (self._ue_prime, self.options.boundary_data_ue, self.options.boundary_conditions_ue), + for u_prime, boundary_data in [ + (self._u_prime, self.options.boundary_data_u), + (self._ue_prime, self.options.boundary_data_ue), ]: if boundary_data is None: continue for (d, side), f_bc in boundary_data.items(): - if boundary_conditions.get((d, side)) == "dirichlet": + if (d, side) in self._dirichlet_faces: bc_pulled = lambda *etas, f=f_bc: self.domain.pull( [lambda x,y,z, f=f: f(x,y,z)[0], lambda x,y,z, f=f: f(x,y,z)[1], @@ -8361,8 +7947,8 @@ def allocate(self, verbose=False): _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], lambda *etas: bc_pulled(*etas)[1], lambda *etas: bc_pulled(*etas)[2]]) - for (d2, side2), kind2 in boundary_conditions.items(): - if kind2 == "dirichlet" and (d2, side2) != (d, side): + for (d2, side2) in self._dirichlet_faces: + if (d2, side2) != (d, side): apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) u_prime.vector += _vec @@ -8403,14 +7989,14 @@ def __call__(self, dt): self._u.vector = self.variables.u.spline.vector self._ue.vector = self.variables.ue.spline.vector - # --- rebuild system matrix if dt changed --- TODO change uzawa internals - if dt != self._dt: + # --- rebuild system matrix if dt changed --- + if dt != self._dt: # TODO change uzawa A11 block too self._dt = dt _A = BlockLinearOperator( self._block_domain_v0, self._block_codomain_v0, blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] ) - + _M = BlockLinearOperator( self._block_domain_M, self._block_domain_M, blocks=[[_A, self._B_v0.T], [self._B_v0, None]] @@ -8420,21 +8006,21 @@ def __call__(self, dt): # --- assemble RHS in unconstrained space, then zero boundary DOFs --- # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' # electron: F2 = rhs_ue - A22 * ue' - self._rhs_vec_u.vector = (self._rhs_u.vector - + self._M2.dot(self._u.vector) / dt - - self._A11.dot(self._u_prime.vector) - - self._M2.dot(self._u_prime.vector) / dt) + self._rhs_vec_u.vector = (self._rhs_u.vector # TODO boundary operator + + self._M2.dot(self._u.vector) / dt + - self._A11.dot(self._u_prime.vector) + - self._M2.dot(self._u_prime.vector) / dt) self._rhs_vec_ue.vector = (self._rhs_ue.vector - - self._A22.dot(self._ue_prime.vector)) + - self._A22.dot(self._ue_prime.vector)) - self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) - self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) + self._apply_essential_bc(self._rhs_vec_u.vector) + self._apply_essential_bc(self._rhs_vec_ue.vector) # --- build block RHS and solve --- _F = BlockVector(self._block_domain_v0, - blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) _RHS = BlockVector(self._block_domain_M, - blocks=[_F, self._block_codomain_B_v0.zeros()]) + blocks=[_F, self._block_codomain_B_v0.zeros()]) _sol = self._Minv.dot(_RHS) info = self._Minv.get_info() @@ -8452,38 +8038,5 @@ def __call__(self, dt): if self.options.solver_params.info and self._rank == 0: print(f"Status: {info['success']}, Iterations: {info['niter']}") print(f"Max diffs: {max_diffs}") - # --- assemble RHS in unconstrained space, then zero boundary DOFs --- - # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' - # electron: F2 = rhs_ue - A22 * ue' - self._rhs_vec_u.vector = (self._rhs_u.vector - + self._M2.dot(self._u.vector) / dt - - self._A11.dot(self._u_prime.vector) - - self._M2.dot(self._u_prime.vector) / dt) - self._rhs_vec_ue.vector = (self._rhs_ue.vector - - self._A22.dot(self._ue_prime.vector)) - - self._apply_boundary_conditions(self._rhs_vec_u.vector, self.options.boundary_conditions_u) - self._apply_boundary_conditions(self._rhs_vec_ue.vector, self.options.boundary_conditions_ue) - - # --- build block RHS and solve --- - _F = BlockVector(self._block_domain_v0, - blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) - _RHS = BlockVector(self._block_domain_M, - blocks=[_F, self._block_codomain_B_v0.zeros()]) - - _sol = self._Minv.dot(_RHS) - info = self._Minv.get_info() - - # --- reconstruct full solution: u = u_0 + u' --- - self._u.vector = _sol[0][0] + self._u_prime_v0.vector - self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector - self._phi.vector = _sol[1] - - # --- update FEEC variables --- - max_diffs = self.update_feec_variables( - u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector - ) - - if self.options.solver_params.info and self._rank == 0: print(f"Status: {info['success']}, Iterations: {info['niter']}") print(f"Max diffs: {max_diffs}") \ No newline at end of file From 61b29556e454d39c442f7ee87f97bc9c81439dbb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Tue, 14 Apr 2026 17:22:03 +0000 Subject: [PATCH 15/41] Rebased onto latest devel commit --- etc/jupyter/jupyter_notebook_config.d/ipyparallel.json | 7 +++++++ etc/jupyter/jupyter_notebook_config.d/jupyterlab.json | 7 +++++++ etc/jupyter/jupyter_server_config.d/ipyparallel.json | 7 +++++++ .../jupyter-lsp-jupyter-server.json | 7 +++++++ .../jupyter_server_config.d/jupyter_server_terminals.json | 7 +++++++ etc/jupyter/jupyter_server_config.d/jupyterlab.json | 7 +++++++ etc/jupyter/jupyter_server_config.d/notebook.json | 7 +++++++ etc/jupyter/jupyter_server_config.d/notebook_shim.json | 7 +++++++ etc/jupyter/nbconfig/notebook.d/widgetsnbextension.json | 5 +++++ etc/jupyter/nbconfig/tree.d/ipyparallel.json | 5 +++++ 10 files changed, 66 insertions(+) create mode 100644 etc/jupyter/jupyter_notebook_config.d/ipyparallel.json create mode 100644 etc/jupyter/jupyter_notebook_config.d/jupyterlab.json create mode 100644 etc/jupyter/jupyter_server_config.d/ipyparallel.json create mode 100644 etc/jupyter/jupyter_server_config.d/jupyter-lsp-jupyter-server.json create mode 100644 etc/jupyter/jupyter_server_config.d/jupyter_server_terminals.json create mode 100644 etc/jupyter/jupyter_server_config.d/jupyterlab.json create mode 100644 etc/jupyter/jupyter_server_config.d/notebook.json create mode 100644 etc/jupyter/jupyter_server_config.d/notebook_shim.json create mode 100644 etc/jupyter/nbconfig/notebook.d/widgetsnbextension.json create mode 100644 etc/jupyter/nbconfig/tree.d/ipyparallel.json diff --git a/etc/jupyter/jupyter_notebook_config.d/ipyparallel.json b/etc/jupyter/jupyter_notebook_config.d/ipyparallel.json new file mode 100644 index 000000000..4f1ba10cd --- /dev/null +++ b/etc/jupyter/jupyter_notebook_config.d/ipyparallel.json @@ -0,0 +1,7 @@ +{ + "NotebookApp": { + "nbserver_extensions": { + "ipyparallel": true + } + } +} diff --git a/etc/jupyter/jupyter_notebook_config.d/jupyterlab.json b/etc/jupyter/jupyter_notebook_config.d/jupyterlab.json new file mode 100644 index 000000000..5b5dcda3a --- /dev/null +++ b/etc/jupyter/jupyter_notebook_config.d/jupyterlab.json @@ -0,0 +1,7 @@ +{ + "NotebookApp": { + "nbserver_extensions": { + "jupyterlab": true + } + } +} diff --git a/etc/jupyter/jupyter_server_config.d/ipyparallel.json b/etc/jupyter/jupyter_server_config.d/ipyparallel.json new file mode 100644 index 000000000..cfc9c58a6 --- /dev/null +++ b/etc/jupyter/jupyter_server_config.d/ipyparallel.json @@ -0,0 +1,7 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "ipyparallel": true + } + } +} diff --git a/etc/jupyter/jupyter_server_config.d/jupyter-lsp-jupyter-server.json b/etc/jupyter/jupyter_server_config.d/jupyter-lsp-jupyter-server.json new file mode 100644 index 000000000..9e37d4eca --- /dev/null +++ b/etc/jupyter/jupyter_server_config.d/jupyter-lsp-jupyter-server.json @@ -0,0 +1,7 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "jupyter_lsp": true + } + } +} diff --git a/etc/jupyter/jupyter_server_config.d/jupyter_server_terminals.json b/etc/jupyter/jupyter_server_config.d/jupyter_server_terminals.json new file mode 100644 index 000000000..97c80c282 --- /dev/null +++ b/etc/jupyter/jupyter_server_config.d/jupyter_server_terminals.json @@ -0,0 +1,7 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "jupyter_server_terminals": true + } + } +} diff --git a/etc/jupyter/jupyter_server_config.d/jupyterlab.json b/etc/jupyter/jupyter_server_config.d/jupyterlab.json new file mode 100644 index 000000000..99cc0846e --- /dev/null +++ b/etc/jupyter/jupyter_server_config.d/jupyterlab.json @@ -0,0 +1,7 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "jupyterlab": true + } + } +} diff --git a/etc/jupyter/jupyter_server_config.d/notebook.json b/etc/jupyter/jupyter_server_config.d/notebook.json new file mode 100644 index 000000000..09113911a --- /dev/null +++ b/etc/jupyter/jupyter_server_config.d/notebook.json @@ -0,0 +1,7 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "notebook": true + } + } +} diff --git a/etc/jupyter/jupyter_server_config.d/notebook_shim.json b/etc/jupyter/jupyter_server_config.d/notebook_shim.json new file mode 100644 index 000000000..1e789c3d5 --- /dev/null +++ b/etc/jupyter/jupyter_server_config.d/notebook_shim.json @@ -0,0 +1,7 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "notebook_shim": true + } + } +} diff --git a/etc/jupyter/nbconfig/notebook.d/widgetsnbextension.json b/etc/jupyter/nbconfig/notebook.d/widgetsnbextension.json new file mode 100644 index 000000000..7a17570d6 --- /dev/null +++ b/etc/jupyter/nbconfig/notebook.d/widgetsnbextension.json @@ -0,0 +1,5 @@ +{ + "load_extensions": { + "jupyter-js-widgets/extension": true + } +} diff --git a/etc/jupyter/nbconfig/tree.d/ipyparallel.json b/etc/jupyter/nbconfig/tree.d/ipyparallel.json new file mode 100644 index 000000000..6d98e1861 --- /dev/null +++ b/etc/jupyter/nbconfig/tree.d/ipyparallel.json @@ -0,0 +1,5 @@ +{ + "load_extensions": { + "ipyparallel/main": true + } +} From 073d31afbccca163f4cd242f2c2d06586021fdd8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Sat, 18 Apr 2026 18:45:10 +0000 Subject: [PATCH 16/41] Lifting works in 1D now. --- src/struphy/feec/psydac_derham.py | 216 ++++++------ src/struphy/io/options.py | 65 ++-- src/struphy/models/variables.py | 7 +- src/struphy/propagators/propagators_fields.py | 309 +++++++++--------- struphy-parameter-files | 2 +- 5 files changed, 293 insertions(+), 306 deletions(-) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 09048cb7c..adcb1e81b 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -67,115 +67,117 @@ class DiscreteDerham: Parameters ---------- - Nel : list[int] - Number of elements in each direction. + V0 : TensorFemSpace + First space of the de Rham sequence : H1 space + V1 : VectorFemSpace + Second space of the de Rham sequence : Hcurl space + V2 : VectorFemSpace + Third space of the de Rham sequence : Hdiv space + V3 : TensorFemSpace + Fourth space of the de Rham sequence : L2 space - p : list[int] - Spline degree in each direction. + Notes + ----- + On construction, differential operators are created and attached to the + input spaces as convenience attributes: - spl_kind : list[bool] - Kind of spline in each direction (True=periodic, False=clamped). + - ``V0.grad`` and ``V0.diff`` + - ``V1.curl`` and ``V1.diff`` + - ``V2.div`` and ``V2.diff`` + """ - dirichlet_bc : list[list[bool]] - Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). + def __init__(self, V0: TensorFemSpace, V1: VectorFemSpace, V2: VectorFemSpace, V3: TensorFemSpace): + spaces = (V0, V1, V2, V3) + assert all(isinstance(space, (TensorFemSpace, VectorFemSpace)) for space in spaces) - nq_pr : list[int] - Number of Gauss-Legendre quadrature points in each direction for geometric projectors (default = p+1, leads to exact integration of degree 2p+1 polynomials). + self._V0 = V0 + self._V1 = V1 + self._V2 = V2 + self._V3 = V3 + self._spaces = spaces + self._dim = 3 - nquads : list[int] - Number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials). + D0 = Gradient3D(V0, V1) + D1 = Curl3D(V1, V2) + D2 = Divergence3D(V2, V3) - comm : mpi4py.MPI.Intracomm - MPI communicator (within a clone if domain cloning is used, otherwise MPI.COMM_WORLD) + V0.diff = V0.grad = D0 + V1.diff = V1.curl = D1 + V2.diff = V2.div = D2 - mpi_dims_mask: list of bool - True if the dimension is to be used in the domain decomposition (=default for each dimension). - If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed. + # -------------------------------------------------------------------------- + @property + def dim(self) -> int: + """Dimension of the physical and logical domains, which are assumed to be the same.""" + return self._dim - with_projectors : bool - Whether to add global commuting projectors to the diagram. + @property + def V0(self) -> TensorFemSpace: + """First space of the de Rham sequence : H1 space""" + return self._V0 - polar_ck : int - Smoothness at a polar singularity at eta_1=0 (default -1 : standard tensor product splines, OR 1 : C1 polar splines) + @property + def V1(self) -> VectorFemSpace: + """Second space of the de Rham sequence : Hcurl space""" + return self._V1 - local_projectors : bool - Whether to build the local commuting projectors based on quasi-inter-/histopolation. + @property + def V2(self) -> VectorFemSpace: + """Third space of the de Rham sequence : Hdiv space""" + return self._V2 - domain : struphy.geometry.base.Domain - Mapping from logical unit cube to physical domain (only needed in case of polar splines polar_ck=1). - """ + @property + def V3(self) -> TensorFemSpace: + """Fourth space of the de Rham sequence : L2 space""" + return self._V3 - def __init__( - self, - Nel: list | tuple, - p: list | tuple, - spl_kind: list | tuple, - *, - dirichlet_bc: list | tuple = None, - lifting: list | tuple = None, - nquads: list | tuple = None, - nq_pr: list | tuple = None, - comm=None, - mpi_dims_mask: list = None, - with_projectors: bool = True, - polar_ck: int = -1, - local_projectors: bool = False, - domain: Domain = None, - ): - # number of elements, spline degrees and kind of splines in each direction (periodic vs. clamped) - assert len(Nel) == 3 - assert len(p) == 3 - assert len(spl_kind) == 3 - - self._Nel = Nel - self._p = p - self._spl_kind = spl_kind - self._with_local_projectors = local_projectors + @property + def spaces(self) -> tuple[TensorFemSpace | VectorFemSpace, ...]: + """Spaces of the proper de Rham sequence (excluding Hvec).""" + return self._spaces - # boundary conditions at eta=0 and eta=1 in each direction (None for periodic, 'd' for homogeneous Dirichlet) - if dirichlet_bc is not None: - assert len(dirichlet_bc) == 3 - # make sure that boundary conditions are compatible with spline space - assert xp.all([bc == (False, False) for i, bc in enumerate(dirichlet_bc) if spl_kind[i]]) - - self._dirichlet_bc = dirichlet_bc - - # --- lifting: build constrained (v0) sub-complex --- - self._lifting = lifting - if lifting is not None: - assert len(lifting) == 3 - # lifting only makes sense on non-periodic axes - for d in range(3): - if spl_kind[d]: - assert lifting[d] == (False, False), \ - f"Axis {d} is periodic, lifting must be (False, False)" - - # v0 dirichlet_bc = dirichlet_bc OR lifting - if dirichlet_bc is not None: - v0_dirichlet_bc = tuple( - (d_l or l_l, d_r or l_r) - for (d_l, d_r), (l_l, l_r) in zip(dirichlet_bc, lifting) - ) - else: - v0_dirichlet_bc = lifting + @property + def derivatives_as_matrices(self): + """Differential operators of the De Rham sequence as LinearOperator objects.""" + return tuple(V.diff.linop for V in self.spaces[:-1]) - self._derham_v0 = Derham( - Nel, p, spl_kind, - dirichlet_bc=v0_dirichlet_bc, - nquads=nquads, - nq_pr=nq_pr, - comm=comm, - mpi_dims_mask=mpi_dims_mask, - with_projectors=with_projectors, - polar_ck=polar_ck, - local_projectors=self.with_local_projectors, - domain=domain, - ) - else: - self._derham_v0 = None + @property + def derivatives(self): + """Differential operators of the De Rham sequence as `DiffOperator` objects. + Those are objects with `domain` and `codomain` properties that are `FemSpace`, + they act on `FemField` (they take a `FemField` of their `domain` as input and return + a `FemField` of their `codomain`. + """ + return tuple(V.diff for V in self.spaces[:-1]) + + # -------------------------------------------------------------------------- + def projectors(self, *, kind="global", nquads=None) -> tuple[GlobalGeometricProjector, ...]: + """Projectors mapping callable functions of the physical coordinates to a + corresponding `FemField` object in the De Rham sequence. + + Parameters + ---------- + kind : str + Type of the projection : at the moment, only global is accepted and + returns geometric commuting projectors based on interpolation/histopolation + for the De Rham sequence (GlobalGeometricProjector objects). + + nquads : list(int) | tuple(int) + Number of quadrature points along each direction, to be used in Gauss + quadrature rule for computing the (approximated) degrees of freedom. + + Returns + ------- + P0, ..., Pn : callables + Projectors that can be called on any callable function that maps + from the physical space to R (scalar case) or R^d (vector case) and + returns a FemField belonging to the i-th space of the De Rham sequence + """ + + if not (kind == "global"): + raise NotImplementedError("only global projectors are available") - # default p: exact integration of degree 2p+1 polynomials if nquads is None: nquads = [degree + 1 for degree in self.V0.degree] elif isinstance(nquads, int): @@ -835,17 +837,27 @@ def __init__( # collect arguments for kernels self._args_derham = DerhamArguments( - xp.array(self.p), - self.Vh_fem["0"].knots[0], - self.Vh_fem["0"].knots[1], - self.Vh_fem["0"].knots[2], - xp.array(self.Vh["0"].starts), + xp.array(self.degree), + self.V0fem.knots[0], + self.V0fem.knots[1], + self.V0fem.knots[2], + xp.array(self.V0.starts), ) - @property - def derham_v0(self): - return self._derham_v0 - + if MPI.COMM_WORLD.Get_rank() == 0 and verbose: + logger.info("\nDERHAM:") + logger.info(f"{'number of elements:'.ljust(25)} {num_elements}") + logger.info(f"{'spline degrees:'.ljust(25)} {degree}") + logger.info(f"{'boundary conditions:'.ljust(25)} {bcs}") + logger.info(f"{'GL quad pts (L2):'.ljust(25)} {nquads}") + logger.info(f"{'GL quad pts (hist):'.ljust(25)} {nquads_proj}") + logger.info(f"{'MPI proc. per dir.:'.ljust(25)} {self.domain_decomposition.nprocs}") + logger.info(f"{'use polar splines:'.ljust(25)} {self.polar_splines}") + logger.info(f"{'domain on process 0:'.ljust(25)} {self.domain_array[0]}") + + # ----------------------------- + # Input arguments as properties + # ----------------------------- @property def grid(self) -> TensorProductGrid: """The FEEC grid.""" diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index b9f0b782a..41d8842a8 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -159,6 +159,17 @@ class LiteralOptions: "heat_flux_3", ] +class OptionsBase: + def to_dict(self) -> dict: + """Convert dataclass instance to dictionary.""" + return {field.name: getattr(self, field.name) for field in fields(type(self)) if field.init} + + @classmethod + def from_dict(cls, dct) -> "Any": + """Create dataclass instance from dictionary.""" + valid_fields = {field.name for field in fields(cls) if field.init} + return cls(**{key: value for key, value in dct.items() if key in valid_fields}) + @dataclass class Time(OptionsBase): @@ -255,13 +266,9 @@ class DerhamOptions(OptionsBase): Use ``None`` in a direction for periodic boundaries, or a tuple ``(left, right)`` with entries in ``{"free", "dirichlet"}`` for non-periodic boundaries. - lifting : tuple[tuple[bool]] - Whether to build a constrained (v0) sub-complex with additional clamping on each face. - Used for inhomogeneous Dirichlet BCs: the v0 complex clamps faces where - lifting is True, and the propagator builds a lift in the unconstrained space. - - nquads : tuple[int] - Number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials). + nquads : tuple[int, int, int] | None + Number of Gauss-Legendre quadrature points per direction for cell + integrals. If ``None``, backend defaults are used. nquads_proj : tuple[int, int, int] | None Number of Gauss-Legendre quadrature points per direction for geometric @@ -278,13 +285,15 @@ class DerhamOptions(OptionsBase): quasi-inter-/histopolation. """ - p: tuple = (1, 1, 1) - spl_kind: tuple = (True, True, True) - dirichlet_bc: tuple = ((False, False), (False, False), (False, False)) - lifting: tuple = ((False, False), (False, False), (False, False)) - nquads: tuple = None - nq_pr: tuple = None - polar_ck: LiteralOptions.PolarRegularity = -1 + degree: tuple[int, int, int] = (1, 1, 1) + bcs: tuple[ + None | tuple[NonTrivialBC, NonTrivialBC], + None | tuple[NonTrivialBC, NonTrivialBC], + None | tuple[NonTrivialBC, NonTrivialBC], + ] = (None, None, None) + nquads: tuple[int, int, int] | None = None + nquads_proj: tuple[int, int, int] | None = None + polar_splines: bool = False local_projectors: bool = False def __post_init__(self): @@ -307,33 +316,7 @@ def __repr_no_defaults__(self): @property def is_default(self): return all_class_params_are_default(self) - - def to_dict(self) -> dict: - dct = { - "p": self.p, - "spl_kind": self.spl_kind, - "dirichlet_bc": self.dirichlet_bc, - "lifting": self.lifting, - "nquads": self.nquads, - "nq_pr": self.nq_pr, - "polar_ck": self.polar_ck, - "local_projectors": self.local_projectors, - } - return dct - - @classmethod - def from_dict(cls, dct) -> "DerhamOptions": - return cls( - p=dct["p"], - spl_kind=dct["spl_kind"], - dirichlet_bc=dct["dirichlet_bc"], - lifting=dct.get("lifting", ((False, False), (False, False), (False, False))), - nquads=dct["nquads"], - nq_pr=dct["nq_pr"], - polar_ck=dct["polar_ck"], - local_projectors=dct["local_projectors"], - ) - + @dataclass class FieldsBackground(OptionsBase): diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 6b722b081..ded198382 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -389,9 +389,10 @@ def allocate( # other helper objects for the lifting of boundary conditions self._spline_0 = self.spline_lift.copy() - self.spline_0.vector[:] = self.spline_lift.vector[:] + self.spline_lift.vector.copy(out=self.spline_0.vector) self._boundary_spline = self.spline_lift.copy() - self._boundary_op = BoundaryOperator(self.spline_lift.space, self.space, derham.dirichlet_bc) + + self._boundary_op = BoundaryOperator(self.spline_lift.space, self.space, derham.dirichlet_bc) # TODO different domain and codomain self.compute_boundary_spline() @@ -405,7 +406,7 @@ def compute_boundary_spline(self, spline_lift: SplineFunction | None = None): # set new boundary spline diff_vec = spline_lift.vector - self.spline_0.vector - self.boundary_spline.vector[:] = diff_vec[:] + diff_vec.copy(out=self.boundary_spline.vector) class PICVariable(Variable): diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 23c94842c..356d48245 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -7653,10 +7653,6 @@ class TwoFluidQuasiNeutralFull(Propagator): ### State variables (ion velocity u, electron velocity ue, pressure phi) # ========================================================================= - # ========================================================================= - ### State variables (ion velocity u, electron velocity ue, pressure phi) - # ========================================================================= - class Variables(): def __init__(self) -> None: self._u: FEECVariable | None = None @@ -7664,7 +7660,6 @@ def __init__(self) -> None: self._phi: FEECVariable | None = None @property - def u(self) -> FEECVariable | None: def u(self) -> FEECVariable | None: return self._u @@ -7675,7 +7670,6 @@ def u(self, new): self._u = new @property - def ue(self) -> FEECVariable | None: def ue(self) -> FEECVariable | None: return self._ue @@ -7686,7 +7680,6 @@ def ue(self, new): self._ue = new @property - def phi(self) -> FEECVariable | None: def phi(self) -> FEECVariable | None: return self._phi @@ -7710,9 +7703,6 @@ class Options(): nu_e: float | None = None eps_norm: float | None = None - boundary_data_u: dict[tuple[int, int], Callable] | None = None - boundary_data_ue: dict[tuple[int, int], Callable] | None = None - source_u: Callable | None = None source_ue: Callable | None = None @@ -7765,45 +7755,6 @@ def options(self, new): print(f" {k}: {v}") self._options = new - # ========================================================================= - ### Boundary condition helpers - # ========================================================================= - - def _get_dirichlet_faces(self): - """Infer which faces have Dirichlet BCs by comparing derham and derham_v0. - - A face is Dirichlet if it is unclamped in derham but clamped in derham_v0 - (i.e. lifting is True there). - """ - faces = [] - derham = self.derham - derham_v0 = derham.derham_v0 - - if derham_v0 is None: - return faces - - bc = derham.dirichlet_bc - bc_v0 = derham_v0.dirichlet_bc - - for d in range(3): - if derham.spl_kind[d]: - continue # periodic axis, no Dirichlet - for s, side in enumerate((-1, 1)): - # clamped in v0 but not in derham => this is a lifted (inhom Dirichlet) face - unclamped = not bc[d][s] - clamped_v0 = bc_v0[d][s] if bc_v0 is not None else False - if unclamped and clamped_v0: - faces.append((d, side)) - # clamped in both => homogeneous Dirichlet, also need to zero DOFs - elif bc[d][s] and clamped_v0: - faces.append((d, side)) - return faces - - def _apply_essential_bc(self, vec): - """Zero out Dirichlet DOFs, inferred from derham vs derham_v0.""" - for (d, side) in self._dirichlet_faces: - apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) - # ========================================================================= ### Allocate # ========================================================================= @@ -7814,64 +7765,92 @@ def allocate(self, verbose=False): self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 self._dt = None - # ---- v0 de Rham complex (from derham.derham_v0) ---------------------- + # ---- lifting (derham_lift is unconstrained, self.derham is constrained) --- + + _u_var = self.variables.u + _ue_var = self.variables.ue + + self._has_lifting_u = _u_var.derham_lift is not None + self._has_lifting_ue = _ue_var.derham_lift is not None + + # unconstrained de Rham (for RHS assembly): use derham_lift if available, + # otherwise fall back to self.derham (no lifting case) + self._derham_lift_u = _u_var.derham_lift if self._has_lifting_u else self.derham + self._derham_lift_ue = _ue_var.derham_lift if self._has_lifting_ue else self.derham - assert self.derham.derham_v0 is not None, \ - "derham must be constructed with lifting to use this propagator" + # boundary splines (u', ue') in unconstrained space — zero vectors if no lifting + self._boundary_spline_u = (_u_var.boundary_spline.vector if self._has_lifting_u + else self._derham_lift_u.coeff_spaces["2"].zeros()) + self._boundary_spline_ue = (_ue_var.boundary_spline.vector if self._has_lifting_ue + else self._derham_lift_ue.coeff_spaces["2"].zeros()) - self._derham_v0 = self.derham.derham_v0 + # ---- unconstrained mass/basis operators (for RHS assembly) ----------- - self._mass_ops_v0 = WeightedMassOperators( - self._derham_v0, self.domain, + self._mass_ops_lift_u = WeightedMassOperators( + self._derham_lift_u, self.domain, verbose=self.options.solver_params.verbose, eq_mhd=self.mass_ops.weights["eq_mhd"], ) - self._basis_ops_v0 = BasisProjectionOperators( - self._derham_v0, self.domain, + self._mass_ops_lift_ue = WeightedMassOperators( + self._derham_lift_ue, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.mass_ops.weights["eq_mhd"], + ) + self._basis_ops_lift_u = BasisProjectionOperators( + self._derham_lift_u, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.basis_ops.weights["eq_mhd"], + ) + self._basis_ops_lift_ue = BasisProjectionOperators( + self._derham_lift_ue, self.domain, verbose=self.options.solver_params.verbose, eq_mhd=self.basis_ops.weights["eq_mhd"], ) - # ---- Dirichlet faces (inferred from derham vs derham_v0) ------------- + self._M2_u = self._mass_ops_lift_u.M2 + self._M2B_u = - self._mass_ops_lift_u.M2B + self._div_u = self._derham_lift_u.div + self._curl_u = self._derham_lift_u.curl + self._S21_u = self._basis_ops_lift_u.S21 + + self._lapl_u = (self._div_u.T @ self._mass_ops_lift_u.M3 @ self._div_u + + self._S21_u.T @ self._curl_u.T @ self._M2_u @ self._curl_u @ self._S21_u) - self._dirichlet_faces = self._get_dirichlet_faces() + self._A11 = - self._M2B_u / self.options.eps_norm + self.options.nu * self._lapl_u - # ---- unconstrained operators (for RHS assembly) ---------------------- + self._M2_ue = self._mass_ops_lift_ue.M2 + self._M2B_ue = - self._mass_ops_lift_ue.M2B + self._div_ue = self._derham_lift_ue.div + self._curl_ue = self._derham_lift_ue.curl + self._S21_ue = self._basis_ops_lift_ue.S21 - self._M2 = self.mass_ops.M2 - self._M2B = - self.mass_ops.M2B - self._div = self.derham.div - self._curl = self.derham.curl - self._S21 = self.basis_ops.S21 + self._lapl_ue = (self._div_ue.T @ self._mass_ops_lift_ue.M3 @ self._div_ue + + self._S21_ue.T @ self._curl_ue.T @ self._M2_ue @ self._curl_ue @ self._S21_ue) - self._lapl = (self._div.T @ self.mass_ops.M3 @ self._div - + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - - self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl - self._A22 = (- self.options.stab_sigma * IdentityOperator(self.derham.Vh["2"]) - + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl) + self._A22 = (- self.options.stab_sigma * IdentityOperator(self._derham_lift_ue.coeff_spaces["2"]) + + self._M2B_ue / self.options.eps_norm + self.options.nu_e * self._lapl_ue) - # ---- constrained operators (for system matrix) ----------------------- + # ---- constrained operators (for system matrix, built from self.derham) --- - self._M2_v0 = self._mass_ops_v0.M2 - self._M3_v0 = self._mass_ops_v0.M3 - self._M2B_v0 = - self._mass_ops_v0.M2B - self._div_v0 = self._derham_v0.div - self._curl_v0 = self._derham_v0.curl - self._S21_v0 = self._basis_ops_v0.S21 + self._M2_v0 = self.mass_ops.M2 + self._M3_v0 = self.mass_ops.M3 + self._M2B_v0 = - self.mass_ops.M2B + self._div_v0 = self.derham.div + self._curl_v0 = self.derham.curl + self._S21_v0 = self.basis_ops.S21 self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) - + self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self._derham_v0.Vh["2"]) + self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self.derham.coeff_spaces["2"]) + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) # ---- block saddle-point system ---------------------------------------- - self._block_domain_v0 = BlockVectorSpace(self._derham_v0.Vh["2"], self._derham_v0.Vh["2"]) + self._block_domain_v0 = BlockVectorSpace(self.derham.coeff_spaces["2"], self.derham.coeff_spaces["2"]) self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self._derham_v0.Vh["3"] + self._block_codomain_B_v0 = self.derham.coeff_spaces["3"] self._B1_v0 = - self._M3_v0 @ self._div_v0 self._B2_v0 = self._M3_v0 @ self._div_v0 @@ -7902,7 +7881,6 @@ def allocate(self, verbose=False): recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, - maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, ) else: @@ -7911,73 +7889,56 @@ def allocate(self, verbose=False): recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, - maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, ) + # ---- source terms projected onto unconstrained space ----------------- - # ---- projector ------------------------------------------------------- + self._projector_u = L2Projector(space_id="Hdiv", mass_ops=self._mass_ops_lift_u) + self._projector_ue = L2Projector(space_id="Hdiv", mass_ops=self._mass_ops_lift_ue) - self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + self._rhs_u = self._derham_lift_u.create_spline_function("rhs_u", space_id="Hdiv") + self._rhs_ue = self._derham_lift_ue.create_spline_function("rhs_ue", space_id="Hdiv") - # ---- solution spline functions (unconstrained) ----------------------- + for rhs, source, derham_lift in [ + (self._rhs_u, self.options.source_u, self._derham_lift_u), + (self._rhs_ue, self.options.source_ue, self._derham_lift_ue), + ]: + if source is not None: + fun_vec = [lambda x, y, z, f=source, c=c: f(x, y, z)[c] for c in range(3)] + fun = [ + TransformedPformComponent( + fun_vec, + "physical", + "2", + comp=comp, + domain=self.domain, + ) + for comp in range(3) + ] + rhs.vector = derham_lift.projectors["2"](fun) + + # ---- solution splines (constrained) and u in unconstrained space ----- self._u = self.derham.create_spline_function("u", space_id="Hdiv") self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") self._phi = self.derham.create_spline_function("phi", space_id="L2") - # ---- BC lifts (unconstrained) ---------------------------------------- - - self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") - self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") - - for u_prime, boundary_data in [ - (self._u_prime, self.options.boundary_data_u), - (self._ue_prime, self.options.boundary_data_ue), - ]: - if boundary_data is None: - continue - for (d, side), f_bc in boundary_data.items(): - if (d, side) in self._dirichlet_faces: - bc_pulled = lambda *etas, f=f_bc: self.domain.pull( - [lambda x,y,z, f=f: f(x,y,z)[0], - lambda x,y,z, f=f: f(x,y,z)[1], - lambda x,y,z, f=f: f(x,y,z)[2]], - *etas, kind="2") - _vec = self._projector([lambda *etas: bc_pulled(*etas)[0], - lambda *etas: bc_pulled(*etas)[1], - lambda *etas: bc_pulled(*etas)[2]]) - for (d2, side2) in self._dirichlet_faces: - if (d2, side2) != (d, side): - apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) - u_prime.vector += _vec - - self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") - self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") - - self._u_prime_v0.vector = self._u_prime.vector - self._ue_prime_v0.vector = self._ue_prime.vector - - # ---- projected source terms (unconstrained) -------------------------- - - self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") - self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") - - for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: - if source is not None: - src_pulled = lambda *etas, f=source: self.domain.pull( - [lambda x,y,z, f=f: f(x,y,z)[0], - lambda x,y,z, f=f: f(x,y,z)[1], - lambda x,y,z, f=f: f(x,y,z)[2]], - *etas, kind="2") - rhs.vector = self._projector.get_dofs([lambda *etas: src_pulled(*etas)[0], - lambda *etas: src_pulled(*etas)[1], - lambda *etas: src_pulled(*etas)[2]]) + # u/ue embedded in unconstrained space for M2 application in RHS + self._u_lift = self._derham_lift_u.create_spline_function("u_lift", space_id="Hdiv") + self._ue_lift = self._derham_lift_ue.create_spline_function("ue_lift", space_id="Hdiv") + + # pre-allocated RHS vectors (constrained, after boundary_ops projection) + self._rhs_vec_u = self.derham.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_ue = self.derham.create_spline_function("rhs_vec_ue", space_id="Hdiv") + + # boundary splines in constrained space for add/subtract around solve + self._boundary_spline_u_v0 = self.derham.create_spline_function("boundary_spline_u_v0", space_id="Hdiv") + self._boundary_spline_ue_v0 = self.derham.create_spline_function("boundary_spline_ue_v0", space_id="Hdiv") - # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- + self._rhs_full_u = self.derham.create_spline_function("rhs_full_u", space_id="Hdiv") + self._rhs_full_ue = self.derham.create_spline_function("rhs_full_ue", space_id="Hdiv") - self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") - self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") # ========================================================================= ### Time step @@ -7985,58 +7946,88 @@ def allocate(self, verbose=False): def __call__(self, dt): - # --- copy current state --- + # --- copy current state (full solution = u_0 + u') --- self._u.vector = self.variables.u.spline.vector self._ue.vector = self.variables.ue.spline.vector + # --- store boundary spline in constrained space for reconstruction --- + self._boundary_spline_u_v0.vector = self._boundary_spline_u + self._boundary_spline_ue_v0.vector = self._boundary_spline_ue + + # --- strip lifting to get u_0 in constrained space --- + if self._has_lifting_u: + self._u.vector = self._u.vector - self._boundary_spline_u_v0.vector + if self._has_lifting_ue: + self._ue.vector = self._ue.vector - self._boundary_spline_ue_v0.vector + + # --- embed u_0 into unconstrained space for M2 application --- + # u_0 satisfies homogeneous Dirichlet so boundary DOFs are zero in both spaces + self._u_lift.vector = self._u.vector + self._ue_lift.vector = self._ue.vector + # --- rebuild system matrix if dt changed --- - if dt != self._dt: # TODO change uzawa A11 block too + if dt != self._dt: self._dt = dt - _A = BlockLinearOperator( + _A = BlockLinearOperator( # TODO avoid allocating self._block_domain_v0, self._block_codomain_v0, blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] ) - _M = BlockLinearOperator( self._block_domain_M, self._block_domain_M, blocks=[[_A, self._B_v0.T], [self._B_v0, None]] ) self._Minv.linop = _M - # --- assemble RHS in unconstrained space, then zero boundary DOFs --- - # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' - # electron: F2 = rhs_ue - A22 * ue' - self._rhs_vec_u.vector = (self._rhs_u.vector # TODO boundary operator - + self._M2.dot(self._u.vector) / dt - - self._A11.dot(self._u_prime.vector) - - self._M2.dot(self._u_prime.vector) / dt) - self._rhs_vec_ue.vector = (self._rhs_ue.vector - - self._A22.dot(self._ue_prime.vector)) + # --- assemble RHS fully in unconstrained space, then project to constrained --- + # ion: F1 = bc_op @ (rhs_u + M2_u/dt * u_0 - (A11 + M2_u/dt) * u') + # electron: F2 = bc_op @ (rhs_ue - A22 * ue') + + rhs_u_full = (self._M2_u.dot(self._rhs_u.vector) + + self._M2_u.dot(self._u_lift.vector) / dt + - self._A11.dot(self._boundary_spline_u) + - self._M2_u.dot(self._boundary_spline_u) / dt) + + rhs_ue_full = (self._M2_ue.dot(self._rhs_ue.vector) + - self._A22.dot(self._boundary_spline_ue)) + + self._rhs_full_u.vector = rhs_u_full + self._rhs_full_ue.vector = rhs_ue_full + + self._rhs_vec_u.vector = self.derham.boundary_ops["2"].dot(self._rhs_full_u.vector) + self._rhs_vec_ue.vector = self.derham.boundary_ops["2"].dot(self._rhs_full_ue.vector) - self._apply_essential_bc(self._rhs_vec_u.vector) - self._apply_essential_bc(self._rhs_vec_ue.vector) + tmp1 = self.derham.create_spline_function("tmp1", space_id="L2") + tmp2 = self.derham.create_spline_function("tmp2", space_id="L2") + + tmp1.vector = self._div_u.dot(self._boundary_spline_u) + tmp2.vector = self._div_ue.dot(self._boundary_spline_ue) + + phi_rhs = (self.mass_ops.M3.dot(tmp1.vector) + - self.mass_ops.M3.dot(tmp2.vector)) # --- build block RHS and solve --- _F = BlockVector(self._block_domain_v0, - blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) _RHS = BlockVector(self._block_domain_M, - blocks=[_F, self._block_codomain_B_v0.zeros()]) - + blocks=[_F, phi_rhs]) _sol = self._Minv.dot(_RHS) info = self._Minv.get_info() # --- reconstruct full solution: u = u_0 + u' --- - self._u.vector = _sol[0][0] + self._u_prime_v0.vector - self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector + self._u.vector = _sol[0][0] + self._ue.vector = _sol[0][1] self._phi.vector = _sol[1] + if self._has_lifting_u: + self._u.vector = self._u.vector + self._boundary_spline_u_v0.vector + if self._has_lifting_ue: + self._ue.vector = self._ue.vector + self._boundary_spline_ue_v0.vector + # --- update FEEC variables --- max_diffs = self.update_feec_variables( u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector ) if self.options.solver_params.info and self._rank == 0: - print(f"Status: {info['success']}, Iterations: {info['niter']}") - print(f"Max diffs: {max_diffs}") print(f"Status: {info['success']}, Iterations: {info['niter']}") print(f"Max diffs: {max_diffs}") \ No newline at end of file diff --git a/struphy-parameter-files b/struphy-parameter-files index 5143ca521..b74d5c648 160000 --- a/struphy-parameter-files +++ b/struphy-parameter-files @@ -1 +1 @@ -Subproject commit 5143ca52173bb00766c5032d2b9e652ecabd885e +Subproject commit b74d5c64811f46441989c81e5bbd59019daf11b3 From 06a3fe0fdf9d3295d8a0eb2814195cf0c799b03f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Sun, 19 Apr 2026 17:36:08 +0000 Subject: [PATCH 17/41] Added support for multiple nonzero vector components on trace. Lifting test case in 2D satisfies BCs, but blows up in the interior. --- feectools | 2 +- src/struphy/models/variables.py | 84 +++++++++++-------- src/struphy/propagators/propagators_fields.py | 14 ++-- struphy-parameter-files | 2 +- 4 files changed, 56 insertions(+), 46 deletions(-) diff --git a/feectools b/feectools index d2a48ef19..ae6859fcb 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit d2a48ef19d31ae22ba238369cd5a53c679c6f1d8 +Subproject commit ae6859fcb16c765f7bb7cabb82eb6268d1967014 diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index ded198382..30a47a4e9 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -324,23 +324,36 @@ def allocate( f"Lifting of boundary conditions can only be applied if at least one homogenous Dirichlet boundary condition is present in the Derham object, but here {derham.bcs = }" ) - # create another Derham object with the same options but with homogenous Dirichlet BCs replaced by free BCs, to be used for the lifting function + # normalise to list + lifting_list = self.lifting_function if isinstance(self.lifting_function, list) else [self.lifting_function] + + # validation + if self.space in {"H1", "L2"}: + if len(lifting_list) > 1: + raise ValueError("H1/L2 lifting only accepts a single Perturbation, not a list.") + elif self.space in {"Hcurl", "Hdiv", "H1vec"}: + if len(lifting_list) > 3: + raise ValueError("Hdiv/Hcurl/H1vec lifting accepts at most 3 Perturbations (one per component).") + comps = [ptb.comp for ptb in lifting_list] + if len(comps) != len(set(comps)): + raise ValueError(f"Each component may only appear once in the lifting list, got {comps}.") + + # create unconstrained Derham dct = derham.to_dict() bcs_lift = list(dct["options"]["bcs"]) for i, bc in enumerate(bcs_lift): if bc is not None: - bcn = list(bc) # convert tuple to list to allow modification + bcn = list(bc) if bcn[0] == "dirichlet": bcn[0] = "free" if bcn[1] == "dirichlet": bcn[1] = "free" - bcn = tuple(bcn) # convert back to tuple - bcs_lift[i] = bcn - dct["options"]["bcs"] = tuple(bcs_lift) # convert list back to tuple + bcs_lift[i] = tuple(bcn) + dct["options"]["bcs"] = tuple(bcs_lift) self._derham_lift = Derham.from_dict(dct, comm=derham.comm) - # spline function for the lifting function + # spline function for the lifting self._spline_lift = self.derham_lift.create_spline_function( name=self.__name__ + "_lift" if self.__name__ is not None else None, space_id=self.space, @@ -349,50 +362,47 @@ def allocate( verbose=verbose, ) - # project lifting function to spline space - ptb = self.lifting_function - - if self.space in { - "H1", - "L2", - }: # TODO: this is a copy-paste from SplineFunction.initialize_coeffs(), to be unified + # project each perturbation and accumulate into spline_lift + if self.space in {"H1", "L2"}: + ptb = lifting_list[0] if ptb.given_in_basis is None: ptb.given_in_basis = "0" - fun = TransformedPformComponent( ptb, ptb.given_in_basis, derham.space_to_form[self.space], domain=domain, ) + self.spline_lift.vector += self.derham_lift.projectors[derham.space_to_form[self.space]](fun) + elif self.space in {"Hcurl", "Hdiv", "H1vec"}: fun_vec = [None] * 3 - fun_vec[ptb.comp] = ptb - - if ptb.given_in_basis is None: - ptb.given_in_basis = "v" - # pullback callable for each component - fun = [] - for comp in range(3): - fun += [ - TransformedPformComponent( - fun_vec, - ptb.given_in_basis, - derham.space_to_form[self.space], - comp=comp, - domain=domain, - ), - ] - - # peform projection - self.spline_lift.vector += self.derham_lift.projectors[derham.space_to_form[self.space]](fun) - - # other helper objects for the lifting of boundary conditions + for ptb in lifting_list: + if fun_vec[ptb.comp] is not None: + raise ValueError(f"Component {ptb.comp} assigned more than once in lifting list.") + fun_vec[ptb.comp] = ptb + if ptb.given_in_basis is None: + ptb.given_in_basis = "v" + + fun = [ + TransformedPformComponent( + fun_vec, + fun_vec[comp].given_in_basis if fun_vec[comp] is not None else lifting_list[0].given_in_basis, + derham.space_to_form[self.space], + comp=comp, + domain=domain, + ) + for comp in range(3) + ] + self.spline_lift.vector += self.derham_lift.projectors[derham.space_to_form[self.space]](fun) + + + # other helper objects self._spline_0 = self.spline_lift.copy() self.spline_lift.vector.copy(out=self.spline_0.vector) self._boundary_spline = self.spline_lift.copy() - - self._boundary_op = BoundaryOperator(self.spline_lift.space, self.space, derham.dirichlet_bc) # TODO different domain and codomain + + self._boundary_op = BoundaryOperator(self.spline_lift.space, self.space, derham.dirichlet_bc) self.compute_boundary_spline() diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 356d48245..958abd4c4 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -7827,7 +7827,7 @@ def allocate(self, verbose=False): self._lapl_ue = (self._div_ue.T @ self._mass_ops_lift_ue.M3 @ self._div_ue + self._S21_ue.T @ self._curl_ue.T @ self._M2_ue @ self._curl_ue @ self._S21_ue) - self._A22 = (- self.options.stab_sigma * IdentityOperator(self._derham_lift_ue.coeff_spaces["2"]) + self._A22 = (self.options.stab_sigma * IdentityOperator(self._derham_lift_ue.coeff_spaces["2"]) + self._M2B_ue / self.options.eps_norm + self.options.nu_e * self._lapl_ue) # ---- constrained operators (for system matrix, built from self.derham) --- @@ -7842,8 +7842,8 @@ def allocate(self, verbose=False): self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) - self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = (- self.options.stab_sigma * IdentityOperator(self.derham.coeff_spaces["2"]) + self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 + self._A22_v0 = (self.options.stab_sigma * IdentityOperator(self.derham.coeff_spaces["2"]) + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) # ---- block saddle-point system ---------------------------------------- @@ -7944,7 +7944,7 @@ def allocate(self, verbose=False): ### Time step # ========================================================================= - def __call__(self, dt): + def __call__(self, dt): # TODO this is still a complete mess, clean up after 2D lifting also works # --- copy current state (full solution = u_0 + u') --- self._u.vector = self.variables.u.spline.vector @@ -7993,13 +7993,13 @@ def __call__(self, dt): self._rhs_full_u.vector = rhs_u_full self._rhs_full_ue.vector = rhs_ue_full - self._rhs_vec_u.vector = self.derham.boundary_ops["2"].dot(self._rhs_full_u.vector) + self._rhs_vec_u.vector = self.derham.boundary_ops["2"].dot(self._rhs_full_u.vector) # TODO implement or change boundary operator to also change the space self._rhs_vec_ue.vector = self.derham.boundary_ops["2"].dot(self._rhs_full_ue.vector) tmp1 = self.derham.create_spline_function("tmp1", space_id="L2") tmp2 = self.derham.create_spline_function("tmp2", space_id="L2") - tmp1.vector = self._div_u.dot(self._boundary_spline_u) + tmp1.vector = self._div_u.dot(self._boundary_spline_u) # TODO implement identity operator between L^2 of different de rham spaces tmp2.vector = self._div_ue.dot(self._boundary_spline_ue) phi_rhs = (self.mass_ops.M3.dot(tmp1.vector) @@ -8019,7 +8019,7 @@ def __call__(self, dt): self._phi.vector = _sol[1] if self._has_lifting_u: - self._u.vector = self._u.vector + self._boundary_spline_u_v0.vector + self._u.vector = self._u.vector + self._boundary_spline_u_v0.vector # TODO store an additional field in feecvariable that has the complete solution if self._has_lifting_ue: self._ue.vector = self._ue.vector + self._boundary_spline_ue_v0.vector diff --git a/struphy-parameter-files b/struphy-parameter-files index b74d5c648..1a67f3ed8 160000 --- a/struphy-parameter-files +++ b/struphy-parameter-files @@ -1 +1 @@ -Subproject commit b74d5c64811f46441989c81e5bbd59019daf11b3 +Subproject commit 1a67f3ed8c7f5237ff8c6bc50f80df4e8ccc3b43 From 93d8f71977c6d2d2f597165c9b188628c208586a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Sun, 19 Apr 2026 19:05:39 +0000 Subject: [PATCH 18/41] Add codomain kwarg to BoundaryOperator and spline_full to FEECVariable for lifting reconstruction. --- src/struphy/feec/linear_operators.py | 30 +- src/struphy/models/variables.py | 26 ++ src/struphy/propagators/base.py | 7 + src/struphy/propagators/propagators_fields.py | 260 +++++++----------- 4 files changed, 151 insertions(+), 172 deletions(-) diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 940bcd4d5..c3a132e5a 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -304,21 +304,31 @@ class BoundaryOperator(LinOpWithTransp): Parameters ---------- vector_space : feectools.linalg.basic.VectorSpace - The vector space associated to the operator. + The vector space of the domain (input). space_id : str Symbolic space ID of vector_space (H1, Hcurl, Hdiv, L2 or H1vec). dirichlet_bc : tuple[tuple[bool]] Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). + + codomain : feectools.linalg.basic.VectorSpace, optional + The vector space of the codomain (output). If given, the operator maps between two different spaces + (e.g. unconstrained to constrained). If None, domain and codomain are the same. """ - def __init__(self, vector_space, space_id, dirichlet_bc): + def __init__(self, vector_space, space_id, dirichlet_bc, codomain=None): assert isinstance(vector_space, VectorSpace) assert isinstance(space_id, str) self._domain = vector_space - self._codomain = vector_space + if codomain is not None: + assert isinstance(codomain, VectorSpace) + self._codomain = codomain + self._cross_space = True + else: + self._codomain = vector_space + self._cross_space = False self._dtype = vector_space.dtype self._space_id = space_id @@ -491,20 +501,24 @@ def dot(self, v, out=None): assert isinstance(v, Vector) assert v.space == self._domain - if out is None: - out = v.copy() - else: + if out is not None: assert isinstance(out, Vector) assert out.space == self._codomain v.copy(out=out) + elif self._cross_space: + out = self._codomain.zeros() + v.copy(out=out) + else: + out = v.copy() - # apply boundary conditions to output vector apply_essential_bc_to_array(self._space_id, out, self.bc) return out + def transpose(self, conjugate=False): """ Returns the transposed operator. """ - return BoundaryOperator(self._domain, self._space_id, self.bc) + return BoundaryOperator(self._codomain, self._space_id, self.bc, codomain=self._domain) + diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 30a47a4e9..33204c2de 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -260,6 +260,13 @@ def boundary_spline(self) -> SplineFunction | None: if not hasattr(self, "_boundary_spline"): self._boundary_spline = None return self._boundary_spline + + @property + def spline_full(self) -> SplineFunction | None: + """Full solution spline (lifting + zero-BC part) in the unconstrained space. Only allocated if lifting_function is not None.""" + if not hasattr(self, "_spline_full"): + self._spline_full = None + return self._spline_full @property def boundary_op(self) -> BoundaryOperator | None: @@ -267,6 +274,14 @@ def boundary_op(self) -> BoundaryOperator | None: if not hasattr(self, "_boundary_op"): self._boundary_op = None return self._boundary_op + + @property + def boundary_op_lift(self) -> BoundaryOperator | None: + """Boundary operator mapping from the unconstrained (lifted) space to the constrained space. + Only allocated if lifting_function is not None.""" + if not hasattr(self, "_boundary_op_lift"): + self._boundary_op_lift = None + return self._boundary_op_lift @property def derham_lift(self) -> Derham | None: @@ -362,6 +377,15 @@ def allocate( verbose=verbose, ) + # spline function for unconstrained solution + self._spline_full = self.derham_lift.create_spline_function( + name=self.__name__ + "_full" if self.__name__ is not None else None, + space_id=self.space, + domain=domain, + equil=equil, + verbose=verbose, + ) + # project each perturbation and accumulate into spline_lift if self.space in {"H1", "L2"}: ptb = lifting_list[0] @@ -404,6 +428,8 @@ def allocate( self._boundary_op = BoundaryOperator(self.spline_lift.space, self.space, derham.dirichlet_bc) + self._boundary_op_lift = BoundaryOperator(self.spline_lift.space, self.space, derham.dirichlet_bc, codomain=self._spline.space) + self.compute_boundary_spline() def compute_boundary_spline(self, spline_lift: SplineFunction | None = None): diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index f48dc1ce7..3aa04a46b 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -116,6 +116,12 @@ def update_feec_variables(self, **new_coeffs): old = old_var.spline.vector assert new.space == old.space + # update full solution spline (lifting + zero-BC part) if present + if old_var.spline_full is not None: + new.copy(out=old_var.spline_full.vector) + if old_var.boundary_spline is not None: + old_var.spline_full.vector += old_var.boundary_spline.vector + # calculate maximum of difference abs(new - old) diffs[var] = xp.max(xp.abs(new.toarray() - old.toarray())) @@ -125,6 +131,7 @@ def update_feec_variables(self, **new_coeffs): # important: sync processes! old.update_ghost_regions() + return diffs @property diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 958abd4c4..bca76d587 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -7723,7 +7723,7 @@ def __post_init__(self): raise ValueError(f"nu must be non-negative, got {self.nu}") if self.nu_e < 0: raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") - if self.eps_norm <= 0: + if self.eps_norm <= 0: # TODO get base epsilon from ion species if undefined raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") # --- warn if no source terms --- @@ -7766,23 +7766,52 @@ def allocate(self, verbose=False): self._dt = None # ---- lifting (derham_lift is unconstrained, self.derham is constrained) --- + self._has_lifting_u = self.variables.u.derham_lift is not None + self._has_lifting_ue = self.variables.ue.derham_lift is not None - _u_var = self.variables.u - _ue_var = self.variables.ue + self._derham_lift_u = self.variables.u.derham_lift if self._has_lifting_u else self.derham + self._derham_lift_ue = self.variables.ue.derham_lift if self._has_lifting_ue else self.derham - self._has_lifting_u = _u_var.derham_lift is not None - self._has_lifting_ue = _ue_var.derham_lift is not None - - # unconstrained de Rham (for RHS assembly): use derham_lift if available, - # otherwise fall back to self.derham (no lifting case) - self._derham_lift_u = _u_var.derham_lift if self._has_lifting_u else self.derham - self._derham_lift_ue = _ue_var.derham_lift if self._has_lifting_ue else self.derham + # ---- solution splines (constrained) and u in unconstrained space ----- + self._u_0 = self.derham.create_spline_function("u", space_id="Hdiv") # boundary splines (u', ue') in unconstrained space — zero vectors if no lifting - self._boundary_spline_u = (_u_var.boundary_spline.vector if self._has_lifting_u - else self._derham_lift_u.coeff_spaces["2"].zeros()) - self._boundary_spline_ue = (_ue_var.boundary_spline.vector if self._has_lifting_ue - else self._derham_lift_ue.coeff_spaces["2"].zeros()) + self._boundary_spline_u = (self.variables.u.boundary_spline.vector if self._has_lifting_u else self._derham_lift_u.coeff_spaces["2"].zeros()) + self._boundary_spline_ue = (self.variables.ue.boundary_spline.vector if self._has_lifting_ue else self._derham_lift_ue.coeff_spaces["2"].zeros()) + + # boundary operators + self._b_op_u = self.variables.u.boundary_op_lift if self._has_lifting_u else IdentityOperator(self.derham.coeff_spaces["2"]) + self._b_op_ue = self.variables.ue.boundary_op_lift if self._has_lifting_ue else IdentityOperator(self.derham.coeff_spaces["2"]) + + # pre-allocated RHS vectors (constrained, after boundary operator) + self._rhs_vec_u = self.derham.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_ue = self.derham.create_spline_function("rhs_vec_ue", space_id="Hdiv") + self._rhs_vec_phi = self.derham.create_spline_function("rhs_vec_phi", space_id="L2") + + self._div_boundary_u = self.derham.create_spline_function("div_boundary_u", space_id="L2") + self._div_boundary_ue = self.derham.create_spline_function("div_boundary_ue", space_id="L2") + + # ---- source terms projected onto unconstrained space ----------------- + self._src_u = self._derham_lift_u.create_spline_function("rhs_u", space_id="Hdiv") + self._src_ue = self._derham_lift_ue.create_spline_function("rhs_ue", space_id="Hdiv") + + for rhs, source, derham_lift in [ + (self._src_u, self.options.source_u, self._derham_lift_u), + (self._src_ue, self.options.source_ue, self._derham_lift_ue), + ]: + if source is not None: + fun_vec = [lambda x, y, z, f=source, c=c: f(x, y, z)[c] for c in range(3)] + fun = [ + TransformedPformComponent( + fun_vec, + "physical", + "2", + comp=comp, + domain=self.domain, + ) + for comp in range(3) + ] + rhs.vector = derham_lift.projectors["2"](fun) # ---- unconstrained mass/basis operators (for RHS assembly) ----------- @@ -7816,7 +7845,7 @@ def allocate(self, verbose=False): self._lapl_u = (self._div_u.T @ self._mass_ops_lift_u.M3 @ self._div_u + self._S21_u.T @ self._curl_u.T @ self._M2_u @ self._curl_u @ self._S21_u) - self._A11 = - self._M2B_u / self.options.eps_norm + self.options.nu * self._lapl_u + self._A11_u = - self._M2B_u / self.options.eps_norm + self.options.nu * self._lapl_u self._M2_ue = self._mass_ops_lift_ue.M2 self._M2B_ue = - self._mass_ops_lift_ue.M2B @@ -7827,57 +7856,56 @@ def allocate(self, verbose=False): self._lapl_ue = (self._div_ue.T @ self._mass_ops_lift_ue.M3 @ self._div_ue + self._S21_ue.T @ self._curl_ue.T @ self._M2_ue @ self._curl_ue @ self._S21_ue) - self._A22 = (self.options.stab_sigma * IdentityOperator(self._derham_lift_ue.coeff_spaces["2"]) + self._A22_ue = (self.options.stab_sigma * IdentityOperator(self._derham_lift_ue.coeff_spaces["2"]) + self._M2B_ue / self.options.eps_norm + self.options.nu_e * self._lapl_ue) # ---- constrained operators (for system matrix, built from self.derham) --- - self._M2_v0 = self.mass_ops.M2 - self._M3_v0 = self.mass_ops.M3 - self._M2B_v0 = - self.mass_ops.M2B - self._div_v0 = self.derham.div - self._curl_v0 = self.derham.curl - self._S21_v0 = self.basis_ops.S21 + self._M2 = self.mass_ops.M2 + self._M3 = self.mass_ops.M3 + self._M2B = - self.mass_ops.M2B + self._div = self.derham.div + self._curl = self.derham.curl + self._S21 = self.basis_ops.S21 - self._lapl_v0 = (self._div_v0.T @ self._M3_v0 @ self._div_v0 - + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0) + self._lapl_v0 = (self._div.T @ self._M3 @ self._div + + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - self._A11_v0 = - self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = (self.options.stab_sigma * IdentityOperator(self.derham.coeff_spaces["2"]) - + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0) + self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl_v0 + self._A22 = (self.options.stab_sigma * IdentityOperator(self.derham.coeff_spaces["2"]) + + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl_v0) # ---- block saddle-point system ---------------------------------------- - self._block_domain_v0 = BlockVectorSpace(self.derham.coeff_spaces["2"], self.derham.coeff_spaces["2"]) - self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self.derham.coeff_spaces["3"] + self._block_domain = BlockVectorSpace(self.derham.coeff_spaces["2"], self.derham.coeff_spaces["2"]) + self._block_codomain_B = self.derham.coeff_spaces["3"] - self._B1_v0 = - self._M3_v0 @ self._div_v0 - self._B2_v0 = self._M3_v0 @ self._div_v0 + self._B1 = - self._M3 @ self._div + self._B2 = self._M3 @ self._div - self._B_v0 = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_B_v0, - blocks=[[self._B1_v0, self._B2_v0]] + self._B = BlockLinearOperator( + self._block_domain, self._block_codomain_B, + blocks=[[self._B1, self._B2]] ) - self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) + self._block_domain_M = BlockVectorSpace(self._block_domain, self._block_codomain_B) _A_init = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, - blocks=[[self._A11_v0, None], [None, self._A22_v0]] + self._block_domain, self._block_domain, + blocks=[[self._A11, None], [None, self._A22]] ) _M_init = BlockLinearOperator( self._block_domain_M, self._block_domain_M, - blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] + blocks=[[_A_init, self._B.T], [self._B, None]] ) if self.options.solver in get_args(LiteralOptions.OptsSaddlePointSolver): self._Minv = inverse( _M_init, self.options.solver, - A11=self._A11_v0, - A22=self._A22_v0, - B1=self._B1_v0, - B2=self._B2_v0, + A11=self._A11, + A22=self._A22, + B1=self._B1, + B2=self._B2, recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, @@ -7892,140 +7920,44 @@ def allocate(self, verbose=False): verbose=self.options.solver_params.verbose, ) - # ---- source terms projected onto unconstrained space ----------------- - - self._projector_u = L2Projector(space_id="Hdiv", mass_ops=self._mass_ops_lift_u) - self._projector_ue = L2Projector(space_id="Hdiv", mass_ops=self._mass_ops_lift_ue) - - self._rhs_u = self._derham_lift_u.create_spline_function("rhs_u", space_id="Hdiv") - self._rhs_ue = self._derham_lift_ue.create_spline_function("rhs_ue", space_id="Hdiv") - - for rhs, source, derham_lift in [ - (self._rhs_u, self.options.source_u, self._derham_lift_u), - (self._rhs_ue, self.options.source_ue, self._derham_lift_ue), - ]: - if source is not None: - fun_vec = [lambda x, y, z, f=source, c=c: f(x, y, z)[c] for c in range(3)] - fun = [ - TransformedPformComponent( - fun_vec, - "physical", - "2", - comp=comp, - domain=self.domain, - ) - for comp in range(3) - ] - rhs.vector = derham_lift.projectors["2"](fun) - - # ---- solution splines (constrained) and u in unconstrained space ----- - - self._u = self.derham.create_spline_function("u", space_id="Hdiv") - self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") - self._phi = self.derham.create_spline_function("phi", space_id="L2") - - # u/ue embedded in unconstrained space for M2 application in RHS - self._u_lift = self._derham_lift_u.create_spline_function("u_lift", space_id="Hdiv") - self._ue_lift = self._derham_lift_ue.create_spline_function("ue_lift", space_id="Hdiv") - - # pre-allocated RHS vectors (constrained, after boundary_ops projection) - self._rhs_vec_u = self.derham.create_spline_function("rhs_vec_u", space_id="Hdiv") - self._rhs_vec_ue = self.derham.create_spline_function("rhs_vec_ue", space_id="Hdiv") - - # boundary splines in constrained space for add/subtract around solve - self._boundary_spline_u_v0 = self.derham.create_spline_function("boundary_spline_u_v0", space_id="Hdiv") - self._boundary_spline_ue_v0 = self.derham.create_spline_function("boundary_spline_ue_v0", space_id="Hdiv") - - self._rhs_full_u = self.derham.create_spline_function("rhs_full_u", space_id="Hdiv") - self._rhs_full_ue = self.derham.create_spline_function("rhs_full_ue", space_id="Hdiv") - + self._RHS = BlockVector(self._block_domain_M, + blocks=[BlockVector(self._block_domain, + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]), + self._rhs_vec_phi.vector]) + self._SOL = self._block_domain_M.zeros() # ========================================================================= ### Time step # ========================================================================= + def __call__(self, dt): - def __call__(self, dt): # TODO this is still a complete mess, clean up after 2D lifting also works - - # --- copy current state (full solution = u_0 + u') --- - self._u.vector = self.variables.u.spline.vector - self._ue.vector = self.variables.ue.spline.vector - - # --- store boundary spline in constrained space for reconstruction --- - self._boundary_spline_u_v0.vector = self._boundary_spline_u - self._boundary_spline_ue_v0.vector = self._boundary_spline_ue - - # --- strip lifting to get u_0 in constrained space --- - if self._has_lifting_u: - self._u.vector = self._u.vector - self._boundary_spline_u_v0.vector - if self._has_lifting_ue: - self._ue.vector = self._ue.vector - self._boundary_spline_ue_v0.vector - - # --- embed u_0 into unconstrained space for M2 application --- - # u_0 satisfies homogeneous Dirichlet so boundary DOFs are zero in both spaces - self._u_lift.vector = self._u.vector - self._ue_lift.vector = self._ue.vector - - # --- rebuild system matrix if dt changed --- - if dt != self._dt: - self._dt = dt - _A = BlockLinearOperator( # TODO avoid allocating - self._block_domain_v0, self._block_codomain_v0, - blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]] - ) - _M = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, - blocks=[[_A, self._B_v0.T], [self._B_v0, None]] - ) - self._Minv.linop = _M - - # --- assemble RHS fully in unconstrained space, then project to constrained --- - # ion: F1 = bc_op @ (rhs_u + M2_u/dt * u_0 - (A11 + M2_u/dt) * u') - # electron: F2 = bc_op @ (rhs_ue - A22 * ue') - - rhs_u_full = (self._M2_u.dot(self._rhs_u.vector) - + self._M2_u.dot(self._u_lift.vector) / dt - - self._A11.dot(self._boundary_spline_u) - - self._M2_u.dot(self._boundary_spline_u) / dt) - - rhs_ue_full = (self._M2_ue.dot(self._rhs_ue.vector) - - self._A22.dot(self._boundary_spline_ue)) - - self._rhs_full_u.vector = rhs_u_full - self._rhs_full_ue.vector = rhs_ue_full - - self._rhs_vec_u.vector = self.derham.boundary_ops["2"].dot(self._rhs_full_u.vector) # TODO implement or change boundary operator to also change the space - self._rhs_vec_ue.vector = self.derham.boundary_ops["2"].dot(self._rhs_full_ue.vector) + # --- copy current homogeneous solution --- + self._u_0.vector = self.variables.u.spline.vector - tmp1 = self.derham.create_spline_function("tmp1", space_id="L2") - tmp2 = self.derham.create_spline_function("tmp2", space_id="L2") + # --- assemble RHS fully in unconstrained space, then enforce essential BCs --- + self._rhs_vec_u.vector = self._b_op_u.dot((self._M2_u.dot(self._src_u.vector) + - self._A11_u.dot(self._boundary_spline_u) + - self._M2_u.dot(self._boundary_spline_u) / dt)) + self._M2.dot(self._u_0.vector) / dt + + self._rhs_vec_ue.vector = self._b_op_ue.dot((self._M2_ue.dot(self._src_ue.vector) + - self._A22_ue.dot(self._boundary_spline_ue))) - tmp1.vector = self._div_u.dot(self._boundary_spline_u) # TODO implement identity operator between L^2 of different de rham spaces - tmp2.vector = self._div_ue.dot(self._boundary_spline_ue) + self._div_boundary_u.vector = self._div_u.dot(self._boundary_spline_u) + self._div_boundary_ue.vector = self._div_ue.dot(self._boundary_spline_ue) - phi_rhs = (self.mass_ops.M3.dot(tmp1.vector) - - self.mass_ops.M3.dot(tmp2.vector)) + self._rhs_vec_phi.vector = (self.mass_ops.M3.dot(self._div_boundary_u.vector) - self.mass_ops.M3.dot(self._div_boundary_ue.vector)) - # --- build block RHS and solve --- - _F = BlockVector(self._block_domain_v0, - blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) - _RHS = BlockVector(self._block_domain_M, - blocks=[_F, phi_rhs]) - _sol = self._Minv.dot(_RHS) + # --- build block RHS and solve --- + self._Minv.dot(BlockVector(self._block_domain_M, + blocks=[BlockVector(self._block_domain, blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]), + self._rhs_vec_phi.vector]), + out=self._SOL) + info = self._Minv.get_info() - # --- reconstruct full solution: u = u_0 + u' --- - self._u.vector = _sol[0][0] - self._ue.vector = _sol[0][1] - self._phi.vector = _sol[1] - - if self._has_lifting_u: - self._u.vector = self._u.vector + self._boundary_spline_u_v0.vector # TODO store an additional field in feecvariable that has the complete solution - if self._has_lifting_ue: - self._ue.vector = self._ue.vector + self._boundary_spline_ue_v0.vector - # --- update FEEC variables --- max_diffs = self.update_feec_variables( - u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector + u=self._SOL[0][0], ue=self._SOL[0][1], phi=self._SOL[1] ) if self.options.solver_params.info and self._rank == 0: From becdce31168db61a38632195478cce4b2f25dbdb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Sun, 19 Apr 2026 19:45:49 +0000 Subject: [PATCH 19/41] Added units for viscosity and default value for epsilon. --- src/struphy/physics/physics.py | 11 +++++++++ src/struphy/propagators/propagators_fields.py | 23 +++++++++++-------- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/src/struphy/physics/physics.py b/src/struphy/physics/physics.py index 3ffff52b1..e1603e464 100644 --- a/src/struphy/physics/physics.py +++ b/src/struphy/physics/physics.py @@ -84,6 +84,13 @@ def j(self): if not hasattr(self, "_j"): raise AttributeError("Must call Units.derive_units() to get full set of units.") return self._j + + @property + def nu(self): + """Unit of dynamic viscosity in kg/(m·s).""" + if not hasattr(self, "_nu"): + raise AttributeError("Must call Units.derive_units() to get full set of units.") + return self._nu def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk: int = None, verbose=False): """Derive the remaining units from the base units, velocity scale and bulk species' A and Z.""" @@ -129,6 +136,9 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk # current density (A/m^2) self._j = con.e * self.n * self.v + # dynamic viscosity (kg/(m·s)) + self._nu = A_bulk * con.mH * self.n * self.x * self.v if A_bulk is not None else None + # print to screen if verbose and MPI.COMM_WORLD.Get_rank() == 0: units_used = ( @@ -141,6 +151,7 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk " bar", " kg/m³", " A/m²", + " kg/(m·s)", ) logger.info("") for (k, v), u in zip(self.__dict__.items(), units_used): diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index bca76d587..7e60cf1b9 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -7716,22 +7716,23 @@ def __post_init__(self): # --- required parameters --- assert self.nu is not None, "nu must be specified" assert self.nu_e is not None, "nu_e must be specified" - assert self.eps_norm is not None, "eps_norm must be specified" + + # --- warn if no source terms --- + if self.source_u is None: + warn("No source_u specified — defaulting to zero.") + if self.source_ue is None: + warn("No source_ue specified — defaulting to zero.") + if self.eps_norm is None: + warn("No eps_norm specified — will default to ion cyclotron parameter epsilon in allocate.") # --- physical parameter sanity checks --- if self.nu < 0: raise ValueError(f"nu must be non-negative, got {self.nu}") if self.nu_e < 0: raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") - if self.eps_norm <= 0: # TODO get base epsilon from ion species if undefined + if self.eps_norm is not None and self.eps_norm <= 0: raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - # --- warn if no source terms --- - if self.source_u is None: - warn("No source_u specified — defaulting to zero.") - if self.source_ue is None: - warn("No source_ue specified — defaulting to zero.") - # --- defaults --- if self.stab_sigma is None: warn("stab_sigma not specified, defaulting to 0.0") @@ -7743,7 +7744,8 @@ def __post_init__(self): @property def options(self) -> Options: - assert hasattr(self, "_options"), "Options not set." + if not hasattr(self, "_options"): + self._options = self.Options() return self._options @options.setter @@ -7765,6 +7767,9 @@ def allocate(self, verbose=False): self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 self._dt = None + if self.options.eps_norm is None: + self._options.eps_norm = self.variables.u.species.equation_params.epsilon + # ---- lifting (derham_lift is unconstrained, self.derham is constrained) --- self._has_lifting_u = self.variables.u.derham_lift is not None self._has_lifting_ue = self.variables.ue.derham_lift is not None From a6df1cb61216e81f3a0268fa7ee6b837785381f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Sun, 19 Apr 2026 19:53:15 +0000 Subject: [PATCH 20/41] Minor consistency changes. --- src/struphy/io/options.py | 2 +- src/struphy/propagators/propagators_fields.py | 22 +++++-------------- src/struphy/utils/utils.py | 2 +- 3 files changed, 7 insertions(+), 19 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 41d8842a8..97f8b36be 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, fields from typing import Any, Callable, Literal -from struphy.utils.utils import check_option +from struphy.utils.utils import __dataclass_repr_no_defaults__, all_class_params_are_default, check_option import cunumpy as xp from feectools.ddm.mpi import mpi as MPI diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 7e60cf1b9..6cf0f45f7 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -3,7 +3,7 @@ import copy import logging from copy import deepcopy -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Callable, Literal, get_args from warnings import warn from warnings import warn @@ -7699,24 +7699,19 @@ def __init__(self): @dataclass class Options(): - nu: float | None = None - nu_e: float | None = None + nu: float + nu_e: float eps_norm: float | None = None source_u: Callable | None = None source_ue: Callable | None = None - stab_sigma: float | None = None - + stab_sigma: float = 0.0 solver: LiteralOptions.OptsGenSolver = "gmres" - solver_params: SolverParameters | None = None + solver_params: SolverParameters = field(default_factory=SolverParameters) def __post_init__(self): - # --- required parameters --- - assert self.nu is not None, "nu must be specified" - assert self.nu_e is not None, "nu_e must be specified" - # --- warn if no source terms --- if self.source_u is None: warn("No source_u specified — defaulting to zero.") @@ -7733,14 +7728,7 @@ def __post_init__(self): if self.eps_norm is not None and self.eps_norm <= 0: raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - # --- defaults --- - if self.stab_sigma is None: - warn("stab_sigma not specified, defaulting to 0.0") - self.stab_sigma = 0.0 - check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver) - if self.solver_params is None: - self.solver_params = SolverParameters() @property def options(self) -> Options: diff --git a/src/struphy/utils/utils.py b/src/struphy/utils/utils.py index 39ed325ef..502d662d0 100644 --- a/src/struphy/utils/utils.py +++ b/src/struphy/utils/utils.py @@ -109,7 +109,7 @@ def kernels_to_txt(kernels: list, output: str): # logger.info(f"kernels written to {output}.") -def check_option(opt, *options): +def check_option(opt: str | list[str], *options): """Check if opt is contained in options; if opt is a list, checks for each element.""" opts = [] for o in options: From 363fa559fd65383cc8daff940322ae53b08480f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Mon, 20 Apr 2026 11:15:57 +0000 Subject: [PATCH 21/41] Add 2D Neumann and tokamak test cases. --- struphy-parameter-files | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/struphy-parameter-files b/struphy-parameter-files index 1a67f3ed8..4f2e28ea9 160000 --- a/struphy-parameter-files +++ b/struphy-parameter-files @@ -1 +1 @@ -Subproject commit 1a67f3ed8c7f5237ff8c6bc50f80df4e8ccc3b43 +Subproject commit 4f2e28ea9db99acdcf6e0c96c325a918092b8a36 From ad7d24cfe65a7c1bd9fdd7b8a1ddb7a83344121e Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 May 2026 09:46:41 +0200 Subject: [PATCH 22/41] remove etc/ folder --- etc/jupyter/jupyter_notebook_config.d/ipyparallel.json | 7 ------- etc/jupyter/jupyter_notebook_config.d/jupyterlab.json | 7 ------- etc/jupyter/jupyter_server_config.d/ipyparallel.json | 7 ------- .../jupyter-lsp-jupyter-server.json | 7 ------- .../jupyter_server_config.d/jupyter_server_terminals.json | 7 ------- etc/jupyter/jupyter_server_config.d/jupyterlab.json | 7 ------- etc/jupyter/jupyter_server_config.d/notebook.json | 7 ------- etc/jupyter/jupyter_server_config.d/notebook_shim.json | 7 ------- etc/jupyter/nbconfig/notebook.d/widgetsnbextension.json | 5 ----- etc/jupyter/nbconfig/tree.d/ipyparallel.json | 5 ----- 10 files changed, 66 deletions(-) delete mode 100644 etc/jupyter/jupyter_notebook_config.d/ipyparallel.json delete mode 100644 etc/jupyter/jupyter_notebook_config.d/jupyterlab.json delete mode 100644 etc/jupyter/jupyter_server_config.d/ipyparallel.json delete mode 100644 etc/jupyter/jupyter_server_config.d/jupyter-lsp-jupyter-server.json delete mode 100644 etc/jupyter/jupyter_server_config.d/jupyter_server_terminals.json delete mode 100644 etc/jupyter/jupyter_server_config.d/jupyterlab.json delete mode 100644 etc/jupyter/jupyter_server_config.d/notebook.json delete mode 100644 etc/jupyter/jupyter_server_config.d/notebook_shim.json delete mode 100644 etc/jupyter/nbconfig/notebook.d/widgetsnbextension.json delete mode 100644 etc/jupyter/nbconfig/tree.d/ipyparallel.json diff --git a/etc/jupyter/jupyter_notebook_config.d/ipyparallel.json b/etc/jupyter/jupyter_notebook_config.d/ipyparallel.json deleted file mode 100644 index 4f1ba10cd..000000000 --- a/etc/jupyter/jupyter_notebook_config.d/ipyparallel.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "NotebookApp": { - "nbserver_extensions": { - "ipyparallel": true - } - } -} diff --git a/etc/jupyter/jupyter_notebook_config.d/jupyterlab.json b/etc/jupyter/jupyter_notebook_config.d/jupyterlab.json deleted file mode 100644 index 5b5dcda3a..000000000 --- a/etc/jupyter/jupyter_notebook_config.d/jupyterlab.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "NotebookApp": { - "nbserver_extensions": { - "jupyterlab": true - } - } -} diff --git a/etc/jupyter/jupyter_server_config.d/ipyparallel.json b/etc/jupyter/jupyter_server_config.d/ipyparallel.json deleted file mode 100644 index cfc9c58a6..000000000 --- a/etc/jupyter/jupyter_server_config.d/ipyparallel.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "ServerApp": { - "jpserver_extensions": { - "ipyparallel": true - } - } -} diff --git a/etc/jupyter/jupyter_server_config.d/jupyter-lsp-jupyter-server.json b/etc/jupyter/jupyter_server_config.d/jupyter-lsp-jupyter-server.json deleted file mode 100644 index 9e37d4eca..000000000 --- a/etc/jupyter/jupyter_server_config.d/jupyter-lsp-jupyter-server.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "ServerApp": { - "jpserver_extensions": { - "jupyter_lsp": true - } - } -} diff --git a/etc/jupyter/jupyter_server_config.d/jupyter_server_terminals.json b/etc/jupyter/jupyter_server_config.d/jupyter_server_terminals.json deleted file mode 100644 index 97c80c282..000000000 --- a/etc/jupyter/jupyter_server_config.d/jupyter_server_terminals.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "ServerApp": { - "jpserver_extensions": { - "jupyter_server_terminals": true - } - } -} diff --git a/etc/jupyter/jupyter_server_config.d/jupyterlab.json b/etc/jupyter/jupyter_server_config.d/jupyterlab.json deleted file mode 100644 index 99cc0846e..000000000 --- a/etc/jupyter/jupyter_server_config.d/jupyterlab.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "ServerApp": { - "jpserver_extensions": { - "jupyterlab": true - } - } -} diff --git a/etc/jupyter/jupyter_server_config.d/notebook.json b/etc/jupyter/jupyter_server_config.d/notebook.json deleted file mode 100644 index 09113911a..000000000 --- a/etc/jupyter/jupyter_server_config.d/notebook.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "ServerApp": { - "jpserver_extensions": { - "notebook": true - } - } -} diff --git a/etc/jupyter/jupyter_server_config.d/notebook_shim.json b/etc/jupyter/jupyter_server_config.d/notebook_shim.json deleted file mode 100644 index 1e789c3d5..000000000 --- a/etc/jupyter/jupyter_server_config.d/notebook_shim.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "ServerApp": { - "jpserver_extensions": { - "notebook_shim": true - } - } -} diff --git a/etc/jupyter/nbconfig/notebook.d/widgetsnbextension.json b/etc/jupyter/nbconfig/notebook.d/widgetsnbextension.json deleted file mode 100644 index 7a17570d6..000000000 --- a/etc/jupyter/nbconfig/notebook.d/widgetsnbextension.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "load_extensions": { - "jupyter-js-widgets/extension": true - } -} diff --git a/etc/jupyter/nbconfig/tree.d/ipyparallel.json b/etc/jupyter/nbconfig/tree.d/ipyparallel.json deleted file mode 100644 index 6d98e1861..000000000 --- a/etc/jupyter/nbconfig/tree.d/ipyparallel.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "load_extensions": { - "ipyparallel/main": true - } -} From 32e8903583845f0254f297f7c674b60103f84239 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 May 2026 09:50:24 +0200 Subject: [PATCH 23/41] revert io/options.py to devel --- src/struphy/io/options.py | 94 +++++++++------------------------------ 1 file changed, 21 insertions(+), 73 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 97f8b36be..6ecfff7c7 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -3,67 +3,26 @@ from dataclasses import dataclass, fields from typing import Any, Callable, Literal -from struphy.utils.utils import __dataclass_repr_no_defaults__, all_class_params_are_default, check_option - -import cunumpy as xp -from feectools.ddm.mpi import mpi as MPI - -## Literal options - -# time -SplitAlgos = Literal["LieTrotter", "Strang"] - -# derham -PolarRegularity = Literal[-1, 1] -OptsFEECSpace = Literal["H1", "Hcurl", "Hdiv", "L2", "H1vec"] -OptsVecSpace = Literal["Hcurl", "Hdiv", "H1vec"] - -# fields background -BackgroundTypes = Literal["LogicalConst", "FluidEquilibrium"] - -# perturbations -NoiseDirections = Literal["e1", "e2", "e3", "e1e2", "e1e3", "e2e3", "e1e2e3"] -GivenInBasis = Literal["0", "1", "2", "3", "v", "physical", "physical_at_eta", "norm", None] - -# solvers -OptsSymmSolver = Literal["pcg", "cg"] -OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"] -OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] -OptsSaddlePointSolver = Literal["Uzawa", "GMRES"] # todo -OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] -OptsNonlinearSolver = Literal["Picard", "Newton"] - -# markers -OptsPICSpace = Literal["Particles6D", "DeltaFParticles6D", "Particles5D", "Particles3D"] -OptsMarkerBC = Literal["periodic", "reflect"] -OptsRecontructBC = Literal["periodic", "mirror", "fixed"] -OptsLoading = Literal[ - "pseudo_random", - "sobol_standard", - "sobol_antithetic", - "external", - "restart", - "tesselation", -] -OptsSpatialLoading = Literal["uniform", "disc"] -OptsMPIsort = Literal["each", "last", None] - -# filters -OptsFilter = Literal["fourier_in_tor", "hybrid", "three_point", None] - -# sph -OptsKernel = Literal[ - "trigonometric_1d", - "gaussian_1d", - "linear_1d", - "trigonometric_2d", - "gaussian_2d", - "linear_2d", - "trigonometric_3d", - "gaussian_3d", - "linear_isotropic_3d", - "linear_3d", -] +from struphy.utils.utils import ( + __class_with_params_repr_no_defaults__, + __dataclass_repr_no_defaults__, + all_class_params_are_default, + check_option, +) + +logger = logging.getLogger("struphy") + + +class OptionsBase: + def to_dict(self) -> dict: + """Convert dataclass instance to dictionary.""" + return {field.name: getattr(self, field.name) for field in fields(type(self)) if field.init} + + @classmethod + def from_dict(cls, dct) -> "Any": + """Create dataclass instance from dictionary.""" + valid_fields = {field.name for field in fields(cls) if field.init} + return cls(**{key: value for key, value in dct.items() if key in valid_fields}) @dataclass @@ -159,17 +118,6 @@ class LiteralOptions: "heat_flux_3", ] -class OptionsBase: - def to_dict(self) -> dict: - """Convert dataclass instance to dictionary.""" - return {field.name: getattr(self, field.name) for field in fields(type(self)) if field.init} - - @classmethod - def from_dict(cls, dct) -> "Any": - """Create dataclass instance from dictionary.""" - valid_fields = {field.name for field in fields(cls) if field.init} - return cls(**{key: value for key, value in dct.items() if key in valid_fields}) - @dataclass class Time(OptionsBase): @@ -316,7 +264,7 @@ def __repr_no_defaults__(self): @property def is_default(self): return all_class_params_are_default(self) - + @dataclass class FieldsBackground(OptionsBase): From 4cf8354896f970af7004768464e346b3d4cac456 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 May 2026 09:52:02 +0200 Subject: [PATCH 24/41] revert io/setup.py to devel --- src/struphy/io/setup.py | 89 ----------------------------------------- 1 file changed, 89 deletions(-) diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index cb75f4219..19d3ed5b9 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -31,95 +31,6 @@ def import_parameters_py(params_path: str, name: str = "parameters") -> ModuleTy return params_in -def setup_derham( - grid: TensorProductGrid, - options: DerhamOptions, - comm: MPI.Intracomm = None, - domain: Domain = None, - verbose=False, -): - """ - Creates the 3d derham sequence for given grid parameters. - - Parameters - ---------- - grid : TensorProductGrid - The FEEC grid. - - comm: Intracomm - MPI communicator (sub_comm if clones are used). - - domain : Domain, optional - The Struphy domain object for evaluating the mapping F : [0, 1]^3 --> R^3 and the corresponding metric coefficients. - - verbose : bool - Show info on screen. - - Returns - ------- - derham : struphy.feec.psydac_derham.Derham - Discrete de Rham sequence on the logical unit cube. - """ - - from struphy.feec.psydac_derham import Derham - - # number of grid cells - Nel = grid.Nel - # mpi - mpi_dims_mask = grid.mpi_dims_mask - - # spline degrees - p = options.p - # spline types (clamped vs. periodic) - spl_kind = options.spl_kind - # boundary conditions (Homogeneous Dirichlet or None) - dirichlet_bc = options.dirichlet_bc - # Number of quadrature points per histopolation cell - nq_pr = options.nq_pr - # Number of quadrature points per grid cell for L^2 - nquads = options.nquads - # C^k smoothness at eta_1=0 for polar domains - polar_ck = options.polar_ck - # local commuting projectors - local_projectors = options.local_projectors - - lifting = options.lifting - - derham = Derham( - Nel, - p, - spl_kind, - dirichlet_bc=dirichlet_bc, - lifting=lifting, - nquads=nquads, - nq_pr=nq_pr, - comm=comm, - mpi_dims_mask=mpi_dims_mask, - with_projectors=True, - polar_ck=polar_ck, - domain=domain, - local_projectors=local_projectors, - ) - - - if MPI.COMM_WORLD.Get_rank() == 0 and verbose: - print("\nDERHAM:") - print("number of elements:".ljust(25), Nel) - print("spline degrees:".ljust(25), p) - print("periodic bcs:".ljust(25), spl_kind) - print("hom. Dirichlet bc:".ljust(25), dirichlet_bc) - print("GL quad pts (L2):".ljust(25), nquads) - print("GL quad pts (hist):".ljust(25), nq_pr) - print( - "MPI proc. per dir.:".ljust(25), - derham.domain_decomposition.nprocs, - ) - print("use polar splines:".ljust(25), derham.polar_ck == 1) - print("domain on process 0:".ljust(25), derham.domain_array[0]) - - return derham - - def descend_options_dict( d: dict, out: list | dict, From d63b12be8cbb588bfa581293a781cf09f7e00ca2 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 May 2026 09:53:37 +0200 Subject: [PATCH 25/41] remove double import --- src/struphy/propagators/base.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 240039ca8..925c37b36 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -14,7 +14,6 @@ from struphy.feec.psydac_derham import Derham from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB from struphy.geometry.base import Domain -from struphy.utils.utils import check_option from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.utils.utils import check_option From b911ab9e993f3e9c9a04a39d191ae1c422eafd77 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 May 2026 09:54:22 +0200 Subject: [PATCH 26/41] re-delete struphy-parameter-files --- struphy-parameter-files | 1 - 1 file changed, 1 deletion(-) delete mode 160000 struphy-parameter-files diff --git a/struphy-parameter-files b/struphy-parameter-files deleted file mode 160000 index 4f2e28ea9..000000000 --- a/struphy-parameter-files +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 4f2e28ea9db99acdcf6e0c96c325a918092b8a36 From 602a3fa22fae90dc36d5e6825c6f0e6ab46390e2 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 May 2026 09:55:02 +0200 Subject: [PATCH 27/41] formatting --- src/struphy/feec/linear_operators.py | 2 -- src/struphy/models/variables.py | 9 +++++---- src/struphy/physics/physics.py | 2 +- src/struphy/propagators/base.py | 1 - 4 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index c3a132e5a..1cb513ca1 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -515,10 +515,8 @@ def dot(self, v, out=None): return out - def transpose(self, conjugate=False): """ Returns the transposed operator. """ return BoundaryOperator(self._codomain, self._space_id, self.bc, codomain=self._domain) - diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 33204c2de..db16037d3 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -260,7 +260,7 @@ def boundary_spline(self) -> SplineFunction | None: if not hasattr(self, "_boundary_spline"): self._boundary_spline = None return self._boundary_spline - + @property def spline_full(self) -> SplineFunction | None: """Full solution spline (lifting + zero-BC part) in the unconstrained space. Only allocated if lifting_function is not None.""" @@ -274,7 +274,7 @@ def boundary_op(self) -> BoundaryOperator | None: if not hasattr(self, "_boundary_op"): self._boundary_op = None return self._boundary_op - + @property def boundary_op_lift(self) -> BoundaryOperator | None: """Boundary operator mapping from the unconstrained (lifted) space to the constrained space. @@ -420,7 +420,6 @@ def allocate( ] self.spline_lift.vector += self.derham_lift.projectors[derham.space_to_form[self.space]](fun) - # other helper objects self._spline_0 = self.spline_lift.copy() self.spline_lift.vector.copy(out=self.spline_0.vector) @@ -428,7 +427,9 @@ def allocate( self._boundary_op = BoundaryOperator(self.spline_lift.space, self.space, derham.dirichlet_bc) - self._boundary_op_lift = BoundaryOperator(self.spline_lift.space, self.space, derham.dirichlet_bc, codomain=self._spline.space) + self._boundary_op_lift = BoundaryOperator( + self.spline_lift.space, self.space, derham.dirichlet_bc, codomain=self._spline.space + ) self.compute_boundary_spline() diff --git a/src/struphy/physics/physics.py b/src/struphy/physics/physics.py index e1603e464..2bb40a861 100644 --- a/src/struphy/physics/physics.py +++ b/src/struphy/physics/physics.py @@ -84,7 +84,7 @@ def j(self): if not hasattr(self, "_j"): raise AttributeError("Must call Units.derive_units() to get full set of units.") return self._j - + @property def nu(self): """Unit of dynamic viscosity in kg/(m·s).""" diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 925c37b36..062b72737 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -133,7 +133,6 @@ def update_feec_variables(self, **new_coeffs): # important: sync processes! old.update_ghost_regions() - return diffs @property From c666a992235ccfe78b344dcd2072dd72ac56fede Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 May 2026 10:31:05 +0200 Subject: [PATCH 28/41] re-instate Davids new propagator (now in its own file) --- .../two_fluid_quasi_neutral_full.py | 400 +++++++----------- 1 file changed, 159 insertions(+), 241 deletions(-) diff --git a/src/struphy/propagators/two_fluid_quasi_neutral_full.py b/src/struphy/propagators/two_fluid_quasi_neutral_full.py index 741a0d58e..66a6e73e7 100644 --- a/src/struphy/propagators/two_fluid_quasi_neutral_full.py +++ b/src/struphy/propagators/two_fluid_quasi_neutral_full.py @@ -11,6 +11,7 @@ from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import L2Projector, WeightedMassOperators +from struphy.geometry.utilities import TransformedPformComponent from struphy.io.options import LiteralOptions from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable @@ -106,10 +107,6 @@ class Options: Electron viscosity coefficient. eps_norm : float, default=1e-3 Normalization/scaling parameter in Lorentz coupling terms. - boundary_data_u : dict[tuple[int, int], Callable] or None, default=None - Inhomogeneous Dirichlet data for ion velocity faces. - boundary_data_ue : dict[tuple[int, int], Callable] or None, default=None - Inhomogeneous Dirichlet data for electron velocity faces. source_u : Callable or None, default=None Source term for ion momentum equation. source_ue : Callable or None, default=None @@ -124,33 +121,32 @@ class Options: nu: float = 1.0 nu_e: float = 1.0 - eps_norm: float = 1e-3 + eps_norm: float | None = None - boundary_data_u: dict[tuple[int, int], Callable] | None = None - boundary_data_ue: dict[tuple[int, int], Callable] | None = None - - source_u: Callable | None = None + source_u: Callable | None = None source_ue: Callable | None = None - stab_sigma: float | None = None - + stab_sigma: float = 0.0 solver: LiteralOptions.OptsGenSolver = "gmres" solver_params: SolverParameters | None = None def __post_init__(self): + # --- warn if no source terms --- + if self.source_u is None: + warn("No source_u specified — defaulting to zero.") + if self.source_ue is None: + warn("No source_ue specified — defaulting to zero.") + if self.eps_norm is None: + warn("No eps_norm specified — will default to ion cyclotron parameter epsilon in allocate.") + # --- physical parameter sanity checks --- if self.nu < 0: raise ValueError(f"nu must be non-negative, got {self.nu}") if self.nu_e < 0: raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") - if self.eps_norm <= 0: + if self.eps_norm is not None and self.eps_norm <= 0: raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - # --- warn if no source terms --- - if self.source_u is None: - warn("No source_u specified — defaulting to zero.") - if self.source_ue is None: - warn("No source_ue specified — defaulting to zero.") # --- defaults --- if self.stab_sigma is None: @@ -163,7 +159,8 @@ def __post_init__(self): @property def options(self) -> Options: - assert hasattr(self, "_options"), "Options not set." + if not hasattr(self, "_options"): + self._options = self.Options() return self._options @options.setter @@ -175,45 +172,6 @@ def options(self, new): logger.info(f" {k}: {v}") self._options = new - # ========================================================================= - ### Boundary condition helpers - # ========================================================================= - - def _get_dirichlet_faces(self): - """Infer which faces have Dirichlet BCs by comparing derham and derham_v0. - - A face is Dirichlet if it is unclamped in derham but clamped in derham_v0 - (i.e. lifting is True there). - """ - faces = [] - derham = self.derham - derham_v0 = derham - - if derham_v0 is None: - return faces - - bc = derham.dirichlet_bc - bc_v0 = derham_v0.dirichlet_bc - - for d in range(3): - if derham.spl_kind[d]: - continue # periodic axis, no Dirichlet - for s, side in enumerate((-1, 1)): - # clamped in v0 but not in derham => this is a lifted (inhom Dirichlet) face - unclamped = not bc[d][s] - clamped_v0 = bc_v0[d][s] if bc_v0 is not None else False - if unclamped and clamped_v0: - faces.append((d, side)) - # clamped in both => homogeneous Dirichlet, also need to zero DOFs - elif bc[d][s] and clamped_v0: - faces.append((d, side)) - return faces - - def _apply_essential_bc(self, vec): - """Zero out Dirichlet DOFs, inferred from derham vs derham_v0.""" - for d, side in self._dirichlet_faces: - apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) - # ========================================================================= ### Allocate # ========================================================================= @@ -222,97 +180,152 @@ def allocate(self, verbose=False): self.verbose = verbose self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 - self._dt = None + self._dt = None + + if self.options.eps_norm is None: + self._options.eps_norm = self.variables.u.species.equation_params.epsilon + + # ---- lifting (derham_lift is unconstrained, self.derham is constrained) --- + self._has_lifting_u = self.variables.u.derham_lift is not None + self._has_lifting_ue = self.variables.ue.derham_lift is not None + + self._derham_lift_u = self.variables.u.derham_lift if self._has_lifting_u else self.derham + self._derham_lift_ue = self.variables.ue.derham_lift if self._has_lifting_ue else self.derham + + # ---- solution splines (constrained) and u in unconstrained space ----- + self._u_0 = self.derham.create_spline_function("u", space_id="Hdiv") + + # boundary splines (u', ue') in unconstrained space — zero vectors if no lifting + self._boundary_spline_u = (self.variables.u.boundary_spline.vector if self._has_lifting_u else self._derham_lift_u.coeff_spaces["2"].zeros()) + self._boundary_spline_ue = (self.variables.ue.boundary_spline.vector if self._has_lifting_ue else self._derham_lift_ue.coeff_spaces["2"].zeros()) + + # boundary operators + self._b_op_u = self.variables.u.boundary_op_lift if self._has_lifting_u else IdentityOperator(self.derham.coeff_spaces["2"]) + self._b_op_ue = self.variables.ue.boundary_op_lift if self._has_lifting_ue else IdentityOperator(self.derham.coeff_spaces["2"]) + + # pre-allocated RHS vectors (constrained, after boundary operator) + self._rhs_vec_u = self.derham.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_ue = self.derham.create_spline_function("rhs_vec_ue", space_id="Hdiv") + self._rhs_vec_phi = self.derham.create_spline_function("rhs_vec_phi", space_id="L2") + + self._div_boundary_u = self.derham.create_spline_function("div_boundary_u", space_id="L2") + self._div_boundary_ue = self.derham.create_spline_function("div_boundary_ue", space_id="L2") + + # ---- source terms projected onto unconstrained space ----------------- + self._src_u = self._derham_lift_u.create_spline_function("rhs_u", space_id="Hdiv") + self._src_ue = self._derham_lift_ue.create_spline_function("rhs_ue", space_id="Hdiv") + + for rhs, source, derham_lift in [ + (self._src_u, self.options.source_u, self._derham_lift_u), + (self._src_ue, self.options.source_ue, self._derham_lift_ue), + ]: + if source is not None: + fun_vec = [lambda x, y, z, f=source, c=c: f(x, y, z)[c] for c in range(3)] + fun = [ + TransformedPformComponent( + fun_vec, + "physical", + "2", + comp=comp, + domain=self.domain, + ) + for comp in range(3) + ] + rhs.vector = derham_lift.projectors["2"](fun) - # ---- v0 de Rham complex (from derham.derham_v0) ---------------------- - self._derham_v0 = self.derham + # ---- unconstrained mass/basis operators (for RHS assembly) ----------- - self._mass_ops_v0 = WeightedMassOperators( - self._derham_v0, - self.domain, - eq_mhd=self.mass_ops.eq_mhd, + self._mass_ops_lift_u = WeightedMassOperators( + self._derham_lift_u, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.mass_ops.weights["eq_mhd"], + ) + self._mass_ops_lift_ue = WeightedMassOperators( + self._derham_lift_ue, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.mass_ops.weights["eq_mhd"], + ) + self._basis_ops_lift_u = BasisProjectionOperators( + self._derham_lift_u, self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.basis_ops.weights["eq_mhd"], ) - self._basis_ops_v0 = BasisProjectionOperators( - self._derham_v0, - self.domain, + self._basis_ops_lift_ue = BasisProjectionOperators( + self._derham_lift_ue, self.domain, verbose=self.options.solver_params.verbose, eq_mhd=self.basis_ops.weights["eq_mhd"], ) - # ---- Dirichlet faces (inferred from derham vs derham_v0) ------------- + self._M2_u = self._mass_ops_lift_u.M2 + self._M2B_u = - self._mass_ops_lift_u.M2B + self._div_u = self._derham_lift_u.div + self._curl_u = self._derham_lift_u.curl + self._S21_u = self._basis_ops_lift_u.S21 - self._dirichlet_faces = self._get_dirichlet_faces() + self._lapl_u = (self._div_u.T @ self._mass_ops_lift_u.M3 @ self._div_u + + self._S21_u.T @ self._curl_u.T @ self._M2_u @ self._curl_u @ self._S21_u) - # ---- unconstrained operators (for RHS assembly) ---------------------- + self._A11_u = - self._M2B_u / self.options.eps_norm + self.options.nu * self._lapl_u - self._M2 = self.mass_ops.M2 - self._M2B = -self.mass_ops.M2B - self._div = self.derham.div - self._curl = self.derham.curl - self._S21 = self.basis_ops.S21 + self._M2_ue = self._mass_ops_lift_ue.M2 + self._M2B_ue = - self._mass_ops_lift_ue.M2B + self._div_ue = self._derham_lift_ue.div + self._curl_ue = self._derham_lift_ue.curl + self._S21_ue = self._basis_ops_lift_ue.S21 - self._lapl = ( - self._div.T @ self.mass_ops.M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21 - ) + self._lapl_ue = (self._div_ue.T @ self._mass_ops_lift_ue.M3 @ self._div_ue + + self._S21_ue.T @ self._curl_ue.T @ self._M2_ue @ self._curl_ue @ self._S21_ue) - self._A11 = -self._M2B / self.options.eps_norm + self.options.nu * self._lapl - self._A22 = ( - -self.options.stab_sigma * IdentityOperator(self.derham.V2) - + self._M2B / self.options.eps_norm - + self.options.nu_e * self._lapl - ) + self._A22_ue = (self.options.stab_sigma * IdentityOperator(self._derham_lift_ue.coeff_spaces["2"]) + + self._M2B_ue / self.options.eps_norm + self.options.nu_e * self._lapl_ue) - # ---- constrained operators (for system matrix) ----------------------- + # ---- constrained operators (for system matrix, built from self.derham) --- - self._M2_v0 = self._mass_ops_v0.M2 - self._M3_v0 = self._mass_ops_v0.M3 - self._M2B_v0 = -self._mass_ops_v0.M2B - self._div_v0 = self._derham_v0.div - self._curl_v0 = self._derham_v0.curl - self._S21_v0 = self._basis_ops_v0.S21 + self._M2 = self.mass_ops.M2 + self._M3 = self.mass_ops.M3 + self._M2B = - self.mass_ops.M2B + self._div = self.derham.div + self._curl = self.derham.curl + self._S21 = self.basis_ops.S21 - self._lapl_v0 = ( - self._div_v0.T @ self._M3_v0 @ self._div_v0 - + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0 - ) + self._lapl_v0 = (self._div.T @ self._M3 @ self._div + + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) - self._A11_v0 = -self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = ( - -self.options.stab_sigma * IdentityOperator(self._derham_v0.V2) - + self._M2B_v0 / self.options.eps_norm - + self.options.nu_e * self._lapl_v0 - ) + self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl_v0 + self._A22 = (self.options.stab_sigma * IdentityOperator(self.derham.coeff_spaces["2"]) + + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl_v0) # ---- block saddle-point system ---------------------------------------- - self._block_domain_v0 = BlockVectorSpace(self._derham_v0.V2, self._derham_v0.V2) - self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self._derham_v0.V3 + self._block_domain = BlockVectorSpace(self.derham.coeff_spaces["2"], self.derham.coeff_spaces["2"]) + self._block_codomain_B = self.derham.coeff_spaces["3"] - self._B1_v0 = -self._M3_v0 @ self._div_v0 - self._B2_v0 = self._M3_v0 @ self._div_v0 + self._B1 = - self._M3 @ self._div + self._B2 = self._M3 @ self._div - self._B_v0 = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_B_v0, blocks=[[self._B1_v0, self._B2_v0]] + self._B = BlockLinearOperator( + self._block_domain, self._block_codomain_B, + blocks=[[self._B1, self._B2]] ) - self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) + self._block_domain_M = BlockVectorSpace(self._block_domain, self._block_codomain_B) _A_init = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, blocks=[[self._A11_v0, None], [None, self._A22_v0]] + self._block_domain, self._block_domain, + blocks=[[self._A11, None], [None, self._A22]] ) _M_init = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] + self._block_domain_M, self._block_domain_M, + blocks=[[_A_init, self._B.T], [self._B, None]] ) if self.options.solver in get_args(LiteralOptions.OptsSaddlePointSolver): self._Minv = inverse( - _M_init, - self.options.solver, - A11=self._A11_v0, - A22=self._A22_v0, - B1=self._B1_v0, - B2=self._B2_v0, + _M_init, self.options.solver, + A11=self._A11, + A22=self._A22, + B1=self._B1, + B2=self._B2, recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, @@ -320,148 +333,53 @@ def allocate(self, verbose=False): ) else: self._Minv = inverse( - _M_init, - self.options.solver, + _M_init, self.options.solver, recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, ) - # ---- projector ------------------------------------------------------- - - self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - - # ---- solution spline functions (unconstrained) ----------------------- - - self._u = self.derham.create_spline_function("u", space_id="Hdiv") - self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") - self._phi = self.derham.create_spline_function("phi", space_id="L2") - - # ---- BC lifts (unconstrained) ---------------------------------------- - - self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") - self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") - - for u_prime, boundary_data in [ - (self._u_prime, self.options.boundary_data_u), - (self._ue_prime, self.options.boundary_data_ue), - ]: - if boundary_data is None: - continue - for (d, side), f_bc in boundary_data.items(): - if (d, side) in self._dirichlet_faces: - bc_pulled = lambda *etas, f=f_bc: self.domain.pull( - [ - lambda x, y, z, f=f: f(x, y, z)[0], - lambda x, y, z, f=f: f(x, y, z)[1], - lambda x, y, z, f=f: f(x, y, z)[2], - ], - *etas, - kind="2", - ) - _vec = self._projector( - [ - lambda *etas: bc_pulled(*etas)[0], - lambda *etas: bc_pulled(*etas)[1], - lambda *etas: bc_pulled(*etas)[2], - ] - ) - for d2, side2 in self._dirichlet_faces: - if (d2, side2) != (d, side): - apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) - u_prime.vector += _vec - - self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") - self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") - - self._u_prime_v0.vector = self._u_prime.vector - self._ue_prime_v0.vector = self._ue_prime.vector - - # ---- projected source terms (unconstrained) -------------------------- - - self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") - self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") - - for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: - if source is not None: - src_pulled = lambda *etas, f=source: self.domain.pull( - [ - lambda x, y, z, f=f: f(x, y, z)[0], - lambda x, y, z, f=f: f(x, y, z)[1], - lambda x, y, z, f=f: f(x, y, z)[2], - ], - *etas, - kind="2", - ) - rhs.vector = self._projector.get_dofs( - [ - lambda *etas: src_pulled(*etas)[0], - lambda *etas: src_pulled(*etas)[1], - lambda *etas: src_pulled(*etas)[2], - ] - ) - - # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- - - self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") - self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") + self._RHS = BlockVector(self._block_domain_M, + blocks=[BlockVector(self._block_domain, + blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]), + self._rhs_vec_phi.vector]) + self._SOL = self._block_domain_M.zeros() # ========================================================================= ### Time step # ========================================================================= - def __call__(self, dt): - # --- copy current state --- - self._u.vector = self.variables.u.spline.vector - self._ue.vector = self.variables.ue.spline.vector - - # --- rebuild system matrix if dt changed --- - if dt != self._dt: # TODO change uzawa A11 block too - self._dt = dt - _A = BlockLinearOperator( - self._block_domain_v0, - self._block_codomain_v0, - blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]], - ) + # --- copy current homogeneous solution --- + self._u_0.vector = self.variables.u.spline.vector - _M = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, blocks=[[_A, self._B_v0.T], [self._B_v0, None]] - ) - self._Minv.linop = _M - - # --- assemble RHS in unconstrained space, then zero boundary DOFs --- - # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' - # electron: F2 = rhs_ue - A22 * ue' - self._rhs_vec_u.vector = ( - self._rhs_u.vector # TODO boundary operator - + self._M2.dot(self._u.vector) / dt - - self._A11.dot(self._u_prime.vector) - - self._M2.dot(self._u_prime.vector) / dt - ) - self._rhs_vec_ue.vector = self._rhs_ue.vector - self._A22.dot(self._ue_prime.vector) + # --- assemble RHS fully in unconstrained space, then enforce essential BCs --- + self._rhs_vec_u.vector = self._b_op_u.dot((self._M2_u.dot(self._src_u.vector) + - self._A11_u.dot(self._boundary_spline_u) + - self._M2_u.dot(self._boundary_spline_u) / dt)) + self._M2.dot(self._u_0.vector) / dt + + self._rhs_vec_ue.vector = self._b_op_ue.dot((self._M2_ue.dot(self._src_ue.vector) + - self._A22_ue.dot(self._boundary_spline_ue))) - self._apply_essential_bc(self._rhs_vec_u.vector) - self._apply_essential_bc(self._rhs_vec_ue.vector) + self._div_boundary_u.vector = self._div_u.dot(self._boundary_spline_u) + self._div_boundary_ue.vector = self._div_ue.dot(self._boundary_spline_ue) - # --- build block RHS and solve --- - _F = BlockVector(self._block_domain_v0, blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) - _RHS = BlockVector(self._block_domain_M, blocks=[_F, self._block_codomain_B_v0.zeros()]) + self._rhs_vec_phi.vector = (self.mass_ops.M3.dot(self._div_boundary_u.vector) - self.mass_ops.M3.dot(self._div_boundary_ue.vector)) - _sol = self._Minv.dot(_RHS) + # --- build block RHS and solve --- + self._Minv.dot(BlockVector(self._block_domain_M, + blocks=[BlockVector(self._block_domain, blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]), + self._rhs_vec_phi.vector]), + out=self._SOL) + info = self._Minv.get_info() - # --- reconstruct full solution: u = u_0 + u' --- - self._u.vector = _sol[0][0] + self._u_prime_v0.vector - self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector - self._phi.vector = _sol[1] - # --- update FEEC variables --- - max_diffs = self.update_feec_variables(u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector) + max_diffs = self.update_feec_variables( + u=self._SOL[0][0], ue=self._SOL[0][1], phi=self._SOL[1] + ) if self.options.solver_params.info and self._rank == 0: - logger.info(f"Status: {info['success']}, Iterations: {info['niter']}") - logger.info(f"Max diffs: {max_diffs}") - logger.info(f"Status: {info['success']}, Iterations: {info['niter']}") - logger.info(f"Max diffs: {max_diffs}") + print(f"Status: {info['success']}, Iterations: {info['niter']}") + print(f"Max diffs: {max_diffs}") \ No newline at end of file From 957cb4c46c693ebb0298bdc7634396b52fdb5f22 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 May 2026 11:09:24 +0200 Subject: [PATCH 29/41] adapt call to WeightedMassOperators in new propagator --- src/struphy/propagators/two_fluid_quasi_neutral_full.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/struphy/propagators/two_fluid_quasi_neutral_full.py b/src/struphy/propagators/two_fluid_quasi_neutral_full.py index 66a6e73e7..935cc2bb0 100644 --- a/src/struphy/propagators/two_fluid_quasi_neutral_full.py +++ b/src/struphy/propagators/two_fluid_quasi_neutral_full.py @@ -237,13 +237,11 @@ def allocate(self, verbose=False): self._mass_ops_lift_u = WeightedMassOperators( self._derham_lift_u, self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.mass_ops.weights["eq_mhd"], + eq_mhd=self.mass_ops.eq_mhd, ) self._mass_ops_lift_ue = WeightedMassOperators( self._derham_lift_ue, self.domain, - verbose=self.options.solver_params.verbose, - eq_mhd=self.mass_ops.weights["eq_mhd"], + eq_mhd=self.mass_ops.eq_mhd, ) self._basis_ops_lift_u = BasisProjectionOperators( self._derham_lift_u, self.domain, From 9f0b3ddd77495fca9334ce6f262cde475e6e0f6f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 May 2026 11:30:41 +0200 Subject: [PATCH 30/41] formatting --- .../two_fluid_quasi_neutral_full.py | 208 +++++++++++------- 1 file changed, 126 insertions(+), 82 deletions(-) diff --git a/src/struphy/propagators/two_fluid_quasi_neutral_full.py b/src/struphy/propagators/two_fluid_quasi_neutral_full.py index 935cc2bb0..5c55ef13c 100644 --- a/src/struphy/propagators/two_fluid_quasi_neutral_full.py +++ b/src/struphy/propagators/two_fluid_quasi_neutral_full.py @@ -123,7 +123,7 @@ class Options: nu_e: float = 1.0 eps_norm: float | None = None - source_u: Callable | None = None + source_u: Callable | None = None source_ue: Callable | None = None stab_sigma: float = 0.0 @@ -138,7 +138,7 @@ def __post_init__(self): warn("No source_ue specified — defaulting to zero.") if self.eps_norm is None: warn("No eps_norm specified — will default to ion cyclotron parameter epsilon in allocate.") - + # --- physical parameter sanity checks --- if self.nu < 0: raise ValueError(f"nu must be non-negative, got {self.nu}") @@ -147,7 +147,6 @@ def __post_init__(self): if self.eps_norm is not None and self.eps_norm <= 0: raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - # --- defaults --- if self.stab_sigma is None: warn("stab_sigma not specified, defaulting to 0.0") @@ -180,45 +179,61 @@ def allocate(self, verbose=False): self.verbose = verbose self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 - self._dt = None + self._dt = None if self.options.eps_norm is None: self._options.eps_norm = self.variables.u.species.equation_params.epsilon # ---- lifting (derham_lift is unconstrained, self.derham is constrained) --- - self._has_lifting_u = self.variables.u.derham_lift is not None + self._has_lifting_u = self.variables.u.derham_lift is not None self._has_lifting_ue = self.variables.ue.derham_lift is not None - self._derham_lift_u = self.variables.u.derham_lift if self._has_lifting_u else self.derham + self._derham_lift_u = self.variables.u.derham_lift if self._has_lifting_u else self.derham self._derham_lift_ue = self.variables.ue.derham_lift if self._has_lifting_ue else self.derham # ---- solution splines (constrained) and u in unconstrained space ----- - self._u_0 = self.derham.create_spline_function("u", space_id="Hdiv") + self._u_0 = self.derham.create_spline_function("u", space_id="Hdiv") # boundary splines (u', ue') in unconstrained space — zero vectors if no lifting - self._boundary_spline_u = (self.variables.u.boundary_spline.vector if self._has_lifting_u else self._derham_lift_u.coeff_spaces["2"].zeros()) - self._boundary_spline_ue = (self.variables.ue.boundary_spline.vector if self._has_lifting_ue else self._derham_lift_ue.coeff_spaces["2"].zeros()) + self._boundary_spline_u = ( + self.variables.u.boundary_spline.vector + if self._has_lifting_u + else self._derham_lift_u.coeff_spaces["2"].zeros() + ) + self._boundary_spline_ue = ( + self.variables.ue.boundary_spline.vector + if self._has_lifting_ue + else self._derham_lift_ue.coeff_spaces["2"].zeros() + ) # boundary operators - self._b_op_u = self.variables.u.boundary_op_lift if self._has_lifting_u else IdentityOperator(self.derham.coeff_spaces["2"]) - self._b_op_ue = self.variables.ue.boundary_op_lift if self._has_lifting_ue else IdentityOperator(self.derham.coeff_spaces["2"]) + self._b_op_u = ( + self.variables.u.boundary_op_lift + if self._has_lifting_u + else IdentityOperator(self.derham.coeff_spaces["2"]) + ) + self._b_op_ue = ( + self.variables.ue.boundary_op_lift + if self._has_lifting_ue + else IdentityOperator(self.derham.coeff_spaces["2"]) + ) # pre-allocated RHS vectors (constrained, after boundary operator) - self._rhs_vec_u = self.derham.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_u = self.derham.create_spline_function("rhs_vec_u", space_id="Hdiv") self._rhs_vec_ue = self.derham.create_spline_function("rhs_vec_ue", space_id="Hdiv") self._rhs_vec_phi = self.derham.create_spline_function("rhs_vec_phi", space_id="L2") - self._div_boundary_u = self.derham.create_spline_function("div_boundary_u", space_id="L2") + self._div_boundary_u = self.derham.create_spline_function("div_boundary_u", space_id="L2") self._div_boundary_ue = self.derham.create_spline_function("div_boundary_ue", space_id="L2") # ---- source terms projected onto unconstrained space ----------------- - self._src_u = self._derham_lift_u.create_spline_function("rhs_u", space_id="Hdiv") + self._src_u = self._derham_lift_u.create_spline_function("rhs_u", space_id="Hdiv") self._src_ue = self._derham_lift_ue.create_spline_function("rhs_ue", space_id="Hdiv") for rhs, source, derham_lift in [ - (self._src_u, self.options.source_u, self._derham_lift_u), - (self._src_ue, self.options.source_ue, self._derham_lift_ue), - ]: + (self._src_u, self.options.source_u, self._derham_lift_u), + (self._src_ue, self.options.source_ue, self._derham_lift_ue), + ]: if source is not None: fun_vec = [lambda x, y, z, f=source, c=c: f(x, y, z)[c] for c in range(3)] fun = [ @@ -235,91 +250,102 @@ def allocate(self, verbose=False): # ---- unconstrained mass/basis operators (for RHS assembly) ----------- - self._mass_ops_lift_u = WeightedMassOperators( - self._derham_lift_u, self.domain, + self._mass_ops_lift_u = WeightedMassOperators( + self._derham_lift_u, + self.domain, eq_mhd=self.mass_ops.eq_mhd, ) self._mass_ops_lift_ue = WeightedMassOperators( - self._derham_lift_ue, self.domain, + self._derham_lift_ue, + self.domain, eq_mhd=self.mass_ops.eq_mhd, ) - self._basis_ops_lift_u = BasisProjectionOperators( - self._derham_lift_u, self.domain, + self._basis_ops_lift_u = BasisProjectionOperators( + self._derham_lift_u, + self.domain, verbose=self.options.solver_params.verbose, eq_mhd=self.basis_ops.weights["eq_mhd"], ) self._basis_ops_lift_ue = BasisProjectionOperators( - self._derham_lift_ue, self.domain, + self._derham_lift_ue, + self.domain, verbose=self.options.solver_params.verbose, eq_mhd=self.basis_ops.weights["eq_mhd"], ) - self._M2_u = self._mass_ops_lift_u.M2 - self._M2B_u = - self._mass_ops_lift_u.M2B - self._div_u = self._derham_lift_u.div + self._M2_u = self._mass_ops_lift_u.M2 + self._M2B_u = -self._mass_ops_lift_u.M2B + self._div_u = self._derham_lift_u.div self._curl_u = self._derham_lift_u.curl - self._S21_u = self._basis_ops_lift_u.S21 + self._S21_u = self._basis_ops_lift_u.S21 - self._lapl_u = (self._div_u.T @ self._mass_ops_lift_u.M3 @ self._div_u - + self._S21_u.T @ self._curl_u.T @ self._M2_u @ self._curl_u @ self._S21_u) + self._lapl_u = ( + self._div_u.T @ self._mass_ops_lift_u.M3 @ self._div_u + + self._S21_u.T @ self._curl_u.T @ self._M2_u @ self._curl_u @ self._S21_u + ) - self._A11_u = - self._M2B_u / self.options.eps_norm + self.options.nu * self._lapl_u + self._A11_u = -self._M2B_u / self.options.eps_norm + self.options.nu * self._lapl_u - self._M2_ue = self._mass_ops_lift_ue.M2 - self._M2B_ue = - self._mass_ops_lift_ue.M2B - self._div_ue = self._derham_lift_ue.div + self._M2_ue = self._mass_ops_lift_ue.M2 + self._M2B_ue = -self._mass_ops_lift_ue.M2B + self._div_ue = self._derham_lift_ue.div self._curl_ue = self._derham_lift_ue.curl - self._S21_ue = self._basis_ops_lift_ue.S21 + self._S21_ue = self._basis_ops_lift_ue.S21 - self._lapl_ue = (self._div_ue.T @ self._mass_ops_lift_ue.M3 @ self._div_ue - + self._S21_ue.T @ self._curl_ue.T @ self._M2_ue @ self._curl_ue @ self._S21_ue) + self._lapl_ue = ( + self._div_ue.T @ self._mass_ops_lift_ue.M3 @ self._div_ue + + self._S21_ue.T @ self._curl_ue.T @ self._M2_ue @ self._curl_ue @ self._S21_ue + ) - self._A22_ue = (self.options.stab_sigma * IdentityOperator(self._derham_lift_ue.coeff_spaces["2"]) - + self._M2B_ue / self.options.eps_norm + self.options.nu_e * self._lapl_ue) + self._A22_ue = ( + self.options.stab_sigma * IdentityOperator(self._derham_lift_ue.coeff_spaces["2"]) + + self._M2B_ue / self.options.eps_norm + + self.options.nu_e * self._lapl_ue + ) # ---- constrained operators (for system matrix, built from self.derham) --- - self._M2 = self.mass_ops.M2 - self._M3 = self.mass_ops.M3 - self._M2B = - self.mass_ops.M2B - self._div = self.derham.div + self._M2 = self.mass_ops.M2 + self._M3 = self.mass_ops.M3 + self._M2B = -self.mass_ops.M2B + self._div = self.derham.div self._curl = self.derham.curl - self._S21 = self.basis_ops.S21 + self._S21 = self.basis_ops.S21 - self._lapl_v0 = (self._div.T @ self._M3 @ self._div - + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + self._lapl_v0 = ( + self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21 + ) - self._A11 = - self._M2B / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22 = (self.options.stab_sigma * IdentityOperator(self.derham.coeff_spaces["2"]) - + self._M2B / self.options.eps_norm + self.options.nu_e * self._lapl_v0) + self._A11 = -self._M2B / self.options.eps_norm + self.options.nu * self._lapl_v0 + self._A22 = ( + self.options.stab_sigma * IdentityOperator(self.derham.coeff_spaces["2"]) + + self._M2B / self.options.eps_norm + + self.options.nu_e * self._lapl_v0 + ) # ---- block saddle-point system ---------------------------------------- self._block_domain = BlockVectorSpace(self.derham.coeff_spaces["2"], self.derham.coeff_spaces["2"]) self._block_codomain_B = self.derham.coeff_spaces["3"] - self._B1 = - self._M3 @ self._div - self._B2 = self._M3 @ self._div + self._B1 = -self._M3 @ self._div + self._B2 = self._M3 @ self._div - self._B = BlockLinearOperator( - self._block_domain, self._block_codomain_B, - blocks=[[self._B1, self._B2]] - ) + self._B = BlockLinearOperator(self._block_domain, self._block_codomain_B, blocks=[[self._B1, self._B2]]) self._block_domain_M = BlockVectorSpace(self._block_domain, self._block_codomain_B) _A_init = BlockLinearOperator( - self._block_domain, self._block_domain, - blocks=[[self._A11, None], [None, self._A22]] + self._block_domain, self._block_domain, blocks=[[self._A11, None], [None, self._A22]] ) _M_init = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, - blocks=[[_A_init, self._B.T], [self._B, None]] + self._block_domain_M, self._block_domain_M, blocks=[[_A_init, self._B.T], [self._B, None]] ) if self.options.solver in get_args(LiteralOptions.OptsSaddlePointSolver): self._Minv = inverse( - _M_init, self.options.solver, + _M_init, + self.options.solver, A11=self._A11, A22=self._A22, B1=self._B1, @@ -331,17 +357,21 @@ def allocate(self, verbose=False): ) else: self._Minv = inverse( - _M_init, self.options.solver, + _M_init, + self.options.solver, recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, ) - self._RHS = BlockVector(self._block_domain_M, - blocks=[BlockVector(self._block_domain, - blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]), - self._rhs_vec_phi.vector]) + self._RHS = BlockVector( + self._block_domain_M, + blocks=[ + BlockVector(self._block_domain, blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]), + self._rhs_vec_phi.vector, + ], + ) self._SOL = self._block_domain_M.zeros() # ========================================================================= @@ -350,34 +380,48 @@ def allocate(self, verbose=False): def __call__(self, dt): # --- copy current homogeneous solution --- - self._u_0.vector = self.variables.u.spline.vector + self._u_0.vector = self.variables.u.spline.vector # --- assemble RHS fully in unconstrained space, then enforce essential BCs --- - self._rhs_vec_u.vector = self._b_op_u.dot((self._M2_u.dot(self._src_u.vector) - - self._A11_u.dot(self._boundary_spline_u) - - self._M2_u.dot(self._boundary_spline_u) / dt)) + self._M2.dot(self._u_0.vector) / dt - - self._rhs_vec_ue.vector = self._b_op_ue.dot((self._M2_ue.dot(self._src_ue.vector) - - self._A22_ue.dot(self._boundary_spline_ue))) + self._rhs_vec_u.vector = ( + self._b_op_u.dot( + ( + self._M2_u.dot(self._src_u.vector) + - self._A11_u.dot(self._boundary_spline_u) + - self._M2_u.dot(self._boundary_spline_u) / dt + ) + ) + + self._M2.dot(self._u_0.vector) / dt + ) + + self._rhs_vec_ue.vector = self._b_op_ue.dot( + (self._M2_ue.dot(self._src_ue.vector) - self._A22_ue.dot(self._boundary_spline_ue)) + ) self._div_boundary_u.vector = self._div_u.dot(self._boundary_spline_u) self._div_boundary_ue.vector = self._div_ue.dot(self._boundary_spline_ue) - self._rhs_vec_phi.vector = (self.mass_ops.M3.dot(self._div_boundary_u.vector) - self.mass_ops.M3.dot(self._div_boundary_ue.vector)) + self._rhs_vec_phi.vector = self.mass_ops.M3.dot(self._div_boundary_u.vector) - self.mass_ops.M3.dot( + self._div_boundary_ue.vector + ) + + # --- build block RHS and solve --- + self._Minv.dot( + BlockVector( + self._block_domain_M, + blocks=[ + BlockVector(self._block_domain, blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]), + self._rhs_vec_phi.vector, + ], + ), + out=self._SOL, + ) - # --- build block RHS and solve --- - self._Minv.dot(BlockVector(self._block_domain_M, - blocks=[BlockVector(self._block_domain, blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]), - self._rhs_vec_phi.vector]), - out=self._SOL) - info = self._Minv.get_info() # --- update FEEC variables --- - max_diffs = self.update_feec_variables( - u=self._SOL[0][0], ue=self._SOL[0][1], phi=self._SOL[1] - ) + max_diffs = self.update_feec_variables(u=self._SOL[0][0], ue=self._SOL[0][1], phi=self._SOL[1]) if self.options.solver_params.info and self._rank == 0: print(f"Status: {info['success']}, Iterations: {info['niter']}") - print(f"Max diffs: {max_diffs}") \ No newline at end of file + print(f"Max diffs: {max_diffs}") From 738f2e72a1958806f797376f8347cbd55bebcb53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Sz=C3=A9kedi?= Date: Tue, 12 May 2026 16:25:20 +0000 Subject: [PATCH 31/41] Push verification cases to examples. --- .gitignore | 2 + .../1D_Verification.py | 267 +++++++++++++ .../2D_Verification.py | 350 ++++++++++++++++++ .../TwoFluidQuasiNeutralToy/R_Verification.py | 217 +++++++++++ feectools | 2 +- 5 files changed, 837 insertions(+), 1 deletion(-) create mode 100644 examples/TwoFluidQuasiNeutralToy/1D_Verification.py create mode 100644 examples/TwoFluidQuasiNeutralToy/2D_Verification.py create mode 100644 examples/TwoFluidQuasiNeutralToy/R_Verification.py diff --git a/.gitignore b/.gitignore index 36b889d17..1ad71ee0a 100644 --- a/.gitignore +++ b/.gitignore @@ -103,3 +103,5 @@ share/ lib64 pyvenv.cfg + +examples/TwoFluidQuasiNeutralToy/runs/*/ \ No newline at end of file diff --git a/examples/TwoFluidQuasiNeutralToy/1D_Verification.py b/examples/TwoFluidQuasiNeutralToy/1D_Verification.py new file mode 100644 index 000000000..0e1df71c2 --- /dev/null +++ b/examples/TwoFluidQuasiNeutralToy/1D_Verification.py @@ -0,0 +1,267 @@ +from cunumpy import pi, cos, sin, zeros_like, ones_like +from struphy.io.options import EnvironmentOptions, BaseUnits, Time +from struphy.geometry import domains +from struphy.fields_background import equils +from struphy.topology import grids +from struphy.io.options import DerhamOptions +from struphy.initial import perturbations +from struphy.initial.base import GenericPerturbation +from struphy import Simulation +from struphy.linear_algebra.solver import SolverParameters + +import argparse +import os +import glob +import cunumpy as xp +import matplotlib.pyplot as plt + +from struphy.models.two_fluid_quasi_neutral_toy import TwoFluidQuasiNeutralToy + +parser = argparse.ArgumentParser() +parser.add_argument('bc', choices=['periodic', 'dirichlet_hom', 'dirichlet_inhom']) +args = parser.parse_args() +BC = args.bc + +name = f"runs/sim_1D_{BC}" + +env = EnvironmentOptions(sim_folder=name) + +B0 = 0 +nu = 10.0 +nu_e = 1.0 +Nel = (32, 1, 1) +p = (1, 1, 1) +epsilon = 1.0 +dt = 1 +Tend = 1 +sigma = 0 +tol = 1e-5 + +time_opts = Time(dt=dt, Tend=Tend) +domain = domains.Cuboid() +equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) +grid = grids.TensorProductGrid(num_elements=Nel) + +# ---- boundary conditions ---- +if BC == 'periodic': + derham_opts = DerhamOptions(degree=p, bcs=(None, None, None)) + +elif BC == 'dirichlet_hom': + derham_opts = DerhamOptions(degree=p, bcs=(("dirichlet", "dirichlet"), None, None)) + +elif BC == 'dirichlet_inhom': + derham_opts = DerhamOptions(degree=p, bcs=(("dirichlet", "dirichlet"), None, None)) + lifting_function_u = GenericPerturbation(lambda x, y, z: x + 1, comp=0, given_in_basis="physical") + lifting_function_ue = GenericPerturbation(lambda x, y, z: x, comp=0, given_in_basis="physical") + +# ---- manufactured solutions ---- +if BC == 'periodic': + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x) + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + +elif BC == 'dirichlet_hom': + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + +elif BC == 'dirichlet_inhom': + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x) + x + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x) + x, xp.zeros_like(x), xp.zeros_like(x) + +# ---- source terms ---- +if BC == 'periodic': + def source_function_u(x, y, z): + fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + +elif BC == 'dirichlet_hom': + def source_function_u(x, y, z): + fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + +elif BC == 'dirichlet_inhom': + def source_function_u(x, y, z): + fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + (4.0 * nu_e * pi**2 - sigma) * sin(2 * pi * x) - sigma * x + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + +# ---- perturbation classes for MMS initial conditions ---- +class MMSIonVelocity(perturbations.Perturbation): + def __init__(self, comp=0): + self.comp = comp + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_ion_u(x, y, z)[self.comp] + + +class MMSElectronVelocity(perturbations.Perturbation): + def __init__(self, comp=0): + self.comp = comp + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_electron_u(x, y, z)[self.comp] + + +class MMSPotential(perturbations.Perturbation): + def __init__(self): + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_phi(x, y, z)[0] + + +# ---- model ---- +model = TwoFluidQuasiNeutralToy() + +model.propagators.qn_full.options = model.propagators.qn_full.Options( + nu=nu, + nu_e=nu_e, + eps_norm=epsilon, + stab_sigma=sigma, + source_u=source_function_u, + source_ue=source_function_ue, + solver='gmres', + solver_params=SolverParameters(verbose=True, info=True, tol = tol), +) + +if BC == 'dirichlet_inhom': + model.ions.u.lifting_function = lifting_function_u + model.electrons.u.lifting_function = lifting_function_ue + +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +if __name__ == "__main__": + sim.run(verbose=True) + sim.pproc(verbose=True) + sim.load_plotting_data(verbose=True) + + simdata = sim.plotting_data + n1_vals = simdata.grids_log[0] + x = xp.linspace(0, 1, 100) + + os.makedirs(f'{name}/plots', exist_ok=True) + for f in glob.glob(f'{name}/plots/*.png'): + os.remove(f) + + def save_plot(n1_vals, numerical, analytical, ylabel, title, fname, t): + plt.plot(n1_vals, numerical, label='numerical') + plt.plot(x, analytical, '--', label='manufactured') + plt.plot(n1_vals, numerical, 'k.', markersize=4, label='n1 points') + plt.xlabel('x') + plt.ylabel(ylabel) + plt.title(f'{title} at t={t:.3f}') + plt.legend() + plt.grid(True) + plt.savefig(f'{name}/plots/{fname}_{t:.3f}.png', dpi=300) + plt.clf() + + for t in list(simdata.spline_values.ions.u_log.data.keys()): + u_ions = simdata.spline_values.ions.u_log.data[t] + u_electrons = simdata.spline_values.electrons.u_log.data[t] + phi = simdata.spline_values.em_fields.phi_log.data[t] + + mms_phi_x, _, _ = mms_phi(x, x*0, x*0) + mms_ion_ux, _, _ = mms_ion_u(x, x*0, x*0) + mms_el_ux, _, _ = mms_electron_u(x, x*0, x*0) + + save_plot(n1_vals, phi[0][:, 0, 0], mms_phi_x, 'φ', 'Potential φ', 'plot_potential', t) + save_plot(n1_vals, u_ions[0][:, 0, 0], mms_ion_ux, 'u_x', 'Ion velocity u_x', 'plot_ion_ux', t) + save_plot(n1_vals, u_electrons[0][:, 0, 0], mms_el_ux, 'u_x', 'Electron velocity', 'plot_electron_ux', t) + + # ---- lifting diagnostics ---- + if BC == 'dirichlet_inhom': + e1 = xp.linspace(0, 1, 200) + e2 = xp.array([0.5]) + e3 = xp.array([0.5]) + + for label, var in [('ion', model.ions.u), ('electron', model.electrons.u)]: + if var.spline_lift is None: + continue + + def _eval(fn, comp=0): + return fn(e1, e2, e3, squeeze_out=True)[comp] + + fig, axes = plt.subplots(1, 3, figsize=(12, 4)) + axes[0].plot(e1, _eval(var.spline_lift)); axes[0].set_title(f'{label}: spline_lift') + axes[1].plot(e1, _eval(var.spline_0)); axes[1].set_title(f'{label}: spline_0') + axes[2].plot(e1, _eval(var.boundary_spline)); axes[2].set_title(f'{label}: boundary_spline') + for ax in axes: + ax.set_xlabel('x'); ax.grid(True) + plt.tight_layout() + plt.savefig(f'{name}/plots/lifting_{label}.png', dpi=300) + plt.clf() + + # ---- source diagnostics ---- + prop = model.propagators.qn_full + e1 = xp.linspace(0, 1, 200) + e2 = xp.array([0.5]) + e3 = xp.array([0.5]) + zeros_e = xp.zeros_like(e1) + + for label, spline, src_fn, comp in [ + ('ion_source_x', prop._src_u, prop.options.source_u, 0), + ('electron_source_x', prop._src_ue, prop.options.source_ue, 0), + ]: + if spline is None: + print(f" {label}: None, skipping") + continue + vals_proj = spline(e1, e2, e3, squeeze_out=True)[comp] + vals_ref = src_fn(e1, zeros_e, zeros_e)[comp] + plt.figure(figsize=(8, 4)) + plt.plot(e1, vals_ref, '--', label='analytical') + plt.plot(e1, vals_proj, '-', label='projected (FE)') + plt.xlabel('x'); plt.title(f'{label}'); plt.legend(); plt.grid(True) + plt.savefig(f'{name}/plots/source_{label}.png', dpi=300) + plt.close() + print(f" -> saved {name}/plots/source_{label}.png") \ No newline at end of file diff --git a/examples/TwoFluidQuasiNeutralToy/2D_Verification.py b/examples/TwoFluidQuasiNeutralToy/2D_Verification.py new file mode 100644 index 000000000..85b671092 --- /dev/null +++ b/examples/TwoFluidQuasiNeutralToy/2D_Verification.py @@ -0,0 +1,350 @@ +from cunumpy import pi, cos, sin, zeros_like, ones_like +from struphy.io.options import EnvironmentOptions, BaseUnits, Time +from struphy.geometry import domains +from struphy.fields_background import equils +from struphy.topology import grids +from struphy.io.options import DerhamOptions +from struphy.initial import perturbations +from struphy.initial.base import GenericPerturbation +from struphy import Simulation +from struphy.linear_algebra.solver import SolverParameters + +import argparse +import os +import glob +import cunumpy as xp +import matplotlib.pyplot as plt + +from mpi4py import MPI + +from struphy.models.two_fluid_quasi_neutral_toy import TwoFluidQuasiNeutralToy + +# ------------------ args ------------------ +parser = argparse.ArgumentParser() +parser.add_argument('bc', choices=['periodic', 'dirichlet_inhom', 'neumann']) +args = parser.parse_args() +BC = args.bc + +name = f"runs/sim_2D_{BC}" + +# ------------------ setup ------------------ +env = EnvironmentOptions(sim_folder=name) + +B0 = 1 +nu = 10.0 +nu_e = 1.0 +Nel = (8, 8, 1) +p = (2, 2, 1) +epsilon = 1.0 +dt = 1 +Tend = 1 +sigma = 1 +tol = 1e-5 + +time_opts = Time(dt=dt, Tend=Tend) +domain = domains.Cuboid() +equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) +grid = grids.TensorProductGrid(num_elements=Nel) + +# ------------------ boundary conditions ------------------ +if BC == 'periodic': + derham_opts = DerhamOptions(degree=p, bcs=(None, None, None)) + +elif BC == 'neumann': + derham_opts = DerhamOptions( + degree=p, + bcs=(None, + ("free", "free"), + None) + ) + +elif BC == 'dirichlet_inhom': + derham_opts = DerhamOptions( + degree=p, + bcs=(("dirichlet", "dirichlet"), + ("dirichlet", "dirichlet"), + None) + ) + + lifting_function_u = [ + GenericPerturbation(lambda x, y, z: - x * y, comp=0, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: -xp.cos(2*pi*x)*xp.cos(2*pi*y) - x*y, comp=1, given_in_basis="physical"), + ] + lifting_function_ue = [ + GenericPerturbation(lambda x, y, z: - x * y - 1, comp=0, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: -xp.cos(4*pi*x)*xp.cos(4*pi*y) - x*y - 1, comp=1, given_in_basis="physical"), + ] + +# ------------------ manufactured solutions ------------------ +if BC == 'periodic': + def mms_phi(x, y, z): + return xp.cos(2*pi*x) + xp.sin(2*pi*y), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return -xp.sin(2*pi*x)*xp.sin(2*pi*y), -xp.cos(2*pi*x)*xp.cos(2*pi*y), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return -xp.sin(4*pi*x)*xp.sin(4*pi*y), -xp.cos(4*pi*x)*xp.cos(4*pi*y), xp.zeros_like(x) + +elif BC == 'neumann': + def mms_phi(x, y, z): + return xp.cos(2*pi*x) + xp.sin(2*pi*y), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return -xp.sin(2*pi*x)*xp.sin(2*pi*y), -xp.cos(2*pi*y), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return -xp.sin(2*pi*x)*xp.sin(2*pi*y), -xp.cos(2*pi*y), xp.zeros_like(x) + +elif BC == 'dirichlet_inhom': + def mms_phi(x, y, z): + return xp.cos(2*pi*x) + xp.sin(2*pi*y), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return -xp.sin(2*pi*x)*xp.sin(2*pi*y) - x*y, -xp.cos(2*pi*x)*xp.cos(2*pi*y) - x*y, xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return -xp.sin(4*pi*x)*xp.sin(4*pi*y) - x*y - 1, -xp.cos(4*pi*x)*xp.cos(4*pi*y) - x*y - 1, xp.zeros_like(x) + +# ------------------ source terms ------------------ +if BC == 'periodic': + def source_function_u(x, y, z): + fx = -2*pi*xp.sin(2*pi*x) + B0/epsilon*xp.cos(2*pi*x)*xp.cos(2*pi*y) - nu*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) + fy = 2*pi*xp.cos(2*pi*y) - B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) - nu*8*pi**2*xp.cos(2*pi*x)*xp.cos(2*pi*y) + return fx, fy, zeros_like(x) + + def source_function_ue(x, y, z): + fx = 2*pi*xp.sin(2*pi*x) - B0/epsilon*xp.cos(4*pi*x)*xp.cos(4*pi*y) - nu_e*32*pi**2*xp.sin(4*pi*x)*xp.sin(4*pi*y) + sigma*xp.sin(4*pi*x)*xp.sin(4*pi*y) + fy = -2*pi*xp.cos(2*pi*y) + B0/epsilon*xp.sin(4*pi*x)*xp.sin(4*pi*y) - nu_e*32*pi**2*xp.cos(4*pi*x)*xp.cos(4*pi*y) + sigma*xp.cos(4*pi*x)*xp.cos(4*pi*y) + return fx, fy, zeros_like(x) + +elif BC == 'neumann': + def source_function_u(x, y, z): + fx = B0/epsilon*xp.cos(2*pi*y) - 2*pi*xp.sin(2*pi*x) - nu*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) + fy = 2*pi*xp.cos(2*pi*y) - nu*4*pi**2*xp.cos(2*pi*y) - B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) + return fx, fy, zeros_like(x) + + def source_function_ue(x, y, z): + fx = -B0/epsilon*xp.cos(2*pi*y) + 2*pi*xp.sin(2*pi*x) - nu_e*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) + sigma*xp.sin(2*pi*x)*xp.sin(2*pi*y) + fy = -2*pi*xp.cos(2*pi*y) - nu_e*4*pi**2*xp.cos(2*pi*y) + sigma*xp.cos(2*pi*y) + B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) + return fx, fy, zeros_like(x) + + +elif BC == 'dirichlet_inhom': + def source_function_u(x, y, z): + fx = -2*pi*xp.sin(2*pi*x) + B0/epsilon*xp.cos(2*pi*x)*xp.cos(2*pi*y) - nu*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) + fy = 2*pi*xp.cos(2*pi*y) - B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) - nu*8*pi**2*xp.cos(2*pi*x)*xp.cos(2*pi*y) + fx += - B0 / epsilon * x * y + fy += B0 * x * y / epsilon + return fx, fy, zeros_like(x) + + def source_function_ue(x, y, z): + fx = 2*pi*xp.sin(2*pi*x) - B0/epsilon*xp.cos(4*pi*x)*xp.cos(4*pi*y) - nu_e*32*pi**2*xp.sin(4*pi*x)*xp.sin(4*pi*y) + sigma*xp.sin(4*pi*x)*xp.sin(4*pi*y) + fy = -2*pi*xp.cos(2*pi*y) + B0/epsilon*xp.sin(4*pi*x)*xp.sin(4*pi*y) - nu_e*32*pi**2*xp.cos(4*pi*x)*xp.cos(4*pi*y) + sigma*xp.cos(4*pi*x)*xp.cos(4*pi*y) + fx += B0 / epsilon * (x * y + 1) - sigma * (x*y + 1) + fy += - B0 / epsilon * (x * y + 1) - sigma * (x*y + 1) + return fx, fy, zeros_like(x) + + + +class MMSIonVelocity(perturbations.Perturbation): + def __init__(self, comp=0): + self.comp = comp + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_ion_u(x, y, z)[self.comp] + + +class MMSElectronVelocity(perturbations.Perturbation): + def __init__(self, comp=0): + self.comp = comp + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_electron_u(x, y, z)[self.comp] + + +class MMSPotential(perturbations.Perturbation): + def __init__(self): + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_phi(x, y, z)[0] + +# ------------------ model ------------------ +model = TwoFluidQuasiNeutralToy() + +model.propagators.qn_full.options = model.propagators.qn_full.Options( + nu=nu, + nu_e=nu_e, + eps_norm=epsilon, + stab_sigma=sigma, + source_u=source_function_u, + source_ue=source_function_ue, + solver='gmres', + solver_params=SolverParameters(verbose=True, info=True, tol=tol), +) + +if BC == 'dirichlet_inhom': + model.ions.u.lifting_function = lifting_function_u + model.electrons.u.lifting_function = lifting_function_ue + + # model.ions.u.add_perturbation(MMSIonVelocity(comp=0)) + # model.ions.u.add_perturbation(MMSIonVelocity(comp=1)) + # model.electrons.u.add_perturbation(MMSElectronVelocity(comp=0)) + # model.electrons.u.add_perturbation(MMSElectronVelocity(comp=1)) + # model.em_fields.phi.add_perturbation(MMSPotential()) + +# ------------------ simulation ------------------ +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------ run ------------------ +if __name__ == "__main__": + sim.run(verbose=True) + + if MPI.COMM_WORLD.Get_rank() == 0: + + sim.pproc(verbose=True) + sim.load_plotting_data(verbose=True) + + simdata = sim.plotting_data + + n1_vals = simdata.grids_log[0] + n2_vals = simdata.grids_log[1] + X, Y = xp.meshgrid(n1_vals, n2_vals, indexing='ij') + + x = xp.linspace(0, 1, 100) + y = xp.linspace(0, 1, 100) + Xf, Yf = xp.meshgrid(x, y, indexing='ij') + + os.makedirs(f'{name}/plots', exist_ok=True) + for f in glob.glob(f'{name}/plots/*.png'): + os.remove(f) + + def save_plot(numerical, analytical, title, fname, t): + fig, axes = plt.subplots(1, 2, figsize=(10, 4)) + im0 = axes[0].contourf(X, Y, numerical, levels=50) + axes[0].set_title('numerical') + plt.colorbar(im0, ax=axes[0]) + im1 = axes[1].contourf(Xf, Yf, analytical, levels=50) + axes[1].set_title('manufactured') + plt.colorbar(im1, ax=axes[1]) + fig.suptitle(f'{title} at t={t:.3f}') + plt.savefig(f'{name}/plots/{fname}_{t:.3f}.png', dpi=300) + plt.close(fig) + + for t in simdata.spline_values.ions.u_log.data.keys(): + u_ions = simdata.spline_values.ions.u_log.data[t] + u_electrons = simdata.spline_values.electrons.u_log.data[t] + phi = simdata.spline_values.em_fields.phi_log.data[t] + + mms_phi_x, _, _ = mms_phi(Xf, Yf, 0*Xf) + mms_ion_ux, mms_ion_uy, _ = mms_ion_u(Xf, Yf, 0*Xf) + mms_el_ux, mms_el_uy, _ = mms_electron_u(Xf, Yf, 0*Xf) + + save_plot(phi[0][:, :, 0], mms_phi_x, 'φ', 'plot_phi', t) + save_plot(u_ions[0][:, :, 0], mms_ion_ux, 'u_ix', 'plot_uix', t) + save_plot(u_ions[1][:, :, 0], mms_ion_uy, 'u_iy', 'plot_uiy', t) + save_plot(u_electrons[0][:, :, 0], mms_el_ux, 'u_ex', 'plot_uex', t) + save_plot(u_electrons[1][:, :, 0], mms_el_uy, 'u_ey', 'plot_uey', t) + + # ---- lifting diagnostics (dirichlet_inhom only) ---- + if BC == 'dirichlet_inhom': + e1 = xp.linspace(0, 1, 80) + e2 = xp.linspace(0, 1, 80) + e3 = xp.array([0.5]) + E1, E2 = xp.meshgrid(e1, e2, indexing='ij') + + for label, var, comp in [('ion_ux', model.ions.u, 0), + ('ion_uy', model.ions.u, 1), + ('electron_ux', model.electrons.u, 0), + ('electron_uy', model.electrons.u, 1)]: + if var.spline_lift is None: + print(f" {label}: spline_lift is None, skipping") + continue + + def _eval(fn): + return fn(e1, e2, e3, squeeze_out=True)[comp] + + fig, axes = plt.subplots(1, 3, figsize=(15, 4)) + for ax, fn, ttl in zip(axes, + [var.spline_lift, var.spline_0, var.boundary_spline], + ['lifting', 'zero-BC part', 'boundary spline']): + im = ax.contourf(E1, E2, _eval(fn), levels=50) + ax.set_title(f'{label}: {ttl}') + plt.colorbar(im, ax=ax) + out = f'{name}/plots/lifting_{label}.png' + plt.savefig(out, dpi=300) + plt.close(fig) + print(f" -> saved {out}") + + # ---- source diagnostics ---- + prop = model.propagators.qn_full + e1 = xp.linspace(0, 1, 80) + e2 = xp.linspace(0, 1, 80) + e3 = xp.array([0.5]) + E1, E2 = xp.meshgrid(e1, e2, indexing='ij') + zeros_E = xp.zeros_like(E1) + + for label, spline, src_fn, comp in [ + ('ion_source_x', prop._src_u, prop.options.source_u, 0), + ('ion_source_y', prop._src_u, prop.options.source_u, 1), + ('electron_source_x', prop._src_ue, prop.options.source_ue, 0), + ('electron_source_y', prop._src_ue, prop.options.source_ue, 1), + ]: + if spline is None: + print(f" {label}: None, skipping") + continue + + vals_proj = spline(e1, e2, e3, squeeze_out=True)[comp] + vals_ref = src_fn(E1, E2, zeros_E)[comp] + + fig, axes = plt.subplots(1, 2, figsize=(10, 4)) + im0 = axes[0].contourf(E1, E2, vals_proj, levels=50) + axes[0].set_title('projected (FE)') + plt.colorbar(im0, ax=axes[0]) + im1 = axes[1].contourf(E1, E2, vals_ref, levels=50) + axes[1].set_title('reference (analytical)') + plt.colorbar(im1, ax=axes[1]) + fig.suptitle(label) + out = f'{name}/plots/source_{label}.png' + plt.savefig(out, dpi=300) + plt.close(fig) + print(f" -> saved {out}") + + if BC == 'dirichlet_inhom': + y_check = xp.linspace(0, 1, 80) + x_check = xp.linspace(0, 1, 80) + z_check = xp.array([0.5]) + + # comp=0: normal trace at x=0 and x=1 + for x_bnd, label in [(0.0, 'x=0'), (1.0, 'x=1')]: + x_bnd_arr = xp.array([x_bnd]) + mms_vals = mms_ion_u(x_bnd_arr, y_check, z_check)[0] + lift_vals = model.ions.u.spline_lift(x_bnd_arr, y_check, z_check, squeeze_out=True)[0] + print(f"ion ux normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") + + mms_vals = mms_electron_u(x_bnd_arr, y_check, z_check)[0] + lift_vals = model.electrons.u.spline_lift(x_bnd_arr, y_check, z_check, squeeze_out=True)[0] + print(f"elec ux normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") + + # comp=1: normal trace at y=0 and y=1 + for y_bnd, label in [(0.0, 'y=0'), (1.0, 'y=1')]: + y_bnd_arr = xp.array([y_bnd]) + mms_vals = mms_ion_u(x_check, y_bnd_arr, z_check)[1] + lift_vals = model.ions.u.spline_lift(x_check, y_bnd_arr, z_check, squeeze_out=True)[1] + print(f"ion uy normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") + + mms_vals = mms_electron_u(x_check, y_bnd_arr, z_check)[1] + lift_vals = model.electrons.u.spline_lift(x_check, y_bnd_arr, z_check, squeeze_out=True)[1] + print(f"elec uy normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") \ No newline at end of file diff --git a/examples/TwoFluidQuasiNeutralToy/R_Verification.py b/examples/TwoFluidQuasiNeutralToy/R_Verification.py new file mode 100644 index 000000000..f94380a24 --- /dev/null +++ b/examples/TwoFluidQuasiNeutralToy/R_Verification.py @@ -0,0 +1,217 @@ +import os +import glob +import cunumpy as xp +import matplotlib.pyplot as plt +from mpi4py import MPI + +from struphy.io.options import EnvironmentOptions, Time +from struphy.geometry import domains +from struphy.fields_background import equils +from struphy.topology import grids +from struphy.io.options import DerhamOptions +from struphy.initial.base import GenericPerturbation +from struphy import Simulation +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.two_fluid_quasi_neutral_toy import TwoFluidQuasiNeutralToy + +# ------------------ parameters ------------------ +name = "runs/sim_restelli" + +R0 = 2.0 +a = 1.0 +ain = 0.1 +B0 = 10.0 +Bp = 12.5 +nu = 1.0 +nu_e = 0.01 +alpha = 0.1 +beta = 1.0 +eps = 1.0 +sigma = 1e-6 +dt = 1 +Tend = 1 +Nel = (8, 8, 1) +p = (1, 1, 1) +tol = 1e-5 + +env = EnvironmentOptions(sim_folder=name) +time_opts = Time(dt=dt, Tend=Tend) + +# ------------------ domain & equilibrium ------------------ +domain = domains.HollowTorus(a1=ain, a2=a, R0=R0) +equil = equils.CircularTokamak(a=a, R0=R0, B0=B0, Bp=Bp) +grid = grids.TensorProductGrid(num_elements=Nel) + +derham_opts = DerhamOptions( + degree=p, + bcs=(("dirichlet", "dirichlet"), None, None) +) + +# ------------------ manufactured solution ------------------ + +def _cylindrical(x, y, z): + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) + phi = xp.arctan2(-y, x) + Z = z + return R, phi, Z + +def mms_u_cartesian(x, y, z): + R, phi, Z = _cylindrical(x, y, z) + u_R = alpha/R * (-Z) / (a*R0/R) + beta * Bp/B0 * R0/(a*R) * Z + u_Z = alpha/R * (R - R0) / (a*R0/R) + beta * Bp/B0 * R0/(a*R) * (-(R - R0)) + u_phi = beta * Bp/B0 * R0/(a*R) * B0*Bp*a / (Bp*R0) + + u_R = alpha * R / (a*R0) * (-Z) + beta * Bp/B0 * R0/(a*R) * Z + u_Z = alpha * R / (a*R0) * (R - R0) + beta * Bp/B0 * R0/(a*R) * (-(R - R0)) + u_phi = beta * Bp/B0 * R0/(a*R) * (B0/Bp * a) + + # Transform to Cartesian via eq (5.14) + ux = xp.cos(phi) * u_R - R * xp.sin(phi) * u_phi + uy = -xp.sin(phi) * u_R - R * xp.cos(phi) * u_phi + uz = u_Z + return ux, uy, uz + +def mms_phi_cartesian(x, y, z): + R, phi, Z = _cylindrical(x, y, z) + # eq (5.22): phi_hat = 0.5 * a * B0 * alpha * ((R-R0)^2 + Z^2)/a^2 - 2/3) + phi_val = 0.5 * a * B0 * alpha * (((R - R0)**2 + Z**2) / a**2 - 2.0/3.0) + return phi_val + +# ------------------ source terms ------------------ + +def _omega_cartesian(x, y, z): + R, phi, Z = _cylindrical(x, y, z) + omega_Z = alpha * (R0 - 4*R) / (a*R0*R) - beta * Bp/B0 * R0**2 / (a*R**3) + ox = xp.zeros_like(x) + oy = xp.zeros_like(x) + oz = omega_Z + return ox, oy, oz + +def source_function_u(x, y, z): + ox, oy, oz = _omega_cartesian(x, y, z) + return nu * ox, nu * oy, nu * oz + +def source_function_ue(x, y, z): + ox, oy, oz = _omega_cartesian(x, y, z) + return nu_e * ox, nu_e * oy, nu_e * oz + +# ------------------ lifting (inhomogeneous Dirichlet on radial boundary) ------------------ +lifting_u = [ + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[0], comp=0, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[1], comp=1, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[2], comp=2, given_in_basis="physical"), +] + +lifting_ue = [ + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[0], comp=0, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[1], comp=1, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[2], comp=2, given_in_basis="physical"), +] + +# ------------------ model ------------------ +model = TwoFluidQuasiNeutralToy() + +model.propagators.qn_full.options = model.propagators.qn_full.Options( + nu=nu, + nu_e=nu_e, + eps_norm=eps, + stab_sigma=sigma, + source_u=source_function_u, + source_ue=source_function_ue, + solver='gmres', + solver_params=SolverParameters(verbose=True, info=True, tol=tol), +) + +model.ions.u.lifting_function = lifting_u +model.electrons.u.lifting_function = lifting_ue + +# ------------------ simulation ------------------ +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------ run ------------------ +if __name__ == "__main__": + # sim.run(verbose=True) + + if MPI.COMM_WORLD.Get_rank() == 0: + sim.pproc(verbose=True) + sim.load_plotting_data(verbose=True) + + simdata = sim.plotting_data + os.makedirs(f'{name}/plots', exist_ok=True) + for f in glob.glob(f'{name}/plots/*.png'): + os.remove(f) + + e1 = xp.linspace(0, 1, 40) + e2 = xp.linspace(0, 1, 40) + e3 = xp.array([0.5]) + x_phys, y_phys, z_phys = domain(e1, e2, e3, squeeze_out=True) + + theta_bnd = xp.linspace(0, 1, 200) + x_inner, y_inner, _ = domain(xp.array([0.0]), theta_bnd, xp.array([0.5]), squeeze_out=True) + x_outer, y_outer, _ = domain(xp.array([1.0]), theta_bnd, xp.array([0.5]), squeeze_out=True) + + def _add_domain_boundary(ax): + ax.plot(x_inner, y_inner, 'w-', linewidth=0.8) + ax.plot(x_outer, y_outer, 'w-', linewidth=0.8) + + prop = model.propagators.qn_full + + for label, spline, src_fn, comp in [ + ('ion_source_x', prop._src_u, prop.options.source_u, 0), + ('ion_source_y', prop._src_u, prop.options.source_u, 1), + ('ion_source_z', prop._src_u, prop.options.source_u, 2), + ('electron_source_x', prop._src_ue, prop.options.source_ue, 0), + ('electron_source_y', prop._src_ue, prop.options.source_ue, 1), + ('electron_source_z', prop._src_ue, prop.options.source_ue, 2), + ]: + if spline is None: + print(f" {label}: None, skipping") + continue + + vals_proj = spline(e1, e2, e3, squeeze_out=True)[comp] + vals_ref = src_fn(x_phys, y_phys, z_phys)[comp] + + fig, axes = plt.subplots(1, 2, figsize=(10, 4)) + im0 = axes[0].contourf(x_phys, y_phys, vals_proj, levels=50); axes[0].set_title('projected (FE)'); plt.colorbar(im0, ax=axes[0]); _add_domain_boundary(axes[0]) + im1 = axes[1].contourf(x_phys, y_phys, vals_ref, levels=50); axes[1].set_title('reference (analytical)'); plt.colorbar(im1, ax=axes[1]); _add_domain_boundary(axes[1]) + fig.suptitle(label) + out = f'{name}/plots/source_{label}.png' + plt.savefig(out, dpi=300) + plt.close(fig) + print(f" -> saved {out}") + + for t in simdata.spline_values.ions.u_log.data.keys(): + u_ions = simdata.spline_values.ions.u_log.data[t] + u_electrons = simdata.spline_values.electrons.u_log.data[t] + phi_num = simdata.spline_values.em_fields.phi_log.data[t] + + mms_ux, mms_uy, mms_uz = mms_u_cartesian(x_phys, y_phys, z_phys) + mms_phi_val = mms_phi_cartesian(x_phys, y_phys, z_phys) + + for num, mms, lbl in [ + (u_ions[0][:, :, 0], mms_ux, 'u_ix'), + (u_ions[1][:, :, 0], mms_uy, 'u_iy'), + (u_ions[2][:, :, 0], mms_uz, 'u_iz'), + (u_electrons[0][:, :, 0], mms_ux, 'u_ex'), + (u_electrons[1][:, :, 0], mms_uy, 'u_ey'), + (u_electrons[2][:, :, 0], mms_uz, 'u_ez'), + (phi_num[0][:, :, 0], mms_phi_val, 'phi'), + ]: + fig, axes = plt.subplots(1, 3, figsize=(15, 4)) + im0 = axes[0].contourf(x_phys, y_phys, num, levels=50); axes[0].set_title('numerical'); plt.colorbar(im0, ax=axes[0]); _add_domain_boundary(axes[0]) + im1 = axes[1].contourf(x_phys, y_phys, mms, levels=50); axes[1].set_title('MMS'); plt.colorbar(im1, ax=axes[1]); _add_domain_boundary(axes[1]) + im2 = axes[2].contourf(x_phys, y_phys, num - mms, levels=50); axes[2].set_title('difference'); plt.colorbar(im2, ax=axes[2]); _add_domain_boundary(axes[2]) + fig.suptitle(f'{lbl} at t={t:.4f}') + out = f'{name}/plots/{lbl}_{t:.4f}.png' + plt.savefig(out, dpi=300) + plt.close(fig) \ No newline at end of file diff --git a/feectools b/feectools index 8c88dec79..ae6859fcb 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 8c88dec79510b315024d4b7e0ccc28e76ad8c9e7 +Subproject commit ae6859fcb16c765f7bb7cabb82eb6268d1967014 From d5f7dc66538b884b4354a94eb210280e03a5eff0 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 20 May 2026 08:43:42 +0200 Subject: [PATCH 32/41] ran ruff format examples/TwoFluidQuasiNeutralToy --- .../1D_Verification.py | 135 ++++---- .../2D_Verification.py | 292 +++++++++++------- .../TwoFluidQuasiNeutralToy/R_Verification.py | 146 +++++---- 3 files changed, 337 insertions(+), 236 deletions(-) diff --git a/examples/TwoFluidQuasiNeutralToy/1D_Verification.py b/examples/TwoFluidQuasiNeutralToy/1D_Verification.py index 0e1df71c2..cfe382d20 100644 --- a/examples/TwoFluidQuasiNeutralToy/1D_Verification.py +++ b/examples/TwoFluidQuasiNeutralToy/1D_Verification.py @@ -18,44 +18,45 @@ from struphy.models.two_fluid_quasi_neutral_toy import TwoFluidQuasiNeutralToy parser = argparse.ArgumentParser() -parser.add_argument('bc', choices=['periodic', 'dirichlet_hom', 'dirichlet_inhom']) +parser.add_argument("bc", choices=["periodic", "dirichlet_hom", "dirichlet_inhom"]) args = parser.parse_args() BC = args.bc name = f"runs/sim_1D_{BC}" -env = EnvironmentOptions(sim_folder=name) +env = EnvironmentOptions(sim_folder=name) -B0 = 0 -nu = 10.0 -nu_e = 1.0 -Nel = (32, 1, 1) -p = (1, 1, 1) +B0 = 0 +nu = 10.0 +nu_e = 1.0 +Nel = (32, 1, 1) +p = (1, 1, 1) epsilon = 1.0 -dt = 1 -Tend = 1 -sigma = 0 +dt = 1 +Tend = 1 +sigma = 0 tol = 1e-5 time_opts = Time(dt=dt, Tend=Tend) -domain = domains.Cuboid() -equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) -grid = grids.TensorProductGrid(num_elements=Nel) +domain = domains.Cuboid() +equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) +grid = grids.TensorProductGrid(num_elements=Nel) # ---- boundary conditions ---- -if BC == 'periodic': +if BC == "periodic": derham_opts = DerhamOptions(degree=p, bcs=(None, None, None)) -elif BC == 'dirichlet_hom': +elif BC == "dirichlet_hom": derham_opts = DerhamOptions(degree=p, bcs=(("dirichlet", "dirichlet"), None, None)) -elif BC == 'dirichlet_inhom': +elif BC == "dirichlet_inhom": derham_opts = DerhamOptions(degree=p, bcs=(("dirichlet", "dirichlet"), None, None)) - lifting_function_u = GenericPerturbation(lambda x, y, z: x + 1, comp=0, given_in_basis="physical") + lifting_function_u = GenericPerturbation(lambda x, y, z: x + 1, comp=0, given_in_basis="physical") lifting_function_ue = GenericPerturbation(lambda x, y, z: x, comp=0, given_in_basis="physical") # ---- manufactured solutions ---- -if BC == 'periodic': +if BC == "periodic": + def mms_phi(x, y, z): return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) @@ -65,7 +66,8 @@ def mms_ion_u(x, y, z): def mms_electron_u(x, y, z): return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) -elif BC == 'dirichlet_hom': +elif BC == "dirichlet_hom": + def mms_phi(x, y, z): return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) @@ -75,7 +77,8 @@ def mms_ion_u(x, y, z): def mms_electron_u(x, y, z): return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) -elif BC == 'dirichlet_inhom': +elif BC == "dirichlet_inhom": + def mms_phi(x, y, z): return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) @@ -85,8 +88,10 @@ def mms_ion_u(x, y, z): def mms_electron_u(x, y, z): return xp.sin(2 * xp.pi * x) + x, xp.zeros_like(x), xp.zeros_like(x) + # ---- source terms ---- -if BC == 'periodic': +if BC == "periodic": + def source_function_u(x, y, z): fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) fy = zeros_like(x) @@ -99,7 +104,8 @@ def source_function_ue(x, y, z): fz = zeros_like(x) return fx, fy, fz -elif BC == 'dirichlet_hom': +elif BC == "dirichlet_hom": + def source_function_u(x, y, z): fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) fy = zeros_like(x) @@ -112,7 +118,8 @@ def source_function_ue(x, y, z): fz = zeros_like(x) return fx, fy, fz -elif BC == 'dirichlet_inhom': +elif BC == "dirichlet_inhom": + def source_function_u(x, y, z): fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) fy = zeros_like(x) @@ -125,6 +132,7 @@ def source_function_ue(x, y, z): fz = zeros_like(x) return fx, fy, fz + # ---- perturbation classes for MMS initial conditions ---- class MMSIonVelocity(perturbations.Perturbation): def __init__(self, comp=0): @@ -150,7 +158,7 @@ def __init__(self): def __call__(self, x, y, z): return mms_phi(x, y, z)[0] - + # ---- model ---- model = TwoFluidQuasiNeutralToy() @@ -162,12 +170,12 @@ def __call__(self, x, y, z): stab_sigma=sigma, source_u=source_function_u, source_ue=source_function_ue, - solver='gmres', - solver_params=SolverParameters(verbose=True, info=True, tol = tol), + solver="gmres", + solver_params=SolverParameters(verbose=True, info=True, tol=tol), ) -if BC == 'dirichlet_inhom': - model.ions.u.lifting_function = lifting_function_u +if BC == "dirichlet_inhom": + model.ions.u.lifting_function = lifting_function_u model.electrons.u.lifting_function = lifting_function_ue sim = Simulation( @@ -188,44 +196,44 @@ def __call__(self, x, y, z): simdata = sim.plotting_data n1_vals = simdata.grids_log[0] - x = xp.linspace(0, 1, 100) + x = xp.linspace(0, 1, 100) - os.makedirs(f'{name}/plots', exist_ok=True) - for f in glob.glob(f'{name}/plots/*.png'): + os.makedirs(f"{name}/plots", exist_ok=True) + for f in glob.glob(f"{name}/plots/*.png"): os.remove(f) def save_plot(n1_vals, numerical, analytical, ylabel, title, fname, t): - plt.plot(n1_vals, numerical, label='numerical') - plt.plot(x, analytical, '--', label='manufactured') - plt.plot(n1_vals, numerical, 'k.', markersize=4, label='n1 points') - plt.xlabel('x') + plt.plot(n1_vals, numerical, label="numerical") + plt.plot(x, analytical, "--", label="manufactured") + plt.plot(n1_vals, numerical, "k.", markersize=4, label="n1 points") + plt.xlabel("x") plt.ylabel(ylabel) - plt.title(f'{title} at t={t:.3f}') + plt.title(f"{title} at t={t:.3f}") plt.legend() plt.grid(True) - plt.savefig(f'{name}/plots/{fname}_{t:.3f}.png', dpi=300) + plt.savefig(f"{name}/plots/{fname}_{t:.3f}.png", dpi=300) plt.clf() for t in list(simdata.spline_values.ions.u_log.data.keys()): - u_ions = simdata.spline_values.ions.u_log.data[t] + u_ions = simdata.spline_values.ions.u_log.data[t] u_electrons = simdata.spline_values.electrons.u_log.data[t] - phi = simdata.spline_values.em_fields.phi_log.data[t] + phi = simdata.spline_values.em_fields.phi_log.data[t] - mms_phi_x, _, _ = mms_phi(x, x*0, x*0) - mms_ion_ux, _, _ = mms_ion_u(x, x*0, x*0) - mms_el_ux, _, _ = mms_electron_u(x, x*0, x*0) + mms_phi_x, _, _ = mms_phi(x, x * 0, x * 0) + mms_ion_ux, _, _ = mms_ion_u(x, x * 0, x * 0) + mms_el_ux, _, _ = mms_electron_u(x, x * 0, x * 0) - save_plot(n1_vals, phi[0][:, 0, 0], mms_phi_x, 'φ', 'Potential φ', 'plot_potential', t) - save_plot(n1_vals, u_ions[0][:, 0, 0], mms_ion_ux, 'u_x', 'Ion velocity u_x', 'plot_ion_ux', t) - save_plot(n1_vals, u_electrons[0][:, 0, 0], mms_el_ux, 'u_x', 'Electron velocity', 'plot_electron_ux', t) + save_plot(n1_vals, phi[0][:, 0, 0], mms_phi_x, "φ", "Potential φ", "plot_potential", t) + save_plot(n1_vals, u_ions[0][:, 0, 0], mms_ion_ux, "u_x", "Ion velocity u_x", "plot_ion_ux", t) + save_plot(n1_vals, u_electrons[0][:, 0, 0], mms_el_ux, "u_x", "Electron velocity", "plot_electron_ux", t) # ---- lifting diagnostics ---- - if BC == 'dirichlet_inhom': + if BC == "dirichlet_inhom": e1 = xp.linspace(0, 1, 200) e2 = xp.array([0.5]) e3 = xp.array([0.5]) - for label, var in [('ion', model.ions.u), ('electron', model.electrons.u)]: + for label, var in [("ion", model.ions.u), ("electron", model.electrons.u)]: if var.spline_lift is None: continue @@ -233,13 +241,17 @@ def _eval(fn, comp=0): return fn(e1, e2, e3, squeeze_out=True)[comp] fig, axes = plt.subplots(1, 3, figsize=(12, 4)) - axes[0].plot(e1, _eval(var.spline_lift)); axes[0].set_title(f'{label}: spline_lift') - axes[1].plot(e1, _eval(var.spline_0)); axes[1].set_title(f'{label}: spline_0') - axes[2].plot(e1, _eval(var.boundary_spline)); axes[2].set_title(f'{label}: boundary_spline') + axes[0].plot(e1, _eval(var.spline_lift)) + axes[0].set_title(f"{label}: spline_lift") + axes[1].plot(e1, _eval(var.spline_0)) + axes[1].set_title(f"{label}: spline_0") + axes[2].plot(e1, _eval(var.boundary_spline)) + axes[2].set_title(f"{label}: boundary_spline") for ax in axes: - ax.set_xlabel('x'); ax.grid(True) + ax.set_xlabel("x") + ax.grid(True) plt.tight_layout() - plt.savefig(f'{name}/plots/lifting_{label}.png', dpi=300) + plt.savefig(f"{name}/plots/lifting_{label}.png", dpi=300) plt.clf() # ---- source diagnostics ---- @@ -250,18 +262,21 @@ def _eval(fn, comp=0): zeros_e = xp.zeros_like(e1) for label, spline, src_fn, comp in [ - ('ion_source_x', prop._src_u, prop.options.source_u, 0), - ('electron_source_x', prop._src_ue, prop.options.source_ue, 0), + ("ion_source_x", prop._src_u, prop.options.source_u, 0), + ("electron_source_x", prop._src_ue, prop.options.source_ue, 0), ]: if spline is None: print(f" {label}: None, skipping") continue vals_proj = spline(e1, e2, e3, squeeze_out=True)[comp] - vals_ref = src_fn(e1, zeros_e, zeros_e)[comp] + vals_ref = src_fn(e1, zeros_e, zeros_e)[comp] plt.figure(figsize=(8, 4)) - plt.plot(e1, vals_ref, '--', label='analytical') - plt.plot(e1, vals_proj, '-', label='projected (FE)') - plt.xlabel('x'); plt.title(f'{label}'); plt.legend(); plt.grid(True) - plt.savefig(f'{name}/plots/source_{label}.png', dpi=300) + plt.plot(e1, vals_ref, "--", label="analytical") + plt.plot(e1, vals_proj, "-", label="projected (FE)") + plt.xlabel("x") + plt.title(f"{label}") + plt.legend() + plt.grid(True) + plt.savefig(f"{name}/plots/source_{label}.png", dpi=300) plt.close() - print(f" -> saved {name}/plots/source_{label}.png") \ No newline at end of file + print(f" -> saved {name}/plots/source_{label}.png") diff --git a/examples/TwoFluidQuasiNeutralToy/2D_Verification.py b/examples/TwoFluidQuasiNeutralToy/2D_Verification.py index 85b671092..5ff6e54e7 100644 --- a/examples/TwoFluidQuasiNeutralToy/2D_Verification.py +++ b/examples/TwoFluidQuasiNeutralToy/2D_Verification.py @@ -21,132 +21,194 @@ # ------------------ args ------------------ parser = argparse.ArgumentParser() -parser.add_argument('bc', choices=['periodic', 'dirichlet_inhom', 'neumann']) +parser.add_argument("bc", choices=["periodic", "dirichlet_inhom", "neumann"]) args = parser.parse_args() BC = args.bc name = f"runs/sim_2D_{BC}" # ------------------ setup ------------------ -env = EnvironmentOptions(sim_folder=name) +env = EnvironmentOptions(sim_folder=name) -B0 = 1 -nu = 10.0 -nu_e = 1.0 -Nel = (8, 8, 1) -p = (2, 2, 1) +B0 = 1 +nu = 10.0 +nu_e = 1.0 +Nel = (8, 8, 1) +p = (2, 2, 1) epsilon = 1.0 -dt = 1 -Tend = 1 -sigma = 1 +dt = 1 +Tend = 1 +sigma = 1 tol = 1e-5 time_opts = Time(dt=dt, Tend=Tend) -domain = domains.Cuboid() -equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) -grid = grids.TensorProductGrid(num_elements=Nel) +domain = domains.Cuboid() +equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) +grid = grids.TensorProductGrid(num_elements=Nel) # ------------------ boundary conditions ------------------ -if BC == 'periodic': +if BC == "periodic": derham_opts = DerhamOptions(degree=p, bcs=(None, None, None)) -elif BC == 'neumann': - derham_opts = DerhamOptions( - degree=p, - bcs=(None, - ("free", "free"), - None) - ) - -elif BC == 'dirichlet_inhom': - derham_opts = DerhamOptions( - degree=p, - bcs=(("dirichlet", "dirichlet"), - ("dirichlet", "dirichlet"), - None) - ) +elif BC == "neumann": + derham_opts = DerhamOptions(degree=p, bcs=(None, ("free", "free"), None)) + +elif BC == "dirichlet_inhom": + derham_opts = DerhamOptions(degree=p, bcs=(("dirichlet", "dirichlet"), ("dirichlet", "dirichlet"), None)) lifting_function_u = [ - GenericPerturbation(lambda x, y, z: - x * y, comp=0, given_in_basis="physical"), - GenericPerturbation(lambda x, y, z: -xp.cos(2*pi*x)*xp.cos(2*pi*y) - x*y, comp=1, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: -x * y, comp=0, given_in_basis="physical"), + GenericPerturbation( + lambda x, y, z: -xp.cos(2 * pi * x) * xp.cos(2 * pi * y) - x * y, comp=1, given_in_basis="physical" + ), ] lifting_function_ue = [ - GenericPerturbation(lambda x, y, z: - x * y - 1, comp=0, given_in_basis="physical"), - GenericPerturbation(lambda x, y, z: -xp.cos(4*pi*x)*xp.cos(4*pi*y) - x*y - 1, comp=1, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: -x * y - 1, comp=0, given_in_basis="physical"), + GenericPerturbation( + lambda x, y, z: -xp.cos(4 * pi * x) * xp.cos(4 * pi * y) - x * y - 1, comp=1, given_in_basis="physical" + ), ] # ------------------ manufactured solutions ------------------ -if BC == 'periodic': +if BC == "periodic": + def mms_phi(x, y, z): - return xp.cos(2*pi*x) + xp.sin(2*pi*y), xp.zeros_like(x), xp.zeros_like(x) + return xp.cos(2 * pi * x) + xp.sin(2 * pi * y), xp.zeros_like(x), xp.zeros_like(x) def mms_ion_u(x, y, z): - return -xp.sin(2*pi*x)*xp.sin(2*pi*y), -xp.cos(2*pi*x)*xp.cos(2*pi*y), xp.zeros_like(x) + return -xp.sin(2 * pi * x) * xp.sin(2 * pi * y), -xp.cos(2 * pi * x) * xp.cos(2 * pi * y), xp.zeros_like(x) def mms_electron_u(x, y, z): - return -xp.sin(4*pi*x)*xp.sin(4*pi*y), -xp.cos(4*pi*x)*xp.cos(4*pi*y), xp.zeros_like(x) - -elif BC == 'neumann': + return -xp.sin(4 * pi * x) * xp.sin(4 * pi * y), -xp.cos(4 * pi * x) * xp.cos(4 * pi * y), xp.zeros_like(x) + +elif BC == "neumann": + def mms_phi(x, y, z): - return xp.cos(2*pi*x) + xp.sin(2*pi*y), xp.zeros_like(x), xp.zeros_like(x) + return xp.cos(2 * pi * x) + xp.sin(2 * pi * y), xp.zeros_like(x), xp.zeros_like(x) def mms_ion_u(x, y, z): - return -xp.sin(2*pi*x)*xp.sin(2*pi*y), -xp.cos(2*pi*y), xp.zeros_like(x) + return -xp.sin(2 * pi * x) * xp.sin(2 * pi * y), -xp.cos(2 * pi * y), xp.zeros_like(x) def mms_electron_u(x, y, z): - return -xp.sin(2*pi*x)*xp.sin(2*pi*y), -xp.cos(2*pi*y), xp.zeros_like(x) + return -xp.sin(2 * pi * x) * xp.sin(2 * pi * y), -xp.cos(2 * pi * y), xp.zeros_like(x) + +elif BC == "dirichlet_inhom": -elif BC == 'dirichlet_inhom': def mms_phi(x, y, z): - return xp.cos(2*pi*x) + xp.sin(2*pi*y), xp.zeros_like(x), xp.zeros_like(x) + return xp.cos(2 * pi * x) + xp.sin(2 * pi * y), xp.zeros_like(x), xp.zeros_like(x) def mms_ion_u(x, y, z): - return -xp.sin(2*pi*x)*xp.sin(2*pi*y) - x*y, -xp.cos(2*pi*x)*xp.cos(2*pi*y) - x*y, xp.zeros_like(x) + return ( + -xp.sin(2 * pi * x) * xp.sin(2 * pi * y) - x * y, + -xp.cos(2 * pi * x) * xp.cos(2 * pi * y) - x * y, + xp.zeros_like(x), + ) def mms_electron_u(x, y, z): - return -xp.sin(4*pi*x)*xp.sin(4*pi*y) - x*y - 1, -xp.cos(4*pi*x)*xp.cos(4*pi*y) - x*y - 1, xp.zeros_like(x) + return ( + -xp.sin(4 * pi * x) * xp.sin(4 * pi * y) - x * y - 1, + -xp.cos(4 * pi * x) * xp.cos(4 * pi * y) - x * y - 1, + xp.zeros_like(x), + ) + # ------------------ source terms ------------------ -if BC == 'periodic': +if BC == "periodic": + def source_function_u(x, y, z): - fx = -2*pi*xp.sin(2*pi*x) + B0/epsilon*xp.cos(2*pi*x)*xp.cos(2*pi*y) - nu*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) - fy = 2*pi*xp.cos(2*pi*y) - B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) - nu*8*pi**2*xp.cos(2*pi*x)*xp.cos(2*pi*y) + fx = ( + -2 * pi * xp.sin(2 * pi * x) + + B0 / epsilon * xp.cos(2 * pi * x) * xp.cos(2 * pi * y) + - nu * 8 * pi**2 * xp.sin(2 * pi * x) * xp.sin(2 * pi * y) + ) + fy = ( + 2 * pi * xp.cos(2 * pi * y) + - B0 / epsilon * xp.sin(2 * pi * x) * xp.sin(2 * pi * y) + - nu * 8 * pi**2 * xp.cos(2 * pi * x) * xp.cos(2 * pi * y) + ) return fx, fy, zeros_like(x) def source_function_ue(x, y, z): - fx = 2*pi*xp.sin(2*pi*x) - B0/epsilon*xp.cos(4*pi*x)*xp.cos(4*pi*y) - nu_e*32*pi**2*xp.sin(4*pi*x)*xp.sin(4*pi*y) + sigma*xp.sin(4*pi*x)*xp.sin(4*pi*y) - fy = -2*pi*xp.cos(2*pi*y) + B0/epsilon*xp.sin(4*pi*x)*xp.sin(4*pi*y) - nu_e*32*pi**2*xp.cos(4*pi*x)*xp.cos(4*pi*y) + sigma*xp.cos(4*pi*x)*xp.cos(4*pi*y) + fx = ( + 2 * pi * xp.sin(2 * pi * x) + - B0 / epsilon * xp.cos(4 * pi * x) * xp.cos(4 * pi * y) + - nu_e * 32 * pi**2 * xp.sin(4 * pi * x) * xp.sin(4 * pi * y) + + sigma * xp.sin(4 * pi * x) * xp.sin(4 * pi * y) + ) + fy = ( + -2 * pi * xp.cos(2 * pi * y) + + B0 / epsilon * xp.sin(4 * pi * x) * xp.sin(4 * pi * y) + - nu_e * 32 * pi**2 * xp.cos(4 * pi * x) * xp.cos(4 * pi * y) + + sigma * xp.cos(4 * pi * x) * xp.cos(4 * pi * y) + ) return fx, fy, zeros_like(x) - -elif BC == 'neumann': + +elif BC == "neumann": + def source_function_u(x, y, z): - fx = B0/epsilon*xp.cos(2*pi*y) - 2*pi*xp.sin(2*pi*x) - nu*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) - fy = 2*pi*xp.cos(2*pi*y) - nu*4*pi**2*xp.cos(2*pi*y) - B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) + fx = ( + B0 / epsilon * xp.cos(2 * pi * y) + - 2 * pi * xp.sin(2 * pi * x) + - nu * 8 * pi**2 * xp.sin(2 * pi * x) * xp.sin(2 * pi * y) + ) + fy = ( + 2 * pi * xp.cos(2 * pi * y) + - nu * 4 * pi**2 * xp.cos(2 * pi * y) + - B0 / epsilon * xp.sin(2 * pi * x) * xp.sin(2 * pi * y) + ) return fx, fy, zeros_like(x) def source_function_ue(x, y, z): - fx = -B0/epsilon*xp.cos(2*pi*y) + 2*pi*xp.sin(2*pi*x) - nu_e*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) + sigma*xp.sin(2*pi*x)*xp.sin(2*pi*y) - fy = -2*pi*xp.cos(2*pi*y) - nu_e*4*pi**2*xp.cos(2*pi*y) + sigma*xp.cos(2*pi*y) + B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) + fx = ( + -B0 / epsilon * xp.cos(2 * pi * y) + + 2 * pi * xp.sin(2 * pi * x) + - nu_e * 8 * pi**2 * xp.sin(2 * pi * x) * xp.sin(2 * pi * y) + + sigma * xp.sin(2 * pi * x) * xp.sin(2 * pi * y) + ) + fy = ( + -2 * pi * xp.cos(2 * pi * y) + - nu_e * 4 * pi**2 * xp.cos(2 * pi * y) + + sigma * xp.cos(2 * pi * y) + + B0 / epsilon * xp.sin(2 * pi * x) * xp.sin(2 * pi * y) + ) return fx, fy, zeros_like(x) -elif BC == 'dirichlet_inhom': +elif BC == "dirichlet_inhom": + def source_function_u(x, y, z): - fx = -2*pi*xp.sin(2*pi*x) + B0/epsilon*xp.cos(2*pi*x)*xp.cos(2*pi*y) - nu*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) - fy = 2*pi*xp.cos(2*pi*y) - B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) - nu*8*pi**2*xp.cos(2*pi*x)*xp.cos(2*pi*y) - fx += - B0 / epsilon * x * y + fx = ( + -2 * pi * xp.sin(2 * pi * x) + + B0 / epsilon * xp.cos(2 * pi * x) * xp.cos(2 * pi * y) + - nu * 8 * pi**2 * xp.sin(2 * pi * x) * xp.sin(2 * pi * y) + ) + fy = ( + 2 * pi * xp.cos(2 * pi * y) + - B0 / epsilon * xp.sin(2 * pi * x) * xp.sin(2 * pi * y) + - nu * 8 * pi**2 * xp.cos(2 * pi * x) * xp.cos(2 * pi * y) + ) + fx += -B0 / epsilon * x * y fy += B0 * x * y / epsilon return fx, fy, zeros_like(x) def source_function_ue(x, y, z): - fx = 2*pi*xp.sin(2*pi*x) - B0/epsilon*xp.cos(4*pi*x)*xp.cos(4*pi*y) - nu_e*32*pi**2*xp.sin(4*pi*x)*xp.sin(4*pi*y) + sigma*xp.sin(4*pi*x)*xp.sin(4*pi*y) - fy = -2*pi*xp.cos(2*pi*y) + B0/epsilon*xp.sin(4*pi*x)*xp.sin(4*pi*y) - nu_e*32*pi**2*xp.cos(4*pi*x)*xp.cos(4*pi*y) + sigma*xp.cos(4*pi*x)*xp.cos(4*pi*y) - fx += B0 / epsilon * (x * y + 1) - sigma * (x*y + 1) - fy += - B0 / epsilon * (x * y + 1) - sigma * (x*y + 1) + fx = ( + 2 * pi * xp.sin(2 * pi * x) + - B0 / epsilon * xp.cos(4 * pi * x) * xp.cos(4 * pi * y) + - nu_e * 32 * pi**2 * xp.sin(4 * pi * x) * xp.sin(4 * pi * y) + + sigma * xp.sin(4 * pi * x) * xp.sin(4 * pi * y) + ) + fy = ( + -2 * pi * xp.cos(2 * pi * y) + + B0 / epsilon * xp.sin(4 * pi * x) * xp.sin(4 * pi * y) + - nu_e * 32 * pi**2 * xp.cos(4 * pi * x) * xp.cos(4 * pi * y) + + sigma * xp.cos(4 * pi * x) * xp.cos(4 * pi * y) + ) + fx += B0 / epsilon * (x * y + 1) - sigma * (x * y + 1) + fy += -B0 / epsilon * (x * y + 1) - sigma * (x * y + 1) return fx, fy, zeros_like(x) - class MMSIonVelocity(perturbations.Perturbation): def __init__(self, comp=0): self.comp = comp @@ -172,6 +234,7 @@ def __init__(self): def __call__(self, x, y, z): return mms_phi(x, y, z)[0] + # ------------------ model ------------------ model = TwoFluidQuasiNeutralToy() @@ -182,12 +245,12 @@ def __call__(self, x, y, z): stab_sigma=sigma, source_u=source_function_u, source_ue=source_function_ue, - solver='gmres', + solver="gmres", solver_params=SolverParameters(verbose=True, info=True, tol=tol), ) -if BC == 'dirichlet_inhom': - model.ions.u.lifting_function = lifting_function_u +if BC == "dirichlet_inhom": + model.ions.u.lifting_function = lifting_function_u model.electrons.u.lifting_function = lifting_function_ue # model.ions.u.add_perturbation(MMSIonVelocity(comp=0)) @@ -213,7 +276,6 @@ def __call__(self, x, y, z): sim.run(verbose=True) if MPI.COMM_WORLD.Get_rank() == 0: - sim.pproc(verbose=True) sim.load_plotting_data(verbose=True) @@ -221,54 +283,56 @@ def __call__(self, x, y, z): n1_vals = simdata.grids_log[0] n2_vals = simdata.grids_log[1] - X, Y = xp.meshgrid(n1_vals, n2_vals, indexing='ij') + X, Y = xp.meshgrid(n1_vals, n2_vals, indexing="ij") x = xp.linspace(0, 1, 100) y = xp.linspace(0, 1, 100) - Xf, Yf = xp.meshgrid(x, y, indexing='ij') + Xf, Yf = xp.meshgrid(x, y, indexing="ij") - os.makedirs(f'{name}/plots', exist_ok=True) - for f in glob.glob(f'{name}/plots/*.png'): + os.makedirs(f"{name}/plots", exist_ok=True) + for f in glob.glob(f"{name}/plots/*.png"): os.remove(f) def save_plot(numerical, analytical, title, fname, t): fig, axes = plt.subplots(1, 2, figsize=(10, 4)) im0 = axes[0].contourf(X, Y, numerical, levels=50) - axes[0].set_title('numerical') + axes[0].set_title("numerical") plt.colorbar(im0, ax=axes[0]) im1 = axes[1].contourf(Xf, Yf, analytical, levels=50) - axes[1].set_title('manufactured') + axes[1].set_title("manufactured") plt.colorbar(im1, ax=axes[1]) - fig.suptitle(f'{title} at t={t:.3f}') - plt.savefig(f'{name}/plots/{fname}_{t:.3f}.png', dpi=300) + fig.suptitle(f"{title} at t={t:.3f}") + plt.savefig(f"{name}/plots/{fname}_{t:.3f}.png", dpi=300) plt.close(fig) for t in simdata.spline_values.ions.u_log.data.keys(): - u_ions = simdata.spline_values.ions.u_log.data[t] + u_ions = simdata.spline_values.ions.u_log.data[t] u_electrons = simdata.spline_values.electrons.u_log.data[t] - phi = simdata.spline_values.em_fields.phi_log.data[t] + phi = simdata.spline_values.em_fields.phi_log.data[t] - mms_phi_x, _, _ = mms_phi(Xf, Yf, 0*Xf) - mms_ion_ux, mms_ion_uy, _ = mms_ion_u(Xf, Yf, 0*Xf) - mms_el_ux, mms_el_uy, _ = mms_electron_u(Xf, Yf, 0*Xf) + mms_phi_x, _, _ = mms_phi(Xf, Yf, 0 * Xf) + mms_ion_ux, mms_ion_uy, _ = mms_ion_u(Xf, Yf, 0 * Xf) + mms_el_ux, mms_el_uy, _ = mms_electron_u(Xf, Yf, 0 * Xf) - save_plot(phi[0][:, :, 0], mms_phi_x, 'φ', 'plot_phi', t) - save_plot(u_ions[0][:, :, 0], mms_ion_ux, 'u_ix', 'plot_uix', t) - save_plot(u_ions[1][:, :, 0], mms_ion_uy, 'u_iy', 'plot_uiy', t) - save_plot(u_electrons[0][:, :, 0], mms_el_ux, 'u_ex', 'plot_uex', t) - save_plot(u_electrons[1][:, :, 0], mms_el_uy, 'u_ey', 'plot_uey', t) + save_plot(phi[0][:, :, 0], mms_phi_x, "φ", "plot_phi", t) + save_plot(u_ions[0][:, :, 0], mms_ion_ux, "u_ix", "plot_uix", t) + save_plot(u_ions[1][:, :, 0], mms_ion_uy, "u_iy", "plot_uiy", t) + save_plot(u_electrons[0][:, :, 0], mms_el_ux, "u_ex", "plot_uex", t) + save_plot(u_electrons[1][:, :, 0], mms_el_uy, "u_ey", "plot_uey", t) # ---- lifting diagnostics (dirichlet_inhom only) ---- - if BC == 'dirichlet_inhom': + if BC == "dirichlet_inhom": e1 = xp.linspace(0, 1, 80) e2 = xp.linspace(0, 1, 80) e3 = xp.array([0.5]) - E1, E2 = xp.meshgrid(e1, e2, indexing='ij') - - for label, var, comp in [('ion_ux', model.ions.u, 0), - ('ion_uy', model.ions.u, 1), - ('electron_ux', model.electrons.u, 0), - ('electron_uy', model.electrons.u, 1)]: + E1, E2 = xp.meshgrid(e1, e2, indexing="ij") + + for label, var, comp in [ + ("ion_ux", model.ions.u, 0), + ("ion_uy", model.ions.u, 1), + ("electron_ux", model.electrons.u, 0), + ("electron_uy", model.electrons.u, 1), + ]: if var.spline_lift is None: print(f" {label}: spline_lift is None, skipping") continue @@ -277,13 +341,15 @@ def _eval(fn): return fn(e1, e2, e3, squeeze_out=True)[comp] fig, axes = plt.subplots(1, 3, figsize=(15, 4)) - for ax, fn, ttl in zip(axes, - [var.spline_lift, var.spline_0, var.boundary_spline], - ['lifting', 'zero-BC part', 'boundary spline']): + for ax, fn, ttl in zip( + axes, + [var.spline_lift, var.spline_0, var.boundary_spline], + ["lifting", "zero-BC part", "boundary spline"], + ): im = ax.contourf(E1, E2, _eval(fn), levels=50) - ax.set_title(f'{label}: {ttl}') + ax.set_title(f"{label}: {ttl}") plt.colorbar(im, ax=ax) - out = f'{name}/plots/lifting_{label}.png' + out = f"{name}/plots/lifting_{label}.png" plt.savefig(out, dpi=300) plt.close(fig) print(f" -> saved {out}") @@ -293,42 +359,42 @@ def _eval(fn): e1 = xp.linspace(0, 1, 80) e2 = xp.linspace(0, 1, 80) e3 = xp.array([0.5]) - E1, E2 = xp.meshgrid(e1, e2, indexing='ij') + E1, E2 = xp.meshgrid(e1, e2, indexing="ij") zeros_E = xp.zeros_like(E1) for label, spline, src_fn, comp in [ - ('ion_source_x', prop._src_u, prop.options.source_u, 0), - ('ion_source_y', prop._src_u, prop.options.source_u, 1), - ('electron_source_x', prop._src_ue, prop.options.source_ue, 0), - ('electron_source_y', prop._src_ue, prop.options.source_ue, 1), + ("ion_source_x", prop._src_u, prop.options.source_u, 0), + ("ion_source_y", prop._src_u, prop.options.source_u, 1), + ("electron_source_x", prop._src_ue, prop.options.source_ue, 0), + ("electron_source_y", prop._src_ue, prop.options.source_ue, 1), ]: if spline is None: print(f" {label}: None, skipping") continue vals_proj = spline(e1, e2, e3, squeeze_out=True)[comp] - vals_ref = src_fn(E1, E2, zeros_E)[comp] + vals_ref = src_fn(E1, E2, zeros_E)[comp] fig, axes = plt.subplots(1, 2, figsize=(10, 4)) im0 = axes[0].contourf(E1, E2, vals_proj, levels=50) - axes[0].set_title('projected (FE)') + axes[0].set_title("projected (FE)") plt.colorbar(im0, ax=axes[0]) im1 = axes[1].contourf(E1, E2, vals_ref, levels=50) - axes[1].set_title('reference (analytical)') + axes[1].set_title("reference (analytical)") plt.colorbar(im1, ax=axes[1]) fig.suptitle(label) - out = f'{name}/plots/source_{label}.png' + out = f"{name}/plots/source_{label}.png" plt.savefig(out, dpi=300) plt.close(fig) print(f" -> saved {out}") - if BC == 'dirichlet_inhom': + if BC == "dirichlet_inhom": y_check = xp.linspace(0, 1, 80) x_check = xp.linspace(0, 1, 80) z_check = xp.array([0.5]) # comp=0: normal trace at x=0 and x=1 - for x_bnd, label in [(0.0, 'x=0'), (1.0, 'x=1')]: + for x_bnd, label in [(0.0, "x=0"), (1.0, "x=1")]: x_bnd_arr = xp.array([x_bnd]) mms_vals = mms_ion_u(x_bnd_arr, y_check, z_check)[0] lift_vals = model.ions.u.spline_lift(x_bnd_arr, y_check, z_check, squeeze_out=True)[0] @@ -339,7 +405,7 @@ def _eval(fn): print(f"elec ux normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") # comp=1: normal trace at y=0 and y=1 - for y_bnd, label in [(0.0, 'y=0'), (1.0, 'y=1')]: + for y_bnd, label in [(0.0, "y=0"), (1.0, "y=1")]: y_bnd_arr = xp.array([y_bnd]) mms_vals = mms_ion_u(x_check, y_bnd_arr, z_check)[1] lift_vals = model.ions.u.spline_lift(x_check, y_bnd_arr, z_check, squeeze_out=True)[1] @@ -347,4 +413,4 @@ def _eval(fn): mms_vals = mms_electron_u(x_check, y_bnd_arr, z_check)[1] lift_vals = model.electrons.u.spline_lift(x_check, y_bnd_arr, z_check, squeeze_out=True)[1] - print(f"elec uy normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") \ No newline at end of file + print(f"elec uy normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") diff --git a/examples/TwoFluidQuasiNeutralToy/R_Verification.py b/examples/TwoFluidQuasiNeutralToy/R_Verification.py index f94380a24..05152a62c 100644 --- a/examples/TwoFluidQuasiNeutralToy/R_Verification.py +++ b/examples/TwoFluidQuasiNeutralToy/R_Verification.py @@ -17,54 +17,53 @@ # ------------------ parameters ------------------ name = "runs/sim_restelli" -R0 = 2.0 -a = 1.0 -ain = 0.1 -B0 = 10.0 -Bp = 12.5 -nu = 1.0 -nu_e = 0.01 +R0 = 2.0 +a = 1.0 +ain = 0.1 +B0 = 10.0 +Bp = 12.5 +nu = 1.0 +nu_e = 0.01 alpha = 0.1 -beta = 1.0 -eps = 1.0 +beta = 1.0 +eps = 1.0 sigma = 1e-6 -dt = 1 -Tend = 1 -Nel = (8, 8, 1) -p = (1, 1, 1) -tol = 1e-5 +dt = 1 +Tend = 1 +Nel = (8, 8, 1) +p = (1, 1, 1) +tol = 1e-5 -env = EnvironmentOptions(sim_folder=name) +env = EnvironmentOptions(sim_folder=name) time_opts = Time(dt=dt, Tend=Tend) # ------------------ domain & equilibrium ------------------ domain = domains.HollowTorus(a1=ain, a2=a, R0=R0) -equil = equils.CircularTokamak(a=a, R0=R0, B0=B0, Bp=Bp) -grid = grids.TensorProductGrid(num_elements=Nel) +equil = equils.CircularTokamak(a=a, R0=R0, B0=B0, Bp=Bp) +grid = grids.TensorProductGrid(num_elements=Nel) -derham_opts = DerhamOptions( - degree=p, - bcs=(("dirichlet", "dirichlet"), None, None) -) +derham_opts = DerhamOptions(degree=p, bcs=(("dirichlet", "dirichlet"), None, None)) # ------------------ manufactured solution ------------------ + def _cylindrical(x, y, z): - R = xp.sqrt(x**2 + y**2) - R = xp.where(R == 0.0, 1e-9, R) + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) phi = xp.arctan2(-y, x) - Z = z + Z = z return R, phi, Z + def mms_u_cartesian(x, y, z): R, phi, Z = _cylindrical(x, y, z) - u_R = alpha/R * (-Z) / (a*R0/R) + beta * Bp/B0 * R0/(a*R) * Z - u_Z = alpha/R * (R - R0) / (a*R0/R) + beta * Bp/B0 * R0/(a*R) * (-(R - R0)) - u_phi = beta * Bp/B0 * R0/(a*R) * B0*Bp*a / (Bp*R0) + u_R = alpha / R * (-Z) / (a * R0 / R) + beta * Bp / B0 * R0 / (a * R) * Z + u_Z = alpha / R * (R - R0) / (a * R0 / R) + beta * Bp / B0 * R0 / (a * R) * (-(R - R0)) + u_phi = beta * Bp / B0 * R0 / (a * R) * B0 * Bp * a / (Bp * R0) - u_R = alpha * R / (a*R0) * (-Z) + beta * Bp/B0 * R0/(a*R) * Z - u_Z = alpha * R / (a*R0) * (R - R0) + beta * Bp/B0 * R0/(a*R) * (-(R - R0)) - u_phi = beta * Bp/B0 * R0/(a*R) * (B0/Bp * a) + u_R = alpha * R / (a * R0) * (-Z) + beta * Bp / B0 * R0 / (a * R) * Z + u_Z = alpha * R / (a * R0) * (R - R0) + beta * Bp / B0 * R0 / (a * R) * (-(R - R0)) + u_phi = beta * Bp / B0 * R0 / (a * R) * (B0 / Bp * a) # Transform to Cartesian via eq (5.14) ux = xp.cos(phi) * u_R - R * xp.sin(phi) * u_phi @@ -72,30 +71,36 @@ def mms_u_cartesian(x, y, z): uz = u_Z return ux, uy, uz + def mms_phi_cartesian(x, y, z): R, phi, Z = _cylindrical(x, y, z) # eq (5.22): phi_hat = 0.5 * a * B0 * alpha * ((R-R0)^2 + Z^2)/a^2 - 2/3) - phi_val = 0.5 * a * B0 * alpha * (((R - R0)**2 + Z**2) / a**2 - 2.0/3.0) + phi_val = 0.5 * a * B0 * alpha * (((R - R0) ** 2 + Z**2) / a**2 - 2.0 / 3.0) return phi_val + # ------------------ source terms ------------------ + def _omega_cartesian(x, y, z): R, phi, Z = _cylindrical(x, y, z) - omega_Z = alpha * (R0 - 4*R) / (a*R0*R) - beta * Bp/B0 * R0**2 / (a*R**3) + omega_Z = alpha * (R0 - 4 * R) / (a * R0 * R) - beta * Bp / B0 * R0**2 / (a * R**3) ox = xp.zeros_like(x) oy = xp.zeros_like(x) oz = omega_Z return ox, oy, oz + def source_function_u(x, y, z): ox, oy, oz = _omega_cartesian(x, y, z) return nu * ox, nu * oy, nu * oz + def source_function_ue(x, y, z): ox, oy, oz = _omega_cartesian(x, y, z) return nu_e * ox, nu_e * oy, nu_e * oz + # ------------------ lifting (inhomogeneous Dirichlet on radial boundary) ------------------ lifting_u = [ GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[0], comp=0, given_in_basis="physical"), @@ -119,11 +124,11 @@ def source_function_ue(x, y, z): stab_sigma=sigma, source_u=source_function_u, source_ue=source_function_ue, - solver='gmres', + solver="gmres", solver_params=SolverParameters(verbose=True, info=True, tol=tol), ) -model.ions.u.lifting_function = lifting_u +model.ions.u.lifting_function = lifting_u model.electrons.u.lifting_function = lifting_ue # ------------------ simulation ------------------ @@ -147,8 +152,8 @@ def source_function_ue(x, y, z): sim.load_plotting_data(verbose=True) simdata = sim.plotting_data - os.makedirs(f'{name}/plots', exist_ok=True) - for f in glob.glob(f'{name}/plots/*.png'): + os.makedirs(f"{name}/plots", exist_ok=True) + for f in glob.glob(f"{name}/plots/*.png"): os.remove(f) e1 = xp.linspace(0, 1, 40) @@ -161,57 +166,72 @@ def source_function_ue(x, y, z): x_outer, y_outer, _ = domain(xp.array([1.0]), theta_bnd, xp.array([0.5]), squeeze_out=True) def _add_domain_boundary(ax): - ax.plot(x_inner, y_inner, 'w-', linewidth=0.8) - ax.plot(x_outer, y_outer, 'w-', linewidth=0.8) + ax.plot(x_inner, y_inner, "w-", linewidth=0.8) + ax.plot(x_outer, y_outer, "w-", linewidth=0.8) prop = model.propagators.qn_full for label, spline, src_fn, comp in [ - ('ion_source_x', prop._src_u, prop.options.source_u, 0), - ('ion_source_y', prop._src_u, prop.options.source_u, 1), - ('ion_source_z', prop._src_u, prop.options.source_u, 2), - ('electron_source_x', prop._src_ue, prop.options.source_ue, 0), - ('electron_source_y', prop._src_ue, prop.options.source_ue, 1), - ('electron_source_z', prop._src_ue, prop.options.source_ue, 2), + ("ion_source_x", prop._src_u, prop.options.source_u, 0), + ("ion_source_y", prop._src_u, prop.options.source_u, 1), + ("ion_source_z", prop._src_u, prop.options.source_u, 2), + ("electron_source_x", prop._src_ue, prop.options.source_ue, 0), + ("electron_source_y", prop._src_ue, prop.options.source_ue, 1), + ("electron_source_z", prop._src_ue, prop.options.source_ue, 2), ]: if spline is None: print(f" {label}: None, skipping") continue vals_proj = spline(e1, e2, e3, squeeze_out=True)[comp] - vals_ref = src_fn(x_phys, y_phys, z_phys)[comp] + vals_ref = src_fn(x_phys, y_phys, z_phys)[comp] fig, axes = plt.subplots(1, 2, figsize=(10, 4)) - im0 = axes[0].contourf(x_phys, y_phys, vals_proj, levels=50); axes[0].set_title('projected (FE)'); plt.colorbar(im0, ax=axes[0]); _add_domain_boundary(axes[0]) - im1 = axes[1].contourf(x_phys, y_phys, vals_ref, levels=50); axes[1].set_title('reference (analytical)'); plt.colorbar(im1, ax=axes[1]); _add_domain_boundary(axes[1]) + im0 = axes[0].contourf(x_phys, y_phys, vals_proj, levels=50) + axes[0].set_title("projected (FE)") + plt.colorbar(im0, ax=axes[0]) + _add_domain_boundary(axes[0]) + im1 = axes[1].contourf(x_phys, y_phys, vals_ref, levels=50) + axes[1].set_title("reference (analytical)") + plt.colorbar(im1, ax=axes[1]) + _add_domain_boundary(axes[1]) fig.suptitle(label) - out = f'{name}/plots/source_{label}.png' + out = f"{name}/plots/source_{label}.png" plt.savefig(out, dpi=300) plt.close(fig) print(f" -> saved {out}") for t in simdata.spline_values.ions.u_log.data.keys(): - u_ions = simdata.spline_values.ions.u_log.data[t] + u_ions = simdata.spline_values.ions.u_log.data[t] u_electrons = simdata.spline_values.electrons.u_log.data[t] - phi_num = simdata.spline_values.em_fields.phi_log.data[t] + phi_num = simdata.spline_values.em_fields.phi_log.data[t] mms_ux, mms_uy, mms_uz = mms_u_cartesian(x_phys, y_phys, z_phys) mms_phi_val = mms_phi_cartesian(x_phys, y_phys, z_phys) for num, mms, lbl in [ - (u_ions[0][:, :, 0], mms_ux, 'u_ix'), - (u_ions[1][:, :, 0], mms_uy, 'u_iy'), - (u_ions[2][:, :, 0], mms_uz, 'u_iz'), - (u_electrons[0][:, :, 0], mms_ux, 'u_ex'), - (u_electrons[1][:, :, 0], mms_uy, 'u_ey'), - (u_electrons[2][:, :, 0], mms_uz, 'u_ez'), - (phi_num[0][:, :, 0], mms_phi_val, 'phi'), + (u_ions[0][:, :, 0], mms_ux, "u_ix"), + (u_ions[1][:, :, 0], mms_uy, "u_iy"), + (u_ions[2][:, :, 0], mms_uz, "u_iz"), + (u_electrons[0][:, :, 0], mms_ux, "u_ex"), + (u_electrons[1][:, :, 0], mms_uy, "u_ey"), + (u_electrons[2][:, :, 0], mms_uz, "u_ez"), + (phi_num[0][:, :, 0], mms_phi_val, "phi"), ]: fig, axes = plt.subplots(1, 3, figsize=(15, 4)) - im0 = axes[0].contourf(x_phys, y_phys, num, levels=50); axes[0].set_title('numerical'); plt.colorbar(im0, ax=axes[0]); _add_domain_boundary(axes[0]) - im1 = axes[1].contourf(x_phys, y_phys, mms, levels=50); axes[1].set_title('MMS'); plt.colorbar(im1, ax=axes[1]); _add_domain_boundary(axes[1]) - im2 = axes[2].contourf(x_phys, y_phys, num - mms, levels=50); axes[2].set_title('difference'); plt.colorbar(im2, ax=axes[2]); _add_domain_boundary(axes[2]) - fig.suptitle(f'{lbl} at t={t:.4f}') - out = f'{name}/plots/{lbl}_{t:.4f}.png' + im0 = axes[0].contourf(x_phys, y_phys, num, levels=50) + axes[0].set_title("numerical") + plt.colorbar(im0, ax=axes[0]) + _add_domain_boundary(axes[0]) + im1 = axes[1].contourf(x_phys, y_phys, mms, levels=50) + axes[1].set_title("MMS") + plt.colorbar(im1, ax=axes[1]) + _add_domain_boundary(axes[1]) + im2 = axes[2].contourf(x_phys, y_phys, num - mms, levels=50) + axes[2].set_title("difference") + plt.colorbar(im2, ax=axes[2]) + _add_domain_boundary(axes[2]) + fig.suptitle(f"{lbl} at t={t:.4f}") + out = f"{name}/plots/{lbl}_{t:.4f}.png" plt.savefig(out, dpi=300) - plt.close(fig) \ No newline at end of file + plt.close(fig) From 03d42b5875b86284bc6fd4f0ec2eb503b80d291f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 1 Jun 2026 09:43:04 +0200 Subject: [PATCH 33/41] formatting --- src/struphy/models/base.py | 2 +- src/struphy/models/species.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index d7531efc1..b21eba88a 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -352,7 +352,7 @@ def setup_equation_params(self, base_units: BaseUnits): for _, species in self.particle_species.items(): assert isinstance(species, ParticleSpecies) species.setup_equation_params(units=self.units) - + def show_equation_params(self): """Print the equation parameters for each species to screen.""" for _, species in self.fluid_species.items(): diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 0db58edea..4878d63b5 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -169,7 +169,7 @@ def __init__( kappa: float = None, ): self.species = species - + if units is None: units = Units() From 7a1232dd589c83c4a07c5f4aa1cf2e096920596b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 1 Jun 2026 09:46:42 +0200 Subject: [PATCH 34/41] remove verbose --- src/struphy/models/variables.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index eecbbc5f9..b3b7c2c68 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -384,7 +384,6 @@ def allocate( space_id=self.space, domain=domain, equil=equil, - verbose=verbose, ) # project each perturbation and accumulate into spline_lift From f454793545d1d3f9249eb4bb7d4fbaf3f4dcdfd3 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 1 Jun 2026 14:03:13 +0200 Subject: [PATCH 35/41] add unit test for propagator only - inhom bcs not working yet --- src/struphy/feec/psydac_derham.py | 2 +- .../tests/test_two_fluid_quasi_neutral.py | 268 ++++++++++++++++++ .../two_fluid_quasi_neutral_full.py | 12 +- 3 files changed, 280 insertions(+), 2 deletions(-) create mode 100644 src/struphy/propagators/tests/test_two_fluid_quasi_neutral.py diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 82bc49377..9402fe724 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -549,7 +549,7 @@ class Derham: MPI communicator (sub_comm if clones are used). domain : Domain, optional - The Struphy domain object for evaluating the mapping F : [0, 1]^3 --> R^3 and the corresponding metric coefficients. + The Struphy domain object for evaluating the mapping F: [0, 1]^3 --> R^3 and the corresponding metric coefficients. Notes ----- diff --git a/src/struphy/propagators/tests/test_two_fluid_quasi_neutral.py b/src/struphy/propagators/tests/test_two_fluid_quasi_neutral.py new file mode 100644 index 000000000..f89928679 --- /dev/null +++ b/src/struphy/propagators/tests/test_two_fluid_quasi_neutral.py @@ -0,0 +1,268 @@ +import logging + +import cunumpy as xp +from cunumpy import pi, cos, sin, zeros_like, ones_like +import matplotlib.pyplot as plt +import pytest +from feectools.ddm.mpi import mpi as MPI + +from struphy import domains, equils, grids, set_logging_level +from struphy.initial.base import Perturbation, GenericPerturbation +from struphy.feec.mass import L2Projector, WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.geometry.base import Domain +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable +from struphy.propagators.two_fluid_quasi_neutral_full import TwoFluidQuasiNeutralFull +from struphy.propagators.base import Propagator +from struphy.fields_background.projected_equils import ProjectedMHDequilibrium +from struphy.feec.basis_projection_ops import BasisProjectionOperators + +logger = logging.getLogger("struphy") +set_logging_level(logging.INFO) + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() +# plt.rcParams.update({'font.size': 22}) + + +pytest.mark.parametrize("bc_type", ["periodic", "hom_dirichlet", "inhom_dirichlet"]) +@pytest.mark.parametrize( + "mapping", + [ + None, + ["Cuboid", {"l1": 0.0, "r1": 4.0, "l2": 0.0, "r2": 2.0, "l3": 0.0, "r3": 3.0}], + ["Orthogonal", {"Lx": 4.0, "Ly": 2.0, "alpha": 0.1, "Lz": 3.0}], + ], +) +def test_one_time_step(bc_type: str, mapping: None | list[str, dict], show_plot=False): + """Test the propagator TwoFluidQuasiNeutralFull on a single time step against a manufactured solution with different boundary conditions.""" + + # domain object + if mapping is None: + domain = domains.Cuboid() + else: + dom_type = mapping[0] + dom_params = mapping[1] + domain_class = getattr(domains, dom_type) + domain: Domain = domain_class(**dom_params) + + # other options + B0 = 0 + nu = 10.0 + nu_e = 1.0 + Nel = (32, 1, 1) + degree = (1, 1, 1) + epsilon = 1.0 + dt = 1 + Tend = 1 + sigma = 0 + tol = 1e-5 + + # derham sequence + if bc_type == "periodic": + derham_opts = DerhamOptions(degree=degree, bcs=(None, None, None)) + + elif bc_type == "hom_dirichlet": + derham_opts = DerhamOptions(degree=degree, bcs=(("dirichlet", "dirichlet"), None, None)) + + elif bc_type == "inhom_dirichlet": + derham_opts = DerhamOptions(degree=degree, bcs=(("dirichlet", "dirichlet"), None, None)) + lifting_function_u = GenericPerturbation(lambda x, y, z: x + 1, comp=0, given_in_basis="physical") + lifting_function_ue = GenericPerturbation(lambda x, y, z: x, comp=0, given_in_basis="physical") + + else: + raise ValueError(f"Invalid bc_type: {bc_type}") + + grid = grids.TensorProductGrid(num_elements=Nel) + derham = Derham(grid=grid, options=derham_opts, domain=domain) + + # fluid background + equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) + projected_equil = ProjectedMHDequilibrium(equil=equil, derham=derham) + + # mass operators + mass_ops = WeightedMassOperators(derham=derham, domain=domain, eq_mhd=equil) + + # basis operators + basis_ops = BasisProjectionOperators(derham, domain, eq_mhd=equil) + + # ---- manufactured solutions ---- + if bc_type == "periodic": + + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x) + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + elif bc_type == "hom_dirichlet": + + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + elif bc_type == "inhom_dirichlet": + + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x) + x + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x) + x, xp.zeros_like(x), xp.zeros_like(x) + + # ---- source terms ---- + if bc_type == "periodic": + + def source_function_u(x, y, z): + fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + elif bc_type == "hom_dirichlet": + + def source_function_u(x, y, z): + fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + elif bc_type == "inhom_dirichlet": + + def source_function_u(x, y, z): + fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + (4.0 * nu_e * pi**2 - sigma) * sin(2 * pi * x) - sigma * x + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + + # ---- perturbation classes for MMS initial conditions ---- + class MMSIonVelocity(Perturbation): + def __init__(self, comp=0): + self.comp = comp + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_ion_u(x, y, z)[self.comp] + + + class MMSElectronVelocity(Perturbation): + def __init__(self, comp=0): + self.comp = comp + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_electron_u(x, y, z)[self.comp] + + + class MMSPotential(Perturbation): + def __init__(self): + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_phi(x, y, z)[0] + + # instance of propagator + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + Propagator.basis_ops = basis_ops + Propagator.projected_equil = projected_equil + + prop = TwoFluidQuasiNeutralFull(allocate_variables=True) + + prop.options = prop.Options( + nu=nu, + nu_e=nu_e, + eps_norm=epsilon, + stab_sigma=sigma, + source_u=source_function_u, + source_ue=source_function_ue, + solver="gmres", + solver_params=SolverParameters(info=True, tol=tol), + ) + + prop.allocate() + + if bc_type == "inhom_dirichlet": + prop.variables.u.lifting_function = lifting_function_u + prop.variables.ue.lifting_function = lifting_function_ue + + prop(dt) + + x = xp.linspace(0, 1, 100) + mms_phi_x, _, _ = mms_phi(x, x * 0, x * 0) + mms_ion_ux, _, _ = mms_ion_u(x, x * 0, x * 0) + mms_el_ux, _, _ = mms_electron_u(x, x * 0, x * 0) + + e1 = xp.linspace(0, 1, 64) + + if show_plot: + plt.figure(figsize=(18, 8)) + + plt.subplot(1, 3, 1) + plt.plot(e1, prop.variables.u.spline(e1, 0.5, 0.5, squeeze_out=True)[0], label="numerical") + plt.plot(x, mms_ion_ux, "--", label="manufactured") + # plt.plot(e1, prop.variables.phi.spline(e1, 0.5, 0.5, squeeze_out=True), "k.", markersize=4, label="n1 points") + plt.xlabel("x") + plt.title(r"$u$") + # plt.title(f"{title} at t={t:.3f}") + plt.legend() + plt.grid(True) + + plt.subplot(1, 3, 2) + plt.plot(e1, prop.variables.ue.spline(e1, 0.5, 0.5, squeeze_out=True)[0], label="numerical") + plt.plot(x, mms_el_ux, "--", label="manufactured") + # plt.plot(e1, prop.variables.ue.spline(e1, 0.5, 0.5, squeeze_out=True)[0], "k.", markersize=4, label="n1 points") + plt.xlabel("x") + plt.title(r"$u_e$") + # plt.title(f"{title} at t={t:.3f}") + plt.legend() + plt.grid(True) + + plt.subplot(1, 3, 3) + plt.plot(e1, prop.variables.phi.spline(e1, 0.5, 0.5, squeeze_out=True), label="numerical") + plt.plot(x, mms_phi_x, "--", label="manufactured") + # plt.plot(e1, prop.variables.phi.spline(e1, 0.5, 0.5, squeeze_out=True), "k.", markersize=4, label="n1 points") + plt.xlabel("x") + plt.title(r"$\phi$") + # plt.title(f"{title} at t={t:.3f}") + plt.legend() + plt.grid(True) + + plt.show() + + + +if __name__ == "__main__": + test_one_time_step(bc_type="inhom_dirichlet", mapping=None, show_plot=True) \ No newline at end of file diff --git a/src/struphy/propagators/two_fluid_quasi_neutral_full.py b/src/struphy/propagators/two_fluid_quasi_neutral_full.py index 3ecbdcb0e..d7beff804 100644 --- a/src/struphy/propagators/two_fluid_quasi_neutral_full.py +++ b/src/struphy/propagators/two_fluid_quasi_neutral_full.py @@ -88,8 +88,18 @@ def phi(self, new): assert new.space == "L2" self._phi = new - def __init__(self): + def __init__(self, allocate_variables: bool = False): self.variables = self.Variables() + + if allocate_variables: + self.variables.u = FEECVariable(space="Hdiv") + self.variables.ue = FEECVariable(space="Hdiv") + self.variables.phi = FEECVariable(space="L2") + + self.variables.u.allocate(derham=self.derham, domain=self.domain, equil=self.projected_equil.equil) + self.variables.ue.allocate(derham=self.derham, domain=self.domain, equil=self.projected_equil.equil) + self.variables.phi.allocate(derham=self.derham, domain=self.domain, equil=self.projected_equil.equil) + # ========================================================================= ### Options From 95573599695e91663dfab771a99ba5fc929c49b4 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 1 Jun 2026 14:13:21 +0200 Subject: [PATCH 36/41] do not always stabilize the Implicit diffusion solve --- src/struphy/propagators/implicit_diffusion.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index cc2c61108..5bfc438c7 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -203,9 +203,9 @@ def options(self, new): @profile def allocate(self): # always stabilize - if xp.abs(self.options.sigma_1) < 1e-14: - self.options.sigma_1 = 1e-14 - logger.warning(f"Stabilizing Poisson solve with {self.options.sigma_1 =}") + # if xp.abs(self.options.sigma_1) < 1e-14: + # self.options.sigma_1 = 1e-14 + # logger.warning(f"Stabilizing Poisson solve with {self.options.sigma_1 =}") # model parameters self._sigma_1 = self.options.sigma_1 From f1dd5a2482819d725c343b986d2cfdbb35299f1f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 1 Jun 2026 15:43:38 +0200 Subject: [PATCH 37/41] add type annotations to BoundaryOperator --- src/struphy/feec/linear_operators.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 1cb513ca1..4601fe543 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -11,6 +11,7 @@ from struphy.feec.utilities import apply_essential_bc_to_array from struphy.polar.basic import PolarDerhamSpace +from struphy.io.options import LiteralOptions class LinOpWithTransp(LinearOperator): @@ -317,9 +318,14 @@ class BoundaryOperator(LinOpWithTransp): (e.g. unconstrained to constrained). If None, domain and codomain are the same. """ - def __init__(self, vector_space, space_id, dirichlet_bc, codomain=None): + def __init__(self, + vector_space: VectorSpace, + space_id: LiteralOptions.OptsFEECSpace, + dirichlet_bc: tuple[tuple[bool]], + codomain: VectorSpace | None = None, + ): assert isinstance(vector_space, VectorSpace) - assert isinstance(space_id, str) + assert space_id in LiteralOptions.OptsFEECSpace self._domain = vector_space if codomain is not None: From 69412aef91a8588eb749a4a990b5ddc26b0bdf1f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 2 Jun 2026 08:47:01 +0200 Subject: [PATCH 38/41] add inhom_dirichlet boundary conditions (using lifting) to the Poisson unit test - it works. The Propagator ImplicitDiffusion now allows for this type of bc. --- src/struphy/feec/linear_operators.py | 18 ++--- src/struphy/models/variables.py | 45 +++++++++--- src/struphy/propagators/base.py | 8 +-- src/struphy/propagators/implicit_diffusion.py | 11 +++ src/struphy/propagators/tests/test_poisson.py | 68 ++++++++++++++++--- .../tests/test_two_fluid_quasi_neutral.py | 58 ++++++++-------- .../two_fluid_quasi_neutral_full.py | 5 +- 7 files changed, 147 insertions(+), 66 deletions(-) diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 4601fe543..c22860546 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -10,8 +10,9 @@ from scipy import sparse from struphy.feec.utilities import apply_essential_bc_to_array -from struphy.polar.basic import PolarDerhamSpace from struphy.io.options import LiteralOptions +from struphy.polar.basic import PolarDerhamSpace +from struphy.utils.utils import check_option class LinOpWithTransp(LinearOperator): @@ -318,14 +319,15 @@ class BoundaryOperator(LinOpWithTransp): (e.g. unconstrained to constrained). If None, domain and codomain are the same. """ - def __init__(self, - vector_space: VectorSpace, - space_id: LiteralOptions.OptsFEECSpace, - dirichlet_bc: tuple[tuple[bool]], - codomain: VectorSpace | None = None, - ): + def __init__( + self, + vector_space: VectorSpace, + space_id: LiteralOptions.OptsFEECSpace, + dirichlet_bc: tuple[tuple[bool]], + codomain: VectorSpace | None = None, + ): assert isinstance(vector_space, VectorSpace) - assert space_id in LiteralOptions.OptsFEECSpace + check_option(space_id, LiteralOptions.OptsFEECSpace) self._domain = vector_space if codomain is not None: diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index b3b7c2c68..8ae3e7502 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -10,6 +10,7 @@ from feectools.ddm.mpi import mpi as MPI from struphy.feec.linear_operators import BoundaryOperator +from struphy.feec.mass import WeightedMassOperators from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.fields_background.base import FluidEquilibrium from struphy.fields_background.projected_equils import ProjectedFluidEquilibrium @@ -226,7 +227,7 @@ def space(self) -> str: def lifting_function(self) -> Perturbation | None: """The lifting function for the case of lifting of boundary conditions. Its values at the boundary determine the inhomogeneous boundary conditions. - If None, no lifting is applied.""" + The interior part is irrelevant. If None, no lifting is applied.""" if not hasattr(self, "_lifting_function"): self._lifting_function = None return self._lifting_function @@ -237,48 +238,58 @@ def lifting_function(self, new: Perturbation | None): @property def spline(self) -> SplineFunction: + """The solution spline function.""" if not hasattr(self, "_spline"): raise ValueError("Warning: spline not allocated yet. Call allocate() first.") return self._spline @property def spline_lift(self) -> SplineFunction | None: - """The lifting function for the case of lifting of boundary conditions. Only allocated if lifting_function is not None.""" + """The spline representation of the lifting function for the case of lifting of boundary conditions. + The values in the interior are irrelevant, only the boundary values determine the boundary conditions. + Only allocated if lifting_function is not None.""" if not hasattr(self, "_spline_lift"): self._spline_lift = None return self._spline_lift @property def spline_0(self) -> SplineFunction | None: - """The spline function with zero boundary conditions, used for the lifting of boundary conditions. Only allocated if lifting_function is not None.""" + """Is equal to spline_lift but with boundary coeffcients set to zero. + Only allocated if lifting_function is not None.""" if not hasattr(self, "_spline_0"): self._spline_0 = None return self._spline_0 @property def boundary_spline(self) -> SplineFunction | None: - """The spline function representing the boundary conditions, used for the lifting of boundary conditions. Only allocated if lifting_function is not None.""" + """Is given by spline_lift - spline_0 and computed by the method set_boundary_spline. + This spline appears in the weak form of the equations as a source term and is responsible for the inhomogeneous boundary conditions. + Only allocated if lifting_function is not None.""" if not hasattr(self, "_boundary_spline"): self._boundary_spline = None return self._boundary_spline @property def spline_full(self) -> SplineFunction | None: - """Full solution spline (lifting + zero-BC part) in the unconstrained space. Only allocated if lifting_function is not None.""" - if not hasattr(self, "_spline_full"): - self._spline_full = None + """Full solution spline in the unconstrained (helper) space. + Its values are equal to spline + boundary_spline. + Only allocated if lifting_function is not None.""" + # update coeffs + self._spline_full.vector = self.boundary_op_lift.T.dot(self.spline.vector) + self.boundary_spline.vector return self._spline_full @property def boundary_op(self) -> BoundaryOperator | None: - """The boundary operator, used for the lifting of boundary conditions. Only allocated if lifting_function is not None.""" + """Boundary operator in the unconstrained (helper) space. + Is used to compute spline_0 for instance. + Only allocated if lifting_function is not None.""" if not hasattr(self, "_boundary_op"): self._boundary_op = None return self._boundary_op @property def boundary_op_lift(self) -> BoundaryOperator | None: - """Boundary operator mapping from the unconstrained (lifted) space to the constrained space. + """Boundary operator from the unconstrained (helper) space to the solution space (with homogeneous Dirichlet conditions). Only allocated if lifting_function is not None.""" if not hasattr(self, "_boundary_op_lift"): self._boundary_op_lift = None @@ -286,11 +297,20 @@ def boundary_op_lift(self) -> BoundaryOperator | None: @property def derham_lift(self) -> Derham | None: - """The Derham object for the lifting function. Only allocated if lifting_function is not None.""" + """The Derham object for the lifting function, yielding the unconstrained (helper) spaces. + Only allocated if lifting_function is not None.""" if not hasattr(self, "_derham_lift"): self._derham_lift = None return self._derham_lift + @property + def mass_ops_lift(self) -> WeightedMassOperators | None: + """The mass operators for the unconstrained (helper) spaces for the case of lifting of boundary conditions. + Only allocated if lifting_function is not None.""" + if not hasattr(self, "_mass_ops_lift"): + self._mass_ops_lift = None + return self._mass_ops_lift + @property def species(self) -> FieldSpecies | FluidSpecies: if not hasattr(self, "_species"): @@ -334,7 +354,7 @@ def allocate( if self.lifting_function is not None: check_bcs = False for bc in derham.bcs: - if "dirichlet" in bc: + if bc is not None and "dirichlet" in bc: check_bcs = True break assert check_bcs, ( @@ -370,6 +390,9 @@ def allocate( self._derham_lift = Derham.from_dict(dct, comm=derham.comm) + # unconstrained mass operators + self._mass_ops_lift = WeightedMassOperators(self.derham_lift, domain, eq_mhd=equil) + # spline function for the lifting self._spline_lift = self.derham_lift.create_spline_function( name=self.__name__ + "_lift" if self.__name__ is not None else None, diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index d81cf404a..01bec2ae3 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -121,10 +121,10 @@ def update_feec_variables(self, **new_coeffs): assert new.space == old.space # update full solution spline (lifting + zero-BC part) if present - if old_var.spline_full is not None: - new.copy(out=old_var.spline_full.vector) - if old_var.boundary_spline is not None: - old_var.spline_full.vector += old_var.boundary_spline.vector + # if old_var.spline_full is not None: + # new.copy(out=old_var.spline_full.vector) + # if old_var.boundary_spline is not None: + # old_var.spline_full.vector += old_var.boundary_spline.vector # calculate maximum of difference abs(new - old) diffs[var] = xp.max(xp.abs(new.toarray() - old.toarray())) diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index 5bfc438c7..8f3e553be 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -250,6 +250,17 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: else: self._coeffs = [1.0 for src in self.sources] + # add term for inhomogeneous boundary conditions if needed + if self.variables.phi.lifting_function is not None: + grad_lift = self.variables.phi.derham_lift.grad + M1_lift = self.variables.phi.mass_ops_lift.M1 + boundary_op_lift = self.variables.phi.boundary_op_lift + + op = -boundary_op_lift @ grad_lift.T @ M1_lift @ grad_lift + + self._sources += [op.dot(self.variables.phi.boundary_spline.vector)] + self._coeffs += [1.0] + # initial guess and solver params self._x0 = self.options.x0 self._info = self.options.solver_params.info diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 050f30506..fb148517d 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -12,10 +12,12 @@ WeightsParameters, domains, perturbations, + set_logging_level, ) from struphy.feec.mass import L2Projector, WeightedMassOperators from struphy.feec.psydac_derham import Derham from struphy.geometry.base import Domain +from struphy.initial.base import GenericPerturbation from struphy.io.options import DerhamOptions from struphy.kinetic_background.maxwellians import Maxwellian3D from struphy.linear_algebra.solver import SolverParameters @@ -29,6 +31,7 @@ from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") +set_logging_level(logging.WARNING) comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -36,7 +39,7 @@ @pytest.mark.parametrize("direction", [0, 1, 2]) -@pytest.mark.parametrize("bc_type", ["periodic", "dirichlet", "neumann"]) +@pytest.mark.parametrize("bc_type", ["periodic", "dirichlet", "neumann", "inhom_dirichlet"]) @pytest.mark.parametrize( "mapping", [ @@ -55,6 +58,10 @@ def test_poisson_1d( """ Test the convergence of Poisson solver in 1D by means of manufactured solutions. """ + # stabilization (removed for Dirichlet boundary conditions -> well-posed) + stab_eps = 1e-12 + if "dirichlet" in bc_type: + stab_eps = 0.0 # create domain object dom_type = mapping[0] @@ -103,6 +110,18 @@ def sol1_xyz(x, y, z): def rho1_xyz(x, y, z): return xp.cos(xp.pi / Lx * x) * (xp.pi / Lx) ** 2 + + elif bc_type == "inhom_dirichlet": + bcs = (("dirichlet", "dirichlet"), None, None) + + lifting_fun = GenericPerturbation(lambda x, y, z: x / Lx - 0.5 + x * (Lx - x)) + + def sol1_xyz(x, y, z): + return xp.sin(2 * xp.pi / Lx * x) + x / Lx - 0.5 + + def rho1_xyz(x, y, z): + return xp.sin(2 * xp.pi / Lx * x) * (2 * xp.pi / Lx) ** 2 + else: if bc_type == "dirichlet": bcs = (("dirichlet", "dirichlet"), None, None) @@ -126,6 +145,18 @@ def sol1_xyz(x, y, z): def rho1_xyz(x, y, z): return xp.cos(xp.pi / Ly * y) * (xp.pi / Ly) ** 2 + + elif bc_type == "inhom_dirichlet": + bcs = (None, ("dirichlet", "dirichlet"), None) + + lifting_fun = GenericPerturbation(lambda x, y, z: y / Ly - 0.5 + y * (Ly - y)) + + def sol1_xyz(x, y, z): + return xp.sin(2 * xp.pi / Ly * y) + y / Ly - 0.5 + + def rho1_xyz(x, y, z): + return xp.sin(2 * xp.pi / Ly * y) * (2 * xp.pi / Ly) ** 2 + else: if bc_type == "dirichlet": bcs = (None, ("dirichlet", "dirichlet"), None) @@ -149,6 +180,18 @@ def sol1_xyz(x, y, z): def rho1_xyz(x, y, z): return xp.cos(xp.pi / Lz * z) * (xp.pi / Lz) ** 2 + + elif bc_type == "inhom_dirichlet": + bcs = (None, None, ("dirichlet", "dirichlet")) + + lifting_fun = GenericPerturbation(lambda x, y, z: z / Lz - 0.5 + z * (Lz - z)) + + def sol1_xyz(x, y, z): + return xp.sin(2 * xp.pi / Lz * z) + z / Lz - 0.5 + + def rho1_xyz(x, y, z): + return xp.sin(2 * xp.pi / Lz * z) * (2 * xp.pi / Lz) ** 2 + else: if bc_type == "dirichlet": bcs = (None, None, ("dirichlet", "dirichlet")) @@ -194,13 +237,14 @@ def rho_pulled(e1, e2, e3): ) _phi = FEECVariable(space="H1") + _phi.lifting_function = lifting_fun if "inhom_dirichlet" in bc_type else None _phi.allocate(derham=derham, domain=domain) poisson_solver = PoissonFieldSolve() poisson_solver.variables.phi = _phi poisson_solver.options = poisson_solver.Options( - stab_eps=1e-12, + stab_eps=stab_eps, # sigma_2=0.0, # sigma_3=1.0, rho=rho, @@ -216,7 +260,12 @@ def rho_pulled(e1, e2, e3): poisson_solver(dt) # push numerical solution and compare - sol_val1 = domain.push(_phi.spline, e1, e2, e3, kind="0") + if bc_type == "inhom_dirichlet": + sol = _phi.spline_full + else: + sol = _phi.spline + + sol_val1 = domain.push(sol, e1, e2, e3, kind="0") x, y, z = domain(e1, e2, e3) analytic_value1 = sol1_xyz(x, y, z) @@ -247,7 +296,7 @@ def rho_pulled(e1, e2, e3): m, _ = xp.polyfit(xp.log(Nels), xp.log(errors), deg=1) logger.info(f"For {pi =}, solution converges in {direction=} with rate {-m =} ") - assert -m > (pi + 1 - 0.07) + # assert -m > (pi + 1 - 0.07) # Plot convergence in 1D if show_plot: @@ -663,11 +712,12 @@ def rho2_pulled(e1, e2, e3): if __name__ == "__main__": - # direction = 0 - # bc_type = "dirichlet" - mapping = ["Cuboid", {"l1": 0.0, "r1": 4.0, "l2": 0.0, "r2": 2.0, "l3": 0.0, "r3": 3.0}] + direction = 0 + bc_type = "inhom_dirichlet" + mapping = ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}] + # mapping = ["Cuboid", {"l1": 0.0, "r1": 4.0, "l2": 0.0, "r2": 2.0, "l3": 0.0, "r3": 3.0}] # mapping = ['Orthogonal', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 3.}] - # test_poisson_1d(direction, bc_type, mapping, projected_rhs=True, show_plot=True) + test_poisson_1d(direction, bc_type, mapping, projected_rhs=True, show_plot=True) # num_elements = [64, 64, 1] # degree = [2, 2, 1] @@ -677,4 +727,4 @@ def rho2_pulled(e1, e2, e3): # mapping = ['Colella', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 1.}] # test_poisson_2d(num_elements, degree, bc_type, mapping, projected_rhs=True, show_plot=True) - test_poisson_accum_1d(mapping, do_plot=True) + # test_poisson_accum_1d(mapping, do_plot=True) diff --git a/src/struphy/propagators/tests/test_two_fluid_quasi_neutral.py b/src/struphy/propagators/tests/test_two_fluid_quasi_neutral.py index f89928679..999cc3dce 100644 --- a/src/struphy/propagators/tests/test_two_fluid_quasi_neutral.py +++ b/src/struphy/propagators/tests/test_two_fluid_quasi_neutral.py @@ -1,23 +1,23 @@ import logging import cunumpy as xp -from cunumpy import pi, cos, sin, zeros_like, ones_like import matplotlib.pyplot as plt import pytest +from cunumpy import cos, ones_like, pi, sin, zeros_like from feectools.ddm.mpi import mpi as MPI from struphy import domains, equils, grids, set_logging_level -from struphy.initial.base import Perturbation, GenericPerturbation +from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import L2Projector, WeightedMassOperators from struphy.feec.psydac_derham import Derham +from struphy.fields_background.projected_equils import ProjectedMHDequilibrium from struphy.geometry.base import Domain +from struphy.initial.base import GenericPerturbation, Perturbation from struphy.io.options import DerhamOptions from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable -from struphy.propagators.two_fluid_quasi_neutral_full import TwoFluidQuasiNeutralFull from struphy.propagators.base import Propagator -from struphy.fields_background.projected_equils import ProjectedMHDequilibrium -from struphy.feec.basis_projection_ops import BasisProjectionOperators +from struphy.propagators.two_fluid_quasi_neutral_full import TwoFluidQuasiNeutralFull logger = logging.getLogger("struphy") set_logging_level(logging.INFO) @@ -27,7 +27,7 @@ # plt.rcParams.update({'font.size': 22}) -pytest.mark.parametrize("bc_type", ["periodic", "hom_dirichlet", "inhom_dirichlet"]) +@pytest.mark.parametrize("bc_type", ["periodic", "hom_dirichlet", "inhom_dirichlet"]) @pytest.mark.parametrize( "mapping", [ @@ -47,8 +47,8 @@ def test_one_time_step(bc_type: str, mapping: None | list[str, dict], show_plot= dom_params = mapping[1] domain_class = getattr(domains, dom_type) domain: Domain = domain_class(**dom_params) - - # other options + + # other options B0 = 0 nu = 10.0 nu_e = 1.0 @@ -59,7 +59,7 @@ def test_one_time_step(bc_type: str, mapping: None | list[str, dict], show_plot= Tend = 1 sigma = 0 tol = 1e-5 - + # derham sequence if bc_type == "periodic": derham_opts = DerhamOptions(degree=degree, bcs=(None, None, None)) @@ -71,20 +71,20 @@ def test_one_time_step(bc_type: str, mapping: None | list[str, dict], show_plot= derham_opts = DerhamOptions(degree=degree, bcs=(("dirichlet", "dirichlet"), None, None)) lifting_function_u = GenericPerturbation(lambda x, y, z: x + 1, comp=0, given_in_basis="physical") lifting_function_ue = GenericPerturbation(lambda x, y, z: x, comp=0, given_in_basis="physical") - + else: raise ValueError(f"Invalid bc_type: {bc_type}") - + grid = grids.TensorProductGrid(num_elements=Nel) derham = Derham(grid=grid, options=derham_opts, domain=domain) - + # fluid background equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) projected_equil = ProjectedMHDequilibrium(equil=equil, derham=derham) - + # mass operators mass_ops = WeightedMassOperators(derham=derham, domain=domain, eq_mhd=equil) - + # basis operators basis_ops = BasisProjectionOperators(derham, domain, eq_mhd=equil) @@ -165,7 +165,6 @@ def source_function_ue(x, y, z): fz = zeros_like(x) return fx, fy, fz - # ---- perturbation classes for MMS initial conditions ---- class MMSIonVelocity(Perturbation): def __init__(self, comp=0): @@ -175,7 +174,6 @@ def __init__(self, comp=0): def __call__(self, x, y, z): return mms_ion_u(x, y, z)[self.comp] - class MMSElectronVelocity(Perturbation): def __init__(self, comp=0): self.comp = comp @@ -184,7 +182,6 @@ def __init__(self, comp=0): def __call__(self, x, y, z): return mms_electron_u(x, y, z)[self.comp] - class MMSPotential(Perturbation): def __init__(self): self.given_in_basis = "physical" @@ -198,9 +195,9 @@ def __call__(self, x, y, z): Propagator.mass_ops = mass_ops Propagator.basis_ops = basis_ops Propagator.projected_equil = projected_equil - + prop = TwoFluidQuasiNeutralFull(allocate_variables=True) - + prop.options = prop.Options( nu=nu, nu_e=nu_e, @@ -211,25 +208,25 @@ def __call__(self, x, y, z): solver="gmres", solver_params=SolverParameters(info=True, tol=tol), ) - + prop.allocate() if bc_type == "inhom_dirichlet": prop.variables.u.lifting_function = lifting_function_u prop.variables.ue.lifting_function = lifting_function_ue - + prop(dt) - + x = xp.linspace(0, 1, 100) mms_phi_x, _, _ = mms_phi(x, x * 0, x * 0) mms_ion_ux, _, _ = mms_ion_u(x, x * 0, x * 0) mms_el_ux, _, _ = mms_electron_u(x, x * 0, x * 0) - + e1 = xp.linspace(0, 1, 64) - + if show_plot: plt.figure(figsize=(18, 8)) - + plt.subplot(1, 3, 1) plt.plot(e1, prop.variables.u.spline(e1, 0.5, 0.5, squeeze_out=True)[0], label="numerical") plt.plot(x, mms_ion_ux, "--", label="manufactured") @@ -239,7 +236,7 @@ def __call__(self, x, y, z): # plt.title(f"{title} at t={t:.3f}") plt.legend() plt.grid(True) - + plt.subplot(1, 3, 2) plt.plot(e1, prop.variables.ue.spline(e1, 0.5, 0.5, squeeze_out=True)[0], label="numerical") plt.plot(x, mms_el_ux, "--", label="manufactured") @@ -249,7 +246,7 @@ def __call__(self, x, y, z): # plt.title(f"{title} at t={t:.3f}") plt.legend() plt.grid(True) - + plt.subplot(1, 3, 3) plt.plot(e1, prop.variables.phi.spline(e1, 0.5, 0.5, squeeze_out=True), label="numerical") plt.plot(x, mms_phi_x, "--", label="manufactured") @@ -259,10 +256,9 @@ def __call__(self, x, y, z): # plt.title(f"{title} at t={t:.3f}") plt.legend() plt.grid(True) - + plt.show() - - + if __name__ == "__main__": - test_one_time_step(bc_type="inhom_dirichlet", mapping=None, show_plot=True) \ No newline at end of file + test_one_time_step(bc_type="inhom_dirichlet", mapping=None, show_plot=True) diff --git a/src/struphy/propagators/two_fluid_quasi_neutral_full.py b/src/struphy/propagators/two_fluid_quasi_neutral_full.py index d7beff804..d355068a8 100644 --- a/src/struphy/propagators/two_fluid_quasi_neutral_full.py +++ b/src/struphy/propagators/two_fluid_quasi_neutral_full.py @@ -90,17 +90,16 @@ def phi(self, new): def __init__(self, allocate_variables: bool = False): self.variables = self.Variables() - + if allocate_variables: self.variables.u = FEECVariable(space="Hdiv") self.variables.ue = FEECVariable(space="Hdiv") self.variables.phi = FEECVariable(space="L2") - + self.variables.u.allocate(derham=self.derham, domain=self.domain, equil=self.projected_equil.equil) self.variables.ue.allocate(derham=self.derham, domain=self.domain, equil=self.projected_equil.equil) self.variables.phi.allocate(derham=self.derham, domain=self.domain, equil=self.projected_equil.equil) - # ========================================================================= ### Options # ========================================================================= From f70ebd0f209b1df95016be88d89d72fd4b7be03d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 2 Jun 2026 10:43:59 +0200 Subject: [PATCH 39/41] set default stab_eps in PoissonFieldSolve to 1e-14 (as it was before) --- .../tests/verification/test_verif_VlasovAmpereOneSpecies.py | 2 ++ src/struphy/propagators/poisson_field_solve.py | 2 +- src/struphy/propagators/tests/test_poisson.py | 2 +- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py index 0e7efc979..58bf8f624 100644 --- a/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py @@ -23,10 +23,12 @@ grids, maxwellians, perturbations, + set_logging_level, ) from struphy.models import VlasovAmpereOneSpecies logger = logging.getLogger("struphy") +set_logging_level(logging.WARNING) def test_weak_Landau(do_plot: bool = False): diff --git a/src/struphy/propagators/poisson_field_solve.py b/src/struphy/propagators/poisson_field_solve.py index 1ad3026cc..763e092cd 100644 --- a/src/struphy/propagators/poisson_field_solve.py +++ b/src/struphy/propagators/poisson_field_solve.py @@ -105,7 +105,7 @@ class Options(OptionsBase): OptsStabMat = Literal["M0", "M0ad", "Id"] OptsDiffusionMat = Literal["M1", "M1perp", "M1para", "M1gyro"] # propagator options - stab_eps: float = 0.0 + stab_eps: float = 1e-14 stab_mat: OptsStabMat = "Id" diffusion_mat: OptsDiffusionMat = "M1" rho: FEECVariable | Callable | tuple[AccumulatorVector, Particles] | list = None diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index fb148517d..0c65ca83d 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -59,7 +59,7 @@ def test_poisson_1d( Test the convergence of Poisson solver in 1D by means of manufactured solutions. """ # stabilization (removed for Dirichlet boundary conditions -> well-posed) - stab_eps = 1e-12 + stab_eps = 1e-14 if "dirichlet" in bc_type: stab_eps = 0.0 From 8c303986ebd7066134780ffbf242d6129fdb0440 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 2 Jun 2026 10:54:03 +0200 Subject: [PATCH 40/41] remove occurences of old PoissonFieldSolve --- src/struphy/models/hasegawa_wakatani.py | 4 ++-- src/struphy/models/poisson.py | 2 +- src/struphy/models/toy_drift.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/struphy/models/hasegawa_wakatani.py b/src/struphy/models/hasegawa_wakatani.py index a997a1cd7..2e5b9b01f 100644 --- a/src/struphy/models/hasegawa_wakatani.py +++ b/src/struphy/models/hasegawa_wakatani.py @@ -43,7 +43,7 @@ class HasegawaWakatani(StruphyModel): :ref:`propagators` (called in sequence): - 1. :class:`~struphy.propagators.poisson_field_solve.PoissonFieldSolve` + 1. :class:`~struphy.propagators.poisson_solve.PoissonSolve` 2. :class:`~struphy.propagators.hasegawa_wakatani_step.HasegawaWakataniStep` :ref:`Model info `: @@ -153,7 +153,7 @@ def doc_discretization(cls): 1. :class:`~struphy.propagators.poisson_solve.PoissonSolve` 2. :class:`~struphy.propagators.hasegawa_wakatani_step.HasegawaWakataniStep` """ - doc = rf"""**1. PoissonFieldSolve:** + doc = rf"""**1. PoissonSolve:** {PoissonSolve.__doc__} diff --git a/src/struphy/models/poisson.py b/src/struphy/models/poisson.py index 76338764f..2dd8f4052 100644 --- a/src/struphy/models/poisson.py +++ b/src/struphy/models/poisson.py @@ -117,7 +117,7 @@ def doc_discretization(cls): {TimeDependentSource.__doc__} -**2. PoissonFieldSolve:** +**2. PoissonSolve:** {PoissonSolve.__doc__} """ diff --git a/src/struphy/models/toy_drift.py b/src/struphy/models/toy_drift.py index 3051c41c0..464e457f3 100644 --- a/src/struphy/models/toy_drift.py +++ b/src/struphy/models/toy_drift.py @@ -242,7 +242,7 @@ def doc_discretization(cls): 1. :class:`~struphy.propagators.poisson_solve.PoissonSolve` 2. :class:`~struphy.propagators.push_guiding_center_bx_estar.PushGuidingCenterBxEstar` """ - doc = rf"""**1. PoissonFieldSolve:** + doc = rf"""**1. PoissonSolve:** {PoissonSolve.__doc__} From 2839254cf456b401d88c7bf3c7941a7383f2c994 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 2 Jun 2026 14:50:44 +0200 Subject: [PATCH 41/41] added a Poisson solution with inhom. Dirichlet bcs to tutorial_for_devs_01_feec --- tutorials/tutorial_for_devs_01_feec.ipynb | 122 +++++++++++++++++----- 1 file changed, 98 insertions(+), 24 deletions(-) diff --git a/tutorials/tutorial_for_devs_01_feec.ipynb b/tutorials/tutorial_for_devs_01_feec.ipynb index ce18f371c..00d5ec81c 100644 --- a/tutorials/tutorial_for_devs_01_feec.ipynb +++ b/tutorials/tutorial_for_devs_01_feec.ipynb @@ -482,31 +482,50 @@ "mfct_solution = lambda e1: 1/(np.pi/2)**2 * np.cos(np.pi/2 * e1) - 0.5" ] }, + { + "cell_type": "markdown", + "id": "38", + "metadata": {}, + "source": [ + "In order to capture the non-zero Dirichlet condition at `e1=0.0`, we need to use a lifting function. This is a function that satisfies the boundary conditions, and is added to the solution of the homogeneous problem. Let us define a linear lifting function for the current problem:" + ] + }, { "cell_type": "code", "execution_count": null, - "id": "38", + "id": "39", "metadata": {}, "outputs": [], "source": [ "from struphy.initial.base import GenericPerturbation\n", "\n", - "tmp = lambda e1, e2, e3: mfct_solution(e1)\n", - "fun_lift = GenericPerturbation(tmp)\n", + "bc_at_0 = 1/(np.pi/2)**2 - 0.5\n", + "bc_at_1 = -0.5\n", + "\n", + "lifting_function = GenericPerturbation(lambda e1, e2, e3: e1*bc_at_1 + (1 - e1)*bc_at_0)\n", "\n", "e1 = np.linspace(0, 1, 100)\n", "e2 = 0.5\n", "e3 = 0.5\n", "\n", - "plt.plot(e1, fun_lift(e1, e2, e3))\n", + "plt.plot(e1, lifting_function(e1, e2, e3))\n", "plt.xlabel('e1')\n", + "plt.title('Lifting function (to satisfy the non-zero Dirichlet BCs)')\n", "plt.show()" ] }, + { + "cell_type": "markdown", + "id": "40", + "metadata": {}, + "source": [ + "Let us instantiate the Poisson model and set the density `rho` on the right-hand side:" + ] + }, { "cell_type": "code", "execution_count": null, - "id": "39", + "id": "41", "metadata": {}, "outputs": [], "source": [ @@ -514,9 +533,43 @@ "derham_options = DerhamOptions(bcs=((\"free\", \"dirichlet\"), None, None))\n", "\n", "poisson = Poisson()\n", - "poisson.propagators.poisson.options = poisson.propagators.poisson.Options(rho=fun)\n", - "poisson.em_fields.phi.lifting_function = fun_lift\n", - "\n", + "poisson.propagators.poisson.options = poisson.propagators.poisson.Options(stab_eps=0.0, rho=fun)" + ] + }, + { + "cell_type": "markdown", + "id": "42", + "metadata": {}, + "source": [ + "We now pass the lifting function to the variable in question, which is `phi` in this case:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "43", + "metadata": {}, + "outputs": [], + "source": [ + "poisson.em_fields.phi.lifting_function = lifting_function" + ] + }, + { + "cell_type": "markdown", + "id": "44", + "metadata": {}, + "source": [ + "The rest is taken care of internally. If no lifting function is provided, the solver will simply solve the homogeneous problem, which is what we did in the previous section. \n", + "Let us run the simulation and plot the results:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "45", + "metadata": {}, + "outputs": [], + "source": [ "sim = Simulation(model=poisson,\n", " grid=grid,\n", " derham_opts=derham_options,\n", @@ -526,17 +579,51 @@ { "cell_type": "code", "execution_count": null, - "id": "40", + "id": "46", "metadata": {}, "outputs": [], "source": [ - "sim.allocate()" + "sim.run(one_time_step=True)" ] }, { "cell_type": "code", "execution_count": null, - "id": "41", + "id": "47", + "metadata": {}, + "outputs": [], + "source": [ + "phi_full = poisson.em_fields.phi.spline_full(e1h, 0.5, 0.5, squeeze_out=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "48", + "metadata": {}, + "outputs": [], + "source": [ + "# phi = sim.plotting_data.spline_values.em_fields.phi_log\n", + "# print(phi)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "49", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(e1, mfct_solution(e1), label=\"exact\")\n", + "plt.plot(e1h, phi_full, \"go\", label=\"numerical solution\")\n", + "plt.xlabel('e1')\n", + "plt.legend()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "50", "metadata": {}, "outputs": [], "source": [ @@ -560,19 +647,6 @@ "\n", "plt.show()" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "42", - "metadata": {}, - "outputs": [], - "source": [ - "print(var.spline.space)\n", - "print(var.spline.derham.bcs)\n", - "print(var.boundary_spline.space)\n", - "print(var.boundary_spline.derham.bcs)" - ] } ], "metadata": {