diff --git a/examples/ShearAlfven/itpa_tae_benchmark/params_ShearAlfven.py b/examples/ShearAlfven/itpa_tae_benchmark/params_ShearAlfven.py index efc6d07e6..b113e5b7c 100644 --- a/examples/ShearAlfven/itpa_tae_benchmark/params_ShearAlfven.py +++ b/examples/ShearAlfven/itpa_tae_benchmark/params_ShearAlfven.py @@ -182,7 +182,7 @@ derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(("dirichlet", "dirichlet"), None, None)) # Profiling options -profiling_opts = ProfilingOptions(use_line_profiler=True) +profiling_opts = ProfilingOptions(use_line_profiler=False) # Simulation object sim = Simulation( @@ -238,4 +238,4 @@ model.mhd.velocity.add_perturbation(perturbation=perturbation_poloidal) if __name__ == "__main__": - sim.run() + sim.run(profiling_activated=True) diff --git a/pyproject.toml b/pyproject.toml index 3a0f57031..fcc8950b7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,7 +50,7 @@ dependencies = [ "pytest-testmon<=2.2.0", "ruff==0.16.10", "line_profiler<=5.0.2", - "scope-profiler>=0.5.0, <=0.7.2", + "scope-profiler>=0.5.0, <=0.7.3", ] [project.license] diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 143416980..1f135b8d2 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -28,6 +28,7 @@ from line_profiler import profile from maybempi import MPI, SerialComm from scipy import sparse +from scope_profiler import ProfileManager from struphy.feec.kronecker_block import KroneckerBlockOperator from struphy.feec.linear_operators import BoundaryOperator @@ -112,12 +113,14 @@ def __init__( self._bc = _boundary_conditions(mass_operator, apply_bc) # approximation of the core operator and its exact inverse - self._matrix, self._solver = self._build_approximation(self._bc) + with ProfileManager.profile_region("build_approximation", functions=[self._build_approximation]): + self._matrix, self._solver = self._build_approximation(self._bc) # diagonal scaling S (a sequence of diagonal operators, applied in this order before the # solver and in reverse order after it) - self._scaling = self._build_scaling() if diagonal_scaling else () - self._tmp_core = tuple(self._matrix.codomain.zeros() for _ in range(2)) if self._scaling else () + with ProfileManager.profile_region("build_scaling", functions=[self._build_scaling]): + self._scaling = self._build_scaling() if diagonal_scaling else () + self._tmp_core = tuple(self._matrix.codomain.zeros() for _ in range(2)) if self._scaling else () # operator to be inverted (with boundary operators B, E if apply_bc), needed in solve self._M, self._mass_index, self._tmp_vectors = _operator_to_invert(mass_operator, self._bc is not None) @@ -194,7 +197,8 @@ 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()) + with ProfileManager.profile_region("local_diagonal of core", functions=[_local_diagonal, self._core_diagonal]): + D = _local_diagonal(self._core_diagonal()) 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),) @@ -738,36 +742,37 @@ def __init__( ops = list(derivative) if isinstance(derivative, (list, tuple)) else [derivative] self._general = all(isinstance(op, KroneckerBlockOperator) for op in ops) if self._general: - # general operators L_i: no diagonal scaling, kernel correction or averaged stabilization - if stab_average is not None: - raise NotImplementedError("stab_average is not supported for a general operator.") - diagonal_scaling = False - kernel_correction = False - assert len({op.domain_space for op in ops}) == 1, "The operators L_i must have the same domain." - mids = list(mid) if isinstance(mid, (list, tuple)) else [mid] * len(ops) - assert len(mids) == len(ops), "Give one mid per operator L_i." - form_dom = derham.space_to_form[ops[0].domain_space] - form_mid = derham.space_to_form[ops[0].codomain_space] - # pairs (L_i, terms of mid_i) - self._general_pairs = [ - ( - op, - _mass_terms( - _resolve_mass_operator( - mass_ops, m, getattr(mass_ops, f"M{derham.space_to_form[op.codomain_space]}") - ) - ), - ) - for op, m in zip(ops, mids) - ] - for op, mid_terms in self._general_pairs: - ref = getattr(mass_ops, f"M{derham.space_to_form[op.codomain_space]}") - for _, t in mid_terms: - assert t.domain == ref.domain and t.codomain == ref.codomain, ( - f"Each mid must have the domain and codomain of the mass operator of the codomain " - f"({op.codomain_space}) of its operator." + with ProfileManager.profile_region("StiffnessPreconditioner: general L_i"): + # general operators L_i: no diagonal scaling, kernel correction or averaged stabilization + if stab_average is not None: + raise NotImplementedError("stab_average is not supported for a general operator.") + diagonal_scaling = False + kernel_correction = False + assert len({op.domain_space for op in ops}) == 1, "The operators L_i must have the same domain." + mids = list(mid) if isinstance(mid, (list, tuple)) else [mid] * len(ops) + assert len(mids) == len(ops), "Give one mid per operator L_i." + form_dom = derham.space_to_form[ops[0].domain_space] + form_mid = derham.space_to_form[ops[0].codomain_space] + # pairs (L_i, terms of mid_i) + self._general_pairs = [ + ( + op, + _mass_terms( + _resolve_mass_operator( + mass_ops, m, getattr(mass_ops, f"M{derham.space_to_form[op.codomain_space]}") + ) + ), ) - mid = mids[0] + for op, m in zip(ops, mids) + ] + for op, mid_terms in self._general_pairs: + ref = getattr(mass_ops, f"M{derham.space_to_form[op.codomain_space]}") + for _, t in mid_terms: + assert t.domain == ref.domain and t.codomain == ref.codomain, ( + f"Each mid must have the domain and codomain of the mass operator of the codomain " + f"({op.codomain_space}) of its operator." + ) + mid = mids[0] elif len(ops) > 1 or not isinstance(derivative, str): raise TypeError("derivative must be 'grad', 'curl', 'div', a KroneckerBlockOperator or a list of these.") else: @@ -789,18 +794,25 @@ def __init__( self._weights = weights # matrix in the middle and stabilization, as sums of weighted mass operators - mass_mid_default = getattr(mass_ops, f"M{form_mid}") - 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{form_dom}") - 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")): + with ProfileManager.profile_region("StiffnessPreconditioner: resolve mass operators"): + # the default M{form_mid} is resolved (and assembled) only if needed + self._mid = mid + self._mid_terms = _mass_terms( + _resolve_mass_operator(mass_ops, f"M{form_mid}" if mid is None else mid, None) + ) + self._stab = stab + self._stab_is_identity = isinstance(stab, str) and stab == "Id" + mass_k = getattr(mass_ops, f"M{form_dom}") + self._stab_terms = ( + None if self._stab_is_identity else _mass_terms(_resolve_mass_operator(mass_ops, stab, mass_k)) + ) + mid_space = derham.coeff_spaces[form_mid] + for terms, ref_dom, ref_cod, name in ( + (self._mid_terms, mid_space, mid_space, "mid"), + (self._stab_terms or [], mass_k.domain, mass_k.codomain, "stab"), + ): for _, op in terms: - assert op.domain == ref.domain and op.codomain == ref.codomain, ( + assert op.domain == ref_dom and op.codomain == ref_cod, ( f"{name} must have the domain and codomain of M{form_mid if name == 'mid' else form_dom}." ) @@ -813,25 +825,29 @@ def __init__( # tensor-product derivative without boundary operators (no polar splines), or the general operator self._d = derivative if self._general else getattr(derham, f"{derivative}_bcfree") - super().__init__(mass_k, apply_bc=apply_bc, diagonal_scaling=diagonal_scaling) + with ProfileManager.profile_region( + "KroneckerPreconditioner init", functions=[KroneckerPreconditioner.__init__] + ): + 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, # 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" - d_prev = getattr(derham, prev) - P_prev = StiffnessPreconditioner( - mass_ops, - prev, - apply_bc=apply_bc, - 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() + with ProfileManager.profile_region("StiffnessPreconditioner: kernel correction"): + prev = "grad" if derivative == "curl" else "curl" + d_prev = getattr(derham, prev) + P_prev = StiffnessPreconditioner( + mass_ops, + prev, + apply_bc=apply_bc, + 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() @property def derivative(self) -> str | KroneckerBlockOperator: @@ -1116,10 +1132,20 @@ def _build_general_approximation(self, bc): 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 + # which neither the stencil data nor probing 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(with_average=False), space)) + diags = _stiffness_local_diagonal(self._d, self._mid_terms) + if diags is None: + # mass operators not given by stencil matrices: probe the core operator + return _diagonal_operator(space, _probe_diagonal(self._core(with_average=False), space)) + if self._sigma != 0.0: + if self._stab_is_identity: + stab = [xp.ones_like(d) for d in diags] + else: + stab = _local_diagonal(_terms_matrix(self._stab_terms)) + diags = [d + self._sigma * sd for d, sd in zip(diags, stab)] + return _diagonal_operator(space, diags) def _scaling_matrix(self) -> LinearOperator: return self._matrix_without_average @@ -1638,29 +1664,34 @@ def __init__( # smoother and gradient correction if smoother == "block_jacobi": - # block Jacobi plus the correction sigma^{-1} G P_0 G^T on the gradients - self._S = StiffnessPreconditioner( - mass_ops, "curl", sigma=sigma, mid=mid, stab=stab, kernel_correction=True, **kron_kwargs - ) - self._grad_correction = None + with ProfileManager.profile_region("block_jacobi smoother", functions=[StiffnessPreconditioner.__init__]): + # block Jacobi plus the correction sigma^{-1} G P_0 G^T on the gradients + self._S = StiffnessPreconditioner( + mass_ops, "curl", sigma=sigma, mid=mid, stab=stab, kernel_correction=True, **kron_kwargs + ) + self._grad_correction = None else: - assert A is not None, "smoother='jacobi' needs the operator A." - assert A.domain is self._space and A.codomain is self._space - self._S = inverse_diagonal(DiagonalComputer(derham.degree)(A)) - P_0 = StiffnessPreconditioner(mass_ops, "grad", mid=None if stab == "Id" else stab, **kron_kwargs) - self._grad_correction = (1.0 / sigma) * (derham.grad @ P_0 @ derham.grad.T) + with ProfileManager.profile_region("jacobi smoother", functions=[StiffnessPreconditioner.__init__]): + assert A is not None, "smoother='jacobi' needs the operator A." + assert A.domain is self._space and A.codomain is self._space + self._S = inverse_diagonal(DiagonalComputer(derham.degree)(A)) + P_0 = StiffnessPreconditioner(mass_ops, "grad", mid=None if stab == "Id" else stab, **kron_kwargs) + self._grad_correction = (1.0 / sigma) * (derham.grad @ P_0 @ derham.grad.T) # vector auxiliary space: Pi P_v Pi^T self._Pi = None if aux_weight != 0.0: - identity = [[(lambda e1, e2, e3: 1.0 + 0.0 * e1) if m == n else None for n in range(3)] for m in range(3)] - basis_ops = BasisProjectionOperators(derham, mass_ops.domain) - self._Pi = basis_ops.create_basis_op(identity, "H1vec", "Hcurl", assemble=True, name="Pi_v1") - self._PiT = self._Pi.T - self._P_v = StiffnessPreconditioner(mass_ops, "grad", sigma=sigma, **kron_kwargs) - self._aux = derham.coeff_spaces["v"] - self._tmp_v = (self._aux.zeros(), self._aux.zeros()) - self._tmp_0 = (self._P_v.domain.zeros(), self._P_v.codomain.zeros()) + with ProfileManager.profile_region("auxiliary_space correction"): + identity = [ + [(lambda e1, e2, e3: 1.0 + 0.0 * e1) if m == n else None for n in range(3)] for m in range(3) + ] + basis_ops = BasisProjectionOperators(derham, mass_ops.domain) + self._Pi = basis_ops.create_basis_op(identity, "H1vec", "Hcurl", assemble=True, name="Pi_v1") + self._PiT = self._Pi.T + self._P_v = StiffnessPreconditioner(mass_ops, "grad", sigma=sigma, **kron_kwargs) + self._aux = derham.coeff_spaces["v"] + self._tmp_v = (self._aux.zeros(), self._aux.zeros()) + self._tmp_0 = (self._P_v.domain.zeros(), self._P_v.codomain.zeros()) self._tmp = self._space.zeros() @@ -2617,6 +2648,129 @@ def _probe_diagonal(op: LinearOperator, space: VectorSpace) -> list[xp.ndarray]: return diags +def _stiffness_local_diagonal( + derivative: LinearOperator, mid_terms: list[tuple[float, WeightedMassOperator]] +) -> list[xp.ndarray] | None: + """ + Local diagonal of ``d^T M d`` (one array per block of the domain of ``d``), ``M`` the sum of ``mid_terms`` and + ``d = derivative`` a tensor-product derivative without boundary operators, computed exactly from the stencil + data of ``M`` instead of probing (no operator applications). + + Each block ``(a, c)`` of ``d`` is a Kronecker product of 1d factors, so its column ``i`` has a few non-zeros + ``D^a[i + o, i]`` (``o`` in {-1, 0} in the differentiated direction, ``o = 0`` in the others), and + + diag(d^T M d)_i = sum_{a, b} sum_{o, o'} D^a[i + o, i] D^b[i + o', i] M^{ab}[i + o, i + o']. + + The band entries ``M^{ab}[r, r + o' - o]`` (including the entries that wrap onto it in a periodic direction with + few points) are taken at the owned rows ``r`` and read at ``r = i + o`` through the ghost regions of a vector, + since the row can be owned by the neighbouring process; this is collective. Returns None if a block of ``M`` is + not a StencilMatrix (e.g. matrix-free mass operators). + """ + + def block(op, r, c): + if isinstance(op, BlockLinearOperator): + return op[r, c] + return op if r == c == 0 else None + + n_rows = derivative.n_block_rows if isinstance(derivative, BlockLinearOperator) else 1 + n_cols = derivative.n_block_cols if isinstance(derivative, BlockLinearOperator) else 1 + for _, t in mid_terms: + for a, b in itertools.product(range(n_rows), repeat=2): + M = block(t.matrix, a, b) + if M is not None and not isinstance(M, StencilMatrix): + return None + + dom = derivative.domain + dom_spaces = dom.spaces if isinstance(dom, BlockVectorSpace) else (dom,) + assert len(dom_spaces) == n_cols + + # 1d factors of the blocks (a, c): per direction, {o: D[i + o, i] at the owned indices i of the domain} + factors = {} + for c, V in enumerate(dom_spaces): + for a in range(n_rows): + op = block(derivative, a, c) + if op is None: + continue + assert isinstance(op, DirectionalDerivativeOperator) and not op.transposed + per_dir = [] + for D_1d, s, e, periodic in zip(op.tokronstencil().mats, V.starts, V.ends, V.periods): + # the sign of a negated block is included in its 1d factor + D = np.asarray(xp.to_numpy(D_1d.toarray()), dtype=float) + n_out, n_in = D.shape + offsets = sorted({int(j - i) for j, i in zip(*np.nonzero(D))}) + if periodic: + offsets = sorted({(o + n_in // 2) % n_in - n_in // 2 for o in offsets}) + i = np.arange(s, e + 1) + F = {} + for o in offsets: + rows = i + o + if periodic: + vals = D[rows % n_out, i] + else: + valid = (rows >= 0) & (rows < n_out) + vals = np.zeros(i.size) + vals[valid] = D[rows[valid], i[valid]] + F[o] = xp.asarray(vals) + per_dir.append(F) + factors[a, c] = per_dir + + # band of M^{ab} at the column offset k (owned rows plus ghost regions), summed over the terms + bands = {} + + def band(a: int, b: int, k: tuple[int, ...]) -> StencilVector | None: + if (a, b, k) not in bands: + vec = None + for coeff, t in mid_terms: + M = block(t.matrix, a, b) + if M is None: + continue + W = M.codomain + assert all(m == 1 for m in W.shifts) and all(m == 1 for m in M.domain.shifts), "Only shifts 1." + if vec is None: + vec = W.zeros() + rows = tuple(slice(p, p + e - s + 1) for p, s, e in zip(W.pads, W.starts, W.ends)) + aliases = [ + [kk for kk in range(-p, p + 1) if (kk - kd) % n == 0] if periodic else [kd] * (abs(kd) <= p) + for kd, n, p, periodic in zip(k, W.npts, M.pads, W.periods) + ] + for ks in itertools.product(*aliases): + vec._data[rows] += coeff * M._data[rows + tuple(p + kk for p, kk in zip(M.pads, ks))] + if vec is not None: + vec.ghost_regions_in_sync = False + vec.update_ghost_regions() + bands[a, b, k] = vec + return bands[a, b, k] + + diags = [] + for c, V in enumerate(dom_spaces): + diag = xp.zeros(tuple(e - s + 1 for s, e in zip(V.starts, V.ends)), dtype=float) + rs = [a for a in range(n_rows) if (a, c) in factors] + for a, b in itertools.product(rs, rs): + Fa, Fb = factors[a, c], factors[b, c] + for oa in itertools.product(*Fa): + for ob in itertools.product(*Fb): + vec = band(a, b, tuple(y - x for x, y in zip(oa, ob))) + if vec is None: + continue + # rows i + oa of the band, from the ghost regions where not owned + W = vec.space + read = tuple( + slice(s + o - sw + p, e + o - sw + p + 1) + for s, e, o, sw, p in zip(V.starts, V.ends, oa, W.starts, W.pads) + ) + assert all(sl.start >= 0 and sl.stop <= n for sl, n in zip(read, vec._data.shape)) + w = reduce( + xp.multiply, + [ + (Fa[d][oa[d]] * Fb[d][ob[d]]).reshape([-1 if e == d else 1 for e in range(len(oa))]) + for d in range(len(oa)) + ], + ) + diag += w * vec._data[read] + diags.append(diag) + return diags + + class FFTSolver(BandedSolver): """ Solve the equation Ax = b for x, assuming A is a circulant matrix. diff --git a/src/struphy/feec/tests/test_stiffness_preconditioner.py b/src/struphy/feec/tests/test_stiffness_preconditioner.py index 4bc498197..7f4c05992 100644 --- a/src/struphy/feec/tests/test_stiffness_preconditioner.py +++ b/src/struphy/feec/tests/test_stiffness_preconditioner.py @@ -74,6 +74,53 @@ def test_stiffness_approximation_unit_cube(derivative, sigma, num_elements): assert _max_abs(diffs, comm) < 1e-12 * _max_abs(_local_diagonal(P.matrix), comm) +@pytest.mark.parametrize("derivative", ["grad", "curl", "div"]) +@pytest.mark.parametrize( + "num_elements, degree, bcs", + [ + ([6, 5, 4], [2, 3, 2], (("dirichlet", "free"), None, None)), + # periodic directions with fewer points than the band width (entries wrap onto the band) + ([6, 4, 3], [3, 3, 3], (("dirichlet", "dirichlet"), None, None)), + ], +) +def test_stiffness_local_diagonal(derivative, num_elements, degree, bcs): + """The diagonal of d^T M d from the stencil data of M equals the probed one (sum of mass terms, full metric).""" + + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import ( + _mass_terms, + _probe_diagonal, + _stiffness_local_diagonal, + _terms_matrix, + ) + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + comm = MPI.COMM_WORLD + derham = Derham(TensorProductGrid(num_elements=num_elements), DerhamOptions(degree=degree, bcs=bcs), comm=comm) + mass_ops = WeightedMassOperators(derham, domains.HollowTorus(a1=0.1, a2=1.0, R0=3.0, tor_period=3)) + + form = FORMS[derivative] + d = getattr(derham, f"{derivative}_bcfree") + mid = getattr(mass_ops, f"M{int(form) + 1}") + space, metric = {"0": ("Hcurl", ("Ginv", "sqrt_g")), "1": ("Hdiv", ("G", "1/sqrt_g")), "2": ("L2", ("1/sqrt_g",))}[ + form + ] + weighted = mass_ops.create_weighted_mass( + space, space, weights=metric + (lambda e1, e2, e3: 1.0 + e1 * e2,), name="M_test", assemble=True + ) + terms = _mass_terms(mid - 0.3 * weighted) + + diag = _stiffness_local_diagonal(d, terms) + ref = _probe_diagonal(d.T @ _terms_matrix(terms) @ d, d.domain) + assert len(diag) == len(ref) + assert _max_abs([a - b for a, b in zip(diag, ref)], comm) < 1e-12 * _max_abs(ref, comm) + + @pytest.mark.parametrize("mapping", ["Cuboid", "HollowCylinder"]) @pytest.mark.parametrize("sigma", [0.0, 1.0]) def test_stiffness_preconditioner_grad(mapping, sigma): diff --git a/src/struphy/feec/utilities.py b/src/struphy/feec/utilities.py index 125e34782..3045e78c8 100644 --- a/src/struphy/feec/utilities.py +++ b/src/struphy/feec/utilities.py @@ -47,10 +47,10 @@ def __init__(self, *vec_fun): self._funs = vec_fun def __call__(self, e1, e2, e3): + # each component evaluated once (the equilibrium callables can be expensive) + a = [fun(e1, e2, e3) for fun in self._funs] # 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)], - ) + tmp = xp.array([[a[m] * a[n] 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)) @@ -76,24 +76,18 @@ def __init__(self, *vec_fun): assert len(vec_fun) == 3 assert all([callable(fun) for fun in vec_fun]) - self._cross_mask = [ - [0, -1, 1], - [1, 0, -1], - [-1, 1, 0], - ] - - self._funs = [ - [lambda e1, e2, e3: 0 * e1, vec_fun[2], vec_fun[1]], - [vec_fun[2], lambda e1, e2, e3: 0 * e2, vec_fun[0]], - [vec_fun[1], vec_fun[0], lambda e1, e2, e3: 0 * e3], - ] + self._funs = vec_fun def __call__(self, e1, e2, e3): + # each component evaluated once (the equilibrium callables can be expensive) + a1, a2, a3 = (fun(e1, e2, e3) for fun in self._funs) + zero = 0.0 * a1 # array from 2d list gives 3x3 array is in the first two indices tmp = xp.array( [ - [self._cross_mask[m][n] * fun(e1, e2, e3) for n, fun in enumerate(row)] - for m, row in enumerate(self._funs) + [zero, -a3, a2], + [a3, zero, -a1], + [-a2, a1, zero], ], ) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index aa7e399d7..9be6a8cf5 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -101,6 +101,17 @@ def __call__(self, dt: float): Time step size. """ + def prepare_time_step(self, dt: float): + """Set up the data that depends on the time step (e.g. preconditioners) for ``dt``, before the first call + with ``dt``, so that it is part of the setup rather than of the first time step (see + :meth:`~struphy.simulation.sim.Simulation.run`). Optional; does nothing by default. + + Parameters + ---------- + dt : float + Time step size of the first call. + """ + @property def _solve_region(self) -> str: """Name of the profiling region for the linear solve(s) of this propagator.""" diff --git a/src/struphy/propagators/shear_alfven_curl_curl_propagator.py b/src/struphy/propagators/shear_alfven_curl_curl_propagator.py index 80d21555e..d9b594a42 100644 --- a/src/struphy/propagators/shear_alfven_curl_curl_propagator.py +++ b/src/struphy/propagators/shear_alfven_curl_curl_propagator.py @@ -6,6 +6,7 @@ from feectools.linalg.solvers import inverse from line_profiler import profile from maybempi import MPI +from scope_profiler import ProfileManager from struphy.feec import preconditioner from struphy.feec.kronecker_block import basis_projection_operator, derivative_operator @@ -131,7 +132,8 @@ class Options(OptionsBase): precond : LiteralOptions.OptsShearAlfvenCurlCurlPrecond, default="HiptmairXu" Preconditioner of the Schur system: Hiptmair-Xu for the curl-curl operator with ``M2gyro`` in the middle (``"HiptmairXu"``, implicit algorithm with ``solver="pcg"`` only, no polar splines; it depends on the time - step, so it is built at the first time step and rebuilt when the time step changes), or a preconditioner + step, so it is built in the setup (:meth:`prepare_time_step`) and rebuilt when the time step changes), or a + preconditioner of the mass matrix ``M1`` (``"MassMatrixPreconditioner"``, ``"PolarMassMatrixPreconditioner"``). With ``algo="explicit"``, ``"HiptmairXu"`` is replaced by ``"MassMatrixPreconditioner"`` (for ``M1``). aux_correction : bool, default=False @@ -341,25 +343,37 @@ def weight(e1, e2, e3, m, n): (Pi_par, Pi_par.T, [kron_par, kron_curl @ kron_par]), ] + @profile def _build_precond(self, dt: float): """Hiptmair-Xu preconditioner of S ~ dt^2/4 (C^T M2gyro C + 4/dt^2 M1) (sigma = 4/dt^2 with the unscaled M2gyro in the middle, as the vector auxiliary space correction of Hiptmair-Xu assumes a unit-scale curl-curl term), plus the field-aligned auxiliary space corrections N P N^T, P the Kronecker preconditioner of N^T (M1 + dt^2/4 C^T M2gyro C) N on V0.""" sigma = 4.0 / dt**2 - hx = HiptmairXuPreconditioner( - None, self.mass_ops, sigma=sigma, mid="M2gyro", aux_weight=self.options.hx_aux_weight - ) + with ProfileManager.profile_region("init HiptmairXu"): + hx = HiptmairXuPreconditioner( + None, self.mass_ops, sigma=sigma, mid="M2gyro", aux_weight=self.options.hx_aux_weight + ) pc = sigma * hx if self.options.aux_correction: - mid = (dt**2 / 4) * self.mass_ops.M2gyro - corrections = [ - (N, NT, preconditioner.StiffnessPreconditioner(self.mass_ops, ops, sigma=0.0, mid=["M1", mid])) - for N, NT, ops in self._aux_spaces - ] - pc = preconditioner.AuxiliarySpacePreconditioner(pc, corrections) + with ProfileManager.profile_region("add auxiliary space corrections"): + mid = (dt**2 / 4) * self.mass_ops.M2gyro + corrections = [ + (N, NT, preconditioner.StiffnessPreconditioner(self.mass_ops, ops, sigma=0.0, mid=["M1", mid])) + for N, NT, ops in self._aux_spaces + ] + pc = preconditioner.AuxiliarySpacePreconditioner(pc, corrections) return pc + def prepare_time_step(self, dt: float): + """Hiptmair-Xu preconditioner of S ~ dt^2/4 (C^T M2gyro C + 4/dt^2 M1) for the time step ``dt``: built in the + setup, rebuilt when dt changes (sigma = 4/dt^2 with the unscaled M2gyro in the middle, as the vector auxiliary + space correction of Hiptmair-Xu assumes a unit-scale curl-curl term).""" + if self.options.algo == "implicit" and self.options.precond == "HiptmairXu" and dt != self._hx_dt: + with ProfileManager.profile_region("build_HiptmairXu_precond"): + self._schur_solver.precond = self._build_precond(dt) + self._hx_dt = dt + @profile def __call__(self, dt): # current FE coeffs @@ -367,12 +381,8 @@ def __call__(self, dt): bn = self.variables.b.spline.vector if self.options.algo == "implicit": - # Hiptmair-Xu preconditioner of S ~ dt^2/4 (C^T M2gyro C + 4/dt^2 M1), for the current time step: built - # at the first time step, rebuilt when dt changes (sigma = 4/dt^2 with the unscaled M2gyro in the middle, - # as the vector auxiliary space correction of Hiptmair-Xu assumes a unit-scale curl-curl term) - if self.options.precond == "HiptmairXu" and dt != self._hx_dt: - self._schur_solver.precond = self._build_precond(dt) - self._hx_dt = dt + # Hiptmair-Xu preconditioner for the current time step (no-op unless dt changed since the setup) + self.prepare_time_step(dt) # solve for new b coeffs byn = self._B.dot(un, out=self._byn) diff --git a/src/struphy/propagators/shear_alfven_propagator.py b/src/struphy/propagators/shear_alfven_propagator.py index 70d4cd293..ce35cfb52 100644 --- a/src/struphy/propagators/shear_alfven_propagator.py +++ b/src/struphy/propagators/shear_alfven_propagator.py @@ -105,7 +105,7 @@ class Options(OptionsBase): the velocity components, ``u_space="Hdiv"`` and no polar splines), with the auxiliary space correction of the :math:`\mathbf E \times \mathbf B` flows if ``exb_correction`` (see there). Both need the implicit algorithm with ``solver="pcg"``. The Schur operator depends on the time step, so these preconditioners are - built at the first time step (and updated if the time step changes). + built in the setup (:meth:`prepare_time_step`) and updated if the time step changes. The multigrid preconditioner is experimental: on the TAE example (HollowTorus, 8x32x4) it needs more PCG iterations than the mass matrix preconditioner and is much slower (the smoother does not treat the large kernel of the curl term on 2-forms). The Kronecker preconditioner (with ``exb_correction``) needs about 11 @@ -278,6 +278,42 @@ def f2(t, y1, y2, out: BlockVector = out2): self._u_tmp2 = self.variables.u.spline.vector.space.zeros() self._b_tmp1 = self.variables.b.spline.vector.space.zeros() + def prepare_time_step(self, dt: float): + """Preconditioner of the Schur operator for the time step ``dt`` (multigrid or Kronecker): built in the setup, + updated or rebuilt when dt changes.""" + if self.options.algo != "implicit": + return + # the multigrid preconditioner of the Schur operator A - dt^2 B C: updated (re-using the coarse mass and + # derivative operators) when dt changes + if self.options.precond == "MultiGrid" and dt != self._multigrid_dt: + schur = self._schur_solver.A - dt**2 * self._schur_solver.BC + if self._schur_solver.precond is None: + self._schur_solver.precond = MultiGridPreconditioner( + schur, self.derham, self.domain, self.options.multigrid, mass_ops=self.mass_ops + ) + else: + self._schur_solver.precond.update(schur) + self._multigrid_dt = dt + + # the Kronecker preconditioner of S = M + dt^2/4 T^T C^T M2 C T (plus the E x B correction), rebuilt when dt + # changes + if self.options.precond == "StiffnessPreconditioner" and dt != self._schur_precond_dt: + mid = (dt**2 / 4) * self.mass_ops.M2 + pc = preconditioner.StiffnessPreconditioner( + self.mass_ops, self._kron_curl_T, sigma=1.0, mid=mid, stab=self._id_M + ) + if self.options.exb_correction: + # N^T S N = N^T M N + dt^2/4 (C T N)^T M2 (C T N) on V0 + pc_phi = preconditioner.StiffnessPreconditioner( + self.mass_ops, + [self._kron_exb, self._kron_curl_T_exb], + sigma=0.0, + mid=[self._id_M, mid], + ) + pc = preconditioner.AuxiliarySpacePreconditioner(pc, [(self._exb_N, self._exb_NT, pc_phi)]) + self._schur_solver.precond = pc + self._schur_precond_dt = dt + @profile def __call__(self, dt): # current FE coeffs @@ -285,36 +321,8 @@ def __call__(self, dt): bn = self.variables.b.spline.vector if self.options.algo == "implicit": - # the multigrid preconditioner of the Schur operator A - dt^2 B C, for the current time step: built at the - # first time step, updated (re-using the coarse mass and derivative operators) when dt changes - if self.options.precond == "MultiGrid" and dt != self._multigrid_dt: - schur = self._schur_solver.A - dt**2 * self._schur_solver.BC - if self._schur_solver.precond is None: - self._schur_solver.precond = MultiGridPreconditioner( - schur, self.derham, self.domain, self.options.multigrid, mass_ops=self.mass_ops - ) - else: - self._schur_solver.precond.update(schur) - self._multigrid_dt = dt - - # the Kronecker preconditioner of S = M + dt^2/4 T^T C^T M2 C T (plus the E x B correction), rebuilt when dt - # changes - if self.options.precond == "StiffnessPreconditioner" and dt != self._schur_precond_dt: - mid = (dt**2 / 4) * self.mass_ops.M2 - pc = preconditioner.StiffnessPreconditioner( - self.mass_ops, self._kron_curl_T, sigma=1.0, mid=mid, stab=self._id_M - ) - if self.options.exb_correction: - # N^T S N = N^T M N + dt^2/4 (C T N)^T M2 (C T N) on V0 - pc_phi = preconditioner.StiffnessPreconditioner( - self.mass_ops, - [self._kron_exb, self._kron_curl_T_exb], - sigma=0.0, - mid=[self._id_M, mid], - ) - pc = preconditioner.AuxiliarySpacePreconditioner(pc, [(self._exb_N, self._exb_NT, pc_phi)]) - self._schur_solver.precond = pc - self._schur_precond_dt = dt + # preconditioner for the current time step (no-op unless dt changed since the setup) + self.prepare_time_step(dt) # solve for new u coeffs byn = self._B.dot(bn, out=self._byn) diff --git a/src/struphy/propagators/tests/test_shear_alfven_curl_curl.py b/src/struphy/propagators/tests/test_shear_alfven_curl_curl.py index 251735a6c..0e18d78c2 100644 --- a/src/struphy/propagators/tests/test_shear_alfven_curl_curl.py +++ b/src/struphy/propagators/tests/test_shear_alfven_curl_curl.py @@ -101,6 +101,22 @@ def test_energy_dissipation(tmp_path): assert en[-1] > 0.9 * en[0], en +def test_prepare_time_step(tmp_path): + """prepare_time_step (called in the setup of Simulation.run) builds the Hiptmair-Xu preconditioner, which a step + with the same dt reuses and one with another dt rebuilds.""" + options = ShearAlfvenCurlCurlPropagator.Options(solver_params=SolverParameters(tol=1e-13, maxiter=1000)) + model = allocate_model(tmp_path, "prepare", options) + shear_alf = model.propagators.shear_alf + assert shear_alf._hx_dt is None + shear_alf.prepare_time_step(0.5) + assert shear_alf._hx_dt == 0.5 + pc = shear_alf._schur_solver.precond + shear_alf(0.5) + assert shear_alf._schur_solver.precond is pc + shear_alf(0.2) + assert shear_alf._hx_dt == 0.2 and shear_alf._schur_solver.precond is not pc + + def test_formulations_agree(tmp_path): """For a smooth mode, the steps with and without the projection in the curl-curl term differ by the projection error only.""" diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 6567c0d46..9452e81c5 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -727,6 +727,10 @@ def run(self, one_time_step: bool = False, profiling_activated: bool | None = No # CUDA kernels (CuPy backend only): compile them now rather than in the first time step self.compile_cuda_kernels() + + # time step dependent data of the propagators (e.g. preconditioners), rather than in the first step + with ProfileManager.profile_region("setup: prepare time step", functions=[self._prepare_time_step]): + self._prepare_time_step() with ProfileManager.profile_region("setup: run metadata", functions=[self._write_run_metadata]): self._write_run_metadata( one_time_step=one_time_step, @@ -1359,6 +1363,16 @@ def _allocate_variables(self): # self._pointer[key] = val["obj"].vector + def _prepare_time_step(self): + """Call :meth:`~struphy.propagators.base.Propagator.prepare_time_step` of each propagator with the time step + it receives from :meth:`~struphy.models.base.StruphyModel.integrate` (``dt / 2`` for all but the last + propagator with Strang splitting).""" + dt = self.time_opts.dt + props = self.model.prop_list + for n, prop in enumerate(props): + half = self.time_opts.split_algo == "Strang" and n < len(props) - 1 + prop.prepare_time_step(dt / 2 if half else dt) + @profile def _allocate_propagators(self): """Allocate propagators and bind shared FEEC/domain operators.