sage.dsp.multibanding

Frequency-domain multibanding utilities.

This module compresses uniformly sampled rFFT tensors by keeping a variable frequency grid. It is meant as the frequency-domain analogue of time-domain multirate sampling: keep fine resolution where it matters, and use coarser frequency spacing elsewhere.

Attributes

PoolMode

Classes

`FrequencyBand`	A frequency interval sampled with an integer bin stride.
`FrequencyBandLayout`	A concrete multiband layout on an rFFT frequency grid.
`FrequencyMultibandCompressor`	Compress rFFT data with a variable frequency grid.

Functions

`make_dyadic_frequency_bands`(*, f_min, f_max[, ...])	Make bands whose stride doubles with frequency.
`describe_layout`(layout)	Return a compact, notebook-friendly description of a layout.
`make_prior_informed_frequency_bands`(*, m1, m2[, s1z, ...])	Method 1 — analytic chirp-time boundary placement.
`make_empirical_frequency_bands`(*, m_min, m_max[, ...])	Method 2 — simulation-based boundary placement.

Module Contents

PoolMode[source]

class FrequencyBand[source]

A frequency interval sampled with an integer bin stride.

Frequencies use half-open intervals [f_low, f_high) except for the last band after conversion to bins, which naturally includes any remaining bins up to the requested maximum.

f_low: float[source]

f_high: float[source]

stride: int[source]

class FrequencyBandLayout[source]

A concrete multiband layout on an rFFT frequency grid.

sample_rate: float[source]

duration: float[source]

bands: tuple[FrequencyBand, Ellipsis][source]

property n_time: int[source]

Return type:: int

property n_freq: int[source]

Return type:: int

property df: float[source]

Return type:: float

property nyquist: float[source]

Return type:: float

property compressed_length: int[source]

Return type:: int

band_indices(device=None)[source]

Return one integer index tensor per band.

Each index is the start of a complete stride-sized window. Partial trailing windows (where fewer than stride native bins remain before the band boundary) are excluded, matching DINGO’s floor-division decimation exactly.

Parameters:: device (torch.device | None)
Return type:: tuple[torch.Tensor, Ellipsis]

frequencies(device=None)[source]

Return the retained frequencies in Hz.

Parameters:: device (torch.device | None)
Return type:: torch.Tensor

frequency_to_bin(frequency, *, round_up)[source]

Parameters:

frequency (float)
round_up (bool)

Return type:

int

validate_for(fd)[source]

Parameters:: fd (torch.Tensor)
Return type:: None

class FrequencyMultibandCompressor(layout, *, pool='sample')[source]

Bases: torch.nn.Module

Compress rFFT data with a variable frequency grid.

Input tensors must have shape (..., n_freq) and may be any dtype (complex or real). The last axis is treated as the frequency axis; all leading axes are passed through unchanged.

pool="sample" keeps one representative bin per stride block via a single index_select call — zero Python overhead in the hot path. pool="mean" averages each stride-sized block using unfold(...).mean(-1) — two GPU ops per band, no Python bin loop.

The compressor is format-agnostic: complex (B, D, N) tensors and real-packed (B, D, 2, N) tensors both work because all operations act on the last axis only.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

layout (FrequencyBandLayout)
pool (PoolMode)

GRAPH_READY = True[source]

pool = 'sample'[source]

classmethod from_bands(sample_rate, duration, bands, *, pool='sample')[source]

Parameters:

sample_rate (float)
duration (float)
bands (Sequence[FrequencyBand | tuple[float, float, int]])
pool (PoolMode)

Return type:

FrequencyMultibandCompressor

forward(fd)[source]

Compress fd along its last axis.

Parameters:: fd (Tensor, shape (..., N_freq))
Return type:: Tensor, shape (..., N_compressed)

retained_frequencies()[source]

Retained frequency grid on the same device as the module.

Return type:: torch.Tensor

make_dyadic_frequency_bands(*, f_min, f_max, base_stride=1, max_stride=64, first_width=32.0)[source]

Make bands whose stride doubles with frequency.

This is intentionally conservative near f_min and progressively cheaper at high frequency. Pass explicit bands to FrequencyMultibandCompressor when you want a layout matched to a PSD, waveform duration, or detector band.

Parameters:

f_min (float)
f_max (float)
base_stride (int)
max_stride (int)
first_width (float)

Return type:

tuple[FrequencyBand, Ellipsis]

describe_layout(layout)[source]

Return a compact, notebook-friendly description of a layout.

Parameters:: layout (FrequencyBandLayout)
Return type:: list[dict[str, float | int]]

make_prior_informed_frequency_bands(*, m1, m2, s1z=0.0, s2z=0.0, f_min, f_max, duration, n_bins_per_period=32, max_stride=128)[source]

Method 1 — analytic chirp-time boundary placement.

Build a multibanding layout by evaluating the IMRPhenomD chirp time τ(f) = SimIMRPhenomDChirpTime analytically for a single worst-case binary. Band boundaries are the smallest rFFT bin frequencies at which τ(f) drops below the resolution threshold for the next stride level. Boundaries are evaluated exactly on the discrete rFFT grid via binary search — no interpolation is performed.

See also make_empirical_frequency_bands() (Method 2) which instead measures τ(f) empirically from a simulated sample of waveforms.

The resolution criterion at bin k (f_k = k / duration) is:

stride(f_k) <= duration / (n_bins_per_period * tau(f_k))

where tau(f_k) = SimIMRPhenomDChirpTime(f_k) for (m1, m2, s1z, s2z).

Parameters:

m1 (float) – Component masses [M_sun]. Pass the lightest binary in the prior (worst case: longest inspiral, finest required resolution everywhere).
m2 (float) – Component masses [M_sun]. Pass the lightest binary in the prior (worst case: longest inspiral, finest required resolution everywhere).
s1z (float) – Aligned-spin parameters. Higher aligned spin increases tau(f) at every frequency, so pass the prior maximum spin (e.g. 0.99) for the most conservative (finest) layout. This is the same convention as TD multirate sampling — both use lightest mass + highest spin as the worst case.
s2z (float) – Aligned-spin parameters. Higher aligned spin increases tau(f) at every frequency, so pass the prior maximum spin (e.g. 0.99) for the most conservative (finest) layout. This is the same convention as TD multirate sampling — both use lightest mass + highest spin as the worst case.
f_min (float) – Lower frequency edge [Hz].
f_max (float) – Upper frequency edge [Hz] (typically the Nyquist frequency).
duration (float) – Segment duration [s]; sets native resolution Δf = 1 / duration.
n_bins_per_period (int) – Minimum multibanded bins per GW signal oscillation period (DINGO: 32).
max_stride (int) – Maximum stride (power of 2). The final band runs to f_max at this stride regardless of the chirp-time criterion.

Returns:

Ready to pass to FrequencyMultibandCompressor. Use pool="mean" to replicate DINGO’s bin-averaging decimation.

Return type:

tuple[FrequencyBand, …]

make_empirical_frequency_bands(*, m_min, m_max, spin_max=0.99, f_min, f_max, duration, n_samples=1000, n_bins_per_period=32, max_stride=128, seed=42)[source]

Method 2 — simulation-based boundary placement.

Build a multibanding layout by simulating n_samples FD waveforms drawn from the prior, measuring the worst-case local phase gradient |dΨ/df| at each rFFT bin, and converting it to an empirical chirp-time envelope τ_max(f). Band boundaries are then placed by the same criterion as make_prior_informed_frequency_bands() (Method 1), but using τ_max(f) in place of the analytic SimIMRPhenomDChirpTime.

This follows the procedure described in the DINGO multibanding paper: simulate waveforms from the prior and demand that every oscillation period of each signal is covered by at least n_bins_per_period bins in the multibanded domain. The key difference from Method 1 is that τ_max(f) is measured directly from waveform phase rather than from the SPA formula, making Method 2 valid for any approximant including precessing or HOM models where the analytic chirp time is inaccurate.

Waveforms are generated with IMRPhenomD (aligned spin) with t_c = 0 so that the FD phase Ψ(f) is the pure inspiral chirp phase with no linear time-shift term; the phase gradient then directly yields τ(f) without a heterodyne step.

Parameters:

m_min (float) – Component mass range [M_sun]. m1 is drawn from U[m_min, m_max] and m2 from U[m_min, m1] (mass ordering enforced). The analytic worst case (m1 = m2 = m_min, s1z = s2z = spin_max) is always included as the first sample.
m_max (float) – Component mass range [M_sun]. m1 is drawn from U[m_min, m_max] and m2 from U[m_min, m1] (mass ordering enforced). The analytic worst case (m1 = m2 = m_min, s1z = s2z = spin_max) is always included as the first sample.
spin_max (float) – Maximum aligned spin magnitude. s1z, s2z drawn from U[-spin_max, spin_max]. Higher spin increases τ(f); use the prior maximum for the most conservative layout.
f_min (float) – Frequency range [Hz].
f_max (float) – Frequency range [Hz].
duration (float) – Segment duration [s]; native resolution Δf = 1 / duration.
n_samples (int) – Number of IMRPhenomD waveforms to simulate (DINGO paper: 1000). Includes the explicit worst-case binary as sample 0.
n_bins_per_period (int) – Minimum multibanded bins per GW period (DINGO: 32).
max_stride (int) – Maximum stride (power of 2).
seed (int) – RNG seed for reproducibility.

Returns:

Ready to pass to FrequencyMultibandCompressor with pool="mean".

Return type:

tuple[FrequencyBand, …]