[pbrubeck]

Summary

AxisForest.getitem was trying the indices on each of its trees and skipping the trees that they did not match. The trees were caching only successful indexing, so every indexing of a forest was indexing the mismatched trees again. In Firedrake this was happening on every Function.assign or zero on a subset, e.g. each DirichletBC.zero.

AxisForest.getitem is now cached for static indices, with the same strategy as AbstractNonUnitAxisTree.getitem:

 AxisForest.getitem(indices)
+  if indices have loop indices
+    return _getitem_uncached(indices)
+  return _getitem_cached(indices)          # @cached_method()

+AxisForest._getitem_uncached(indices)
   for tree in trees
     try tree.getitem(indices)
     except InvalidIndexTargetException, IncompatibleAxisTargetException: skip

Indexing with loop indices is still not cached, so the FIXME about this bottleneck now names loop indices.

Evidence

test_axis_forest_getitem_is_cached (new) indexes a forest twice, and fails if the second indexing indexes any of its trees:

• Before: fails
  After: passes

DirichletBC.zero for CG2 on a 16×16 mesh (minimum of 200 calls):

	before	after	pyop2
bc.zero(Function)	0.53 ms	0.065 ms	0.006 ms
bc.zero(Cofunction)	0.53 ms	0.069 ms	0.006 ms

Merge Danger

Door: two-way

Blast Radius: small

Only indexing of axis forests with static indices changes. The indexed trees are now kept for the lifetime of the forest, as they already were for the lifetime of each tree.

---

AI assistance: this PR was written with Claude Code (Claude Opus 5.5).

🤖 Generated with Claude Code