unwrap_expression(symbolic_keep_constants) is exponential on nested constitutive models: 563 s for a 10k-node tree; JIT setup 10-40+ min on the notch #823

unwrap_expression(expr, mode="symbolic_keep_constants") is pathologically slow on nested constitutive models. That mode is the one _jacobian_unwrap uses before every consistent-tangent derivative, and prepare_for_cache_key uses it in the JIT. On the Spiegelman notch, the Jacobian setup alone took 10 to more than 40 minutes before any solve.

Measured (notch, refinement 1, ViscoPlastic with Drucker-Prager yield, residual strength, viscosity floor, smooth yield law)

_setup_pointwise_functions takes 616 s with the default law and 1,352 s with the powermean law. With a declared plastic_rate_strengthening it exceeds 40 minutes. With the prototype swapped in, the same setup takes 54 s, or 108 s including the JIT, both under cProfile.

The viscosity is a 130-node graph with 18 named sub-expressions. Unwrapped it becomes a 9,913-node tree, so the output is not what costs 563 s.

Cause

unwrap_expression iterates expr.subs(...) passes to a fixed point. Each pass traverses and re-canonicalises the whole growing tree. In symbolic_keep_constants mode, every atom on every pass also calls _is_truly_constant, which is itself a full nondimensional unwrap of that atom. Nothing is memoised.

Fix

Expand each UW atom once (memoised by identity), using the same _unwrap_atom rule and the same truly-constant predicate (also memoised). Then rebuild each expression in one xreplace pass. The output is structurally identical on the notch in both modes.

The rest of the compile (after the fix, profiled)

• C source generation: 76 s, of which atoms() scans over the expanded trees are 35 s and printing is 43 s.
• debugging_text, a string comment for every kernel, generated whether or not anyone debugs: 23 s.
• The sqrt guard's tree-wide .replace in _jacobian_unwrap: 16 s.
• Constant extraction: 12 s.
• Differentiation: 11 s.

Fixing these is tier 1. Tier 2 would generate code from the shared graph rather than the expanded tree, with each named sub-expression as one C temporary and the Jacobian built by chain rule. That needs a design note.

Found while running the shear-band width scan for PR #794.

Underworld development team with AI support from Claude Code