Rebasing bns model commits from bns-o3-mdc branch onto dev branch

libs/architectures/architectures/supervised.py

@@ -302,13 +302,17 @@ def __init__(
         zero_init_residual: bool = False,
         groups: int = 1,
         width_per_group: int = 64,
+        top_k: int | None = None,
         stride_type: list[Literal["stride", "dilation"]] | None = None,
         norm_layer: NormLayer | None = 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

@@ -1,5 +1,6 @@
 import pytest
 import torch
+import torch.nn.functional as F
 from ml4gw.transforms import Heterodyne
 from utils.augmentation import HeterodyneAugmentor
 from utils.preprocessing import (
@@ -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

@@ -2,10 +2,76 @@
 from typing import Literal
 
 import torch
+import torch.nn.functional as F
 from ml4gw.transforms import Heterodyne
 from torch import Tensor
 
 
+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

@@ -78,3 +78,13 @@ output_buffer_length: 8
 device: "cuda"
 verbose: true
 matmul_precision: "high"
+# optional augmentor arguments
+# augmentor:
+#   class_path: utils.augmentation.HeterodyneAugmentor
+#   init_args:
+#     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

projects/online/online/cli.py

@@ -23,6 +23,27 @@ def build_parser():
         apply_on="parse",
     )
 
+    # TODO: This is a workaround for linking sample_rate and
+    # kernel_length argument between the augmentor and
+    # the global parameters in the config.
+    try:
+        parser.link_arguments(
+            "sample_rate",
+            "augmentor.init_args.sample_rate",
+            apply_on="parse",
+        )
+    except Exception:
+        pass
+
+    try:
+        parser.link_arguments(
+            "kernel_length",
+            "augmentor.init_args.kernel_length",
+            apply_on="parse",
+        )
+    except Exception:
+        pass
+
     return parser
 
 

projects/online/online/main.py

@@ -386,6 +386,7 @@ def main(
     integration_window_length: float,
     astro_event_rate: float,
     data_source: Literal["frames", "arrakis"] = "frames",
+    augmentor: torch.nn.Module | None = None,
     state_channels: list[str] | None = None,
     fftlength: float | None = None,
     highpass: float | None = None,
@@ -450,6 +451,10 @@ def main(
         offline_inference_rate:
             Rate at which inference was performed offline when
             establishing the background and foreground distributions
+        augmentor:
+            Optional augmentor module that performs augmentation
+            of the input data that need to be analyzed by Aframe.
+            If not provided, no augmentation will be performed.
         psd_length:
             Length of PSD estimation window in seconds for PSD
             used to whiten aframe data
@@ -842,6 +847,7 @@ def main(
         batch_size=int(update_size * online_inference_rate),
         fduration=fduration,
         fftlength=fftlength,
+        augmentor=augmentor,
         highpass=highpass,
         lowpass=lowpass,
     ).to(device)