AresGW Documentation

AresGW Documentation#

This page documents the actual code structure used in the ACME AresGW development repository at niksterg/acme-aresgw-dev. It complements the existing notebooks in this book: the notebooks remain the interactive tutorial layer, while this page explains how the Python codebase itself is organized.

Unlike GWG, AresGW is not structured as a small importable package with a single public module. It is organized as a script-driven training and evaluation pipeline built around PyTorch models, HDF5 time-series data, and utility modules for slicing, whitening, waveform generation, inference, and sensitivity evaluation.

Repository Layout#

The top-level repository structure is:

acme-aresgw-dev/
  train.py
  test.py
  evaluate.py
  generate_waveforms_npy.py
  sensitivity_plot.py
  run_on_test.sh
  modules/
  utils/
  generatedata/
  environments/
  evaluation-real/
  doc/

Top-Level Entry Scripts#

  • train.py: trains the AresGW classifier.

  • test.py: runs a trained model on HDF5 detector data and writes clustered triggers.

  • evaluate.py: compares foreground and background triggers against injections and computes false-alarm and sensitivity metrics.

  • generate_waveforms_npy.py: precomputes waveform arrays used for training and validation augmentation.

  • sensitivity_plot.py: plots the evaluation output.

Supporting Directories#

  • modules/: neural-network components and preprocessing layers.

  • utils/: dataset construction, waveform generation, trigger extraction, clustering, and training helpers.

  • generatedata/: scripts for creating the HDF5 training, validation, and test files.

  • environments/: conda environment definitions.

  • evaluation-real/: utilities and example material for real LVK data evaluation.

  • doc/: slides and training-session material.

Pipeline Overview#

The codebase is designed around a staged workflow:

real noise + injection settings
  -> generatedata/generate_data.py and generate_all_sets.py
  -> background_*.hdf, foreground_*.hdf, injections_*.hdf
  -> generate_waveforms_npy.py
  -> injections_train.npy and injections_val.npy
  -> train.py
  -> weights.pt and training_stats.json
  -> test.py
  -> foreground/background trigger files
  -> evaluate.py
  -> sensitivity metrics and evaluation HDF5 output

The notebooks in this Jupyter Book consume the same artifacts at different stages of that pipeline.

Main Entry Scripts#

train.py#

train.py is the central training entry point. It assembles the full model, constructs the datasets, configures the optimizer and scheduler, and runs the training and validation loops.

Its main steps are:

  1. create the output directory and weight path,

  2. build the model as CropWhitenNet(ResNet54Double(), DAIN_Layer(...)),

  3. load validation data from background_val.hdf and injections_val.npy,

  4. load training data from background_train.hdf, injections_train.hdf, and injections_train.npy,

  5. decode the epoch-by-epoch SNR curriculum,

  6. train with reg_BCELoss, Adam, MultiStepLR, and optional warmup,

  7. save weights.pt, training_curves.png, and training_stats.json.

Two details matter for understanding the implementation:

  • positive training samples are created on the fly by injecting precomputed waveforms into background segments,

  • the SNR range of allowed injections is updated during training via the --snr-schedule argument.

test.py#

test.py is the inference entry point used after training. It loads the same model architecture as training, restores saved weights, scans an input HDF5 file in sliding windows, extracts triggers above a threshold, clusters them in time, and writes the result to an output HDF5 file.

The script fixes several deployment-time parameters internally:

  • step_size = 5.1

  • slice_dur = 6.25

  • trigger_threshold = 0.5

  • cluster_threshold = 0.35

  • var = 0.5

The output file contains three datasets:

  • time

  • stat

  • var

These are the trigger products later consumed by evaluate.py and the tutorial notebooks.

evaluate.py#

evaluate.py performs sensitivity evaluation by comparing foreground detections, background detections, and the known injection catalogue.

The script:

  1. identifies which injections are actually contained in the analyzed foreground files,

  2. loads injections and trigger files,

  3. separates true positives from false positives using the trigger time window var,

  4. computes foreground and background false-alarm rates,

  5. computes sensitive volume and sensitive distance,

  6. writes the evaluation products to an HDF5 output file.

The code supports a chirp-mass weighted sensitivity mode when the injection file contains the required information.

generate_waveforms_npy.py#

This script bridges generated injection metadata and training-time augmentation. It loops over the train and val sets, then calls python -m utils.waveform_gen for each one.

Its job is intentionally small: it is an orchestration layer that creates:

  • injections_train.npy

  • injections_val.npy

Those arrays are then memory-mapped by the dataset classes during training and validation.

Data Generation Code#

generatedata/generate_data.py#

This script is the data-generation entry point used by the tutorials. The repository README explicitly notes that it comes from MLGWSC-1. In practice it is the producer of the HDF5 files consumed by the rest of the pipeline.

The important output products are:

  • background_*.hdf

  • foreground_*.hdf

  • injections_*.hdf

generatedata/generate_all_sets.py#

generate_all_sets.py is a convenience wrapper around generate_data.py. It creates train, validation, and test splits by invoking the generator three times with different:

  • output filenames,

  • offsets into the noise file,

  • seeds.

The current local version uses:

  • train: offset 0, seed seed

  • val: offset duration, seed seed + 1

  • test: offset 2 * duration, seed seed + 2

This is one of the places where documenting the local repository matters more than documenting the original upstream AresGW repository.

generatedata/subset_noise.py#

This utility creates a shorter HDF5 subset from a larger segmented noise file. It walks detector groups and segment objects, computes temporal overlap with a requested window, and writes only the intersecting slices to a new file.

Its practical role is to make the ACME tutorials feasible without requiring the full multi-day noise archive.

Model Components#

modules/resnet.py#

This file defines the 1D convolutional classifier backbone.

  • ResBlock: residual 1D building block.

  • ResNet54: narrower classifier variant.

  • ResNet54Double: wider variant used by the current training and testing scripts.

ResNet54Double is the effective classifier used in the ACME workflow.

modules/whiten.py#

This file defines the whitening and deployment wrapper logic.

  • Whiten: estimates PSDs and whitens detector time series in PyTorch.

  • CropWhitenNet: wraps whitening, cropping, optional DAIN normalization, and classifier inference.

CropWhitenNet is the central runtime wrapper around the network. During training it crops around known injection times; during testing it unfolds the input into many overlapping 1.25 s segments and returns both scores and the associated time offsets.

modules/dain.py#

This module provides the DAIN_Layer, an adaptive normalization layer applied after whitening. It performs learned centering, scaling, and gating over feature vectors.

modules/loss.py#

This module defines reg_BCELoss, a regularized binary cross-entropy loss that slightly shifts the network outputs away from exact 0 and 1 values before computing BCE.

Dataset and Evaluation Utilities#

utils/dataset.py#

This file contains the training-time data pipeline.

Slicer#

Slicer reads detector segments from HDF5 files and exposes sliding windows over the raw data. It supports multiprocessing and prefetching.

SlicerDataset#

SlicerDataset pairs background noise segments with precomputed waveforms from .npy files. It uses each background slice once as a negative sample and once as a positive sample after injecting a waveform.

SlicerDatasetSNR#

The training script imports SlicerDatasetSNR, which extends the slicing logic with SNR-based filtering using the injection metadata. This is what lets the training code implement the epoch-dependent --snr-schedule curriculum.

The key design choice in this file is that waveform augmentation is not stored as fixed foreground training tensors. Instead, clean waveforms are precomputed once and injected into fresh noise windows dynamically during training.

utils/eval_utils.py#

This file contains the evaluation-time support code used by test.py.

Important functions:

  • TorchSlicer: dataset-style slicer for evaluation.

  • get_triggers(...): runs the network over sliding windows and collects threshold crossings.

  • get_clusters(...): groups neighboring triggers and centers each cluster at its maximum statistic.

  • find_injection_times(...): identifies which injections lie inside the analyzed foreground files.

  • mchirp(...): chirp-mass helper used in sensitivity calculations.

utils/waveform_gen.py#

This module converts injection metadata into time-domain waveforms. It uses PyCBC waveform generation and detector projection, pads and crops the strain around coalescence, and saves a (n_injections, 2, 2560) NumPy array.

This is the code path behind generate_waveforms_npy.py.

utils/train_utils.py#

This file contains small but important training helpers:

  • WarmUpLR: linear learning-rate warmup scheduler.

  • initialize_xavier(...): initialization helper for conv, batch norm, and linear layers.

  • progress_bar(...): terminal progress display.

Data Products#

The local codebase revolves around a small set of recurring artifacts.

Generated HDF5 Files#

  • background_train.hdf, background_val.hdf, background_test.hdf

  • foreground_train.hdf, foreground_val.hdf, foreground_test.hdf

  • injections_train.hdf, injections_val.hdf, injections_test.hdf

Generated NumPy Files#

  • injections_train.npy

  • injections_val.npy

These contain waveform tensors used for training-time injection.

Training Outputs#

  • weights.pt

  • training_curves.png

  • training_stats.json

Inference and Evaluation Outputs#

  • trigger files from test.py

  • evaluation summaries from evaluate.py

  • sensitivity plots from sensitivity_plot.py

Short API Reference#

Entry point or component

Signature

Key arguments

train.py

train.py --data-dir DIR --output-dir DIR [options]

--data-dir: location of generated HDF5 and .npy files; --train-device: cpu or cuda:*; --slice-dur, --slice-stride: window geometry; --snr-schedule: curriculum; --p-augment: cross-noise augmentation; --resume-from: restore weights

test.py

test.py [--weights PATH] INPUTFILE OUTPUTFILE [--test-device DEVICE]

--weights: model checkpoint; INPUTFILE: foreground/background HDF5 file; OUTPUTFILE: trigger HDF5 file; --test-device: deployment device

evaluate.py

evaluate.py --injection-file FILE --foreground-files FILE... --foreground-events FILE... --background-events FILE... --output-file FILE [--verbose] [--force]

injection catalogue, analyzed foreground files, foreground trigger files, background trigger files, output HDF5 path

generate_waveforms_npy.py

generate_waveforms_npy.py --hdf-filename FILE --data-dir DIR

--hdf-filename: background HDF5 reference file; --data-dir: directory containing injections_train.hdf and injections_val.hdf

generatedata/generate_all_sets.py

generate_all_sets.py --output-dir DIR --seed N --real-noise-path FILE --duration SEC [--verbose] [--force]

split output directory, initial seed, source noise file, duration per split

utils.waveform_gen

python -m utils.waveform_gen --hdf-filename FILE --injections-file FILE --output-npy FILE

source background file, injection metadata file, output waveform array

utils.eval_utils.get_triggers

get_triggers(Network, inputfile, step_size=0.1, trigger_threshold=0.2, device='cpu', verbose=False, dtype=torch.float32, batch_size=512, slicer_cls=TorchSlicer, num_workers=8, whiten=True, slice_length=2048)

trained network, input HDF5 file, step size, trigger threshold, device, whitening choice

utils.eval_utils.get_clusters

get_clusters(triggers, cluster_threshold=0.35, var=0.2)

trigger list, clustering window, time-tolerance value

modules.whiten.CropWhitenNet

CropWhitenNet(net=None, norm=None, deploy=False, m=0.625, l=0.5, f=15.)

classifier backbone, DAIN layer, deployment mode, whitening window and cutoff parameters

modules.resnet.ResNet54Double

ResNet54Double()

main classifier backbone used by the local training and testing scripts

Notable Design Choices In The Local Repo#

  • The current ACME development repo is script-first, not package-first. The CLI entry points are the primary interface.

  • Training uses dynamically injected waveforms rather than a fully materialized foreground training set.

  • Whitening is performed inside the network wrapper, not as a separate offline preprocessing stage.

  • Testing uses deployment logic in CropWhitenNet to unfold longer inputs into overlapping whitened inference windows.

  • The local generate_all_sets.py and subset_noise.py provide ACME-specific workflow improvements compared to the original repository lineage.

Relationship To The Tutorials#

The existing notebooks in this book remain the user-facing walkthroughs:

This documentation page is the code-reference companion to those notebooks.

References#

Contents