Fix GenericFluxModel flux storage — preserve symbolic field dependencies

src/underworld3/constitutive_models.py

@@ -2593,17 +2593,11 @@ 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. 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
         def flux(inner_self):
@@ -2612,26 +2606,35 @@ 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. 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
 
             # 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

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.constitutive_model = uw.constitutive_models.GenericFluxModel(solver.Unknowns)
+    # 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")
+    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}"
+    )