sage.data.waveform.snr

Filename : snr.py Description : Short description of the file

Created on 2026-02-16 11:14:49

__author__ = Narenraju Nagarajan __copyright__ = Copyright 2026, ProjectName __license__ = MIT Licence __version__ = 0.0.1 __maintainer__ = Narenraju Nagarajan __affiliation__ = N/A __email__ = N/A __status__ = [‘inProgress’, ‘Archived’, ‘inUsage’, ‘Debugging’]

GitHub Repository: NULL

Documentation:

snr = OptimalSNREstimator(

psds=fiducial_psds, # (D, F) delta_f=delta_f, f_low=20.0, f_high=1024.0, device=”cuda”

)

rho_net, rho_det = snr(h_batch)

Classes

OptimalSNREstimator	Fast batched optimal matched-filter SNR estimator (equivalent to PyCBC sigmasq).
OptimalSNRRescaler	Rescales a batch of signals to match target SNRs.

Module Contents

class OptimalSNREstimator

Bases: torch.nn.Module

Fast batched optimal matched-filter SNR estimator (equivalent to PyCBC sigmasq).

Computes the optimal (whitened) SNR for a batch of frequency-domain detector-projected waveforms using fiducial PSDs loaded from disk. The integration is performed as:

\[\rho^2 = 4 \Delta f \sum_f \frac{|h(f)|^2}{S_n(f)}\]

for each detector, summed over detectors for the network SNR.

asds

Amplitude spectral densities (sqrt of PSDs) per detector.

Type:

torch.Tensor, shape (1, D, F)

mask

Pre-computed frequency mask for [f_low, f_high] integration band.

Type:

torch.Tensor or None, shape (1, 1, F)

delta_f

Frequency bin spacing in Hz.

Type:

float

Expected input shapes

---------------------

h

Type:

(B, D, F) complex tensor — detector-projected FD waveforms.

device

delta_f

asds

mask = None

forward(h)

Batched optimal SNR for multi-detector frequency-domain waveforms.

Parameters:

• h (complex tensor (B, D, F)) – Detector-projected frequency-domain strain

• psd (real tensor (D, F)) – One-sided PSD for each detector

• delta_f (float) – Frequency spacing

• mask (optional bool tensor (F,)) – Frequency mask for f_low / f_high cutoffs

Returns:

• rho_net ((B, 1)) – Network optimal SNR

• rho_det ((B, D, 1)) – Per-detector optimal SNR

class OptimalSNRRescaler(target_snr_sampler)

Bases: torch.nn.Module

Rescales a batch of signals to match target SNRs.

Parameters:

• snr_estimator – instance of OptimalSNREstimator

• target_snr_sampler (Callable[[int], torch.Tensor]) – callable(batch_size) -> Tensor of target SNRs

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

snr_estimator

target_snr_sampler

forward(signal_batch)

Rescale signals to target SNR.

Parameters:

signal_batch (torch.Tensor) – shape [B, L] or [B, C, L]

Returns:

same shape as input, shape (B, …) scale: (B,) float tensor — per-sample amplitude scale factors

(hf_new = hf_old * scale, so distance_new = distance_old / scale)

Return type:

rescaled_signal_batch