sage.dsp.whiten

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

Created on 2026-01-19 16:26:37

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

GitHub Repository: NULL

Documentation: NULL

Classes

FiducialWhitening

Whiten frequency-domain strain using fixed, detector-specific fiducial PSDs.

Module Contents

class FiducialWhitening(**kwargs)

Bases: torch.nn.Module

Whiten frequency-domain strain using fixed, detector-specific fiducial PSDs.

The whitening kernel is derived once from pre-computed fiducial ASDs and stored as a registered buffer so it moves to the correct device automatically and is included in torch.compile graphs.

Pipeline (per sample)

1. Multiply FD strain by the whitening kernel: X_white = X_fd * whitening where whitening[d, f] = 2 Δf / (√0.5 · ASD[d, f]).

2. Convert back to time domain via inverse real FFT.

3. Strip the corrupted edge samples introduced by the Welch PSD estimation window (padding_nsamples on each side).

The @torch.no_grad() decorator on forward() means this module severs the autograd graph. Adversarial perturbations or any gradient-based optimisation must therefore operate on the output of this module, not on its FD input.

param **kwargs:

Forwarded to nn.Module.__init__.

whitening

Per-detector, per-frequency whitening kernel (registered buffer).

Type:

torch.Tensor, shape (D, F)

corrupted_len

Number of samples removed from each end of the whitened time series.

Type:

int

Input / Output

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

forward(X_fd)

where T_valid = seq_len - 2 * corrupted_len.

Type:

(B, D, F) complex64 → (B, D, T_valid) float32

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

device

seq_len

sample_rate

corrupted_len

delta_f

remove_corrupted(x)

Strip edge samples corrupted by the Welch PSD estimation window.

Parameters:

x (torch.Tensor, shape (B, D, T)) – Whitened time-domain strain (full length, including corrupted ends).

Returns:

Valid central samples only.

Return type:

torch.Tensor, shape (B, D, T - 2 * corrupted_len)

forward(input)

Whiten frequency-domain strain.

Accepts either a raw tensor (legacy path) or a GWBatch (state-tracked path). The behaviour depends on the grid type:

• FD_UNIFORM — whiten → IFFT → strip corrupted edges → return GWBatch with TD_UNIFORM state (real float32, shape (B, D, T_valid)).

• FD_COARSE — whiten at the coarse frequency indices using batch.coarse_indices → return GWBatch with FD_COARSE whitened state (complex, shape (B, D, N_coarse)). No IFFT is applied — the non-uniform grid cannot be IFFTed.

• Raw tensor (no GWBatch) — treated as FD_UNIFORM and the raw whitened TD tensor is returned for backward compatibility.

Parameters:

input (torch.Tensor or GWBatch) – FD strain (B, D, F) complex, or a GWBatch wrapping it.

Returns:

GWBatch when input is a GWBatch; raw float32 tensor otherwise.

Return type:

GWBatch or torch.Tensor