[Shizheng-Wen]

Summary

Transformation caches its geometric quantities (shape gradients, JxW, facet Jacobians, ...) lazily as buffers, but the cache had three problems.

1. Assembly bypassed it. ElementAssembler.__call__ and NodeAssembler.__call__ take shape gradients from Transformation.batch_shape_grad_jxw. That method used the cache only if energy() or the shape_grad/JxW properties had filled it, and never filled it itself. So every plain asm() call inverted every element Jacobian again: one LAPACK call per 2×2/3×3 matrix on the CPU.
2. update_points didn't drop anything, despite its docstring ("reset derived buffers"). Every points= argument goes through it.
   • After energy() had filled the cache, asm(points=new) returned the old geometry's matrix (max error 1.0 on a stretched mesh).
   • FacetAssembler, whose facet quantities always go through the caching properties, did this from its second call on.
   • So did NodeAssembler's compile() fast path, which captures shape_grad/jxw when it is built.
   • The user guide (forms.rst) promises "all paths re-cache the Transformation on new points".
3. Differentiable geometry (points that require grad): the cached tensors kept the first call's autograd graph.
   • A second backward failed with "Trying to backward through the graph a second time".
   • New points were ignored.

Changes

• Cache rules (Transformation):
  • A full-range request (batch_size=-1, the ElementAssembler default) fills the shape-gradient + JxW cache, and later calls reuse it.
  • Chunked requests (batch_size>0, the memory-saving mode, and NodeAssembler's default) are served from the cache when it exists. Otherwise they compute their chunk and keep nothing, as before.
  • update_points drops every geometry-derived buffer, unless it is handed the tensor it already holds, unmodified (identity + _version). Passing the same coordinates again therefore keeps the cache.
  • While the points require grad (with grad mode on), geometry is recomputed on every call, so each call builds its own graph.
  • The Jacobians are no longer cached alongside the shape gradients, since assembly never uses them. jacobian rebuilds them on access, without an inverse. This keeps the cache to what assembly needs.
• small_inverse (tensormesh/element/element.py, internal): closed-form adjugate inverse for D ≤ 3, falling back to torch.linalg.inv above that. It raises ValueError on a singular (degenerate) element, as torch.inverse did. It replaces torch.inverse in eval_shape_grad, in the Nanson scale and in MixedElementAssembler, which recomputes its geometry on every call by design.
• NodeAssembler.compile(): the fast path's rebuild key includes Transformation._geometry_version.
• The batch_size docstrings say what gets cached.

Timings

Laplacian, quadrature_order=2, one Clariden node (Grace CPU / GH200). "New points" passes a fresh coordinate tensor every call, as a moving mesh would.

mesh	device	main: first / repeated / new points	this PR: first / repeated / new points
12,168 P1 triangles	CPU	148 / 146 / 146 ms	2.4 / 0.9 / 2.2 ms
12,168 P1 triangles	GPU	1.8 / 1.3 / 1.4 ms	1.9 / 0.6 / 1.5 ms
48,000 P1 tetrahedra	CPU	780 / 823 / 795 ms	7.6 / 5.4 / 8.5 ms
48,000 P1 tetrahedra	GPU	2.4 / 1.6 / 1.6 ms	1.7 / 0.6 / 1.3 ms

The whole test suite on the same node: 278 s on main, 102 s with this PR.

Tests

tests/element/test_transformation_cache.py (17 tests). Meshes are built with NumPy + meshio, so the tests don't need gmsh.

• A repeated asm() evaluates the shape gradients once. Chunked assembly keeps nothing and agrees with full assembly, both before and after the cache is filled.
• New points are honoured after the cache was filled, by:
  • ElementAssembler, with the cache filled by assembly or by energy();
  • NodeAssembler, chunked and full;
  • NodeAssembler.compile();
  • FacetAssembler.
    Each result is compared with a fresh assembler on the moved mesh.
• The same unchanged points keep the cache. An in-place edit of them drops it.
• Points requiring grad:
  • d(area)/d(stretch) = 1, for two stretches in a row;
  • repeated autograd.grad through the same leaf;
  • backward through an assembled matrix.
• small_inverse against torch.linalg.inv for D = 1…4, gradcheck, and singular input.

Against unmodified main, 6 of the 10 caching tests fail. Among the failures, the autograd test fails with the graph error above. The other 4 pass, because nothing is cached on main to go stale.

The full suite on a Clariden GPU node gives an identical failure set on main and on this branch. All of those failures are environmental: gmsh has no aarch64 build there. The branch adds 17 passes. CI covers the gmsh-dependent tests.

Notes

• Memory: the default full-range assembly now keeps shape_grad ([E, Q, B, D]) and JxW per Transformation between calls. Its peak memory during a call is unchanged. Memory-bound runs that pass batch_size>0 keep nothing, as before.
• Transformation.jacobian is now computed on its own (no inverse) and cached separately. Previously it was a side product of shape_grad. The values are identical.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Dst8zhPLnxVpZMBhR4Pvfs