feat(jax_profiling): hannah ALMA preset for datacube + interferometer delaunay#63
Merged
Merged
Conversation
… delaunay Add a "hannah" instrument preset that mirrors Hannah Stacey's realistic ALMA science-case parameters (n_channels=34, n_visibilities=16984, pixel_scale=0.125", shape_native=(40, 40), mask_radius=2.3"). Promote mask_radius to a per-instrument field in the INSTRUMENTS dicts so it travels with the preset instead of being hardcoded. Gate the full-pipeline cube JIT (Part C) behind CUBE_FULL_JIT=1 - at n_channels=34 the lower+compile alone is ~70s, so the default cube run stays usable while step-by-step Part B JIT (the data feeding the shared-LtWL analysis) still completes in ~3 min. Per-channel hannah regression literal pinned: -204838.07924622478. Cube total at n_channels=34: -6964494.694372 (= 34 x per-channel under identical-channel reuse). Both regression assertions PASS in this commit. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
3 tasks
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
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
Add a
"hannah"instrument preset to thejax_profilingdatacube + interferometer Delaunay profilers that pins Hannah Stacey's realistic ALMA science-case parameters. Step-by-step cube cost at this scale (n_channels=34, 16984 visibilities/channel, 40×40 real-space): 205.92 s per cube evaluation, with an estimated 34.82 s saving available once the shared-Lᵀ W̃ L optimisation lands.Hannah's quoted settings:
n_channels = 34n_visibilities = 16984per channelpixel_scale = 0.125 "/pixelshape_native = (40, 40)mask_radius = 2.3 "Scripts Changed
jax_profiling/dataset_setup/interferometer.py— add"hannah"preset to theINSTRUMENTSdict; threadtransformer_classper preset (NUFFT for ALMA-scale, DFT for sma/alma); dataset auto-simulates todataset/interferometer/hannah/.jax_profiling/jit/interferometer/delaunay.py— add"hannah"to INSTRUMENTS; promotemask_radiusto a per-instrument field; convertEXPECTED_LOG_EVIDENCE_SMAto a per-instrumentEXPECTED_LOG_EVIDENCEdict (hannah pinned at-204838.07924622478); passraise_error_dft_visibilities_limit=Falseso the JAX-traceable DFT path can load >10k-visibility datasets (the JIT path is intentionally DFT until nufftax is wired in).jax_profiling/jit/datacube/delaunay.py— add"hannah"to INSTRUMENTS; promotemask_radius; switch defaultinstrument="hannah"andn_channels=34; gate Part C (full-pipeline cube JIT) behindCUBE_FULL_JIT=1since lower+compile at n_channels=34 is ~70 s on this laptop and was killing the default run; convert regression literal to a per-instrumentEXPECTED_LOG_EVIDENCE_PER_CHANNELdict.jax_profiling/dataset/interferometer/hannah/— new simulated dataset (16984-visibility NUFFT-simulated ALMA cube, 40×40 grid, single-channel reused 34× for the cube profile).jax_profiling/results/jit/interferometer/delaunay_likelihood_summary_hannah_v2026.5.14.2.{json,png}— captured single-channel hannah timings + bar chart.jax_profiling/results/jit/datacube/delaunay_likelihood_summary_hannah_v2026.5.14.2.{json,png}— captured 34-channel cube timings + bar chart.Upstream PR
Interferometer.from_fitsacceptsraise_error_dft_visibilities_limit. This profiler PR consumes that kwarg, so the library PR should merge first.Test Plan
-204838.07924622478PASSED.FitInterferometer.log_evidenceatrtol=1e-4(=-6964494.694372); eager cube regression vs34 × -204838.07924622478PASSED.instrument="sma"remains the default in the single-channel profiler; the existing SMA regression literal is unchanged;n_channels=4for the cube still works when toggled back.🤖 Generated with Claude Code