ml4gw/ml4gw

1. What this repo is

- ml4gw is a PyTorch-based scientific library for gravitational-wave ML workflows, not an application server or standalone pipeline.

- Its architecture is organized around reusable building blocks for signal generation, preprocessing, dataset loading, and neural-network modules.

2. Core data flow

- Typical flow is: raw or simulated time-series data → signal-conditioning transforms → optional augmentation/rescaling → dataset wrappers → neural network models.

- A second major flow is waveform synthesis: astrophysical parameters → waveform generators → injected or standalone strain-like tensors used in training/evaluation.

3. Time-series transforms subsystem

- ml4gw/transforms is the main preprocessing layer.

- It contains operations for decimation, heterodyning, IIR filtering, integration, Pearson correlation, spectral estimation, spectrograms, Q-transforms, spline interpolation, scaling, SNR rescaling, and whitening.

- These modules are the main boundary between raw detector-like signals and model-ready tensors.

4. Waveform generation subsystem

- ml4gw/waveforms generates synthetic gravitational-wave signals.

- It splits into cbc (compact binary coalescence models such as TaylorF2, PhenomD, PhenomP) and adhoc (e.g. sine-Gaussian and ringdown-style signals).

- generator.py and conversion.py suggest a parameterization layer that converts physical parameters into waveform tensors.

5. Data loading subsystem

- ml4gw/dataloading provides training data adapters for different storage/streaming patterns.

- in_memory_dataset.py, hdf5_dataset.py, and chunked_dataset.py imply support for small local datasets, large on-disk datasets, and chunked access for scalable training.

- This is the main I/O boundary between persisted data and the model-training stack.

6. Neural network components

- ml4gw/nn contains reusable model components rather than a single end-to-end model.

- Subpackages include resnet for 1D/2D residual architectures, autoencoder for convolutional autoencoders, ssm/s4d.py for state-space sequence modeling, streaming for online inference helpers, and norm.py for normalization layers.

- These modules are designed to consume preprocessed GW-like tensors produced by the transforms and datasets.

7. Utilities and physics helpers

- ml4gw/utils contains low-level helpers such as interferometer-related utilities and slicing.

- Top-level modules like gw.py, spectral.py, augmentations.py, distributions.py, constants.py, and types.py provide shared physics/math helpers, tensor utilities, and type definitions used across the package.

8. Streaming / online inference boundary

- The nn/streaming package introduces explicit online-processing support through components like snapshotter and online_average.

- This indicates the library is not limited to offline training; it also supports sliding-window or incremental inference over time series.

9. Testing strategy

- tests/ mirrors the package layout almost one-to-one, which makes the architecture easy to map and suggests strong module-level isolation.

- There are dedicated test groups for dataloading, transforms, waveform synthesis, utilities, and each NN subsystem.

- benchmarks/ mirrors major functional areas as well, indicating the project tracks performance for compute-heavy signal-processing and model components.

10. Documentation and examples

- docs/ is a first-class part of the repo, with tutorials, installation docs, and example pages for augmentation, distributions, GW utilities, and transform families.

- The notebooks under docs/tutorials/ act as higher-level integration examples that likely demonstrate the intended end-to-end flow from signals to models.

11. Packaging and release

- The project is packaged as a Python library via pyproject.toml and pinned with uv.lock.

- GitHub Actions handles documentation builds, pre-commit checks, unit tests, and publishing, so the repo’s operational boundary is standard library CI/CD rather than runtime infrastructure.

12. Architecture summary for a diagram

- The cleanest diagram would show three main capability layers: signal generation, signal preprocessing/data access, and ML models/streaming inference.

- Shared utilities and spectral/physics helpers sit beneath those layers, while tests/docs/benchmarks wrap the whole codebase as validation and usability scaffolding.