Consider caching strategy for repeated Poisson bracket calls

[anjor]

Background

In typical KRMHD time-stepping, the Poisson bracket is called repeatedly with the same field pairs (e.g., {φ, A∥} computed at each RK4 substep). Currently, each call recomputes:

1. Four derivative transforms (∂f/∂x, ∂f/∂y, ∂g/∂x, ∂g/∂y)
2. Four inverse FFTs
3. One forward FFT
4. Dealiasing

Potential Optimizations

1. Cache derivatives: If f or g unchanged, reuse ∂f/∂x, ∂f/∂y
2. Batch operations: Compute multiple brackets in parallel (vmap)
3. Lazy evaluation: Use SpectralField2D/3D lazy caching for derivatives
4. JIT compilation: Ensure XLA optimizes repeated calls

Investigation Needed

• Profile typical time-stepping workload (RK4 with 3-4 fields)
• Measure cache hit rate if derivatives were cached
• Benchmark batch bracket evaluation vs sequential
• Check if JAX XLA already optimizes this pattern
• Estimate memory overhead of caching derivatives

Decision Criteria

Only implement if:

• Profiling shows bracket is >50% of runtime (expected)
• Cache provides >20% speedup
• Memory overhead acceptable (<2x current usage)

Related

• Issue Step 4: Poisson Bracket #4 (Poisson bracket implementation)
• Issue Step 8: Time Integrator #8 (Time integration - RK4)
• Issue Add performance benchmark for Poisson bracket #35 (Performance benchmarks)

[anjor]

Status Check - Investigation phase, status unchanged

Current State

• No caching implemented
• No profiling data yet
• Original issue correctly identified this as "investigation needed"

Context Update

Since this issue was created, the GANDALF timestepper (PR #47, #50) has been implemented. We now have real usage patterns to profile.

Actual Poisson bracket usage in gandalf_step():

# In timestepping.py:
bracket1 = poisson_bracket_3d(z_plus, z_minus_lap_perp, ...)  
bracket2 = poisson_bracket_3d(z_plus, z_minus, ...)
# → Two brackets per Elsasser variable
# → Four total brackets per RK2 substep
# → Eight total per full timestep

So yes, there is repeated computation, but fields change every substep, so caching would have limited benefit.

Decision Criteria (from original issue)

Only implement if:

• Profiling shows bracket is >50% of runtime (expected) ✓ Likely true
• Cache provides >20% speedup - ❓ Need profiling
• Memory overhead acceptable (<2x current usage) - ❓ Need measurement

Recommended Next Steps

Phase 1: Profile with Issue #35 (Performance benchmarks)

# What we need to measure:
1. % of time in Poisson bracket vs other operations
2. FFT cost vs everything else
3. Memory usage at 256³ resolution

Phase 2: Decide on caching strategy

• If bracket is >50% and cache helps >20% → implement
• If not → close as "investigated, not worth it"

Alternative optimization: Issue #52 (diagnostics) showed Python loops are bigger bottleneck

• Fix Issue Optimize diagnostics spectrum calculation performance #52 first (10-100× speedup from vectorization)
• Then profile to see if Poisson bracket caching still matters

Related Issues

• Issue Add performance benchmark for Poisson bracket #35 - Needs profiling infrastructure
• Issue Optimize diagnostics spectrum calculation performance #52 - Already found bigger performance issue in diagnostics
• Issue Optimize Elsasser RHS Laplacian computations #41 - Similar Laplacian optimization question

Current Status

Keeping open for investigation, but blocked on Issue #35 (profiling). Recommend:

1. Implement Issue Add performance benchmark for Poisson bracket #35 first (benchmarks)
2. Fix Issue Optimize diagnostics spectrum calculation performance #52 (diagnostics performance)
3. Re-profile with fixed diagnostics
4. Then decide if Poisson bracket caching is worth it

Priority: Low - Wait for profiling data before investing effort.

Keeping open pending profiling data. 🔍