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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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)