diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index c7df6b29f..fa3813c2b 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -12,10 +12,11 @@ from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilDiagonalMatrix, StencilMatrix, StencilVector +from struphy import domains, equils from struphy.feec import mass_kernels from struphy.feec.linear_operators import BoundaryOperator, LinOpWithTransp from struphy.feec.psydac_derham import Derham, SplineFunction -from struphy.feec.utilities import LocalRotationMatrix, get_quad_grids +from struphy.feec.utilities import LocalProjectionMatrix, LocalRotationMatrix, get_quad_grids from struphy.fields_background.base import MHDequilibrium from struphy.fields_background.equils import set_defaults from struphy.geometry.base import Domain @@ -59,6 +60,11 @@ def __init__( self._matrix_free = matrix_free self._eq_mhd = eq_mhd + if self._eq_mhd is None: + self._eq_mhd = equils.HomogenSlab() + if not hasattr(self.eq_mhd, "_domain"): + self._eq_mhd.domain = self._domain + # only for M1 Mac users PSYDAC_BACKEND_GPYCCEL["flags"] = "-O3 -march=native -mtune=native -ffast-math -ffree-line-length-none" @@ -338,7 +344,7 @@ def M1ninv(self): weights=( "Ginv", "sqrt_g", - "1/eq_n0", + lambda *etas: 1 / self.eq_mhd.n0(*etas), ), name="M1ninv", assemble=True, @@ -693,7 +699,7 @@ def M1Bninv(self): rot_B, "Ginv", "sqrt_g", - "1/eq_n0", + lambda *etas: 1 / self.eq_mhd.n0(*etas), ), name="M1Bninv", assemble=True, @@ -701,6 +707,31 @@ def M1Bninv(self): return self._M1Bninv + @auto_convert_docstring + @property + def M1para(self): + r""" + Mass matrix + + .. math:: + + \mathbb M^{1,\parallel}_{(\mu,ijk), (\nu,mno)} = \int \vec{\Lambda}^1_{\mu,ijk} b_0 b_0^\top \vec{\Lambda}^1_{\nu,mno} \sqrt{g} \textnormal{d}\boldsymbol{\eta}. + """ + if not hasattr(self, "_M1para"): + bb = LocalProjectionMatrix(self.eq_mhd.unit_bv_1, self.eq_mhd.unit_bv_2, self.eq_mhd.unit_bv_3) + + self._M1para = self.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + bb, + "sqrt_g", + ), + name="M1para", + assemble=True, + ) + return self._M1para + @auto_convert_docstring @property def M1perp(self): @@ -709,25 +740,79 @@ def M1perp(self): .. math:: - \mathbb M^{1,\perp}_{(\mu,ijk), (\nu,mno)} = \int \vec{\Lambda}^1_{\mu,ijk} DF^{-1} \begin{pmatrix} 1 & 0 & 0 \\ 0 & 1 & 0 \\ 0 & 0 & 0 \end{pmatrix} DF^{-\top} \vec{\Lambda}^1_{\nu,mno} \sqrt{g} \textnormal{d}\boldsymbol{\eta}. + \mathbb M^{1,\perp}_{(\mu,ijk), (\nu,mno)} = \int \vec{\Lambda}^1_{\mu,ijk} \left(G^{-1} - b_0 b_0^\top \right) \vec{\Lambda}^1_{\nu,mno} \sqrt{g} \textnormal{d}\boldsymbol{\eta}. """ if not hasattr(self, "_M1perp"): - D = [[1, 0, 0], [0, 1, 0], [0, 0, 0]] + self._M1perp = self.M1.copy() + self._M1perp -= self.M1para + return self._M1perp - self._M1perp = self.create_weighted_mass( + @auto_convert_docstring + @property + def M1para_MHDeq(self): + r""" + Mass matrix + + .. math:: + + \mathbb M^{1,\parallel}_{(\mu,ijk), (\nu,mno)} = \int \frac{n^0_{\textnormal{eq}}(\boldsymbol{\eta})}{\|B_0(\boldsymbol{\eta})\|^2} \vec{\Lambda}^1_{\mu,ijk} b_0 b_0^\top \vec{\Lambda}^1_{\nu,mno} \sqrt{g} \textnormal{d}\boldsymbol{\eta}. + """ + if not hasattr(self, "_M1para_MHDeq"): + bb = LocalProjectionMatrix(self.eq_mhd.unit_bv_1, self.eq_mhd.unit_bv_2, self.eq_mhd.unit_bv_3) + + self._M1para_MHDeq = self.create_weighted_mass( "Hcurl", "Hcurl", weights=( - "DFinv", - D, - "DFinv", + bb, + lambda *etas: 1 / self.eq_mhd.absB0(*etas) ** 2, + "eq_n0", "sqrt_g", ), - name="M1perp", + name="M1para_MHDeq", assemble=True, ) + return self._M1para_MHDeq - return self._M1perp + @auto_convert_docstring + @property + def M1_MHDeq(self): + r""" + Mass matrix + + .. math:: + + \mathbb M^{1}_{(\mu,ijk), (\nu,mno)} = \int \frac{n^0_{\textnormal{eq}}(\boldsymbol{\eta})}{\|B_0(\boldsymbol{\eta})\|^2} \vec{\Lambda}^1_{\mu,ijk} G^{-1} \vec{\Lambda}^1_{\nu,mno} \sqrt{g} \textnormal{d}\boldsymbol{\eta}. + """ + if not hasattr(self, "_M1_MHDeq"): + self._M1_MHDeq = self.create_weighted_mass( + "Hcurl", + "Hcurl", + weights=( + "Ginv", + lambda *etas: 1 / self.eq_mhd.absB0(*etas) ** 2, + "eq_n0", + "sqrt_g", + ), + name="M1_MHDeq", + assemble=True, + ) + return self._M1_MHDeq + + @auto_convert_docstring + @property + def M1gyro(self): + r""" + Mass matrix + + .. math:: + + \mathbb M^{1,\perp}_{(\mu,ijk), (\nu,mno)} = \int \frac{n^0_{\textnormal{eq}}(\boldsymbol{\eta})}{\|B_0(\boldsymbol{\eta})\|^2} \vec{\Lambda}^1_{\mu,ijk} \left(G^{-1} - b_0 b_0^\top \right) \vec{\Lambda}^1_{\nu,mno} \sqrt{g} \textnormal{d}\boldsymbol{\eta}. + """ + if not hasattr(self, "_M1gyro"): + self._M1gyro = self.M1_MHDeq.copy() + self._M1gyro -= self.M1para_MHDeq + return self._M1gyro @auto_convert_docstring @property @@ -746,7 +831,10 @@ def M0ad(self): self._M0ad = self.create_weighted_mass( "H1", "H1", - weights=("eq_n0", "sqrt_g"), + weights=( + "eq_n0", + "sqrt_g", + ), name="M0ad", assemble=True, ) @@ -755,39 +843,30 @@ def M0ad(self): @auto_convert_docstring @property - def M1gyro(self): + def M0ad_withT(self): r""" Mass matrix .. math:: - \mathbb M^{1,n}_{(\mu,ijk), (\nu,mno)} = \int n^0_{\textnormal{eq}}(\boldsymbol{\eta}) \Lambda^1_{\mu,ijk} G^{-1}_{\mu,\nu} \Lambda^1_{\nu,mno} \sqrt{g} \textnormal{d}\boldsymbol{\eta}, + \mathbb M^0_{ijk, mno} = \int \frac{n^0_{\textnormal{eq}}(\boldsymbol{\eta})}{T^0_{\textnormal{eq}}(\boldsymbol{\eta})} \Lambda^0_{ijk} \Lambda^0_{mno} \sqrt{g} \textnormal{d}\boldsymbol{\eta}. - where :math:`n^0_{\textnormal{eq}}(\boldsymbol{\eta})` is an MHD equilibrium density (0-form). + where :math:`n^0_{\textnormal{eq}}(\boldsymbol{\eta})` and :math:`T^0_{\textnormal{eq}}(\boldsymbol{\eta})` are MHD equilibrium density and electron temperature (0-forms), respectively. """ - - if not hasattr(self, "_M1gyro"): - D = [[1, 0, 0], [0, 1, 0], [0, 0, 0]] - - self._M1gyro = self.create_weighted_mass( - "Hcurl", - "Hcurl", + if not hasattr(self, "_M0ad_withT"): + self._M0ad_withT = self.create_weighted_mass( + "H1", + "H1", weights=( "eq_n0", - "1/eq_absB0", - "1/eq_absB0", - D, - "Ginv", - D, + lambda *etas: 1 / self.eq_mhd.t0(*etas), "sqrt_g", ), - name="M1gyro", + name="M0ad_withT", assemble=True, ) - # 1/eq_absB0**2 written twice instead of square - - return self._M1gyro + return self._M0ad_withT @property def WMM(self): @@ -917,18 +996,8 @@ def create_weighted_mass( for n, f in enumerate(weights): if isinstance(f, str): # determine the callable - if "/" in f: - f_components = f.split("/") - if f_components[-1] == "sqrt_g": - f_call = lambda e1, e2, e3: 1.0 / abs(self.domain.jacobian_det(e1, e2, e3)) - elif f_components[-1] == "eq_n0": - f_call = lambda e1, e2, e3: 1.0 / self.eq_mhd.n0(e1, e2, e3) - elif f_components[-1] == "eq_absB0": - f_call = lambda e1, e2, e3: 1.0 / self.eq_mhd.absB0(e1, e2, e3) - else: - raise NotImplementedError( - f"The option {f} is not available for division ('/') yet.", - ) + if f == "1/sqrt_g": + f_call = lambda e1, e2, e3: 1.0 / abs(self.domain.jacobian_det(e1, e2, e3)) elif "eq_" in f: f_components = f.split("q_") f_call = getattr(self.eq_mhd, f_components[-1]) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index b6201f8e7..be7938457 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -149,6 +149,14 @@ def fun(e): ): # this branch is only entered if comm exists (and thus subcomm has been initialized) if subcomm != MPI.COMM_NULL: subcomm.Allgather(local_fun, fun) + """gathered = subcomm.gather(local_fun, root=selected_ranks[0]) + if rank == selected_ranks[0]: + if gathered is None: + raise RuntimeError("MPI gather failed to return data on root rank") + fun[:] = xp.concatenate(gathered) + assert fun.size == npts, ( + f"Gathered weight size {fun.size} does not match expected {npts}" + )""" comm.Bcast(fun, root=selected_ranks[0]) else: fun[:] = local_fun diff --git a/src/struphy/feec/tests/test_mass_matrices.py b/src/struphy/feec/tests/test_mass_matrices.py index 82b87d097..60827e9a6 100644 --- a/src/struphy/feec/tests/test_mass_matrices.py +++ b/src/struphy/feec/tests/test_mass_matrices.py @@ -22,7 +22,7 @@ def test_mass(num_elements, degree, bcs, map_and_equil, matrix_free, show_plots=False): """Test weighted mass matrices by recovering projected functions from the DeRham complex. - For each mass operator in ``{M0, M1, M2, M3, Mv, M1n, M2n, Mvn, M1ninv, M0ad}``, + For each mass operator in ``{M0, M1, M2, M3, Mv, M1n, M2n, Mvn, M1ninv, M0ad, M0ad_withT}``, the test: 1. Projects known trigonometric right-hand-side functions onto the @@ -32,10 +32,12 @@ def test_mass(num_elements, degree, bcs, map_and_equil, matrix_free, show_plots= it point-wise to the exact function. The density-weighted operators (``M1n``, ``M2n``, ``Mvn``, ``M0ad``) are - tested against ``exact / n0``, and the inverse-density operator + tested against ``exact / n0``, ``M0ad_withT`` is tested against ``exact * t0 / n0``, and the inverse-density operator (``M1ninv``) is tested against ``exact * n0``. """ + from types import MethodType + import cunumpy as xp from feectools.ddm.mpi import mpi as MPI from feectools.linalg.solvers import inverse @@ -82,7 +84,7 @@ def test_mass(num_elements, degree, bcs, map_and_equil, matrix_free, show_plots= # derham object grid = TensorProductGrid(num_elements=num_elements) derham_opts = DerhamOptions(degree=degree, bcs=bcs) - derham = Derham(grid, derham_opts, comm=mpi_comm) + derham = Derham(grid, derham_opts, comm=mpi_comm, domain=domain) logger.debug(f"Rank {mpi_rank} | Local domain : " + str(derham.domain_array[mpi_rank])) @@ -111,6 +113,7 @@ def rhs_2(e1, e2, e3): rhs = {} rhs["M0"] = l2proj_0.get_dofs(rhs_0, apply_bc=True) rhs["M0ad"] = rhs["M0"] + rhs["M0ad_withT"] = rhs["M0"] rhs["M1"] = l2proj_1.get_dofs((rhs_0, rhs_1, rhs_2), apply_bc=True) rhs["M1n"] = rhs["M1"] rhs["M1ninv"] = rhs["M1"] @@ -134,7 +137,7 @@ def rhs_2(e1, e2, e3): elif min(degree) == 2: err_bound = 2.6e-2 - names = ["M0", "M1", "M2", "M3", "Mv", "M1n", "M2n", "Mvn", "M1ninv", "M0ad", "WMM", "WMMnew"] + names = ["M0", "M1", "M2", "M3", "Mv", "M1n", "M2n", "Mvn", "M1ninv", "M0ad", "M0ad_withT", "WMM", "WMMnew"] for name in names: if name == "WMM": intermediate = mass_ops.WMM @@ -155,7 +158,9 @@ def rhs_2(e1, e2, e3): exact = xp.array([rhs_0(ee1, ee2, ee3), rhs_1(ee1, ee2, ee3), rhs_2(ee1, ee2, ee3)]) solver = "cg" - if name in ["M1n", "M2n", "Mvn", "M0ad", "WMM", "WMMnew"]: + if name == "M0ad_withT": + exact *= equil.t0(e1, e2, e3) + if name in ["M1n", "M2n", "Mvn", "M0ad", "WMM", "WMMnew", "M0ad_withT"]: # solve n0 * u = f, where n0 is the equilibrium density exact /= equil.n0(e1, e2, e3) elif name == "M1ninv": diff --git a/src/struphy/feec/utilities.py b/src/struphy/feec/utilities.py index e7a4283dd..fdc20ab39 100644 --- a/src/struphy/feec/utilities.py +++ b/src/struphy/feec/utilities.py @@ -23,6 +23,38 @@ def get_quad_grids( return tuple({q: gag} for q, gag in zip(nquads, space.get_assembly_grids(*nquads))) +class LocalProjectionMatrix: + """For a given triple of callables representing the components of a normalized vector-valued function a(e1, e2, e3), + represents the local projection matrix P defined by P = a a^T at (e1, e2, e3). + + LocalProjectionMatrix(e1, e2, e3) returns a five-dimensional array, with the 3x3 matrix in the last two indices. + + This can then be used with the following numpy functions: + * matvec for matrix-vector multiplication in the last indices + * @ for matrix-matrix multiplication in the last two indices + + Parameters + ---------- + *vec_fun : list + Three callables that represent the components of the vector-valued function a. + """ + + def __init__(self, *vec_fun): + assert len(vec_fun) == 3 + assert all([callable(fun) for fun in vec_fun]) + + self._funs = vec_fun + + def __call__(self, e1, e2, e3): + # array from 2d list gives 3x3 array is in the first two indices + tmp = xp.array( + [[self._funs[m](e1, e2, e3) * self._funs[n](e1, e2, e3) for n in range(3)] for m in range(3)], + ) + + # numpy operates on the last two indices with @ + return xp.transpose(tmp, axes=(2, 3, 4, 0, 1)) + + class LocalRotationMatrix: """For a given triple of callables representing the components of a vector-valued function a(e1, e2, e3), represents the local rotation matrix R defined by Rv = a x v at (e1, e2, e3) for any vector v in R^3. diff --git a/src/struphy/propagators/poisson_field_solve.py b/src/struphy/propagators/poisson_field_solve.py index 2ddbab220..64406ce4f 100644 --- a/src/struphy/propagators/poisson_field_solve.py +++ b/src/struphy/propagators/poisson_field_solve.py @@ -48,6 +48,14 @@ class Options: - ``"M0ad"``: adiabatic-electron weighted 0-form mass operator. - ``"Id"``: identity operator. + diffusion_mat : {"M1", "M1perp", "M1gyro"}, defaults="M1" + Diffusion matrix. + + - ``"M1"``: standard weighted 1-form mass operator. + - ``"M1perp"``: weighted 1-form mass operator perpendicular to magnetic field. + - ``"M1para"``: weighted 1-form mass operator parallele to magnetic field. + - ``"M1gyro"``: weighted 1-form mass operator used in gyrokinetic model. + rho : FEECVariable or Callable or tuple or list, default=None Right-hand side source term(s) of the Poisson problem. Accepted entries are: @@ -92,9 +100,11 @@ class Options: # specific literals OptsStabMat = Literal["M0", "M0ad", "Id"] + OptsDiffusionMat = Literal["M1", "M1perp", "M1para", "M1gyro"] # propagator options stab_eps: float = 0.0 stab_mat: OptsStabMat = "Id" + diffusion_mat: OptsDiffusionMat = "M1" rho: FEECVariable | Callable | tuple[AccumulatorVector, Particles] | list = None rho_coeffs: float | list = None x0: StencilVector = None @@ -105,6 +115,7 @@ class Options: def __post_init__(self): # checks check_option(self.stab_mat, self.OptsStabMat) + check_option(self.diffusion_mat, self.OptsDiffusionMat) check_option(self.solver, LiteralOptions.OptsSymmSolver) check_option(self.precond, LiteralOptions.OptsMassPrecond) @@ -117,7 +128,6 @@ def __post_init__(self): self.sigma_2 = 0.0 self.sigma_3 = 1.0 self.divide_by_dt = False - self.diffusion_mat = "M1" @property def options(self) -> Options: diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index 8f298aad3..25a46cb07 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -5,7 +5,7 @@ import pytest from feectools.ddm.mpi import mpi as MPI -from struphy import domains +from struphy import domains, equils, set_logging_level from struphy.feec.mass import L2Projector, WeightedMassOperators from struphy.feec.psydac_derham import Derham from struphy.geometry.base import Domain @@ -14,9 +14,11 @@ from struphy.models.variables import FEECVariable from struphy.propagators.base import Propagator from struphy.propagators.implicit_diffusion import ImplicitDiffusion +from struphy.propagators.poisson_field_solve import PoissonFieldSolve from struphy.topology.grids import TensorProductGrid logger = logging.getLogger("struphy") +set_logging_level(logging.INFO) comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -453,21 +455,19 @@ def rho2_pulled(e1, e2, e3): assert error2 < 0.023 -@pytest.mark.skip(reason="Not clear if the 2.5d strategy is sound.") @pytest.mark.parametrize("num_elements", [[32, 32, 16]]) @pytest.mark.parametrize("degree", [[1, 1, 1], [2, 2, 1]]) @pytest.mark.parametrize( "mapping", [ ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}], - ["Colella", {"Lx": 1.0, "Ly": 1.0, "alpha": 0.1, "Lz": 1.0}], ], ) -def test_poisson_M1perp_3d_compare_2p5d(num_elements, degree, mapping, show_plot=False): +def test_poisson_M1perp_3d_compare_M1(num_elements, degree, mapping, show_plot=False): """ Test the Poisson solver with M1perp diffusion matrix - by comparing 3d simulation to a loop over 2d simulations. - Dirichlet boundary conditions in eta1. + by comparing 3d simulation using M1perp and M1 diffusion matrices, and integrating over the third direction. + The two analytical solutions must be exactly the same in both cases. """ from time import time @@ -479,36 +479,32 @@ def test_poisson_M1perp_3d_compare_2p5d(num_elements, degree, mapping, show_plot domain_class = getattr(domains, dom_type) domain: Domain = domain_class(**dom_params) - # boundary conditions - bcs = (("dirichlet", "dirichlet"), None, None) + equil = equils.HomogenSlab() + equil.domain = domain # evaluation grid - e1 = xp.linspace(0.0, 1.0, 50) - e2 = xp.linspace(0.0, 1.0, 60) - e3 = xp.linspace(0.0, 1.0, 30) + e1 = xp.linspace(0.0, 1.0, num_elements[0]) + e2 = xp.linspace(0.0, 1.0, num_elements[1]) + e3 = xp.linspace(0.0, 1.0, num_elements[2]) # solution and right-hand side on unit cube def rho(e1, e2, e3): - dd1 = xp.sin(xp.pi * e1) * xp.sin(4 * xp.pi * e2) * xp.cos(2 * xp.pi * e3) * (xp.pi) ** 2 - dd2 = xp.sin(xp.pi * e1) * xp.sin(4 * xp.pi * e2) * xp.cos(2 * xp.pi * e3) * (4 * xp.pi) ** 2 - return dd1 + dd2 + dd1 = xp.sin(xp.pi * e1) * xp.sin(2 * xp.pi * e2) * (1 + xp.cos(2 * xp.pi * e3)) * (xp.pi) ** 2 + return dd1 # create 3d derham object grid = TensorProductGrid(num_elements=num_elements) - derham_opts = DerhamOptions(degree=degree, bcs=bcs) + derham_opts = DerhamOptions(degree=degree, bcs=(None, None, None)) derham = Derham(grid, derham_opts, comm=comm) - mass_ops = WeightedMassOperators(derham, domain) + mass_ops = WeightedMassOperators(derham, domain, eq_mhd=equil) Propagator.derham = derham Propagator.domain = domain Propagator.mass_ops = mass_ops # discrete right-hand sides - l2_proj = L2Projector("H1", mass_ops) - rho_vec = l2_proj.get_dofs(rho, apply_bc=True) - - logger.info(f"{rho_vec[:].shape =}") + # l2_proj = L2Projector("H1", mass_ops) # Create 3d Poisson solver solver_params = SolverParameters( @@ -519,20 +515,34 @@ def rho(e1, e2, e3): recycle=False, ) - _phi = FEECVariable(space="H1") - _phi.allocate(derham=derham, domain=domain) + _phi_M1 = FEECVariable(space="H1") + _phi_M1.allocate(derham=derham, domain=domain) - _phi_2p5d = FEECVariable(space="H1") - _phi_2p5d.allocate(derham=derham, domain=domain) + poisson_solver_M1 = PoissonFieldSolve() + poisson_solver_M1.variables.phi = _phi_M1 - poisson_solver_3d = ImplicitDiffusion() - poisson_solver_3d.variables.phi = _phi + poisson_solver_M1.options = poisson_solver_M1.Options( + stab_eps=1e-8, + diffusion_mat="M1", + rho=rho, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, + ) - poisson_solver_3d.options = poisson_solver_3d.Options( - sigma_1=1e-8, - sigma_2=0.0, - sigma_3=1.0, - divide_by_dt=True, + poisson_solver_M1.allocate() + + s = _phi_M1.spline.starts + e = _phi_M1.spline.ends + + _phi_M1perp = FEECVariable(space="H1") + _phi_M1perp.allocate(derham=derham, domain=domain) + + poisson_solver_M1perp = PoissonFieldSolve() + poisson_solver_M1perp.variables.phi = _phi_M1perp + + poisson_solver_M1perp.options = poisson_solver_M1perp.Options( + stab_eps=1e-8, diffusion_mat="M1perp", rho=rho, solver="pcg", @@ -540,45 +550,170 @@ def rho(e1, e2, e3): solver_params=solver_params, ) - poisson_solver_3d.allocate() + poisson_solver_M1perp.allocate() - s = _phi.spline.starts - e = _phi.spline.ends + # Solve M1 Poisson equation (call propagator with dt=1.) + dt = 1.0 + t0 = time() + poisson_solver_M1(dt) + t1 = time() + logger.info(f"rank {rank}, M1 3d solve time = {t1 - t0}") - # create 2.5d deRham object - num_elements_new = [num_elements[0], num_elements[1], 1] - degree[2] = 1 + # Solve M1perp Poisson equation (call propagator with dt=1.) + t0 = time() + poisson_solver_M1perp(dt) + t1 = time() + logger.info(f"rank {rank}, M1perp 3d solve time = {t1 - t0}") - grid = TensorProductGrid(num_elements=num_elements_new) - derham_opts = DerhamOptions(degree=degree, bcs=bcs) - derham = Derham(grid, derham_opts, comm=comm) + # push numerical solutions + sol_val_M1 = domain.push(_phi_M1.spline, e1, e2, e3, kind="0") + sol_val_M1perp = domain.push(_phi_M1perp.spline, e1, e2, e3, kind="0") + x, y, z = domain(e1, e2, e3) - mass_ops = WeightedMassOperators(derham, domain) + # The two solutions might not be the same: + logger.info(f"max diff: {xp.max(xp.abs(sol_val_M1 - sol_val_M1perp))}") - Propagator.derham = derham - Propagator.mass_ops = mass_ops + # We take the average of both solutions and we compare their value, we should obtain the same: + max_diff_av = xp.max(xp.abs(xp.trapezoid(sol_val_M1 - sol_val_M1perp, e3, axis=2) / (e3[-1] - e3[0]))) + logger.info(f"max diff of the averaged solutions (over e3): {max_diff_av}") + assert max_diff_av < 0.001 + if show_plot and rank == 0: + plt.figure("e1-e2 plane", figsize=(24, 16)) + plt.subplot(2, 3, 1) + plt.title("charge density averaged over e3") + plt.pcolor( + x[:, :, 0], + y[:, :, 0], + xp.transpose(xp.sum(rho(*xp.meshgrid(e1, e2, e3)), axis=2)) / len(e3), + shading="nearest", + ) + plt.colorbar() + ax = plt.gca() + ax.set_aspect("equal", adjustable="box") + plt.subplot(2, 3, 4) + plt.title(f"charge density at e3={e3[len(e3) // 2]:.2f}") + plt.pcolor( + x[:, :, 0], y[:, :, 0], xp.transpose(rho(*xp.meshgrid(e1, e2, e3))[:, :, len(e3) // 2]), shading="nearest" + ) + plt.colorbar() + ax = plt.gca() + ax.set_aspect("equal", adjustable="box") + plt.subplot(2, 3, 2) + plt.title(f"phi at e3={e3[len(e3) // 2]:.2f} with M1 solve") + plt.pcolor(x[:, :, 0], y[:, :, 0], sol_val_M1[:, :, len(e3) // 2]) + plt.colorbar() + ax = plt.gca() + ax.set_aspect("equal", adjustable="box") + plt.subplot(2, 3, 5) + plt.title(f"phi at e3={e3[len(e3) // 2]:.2f} with M1perp solve") + plt.pcolor(x[:, :, 0], y[:, :, 0], sol_val_M1perp[:, :, len(e3) // 2]) + plt.colorbar() + ax = plt.gca() + ax.set_aspect("equal", adjustable="box") + plt.subplot(2, 3, 3) + plt.title("average over e3 of M1 solve") + plt.pcolor(x[:, :, 0], y[:, :, 0], xp.trapezoid(sol_val_M1, e3, axis=2) / (e3[-1] - e3[0])) + plt.colorbar() + ax = plt.gca() + ax.set_aspect("equal", adjustable="box") + plt.subplot(2, 3, 6) + plt.title("average over e3 of M1perp solve") + plt.pcolor(x[:, :, 0], y[:, :, 0], xp.trapezoid(sol_val_M1perp, e3, axis=2) / (e3[-1] - e3[0])) + plt.colorbar() + ax = plt.gca() + ax.set_aspect("equal", adjustable="box") + plt.show() + + +# @pytest.mark.skip(reason="Not clear if the 2.5d strategy is sound.") +@pytest.mark.parametrize("num_elements", [[32, 32, 16]]) +@pytest.mark.parametrize("degree", [[1, 1, 1], [2, 2, 1]]) +@pytest.mark.parametrize( + "mapping", + [ + ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}], + ], +) +def test_poisson_M1perp_3d_compare_2p5d(num_elements, degree, mapping, show_plot=False): + """ + Test the Poisson solver with M1perp diffusion matrix + by comparing 3d simulation to a loop over 2d simulations. + """ - _phi_small = FEECVariable(space="H1") - _phi_small.allocate(derham=derham, domain=domain) + from time import time - poisson_solver_2p5d = ImplicitDiffusion() - poisson_solver_2p5d.variables.phi = _phi_small + # create domain object + dom_type = mapping[0] + dom_params = mapping[1] - poisson_solver_2p5d.options = poisson_solver_2p5d.Options( - sigma_1=1e-8, - sigma_2=0.0, - sigma_3=1.0, - divide_by_dt=True, + domain_class = getattr(domains, dom_type) + domain_3D: Domain = domain_class(**dom_params) + + # boundary conditions + bcs = (None, None, ("free", "free")) + + # evaluation grid + e1 = xp.linspace(0.0, 1.0, num_elements[0]) + e2 = xp.linspace(0.0, 1.0, num_elements[1]) + e3 = xp.linspace(0.0, 1.0, num_elements[2]) + + # solution and right-hand side on unit cube + def rho_physical(x, y, z): + dd1 = xp.sin(2 * xp.pi * x) * xp.sin(xp.pi * y) * xp.cos(xp.pi * z) * (xp.pi) ** 2 + return dd1 + + # create 3d derham object + grid_3D = TensorProductGrid(num_elements=num_elements) + derham_opts_3D = DerhamOptions(degree=degree, bcs=bcs) + derham_3D = Derham(grid_3D, derham_opts_3D, comm=comm) + + mass_ops_3D = WeightedMassOperators(derham_3D, domain_3D, eq_mhd=equils.HomogenSlab(B0x=0.0, B0y=0.0, B0z=1.0)) + + Propagator.derham = derham_3D + Propagator.domain = domain_3D + Propagator.mass_ops = mass_ops_3D + + rho_logical_3D = lambda e1, e2, e3: rho_physical(*domain_3D(e1, e2, e3)) + + # discrete right-hand sides + # l2_proj = L2Projector("H1", mass_ops_3D) + # rho_vec = l2_proj.get_dofs(rho, apply_bc=True) + + # logger.info(f"{rho_vec[:].shape =}") + + # Create 3d Poisson solver + solver_params = SolverParameters( + tol=1.0e-13, + maxiter=3000, + info=True, + verbose=False, + recycle=False, + ) + + _phi = FEECVariable(space="H1") + _phi.allocate(derham=derham_3D, domain=domain_3D) + + _phi_2p5d = FEECVariable(space="H1") + _phi_2p5d.allocate(derham=derham_3D, domain=domain_3D) + + poisson_solver_3d = PoissonFieldSolve() + poisson_solver_3d.variables.phi = _phi + + poisson_solver_3d.options = poisson_solver_3d.Options( + stab_eps=1e-8, diffusion_mat="M1perp", - rho=rho, + rho=rho_logical_3D, solver="pcg", precond="MassMatrixPreconditioner", solver_params=solver_params, ) - poisson_solver_2p5d.allocate() + poisson_solver_3d.allocate() - # Solve Poisson equation (call propagator with dt=1.) + s = _phi.spline.starts + e = _phi.spline.ends + + # Solve 3d Poisson equation (call propagator with dt=1.) dt = 1.0 t0 = time() poisson_solver_3d(dt) @@ -586,9 +721,49 @@ def rho(e1, e2, e3): logger.info(f"rank {rank}, 3d solve time = {t1 - t0}") + # create 2.5d deRham object for each slice in e3, and solve 2d Poisson equations + num_elements_new = [num_elements[0], num_elements[1], 1] + degree_sliced = degree.copy() + degree_sliced[2] = 1 + dom_params_sliced = dom_params.copy() t0 = time() t_inner = 0.0 for n in range(s[2], e[2] + 1): + # We define another domain that maps into a slice of the previous Cuboid domain, slices are perpendicular to the third direction. + dom_params_sliced["l3"] = dom_params["l3"] + (dom_params["r3"] - dom_params["l3"]) / (len(e3) + degree[2]) * n + dom_params_sliced["r3"] = dom_params["l3"] + (dom_params["r3"] - dom_params["l3"]) / (len(e3) + degree[2]) * ( + n + 1 + ) + domain_sliced = domain_class(**dom_params_sliced) + + grid_sliced = TensorProductGrid(num_elements=num_elements_new) + derham_opts_sliced = DerhamOptions(degree=degree_sliced, bcs=bcs) + derham_sliced = Derham(grid_sliced, derham_opts_sliced, comm=comm) + + mass_ops_sliced = WeightedMassOperators(derham_sliced, domain_sliced) + + Propagator.derham = derham_sliced + Propagator.mass_ops = mass_ops_sliced + Propagator.domain = domain_sliced + + rho_logical_sliced = lambda e1, e2, e3: rho_physical(*domain_sliced(e1, e2, e3)) + + _phi_small = FEECVariable(space="H1") + _phi_small.allocate(derham=derham_sliced, domain=domain_sliced) + + poisson_solver_2p5d = PoissonFieldSolve() + poisson_solver_2p5d.variables.phi = _phi_small + + poisson_solver_2p5d.options = poisson_solver_2p5d.Options( + stab_eps=1e-8, + diffusion_mat="M1", + rho=rho_logical_sliced, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, + ) + poisson_solver_2p5d.allocate() + t0i = time() poisson_solver_2p5d(dt) t1i = time() @@ -601,39 +776,42 @@ def rho(e1, e2, e3): logger.info(f"rank {rank}, 2.5d solve time = {t1 - t0}") # push numerical solutions - sol_val = domain.push(_phi.spline, e1, e2, e3, kind="0") - sol_val_2p5d = domain.push(_phi_2p5d.spline, e1, e2, e3, kind="0") - x, y, z = domain(e1, e2, e3) + sol_val = domain_3D.push(_phi.spline, e1, e2, e3, kind="0") + sol_val_2p5d = domain_3D.push(_phi_2p5d.spline, e1, e2, e3, kind="0") + x, y, z = domain_3D(e1, e2, e3) + logger.info(f"mean diff: {xp.mean(xp.abs(sol_val - sol_val_2p5d))}") logger.info(f"max diff: {xp.max(xp.abs(sol_val - sol_val_2p5d))}") - assert xp.max(xp.abs(sol_val - sol_val_2p5d)) < 0.026 + assert xp.max(xp.abs(sol_val - sol_val_2p5d)) < 0.01 if show_plot and rank == 0: plt.figure("e1-e2 plane", figsize=(24, 16)) + plot_id_e3 = [0, len(e3) // 2, len(e3) - 1] + plot_id_e2 = [0, len(e2) // 2, len(e2) - 1] for n in range(3): plt.subplot(2, 3, n + 1) - plt.title(f"e3 = {e3[n * 6]} from 3d solve") - plt.pcolor(x[:, :, n * 6], y[:, :, n * 6], sol_val[:, :, n * 6], vmin=-1.0, vmax=1.0) + plt.title(f"e3 = {e3[plot_id_e3[n]]} from 3d solve") + plt.pcolor(x[:, :, plot_id_e3[n]], y[:, :, plot_id_e3[n]], sol_val[:, :, plot_id_e3[n]]) plt.colorbar() ax = plt.gca() ax.set_aspect("equal", adjustable="box") plt.subplot(2, 3, 4 + n) - plt.title(f"e3 = {e3[n * 6]} from 2.5d solve") - plt.pcolor(x[:, :, n * 6], y[:, :, n * 6], sol_val_2p5d[:, :, n * 6], vmin=-1.0, vmax=1.0) + plt.title(f"e3 = {e3[plot_id_e3[n]]} from 2.5d solve") + plt.pcolor(x[:, :, plot_id_e3[n]], y[:, :, plot_id_e3[n]], sol_val_2p5d[:, :, plot_id_e3[n]]) plt.colorbar() ax = plt.gca() ax.set_aspect("equal", adjustable="box") plt.figure("e1-e3 plane", figsize=(24, 16)) for n in range(3): plt.subplot(2, 3, n + 1) - plt.title(f"e2 = {e2[n * 12]} from 3d solve") - plt.pcolor(x[:, n * 12, :], z[:, n * 12, :], sol_val[:, n * 12, :], vmin=-1.0, vmax=1.0) + plt.title(f"e2 = {e2[plot_id_e2[n]]} from 3d solve") + plt.pcolor(x[:, plot_id_e2[n], :], z[:, plot_id_e2[n], :], sol_val[:, plot_id_e2[n], :]) plt.colorbar() ax = plt.gca() ax.set_aspect("equal", adjustable="box") plt.subplot(2, 3, 4 + n) - plt.title(f"e2 = {e2[n * 12]} from 2.5d solve") - plt.pcolor(x[:, n * 12, :], z[:, n * 12, :], sol_val_2p5d[:, n * 12, :], vmin=-1.0, vmax=1.0) + plt.title(f"e2 = {e2[plot_id_e2[n]]} from 2.5d solve") + plt.pcolor(x[:, plot_id_e2[n], :], z[:, plot_id_e2[n], :], sol_val_2p5d[:, plot_id_e2[n], :]) plt.colorbar() ax = plt.gca() ax.set_aspect("equal", adjustable="box") @@ -645,17 +823,19 @@ def rho(e1, e2, e3): direction = 0 bc_type = "dirichlet" mapping = ["Cuboid", {"l1": 0.0, "r1": 4.0, "l2": 0.0, "r2": 2.0, "l3": 0.0, "r3": 3.0}] - mapping = ["Orthogonal", {"Lx": 4.0, "Ly": 2.0, "alpha": 0.1, "Lz": 3.0}] - test_poisson_M1perp_1d(direction, bc_type, mapping, show_plot=True) - - # num_elements = [64, 64, 1] - # degree = [2, 2, 1] - # bc_type = 'neumann' - # #mapping = ['Cuboid', {'l1': 0., 'r1': 4., 'l2': 0., 'r2': 2., 'l3': 0., 'r3': 3.}] - # mapping = ['Orthogonal', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 1.}] - # test_poisson_M1perp_2d(num_elements, degree, bc_type, mapping, show_plot=True) - - # num_elements = [64, 64, 16] - # degree = [2, 2, 1] - # mapping = ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}] + # mapping = ["Orthogonal", {"Lx": 4.0, "Ly": 2.0, "alpha": 0.1, "Lz": 3.0}] + # test_poisson_M1perp_1d(direction, bc_type, mapping, projected_rhs=True, show_plot=True) + + num_elements = [64, 64, 1] + degree = [2, 2, 1] + bc_type = "neumann" + mapping = ["Cuboid", {"l1": 0.0, "r1": 4.0, "l2": 0.0, "r2": 2.0, "l3": 0.0, "r3": 3.0}] + mapping = ["Orthogonal", {"Lx": 4.0, "Ly": 2.0, "alpha": 0.1, "Lz": 1.0}] + # test_poisson_M1perp_2d(num_elements, degree, bc_type, mapping, projected_rhs=True, show_plot=True) + + num_elements = [50, 50, 50] + degree = [2, 2, 1] + mapping = ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}] + # test_poisson_M1perp_3d_compare_M1(num_elements, degree, mapping, show_plot=True) + num_elements = [50, 50, 50] # test_poisson_M1perp_3d_compare_2p5d(num_elements, degree, mapping, show_plot=True)