Skip to content

feat(jax_profiling): hannah ALMA preset for datacube + interferometer delaunay#63

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/datacube-hannah-preset
May 15, 2026
Merged

feat(jax_profiling): hannah ALMA preset for datacube + interferometer delaunay#63
Jammy2211 merged 1 commit into
mainfrom
feature/datacube-hannah-preset

Conversation

@Jammy2211
Copy link
Copy Markdown
Contributor

Summary

Add a "hannah" instrument preset to the jax_profiling datacube + 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 = 34
  • n_visibilities = 16984 per channel
  • pixel_scale = 0.125 "/pixel
  • shape_native = (40, 40)
  • mask_radius = 2.3 "

Scripts Changed

  • jax_profiling/dataset_setup/interferometer.py — add "hannah" preset to the INSTRUMENTS dict; thread transformer_class per preset (NUFFT for ALMA-scale, DFT for sma/alma); dataset auto-simulates to dataset/interferometer/hannah/.
  • jax_profiling/jit/interferometer/delaunay.py — add "hannah" to INSTRUMENTS; promote mask_radius to a per-instrument field; convert EXPECTED_LOG_EVIDENCE_SMA to a per-instrument EXPECTED_LOG_EVIDENCE dict (hannah pinned at -204838.07924622478); pass raise_error_dft_visibilities_limit=False so 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; promote mask_radius; switch default instrument="hannah" and n_channels=34; gate Part C (full-pipeline cube JIT) behind CUBE_FULL_JIT=1 since lower+compile at n_channels=34 is ~70 s on this laptop and was killing the default run; convert regression literal to a per-instrument EXPECTED_LOG_EVIDENCE_PER_CHANNEL dict.
  • 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

Test Plan

  • Single-channel hannah profiler runs end-to-end: eager-vs-JIT correctness PASSED, per-channel regression pinned at -204838.07924622478 PASSED.
  • 34-channel hannah cube profiler runs end-to-end with default (Part C skipped): per-step recompute matches summed FitInterferometer.log_evidence at rtol=1e-4 (= -6964494.694372); eager cube regression vs 34 × -204838.07924622478 PASSED.
  • Backward-compat: instrument="sma" remains the default in the single-channel profiler; the existing SMA regression literal is unchanged; n_channels=4 for the cube still works when toggled back.

🤖 Generated with Claude Code

… 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>
@Jammy2211 Jammy2211 added the pending-release Pending release build label May 15, 2026
@Jammy2211 Jammy2211 merged commit 31365b0 into main May 15, 2026
@Jammy2211 Jammy2211 deleted the feature/datacube-hannah-preset branch May 15, 2026 14:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release Pending release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant