JIT constants nested in wrapper expressions are live-rampable — reveal before substitute

src/underworld3/utilities/_jitextension.py

@@ -260,32 +260,60 @@ def counts(self):
                 len(self.bd_residual), len(self.bd_jacobian))
 
 
+def _reveal_constants(fn):
+    """Expand non-constant UWexpressions to fixpoint, KEEPING truly-constant
+    atoms symbolic — so constants nested at ANY depth surface as atoms.
+
+    This must run BEFORE the ``constants[]`` substitution: a top-level
+    ``xreplace`` cannot see a constant hidden inside a nested UWexpression
+    (every ``Parameters.*`` value is template-wrapped in one), so nested
+    constants were silently folded to C literals while the manifest still
+    listed them as live — issue #302. The keep-constants predicate here is
+    the same ``_is_truly_constant`` used to build the manifest, so the set
+    of atoms revealed is exactly the set the manifest routes to
+    ``constants[]``.
+    """
+    from underworld3.function.expressions import (
+        unwrap_expression,
+        UWDerivativeExpression,
+    )
+
+    if fn is None:
+        return fn
+    if isinstance(fn, UWDerivativeExpression):
+        fn = fn.doit()
+    if isinstance(fn, sympy.MatrixBase):
+        return fn.applyfunc(
+            lambda e: unwrap_expression(e, mode='symbolic_keep_constants'))
+    return unwrap_expression(fn, mode='symbolic_keep_constants')
+
+
 def prepare_for_cache_key(fn, constants_subs_map):
     """Prepare a single expression for JIT cache hashing.
 
-    Two-phase process:
-    1. Substitute constant UWexpressions with ``_JITConstant`` placeholders
+    Three-phase process (mirrors the codegen lowering in ``_createext`` so
+    the cache key and the generated C agree — issue #302):
+    1. Reveal nested constants (``_reveal_constants``).
+    2. Substitute manifested constants with ``_JITConstant`` placeholders
        so that changing a constant's *value* does not invalidate the cache.
-    2. Unwrap remaining (non-constant) UWexpressions to pure SymPy so the
-       hash is deterministic.
-
-    Parameters
-    ----------
-    fn : sympy expression or None
-        The expression to expand.
-    constants_subs_map : dict or None
-        Mapping from UWexpression symbols to ``_JITConstant`` placeholders.
+    3. Unwrap the remaining UW atoms to pure SymPy so the hash is
+       deterministic.
     """
-    # Phase 1: Substitute constants with _JITConstant placeholders
-    if constants_subs_map and fn is not None:
+    # Phase 1: reveal constants nested inside other UWexpressions. Loud on
+    # failure, exactly like the codegen path — a silent fallback here would
+    # hash constant VALUES into the cache key and force a recompile on
+    # every ramp (adversarial-review finding).
+    fn_structural = _reveal_constants(fn)
+
+    # Phase 2: Substitute constants with _JITConstant placeholders
+    if constants_subs_map and fn_structural is not None:
         try:
-            fn_structural = fn.xreplace(constants_subs_map) if hasattr(fn, "xreplace") else fn
+            if hasattr(fn_structural, "xreplace"):
+                fn_structural = fn_structural.xreplace(constants_subs_map)
         except Exception:
-            fn_structural = fn
-    else:
-        fn_structural = fn
+            pass
 
-    # Phase 2: Unwrap remaining (non-constant) expressions
+    # Phase 3: Unwrap remaining (non-constant) expressions
     return underworld3.function.expressions.unwrap(
         fn_structural, keep_constants=False, return_self=False
     )
@@ -962,18 +990,55 @@ def _basescalar_ccode(self, printer):
         # Save original for debugging
         fn_original = fn
 
-        # Two-phase unwrap:
-        # Phase 1: Substitute constant UWexpressions with _JITConstant symbols
+        # Three-phase lowering (issue #302 — must match prepare_for_cache_key):
+        # Phase 1: reveal constants nested inside other UWexpressions, so the
+        #          substitution below can reach them. A top-level xreplace
+        #          missed constants inside template-wrapped parameters and
+        #          baked them as C literals while the manifest listed them.
+        fn = _reveal_constants(fn)
+
+        # A truly-constant atom the manifest does NOT know about would be
+        # silently folded to a literal in phase 3 — the manifest and the
+        # C source must never disagree (issue #302).
+        from underworld3.function.expressions import UWexpression as _UWexpr
+        if constants_subs_map is not None and hasattr(fn, 'atoms'):
+            unmanifested = [
+                a.name for a in _stable_sorted(fn.atoms(sympy.Symbol))
+                if isinstance(a, _UWexpr)
+                and _is_truly_constant(a, _UWexpr)
+                and a not in constants_subs_map
+            ]
+            if unmanifested:
+                raise RuntimeError(
+                    f"JIT constants manifest is incomplete: constant expression(s) "
+                    f"{unmanifested} appear in a kernel but have no constants[] "
+                    f"slot — they would be baked into the C source (issue #302)."
+                )
+
+        # Phase 2: Substitute constant UWexpressions with _JITConstant symbols
         #          These survive into C code as constants[i]
         if constants_subs_map and fn is not None:
             try:
                 fn = fn.xreplace(constants_subs_map) if hasattr(fn, 'xreplace') else fn
             except Exception:
                 pass
 
-        # Phase 2: Unwrap remaining non-constant UWexpressions to numerical values
+        # Phase 3: Unwrap remaining non-constant UWexpressions to numerical values
         fn = underworld3.function.expressions.unwrap(fn, keep_constants=False, return_self=False)
 
+        # A manifested constant surviving to here bypassed its constants[]
+        # slot and is about to be baked — refuse rather than freeze the
+        # parameter silently (issue #302).
+        if constants_subs_map and hasattr(fn, 'atoms'):
+            baked = [a.name for a in _stable_sorted(fn.atoms(sympy.Symbol))
+                     if a in constants_subs_map]
+            if baked:
+                raise RuntimeError(
+                    f"Manifested constant(s) {baked} were not routed through "
+                    f"constants[] and would be baked into the C source "
+                    f"(issue #302)."
+                )
+
         if isinstance(fn, sympy.vector.Vector):
             fn = fn.to_matrix(mesh.N)[0 : mesh.dim, 0]
         elif isinstance(fn, sympy.vector.Dyadic):

tests/test_0103_jit_rampable_constants.py

@@ -0,0 +1,85 @@
+"""Live-rampable JIT constants — issue #302.
+
+A constant UWexpression listed in ``solver.constants_manifest`` must be
+readable from its ``constants[]`` slot in the compiled kernels, never baked
+as a C literal. The old two-phase lowering substituted constants[] slots
+with a TOP-LEVEL xreplace before unwrapping, so any constant nested inside
+another UWexpression — which includes EVERY ``Parameters.*`` value, since
+the setter template-wraps it — was silently folded to a literal: ramping
+its ``.sym`` between solves had no effect until a full rebuild.
+"""
+
+import numpy as np
+import pytest
+import underworld3 as uw
+
+pytestmark = [pytest.mark.level_1, pytest.mark.tier_b]
+
+
+def _build():
+    mesh = uw.meshing.StructuredQuadBox(
+        elementRes=(8, 8), minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0))
+    T = uw.discretisation.MeshVariable("T302", mesh, 1, degree=2)
+
+    k_direct = uw.expression(r"k_d", 1.0, "constant at parameter top level")
+    k_nested = uw.expression(r"k_n302", 1.0, "constant nested in a wrapper")
+    # Non-constant wrapper: the collection recurses into it and manifests
+    # k_nested, but a top-level substitution cannot see inside it — the
+    # baked-constant topology from issue #302.
+    wrapper = uw.expression(
+        r"\eta_{w302}", k_nested * 2.0 + 0.05 * T.sym[0] ** 2, "wrapper")
+
+    poisson = uw.systems.Poisson(mesh, u_Field=T)
+    poisson.constitutive_model = uw.constitutive_models.DiffusionModel
+    poisson.constitutive_model.Parameters.diffusivity = (
+        k_direct * (1.0 + 0.1 * T.sym[0] ** 2) + wrapper)
+    poisson.f = 1.0
+    poisson.add_dirichlet_bc(0.0, "Top")
+    poisson.add_dirichlet_bc(0.0, "Bottom")
+    return poisson, T, k_direct, k_nested
+
+
+def _mean(poisson, T):
+    poisson.solve()
+    return float(np.asarray(T.data)[:, 0].mean())
+
+
+def test_manifested_constants_ramp_without_rebuild():
+    poisson, T, k_direct, k_nested = _build()
+    base = _mean(poisson, T)
+
+    names = [e.name for _i, e in poisson.constants_manifest]
+    assert "k_d" in names and "k_n302" in names
+
+    # Ramp each constant via .sym alone — REAL solves, no rebuild between.
+    k_direct.sym = 10.0
+    direct_ramped = _mean(poisson, T)
+    k_direct.sym = 1.0
+
+    k_nested.sym = 10.0
+    nested_ramped = _mean(poisson, T)
+    k_nested.sym = 1.0
+
+    assert abs(direct_ramped - base) > 1e-8, (
+        "top-level constant did not ramp without a rebuild")
+    assert abs(nested_ramped - base) > 1e-8, (
+        "NESTED constant did not ramp without a rebuild (issue #302)")
+
+    # Ramping back restores the base solution exactly — no hidden state.
+    assert np.isclose(_mean(poisson, T), base, rtol=1e-12)
+
+
+def test_ramped_solution_matches_rebuilt_solution():
+    # The live-ramped answer must equal the answer from a full rebuild at
+    # the same parameter value — the manifest and the C source agree.
+    poisson, T, k_direct, k_nested = _build()
+    _mean(poisson, T)
+    k_nested.sym = 10.0
+    ramped = _mean(poisson, T)
+
+    poisson2, T2, _kd2, k_nested2 = _build()
+    k_nested2.sym = 10.0
+    rebuilt = _mean(poisson2, T2)
+
+    assert np.isclose(ramped, rebuilt, rtol=1e-10), (
+        f"ramped {ramped} != rebuilt {rebuilt}")