Add top-k support to heterodyne functionality

libs/architectures/architectures/supervised.py

@@ -301,13 +301,17 @@ def __init__(
         zero_init_residual: bool = False,
         groups: int = 1,
         width_per_group: int = 64,
+        top_k: int | None = None,
         stride_type: Optional[list[Literal["stride", "dilation"]]] = None,
         norm_layer: Optional[NormLayer] = None,
         **kwargs,
     ) -> None:
         super().__init__()
+
+        num_channels = top_k if top_k is not None else num_chirp_masses
+
         self.time_domain_resnet = ResNet1D(
-            in_channels=num_ifos * num_chirp_masses,
+            in_channels=num_ifos * num_channels,
             layers=layers,
             classes=1,
             kernel_size=kernel_size,

libs/utils/tests/test_preprocessing.py

@@ -8,6 +8,7 @@
 from utils.augmentation import HeterodyneAugmentor
 from ml4gw.transforms import Heterodyne
 import torch
+import torch.nn.functional as F
 import pytest
 
 
@@ -302,6 +303,84 @@ def test_with_heterodyne_augmentor(self):
         kernels_aug = whitener_aug(x)
         assert torch.allclose(kernels_aug, kernels)
 
+    def test_with_heterodyne_augmentor_top_k(self):
+
+        heterodyne_augmentor = HeterodyneAugmentor(
+            sample_rate=self.sample_rate,
+            kernel_length=self.kernel_length,
+            chirp_mass_low=1.0,
+            chirp_mass_high=2.5,
+            num_chirp_masses=10,
+            chirp_mass_spacing="log",
+            keep_last_n_seconds=1.5,
+            top_k=5,
+        )
+
+        whitener = BatchWhitener(
+            kernel_length=self.kernel_length,
+            sample_rate=self.sample_rate,
+            inference_sampling_rate=self.inference_sampling_rate,
+            batch_size=self.batch_size,
+            fduration=self.fduration,
+            fftlength=self.fftlength,
+        )
+
+        whitener_aug = BatchWhitener(
+            kernel_length=self.kernel_length,
+            sample_rate=self.sample_rate,
+            inference_sampling_rate=self.inference_sampling_rate,
+            batch_size=self.batch_size,
+            fduration=self.fduration,
+            fftlength=self.fftlength,
+            augmentor=heterodyne_augmentor,
+        )
+
+        channels = 2
+        total_samples = int(
+            (
+                (self.batch_size - 1) * whitener_aug.stride_size
+                + whitener_aug.kernel_size
+            )
+            * 2
+        )
+
+        x = torch.randn(channels, total_samples)
+
+        heterodyne = Heterodyne(
+            sample_rate=self.sample_rate,
+            kernel_length=self.kernel_length,
+            chirp_mass=heterodyne_augmentor.chirp_mass_grid,
+            return_type="time",
+        )
+
+        kernels = whitener_aug(x)
+
+        kernels_heterodyned = heterodyne(whitener(x))
+        _B, _C, _M, _T = kernels_heterodyned.shape
+        avgpool = F.avg_pool1d(
+            torch.abs(kernels_heterodyned.reshape(_B, _C * _M, _T)),
+            31,
+            stride=1,
+            padding=15,
+        ).reshape(_B, _C, _M, _T)
+        avgpool_snr = torch.sqrt(
+            (avgpool[..., -int(1.5 * self.sample_rate) :] ** 2).sum(dim=1)
+        )
+        vals = torch.max(avgpool_snr, dim=-1).values
+        idx = torch.topk(vals, k=5, dim=-1).indices
+        idx_expand = idx[:, None, :, None].expand(-1, _C, -1, _T)
+        kernels_heterodyned = torch.gather(
+            kernels_heterodyned, dim=2, index=idx_expand
+        )
+        kernels_heterodyned = kernels_heterodyned.reshape(_B, _C * 5, _T)
+        kernels_heterodyned = kernels_heterodyned[
+            ..., -int(1.5 * self.sample_rate) :
+        ]
+        kernels = kernels_heterodyned
+
+        kernels_aug = whitener_aug(x)
+        assert torch.allclose(kernels_aug, kernels)
+
 
 class TestMultiModalPreprocessor:
     """Test suite for MultiModalPreprocessor module."""

libs/utils/utils/augmentation.py

@@ -1,11 +1,77 @@
 import math
 import torch
+import torch.nn.functional as F
 from torch import Tensor
 from typing import Literal
 
 from ml4gw.transforms import Heterodyne
 
 
+def select_top_k(
+    X: torch.Tensor,
+    top_k: int,
+    kernel_size: int = 31,
+    stride: int = 1,
+    padding: int = 15,
+    keep_last_n_samples: int | None = None,
+) -> torch.Tensor:
+    """
+    Select top `k` chirp mass channels from `M` heterodyned timeseries.
+
+    Args:
+        X (Tensor):
+            Input tensor of shape (B, C, M, T) where B is the batch
+            size, C is the number of channels, M is the number of chirp
+            mass channels, and T is the number of time samples.
+        top_k (int):
+            Number of chirp mass channels to retain.
+        kernel_size (int):
+            Size of the running average (average pooling) window used to
+            smooth the absolute value of the input before computing the
+            selection statistic.
+        stride (int):
+            Stride of the average pooling operation.
+        padding (int):
+            Zero-padding applied to both sides of the timeseries before
+            the average pooling operation.
+        keep_last_n_samples (int, optional):
+            If provided, only the final `n` samples are used when
+            computing the statistic used to select the top k chirp mass
+            channels. If `None`, all samples are used.
+
+    Returns:
+        Tensor:
+            Output tensor of shape (B, C * top_k, T) containing the
+            selected top k chirp mass channels for each batch and channel.
+    """
+
+    B, C, M, T = X.shape
+
+    avgpool = torch.stack(
+        [
+            F.avg_pool1d(
+                torch.abs(x.reshape(C * M, T)),
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+            )
+            for x in X
+        ]
+    ).reshape(B, C, M, T)
+
+    if keep_last_n_samples is not None:
+        avgpool_snr = torch.sqrt(
+            (avgpool[..., -keep_last_n_samples:] ** 2).sum(dim=1)
+        )
+    else:
+        avgpool_snr = torch.sqrt((avgpool**2).sum(dim=1))
+
+    vals = avgpool_snr.max(dim=-1).values
+    idx = torch.topk(vals, k=top_k, dim=-1).indices
+    idx = idx[:, None, :, None].expand(B, C, top_k, T)
+    return torch.gather(X, dim=2, index=idx).reshape(B, C * top_k, T)
+
+
 class HeterodyneAugmentor(torch.nn.Module):
     """
     Apply a heterodyne transform over a grid of chirp masses to a batch
@@ -26,6 +92,10 @@ class HeterodyneAugmentor(torch.nn.Module):
         keep_last_n_seconds (float):
             If provided, only the last `n` seconds of the kernel_length are
             returned. Otherwise, the full kernel_length is returned.
+        top_k (int):
+            If provided, the top `k` chirp mass channels are selected based on
+            their absolute amplitude, corresponding to the `k` matching chirp
+            masses with. If `None`, all chirp mass channels are returned.
     Shape:
         Input: (batch_size, channels, time)
         Output: (batch_size, channels * num_chirp_masses, time_out)
@@ -58,12 +128,13 @@ def __init__(
         chirp_mass_high: float = 2.5,
         num_chirp_masses: int = 100,
         chirp_mass_spacing: Literal["linear", "log"] = "log",
-        keep_last_n_seconds: float = None,
+        keep_last_n_seconds: float | None = None,
+        top_k: int | None = None,
     ):
         super().__init__()
         self.sample_rate = sample_rate
         self.kernel_length = kernel_length
-        self.keep_last_n_seconds = keep_last_n_seconds
+        self.top_k = top_k
         self.num_chirp_masses = num_chirp_masses
         self.keep_last_n_seconds = keep_last_n_seconds
 
@@ -78,6 +149,8 @@ def __init__(
             self.keep_last_n_samples = int(
                 self.keep_last_n_seconds * sample_rate
             )
+        else:
+            self.keep_last_n_samples = None
 
         self.heterodyne_transform = Heterodyne(
             sample_rate=sample_rate,
@@ -120,12 +193,22 @@ def forward(self, x: Tensor) -> Tensor:
                 or determined by `keep_last_n_seconds`.
         """
         _B, _C, _T = x.shape
-        x_heterodyned = torch.empty((_B, _C * self.num_chirp_masses, _T))
+        if self.top_k is not None:
+            x_heterodyned = torch.empty((_B, _C * self.top_k, _T))
+        else:
+            x_heterodyned = torch.empty((_B, _C * self.num_chirp_masses, _T))
         # Heterodyne the whitened timeseries
         x = self.heterodyne_transform(x)
-        # Reshaping x from (batch_size, channels, num_chirp_mass, kernel_size)
-        # to (batch_size, channels x num_chirp_mass, kernel_size)
-        x = x.reshape(_B, _C * self.num_chirp_masses, _T)
+        if self.top_k is not None:
+            # Select the top k chirp mass channels
+            x = select_top_k(
+                x, self.top_k, keep_last_n_samples=self.keep_last_n_samples
+            )
+        else:
+            # Reshaping x from
+            # (batch_size, channels, num_chirp_mass, kernel_size)
+            # to (batch_size, channels x num_chirp_mass, kernel_size)
+            x = x.reshape(_B, _C * self.num_chirp_masses, _T)
         x_heterodyned[:, :, :] = x
         # Returning the desired length of heterodyned strain in the
         # time dimension

projects/export/export_heterodyne.yaml

@@ -9,7 +9,7 @@ weights:
 batch_file:
 repository_directory:
 num_ifos: 2
-kernel_length: 4.0
+kernel_length: 8.0
 inference_sampling_rate: 4
 sample_rate: 2048
 batch_size: 128
@@ -18,7 +18,7 @@ psd_length: 64
 preprocessor:
   class_path: utils.preprocessing.BatchWhitener
   init_args:
-    kernel_length: 4.0
+    kernel_length: 8.0
     sample_rate: 2048
     inference_sampling_rate: 4
     batch_size: 128
@@ -28,12 +28,13 @@ preprocessor:
       class_path: utils.augmentation.HeterodyneAugmentor
       init_args:
         sample_rate: 2048
-        kernel_length: 4.0
+        kernel_length: 8.0
         chirp_mass_low: 1.0
         chirp_mass_high: 2.5 
         num_chirp_masses: 100
         chirp_mass_spacing: "log" 
         keep_last_n_seconds: null
+        top_k: 10
     highpass: 32.0
 streams_per_gpu: 12
 # num_outputs:

projects/online/config.yaml

@@ -52,13 +52,15 @@ state_channels: ["H1:GDS-CALIB_STATE_VECTOR", "L1:GDS-CALIB_STATE_VECTOR", "V1:D
 data_source: "frames"
 sample_rate: 2048
 astro_event_rate: 89
-kernel_length: 20
+kernel_length: 8
 online_inference_rate: 512
 offline_inference_rate: 16
-schedule: [[0, 16, 512], [16, 20, 2048]]
-split: True
-q: 45.6
-spectrogram: [64, 128]
+chirp_mass_low: 1.0
+chirp_mass_high: 2.5
+num_chirp_masses: 100
+chirp_mass_spacing: "log"
+keep_last_n_seconds: 1.5
+top_k: 10
 inference_params: ["chirp_mass", "mass_ratio", "luminosity_distance", "phic", "inclination", "dec", "psi", "phi"]
 psd_length: 64
 amplfi_psd_length: 10