General boundary_flux primitive: Consistent Boundary Flux for any solver (heat flux / Nusselt, traction)

src/underworld3/cython/petsc_generic_snes_solvers.pyx

@@ -1819,6 +1819,114 @@ class SolverBaseClass(uw_object):
 
         return
 
+    def _assemble_volume_reaction(self, time=None, verbose=False):
+        """RAW per-rank FEM VOLUME residual (no boundary terms) in the DM-local layout,
+        as a numpy array.
+
+        At an essential-BC (Dirichlet) node this residual IS the consistent boundary
+        reaction — the integrated nodal flux :math:`\\int_\\Gamma (F\\cdot\\hat n)\\phi_i`
+        (heat flux for a scalar diffusion solve, traction for Stokes). Interior nodes are
+        ~0. General across scalar / vector / Stokes solvers: the current solution is
+        gathered from ``self.fields`` when present (Stokes) else the single
+        ``Unknowns.u`` field.
+
+        NOTE: the returned array is NOT globally assembled. The DM has overlap=0, so each
+        rank computes only its OWNED cells' contribution; a boundary node shared across a
+        partition cut therefore holds only this rank's PARTIAL reaction. The complete
+        reaction is assembled by the caller (``utilities.boundary_flux._desmear``) by
+        SUMMING each rank's partial by coordinate — not by a hand-rolled localToGlobal.
+        """
+        cdef DM dm
+        cdef Vec xvec
+        cdef Vec fvec
+        cdef DM _time_dm_reaction
+        cdef PetscFormKey key
+        cdef IS ccell_is
+        cdef PetscReal residual_time = 0.0
+        cdef PetscReal implicit_form_time = <PetscReal>-1.7976931348623157e308
+
+        self._build(verbose, False, None)
+
+        if time is not None:
+            if hasattr(time, 'magnitude') or hasattr(time, '_pint_qty'):
+                t_nd = float(uw.non_dimensionalise(time))
+            else:
+                t_nd = float(time)
+            _time_dm_reaction = self.dm
+            UW_DMSetTime(_time_dm_reaction.dm, <PetscReal>t_nd)
+
+        self.mesh.update_lvec()
+        self.dm.setAuxiliaryVec(self.mesh.lvec, None)
+        self._update_constants()
+
+        gvec = self.dm.getGlobalVec()
+        xlocal = self.dm.getLocalVec()
+        flocal = self.dm.getLocalVec()
+        gvec.setArray(0.0)
+        xlocal.setArray(0.0)
+        flocal.setArray(0.0)
+
+        try:
+            # gather the current solution into the global vector (field-structure agnostic)
+            if getattr(self, "fields", None):
+                for name, var in self.fields.items():
+                    sgvec = gvec.getSubVector(self._subdict[name][0])
+                    self._subdict[name][1].localToGlobal(var.vec, sgvec)
+                    gvec.restoreSubVector(self._subdict[name][0], sgvec)
+            else:
+                _names, _iss, _subdms = self.dm.createFieldDecomposition()
+                sgvec = gvec.getSubVector(_iss[0])
+                _subdms[0].localToGlobal(self.Unknowns.u.vec, sgvec)
+                gvec.restoreSubVector(_iss[0], sgvec)
+
+            self.dm.globalToLocal(gvec, xlocal)
+
+            dm = self.dm
+            xvec = xlocal
+            fvec = flocal
+            CHKERRQ(DMPlexSNESComputeResidualFEM(dm.dm, xvec.vec, fvec.vec, NULL))
+
+            # Return the RAW local residual: each rank has computed its OWNED cells'
+            # contribution to its local nodes (the DM has overlap=0, so a boundary node
+            # shared across a partition cut holds only this rank's PARTIAL contribution).
+            # The caller assembles the complete reaction by summing these partials across
+            # ranks by coordinate (boundary_flux._desmear), consistent with the boundary-
+            # mass gather — this reproduces the rock-solid volume integral at cut nodes.
+            return np.array(flocal.array, copy=True)
+        finally:
+            self.dm.restoreLocalVec(flocal)
+            self.dm.restoreLocalVec(xlocal)
+            self.dm.restoreGlobalVec(gvec)
+
+    def boundary_flux(self, boundary, mass="lumped", remove_mean=False, normal=None):
+        r"""Consistent boundary flux on ``boundary``, recovered from the essential-BC
+        reaction of the last solve (the Consistent Boundary Flux method).
+
+        Returns ``(xs, flux)`` with one entry per boundary node on this rank: for a
+        **scalar** solver the outward normal flux :math:`F\cdot\hat n` (e.g. surface heat
+        flux :math:`-k\,\partial T/\partial n`, whose boundary mean is the Nusselt
+        number); for a **vector** solver the traction :math:`\sigma\cdot\hat n` (pass
+        ``normal`` to get the scalar normal component :math:`\hat n\cdot\sigma\cdot\hat n`).
+
+        ``mass`` de-smears the nodal reaction with the ``"lumped"`` (diagonal, monotone —
+        no overshoot at a flux jump) or ``"consistent"`` boundary mass. ``remove_mean``
+        subtracts the boundary mean — leave ``False`` for a physical flux (the mean is
+        the Nusselt number); ``True`` gives a gauge-free field (e.g. dynamic topography).
+        Parallel-safe and partition-independent."""
+        from underworld3.utilities.boundary_flux import boundary_flux as _bf
+        return _bf(self, boundary, mass=mass, remove_mean=remove_mean, normal=normal)
+
+    def boundary_flux_field(self, boundary, field, mass="lumped",
+                            remove_mean=False, scale=1.0, normal=None):
+        r"""Write the consistent boundary flux (see :meth:`boundary_flux`) onto a scalar
+        MeshVariable ``field`` at the boundary nodes (interior untouched), multiplied by
+        ``scale``. This is the field hand-off for downstream machinery (surface heat
+        flux for coupling, or — with ``remove_mean=True`` and ``scale=-1/(\Delta\rho g)``
+        — dynamic topography). Returns ``field``."""
+        from underworld3.utilities.boundary_flux import boundary_flux_to_field as _bff
+        return _bff(self, boundary, field, mass=mass, remove_mean=remove_mean,
+                    scale=scale, normal=normal)
+
 ## Specific to dimensionality
 
 

src/underworld3/utilities/boundary_flux.py

@@ -0,0 +1,255 @@
+"""Consistent Boundary Flux (CBF) recovery — general across solvers.
+
+The residual of the assembled interior (volume) FEM problem, read at an essential-BC
+boundary node, is the *consistent* nodal flux there (Gresho et al.). De-smearing that
+nodal reaction with the boundary mass gives a pointwise surface flux:
+
+  * scalar diffusion / advection-diffusion  →  surface heat flux  -k dT/dn  (Nusselt),
+  * Stokes                                   →  boundary traction sigma.n   (sigma_nn).
+
+This module holds the parts that do not depend on the equation: the boundary-node
+gathering, the boundary-mass de-smear (lumped / consistent), and the field hand-off.
+The equation-specific bit — extracting the nodal reaction — is the solver method
+``_assemble_volume_reaction``, which returns each rank's RAW (per-rank) volume FEM
+residual; the complete reaction at a boundary node shared across a partition cut is then
+assembled here in ``_desmear`` by SUMMING each rank's partial contribution by coordinate
+(the same rock-solid gather used for the boundary mass — no hand-rolled global assembly).
+
+``mass="lumped"`` (default) uses the diagonal boundary mass: being an M-matrix it cannot
+overshoot where the flux jumps (no Gibbs wiggle) and is a purely local division.
+``remove_mean=False`` (default) keeps the physical mean flux (the Nusselt number);
+set ``remove_mean=True`` for a gauge-free field (e.g. dynamic topography).
+"""
+import numpy as np
+from mpi4py import MPI
+
+
+def _key(c, dim):
+    return tuple(round(float(t), 9) for t in np.asarray(c).ravel()[:dim])
+
+
+def _boundary_stratum_is(dm, mesh, boundary):
+    """Facet stratum IS for ``boundary`` via the CONSOLIDATED ``UW_Boundaries`` label
+    (boundaries are distinguished by value on this one label — the per-boundary labels
+    named like the boundary do not survive mesh adaptation; only ``UW_Boundaries`` is
+    rebuilt). Returns None if this rank owns no part of the boundary. Raises a clear
+    error for an unknown boundary name."""
+    match = [b.value for b in mesh.boundaries if b.name == boundary]
+    if not match:
+        raise ValueError(
+            f"Unknown boundary {boundary!r}; known: {[b.name for b in mesh.boundaries]}")
+    label = dm.getLabel("UW_Boundaries")
+    if label is None:
+        return None
+    return label.getStratumIS(match[0])
+
+
+def _point_coord(dm, dim, cvec, csec, v0, v1, q):
+    """Coordinate of a DMPlex point (vertex → its coord; higher point → mean of its
+    closure vertices)."""
+    if v0 <= q < v1:
+        return cvec[csec.getOffset(q) // dim]
+    clo = dm.getTransitiveClosure(q)[0]
+    verts = [int(c) for c in clo if v0 <= c < v1]
+    return np.mean([cvec[csec.getOffset(v) // dim] for v in verts], axis=0)
+
+
+def _boundary_field_nodes(solver, boundary, field_id=0):
+    """DMPlex points carrying `field_id` DOFs on `boundary`, with their coordinates.
+    Parallel-safe: a rank owning no part of the boundary gets a NULL stratum IS
+    (guarded); ghost/shared nodes are included and their partial per-rank reactions are
+    summed by coordinate in ``_desmear`` to form the complete reaction."""
+    dm = solver.dm
+    dim = solver.mesh.dim
+    lsec = dm.getLocalSection()
+    csec = dm.getCoordinateSection()
+    cvec = np.asarray(dm.getCoordinatesLocal().array).reshape(-1, dim)
+    v0, v1 = dm.getDepthStratum(0)
+    fS, fE = dm.getHeightStratum(1)
+    sis = _boundary_stratum_is(dm, solver.mesh, boundary)
+    if sis is None or sis.handle == 0:
+        return [], lsec, csec, cvec, v0, v1
+    facets = [int(z) for z in sis.getIndices()]
+    seen = set(); out = []
+    for f in facets:
+        if not (fS <= f < fE):
+            continue
+        for q in (int(c) for c in dm.getTransitiveClosure(f)[0]):
+            if q in seen or lsec.getFieldDof(q, field_id) <= 0:
+                continue
+            seen.add(q)
+            out.append((q, _point_coord(dm, dim, cvec, csec, v0, v1, q)))
+    return out, lsec, csec, cvec, v0, v1
+
+
+def _node_normals(solver, boundary, normal, nodes, dm, dim, cvec, csec, v0, v1):
+    """Per-node outward unit normal (only needed to project a vector reaction).
+    ``normal`` is None (geometric facet normal), a sympy 1×dim Matrix (analytic,
+    lambdified), or a constant (dim,) vector."""
+    interior_ref = cvec.mean(axis=0)
+    sym_fn = const = None
+    if normal is not None:
+        try:
+            import sympy
+            if isinstance(normal, sympy.Matrix):
+                sym_fn = sympy.lambdify(list(solver.mesh.X),
+                                        [normal[0, k] for k in range(dim)], "numpy")
+        except Exception:
+            sym_fn = None
+        if sym_fn is None:
+            const = np.asarray(normal, dtype=float).ravel()
+    nmap = {}
+    coord = {q: c for q, c in nodes}
+    if normal is None:
+        # accumulate area-weighted facet normals to the closure nodes
+        sis = _boundary_stratum_is(dm, solver.mesh, boundary)
+        facets = [] if (sis is None or sis.handle == 0) else [int(z) for z in sis.getIndices()]
+        fS, fE = dm.getHeightStratum(1)
+        acc = {}
+        for f in facets:
+            if not (fS <= f < fE):
+                continue
+            _, cent, nrm = dm.computeCellGeometryFVM(f)
+            ne = np.asarray(nrm, float); ne = ne / (np.linalg.norm(ne) + 1e-30)
+            if np.dot(ne, np.asarray(cent) - interior_ref) < 0:
+                ne = -ne
+            for q in (int(c) for c in dm.getTransitiveClosure(f)[0]):
+                if q in coord:
+                    acc[q] = acc.get(q, np.zeros(dim)) + ne
+        for q in coord:
+            nn = acc.get(q, np.zeros(dim))
+            nmap[q] = nn / (np.linalg.norm(nn) + 1e-30)
+    else:
+        for q, c in nodes:
+            ne = np.asarray(sym_fn(*c), float).ravel() if sym_fn is not None else const.copy()
+            nmap[q] = ne / (np.linalg.norm(ne) + 1e-30)
+    return nmap
+
+
+def _desmear(solver, boundary, xs, R, mass, remove_mean):
+    """De-smear per-node reaction loads R (aligned with xs) into a pointwise flux via the
+    boundary mass, assembled globally by a coordinate-keyed allgather so every rank forms
+    the identical system. Returns the flux at this rank's local nodes (xs order)."""
+    dm = solver.dm; dim = solver.mesh.dim; comm = dm.comm.tompi4py()
+    csec = dm.getCoordinateSection()
+    cvec = np.asarray(dm.getCoordinatesLocal().array).reshape(-1, dim)
+    v0, v1 = dm.getDepthStratum(0)
+    if dim != 2:
+        # no line-mass geometry yet in 3D → global-mean lumped fallback
+        tot = comm.allreduce(float(np.sum(R)), op=MPI.SUM)
+        cnt = comm.allreduce(int(len(R)), op=MPI.SUM)
+        m = tot / max(cnt, 1)
+        return np.asarray(R) - (m if remove_mean else 0.0)
+
+    e0, e1 = dm.getDepthStratum(1)
+    def vcoord(q): return cvec[csec.getOffset(q) // dim]
+    nodeR = {_key(x, dim): float(r) for x, r in zip(xs, R)}
+    sis = _boundary_stratum_is(dm, solver.mesh, boundary)
+    strat = [] if (sis is None or sis.handle == 0) else [int(z) for z in sis.getIndices()]
+    local_elems = []
+    for e in [q for q in strat if e0 <= q < e1]:
+        a, b = (int(c) for c in dm.getCone(e))
+        cmid = _point_coord(dm, dim, cvec, csec, v0, v1, e)
+        h = float(np.hypot(*(vcoord(b) - vcoord(a))))
+        local_elems.append((_key(vcoord(a), dim), _key(cmid, dim), _key(vcoord(b), dim), h))
+
+    # SUM the nodal reaction across ranks by coordinate: with overlap=0 a boundary node
+    # shared across a partition cut holds only each rank's partial cell contribution, so
+    # summing them assembles the complete reaction (matches the rock-solid volume integral).
+    R_by = {}
+    for d in comm.allgather(nodeR):
+        for k, v in d.items():
+            R_by[k] = R_by.get(k, 0.0) + v
+    uniq = {}
+    for lst in comm.allgather(local_elems):
+        for (ka, km, kb, h) in lst:
+            uniq[(ka, km, kb)] = h
+    keys = sorted(R_by.keys()); gi = {k: i for i, k in enumerate(keys)}
+    n = len(keys); Rg = np.zeros(n)
+    for k, i in gi.items():
+        Rg[i] = R_by[k]
+    if mass == "lumped":
+        mL = np.zeros(n)
+        for (ka, km, kb), h in uniq.items():
+            mL[gi[ka]] += h / 6.0; mL[gi[km]] += 2.0 * h / 3.0; mL[gi[kb]] += h / 6.0
+        sig = Rg / mL
+    else:
+        # consistent P2 line mass — a dense (n×n) solve in the number of boundary nodes
+        # (O(n^3)); fine for a 1D boundary (n ~ resolution) but prefer the default lumped
+        # (O(n), monotone) for very large boundaries.
+        M = np.zeros((n, n))
+        Me = np.array([[4., 2, -1], [2, 16, 2], [-1, 2, 4]])
+        for (ka, km, kb), h in uniq.items():
+            tri = [gi[ka], gi[km], gi[kb]]; Mh = (h / 30.0) * Me
+            for ii in range(3):
+                for jj in range(3):
+                    M[tri[ii], tri[jj]] += Mh[ii, jj]
+        sig = np.linalg.solve(M, Rg)
+    if remove_mean:
+        sig = sig - sig.mean()
+    return np.array([sig[gi[_key(x, dim)]] for x in xs])
+
+
+def boundary_flux(solver, boundary, mass="lumped", remove_mean=False, normal=None):
+    """See ``SolverBaseClass.boundary_flux``. Returns ``(xs, flux)`` for this rank's
+    boundary nodes; scalar solver → normal flux, vector solver → traction (or its normal
+    component if ``normal`` is given)."""
+    dm = solver.dm; dim = solver.mesh.dim
+    ra = np.asarray(solver._assemble_volume_reaction()).ravel()
+    nodes, lsec, csec, cvec, v0, v1 = _boundary_field_nodes(solver, boundary, field_id=0)
+    ncomp = lsec.getFieldComponents(0)
+    xs = np.array([c for _q, c in nodes]) if nodes else np.zeros((0, dim))
+
+    if ncomp == 1:
+        R = np.array([ra[lsec.getFieldOffset(q, 0)] for q, _c in nodes]) if nodes else np.zeros(0)
+        flux = _desmear(solver, boundary, xs, R, mass, remove_mean)
+        return xs, flux
+
+    # vector reaction (traction sigma.n at each node)
+    Rvec = np.array([ra[lsec.getFieldOffset(q, 0):lsec.getFieldOffset(q, 0) + ncomp]
+                     for q, _c in nodes]) if nodes else np.zeros((0, ncomp))
+    if normal is not None:
+        # scalar NORMAL component sigma_nn = n.(sigma.n)
+        nmap = _node_normals(solver, boundary, normal, nodes, dm, dim, cvec, csec, v0, v1)
+        Rn = np.array([float(np.dot(nmap[q], Rvec[i])) for i, (q, _c) in enumerate(nodes)]) \
+            if nodes else np.zeros(0)
+        return xs, _desmear(solver, boundary, xs, Rn, mass, remove_mean)
+    # full traction vector: de-smear each component independently
+    cols = [_desmear(solver, boundary, xs, Rvec[:, k] if len(Rvec) else np.zeros(0),
+                     mass, remove_mean) for k in range(ncomp)]
+    return xs, (np.column_stack(cols) if nodes else np.zeros((0, ncomp)))
+
+
+def boundary_flux_to_field(solver, boundary, field, mass="lumped",
+                           remove_mean=False, scale=1.0, normal=None):
+    """See ``SolverBaseClass.boundary_flux_field``. Writes ``scale * flux`` onto the
+    scalar MeshVariable ``field`` at the boundary nodes (interior untouched)."""
+    dim = solver.mesh.dim
+    xs, flux = boundary_flux(solver, boundary, mass=mass, remove_mean=remove_mean, normal=normal)
+    flux = np.asarray(flux)
+    # a SCALAR field can only hold a scalar flux — a vector solver returns a per-node
+    # traction VECTOR unless a `normal` is given to project it. Fail fast rather than
+    # silently pairing the flattened vector with the nodes.
+    if flux.ndim > 1 and flux.shape[1] != 1:
+        raise ValueError(
+            "boundary_flux_field target is a scalar MeshVariable but boundary_flux "
+            "returned a vector (traction). Pass normal= to project onto the normal "
+            "component, or use boundary_flux() directly for the full vector.")
+    fmap = {_key(x, dim): scale * float(f) for x, f in zip(np.asarray(xs), flux.ravel())}
+    fc = np.asarray(field.coords)
+    # Write the field ONCE from a local copy: a per-node write to var.data fires the
+    # write-callback each time, and the boundary-node count differs per rank (a rank may
+    # own none of the boundary), so per-node writes would desync the callback and hang.
+    newdata = np.asarray(field.data).copy()
+    for i in range(fc.shape[0]):
+        v = fmap.get(_key(fc[i], dim))
+        if v is not None:
+            newdata[i, 0] = v
+    field.data[...] = newdata
+    base = getattr(field, "_base_var", field)
+    if hasattr(base, "_sync_lvec_to_gvec"):
+        base._sync_lvec_to_gvec()
+    if hasattr(base, "_canonical_data"):
+        base._canonical_data = None
+    solver.mesh._stale_lvec = True
+    return field