Merge pull request #87 from underworldcode/feature/petscds-constants

Route constant solver parameters through PetscDS constants[] array

Author: lmoresi

docs/beginner/parameters.md

@@ -251,6 +251,47 @@ mpirun -np 256 python convection.py \
     -uw_max_iterations 100
 ```
 
+## Using Parameters with Solvers: Expressions
+
+When passing parameters to solver constitutive models, wrap them in
+`uw.expression()`. This creates a named symbolic container that the JIT
+compiler can update efficiently — changing the value between time steps
+does **not** trigger recompilation of the C extension.
+
+```python
+import underworld3 as uw
+
+# Define parameters
+VISCOSITY = 1e21      # Named constant
+MODULUS = 1e10        # Named constant
+DT = 0.01            # Timestep
+
+params = uw.Params(
+    uw_viscosity = VISCOSITY,
+    uw_modulus = MODULUS,
+)
+
+# Wrap in expressions for solver use
+eta = uw.expression("eta", params.uw_viscosity)
+mu = uw.expression("mu", params.uw_modulus)
+dt_e = uw.expression("dt_e", DT)
+
+# Pass expressions to constitutive model
+stokes.constitutive_model.Parameters.shear_viscosity_0 = eta
+stokes.constitutive_model.Parameters.shear_modulus = mu
+stokes.constitutive_model.Parameters.dt_elastic = dt_e
+
+# Time-stepping: change dt without recompilation
+for step in range(100):
+    dt_e.sym = compute_new_timestep()  # Updates value, ~0ms
+    stokes.solve()                      # No JIT rebuild needed
+```
+
+Without expressions, changing a solver parameter between steps requires
+setting `_force_setup=True` on the solve call, which triggers a full JIT
+recompilation (~5–15 seconds). With expressions, parameter updates go
+through PETSc's `constants[]` array and cost essentially nothing.
+
 ## Angle Units
 
 Angles work naturally - you can define in degrees and provide radians (or vice versa):

docs/developer/UW3_Developers_MathematicalObjects.md

@@ -428,13 +428,13 @@ The JIT compilation system needs to:
 1. Identify SymPy Function atoms in expressions
 2. Map them to PETSc vector components
 3. Generate C code with appropriate substitutions
+4. Allow constant parameters to change without recompilation
 
 ## The Solution: Transparent SymPy Objects
 
 The mathematical object system preserves JIT compatibility by ensuring all operations return pure SymPy objects:
 
 ```python
-
 # User writes natural syntax
 momentum = density * velocity
 
@@ -447,23 +447,50 @@ atoms = momentum.atoms(sympy.Function)  # Finds V_0, V_1
 # Maps to velocity.fn for PETSc substitution
 ```
 
-## Expression Unwrapping
+## Expression Unwrapping and Constants
 
-The `unwrap()` function resolves nested expressions before compilation:
+The JIT compiler performs a **two-phase unwrap** on expressions:
 
-```python
+1. **Phase 1 — Constants extraction**: `UWexpression` atoms that resolve to
+   pure numbers (no spatial/field dependencies) are replaced with
+   `_JITConstant` symbols that render as `constants[i]` in C code.
+
+2. **Phase 2 — Full unwrap**: Remaining `UWexpression` atoms are expanded
+   to their numerical values and baked into the C code.
 
+```python
 # Expression with nested UWexpressions
 complex_expr = alpha * (temperature - T0) * velocity
 
-# unwrap() substitutes all UWexpression.sym values
-unwrapped = unwrap(complex_expr)
-# Result: 2e-5 * (T(x,y,z) - 293) * Matrix([[V_0(x,y,z)], [V_1(x,y,z)]])
+# Phase 1: alpha and T0 are constants → constants[0], constants[1]
+# Phase 2: temperature and velocity are field variables → petsc_a[], petsc_u[]
+
+# Generated C code:
+#   result = constants[0] * (petsc_a[0] - constants[1]) * petsc_u[0];
+```
+
+This means **changing `alpha` or `T0` between solves does not require
+recompilation** — only `PetscDSSetConstants()` is called (~0ms vs ~10s).
+
+## Why Use UWexpressions for Solver Parameters
+
+**UWexpressions are the preferred way to define solver parameters.**
+Using raw numbers forces recompilation when values change:
 
-# JIT compilation proceeds normally
-compiled = uw.systems.compile(unwrapped)
+```python
+# PREFERRED — expression parameter, efficient for time-stepping
+eta = uw.expression("eta", 1e21)
+stokes.constitutive_model.Parameters.shear_viscosity_0 = eta
+# Changing eta.sym later → no recompilation
+
+# AVOID for time-varying parameters — raw number, requires rebuild
+stokes.constitutive_model.Parameters.shear_viscosity_0 = 1e21
+# Changing this later → full JIT rebuild (~10s)
 ```
 
+This is particularly important for viscoelastic and Navier-Stokes
+solvers where parameters like `dt_elastic` change every time step.
+
 # Migration from Legacy Patterns
 
 ## Automatic Migration Patterns

docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md

@@ -610,16 +610,34 @@ def test_swarm_functionality():
 
 ### JIT Compilation Issues
 
-If you see generated C code with symbolic expressions instead of numbers:
+**Slow time-stepping loops**: If each `solver.solve()` call takes 10+ seconds
+in a time-stepping loop, the solver is probably recompiling the JIT extension
+every step. Use `UWexpression` objects for any parameter that changes between
+steps:
+
+```python
+# FAST — expression parameter, no recompilation on change
+dt_e = uw.expression("dt_e", 0.01)
+model.Parameters.dt_elastic = dt_e
+
+for step in range(100):
+    dt_e.sym = compute_timestep()  # Updates constants[], ~0ms
+    solver.solve()                  # No JIT rebuild
+```
+
+**Symbolic names in generated C code**: If you see LaTeX-like names in the
+generated C code (e.g., `\eta` instead of a number or `constants[i]`):
 
 ```text
-// ERROR symptom in generated code:
-out[0] = 1.0/{ \eta \hspace{ 0.0006pt } };  // Should be numeric!
+// ERROR symptom — expression not unwrapped:
+out[0] = 1.0/{ \eta \hspace{ 0.0006pt } };
 ```
 
-**Cause**: `unwrap(fn, keep_constants=False)` not properly unwrapping constants.
+**Cause**: A `UWexpression` was not properly detected as constant or unwrapped.
 
-**Solution**: Check that constants (like UWQuantity) are being unwrapped to numeric values.
+**Solution**: Check that the expression resolves to a pure number when fully
+unwrapped. Composite expressions containing mesh variables or coordinates
+cannot be routed through `constants[]`.
 
 ### PETSc DM Errors with Swarms