ml4gw 0.2.0__tar.gz → 0.4.0__tar.gz

ml4gw-0.2.0/README.md → ml4gw-0.4.0/PKG-INFO

@@ -1,3 +1,22 @@
+Metadata-Version: 2.1
+Name: ml4gw
+Version: 0.4.0
+Summary: Tools for training torch models on gravitational wave data
+Author: Alec Gunny
+Author-email: alec.gunny@ligo.org
+Requires-Python: >=3.8,<3.12
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: torch (>=1.10,<2.0) ; python_version >= "3.8" and python_version < "3.11"
+Requires-Dist: torch (>=2.0,<3.0) ; python_version >= "3.11"
+Requires-Dist: torchaudio (>=0.13,<0.14) ; python_version >= "3.8" and python_version < "3.11"
+Requires-Dist: torchaudio (>=2.0,<3.0) ; python_version >= "3.11"
+Requires-Dist: torchtyping (>=0.1,<0.2)
+Description-Content-Type: text/markdown
+
 # ML4GW
 
 Torch utilities for training neural networks in gravitational wave physics applications.
@@ -21,8 +40,8 @@ pip install ml4gw torch==1.12.0 --extra-index-url=https://download.pytorch.org/w
 
 ```toml
 [tool.poetry.dependencies]
-python = "^3.8"  # python versions 3.8-3.10 are supported
-ml4gw = "^0.1.0"
+python = "^3.8"  # python versions 3.8-3.11 are supported
+ml4gw = "^0.3.0"
 ```
 
 To build against a specific PyTorch/CUDA combination, consult the PyTorch installation documentation above and specify the `extra-index-url` via the `tool.poetry.source` table in your `pyproject.toml`. For example, to build against CUDA 11.6, you would do something like:
@@ -30,7 +49,7 @@ To build against a specific PyTorch/CUDA combination, consult the PyTorch instal
 ```toml
 [tool.poetry.dependencies]
 python = "^3.8"
-ml4gw = "^0.1.0"
+ml4gw = "^0.3.0"
 torch = {version = "^1.12", source = "torch"}
 
 [[tool.poetry.source]]
@@ -40,6 +59,8 @@ secondary = true
 default = false
 ```
 
+Note: if you are building against CUDA 11.6 or 11.7, make sure that you are using python 3.8, 3.9, or 3.10. Python 3.11 is incompatible with `torchaudio` 0.13, and the following `torchaudio` version is incompatible with CUDA 11.7 and earlier.
+
 ## Use cases
 This library provided utilities for both data iteration and transformation via dataloaders defined in `ml4gw/dataloading` and transform layers exposed in `ml4gw/transforms`. Lower level functions and utilies are defined at the top level of the library and in the `utils` library.
 
@@ -128,3 +149,7 @@ We encourage users who encounter these difficulties to file issues on GitHub, an
 We also strongly encourage ML users in the GW physics space to try their hand at working on these issues and joining on as collaborators!
 For more information about how to get involved, feel free to reach out to [ml4gw@ligo.mit.edu](mailto:ml4gw@ligo.mit.edu) .
 By bringing in new users with new use cases, we hope to develop this library into a truly general-purpose tool which makes DL more accessible for gravitational wave physicists everywhere.
+
+## Funding
+We are grateful for the support of the U.S. National Science Foundation (NSF) Harnessing the Data Revolution (HDR) Institute for <a href="https://a3d3.ai">Accelerating AI Algorithms for Data Driven Discovery (A3D3)</a> under Cooperative Agreement No. <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=2117997">PHY-2117997</a>.
+

ml4gw-0.2.0/PKG-INFO → ml4gw-0.4.0/README.md

@@ -1,19 +1,3 @@
-Metadata-Version: 2.1
-Name: ml4gw
-Version: 0.2.0
-Summary: Tools for training torch models on gravitational wave data
-Author: Alec Gunny
-Author-email: alec.gunny@ligo.org
-Requires-Python: >=3.8,<4.0
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: torch (>=1.10,<2.0)
-Requires-Dist: torchtyping (>=0.1,<0.2)
-Description-Content-Type: text/markdown
-
 # ML4GW
 
 Torch utilities for training neural networks in gravitational wave physics applications.
@@ -37,8 +21,8 @@ pip install ml4gw torch==1.12.0 --extra-index-url=https://download.pytorch.org/w
 
 ```toml
 [tool.poetry.dependencies]
-python = "^3.8"  # python versions 3.8-3.10 are supported
-ml4gw = "^0.1.0"
+python = "^3.8"  # python versions 3.8-3.11 are supported
+ml4gw = "^0.3.0"
 ```
 
 To build against a specific PyTorch/CUDA combination, consult the PyTorch installation documentation above and specify the `extra-index-url` via the `tool.poetry.source` table in your `pyproject.toml`. For example, to build against CUDA 11.6, you would do something like:
@@ -46,7 +30,7 @@ To build against a specific PyTorch/CUDA combination, consult the PyTorch instal
 ```toml
 [tool.poetry.dependencies]
 python = "^3.8"
-ml4gw = "^0.1.0"
+ml4gw = "^0.3.0"
 torch = {version = "^1.12", source = "torch"}
 
 [[tool.poetry.source]]
@@ -56,6 +40,8 @@ secondary = true
 default = false
 ```
 
+Note: if you are building against CUDA 11.6 or 11.7, make sure that you are using python 3.8, 3.9, or 3.10. Python 3.11 is incompatible with `torchaudio` 0.13, and the following `torchaudio` version is incompatible with CUDA 11.7 and earlier.
+
 ## Use cases
 This library provided utilities for both data iteration and transformation via dataloaders defined in `ml4gw/dataloading` and transform layers exposed in `ml4gw/transforms`. Lower level functions and utilies are defined at the top level of the library and in the `utils` library.
 
@@ -145,3 +131,5 @@ We also strongly encourage ML users in the GW physics space to try their hand at
 For more information about how to get involved, feel free to reach out to [ml4gw@ligo.mit.edu](mailto:ml4gw@ligo.mit.edu) .
 By bringing in new users with new use cases, we hope to develop this library into a truly general-purpose tool which makes DL more accessible for gravitational wave physicists everywhere.
 
+## Funding
+We are grateful for the support of the U.S. National Science Foundation (NSF) Harnessing the Data Revolution (HDR) Institute for <a href="https://a3d3.ai">Accelerating AI Algorithms for Data Driven Discovery (A3D3)</a> under Cooperative Agreement No. <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=2117997">PHY-2117997</a>.

ml4gw-0.4.0/ml4gw/augmentations.py

@@ -0,0 +1,43 @@
+import torch
+
+
+class SignalInverter(torch.nn.Module):
+    """
+    Takes a tensor of timeseries of arbitrary dimension
+    and randomly inverts (i.e. h(t) -> -h(t))
+    each timeseries with probability `prob`.
+
+    Args:
+        prob:
+            Probability that a timeseries is inverted
+    """
+
+    def __init__(self, prob: float = 0.5):
+        super().__init__()
+        self.prob = prob
+
+    def forward(self, X):
+        mask = torch.rand(size=X.shape[:-1]) < self.prob
+        X[mask] *= -1
+        return X
+
+
+class SignalReverser(torch.nn.Module):
+    """
+    Takes a tensor of timeseries of arbitrary dimension
+    and randomly reverses (i.e. h(t) -> h(-t))
+    each timeseries with probability `prob`.
+
+    Args:
+        prob:
+            Probability that a kernel is reversed
+    """
+
+    def __init__(self, prob: float = 0.5):
+        super().__init__()
+        self.prob = prob
+
+    def forward(self, X):
+        mask = torch.rand(size=X.shape[:-1]) < self.prob
+        X[mask] = X[mask].flip(-1)
+        return X

ml4gw-0.4.0/ml4gw/dataloading/__init__.py

@@ -0,0 +1,3 @@
+from .chunked_dataset import ChunkedTimeSeriesDataset
+from .hdf5_dataset import Hdf5TimeSeriesDataset
+from .in_memory_dataset import InMemoryDataset

ml4gw-0.4.0/ml4gw/dataloading/chunked_dataset.py

@@ -0,0 +1,134 @@
+from collections.abc import Iterable
+
+import torch
+
+
+class ChunkedTimeSeriesDataset(torch.utils.data.IterableDataset):
+    """
+    Wrapper dataset that will loop through chunks of timeseries
+    data produced by another iterable and sample windows from
+    these chunks.
+
+    Args:
+        chunk_it:
+            Iterator which will produce chunks of timeseries
+            data to sample windows from. Should have shape
+            `(N, C, T)`, where `N` is the number of chunks
+            to sample from, `C` is the number of channels,
+            and `T` is the number of samples along the
+            time dimension for each chunk.
+        kernel_size:
+            Size of windows to be sampled from each chunk.
+            Should be less than the size of each chunk
+            along the time dimension.
+        batch_size:
+            Number of windows to sample at each iteration
+        batches_per_chunk:
+            Number of batches of windows to sample from
+            each chunk before moving on to the next one.
+            Sampling fewer batches from each chunk means
+            a lower likelihood of sampling duplicate windows,
+            but an increase in chunk-loading overhead.
+        coincident:
+            Whether the windows sampled from individual
+            channels in each batch element should be
+            sampled coincidentally, i.e. consisting of
+            the same timesteps, or whether each window
+            should be sample independently from the others.
+        device:
+            Which device chunks should be moved to upon loading.
+    """
+
+    def __init__(
+        self,
+        chunk_it: Iterable,
+        kernel_size: float,
+        batch_size: int,
+        batches_per_chunk: int,
+        coincident: bool = True,
+        device: str = "cpu",
+    ) -> None:
+        self.chunk_it = chunk_it
+        self.kernel_size = kernel_size
+        self.batch_size = batch_size
+        self.batches_per_chunk = batches_per_chunk
+        self.coincident = coincident
+        self.device = device
+
+    def __len__(self):
+        return len(self.chunk_it) * self.batches_per_chunk
+
+    def __iter__(self):
+        it = iter(self.chunk_it)
+        chunk = next(it)
+        num_chunks, num_channels, chunk_size = chunk.shape
+
+        # if we're sampling coincidentally, we only need
+        # to sample indices on a per-batch-element basis.
+        # Otherwise, we'll need indices for both each
+        # batch sample _and_ each channel with each sample
+        if self.coincident:
+            sample_size = (self.batch_size,)
+        else:
+            sample_size = (self.batch_size, num_channels)
+
+        # slice kernels out a flattened chunk tensor
+        # index-for-index. We'll account for batch/
+        # channel indices by introducing offsets later on
+        idx = torch.arange(self.kernel_size, device=self.device)
+        idx = idx.view(1, 1, -1)
+        idx = idx.repeat(self.batch_size, num_channels, 1)
+
+        # this will just be a set of aranged channel indices
+        # repeated to offset the kernel indices in the
+        # flattened chunk tensor
+        channel_idx = torch.arange(num_channels, device=self.device)
+        channel_idx = channel_idx.view(1, -1, 1)
+        channel_idx = channel_idx.repeat(self.batch_size, 1, self.kernel_size)
+        idx += channel_idx * chunk_size
+
+        while True:
+            # record the number of rows in the chunk, then
+            # flatten it to make it easier to slice
+            if chunk_size < self.kernel_size:
+                raise ValueError(
+                    "Can't sample kernels of size {} from chunk "
+                    "with size {}".format(self.kernel_size, chunk_size)
+                )
+            chunk = chunk.reshape(-1)
+
+            # generate batches from the current chunk
+            for _ in range(self.batches_per_chunk):
+                # first sample the indices of which chunk elements
+                # we're going to read batch elements from
+                chunk_idx = torch.randint(
+                    0, num_chunks, size=sample_size, device=self.device
+                )
+
+                # account for the offset this batch element
+                # introduces in the flattened array
+                chunk_idx *= num_channels * chunk_size
+                chunk_idx = chunk_idx.view(self.batch_size, -1, 1)
+                chunk_idx = chunk_idx + idx
+
+                # now sample the start index within each chunk
+                # element we're going to grab our time windows from
+                time_idx = torch.randint(
+                    0,
+                    chunk_size - self.kernel_size,
+                    size=sample_size,
+                    device=self.device,
+                )
+                time_idx = time_idx.view(self.batch_size, -1, 1)
+
+                # there's no additional offset factor to account for here
+                chunk_idx += time_idx
+
+                # now slice this 3D tensor from our flattened chunk
+                yield chunk[chunk_idx]
+
+            try:
+                chunk = next(it)
+            except StopIteration:
+                break
+            num_chunks, num_channels, chunk_size = chunk.shape

ml4gw-0.4.0/ml4gw/dataloading/hdf5_dataset.py

@@ -0,0 +1,176 @@
+import warnings
+from typing import Sequence, Union
+
+import h5py
+import numpy as np
+import torch
+
+from ml4gw.types import WaveformTensor
+
+
+class ContiguousHdf5Warning(Warning):
+    pass
+
+
+class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
+    """
+    Iterable dataset that samples and loads windows of
+    timeseries data uniformly from a set of HDF5 files.
+    It is _strongly_ recommended that these files have been
+    written using [chunked storage]
+    (https://docs.h5py.org/en/stable/high/dataset.html#chunked-storage).
+    This has shown to produce increases in read-time speeds
+    of over an order of magnitude.
+
+    Args:
+        fnames:
+            Paths to HDF5 files from which to sample data.
+        channels:
+            Datasets to read from the indicated files, which
+            will be stacked along dim 1 of the generated batches
+            during iteration.
+        kernel_size:
+            Size of the windows to read, in number of samples.
+            This will be the size of the last dimension of the
+            generated batches.
+        batch_size:
+            Number of windows to sample at each iteration.
+        batches_per_epoch:
+            Number of batches to generate during each call
+            to `__iter__`.
+        coincident:
+            Whether windows for each channel in a given batch
+            element should be sampled coincidentally, i.e.
+            corresponding to the same time indices from the
+            same files, or should be sampled independently.
+            For the latter case, users can either specify
+            `False`, which will sample filenames independently
+            for each channel, or `"files"`, which will sample
+            windows independently within a given file for each
+            channel. The latter setting limits the amount of
+            entropy in the effective dataset, but can provide
+            over 2x improvement in total throughput.
+    """
+
+    def __init__(
+        self,
+        fnames: Sequence[str],
+        channels: Sequence[str],
+        kernel_size: int,
+        batch_size: int,
+        batches_per_epoch: int,
+        coincident: Union[bool, str],
+    ) -> None:
+        if not isinstance(coincident, bool) and coincident != "files":
+            raise ValueError(
+                "coincident must be either a boolean or 'files', "
+                "got unrecognized value {}".format(coincident)
+            )
+
+        self.fnames = fnames
+        self.channels = channels
+        self.num_channels = len(channels)
+        self.kernel_size = kernel_size
+        self.batch_size = batch_size
+        self.batches_per_epoch = batches_per_epoch
+        self.coincident = coincident
+
+        self.sizes = {}
+        for fname in self.fnames:
+            with h5py.File(fname, "r") as f:
+                dset = f[channels[0]]
+                if dset.chunks is None:
+                    warnings.warn(
+                        "File {} contains datasets that were generated "
+                        "without using chunked storage. This can have "
+                        "severe performance impacts at data loading time. "
+                        "If you need faster loading, try re-generating "
+                        "your datset with chunked storage turned on.".format(
+                            fname
+                        ),
+                        category=ContiguousHdf5Warning,
+                    )
+
+                self.sizes[fname] = len(dset)
+        total = sum(self.sizes.values())
+        self.probs = np.array([i / total for i in self.sizes.values()])
+
+    def __len__(self) -> int:
+        return self.batches_per_epoch
+
+    def sample_fnames(self, size) -> np.ndarray:
+        return np.random.choice(
+            self.fnames,
+            p=self.probs,
+            size=size,
+            replace=True,
+        )
+
+    def sample_batch(self) -> WaveformTensor:
+        """
+        Sample a single batch of multichannel timeseries
+        """
+
+        # allocate memory up front
+        x = np.zeros((self.batch_size, len(self.channels), self.kernel_size))
+
+        # sample filenames, but only loop through each unique
+        # filename once to avoid unnecessary I/O overhead
+        if self.coincident is not False:
+            size = (self.batch_size,)
+        else:
+            size = (self.batch_size, self.num_channels)
+        fnames = self.sample_fnames(size)
+
+        unique_fnames, inv, counts = np.unique(
+            fnames, return_inverse=True, return_counts=True
+        )
+        for i, (fname, count) in enumerate(zip(unique_fnames, counts)):
+            size = self.sizes[fname]
+            max_idx = size - self.kernel_size
+
+            # figure out which batch indices should be
+            # sampled from the current filename
+            indices = np.where(inv == i)[0]
+
+            # when sampling coincidentally either fully
+            # or at the file level, all channels will
+            # correspond to the same file
+            if self.coincident is not False:
+                batch_indices = np.repeat(indices, self.num_channels)
+                channel_indices = np.arange(self.num_channels)
+                channel_indices = np.concatenate([channel_indices] * count)
+            else:
+                batch_indices = indices // self.num_channels
+                channel_indices = indices % self.num_channels
+
+            # if we're sampling fully coincidentally, each
+            # channel will be the same in each file
+            if self.coincident is True:
+                idx = np.random.randint(max_idx, size=count)
+                idx = np.repeat(idx, self.num_channels)
+            else:
+                # otherwise, every channel will be different
+                # for the given file
+                idx = np.random.randint(max_idx, size=len(batch_indices))
+
+            # open the file and sample a different set of
+            # kernels for each batch element it occupies
+            with h5py.File(fname, "r") as f:
+                for b, c, i in zip(batch_indices, channel_indices, idx):
+                    x[b, c] = f[self.channels[c]][i : i + self.kernel_size]
+        return torch.Tensor(x)
+
+    def __iter__(self) -> torch.Tensor:
+        worker_info = torch.utils.data.get_worker_info()
+        if worker_info is None:
+            num_batches = self.batches_per_epoch
+        else:
+            num_batches, remainder = divmod(
+                self.batches_per_epoch, worker_info.num_workers
+            )
+            if worker_info.id < remainder:
+                num_batches += 1
+
+        for _ in range(num_batches):
+            yield self.sample_batch()

ml4gw-0.4.0/ml4gw/nn/autoencoder/__init__.py

@@ -0,0 +1,3 @@
+from .base import Autoencoder
+from .convolutional import ConvolutionalAutoencoder
+from .skip_connection import AddSkipConnect, ConcatSkipConnect, SkipConnection

ml4gw-0.4.0/ml4gw/nn/autoencoder/base.py

@@ -0,0 +1,89 @@
+from collections.abc import Sequence
+from typing import Optional
+
+import torch
+
+from ml4gw.nn.autoencoder.skip_connection import SkipConnection
+
+
+class Autoencoder(torch.nn.Module):
+    """
+    Base autoencoder class that defines some of the
+    basic methods and functionality. Autoencoders are
+    defined here as a set of sequential blocks that
+    have an `encode` method, which acts on the input
+    data to the autoencoder, and a `decode` method, which
+    acts on the encoded vector generated by the `encode`
+    method. `forward` just runs these steps one after the
+    other. Although it isn't explicitly enforced, a good
+    rule of thumb is that the ouput of a block's `decode`
+    method should have the same shape as the _input_ of its
+    `encode` method.
+
+    Accepts a `skip_connection` argument that defines how to
+    combine information from the input of one block's `encode`
+    layer with the output to its `decode`layer. See `skip_connections.py`
+    for more info about what these classes are expected to contain
+    and how they operate.
+    """
+
+    def __init__(self, skip_connection: Optional[SkipConnection] = None):
+        super().__init__()
+        self.skip_connection = skip_connection
+        self.blocks = torch.nn.ModuleList()
+
+    def encode(self, *X: torch.Tensor, return_states: bool = False):
+        states = []
+        for block in self.blocks:
+            if isinstance(X, tuple):
+                X = block.encode(*X)
+            else:
+                X = block.encode(X)
+            states.append(X)
+
+        # don't need to return the last
+        # state, since that's just equal
+        # to the output of this layer
+        if return_states:
+            return X, states[:-1]
+        return X
+
+    def decode(self, *X, states: Optional[Sequence[torch.Tensor]] = None):
+        if self.skip_connection is not None and states is None:
+            raise ValueError(
+                "Must pass intermediate states when autoencoder "
+                "has a skip connection function specified"
+            )
+        elif states is not None:
+            if len(states) != len(self.blocks) - 1:
+                raise ValueError(
+                    "Passed {} intermediate states, expected {}".format(
+                        len(states), len(self.blocks) - 1
+                    )
+                )
+
+            # Don't skip connect the output layer
+            states = states[::-1] + [None]
+
+        for i, block in enumerate(self.blocks[::-1]):
+            if isinstance(X, tuple):
+                X = block.decode(*X)
+            else:
+                X = block.decode(X)
+
+            state = states[-i - 1]
+            if state is not None:
+                X = self.skip_connection(X, state)
+        return X
+
+    def forward(self, *X):
+        return_states = self.skip_connection is not None
+        X = self.encode(*X, return_states=return_states)
+        if return_states:
+            *X, states = X
+        else:
+            states = None
+
+        if isinstance(X, torch.Tensor):
+            X = (X,)
+        return self.decode(*X, states=states)