diff --git a/docs/developer/design/eulerian-supg-transport.md b/docs/developer/design/eulerian-supg-transport.md index 39146431..0ec9bbf0 100644 --- a/docs/developer/design/eulerian-supg-transport.md +++ b/docs/developer/design/eulerian-supg-transport.md @@ -742,6 +742,49 @@ semi-Lagrangian Stokes stress history (`DFDt` on a viscoelastic Stokes solve) is untouched: it advects a stress that is not an unknown of the solve, which is a different job from the one the contract describes. +### The histories share one characteristic trace + +The composable terms are the right design, and they exposed a cost: two histories on +the same nodes each traced their own characteristics. A profile of the nodal +`AdvDiffusionSLCN` step (2026-09-08) counted, per step, two managers (the value and the +flux history) each recording its field, tracing the same nodes with the same velocity +(two evaluations), sampling its history and caching the same velocity level: ten +evaluate-class calls and a projection where the work is one trace, two samples and a +copy. A hidden `simplify` in the parallel evaluator, run on every evaluation of an +expression holding a mesh variable, had been adding 14 of 25 seconds on top. + +The mathematics stays with the term. What the terms now share is the evaluation behind +`update_pre_solve`: a `CharacteristicTrace` (`systems/ddt.py`) owned by the solver +holds the departure points per node set and segment structure for the current step, and +the velocity levels $v^{n-1}, v^{n-2}, \dots$ cached by evaluation at the true nodes. Each +history asks it for "the feet of my nodes through these segments" and samples its own +field; the second history on the same nodes, and the older slots of a one-segment +history, are served from the cache. The solver delimits the step (`begin_step`, +`finish_step`, which records the velocity used this step); a manager used on its own +owns a private trace and delimits its own steps, so nothing changes for standalone use. +`share_characteristics(DuDt, DFDt)` attaches one trace to every manager that follows +the same velocity on the same mesh; the Navier-Stokes SLCN solver shares the velocity +levels this way even though its stress history lives on different nodes. + +A history is skipped only when it is a derived quantity the weak form does not read: +the diffusive flux history under BDF or theta = 1, whose old-level weight is zero. It is +never a state history. The flux history of a viscoelastic solve is state (the stress at +the old time cannot be rebuilt from the present velocity gradient and rheology), and +Crank-Nicolson is what keeps the elastic response undamped, so that history is always +carried. + +Measured on the rotating Gaussian (h = 0.1, C = 0.25, one revolution, answer identical +to every printed digit, L2 5.142e-2, peak 0.6859): + +| | nodal SLCN | integration-point SLCN | +|---|---|---| +| before | 1350 ms per step | 2250 | +| `simplify` forwarded | 510 | 1250 | +| shared trace, unread flux skipped | 250 | 660 | + +The cost ratio against SUPG in the convection comparison above was measured before +both fixes and should be re-read with that in mind. + ## What the timestep estimate means The cell-crossing time is not a stability limit for either scheme and says diff --git a/docs/developer/subsystems/integration-point-variables.md b/docs/developer/subsystems/integration-point-variables.md index 542ec9f7..d8ee4cfc 100644 --- a/docs/developer/subsystems/integration-point-variables.md +++ b/docs/developer/subsystems/integration-point-variables.md @@ -260,3 +260,142 @@ integration-point scheme loses 45 times less of it. The trace-back samples twelve points per cell rather than the P2 nodes, and the snapshot evaluation at the moving feet misses the locator cache every step, so the update costs about three times the nodal one. + +## Swarm proxy at the integration points + +A swarm variable normally reaches the weak form through a nodal proxy: +the particle field is reconstructed at the proxy's nodes from the nearest +particles and the assembler interpolates it to the integration points with +the basis. With `proxy_location="integration_points"` the proxy is an +integration-point variable, reconstructed from the nearest particles at +every integration point and read there directly. That is the +Ellipsis / Underworld PIC-LIP mapping: material properties are sampled +from the particles around each integration point, not smoothed to the +nodes and back. + +```python +swarm = uw.swarm.Swarm(mesh) +M = uw.swarm.SwarmVariable("M", swarm, 1, proxy_location="integration_points") +swarm.populate(fill_param=3) +M.data[:, 0] = ... # per particle +stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_0 * M.sym[0] + eta_1 * (1 - M.sym[0]) +``` + +The reconstruction itself is unchanged (a linear-exact RBF over the nearest +particles, `rbf_interpolate`); only its target moved. A particle-carried +material step is reproduced at the integration points with less than half +the L2 error of the nodal proxy, and is exactly 0 or 1 one cell away from +the interface (`tests/test_0067_integration_point_proxy.py`). +`proxy_degree` and `proxy_continuous` are ignored for this proxy; the proxy +has no gradient, so a derivative of the swarm variable's symbol is refused. +Vector and tensor swarm variables get a multi-component proxy. + +`Lagrangian_Swarm(..., proxy_location="integration_points")` applies the +same to the fully Lagrangian history: the slots carried on the particles +are reconstructed at the integration points and the weak form reads them +there, with no nodal history field. This is the Lagrangian option for large +particle swarms, where the particles carry the state and the mesh only +integrates it. + +## Swarm proxy as a polynomial per cell + +`proxy_location="cells"` is the third target. The proxy is a discontinuous +mesh variable of `proxy_degree`, and every cell holds the least-squares +polynomial through the particles that cell holds +(`utilities/cell_polynomial_projection.py`). The assembler reads it at the +integration points through the ordinary basis, so: + +- a polynomial particle field up to `proxy_degree` is reproduced exactly; +- the value at the rule is a polynomial on the mesh cell, so the default + rule integrates it exactly and the oversampling guard of the + integration-point history does not apply; +- a material step on a cell edge is exactly 0 or 1 on either side, with no + overshoot (the RBF reconstruction overshoots a step by up to 14%); +- the proxy has a gradient, so Crank-Nicolson and the Adams-Moulton flux + of the history work; +- each rank fits its own cells from its own particles: no neighbour search + across ranks, no halo particles. + +A cell with fewer particles than the basis size plus two takes a linear +fit to the particles nearest its centroid, which is the RBF's +neighbourhood; linear, because a higher-degree polynomial extrapolated +from a distant neighbourhood is unbounded (a P2 extrapolation into the +emptied corner cells of a rotating box reached twice the field maximum). +That cell is consistent to first order but no longer a cell-local fit, and +a light swarm degrades the same way the RBF proxy does. A cell with no +particles keeps its previous proxy value: no particles is no information, +and that holds until the swarm is repopulated. The threshold and patch +size are `nmin` and `patch_nnn` on `CellPolynomialProjector.fit`. + +```python +M = uw.swarm.SwarmVariable("M", swarm, 1, proxy_location="cells", proxy_degree=2) +``` + +### Why a least-squares fit and not a conservative transfer + +The conservative particle-to-mesh transfer solves the rule mass matrix +against the particle moments, `M u = sum_p V_p phi(x_p) psi_p` (PETSc's +`DMSwarmProjectFields` on a Plex does exactly this, with unit weights). It +hands the mesh the particle sums exactly, but its nodal values carry the +error of the particle "quadrature", which scales with the field value over +the square root of the particle count, not with the field's variation over +the cell. Measured on `UnstructuredSimplexBox(cellSize=0.05)` with the +`populate()` lattice jittered by 30% of the particle spacing, L2 error at +the integration points (`~/+Simulations/integration_point_proxy`, +2026-09-08): + +| field, fill (particles per cell) | RBF at the points | conservative P1 | conservative P2 | least squares P1 | least squares P2 | +|---|---|---|---|---|---| +| linear, 3 | 7e-16 | 3.5 | 5.8 | 2e-10 | 6e-10 | +| linear, 21 | 6e-16 | 0.89 | 2.6 | 9e-16 | 1e-15 | +| Gaussian (width 0.1), 3 | 1.1e-3 | 1.6e-1 | 2.7e-1 | 2.8e-3 | 4.4e-4 | +| Gaussian, 10 | 3.7e-4 | 6.8e-2 | 1.5e-1 | 1.7e-3 | 1.2e-4 | +| Gaussian, 21 | 1.2e-4 | 3.8e-2 | 1.3e-1 | 1.6e-3 | 5.5e-5 | +| Gaussian, 1 per cell on average (34% of cells empty) | 3.9e-3 | 2.5e-1 | 3.8e-1 | 8.9e-3 | 3.1e-3 | + +Moment-matched particle weights (chosen so constants transfer exactly) +repair the conservative transfer's constant mode and cut the linear error +fifty-fold, but go negative on thin cells and still trail the fit by two +orders of magnitude on the Gaussian. Conservation and light sampling are in +tension: what makes a light swarm usable is polynomial reproduction with a +support that widens when the cell is thin, and the fit degree has to reach +the mesh degree to profit from particle density (the P1 fit is limited by +cell size and is worse than the RBF; the P2 fit beats the RBF three times +over at ten particles per cell and at every density tested). The cell-mean +of the fit matches the particle mean of the cell exactly; the integral of +the fitted field differs from the particle sum by the particle-quadrature +error, 1e-4 relative at ten particles per cell. + +The refresh costs about the same as the RBF path: at ten particles per +cell on 944 cells, 8 ms (locate 6 ms, fit 2 ms) against 22 ms for the RBF +proxy with its kd-tree rebuilt. + +### As the transport term of an advection-diffusion solve + +`Lagrangian_Swarm(..., proxy_location="cells")` composed into +`uw.systems.AdvDiffusion` gives a particle-in-cell transport scheme: the +particles are advected (`swarm.advection`), each history slot is fitted +to the cells, the mesh solves the diffusion against that history, and the +particles re-read the solution (`particle_update="pic"`, the default; +`"flip"` adds the mesh increment instead and is kept for the MPM line of +work, it accumulates the projection increments). The history is sampled +at the particles the first time the swarm moves, through the swarm's +pre-advection hook; sampled at the first solve instead, it would see the +landed positions and lose a step. + +Rotating Gaussian (sigma 0.1 at radius 0.4, one revolution, h = 0.1, P2, +C = 0.25, `~/+Simulations/integration_point_proxy/scripts/transport_gaussian.py`), +L2 error of the mesh field against the exact solution: + +| Pe_h | nodal SLCN | integration-point SLCN | PIC, cells P2, 10 particles per cell | PIC, 21 per cell | +|---|---|---|---|---| +| infinite | 5.1e-2 (peak 0.69) | 9.1e-3 (peak 0.99) | 1.6e-2 (peak 0.86) | 5.2e-3 (peak 0.995) | +| 400 | 4.3e-2 | 6.8e-3 | 1.3e-2 | 4.1e-3 | +| 100 | 2.7e-2 | 3.4e-3 | 8.2e-3 | 2.4e-3 | +| ms per step | 510 | 1250 | 140 to 170 | 180 to 250 | + +The particle scheme's error is the per-step re-projection (fit, then +Galerkin projection, then read-back) and falls with particle density; at +21 particles per cell it is below the integration-point history at a +fifth of the cost. At C = 2 all three schemes are limited by the midpoint +RK2 trajectory (half a radian per step, 4% phase lead), not by transport. diff --git a/src/underworld3/cython/petsc_quadrature_fe.pyx b/src/underworld3/cython/petsc_quadrature_fe.pyx index efbbd553..d85d56c5 100644 --- a/src/underworld3/cython/petsc_quadrature_fe.pyx +++ b/src/underworld3/cython/petsc_quadrature_fe.pyx @@ -30,7 +30,7 @@ from petsc4py import PETSc from petsc4py.PETSc cimport FE, PetscFE, Quad, PetscQuadrature, DM, PetscDM from petsc4py.PETSc cimport PetscSpace, PetscDualSpace, PetscObject, MPI_Comm from petsc4py.PETSc cimport CHKERR as CHKERRQ -from underworld3.cython.petsc_types cimport PetscInt, PetscReal, PetscErrorCode +from underworld3.cython.petsc_types cimport PetscInt, PetscReal, PetscErrorCode, PetscBool import numpy as np @@ -62,6 +62,8 @@ cdef extern from "petsc.h" nogil: PetscErrorCode PetscQuadratureDestroy(PetscQuadrature*) PetscErrorCode PetscFECreateFromSpaces(PetscSpace, PetscDualSpace, PetscQuadrature, PetscQuadrature, PetscFE*) + PetscErrorCode PetscFECreateVector(PetscFE, PetscInt, PetscBool, PetscBool, PetscFE*) + PetscErrorCode PetscFEDestroy(PetscFE*) PetscErrorCode PetscObjectReference(PetscObject) PetscErrorCode PetscObjectSetName(PetscObject, const char*) PetscErrorCode PetscMalloc(size_t, void**) @@ -90,8 +92,12 @@ cdef extern from "uw_delta_space.h" nogil: CHKERRQ(UWDeltaSpaceRegister()) -def create_delta_fe(Quad quad, int polytope, name="quadrature_point_fe"): - r"""Build the scalar quadrature-point element on ``quad``. +def create_delta_fe(Quad quad, int polytope, name="quadrature_point_fe", int num_components=1): + r"""Build the quadrature-point element on ``quad``. + + ``num_components > 1`` wraps the scalar element with ``PetscFECreateVector`` + (interleaved basis and components): the dofs of a cell are point-major, + component-minor, so a local vector reshapes to ``(ncells * Nq, Nc)``. Parameters ---------- @@ -107,7 +113,8 @@ def create_delta_fe(Quad quad, int polytope, name="quadrature_point_fe"): Returns ------- petsc4py.PETSc.FE - Element of dimension ``Nq`` (points in the rule), one component, + Element of ``Nq * num_components`` basis functions (``Nq`` points in + the rule, ``num_components`` interleaved components, one by default), with ``quad`` as its cell quadrature and no face quadrature. """ cdef PetscInt qdim = 0, qNc = 0, Nq = 0, i, d @@ -120,7 +127,10 @@ def create_delta_fe(Quad quad, int polytope, name="quadrature_point_fe"): cdef PetscDualSpace Q = NULL cdef PetscDM refcell = NULL cdef PetscFE cfe = NULL + cdef PetscFE vfe = NULL cdef FE pyfe + if num_components < 1: + raise ValueError("num_components must be >= 1") CHKERRQ(PetscQuadratureGetData(quad.quad, &qdim, &qNc, &Nq, &points, &weights)) if qNc != 1: @@ -161,6 +171,10 @@ def create_delta_fe(Quad quad, int polytope, name="quadrature_point_fe"): # caller's Quad alive by taking a reference first. No face quadrature. CHKERRQ(PetscObjectReference(quad.quad)) CHKERRQ(PetscFECreateFromSpaces(P, Q, quad.quad, NULL, &cfe)) + if num_components > 1: + CHKERRQ(PetscFECreateVector(cfe, num_components, 1, 1, &vfe)) + CHKERRQ(PetscFEDestroy(&cfe)) # the vector element holds its own reference + cfe = vfe CHKERRQ(PetscObjectSetName(cfe, name.encode())) pyfe = FE() @@ -231,3 +245,79 @@ def cell_quadrature_points(DM dm, Quad quad): for d in range(cdim): ov[c - cStart, q, d] = v[q * cdim + d] return out + + +def tabulate_with_derivatives(FE fe, points): + r"""Tabulate ``fe``'s basis and its reference gradient at ``points``. + + Returns ``(B, D)`` shaped ``(Np, Nb, Nc)`` and ``(Np, Nb, Nc, dim)``; + ``D`` is the gradient with respect to the reference coordinates, so a + physical gradient is ``invJ^T D``. + """ + cdef PetscTabulation T = NULL + cdef PetscInt Np, Nb, Nc, cdim, p, b, c, d + pts = np.ascontiguousarray(points, dtype=np.float64) + if pts.ndim != 2: + raise ValueError("points must be (Np, dim)") + cdef double[:, ::1] pv = pts + Np = pts.shape[0] + CHKERRQ(PetscFECreateTabulation(fe.fe, 1, Np, &pv[0, 0], 1, &T)) + Nb = T.Nb + Nc = T.Nc + cdim = T.cdim + B = np.empty((Np, Nb, Nc), dtype=np.float64) + D = np.empty((Np, Nb, Nc, cdim), dtype=np.float64) + cdef double[:, :, ::1] bv = B + cdef double[:, :, :, ::1] dv = D + for p in range(Np): + for b in range(Nb): + for c in range(Nc): + bv[p, b, c] = T.T[0][(p * Nb + b) * Nc + c] + for d in range(cdim): + dv[p, b, c, d] = T.T[1][((p * Nb + b) * Nc + c) * cdim + d] + CHKERRQ(PetscTabulationDestroy(&T)) + return B, D + + +def cell_affine_maps(DM dm): + r"""Affine reference map of every local cell. + + Returns ``(v0, invJ, detJ)`` shaped ``(ncells, cdim)``, ``(ncells, cdim, + cdim)`` and ``(ncells,)`` in local cell order, from + ``DMPlexComputeCellGeometryFEM`` with no rule (the affine map). The + reference coordinate of a physical point ``x`` in cell ``c`` is + ``invJ[c] @ (x - v0[c]) - 1`` in PETSc's ``[-1, 1]`` reference frame + (``v0`` is the image of the reference corner ``(-1, ..., -1)``), which is + the frame :func:`tabulate` expects. + """ + cdef PetscInt cStart = 0, cEnd = 0, cdim = 0, c, d, e + cdef PetscReal *v = NULL + cdef PetscReal *J = NULL + cdef PetscReal *invJ = NULL + cdef PetscReal detJ = 0.0 + CHKERRQ(DMPlexGetHeightStratum(dm.dm, 0, &cStart, &cEnd)) + CHKERRQ(DMGetCoordinateDim(dm.dm, &cdim)) + ncells = cEnd - cStart + v0 = np.empty((ncells, cdim), dtype=np.float64) + iJ = np.empty((ncells, cdim, cdim), dtype=np.float64) + dJ = np.empty((ncells,), dtype=np.float64) + cdef double[:, ::1] v0v = v0 + cdef double[:, :, ::1] iJv = iJ + cdef double[::1] dJv = dJ + vbuf = np.empty(cdim, dtype=np.float64) + Jbuf = np.empty(cdim * cdim, dtype=np.float64) + iJbuf = np.empty(cdim * cdim, dtype=np.float64) + cdef double[::1] vv = vbuf + cdef double[::1] Jv = Jbuf + cdef double[::1] iJvb = iJbuf + if ncells == 0: + return v0, iJ, dJ + v = &vv[0]; J = &Jv[0]; invJ = &iJvb[0] + for c in range(cStart, cEnd): + CHKERRQ(DMPlexComputeCellGeometryFEM(dm.dm, c, NULL, v, J, invJ, &detJ)) + dJv[c - cStart] = detJ + for d in range(cdim): + v0v[c - cStart, d] = v[d] + for e in range(cdim): + iJv[c - cStart, d, e] = invJ[d * cdim + e] + return v0, iJ, dJ diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index 4a277851..cbea9a68 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -3525,7 +3525,9 @@ class _BaseIntegrationPointVariable(_BaseMeshVariable): ``coords`` are the physical integration points in the same order. Derivatives of the symbol are meaningless (the tabulated gradient is zero) - and the JIT refuses them. Scalar components only for now. + and the JIT refuses them. Any number of components: the element is the + scalar delta element wrapped as a vector element, dofs point-major and + component-minor within a cell. """ is_integration_point = True @@ -3551,15 +3553,11 @@ def _basis_key(self): def _create_petsc_fe(self, dim, prefix): from underworld3.cython.petsc_quadrature_fe import create_delta_fe - if self.num_components != 1: - raise NotImplementedError( - "IntegrationPointVariable: scalar components only for now - " - "use one variable per component" - ) cStart, _ = self.mesh.dm.getHeightStratum(0) - fe = create_delta_fe(self.mesh.integration_rule, self.mesh.dm.getCellType(cStart), - name=f"{prefix}integration_point_fe") - return fe + return create_delta_fe( + self.mesh.integration_rule, self.mesh.dm.getCellType(cStart), + name=f"{prefix}integration_point_fe", num_components=self.num_components, + ) # -- geometry --------------------------------------------------------------- diff --git a/src/underworld3/discretisation/enhanced_variables.py b/src/underworld3/discretisation/enhanced_variables.py index 385f60cc..7c3b8e3c 100644 --- a/src/underworld3/discretisation/enhanced_variables.py +++ b/src/underworld3/discretisation/enhanced_variables.py @@ -66,6 +66,7 @@ class EnhancedMeshVariable(DimensionalityMixin, MathematicalMixin): # The storage class this wrapper delegates to; IntegrationPointVariable # swaps in the quadrature-point element. _base_variable_class = _BaseMeshVariable + is_integration_point = False def __new__(cls, varname, mesh, *args, **kwargs): """Custom __new__ to ensure proper initialization and registration.""" @@ -947,7 +948,8 @@ class IntegrationPointVariable(EnhancedMeshVariable): nearest-integration-point partition of each cell; ``uw.function.evaluate`` returns that, so a query agrees with what the assembler used at the same point. Derivatives of the symbol are refused by the JIT (the gradient is - identically zero). Scalar components only for now. + identically zero). Vector and tensor variables are supported (one dof + per component per point). Examples -------- diff --git a/src/underworld3/function/_function.pyx b/src/underworld3/function/_function.pyx index 6a144f25..3fb4e916 100644 --- a/src/underworld3/function/_function.pyx +++ b/src/underworld3/function/_function.pyx @@ -361,7 +361,7 @@ def global_evaluate_nd( expr, coords=None, coord_sys=None, other_arguments=None, - simplify=True, + simplify=False, verbose=False, evalf=False, rbf=False, @@ -514,7 +514,10 @@ def global_evaluate_nd( expr, evaluation_swarm.migrate(remove_sent_points=True, delete_lost_points=False) local_coords = evaluation_swarm._particle_coordinates.array[...].reshape(-1,evaluation_swarm.cdim) - values, extrapolated = evaluate_nd(expr, local_coords, rbf=rbf, evalf=evalf, verbose=verbose, check_extrapolated=True,) + # Forward `simplify`: without it the local evaluator's default (True) ran + # sympy.simplify on every call for any expression holding a mesh variable + # (14 of 25 s in a semi-Lagrangian step with a tanh velocity, 2026-09-08). + values, extrapolated = evaluate_nd(expr, local_coords, rbf=rbf, evalf=evalf, verbose=verbose, check_extrapolated=True, simplify=simplify,) if local_coords.shape[0] > 0: data_container.array[...] = values[...] @@ -618,7 +621,7 @@ def global_evaluate_nd( expr, # This rank's local rbf extrapolation of the global set. NON-collective # value path — see DEADLOCK SAFETY above (must be rbf=True, never FE). ext_vals, ext_flag = evaluate_nd( - expr, all_ext, rbf=True, evalf=False, verbose=False, + expr, all_ext, rbf=True, evalf=False, verbose=False, simplify=simplify, check_extrapolated=True,) ext_vals = np.ascontiguousarray( np.asarray(ext_vals, dtype=np.float64).reshape((n_ext_total,) + expr_shape)) @@ -916,7 +919,7 @@ def evaluate_nd( expr, coords=None, coord_sys=None, other_arguments=None, - simplify=True, + simplify=False, verbose=False, evalf=False, rbf=False, @@ -1113,7 +1116,7 @@ def petsc_interpolate( expr, coord_sys=None, mesh=None, other_arguments=None, - simplify=True, + simplify=False, verbose=False, cell_hints=None, ): """ @@ -1516,7 +1519,7 @@ def rbf_evaluate( expr, mesh=None, other_arguments=None, verbose=False, - simplify=True,): + simplify=False,): """ Evaluate a given expression at a list of coordinates. diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index 9788c447..ca6a93fb 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -143,6 +143,27 @@ class SwarmVariable(DimensionalityMixin, MathematicalMixin, Stateful, uw_object) If None, inferred from ``size``. dtype : type, default=float Data type for storage (float or int). + proxy_location : {"nodes", "integration_points", "cells"}, default="nodes" + Where the proxy lives. ``"nodes"``: a nodal mesh variable of + ``proxy_degree`` / ``proxy_continuous``, reconstructed from the + particles at its nodes and interpolated by the basis wherever the + weak form reads it. ``"integration_points"``: an + :class:`~underworld3.discretisation.IntegrationPointVariable` + reconstructed from the nearest particles at every integration point + and read there directly, with no second interpolation (the + Ellipsis / Underworld PIC-LIP mapping); a material interface keeps + its sub-cell position, and the proxy has no gradient (a derivative + of its symbol is refused). ``proxy_degree`` / ``proxy_continuous`` + are ignored in that case. ``"cells"``: a discontinuous mesh variable + of ``proxy_degree`` holding, in every cell, the least-squares + polynomial through the particles that cell holds (a thin cell takes + a linear fit to the particles nearest its centroid, an empty cell + keeps its previous value). Exact for + polynomial particle fields up to ``proxy_degree``, integrated exactly + by the default rule, sharp at cell edges, with a gradient, and no + neighbour search across ranks; see + :class:`~underworld3.utilities.cell_polynomial_projection.CellPolynomialProjector`. + ``proxy_continuous`` is ignored (the proxy is discontinuous). proxy_degree : int, default=1 Polynomial degree for the mesh proxy variable. proxy_continuous : bool, default=True @@ -192,6 +213,7 @@ def __init__( dtype=float, proxy_degree=1, proxy_continuous=True, + proxy_location="nodes", _register=True, _proxy=True, varsymbol=None, @@ -376,6 +398,13 @@ def __init__( self._vtype = vtype self._proxy_degree = proxy_degree self._proxy_continuous = proxy_continuous + if proxy_location not in ("nodes", "integration_points", "cells"): + raise ValueError( + "proxy_location must be 'nodes', 'integration_points' or 'cells', " + f"not {proxy_location!r}" + ) + self._cell_projector = None + self._proxy_location = proxy_location self._create_proxy_variable() # Inert: kept for backward compatibility with the removed @@ -1078,22 +1107,50 @@ def _create_proxy_variable(self): # var stale via _mark_reinit_stale; we wire that callback # below to set ``self._proxy_stale = True`` so the next # access re-projects. - self._meshVar = uw.discretisation.MeshVariable( - "proxy_" + self.clean_name, - self.swarm.mesh, - self.shape, - self._vtype, - degree=self._proxy_degree, - continuous=self._proxy_continuous, - varsymbol=r"\left<" + self.symbol + r"\right>", - remesh_policy="reinit", - # The proxy is what `var.sym` resolves to, so it advertises - # the same units as the variable it stands for. Without this, - # evaluating a proxied symbol returned the NON-DIMENSIONAL - # number with no units attached, as though it were the answer - # (issue #439). Stored data stays non-dimensional either way. - units=self._units, - ) + if getattr(self, "_proxy_location", "nodes") == "integration_points": + # Particles -> integration points directly: the assembler reads + # the stored values at the rule with no basis interpolation. + self._meshVar = uw.discretisation.IntegrationPointVariable( + "proxy_" + self.clean_name, + self.swarm.mesh, + self.shape, + self._vtype, + varsymbol=r"\left<" + self.symbol + r"\right>_q", + remesh_policy="reinit", + units=self._units, + ) + elif getattr(self, "_proxy_location", "nodes") == "cells": + # Particles -> a polynomial per cell (least squares); read by + # the assembler through the ordinary discontinuous basis. + self._meshVar = uw.discretisation.MeshVariable( + "proxy_" + self.clean_name, + self.swarm.mesh, + self.shape, + self._vtype, + degree=self._proxy_degree, + continuous=False, + varsymbol=r"\left<" + self.symbol + r"\right>_c", + remesh_policy="reinit", + units=self._units, + ) + self._cell_projector = None + else: + self._meshVar = uw.discretisation.MeshVariable( + "proxy_" + self.clean_name, + self.swarm.mesh, + self.shape, + self._vtype, + degree=self._proxy_degree, + continuous=self._proxy_continuous, + varsymbol=r"\left<" + self.symbol + r"\right>", + remesh_policy="reinit", + # The proxy is what `var.sym` resolves to, so it advertises + # the same units as the variable it stands for. Without this, + # evaluating a proxied symbol returned the NON-DIMENSIONAL + # number with no units attached, as though it were the answer + # (issue #439). Stored data stays non-dimensional either way. + units=self._units, + ) # The remesh helper calls this on REINIT vars after an # adapt. Bound here so the closure captures ``self`` (the # SwarmVariable) rather than the proxy MeshVariable. @@ -1158,13 +1215,66 @@ def _update_proxy_if_stale(self): try: self._updating_proxy = True - self._rbf_to_meshVar(self._meshVar) + if getattr(self, "_proxy_location", "nodes") == "cells": + self._cells_to_meshVar(self._meshVar) + else: + self._rbf_to_meshVar(self._meshVar) self._proxy_stale = False # Mark as fresh finally: self._updating_proxy = False return + def _cells_to_meshVar(self, meshVar): + """Refresh a ``proxy_location="cells"`` proxy: a least-squares polynomial + per cell through the particles it holds (see + :class:`~underworld3.utilities.cell_polynomial_projection.CellPolynomialProjector`). + + Rank-local: each rank fits the cells it holds from the particles it + holds; the proxy's own ghost synchronisation delivers owned values + to the neighbours. The starved-rank guard and the collective + read-then-write sequence mirror :meth:`_rbf_to_meshVar`. + """ + from underworld3.utilities.cell_polynomial_projection import CellPolynomialProjector + + if meshVar.mesh != self.swarm.mesh: + if hasattr(self, "_meshVar") and meshVar is self._meshVar: + self._create_proxy_variable() + meshVar = self._meshVar + else: + raise RuntimeError("Cannot map a swarm to a different mesh") + + current_values = np.array(meshVar.data[...], copy=True) + + if self.swarm.local_size <= 1: + if self.swarm._population_generation > 0: + import warnings + + warnings.warn( + f"Swarm proxy update: rank {uw.mpi.rank} holds " + f"{max(self.swarm.local_size, 0)} particles; proxy variable " + f"'{getattr(meshVar, 'clean_name', meshVar.name)}' left " + "unchanged on this rank.", + stacklevel=2, + ) + Values = current_values + else: + projector = self._cell_projector + if ( + projector is None + or projector.var is not meshVar + or projector.mesh_version != self.swarm.mesh._mesh_version + ): + projector = CellPolynomialProjector(meshVar) + self._cell_projector = projector + raw_data = self.unpack_raw_data_from_petsc(squeeze=False) + Values = projector.fit( + self.swarm._particle_coordinates.data, raw_data, old=current_values + ) + + meshVar.data[...] = Values[...] + return + # Maybe rbf_interpolate for this one and meshVar is a special case def _rbf_to_meshVar(self, meshVar, nnn=None, verbose=False, order=1, monotone=False): @@ -1473,6 +1583,7 @@ def _object_viewer(self): > symbol: ${self.symbol}$\n > shape: ${self.shape}$\n > proxy: ${self._proxy}$\n + > proxy_location: `{self._proxy_location}`\n > proxy_degree: ${self._proxy_degree}$\n > proxy_continuous: `{self._proxy_continuous}`\n > type: `{self.vtype.name}`""" @@ -2880,6 +2991,10 @@ def __init__(self, mesh, recycle_rate=0, verbose=False, clip_to_mesh=True): ) self._X0_uninitialised = True + # Callables run at the top of advection(), before any particle moves: + # a Lagrangian history registers its first sampling here so it sees + # the field at the launch positions, not at the landing ones. + self._pre_advection_hooks = [] self._index = None # Particle -> proxy-node transfer operators, keyed by geometry and # stencil and shared by every proxied variable of this swarm. Entries @@ -4917,6 +5032,9 @@ def advection( if uw.mpi.rank == 0 and self.verbose: print(f"Substepping {substeps} / {abs(delta_t) / dt_limit}, {delta_t} ") + for hook in list(getattr(self, "_pre_advection_hooks", ())): + hook() + # X0 holds the particle location at the start of advection # This is needed because the particles may be migrated off-proc # during timestepping. Probably not needed - use global evaluation instead diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index fed74ea7..180e4346 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -753,22 +753,25 @@ def _velocity_degree(self): degs = [fn.meshvar().degree for fn in varfns] return max(degs) if degs else 2 - def _make_velocity_level(self, tag): - """A cached velocity level: ``V_fn`` EVALUATED at the true nodes of a - vector field (no nudge; the evaluator is exact at node coordinates on - simplex, quad and annulus meshes). Caching by evaluation, rather than - by substituting snapshots of the mesh variables into the expression, - is what captures everything ``V_fn`` depends on at that time: the - variables, constants that ramp, swarm proxies, the mesh geometry.""" - snap = uw.discretisation.MeshVariable( - f"vcache_{tag}_{self.instance_number}", self.mesh, self.mesh.dim, - degree=self._velocity_degree(), continuous=True, - varsymbol=rf"{{ V^{{ ({tag}) }}_{{ [{self.instance_number}] }} }}", - units=self._velocity_units(), # same frame as V_fn, so 1.5 v - 0.5 v_prev is consistent - ) - snap.remesh_policy = RemeshPolicy.CARRY - snap._remesh_managed_by = self - return {"var": snap, "expr": snap.sym} + def attach_characteristics(self, trace): + """Trace along ``trace`` (a :class:`CharacteristicTrace` on this + mesh and velocity) instead of a private one. The owner of the trace + (a solver) delimits the steps with ``begin_step`` / ``finish_step``.""" + self._characteristics = trace + self._owns_characteristics = False + + @property + def characteristics(self): + """The :class:`CharacteristicTrace` this history samples from; a + private one is created on first use when no solver shared one.""" + tr = getattr(self, "_characteristics", None) + if tr is None: + tr = CharacteristicTrace( + self.mesh, self.V_fn, midtime_velocity=getattr(self, "midtime_velocity", True) + ) + self._characteristics = tr + self._owns_characteristics = True + return tr def _velocity_units(self): """Units of ``V_fn`` under an active units model, else None.""" @@ -777,17 +780,6 @@ def _velocity_units(self): units = None return units - def _copy_velocity_level(self, dst, src=None): - """``dst`` <- ``src`` (another level) or, with ``src=None``, ``V_fn`` - evaluated now at ``dst``'s nodes (reduced to the non-dimensional - frame ``.data`` stores, issue #267).""" - if src is None: - vals = uw.function.evaluate(self._V_matrix(), np.asarray(dst["var"].coords_nd)) - vals = _to_nondim_ndarray(vals, units=self._velocity_units()) - dst["var"].data[...] = np.asarray(vals).reshape(-1, self.mesh.dim) - else: - dst["var"].data[...] = src["var"].data[...] - def bdf(self, order: Optional[int] = None): r"""Backward differentiation approximation of the time-derivative of :math:`\psi`. @@ -1846,6 +1838,247 @@ def _object_viewer(self): display(Latex(rf"$\quad$ integrator: {self.integrator}, tau shape: {self.tau_shape}")) +class CharacteristicTrace: + r"""Departure points and cached velocity levels for one advecting + velocity on one mesh, shared by every history that follows it. + + The mathematics of a history term stays with the term: its symbols, + its BDF / Adams-Moulton stencils, its view. What the terms share is the + evaluation behind them: the characteristic traced back from a node set, + and the velocity levels :math:`v^{n-1}, v^{n-2}, \dots` cached by + EVALUATION at the true nodes, so any expression for ``V_fn`` (a + variable, ``-v``, ``v/2``, ``c(t)\,v``) is carried as it was then. + + A solver creates one trace per advecting velocity and attaches it to + each history manager (:meth:`_DDtBase.attach_characteristics`); a + manager used on its own owns a private one. Within a step every request + for the same node set and the same segment structure is served from the + cache, so two histories on the same nodes cost one trace, and the older + slots of a one-segment history cost nothing beyond the first. + + Steps are delimited by :meth:`begin_step` (clears the cache) and + :meth:`finish_step` (records the velocity used this step as + :math:`v^{n-1}`), called by the solver that owns the trace, or by the + manager when the trace is private. + + A segment is ``("first", 0, dt)``: from the launch points, start velocity + :math:`v^n` (``V_fn`` live), mid-time velocity + :math:`\tfrac32 v^n - \tfrac12 v^{n-1}` (``V_fn`` alone until a level is + recorded, or with ``midtime_velocity=False``); or ``("older", k, dt)``: + one more step back through cached levels, start :math:`v^{n-k}`, mid + :math:`\tfrac12 (v^{n-k} + v^{n-k-1})`. A chain of segments is cached by + prefix, so slot ``k`` of a multi-segment history extends slot ``k-1``. + """ + + _next_instance = 0 + + def __init__(self, mesh, V_fn, midtime_velocity=True): + self.mesh = mesh + self.V_fn = V_fn + self.midtime_velocity = midtime_velocity + self.instance_number = CharacteristicTrace._next_instance + CharacteristicTrace._next_instance += 1 + self._levels = [None] # index k >= 1 holds v^{n-k} + self._levels_valid = 0 # levels 1..valid hold a recorded velocity + self._cache = {} + self._step = 0 + self._dt = None + self.n_velocity_evaluations = 0 + self.n_cache_hits = 0 + + # -- the velocity --------------------------------------------------------- + + def V_matrix(self): + V = self.V_fn + if hasattr(V, "sym") and not isinstance(V, sympy.Basic): + return sympy.Matrix(V.sym) + return sympy.Matrix(V) + + def velocity_degree(self): + _, varfns, _ = uw.function.expressions.mesh_vars_in_expression(self.V_matrix()) + degs = [fn.meshvar().degree for fn in varfns] + return max(degs) if degs else 2 + + def velocity_units(self): + units = uw.get_units(self.V_matrix()) + if units is not None and not uw.get_default_model().has_units(): + units = None + return units + + def _new_level(self, k): + snap = uw.discretisation.MeshVariable( + f"vtrace_{self.instance_number}_n{k}", self.mesh, self.mesh.dim, + degree=self.velocity_degree(), continuous=True, + varsymbol=rf"{{ V^{{ (n-{k}) }}_{{ [{self.instance_number}] }} }}", + units=self.velocity_units(), + ) + snap.remesh_policy = RemeshPolicy.CARRY + snap._remesh_managed_by = self + return {"var": snap, "expr": snap.sym} + + def ensure_levels(self, n): + """Make cached levels ``1 .. n-1`` exist (level 0 is ``V_fn`` live).""" + while len(self._levels) < n: + self._levels.append(self._new_level(len(self._levels))) + + def level_expr(self, k): + if k == 0: + return self.V_matrix() + self.ensure_levels(k + 1) + return self._levels[k]["expr"] + + def level_valid(self, k): + return k == 0 or k <= self._levels_valid + + def _evaluate_into(self, level): + vals = uw.function.evaluate(self.V_matrix(), np.asarray(level["var"].coords_nd)) + vals = _to_nondim_ndarray(vals, units=self.velocity_units()) + level["var"].data[...] = np.asarray(vals).reshape(-1, self.mesh.dim) + + def initialise_levels(self, n): + """Every cached level ``1 .. n-1`` <- ``V_fn`` now (a history starting + from rest in time: the older velocities are the current one).""" + self.ensure_levels(n) + if n > 1: + self._evaluate_into(self._levels[1]) + for k in range(2, n): + self._levels[k]["var"].data[...] = self._levels[1]["var"].data[...] + self._levels_valid = max(self._levels_valid, n - 1) + + def record_velocity(self): + """Shift the cached levels and record ``V_fn`` now as :math:`v^{n-1}`.""" + if len(self._levels) < 2: + return + for k in range(len(self._levels) - 1, 1, -1): + self._levels[k]["var"].data[...] = self._levels[k - 1]["var"].data[...] + self._evaluate_into(self._levels[1]) + self._levels_valid = min(self._levels_valid + 1, len(self._levels) - 1) + + def midtime_expr(self): + r""":math:`\tfrac32 v^n - \tfrac12 v^{n-1}`, or :math:`v^n` alone.""" + if not self.midtime_velocity: + return self.V_matrix() + self.ensure_levels(2) + if not self.level_valid(1): + return self.V_matrix() + return self.V_matrix() * sympy.Rational(3, 2) - self.level_expr(1) * sympy.Rational(1, 2) + + def velocity_at(self, expr, coords, use_global=False, evalf=False, + subtract_v_mesh=False, v_mesh_var=None): + """``expr`` (a velocity expression) at ``coords``, as a plain + non-dimensional ``(N, dim)`` array. Node points are rank-local and + use ``evaluate``; points that may have left the partition route + through ``global_evaluate``. With ``subtract_v_mesh`` the mesh + velocity sampled at the same points is removed (ALE), after the + evaluation so it inherits the unit treatment of ``V_fn``.""" + self.n_velocity_evaluations += 1 + if use_global: + v_result = uw.function.global_evaluate(expr, coords, evalf=evalf) + if subtract_v_mesh: + v_result = v_result - uw.function.global_evaluate(v_mesh_var.sym, coords, evalf=evalf) + else: + v_result = uw.function.evaluate(expr, coords) + if subtract_v_mesh: + v_result = v_result - uw.function.evaluate(v_mesh_var.sym, coords) + if isinstance(v_result, UnitAwareArray): + v_at_pts = v_result[:, 0, :] + if not isinstance(v_at_pts, UnitAwareArray): + v_at_pts = UnitAwareArray(v_at_pts, units=v_result.units) + else: + v_at_pts = np.asarray(v_result) + v_at_pts = v_at_pts[:, 0, :] if v_at_pts.ndim == 3 else v_at_pts + out = _to_nondim_ndarray(v_at_pts, units=uw.get_units(self.V_fn)) + return np.asarray(out).reshape(coords.shape[0], self.mesh.dim) + + # -- the steps and the trace ----------------------------------------------- + + def begin_step(self, dt): + self._step += 1 + self._dt = dt + self._cache.clear() + + def finish_step(self): + """The velocity used this step becomes :math:`v^{n-1}`.""" + self.record_velocity() + + def _segment_exprs(self, seg): + kind, k, _ = seg + if kind == "first": + return self.V_matrix(), self.midtime_expr() + half = sympy.Rational(1, 2) + return self.level_expr(k - 1), (self.level_expr(k - 1) + self.level_expr(k)) * half + + def departure_points(self, key, X0, segments, evalf=False, X_eval=None, + clamp_final=True, subtract_v_mesh=False, v_mesh_var=None): + r"""Trace ``X0`` back through ``segments`` (RK2 midpoint each): + ``x_mid = x - dt/2 v_start(x)``, ``x_dep = x - dt v_mid(x_mid)``. + + ``key`` names the launch node set (a variable's ``_basis_key`` plus + a tag for the nudge); ``X_eval`` are the points where the first + segment's start velocity is evaluated when they differ from ``X0`` + (the centroid-nudged nodes of the nodal history). Midpoints are + clamped to the domain; the last point is clamped unless + ``clamp_final`` is False (old-frame reach-back). + """ + clamp = self.mesh.return_coords_to_bounds + vm = id(v_mesh_var) if subtract_v_mesh else None + X = np.asarray(X0) + for j in range(len(segments)): + last = j == len(segments) - 1 + clamp_this = clamp_final or not last + ckey = (key, tuple(segments[: j + 1]), clamp_this, subtract_v_mesh, vm) + hit = self._cache.get(ckey) + if hit is not None: + self.n_cache_hits += 1 + X = hit + continue + kind, k, dt = segments[j] + v_start, v_mid = self._segment_exprs(segments[j]) + X_start = X_eval if (j == 0 and X_eval is not None) else X + v0 = self.velocity_at(v_start, X_start, use_global=j > 0, evalf=evalf, + subtract_v_mesh=subtract_v_mesh, v_mesh_var=v_mesh_var) + Xm = X - v0 * (0.5 * dt) + if clamp is not None: + Xm = clamp(Xm) + vmid = self.velocity_at(v_mid, Xm, use_global=True, evalf=evalf, + subtract_v_mesh=subtract_v_mesh, v_mesh_var=v_mesh_var) + X = X - vmid * dt + if clamp is not None and clamp_this: + X = clamp(X) + self._cache[ckey] = X + return X + + +def share_characteristics(*managers, midtime_velocity=None): + """One :class:`CharacteristicTrace` for every manager that traces along + the same velocity on the same mesh; managers that do not trace (Eulerian, + swarm, symbolic) are left alone. Returns the trace, or None.""" + tracers = [m for m in managers + if m is not None and hasattr(m, "attach_characteristics") and getattr(m, "V_fn", None) is not None] + if not tracers: + return None + first = tracers[0] + same = [m for m in tracers + if m.mesh is first.mesh and (m.V_fn is first.V_fn or sympy.Matrix(_matrix_of(m.V_fn)) == sympy.Matrix(_matrix_of(first.V_fn)))] + if midtime_velocity is None: + midtime_velocity = getattr(first, "midtime_velocity", True) + trace = CharacteristicTrace(first.mesh, first.V_fn, midtime_velocity=midtime_velocity) + for m in same: + m.attach_characteristics(trace) + return trace + + +def _basis_key_of(var): + """A variable's node-set key (the enhanced wrapper hides underscore names).""" + return getattr(var, "_base_var", var)._basis_key + + +def _matrix_of(V): + if hasattr(V, "sym") and not isinstance(V, sympy.Basic): + return V.sym + return V + + class SemiLagrangian(_DDtBase): r""" Semi-Lagrangian history manager. @@ -2700,21 +2933,15 @@ def _midtime_velocity_expr(self): carried as it was then.""" if not getattr(self, "midtime_velocity", True): return None - level = getattr(self, "_v_prev_level", None) - if level is None or not getattr(self, "_v_prev_valid", False): - return None - return self._V_matrix() * sympy.Rational(3, 2) - level["expr"] * sympy.Rational(1, 2) + return self.characteristics.midtime_expr() def _record_velocity_history(self): """Cache ``V_fn`` evaluated at the true nodes as v^{n-1} for the next step. Evaluating at NUDGED nodes left a 0.001 h |grad v| bias that the extrapolation fed into every trace and moved the Blankenbach 1a wall Nusselt number by 0.9 %.""" - if getattr(self, "_v_prev_level", None) is None: - self._v_prev_level = self._make_velocity_level("n-1") - self._v_prev_valid = False - self._copy_velocity_level(self._v_prev_level) - self._v_prev_valid = True + if getattr(self, "_owns_characteristics", True): + self.characteristics.finish_step() def _centroid_shifted_var_coords(self, var): """ND node coordinates of ``var`` nudged 0.1 % toward their cell @@ -2763,31 +2990,10 @@ def _velocity_nd_at( subtraction inherits the same unit treatment as ``V_fn``. """ fn = self._V_matrix() if expr is None else expr - if use_global: - v_result = uw.function.global_evaluate(fn, coords, evalf=evalf) - if subtract_v_mesh: - v_mesh = uw.function.global_evaluate( - self._v_mesh_var.sym, coords, evalf=evalf - ) - v_result = v_result - v_mesh - else: - v_result = uw.function.evaluate(fn, coords) - if subtract_v_mesh: - v_mesh = uw.function.evaluate(self._v_mesh_var.sym, coords) - v_result = v_result - v_mesh - - # Slicing can drop the UnitAwareArray wrapper — rewrap before the - # ND reduction so the units are not silently lost. - if isinstance(v_result, UnitAwareArray): - v_at_pts = v_result[:, 0, :] - if not isinstance(v_at_pts, UnitAwareArray): - v_at_pts = UnitAwareArray(v_at_pts, units=v_result.units) - else: - v_at_pts = v_result[:, 0, :] - - # Non-dimensionalise to the DM/ND space: the trace-back arithmetic - # and the subsequent point-location both work in ND coordinates. - return _to_nondim_ndarray(v_at_pts, units=uw.get_units(self.V_fn)) + return self.characteristics.velocity_at( + fn, coords, use_global=use_global, evalf=evalf, + subtract_v_mesh=subtract_v_mesh, v_mesh_var=getattr(self, "_v_mesh_var", None), + ) def update( self, @@ -3008,52 +3214,21 @@ def _trace_departure_points( any foot that falls outside the old mesh, matching the validated prototype, which omits this clamp). """ - # Use shifted ND coords to avoid quad mesh boundary issues - # (node_coords_nd is slightly shifted toward cell centroids — - # see _centroid_shifted_node_coords) - v_at_node_pts = self._velocity_nd_at( - node_coords_nd, subtract_v_mesh=subtract_v_mesh - ) - - # Departure point in the mesh's ND (DM) coordinate space. coords_nd is - # the ND reduction of the (possibly dimensional) node coordinates — - # identical to .coords for a non-units model, and the DM-space values - # (0..L_model) when units are active, matching what global_evaluate / - # the DM point-location expect. See #267. - coords = np.asarray(self.psi_star[i].coords_nd) - - # CRITICAL (2025-11-27): Multiply velocity FIRST so UnitAwareArray.__mul__ handles it. - # If we do `dt_for_calc * v_at_node_pts`, Pint handles it and loses UnitAwareArray units. - mid_pt_coords = coords - v_at_node_pts * (0.5 * dt_for_calc) - - # Clamp midpoint coordinates to the domain boundary - if self.mesh.return_coords_to_bounds is not None: - mid_pt_coords = self.mesh.return_coords_to_bounds(mid_pt_coords) - - # Mid-point velocities may lie off-rank, so route through - # global_evaluate (with evalf forwarded), unlike the on-node - # evaluation above. The mid-point velocity is taken at the mid - # TIME, t^{n+1/2}, by extrapolation from the two most recent - # velocity fields, 1.5 v^n - 0.5 v^{n-1}; with v^n alone the - # trace is only first order in an unsteady flow. On the first - # step (no previous velocity) v^n is used. - v_at_mid_pts = self._velocity_nd_at( - mid_pt_coords, - use_global=True, + # One RK2 segment from the true nodes, the start velocity taken at + # the centroid-nudged coordinates (node_coords_nd). Served from the + # shared trace: a second history on the same nodes, or an older slot + # of this one, reuses the departure points computed here. + return self.characteristics.departure_points( + (_basis_key_of(self.psi_star[i]), "nudged"), + np.asarray(self.psi_star[i].coords_nd), + (("first", 0, dt_for_calc),), evalf=evalf, + X_eval=node_coords_nd, + clamp_final=not oldframe_active, subtract_v_mesh=subtract_v_mesh, - expr=self._midtime_velocity_expr(), + v_mesh_var=getattr(self, "_v_mesh_var", None), ) - # Upstream (departure) coordinates: current position - velocity * timestep - end_pt_coords = coords - v_at_mid_pts * dt_for_calc - - if (self.mesh.return_coords_to_bounds is not None - and not oldframe_active): - end_pt_coords = self.mesh.return_coords_to_bounds(end_pt_coords) - - return end_pt_coords - def _sample_history_at_departure( self, i, end_pt_coords, evalf, monotone_mode, oldframe_active, oldframe_X ): @@ -3169,6 +3344,13 @@ def update_pre_solve( if not self._history_initialised: self.initialise_history() + # A private trace delimits its own step; a shared one is delimited + # by the solver that owns it (begin before the first manager, finish + # after the last, so every history sees the same velocity). + trace = self.characteristics + if self._owns_characteristics: + trace.begin_step(dt) + # Old-frame reach-back (mutually exclusive with the ALE pulse: # ``on_remesh`` stashes ``_oldframe_X`` INSTEAD of a v_mesh disp, # so ``_ale_active`` is False below whenever this is True). When @@ -3718,6 +3900,14 @@ class Lagrangian_Swarm(_DDtBase): Order of time integration (1-3) (default ``1``). smoothing : float, optional Smoothing parameter (default ``0.0``). + proxy_location : {"nodes", "integration_points", "cells"}, optional + Where each history slot's proxy lives; ``"integration_points"`` + reconstructs the particle history at the integration points and the + weak form reads it there directly, with no nodal proxy and no basis + interpolation (the Ellipsis / Underworld PIC-LIP mapping); + ``"cells"`` fits a least-squares polynomial of degree ``degree`` per + cell, exact for polynomial histories and integrated exactly by the + default rule. step_averaging : int, optional Number of steps for history averaging (default ``2``). @@ -3769,10 +3959,13 @@ def __init__( continuous: bool, varsymbol: Optional[str] = r"u", verbose: Optional[bool] = False, - bcs=[], + bcs=None, order=1, smoothing=0.0, step_averaging=2, + proxy_location="nodes", + particle_update="pic", + residual_retention=1.0, ): super().__init__() @@ -3782,9 +3975,45 @@ def __init__( self.verbose = verbose self.order = order self.step_averaging = step_averaging + if particle_update not in ("pic", "flip"): + raise ValueError(f"particle_update must be 'pic' or 'flip', not {particle_update!r}") + # "pic": after a solve every particle takes the mesh solution at its + # position (blended over step_averaging steps), so sub-cell particle + # detail is re-projected away each step. "flip": the particle keeps + # its own value and adds the mesh INCREMENT, solution minus the proxy + # the mesh saw, evaluated at the particle; the sub-cell residual + # survives, scaled by residual_retention (1 = FLIP, 0 = PIC; set it + # to exp(-kappa dt pi^2 / l^2) to let a diffusing residual decay). + self.particle_update = particle_update + self.residual_retention = residual_retention + # "integration_points": each slot's proxy is an IntegrationPointVariable + # reconstructed from the particles at the rule and read there directly + # (the Ellipsis / Underworld PIC-LIP mapping); no nodal proxy. + # "cells": each slot's proxy is a least-squares polynomial per cell + # (discontinuous, degree `degree`), exact for polynomial histories + # and integrated exactly by the default rule. + self.proxy_location = proxy_location self._init_history_tracking(order) + # Sample the history before the particles first move. Left to the + # first update_pre_solve, the sampling happens AFTER the user's + # swarm.advection() and the first step transports nothing (a + # one-step lag, measured as 0.05 of displacement on the rotating + # Gaussian, 2026-09-08). Weak reference: the swarm must not own us. + import weakref + + _self = weakref.ref(self) + + def _initialise_before_first_move(): + mgr = _self() + if mgr is not None and not mgr._history_initialised: + mgr.initialise_history() + + hooks = getattr(swarm, "_pre_advection_hooks", None) + if hooks is not None: + hooks.append(_initialise_before_first_move) + psi_star = [] self.psi_star = psi_star @@ -3796,6 +4025,7 @@ def __init__( vtype=vtype, proxy_degree=degree, proxy_continuous=continuous, + proxy_location=proxy_location, varsymbol=rf"{varsymbol}^{{ {'*'*(i+1)} }}", ) ) @@ -3900,6 +4130,31 @@ def update_pre_solve( return + def _proxy_values_at_particles(self, slot, coords, evalf): + """The slot's proxy evaluated at the particles, shaped like ``slot.data``. + + A ``"cells"`` proxy is read through its own fitted polynomials (exact, + no locator round trip); any other proxy through ``evaluate`` of the + proxy mesh variable's symbol. + """ + # The proxy refreshes lazily on access through the SWARM variable's + # symbol; reading its mesh variable directly bypasses that, so refresh + # first on every path, else the residual is taken against a stale fit. + slot._update_proxy_if_stale() + projector = getattr(slot, "_cell_projector", None) + if projector is not None and getattr(slot, "_proxy_location", None) == "cells": + vals = projector.interpolate(np.asarray(slot._meshVar.data), coords) + return np.nan_to_num(vals) + mv = slot._meshVar + out = np.empty((coords.shape[0], slot.data.shape[1])) + for i in range(slot.shape[0]): + for j in range(slot.shape[1]): + ij = slot._data_layout(i, j) + out[:, ij] = np.asarray( + uw.function.evaluate(mv.sym[i, j], coords, evalf=evalf) + ).reshape(-1) + return out + def update_post_solve( self, dt: float, @@ -3931,9 +4186,13 @@ def update_post_solve( phi = 1 / self.step_averaging psi_star_0 = self.psi_star[0] + coords = np.asarray(self.swarm._particle_coordinates.data) + if self.particle_update == "flip": + # The proxy the mesh saw during this solve, at the particles: the + # residual psi_p - proxy(x_p) is what the mesh never resolved. + proxy_at_p = self._proxy_values_at_particles(psi_star_0, coords, evalf) # Blend the freshly-evaluated psi into slot 0 component-by-component # through the canonical (N, components) storage (audit SWARM-06). - coords = np.asarray(self.swarm._particle_coordinates.data) for i in range(psi_star_0.shape[0]): for j in range(psi_star_0.shape[1]): ij = psi_star_0._data_layout(i, j) @@ -3944,9 +4203,13 @@ def update_post_solve( evalf=evalf, ) ).reshape(-1) - psi_star_0.data[:, ij] = ( - phi * updated_psi + (1 - phi) * psi_star_0.data[:, ij] - ) + if self.particle_update == "flip": + residual = np.asarray(psi_star_0.data[:, ij]) - proxy_at_p[:, ij] + psi_star_0.data[:, ij] = updated_psi + self.residual_retention * residual + else: + psi_star_0.data[:, ij] = ( + phi * updated_psi + (1 - phi) * psi_star_0.data[:, ij] + ) if self._n_solves_completed < self.order: self._n_solves_completed += 1 @@ -4064,8 +4327,7 @@ def __init__( # velocity is extrapolated from v^n and v^{n-1}. Each level caches # V_fn evaluated at the nodes at that time, so V_fn may be any # expression (variables, ramping constants, swarm proxies). - self._n_v = max(order, 2) - self.v_levels = [self._make_velocity_level(f"n-{k}") for k in range(self._n_v)] + self._n_v = max(order, 2) # velocity levels the segments read self._init_coefficient_expressions(order, self.theta, with_exp=False) def spatial_weights(self): @@ -4174,29 +4436,6 @@ def _record_current(self): else: vals = uw.function.evaluate(self.psi_fn[0], self._nudged_node_coords(ps)) ps.data[:, 0] = np.asarray(_to_nondim_ndarray(vals, units=self._psi_units)).reshape(-1) - self._copy_velocity_level(self.v_levels[0]) - - def _velocity_at(self, v_sym, coords, evalf): - v = uw.function.global_evaluate(v_sym, coords, evalf=evalf) - v = np.asarray(_to_nondim_ndarray(v, units=self._velocity_units())) - if v.ndim == 3: - v = v[:, 0, :] - return v.reshape(coords.shape[0], self.mesh.dim) - - def _trace_segment(self, X, v_start_sym, v_mid_sym, dt, evalf): - r"""One RK2 (midpoint) segment of the characteristic, backwards: - ``x_mid = x - dt/2 v_start(x)``, ``x_dep = x - dt v_mid(x_mid)``, - with ``v_mid`` the velocity at the segment's mid TIME.""" - clamp = self.mesh.return_coords_to_bounds - v0 = self._velocity_at(v_start_sym, X, evalf) - Xm = X - 0.5 * dt * v0 - if clamp is not None: - Xm = clamp(Xm) - vm = self._velocity_at(v_mid_sym, Xm, evalf) - Xd = X - dt * vm - if clamp is not None: - Xd = clamp(Xd) - return Xd def _segment_dt(self, j, dt): """Length of segment ``j`` (0 = the current step).""" @@ -4208,25 +4447,17 @@ def _segment_dt(self, j, dt): def _fill_slots(self, dt, evalf): """Trace back from the integration points and sample the snapshots.""" X0 = np.asarray(self.psi_star[0].coords_nd) - X = X0.copy() - half = sympy.Rational(1, 2) + trace = self.characteristics + trace.ensure_levels(self._n_v) + key = (_basis_key_of(self.psi_star[0]), "true") + segments = [] for k in range(self.order): - # Segment k extends the trace from slot k-1's feet, so the - # feet for slot k are those of slot k-1 traced one more step. - # Segment k runs from t^{n+1-k} back to t^{n-k}. Its mid-time - # velocity: for k=0 extrapolated, 1.5 v^n - 0.5 v^{n-1} (v^{n+1} - # is not known yet); for k>=1 both ends are known, so the - # average of v^{n+1-k} and v^{n-k}. The first stage, which - # only places the mid-point, uses the velocity at the - # segment's start time. - V = [lvl["expr"] for lvl in self.v_levels] - if k == 0: - v_start = V[0] - v_mid = V[0] * sympy.Rational(3, 2) - V[1] * half - else: - v_start = V[k - 1] - v_mid = (V[k - 1] + V[k]) * half - X = self._trace_segment(X, v_start, v_mid, self._segment_dt(k, dt), evalf) + # Slot k's feet are slot k-1's traced one more step back. The + # trace caches by prefix, so this extends the previous chain by + # one segment; the velocities per segment are described there. + segments.append(("first", 0, self._segment_dt(0, dt)) if k == 0 + else ("older", k, self._segment_dt(k, dt))) + X = trace.departure_points(key, X0, tuple(segments), evalf=evalf) vals = uw.function.global_evaluate( self.psi_snap[k].sym[0], X, evalf=evalf, monotone=self.monotone_mode ) @@ -4239,8 +4470,7 @@ def initialise_history(self): self._record_current() for k in range(1, self.order): self.psi_snap[k].data[...] = self.psi_snap[0].data[...] - for k in range(1, self._n_v): - self._copy_velocity_level(self.v_levels[k], self.v_levels[0]) + self.characteristics.initialise_levels(self._n_v) X = np.asarray(self.psi_star[0].coords_nd) vals = uw.function.evaluate(self.psi_snap[0].sym[0], X) vals = np.asarray(_to_nondim_ndarray(vals, units=self._psi_units)).reshape(-1) @@ -4256,10 +4486,13 @@ def update_pre_solve(self, dt, evalf=False, verbose=False, **_ignored): _update_am_values(self._am_coeffs, self.effective_order, self.theta) for k in range(self.order - 1, 0, -1): self.psi_snap[k].data[...] = self.psi_snap[k - 1].data[...] - for k in range(self._n_v - 1, 0, -1): - self._copy_velocity_level(self.v_levels[k], self.v_levels[k - 1]) + trace = self.characteristics + if self._owns_characteristics: + trace.begin_step(dt) self._record_current() self._fill_slots(dt, evalf) + if self._owns_characteristics: + trace.finish_step() def update(self, dt, evalf=False, verbose=False, **kwargs): self.update_pre_solve(dt, evalf=evalf, verbose=verbose, **kwargs) diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 524a411b..24351f47 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -4278,8 +4278,28 @@ def __init__( self.u.remesh_policy = RemeshPolicy.CARRY self.u._remesh_managed_by = self.Unknowns.DuDt + # The value and flux histories follow the same velocity from the + # same nodes: one characteristic trace serves both (each term keeps + # its own mathematics; only the departure points and the cached + # velocity levels are shared). The solver delimits the steps. + self._characteristics = uw.systems.ddt.share_characteristics( + self.Unknowns.DuDt, self.Unknowns.DFDt + ) + return + def _flux_history_is_read(self): + """Whether the weak form reads the old-level flux this step: only the + theta rule with theta < 1 does (BDF orders and theta = 1 weight the + old levels by zero). An unread flux history is a derived quantity + (the flux of the value history) and is not traced or sampled; when + theta later drops below 1 it initialises from the field then.""" + d = self.Unknowns.DFDt + if getattr(d, "integrator", "am") == "bdf": + return False + theta = getattr(d, "theta", 0.5) + return theta is None or float(theta) < 1.0 + @property def F0(self): """Pointwise source term including time derivative.""" @@ -4515,8 +4535,14 @@ def solve( # Update History / Flux History terms # SemiLagrange and Lagrange may have different sequencing. + trace = getattr(self, "_characteristics", None) + if trace is not None: + trace.begin_step(timestep) self.DuDt.update_pre_solve(timestep, verbose=verbose, evalf=_evalf) - self.DFDt.update_pre_solve(timestep, verbose=verbose, evalf=_evalf) + if self._flux_history_is_read(): + self.DFDt.update_pre_solve(timestep, verbose=verbose, evalf=_evalf) + if trace is not None: + trace.finish_step() super().solve(zero_init_guess, _force_setup, divergence_retries=divergence_retries) @@ -5250,9 +5276,18 @@ def solve( if uw.mpi.rank == 0 and verbose: print(f"NS solver - pre-solve DuDt update", flush=True) - # Update SemiLagrange Flux terms + # Update SemiLagrange Flux terms. The velocity and stress histories + # follow the same velocity: one characteristic trace (velocity + # levels shared; departure points too where the node sets match). + if not hasattr(self, "_characteristics"): + self._characteristics = uw.systems.ddt.share_characteristics(self.DuDt, self.DFDt) + trace = self._characteristics + if trace is not None: + trace.begin_step(timestep) self.DuDt.update_pre_solve(timestep, verbose=verbose, evalf=_evalf) self.DFDt.update_pre_solve(timestep, verbose=verbose, evalf=_evalf) + if trace is not None: + trace.finish_step() # Override AM coefficients if flux_order is explicitly set if self._flux_order is not None: diff --git a/src/underworld3/utilities/cell_polynomial_projection.py b/src/underworld3/utilities/cell_polynomial_projection.py new file mode 100644 index 00000000..4fa46067 --- /dev/null +++ b/src/underworld3/utilities/cell_polynomial_projection.py @@ -0,0 +1,191 @@ +"""Cell-by-cell polynomial fit of particle data. + +This is the reconstruction behind ``SwarmVariable(proxy_location="cells")``: +every cell of a discontinuous mesh variable receives the least-squares +polynomial through the particles it holds. The result is exact for +polynomial particle fields up to the proxy degree, is a polynomial on the +mesh cell (so the default integration rule integrates it exactly, no +oversampling guard), keeps a material step sharp at cell edges, has a +gradient, and needs no neighbour search across ranks: a rank fits its own +cells from its own particles. + +A cell with too few particles for a well-posed fit (fewer than the basis +size plus two) receives a LINEAR fit to the particles nearest its centroid, +which is what the RBF proxy's neighbourhood does; linear because a +higher-degree polynomial extrapolated from a distant neighbourhood is +unbounded (measured: a P2 extrapolation into the emptied corner cells of a +rotating box reached twice the field maximum). That cell is consistent to +first order but not a cell-local fit. A cell with no particles at all +keeps its previous proxy value when one is supplied: no particles is no +information, and holding the old value is the honest default until the +swarm is repopulated. + +Why least squares rather than moment matching: the conservative +particle-to-mesh transfer (rule mass on the left, particle moments on the +right, as PETSc's ``DMSwarmProjectFields`` does) conserves the particle +sums exactly but its nodal values carry the Monte-Carlo error of the +particle "quadrature", which scales with the field value over the square +root of the particle count. Measured on a linear field with 21 jittered +particles per cell its error was 25% of the field; the least-squares fit +was exact (`~/+Simulations/integration_point_proxy`, 2026-09-08). +""" + +from __future__ import annotations + +import numpy as np + +import underworld3 as uw +from underworld3.cython.petsc_quadrature_fe import cell_affine_maps, tabulate + + +class CellPolynomialProjector: + """Least-squares fit of particle values onto a discontinuous mesh variable, cell by cell. + + Parameters + ---------- + meshVar : + A discontinuous (``continuous=False``) mesh variable of any degree and + component count. Its element and the mesh's affine cell maps are + tabulated once here; rebuild the projector when the mesh moves + (``mesh_version`` records the mesh version it was built for). + """ + + def __init__(self, meshVar): + mesh = meshVar.mesh + if getattr(meshVar, "continuous", True) or getattr(meshVar, "is_integration_point", False): + raise ValueError("CellPolynomialProjector needs a discontinuous (cell-local) mesh variable") + self.mesh = mesh + self.var = meshVar + self.mesh_version = mesh._mesh_version + self.dim = mesh.dim + self.num_components = meshVar.num_components + self.fe = mesh.dm.getField(meshVar.field_id)[0] + self.v0, self.invJ, self.detJ = cell_affine_maps(mesh.dm) + self.ncells = self.detJ.shape[0] + probe = tabulate(self.fe, np.zeros((1, self.dim))) + self.Nb = probe.shape[1] // self.num_components # scalar basis size + # Reference-cell centroid, mapped: xi_c + 1 = 2 / (dim + 1) on every axis. + J = np.linalg.inv(self.invJ) if self.ncells else self.invJ + self.centroids = self.v0 + np.einsum( + "cij,j->ci", J, np.full(self.dim, 2.0 / (self.dim + 1)) + ) + self._check_layout() + # Reference coordinates of the cell's dof nodes (the same in every + # cell), for evaluating a thin cell's linear fit at its nodes. + X = np.asarray(meshVar.coords_nd) + self.xi_dof = ( + self.reference_coords(X[: self.Nb], np.zeros(self.Nb, dtype=np.int64)) + if self.ncells else np.zeros((self.Nb, self.dim)) + ) + + # -- geometry ----------------------------------------------------------- + + def _scalar_basis(self, xi): + """Scalar basis values at reference points, shape (Np, Nb).""" + if xi.shape[0] == 0: + return np.zeros((0, self.Nb)) + T = tabulate(self.fe, xi) # (Np, Nb*Nc, Nc), interleaved + return T[:, 0 :: self.num_components, 0] + + def _check_layout(self): + """The discontinuous local vector is cell-major with ``Nb`` dofs per cell in basis order.""" + X = np.asarray(self.var.coords_nd) + if X.shape[0] != self.ncells * self.Nb: + raise RuntimeError( + f"unexpected dof count for {self.var.clean_name}: {X.shape[0]} != {self.ncells} x {self.Nb}" + ) + if self.ncells == 0: + return + cells = np.repeat(np.arange(self.ncells), self.Nb) + xi = self.reference_coords(X, cells) + B = self._scalar_basis(xi).reshape(self.ncells, self.Nb, self.Nb) + if not np.allclose(B, np.eye(self.Nb)[None], atol=1e-9): + raise RuntimeError("discontinuous dof layout is not cell-major; cannot fit cell by cell") + + def reference_coords(self, coords, cells): + """Reference coordinates (PETSc's [-1, 1] frame) of points in their cells.""" + return np.einsum("cij,cj->ci", self.invJ[cells], coords - self.v0[cells]) - 1.0 + + def locate(self, coords): + """Owning local cell of each point (-1 when not on this rank) and its reference coordinates.""" + coords = np.asarray(coords, dtype=np.float64) + cells = np.asarray(self.mesh._robust_owning_cells(coords), dtype=np.int64) + ok = cells >= 0 + xi = np.zeros_like(coords) + if ok.any(): + xi[ok] = self.reference_coords(coords[ok], cells[ok]) + return cells, ok, xi + + # -- the fit ------------------------------------------------------------ + + def fit(self, coords, values, nmin=None, patch_nnn=None, old=None): + """Fit every cell; returns nodal values shaped like ``meshVar.data``. + + Parameters + ---------- + coords : (N, dim) particle coordinates (non-dimensional). + values : (N, num_components) particle values. + nmin : particles a cell needs for its own fit (default basis size + 2). + patch_nnn : particles nearest the centroid used for a thin cell's + linear fit (default twice the basis size + 2). + old : current proxy values, shaped like ``meshVar.data``; after the + first fit a cell with no particles keeps them (on the first fit, + or without ``old``, it takes the linear patch fit). + """ + coords = np.asarray(coords, dtype=np.float64).reshape(-1, self.dim) + values = np.asarray(values, dtype=np.float64).reshape(coords.shape[0], -1) + nc = values.shape[1] + cells, ok, xi = self.locate(coords) + c = cells[ok] + B = self._scalar_basis(xi[ok]) # (Np, Nb) + psi = values[ok] # (Np, nc) + npc = np.bincount(c, minlength=self.ncells) + + G = np.zeros((self.ncells, self.Nb, self.Nb)) + R = np.zeros((self.ncells, self.Nb, nc)) + np.add.at(G, c, B[:, :, None] * B[:, None, :]) + np.add.at(R, c, B[:, :, None] * psi[:, None, :]) + + U = np.zeros((self.ncells, self.Nb, nc)) + nmin = nmin or self.Nb + 2 + dense = npc >= nmin + if dense.any(): + ridge = 1e-10 * np.trace(G[dense], axis1=1, axis2=2)[:, None, None] / self.Nb + U[dense] = np.linalg.solve(G[dense] + ridge * np.eye(self.Nb)[None], R[dense]) + + self.n_empty = int((npc == 0).sum()) + held = np.zeros(self.ncells, dtype=bool) + if old is not None and getattr(self, "_has_fit", False): + held = npc == 0 + U[held] = np.asarray(old, dtype=np.float64).reshape(self.ncells, self.Nb, nc)[held] + thin = np.nonzero(~dense & ~held)[0] + self.n_thin = int(thin.shape[0]) + if thin.shape[0] > 0 and c.shape[0] > 0: + # Linear fit (monomials 1, xi_1, ..., xi_dim in the cell's frame) + # to the nearest particles, evaluated at the cell's dof nodes. + Xp = coords[ok] + nnn = min(patch_nnn or 2 * self.Nb + 2, Xp.shape[0]) + tree = uw.kdtree.KDTree(Xp) + _, idx = tree.query(self.centroids[thin], k=nnn) + idx = np.asarray(idx).reshape(thin.shape[0], nnn) + xp = Xp[idx] # (nthin, nnn, dim) + xit = np.einsum("cij,cpj->cpi", self.invJ[thin], xp - self.v0[thin][:, None, :]) - 1.0 + A = np.concatenate([np.ones((thin.shape[0], nnn, 1)), xit], axis=2) # (nthin, nnn, dim+1) + Gt = np.einsum("cpa,cpb->cab", A, A) + Rt = np.einsum("cpa,cpk->cak", A, psi[idx]) + ridge = 1e-10 * np.trace(Gt, axis1=1, axis2=2)[:, None, None] / (self.dim + 1) + 1e-30 + coef = np.linalg.solve(Gt + ridge * np.eye(self.dim + 1)[None], Rt) # (nthin, dim+1, nc) + Adof = np.concatenate([np.ones((self.Nb, 1)), self.xi_dof], axis=1) # (Nb, dim+1) + U[thin] = np.einsum("ba,cak->cbk", Adof, coef) + + self._has_fit = True + return U.reshape(self.ncells * self.Nb, nc) + + def interpolate(self, U, coords): + """The fitted polynomials evaluated at points (NaN off-rank): the FLIP read-back.""" + U = np.asarray(U).reshape(self.ncells, self.Nb, -1) + cells, ok, xi = self.locate(coords) + out = np.full((coords.shape[0], U.shape[2]), np.nan) + B = self._scalar_basis(xi[ok]) + out[ok] = np.einsum("pb,pbk->pk", B, U[cells[ok]]) + return out diff --git a/tests/test_0064_quadrature_point_fe.py b/tests/test_0064_quadrature_point_fe.py index e989750a..111cb281 100644 --- a/tests/test_0064_quadrature_point_fe.py +++ b/tests/test_0064_quadrature_point_fe.py @@ -100,3 +100,52 @@ def test_rule_is_the_mesh_rule(): fe = PETSc.FE().createDefault(2, 1, True, 2, f"p{degree}_", PETSc.COMM_SELF) pts = np.array(fe.getQuadrature().getData()[0]).reshape(-1, 2) assert len(pts) == 6 + + +@pytest.mark.parametrize("dim,simplex", CELLS, ids=IDS) +def test_vector_element_is_the_identity_point_major(dim, simplex): + """Two components: basis (p, c) is the delta at point p times e_c, ordered + point-major, component-minor, so a cell's dofs reshape to (Nq, Nc).""" + _, quad, pts = _rule(dim, simplex, 2) + _, polytope = _box(dim, simplex) + Nc = 2 + fe = create_delta_fe(quad, polytope, num_components=Nc) + Nq = len(pts) + assert fe.getDimension() == Nq * Nc + assert fe.getNumComponents() == Nc + B = tabulate(fe, pts) # (Np, Nb, Nc) + expected = np.zeros((Nq, Nq * Nc, Nc)) + for p in range(Nq): + for c in range(Nc): + expected[p, p * Nc + c, c] = 1.0 + assert np.array_equal(B, expected) + assert np.all(tabulate(fe, pts[:2] + 0.05) == 0.0) + + +def test_cell_affine_maps_and_derivative_tabulation(): + """The affine map returns the reference frame the tabulation expects, and + the derivative tabulation differentiates a discontinuous P1 field exactly.""" + import numpy as np + import underworld3 as uw + from underworld3.cython.petsc_quadrature_fe import ( + cell_affine_maps, tabulate, tabulate_with_derivatives, + ) + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.25, qdegree=2) + var = uw.discretisation.MeshVariable("dg1", mesh, 1, degree=1, continuous=False) + fe = mesh.dm.getField(var.field_id)[0] + v0, invJ, detJ = cell_affine_maps(mesh.dm) + ncells = detJ.shape[0] + X = np.asarray(var.coords_nd) + cells = np.repeat(np.arange(ncells), 3) + xi = np.einsum("cij,cj->ci", invJ[cells], X - v0[cells]) - 1.0 + B, D = tabulate_with_derivatives(fe, xi) + assert np.allclose(B[:, :, 0].reshape(ncells, 3, 3), np.eye(3)[None], atol=1e-10) + assert np.allclose(tabulate(fe, xi), B) + # d/dx of the nodal interpolant of (2x + 3y): physical gradient = invJ^T D + vals = (2.0 * X[:, 0] + 3.0 * X[:, 1]).reshape(ncells, 3) + grad_ref = np.einsum("cnbd,cb->cnd", D[:, :, 0, :].reshape(ncells, 3, 3, 2), vals) + grad = np.einsum("cji,cnj->cni", invJ, grad_ref) + assert np.allclose(grad, [2.0, 3.0], atol=1e-10) + # cell volumes from detJ: the reference triangle has area 2 + assert np.isclose(2.0 * detJ.sum(), 1.0, atol=1e-12) or uw.mpi.size > 1 diff --git a/tests/test_0065_integration_point_variable.py b/tests/test_0065_integration_point_variable.py index 4f334436..8d9bef82 100644 --- a/tests/test_0065_integration_point_variable.py +++ b/tests/test_0065_integration_point_variable.py @@ -158,7 +158,33 @@ def test_other_rule_is_refused(): mesh._verify_integration_rule(other) -def test_vector_components_not_yet_supported(): +def test_vector_variable_layout_projection_and_evaluate(): + """A two-component integration-point variable: (ncells*Nq, 2) layout, + each component reproduced exactly by a P2 projection of P2 point data, + and evaluate exact at its own points.""" mesh = _mesh("triangle") - with pytest.raises(NotImplementedError): - uw.discretisation.IntegrationPointVariable("v", mesh, num_components=2) + x, y = mesh.X + v = uw.discretisation.IntegrationPointVariable("v", mesh, num_components=2) + Nq = len(np.asarray(mesh.integration_rule.getData()[1])) + n = _ncells(mesh) + assert v.data.shape == (n * Nq, 2) + assert v.cell_data.shape == (n, Nq, 2) + X = np.asarray(v.coords) + f0 = lambda X: 1.0 + 2.0 * X[:, 0] - 3.0 * X[:, 1] + 0.5 * X[:, 0] ** 2 + f1 = lambda X: -2.0 + X[:, 0] * X[:, 1] + X[:, 1] ** 2 + v.data[:, 0] = f0(X) + v.data[:, 1] = f1(X) + + own = np.asarray(uw.function.evaluate(v.sym, X)).reshape(-1, 2) + assert np.allclose(own, np.asarray(v.data), rtol=0, atol=1e-14) + + T = uw.discretisation.MeshVariable("Tv", mesh, 2, degree=2) + proj = uw.systems.solvers.SNES_Vector_Projection(mesh, T) + proj.uw_function = v.sym + proj.smoothing = 0.0 + proj.petsc_options["ksp_rtol"] = 1e-13 + proj.petsc_options["snes_rtol"] = 1e-13 + proj.solve() + Xt = np.asarray(T.coords) + assert np.abs(T.data[:, 0] - f0(Xt)).max() < 1e-9 + assert np.abs(T.data[:, 1] - f1(Xt)).max() < 1e-9 diff --git a/tests/test_0066_integration_point_slcn.py b/tests/test_0066_integration_point_slcn.py index 18520e44..27f18757 100644 --- a/tests/test_0066_integration_point_slcn.py +++ b/tests/test_0066_integration_point_slcn.py @@ -214,3 +214,53 @@ def run(solver_cls, kwargs): # Same history, same time derivative; the solvers differ only in how the # (negligible) diffusion is applied, so the fields agree closely. assert np.abs(T_composed - T_slcn).max() < 5e-3 + + +def test_value_and_flux_histories_share_one_characteristic_trace(): + """The SLCN solver's value and flux histories follow the same velocity + from the same nodes: one trace per step (two velocity evaluations for + the RK2 segment), the flux history served from the cache. With theta=1 + the old-level flux is never read, so the flux history is not traced at + all and the trace records one velocity level per step.""" + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.2, qdegree=2) + x, y = mesh.X + V = sympy.Matrix([[-y, x]]) + for theta, hits in ((0.5, 1), (1.0, 0)): + T = uw.discretisation.MeshVariable(f"T{int(theta * 10)}", mesh, 1, degree=2) + T.data[:, 0] = np.exp(-((np.asarray(T.coords) - 0.5) ** 2).sum(1) / 0.02) + adv = uw.systems.AdvDiffusionSLCN(mesh, u_Field=T, V_fn=V, order=1, theta=theta) + adv.constitutive_model = uw.constitutive_models.DiffusionModel + adv.constitutive_model.Parameters.diffusivity = 1e-3 + for b in ("Left", "Right", "Top", "Bottom"): + adv.add_dirichlet_bc(0.0, b) + trace = adv._characteristics + assert trace is not None + assert adv.DuDt.characteristics is trace and adv.DFDt.characteristics is trace + assert not adv.DuDt._owns_characteristics + adv.solve(timestep=0.05) + n0, h0 = trace.n_velocity_evaluations, trace.n_cache_hits + adv.solve(timestep=0.05) + # per step: one RK2 segment = 2 velocity evaluations, plus the one + # evaluation that records v^{n-1} at the nodes + assert trace.n_velocity_evaluations - n0 == 2, trace.n_velocity_evaluations - n0 + assert trace.n_cache_hits - h0 == hits + assert adv._flux_history_is_read() == (theta < 1.0) + if theta == 1.0: + assert not adv.DFDt._history_initialised + + +def test_private_trace_when_a_manager_stands_alone(): + """A manager used without a solver owns its trace and delimits its own + steps; the mid-time velocity becomes available after the first step.""" + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.25, qdegree=2) + x, y = mesh.X + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + T.data[:, 0] = np.asarray(T.coords)[:, 0] + D = uw.systems.ddt.SemiLagrangian(mesh, T.sym, sympy.Matrix([[1.0, 0.0]]), vtype=uw.VarType.SCALAR, degree=2, continuous=True) + tr = D.characteristics + assert D._owns_characteristics + assert tr.midtime_expr() == tr.V_matrix() # nothing recorded yet + D.update_pre_solve(0.1) + assert tr.level_valid(1) + assert tr.midtime_expr() != tr.V_matrix() # 1.5 v^n - 0.5 v^{n-1} + assert tr.n_velocity_evaluations == 2 diff --git a/tests/test_0067_integration_point_proxy.py b/tests/test_0067_integration_point_proxy.py new file mode 100644 index 00000000..9399f3db --- /dev/null +++ b/tests/test_0067_integration_point_proxy.py @@ -0,0 +1,311 @@ +"""Swarm proxy directly at the integration points. + +``SwarmVariable(..., proxy_location="integration_points")`` reconstructs the +particle field at the mesh integration points into an +``IntegrationPointVariable`` and the weak form reads it there with no basis +interpolation (the Ellipsis / Underworld PIC-LIP mapping). Checked here: + +- a material step carried by particles is reproduced far more sharply at + the integration points than through the nodal proxy; +- the proxy has no gradient and the JIT refuses one; +- ``Lagrangian_Swarm`` accepts the option and its history slots become + integration-point proxies that reproduce a linear field exactly; +- a symmetric-tensor swarm variable gets a multi-component proxy. +""" + +import numpy as np +import pytest +import sympy + +import underworld3 as uw + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +def _mesh(): + return uw.meshing.UnstructuredSimplexBox(cellSize=0.1, qdegree=2) + + +def test_material_step_is_sharper_at_integration_points(): + mesh = _mesh() + x, y = mesh.X + swarm = uw.swarm.Swarm(mesh) + M_n = uw.swarm.SwarmVariable("Mn", swarm, 1, proxy_degree=1, proxy_location="nodes") + M_q = uw.swarm.SwarmVariable("Mq", swarm, 1, proxy_location="integration_points") + swarm.populate(fill_param=3) + assert M_q._meshVar.is_integration_point and not M_n._meshVar.is_integration_point + + X = np.asarray(swarm._particle_coordinates.data) + step = (X[:, 0] < 0.5).astype(float) + with uw.synchronised_array_update(): + M_n.data[:, 0] = step + M_q.data[:, 0] = step + + # The exact step at the integration points, as an integration-point field. + S = uw.discretisation.IntegrationPointVariable("S", mesh) + S.data[:, 0] = (np.asarray(S.coords)[:, 0] < 0.5).astype(float) + + err_n = uw.maths.Integral(mesh, (M_n.sym[0] - S.sym[0]) ** 2).evaluate() + err_q = uw.maths.Integral(mesh, (M_q.sym[0] - S.sym[0]) ** 2).evaluate() + # Both proxies are nonzero-error (the interface sits between particles), + # and the integration-point proxy is sharper by a clear margin. + assert err_q > 0.0 and err_n > 0.0 + assert err_q < 0.5 * err_n, (err_q, err_n) + + # The integration-point proxy carries the step within one cell: away + # from the interface it is exactly 0 or 1 at every integration point. + Xq = np.asarray(M_q._meshVar.coords) + far = np.abs(Xq[:, 0] - 0.5) > 0.1 + vals = np.asarray(uw.function.evaluate(M_q.sym[0], Xq)).reshape(-1) # own points, through the symbol + assert np.allclose(vals[far], (Xq[far, 0] < 0.5).astype(float), rtol=0, atol=1e-14) + + # No gradient: the JIT refuses a derivative of the proxy symbol. + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) + proj = uw.systems.solvers.SNES_Projection(mesh, T) + proj.uw_function = M_q.sym[0].diff(x) + with pytest.raises(RuntimeError, match="integration-point"): + proj.solve() + + +def test_lagrangian_swarm_history_at_integration_points(): + mesh = _mesh() + x, y = mesh.X + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) + T.data[:, 0] = np.asarray(T.coords) @ np.array([1.0, 2.0]) # x + 2y + swarm = uw.swarm.Swarm(mesh) + lag = uw.systems.ddt.Lagrangian_Swarm( + swarm=swarm, psi_fn=T.sym, vtype=uw.VarType.SCALAR, degree=1, continuous=True, + order=1, proxy_location="integration_points", + ) + swarm.populate(fill_param=3) + lag.update_pre_solve(dt=0.1) + + slot = lag.psi_star[0] + assert slot._meshVar.is_integration_point + # Particle values are the field at the particles ... + Xp = np.asarray(swarm._particle_coordinates.data) + assert np.allclose(np.asarray(slot.data)[:, 0], Xp[:, 0] + 2.0 * Xp[:, 1], atol=1e-10) + # ... and the proxy at the integration points reproduces the linear field + # (order-1 RBF is linear-exact), read straight from the slot's symbol. + Xq = np.asarray(slot._meshVar.coords) + got = np.asarray(uw.function.evaluate(slot.sym[0], Xq)).reshape(-1) + assert np.allclose(got, Xq[:, 0] + 2.0 * Xq[:, 1], atol=1e-8) + # bdf() is built from the integration-point symbols. + assert slot._meshVar.sym[0] in lag.bdf()[0].free_symbols or lag.bdf()[0].has(slot._meshVar.sym[0]) + + +def test_tensor_swarm_variable_gets_a_multicomponent_proxy(): + mesh = _mesh() + swarm = uw.swarm.Swarm(mesh) + S = uw.swarm.SwarmVariable("S", swarm, vtype=uw.VarType.SYM_TENSOR, proxy_location="integration_points") + swarm.populate(fill_param=2) + Nq = len(np.asarray(mesh.integration_rule.getData()[1])) + c0, c1 = mesh.dm.getHeightStratum(0) + assert S._meshVar.is_integration_point + assert S._meshVar.data.shape == ((c1 - c0) * Nq, 3) + with uw.synchronised_array_update(): + S.data[:, 0] = 1.0; S.data[:, 1] = 2.0; S.data[:, 2] = 3.0 + # Read through the symbol (which refreshes the proxy): the symmetric + # tensor [[c0, c2], [c2, c1]] at every integration point. + Xq = np.asarray(S._meshVar.coords) + vals = np.asarray(uw.function.evaluate(S.sym, Xq)).reshape(len(Xq), 2, 2) + assert np.allclose(vals, [[1.0, 3.0], [3.0, 2.0]]) + + +@pytest.mark.level_2 +def test_layered_couette_interface_on_edges_is_exact_at_integration_points(): + """Two viscosity layers (1 and 1e3) carried by particles, interface on mesh + edges, top-driven layer flow. The exact velocity is piecewise linear and + lies in the P2 space, so the only error is the proxy's representation of + the step: the integration-point proxy is exact to solver tolerance, the + nodal proxy (a node on the interface averages both materials) is not.""" + h, eta_top = 0.5, 1.0e3 + results = {} + for proxy in ("nodes", "integration_points"): + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.1, qdegree=2, regular=True) + v = uw.discretisation.MeshVariable(f"v_{proxy}", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable(f"p_{proxy}", mesh, 1, degree=1) + swarm = uw.swarm.Swarm(mesh) + if proxy == "nodes": + M = uw.swarm.SwarmVariable("M", swarm, 1, proxy_degree=1) + else: + M = uw.swarm.SwarmVariable("M", swarm, 1, proxy_location="integration_points") + swarm.populate(fill_param=3) + X = np.asarray(swarm._particle_coordinates.data) + with uw.synchronised_array_update(): + M.data[:, 0] = (X[:, 1] > h).astype(float) + eta = 1.0 + (eta_top - 1.0) * sympy.Max(0, sympy.Min(1, M.sym[0])) + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = eta + stokes.add_dirichlet_bc((1.0, 0.0), "Top") + stokes.add_dirichlet_bc((0.0, 0.0), "Bottom") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Left") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Right") + stokes.tolerance = 1e-8 + stokes.solve() + A = 1.0 / (h + (1.0 - h) / eta_top) + Xv = np.asarray(v.coords) + vx_exact = np.where(Xv[:, 1] < h, A * Xv[:, 1], A * h + A / eta_top * (Xv[:, 1] - h)) + results[proxy] = np.abs(np.asarray(v.data[:, 0]) - vx_exact).max() + assert results["integration_points"] < 1e-5 + assert results["nodes"] > 1e-2 # the control: the nodal proxy smears the step + + +# --------------------------------------------------------------------------- +# proxy_location="cells": a least-squares polynomial per cell +# --------------------------------------------------------------------------- + + +def _jittered_swarm(mesh, fill, seed=0, **var_kwargs): + """A populated swarm with its lattice jittered so cells hold uneven particle counts.""" + swarm = uw.swarm.Swarm(mesh) + M = uw.swarm.SwarmVariable("M", swarm, 1, **var_kwargs) + swarm.populate(fill_param=fill) + rng = np.random.default_rng(seed + uw.mpi.rank) + X = np.array(swarm._particle_coordinates.data, copy=True) + X += rng.uniform(-1, 1, X.shape) * 0.3 * 0.1 / fill + X = np.clip(X, 1e-6, 1 - 1e-6) + with uw.synchronised_array_update(): + swarm._particle_coordinates.data[...] = X + swarm.migrate() + return swarm, M + + +def test_cells_proxy_reproduces_polynomials_and_has_a_gradient(): + mesh = _mesh() + x, y = mesh.X + swarm, M = _jittered_swarm(mesh, 4, proxy_location="cells", proxy_degree=2) + assert not M._meshVar.continuous and M._meshVar.degree == 2 + X = np.asarray(swarm._particle_coordinates.data) + quad = 1.0 + 2.0 * X[:, 0] - 3.0 * X[:, 1] + 4.0 * X[:, 0] * X[:, 1] - X[:, 1] ** 2 + with uw.synchronised_array_update(): + M.data[:, 0] = quad + quad_sym = 1 + 2 * x - 3 * y + 4 * x * y - y ** 2 + # Read through the weak form: the assembler evaluates the discontinuous + # proxy at the rule, where the fit must be exact. + err = uw.maths.Integral(mesh, (M.sym[0] - quad_sym) ** 2).evaluate() + assert err < 1e-14, err + # Fifteen particles per cell: no cell dropped below Nb + 2 = 8 after the + # jitter, so every cell took its own P2 fit. + assert M._cell_projector.n_thin == 0 + # The proxy has a gradient (unlike the integration-point proxy). + gerr = uw.maths.Integral(mesh, (M.sym[0].diff(x) - (2 + 4 * y)) ** 2).evaluate() + assert gerr < 1e-12, gerr + + +def test_cells_proxy_light_swarm_stays_linear_exact(): + """Three particles per cell (below the degree-2 fit's own threshold): every + cell takes the linear patch fit and a linear field is still exact.""" + mesh = _mesh() + x, y = mesh.X + swarm, M = _jittered_swarm(mesh, 1, proxy_location="cells", proxy_degree=2) + X = np.asarray(swarm._particle_coordinates.data) + with uw.synchronised_array_update(): + M.data[:, 0] = 1.0 + 2.0 * X[:, 0] + 3.0 * X[:, 1] + err = uw.maths.Integral(mesh, (M.sym[0] - (1 + 2 * x + 3 * y)) ** 2).evaluate() + assert err < 1e-14, err + assert M._cell_projector.n_thin > 0 + + +def test_cells_proxy_material_step_is_sharp_and_bounded_at_cell_edges(): + """A step on x = 0.5 with the mesh regular so cell edges lie on it: the + per-cell fit is exactly 0 or 1 in every cell (no overshoot, unlike the RBF).""" + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.1, qdegree=2, regular=True) + x, y = mesh.X + swarm = uw.swarm.Swarm(mesh) + M = uw.swarm.SwarmVariable("M", swarm, 1, proxy_location="cells", proxy_degree=1) + swarm.populate(fill_param=3) + X = np.asarray(swarm._particle_coordinates.data) + with uw.synchronised_array_update(): + M.data[:, 0] = (X[:, 0] < 0.5).astype(float) + S = uw.discretisation.IntegrationPointVariable("S", mesh) + S.data[:, 0] = (np.asarray(S.coords)[:, 0] < 0.5).astype(float) + err = uw.maths.Integral(mesh, (M.sym[0] - S.sym[0]) ** 2).evaluate() + assert err < 1e-14, err + + +def test_cells_proxy_vector_variable_and_lagrangian_swarm(): + mesh = _mesh() + swarm = uw.swarm.Swarm(mesh) + V = uw.swarm.SwarmVariable("V", swarm, vtype=uw.VarType.VECTOR, proxy_location="cells", proxy_degree=1) + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) + T.data[:, 0] = np.asarray(T.coords) @ np.array([1.0, 2.0]) + lag = uw.systems.ddt.Lagrangian_Swarm( + swarm=swarm, psi_fn=T.sym, vtype=uw.VarType.SCALAR, degree=1, continuous=False, + order=1, proxy_location="cells", + ) + swarm.populate(fill_param=3) + X = np.asarray(swarm._particle_coordinates.data) + with uw.synchronised_array_update(): + V.data[:, 0] = X[:, 0] + V.data[:, 1] = 2.0 * X[:, 1] + x, y = mesh.X + err = uw.maths.Integral(mesh, (V.sym[0] - x) ** 2 + (V.sym[1] - 2 * y) ** 2).evaluate() + assert err < 1e-14, err + lag.update_pre_solve(dt=0.1) + slot = lag.psi_star[0] + assert not slot._meshVar.continuous + err = uw.maths.Integral(mesh, (slot.sym[0] - (x + 2 * y)) ** 2).evaluate() + assert err < 1e-14, err + + +def test_lagrangian_swarm_flip_update_keeps_particle_values_for_a_resolved_field(): + """FLIP read-back: the particle takes the mesh INCREMENT (solution minus + the proxy the mesh saw). For a field the proxy resolves exactly the + increment is zero, so the particle values are untouched by a solve that + reproduces the field; PIC would re-sample them.""" + mesh = _mesh() + x, y = mesh.X + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + T.data[:, 0] = np.asarray(T.coords) @ np.array([1.0, 2.0]) + 0.5 + swarm = uw.swarm.Swarm(mesh) + lag = uw.systems.ddt.Lagrangian_Swarm( + swarm=swarm, psi_fn=T.sym, vtype=uw.VarType.SCALAR, degree=2, continuous=False, + order=1, step_averaging=1, proxy_location="cells", particle_update="flip", + ) + swarm.populate(fill_param=3) + lag.update_pre_solve(dt=0.1) + Xp = np.asarray(swarm._particle_coordinates.data) + before = np.array(lag.psi_star[0].data[:, 0], copy=True) + assert np.allclose(before, Xp @ np.array([1.0, 2.0]) + 0.5, atol=1e-10) + # Perturb the particle values by a sub-cell "residual" the P2 proxy cannot + # hold, then run the post-solve update with the mesh field unchanged. + rng = np.random.default_rng(1) + noise = 1e-3 * rng.standard_normal(before.shape[0]) + with uw.synchronised_array_update(): + lag.psi_star[0].data[:, 0] = before + noise + lag.update_pre_solve(dt=0.1) # proxy refits from the noisy particles + lag.update_post_solve(dt=0.1) + after = np.asarray(lag.psi_star[0].data[:, 0]) + # solution (T, exact linear) - proxy (linear + fit of the noise): the + # residual survives up to the part of the noise the fit absorbed. + assert np.abs(after - (before + noise)).max() < 5e-3 + assert np.abs(after - before).max() > 1e-4 # PIC would have given `before` back + with pytest.raises(ValueError, match="particle_update"): + uw.systems.ddt.Lagrangian_Swarm(swarm=swarm, psi_fn=T.sym, vtype=uw.VarType.SCALAR, + degree=1, continuous=False, proxy_location="cells", particle_update="xx") + + +def test_lagrangian_swarm_history_is_sampled_before_the_first_move(): + """The first advection samples the history at the launch positions. Left + to the first solve, the sampling would see the landed positions and the + first step would transport nothing (a one-step lag).""" + mesh = _mesh() + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + T.data[:, 0] = np.asarray(T.coords)[:, 0] # psi = x + swarm = uw.swarm.Swarm(mesh) + lag = uw.systems.ddt.Lagrangian_Swarm( + swarm=swarm, psi_fn=T.sym, vtype=uw.VarType.SCALAR, degree=2, continuous=False, + order=1, proxy_location="cells", + ) + swarm.populate(fill_param=2) + assert not lag._history_initialised + X_before = np.array(swarm._particle_coordinates.data, copy=True) + swarm.advection(sympy.Matrix([[0.1, 0.0]]), 0.5, order=2) # every particle moves +0.05 in x + X_after = np.asarray(swarm._particle_coordinates.data) + assert lag._history_initialised + kept = np.abs(X_after[:, 0] - X_before[:, 0] - 0.05) < 1e-12 # particles not returned to bounds + vals = np.asarray(lag.psi_star[0].data[:, 0]) + assert np.allclose(vals[kept], X_before[kept, 0], atol=1e-10) # launch positions ... + assert not np.allclose(vals[kept], X_after[kept, 0], atol=1e-3) # ... not landing positions