Skip to content

Add heterodyning transforms for BNS chirp mass GNPE#355

Open
stephengreen wants to merge 2 commits into
bnsfrom
bns-heterodyne-transforms
Open

Add heterodyning transforms for BNS chirp mass GNPE#355
stephengreen wants to merge 2 commits into
bnsfrom
bns-heterodyne-transforms

Conversation

@stephengreen

Copy link
Copy Markdown
Member

Summary

  • Add factor_fiducial_waveform() and HeterodynePhase to waveform_transforms.py for phase heterodyning
  • Add GNPEChirp(GNPEBase) to gnpe_transforms.py for chirp mass GNPE conditioning
  • Add test suite in tests/gw/transforms/test_heterodyne.py (12 unit tests + 1 integration test)

Based on Max Dax's implementation on bns_add_dingo_pipe_max.

Note

A future PR will update ApplySVD, WhitenFixedASD, WaveformDataset, and WaveformGenerator to use a wrapped dict interface so HeterodynePhase can be used in decompression/compression chains.

Test plan

  • 12 unit tests for factor_fiducial_waveform, HeterodynePhase, and GNPEChirp
  • Integration test with real BNS waveforms (IMRPhenomPv2_NRTidal)
  • Full transform test suite passes (33/33), no regressions

Add factor_fiducial_waveform() and HeterodynePhase transform for
phase heterodyning, and GNPEChirp transform for chirp mass GNPE
conditioning. These enable removing the dominant chirp phase evolution
from BNS waveforms, which is critical for neural network training.

Based on Max Dax's implementation on bns_add_dingo_pipe_max.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@stephengreen stephengreen requested a review from max-dax February 24, 2026 23:07
ApplySVD and WhitenFixedASD now operate on sample["waveform"] instead of
flat dicts, matching the interface used by all other transforms. This
enables HeterodynePhase (which needs sample["parameters"]) to be part of
the compression/decompression chain.

Also wires in HeterodynePhase to compression (generate_dataset.py) and
decompression (waveform_dataset.py), and wraps data as sample dicts in
WaveformGenerator and WaveformDataset before passing to transforms.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@stephengreen

Copy link
Copy Markdown
Member Author

Added a second commit that updates the compression transform interface: ApplySVD and WhitenFixedASD now operate on sample dicts (sample["waveform"]) instead of flat dicts. This enables HeterodynePhase to participate in the compression/decompression chain (it needs access to sample["parameters"] for chirp_mass). Also wires HeterodynePhase into generate_dataset.py and waveform_dataset.py, and wraps data in WaveformGenerator before passing to transforms. All 101 tests pass.

stephengreen added a commit that referenced this pull request Jul 7, 2026
From PR #355: factor_fiducial_waveform and HeterodynePhase (leading-order
chirp heterodyning with an optional 1PN correction; a fixed-parameters
mode for inference-time use) and the GNPEChirp proxy transform, with
their tests -- including a slow integration test on real NRTidal
waveforms (roundtrip mismatch < 1e-14, an order of magnitude fewer zero
crossings after heterodyning).

From PR #354: d_inner_h_complex_multi gains return_rho2opt, returning
(h|h) alongside (d|h) so the full likelihood decomposition
log L = log_Zn + Re[(d|h) e^(2i phase)] - (h|h)/2 is available from a
single waveform generation (the basis for the chirp-mass sweep), with
its tests. The synthetic-phase profile path is unchanged (re-verified
bit-exact).

PR #354's sampler extensions (fixed_context_parameters,
_preprocess_data) are deliberately not ported: the composed spine
already expresses them as a DeltaFactor pin and the context's one-time
prepared_data. The remaining #355 changes (the sample-dict interface for
the dataset-compression transforms) are training-side and come later as
a package.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant