From 9fe39add402acb661619105cf348d1796e7543c1 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Sat, 4 Jul 2026 08:58:09 +1000 Subject: [PATCH] fix(ddt): Lagrangian history writes use the modern data layout (BF-04) The history-update blocks in Lagrangian and Lagrangian_Swarm wrote per-component particle history as psi_star_0[i, j].data[:] = ... Indexing a SwarmVariable returns a symbolic component (via MathematicalMixin.__getitem__, an UnderworldAppliedFunction) which has no .data attribute, so the first update_pre_solve of any user- configured Lagrangian history raised AttributeError at ddt.py:3197 and the advertised feature was unusable (audit finding SWARM-06, runtime-confirmed). Component writes now go through the canonical (N, components) storage via psi_star_0.data[:, var._data_layout(i, j)], particle coordinates are read from swarm._particle_coordinates.data instead of the deprecated swarm.data, and the deprecated swarm.access() context wrappers around these blocks are removed. The four sites (previously ddt.py:3197, 3281, 3549, 3638-3639) now sit directly on the modern data-access pattern, so Wave B migration has nothing left to do there. Adds advecting-history regression tests for both classes (tests/test_0857_lagrangian_ddt_advecting_history.py): scalar and vector Lagrangian histories through initialise/advect/shift, and Lagrangian_Swarm step-averaged blending across a user-driven advection. All three fail with the original AttributeError against the old write path. Underworld development team with AI support from Claude Code --- src/underworld3/systems/ddt.py | 96 ++++---- ...t_0857_lagrangian_ddt_advecting_history.py | 217 ++++++++++++++++++ 2 files changed, 274 insertions(+), 39 deletions(-) create mode 100644 tests/test_0857_lagrangian_ddt_advecting_history.py diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index 662acab8..624d488a 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -3187,19 +3187,24 @@ def initialise_history(self): be called manually after setting initial conditions. """ psi_star_0 = self.psi_star[0] - with self.swarm.access(psi_star_0): - for i in range(psi_star_0.shape[0]): - for j in range(psi_star_0.shape[1]): - updated_psi = uw.function.evaluate( - self.psi_fn[i, j], - self.swarm.data, - ) - psi_star_0[i, j].data[:] = updated_psi + # Component-wise write through the canonical (N, components) storage. + # Indexing the SwarmVariable itself (``psi_star_0[i, j]``) returns a + # *symbolic* component with no ``.data`` — the modern component + # address is ``.data[:, var._data_layout(i, j)]`` (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]): + updated_psi = uw.function.evaluate( + self.psi_fn[i, j], + coords, + ) + psi_star_0.data[:, psi_star_0._data_layout(i, j)] = np.asarray( + updated_psi + ).reshape(-1) # Copy to all other history slots for k in range(1, self.order): - with self.swarm.access(self.psi_star[k]): - self.psi_star[k].data[...] = psi_star_0.data[...] + self.psi_star[k].data[...] = psi_star_0.data[...] self._history_initialised = True return @@ -3270,15 +3275,19 @@ def update_post_solve( # Now update the swarm variable psi_star_0 = self.psi_star[0] - with self.swarm.access(psi_star_0): - for i in range(psi_star_0.shape[0]): - for j in range(psi_star_0.shape[1]): - updated_psi = uw.function.evaluate( - self.psi_fn[i, j], - self.swarm.data, - evalf=evalf, - ) - psi_star_0[i, j].data[:] = updated_psi + # Grab the current psi values at the (pre-advection) particle + # positions via the canonical component 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]): + updated_psi = uw.function.evaluate( + self.psi_fn[i, j], + coords, + evalf=evalf, + ) + psi_star_0.data[:, psi_star_0._data_layout(i, j)] = np.asarray( + updated_psi + ).reshape(-1) # Now update the swarm locations @@ -3539,19 +3548,24 @@ def initialise_history(self): be called manually after setting initial conditions. """ psi_star_0 = self.psi_star[0] - with self.swarm.access(psi_star_0): - for i in range(psi_star_0.shape[0]): - for j in range(psi_star_0.shape[1]): - updated_psi = uw.function.evaluate( - self.psi_fn[i, j], - self.swarm.data, - ) - psi_star_0[i, j].data[:] = updated_psi + # Component-wise write through the canonical (N, components) storage. + # Indexing the SwarmVariable itself (``psi_star_0[i, j]``) returns a + # *symbolic* component with no ``.data`` — the modern component + # address is ``.data[:, var._data_layout(i, j)]`` (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]): + updated_psi = uw.function.evaluate( + self.psi_fn[i, j], + coords, + ) + psi_star_0.data[:, psi_star_0._data_layout(i, j)] = np.asarray( + updated_psi + ).reshape(-1) # Copy to all other history slots for k in range(1, self.order): - with self.swarm.access(self.psi_star[k]): - self.psi_star[k].data[...] = psi_star_0.data[...] + self.psi_star[k].data[...] = psi_star_0.data[...] self._history_initialised = True return @@ -3621,23 +3635,27 @@ def update_post_solve( ) print(f"Lagrange swarm copying {i-1} to {i}", flush=True) - with self.swarm.access(self.psi_star[i]): - self.psi_star[i].data[...] = self.psi_star[i - 1].data[...] + self.psi_star[i].data[...] = self.psi_star[i - 1].data[...] phi = 1 / self.step_averaging psi_star_0 = self.psi_star[0] - with self.swarm.access(psi_star_0): - for i in range(psi_star_0.shape[0]): - for j in range(psi_star_0.shape[1]): - updated_psi = uw.function.evaluate( + # 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) + updated_psi = np.asarray( + uw.function.evaluate( self.psi_fn[i, j], - self.swarm.data, + coords, evalf=evalf, ) - psi_star_0[i, j].data[:] = ( - phi * updated_psi + (1 - phi) * psi_star_0[i, j].data[:] - ) + ).reshape(-1) + 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 diff --git a/tests/test_0857_lagrangian_ddt_advecting_history.py b/tests/test_0857_lagrangian_ddt_advecting_history.py new file mode 100644 index 00000000..8e94d49e --- /dev/null +++ b/tests/test_0857_lagrangian_ddt_advecting_history.py @@ -0,0 +1,217 @@ +"""Advecting-history regression tests for the swarm-based Lagrangian DDt +classes (``Lagrangian`` and ``Lagrangian_Swarm``). + +Audit finding SWARM-06 (BF-04, docs/reviews/2026-07): the history-update +blocks in both classes wrote components as ``psi_star_0[i, j].data[:] = ...``. +Indexing a SwarmVariable returns a *symbolic* component (an +UnderworldAppliedFunction via ``MathematicalMixin.__getitem__``) which has no +``.data`` — so every first ``update_pre_solve`` raised AttributeError and the +Lagrangian history machinery was unusable. The fix routes component writes +through the canonical storage: ``psi_star_0.data[:, var._data_layout(i, j)]``. + +These tests actually *advect* the history (the pre-existing test_0007 +coverage only round-tripped ``.state`` with ``_history_initialised`` set +manually) and fail with AttributeError against the old write path. + +All expected values are expressed as functions of the particles' *current* +coordinates (minus the known constant-velocity displacement where relevant) +so the assertions are robust to any particle reordering during migration. +""" + +import numpy as np +import pytest +import sympy + +import underworld3 as uw +from underworld3.systems import ddt as ddt_module + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + +# Constant velocity: RK advection is exact, displacement is known a priori. +VX, VY = 0.30, 0.20 +DT = 0.01 + + +def _make_mesh(): + return uw.meshing.StructuredQuadBox( + elementRes=(8, 8), minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0) + ) + + +def _linear_scalar_field(mesh, name, a, b): + """P1 MeshVariable holding the (exactly representable) field a*x + b*y.""" + T = uw.discretisation.MeshVariable(name, mesh, 1, degree=1) + coords = np.asarray(T.coords) + T.data[:, 0] = a * coords[:, 0] + b * coords[:, 1] + return T + + +def _set_linear(T, a, b): + coords = np.asarray(T.coords) + T.data[:, 0] = a * coords[:, 0] + b * coords[:, 1] + + +def _constant_velocity(mesh): + V = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=1) + V.data[:, 0] = VX + V.data[:, 1] = VY + return V + + +def test_lagrangian_scalar_advecting_history(): + """Lagrangian (own swarm): initialise, advect, and verify the shifted + history holds psi at each particle's pre-advection position.""" + mesh = _make_mesh() + T = _linear_scalar_field(mesh, "T", 1.0, 2.0) # T1 = x + 2y + V = _constant_velocity(mesh) + + lag_ddt = ddt_module.Lagrangian( + mesh=mesh, + psi_fn=T.sym, + V_fn=V.sym, + vtype=uw.VarType.SCALAR, + degree=1, + continuous=True, + order=2, + fill_param=2, + ) + + # --- first update: auto-initialise (old write path raised here) --- + lag_ddt.update_pre_solve(dt=DT) + assert lag_ddt._history_initialised + + coords0 = np.asarray(lag_ddt.swarm._particle_coordinates.data).copy() + expected_T1 = coords0[:, 0] + 2.0 * coords0[:, 1] + got0 = np.asarray(lag_ddt.psi_star[0].data)[:, 0] + got1 = np.asarray(lag_ddt.psi_star[1].data)[:, 0] + assert np.allclose(got0, expected_T1, atol=1e-10) + # all history slots initialised identically + assert np.allclose(got1, expected_T1, atol=1e-10) + + # --- post-solve: field changes, history shifts, swarm advects --- + _set_linear(T, 3.0, -1.0) # T2 = 3x - y + lag_ddt.update_post_solve(dt=DT) + + coords1 = np.asarray(lag_ddt.swarm._particle_coordinates.data) + # the swarm actually moved (displacement DT * V for interior particles) + disp = np.linalg.norm(coords1.mean(axis=0) - coords0.mean(axis=0)) + assert disp > 0.5 * DT * np.hypot(VX, VY) + + # Each particle's pre-advection position (constant V => exact rewind). + x0 = coords1[:, 0] - DT * VX + y0 = coords1[:, 1] - DT * VY + # ignore particles clamped back into the domain by the bounds restore + interior = ( + (x0 > 1e-6) & (x0 < 1 - DT * VX - 1e-6) & (y0 > 1e-6) & (y0 < 1 - DT * VY - 1e-6) + ) + assert interior.sum() > 0.9 * interior.size + + # slot 0 = T2 sampled at the pre-advection position; slot 1 = old slot 0 + got0 = np.asarray(lag_ddt.psi_star[0].data)[:, 0] + got1 = np.asarray(lag_ddt.psi_star[1].data)[:, 0] + assert np.allclose(got0[interior], 3.0 * x0[interior] - y0[interior], atol=1e-8) + assert np.allclose(got1[interior], x0[interior] + 2.0 * y0[interior], atol=1e-8) + + +def test_lagrangian_vector_component_layout(): + """Vector psi: each component must land in its own canonical-storage + column (exercises the _data_layout(i, j) mapping, not just scalars).""" + mesh = _make_mesh() + W = uw.discretisation.MeshVariable("W", mesh, mesh.dim, degree=1) + wc = np.asarray(W.coords) + W.data[:, 0] = wc[:, 0] + 2.0 * wc[:, 1] # W_x = x + 2y + W.data[:, 1] = 3.0 * wc[:, 0] - wc[:, 1] # W_y = 3x - y + V = _constant_velocity(mesh) + + lag_ddt = ddt_module.Lagrangian( + mesh=mesh, + psi_fn=W.sym, + V_fn=V.sym, + vtype=uw.VarType.VECTOR, + degree=1, + continuous=True, + order=1, + fill_param=2, + ) + + lag_ddt.update_pre_solve(dt=DT) + + coords = np.asarray(lag_ddt.swarm._particle_coordinates.data) + got = np.asarray(lag_ddt.psi_star[0].data) + assert got.shape[1] == mesh.dim + assert np.allclose(got[:, 0], coords[:, 0] + 2.0 * coords[:, 1], atol=1e-10) + assert np.allclose(got[:, 1], 3.0 * coords[:, 0] - coords[:, 1], atol=1e-10) + + +def test_lagrangian_swarm_advecting_history_with_step_averaging(): + """Lagrangian_Swarm (user swarm): initialise, blend with step_averaging, + then advect the user swarm and blend again at the new positions.""" + mesh = _make_mesh() + T = _linear_scalar_field(mesh, "T", 1.0, 2.0) # T1 = x + 2y + V = _constant_velocity(mesh) + + swarm = uw.swarm.Swarm(mesh) + lag_ddt = ddt_module.Lagrangian_Swarm( + swarm=swarm, + psi_fn=T.sym, + vtype=uw.VarType.SCALAR, + degree=1, + continuous=True, + order=2, + step_averaging=2, # phi = 1/2 blending + ) + swarm.populate(fill_param=2) + + # --- first update: auto-initialise (old write path raised here) --- + lag_ddt.update_pre_solve(dt=DT) + coords0 = np.asarray(swarm._particle_coordinates.data).copy() + T1 = lambda x, y: x + 2.0 * y + got0 = np.asarray(lag_ddt.psi_star[0].data)[:, 0] + assert np.allclose(got0, T1(coords0[:, 0], coords0[:, 1]), atol=1e-10) + + # --- post-solve without advection: blended update, shifted chain --- + _set_linear(T, 3.0, -1.0) # T2 = 3x - y + T2 = lambda x, y: 3.0 * x - y + lag_ddt.update_post_solve(dt=DT) + + # positions unchanged (Lagrangian_Swarm never advects the swarm itself) + coords_after = np.asarray(swarm._particle_coordinates.data) + assert np.allclose(coords_after, coords0) + + x, y = coords0[:, 0], coords0[:, 1] + blend1 = 0.5 * T2(x, y) + 0.5 * T1(x, y) + got0 = np.asarray(lag_ddt.psi_star[0].data)[:, 0] + got1 = np.asarray(lag_ddt.psi_star[1].data)[:, 0] + assert np.allclose(got0, blend1, atol=1e-10) + assert np.allclose(got1, T1(x, y), atol=1e-10) # chain copy of old slot 0 + + # --- user advects the swarm, then another blended update --- + swarm.advection( + V.sym, delta_t=DT, restore_points_to_domain_func=mesh.return_coords_to_bounds + ) + _set_linear(T, -1.0, 4.0) # T3 = -x + 4y + T3 = lambda x, y: -x + 4.0 * y + lag_ddt.update_post_solve(dt=DT) + + coords1 = np.asarray(swarm._particle_coordinates.data) + # pre-advection positions (constant V => exact rewind); skip clamped points + x0 = coords1[:, 0] - DT * VX + y0 = coords1[:, 1] - DT * VY + interior = ( + (x0 > 1e-6) & (x0 < 1 - DT * VX - 1e-6) & (y0 > 1e-6) & (y0 < 1 - DT * VY - 1e-6) + ) + assert interior.sum() > 0.9 * interior.size + + # slot 0 = phi*T3(current pos) + (1-phi)*(previous per-particle blend), + # where the previous blend is a known function of the pre-advection pos. + expected0 = 0.5 * T3(coords1[:, 0], coords1[:, 1]) + 0.5 * ( + 0.5 * T2(x0, y0) + 0.5 * T1(x0, y0) + ) + got0 = np.asarray(lag_ddt.psi_star[0].data)[:, 0] + got1 = np.asarray(lag_ddt.psi_star[1].data)[:, 0] + assert np.allclose(got0[interior], expected0[interior], atol=1e-8) + # slot 1 = the previous blend, carried with the particle + assert np.allclose( + got1[interior], 0.5 * T2(x0[interior], y0[interior]) + 0.5 * T1(x0[interior], y0[interior]), + atol=1e-8, + )