From 02efb64e87ed5a012d4f564c039dbc6929ed30e0 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:48:38 +0800 Subject: [PATCH 1/3] =?UTF-8?q?Fix=20GenericFluxModel=20flux=20storage=20?= =?UTF-8?q?=E2=80=94=20preserve=20symbolic=20field=20dependencies?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GenericFluxModel.Parameters.flux setter (and __init__) wrapped each flux component in a UWexpression via validate_parameters(), creating intermediate symbols (q_0, q_1) that the JIT treated as independent constants, losing the dependency on the unknown field variable and producing a singular matrix. Fix: store raw sympy expressions directly, bypassing the UWexpression wrapper. Includes test test_1060_generic_flux_model_update.py verifying both the _force_setup and _solver-link propagation paths. Underworld development team with AI support from Claude Code --- src/underworld3/constitutive_models.py | 40 +++---- tests/test_1060_generic_flux_model_update.py | 119 +++++++++++++++++++ 2 files changed, 135 insertions(+), 24 deletions(-) create mode 100644 tests/test_1060_generic_flux_model_update.py diff --git a/src/underworld3/constitutive_models.py b/src/underworld3/constitutive_models.py index 35f34a64..1f911338 100644 --- a/src/underworld3/constitutive_models.py +++ b/src/underworld3/constitutive_models.py @@ -2593,17 +2593,10 @@ class _Parameters(_ParameterBase): def __init__(inner_self, _owning_model): inner_self._owning_model = _owning_model - default_flux = sympy.zeros(_owning_model.dim, 1) - elements = [default_flux[i] for i in range(_owning_model.dim)] - validated = [] - for i, v in enumerate(elements): - flux_component = validate_parameters( - rf"q_{{{i}}}", v, f"Flux component in x_{i}", allow_number=True - ) - if flux_component is not None: - validated.append(flux_component) - - inner_self._flux = sympy.Matrix(validated) + # Raw sympy zero vector — NOT wrapped in UWexpressions, so the + # JIT sees the symbolic content directly (field variable dependencies). + # See the flux setter for the rationale. + inner_self._flux = sympy.zeros(_owning_model.dim, 1) @property def flux(inner_self): @@ -2612,26 +2605,25 @@ def flux(inner_self): @flux.setter def flux(inner_self, value: sympy.Matrix): - """Set the flux expression (must be a vector of length dim).""" + """Set the flux expression (must be a vector of length dim). + + Stores the raw sympy expression directly to preserve symbolic + dependencies (field variables, mesh coordinates) for the JIT. + Previously wrapped each component in a UWexpression via + ``validate_parameters``, which created independent Symbols that + the JIT treated as constants — losing the dependency on the + unknown field (e.g. ``u.sym``) and producing a singular matrix. + """ dim = inner_self._owning_model.dim # Accept shape (dim, 1) or (1, dim) if value.shape not in [(dim, 1), (1, dim)]: raise ValueError( - f"Flux must be a symbolic vector of length {dim}. " f"Got shape {value.shape}." - ) - - # Flatten and validate - elements = [value[i] for i in range(dim)] - validated = [] - for i, v in enumerate(elements): - flux_component = validate_parameters( - rf"q_{{{i}}}", v, f"Flux component in x_{i}", allow_number=True + f"Flux must be a symbolic vector of length {dim}. " + f"Got shape {value.shape}." ) - if flux_component is not None: - validated.append(flux_component) - inner_self._flux = sympy.Matrix(validated).reshape(dim, 1) + inner_self._flux = sympy.Matrix(value).reshape(dim, 1) inner_self._reset() @property diff --git a/tests/test_1060_generic_flux_model_update.py b/tests/test_1060_generic_flux_model_update.py new file mode 100644 index 00000000..74c54637 --- /dev/null +++ b/tests/test_1060_generic_flux_model_update.py @@ -0,0 +1,119 @@ +"""Test that `GenericFluxModel.Parameters.flux` properly propagates to the solver. + +Verifies that the flux expression is stored as raw sympy (not wrapped in +UWexpression) so the JIT sees field-variable dependencies correctly. + +Problem: -div(k·grad(phi)) = f with phi=0 on all walls, f=1. +Solve with k=1, then change k → 10 via Parameters.flux. +Solution scales as 1/k. +""" + +import sympy +import pytest +import numpy as np +import underworld3 as uw +from underworld3.systems import Poisson + + +@pytest.fixture(scope="function") +def mesh(): + uw.reset_default_model() + return uw.meshing.StructuredQuadBox(elementRes=(4, 4)) + + +def _setup_solver(u, mesh, flux_expr): + """Helper: configure Poisson with GenericFluxModel and direct solve.""" + solver = Poisson(mesh, u_Field=u) + solver.petsc_options["ksp_type"] = "preonly" + solver.petsc_options["pc_type"] = "lu" + solver.petsc_options["pc_factor_mat_solver_type"] = "mumps" + solver.constitutive_model = uw.constitutive_models.GenericFluxModel(solver.Unknowns) + solver.constitutive_model.Parameters._solver = solver + solver.constitutive_model.Parameters.flux = flux_expr + solver.f = 1.0 + solver.add_essential_bc(0.0, "Bottom") + solver.add_essential_bc(0.0, "Top") + solver.add_essential_bc(0.0, "Left") + solver.add_essential_bc(0.0, "Right") + return solver + + +@pytest.mark.tier_a +@pytest.mark.level_1 +def test_flux_update_via_force_setup(mesh): + """Change flux between solves using _force_setup=True.""" + u = uw.discretisation.MeshVariable("u", mesh, 1, degree=2) + x, y = mesh.X + + grad = sympy.Matrix([u.sym.diff(x)[0], u.sym.diff(y)[0]]) + solver = _setup_solver(u, mesh, grad) + + solver.solve(_force_setup=True, zero_init_guess=True) + peak_a = float(np.max(np.abs(u.data))) + assert peak_a > 1e-10, f"Solution is zero (k=1)" + + # Change flux: k=10 + solver.constitutive_model.Parameters.flux = \ + sympy.Matrix([10 * u.sym.diff(x)[0], 10 * u.sym.diff(y)[0]]) + solver.solve(_force_setup=True, zero_init_guess=False) + peak_b = float(np.max(np.abs(u.data))) + assert peak_b > 1e-10, f"Solution is zero (k=10)" + + ratio = peak_a / peak_b if peak_b > 0 else 0 + print(f" force_setup: peak_a={peak_a:.6f}, peak_b={peak_b:.6f}, ratio={ratio:.3f}") + assert abs(ratio - 10.0) < 2.0 + + +@pytest.mark.tier_a +@pytest.mark.level_1 +def test_flux_update_via_solver_link(mesh): + """Change flux WITHOUT _force_setup, relying on _solver link + _reset().""" + u = uw.discretisation.MeshVariable("u", mesh, 1, degree=2) + x, y = mesh.X + + grad = sympy.Matrix([u.sym.diff(x)[0], u.sym.diff(y)[0]]) + solver = _setup_solver(u, mesh, grad) + + solver.solve(_force_setup=True, zero_init_guess=True) + peak_a = float(np.max(np.abs(u.data))) + assert peak_a > 1e-10 + + # Change flux via _solver link (no _force_setup) + solver.constitutive_model.Parameters.flux = \ + sympy.Matrix([10 * u.sym.diff(x)[0], 10 * u.sym.diff(y)[0]]) + + assert not solver.constitutive_model._solver_is_setup + assert solver._needs_function_rewire + + solver.solve(zero_init_guess=False) + peak_b = float(np.max(np.abs(u.data))) + assert peak_b > 1e-10 + + ratio = peak_a / peak_b if peak_b > 0 else 0 + print(f" solver_link: peak_a={peak_a:.6f}, peak_b={peak_b:.6f}, ratio={ratio:.3f}") + assert abs(ratio - 10.0) < 2.0 + + +@pytest.mark.tier_a +@pytest.mark.level_1 +def test_flux_stores_raw_sympy(mesh): + """Flux must NOT be wrapped in UWexpression — JIT must see field variables.""" + u = uw.discretisation.MeshVariable("u", mesh, 1, degree=2) + x, y = mesh.X + + solver = Poisson(mesh, u_Field=u) + solver.constitutive_model = uw.constitutive_models.GenericFluxModel(solver.Unknowns) + solver.constitutive_model.Parameters.flux = \ + sympy.Matrix([u.sym.diff(x)[0], u.sym.diff(y)[0]]) + + flux = solver.constitutive_model.Parameters.flux + + # Check it's a plain sympy Matrix, not a UWexpression + assert isinstance(flux, sympy.Matrix), f"Expected sympy.Matrix, got {type(flux)}" + # The elements should be raw sympy, not UWexpressions + # UWexpressions show as "q_0 = ..." in str representation + flux_str = str(flux) + assert "q_" not in flux_str, ( + f"Flux elements are wrapped in UWexpressions. " + f"Got: {flux_str}" + ) From f857ca106f482704c489b1d8810eb0a2323984e6 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 28 Jul 2026 17:39:21 +1000 Subject: [PATCH 2/3] review response (wording + test hygiene): the mechanism is the Picard tangent, not the JIT MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit From the adversarial review on the PR thread: - The setter docstring and __init__ comment claimed the JIT treated the wrapped flux components as independent constants. Probed false: the constants manifest was empty, the residual was never corrupted, and the identical pre-fix code solves under consistent_jacobian=True. The true mechanism is that the DEFAULT (Picard) tangent differentiates the flux without unwrapping (frozen-coefficient semantics) — an opaque wrapper holding the ENTIRE flux has zero derivative w.r.t. the unknown and its gradient, so G0-G3 vanish and the operator is structurally singular. Wording corrected so the record does not indict the JIT constants machinery (#416, exonerated). - Documented the side effect: with nothing left to freeze, GenericFluxModel's default tangent is effectively full Newton. - Test hygiene: removed the hand-wiring of Parameters._solver (the constitutive_model setter does it — hand-wiring masks a regression in that path) and the MUMPS solver pin (house policy avoids MUMPS; plain LU is used by the review probes and passes identically). Fix itself unchanged; all review probes pass (see thread). Underworld development team with AI support from Claude Code --- src/underworld3/constitutive_models.py | 29 ++++++++++++++++++-------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/src/underworld3/constitutive_models.py b/src/underworld3/constitutive_models.py index 1f911338..79f960fb 100644 --- a/src/underworld3/constitutive_models.py +++ b/src/underworld3/constitutive_models.py @@ -2593,9 +2593,10 @@ class _Parameters(_ParameterBase): def __init__(inner_self, _owning_model): inner_self._owning_model = _owning_model - # Raw sympy zero vector — NOT wrapped in UWexpressions, so the - # JIT sees the symbolic content directly (field variable dependencies). - # See the flux setter for the rationale. + # Raw sympy zero vector — NOT wrapped in UWexpressions. The default + # (Picard) tangent differentiates the flux WITHOUT unwrapping, so a + # wrapped whole-flux component has zero derivative w.r.t. the unknown + # and its gradient. See the flux setter for the full rationale. inner_self._flux = sympy.zeros(_owning_model.dim, 1) @property @@ -2607,12 +2608,22 @@ def flux(inner_self): def flux(inner_self, value: sympy.Matrix): """Set the flux expression (must be a vector of length dim). - Stores the raw sympy expression directly to preserve symbolic - dependencies (field variables, mesh coordinates) for the JIT. - Previously wrapped each component in a UWexpression via - ``validate_parameters``, which created independent Symbols that - the JIT treated as constants — losing the dependency on the - unknown field (e.g. ``u.sym``) and producing a singular matrix. + Stores the raw sympy expression directly. Previously each component + was wrapped in a UWexpression via ``validate_parameters`` — and the + default (Picard) tangent differentiates the flux WITHOUT unwrapping + (frozen-coefficient semantics; unwrap-before-differentiate runs only + on the consistent-Newton path). An opaque wrapper holding the ENTIRE + flux therefore has zero derivative w.r.t. the unknown and its + gradient: G0–G3 vanish and the operator is structurally singular. + The residual was never affected (the JIT's constants machinery + reveals non-constant wrappers correctly); the same pre-fix model + solved under ``consistent_jacobian=True``. Coefficient-level + wrapping in other models is safe — d(k*grad u)/d(grad u) = k — the + hazard is specific to wrapping a whole flux term. + + Side effect of raw storage: with nothing left to freeze, this + model's default tangent is effectively full Newton (the + Picard/Newton switch is a no-op for GenericFluxModel). """ dim = inner_self._owning_model.dim From 524465df8df4778481172183a5b3620e59f7f330 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 28 Jul 2026 17:39:55 +1000 Subject: [PATCH 3/3] review response (test hygiene): drop hand-wired Parameters._solver and the MUMPS pin MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The constitutive_model setter wires Parameters._solver — hand-wiring masked a regression in that path and weakened test_flux_update_via_solver_link. Plain LU replaces the MUMPS pin (house policy; review probes used LU throughout). Underworld development team with AI support from Claude Code --- tests/test_1060_generic_flux_model_update.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_1060_generic_flux_model_update.py b/tests/test_1060_generic_flux_model_update.py index 74c54637..b5843e5f 100644 --- a/tests/test_1060_generic_flux_model_update.py +++ b/tests/test_1060_generic_flux_model_update.py @@ -26,9 +26,9 @@ def _setup_solver(u, mesh, flux_expr): solver = Poisson(mesh, u_Field=u) solver.petsc_options["ksp_type"] = "preonly" solver.petsc_options["pc_type"] = "lu" - solver.petsc_options["pc_factor_mat_solver_type"] = "mumps" solver.constitutive_model = uw.constitutive_models.GenericFluxModel(solver.Unknowns) - solver.constitutive_model.Parameters._solver = solver + # NOTE: the constitutive_model setter wires Parameters._solver automatically; + # hand-wiring it here would mask a regression in that path. solver.constitutive_model.Parameters.flux = flux_expr solver.f = 1.0 solver.add_essential_bc(0.0, "Bottom")