perf: enable apply_sparse_operator on Delaunay JAX profilers#65
Merged
Conversation
Wire the now-unblocked apply_sparse_operator(use_jax=True) path into the two Delaunay JAX profilers so per-fit curvature assembly uses the FFT-based sparse precision-matrix preload instead of dense DFT per source pixel. Unblocked by PyAutoLabs/PyAutoArray#316 (Pmax > 1 extent-indexing fix); the path was previously guarded with NotImplementedError on Delaunay by PR #315. - jax_profiling/jit/interferometer/delaunay.py: new "apply_sparse_operator" timer section right after dataset_load. - jax_profiling/jit/datacube/delaunay.py: chain .apply_sparse_operator(use_jax=True, show_progress=False) onto the per-channel from_fits in the dataset_list comprehension. Validation: - jit/interferometer/delaunay.py: end-to-end run, eager log_likelihood = -3152.876, figure_of_merit = -3167.526, all 11 per-step JIT stages completed. - jit/datacube/delaunay.py: end-to-end run on the hannah 34-channel cube, cube log_evidence = -6964494.694 (regression assertion PASSED), all 8 per-step JIT stages completed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Enables
apply_sparse_operator(use_jax=True)on the two Delaunay JAX profilers injax_profiling/jit/. Per-fit curvature assembly now uses the FFT-based sparse precision-matrix preload instead of dense DFT per source pixel.This was blocked by the
NotImplementedErrorguard on Delaunay added in PyAutoLabs/PyAutoArray#315 and replaced by the actual math fix in PyAutoLabs/PyAutoArray#316.Scripts Changed
jax_profiling/jit/interferometer/delaunay.py— new"apply_sparse_operator"timer section right afterdataset_load, callingdataset.apply_sparse_operator(use_jax=True, show_progress=True). The 11 per-step JIT-profiled stages downstream now operate on the sparse-curvature path.jax_profiling/jit/datacube/delaunay.py— chains.apply_sparse_operator(use_jax=True, show_progress=False)onto each per-channelfrom_fitsin thedataset_listcomprehension. 34-channel per-channel precision-matrix preload.Upstream PR
Test plan
python jax_profiling/jit/interferometer/delaunay.py— end-to-end. Eagerlog_likelihood = -3152.876,figure_of_merit = -3167.526. All 11 per-step JIT stages completed. Results JSON + PNG written.python jax_profiling/jit/datacube/delaunay.py— end-to-end on the hannah 34-channel cube. Eager cubelog_evidence = -6964494.694(regression assertion PASSED — agrees with the existing hardcoded reference, confirming the upstream library fix preserves correctness on the eager numpy path too). All 8 per-step JIT stages completed. Results JSON + PNG written.🤖 Generated with Claude Code