From 91ce5ca275f95ea32a9f1e532dd715077d87c806 Mon Sep 17 00:00:00 2001 From: Monoclod Date: Fri, 12 Jun 2026 16:25:34 +0200 Subject: [PATCH 1/6] Added new propagator for the full perturbation system --- .../propagators/perturbation_system_full.py | 663 ++++++++++++++++++ 1 file changed, 663 insertions(+) create mode 100644 src/struphy/propagators/perturbation_system_full.py diff --git a/src/struphy/propagators/perturbation_system_full.py b/src/struphy/propagators/perturbation_system_full.py new file mode 100644 index 000000000..1abda4e00 --- /dev/null +++ b/src/struphy/propagators/perturbation_system_full.py @@ -0,0 +1,663 @@ +import logging +from dataclasses import dataclass +from typing import Callable, get_args +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 +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 L2Projector, WeightedMassOperators +from struphy.feec.utilities import LocalRotationMatrix +from struphy.io.options import LiteralOptions, OptionsBase +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 + +logger = logging.getLogger("struphy") + + +class ColdPlasmaPerturbation(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 (electron density rhosin and rhocos, electron velocity usin and ucos, electric field Esin and Ecos, magnetic field Bsin and Bcos) + # ========================================================================= + + class Variables: + """Container for variables advanced by :class:`ColdPlasmaPerturbation`. + + Attributes + ---------- + rhosin : FEECVariable or None + Sine component of electron density variable in ``"H1"`` space. + rhocos : FEECVariable or None + Cosine component of electron density variable in ``"H1"`` space. + usin : FEECVariable or None + Sine component of the electron velocity variable in ``"Hcurl"`` space. + ucos : FEECVariable or None + Cosine component of the electron velocity variable in ``"Hcurl"`` space. + Esin : FEECVariable or None + Sine component of the electric field variable in ``"Hcurl"`` space. + Ecos : FEECVariable or None + Cosine component of the electric field variable in ``"Hcurl"`` space. + Bsin : FEECVariable or None + Sine component of the magnetic field variable in ``"Hdiv"`` space. + Bcos : FEECVariable or None + Cosine component of the magnetic field variable in ``"Hdiv"`` space. + """ + + def __init__(self) -> None: + self._rhosin: FEECVariable | None = None + self._rhocos: FEECVariable | None = None + self._usin: FEECVariable | None = None + self._ucos: FEECVariable | None = None + self._Esin: FEECVariable | None = None + self._Ecos: FEECVariable | None = None + self._Bsin: FEECVariable | None = None + self._Bcos: FEECVariable | None = None + + @property + def rhosin(self) -> FEECVariable | None: + return self._rhosin + + @rhosin.setter + def rhosin(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1" + self._rhosin = new + + @property + def rhocos(self) -> FEECVariable | None: + return self._rhocos + + @rhocos.setter + def rhocos(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1" + self._rhocos = new + + @property + def usin(self) -> FEECVariable | None: + return self._usin + + @usin.setter + def usin(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._usin = new + + @property + def ucos(self) -> FEECVariable | None: + return self._ucos + + @ucos.setter + def ucos(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._ucos = new + + @property + def Esin(self) -> FEECVariable | None: + return self._Esin + + @Esin.setter + def Esin(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._Esin = new + + @property + def Ecos(self) -> FEECVariable | None: + return self._Ecos + + @Ecos.setter + def Ecos(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._Ecos = new + + @property + def Bsin(self) -> FEECVariable | None: + return self._Bsin + + @Bsin.setter + def Bsin(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._Bsin = new + + @property + def Bcos(self) -> FEECVariable | None: + return self._Bcos + + @ucos.setter + def Bcos(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._Bcos = new + + + def __init__(self): + self.variables = self.Variables() + + # ========================================================================= + ### Options + # ========================================================================= + + @dataclass(repr=False) + class Options(OptionsBase): + """Configuration options for :class:`ColdPlasmaPerturbation`. + + Parameters + ---------- + J : FEECVariable in ``"Hcurl"`` or list + Cosine component of the source term. + omega : float, default=1.0 + Source term oscillation frequency. + curlcurl_lambda : float, default=1.0 + Coefficient in the curl-curl operator. + mass : float, default=1.0 + Electron mass in relative unis. + mu : Callable or float, default=1.0 + Electron viscosity coefficient. + nu : Callable or float, default=1.0 + Electron-Ion collision frequency. + rhobar : FEECVariable in ``"H1"`` or Callable or float, default=1.0 + Average electron mass density. + theta : FEECVariable in ``"H1"`` or Callable or float, default=1.0 + Average electron temperature. + Ebar : FEECVariable in ``"Hcurl"`` or list + Average electrostatic field. + Esin0 : StencilVector, default=None + Initial Esin guess for the iterative linear solver. + Ecos0 : StencilVector, default=None + Initial Ecos guess for the iterative linear solver. + solver : LiteralOptions.OptsGenSolver, default="gmres" + Linear/saddle-point solver used for the global system. + solver_params : SolverParameters or None, default=None + Solver controls. + """ + + J: FEECVariable | list + omega: float = 1.0 + curlcurl_lambda: float = 1.0 + mass: float = 1.0 + mu: Callable | float = 1.0 + nu: Callable | float = 1.0 + rhobar: FEECVariable | Callable | float = 1.0 + theta: FEECVariable | Callable | float = 1.0 + Ebar: FEECVariable | list + + Esin0: FEECVariable | StencilVector = None + Ecos0: FEECVariable | StencilVector = None + + solver: LiteralOptions.OptsGenSolver = "gmres" + solver_params: SolverParameters | None = None + + def __post_init__(self): + # input format correctness + assert self.J is not None + if (not isinstance(self.J, (FEECVariable, list))): + raise TypeError(f"J must be either a Hcurl FEECVariable or list of Callables, got {type(self.J)}") + if isinstance(self.J, FEECVariable): + assert self.J.space == "Hcurl" + if isinstance(self.J,list): + assert len(self.J) == 3 + for ji in self.J: + assert isinstance(ji, Callable) + + if (self.rhobar is not None) and (not isinstance(self.rhobar, (FEECVariable, Callable, float))): + raise TypeError(f"rhobar must be either a H1 FEECVariable or a Callable or a float, got {type(self.rhobar)}") + if isinstance(rhobar, FEECVariable): + assert rhobar.space == "H1" + + if (self.theta is not None) and (not isinstance(self.theta, (FEECVariable, Callable, float))): + raise TypeError(f"theta must be either a H1 FEECVariable or a Callable or a float, got {type(self.theta)}") + if isinstance(self.theta, FEECVariable): + assert self.theta.space == "H1" + + assert self.Ebar is not None + if (not isinstance(self.Ebar,(FEECVariable, list))): + raise TypeError(f"Ebar must be either a Hcurl FEECVariable or list of Callables, got {type(self.Ebar)}") + if isinstance(self.Ebar, FEECVariable): + assert self.Ebar.space == "Hcurl" + if isinstance(self.Ebar,list): + assert len(self.Ebar) == 3 + for ei in self.Ebar: + assert isinstance(ei, Callable) + + # --- physical parameter sanity checks --- + if self.omega <= 0: + raise ValueError(f"omega must be positive, got {self.omega}") + if self.curlcurl_lambda <= 0: + raise ValueError(f"curlcurl_lambda must be positive, got {self.curlcurl_lambda}") + if self.mass <= 0: + raise ValueError(f"mass must be positive, got {self.mass}") + if isinstance(mu, float) and self.mu < 0: + raise ValueError(f"mu must be non-negative, got {self.mu}") + if isinstance(nu, float) and self.nu < 0: + raise ValueError(f"nu must be non-negative, got {self.nu}") + + # --- E initial guess correct space check --- + if isinstance(self.Esin0, FEECVariable): + assert self.Esin0.space == "Hcurl" + if isinstance(self.Ecos0, FEECVariable): + assert self.Ecos0.space == "Hcurl" + + 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) + self._options = new + logger.info(f"\nNew options for propagator '{self.__class__.__name__}':\n{self._options}") + + # ========================================================================= + ### Allocate + # ========================================================================= + + def allocate(self): + + # ---- source term vector (for RHS assembly) --------------------------- + + self._j: StencilVector + + if isinstance(self._options.J,FEECVariable): + self._j = self._options.J.spline.vector + else: + self._j = self.derham.P1(self._options.J) # works if J is a list of Callables + + # ---- unconstrained operators (for RHS assembly) ---------------------- + + self._M0 = self.mass_ops.M0 + self._M1 = self.mass_ops.M1 + self._M2 = self.mass_ops.M2 + self._grad = self.derham.grad + self._curl = self.derham.curl + self._div = self.derham.div + + self._M0inv = inverse( + self._M0, + "pcg", + pc="MassMatrixPreconditioner", + tol=self._options.solver_params.tol, + maxiter=self._options.solver_params.maxiter, + verbose = False, + recycle = self._options.solver_params.recycle, + ) + + self._M1mu: WeightedMassOperators + self._M2mu: WeightedMassOperators + self._M3mu: WeightedMassOperators + + if isinstance(self._options.mu, float): + self._M1mu = self._options.mu * self._M1 + self._M2mu = self._options.mu * self._M2 + self._M3mu = self._options.mu * self.mass_ops.M3 + else: + self._M1mu = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + self._options.mu, + ), + name = "M1mu", + assemble = True, + ) + + self._M2mu = self.mass_ops.create_weighted_mass( + "Hdiv", + "Hdiv", + weights=( + "G", + "1/sqrt_g", + self._options.mu, + ), + name = "M2mu", + assemble = True, + ) + + self._M3mu = self.mass_ops.create_weighted_mass( + "L2", + "L2", + weights=( + "1/sqrt_g", + self._options.mu, + ), + name = "M3mu", + assemble = True, + ) + + + self._M1rho: WeightedMassOperators + self._M1xrhoB: WeightedMassOperators + + rot_B = LocalRotationMatrix( + self.eq_mhd.b2_1, + self.eq_mhd.b2_2, + self.eq_mhd.b2_3, + ) + + rhoB1: Callable + rhoB2: Callable + rhoB3: Callable + + if isinstance(self._options.rhobar, float): + self._M1rho = self._options.rhobar * self.mass_ops.M1 + self._M1xrhoB = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + rot_B, + "Ginv", + "sqrt_g", + ), + name = "M1_xrhoB", + assemble = True, + ) + self._M1xrhoB *= self._options.rhobar + + if isinstance(self._options.rhobar, Callable): + self._M1rho = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + self._options.rhobar, + ), + name = "M1rho", + assemble = True, + ) + + self._M1xrhoB = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + rot_B, + self._options.rhobar, + "Ginv", + "sqrt_g", + ), + name = "M1_xrhoB", + assemble = True, + ) + + if isinstance(self._options.rhobar,FEECVariable): + self._M1rho = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + self._options.rhobar.spline, + ), + name = "M1rho", + assemble = True, + ) + + self._M1xrhoB = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + rot_B, + self._options.rhobar.spline, + "Ginv", + "sqrt_g", + ), + name = "M1_xrhoB", + assemble = True, + ) + + self._M1rho_inv = inverse( + self._M1rho, + "pcg", + pc="MassMatrixPreconditioner", + tol=self._options.solver_params.tol, + maxiter=self._options.solver_params.maxiter, + verbose=False, + recycle=self._options.solver_params.recycle, + ) + + + self._M1nurho: WeightedMassOperators + + if isinstance(self._options.nu, float): + self._M1nurho = self._options.nu * self._M1rho + + if isinstance(self._options.nu, Callable): + nurho: Callable + if isinstance(self._options.rhobar, float): + nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar + if isinstance(self._options.rhobar,Callable): + nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar(*etas) + if isinstance(self._options.rhobar,FEECVariable): + nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar.spline(*etas) + + self._M1nurho = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + nurho, + ), + name = "M1nurho", + assemble = True, + ) + + + self._P00theta: BasisProjectionOperators + + if isinstance(self._options.theta, float): + self._P00theta = self._options.theta * IdentityOperator(self.derham.V0) + + if isinstance(self._options.theta, Callable): + self._P00theta = self.basis_ops.create_basis_op( + [[self._options.theta]], + "H1", + "H1", + assemble = True, + name = "P00theta", + ) + + if isinstance(self._options.theta, FEECVariable): + self._P00theta = self.basis_ops.create_basis_op( + [[self._options.theta.spline]], + "H1", + "H1", + assemble = True, + name = "P00theta", + ) + + self._P01Ebar: BasisProjectionOperators + + if isinstance(self._options.Ebar, list): + self._P01Ebar = self.basis_ops.create_basis_op( + [[self._options.Ebar[0]],[self._options.Ebar[1]],[self._options.Ebar[2]]], + "H1", + "Hcurl", + assemble = True, + name = "P01Ebar", + ) + + if isinstance(self._options.Ebar, FEECVariable): + Ebar1 = lambda *etas: self._options.Ebar.spline(etas)[0] + Ebar2 = lambda *etas: self._options.Ebar.spline(etas)[1] + Ebar3 = lambda *etas: self._options.Ebar.spline(etas)[2] + + self._P01Ebar = self.basis_ops.create_basis_op( + [[Ebar1],[Ebar2],[Ebar3]], + "H1", + "Hcurl", + assemble = True, + name = "P01Ebar", + ) + + + self._P12 = self.basis_ops.U1 + + + ones = lambda *etas: 1.0 + 0 * etas + zeroes = lambda *etas: 0 * etas + + self._O1 = self.basis_ops.create_basis_op( + [[ones, zeroes, zeroes]], + "Hcurl", + "H1", + assemble = True, + name = "O1", + ) + self._O2 = self.basis_ops.create_basis_op( + [[zeroes, ones, zeroes]], + "Hcurl", + "H1", + assemble = True, + name = "O1", + ) + self._O3 = self.basis_ops.create_basis_op( + [[zeroes, zeroes, ones]], + "Hcurl", + "H1", + assemble = True, + name = "O1", + ) + + + self._Acurlcurl = self._curl.T @ self._M2 @ self._curl - self._options.curlcurl_lambda * self._M1 + + self._divPi = - self._curl.T @ self._M2mu @ self._curl \ + - 2/3 * self._P12.T @ self._div.T @ self._M3mu @ self._div @ self._P12 \ + + 2 * self._O1.T @ self._grad.T @ self._M1mu @ self._grad @ self._O1 \ + + 2 * self._O2.T @ self._grad.T @ self._M1mu @ self._grad @ self._O2 \ + + 2 * self._O3.T @ self._grad.T @ self._M1mu @ self._grad @ self._O3 + + self._A = self._M1rho + self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) @ self._M0inv @ self._grad.T @ self._M1rho \ + / (self._options.mass * self._options.omega * self._options.omega) + + self._B = (self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho) / self._options.omega + + # ---- block saddle-point system ---------------------------------------- + + self._block_domain = BlockVectorSpace(self.derham.V1, self.derham.V1) + self._block_codomain = self._block_domain + + self._coupled_equations_matrix = BlockLinearOperator( + self._block_domain, self._block_codomain, blocks=[[self._A, -self._B], [self._B, self._A]] + ) + + M1rhoinv_j = self._M1rho_inv.solve(self._j) + A_j = self._A.dot(M1rhoinv_j) + minusB_j = - self.B.dot(M1rhoinv_j) + + self._calEsin0: StencilVector = None + self._calEcos0: StencilVector = None + + # --- copy current state --- + Esin0 = self.variables.Esin.spline.vector + Ecos0 = self.variables.Ecos.spline.vector + + self._calEsin0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Esin0)) + self._calEcos0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Ecos0)) + + self._calE0 = BlockVector(self._block_domain, blocks=[self._calEsin0, self._calEcos0]) + + self._coupled_equations_matrix_inverse = inverse( + self._coupled_equations_matrix, + solver="gmres", + x0=self._calE0, + tol=self._options.solver_params.tol, + maxiter=self._options.solver_params.maxiter, + verbose=False, + recycle=True, + ) + + # --- build inverses of the curl-curl matrices with good initial guesses --- + self._Acurlcurl_inv_sin = inverse( + self._Acurlcurl, + solver="pcg", + x0=Esin0, + tol=self._options.solver_params.tol, + maxiter=self._options.solver_params.maxiter, + verbose=False, + recycle=True, + ) + + self._Acurlcurl_inv_cos = inverse( + self._Acurlcurl, + solver="pcg", + x0=Ecos0, + tol=self._options.solver_params.tol, + maxiter=self._options.solver_params.maxiter, + verbose=False, + recycle=True, + ) + + + # ========================================================================= + ### Equation solve + # ========================================================================= + + def __call__(self, dt): + + # --- calculate auxilliary vectors --- + _calE = self._coupled_equations_matrix_inverse.solve(self._calE_RHS) + + _calEsin = _calE[0] + _calEcos = _calE[1] + + _m_curlcurlEsin = self._M1rho.dot(_calEsin) + _m_curlcurlEcos = self._M1rho.dot(_calEcos) + + _M1rho_usin = _m_curlcurlEcos.copy() + _M1rho_ucos = self._options.mass * self._j - _m_curlcurlEsin + + # --- calculate solutions --- + self._Esin.vector = self._Acurlcurl_inv_sin.solve(_m_curlcurlEsin / self._options.mass) + self._Ecos.vector = self._Acurlcurl_inv_cos.solve(_m_curlcurlEcos / self._options.mass) + + self._Bsin.vector = - self._curl.dot(self._Ecos.vector) / self._options.omega + self._Bcos.vector = self._curl.dot(self._Esin.vector) / self._options.omega + + self._usin.vector = self._M1rho_inv.solve(_M1rho_usin) + self._ucos.vector = self._M1rho_inv.solve(_M1rho_ucos) + + self._rhosin.vector = self._M0inv.solve(self._grad.T.dot(_M1rho_ucos)) / self._options.omega + self._rhocos.vector = - self._M0inv.solve(self._grad.T.dot(_M1rho_usin)) / self._options.omega + + # --- update FEEC variables --- + self.update_feec_variables( + rhosin=self._rhosin.vector, rhocos=self._rhocos.vector, + usin=self._usin.vector, ucos=self._ucos.vector, + Esin=self._Esin.vector, Ecos=self._Ecos.vector, + Bsin=self._Bsin.vector, Bcos=self._Bcos.vector) + From 94c761edc6c4de2229a53d3dd0916b7be5218fc1 Mon Sep 17 00:00:00 2001 From: Monoclod Date: Fri, 19 Jun 2026 16:13:32 +0200 Subject: [PATCH 2/6] Updated the design of the solver --- .../propagators/perturbation_system_full.py | 145 +++++++++++++----- 1 file changed, 108 insertions(+), 37 deletions(-) diff --git a/src/struphy/propagators/perturbation_system_full.py b/src/struphy/propagators/perturbation_system_full.py index 1abda4e00..9e1cf41b0 100644 --- a/src/struphy/propagators/perturbation_system_full.py +++ b/src/struphy/propagators/perturbation_system_full.py @@ -561,67 +561,138 @@ def allocate(self): + 2 * self._O2.T @ self._grad.T @ self._M1mu @ self._grad @ self._O2 \ + 2 * self._O3.T @ self._grad.T @ self._M1mu @ self._grad @ self._O3 - self._A = self._M1rho + self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) @ self._M0inv @ self._grad.T @ self._M1rho \ - / (self._options.mass * self._options.omega * self._options.omega) + # self._A = self._M1rho + self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) @ self._M0inv @ self._grad.T @ self._M1rho \ + # / (self._options.mass * self._options.omega * self._options.omega) - self._B = (self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho) / self._options.omega + # self._B = (self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho) / self._options.omega # ---- block saddle-point system ---------------------------------------- - self._block_domain = BlockVectorSpace(self.derham.V1, self.derham.V1) - self._block_codomain = self._block_domain + self._block_V0 = BlockVectorSpace(self.derham.V0, self.derham.V0) + self._block_V1 = BlockVectorSpace(self.derham.V1, self.derham.V1) + self._block_V2 = BlockVectorSpace(self.derham.V2, self.derham.V2) - self._coupled_equations_matrix = BlockLinearOperator( - self._block_domain, self._block_codomain, blocks=[[self._A, -self._B], [self._B, self._A]] + self._block_source = BlockVector(self._block_V1, blocks=[self._M1.dot(self._j), None]) + + self._block_M = BlockLinearOperator( + self._block_V0, self._block_V0, blocks=[[self._M0, None], [None, self._M0]] + ) + + self._block_Divergence = BlockLinearOperator( + self._block_V1, self._block_V0, blocks=[[None, - self._grad.T @ self._M1rho], [self._grad.T @ self._M1rho, None]] ) - M1rhoinv_j = self._M1rho_inv.solve(self._j) - A_j = self._A.dot(M1rhoinv_j) - minusB_j = - self.B.dot(M1rhoinv_j) + self._block_Acurlcurl = BlockLinearOperator( + self.block_V, self._block_V1, blocks=[[self._Acurlcurl, None], [None, self._Acurlcurl]] + ) - self._calEsin0: StencilVector = None - self._calEcos0: StencilVector = None + self._block_B = BlockLinearOperator( + self._block_V1, self._block_V1, blocks=[[None, self._M1rho / self._options.mass], [- self._M1rho / self._options.mass, None]] + ) - # --- copy current state --- - Esin0 = self.variables.Esin.spline.vector - Ecos0 = self.variables.Ecos.spline.vector + self._block_P = BlockLinearOperator( + self._block_V0, self._block_V1, + blocks=[[None, self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) / self._options.mass], + [self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) / self._options.mass, None]] + ) - self._calEsin0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Esin0)) - self._calEcos0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Ecos0)) + self._block_Q = BlockLinearOperator( + self._block_V1, self._block_V1, + blocks=[[self._options.omega * self._M1rho, self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho], + [self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho, - self._options.omega * self._M1rho]] + ) - self._calE0 = BlockVector(self._block_domain, blocks=[self._calEsin0, self._calEcos0]) + self._block_R = BlockLinearOperator( + self._block_V1, self._block_V1, blocks=[[None, self._M1rho / self._options.mass], [self._M1rho / self._options.mass, None]] + ) - self._coupled_equations_matrix_inverse = inverse( - self._coupled_equations_matrix, - solver="gmres", - x0=self._calE0, + self._block_Minv = inverse( + self._block_M, + "pcg", + pc="MassMatrixPreconditioner", tol=self._options.solver_params.tol, maxiter=self._options.solver_params.maxiter, - verbose=False, - recycle=True, + verbose = False, + recycle = self._options.solver_params.recycle, ) - # --- build inverses of the curl-curl matrices with good initial guesses --- - self._Acurlcurl_inv_sin = inverse( - self._Acurlcurl, - solver="pcg", - x0=Esin0, + self._block_rhomatrix = - self._block_Minv @ self._block_Divergence / self._options.omega + + self._block_umatrix = - self._block_P @ self._block_rhomatrix - self._block_Q + + self._block_umatrix_inv = inverse( + self._block_umatrix, + "gmres", + pc=None, tol=self._options.solver_params.tol, maxiter=self._options.solver_params.maxiter, - verbose=False, - recycle=True, + verbose = False, + recycle = self._options.solver_params.recycle, ) - self._Acurlcurl_inv_cos = inverse( - self._Acurlcurl, - solver="pcg", - x0=Ecos0, + self._block_Ematrix = self._block_Acurlcurl / self._options.omega + self._block_B @ self._block_umatrix_inv @ self._block_R + + self._block_Ematrix_inv = inverse( + self._block_Ematrix, + "gmres", + pc=None, tol=self._options.solver_params.tol, maxiter=self._options.solver_params.maxiter, - verbose=False, - recycle=True, + verbose = False, + recycle = self._options.solver_params.recycle, ) + # self._coupled_equations_matrix = BlockLinearOperator( + # self._block_domain, self._block_codomain, blocks=[[self._A, -self._B], [self._B, self._A]] + # ) + + # M1rhoinv_j = self._M1rho_inv.solve(self._j) + # A_j = self._A.dot(M1rhoinv_j) + # minusB_j = - self.B.dot(M1rhoinv_j) + + # self._calEsin0: StencilVector = None + # self._calEcos0: StencilVector = None + + # # --- copy current state --- + # Esin0 = self.variables.Esin.spline.vector + # Ecos0 = self.variables.Ecos.spline.vector + + # self._calEsin0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Esin0)) + # self._calEcos0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Ecos0)) + + # self._calE0 = BlockVector(self._block_domain, blocks=[self._calEsin0, self._calEcos0]) + + # self._coupled_equations_matrix_inverse = inverse( + # self._coupled_equations_matrix, + # solver="gmres", + # x0=self._calE0, + # tol=self._options.solver_params.tol, + # maxiter=self._options.solver_params.maxiter, + # verbose=False, + # recycle=True, + # ) + + # # --- build inverses of the curl-curl matrices with good initial guesses --- + # self._Acurlcurl_inv_sin = inverse( + # self._Acurlcurl, + # solver="pcg", + # x0=Esin0, + # tol=self._options.solver_params.tol, + # maxiter=self._options.solver_params.maxiter, + # verbose=False, + # recycle=True, + # ) + + # self._Acurlcurl_inv_cos = inverse( + # self._Acurlcurl, + # solver="pcg", + # x0=Ecos0, + # tol=self._options.solver_params.tol, + # maxiter=self._options.solver_params.maxiter, + # verbose=False, + # recycle=True, + # ) + # ========================================================================= ### Equation solve From 3effcc38279780e9f9cda37955a028aa92dc99eb Mon Sep 17 00:00:00 2001 From: Monoclod Date: Fri, 19 Jun 2026 17:08:18 +0200 Subject: [PATCH 3/6] Renamed propagator --- .../propagators/perturbation_system_cold.py | 734 ++++++++++++++++++ 1 file changed, 734 insertions(+) create mode 100644 src/struphy/propagators/perturbation_system_cold.py diff --git a/src/struphy/propagators/perturbation_system_cold.py b/src/struphy/propagators/perturbation_system_cold.py new file mode 100644 index 000000000..9e1cf41b0 --- /dev/null +++ b/src/struphy/propagators/perturbation_system_cold.py @@ -0,0 +1,734 @@ +import logging +from dataclasses import dataclass +from typing import Callable, get_args +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 +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 L2Projector, WeightedMassOperators +from struphy.feec.utilities import LocalRotationMatrix +from struphy.io.options import LiteralOptions, OptionsBase +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 + +logger = logging.getLogger("struphy") + + +class ColdPlasmaPerturbation(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 (electron density rhosin and rhocos, electron velocity usin and ucos, electric field Esin and Ecos, magnetic field Bsin and Bcos) + # ========================================================================= + + class Variables: + """Container for variables advanced by :class:`ColdPlasmaPerturbation`. + + Attributes + ---------- + rhosin : FEECVariable or None + Sine component of electron density variable in ``"H1"`` space. + rhocos : FEECVariable or None + Cosine component of electron density variable in ``"H1"`` space. + usin : FEECVariable or None + Sine component of the electron velocity variable in ``"Hcurl"`` space. + ucos : FEECVariable or None + Cosine component of the electron velocity variable in ``"Hcurl"`` space. + Esin : FEECVariable or None + Sine component of the electric field variable in ``"Hcurl"`` space. + Ecos : FEECVariable or None + Cosine component of the electric field variable in ``"Hcurl"`` space. + Bsin : FEECVariable or None + Sine component of the magnetic field variable in ``"Hdiv"`` space. + Bcos : FEECVariable or None + Cosine component of the magnetic field variable in ``"Hdiv"`` space. + """ + + def __init__(self) -> None: + self._rhosin: FEECVariable | None = None + self._rhocos: FEECVariable | None = None + self._usin: FEECVariable | None = None + self._ucos: FEECVariable | None = None + self._Esin: FEECVariable | None = None + self._Ecos: FEECVariable | None = None + self._Bsin: FEECVariable | None = None + self._Bcos: FEECVariable | None = None + + @property + def rhosin(self) -> FEECVariable | None: + return self._rhosin + + @rhosin.setter + def rhosin(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1" + self._rhosin = new + + @property + def rhocos(self) -> FEECVariable | None: + return self._rhocos + + @rhocos.setter + def rhocos(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1" + self._rhocos = new + + @property + def usin(self) -> FEECVariable | None: + return self._usin + + @usin.setter + def usin(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._usin = new + + @property + def ucos(self) -> FEECVariable | None: + return self._ucos + + @ucos.setter + def ucos(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._ucos = new + + @property + def Esin(self) -> FEECVariable | None: + return self._Esin + + @Esin.setter + def Esin(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._Esin = new + + @property + def Ecos(self) -> FEECVariable | None: + return self._Ecos + + @Ecos.setter + def Ecos(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._Ecos = new + + @property + def Bsin(self) -> FEECVariable | None: + return self._Bsin + + @Bsin.setter + def Bsin(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._Bsin = new + + @property + def Bcos(self) -> FEECVariable | None: + return self._Bcos + + @ucos.setter + def Bcos(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._Bcos = new + + + def __init__(self): + self.variables = self.Variables() + + # ========================================================================= + ### Options + # ========================================================================= + + @dataclass(repr=False) + class Options(OptionsBase): + """Configuration options for :class:`ColdPlasmaPerturbation`. + + Parameters + ---------- + J : FEECVariable in ``"Hcurl"`` or list + Cosine component of the source term. + omega : float, default=1.0 + Source term oscillation frequency. + curlcurl_lambda : float, default=1.0 + Coefficient in the curl-curl operator. + mass : float, default=1.0 + Electron mass in relative unis. + mu : Callable or float, default=1.0 + Electron viscosity coefficient. + nu : Callable or float, default=1.0 + Electron-Ion collision frequency. + rhobar : FEECVariable in ``"H1"`` or Callable or float, default=1.0 + Average electron mass density. + theta : FEECVariable in ``"H1"`` or Callable or float, default=1.0 + Average electron temperature. + Ebar : FEECVariable in ``"Hcurl"`` or list + Average electrostatic field. + Esin0 : StencilVector, default=None + Initial Esin guess for the iterative linear solver. + Ecos0 : StencilVector, default=None + Initial Ecos guess for the iterative linear solver. + solver : LiteralOptions.OptsGenSolver, default="gmres" + Linear/saddle-point solver used for the global system. + solver_params : SolverParameters or None, default=None + Solver controls. + """ + + J: FEECVariable | list + omega: float = 1.0 + curlcurl_lambda: float = 1.0 + mass: float = 1.0 + mu: Callable | float = 1.0 + nu: Callable | float = 1.0 + rhobar: FEECVariable | Callable | float = 1.0 + theta: FEECVariable | Callable | float = 1.0 + Ebar: FEECVariable | list + + Esin0: FEECVariable | StencilVector = None + Ecos0: FEECVariable | StencilVector = None + + solver: LiteralOptions.OptsGenSolver = "gmres" + solver_params: SolverParameters | None = None + + def __post_init__(self): + # input format correctness + assert self.J is not None + if (not isinstance(self.J, (FEECVariable, list))): + raise TypeError(f"J must be either a Hcurl FEECVariable or list of Callables, got {type(self.J)}") + if isinstance(self.J, FEECVariable): + assert self.J.space == "Hcurl" + if isinstance(self.J,list): + assert len(self.J) == 3 + for ji in self.J: + assert isinstance(ji, Callable) + + if (self.rhobar is not None) and (not isinstance(self.rhobar, (FEECVariable, Callable, float))): + raise TypeError(f"rhobar must be either a H1 FEECVariable or a Callable or a float, got {type(self.rhobar)}") + if isinstance(rhobar, FEECVariable): + assert rhobar.space == "H1" + + if (self.theta is not None) and (not isinstance(self.theta, (FEECVariable, Callable, float))): + raise TypeError(f"theta must be either a H1 FEECVariable or a Callable or a float, got {type(self.theta)}") + if isinstance(self.theta, FEECVariable): + assert self.theta.space == "H1" + + assert self.Ebar is not None + if (not isinstance(self.Ebar,(FEECVariable, list))): + raise TypeError(f"Ebar must be either a Hcurl FEECVariable or list of Callables, got {type(self.Ebar)}") + if isinstance(self.Ebar, FEECVariable): + assert self.Ebar.space == "Hcurl" + if isinstance(self.Ebar,list): + assert len(self.Ebar) == 3 + for ei in self.Ebar: + assert isinstance(ei, Callable) + + # --- physical parameter sanity checks --- + if self.omega <= 0: + raise ValueError(f"omega must be positive, got {self.omega}") + if self.curlcurl_lambda <= 0: + raise ValueError(f"curlcurl_lambda must be positive, got {self.curlcurl_lambda}") + if self.mass <= 0: + raise ValueError(f"mass must be positive, got {self.mass}") + if isinstance(mu, float) and self.mu < 0: + raise ValueError(f"mu must be non-negative, got {self.mu}") + if isinstance(nu, float) and self.nu < 0: + raise ValueError(f"nu must be non-negative, got {self.nu}") + + # --- E initial guess correct space check --- + if isinstance(self.Esin0, FEECVariable): + assert self.Esin0.space == "Hcurl" + if isinstance(self.Ecos0, FEECVariable): + assert self.Ecos0.space == "Hcurl" + + 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) + self._options = new + logger.info(f"\nNew options for propagator '{self.__class__.__name__}':\n{self._options}") + + # ========================================================================= + ### Allocate + # ========================================================================= + + def allocate(self): + + # ---- source term vector (for RHS assembly) --------------------------- + + self._j: StencilVector + + if isinstance(self._options.J,FEECVariable): + self._j = self._options.J.spline.vector + else: + self._j = self.derham.P1(self._options.J) # works if J is a list of Callables + + # ---- unconstrained operators (for RHS assembly) ---------------------- + + self._M0 = self.mass_ops.M0 + self._M1 = self.mass_ops.M1 + self._M2 = self.mass_ops.M2 + self._grad = self.derham.grad + self._curl = self.derham.curl + self._div = self.derham.div + + self._M0inv = inverse( + self._M0, + "pcg", + pc="MassMatrixPreconditioner", + tol=self._options.solver_params.tol, + maxiter=self._options.solver_params.maxiter, + verbose = False, + recycle = self._options.solver_params.recycle, + ) + + self._M1mu: WeightedMassOperators + self._M2mu: WeightedMassOperators + self._M3mu: WeightedMassOperators + + if isinstance(self._options.mu, float): + self._M1mu = self._options.mu * self._M1 + self._M2mu = self._options.mu * self._M2 + self._M3mu = self._options.mu * self.mass_ops.M3 + else: + self._M1mu = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + self._options.mu, + ), + name = "M1mu", + assemble = True, + ) + + self._M2mu = self.mass_ops.create_weighted_mass( + "Hdiv", + "Hdiv", + weights=( + "G", + "1/sqrt_g", + self._options.mu, + ), + name = "M2mu", + assemble = True, + ) + + self._M3mu = self.mass_ops.create_weighted_mass( + "L2", + "L2", + weights=( + "1/sqrt_g", + self._options.mu, + ), + name = "M3mu", + assemble = True, + ) + + + self._M1rho: WeightedMassOperators + self._M1xrhoB: WeightedMassOperators + + rot_B = LocalRotationMatrix( + self.eq_mhd.b2_1, + self.eq_mhd.b2_2, + self.eq_mhd.b2_3, + ) + + rhoB1: Callable + rhoB2: Callable + rhoB3: Callable + + if isinstance(self._options.rhobar, float): + self._M1rho = self._options.rhobar * self.mass_ops.M1 + self._M1xrhoB = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + rot_B, + "Ginv", + "sqrt_g", + ), + name = "M1_xrhoB", + assemble = True, + ) + self._M1xrhoB *= self._options.rhobar + + if isinstance(self._options.rhobar, Callable): + self._M1rho = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + self._options.rhobar, + ), + name = "M1rho", + assemble = True, + ) + + self._M1xrhoB = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + rot_B, + self._options.rhobar, + "Ginv", + "sqrt_g", + ), + name = "M1_xrhoB", + assemble = True, + ) + + if isinstance(self._options.rhobar,FEECVariable): + self._M1rho = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + self._options.rhobar.spline, + ), + name = "M1rho", + assemble = True, + ) + + self._M1xrhoB = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + rot_B, + self._options.rhobar.spline, + "Ginv", + "sqrt_g", + ), + name = "M1_xrhoB", + assemble = True, + ) + + self._M1rho_inv = inverse( + self._M1rho, + "pcg", + pc="MassMatrixPreconditioner", + tol=self._options.solver_params.tol, + maxiter=self._options.solver_params.maxiter, + verbose=False, + recycle=self._options.solver_params.recycle, + ) + + + self._M1nurho: WeightedMassOperators + + if isinstance(self._options.nu, float): + self._M1nurho = self._options.nu * self._M1rho + + if isinstance(self._options.nu, Callable): + nurho: Callable + if isinstance(self._options.rhobar, float): + nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar + if isinstance(self._options.rhobar,Callable): + nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar(*etas) + if isinstance(self._options.rhobar,FEECVariable): + nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar.spline(*etas) + + self._M1nurho = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + nurho, + ), + name = "M1nurho", + assemble = True, + ) + + + self._P00theta: BasisProjectionOperators + + if isinstance(self._options.theta, float): + self._P00theta = self._options.theta * IdentityOperator(self.derham.V0) + + if isinstance(self._options.theta, Callable): + self._P00theta = self.basis_ops.create_basis_op( + [[self._options.theta]], + "H1", + "H1", + assemble = True, + name = "P00theta", + ) + + if isinstance(self._options.theta, FEECVariable): + self._P00theta = self.basis_ops.create_basis_op( + [[self._options.theta.spline]], + "H1", + "H1", + assemble = True, + name = "P00theta", + ) + + self._P01Ebar: BasisProjectionOperators + + if isinstance(self._options.Ebar, list): + self._P01Ebar = self.basis_ops.create_basis_op( + [[self._options.Ebar[0]],[self._options.Ebar[1]],[self._options.Ebar[2]]], + "H1", + "Hcurl", + assemble = True, + name = "P01Ebar", + ) + + if isinstance(self._options.Ebar, FEECVariable): + Ebar1 = lambda *etas: self._options.Ebar.spline(etas)[0] + Ebar2 = lambda *etas: self._options.Ebar.spline(etas)[1] + Ebar3 = lambda *etas: self._options.Ebar.spline(etas)[2] + + self._P01Ebar = self.basis_ops.create_basis_op( + [[Ebar1],[Ebar2],[Ebar3]], + "H1", + "Hcurl", + assemble = True, + name = "P01Ebar", + ) + + + self._P12 = self.basis_ops.U1 + + + ones = lambda *etas: 1.0 + 0 * etas + zeroes = lambda *etas: 0 * etas + + self._O1 = self.basis_ops.create_basis_op( + [[ones, zeroes, zeroes]], + "Hcurl", + "H1", + assemble = True, + name = "O1", + ) + self._O2 = self.basis_ops.create_basis_op( + [[zeroes, ones, zeroes]], + "Hcurl", + "H1", + assemble = True, + name = "O1", + ) + self._O3 = self.basis_ops.create_basis_op( + [[zeroes, zeroes, ones]], + "Hcurl", + "H1", + assemble = True, + name = "O1", + ) + + + self._Acurlcurl = self._curl.T @ self._M2 @ self._curl - self._options.curlcurl_lambda * self._M1 + + self._divPi = - self._curl.T @ self._M2mu @ self._curl \ + - 2/3 * self._P12.T @ self._div.T @ self._M3mu @ self._div @ self._P12 \ + + 2 * self._O1.T @ self._grad.T @ self._M1mu @ self._grad @ self._O1 \ + + 2 * self._O2.T @ self._grad.T @ self._M1mu @ self._grad @ self._O2 \ + + 2 * self._O3.T @ self._grad.T @ self._M1mu @ self._grad @ self._O3 + + # self._A = self._M1rho + self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) @ self._M0inv @ self._grad.T @ self._M1rho \ + # / (self._options.mass * self._options.omega * self._options.omega) + + # self._B = (self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho) / self._options.omega + + # ---- block saddle-point system ---------------------------------------- + + self._block_V0 = BlockVectorSpace(self.derham.V0, self.derham.V0) + self._block_V1 = BlockVectorSpace(self.derham.V1, self.derham.V1) + self._block_V2 = BlockVectorSpace(self.derham.V2, self.derham.V2) + + self._block_source = BlockVector(self._block_V1, blocks=[self._M1.dot(self._j), None]) + + self._block_M = BlockLinearOperator( + self._block_V0, self._block_V0, blocks=[[self._M0, None], [None, self._M0]] + ) + + self._block_Divergence = BlockLinearOperator( + self._block_V1, self._block_V0, blocks=[[None, - self._grad.T @ self._M1rho], [self._grad.T @ self._M1rho, None]] + ) + + self._block_Acurlcurl = BlockLinearOperator( + self.block_V, self._block_V1, blocks=[[self._Acurlcurl, None], [None, self._Acurlcurl]] + ) + + self._block_B = BlockLinearOperator( + self._block_V1, self._block_V1, blocks=[[None, self._M1rho / self._options.mass], [- self._M1rho / self._options.mass, None]] + ) + + self._block_P = BlockLinearOperator( + self._block_V0, self._block_V1, + blocks=[[None, self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) / self._options.mass], + [self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) / self._options.mass, None]] + ) + + self._block_Q = BlockLinearOperator( + self._block_V1, self._block_V1, + blocks=[[self._options.omega * self._M1rho, self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho], + [self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho, - self._options.omega * self._M1rho]] + ) + + self._block_R = BlockLinearOperator( + self._block_V1, self._block_V1, blocks=[[None, self._M1rho / self._options.mass], [self._M1rho / self._options.mass, None]] + ) + + self._block_Minv = inverse( + self._block_M, + "pcg", + pc="MassMatrixPreconditioner", + tol=self._options.solver_params.tol, + maxiter=self._options.solver_params.maxiter, + verbose = False, + recycle = self._options.solver_params.recycle, + ) + + self._block_rhomatrix = - self._block_Minv @ self._block_Divergence / self._options.omega + + self._block_umatrix = - self._block_P @ self._block_rhomatrix - self._block_Q + + self._block_umatrix_inv = inverse( + self._block_umatrix, + "gmres", + pc=None, + tol=self._options.solver_params.tol, + maxiter=self._options.solver_params.maxiter, + verbose = False, + recycle = self._options.solver_params.recycle, + ) + + self._block_Ematrix = self._block_Acurlcurl / self._options.omega + self._block_B @ self._block_umatrix_inv @ self._block_R + + self._block_Ematrix_inv = inverse( + self._block_Ematrix, + "gmres", + pc=None, + tol=self._options.solver_params.tol, + maxiter=self._options.solver_params.maxiter, + verbose = False, + recycle = self._options.solver_params.recycle, + ) + + # self._coupled_equations_matrix = BlockLinearOperator( + # self._block_domain, self._block_codomain, blocks=[[self._A, -self._B], [self._B, self._A]] + # ) + + # M1rhoinv_j = self._M1rho_inv.solve(self._j) + # A_j = self._A.dot(M1rhoinv_j) + # minusB_j = - self.B.dot(M1rhoinv_j) + + # self._calEsin0: StencilVector = None + # self._calEcos0: StencilVector = None + + # # --- copy current state --- + # Esin0 = self.variables.Esin.spline.vector + # Ecos0 = self.variables.Ecos.spline.vector + + # self._calEsin0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Esin0)) + # self._calEcos0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Ecos0)) + + # self._calE0 = BlockVector(self._block_domain, blocks=[self._calEsin0, self._calEcos0]) + + # self._coupled_equations_matrix_inverse = inverse( + # self._coupled_equations_matrix, + # solver="gmres", + # x0=self._calE0, + # tol=self._options.solver_params.tol, + # maxiter=self._options.solver_params.maxiter, + # verbose=False, + # recycle=True, + # ) + + # # --- build inverses of the curl-curl matrices with good initial guesses --- + # self._Acurlcurl_inv_sin = inverse( + # self._Acurlcurl, + # solver="pcg", + # x0=Esin0, + # tol=self._options.solver_params.tol, + # maxiter=self._options.solver_params.maxiter, + # verbose=False, + # recycle=True, + # ) + + # self._Acurlcurl_inv_cos = inverse( + # self._Acurlcurl, + # solver="pcg", + # x0=Ecos0, + # tol=self._options.solver_params.tol, + # maxiter=self._options.solver_params.maxiter, + # verbose=False, + # recycle=True, + # ) + + + # ========================================================================= + ### Equation solve + # ========================================================================= + + def __call__(self, dt): + + # --- calculate auxilliary vectors --- + _calE = self._coupled_equations_matrix_inverse.solve(self._calE_RHS) + + _calEsin = _calE[0] + _calEcos = _calE[1] + + _m_curlcurlEsin = self._M1rho.dot(_calEsin) + _m_curlcurlEcos = self._M1rho.dot(_calEcos) + + _M1rho_usin = _m_curlcurlEcos.copy() + _M1rho_ucos = self._options.mass * self._j - _m_curlcurlEsin + + # --- calculate solutions --- + self._Esin.vector = self._Acurlcurl_inv_sin.solve(_m_curlcurlEsin / self._options.mass) + self._Ecos.vector = self._Acurlcurl_inv_cos.solve(_m_curlcurlEcos / self._options.mass) + + self._Bsin.vector = - self._curl.dot(self._Ecos.vector) / self._options.omega + self._Bcos.vector = self._curl.dot(self._Esin.vector) / self._options.omega + + self._usin.vector = self._M1rho_inv.solve(_M1rho_usin) + self._ucos.vector = self._M1rho_inv.solve(_M1rho_ucos) + + self._rhosin.vector = self._M0inv.solve(self._grad.T.dot(_M1rho_ucos)) / self._options.omega + self._rhocos.vector = - self._M0inv.solve(self._grad.T.dot(_M1rho_usin)) / self._options.omega + + # --- update FEEC variables --- + self.update_feec_variables( + rhosin=self._rhosin.vector, rhocos=self._rhocos.vector, + usin=self._usin.vector, ucos=self._ucos.vector, + Esin=self._Esin.vector, Ecos=self._Ecos.vector, + Bsin=self._Bsin.vector, Bcos=self._Bcos.vector) + From 4fb1032652ec8352d063d3ce512ae1051bff16a6 Mon Sep 17 00:00:00 2001 From: Monoclod Date: Fri, 19 Jun 2026 17:37:08 +0200 Subject: [PATCH 4/6] small cleanup --- .../propagators/perturbation_system_cold.py | 101 +-- .../propagators/perturbation_system_full.py | 734 ------------------ 2 files changed, 54 insertions(+), 781 deletions(-) delete mode 100644 src/struphy/propagators/perturbation_system_full.py diff --git a/src/struphy/propagators/perturbation_system_cold.py b/src/struphy/propagators/perturbation_system_cold.py index 9e1cf41b0..de26bbd15 100644 --- a/src/struphy/propagators/perturbation_system_cold.py +++ b/src/struphy/propagators/perturbation_system_cold.py @@ -185,10 +185,6 @@ class Options(OptionsBase): Average electron temperature. Ebar : FEECVariable in ``"Hcurl"`` or list Average electrostatic field. - Esin0 : StencilVector, default=None - Initial Esin guess for the iterative linear solver. - Ecos0 : StencilVector, default=None - Initial Ecos guess for the iterative linear solver. solver : LiteralOptions.OptsGenSolver, default="gmres" Linear/saddle-point solver used for the global system. solver_params : SolverParameters or None, default=None @@ -204,9 +200,6 @@ class Options(OptionsBase): rhobar: FEECVariable | Callable | float = 1.0 theta: FEECVariable | Callable | float = 1.0 Ebar: FEECVariable | list - - Esin0: FEECVariable | StencilVector = None - Ecos0: FEECVariable | StencilVector = None solver: LiteralOptions.OptsGenSolver = "gmres" solver_params: SolverParameters | None = None @@ -255,12 +248,6 @@ def __post_init__(self): if isinstance(nu, float) and self.nu < 0: raise ValueError(f"nu must be non-negative, got {self.nu}") - # --- E initial guess correct space check --- - if isinstance(self.Esin0, FEECVariable): - assert self.Esin0.space == "Hcurl" - if isinstance(self.Ecos0, FEECVariable): - assert self.Ecos0.space == "Hcurl" - check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver) if self.solver_params is None: self.solver_params = SolverParameters() @@ -284,7 +271,7 @@ def allocate(self): # ---- source term vector (for RHS assembly) --------------------------- - self._j: StencilVector + self._j: StencilVector = None if isinstance(self._options.J,FEECVariable): self._j = self._options.J.spline.vector @@ -310,9 +297,9 @@ def allocate(self): recycle = self._options.solver_params.recycle, ) - self._M1mu: WeightedMassOperators - self._M2mu: WeightedMassOperators - self._M3mu: WeightedMassOperators + self._M1mu: WeightedMassOperators = None + self._M2mu: WeightedMassOperators = None + self._M3mu: WeightedMassOperators = None if isinstance(self._options.mu, float): self._M1mu = self._options.mu * self._M1 @@ -355,8 +342,8 @@ def allocate(self): ) - self._M1rho: WeightedMassOperators - self._M1xrhoB: WeightedMassOperators + self._M1rho: WeightedMassOperators = None + self._M1xrhoB: WeightedMassOperators = None rot_B = LocalRotationMatrix( self.eq_mhd.b2_1, @@ -364,9 +351,9 @@ def allocate(self): self.eq_mhd.b2_3, ) - rhoB1: Callable - rhoB2: Callable - rhoB3: Callable + rhoB1: Callable = None + rhoB2: Callable = None + rhoB3: Callable = None if isinstance(self._options.rhobar, float): self._M1rho = self._options.rhobar * self.mass_ops.M1 @@ -449,7 +436,7 @@ def allocate(self): ) - self._M1nurho: WeightedMassOperators + self._M1nurho: WeightedMassOperators = None if isinstance(self._options.nu, float): self._M1nurho = self._options.nu * self._M1rho @@ -476,7 +463,7 @@ def allocate(self): ) - self._P00theta: BasisProjectionOperators + self._P00theta: BasisProjectionOperators = None if isinstance(self._options.theta, float): self._P00theta = self._options.theta * IdentityOperator(self.derham.V0) @@ -499,7 +486,7 @@ def allocate(self): name = "P00theta", ) - self._P01Ebar: BasisProjectionOperators + self._P01Ebar: BasisProjectionOperators = None if isinstance(self._options.Ebar, list): self._P01Ebar = self.basis_ops.create_basis_op( @@ -582,6 +569,10 @@ def allocate(self): self._block_V1, self._block_V0, blocks=[[None, - self._grad.T @ self._M1rho], [self._grad.T @ self._M1rho, None]] ) + self._block_curl = BlockLinearOperator( + self._block_V1, self._block_V2, blocks=[[None, - self._curl], [self._curl, None]] + ) + self._block_Acurlcurl = BlockLinearOperator( self.block_V, self._block_V1, blocks=[[self._Acurlcurl, None], [None, self._Acurlcurl]] ) @@ -701,34 +692,50 @@ def allocate(self): def __call__(self, dt): # --- calculate auxilliary vectors --- - _calE = self._coupled_equations_matrix_inverse.solve(self._calE_RHS) + block_E = self._block_Ematrix_inv.solve(self._block_source) - _calEsin = _calE[0] - _calEcos = _calE[1] + block_B = self._block_curl.dot(block_E) / self._options.omega - _m_curlcurlEsin = self._M1rho.dot(_calEsin) - _m_curlcurlEcos = self._M1rho.dot(_calEcos) + block_u = self._block_umatrix_inv.solve(self._block_R.dot(block_E)) - _M1rho_usin = _m_curlcurlEcos.copy() - _M1rho_ucos = self._options.mass * self._j - _m_curlcurlEsin + block_rho = self._block_rhomatrix.dot(block_u) - # --- calculate solutions --- - self._Esin.vector = self._Acurlcurl_inv_sin.solve(_m_curlcurlEsin / self._options.mass) - self._Ecos.vector = self._Acurlcurl_inv_cos.solve(_m_curlcurlEcos / self._options.mass) + # --- update FEEC variables --- + self.update_feec_variables( + rhosin=block_rho[0], rhocos=block_rho[1], + usin=block_u[0], ucos=block_u[1], + Esin=block_E[0], Ecos=block_E[1], + Bsin=block_B[0], Bcos=Block_B[1]) - self._Bsin.vector = - self._curl.dot(self._Ecos.vector) / self._options.omega - self._Bcos.vector = self._curl.dot(self._Esin.vector) / self._options.omega - self._usin.vector = self._M1rho_inv.solve(_M1rho_usin) - self._ucos.vector = self._M1rho_inv.solve(_M1rho_ucos) + # _calE = self._coupled_equations_matrix_inverse.solve(self._calE_RHS) - self._rhosin.vector = self._M0inv.solve(self._grad.T.dot(_M1rho_ucos)) / self._options.omega - self._rhocos.vector = - self._M0inv.solve(self._grad.T.dot(_M1rho_usin)) / self._options.omega + # _calEsin = _calE[0] + # _calEcos = _calE[1] - # --- update FEEC variables --- - self.update_feec_variables( - rhosin=self._rhosin.vector, rhocos=self._rhocos.vector, - usin=self._usin.vector, ucos=self._ucos.vector, - Esin=self._Esin.vector, Ecos=self._Ecos.vector, - Bsin=self._Bsin.vector, Bcos=self._Bcos.vector) + # _m_curlcurlEsin = self._M1rho.dot(_calEsin) + # _m_curlcurlEcos = self._M1rho.dot(_calEcos) + + # _M1rho_usin = _m_curlcurlEcos.copy() + # _M1rho_ucos = self._options.mass * self._j - _m_curlcurlEsin + + # # --- calculate solutions --- + # self._Esin.vector = self._Acurlcurl_inv_sin.solve(_m_curlcurlEsin / self._options.mass) + # self._Ecos.vector = self._Acurlcurl_inv_cos.solve(_m_curlcurlEcos / self._options.mass) + + # self._Bsin.vector = - self._curl.dot(self._Ecos.vector) / self._options.omega + # self._Bcos.vector = self._curl.dot(self._Esin.vector) / self._options.omega + + # self._usin.vector = self._M1rho_inv.solve(_M1rho_usin) + # self._ucos.vector = self._M1rho_inv.solve(_M1rho_ucos) + + # self._rhosin.vector = self._M0inv.solve(self._grad.T.dot(_M1rho_ucos)) / self._options.omega + # self._rhocos.vector = - self._M0inv.solve(self._grad.T.dot(_M1rho_usin)) / self._options.omega + + # # --- update FEEC variables --- + # self.update_feec_variables( + # rhosin=self._rhosin.vector, rhocos=self._rhocos.vector, + # usin=self._usin.vector, ucos=self._ucos.vector, + # Esin=self._Esin.vector, Ecos=self._Ecos.vector, + # Bsin=self._Bsin.vector, Bcos=self._Bcos.vector) diff --git a/src/struphy/propagators/perturbation_system_full.py b/src/struphy/propagators/perturbation_system_full.py deleted file mode 100644 index 9e1cf41b0..000000000 --- a/src/struphy/propagators/perturbation_system_full.py +++ /dev/null @@ -1,734 +0,0 @@ -import logging -from dataclasses import dataclass -from typing import Callable, get_args -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 -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 L2Projector, WeightedMassOperators -from struphy.feec.utilities import LocalRotationMatrix -from struphy.io.options import LiteralOptions, OptionsBase -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 - -logger = logging.getLogger("struphy") - - -class ColdPlasmaPerturbation(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 (electron density rhosin and rhocos, electron velocity usin and ucos, electric field Esin and Ecos, magnetic field Bsin and Bcos) - # ========================================================================= - - class Variables: - """Container for variables advanced by :class:`ColdPlasmaPerturbation`. - - Attributes - ---------- - rhosin : FEECVariable or None - Sine component of electron density variable in ``"H1"`` space. - rhocos : FEECVariable or None - Cosine component of electron density variable in ``"H1"`` space. - usin : FEECVariable or None - Sine component of the electron velocity variable in ``"Hcurl"`` space. - ucos : FEECVariable or None - Cosine component of the electron velocity variable in ``"Hcurl"`` space. - Esin : FEECVariable or None - Sine component of the electric field variable in ``"Hcurl"`` space. - Ecos : FEECVariable or None - Cosine component of the electric field variable in ``"Hcurl"`` space. - Bsin : FEECVariable or None - Sine component of the magnetic field variable in ``"Hdiv"`` space. - Bcos : FEECVariable or None - Cosine component of the magnetic field variable in ``"Hdiv"`` space. - """ - - def __init__(self) -> None: - self._rhosin: FEECVariable | None = None - self._rhocos: FEECVariable | None = None - self._usin: FEECVariable | None = None - self._ucos: FEECVariable | None = None - self._Esin: FEECVariable | None = None - self._Ecos: FEECVariable | None = None - self._Bsin: FEECVariable | None = None - self._Bcos: FEECVariable | None = None - - @property - def rhosin(self) -> FEECVariable | None: - return self._rhosin - - @rhosin.setter - def rhosin(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "H1" - self._rhosin = new - - @property - def rhocos(self) -> FEECVariable | None: - return self._rhocos - - @rhocos.setter - def rhocos(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "H1" - self._rhocos = new - - @property - def usin(self) -> FEECVariable | None: - return self._usin - - @usin.setter - def usin(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hcurl" - self._usin = new - - @property - def ucos(self) -> FEECVariable | None: - return self._ucos - - @ucos.setter - def ucos(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hcurl" - self._ucos = new - - @property - def Esin(self) -> FEECVariable | None: - return self._Esin - - @Esin.setter - def Esin(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hcurl" - self._Esin = new - - @property - def Ecos(self) -> FEECVariable | None: - return self._Ecos - - @Ecos.setter - def Ecos(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hcurl" - self._Ecos = new - - @property - def Bsin(self) -> FEECVariable | None: - return self._Bsin - - @Bsin.setter - def Bsin(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hdiv" - self._Bsin = new - - @property - def Bcos(self) -> FEECVariable | None: - return self._Bcos - - @ucos.setter - def Bcos(self, new): - assert isinstance(new, FEECVariable) - assert new.space == "Hdiv" - self._Bcos = new - - - def __init__(self): - self.variables = self.Variables() - - # ========================================================================= - ### Options - # ========================================================================= - - @dataclass(repr=False) - class Options(OptionsBase): - """Configuration options for :class:`ColdPlasmaPerturbation`. - - Parameters - ---------- - J : FEECVariable in ``"Hcurl"`` or list - Cosine component of the source term. - omega : float, default=1.0 - Source term oscillation frequency. - curlcurl_lambda : float, default=1.0 - Coefficient in the curl-curl operator. - mass : float, default=1.0 - Electron mass in relative unis. - mu : Callable or float, default=1.0 - Electron viscosity coefficient. - nu : Callable or float, default=1.0 - Electron-Ion collision frequency. - rhobar : FEECVariable in ``"H1"`` or Callable or float, default=1.0 - Average electron mass density. - theta : FEECVariable in ``"H1"`` or Callable or float, default=1.0 - Average electron temperature. - Ebar : FEECVariable in ``"Hcurl"`` or list - Average electrostatic field. - Esin0 : StencilVector, default=None - Initial Esin guess for the iterative linear solver. - Ecos0 : StencilVector, default=None - Initial Ecos guess for the iterative linear solver. - solver : LiteralOptions.OptsGenSolver, default="gmres" - Linear/saddle-point solver used for the global system. - solver_params : SolverParameters or None, default=None - Solver controls. - """ - - J: FEECVariable | list - omega: float = 1.0 - curlcurl_lambda: float = 1.0 - mass: float = 1.0 - mu: Callable | float = 1.0 - nu: Callable | float = 1.0 - rhobar: FEECVariable | Callable | float = 1.0 - theta: FEECVariable | Callable | float = 1.0 - Ebar: FEECVariable | list - - Esin0: FEECVariable | StencilVector = None - Ecos0: FEECVariable | StencilVector = None - - solver: LiteralOptions.OptsGenSolver = "gmres" - solver_params: SolverParameters | None = None - - def __post_init__(self): - # input format correctness - assert self.J is not None - if (not isinstance(self.J, (FEECVariable, list))): - raise TypeError(f"J must be either a Hcurl FEECVariable or list of Callables, got {type(self.J)}") - if isinstance(self.J, FEECVariable): - assert self.J.space == "Hcurl" - if isinstance(self.J,list): - assert len(self.J) == 3 - for ji in self.J: - assert isinstance(ji, Callable) - - if (self.rhobar is not None) and (not isinstance(self.rhobar, (FEECVariable, Callable, float))): - raise TypeError(f"rhobar must be either a H1 FEECVariable or a Callable or a float, got {type(self.rhobar)}") - if isinstance(rhobar, FEECVariable): - assert rhobar.space == "H1" - - if (self.theta is not None) and (not isinstance(self.theta, (FEECVariable, Callable, float))): - raise TypeError(f"theta must be either a H1 FEECVariable or a Callable or a float, got {type(self.theta)}") - if isinstance(self.theta, FEECVariable): - assert self.theta.space == "H1" - - assert self.Ebar is not None - if (not isinstance(self.Ebar,(FEECVariable, list))): - raise TypeError(f"Ebar must be either a Hcurl FEECVariable or list of Callables, got {type(self.Ebar)}") - if isinstance(self.Ebar, FEECVariable): - assert self.Ebar.space == "Hcurl" - if isinstance(self.Ebar,list): - assert len(self.Ebar) == 3 - for ei in self.Ebar: - assert isinstance(ei, Callable) - - # --- physical parameter sanity checks --- - if self.omega <= 0: - raise ValueError(f"omega must be positive, got {self.omega}") - if self.curlcurl_lambda <= 0: - raise ValueError(f"curlcurl_lambda must be positive, got {self.curlcurl_lambda}") - if self.mass <= 0: - raise ValueError(f"mass must be positive, got {self.mass}") - if isinstance(mu, float) and self.mu < 0: - raise ValueError(f"mu must be non-negative, got {self.mu}") - if isinstance(nu, float) and self.nu < 0: - raise ValueError(f"nu must be non-negative, got {self.nu}") - - # --- E initial guess correct space check --- - if isinstance(self.Esin0, FEECVariable): - assert self.Esin0.space == "Hcurl" - if isinstance(self.Ecos0, FEECVariable): - assert self.Ecos0.space == "Hcurl" - - 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) - self._options = new - logger.info(f"\nNew options for propagator '{self.__class__.__name__}':\n{self._options}") - - # ========================================================================= - ### Allocate - # ========================================================================= - - def allocate(self): - - # ---- source term vector (for RHS assembly) --------------------------- - - self._j: StencilVector - - if isinstance(self._options.J,FEECVariable): - self._j = self._options.J.spline.vector - else: - self._j = self.derham.P1(self._options.J) # works if J is a list of Callables - - # ---- unconstrained operators (for RHS assembly) ---------------------- - - self._M0 = self.mass_ops.M0 - self._M1 = self.mass_ops.M1 - self._M2 = self.mass_ops.M2 - self._grad = self.derham.grad - self._curl = self.derham.curl - self._div = self.derham.div - - self._M0inv = inverse( - self._M0, - "pcg", - pc="MassMatrixPreconditioner", - tol=self._options.solver_params.tol, - maxiter=self._options.solver_params.maxiter, - verbose = False, - recycle = self._options.solver_params.recycle, - ) - - self._M1mu: WeightedMassOperators - self._M2mu: WeightedMassOperators - self._M3mu: WeightedMassOperators - - if isinstance(self._options.mu, float): - self._M1mu = self._options.mu * self._M1 - self._M2mu = self._options.mu * self._M2 - self._M3mu = self._options.mu * self.mass_ops.M3 - else: - self._M1mu = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - "Ginv", - "sqrt_g", - self._options.mu, - ), - name = "M1mu", - assemble = True, - ) - - self._M2mu = self.mass_ops.create_weighted_mass( - "Hdiv", - "Hdiv", - weights=( - "G", - "1/sqrt_g", - self._options.mu, - ), - name = "M2mu", - assemble = True, - ) - - self._M3mu = self.mass_ops.create_weighted_mass( - "L2", - "L2", - weights=( - "1/sqrt_g", - self._options.mu, - ), - name = "M3mu", - assemble = True, - ) - - - self._M1rho: WeightedMassOperators - self._M1xrhoB: WeightedMassOperators - - rot_B = LocalRotationMatrix( - self.eq_mhd.b2_1, - self.eq_mhd.b2_2, - self.eq_mhd.b2_3, - ) - - rhoB1: Callable - rhoB2: Callable - rhoB3: Callable - - if isinstance(self._options.rhobar, float): - self._M1rho = self._options.rhobar * self.mass_ops.M1 - self._M1xrhoB = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - "Ginv", - rot_B, - "Ginv", - "sqrt_g", - ), - name = "M1_xrhoB", - assemble = True, - ) - self._M1xrhoB *= self._options.rhobar - - if isinstance(self._options.rhobar, Callable): - self._M1rho = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - "Ginv", - "sqrt_g", - self._options.rhobar, - ), - name = "M1rho", - assemble = True, - ) - - self._M1xrhoB = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - "Ginv", - rot_B, - self._options.rhobar, - "Ginv", - "sqrt_g", - ), - name = "M1_xrhoB", - assemble = True, - ) - - if isinstance(self._options.rhobar,FEECVariable): - self._M1rho = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - "Ginv", - "sqrt_g", - self._options.rhobar.spline, - ), - name = "M1rho", - assemble = True, - ) - - self._M1xrhoB = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - "Ginv", - rot_B, - self._options.rhobar.spline, - "Ginv", - "sqrt_g", - ), - name = "M1_xrhoB", - assemble = True, - ) - - self._M1rho_inv = inverse( - self._M1rho, - "pcg", - pc="MassMatrixPreconditioner", - tol=self._options.solver_params.tol, - maxiter=self._options.solver_params.maxiter, - verbose=False, - recycle=self._options.solver_params.recycle, - ) - - - self._M1nurho: WeightedMassOperators - - if isinstance(self._options.nu, float): - self._M1nurho = self._options.nu * self._M1rho - - if isinstance(self._options.nu, Callable): - nurho: Callable - if isinstance(self._options.rhobar, float): - nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar - if isinstance(self._options.rhobar,Callable): - nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar(*etas) - if isinstance(self._options.rhobar,FEECVariable): - nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar.spline(*etas) - - self._M1nurho = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - "Ginv", - "sqrt_g", - nurho, - ), - name = "M1nurho", - assemble = True, - ) - - - self._P00theta: BasisProjectionOperators - - if isinstance(self._options.theta, float): - self._P00theta = self._options.theta * IdentityOperator(self.derham.V0) - - if isinstance(self._options.theta, Callable): - self._P00theta = self.basis_ops.create_basis_op( - [[self._options.theta]], - "H1", - "H1", - assemble = True, - name = "P00theta", - ) - - if isinstance(self._options.theta, FEECVariable): - self._P00theta = self.basis_ops.create_basis_op( - [[self._options.theta.spline]], - "H1", - "H1", - assemble = True, - name = "P00theta", - ) - - self._P01Ebar: BasisProjectionOperators - - if isinstance(self._options.Ebar, list): - self._P01Ebar = self.basis_ops.create_basis_op( - [[self._options.Ebar[0]],[self._options.Ebar[1]],[self._options.Ebar[2]]], - "H1", - "Hcurl", - assemble = True, - name = "P01Ebar", - ) - - if isinstance(self._options.Ebar, FEECVariable): - Ebar1 = lambda *etas: self._options.Ebar.spline(etas)[0] - Ebar2 = lambda *etas: self._options.Ebar.spline(etas)[1] - Ebar3 = lambda *etas: self._options.Ebar.spline(etas)[2] - - self._P01Ebar = self.basis_ops.create_basis_op( - [[Ebar1],[Ebar2],[Ebar3]], - "H1", - "Hcurl", - assemble = True, - name = "P01Ebar", - ) - - - self._P12 = self.basis_ops.U1 - - - ones = lambda *etas: 1.0 + 0 * etas - zeroes = lambda *etas: 0 * etas - - self._O1 = self.basis_ops.create_basis_op( - [[ones, zeroes, zeroes]], - "Hcurl", - "H1", - assemble = True, - name = "O1", - ) - self._O2 = self.basis_ops.create_basis_op( - [[zeroes, ones, zeroes]], - "Hcurl", - "H1", - assemble = True, - name = "O1", - ) - self._O3 = self.basis_ops.create_basis_op( - [[zeroes, zeroes, ones]], - "Hcurl", - "H1", - assemble = True, - name = "O1", - ) - - - self._Acurlcurl = self._curl.T @ self._M2 @ self._curl - self._options.curlcurl_lambda * self._M1 - - self._divPi = - self._curl.T @ self._M2mu @ self._curl \ - - 2/3 * self._P12.T @ self._div.T @ self._M3mu @ self._div @ self._P12 \ - + 2 * self._O1.T @ self._grad.T @ self._M1mu @ self._grad @ self._O1 \ - + 2 * self._O2.T @ self._grad.T @ self._M1mu @ self._grad @ self._O2 \ - + 2 * self._O3.T @ self._grad.T @ self._M1mu @ self._grad @ self._O3 - - # self._A = self._M1rho + self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) @ self._M0inv @ self._grad.T @ self._M1rho \ - # / (self._options.mass * self._options.omega * self._options.omega) - - # self._B = (self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho) / self._options.omega - - # ---- block saddle-point system ---------------------------------------- - - self._block_V0 = BlockVectorSpace(self.derham.V0, self.derham.V0) - self._block_V1 = BlockVectorSpace(self.derham.V1, self.derham.V1) - self._block_V2 = BlockVectorSpace(self.derham.V2, self.derham.V2) - - self._block_source = BlockVector(self._block_V1, blocks=[self._M1.dot(self._j), None]) - - self._block_M = BlockLinearOperator( - self._block_V0, self._block_V0, blocks=[[self._M0, None], [None, self._M0]] - ) - - self._block_Divergence = BlockLinearOperator( - self._block_V1, self._block_V0, blocks=[[None, - self._grad.T @ self._M1rho], [self._grad.T @ self._M1rho, None]] - ) - - self._block_Acurlcurl = BlockLinearOperator( - self.block_V, self._block_V1, blocks=[[self._Acurlcurl, None], [None, self._Acurlcurl]] - ) - - self._block_B = BlockLinearOperator( - self._block_V1, self._block_V1, blocks=[[None, self._M1rho / self._options.mass], [- self._M1rho / self._options.mass, None]] - ) - - self._block_P = BlockLinearOperator( - self._block_V0, self._block_V1, - blocks=[[None, self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) / self._options.mass], - [self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) / self._options.mass, None]] - ) - - self._block_Q = BlockLinearOperator( - self._block_V1, self._block_V1, - blocks=[[self._options.omega * self._M1rho, self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho], - [self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho, - self._options.omega * self._M1rho]] - ) - - self._block_R = BlockLinearOperator( - self._block_V1, self._block_V1, blocks=[[None, self._M1rho / self._options.mass], [self._M1rho / self._options.mass, None]] - ) - - self._block_Minv = inverse( - self._block_M, - "pcg", - pc="MassMatrixPreconditioner", - tol=self._options.solver_params.tol, - maxiter=self._options.solver_params.maxiter, - verbose = False, - recycle = self._options.solver_params.recycle, - ) - - self._block_rhomatrix = - self._block_Minv @ self._block_Divergence / self._options.omega - - self._block_umatrix = - self._block_P @ self._block_rhomatrix - self._block_Q - - self._block_umatrix_inv = inverse( - self._block_umatrix, - "gmres", - pc=None, - tol=self._options.solver_params.tol, - maxiter=self._options.solver_params.maxiter, - verbose = False, - recycle = self._options.solver_params.recycle, - ) - - self._block_Ematrix = self._block_Acurlcurl / self._options.omega + self._block_B @ self._block_umatrix_inv @ self._block_R - - self._block_Ematrix_inv = inverse( - self._block_Ematrix, - "gmres", - pc=None, - tol=self._options.solver_params.tol, - maxiter=self._options.solver_params.maxiter, - verbose = False, - recycle = self._options.solver_params.recycle, - ) - - # self._coupled_equations_matrix = BlockLinearOperator( - # self._block_domain, self._block_codomain, blocks=[[self._A, -self._B], [self._B, self._A]] - # ) - - # M1rhoinv_j = self._M1rho_inv.solve(self._j) - # A_j = self._A.dot(M1rhoinv_j) - # minusB_j = - self.B.dot(M1rhoinv_j) - - # self._calEsin0: StencilVector = None - # self._calEcos0: StencilVector = None - - # # --- copy current state --- - # Esin0 = self.variables.Esin.spline.vector - # Ecos0 = self.variables.Ecos.spline.vector - - # self._calEsin0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Esin0)) - # self._calEcos0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Ecos0)) - - # self._calE0 = BlockVector(self._block_domain, blocks=[self._calEsin0, self._calEcos0]) - - # self._coupled_equations_matrix_inverse = inverse( - # self._coupled_equations_matrix, - # solver="gmres", - # x0=self._calE0, - # tol=self._options.solver_params.tol, - # maxiter=self._options.solver_params.maxiter, - # verbose=False, - # recycle=True, - # ) - - # # --- build inverses of the curl-curl matrices with good initial guesses --- - # self._Acurlcurl_inv_sin = inverse( - # self._Acurlcurl, - # solver="pcg", - # x0=Esin0, - # tol=self._options.solver_params.tol, - # maxiter=self._options.solver_params.maxiter, - # verbose=False, - # recycle=True, - # ) - - # self._Acurlcurl_inv_cos = inverse( - # self._Acurlcurl, - # solver="pcg", - # x0=Ecos0, - # tol=self._options.solver_params.tol, - # maxiter=self._options.solver_params.maxiter, - # verbose=False, - # recycle=True, - # ) - - - # ========================================================================= - ### Equation solve - # ========================================================================= - - def __call__(self, dt): - - # --- calculate auxilliary vectors --- - _calE = self._coupled_equations_matrix_inverse.solve(self._calE_RHS) - - _calEsin = _calE[0] - _calEcos = _calE[1] - - _m_curlcurlEsin = self._M1rho.dot(_calEsin) - _m_curlcurlEcos = self._M1rho.dot(_calEcos) - - _M1rho_usin = _m_curlcurlEcos.copy() - _M1rho_ucos = self._options.mass * self._j - _m_curlcurlEsin - - # --- calculate solutions --- - self._Esin.vector = self._Acurlcurl_inv_sin.solve(_m_curlcurlEsin / self._options.mass) - self._Ecos.vector = self._Acurlcurl_inv_cos.solve(_m_curlcurlEcos / self._options.mass) - - self._Bsin.vector = - self._curl.dot(self._Ecos.vector) / self._options.omega - self._Bcos.vector = self._curl.dot(self._Esin.vector) / self._options.omega - - self._usin.vector = self._M1rho_inv.solve(_M1rho_usin) - self._ucos.vector = self._M1rho_inv.solve(_M1rho_ucos) - - self._rhosin.vector = self._M0inv.solve(self._grad.T.dot(_M1rho_ucos)) / self._options.omega - self._rhocos.vector = - self._M0inv.solve(self._grad.T.dot(_M1rho_usin)) / self._options.omega - - # --- update FEEC variables --- - self.update_feec_variables( - rhosin=self._rhosin.vector, rhocos=self._rhocos.vector, - usin=self._usin.vector, ucos=self._ucos.vector, - Esin=self._Esin.vector, Ecos=self._Ecos.vector, - Bsin=self._Bsin.vector, Bcos=self._Bcos.vector) - From 2982839ba8c4e9b3e586424f30fcae65c43349f5 Mon Sep 17 00:00:00 2001 From: Monoclod Date: Wed, 1 Jul 2026 15:24:23 +0200 Subject: [PATCH 5/6] Fixed error in the rotation matrix --- .../propagators/perturbation_system_cold.py | 152 ++++++++---------- 1 file changed, 65 insertions(+), 87 deletions(-) diff --git a/src/struphy/propagators/perturbation_system_cold.py b/src/struphy/propagators/perturbation_system_cold.py index de26bbd15..8e6cbd171 100644 --- a/src/struphy/propagators/perturbation_system_cold.py +++ b/src/struphy/propagators/perturbation_system_cold.py @@ -191,7 +191,7 @@ class Options(OptionsBase): Solver controls. """ - J: FEECVariable | list + J: FEECVariable | list = None omega: float = 1.0 curlcurl_lambda: float = 1.0 mass: float = 1.0 @@ -199,7 +199,7 @@ class Options(OptionsBase): nu: Callable | float = 1.0 rhobar: FEECVariable | Callable | float = 1.0 theta: FEECVariable | Callable | float = 1.0 - Ebar: FEECVariable | list + Ebar: FEECVariable | list = None solver: LiteralOptions.OptsGenSolver = "gmres" solver_params: SolverParameters | None = None @@ -218,7 +218,7 @@ def __post_init__(self): if (self.rhobar is not None) and (not isinstance(self.rhobar, (FEECVariable, Callable, float))): raise TypeError(f"rhobar must be either a H1 FEECVariable or a Callable or a float, got {type(self.rhobar)}") - if isinstance(rhobar, FEECVariable): + if isinstance(self.rhobar, FEECVariable): assert rhobar.space == "H1" if (self.theta is not None) and (not isinstance(self.theta, (FEECVariable, Callable, float))): @@ -287,10 +287,12 @@ def allocate(self): self._curl = self.derham.curl self._div = self.derham.div + _M0preconditioner = MassMatrixPreconditioner(self._M0) + self._M0inv = inverse( self._M0, "pcg", - pc="MassMatrixPreconditioner", + pc=_M0preconditioner, tol=self._options.solver_params.tol, maxiter=self._options.solver_params.maxiter, verbose = False, @@ -346,14 +348,10 @@ def allocate(self): self._M1xrhoB: WeightedMassOperators = None rot_B = LocalRotationMatrix( - self.eq_mhd.b2_1, - self.eq_mhd.b2_2, - self.eq_mhd.b2_3, + self.projected_equil.equil.b1_1, + self.projected_equil.equil.b1_2, + self.projected_equil.equil.b1_3, ) - - rhoB1: Callable = None - rhoB2: Callable = None - rhoB3: Callable = None if isinstance(self._options.rhobar, float): self._M1rho = self._options.rhobar * self.mass_ops.M1 @@ -361,10 +359,7 @@ def allocate(self): "Hcurl", "Hcurl", weights=( - "Ginv", rot_B, - "Ginv", - "sqrt_g", ), name = "M1_xrhoB", assemble = True, @@ -388,14 +383,11 @@ def allocate(self): "Hcurl", "Hcurl", weights=( - "Ginv", rot_B, self._options.rhobar, - "Ginv", - "sqrt_g", - ), - name = "M1_xrhoB", - assemble = True, + ), + name = "M1_xrhoB", + assemble = True, ) if isinstance(self._options.rhobar,FEECVariable): @@ -415,20 +407,19 @@ def allocate(self): "Hcurl", "Hcurl", weights=( - "Ginv", rot_B, self._options.rhobar.spline, - "Ginv", - "sqrt_g", ), name = "M1_xrhoB", assemble = True, ) + _M1rhopreconditioner = MassMatrixPreconditioner(self._M1rho) + self._M1rho_inv = inverse( self._M1rho, "pcg", - pc="MassMatrixPreconditioner", + pc=_M1rhopreconditioner, tol=self._options.solver_params.tol, maxiter=self._options.solver_params.maxiter, verbose=False, @@ -438,11 +429,13 @@ def allocate(self): self._M1nurho: WeightedMassOperators = None + assert isinstance(self._options.nu, (Callable, float)) + if isinstance(self._options.nu, float): self._M1nurho = self._options.nu * self._M1rho if isinstance(self._options.nu, Callable): - nurho: Callable + nurho: Callable = None if isinstance(self._options.rhobar, float): nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar if isinstance(self._options.rhobar,Callable): @@ -498,9 +491,9 @@ def allocate(self): ) if isinstance(self._options.Ebar, FEECVariable): - Ebar1 = lambda *etas: self._options.Ebar.spline(etas)[0] - Ebar2 = lambda *etas: self._options.Ebar.spline(etas)[1] - Ebar3 = lambda *etas: self._options.Ebar.spline(etas)[2] + Ebar1 = lambda *etas: self._options.Ebar.spline(*etas)[0] + Ebar2 = lambda *etas: self._options.Ebar.spline(*etas)[1] + Ebar3 = lambda *etas: self._options.Ebar.spline(*etas)[2] self._P01Ebar = self.basis_ops.create_basis_op( [[Ebar1],[Ebar2],[Ebar3]], @@ -511,11 +504,19 @@ def allocate(self): ) - self._P12 = self.basis_ops.U1 + ones = lambda e1, e2, e3: 1.0 + 0.*(e1 + e2 + e3) + zeroes = lambda e1, e2, e3: 0.*(e1 + e2 + e3) + self._P12 = self.basis_ops.create_basis_op( + [[ones, zeroes, zeroes], + [zeroes, ones, zeroes], + [zeroes, zeroes, ones]], + "Hcurl", + "Hdiv", + assemble = True, + name = "P12", + ) - ones = lambda *etas: 1.0 + 0 * etas - zeroes = lambda *etas: 0 * etas self._O1 = self.basis_ops.create_basis_op( [[ones, zeroes, zeroes]], @@ -529,14 +530,14 @@ def allocate(self): "Hcurl", "H1", assemble = True, - name = "O1", + name = "O2", ) self._O3 = self.basis_ops.create_basis_op( [[zeroes, zeroes, ones]], "Hcurl", "H1", assemble = True, - name = "O1", + name = "O3", ) @@ -547,11 +548,6 @@ def allocate(self): + 2 * self._O1.T @ self._grad.T @ self._M1mu @ self._grad @ self._O1 \ + 2 * self._O2.T @ self._grad.T @ self._M1mu @ self._grad @ self._O2 \ + 2 * self._O3.T @ self._grad.T @ self._M1mu @ self._grad @ self._O3 - - # self._A = self._M1rho + self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) @ self._M0inv @ self._grad.T @ self._M1rho \ - # / (self._options.mass * self._options.omega * self._options.omega) - - # self._B = (self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho) / self._options.omega # ---- block saddle-point system ---------------------------------------- @@ -578,18 +574,18 @@ def allocate(self): ) self._block_B = BlockLinearOperator( - self._block_V1, self._block_V1, blocks=[[None, self._M1rho / self._options.mass], [- self._M1rho / self._options.mass, None]] + self._block_V1, self._block_V1, blocks=[[None, - self._M1rho / self._options.mass], [self._M1rho / self._options.mass, None]] ) self._block_P = BlockLinearOperator( self._block_V0, self._block_V1, - blocks=[[None, self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) / self._options.mass], - [self._M1 @ (self._grad @ self._P0theta + self._P01Ebar) / self._options.mass, None]] + blocks=[[None, self._M1 @ (self._grad @ self._P00theta + self._P01Ebar) / self._options.mass], + [self._M1 @ (self._grad @ self._P00theta + self._P01Ebar) / self._options.mass, None]] ) self._block_Q = BlockLinearOperator( self._block_V1, self._block_V1, - blocks=[[self._options.omega * self._M1rho, self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho], + blocks=[[self._options.omega * self._M1rho, self._divPi - (self._M1xrhoB / self._options.mass) + self._M1nurho], [self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho, - self._options.omega * self._M1rho]] ) @@ -597,36 +593,48 @@ def allocate(self): self._block_V1, self._block_V1, blocks=[[None, self._M1rho / self._options.mass], [self._M1rho / self._options.mass, None]] ) - self._block_Minv = inverse( - self._block_M, - "pcg", - pc="MassMatrixPreconditioner", - tol=self._options.solver_params.tol, - maxiter=self._options.solver_params.maxiter, - verbose = False, - recycle = self._options.solver_params.recycle, + self._block_Minv = BlockLinearOperator( + self._block_V0, self._block_V0, blocks=[[self._M0inv, None], [None, self._M0inv]] ) self._block_rhomatrix = - self._block_Minv @ self._block_Divergence / self._options.omega - self._block_umatrix = - self._block_P @ self._block_rhomatrix - self._block_Q + self._block_umatrix = self._block_Q + self._block_P @ self._block_rhomatrix + + # construction of the initial guess + + _block_Ematrix_inv_approx = BlockLinearOperator( + self._block_V1, self._block_V1, + blocks=[[-self._options.omega * self._M1rho_inv / (self._options.curlcurl_lambda + 1/(self._options.mass*self._options.mass)), None],\ + [None, -self._options.omega * self._M1rho_inv / (self._options.curlcurl_lambda - 1/(self._options.mass*self._options.mass))]] + ) + + _block_E_initialguess = _block_Ematrix_inv_approx.dot(self._block_source) + + _block_umatrix_inv_approx = BlockLinearOperator( + self._block_V1, self._block_V1, + blocks=[[-IdentityOperator(self.derham.V1)/(self._options.mass * self._options.omega), None],\ + [None, IdentityOperator(self.derham.V1)/(self._options.mass * self._options.omega)]] + ) + + _block_u_initialguess = _block_umatrix_inv_approx.dot(_block_E_initialguess) self._block_umatrix_inv = inverse( self._block_umatrix, "gmres", - pc=None, + x0=_block_u_initialguess, tol=self._options.solver_params.tol, maxiter=self._options.solver_params.maxiter, verbose = False, recycle = self._options.solver_params.recycle, ) - self._block_Ematrix = self._block_Acurlcurl / self._options.omega + self._block_B @ self._block_umatrix_inv @ self._block_R + self._block_Ematrix = self._block_Acurlcurl / self._options.omega - self._block_B @ self._block_umatrix_inv @ self._block_R self._block_Ematrix_inv = inverse( self._block_Ematrix, "gmres", - pc=None, + x0=_block_E_initialguess, tol=self._options.solver_params.tol, maxiter=self._options.solver_params.maxiter, verbose = False, @@ -696,7 +704,9 @@ def __call__(self, dt): block_B = self._block_curl.dot(block_E) / self._options.omega - block_u = self._block_umatrix_inv.solve(self._block_R.dot(block_E)) + tmp = - self._block_R.dot(block_E) + + block_u = self._block_umatrix_inv.solve(tmp) block_rho = self._block_rhomatrix.dot(block_u) @@ -707,35 +717,3 @@ def __call__(self, dt): Esin=block_E[0], Ecos=block_E[1], Bsin=block_B[0], Bcos=Block_B[1]) - - # _calE = self._coupled_equations_matrix_inverse.solve(self._calE_RHS) - - # _calEsin = _calE[0] - # _calEcos = _calE[1] - - # _m_curlcurlEsin = self._M1rho.dot(_calEsin) - # _m_curlcurlEcos = self._M1rho.dot(_calEcos) - - # _M1rho_usin = _m_curlcurlEcos.copy() - # _M1rho_ucos = self._options.mass * self._j - _m_curlcurlEsin - - # # --- calculate solutions --- - # self._Esin.vector = self._Acurlcurl_inv_sin.solve(_m_curlcurlEsin / self._options.mass) - # self._Ecos.vector = self._Acurlcurl_inv_cos.solve(_m_curlcurlEcos / self._options.mass) - - # self._Bsin.vector = - self._curl.dot(self._Ecos.vector) / self._options.omega - # self._Bcos.vector = self._curl.dot(self._Esin.vector) / self._options.omega - - # self._usin.vector = self._M1rho_inv.solve(_M1rho_usin) - # self._ucos.vector = self._M1rho_inv.solve(_M1rho_ucos) - - # self._rhosin.vector = self._M0inv.solve(self._grad.T.dot(_M1rho_ucos)) / self._options.omega - # self._rhocos.vector = - self._M0inv.solve(self._grad.T.dot(_M1rho_usin)) / self._options.omega - - # # --- update FEEC variables --- - # self.update_feec_variables( - # rhosin=self._rhosin.vector, rhocos=self._rhocos.vector, - # usin=self._usin.vector, ucos=self._ucos.vector, - # Esin=self._Esin.vector, Ecos=self._Ecos.vector, - # Bsin=self._Bsin.vector, Bcos=self._Bcos.vector) - From 140e8e417d8cca10271fdfde78319a40e1606cf6 Mon Sep 17 00:00:00 2001 From: Monoclod Date: Fri, 31 Jul 2026 13:29:24 +0200 Subject: [PATCH 6/6] Overhauled propagator and added a rough test --- .../propagators/perturbation_system_cold.py | 498 +++--- .../tests/test_dispersion_relation.py | 1418 +++++++++++++++++ 2 files changed, 1677 insertions(+), 239 deletions(-) create mode 100644 src/struphy/propagators/tests/test_dispersion_relation.py diff --git a/src/struphy/propagators/perturbation_system_cold.py b/src/struphy/propagators/perturbation_system_cold.py index 8e6cbd171..356929b9d 100644 --- a/src/struphy/propagators/perturbation_system_cold.py +++ b/src/struphy/propagators/perturbation_system_cold.py @@ -5,10 +5,11 @@ from feectools.api.essential_bc import apply_essential_bc_stencil from feectools.ddm.mpi import mpi as MPI -from feectools.linalg.basic import IdentityOperator +from feectools.linalg.basic import IdentityOperator, ZeroOperator from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace from feectools.linalg.solvers import inverse +from struphy.feec.preconditioner import MassMatrixPreconditioner from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import L2Projector, WeightedMassOperators from struphy.feec.utilities import LocalRotationMatrix @@ -22,18 +23,10 @@ class ColdPlasmaPerturbation(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. + r""":ref:`FEEC ` discretization of a linearized cold plasma fluid system perturbed by a source of frequency :math:`\omega`. + The state variables are the first-order components of the oscillations in electron density, electron velocity, electric field and magnetic field. + Their oscillation-average counterparts (making up the plasma bulk) are passed as parameters of the solver. + Each variable is split into a real (cosine) and imaginary (sine) part to represent the complex quantity fully. """ # ========================================================================= @@ -147,7 +140,7 @@ def Bsin(self, new): def Bcos(self) -> FEECVariable | None: return self._Bcos - @ucos.setter + @Bcos.setter def Bcos(self, new): assert isinstance(new, FEECVariable) assert new.space == "Hdiv" @@ -171,8 +164,10 @@ class Options(OptionsBase): Cosine component of the source term. omega : float, default=1.0 Source term oscillation frequency. - curlcurl_lambda : float, default=1.0 - Coefficient in the curl-curl operator. + c0 : float, default=1.0 + First coefficient in the curl-curl operator. + c1 : float, default=1.0 + Second coefficient in the curl-curl operator. mass : float, default=1.0 Electron mass in relative unis. mu : Callable or float, default=1.0 @@ -193,7 +188,8 @@ class Options(OptionsBase): J: FEECVariable | list = None omega: float = 1.0 - curlcurl_lambda: float = 1.0 + c0: float = 1.0 + c1: float = 1.0 mass: float = 1.0 mu: Callable | float = 1.0 nu: Callable | float = 1.0 @@ -237,18 +233,28 @@ def __post_init__(self): assert isinstance(ei, Callable) # --- physical parameter sanity checks --- + if not isinstance(self.omega, float): + raise TypeError(f"omega must be a float, recieved {type(self.omega)}") if self.omega <= 0: raise ValueError(f"omega must be positive, got {self.omega}") - if self.curlcurl_lambda <= 0: - raise ValueError(f"curlcurl_lambda must be positive, got {self.curlcurl_lambda}") + if not isinstance(self.c0, float): + raise TypeError(f"c0 must be a float, recieved {type(self.c0)}") + if self.c0 <= 0: + raise ValueError(f"c0 must be positive, got {self.c0}") + if not isinstance(self.c1, float): + raise TypeError(f"c1 must be a float, recieved {type(self.c1)}") + if self.c1 <= 0: + raise ValueError(f"c1 must be positive, got {self.c1}") + if not isinstance(self.mass, float): + raise TypeError(f"mass must be a float, recieved {type(self.mass)}") if self.mass <= 0: raise ValueError(f"mass must be positive, got {self.mass}") - if isinstance(mu, float) and self.mu < 0: + if isinstance(self.mu, float) and self.mu < 0: raise ValueError(f"mu must be non-negative, got {self.mu}") - if isinstance(nu, float) and self.nu < 0: + if isinstance(self.nu, float) and self.nu < 0: raise ValueError(f"nu must be non-negative, got {self.nu}") - check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver) + check_option(self.solver, LiteralOptions.OptsGenSolver, LiteralOptions.OptsSaddlePointSolver, LiteralOptions.OptsDirectSolver) if self.solver_params is None: self.solver_params = SolverParameters() @@ -277,6 +283,11 @@ def allocate(self): self._j = self._options.J.spline.vector else: self._j = self.derham.P1(self._options.J) # works if J is a list of Callables + + zeroes = lambda x,y,z: 0. * (x+y+z) + self._zerovectorV1 = self.derham.P1([zeroes,zeroes,zeroes]) # need this for RHS assembly + self._zerovectorV2 = self.derham.P2([zeroes,zeroes,zeroes]) + self._zerofield = self.derham.P0(zeroes) # ---- unconstrained operators (for RHS assembly) ---------------------- @@ -293,20 +304,37 @@ def allocate(self): self._M0, "pcg", pc=_M0preconditioner, - tol=self._options.solver_params.tol, + tol=1e-12, maxiter=self._options.solver_params.maxiter, verbose = False, recycle = self._options.solver_params.recycle, ) + _M1preconditioner = MassMatrixPreconditioner(self.mass_ops.M1) + + self._M1inv = inverse( + self._M1, + "pcg", + pc=_M1preconditioner, + tol=1e-12, + maxiter=self._options.solver_params.maxiter, + verbose=False, + recycle=self._options.solver_params.recycle, + ) + self._M1mu: WeightedMassOperators = None self._M2mu: WeightedMassOperators = None self._M3mu: WeightedMassOperators = None if isinstance(self._options.mu, float): - self._M1mu = self._options.mu * self._M1 - self._M2mu = self._options.mu * self._M2 - self._M3mu = self._options.mu * self.mass_ops.M3 + if self._options.mu == 0.: + self._M1mu = ZeroOperator(self.derham.V1,self.derham.V1) + self._M2mu = ZeroOperator(self.derham.V2,self.derham.V2) + self._M3mu = ZeroOperator(self.derham.V3,self.derham.V3) + else: + self._M1mu = self._options.mu * self._M1 + self._M2mu = self._options.mu * self._M2 + self._M3mu = self._options.mu * self.mass_ops.M3 else: self._M1mu = self.mass_ops.create_weighted_mass( "Hcurl", @@ -354,29 +382,44 @@ def allocate(self): ) if isinstance(self._options.rhobar, float): - self._M1rho = self._options.rhobar * self.mass_ops.M1 - self._M1xrhoB = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - rot_B, + if self._options.rhobar == 0.: + self._M1rho = ZeroOperator(self.derham.V1,self.derham.V1) + self._M1xrhoB = ZeroOperator(self.derham.V1,self.derham.V1) + else: + rhoarray = lambda e1,e2,e3: self._options.rhobar + 0.*(e1+e2+e3) + self._M1rho = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + rhoarray, + ), + name = "M1rho", + assemble = True, + ) + self._M1xrhoB = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + rot_B, + rhoarray, ), name = "M1_xrhoB", assemble = True, - ) - self._M1xrhoB *= self._options.rhobar + ) if isinstance(self._options.rhobar, Callable): self._M1rho = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - "Ginv", - "sqrt_g", - self._options.rhobar, - ), - name = "M1rho", - assemble = True, + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + self._options.rhobar, + ), + name = "M1rho", + assemble = True, ) self._M1xrhoB = self.mass_ops.create_weighted_mass( @@ -392,16 +435,16 @@ def allocate(self): if isinstance(self._options.rhobar,FEECVariable): self._M1rho = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - "Ginv", - "sqrt_g", - self._options.rhobar.spline, - ), - name = "M1rho", - assemble = True, - ) + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + self._options.rhobar.spline, + ), + name = "M1rho", + assemble = True, + ) self._M1xrhoB = self.mass_ops.create_weighted_mass( "Hcurl", @@ -409,57 +452,66 @@ def allocate(self): weights=( rot_B, self._options.rhobar.spline, - ), - name = "M1_xrhoB", - assemble = True, + ), + name = "M1_xrhoB", + assemble = True, ) - _M1rhopreconditioner = MassMatrixPreconditioner(self._M1rho) - - self._M1rho_inv = inverse( - self._M1rho, - "pcg", - pc=_M1rhopreconditioner, - tol=self._options.solver_params.tol, - maxiter=self._options.solver_params.maxiter, - verbose=False, - recycle=self._options.solver_params.recycle, - ) - + # if not isinstance(self._options.rhobar,float): + # _M1rhopreconditioner = MassMatrixPreconditioner(self._M1rho) + + # self._M1rho_inv = inverse( + # self._M1rho, + # "pcg", + # pc=_M1rhopreconditioner, + # tol=1e-12, + # maxiter=self._options.solver_params.maxiter, + # verbose=False, + # recycle=self._options.solver_params.recycle, + # ) self._M1nurho: WeightedMassOperators = None assert isinstance(self._options.nu, (Callable, float)) - if isinstance(self._options.nu, float): - self._M1nurho = self._options.nu * self._M1rho - - if isinstance(self._options.nu, Callable): - nurho: Callable = None - if isinstance(self._options.rhobar, float): - nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar - if isinstance(self._options.rhobar,Callable): - nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar(*etas) - if isinstance(self._options.rhobar,FEECVariable): - nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar.spline(*etas) - - self._M1nurho = self.mass_ops.create_weighted_mass( - "Hcurl", - "Hcurl", - weights=( - "Ginv", - "sqrt_g", - nurho, - ), - name = "M1nurho", - assemble = True, - ) + self._M1nurho = ZeroOperator(self.derham.V1,self.derham.V1) + + if self._options.rhobar != 0.: + if isinstance(self._options.nu, float): + if self._options.nu == 0.: + self._M1nurho = ZeroOperator(self.derham.V1,self.derham.V1) + else: + self._M1nurho = self._options.nu * self._M1rho + + if isinstance(self._options.nu, Callable): + nurho: Callable = None + if isinstance(self._options.rhobar, float): + nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar + if isinstance(self._options.rhobar,Callable): + nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar(*etas) + if isinstance(self._options.rhobar,FEECVariable): + nurho = lambda *etas: self._options.nu(*etas) * self._options.rhobar.spline(*etas) + + self._M1nurho = self.mass_ops.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + "sqrt_g", + nurho, + ), + name = "M1nurho", + assemble = True, + ) self._P00theta: BasisProjectionOperators = None if isinstance(self._options.theta, float): - self._P00theta = self._options.theta * IdentityOperator(self.derham.V0) + if self._options.theta == 0.: + self._P00theta = ZeroOperator(self.derham.V0,self.derham.V0) + else: + self._P00theta = self._options.theta * IdentityOperator(self.derham.V0) if isinstance(self._options.theta, Callable): self._P00theta = self.basis_ops.create_basis_op( @@ -507,191 +559,154 @@ def allocate(self): ones = lambda e1, e2, e3: 1.0 + 0.*(e1 + e2 + e3) zeroes = lambda e1, e2, e3: 0.*(e1 + e2 + e3) - self._P12 = self.basis_ops.create_basis_op( - [[ones, zeroes, zeroes], - [zeroes, ones, zeroes], - [zeroes, zeroes, ones]], - "Hcurl", - "Hdiv", - assemble = True, - name = "P12", - ) + # self._P12 = self.basis_ops.create_basis_op( + # [[ones, zeroes, zeroes], + # [zeroes, ones, zeroes], + # [zeroes, zeroes, ones]], + # "Hcurl", + # "Hdiv", + # assemble = True, + # name = "P12", + # ) - self._O1 = self.basis_ops.create_basis_op( - [[ones, zeroes, zeroes]], - "Hcurl", - "H1", - assemble = True, - name = "O1", - ) - self._O2 = self.basis_ops.create_basis_op( - [[zeroes, ones, zeroes]], - "Hcurl", - "H1", - assemble = True, - name = "O2", - ) - self._O3 = self.basis_ops.create_basis_op( - [[zeroes, zeroes, ones]], - "Hcurl", - "H1", - assemble = True, - name = "O3", - ) + # self._O1 = self.basis_ops.create_basis_op( + # [[ones, zeroes, zeroes]], + # "Hcurl", + # "H1", + # assemble = True, + # name = "O1", + # ) + # self._O2 = self.basis_ops.create_basis_op( + # [[zeroes, ones, zeroes]], + # "Hcurl", + # "H1", + # assemble = True, + # name = "O2", + # ) + # self._O3 = self.basis_ops.create_basis_op( + # [[zeroes, zeroes, ones]], + # "Hcurl", + # "H1", + # assemble = True, + # name = "O3", + # ) - self._Acurlcurl = self._curl.T @ self._M2 @ self._curl - self._options.curlcurl_lambda * self._M1 + self._Acurlcurl = self._options.c0 * self._curl.T @ self._M2 @ self._curl - self._options.c1 * self._M1 - self._divPi = - self._curl.T @ self._M2mu @ self._curl \ - - 2/3 * self._P12.T @ self._div.T @ self._M3mu @ self._div @ self._P12 \ - + 2 * self._O1.T @ self._grad.T @ self._M1mu @ self._grad @ self._O1 \ - + 2 * self._O2.T @ self._grad.T @ self._M1mu @ self._grad @ self._O2 \ - + 2 * self._O3.T @ self._grad.T @ self._M1mu @ self._grad @ self._O3 + # self._divPi = - self._curl.T @ self._M2mu @ self._curl \ + # - 2/3 * self._P12.T @ self._div.T @ self._M3mu @ self._div @ self._P12 \ + # + 2 * self._O1.T @ self._grad.T @ self._M1mu @ self._grad @ self._O1 \ + # + 2 * self._O2.T @ self._grad.T @ self._M1mu @ self._grad @ self._O2 \ + # + 2 * self._O3.T @ self._grad.T @ self._M1mu @ self._grad @ self._O3 + self._divPi = - self._curl.T @ self._M2mu @ self._curl + 4/3 * self._M1mu @ self._grad @ self._M0inv @ self._grad.T @ self._M1 - # ---- block saddle-point system ---------------------------------------- + # ---- block Schur solve system ---------------------------------------- - self._block_V0 = BlockVectorSpace(self.derham.V0, self.derham.V0) - self._block_V1 = BlockVectorSpace(self.derham.V1, self.derham.V1) - self._block_V2 = BlockVectorSpace(self.derham.V2, self.derham.V2) + self._V0squared = BlockVectorSpace(self.derham.V0, self.derham.V0) + self._V1squared = BlockVectorSpace(self.derham.V1, self.derham.V1) + self._V2squared = BlockVectorSpace(self.derham.V2, self.derham.V2) - self._block_source = BlockVector(self._block_V1, blocks=[self._M1.dot(self._j), None]) + self._space_block1 = BlockVectorSpace(self._V0squared, self._V1squared) + self._space_block2 = BlockVectorSpace(self._V1squared, self._V2squared) - self._block_M = BlockLinearOperator( - self._block_V0, self._block_V0, blocks=[[self._M0, None], [None, self._M0]] + self._source = BlockVector(self._V1squared, blocks=[self._zerovectorV1, self._options.omega * self._M1.dot(self._j)]) + self._zerovectorV2squared = BlockVector(self._V2squared, blocks=[self._zerovectorV2, self._zerovectorV2]) + self._block_source = BlockVector(self._space_block2, blocks=[self._source, self._zerovectorV2squared]) + + self._block_M0 = BlockLinearOperator( + self._V0squared, self._V0squared, blocks=[[None, - self._M0], [self._M0, None]] + ) + + self._block_M1 = BlockLinearOperator( + self._V1squared, self._V1squared, blocks=[[None, - self._M1], [self._M1, None]] ) self._block_Divergence = BlockLinearOperator( - self._block_V1, self._block_V0, blocks=[[None, - self._grad.T @ self._M1rho], [self._grad.T @ self._M1rho, None]] + self._V1squared, self._V0squared, blocks=[[self._grad.T @ self._M1rho, None], [None, self._grad.T @ self._M1rho]] ) self._block_curl = BlockLinearOperator( - self._block_V1, self._block_V2, blocks=[[None, - self._curl], [self._curl, None]] + self._V1squared, self._V2squared, blocks=[[self._curl, None], [None, self._curl]] ) - self._block_Acurlcurl = BlockLinearOperator( - self.block_V, self._block_V1, blocks=[[self._Acurlcurl, None], [None, self._Acurlcurl]] + self._block_curlV2 = BlockLinearOperator( + self._V2squared, self._V1squared, blocks=[[self._curl.T @ self._M2, None], [None, self._curl.T @ self._M2]] ) - self._block_B = BlockLinearOperator( - self._block_V1, self._block_V1, blocks=[[None, - self._M1rho / self._options.mass], [self._M1rho / self._options.mass, None]] + self._block_ImV2 = BlockLinearOperator( + self._V2squared, self._V2squared, blocks=[[None, - IdentityOperator(self.derham.V2)], [IdentityOperator(self.derham.V2), None]] ) - self._block_P = BlockLinearOperator( - self._block_V0, self._block_V1, - blocks=[[None, self._M1 @ (self._grad @ self._P00theta + self._P01Ebar) / self._options.mass], - [self._M1 @ (self._grad @ self._P00theta + self._P01Ebar) / self._options.mass, None]] + self._block_Acurlcurl = BlockLinearOperator( + self._V1squared, self._V1squared, blocks=[[self._Acurlcurl, None], [None, self._Acurlcurl]] ) - self._block_Q = BlockLinearOperator( - self._block_V1, self._block_V1, - blocks=[[self._options.omega * self._M1rho, self._divPi - (self._M1xrhoB / self._options.mass) + self._M1nurho], - [self._divPi - self._M1xrhoB / self._options.mass + self._M1nurho, - self._options.omega * self._M1rho]] + self._block_M1rho = BlockLinearOperator( + self._V1squared, self._V1squared, blocks=[[self._M1rho, None], [None, self._M1rho]] ) - self._block_R = BlockLinearOperator( - self._block_V1, self._block_V1, blocks=[[None, self._M1rho / self._options.mass], [self._M1rho / self._options.mass, None]] + self._block_iM1rho = BlockLinearOperator( + self._V1squared, self._V1squared, blocks=[[None, - self._M1rho], [self._M1rho, None]] ) - self._block_Minv = BlockLinearOperator( - self._block_V0, self._block_V0, blocks=[[self._M0inv, None], [None, self._M0inv]] + self._block_P = BlockLinearOperator( + self._V0squared, self._V1squared, + blocks=[[self._M1 @ (self._grad @ self._P00theta + self._P01Ebar) / self._options.mass, None], + [None, self._M1 @ (self._grad @ self._P00theta + self._P01Ebar) / self._options.mass]] ) - self._block_rhomatrix = - self._block_Minv @ self._block_Divergence / self._options.omega - - self._block_umatrix = self._block_Q + self._block_P @ self._block_rhomatrix - - # construction of the initial guess + self._block_Q = BlockLinearOperator( + self._V1squared, self._V1squared, + blocks=[[self._divPi - (self._M1xrhoB / self._options.mass) + self._M1nurho, self._options.omega * self._M1rho], + [- self._options.omega * self._M1rho, self._divPi - (self._M1xrhoB / self._options.mass) + self._M1nurho]] + ) - _block_Ematrix_inv_approx = BlockLinearOperator( - self._block_V1, self._block_V1, - blocks=[[-self._options.omega * self._M1rho_inv / (self._options.curlcurl_lambda + 1/(self._options.mass*self._options.mass)), None],\ - [None, -self._options.omega * self._M1rho_inv / (self._options.curlcurl_lambda - 1/(self._options.mass*self._options.mass))]] + # constru + self._block_A = BlockLinearOperator( + self._space_block1, self._space_block1, + blocks=[[self._options.omega * self._block_M0, self._block_Divergence], [self._block_P, self._block_Q]] ) - _block_E_initialguess = _block_Ematrix_inv_approx.dot(self._block_source) + self._block_B = BlockLinearOperator( + self._space_block2, self._space_block1, + blocks=[[None, None], [self._block_M1rho / self._options.mass, None]] + ) - _block_umatrix_inv_approx = BlockLinearOperator( - self._block_V1, self._block_V1, - blocks=[[-IdentityOperator(self.derham.V1)/(self._options.mass * self._options.omega), None],\ - [None, IdentityOperator(self.derham.V1)/(self._options.mass * self._options.omega)]] + self._block_D = BlockLinearOperator( + self._space_block1, self._space_block2, + blocks=[[None, self._options.omega * self._block_iM1rho / self.options.mass], [None, None]] ) - _block_u_initialguess = _block_umatrix_inv_approx.dot(_block_E_initialguess) + self._block_C = BlockLinearOperator( + self._space_block2, self._space_block2, + blocks=[[self._block_Acurlcurl, None], + [self._block_curl, - self._options.omega * self._block_ImV2]] + ) - self._block_umatrix_inv = inverse( - self._block_umatrix, + self._block_Cinv = inverse( + self._block_C, "gmres", - x0=_block_u_initialguess, - tol=self._options.solver_params.tol, + x0=None, + tol=1e-10, maxiter=self._options.solver_params.maxiter, - verbose = False, + verbose = True, recycle = self._options.solver_params.recycle, ) - self._block_Ematrix = self._block_Acurlcurl / self._options.omega - self._block_B @ self._block_umatrix_inv @ self._block_R + self._block_A_schur = self._block_A - self._block_B @ self._block_Cinv @ self._block_D - self._block_Ematrix_inv = inverse( - self._block_Ematrix, + self._block_A_schur_inv = inverse( + self._block_A_schur, "gmres", - x0=_block_E_initialguess, + x0=None, tol=self._options.solver_params.tol, maxiter=self._options.solver_params.maxiter, - verbose = False, + verbose = True, recycle = self._options.solver_params.recycle, ) - # self._coupled_equations_matrix = BlockLinearOperator( - # self._block_domain, self._block_codomain, blocks=[[self._A, -self._B], [self._B, self._A]] - # ) - - # M1rhoinv_j = self._M1rho_inv.solve(self._j) - # A_j = self._A.dot(M1rhoinv_j) - # minusB_j = - self.B.dot(M1rhoinv_j) - - # self._calEsin0: StencilVector = None - # self._calEcos0: StencilVector = None - - # # --- copy current state --- - # Esin0 = self.variables.Esin.spline.vector - # Ecos0 = self.variables.Ecos.spline.vector - - # self._calEsin0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Esin0)) - # self._calEcos0 = self._M1rho_inv.dot(self._options.mass * self._Acurlcurl.dot(Ecos0)) - - # self._calE0 = BlockVector(self._block_domain, blocks=[self._calEsin0, self._calEcos0]) - - # self._coupled_equations_matrix_inverse = inverse( - # self._coupled_equations_matrix, - # solver="gmres", - # x0=self._calE0, - # tol=self._options.solver_params.tol, - # maxiter=self._options.solver_params.maxiter, - # verbose=False, - # recycle=True, - # ) - - # # --- build inverses of the curl-curl matrices with good initial guesses --- - # self._Acurlcurl_inv_sin = inverse( - # self._Acurlcurl, - # solver="pcg", - # x0=Esin0, - # tol=self._options.solver_params.tol, - # maxiter=self._options.solver_params.maxiter, - # verbose=False, - # recycle=True, - # ) - - # self._Acurlcurl_inv_cos = inverse( - # self._Acurlcurl, - # solver="pcg", - # x0=Ecos0, - # tol=self._options.solver_params.tol, - # maxiter=self._options.solver_params.maxiter, - # verbose=False, - # recycle=True, - # ) - # ========================================================================= ### Equation solve @@ -699,21 +714,26 @@ def allocate(self): def __call__(self, dt): - # --- calculate auxilliary vectors --- - block_E = self._block_Ematrix_inv.solve(self._block_source) + tmp_z = self._block_Cinv.solve(self._block_source) + + rhs_x1 = - self._block_B.dot(tmp_z) + + tmp_x1 = self._block_A_schur_inv.solve(rhs_x1) + + tmp_x2 = tmp_z - self._block_Cinv.solve(self._block_D.dot(tmp_x1)) - block_B = self._block_curl.dot(block_E) / self._options.omega + comp_rho = tmp_x1[0] - tmp = - self._block_R.dot(block_E) + comp_u = tmp_x1[1] - block_u = self._block_umatrix_inv.solve(tmp) + comp_E = tmp_x2[0] - block_rho = self._block_rhomatrix.dot(block_u) + comp_B = tmp_x2[1] # --- update FEEC variables --- self.update_feec_variables( - rhosin=block_rho[0], rhocos=block_rho[1], - usin=block_u[0], ucos=block_u[1], - Esin=block_E[0], Ecos=block_E[1], - Bsin=block_B[0], Bcos=Block_B[1]) + rhosin=comp_rho[1], rhocos=comp_rho[0], + usin=comp_u[1], ucos=comp_u[0], + Esin=comp_E[1], Ecos=comp_E[0], + Bsin=comp_B[1], Bcos=comp_B[0]) diff --git a/src/struphy/propagators/tests/test_dispersion_relation.py b/src/struphy/propagators/tests/test_dispersion_relation.py new file mode 100644 index 000000000..b8caf6fee --- /dev/null +++ b/src/struphy/propagators/tests/test_dispersion_relation.py @@ -0,0 +1,1418 @@ +import logging + +import numpy as np +import cunumpy as xp +from matplotlib import pyplot as plt +from feectools.ddm.mpi import mpi as MPI + +from struphy import ( + BinningPlot, + BoundaryParameters, + LoadingParameters, + WeightsParameters, + domains, + equils, + perturbations, + set_logging_level, +) +from struphy.feec.mass import L2Projector, WeightedMassOperators +from struphy.feec.basis_projection_ops import BasisProjectionOperators +from struphy.feec.psydac_derham import Derham +from struphy.geometry.base import Domain +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable +from struphy.propagators.base import Propagator +from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB +from struphy.propagators.perturbation_system_cold import ColdPlasmaPerturbation +from struphy.topology.grids import TensorProductGrid +from struphy.utils.pyccel import Pyccelkernel + +logger = logging.getLogger("struphy") +set_logging_level(logging.DEBUG) + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() +plt.rcParams.update({"font.size": 22}) + +tol: float = 1e-13 + +m_e: float = 9.1094e-31 # kg +e: float = 1.51874e-14 # kg^1/2 m^3/2 s^{-1} +c: float = 299792458 # m/s +mu0: float = 4*np.pi / (c**2) +kB: float = 1.380649e-23 # kg m^2 s^{-2} K^{-1} + +N: float = 1e18 # m^{-3} + +kT = 10*1.602176621e-19 #J +T: float = kT / kB + +eratio: float = 1.602e-19 / e # e SI over e Gauss +B: float = 1e-4 # Tesla +B_gauss: float = B * c * eratio + +omega_pe_scale = np.sqrt((4*np.pi*N*(e**2)/(m_e))) + +debye: float = np.sqrt(kT / (4*np.pi * N * (e**2))) + +L = c / omega_pe_scale + +Vcyclo: float = L*e*B_gauss/(2*np.pi*m_e * c) # cyclotron speed +Valfven: float = B_gauss / np.sqrt(m_e*N*mu0*c) # Alfven speed + +V: float = np.sqrt(kT / m_e) + +t: float = L/V +print(f"{t=}") + +E = kT / (e*L) + +B = E * c / V + +alpha: float = (debye / L)**2 + +c_normalized = c/V +print(f"{c_normalized=}") + +omega: float = 100000000. * t +print(f"{omega=}") +rhobar: float = 0. +mass: float = 1. +theta: float = 0. # 200. +zeta: float = np.pi * 3/5 # between 0 and pi/2 +B0: float = 5. +coszeta = np.cos(zeta) +sinzeta = np.sin(zeta) + +B_x = 0. # B0*sinzeta +B_y = 0. +B_z = 0. # B0*coszeta + +omega_pe: float = np.sqrt(rhobar) / mass +print(f"{omega_pe=}") + +omega_pe_normalized: float = omega_pe * omega_pe_scale * t +print(f"{omega_pe_normalized=}") + +d_omega = np.sqrt(np.abs(omega**2 - omega_pe_normalized**2)) +k_light = d_omega / c_normalized + +# c_sound_normalized = np.sqrt(theta / mass) # only activate k_sound and anything related to it when theta is nonzero +# k_sound = d_omega / c_sound_normalized + +c0: float = (c_normalized**2) * alpha +c1: float = (omega**2) * alpha + +print(f"{k_light/L=}") +# print(f"{k_sound/L=}") + +def test_dispersion_relation_1d(): + ksquared = lambda k1,k2: ((k1**2)+(k2**2))*id3x3 + + ktensork = lambda k1,k2: np.array([[k1**2, k1*k2, 0], + [k2*k1, k2**2, 0], + [0, 0, 0]]) + + rotB = np.array([[0, -B_z, B_y], + [B_z, 0, -B_x], + [-B_y, B_x, 0]]) + + id3x3 = np.identity(3) + + def matrix(k1,k2): + mat = np.zeros((6,6), dtype="complex") + + mat[:3,:3] = id3x3 - theta / (mass * (omega**2)) * ktensork(k1,k2) - 1j / (mass*c_normalized*omega) * rotB + mat[:3,3:6] = 1j * np.sqrt(4*np.pi) * omega_pe * id3x3 + mat[3:6,:3] = - 1j * omega_pe / (np.sqrt(4*np.pi)*alpha) * id3x3 + mat[3:6,3:6] = (omega**2) * id3x3 - (c_normalized**2) * (ksquared(k1,k2) - ktensork(k1,k2)) + + return mat + + determinant = lambda k1, k2: np.linalg.det(matrix(k1,k2)) + + print(determinant(k_light,0)/(c_normalized**2)) + # print(determinant(k_sound,0)/(c_normalized**2)) + # print(determinant((k_light+k_sound)/2,0)) + + if d_omega == 0.: + kmax = 0.5 + else: + kmax = 1.5 * k_light # np.maximum(k_light, k_sound) + + Nel = 450 + + k = np.linspace(-kmax, kmax, Nel) + + k1, k2 = np.meshgrid(k, k, indexing='ij') + + det = np.zeros((Nel,Nel)) + + for i in range(Nel): + for j in range(Nel): + val = np.real(determinant(k1[i,j],k2[i,j])) + det[i,j] = 0. if np.abs(val) <= tol else val + + det /= c_normalized**2 + + detslice = (det[int(Nel/2),:]+det[int(Nel/2)+1,:])/2 + + # ax = plt.figure().add_subplot(projection='3d') + + # surface = ax.plot_surface(k1,k2,det,linewidth=0) + + # plt.contour(k1, k2, det, levels=[0.]) + + # plt.colorbar() + + plt.figure(1) + plotyscale = np.max(detslice) + plotxscale = 2*kmax/L + plt.title("Determinant of the system of oscillations") + plt.plot(k/L, detslice, label="Horizontal slice of determinant") + plt.axhline(y=0.,color="black") + plt.axvline(x=k_light/L, ymin=0.05, linestyle='--',color="black") + plt.axvline(x=-k_light/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=k_sound/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=-k_sound/L, ymin=0.05, linestyle='--',color="black") + plt.xlabel("Wave vectors [$m^{-1}$]") + plt.text(k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{light}}}}$") + plt.text(-k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{light}}}}$") + # plt.text(k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{therminc}}}}$") + # plt.text(-k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{thermic}}}}$") + plt.legend() + + # plt.show() + # exit() + + p = 3 + Ngrid = 2**8 + + Nfft = 100 + k_cutoff = kmax + dk = k_cutoff / Nfft + maxL: float = 2*np.pi / dk # so that our largest wavenumber value is included + + domain = domains.Cuboid(l1=-maxL/2 ,r1=maxL/2) + equil = equils.HomogenSlab(B0x=B_x, B0y=B_y, B0z=B_z) + equil.domain = domain + + e = np.linspace(0., 1., Nel) + e_x, e_y, e_z = domain(e,0.,0.) # the values the field will be sampled on that will correspond exactly to the k array after the FFT + e_x = e_x[:,0,0] + print(e_x) + cellsize = (maxL / Nel) + e_k = np.linspace(-Nel/2 * dk, (Nel/2-1) * dk, Nel) + + J0: float = 1. + j_physical = lambda x,y,z: J0 * np.sinc(k_cutoff/np.pi * x) + zeroes = lambda x,y,z: 0.*(x+y+z) + print(f"{k_cutoff/np.pi *maxL=}") + zeroes = lambda x,y,z: 0. * (x+y+z) + + j_pulled_1 = lambda e1,e2,e3: domain.pull([j_physical,j_physical,zeroes],e1,e2,e3,kind="1", squeeze_out=False)[0] + j_pulled_2 = lambda e1,e2,e3: domain.pull([j_physical,j_physical,zeroes],e1,e2,e3,kind="1", squeeze_out=False)[1] + j_pulled_3 = lambda e1,e2,e3: domain.pull([j_physical,j_physical,zeroes],e1,e2,e3,kind="1", squeeze_out=False)[2] + + print(f"{np.shape(e)=}") + print(f"{np.shape(e_x)=}") + print(f"{np.shape(j_pulled_1(e,0.,0.)[:,0,0])=}") + print(f"{np.shape(j_physical(e_x,0.,0.))=}") + + jdiff1 = lambda e1,e2,e3: j_pulled_1(e1,e2,e3) / maxL - j_physical(domain(e1,e2,e3)[0],domain(e1,e2,e3)[1],domain(e1,e2,e3)[2]) + jdiff2 = lambda e1,e2,e3: j_pulled_2(e1,e2,e3) - j_physical(domain(e1,e2,e3)[0],domain(e1,e2,e3)[1],domain(e1,e2,e3)[2]) + jdiff3 = lambda e1,e2,e3: j_pulled_3(e1,e2,e3) - zeroes(domain(e1,e2,e3)[0],domain(e1,e2,e3)[1],domain(e1,e2,e3)[2]) + + plt.figure(2) + plt.subplot(1,3,1) + plt.plot(e, jdiff1(e,0.,0.)[:,0,0],label="pulled j1") + plt.legend() + plt.subplot(1,3,2) + plt.plot(e, jdiff2(e,0.,0.)[:,0,0],label="pulled j2") + plt.legend() + plt.subplot(1,3,3) + plt.plot(e, jdiff3(e,0.,0.)[:,0,0],label="pulled j3") + plt.legend() + + plt.figure(3) + plt.subplot(2,3,1) + plt.plot(e, j_physical(e_x,0.,0.),label="physical j1") + plt.legend() + plt.subplot(2,3,2) + plt.plot(e, j_physical(e_x,0.,0.),label="physical j2") + plt.legend() + plt.subplot(2,3,3) + plt.plot(e, zeroes(e_x,0.,0.),label="physical j3") + plt.legend() + plt.subplot(2,3,4) + plt.plot(e, j_pulled_1(e,0.,0.)[:,0,0],label="pulled j1") + plt.legend() + plt.subplot(2,3,5) + plt.plot(e, j_pulled_2(e,0.,0.)[:,0,0],label="pulled j2") + plt.legend() + plt.subplot(2,3,6) + plt.plot(e, j_pulled_3(e,0.,0.)[:,0,0],label="pulled j3") + plt.legend() + + # plt.show() + # exit() + + degree = (p,1,1) + num_elements = (Ngrid,1,1) + bcs = (("dirichlet","dirichlet"), None, None) + + grid = TensorProductGrid(num_elements=num_elements) + derham_opts = DerhamOptions(degree=degree, bcs=bcs) + derham = Derham(grid=grid, options=derham_opts, comm=comm) + projected_equil = ProjectedFluidEquilibriumWithB(equil=equil, derham=derham) + + mass_ops = WeightedMassOperators(derham=derham, domain=domain) + basis_ops = BasisProjectionOperators(derham=derham, domain=domain) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + Propagator.basis_ops = basis_ops + Propagator.projected_equil = projected_equil + + J = FEECVariable(space="Hcurl") + J.allocate(derham=derham, domain=domain) + J.spline.vector = derham.P1([j_pulled_1,j_pulled_2,j_pulled_3]) + + ee1, ee2, ee3 = np.meshgrid(np.linspace(0.,1.,4),np.linspace(0.,1.,5),np.linspace(0.,1.,6), indexing="ij") + # print(f"{np.shape(sincheck(ee1,ee2,ee3))=}") + print(f"{np.shape(j_pulled_1(ee1,ee2,ee3))=}") + print(f"{e=}") + # exit() + + plt.figure(4) + # plt.plot(e, J0.spline(e,0.,0.)[:,0,0],label="projected j1") + # plt.plot(e, j_pulled_1(e,0.,0.)[:,0,0],'x',label="j1") + plt.subplot(3,3,1) + plt.plot(e, j_pulled_1(e,0,0.)[:,0,0] ,label="j1") + plt.legend() + plt.subplot(3,3,2) + plt.plot(e, j_pulled_2(e,0,0.)[:,0,0],label="j2") + plt.legend() + plt.subplot(3,3,3) + plt.plot(e, j_pulled_3(e,0,0.)[:,0,0],label="j3") + plt.legend() + plt.subplot(3,3,4) + plt.plot(e, J.spline(e,0.,0.)[0][:,0,0],label="projected j1") + plt.legend() + plt.subplot(3,3,5) + plt.plot(e, J.spline(e,0.,0.)[1][:,0,0],label="projected j2") + plt.legend() + plt.subplot(3,3,6) + plt.plot(e, J.spline(e,0.,0.)[2][:,0,0],label="projected j3") + plt.legend() + plt.subplot(3,3,7) + plt.plot(e, J.spline(e,0.,0.)[0][:,0,0] - j_pulled_1(e,0,0.)[:,0,0],label="difference j1") + plt.legend() + plt.subplot(3,3,8) + plt.plot(e, J.spline(e,0.,0.)[1][:,0,0] - j_pulled_2(e,0,0.)[:,0,0],label="difference j2") + plt.legend() + plt.subplot(3,3,9) + plt.plot(e, J.spline(e,0.,0.)[2][:,0,0] - j_pulled_3(e,0,0.)[:,0,0],label="difference j3") + # plt.show() + # exit() + + plt.figure(5) + # plt.title("Fourier transform of source term") + plt.subplot(3,1,1) + plt.plot(e_k/L, cellsize * np.fft.fftshift(np.abs(np.fft.fft(domain.push(J.spline,e,0.,0.,kind="1")[0][:,0,0]))),label="FFT of j1") + plt.legend() + plt.subplot(3,1,2) + plt.plot(e_k/L, cellsize * np.fft.fftshift(np.abs(np.fft.fft(domain.push(J.spline,e,0.,0.,kind="1")[1][:,0,0]))),label="FFT of j2") + plt.legend() + plt.subplot(3,1,3) + plt.plot(e_k/L, cellsize * np.fft.fftshift(np.abs(np.fft.fft(domain.push(J.spline,e,0.,0.,kind="1")[2][:,0,0]))),label="FFT of j3") + plt.legend() + # plt.show() + # exit() + + solver_params = SolverParameters( + tol=1e-10, + maxiter=1600, + info=True, + recycle=True, + ) + + _rhosin = FEECVariable(space="H1") + _rhosin.allocate(derham=derham, domain=domain) + + _rhocos = FEECVariable(space="H1") + _rhocos.allocate(derham=derham, domain=domain) + + _usin = FEECVariable(space="Hcurl") + _usin.allocate(derham=derham, domain=domain) + + _ucos = FEECVariable(space="Hcurl") + _ucos.allocate(derham=derham, domain=domain) + + _Esin = FEECVariable(space="Hcurl") + _Esin.allocate(derham=derham, domain=domain) + + _Ecos = FEECVariable(space="Hcurl") + _Ecos.allocate(derham=derham, domain=domain) + + _Bsin = FEECVariable(space="Hdiv") + _Bsin.allocate(derham=derham, domain=domain) + + _Bcos = FEECVariable(space="Hdiv") + _Bcos.allocate(derham=derham, domain=domain) + + solver = ColdPlasmaPerturbation() + solver.variables.rhosin = _rhosin + solver.variables.rhocos = _rhocos + solver.variables.usin = _usin + solver.variables.ucos = _ucos + solver.variables.Esin = _Esin + solver.variables.Ecos = _Ecos + solver.variables.Bsin = _Bsin + solver.variables.Bcos = _Bcos + + solver.options = solver.Options( + J=J, + omega=omega, + c0=c0, + c1=c1, + mass=mass, + mu=0., + nu=0., + rhobar=rhobar, + theta=theta, + Ebar=[zeroes, zeroes, zeroes], + solver="gmres", + solver_params=solver_params, + ) + + solver.allocate() + + dt=1.0 + print("Hi man") + solver(dt) + print("Bye man") + + Esinvalues = domain.push(_Esin.spline, e, 0., 0., kind="1") + Ecosvalues = domain.push(_Ecos.spline, e, 0., 0., kind="1") + + usinvalues = domain.push(_usin.spline, e, 0., 0., kind="1") + ucosvalues = domain.push(_ucos.spline, e, 0., 0., kind="1") + + rhosinvalues = domain.push(_rhosin.spline, e, 0., 0., kind="0") + rhocosvalues = domain.push(_rhocos.spline, e, 0., 0., kind="0") + + print(Esinvalues.shape) + print(Ecosvalues.shape) + + print(usinvalues.shape) + print(ucosvalues.shape) + + print(rhosinvalues.shape) + print(rhocosvalues.shape) + + Esinvalues1 = Esinvalues[0,:,0,0] + Esinvalues2 = Esinvalues[1,:,0,0] + Esinvalues3 = Esinvalues[2,:,0,0] + + Ecosvalues1 = Esinvalues[0,:,0,0] + Ecosvalues2 = Esinvalues[1,:,0,0] + Ecosvalues3 = Esinvalues[2,:,0,0] + + usinvalues1 = usinvalues[0,:,0,0] + usinvalues2 = usinvalues[1,:,0,0] + usinvalues3 = usinvalues[2,:,0,0] + + ucosvalues1 = usinvalues[0,:,0,0] + ucosvalues2 = usinvalues[1,:,0,0] + ucosvalues3 = usinvalues[2,:,0,0] + + Esinvalues1_fft = np.fft.fftshift(np.fft.fft(Esinvalues1)) / cellsize + print("Evaluated FFT of Esin1") + Esinvalues2_fft = np.fft.fftshift(np.fft.fft(Esinvalues2)) / cellsize + print("Evaluated FFT of Esin2") + Esinvalues3_fft = np.fft.fftshift(np.fft.fft(Esinvalues3)) / cellsize + print("Evaluated FFT of Esin3") + + Ecosvalues1_fft = np.fft.fftshift(np.fft.fft(Ecosvalues1)) / cellsize + print("Evaluated FFT of Ecos1") + Ecosvalues2_fft = np.fft.fftshift(np.fft.fft(Ecosvalues2)) / cellsize + print("Evaluated FFT of Ecos2") + Ecosvalues3_fft = np.fft.fftshift(np.fft.fft(Ecosvalues3)) / cellsize + print("Evaluated FFT of Ecos3") + + usinvalues1_fft = np.fft.fftshift(np.fft.fft(usinvalues1)) / cellsize + print("Evaluated FFT of usin1") + usinvalues2_fft = np.fft.fftshift(np.fft.fft(usinvalues2)) / cellsize + print("Evaluated FFT of usin2") + usinvalues3_fft = np.fft.fftshift(np.fft.fft(usinvalues3)) / cellsize + print("Evaluated FFT of usin3") + + ucosvalues1_fft = np.fft.fftshift(np.fft.fft(ucosvalues1)) / cellsize + print("Evaluated FFT of ucos1") + ucosvalues2_fft = np.fft.fftshift(np.fft.fft(ucosvalues2)) / cellsize + print("Evaluated FFT of ucos2") + ucosvalues3_fft = np.fft.fftshift(np.fft.fft(ucosvalues3)) / cellsize + print("Evaluated FFT of ucos3") + + E_abs = np.sqrt(Esinvalues1 * np.conjugate(Esinvalues1) + Esinvalues2 * np.conjugate(Esinvalues2) + Esinvalues3 * np.conjugate(Esinvalues3) \ + + Ecosvalues1 * np.conjugate(Ecosvalues1) + Ecosvalues2 * np.conjugate(Ecosvalues2) + Ecosvalues3 * np.conjugate(Ecosvalues3)) + + E_abs_fft = cellsize * np.sqrt(Esinvalues1_fft * np.conjugate(Esinvalues1_fft) + Esinvalues2_fft * np.conjugate(Esinvalues2_fft) + Esinvalues3_fft * np.conjugate(Esinvalues3_fft) \ + + Ecosvalues1_fft * np.conjugate(Ecosvalues1_fft) + Ecosvalues2_fft * np.conjugate(Ecosvalues2_fft) + Ecosvalues3_fft * np.conjugate(Ecosvalues3_fft)) + print("Evaluated square modulus of FFT of E") + print(np.max(E_abs_fft)) + + u_abs = np.sqrt(usinvalues1 * np.conjugate(usinvalues1) + usinvalues2 * np.conjugate(usinvalues2) + usinvalues3 * np.conjugate(usinvalues3) \ + + ucosvalues1 * np.conjugate(ucosvalues1) + ucosvalues2 * np.conjugate(ucosvalues2) + ucosvalues3 * np.conjugate(ucosvalues3)) + + u_abs_fft = cellsize * np.sqrt(usinvalues1_fft * np.conjugate(usinvalues1_fft) + usinvalues2_fft * np.conjugate(usinvalues2_fft) + usinvalues3_fft * np.conjugate(usinvalues3_fft) \ + + ucosvalues1_fft * np.conjugate(ucosvalues1_fft) + ucosvalues2_fft * np.conjugate(ucosvalues2_fft) + ucosvalues3_fft * np.conjugate(ucosvalues3_fft)) + print("Evaluated square modulus of FFT of u") + print(np.max(u_abs)) + print(np.max(u_abs_fft)) + + # ax = plt.figure().add_subplot(projection='3d') + # ax.plot(k,k,E_abs) + + print(f"{d_omega=}") + print(f"{k_light/L=}") + # print(f"{k_sound/L=}") + + plt.figure(6) + plotyscale = np.max(u_abs_fft*V) + plotxscale = dk * Nel / L + plt.axhline(y=0.,color="black") + plt.axvline(x=k_light/L, ymin=0.05, linestyle='--',color="black") + plt.axvline(x=-k_light/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=k_sound/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=-k_sound/L, ymin=0.05, linestyle='--',color="black") + plt.plot(e_k/L,u_abs_fft*V,label="u") + plt.xlabel("Wave vectors [$m^{-1}$]") + plt.text(k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{light}}}}$") + plt.text(-k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{light}}}}$") + # plt.text(k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{thermic}}}}$") + # plt.text(-k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{thermic}}}}$") + plt.ylabel("Velocity field magnitude") + plt.title(f"Fourier Transform of u with $\\bar{{\\rho}}$={(rhobar*m_e*N*1000):.1f} $g cm^{{-3}}$, $\\bar{{\\theta}}$={(theta*T):.0f} K, \ + $k_{{\\mathrm{{light}}}}$={k_light/L:.3f} $m^{{-1}}$") + plt.legend() + # $k_{{\\mathrm{{thermic}}}}$={k_sound/L:.3f} $m^{{-1}}$ + + plt.figure(7) + plotyscale = np.max(E_abs_fft*E) + plotxscale = dk * Nel / L + plt.axhline(y=0.,color="black") + plt.axvline(x=k_light/L, ymin=0.05, linestyle='--',color="black") + plt.axvline(x=-k_light/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=k_sound/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=-k_sound/L, ymin=0.05, linestyle='--',color="black") + plt.plot(e_k/L,E_abs_fft*E,label="E") + plt.xlabel("Wave vectors [$m^{-1}$]") + plt.text(k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{light}}}}$") + plt.text(-k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{light}}}}$") + # plt.text(k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{thermic}}}}$") + # plt.text(-k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{thermic}}}}$") + plt.ylabel("Electric field strength") + plt.title(f"Fourier Transform of E with $\\bar{{\\rho}}$={(rhobar*m_e*N*1000):.1f} $g cm^{{-3}}$, $\\bar{{\\theta}}$={(theta*T):.0f} K, \ + $k_{{\\mathrm{{light}}}}$={k_light/L:.3f} $m^{{-1}}$") + plt.legend() + # $k_{{\\mathrm{{thermic}}}}$={k_sound/L:.3f} $m^{{-1}}$ + + plt.show() + + +def test_dispersion_relation_2d(): + ksquared = lambda k1,k2: ((k1**2)+(k2**2))*id3x3 + + ktensork = lambda k1,k2: np.array([[k1**2, k1*k2, 0], + [k2*k1, k2**2, 0], + [0, 0, 0]]) + + rotB = np.array([[0, -B_z, B_y], + [B_z, 0, -B_x], + [-B_y, B_x, 0]]) + + id3x3 = np.identity(3) + + def matrix(k1,k2): + mat = np.zeros((6,6), dtype="complex") + + mat[:3,:3] = id3x3 - theta / (mass * (omega**2)) * ktensork(k1,k2) # - 1j * omega_cyclo / omega * rotB + mat[:3,3:6] = 1j * omega_pe_normalized * id3x3 + mat[3:6,:3] = - 1j * omega_pe_normalized * id3x3 + mat[3:6,3:6] = (omega**2) * id3x3 - (c_normalized**2) * (ksquared(k1,k2) - ktensork(k1,k2)) + + return mat + + determinant = lambda k1, k2: np.linalg.det(matrix(k1,k2)) + + print(determinant(k_light,0)/(c_normalized**2)) + # print(determinant(k_sound,0)/(c_normalized**2)) + # print(determinant((k_light+k_sound)/2,0)) + + if d_omega == 0.: + kmax = 0.5 + else: + kmax = 1.6 * k_light # np.maximum(k_light, k_sound) + + Nel = 200 + + k = np.linspace(-kmax, kmax, Nel) + + k1, k2 = np.meshgrid(k, k, indexing='ij') + + det = np.zeros((Nel,Nel)) + + for i in range(Nel): + for j in range(Nel): + val = np.real(determinant(k1[i,j],k2[i,j])) + det[i,j] = 0. if np.abs(val) <= tol else val + + det /= c_normalized**2 + + detslice = (det[int(Nel/2),:]+det[int(Nel/2)+1,:])/2 + + # ax = plt.figure().add_subplot(projection='3d') + + # surface = ax.plot_surface(k1,k2,det,linewidth=0) + + # plt.contour(k1, k2, det, levels=[0.]) + + # plt.colorbar() + + plt.figure(1) + plotyscale = np.max(detslice) + plotxscale = 2*kmax/L + plt.title("Determinant of the system of oscillations") + plt.plot(k/L, detslice, label="Horizontal slice of determinant") + plt.axhline(y=0.,color="black") + plt.axvline(x=k_light/L, ymin=0.05, linestyle='--',color="black") + plt.axvline(x=-k_light/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=k_sound/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=-k_sound/L, ymin=0.05, linestyle='--',color="black") + plt.xlabel("Wave vectors [$m^{-1}$]") + plt.text(k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{light}}}}$") + plt.text(-k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{light}}}}$") + # plt.text(k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{therminc}}}}$") + # plt.text(-k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{thermic}}}}$") + plt.legend() + + + p = 3 + Ngrid = 2**6 + + Nfft = 64 + k_cutoff = kmax + dk = k_cutoff / Nfft + maxL: float = 2*np.pi / dk # so that our largest wavenumber value is included + + domain = domains.Cuboid(l1=-maxL/2 ,r1=maxL/2,l2=-maxL/2 ,r2=maxL/2) + equil = equils.HomogenSlab(B0x=B_x, B0y=B_y, B0z=B_z) + equil.domain = domain + + e = np.linspace(0., 1., Nel) + e_x, e_y, e_z = domain(e,e,0.) # the values the field will be sampled on that will correspond exactly to the k array after the FFT + e_x = e_x[:,:,0] + e_y = e_y[:,:,0] + print(e_x) + print(e_y) + cellsize = (maxL / Nel)**2 + e_k = np.linspace(-Nel/2 * dk, (Nel/2-1) * dk, Nel) + + j_physical = lambda x,y,z: k_cutoff/np.pi * np.sinc(k_cutoff/np.pi * x) * k_cutoff/np.pi * np.sinc(k_cutoff/np.pi * y) + # j_y = lambda x,y,z: np.exp(-((x-0.5)**2)/(0.1**2)) + # j_physical = lambda x,y,z: np.sinc(64/maxL *x) + print(f"{k_cutoff/np.pi *maxL=}") + zeroes = lambda x,y,z: 0. * (x+y+z) + + j_pulled_1 = lambda e1,e2,e3: domain.pull([j_physical,j_physical,j_physical],e1,e2,e3,kind="1", squeeze_out=False)[0] + j_pulled_2 = lambda e1,e2,e3: domain.pull([j_physical,j_physical,j_physical],e1,e2,e3,kind="1", squeeze_out=False)[1] + j_pulled_3 = lambda e1,e2,e3: domain.pull([j_physical,j_physical,j_physical],e1,e2,e3,kind="1", squeeze_out=False)[2] + + print(f"{np.shape(e)=}") + print(f"{np.shape(e_x)=}") + print(f"{np.shape(e_y)=}") + print(f"{np.shape(j_pulled_1(e,e,0.)[:,:,0])=}") + print(f"{np.shape(j_physical(e_x,e_y,0.))=}") + + jdiff1 = lambda e1,e2,e3: j_pulled_1(e1,e2,e3) / maxL - j_physical(domain(e1,e2,e3)[0],domain(e1,e2,e3)[1],domain(e1,e2,e3)[2]) + jdiff2 = lambda e1,e2,e3: j_pulled_2(e1,e2,e3) / maxL - j_physical(domain(e1,e2,e3)[0],domain(e1,e2,e3)[1],domain(e1,e2,e3)[2]) + jdiff3 = lambda e1,e2,e3: j_pulled_3(e1,e2,e3) / maxL - j_physical(domain(e1,e2,e3)[0],domain(e1,e2,e3)[1],domain(e1,e2,e3)[2]) + + plt.figure(2) + plt.subplot(1,3,1) + plt.pcolormesh(e,e, jdiff1(e,e,0.)[:,:,0],label="pulled j1") + plt.colorbar() + plt.subplot(1,3,2) + plt.pcolormesh(e,e, jdiff2(e,e,0.)[:,:,0],label="pulled j2") + plt.colorbar() + plt.subplot(1,3,3) + plt.pcolormesh(e,e, jdiff3(e,e,0.)[:,:,0],label="pulled j3") + plt.colorbar() + + plt.figure(3) + plt.subplot(2,3,1) + plt.pcolormesh(e,e, j_physical(e_x,e_y,0.),label="physical j1") + plt.colorbar() + plt.subplot(2,3,2) + plt.pcolormesh(e,e, j_physical(e_x,e_y,0.),label="physical j2") + plt.colorbar() + plt.subplot(2,3,3) + plt.pcolormesh(e,e, j_physical(e_x,e_y,0.),label="physical j3") + plt.colorbar() + plt.subplot(2,3,4) + plt.pcolormesh(e,e, j_pulled_1(e,e,0.)[:,:,0],label="physical j1") + plt.colorbar() + plt.subplot(2,3,5) + plt.pcolormesh(e,e, j_pulled_2(e,e,0.)[:,:,0],label="physical j2") + plt.colorbar() + plt.subplot(2,3,6) + plt.pcolormesh(e,e, j_pulled_3(e,e,0.)[:,:,0],label="physical j3") + plt.colorbar() + # plt.show() + # exit() + + degree = (p,p,1) + num_elements = (Ngrid,Ngrid,1) + bcs = (("dirichlet","dirichlet"), ("dirichlet","dirichlet"), None) + + grid = TensorProductGrid(num_elements=num_elements) + derham_opts = DerhamOptions(degree=degree, bcs=bcs) + derham = Derham(grid=grid, options=derham_opts, comm=comm) + projected_equil = ProjectedFluidEquilibriumWithB(equil=equil, derham=derham) + + mass_ops = WeightedMassOperators(derham=derham, domain=domain) + basis_ops = BasisProjectionOperators(derham=derham, domain=domain) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + Propagator.basis_ops = basis_ops + Propagator.projected_equil = projected_equil + + J = FEECVariable(space="Hcurl") + J.allocate(derham=derham, domain=domain) + J.spline.vector = derham.P1([j_pulled_1,j_pulled_2,j_pulled_3]) + + ee1, ee2, ee3 = np.meshgrid(np.linspace(0.,1.,4),np.linspace(0.,1.,5),np.linspace(0.,1.,6), indexing="ij") + # print(f"{np.shape(sincheck(ee1,ee2,ee3))=}") + print(f"{np.shape(j_pulled_1(ee1,ee2,ee3))=}") + print(f"{e=}") + + plt.figure(3) + # plt.plot(e, J0.spline(e,0.,0.)[:,0,0],label="projected j1") + # plt.plot(e, j_pulled_1(e,0.,0.)[:,0,0],'x',label="j1") + plt.subplot(2,3,1) + plt.pcolormesh(e,e, j_pulled_1(e,e,0.)[:,:,0],label="j1") + plt.colorbar() + plt.subplot(2,3,2) + plt.pcolormesh(e,e, j_pulled_2(e,e,0.)[:,:,0],label="j2") + plt.colorbar() + plt.subplot(2,3,3) + plt.pcolormesh(e,e, j_pulled_3(e,e,0.)[:,:,0],label="j3") + plt.colorbar() + plt.subplot(2,3,4) + plt.pcolormesh(e,e, J.spline(e,e,0.)[0][:,:,0],label="projected j1") + plt.colorbar() + plt.subplot(2,3,5) + plt.pcolormesh(e,e, J.spline(e,e,0.)[1][:,:,0],label="projected j2") + plt.colorbar() + plt.subplot(2,3,6) + plt.pcolormesh(e,e, J.spline(e,e,0.)[2][:,:,0],label="projected j3") + plt.colorbar() + # plt.show() + # exit() + + plt.figure(4) + # plt.title("Fourier transform of source term") + plt.subplot(3,1,1) + plt.contour(e_k/L, e_k/L, cellsize * np.fft.fftshift(np.abs(np.fft.fft2(domain.push(J.spline,e,e,0.,kind="1")[0][:,:,0]))),levels=50) + plt.colorbar() + plt.subplot(3,1,2) + plt.contour(e_k/L, e_k/L, cellsize * np.fft.fftshift(np.abs(np.fft.fft2(domain.push(J.spline,e,e,0.,kind="1")[1][:,:,0]))),levels=50) + plt.colorbar() + plt.subplot(3,1,3) + plt.contour(e_k/L, e_k/L, cellsize * np.fft.fftshift(np.abs(np.fft.fft2(domain.push(J.spline,e,e,0.,kind="1")[2][:,:,0]))),levels=50) + plt.colorbar() + # plt.show() + # exit() + + solver_params = SolverParameters( + tol=1e-16, + maxiter=3000, + info=True, + recycle=True, + ) + + _rhosin = FEECVariable(space="H1") + _rhosin.allocate(derham=derham, domain=domain) + + _rhocos = FEECVariable(space="H1") + _rhocos.allocate(derham=derham, domain=domain) + + _usin = FEECVariable(space="Hcurl") + _usin.allocate(derham=derham, domain=domain) + + _ucos = FEECVariable(space="Hcurl") + _ucos.allocate(derham=derham, domain=domain) + + _Esin = FEECVariable(space="Hcurl") + _Esin.allocate(derham=derham, domain=domain) + + _Ecos = FEECVariable(space="Hcurl") + _Ecos.allocate(derham=derham, domain=domain) + + _Bsin = FEECVariable(space="Hdiv") + _Bsin.allocate(derham=derham, domain=domain) + + _Bcos = FEECVariable(space="Hdiv") + _Bcos.allocate(derham=derham, domain=domain) + + solver = ColdPlasmaPerturbation() + solver.variables.rhosin = _rhosin + solver.variables.rhocos = _rhocos + solver.variables.usin = _usin + solver.variables.ucos = _ucos + solver.variables.Esin = _Esin + solver.variables.Ecos = _Ecos + solver.variables.Bsin = _Bsin + solver.variables.Bcos = _Bcos + + solver.options = solver.Options( + J=J, + omega=omega, + c0=c0, + c1=c1, + mass=mass, + mu=0., + nu=0., + rhobar=rhobar, + theta=theta, + Ebar=[zeroes, zeroes, zeroes], + solver="gmres", + solver_params=solver_params, + ) + + solver.allocate() + + dt=1.0 + print("Hi man") + solver(dt) + print("Bye man") + + Esinvalues = domain.push(_Esin.spline, e, e, 0., kind="1") + Ecosvalues = domain.push(_Ecos.spline, e, e, 0., kind="1") + + usinvalues = domain.push(_usin.spline, e, e, 0., kind="1") + ucosvalues = domain.push(_ucos.spline, e, e, 0., kind="1") + + rhosinvalues = domain.push(_rhosin.spline, e, e, 0., kind="0") + rhocosvalues = domain.push(_rhocos.spline, e, e, 0., kind="0") + + print(Esinvalues.shape) + print(Ecosvalues.shape) + + print(usinvalues.shape) + print(ucosvalues.shape) + + print(rhosinvalues.shape) + print(rhocosvalues.shape) + + Esinvalues1 = Esinvalues[0,:,:,0] + Esinvalues2 = Esinvalues[1,:,:,0] + Esinvalues3 = Esinvalues[2,:,:,0] + + Ecosvalues1 = Esinvalues[0,:,:,0] + Ecosvalues2 = Esinvalues[1,:,:,0] + Ecosvalues3 = Esinvalues[2,:,:,0] + + usinvalues1 = usinvalues[0,:,:,0] + usinvalues2 = usinvalues[1,:,:,0] + usinvalues3 = usinvalues[2,:,:,0] + + ucosvalues1 = usinvalues[0,:,:,0] + ucosvalues2 = usinvalues[1,:,:,0] + ucosvalues3 = usinvalues[2,:,:,0] + + Esinvalues1_fft = np.fft.fftshift(np.fft.fft2(Esinvalues1)) / cellsize + print("Evaluated FFT of Esin1") + Esinvalues2_fft = np.fft.fftshift(np.fft.fft2(Esinvalues2)) / cellsize + print("Evaluated FFT of Esin2") + Esinvalues3_fft = np.fft.fftshift(np.fft.fft2(Esinvalues3)) / cellsize + print("Evaluated FFT of Esin3") + + Ecosvalues1_fft = np.fft.fftshift(np.fft.fft2(Ecosvalues1)) / cellsize + print("Evaluated FFT of Ecos1") + Ecosvalues2_fft = np.fft.fftshift(np.fft.fft2(Ecosvalues2)) / cellsize + print("Evaluated FFT of Ecos2") + Ecosvalues3_fft = np.fft.fftshift(np.fft.fft2(Ecosvalues3)) / cellsize + print("Evaluated FFT of Ecos3") + + usinvalues1_fft = np.fft.fftshift(np.fft.fft2(usinvalues1)) / cellsize + print("Evaluated FFT of usin1") + usinvalues2_fft = np.fft.fftshift(np.fft.fft2(usinvalues2)) / cellsize + print("Evaluated FFT of usin2") + usinvalues3_fft = np.fft.fftshift(np.fft.fft2(usinvalues3)) / cellsize + print("Evaluated FFT of usin3") + + ucosvalues1_fft = np.fft.fftshift(np.fft.fft2(ucosvalues1)) / cellsize + print("Evaluated FFT of ucos1") + ucosvalues2_fft = np.fft.fftshift(np.fft.fft2(ucosvalues2)) / cellsize + print("Evaluated FFT of ucos2") + ucosvalues3_fft = np.fft.fftshift(np.fft.fft2(ucosvalues3)) / cellsize + print("Evaluated FFT of ucos3") + + E_abs = np.sqrt(Esinvalues1 * np.conjugate(Esinvalues1) + Esinvalues2 * np.conjugate(Esinvalues2) + Esinvalues3 * np.conjugate(Esinvalues3) \ + + Ecosvalues1 * np.conjugate(Ecosvalues1) + Ecosvalues2 * np.conjugate(Ecosvalues2) + Ecosvalues3 * np.conjugate(Ecosvalues3)) + + E_abs_fft = cellsize * np.sqrt(Esinvalues1_fft * np.conjugate(Esinvalues1_fft) + Esinvalues2_fft * np.conjugate(Esinvalues2_fft) + Esinvalues3_fft * np.conjugate(Esinvalues3_fft) \ + + Ecosvalues1_fft * np.conjugate(Ecosvalues1_fft) + Ecosvalues2_fft * np.conjugate(Ecosvalues2_fft) + Ecosvalues3_fft * np.conjugate(Ecosvalues3_fft)) + print("Evaluated square modulus of FFT of E") + print(np.max(E_abs_fft)) + + u_abs = np.sqrt(usinvalues1 * np.conjugate(usinvalues1) + usinvalues2 * np.conjugate(usinvalues2) + usinvalues3 * np.conjugate(usinvalues3) \ + + ucosvalues1 * np.conjugate(ucosvalues1) + ucosvalues2 * np.conjugate(ucosvalues2) + ucosvalues3 * np.conjugate(ucosvalues3)) + + u_abs_fft = cellsize * np.sqrt(usinvalues1_fft * np.conjugate(usinvalues1_fft) + usinvalues2_fft * np.conjugate(usinvalues2_fft) + usinvalues3_fft * np.conjugate(usinvalues3_fft) \ + + ucosvalues1_fft * np.conjugate(ucosvalues1_fft) + ucosvalues2_fft * np.conjugate(ucosvalues2_fft) + ucosvalues3_fft * np.conjugate(ucosvalues3_fft)) + print("Evaluated square modulus of FFT of u") + print(np.max(u_abs)) + print(np.max(u_abs_fft)) + + # ax = plt.figure().add_subplot(projection='3d') + # ax.plot(k,k,E_abs) + + print(f"{d_omega=}") + print(f"{k_light/L=}") + # print(f"{k_sound/L=}") + + plt.figure(5) + plotyscale = np.max(u_abs_fft*V) + plotxscale = dk * Nel / L + plt.axhline(y=0.,color="black") + plt.axvline(x=k_light/L, ymin=0.05, linestyle='--',color="black") + plt.axvline(x=-k_light/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=k_sound/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=-k_sound/L, ymin=0.05, linestyle='--',color="black") + plt.contour(e_k/L,e_k/L,u_abs_fft*V,levels=100,label="u") + plt.xlabel("Wave vectors [$m^{-1}$]") + plt.text(k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{light}}}}$") + plt.text(-k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{light}}}}$") + # plt.text(k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{therminc}}}}$") + # plt.text(-k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{thermic}}}}$") + plt.ylabel("Velocity field magnitude") + plt.title(f"Fourier Transform of u with $\\bar{{\\rho}}$={(rhobar*m_e*N*1000):.1f} $g cm^{{-3}}$, $\\bar{{\\theta}}$={(theta*T):.0f} K, \ + $k_{{\\mathrm{{light}}}}$={k_light/L:.3f} $m^{{-1}}$") + plt.legend() + plt.colorbar() + # $k_{{\\mathrm{{thermic}}}}$={k_sound/L:.3f} $m^{{-1}}$ + + plt.figure(6) + plotyscale = np.max(E_abs_fft*E) + plotxscale = dk * Nel / L + plt.axhline(y=0.,color="black") + plt.axvline(x=k_light/L, ymin=0.05, linestyle='--',color="black") + plt.axvline(x=-k_light/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=k_sound/L, ymin=0.05, linestyle='--',color="black") + # plt.axvline(x=-k_sound/L, ymin=0.05, linestyle='--',color="black") + plt.contour(e_k/L,e_k/L,E_abs_fft*E,levels=100,label="E") + plt.xlabel("Wave vectors [$m^{-1}$]") + plt.text(k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{light}}}}$") + plt.text(-k_light/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{light}}}}$") + # plt.text(k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"+$k_{{\\mathrm{{therminc}}}}$") + # plt.text(-k_sound/L - 0.02*plotxscale,-0.035*plotyscale,f"-$k_{{\\mathrm{{thermic}}}}$") + plt.ylabel("Electric field strength") + plt.title(f"Fourier Transform of E with $\\bar{{\\rho}}$={(rhobar*m_e*N*1000):.1f} $g cm^{{-3}}$, $\\bar{{\\theta}}$={(theta*T):.0f} K, \ + $k_{{\\mathrm{{light}}}}$={k_light/L:.3f} $m^{{-1}}$") + plt.legend() + plt.colorbar() + # $k_{{\\mathrm{{thermic}}}}$={k_sound/L:.3f} $m^{{-1}}$ + + plt.show() + + + + +def test_convergence_1d( + show_plot: bool = False, +): + """Test of the solver on 1d problem by means of manufactured solution""" + + domain: Domain = domains.Cuboid() + equil = equils.HomogenSlab(B0x=B_x, B0y=B_y, B0z=B_z) + equil.domain = domain + + bcs = (None, None, None) + + pmax = 3 + Nmin = 4 + Nmax = 7 + + J0: float = 1.5 # times eNV + + denom_light: float = 1. / (omega**2 - rhobar / (alpha * (mass**2)) - (c_normalized**2) * 16*(xp.pi**2)) + denom_sound: float = 1. / (omega**2 - rhobar / (alpha * (mass**2)) - (theta / mass) * 36*(xp.pi**2)) + + E0: float = omega * J0 * denom_light / alpha + u0: float = J0 * denom_sound / (mass * alpha) + + zeroes = lambda x,y,z: 0.*(x+y+z) + + j_exact_x = lambda x,y,z: J0 * xp.cos(6*xp.pi*x) + j_exact_y = lambda x,y,z: J0 * xp.sin(4*xp.pi*x) + + E_exact_x = lambda x,y,z: - (1 + rhobar / (alpha * (mass**2)) * denom_sound) * J0 / (omega * alpha) * xp.cos(6*xp.pi*x) + E_exact_y = lambda x,y,z: - E0 * xp.sin(4*xp.pi*x) + + u_exact_x = lambda x,y,z: - u0 * xp.cos(6*xp.pi*x) + u_exact_y = lambda x,y,z: - E0 / (mass * omega) * xp.sin(4*xp.pi*x) + + rho_exact = lambda x,y,z: - 6*xp.pi * rhobar * u0 / omega * xp.sin(6*xp.pi*x) + + B_exact_z = lambda x,y,z: - 4*xp.pi / omega * E0 * xp.cos(4*xp.pi*x) + + + # Test over spline degree and grid resolution + Nels = [2**n for n in range(Nmin, Nmax + 1)] + + e1 = xp.linspace(0.0, 1.0, 64) + e2 = 0.0 + e3 = 0.0 + + ee1, ee2, ee3 = xp.meshgrid(e1, e2, e3, indexing="ij") + + for p in range(2, pmax + 1): + errors_Esin = [] + errors_Ecos = [] + errors_usin = [] + errors_ucos = [] + errors_rhosin = [] + errors_rhocos = [] + errors_Bsin = [] + errors_Bcos = [] + errors = [] + h_vec = [] + + for n, Nel in enumerate(Nels): + + degree = (p, 1, 1) + num_elements = (Nel, 1, 1) + + grid = TensorProductGrid(num_elements=num_elements) + derham_opts = DerhamOptions(degree=degree, bcs=bcs) + derham = Derham(grid=grid, options=derham_opts, comm=comm) + projected_equil = ProjectedFluidEquilibriumWithB(equil=equil, derham=derham) + + mass_ops = WeightedMassOperators(derham=derham, domain=domain) + basis_ops = BasisProjectionOperators(derham=derham, domain=domain) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + Propagator.basis_ops = basis_ops + Propagator.projected_equil = projected_equil + + J = FEECVariable(space="Hcurl") + J.allocate(derham=derham, domain=domain) + J.spline.vector = derham.P1([j_exact_x, j_exact_y, zeroes]) + + solver_params = SolverParameters( + tol=1e-16, + maxiter=3000, + info=True, + recycle=False, + ) + + _Esin = FEECVariable(space="Hcurl") + _Esin.allocate(derham=derham, domain=domain) + _Ecos = FEECVariable(space="Hcurl") + _Ecos.allocate(derham=derham, domain=domain) + _usin = FEECVariable(space="Hcurl") + _usin.allocate(derham=derham, domain=domain) + _ucos = FEECVariable(space="Hcurl") + _ucos.allocate(derham=derham, domain=domain) + _rhosin = FEECVariable(space="H1") + _rhosin.allocate(derham=derham, domain=domain) + _rhocos = FEECVariable(space="H1") + _rhocos.allocate(derham=derham, domain=domain) + _Bsin = FEECVariable(space="Hdiv") + _Bsin.allocate(derham=derham, domain=domain) + _Bcos = FEECVariable(space="Hdiv") + _Bcos.allocate(derham=derham, domain=domain) + + solver = ColdPlasmaPerturbation() + solver.variables.Esin = _Esin + solver.variables.Ecos = _Ecos + solver.variables.usin = _usin + solver.variables.ucos = _ucos + solver.variables.rhosin = _rhosin + solver.variables.rhocos = _rhocos + solver.variables.Bsin = _Bsin + solver.variables.Bcos = _Bcos + + solver.options = solver.Options( + J=J, + omega=omega, + c0=c0, + c1=c1, + mass=mass, + mu=0., + nu=0., + rhobar=rhobar, + theta=theta, + Ebar=[zeroes, zeroes, zeroes], + solver="gmres", + solver_params=solver_params, + ) + + solver.allocate() + + dt = 1.0 + solver(dt) + + Esin_calculated = xp.array(_Esin.spline(ee1, ee2, ee3)) + logger.info(f"{Esin_calculated.shape = }") + Esin_analytical = xp.array([E_exact_x(ee1, ee2, ee3), E_exact_y(ee1, ee2, ee3), zeroes(ee1, ee2, ee3)]) + logger.info(f"{Esin_analytical.shape = }") + + Ecos_calculated = xp.array(_Ecos.spline(ee1, ee2, ee3)) + logger.info(f"{Ecos_calculated.shape = }") + Ecos_analytical = xp.array([zeroes(ee1, ee2, ee3), zeroes(ee1, ee2, ee3), zeroes(ee1, ee2, ee3)]) + logger.info(f"{Ecos_analytical.shape = }") + + usin_calculated = xp.array(_usin.spline(ee1, ee2, ee3)) + logger.info(f"{usin_calculated.shape = }") + usin_analytical = xp.array([zeroes(ee1, ee2, ee3), zeroes(ee1, ee2, ee3), zeroes(ee1, ee2, ee3)]) + logger.info(f"{usin_analytical.shape = }") + + ucos_calculated = xp.array(_ucos.spline(ee1, ee2, ee3)) + logger.info(f"{ucos_calculated.shape = }") + print(_ucos.spline(1/xp.e,0,0)) + ucos_analytical = xp.array([u_exact_x(ee1, ee2, ee3), u_exact_y(ee1, ee2, ee3), zeroes(ee1, ee2, ee3)]) + logger.info(f"{ucos_analytical.shape = }") + + Bsin_calculated = xp.array(_Bsin.spline(ee1, ee2, ee3)) + logger.info(f"{Bsin_calculated.shape = }") + Bsin_analytical = xp.array([zeroes(ee1, ee2, ee3), zeroes(ee1, ee2, ee3), zeroes(ee1, ee2, ee3)]) + logger.info(f"{Bsin_analytical.shape = }") + + Bcos_calculated = xp.array(_Bcos.spline(ee1, ee2, ee3)) + logger.info(f"{Bcos_calculated.shape = }") + Bcos_analytical = xp.array([zeroes(ee1, ee2, ee3), zeroes(ee1, ee2, ee3), B_exact_z(ee1, ee2, ee3)]) + logger.info(f"{Bcos_analytical.shape = }") + + rhosin_calculated = xp.array(_rhosin.spline(ee1, ee2, ee3)) + logger.info(f"{rhosin_calculated.shape = }") + rhosin_analytical = rho_exact(ee1, ee2, ee3) + logger.info(f"{rhosin_analytical.shape = }") + + rhocos_calculated = xp.array(_rhocos.spline(ee1, ee2, ee3)) + logger.info(f"{rhocos_calculated.shape = }") + rhocos_analytical = zeroes(ee1, ee2, ee3) + logger.info(f"{rhocos_analytical.shape = }") + + if show_plot: + plt.figure(f"Esin[0] error for degree {p =}, analytical amplitude {(1 + rhobar / (alpha * (mass**2)) * denom_sound) * J0 / (omega * alpha)}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Esin_calculated[0][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Esin_analytical[0][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Esin[1] error for degree {p =}, analytical amplitude {E0}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Esin_calculated[1][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Esin_analytical[1][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Esin[2] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Esin_calculated[2][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Esin_analytical[2][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Ecos[0] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Ecos_calculated[0][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Ecos_analytical[0][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Ecos[1] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Ecos_calculated[1][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Ecos_analytical[1][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Ecos[2] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Ecos_calculated[2][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Ecos_analytical[2][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"usin[0] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, usin_calculated[0][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, usin_analytical[0][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"usin[1] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, usin_calculated[1][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, usin_analytical[1][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"usin[2] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, usin_calculated[2][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, usin_analytical[2][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"ucos[0] error for degree {p =}, analytical amplitude {u0}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, ucos_calculated[0][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, ucos_analytical[0][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"ucos[1] error for degree {p =}, analytical amplitude {E0 / (mass * omega)}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, ucos_calculated[1][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, ucos_analytical[1][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"ucos[2] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, ucos_calculated[2][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, ucos_analytical[2][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"rhosin error for degree {p =}, analytical amplitude {6*xp.pi * rhobar * u0 / omega}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, rhosin_calculated[:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, rhosin_analytical[:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"rhocos error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, rhocos_calculated[:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, rhocos_analytical[:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Bsin[0] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Bsin_calculated[0][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Bsin_analytical[0][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Bsin[1] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Bsin_calculated[1][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Bsin_analytical[1][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Bsin[2] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Bsin_calculated[2][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Bsin_analytical[2][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Bcos[0] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Bcos_calculated[0][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Bcos_analytical[0][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Bcos[1] error for degree {p =}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Bcos_calculated[1][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Bcos_analytical[1][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + plt.figure(f"Bcos[2] error for degree {p =}, analytical amplitude {4*xp.pi / omega * E0}", figsize=(12, 8)) + plt.subplot(2, int((Nmax - Nmin) / 2 + 1), n + 1) + plt.plot(e1, Bcos_calculated[2][:, 0, 0], "o", label=f"{Nel}, numerical") + plt.plot(e1, Bcos_analytical[2][:, 0, 0], "k--", label=f"{Nel}, analytical") + plt.legend() + + # if n == 0: + # plt.show() + # exit() + + error_Esin = xp.max(xp.abs(Esin_calculated - Esin_analytical)) + errors_Esin.append(error_Esin) + error_Ecos = xp.max(xp.abs(Ecos_calculated - Ecos_analytical)) + errors_Ecos.append(error_Ecos) + + error_usin = xp.max(xp.abs(usin_calculated - usin_analytical)) + errors_usin.append(error_usin) + error_ucos = xp.max(xp.abs(ucos_calculated - ucos_analytical)) + errors_ucos.append(error_ucos) + + error_rhosin = xp.max(xp.abs(rhosin_calculated - rhosin_analytical)) + errors_rhosin.append(error_rhosin) + error_rhocos = xp.max(xp.abs(rhocos_calculated - rhocos_analytical)) + errors_rhocos.append(error_rhocos) + + error_Bsin = xp.max(xp.abs(Bsin_calculated - Bsin_analytical)) + errors_Bsin.append(error_Bsin) + error_Bcos = xp.max(xp.abs(Bcos_calculated - Bcos_analytical)) + errors_Bcos.append(error_Ecos) + + error = xp.max([error_Esin,error_Ecos,error_usin,error_ucos,error_rhosin,error_rhocos,error_Bsin,error_Bcos]) + errors.append(error) + + h = 1 / Nel + h_vec.append(h) + + m, _ = xp.polyfit(xp.log(Nels), xp.log(errors), deg=1) + logger.info(f"For {p =}, solution converges with rate {-m =} ") + + if show_plot: + plt.figure(f"Esin Convergence for degree {p =}", figsize=(12, 8)) + plt.title(f"Esin Convergence rate for degree {p =}") + plt.plot(h_vec, errors_Esin, "o", label=f"Calculated Esin error, {m =}") + plt.plot( + h_vec, + [(h ** (p)) / (h_vec[0] ** (p)) * errors[0] for h in h_vec], + "k--", + label="Theoretical Esin error, rate = p", + ) + plt.xscale("log") + plt.yscale("log") + plt.xlabel("Grid spacing h") + plt.ylabel("Error") + plt.legend() + + plt.figure(f"Ecos Convergence for degree {p =}", figsize=(12, 8)) + plt.title(f"Ecos Convergence rate for degree {p =}") + plt.plot(h_vec, errors_Ecos, "o", label=f"Calculated Ecos error, {m =}") + plt.plot( + h_vec, + [(h ** (p)) / (h_vec[0] ** (p)) * errors[0] for h in h_vec], + "k--", + label="Theoretical Ecos error, rate = p", + ) + plt.xscale("log") + plt.yscale("log") + plt.xlabel("Grid spacing h") + plt.ylabel("Error") + plt.legend() + + plt.figure(f"usin Convergence for degree {p =}", figsize=(12, 8)) + plt.title(f"usin Convergence rate for degree {p =}") + plt.plot(h_vec, errors_usin, "o", label=f"Calculated usin error, {m =}") + plt.plot( + h_vec, + [(h ** (p)) / (h_vec[0] ** (p)) * errors[0] for h in h_vec], + "k--", + label="Theoretical usin error, rate = p", + ) + plt.xscale("log") + plt.yscale("log") + plt.xlabel("Grid spacing h") + plt.ylabel("Error") + plt.legend() + + plt.figure(f"ucos Convergence for degree {p =}", figsize=(12, 8)) + plt.title(f"ucos Convergence rate for degree {p =}") + plt.plot(h_vec, errors_Esin, "o", label=f"Calculated ucos error, {m =}") + plt.plot( + h_vec, + [(h ** (p)) / (h_vec[0] ** (p)) * errors[0] for h in h_vec], + "k--", + label="Theoretical ucos error, rate = p", + ) + plt.xscale("log") + plt.yscale("log") + plt.xlabel("Grid spacing h") + plt.ylabel("Error") + plt.legend() + + plt.figure(f"rhosin Convergence for degree {p =}", figsize=(12, 8)) + plt.title(f"rhosin Convergence rate for degree {p =}") + plt.plot(h_vec, errors_rhosin, "o", label=f"Calculated rhosin error, {m =}") + plt.plot( + h_vec, + [(h ** (p + 1)) / (h_vec[0] ** (p + 1)) * errors[0] for h in h_vec], + "k--", + label="Theoretical rhosin error, rate = p + 1", + ) + plt.xscale("log") + plt.yscale("log") + plt.xlabel("Grid spacing h") + plt.ylabel("Error") + plt.legend() + + plt.figure(f"rhocos Convergence for degree {p =}", figsize=(12, 8)) + plt.title(f"rhocos Convergence rate for degree {p =}") + plt.plot(h_vec, errors_rhocos, "o", label=f"Calculated rhocos error, {m =}") + plt.plot( + h_vec, + [(h ** (p + 1)) / (h_vec[0] ** (p + 1)) * errors[0] for h in h_vec], + "k--", + label="Theoretical rhocos error, rate = p + 1", + ) + plt.xscale("log") + plt.yscale("log") + plt.xlabel("Grid spacing h") + plt.ylabel("Error") + plt.legend() + + plt.figure(f"Bsin Convergence for degree {p =}", figsize=(12, 8)) + plt.title(f"Bsin Convergence rate for degree {p =}") + plt.plot(h_vec, errors_Bsin, "o", label=f"Calculated Bsin error, {m =}") + plt.plot( + h_vec, + [(h ** (p - 1)) / (h_vec[0] ** (p - 1)) * errors[0] for h in h_vec], + "k--", + label="Theoretical Esin error, rate = p - 1", + ) + plt.xscale("log") + plt.yscale("log") + plt.xlabel("Grid spacing h") + plt.ylabel("Error") + plt.legend() + + plt.figure(f"Bcos Convergence for degree {p =}", figsize=(12, 8)) + plt.title(f"Bcos Convergence rate for degree {p =}") + plt.plot(h_vec, errors_Bcos, "o", label=f"Calculated Bcos error, {m =}") + plt.plot( + h_vec, + [(h ** (p - 1)) / (h_vec[0] ** (p - 1)) * errors[0] for h in h_vec], + "k--", + label="Theoretical Bcos error, rate = p - 1", + ) + plt.xscale("log") + plt.yscale("log") + plt.xlabel("Grid spacing h") + plt.ylabel("Error") + plt.legend() + + plt.figure(f"Convergence for degree {p =}", figsize=(12, 8)) + plt.title(f"Convergence rate for degree {p =}") + plt.plot(h_vec, errors, "o", label=f"Calculated error, {m =}") + plt.plot( + h_vec, + [(h ** (p)) / (h_vec[0] ** (p)) * errors[0] for h in h_vec], + "k--", + label="Theoretical error, rate = p", + ) + plt.xscale("log") + plt.yscale("log") + plt.xlabel("Grid spacing h") + plt.ylabel("Error") + plt.legend() + + plt.figure("Difference between Esin and ucos") + plt.plot(e1, ucos_calculated[1][:,0,0]-Esin_calculated[1][:, 0, 0]/(mass*omega), '.') + + plt.show() + + tolerance: float = 0.07 + assert -m > (p - 1 - tolerance) + + +if __name__ == "__main__": + test_dispersion_relation_1d() + # test_convergence_1d(show_plot=True) \ No newline at end of file