From f526b7662cb392189cd9a076018ab1dcfb861916 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 9 Oct 2026 14:51:04 +0200 Subject: [PATCH 1/4] StiffnessPreconditioner for PoissonAdiabaticGyrokinetic: general mid/stab, weighted average in the stabilization - StiffnessPreconditioner: mid (default M_{k+1}) and stab (default M_k, or "Id") as names of mass operators or operators (incl. sums such as M1gyro), stab_average ("eta3", "eta2 eta3", grad only); blocks of mid with zero mean have no stiffness term (their 1d stiffness still defines the eigenbasis); the averaged stabilization S (I - A) is approximated by mean-removing terms, exact on the logical cube; the average is left out of the diagonal scaling (not resolvable by probing) - WeightedAverageProjection: S-orthogonal projection onto functions constant in the averaged directions, S (I - P) is symmetric; equals the plain AverageOperator if the weight of S does not depend on the averaged directions; e.g. the flux-surface average with sqrt(g) for M0 - ImplicitDiffusion/PoissonSolve: option stab_average; stab_mat @ (I - P) built in ImplicitDiffusion; the stiffness preconditioner gets diffusion_mat, stab_mat and stab_average; MultiGrid rejected with stab_average - PoissonAdiabaticGyrokinetic: stab_average from which_geometry (cylindrical: eta3, toroidal: eta2 eta3); with the weighted average the toroidal operator is symmetric (it was not with the plain average, pcg was not applicable); MultiGrid rejected - feectools submodule: KroneckerSumSolver with stiffness coefficients and additional Kronecker terms (struphy-hub/feectools#102) - tests: exactness of the new options on the cube, option checks, WeightedAverageProjection (symmetry, projection, plain average), gyrokinetic Poisson with the stiffness preconditioner (cylindrical 34 -> 10, toroidal 109 -> 20 PCG iterations) Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- src/struphy/feec/mass.py | 164 +++++++++ src/struphy/feec/preconditioner.py | 339 ++++++++++++++++-- .../tests/test_stiffness_preconditioner.py | 95 +++++ .../tests/test_weighted_average_projection.py | 56 +++ src/struphy/propagators/implicit_diffusion.py | 35 +- .../poisson_adiabatic_gyrokinetic.py | 34 +- src/struphy/propagators/poisson_solve.py | 10 + .../tests/test_gyrokinetic_poisson.py | 52 +++ 9 files changed, 726 insertions(+), 61 deletions(-) create mode 100644 src/struphy/feec/tests/test_weighted_average_projection.py diff --git a/feectools b/feectools index 5385af968..c2593636d 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 5385af968237c96f7d3ca643a5da4f4e4e31916d +Subproject commit c2593636dade2386881c8f9e0daaada9acc11f50 diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index e56a19af0..0b2fc656a 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -3520,3 +3520,167 @@ def transpose(self, conjugate=False): return AverageOperator( self.derham, self._space, self._direction, transposed=not self._transposed, nquads=self._nquads ) + + +class WeightedAverageProjection(LinearOperator): + r""" + Projection onto functions that are constant in the given directions, orthogonal with respect to + a symmetric positive (semi-)definite operator :math:`\mathbb S` on the 0-form coefficients: + + .. math:: + + \Pi_{\mathbb S} = \mathbb E \, (\mathbb E^\top \mathbb S \, \mathbb E)^{-1} \, \mathbb E^\top \mathbb S \,, + + where :math:`\mathbb E` embeds the coefficients of the remaining directions as constant along the + averaged ones (:math:`\mathbb E^\top` sums over the averaged directions). For a weighted mass matrix + :math:`\mathbb S`, :math:`\Pi_{\mathbb S} \phi` is the spline in the remaining directions closest to + :math:`\phi` in the weighted :math:`L^2` norm, i.e. the weighted average (e.g. the flux-surface + average with the volume element :math:`\sqrt g` for :math:`\mathbb M^0`). Then + + .. math:: + + \mathbb S (\mathbb I - \Pi_{\mathbb S}) = \mathbb S - \mathbb S \mathbb E \, (\mathbb E^\top \mathbb S \, \mathbb E)^{-1} \, \mathbb E^\top \mathbb S + + is symmetric. If the weight of :math:`\mathbb S` does not depend on the averaged directions, + :math:`\Pi_{\mathbb S}` is the plain average of :class:`AverageOperator`. + + The reduced matrix :math:`\mathbb E^\top \mathbb S \, \mathbb E` (on the remaining directions) is + assembled once by probing :math:`\mathbb S` with colored vectors and factorized on every process; + one application costs one application of :math:`\mathbb S` and a sparse solve. + + Parameters + ---------- + derham : Derham + Discrete de Rham sequence (0-forms, no polar splines). + + S : LinearOperator + Symmetric operator on the 0-form coefficients (e.g. a weighted mass operator), positive + definite on the functions constant in the averaged directions (up to essential boundary + conditions, whose reduced rows are replaced by the identity). + + directions : tuple[int] + The averaged directions (0, 1 and/or 2). + """ + + def __init__(self, derham: Derham, S: LinearOperator, directions: tuple[int, ...]): + import numpy as np + from scipy import sparse + from scipy.sparse.linalg import splu + + V = derham.V0 + assert S.domain == V and S.codomain == V, "S must act on the 0-form coefficients." + directions = tuple(sorted(set(directions))) + assert directions and all(d in (0, 1, 2) for d in directions) + kept = tuple(d for d in range(3) if d not in directions) + assert kept, "At least one direction must not be averaged." + + self._derham = derham + self._S = S + self._space = V + self._directions = directions + self._kept = kept + self._comm = None if isinstance(derham.comm, (SerialComm, type(None))) else derham.comm + self._local = tuple(slice(p * m, p * m + e - s + 1) for p, m, s, e in zip(V.pads, V.shifts, V.starts, V.ends)) + self._red_shape = tuple(V.npts[d] for d in kept) + self._red_local = tuple(slice(V.starts[d], V.ends[d] + 1) for d in kept) + self._tmp = V.zeros() + self._tmp2 = V.zeros() + + # reduced matrix E^T S E, column by column with colored probes: in each kept direction the + # columns of one color are at least 2 p + 1 apart (also across a periodic boundary) + strides = [] + for d in kept: + n, p = V.npts[d], V.pads[d] + stride = n + if not V.periods[d]: + stride = min(n, 2 * p + 1) + else: + for s in range(2 * p + 1, n + 1): + if n % s == 0: + stride = s + break + strides.append(stride) + + rows, cols, vals = [], [], [] + n_red = int(np.prod(self._red_shape)) + index = np.arange(n_red).reshape(self._red_shape) + for color in np.ndindex(*strides): + probe = np.zeros(self._red_shape) + sel = tuple(slice(c, None, s) for c, s in zip(color, strides)) + probe[sel] = 1.0 + col = self._ET(self._S.dot(self._E(probe, out=self._tmp), out=self._tmp2)) + # entry (b, a) for the probed columns a: rows b within the coupling width of a + for a in index[sel].ravel(): + a_idx = np.unravel_index(a, self._red_shape) + ranges = [] + for k, d in enumerate(kept): + n, p = V.npts[d], V.pads[d] + r = np.arange(a_idx[k] - p, a_idx[k] + p + 1) + r = r % n if V.periods[d] else r[(r >= 0) & (r < n)] + ranges.append(np.unique(r)) + for b_idx in np.ndindex(*[len(r) for r in ranges]): + b = tuple(r[i] for r, i in zip(ranges, b_idx)) + rows.append(index[b]) + cols.append(a) + vals.append(col[b]) + K = sparse.coo_matrix((vals, (rows, cols)), shape=(n_red, n_red)).tocsc() + K.sum_duplicates() + # rows without entries (essential boundary conditions of S): identity + diag = K.diagonal() + empty = np.abs(diag) <= 1e-14 * max(np.abs(diag).max(), 1e-300) + if np.any(empty): + K = K + sparse.diags(empty.astype(float)) + self._K_lu = splu(K.tocsc()) + + def _ET(self, v): + """E^T v: sum over the averaged directions, as global array over the kept directions (all processes).""" + import numpy as np + + loc = xp.to_numpy(v._data[self._local]).sum(axis=self._directions) + out = np.zeros(self._red_shape) + out[self._red_local] = loc + if self._comm is not None: + out = self._comm.allreduce(out, op=MPI.SUM) + return out + + def _E(self, red, out): + """E red: the global array over the kept directions, constant along the averaged ones.""" + loc = red[self._red_local] + shape = [1, 1, 1] + for k, d in enumerate(self._kept): + shape[d] = loc.shape[k] + out._data[self._local] = xp.asarray(loc.reshape(shape)) + out.update_ghost_regions() + return out + + @property + def domain(self): + return self._space + + @property + def codomain(self): + return self._space + + @property + def dtype(self): + return self._space.dtype + + @property + def directions(self) -> tuple[int, ...]: + """The averaged directions.""" + return self._directions + + def dot(self, v, out=None): + if out is None: + out = self._space.zeros() + red = self._K_lu.solve(self._ET(self._S.dot(v, out=self._tmp)).ravel()).reshape(self._red_shape) + return self._E(red, out) + + def transpose(self, conjugate=False): + raise NotImplementedError("Only S @ (I - P) is symmetric; P itself is not.") + + def tosparse(self): + raise NotImplementedError + + def toarray(self): + raise NotImplementedError diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 1d3481b25..6ef305d68 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -13,7 +13,9 @@ from feectools.fem.tensor import TensorFemSpace from feectools.linalg.basic import ( ComposedLinearOperator, + IdentityOperator, LinearOperator, + LinearSolver, ScaledLinearOperator, SumLinearOperator, Vector, @@ -28,7 +30,7 @@ from scipy import sparse from struphy.feec.linear_operators import BoundaryOperator -from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators +from struphy.feec.mass import WeightedAverageProjection, WeightedMassOperator, WeightedMassOperators logger = logging.getLogger("struphy") @@ -128,6 +130,10 @@ def _core_diagonal(self) -> LinearOperator: """Diagonal :math:`D` of the core operator as diagonal operator (for diagonal scaling).""" raise NotImplementedError(f"{type(self).__name__} does not support diagonal scaling.") + def _scaling_matrix(self) -> LinearOperator: + """The approximation whose diagonal :math:`\\hat D` enters the diagonal scaling (default: ``matrix``).""" + return self._matrix + def transpose(self, conjugate: bool = False) -> KroneckerPreconditioner: raise NotImplementedError @@ -171,7 +177,7 @@ def diagonal_scaling(self) -> bool: def _build_scaling(self) -> tuple[LinearOperator, ...]: """Diagonal scaling :math:`S = \\hat D^{1/2} D^{-1/2}` as one diagonal operator.""" D = _local_diagonal(self._core_diagonal()) - D_hat = _local_diagonal(self._matrix) + D_hat = _local_diagonal(self._scaling_matrix()) S = [xp.where(d > 0, xp.sqrt(xp.abs(dh) / xp.where(d > 0, d, 1.0)), 1.0) for d, dh in zip(D, D_hat)] return (_diagonal_operator(self._matrix.domain, S),) @@ -373,6 +379,15 @@ class StiffnessPreconditioner(KroneckerPreconditioner): :math:`\mathbb C^T \mathbb M_2 \mathbb C + \sigma \mathbb M_1` and :math:`\mathbb D^T \mathbb M_3 \mathbb D + \sigma \mathbb M_2`. + The matrix in the middle (``mid``, default :math:`\mathbb M_{k+1}`) and the stabilization + (``stab``, default :math:`\mathbb M_k`) can be other weighted mass operators with the same + domain and codomain (or sums and multiples of them), e.g. ``mid="M1gyro"`` and + ``stab="M0ad_withT"``; ``stab="Id"`` is the identity. For the gradient, ``stab_average`` + replaces the stabilization :math:`\mathbb S` by :math:`\mathbb S (\mathbb I - \mathbb A)`, where + :math:`\mathbb A` is the :math:`\mathbb S`-weighted average over the given directions + (:class:`~struphy.feec.mass.WeightedAverageProjection`), as in the adiabatic gyrokinetic Poisson + equation. + With unit weights (the operator on the logical cube), the mass matrices are block-diagonal Kronecker products of 1d mass matrices :math:`M_d`, and each diagonal block of :math:`A` is a sum of Kronecker products, @@ -410,7 +425,17 @@ class StiffnessPreconditioner(KroneckerPreconditioner): over the other directions, averaged over the diagonal blocks of :math:`\mathbb M_{k+1}`), times the mean of the weight of each term (the block of :math:`\mathbb M_{k+1}` for the stiffness term of each direction, the block of :math:`\mathbb M_k` for the mass term). This is exact - e.g. for constant but anisotropic weights. + e.g. for constant but anisotropic weights. A block of ``mid`` with zero mean (e.g. the + parallel direction of a perpendicular operator) gives no stiffness term in that direction; its + 1d stiffness still defines the eigenbasis of the fast diagonalization. + + The averaged stabilization :math:`\mathbb S (\mathbb I - \mathbb A)` is approximated by + :math:`\bar s \, (M_1 \otimes M_2 \otimes M_3 - F_1 \otimes F_2 \otimes F_3)` with + :math:`F_d = c_d c_d^T / (\mathbb 1^T c_d)`, :math:`c_d = M_d \mathbb 1`, in the averaged + directions (:math:`F_d = M_d` otherwise). In directions with :math:`S_d \mathbb 1 = 0` + (periodic or without essential boundary conditions) the constants are an eigenvector of the + fast diagonalization and this term is exact in its eigenbasis (it removes the averaged modes). + The identity (``stab="Id"``) is replaced by its diagonal in that eigenbasis. Polar splines are not supported yet. @@ -439,6 +464,20 @@ class StiffnessPreconditioner(KroneckerPreconditioner): weights : str ``"average"`` (separable approximation of the weights, see above; default) or ``"unit"`` (operator on the logical cube). + + mid : str | LinearOperator | None + Matrix in the middle: the name of a mass operator of ``mass_ops`` (e.g. ``"M1gyro"``, + ``"M1perp"``) or the operator itself (a WeightedMassOperator, or a sum or multiple of such), + with the domain and codomain of :math:`\mathbb M_{k+1}`. Default :math:`\mathbb M_{k+1}`. + + stab : str | LinearOperator | None + Stabilization: the name of a mass operator of ``mass_ops`` (e.g. ``"M0ad"``, + ``"M0ad_withT"``), the operator itself (as for ``mid``), with the domain and codomain of + :math:`\mathbb M_k`, or ``"Id"`` for the identity. Default :math:`\mathbb M_k`. + + stab_average : str | None + Gradient only: the directions averaged in the stabilization :math:`\mathbb S (\mathbb I - \mathbb A)`, + e.g. ``"eta3"`` or ``"eta2 eta3"``; None (default) for no averaging. """ _FORMS = {"grad": 0, "curl": 1, "div": 2} @@ -452,6 +491,9 @@ def __init__( diagonal_scaling: bool = True, kernel_correction: bool = True, weights: str = "average", + mid: str | LinearOperator | None = None, + stab: str | LinearOperator | None = None, + stab_average: str | None = None, ): assert derivative in self._FORMS, f"derivative must be one of {tuple(self._FORMS)}, got {derivative!r}." assert sigma >= 0.0 @@ -469,13 +511,36 @@ def __init__( self._derivative = derivative self._sigma = sigma self._weights = weights - self._mass_mid = getattr(mass_ops, f"M{k + 1}") + + # matrix in the middle and stabilization, as sums of weighted mass operators + mass_mid_default = getattr(mass_ops, f"M{k + 1}") + self._mid = mid + self._mid_terms = _mass_terms(_resolve_mass_operator(mass_ops, mid, mass_mid_default)) + self._stab = stab + self._stab_is_identity = isinstance(stab, str) and stab == "Id" + mass_k = getattr(mass_ops, f"M{k}") + self._stab_terms = ( + None if self._stab_is_identity else _mass_terms(_resolve_mass_operator(mass_ops, stab, mass_k)) + ) + for terms, ref, name in ((self._mid_terms, mass_mid_default, "mid"), (self._stab_terms or [], mass_k, "stab")): + for _, op in terms: + assert op.domain == ref.domain and op.codomain == ref.codomain, ( + f"{name} must have the domain and codomain of M{k + 1 if name == 'mid' else k}." + ) + + # averaged directions of the stabilization (gradient only) + self._stab_average = stab_average + self._average_dirs = _parse_directions(stab_average) + if self._average_dirs and derivative != "grad": + raise NotImplementedError("stab_average is only implemented for derivative='grad'.") + # tensor-product derivative without boundary operators (no polar splines) self._d = getattr(derham, f"{derivative}_bcfree") - super().__init__(getattr(mass_ops, f"M{k}"), apply_bc=apply_bc, diagonal_scaling=diagonal_scaling) + super().__init__(mass_k, apply_bc=apply_bc, diagonal_scaling=diagonal_scaling) - # correction on the kernel of the derivative: sigma^{-1} d_prev P_prev d_prev^T + # correction on the kernel of the derivative: sigma^{-1} d_prev P_prev d_prev^T, + # where A = sigma * stab; P_prev approximates the inverse of d_prev^T stab d_prev self._kernel_correction = None if kernel_correction: prev = "grad" if derivative == "curl" else "curl" @@ -487,6 +552,7 @@ def __init__( diagonal_scaling=diagonal_scaling, kernel_correction=False, weights=weights, + mid=None if self._stab_is_identity else stab, ) self._kernel_correction = (1.0 / sigma) * (d_prev @ P_prev @ d_prev.T) self._tmp_kernel = self.codomain.zeros() @@ -511,6 +577,21 @@ def kernel_correction(self) -> LinearOperator | None: """The correction on the kernel of the derivative, or None.""" return self._kernel_correction + @property + def mid(self) -> str | LinearOperator | None: + """The matrix in the middle as given (None for the default :math:`\\mathbb M_{k+1}`).""" + return self._mid + + @property + def stab(self) -> str | LinearOperator | None: + """The stabilization as given (None for the default :math:`\\mathbb M_k`, ``"Id"`` for the identity).""" + return self._stab + + @property + def stab_average(self) -> str | None: + """The directions averaged in the stabilization, or None.""" + return self._stab_average + def solve(self, rhs: Vector, out: Vector | None = None) -> Vector: """Apply the preconditioner (block Jacobi part plus kernel correction, if any) to rhs.""" out = super().solve(rhs, out=out) @@ -520,10 +601,19 @@ def solve(self, rhs: Vector, out: Vector | None = None) -> Vector: @property def core_operator(self) -> LinearOperator: - """The core operator :math:`\\mathbb d^T \\mathbb M_{k+1} \\mathbb d + \\sigma \\mathbb M_k` on the tensor-product spaces.""" - A = self._d.T @ self._mass_mid.matrix @ self._d + """The core operator :math:`\\mathbb d^T \\mathbb M_{\\textrm{mid}} \\mathbb d + \\sigma \\mathbb S (\\mathbb I - \\mathbb A)` + on the tensor-product spaces (without the average if ``stab_average`` is None).""" + return self._core(with_average=True) + + def _core(self, with_average: bool) -> LinearOperator: + A = self._d.T @ _terms_matrix(self._mid_terms) @ self._d if self._sigma != 0.0: - A = A + self._sigma * self._mass_operator.matrix + space = self._mass_operator.matrix.domain + S = IdentityOperator(space) if self._stab_is_identity else _terms_matrix(self._stab_terms) + if self._average_dirs and with_average: + avg = WeightedAverageProjection(self._mass_ops.derham, S, self._average_dirs) + S = S @ (IdentityOperator(space) - avg) + A = A + self._sigma * S return A def _build_approximation(self, bc): @@ -535,62 +625,110 @@ def _build_approximation(self, bc): stiff_dirs = _derivative_directions(self._d) # weights: common shape per direction and means per block (unit or separable average) - n_mid = len(self._mass_mid.weights) + derham = self._mass_operator.derham + n_mid = len(self._mid_terms[0][1].weights) n_dom = len(self._mass_operator.weights) if self._weights == "average": - mid_weights = [self._mass_mid.weights[r][r] for r in range(n_mid)] - shapes, mid_means = _separable_weight(mid_weights, self._mass_operator.derham) - _, dom_means = _separable_weight( - [self._mass_operator.weights[c][c] for c in range(n_dom)], self._mass_operator.derham - ) + shapes, mid_means = _separable_weight(_block_weights(self._mid_terms, n_mid), derham) + if self._stab_is_identity: + dom_means = [1.0] * n_dom + else: + _, dom_means = _separable_weight(_block_weights(self._stab_terms, n_dom), derham) else: shapes, mid_means, dom_means = [_ones_1d] * 3, [1.0] * n_mid, [1.0] * n_dom matrixblocks = [] solverblocks = [] + scalingblocks = [] for c, comp in enumerate(comps): coeff_space = femspace.coeff_space if is_scalar else femspace.coeff_space[c] - stiffness, mass, local_mass, local_stiffness = [], [], [], [] + stiffness, coeffs, mass, local_mass, local_stiffness = [], [], [], [], [] + identity, average = [], [] for d in range(3): basis = comp.spaces[d].basis M_d, domain_decomposition = _dense_mass_1d(self._mass_operator, basis, d, shapes[d]) + # unscaled 1d stiffness (defines the eigenbasis also for a zero coefficient) S_d = None if d in stiff_dirs[c]: assert basis == "B", "Only B-spline directions are differentiated." M_D, _ = _dense_mass_1d(self._mass_operator, "M", d, shapes[d]) D_d = _dense_derivative_1d(self._d, stiff_dirs[c][d], c) assert D_d.shape == (M_D.shape[0], M_d.shape[0]) - S_d = mid_means[stiff_dirs[c][d]] * (D_d.T @ M_D @ D_d) + S_d = D_d.T @ M_D @ D_d + coeffs.append(mid_means[stiff_dirs[c][d]] if S_d is not None else 1.0) if bc is not None and basis == "B": M_d = _apply_bc_dense(M_d, bc[d]) S_d = None if S_d is None else _apply_bc_dense(S_d, bc[d]) + # 1d factors of the stabilization terms (None for M_d) + identity.append(np.eye(M_d.shape[0])) + if d in self._average_dirs: + c_d = M_d @ np.ones(M_d.shape[0]) + average.append(np.outer(c_d, c_d) / c_d.sum()) + else: + average.append(None) + stiffness.append(S_d) mass.append(M_d) local_mass.append(_process_local_matrix_1d(xp.asarray(M_d), comp.coeff_space, d, domain_decomposition)) local_stiffness.append( None if S_d is None - else _process_local_matrix_1d(xp.asarray(S_d), comp.coeff_space, d, domain_decomposition) + else _process_local_matrix_1d( + xp.asarray(coeffs[d] * S_d), comp.coeff_space, d, domain_decomposition + ) ) + # stabilization: sigma_c * (M_1 x M_2 x M_3 - F_1 x F_2 x F_3), or the identity sigma_c = self._sigma * dom_means[c] - solverblocks.append(KroneckerSumSolver(coeff_space, stiffness, mass, sigma=sigma_c)) - matrixblocks.append(_kronecker_sum(coeff_space, local_stiffness, local_mass, sigma_c)) + terms = [] + if self._stab_is_identity and sigma_c != 0.0: + terms.append((sigma_c, identity)) + if self._average_dirs and sigma_c != 0.0: + factors = ( + [average[d] if average[d] is not None else identity[d] for d in range(3)] + if self._stab_is_identity + else average + ) + terms.append((-sigma_c, factors)) + sigma_mass = 0.0 if self._stab_is_identity else sigma_c + + solverblocks.append( + KroneckerSumSolver( + coeff_space, stiffness, mass, sigma=sigma_mass, stiffness_coeffs=coeffs, mass_terms=terms + ) + ) + dense_terms = [ + (coeff, [mass[d] if F is None else F for d, F in enumerate(factors)]) for coeff, factors in terms + ] + matrixblocks.append(_kronecker_sum(coeff_space, local_stiffness, local_mass, sigma_mass, dense_terms)) + # the average term is the last one; it is left out in the diagonal scaling + n_scaling = len(dense_terms) - (1 if self._average_dirs and sigma_c != 0.0 else 0) + scalingblocks.append( + _kronecker_sum(coeff_space, local_stiffness, local_mass, sigma_mass, dense_terms[:n_scaling]) + ) if is_scalar: + self._matrix_without_average = scalingblocks[0] return matrixblocks[0], solverblocks[0] + self._matrix_without_average = _block_diagonal(femspace.coeff_space, scalingblocks) return ( _block_diagonal(femspace.coeff_space, matrixblocks), _block_diagonal(femspace.coeff_space, solverblocks), ) def _core_diagonal(self): + # without the average: it couples all degrees of freedom along the averaged directions, + # which probing cannot resolve; its diagonal is small (O(1/n)) and left out on both sides + # (see _scaling_matrix), which keeps the preconditioner exact on the logical cube space = self._matrix.domain - return _diagonal_operator(space, _probe_diagonal(self.core_operator, space)) + return _diagonal_operator(space, _probe_diagonal(self._core(with_average=False), space)) + + def _scaling_matrix(self) -> LinearOperator: + return self._matrix_without_average def transpose(self, conjugate: bool = False) -> StiffnessPreconditioner: """The operator is symmetric, so is the preconditioner.""" @@ -610,7 +748,8 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: Local (rows owned by this process) diagonal of ``op``, one array per block of its domain. Supported are StencilMatrix, StencilDiagonalMatrix, KroneckerStencilMatrix (1d factors), - sums, scalings and BlockLinearOperators of these (missing diagonal blocks give zeros). + _DenseKroneckerProduct, sums, scalings and BlockLinearOperators of these (missing diagonal + blocks give zeros). """ if isinstance(op, BlockLinearOperator): diags = [] @@ -637,6 +776,8 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: for d in diags[1:]: out = xp.multiply.outer(out, d) return [out] + if isinstance(op, _DenseKroneckerProduct): + return [op.local_diagonal()] if isinstance(op, SumLinearOperator): parts = [_local_diagonal(a) for a in op.addends] return [sum(blocks) for blocks in zip(*parts)] @@ -1080,7 +1221,8 @@ def _separable_weight(weights: list, derham) -> tuple[list[xp.ndarray], list[flo Parameters ---------- weights : list - 3d weights (callable, array at the local quadrature points, or None for 1). + Per weight :math:`w_i`: a 3d weight (callable, array at the local quadrature points, or + None for 1), or a list of ``(coefficient, weight)`` terms (an empty list is a zero weight). derham : Derham Discrete de Rham sequence (quadrature grid). @@ -1089,10 +1231,11 @@ def _separable_weight(weights: list, derham) -> tuple[list[xp.ndarray], list[flo ------- shapes : list[xp.ndarray] :math:`\phi_d` at the global 1d quadrature points: the mean over the other directions, - divided by the mean, averaged over the weights (mean 1). + divided by the mean, averaged over the weights with non-zero mean (mean 1). means : list[float] - The means :math:`\bar w_i` of the weights over the logical cube. + The means :math:`\bar w_i` of the weights over the logical cube (0 for weights whose + mean vanishes relative to the largest one). Collective on ``derham.comm`` for array weights. """ @@ -1102,16 +1245,24 @@ def _separable_weight(weights: list, derham) -> tuple[list[xp.ndarray], list[flo profiles, means = [], [] for w in weights: + terms = w if isinstance(w, list) else [(1.0, w)] prof = [] for d in range(3): - p = _reduced_weight_1d(w, 0, d, derham, "average") - prof.append(xp.asarray(p(pts[d]) if callable(p) else p, dtype=float)) - mean = float(xp.sum(prof[0] * wts[0]) / xp.sum(wts[0])) - assert mean > 0.0, "The weights must have a positive mean." + p_d = xp.zeros(pts[d].size, dtype=float) + for coeff, w_t in terms: + p = _reduced_weight_1d(w_t, 0, d, derham, "average") + p_d = p_d + coeff * xp.asarray(p(pts[d]) if callable(p) else p, dtype=float) + prof.append(p_d) profiles.append(prof) - means.append(mean) + means.append(float(xp.sum(prof[0] * wts[0]) / xp.sum(wts[0]))) + + largest = max(abs(m) for m in means) + assert largest > 0.0, "The weights must not all vanish." + means = [0.0 if abs(m) <= 1e-12 * largest else m for m in means] + assert all(m >= 0.0 for m in means), f"The weights must have non-negative means, got {means}." - shapes = [sum(prof[d] / mean for prof, mean in zip(profiles, means)) / len(weights) for d in range(3)] + nonzero = [(prof, m) for prof, m in zip(profiles, means) if m > 0.0] + shapes = [sum(prof[d] / m for prof, m in nonzero) / len(nonzero) for d in range(3)] return shapes, means @@ -1153,10 +1304,12 @@ def _kronecker_sum( stiffness: list[StencilMatrix | None], mass: list[StencilMatrix], sigma: float, + mass_terms: list[tuple[float, list[StencilMatrix | None]]] | None = None, ) -> LinearOperator: """ - The operator :math:`\\sum_d M_1 \\otimes \\dots \\otimes S_d \\otimes \\dots \\otimes M_3 + \\sigma M_1 \\otimes M_2 \\otimes M_3` - as sum of KroneckerStencilMatrix (process-local 1d factors; directions with ``S_d = None`` have no term). + The operator :math:`\\sum_d M_1 \\otimes \\dots \\otimes S_d \\otimes \\dots \\otimes M_3 + \\sigma M_1 \\otimes M_2 \\otimes M_3 + + \\sum_t c_t F_{t,1} \\otimes F_{t,2} \\otimes F_{t,3}` as sum of KroneckerStencilMatrix (process-local + 1d factors; directions with ``S_d = None`` have no term; factors ``F = None`` stand for ``M_d``). """ terms = [] for d, S_d in enumerate(stiffness): @@ -1165,10 +1318,130 @@ def _kronecker_sum( terms.append(KroneckerStencilMatrix(space, space, *factors)) if sigma != 0.0: terms.append(KroneckerStencilMatrix(space, space, *[M_e.copy() for M_e in mass]) * sigma) + for coeff, factors in mass_terms or []: + terms.append(coeff * _DenseKroneckerProduct(space, factors)) assert terms, "The operator has no terms." return reduce(lambda a, b: a + b, terms) +class _DenseKroneckerProduct(LinearOperator): + """ + Kronecker product :math:`F_1 \\otimes F_2 \\otimes F_3` of dense global 1d matrices on a + StencilVectorSpace (e.g. the rank-one factors of an average), applied with a + KroneckerLinearSolver (also along distributed axes). + """ + + class _DenseApply(LinearSolver): + """'Solver' for KroneckerLinearSolver that applies a dense matrix to each right-hand side.""" + + def __init__(self, mat): + self._mat = mat + + @property + def space(self): + return xp.ndarray + + def transpose(self): + return _DenseKroneckerProduct._DenseApply(self._mat.T) + + def solve(self, rhs, out=None): + # rows of rhs are the vectors: out_i = mat @ rhs_i + result = rhs @ self._mat.T + if out is None: + return result + out[...] = result + return out + + def __init__(self, space: StencilVectorSpace, factors: list[np.ndarray]): + assert all(F.shape == (n, n) for F, n in zip(factors, space.npts)) + self._space = space + self._factors = [np.asarray(F, dtype=float) for F in factors] + self._kron = KroneckerLinearSolver( + space, space, [_DenseKroneckerProduct._DenseApply(xp.asarray(F)) for F in self._factors] + ) + + @property + def domain(self) -> StencilVectorSpace: + return self._space + + @property + def codomain(self) -> StencilVectorSpace: + return self._space + + @property + def dtype(self): + return self._space.dtype + + def local_diagonal(self) -> xp.ndarray: + """Diagonal on the rows owned by this process.""" + diags = [ + xp.asarray(np.diag(F)[s : e + 1]) for F, s, e in zip(self._factors, self._space.starts, self._space.ends) + ] + out = diags[0] + for d in diags[1:]: + out = xp.multiply.outer(out, d) + return out + + def dot(self, v: Vector, out: Vector | None = None) -> Vector: + assert v.space == self._space + return self._kron.solve(v, out=out) + + def transpose(self, conjugate: bool = False) -> _DenseKroneckerProduct: + return _DenseKroneckerProduct(self._space, [F.T for F in self._factors]) + + def tosparse(self): + raise NotImplementedError + + def toarray(self): + raise NotImplementedError + + +def _resolve_mass_operator(mass_ops: WeightedMassOperators, spec, default: LinearOperator) -> LinearOperator: + """The mass operator given by ``spec``: None (``default``), a name in ``mass_ops`` or an operator.""" + if spec is None: + return default + if isinstance(spec, str): + if not hasattr(mass_ops, spec): + raise ValueError(f"WeightedMassOperators has no mass operator {spec!r}.") + return getattr(mass_ops, spec) + assert isinstance(spec, LinearOperator), f"Expected a name or a LinearOperator, got {type(spec).__name__}." + return spec + + +def _mass_terms(op: LinearOperator) -> list[tuple[float, WeightedMassOperator]]: + """``op`` (a WeightedMassOperator, or a sum or multiple of such) as list of ``(coefficient, operator)``.""" + if isinstance(op, WeightedMassOperator): + return [(1.0, op)] + if isinstance(op, ScaledLinearOperator): + return [(op.scalar * c, t) for c, t in _mass_terms(op.operator)] + if isinstance(op, SumLinearOperator): + return [term for a in op.addends for term in _mass_terms(a)] + raise TypeError(f"Expected a WeightedMassOperator or a sum/multiple of such, got {type(op).__name__}.") + + +def _terms_matrix(terms: list[tuple[float, WeightedMassOperator]]) -> LinearOperator: + """Sum of the tensor-product matrices (without boundary operators) of the terms.""" + return reduce(lambda a, b: a + b, [t.matrix if c == 1.0 else c * t.matrix for c, t in terms]) + + +def _block_weights(terms: list[tuple[float, WeightedMassOperator]], n: int) -> list[list[tuple[float, object]]]: + """Diagonal blocks of the weights of a sum of mass operators: per block, the ``(coefficient, weight)`` terms + (a stored weight None is a zero block and is left out).""" + return [[(c, t.weights[r][r]) for c, t in terms if t.weights[r][r] is not None] for r in range(n)] + + +def _parse_directions(spec: str | None) -> tuple[int, ...]: + """``"eta2 eta3"`` (or ``"eta2,eta3"``) -> ``(1, 2)``; None or ``""`` -> ``()``.""" + if not spec: + return () + dirs = [] + for token in spec.replace(",", " ").split(): + if token not in ("eta1", "eta2", "eta3"): + raise ValueError(f"Unknown direction {token!r}, expected 'eta1', 'eta2' or 'eta3'.") + dirs.append(int(token[-1]) - 1) + return tuple(sorted(set(dirs))) + + def _probe_diagonal(op: LinearOperator, space: VectorSpace) -> list[xp.ndarray]: """ Local diagonal of ``op`` (one array per block of ``space``) by probing with colored unit vectors. diff --git a/src/struphy/feec/tests/test_stiffness_preconditioner.py b/src/struphy/feec/tests/test_stiffness_preconditioner.py index 5364dffb9..6d7534ab6 100644 --- a/src/struphy/feec/tests/test_stiffness_preconditioner.py +++ b/src/struphy/feec/tests/test_stiffness_preconditioner.py @@ -274,3 +274,98 @@ def test_stiffness_preconditioner_options(): StiffnessPreconditioner(mass_ops, "rot") with pytest.raises(AssertionError): StiffnessPreconditioner(mass_ops, "grad", sigma=-1.0) + + +@pytest.mark.parametrize( + "case", + ["mid_aniso", "mid_perp", "stab_scaled", "average_eta3", "average_eta2_eta3", "mid_perp_average_eta3", "stab_id"], +) +def test_stiffness_preconditioner_mid_stab_average(case): + """PCG for G^T M_mid G + sigma S (I - A) on the unit cube (Dirichlet in eta1): exact (few iterations) with + constant (anisotropic or perpendicular) mid, a scaled stabilization and averages over periodic directions; + the identity stabilization (approximated) needs fewer iterations than no preconditioner.""" + + from feectools.linalg.basic import IdentityOperator + from feectools.linalg.solvers import inverse + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedAverageProjection, WeightedMassOperators + from struphy.feec.preconditioner import StiffnessPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + derham = Derham( + TensorProductGrid(num_elements=[10, 8, 6]), + DerhamOptions(degree=[2, 3, 2], bcs=(("dirichlet", "dirichlet"), None, None)), + comm=MPI.COMM_WORLD, + ) + mass_ops = WeightedMassOperators(derham, domains.Cuboid()) + sigma = 10.0 + + mid, stab, average = None, None, None + if case in ("mid_aniso",): + mid = mass_ops.create_weighted_mass( + "Hcurl", "Hcurl", weights=([[2.0, 0, 0], [0, 1.0, 0], [0, 0, 0.5]], "sqrt_g"), assemble=True + ) + if case in ("mid_perp", "mid_perp_average_eta3"): + mid = mass_ops.create_weighted_mass( + "Hcurl", "Hcurl", weights=([[1.0, 0, 0], [0, 1.0, 0], [0, 0, 0.0]], "sqrt_g"), assemble=True + ) + if case == "stab_scaled": + stab = 3.0 * mass_ops.M0 + if case == "stab_id": + stab = "Id" + if case in ("average_eta3", "mid_perp_average_eta3"): + average = "eta3" + if case == "average_eta2_eta3": + average = "eta2 eta3" + + # the operator + S = mass_ops.M0 if stab is None else (IdentityOperator(derham.V0) if stab == "Id" else stab) + if average: + dirs = [int(t[-1]) - 1 for t in average.split()] + S = S @ (IdentityOperator(derham.V0) - WeightedAverageProjection(derham, S, dirs)) + A = derham.grad.T @ (mass_ops.M1 if mid is None else mid) @ derham.grad + sigma * S + + P = StiffnessPreconditioner(mass_ops, "grad", sigma=sigma, mid=mid, stab=stab, stab_average=average) + assert P.stab_average == average + + b = derham.boundary_ops["0"].dot(create_equal_random_arrays(derham.V0fem, seed=3, flattened=True)[1]) + niter = [] + for pc in (None, P): + inv = inverse(A, "pcg", pc=pc, tol=1e-10, maxiter=1000) if pc else inverse(A, "cg", tol=1e-10, maxiter=1000) + inv.dot(b) + assert inv._info["success"] + niter.append(inv._info["niter"]) + + if case == "stab_id": + assert niter[1] < niter[0] + else: + assert niter[1] <= 3 + + +def test_stiffness_preconditioner_mid_stab_options(): + """Names of mass operators, wrong domains, averages only for grad, unknown directions.""" + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import StiffnessPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + derham = Derham(TensorProductGrid(num_elements=[4, 4, 2]), DerhamOptions(degree=[1, 1, 1], bcs=(None, None, None))) + mass_ops = WeightedMassOperators(derham, domains.Cuboid()) + P = StiffnessPreconditioner(mass_ops, "grad", sigma=0.5, mid="M1", stab="M0", stab_average="eta2, eta3") + assert P.mid == "M1" and P.stab == "M0" and P.stab_average == "eta2, eta3" + with pytest.raises(ValueError): + StiffnessPreconditioner(mass_ops, "grad", sigma=0.5, mid="M1_does_not_exist") + with pytest.raises(AssertionError): + StiffnessPreconditioner(mass_ops, "grad", sigma=0.5, mid="M0") + with pytest.raises(ValueError): + StiffnessPreconditioner(mass_ops, "grad", sigma=0.5, stab_average="eta4") + with pytest.raises(NotImplementedError): + StiffnessPreconditioner(mass_ops, "curl", sigma=0.5, stab_average="eta3") diff --git a/src/struphy/feec/tests/test_weighted_average_projection.py b/src/struphy/feec/tests/test_weighted_average_projection.py new file mode 100644 index 000000000..ffa484cc4 --- /dev/null +++ b/src/struphy/feec/tests/test_weighted_average_projection.py @@ -0,0 +1,56 @@ +import pytest + + +@pytest.mark.parametrize("mapping", ["Cuboid", "HollowCylinder", "HollowTorus"]) +@pytest.mark.parametrize("stab", ["M0", "M0ad_withT"]) +def test_weighted_average_projection(mapping, stab): + """S (I - P) is symmetric, P is a projection, and P is the plain average (AverageOperator) if the + weight of S does not depend on the averaged directions (Cuboid, HollowCylinder averaged over eta3).""" + + from feectools.linalg.basic import IdentityOperator + from maybempi import MPI + + from struphy import domains, equils + from struphy.feec.mass import AverageOperator, WeightedAverageProjection, WeightedMassOperators + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + if mapping == "Cuboid": + domain, dirs = domains.Cuboid(), (1, 2) + elif mapping == "HollowCylinder": + domain, dirs = domains.HollowCylinder(a1=0.1, a2=1.0, Lz=4.0), (2,) + else: + domain, dirs = domains.HollowTorus(a1=0.1, a2=1.0, R0=4.0, tor_period=1), (1, 2) + equil = equils.HomogenSlab(B0x=0.0, B0y=0.0, B0z=1.0) + equil.domain = domain + + derham = Derham( + TensorProductGrid(num_elements=[8, 8, 6]), + DerhamOptions(degree=[2, 2, 2], bcs=(("dirichlet", "dirichlet"), None, None)), + comm=MPI.COMM_WORLD, + ) + mass_ops = WeightedMassOperators(derham, domain, eq_mhd=equil) + S0 = getattr(mass_ops, stab) + P = WeightedAverageProjection(derham, S0, dirs) + assert P.directions == dirs + S = S0 @ (IdentityOperator(derham.V0) - P) + + bc = derham.boundary_ops["0"] + x = bc.dot(create_equal_random_arrays(derham.V0fem, seed=1, flattened=True)[1]) + y = bc.dot(create_equal_random_arrays(derham.V0fem, seed=2, flattened=True)[1]) + + xSy, ySx = float(x.inner(S.dot(y))), float(y.inner(S.dot(x))) + assert abs(xSy - ySx) < 1e-12 * abs(xSy) + + Px = P.dot(x) + dP = P.dot(Px) - Px + assert float(dP.inner(dP)) ** 0.5 < 1e-12 * float(Px.inner(Px)) ** 0.5 + + if mapping != "HollowTorus": + avg = AverageOperator(derham, "H1", dirs[0]) + for d in dirs[1:]: + avg = avg @ AverageOperator(derham, "H1", d) + diff = avg.dot(x) - Px + assert float(diff.inner(diff)) ** 0.5 < 1e-12 * float(Px.inner(Px)) ** 0.5 diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index 8d0c6ce1a..079a4e230 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -10,7 +10,7 @@ from scope_profiler import ProfileManager from struphy.feec import preconditioner -from struphy.feec.mass import L2Projector, WeightedMassOperator +from struphy.feec.mass import L2Projector, WeightedAverageProjection, WeightedMassOperator from struphy.io.options import LiteralOptions, OptionsBase from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions, MultiGridPreconditioner from struphy.linear_algebra.solver import SolverParameters @@ -170,6 +170,13 @@ class Options(OptionsBase): - ``"M0ad"``: adiabatic-electron weighted 0-form mass operator. - ``"Id"``: identity operator. + stab_average : str, default=None + If not None, the stabilization becomes ``stab_mat @ (I - A)``, where ``A`` is the average + over the given logical directions, e.g. ``"eta3"`` or ``"eta2 eta3"``, weighted with + ``stab_mat`` (:class:`~struphy.feec.mass.WeightedAverageProjection`; e.g. the flux-surface + average with the volume element for ``"M0"``), such that the operator stays symmetric. + Not supported by ``precond="MultiGrid"``. + diffusion_mat : {"M1", "M1perp"} or WeightedMassOperator, default="M1" Diffusion metric in the bilinear form ``grad.T @ diffusion_mat @ grad``. @@ -195,10 +202,10 @@ class Options(OptionsBase): approximation of the inverse of the 0-form mass matrix ``M0`` (see :mod:`struphy.feec.preconditioner`). - ``"StiffnessPreconditioner"``: Kronecker approximation of the inverse of - ``grad.T @ M1 @ grad + sigma * M0`` - (:class:`~struphy.feec.preconditioner.StiffnessPreconditioner`), with - ``sigma = sigma_1`` (divided by dt if ``divide_by_dt=True``); any ``stab_mat`` is - approximated by ``M0`` (this keeps periodic or Neumann problems regular). + ``grad.T @ diffusion_mat @ grad + sigma * stab_mat (I - A)`` + (:class:`~struphy.feec.preconditioner.StiffnessPreconditioner` with ``mid=diffusion_mat``, + ``stab=stab_mat`` and ``stab_average``), with ``sigma = sigma_1`` (divided by dt if + ``divide_by_dt=True``). precond_params : dict, default=None Keyword arguments passed to the constructor of the mass-matrix or stiffness @@ -236,6 +243,7 @@ class Options(OptionsBase): sigma_3: float = 1.0 divide_by_dt: bool = False stab_mat: OptsStabMat = "M0" + stab_average: str = None diffusion_mat: OptsDiffusionMat = "M1" x0: StencilVector = None solver: LiteralOptions.OptsSymmSolver = "pcg" @@ -253,6 +261,9 @@ def __post_init__(self): check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) if self.precond is not None: assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'." + preconditioner._parse_directions(self.stab_average) + if self.stab_average and self.precond == "MultiGrid": + raise ValueError("precond='MultiGrid' does not support stab_average (no coarsening of the average).") # defaults if self.precond_params is None: @@ -363,6 +374,12 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: else: stab_mat = getattr(self.mass_ops, self.options.stab_mat) + # stabilization without the (stab_mat-weighted) average over the given directions: stab_mat @ (I - A) + average_dirs = preconditioner._parse_directions(self.options.stab_average) + if average_dirs: + avg = WeightedAverageProjection(self.derham, stab_mat, average_dirs) + stab_mat = stab_mat @ (IdentityOperator(stab_mat.domain, stab_mat.codomain) - avg) + if isinstance(self.options.diffusion_mat, str): diffusion_mat = getattr(self.mass_ops, self.options.diffusion_mat) else: @@ -421,14 +438,14 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: self._tmp_src = phi.space.zeros() def _build_stiffness_pc(self, sig_1: float) -> preconditioner.StiffnessPreconditioner: - """StiffnessPreconditioner for ``sig_1 * stab_mat + grad.T @ M1 @ grad``; the stabilization - is approximated by ``sig_1 * M0`` for any ``stab_mat`` (with ``sigma = 0``, the pseudo-inverse - would drop the constants, on which the operator is not singular for periodic or Neumann - boundary conditions).""" + """StiffnessPreconditioner for ``sig_1 * stab_mat (I - A) + grad.T @ diffusion_mat @ grad``.""" return preconditioner.StiffnessPreconditioner( self.mass_ops, "grad", sigma=sig_1, + mid=self.options.diffusion_mat, + stab=self.options.stab_mat, + stab_average=self.options.stab_average, **self.options.precond_params, ) diff --git a/src/struphy/propagators/poisson_adiabatic_gyrokinetic.py b/src/struphy/propagators/poisson_adiabatic_gyrokinetic.py index fbf6dab48..866ff4659 100644 --- a/src/struphy/propagators/poisson_adiabatic_gyrokinetic.py +++ b/src/struphy/propagators/poisson_adiabatic_gyrokinetic.py @@ -1,10 +1,8 @@ from dataclasses import dataclass from typing import Callable, Literal -from feectools.linalg.basic import IdentityOperator from feectools.linalg.stencil import StencilVector -from struphy.feec.mass import AverageOperator from struphy.io.options import LiteralOptions from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable, PICVariable, SPHVariable @@ -104,7 +102,12 @@ class Options: - ``"Id"``: identity operator. which_geometry: {"cylindrical", "toroidal"}, default="cylindrical" - Geometry of the problem, determines the meaning of `<\phi>` in the stabilization term. + Geometry of the problem, determines the meaning of `<\phi>` in the stabilization term: + the average over ``eta3`` (cylindrical) or over ``eta2`` and ``eta3`` (toroidal), i.e. + ``stab_average`` of :class:`ImplicitDiffusion` is ``"eta3"`` or ``"eta2 eta3"``. The average + is weighted with ``stab_mat`` (for equilibrium density and temperature constant on the + averaged surfaces, the flux-surface average with the volume element), which keeps the + operator symmetric (:class:`~struphy.feec.mass.WeightedAverageProjection`). diffusion_mat : {"M1", "M1perp", "M1gyro"}, defaults="M1gyro" Diffusion matrix. @@ -133,14 +136,13 @@ class Options: precond : LiteralOptions.OptsDiffusionPrecond, default=None Name of the preconditioner, see :class:`ImplicitDiffusion`. Requires ``solver="pcg"``. - ``"MultiGrid"`` is not supported: multigrid would be built for the stabilization without - the average, which is applied after ``ImplicitDiffusion.allocate()``. + ``"StiffnessPreconditioner"`` approximates the operator with ``diffusion_mat``, ``stab_mat`` + and the average; ``"MultiGrid"`` is not supported (no coarsening of the average). precond_params : dict, default=None Keyword arguments passed to the constructor of the mass-matrix or stiffness preconditioner, see :class:`ImplicitDiffusion`. - solver_params : SolverParameters, default=None Iterative-solver controls (for example ``tol``, ``maxiter``, ``verbose``, ``info``, ``recycle``). @@ -177,11 +179,12 @@ def __post_init__(self): check_option(self.stab_mat, self.OptsStabMat) check_option(self.solver, LiteralOptions.OptsSymmSolver) check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) + check_option(self.which_geometry, self.OptsGeometry) if self.precond is not None: assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'." if self.precond == "MultiGrid": raise ValueError( - "precond='MultiGrid' is not supported (it would not see the average in the stabilization)." + "precond='MultiGrid' is not supported (no coarsening of the average in the stabilization)." ) # defaults @@ -190,6 +193,11 @@ def __post_init__(self): if self.solver_params is None: self.solver_params = SolverParameters() + @property + def stab_average(self) -> str: + """Directions averaged in the stabilization, see ``which_geometry``.""" + return {"cylindrical": "eta3", "toroidal": "eta2 eta3"}[self.which_geometry] + def allocate(self): epsilon = self.epsilon Z = self.Z @@ -197,18 +205,8 @@ def allocate(self): self.options.sigma_2 = 0.0 self.options.sigma_3 = 1 / epsilon self.options.divide_by_dt = False + # the stabilization stab_mat @ (I - A) (A from stab_average) is built in ImplicitDiffusion super().allocate() - if self.options.which_geometry == "cylindrical": - average_mat = AverageOperator(self.derham, "H1", 2) - self._stab_mat = self._stab_mat @ ( - IdentityOperator(self._stab_mat.domain, self._stab_mat.codomain) - average_mat - ) - elif self.options.which_geometry == "toroidal": - average_mat1 = AverageOperator(self.derham, "H1", 1) - average_mat2 = AverageOperator(self.derham, "H1", 2) - self._stab_mat = self._stab_mat @ ( - IdentityOperator(self._stab_mat.domain, self._stab_mat.codomain) - (average_mat1 @ average_mat2) - ) @property def options(self) -> Options: diff --git a/src/struphy/propagators/poisson_solve.py b/src/struphy/propagators/poisson_solve.py index 9dd8998cc..324e08d2e 100644 --- a/src/struphy/propagators/poisson_solve.py +++ b/src/struphy/propagators/poisson_solve.py @@ -4,6 +4,7 @@ from feectools.linalg.stencil import StencilVector +from struphy.feec.preconditioner import _parse_directions from struphy.io.options import LiteralOptions, OptionsBase from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions from struphy.linear_algebra.solver import SolverParameters @@ -52,6 +53,11 @@ class Options(OptionsBase): - ``"M0ad"``: adiabatic-electron weighted 0-form mass operator. - ``"Id"``: identity operator. + stab_average : str, default=None + If not None, the stabilization becomes ``stab_mat @ (I - A)``, where ``A`` is the + ``stab_mat``-weighted average over the given logical directions, e.g. ``"eta3"`` or + ``"eta2 eta3"``, see :class:`ImplicitDiffusion`. + diffusion_mat : {"M1", "M1perp", "M1gyro"}, defaults="M1" Diffusion matrix. @@ -104,6 +110,7 @@ class Options(OptionsBase): # propagator options stab_eps: float = 0.0 stab_mat: OptsStabMat = "M0" + stab_average: str = None diffusion_mat: OptsDiffusionMat = "M1" x0: StencilVector = None solver: LiteralOptions.OptsSymmSolver = "pcg" @@ -121,6 +128,9 @@ def __post_init__(self): check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) if self.precond is not None: assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'." + _parse_directions(self.stab_average) + if self.stab_average and self.precond == "MultiGrid": + raise ValueError("precond='MultiGrid' does not support stab_average (no coarsening of the average).") # defaults if self.precond_params is None: diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index 3e7d7414f..706405b2b 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -808,6 +808,58 @@ def rho_physical(x, y, z): plt.show() +@pytest.mark.parametrize("which_geometry", ["cylindrical", "toroidal"]) +def test_poisson_adiabatic_gyrokinetic_stiffness_preconditioner(which_geometry): + """PoissonAdiabaticGyrokinetic (M1gyro, M0ad_withT and the average in the stabilization): the stiffness + preconditioner gives the solution without preconditioner, in fewer iterations; multigrid is rejected.""" + from struphy.propagators.poisson_adiabatic_gyrokinetic import PoissonAdiabaticGyrokinetic + + if which_geometry == "cylindrical": + domain = domains.HollowCylinder(a1=0.1, a2=1.0, Lz=4.0) + else: + domain = domains.HollowTorus(a1=0.1, a2=1.0, R0=4.0, tor_period=1) + equil = equils.HomogenSlab(B0x=0.0, B0y=0.0, B0z=1.0) + equil.domain = domain + + derham = Derham( + TensorProductGrid(num_elements=[12, 12, 8]), + DerhamOptions(degree=[2, 2, 2], bcs=(("dirichlet", "dirichlet"), None, None)), + comm=comm, + ) + mass_ops = WeightedMassOperators(derham, domain, eq_mhd=equil) + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + def rho(e1, e2, e3): + return xp.sin(xp.pi * e1) * (1.0 + xp.cos(2 * xp.pi * e2) * xp.sin(2 * xp.pi * e3)) + + phis, infos = [], [] + for precond in [None, "StiffnessPreconditioner"]: + phi = FEECVariable(space="H1") + phi.allocate(derham=derham, domain=domain) + prop = PoissonAdiabaticGyrokinetic(rho=rho, epsilon=0.1) + prop.variables.phi = phi + prop.options = prop.Options( + which_geometry=which_geometry, + precond=precond, + solver_params=SolverParameters(tol=1e-11, maxiter=3000, recycle=False), + ) + prop.allocate() + prop(1.0) + phis.append(phi.spline.vector.toarray()) + infos.append(prop._solver._info) + + if comm.Get_size() > 1: + phis = [comm.allreduce(p, op=MPI.SUM) for p in phis] + assert all(info["success"] for info in infos) + assert xp.max(xp.abs(phis[0] - phis[1])) < 1e-7 * xp.max(xp.abs(phis[0])) + assert infos[1]["niter"] < infos[0]["niter"] / 2 + + with pytest.raises(ValueError): + PoissonAdiabaticGyrokinetic.Options(precond="MultiGrid") + + if __name__ == "__main__": direction = 0 bc_type = "dirichlet" From f50e096be8eb92f4fe3176d39f5e399f7048a089 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 9 Oct 2026 16:29:46 +0200 Subject: [PATCH 2/4] Bump feectools: pseudo-inverse cutoff of KroneckerSumSolver for signed terms (feectools#102 review) Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index c2593636d..ff6442cda 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit c2593636dade2386881c8f9e0daaada9acc11f50 +Subproject commit ff6442cda515ca920669cd63f784c5736e141ed2 From 4a215c14c24bd417464afeb9d2777a5915b08b04 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 9 Oct 2026 18:26:33 +0200 Subject: [PATCH 3/4] update feectools and pyproject.toml --- feectools | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/feectools b/feectools index ff6442cda..d40c05fda 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit ff6442cda515ca920669cd63f784c5736e141ed2 +Subproject commit d40c05fda4a357504871b282acc76a949c92347b diff --git a/pyproject.toml b/pyproject.toml index d14c35a11..c75808225 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "cunumpy >= 0.6.1, <= 0.6.1", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", - "feectools>=0.7.0, <=0.7.0", + "feectools>=0.8.0, <=0.8.0", "scipy<=1.18.0", "h5py<=3.16.0", "h5netcdf<=1.8.1", From ccc5ef9c24dc3bd28be0037cfbe7be5de02874fd Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 9 Oct 2026 18:48:34 +0200 Subject: [PATCH 4/4] Address #740 review: Euclidean average factors for stab='Id', global average in WeightedAverageProjection - StiffnessPreconditioner: with stab='Id' the averaged factors are 1 1^T / n (the projection is Euclidean), not the M_d-weighted ones; exact on the cube again (2 instead of 44 PCG iterations for stab_average='eta3'). - WeightedAverageProjection: allow averaging over all three directions (global weighted mean, 1x1 reduced system). - Tests: stab_id with eta3 / eta2 eta3 averages, global average (preconditioner and projection). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/mass.py | 7 ++-- src/struphy/feec/preconditioner.py | 10 +++-- .../tests/test_stiffness_preconditioner.py | 30 ++++++++++---- .../tests/test_weighted_average_projection.py | 39 +++++++++++++++++++ 4 files changed, 71 insertions(+), 15 deletions(-) diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 0b2fc656a..dea89f68b 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -3542,7 +3542,9 @@ class WeightedAverageProjection(LinearOperator): \mathbb S (\mathbb I - \Pi_{\mathbb S}) = \mathbb S - \mathbb S \mathbb E \, (\mathbb E^\top \mathbb S \, \mathbb E)^{-1} \, \mathbb E^\top \mathbb S is symmetric. If the weight of :math:`\mathbb S` does not depend on the averaged directions, - :math:`\Pi_{\mathbb S}` is the plain average of :class:`AverageOperator`. + :math:`\Pi_{\mathbb S}` is the plain average of :class:`AverageOperator`. If all three directions + are averaged, :math:`\mathbb E` is the vector of ones and the reduced matrix is the scalar + :math:`\mathbb 1^\top \mathbb S \, \mathbb 1` (the weighted mean over the whole domain). The reduced matrix :math:`\mathbb E^\top \mathbb S \, \mathbb E` (on the remaining directions) is assembled once by probing :math:`\mathbb S` with colored vectors and factorized on every process; @@ -3559,7 +3561,7 @@ class WeightedAverageProjection(LinearOperator): conditions, whose reduced rows are replaced by the identity). directions : tuple[int] - The averaged directions (0, 1 and/or 2). + The averaged directions (0, 1 and/or 2; all three for the global mean). """ def __init__(self, derham: Derham, S: LinearOperator, directions: tuple[int, ...]): @@ -3572,7 +3574,6 @@ def __init__(self, derham: Derham, S: LinearOperator, directions: tuple[int, ... directions = tuple(sorted(set(directions))) assert directions and all(d in (0, 1, 2) for d in directions) kept = tuple(d for d in range(3) if d not in directions) - assert kept, "At least one direction must not be averaged." self._derham = derham self._S = S diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 6ef305d68..d64838dc2 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -435,7 +435,8 @@ class StiffnessPreconditioner(KroneckerPreconditioner): directions (:math:`F_d = M_d` otherwise). In directions with :math:`S_d \mathbb 1 = 0` (periodic or without essential boundary conditions) the constants are an eigenvector of the fast diagonalization and this term is exact in its eigenbasis (it removes the averaged modes). - The identity (``stab="Id"``) is replaced by its diagonal in that eigenbasis. + The identity (``stab="Id"``) is replaced by its diagonal in that eigenbasis; its average is + Euclidean, i.e. :math:`M_d` is replaced by the identity and :math:`c_d = \mathbb 1`. Polar splines are not supported yet. @@ -477,7 +478,7 @@ class StiffnessPreconditioner(KroneckerPreconditioner): stab_average : str | None Gradient only: the directions averaged in the stabilization :math:`\mathbb S (\mathbb I - \mathbb A)`, - e.g. ``"eta3"`` or ``"eta2 eta3"``; None (default) for no averaging. + e.g. ``"eta3"``, ``"eta2 eta3"`` or ``"eta1 eta2 eta3"`` (global mean); None (default) for no averaging. """ _FORMS = {"grad": 0, "curl": 1, "div": 2} @@ -663,10 +664,11 @@ def _build_approximation(self, bc): M_d = _apply_bc_dense(M_d, bc[d]) S_d = None if S_d is None else _apply_bc_dense(S_d, bc[d]) - # 1d factors of the stabilization terms (None for M_d) + # 1d factors of the stabilization terms (None for M_d); the average is weighted with + # M_d, or Euclidean (1 1^T / n) for the identity stabilization identity.append(np.eye(M_d.shape[0])) if d in self._average_dirs: - c_d = M_d @ np.ones(M_d.shape[0]) + c_d = np.ones(M_d.shape[0]) if self._stab_is_identity else M_d @ np.ones(M_d.shape[0]) average.append(np.outer(c_d, c_d) / c_d.sum()) else: average.append(None) diff --git a/src/struphy/feec/tests/test_stiffness_preconditioner.py b/src/struphy/feec/tests/test_stiffness_preconditioner.py index 6d7534ab6..4bc498197 100644 --- a/src/struphy/feec/tests/test_stiffness_preconditioner.py +++ b/src/struphy/feec/tests/test_stiffness_preconditioner.py @@ -278,12 +278,24 @@ def test_stiffness_preconditioner_options(): @pytest.mark.parametrize( "case", - ["mid_aniso", "mid_perp", "stab_scaled", "average_eta3", "average_eta2_eta3", "mid_perp_average_eta3", "stab_id"], + [ + "mid_aniso", + "mid_perp", + "stab_scaled", + "average_eta3", + "average_eta2_eta3", + "mid_perp_average_eta3", + "stab_id", + "stab_id_average_eta3", + "stab_id_average_eta2_eta3", + "average_global", + ], ) def test_stiffness_preconditioner_mid_stab_average(case): """PCG for G^T M_mid G + sigma S (I - A) on the unit cube (Dirichlet in eta1): exact (few iterations) with - constant (anisotropic or perpendicular) mid, a scaled stabilization and averages over periodic directions; - the identity stabilization (approximated) needs fewer iterations than no preconditioner.""" + constant (anisotropic or perpendicular) mid, a scaled or identity stabilization and averages over periodic + directions (Euclidean for the identity); the global average (not exact in the Dirichlet direction eta1) needs + fewer iterations than no preconditioner.""" from feectools.linalg.basic import IdentityOperator from feectools.linalg.solvers import inverse @@ -316,12 +328,14 @@ def test_stiffness_preconditioner_mid_stab_average(case): ) if case == "stab_scaled": stab = 3.0 * mass_ops.M0 - if case == "stab_id": + if case.startswith("stab_id"): stab = "Id" - if case in ("average_eta3", "mid_perp_average_eta3"): + if case in ("average_eta3", "mid_perp_average_eta3", "stab_id_average_eta3"): average = "eta3" - if case == "average_eta2_eta3": + if case in ("average_eta2_eta3", "stab_id_average_eta2_eta3"): average = "eta2 eta3" + if case == "average_global": + average = "eta1 eta2 eta3" # the operator S = mass_ops.M0 if stab is None else (IdentityOperator(derham.V0) if stab == "Id" else stab) @@ -341,8 +355,8 @@ def test_stiffness_preconditioner_mid_stab_average(case): assert inv._info["success"] niter.append(inv._info["niter"]) - if case == "stab_id": - assert niter[1] < niter[0] + if case == "average_global": + assert niter[1] < niter[0] / 4 else: assert niter[1] <= 3 diff --git a/src/struphy/feec/tests/test_weighted_average_projection.py b/src/struphy/feec/tests/test_weighted_average_projection.py index ffa484cc4..2d67df4e8 100644 --- a/src/struphy/feec/tests/test_weighted_average_projection.py +++ b/src/struphy/feec/tests/test_weighted_average_projection.py @@ -54,3 +54,42 @@ def test_weighted_average_projection(mapping, stab): avg = avg @ AverageOperator(derham, "H1", d) diff = avg.dot(x) - Px assert float(diff.inner(diff)) ** 0.5 < 1e-12 * float(Px.inner(Px)) ** 0.5 + + +def test_weighted_average_projection_global(): + """Average over all three directions: P x is the constant (1^T S x) / (1^T S 1), S (I - P) is symmetric.""" + + from feectools.linalg.basic import IdentityOperator + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedAverageProjection, WeightedMassOperators + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + domain = domains.HollowTorus(a1=0.1, a2=1.0, R0=4.0, tor_period=1) + derham = Derham( + TensorProductGrid(num_elements=[8, 8, 6]), + DerhamOptions(degree=[2, 2, 2], bcs=(("dirichlet", "dirichlet"), None, None)), + comm=MPI.COMM_WORLD, + ) + mass_ops = WeightedMassOperators(derham, domain) + S0 = mass_ops.M0 + P = WeightedAverageProjection(derham, S0, (0, 1, 2)) + assert P.directions == (0, 1, 2) + S = S0 @ (IdentityOperator(derham.V0) - P) + + x = create_equal_random_arrays(derham.V0fem, seed=1, flattened=True)[1] + y = create_equal_random_arrays(derham.V0fem, seed=2, flattened=True)[1] + xSy, ySx = float(x.inner(S.dot(y))), float(y.inner(S.dot(x))) + assert abs(xSy - ySx) < 1e-12 * abs(xSy) + + ones = derham.V0.zeros() + ones[:] = 1.0 + mean = float(ones.inner(S0.dot(x))) / float(ones.inner(S0.dot(ones))) + Px = P.dot(x) + diff = Px - mean * ones + assert float(diff.inner(diff)) ** 0.5 < 1e-12 * float(Px.inner(Px)) ** 0.5 + assert abs(float(ones.inner(S.dot(x)))) < 1e-12 * abs(float(ones.inner(S0.dot(x))))