ml4gw 0.7.6__py3-none-any.whl → 0.7.8__py3-none-any.whl

ml4gw/transforms/spectrogram.py

@@ -1,5 +1,4 @@
 import warnings
-from typing import Dict, List
 
 import torch
 import torch.nn.functional as F
@@ -104,7 +103,7 @@ class MultiResolutionSpectrogram(torch.nn.Module):
         self.register_buffer("freq_idxs", freq_idxs)
         self.register_buffer("time_idxs", time_idxs)
 
-    def _check_and_format_kwargs(self, kwargs: Dict[str, List]) -> List:
+    def _check_and_format_kwargs(self, kwargs: dict[str, list]) -> list:
         lengths = sorted(len(v) for v in kwargs.values())
         lengths = list(set(lengths))
 
@@ -127,7 +126,10 @@ class MultiResolutionSpectrogram(torch.nn.Module):
             size = lengths[1]
             kwargs = {k: v * int(size / len(v)) for k, v in kwargs.items()}
 
-        return [dict(zip(kwargs, col)) for col in zip(*kwargs.values())]
+        return [
+            dict(zip(kwargs, col, strict=True))
+            for col in zip(*kwargs.values(), strict=True)
+        ]
 
     def forward(
         self, X: TimeSeries3d
@@ -161,6 +163,7 @@ class MultiResolutionSpectrogram(torch.nn.Module):
             self.right_pad,
             self.top_pad,
             self.bottom_pad,
+            strict=True,
         ):
             padded_specs.append(F.pad(spec, (left, right, top, bottom)))
 

ml4gw/transforms/spline_interpolation.py

@@ -1,131 +1,27 @@
 """
-Adaptation of code from https://github.com/dottormale/Qtransform
+Adaptation of code from https://github.com/dottormale/Qtransform_torch/
 """
 
-from typing import Optional, Tuple
-
 import torch
-import torch.nn.functional as F
 from torch import Tensor
 
 
-class SplineInterpolate(torch.nn.Module):
+class SplineInterpolateBase(torch.nn.Module):
     """
-    Perform 1D or 2D spline interpolation based on De Boor's method.
-    Supports batched, multi-channel inputs, so acceptable data
-    shapes are ``(width)``, ``(height, width)``, ``(batch, width)``,
-    ``(batch, height, width)``, ``(batch, channel, width)``, and
-    ``(batch, channel, height, width)``.
-
-    During initialization of this Module, both the desired input
-    and output coordinate Tensors can be specified to allow
-    pre-computation of the B-spline basis matrices, though the only
-    mandatory argument is the coordinates of the data along the
-    ``width`` dimension. If no argument is given for coordinates along
-    the ``height`` dimension, it is assumed that 1D interpolation is
-    desired.
-
-    Unlike scipy's implementation of spline interpolation, the data
-    to be interpolated is not passed until actually calling the
-    object. This is useful for cases where the input and output
-    coordinates are known in advance, but the data is not, so that
-    the interpolator can be set up ahead of time.
-
-    WARNING: compared to scipy's spline interpolation, this method
-    produces edge artifacts when the output coordinates are near
-    the boundaries of the input coordinates. Therefore, it is
-    recommended to interpolate only to coordinates that are well
-    within the input coordinate range. Unfortunately, the specific
-    definition of "well within" changes based on the size of the
-    data, so some testing may be required to get good results.
-
-    Args:
-        x_in:
-            Coordinates of the width dimension of the data
-        y_in:
-            Coordinates of the height dimension of the data. If not
-            specified, it is assumed the 1D interpolation is desired,
-            and so the default value is a Tensor of length 1
-        kx:
-            Degree of spline interpolation along the width dimension.
-            Default is cubic.
-        ky:
-            Degree of spline interpolation along the height dimension.
-            Default is cubic.
-        sx:
-            Regularization factor to avoid singularities during matrix
-            inversion for interpolation along the width dimension. Not
-            to be confused with the ``s`` parameter in scipy's spline
-            methods, which controls the number of knots.
-        sy:
-            Regularization factor to avoid singularities during matrix
-            inversion for interpolation along the height dimension.
-        x_out:
-            Coordinates for the data to be interpolated to along the
-            width dimension. If not specified during initialization,
-            this must be specified during the object call.
-        y_out:
-            Coordinates for the data to be interpolated to along the
-            height dimension. If not specified during initialization,
-            this must be specified during the object call.
-
+    Base class for spline interpolation.
     """
 
-    def __init__(
-        self,
-        x_in: Tensor,
-        y_in: Tensor = None,
-        kx: int = 3,
-        ky: int = 3,
-        sx: float = 0.001,
-        sy: float = 0.001,
-        x_out: Optional[Tensor] = None,
-        y_out: Optional[Tensor] = None,
-    ):
-        super().__init__()
-        if y_in is None:
-            y_in = Tensor([1])
-        self.kx = kx
-        self.ky = ky
-        self.sx = sx
-        self.sy = sy
-        self.register_buffer("x_in", x_in)
-        self.register_buffer("y_in", y_in)
-        self.register_buffer("x_out", x_out)
-        self.register_buffer("y_out", y_out)
-
-        tx, Bx, BxT_Bx = self._compute_knots_and_basis_matrices(x_in, kx, sx)
-        self.register_buffer("tx", tx)
-        self.register_buffer("Bx", Bx)
-        self.register_buffer("BxT_Bx", BxT_Bx)
-
-        ty, By, ByT_By = self._compute_knots_and_basis_matrices(y_in, ky, sy)
-        self.register_buffer("ty", ty)
-        self.register_buffer("By", By)
-        self.register_buffer("ByT_By", ByT_By)
-
-        if self.x_out is not None:
-            Bx_out = self.bspline_basis_natural(x_out, kx, self.tx)
-            self.register_buffer("Bx_out", Bx_out)
-        if self.y_out is not None:
-            By_out = self.bspline_basis_natural(y_out, ky, self.ty)
-            self.register_buffer("By_out", By_out)
-
     def _compute_knots_and_basis_matrices(self, x, k, s):
-        knots = self.generate_natural_knots(x, k)
+        knots = self.generate_fitpack_knots(x, k)
         basis_matrix = self.bspline_basis_natural(x, k, knots)
         identity = torch.eye(basis_matrix.shape[-1])
         B_T_B = basis_matrix.T @ basis_matrix + s * identity
         return knots, basis_matrix, B_T_B
 
-    def generate_natural_knots(self, x: Tensor, k: int) -> Tensor:
+    def generate_fitpack_knots(self, x: Tensor, k: int) -> Tensor:
         """
-        Generates a natural knot sequence for B-spline interpolation.
-        Natural knot sequence means that 2*k knots are added to the beginning
-        and end of datapoints as replicas of first and last datapoint
-        respectively in order to enforce natural boundary conditions,
-        i.e. second derivative = 0.
-        The other n nodes are placed in correspondece of the data points.
+        Generates a knot sequence for B-spline interpolation
+        in the same way as the FITPACK algorithm used by SciPy.
 
         Args:
             x: Tensor of data point positions.
@@ -134,7 +30,17 @@ class SplineInterpolate(torch.nn.Module):
         Returns:
             Tensor of knot positions.
         """
-        return F.pad(x[None], (k, k), mode="replicate")[0]
+        num_knots = x.shape[-1] + k + 1
+        knots = torch.zeros(num_knots, dtype=x.dtype)
+        knots[: k + 1] = x[0]
+        knots[-(k + 1) :] = x[-1]
+
+        # Interior knots are the rolling average of the data points
+        # excluding the first and last points
+        windows = x[1:-1].unfold(dimension=-1, size=k, step=1)
+        knots[k + 1 : -k - 1] = windows.mean(dim=-1)
+
+        return knots
 
     def compute_L_R(
         self,
@@ -142,7 +48,7 @@ class SplineInterpolate(torch.nn.Module):
         t: Tensor,
         d: int,
         m: int,
-    ) -> Tuple[Tensor, Tensor]:
+    ) -> tuple[Tensor, Tensor]:
         """
         Compute the L and R values for B-spline basis functions.
         L and R are respectively the first and second coefficient multiplying
@@ -233,9 +139,6 @@ class SplineInterpolate(torch.nn.Module):
         Returns:
             Tensor containing the kth-order B-spline basis functions
         """
-
-        if len(x) == 1:
-            return torch.eye(1)
         n = x.shape[0]
         m = t.shape[0] - k - 1
 
@@ -255,6 +158,271 @@ class SplineInterpolate(torch.nn.Module):
 
         return b[:, :, -1]
 
+
+class SplineInterpolate1D(SplineInterpolateBase):
+    """
+    Perform 1D spline interpolation based on De Boor's method.
+    It is allowed to have two spatial dimensions, but the second
+    dimension cannot be interpolated along. To interpolate along both
+    dimensions, use :class:`SplineInterpolate2D`.
+
+    Supports batched, multi-channel inputs, so acceptable data
+    shapes are ``(width)``, ``(height, width)``, ``(batch, width)``,
+    ``(batch, height, width)``, ``(batch, channel, width)``, and
+    ``(batch, channel, height, width)``.
+
+    During initialization of this Module, both the desired input
+    and output coordinate Tensors can be specified to allow
+    pre-computation of the B-spline basis matrices, though the only
+    mandatory argument is the coordinates of the data along the
+    ``width`` dimension.
+
+    Unlike scipy's implementation of spline interpolation, the data
+    to be interpolated is not passed until actually calling the
+    object. This is useful for cases where the input and output
+    coordinates are known in advance, but the data is not, so that
+    the interpolator can be set up ahead of time.
+
+    Args:
+        x_in:
+            Coordinates of the width dimension of the data
+        kx:
+            Degree of spline interpolation along the width dimension.
+            Default is cubic.
+        sx:
+            Regularization factor to avoid singularities during matrix
+            inversion for interpolation along the width dimension. Not
+            to be confused with the ``s`` parameter in scipy's spline
+            methods, which controls the number of knots.
+        x_out:
+            Coordinates for the data to be interpolated to along the
+            width dimension. If not specified during initialization,
+            this must be specified during the object call.
+
+    """
+
+    def __init__(
+        self,
+        x_in: Tensor,
+        kx: int = 3,
+        sx: float = 0.0,
+        x_out: Tensor | None = None,
+    ):
+        super().__init__()
+
+        if len(x_in) < kx + 2:
+            raise ValueError(
+                "Input x-coordinates must have at least kx + 2 points."
+            )
+
+        # Ensure that coordinates are floats
+        x_in = x_in.float()
+        x_out = x_out.float() if x_out is not None else None
+
+        self.kx = kx
+        self.sx = sx
+        self.register_buffer("x_in", x_in)
+        self.register_buffer("x_out", x_out)
+
+        tx, Bx, BxT_Bx = self._compute_knots_and_basis_matrices(x_in, kx, sx)
+        self.register_buffer("tx", tx)
+        self.register_buffer("Bx", Bx)
+        self.register_buffer("BxT_Bx", BxT_Bx)
+
+        if self.x_out is not None:
+            x_clamped = torch.clamp(x_out, tx[kx], tx[-kx - 1])
+            Bx_out = self.bspline_basis_natural(x_clamped, kx, self.tx)
+            self.register_buffer("Bx_out", Bx_out)
+
+    def spline_fit_natural(self, Z):
+        # Adding batch/channel dimension handling
+        # Bx @ Z
+        BxT_Z = torch.einsum("ij,bchj->bchi", self.Bx.T, Z)
+        # (BxT @ Bx)^-1 @ (BxT @ Z) = Bx^-1 @ Z
+        C = torch.linalg.solve(self.BxT_Bx, BxT_Z.unsqueeze(-1))
+        return C.squeeze(-1)
+
+    def evaluate_spline(self, C: Tensor):
+        """
+        Evaluate a bivariate spline on a grid of x and y points.
+
+        Args:
+            C: Coefficient tensor of shape (batch_size, mx, my).
+
+        Returns:
+            Z_interp: Interpolated values at the grid points.
+        """
+        # Perform matrix multiplication using einsum to get Z_interp
+        return torch.einsum("ij,bchj->bchi", self.Bx_out, C)
+
+    def _validate_inputs(self, Z, x_out):
+        if x_out is None and self.x_out is None:
+            raise ValueError(
+                "Output x-coordinates were not specified in either object "
+                "creation or in forward call"
+            )
+
+        dims = len(Z.shape)
+        if dims > 4:
+            raise ValueError("Input data has more than 4 dimensions")
+
+        if Z.shape[-1] != len(self.x_in):
+            raise ValueError(
+                "The spatial dimensions of the data tensor do not match "
+                "the given input dimensions. "
+                f"Expected {len(self.x_in)}, but got {Z.shape[-1]}"
+            )
+
+        # Expand Z to have a batch, channel, and height dimension if needed
+        while len(Z.shape) < 4:
+            Z = Z.unsqueeze(0)
+
+        return Z
+
+    def forward(
+        self,
+        Z: Tensor,
+        x_out: Tensor | None = None,
+    ) -> Tensor:
+        """
+        Compute the interpolated data
+
+        Args:
+            Z:
+                Tensor of data to be interpolated. Must be between 2 and 4
+                dimensions. The shape of the tensor must agree with the
+                input coordinates given on initialization.
+            x_out:
+                Coordinates to interpolate the data to along the width
+                dimension. Overrides any value that was set during
+                initialization.
+
+        Returns:
+            A 4D tensor with shape ``(batch, channel, height, width)``.
+            Depending on the input data shape, many of these dimensions
+            may have length 1.
+        """
+
+        Z = self._validate_inputs(Z, x_out)
+
+        if x_out is not None:
+            x_out = x_out.float()
+            x_clamped = torch.clamp(
+                x_out, self.tx[self.kx], self.tx[-self.kx - 1]
+            )
+            self.Bx_out = self.bspline_basis_natural(
+                x_clamped, self.kx, self.tx
+            )
+
+        coef = self.spline_fit_natural(Z)
+        Z_interp = self.evaluate_spline(coef)
+        return Z_interp
+
+
+class SplineInterpolate2D(SplineInterpolateBase):
+    """
+    Perform 2D spline interpolation based on De Boor's method.
+    Supports batched, multi-channel inputs, so acceptable data
+    shapes are ``(height, width)``, ``(batch, height, width)``,
+    and ``(batch, channel, height, width)``.
+
+    During initialization of this Module, both the desired input
+    and output coordinate Tensors can be specified to allow
+    pre-computation of the B-spline basis matrices, though the only
+    mandatory arguments are the input coordinates.
+
+    Unlike scipy's implementation of spline interpolation, the data
+    to be interpolated is not passed until actually calling the
+    object. This is useful for cases where the input and output
+    coordinates are known in advance, but the data is not, so that
+    the interpolator can be set up ahead of time.
+
+    Args:
+        x_in:
+            Coordinates of the width dimension of the data
+        y_in:
+            Coordinates of the height dimension of the data.
+        kx:
+            Degree of spline interpolation along the width dimension.
+            Default is cubic.
+        ky:
+            Degree of spline interpolation along the height dimension.
+            Default is cubic.
+        sx:
+            Regularization factor to avoid singularities during matrix
+            inversion for interpolation along the width dimension. Not
+            to be confused with the ``s`` parameter in scipy's spline
+            methods, which controls the number of knots.
+        sy:
+            Regularization factor to avoid singularities during matrix
+            inversion for interpolation along the height dimension.
+        x_out:
+            Coordinates for the data to be interpolated to along the
+            width dimension. If not specified during initialization,
+            this must be specified during the object call.
+        y_out:
+            Coordinates for the data to be interpolated to along the
+            height dimension. If not specified during initialization,
+            this must be specified during the object call.
+
+    """
+
+    def __init__(
+        self,
+        x_in: Tensor,
+        y_in: Tensor,
+        kx: int = 3,
+        ky: int = 3,
+        sx: float = 0.0,
+        sy: float = 0.0,
+        x_out: Tensor | None = None,
+        y_out: Tensor | None = None,
+    ):
+        super().__init__()
+
+        if len(x_in) < kx + 2:
+            raise ValueError(
+                "Input x-coordinates must have at least kx + 2 points."
+            )
+        if len(y_in) < ky + 2:
+            raise ValueError(
+                "Input y-coordinates must have at least ky + 2 points."
+            )
+
+        # Ensure that coordinates are floats
+        x_in = x_in.float()
+        y_in = y_in.float()
+        x_out = x_out.float() if x_out is not None else None
+        y_out = y_out.float() if y_out is not None else None
+
+        self.kx = kx
+        self.ky = ky
+        self.sx = sx
+        self.sy = sy
+        self.register_buffer("x_in", x_in)
+        self.register_buffer("y_in", y_in)
+        self.register_buffer("x_out", x_out)
+        self.register_buffer("y_out", y_out)
+
+        tx, Bx, BxT_Bx = self._compute_knots_and_basis_matrices(x_in, kx, sx)
+        self.register_buffer("tx", tx)
+        self.register_buffer("Bx", Bx)
+        self.register_buffer("BxT_Bx", BxT_Bx)
+
+        ty, By, ByT_By = self._compute_knots_and_basis_matrices(y_in, ky, sy)
+        self.register_buffer("ty", ty)
+        self.register_buffer("By", By)
+        self.register_buffer("ByT_By", ByT_By)
+
+        if self.x_out is not None:
+            x_clamped = torch.clamp(x_out, tx[kx], tx[-kx - 1])
+            Bx_out = self.bspline_basis_natural(x_clamped, kx, self.tx)
+            self.register_buffer("Bx_out", Bx_out)
+        if self.y_out is not None:
+            y_clamped = torch.clamp(y_out, ty[ky], ty[-ky - 1])
+            By_out = self.bspline_basis_natural(y_clamped, ky, self.ty)
+            self.register_buffer("By_out", By_out)
+
     def bivariate_spline_fit_natural(self, Z):
         # Adding batch/channel dimension handling
         # ByT @ Z @ BxW
@@ -285,29 +453,16 @@ class SplineInterpolate(torch.nn.Module):
             )
 
         if y_out is None and self.y_out is None:
-            y_out = self.y_in
+            raise ValueError(
+                "Output y-coordinates were not specified in either object "
+                "creation or in forward call"
+            )
 
         dims = len(Z.shape)
         if dims > 4:
             raise ValueError("Input data has more than 4 dimensions")
-
-        if len(self.y_in) > 1 and dims == 1:
-            raise ValueError(
-                "An input y-coordinate array with length greater than 1 "
-                "was given, but the input data is 1-dimensional. Expected "
-                "input data to be at least 2-dimensional"
-            )
-
-        # Expand Z to have 4 dimensions
-        # There are 6 valid input shapes: (w), (b, w), (b, c, w),
-        # (h, w), (b, h, w), and (b, c, h, w).
-
-        # If the input y coordinate array has length 1,
-        # assume the first dimension(s) are batch dimensions
-        # and that no height dimension is included in Z
-        idx = -2 if len(self.y_in) == 1 else -3
-        while len(Z.shape) < 4:
-            Z = Z.unsqueeze(idx)
+        if dims < 2:
+            raise ValueError("Input data has fewer than 2 dimensions")
 
         if Z.shape[-2:] != torch.Size([len(self.y_in), len(self.x_in)]):
             raise ValueError(
@@ -317,24 +472,26 @@ class SplineInterpolate(torch.nn.Module):
                 f"[{Z.shape[-2]}, {Z.shape[-1]}]"
             )
 
+        # Expand Z to have a batch and channel dimension if needed
+        while len(Z.shape) < 4:
+            Z = Z.unsqueeze(0)
+
         return Z, y_out
 
     def forward(
         self,
         Z: Tensor,
-        x_out: Optional[Tensor] = None,
-        y_out: Optional[Tensor] = None,
+        x_out: Tensor | None = None,
+        y_out: Tensor | None = None,
     ) -> Tensor:
         """
         Compute the interpolated data
 
         Args:
             Z:
-                Tensor of data to be interpolated. Must be between 1 and 4
+                Tensor of data to be interpolated. Must be between 2 and 4
                 dimensions. The shape of the tensor must agree with the
-                input coordinates given on initialization. If ``y_in`` was
-                not specified during initialization, it is assumed that
-                Z does not have a height dimension.
+                input coordinates given on initialization.
             x_out:
                 Coordinates to interpolate the data to along the width
                 dimension. Overrides any value that was set during
@@ -353,9 +510,21 @@ class SplineInterpolate(torch.nn.Module):
         Z, y_out = self._validate_inputs(Z, x_out, y_out)
 
         if x_out is not None:
-            self.Bx_out = self.bspline_basis_natural(x_out, self.kx, self.tx)
+            x_out = x_out.float()
+            x_clamped = torch.clamp(
+                x_out, self.tx[self.kx], self.tx[-self.kx - 1]
+            )
+            self.Bx_out = self.bspline_basis_natural(
+                x_clamped, self.kx, self.tx
+            )
         if y_out is not None:
-            self.By_out = self.bspline_basis_natural(y_out, self.ky, self.ty)
+            y_out = y_out.float()
+            y_clamped = torch.clamp(
+                y_out, self.ty[self.ky], self.ty[-self.ky - 1]
+            )
+            self.By_out = self.bspline_basis_natural(
+                y_clamped, self.ky, self.ty
+            )
 
         coef = self.bivariate_spline_fit_natural(Z)
         Z_interp = self.evaluate_bivariate_spline(coef)

ml4gw/transforms/transform.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import torch
 
 from ..spectral import spectral_density
@@ -20,8 +18,8 @@ class FittableTransform(torch.nn.Module):
     def _check_built(self):
         if not self.built:
             raise ValueError(
-                "Must fit parameters of {} transform to data "
-                "before calling forward step".format(self.__class__.__name__)
+                f"Must fit parameters of {self.__class__.__name__} transform "
+                "to data before calling forward step"
             )
 
     def __call__(self, *args, **kwargs):
@@ -47,8 +45,8 @@ class FittableSpectralTransform(FittableTransform):
         x: TimeSeries1to3d,
         sample_rate: float,
         num_freqs: int,
-        fftlength: Optional[float] = None,
-        overlap: Optional[float] = None,
+        fftlength: float | None = None,
+        overlap: float | None = None,
     ) -> FrequencySeries1to3d:
         # if we specified an FFT length, convert
         # the (assumed) time-domain data to the

ml4gw/transforms/waveforms.py

@@ -1,5 +1,3 @@
-from typing import List, Optional
-
 import torch
 from jaxtyping import Float
 from torch import Tensor
@@ -13,7 +11,7 @@ from ..types import BatchTensor
 class WaveformSampler(torch.nn.Module):
     def __init__(
         self,
-        parameters: Optional[Float[Tensor, "batch num_params"]] = None,
+        parameters: Float[Tensor, "batch num_params"] | None = None,
         **polarizations: Float[Tensor, "batch time"],
     ):
         super().__init__()
@@ -24,10 +22,8 @@ class WaveformSampler(torch.nn.Module):
         for polarization, tensor in polarizations.items():
             if num_waveforms is not None and len(tensor) != num_waveforms:
                 raise ValueError(
-                    "Polarization {} has {} waveforms "
-                    "associated with it, expected {}".format(
-                        polarization, len(tensor), num_waveforms
-                    )
+                    f"Polarization {polarization} has {len(tensor)} waveforms "
+                    f"associated with it, expected {num_waveforms}"
                 )
             elif num_waveforms is None:
                 num_waveforms = tensor.shape[0]
@@ -36,10 +32,8 @@ class WaveformSampler(torch.nn.Module):
 
         if parameters is not None and len(parameters) != num_waveforms:
             raise ValueError(
-                "Waveform parameters has {} waveforms "
-                "associated with it, expected {}".format(
-                    len(parameters), num_waveforms
-                )
+                f"Waveform parameters has {len(parameters)} waveforms "
+                f"associated with it, expected {num_waveforms}"
             )
         self.num_waveforms = num_waveforms
         self.parameters = parameters
@@ -48,9 +42,8 @@ class WaveformSampler(torch.nn.Module):
         # TODO: should we allow sampling with replacement?
         if N > self.num_waveforms:
             raise ValueError(
-                "Requested {} waveforms, but only {} are available".format(
-                    N, self.num_waveforms
-                )
+                f"Requested {N} waveforms, but only {self.num_waveforms} are "
+                "available"
             )
         # TODO: do we still really want this behavior here when a
         # user can do this without instantiating a WaveformSampler?
@@ -67,7 +60,7 @@ class WaveformSampler(torch.nn.Module):
 
 
 class WaveformProjector(torch.nn.Module):
-    def __init__(self, ifos: List[str], sample_rate: float):
+    def __init__(self, ifos: list[str], sample_rate: float):
         super().__init__()
         tensors, vertices = gw.get_ifo_geometry(*ifos)
         self.sample_rate = sample_rate