From 1046c9e22f90e425ccf9c5cc14e881c5da5bee34 Mon Sep 17 00:00:00 2001
From: Ben Knight <55677727+bknight1@users.noreply.github.com>
Date: Tue, 14 Jul 2026 14:53:55 +0800
Subject: [PATCH] fix(constitutive_model): make flux_jacobian settable for
 custom Jacobian tangents
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The flux_jacobian property on Constitutive_Model was read-only and
always returned None.  This prevented models from providing a
simplified flux expression for the Jacobian independently of the
exact residual flux.

Now it has a getter/setter backed by an instance attribute.
Default remains None (identical behaviour) — setting it to a
SymPy expression and enabling consistent_jacobian = True on the
solver gives the solver a custom expression to differentiate for
the Jacobian tangent while the residual keeps the exact flux.

Use cases:
- Viscoplastic models with sharp yield kinks: supply a smooth
  constitutive law for the tangent only (Newton avoids
  Min-composite derivative discontinuities).
- Multicomponent diffusion with composition-dependent
  coefficients: supply a constant-coefficient flux so the
  Jacobian gets a clean SPD stiffness matrix (D_ii · ∇·∇)
  instead of the zero-G3 from UWexpression-wrapped coefficients,
  reducing Newton iterations from ~50 to ~3.
---
 src/underworld3/constitutive_models.py | 39 +++++++++++++++++++++-----
 1 file changed, 32 insertions(+), 7 deletions(-)

diff --git a/src/underworld3/constitutive_models.py b/src/underworld3/constitutive_models.py
index 2944ea0c..7d37eecf 100644
--- a/src/underworld3/constitutive_models.py
+++ b/src/underworld3/constitutive_models.py
@@ -305,6 +305,12 @@ def __init__(self, unknowns, material_name: str = None):
         self._class_instance_number = Constitutive_Model._class_instance_counts[class_name]
         Constitutive_Model._class_instance_counts[class_name] += 1
 
+        # Backing attribute for settable flux_jacobian (set to a custom
+        # SymPy expression to override the Jacobian tangent independently
+        # of the residual flux).  None by default, meaning the solver
+        # differentiates the exact flux.
+        self._flux_jacobian = None
+
         self.Unknowns = unknowns
 
         u = self.Unknowns.u
@@ -581,17 +587,36 @@ def flux_jacobian(self):
         exact :attr:`flux` (the Newton fix unwraps it first; a generic Min/Max
         kink-smoothing fallback then rounds any remaining yield kink).
 
-        A model whose flux has a non-smooth yield kink (e.g. hard-``Min``
-        viscoplasticity) may override this to supply a physically-motivated
-        *smooth constitutive law for the tangent only*. The residual still uses
-        the exact :attr:`flux`, so the converged solution satisfies the true
-        constitutive law — only the Newton search direction is smoothed, giving
-        a robust, line-search-friendly tangent without changing the answer.
+        Set this to a custom SymPy expression to override the Jacobian tangent
+        independently of the residual flux.  The residual still uses the exact
+        :attr:`flux`, so the converged solution satisfies the true constitutive
+        law — only the Newton search direction is smoothed, giving a robust,
+        line-search-friendly tangent without changing the answer.
+
+        Use cases:
+
+        * A model whose flux has a non-smooth yield kink (e.g. hard-``Min``
+          viscoplasticity) supplies a physically-motivated *smooth law for the
+          tangent only*.
+
+        * A model whose flux contains composition-dependent coefficients
+          (e.g. multicomponent diffusion) may supply a *constant-coefficient*
+          version to give the solver a clean SPD Jacobian while keeping the
+          exact physics in the residual.
 
         Shape must match :attr:`flux` (the solver substitutes it for ``F1`` when
         forming the velocity-gradient Jacobian blocks).
         """
-        return None
+        return self._flux_jacobian
+
+    @flux_jacobian.setter
+    def flux_jacobian(self, value):
+        """Set a custom Jacobian flux expression (or None to revert)."""
+        self._flux_jacobian = value
+        # Signal the solver to rebuild its pointwise Jacobian function.
+        # Without this, the change is silently ignored until something
+        # else triggers a re-setup.
+        self._solver_is_setup = False
 
     def _reset(self):
         """Flags that the expressions in the consitutive tensor need to be refreshed and also that the