ml4gw 0.7.6__tar.gz → 0.7.8__tar.gz

{ml4gw-0.7.6 → ml4gw-0.7.8}/PKG-INFO

@@ -1,22 +1,24 @@
 Metadata-Version: 2.4
 Name: ml4gw
-Version: 0.7.6
+Version: 0.7.8
 Summary: Tools for training torch models on gravitational wave data
-Author-email: Ethan Marx <emarx@mit.edu>, Will Benoit <benoi090@umn.edu>, Deep Chatterjee <deep1018@mit.edu>, Alec Gunny <alec.gunny@ligo.org>
-License-File: LICENSE
-Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
-Classifier: Programming Language :: Python :: 3.9
+Author-email: Ethan Marx <emarx@mit.edu>, Will Benoit <benoi090@umn.edu>, Deep Chatterjee <deep1018@mit.edu>, Alec Gunny <alec.gunny@ligo.org>, Ravi Kumar <ravi.kumar@ligo.org>
+Maintainer-email: Ethan Marx <emarx@mit.edu>, Will Benoit <benoi090@umn.edu>, Deep Chatterjee <deep1018@mit.edu>
+License-Expression: GPL-3.0-or-later
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
-Requires-Python: <3.13,>=3.9
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: <3.13,>=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
 Requires-Dist: jaxtyping<0.3,>=0.2
+Requires-Dist: torch~=2.0
+Requires-Dist: torchaudio~=2.0
 Requires-Dist: numpy<2.0.0
 Requires-Dist: scipy<1.15,>=1.9.0
-Requires-Dist: torchaudio~=2.0
-Requires-Dist: torch~=2.0
-Description-Content-Type: text/markdown
+Dynamic: license-file
 
 # ML4GW
 ![PyPI - Version](https://img.shields.io/pypi/v/ml4gw)
@@ -29,7 +31,10 @@ Torch utilities for training neural networks in gravitational wave physics appli
 
 ## Documentation
 Please visit our [documentation page](https://ml4gw.github.io/ml4gw/) to see descriptions and examples of the functions and modules available in `ml4gw`.
-We also have an interactive Jupyter notebook that demonstrates much of the core functionality available in the `examples` directory.
+We also have an interactive Jupyter notebook demonstrating much of the core functionality available [here](https://github.com/ML4GW/ml4gw/blob/main/docs/tutorials/ml4gw_tutorial.ipynb).
+To run this notebook, download it from the above link and follow the instructions within it to install the required packages.
+See also the [documentation page](https://ml4gw.github.io/ml4gw/tutorials/ml4gw_tutorial.html) for the tutorial to look
+through it without running the code.
 
 ## Installation
 ### Pip installation
@@ -45,9 +50,25 @@ To build with a specific version of PyTorch/CUDA, please see the PyTorch install
 pip install ml4gw torch==2.5.1--extra-index-url=https://download.pytorch.org/whl/cu118
 ```
 
+### uv installation
+If you want to develop `ml4gw`, you can use [uv](https://docs.astral.sh/uv/getting-started/installation/) to install the project in editable mode.
+For example, after cloning the repository, create a virtualenv using
+```bash
+uv venv --python=3.11
+```
+Then sync the dependencies from the [uv lock file](/uv.lock) using
+```bash
+uv sync --all-extras
+```
+Code changes can be tested using
+```bash
+uv run pytest
+```
+See [contribution guide](/CONTRIBUTING.md) for more details.
+
 ## Contributing
 If you come across errors in the code, have difficulties using this software, or simply find that the current version doesn't cover your use case, please file an issue on our GitHub page, and we'll be happy to offer support.
-We encourage users who encounter these difficulties to file issues on GitHub, and we'll be happy to offer support to extend our coverage to new or improved functionality.
+If you want to add feature, please refer to the [contribution guide](/CONTRIBUTING.md) for more details.
 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 that makes deep learning more accessible for gravitational wave physicists everywhere.

{ml4gw-0.7.6 → ml4gw-0.7.8}/README.md

@@ -9,7 +9,10 @@ Torch utilities for training neural networks in gravitational wave physics appli
 
 ## Documentation
 Please visit our [documentation page](https://ml4gw.github.io/ml4gw/) to see descriptions and examples of the functions and modules available in `ml4gw`.
-We also have an interactive Jupyter notebook that demonstrates much of the core functionality available in the `examples` directory.
+We also have an interactive Jupyter notebook demonstrating much of the core functionality available [here](https://github.com/ML4GW/ml4gw/blob/main/docs/tutorials/ml4gw_tutorial.ipynb).
+To run this notebook, download it from the above link and follow the instructions within it to install the required packages.
+See also the [documentation page](https://ml4gw.github.io/ml4gw/tutorials/ml4gw_tutorial.html) for the tutorial to look
+through it without running the code.
 
 ## Installation
 ### Pip installation
@@ -25,9 +28,25 @@ To build with a specific version of PyTorch/CUDA, please see the PyTorch install
 pip install ml4gw torch==2.5.1--extra-index-url=https://download.pytorch.org/whl/cu118
 ```
 
+### uv installation
+If you want to develop `ml4gw`, you can use [uv](https://docs.astral.sh/uv/getting-started/installation/) to install the project in editable mode.
+For example, after cloning the repository, create a virtualenv using
+```bash
+uv venv --python=3.11
+```
+Then sync the dependencies from the [uv lock file](/uv.lock) using
+```bash
+uv sync --all-extras
+```
+Code changes can be tested using
+```bash
+uv run pytest
+```
+See [contribution guide](/CONTRIBUTING.md) for more details.
+
 ## Contributing
 If you come across errors in the code, have difficulties using this software, or simply find that the current version doesn't cover your use case, please file an issue on our GitHub page, and we'll be happy to offer support.
-We encourage users who encounter these difficulties to file issues on GitHub, and we'll be happy to offer support to extend our coverage to new or improved functionality.
+If you want to add feature, please refer to the [contribution guide](/CONTRIBUTING.md) for more details.
 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 that makes deep learning more accessible for gravitational wave physicists everywhere.

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/augmentations.py

@@ -1,3 +1,8 @@
+"""
+This module contains transformations that may be useful
+for augmenting timeseries data during training
+"""
+
 import torch
 from jaxtyping import Float
 from torch import Tensor

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/dataloading/__init__.py

@@ -1,3 +1,8 @@
+"""
+This module contains tools for efficient in-memory and
+out-of-memory dataloading.
+"""
+
 from .chunked_dataset import ChunkedTimeSeriesDataset
 from .hdf5_dataset import Hdf5TimeSeriesDataset
 from .in_memory_dataset import InMemoryDataset

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/dataloading/chunked_dataset.py

@@ -94,10 +94,8 @@ class ChunkedTimeSeriesDataset(torch.utils.data.IterableDataset):
             # 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)
+                    f"Can't sample kernels of size {self.kernel_size} from "
+                    f"chunk with size {chunk_size}"
                 )
             chunk = chunk.reshape(-1)
 

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/dataloading/hdf5_dataset.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Optional, Sequence, Union
+from collections.abc import Sequence
 
 import h5py
 import numpy as np
@@ -63,13 +63,13 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
         kernel_size: int,
         batch_size: int,
         batches_per_epoch: int,
-        coincident: Union[bool, str],
-        num_files_per_batch: Optional[int] = None,
+        coincident: bool | str,
+        num_files_per_batch: int | None = None,
     ) -> None:
         if not isinstance(coincident, bool) and coincident != "files":
             raise ValueError(
                 "coincident must be either a boolean or 'files', "
-                "got unrecognized value {}".format(coincident)
+                f"got unrecognized value {coincident}"
             )
 
         self.fnames = np.array(fnames)
@@ -94,13 +94,11 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
                 dset = f[channels[0]]
                 if dset.chunks is None:
                     warnings.warn(
-                        "File {} contains datasets that were generated "
+                        f"File {fname} 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 dataset with chunked storage turned on.".format(
-                            fname
-                        ),
+                        "your dataset with chunked storage turned on.",
                         category=ContiguousHdf5Warning,
                         stacklevel=2,
                     )
@@ -153,7 +151,9 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
         unique_fnames, inv, counts = np.unique(
             fnames, return_inverse=True, return_counts=True
         )
-        for i, (fname, count) in enumerate(zip(unique_fnames, counts)):
+        for i, (fname, count) in enumerate(
+            zip(unique_fnames, counts, strict=True)
+        ):
             size = self.sizes[fname]
             max_idx = size - self.kernel_size
 
@@ -185,7 +185,9 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
             # 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):
+                for b, c, i in zip(
+                    batch_indices, channel_indices, idx, strict=True
+                ):
                     x[b, c] = f[self.channels[c]][i : i + self.kernel_size]
         return torch.Tensor(x)
 

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/dataloading/in_memory_dataset.py

@@ -1,5 +1,4 @@
 import itertools
-from typing import Optional, Tuple, Union
 
 import torch
 from jaxtyping import Float
@@ -79,10 +78,10 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
         self,
         X: Float[Tensor, "channels time"],
         kernel_size: int,
-        y: Optional[Float[Tensor, " time"]] = None,
+        y: Float[Tensor, " time"] | None = None,
         batch_size: int = 32,
         stride: int = 1,
-        batches_per_epoch: Optional[int] = None,
+        batches_per_epoch: int | None = None,
         coincident: bool = True,
         shuffle: bool = True,
         device: str = "cpu",
@@ -122,10 +121,9 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
             batch_size * batches_per_epoch
         ):
             raise ValueError(
-                "Number of kernels {} in timeseries insufficient "
-                "to generate {} batches of size {}".format(
-                    self.num_kernels, batch_size, batches_per_epoch
-                )
+                f"Number of kernels {self.num_kernels} in timeseries "
+                f"insufficient to generate {batch_size} batches of size "
+                f"{batches_per_epoch}"
             )
 
         self.batch_size = batch_size
@@ -191,7 +189,9 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
             # indices we'll need rather than having to generate
             # everything.
             idx = [range(self.num_kernels) for _ in range(len(self.X))]
-            idx = zip(range(num_kernels), itertools.product(*idx))
+            idx = zip(
+                range(num_kernels), itertools.product(*idx), strict=False
+            )
             idx = torch.stack([torch.Tensor(i[1]) for i in idx])
             idx = idx.type(torch.int64).to(device)
         elif self.shuffle:
@@ -208,10 +208,10 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
 
     def __iter__(
         self,
-    ) -> Union[
-        Float[Tensor, "batch channel time"],
-        Tuple[Float[Tensor, "batch channel time"], Float[Tensor, " batch"]],
-    ]:
+    ) -> (
+        Float[Tensor, "batch channel time"]
+        | tuple[Float[Tensor, "batch channel time"], Float[Tensor, " batch"]]
+    ):
         indices = self.init_indices()
         for i in range(len(self)):
             # slice the array of _indices_ we'll be using to

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/distributions.py

@@ -6,7 +6,7 @@ from the corresponding distribution.
 """
 
 import math
-from typing import Callable, Optional
+from collections.abc import Callable
 
 import torch
 import torch.distributions as dist
@@ -104,7 +104,7 @@ class LogNormal(dist.LogNormal):
         self,
         mean: float,
         std: float,
-        low: Optional[float] = None,
+        low: float | None = None,
         validate_args=None,
     ):
         self.low = low
@@ -137,7 +137,7 @@ class PowerLaw(dist.TransformedDistribution):
     support = dist.constraints.nonnegative
 
     def __init__(
-        self, minimum: float, maximum: float, index: int, validate_args=None
+        self, minimum: float, maximum: float, index: float, validate_args=None
     ):
         if index == 0:
             raise ValueError("Index of 0 is the same as Uniform")

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/gw.py

@@ -1,6 +1,7 @@
 """
-Tools for manipulating raw gravitational waveforms
-and projecting them onto interferometer responses.
+Tools for manipulating raw gravitational waveforms,
+projecting them onto interferometer responses, and
+calculating SNRs.
 Much of the projection code is an extension of the
 implementation made available in
 `bilby <https://arxiv.org/abs/1811.02042>`_.
@@ -8,8 +9,6 @@ Specifically code from
 `this module <https://github.com/lscsoft/bilby/blob/master/bilby/gw/detector/interferometer.py>`_.
 """  # noqa E501
 
-from typing import List, Tuple, Union
-
 import torch
 from jaxtyping import Float
 from torch import Tensor
@@ -58,7 +57,7 @@ def compute_antenna_responses(
     psi: BatchTensor,
     phi: BatchTensor,
     detector_tensors: NetworkDetectorTensors,
-    modes: List[str],
+    modes: list[str],
 ) -> Float[Tensor, "batch polarizations num_ifos"]:
     """
     Compute the antenna pattern factors of a batch of
@@ -257,7 +256,7 @@ def compute_observed_strain(
 
 def get_ifo_geometry(
     *ifos: str,
-) -> Tuple[NetworkDetectorTensors, NetworkVertices]:
+) -> tuple[NetworkDetectorTensors, NetworkVertices]:
     """
     For a given list of interferometer names, retrieve and
     concatenate the associated detector tensors and vertices
@@ -286,8 +285,8 @@ def compute_ifo_snr(
     responses: WaveformTensor,
     psd: PSDTensor,
     sample_rate: float,
-    highpass: Union[float, Float[Tensor, " frequency"], None] = None,
-    lowpass: Union[float, Float[Tensor, " frequency"], None] = None,
+    highpass: float | Float[Tensor, " frequency"] | None = None,
+    lowpass: float | Float[Tensor, " frequency"] | None = None,
 ) -> Float[Tensor, "batch num_ifos"]:
     """Compute the SNRs of a batch of interferometer responses
 
@@ -367,10 +366,9 @@ def compute_ifo_snr(
             highpass = freqs >= highpass
         elif len(highpass) != integrand.shape[-1]:
             raise ValueError(
-                "Can't apply highpass filter mask with {} frequency bins"
-                "to signal fft with {} frequency bins".format(
-                    len(highpass), integrand.shape[-1]
-                )
+                f"Can't apply highpass filter mask with {len(highpass)} "
+                f"frequency bins to signal fft with {integrand.shape[-1]} "
+                "frequency bins"
             )
         integrand *= highpass.to(integrand.device)
     if lowpass is not None:
@@ -379,10 +377,9 @@ def compute_ifo_snr(
             lowpass = freqs < lowpass
         elif len(lowpass) != integrand.shape[-1]:
             raise ValueError(
-                "Can't apply lowpass filter mask with {} frequency bins"
-                "to signal fft with {} frequency bins".format(
-                    len(lowpass), integrand.shape[-1]
-                )
+                f"Can't apply lowpass filter mask with {len(lowpass)} "
+                f"frequency bins to signal fft with {integrand.shape[-1]} "
+                "frequency bins"
             )
         integrand *= lowpass.to(integrand.device)
 
@@ -410,8 +407,8 @@ def compute_network_snr(
     responses: WaveformTensor,
     psd: PSDTensor,
     sample_rate: float,
-    highpass: Union[float, Float[Tensor, " frequency"], None] = None,
-    lowpass: Union[float, Float[Tensor, " frequency"], None] = None,
+    highpass: float | Float[Tensor, " frequency"] | None = None,
+    lowpass: float | Float[Tensor, " frequency"] | None = None,
 ) -> BatchTensor:
     """
     Compute the total SNR from a gravitational waveform
@@ -467,11 +464,11 @@ def compute_network_snr(
 
 def reweight_snrs(
     responses: WaveformTensor,
-    target_snrs: Union[float, BatchTensor],
+    target_snrs: float | BatchTensor,
     psd: PSDTensor,
     sample_rate: float,
-    highpass: Union[float, Float[Tensor, " frequency"], None] = None,
-    lowpass: Union[float, Float[Tensor, " frequency"], None] = None,
+    highpass: float | Float[Tensor, " frequency"] | None = None,
+    lowpass: float | Float[Tensor, " frequency"] | None = None,
 ) -> WaveformTensor:
     """Scale interferometer responses such that they have a desired SNR
 

ml4gw-0.7.8/ml4gw/nn/__init__.py

@@ -0,0 +1,6 @@
+"""
+This module contains neural network architectures and
+architecture components. These can be a good place
+to get started, rather than defining your own
+architecture from the start.
+"""

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/nn/autoencoder/base.py

@@ -1,5 +1,4 @@
 from collections.abc import Sequence
-from typing import Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -28,16 +27,14 @@ class Autoencoder(torch.nn.Module):
     and how they operate.
     """
 
-    def __init__(
-        self, skip_connection: Optional[SkipConnection] = None
-    ) -> None:
+    def __init__(self, skip_connection: SkipConnection | None = None) -> None:
         super().__init__()
         self.skip_connection = skip_connection
         self.blocks = torch.nn.ModuleList()
 
     def encode(
         self, *X: Tensor, return_states: bool = False
-    ) -> Union[Tensor, Tuple[Tensor, Sequence]]:
+    ) -> Tensor | tuple[Tensor, Sequence]:
         states = []
         for block in self.blocks:
             if isinstance(X, tuple):
@@ -53,7 +50,7 @@ class Autoencoder(torch.nn.Module):
             return X, states[:-1]
         return X
 
-    def decode(self, *X, states: Optional[Sequence[Tensor]] = None) -> Tensor:
+    def decode(self, *X, states: Sequence[Tensor] | None = None) -> Tensor:
         if self.skip_connection is not None and states is None:
             raise ValueError(
                 "Must pass intermediate states when autoencoder "
@@ -62,9 +59,8 @@ class Autoencoder(torch.nn.Module):
         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
-                    )
+                    f"Passed {len(states)} intermediate states, expected "
+                    f"{len(self.blocks) - 1}"
                 )
 
             # Don't skip connect the output layer

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/nn/autoencoder/convolutional.py

@@ -1,5 +1,4 @@
 from collections.abc import Callable, Sequence
-from typing import Optional
 
 import torch
 from torch import Tensor
@@ -21,9 +20,9 @@ class ConvBlock(Autoencoder):
         groups: int = 1,
         activation: torch.nn.Module = torch.nn.ReLU,
         norm: Module = torch.nn.BatchNorm1d,
-        decode_channels: Optional[int] = None,
-        output_activation: Optional[torch.nn.Module] = None,
-        skip_connection: Optional[SkipConnection] = None,
+        decode_channels: int | None = None,
+        output_activation: torch.nn.Module | None = None,
+        skip_connection: SkipConnection | None = None,
     ) -> None:
         super().__init__(skip_connection=None)
 
@@ -98,10 +97,10 @@ class ConvolutionalAutoencoder(Autoencoder):
         stride: int = 1,
         groups: int = 1,
         activation: torch.nn.Module = torch.nn.ReLU,
-        output_activation: Optional[torch.nn.Module] = None,
+        output_activation: torch.nn.Module | None = None,
         norm: Module = torch.nn.BatchNorm1d,
-        decode_channels: Optional[int] = None,
-        skip_connection: Optional[SkipConnection] = None,
+        decode_channels: int | None = None,
+        skip_connection: SkipConnection | None = None,
     ) -> None:
         # TODO: how to do this dynamically? Maybe the base
         # architecture looks for overlapping arguments between
@@ -145,9 +144,7 @@ class ConvolutionalAutoencoder(Autoencoder):
             self.blocks.append(block)
             in_channels = channels * groups
 
-    def decode(
-        self, *X, states=None, input_size: Optional[int] = None
-    ) -> Tensor:
+    def decode(self, *X, states=None, input_size: int | None = None) -> Tensor:
         X = super().decode(*X, states=states)
         if input_size is not None:
             return match_size(X, input_size)

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/nn/autoencoder/skip_connection.py

@@ -35,13 +35,11 @@ class ConcatSkipConnect(SkipConnection):
         rem = num_channels % self.groups
         if rem:
             raise ValueError(
-                "Number of channels in input tensor {} cannot "
-                "be divided evenly into {} groups".format(
-                    num_channels, self.groups
-                )
+                f"Number of channels in input tensor {num_channels} cannot "
+                f"be divided evenly into {self.groups} groups"
             )
 
         X = torch.split(X, self.groups, dim=1)
         state = torch.split(state, self.groups, dim=1)
-        frags = [i for j in zip(X, state) for i in j]
+        frags = [i for j in zip(X, state, strict=True) for i in j]
         return torch.cat(frags, dim=1)

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/nn/norm.py

@@ -1,4 +1,4 @@
-from typing import Callable, Optional
+from collections.abc import Callable
 
 import torch
 from jaxtyping import Float
@@ -16,7 +16,7 @@ class GroupNorm1D(torch.nn.Module):
     def __init__(
         self,
         num_channels: int,
-        num_groups: Optional[int] = None,
+        num_groups: int | None = None,
         eps: float = 1e-5,
     ):
         super().__init__()
@@ -77,7 +77,7 @@ class GroupNorm1DGetter:
     for command-line parameterization with jsonargparse.
     """
 
-    def __init__(self, groups: Optional[int] = None) -> None:
+    def __init__(self, groups: int | None = None) -> None:
         self.groups = groups
 
     def __call__(self, num_channels: int) -> torch.nn.Module:
@@ -96,7 +96,7 @@ class GroupNorm2DGetter:
     for command-line parameterization with jsonargparse.
     """
 
-    def __init__(self, groups: Optional[int] = None) -> None:
+    def __init__(self, groups: int | None = None) -> None:
         self.groups = groups
 
     def __call__(self, num_channels: int) -> torch.nn.Module:

{ml4gw-0.7.6 → ml4gw-0.7.8}/ml4gw/nn/resnet/resnet_1d.py

@@ -7,7 +7,8 @@ where training-time statistics are entirely arbitrary due to
 simulations.
 """
 
-from typing import Callable, List, Literal, Optional
+from collections.abc import Callable
+from typing import Literal
 
 import torch
 import torch.nn as nn
@@ -58,11 +59,11 @@ class BasicBlock(nn.Module):
         planes: int,
         kernel_size: int = 3,
         stride: int = 1,
-        downsample: Optional[nn.Module] = None,
+        downsample: nn.Module | None = None,
         groups: int = 1,
         base_width: int = 64,
         dilation: int = 1,
-        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Callable[..., nn.Module] | None = None,
     ) -> None:
         super().__init__()
         if norm_layer is None:
@@ -123,11 +124,11 @@ class Bottleneck(nn.Module):
         planes: int,
         kernel_size: int = 3,
         stride: int = 1,
-        downsample: Optional[nn.Module] = None,
+        downsample: nn.Module | None = None,
         groups: int = 1,
         base_width: int = 64,
         dilation: int = 1,
-        norm_layer: Optional[NormLayer] = None,
+        norm_layer: NormLayer | None = None,
     ) -> None:
         super().__init__()
         if norm_layer is None:
@@ -231,14 +232,14 @@ class ResNet1D(nn.Module):
     def __init__(
         self,
         in_channels: int,
-        layers: List[int],
+        layers: list[int],
         classes: int,
         kernel_size: int = 3,
         zero_init_residual: bool = False,
         groups: int = 1,
         width_per_group: int = 64,
-        stride_type: Optional[List[Literal["stride", "dilation"]]] = None,
-        norm_layer: Optional[NormLayer] = None,
+        stride_type: list[Literal["stride", "dilation"]] | None = None,
+        norm_layer: NormLayer | None = None,
     ) -> None:
         super().__init__()
 
@@ -257,10 +258,8 @@ class ResNet1D(nn.Module):
             stride_type = ["stride"] * (len(layers) - 1)
         if len(stride_type) != (len(layers) - 1):
             raise ValueError(
-                (
-                    "'stride_type' should be None or a {}-element "
-                    "tuple, got {}"
-                ).format(len(layers) - 1, stride_type)
+                f"'stride_type' should be None or a {len(layers) - 1}-element "
+                f"tuple, got {stride_type}"
             )
 
         self.groups = groups
@@ -289,7 +288,7 @@ class ResNet1D(nn.Module):
         # striding or dilating depending on the stride_type
         # argument)
         residual_layers = [self._make_layer(64, layers[0], kernel_size)]
-        it = zip(layers[1:], stride_type)
+        it = zip(layers[1:], stride_type, strict=True)
         for i, (num_blocks, stride) in enumerate(it):
             block_size = 64 * 2 ** (i + 1)
             layer = self._make_layer(