"""API contract for L2D data and waveform representations.
This module defines the protocol-based contracts for representing gravitational wave
signals in different domains and on different grids.
For conceptual diagrams and shape conventions, see :doc:`overview`.
"""
from typing import Any, Iterator, Literal, Protocol, runtime_checkable
from .array import Array
[docs]
@runtime_checkable
class LinspaceLike(Protocol):
"""Protocol for custom linspace-like objects that can be used as axes in representations.
Classes implementing this protocol replace the standard 1D array axis with more memory efficiency.
"""
@property
def start(self) -> float:
"""Start point of the linspace."""
...
@property
def step(self) -> float:
"""Step size (cadence) of the linspace."""
...
[docs]
def __len__(self) -> int:
"""Return the number of points in the linspace."""
...
Axis = Array[Any, Any] | LinspaceLike
"""`Axis` can be either a standard 1D array following the Python Array API standard,
or a custom linspace-like object that implements the :class:`LinspaceLike`
protocol for memory-efficient axis representation.
"""
Grid1D = tuple[Axis]
Grid2D = tuple[Axis, Axis]
[docs]
@runtime_checkable
class Grid2DSparse(Protocol):
"""Protocol for sparse 2D (time-frequency) grids.
A sparse grid selects a subset of points from the Cartesian product of two axes.
Implements tuple-like indexing to access axes (compatible with dense Grid2D interface)
and provides an ``indices`` property mapping sparse points to dense grid coordinates.
"""
[docs]
def __getitem__(self, __index: int) -> Axis:
"""Get the axis for the specified dimension index (0 for frequency, 1 for time)."""
...
[docs]
def __len__(self) -> Literal[2]:
"""Return the number of dimensions (always 2 for :class:`Grid2DSparse`)."""
...
@property
def indices(self) -> Array[Any, Any]:
"""Array of shape (n_sparse, 2) containing the indices of the non-zero points in the 2D grid."""
...
Grid = Grid1D | Grid2D | Grid2DSparse
"""Physical domain of representations.
It can be one of the following:
- :class:`Grid1D`: A tuple containing a single `Axis`
- :class:`Grid2D`: A tuple containing two `Axis` objects
- :class:`Grid2DSparse`: A protocol for sparse 2D grids,
where the non-zero points are specified by a `indices`
array of shape `(n_sparse, 2)`.
"""
Domain = Literal["time", "frequency", "time-frequency"]
[docs]
class Representation[DomainT: Domain, GridT: Grid, KindT: str | None](Protocol):
"""API contract for representation of gravitational wave signals."""
@property
def entries(self) -> Array[Any, Any]:
"""Multi-dimensional array following the Python Array API standard.
Shape convention: ``(n_batches, n_channels, n_harmonics, n_features, *grid_like)``
- ``n_batches``: Independent signal realizations
- ``n_channels``: Detector channels (e.g., 1 for single, 3 for TDI X, Y, Z)
- ``n_harmonics``: Harmonic modes (1 for single-mode, >1 for multi-mode)
- ``n_features``: Features per grid point (1 for scalar, >1 for multivariate)
- Most common case: time-series of scalar or Fourier coefficients (``n_features=1``)
- Multivariate example: frequency-domain series with both amplitude and phase (``n_features=2``)
- ``*grid_like``: Remaining dimensions determined by grid type
- 1D grid: time-domain series or frequency-domain series ``(n_grid,)``
- 2D dense grid: dense time-frequency representations ``(n_freq, n_time)``
- 2D sparse grid: sparse time-frequency representations ``(n_sparse,)``
where `n_sparse` is the number of non-zero points. This can be read
from `grid.indices` which is of shape ``(n_sparse, 2)``
.. note::
Never squeeze dimensions, even if trivial (e.g., ``n_channels=1``).
The presence of reserved dimensions for channels and harmonics is a design choice
to support efficient cross-channel and cross-harmonic operations where applicable.
It should not be taken to imply that all representations must have multiple channels or harmonics
(many will have just one), nor that it is priledged to populate these dimensions by
all means. In fact, when signals are homogeneous acros harmonics (common in waveform
generation), we should use mapping containers keyed by harmonic mode and valued
by representations with ``n_harmonics=1`` (i.e., ``shape[2] == 1``) rather than
forcing them into a single array with ``n_harmonics > 1`` (see :class:`HarmonicWaveform`
and :class:`HarmonicProjectedWaveform`). Even more so, when signals are homogeneous across channels
(common in detector response and recorded data), mapping containers keyed by channel name and
valued by representations with ``n_channels=1`` (i.e., ``shape[1] == 1``) also provide more
semantic clarity (see :class:`Data`, :class:`TransformedData`, and :class:`ProjectedWaveform`).
"""
...
@property
def grid(self) -> GridT:
"""Grid specification defining axis points.
1D grid for time-domain or frequency-domain series; 2D for time-frequency.
"""
...
@property
def domain(self) -> DomainT:
"""Physical domain of the representation.
One of: ``'time'``, ``'frequency'``, or ``'time-frequency'``.
See :class:`TDRepresentation`, :class:`FDRepresentation`, :class:`DenseTFRepresentation`
and :class:`SparseTFRepresentation`.
"""
...
@property
def kind(self) -> KindT:
"""Optional semantic kind for domain-specific variants.
Examples: ``'wavelet'`` for time-frequency representations.
``None`` for standard representations (e.g., scalar time/frequency series).
"""
...
[docs]
class TDRepresentation[KindT: str | None](
Representation[Literal["time"], Grid1D, KindT], Protocol
):
"""API contract for time-domain representations of gravitational wave signals."""
[docs]
class FDRepresentation[KindT: str | None](
Representation[Literal["frequency"], Grid1D, KindT], Protocol
):
"""API contract for frequency-domain representations of gravitational wave signals."""
[docs]
class DenseTFRepresentation[KindT: str | None](
Representation[Literal["time-frequency"], Grid2D, KindT], Protocol
):
"""API contract for dense time-frequency representations of gravitational wave signals.
The entries array has shape ``(n_batches, n_channels, n_harmonics, n_features, n_freq, n_time)``
representing the full Cartesian product of frequency and time axes.
"""
[docs]
class SparseTFRepresentation[KindT: str | None](
Representation[Literal["time-frequency"], Grid2DSparse, KindT], Protocol
):
"""API contract for sparse time-frequency representations of gravitational wave signals.
For sparse representations, entries are flattened along the time-frequency dimensions.
The actual TF point coordinates are recovered from ``grid.indices``.
The entries array has shape ``(n_batches, n_channels, n_harmonics, n_features, n_sparse)``
where ``n_sparse`` is the number of non-zero time-frequency points.
Example
-------
Sparse time-frequency representation (only 5000 active points out of 100×500=50000)::
# entries shape: (1, 1, 1, 1, 5000)
# grid.indices: (5000, 2) with (freq_idx, time_idx) pairs
# Only 10% of the dense grid is used, enabling memory-efficient storage
"""
class _ChannelMapping[DomainT: Domain, GridT: Grid, KindT: str | None](Protocol):
"""Base protocol for channel-keyed representation mappings.
This is an internal base protocol that defines the common API for containers
that map channel names to :class:`Representation` objects. All channels share
the same domain, grid, and kind.
"""
def __getitem__(self, key: str) -> Representation[DomainT, GridT, KindT]:
"""Get a channel representation by name.
The :class:`Representation` objects returned by this method
must have ``n_channels=1`` (i.e. ``shape[1] == 1``) and ``n_harmonics=1`` (i.e. ``shape[2] == 1``).
.. note::
Though not strictly required, implementations are encouraged
to return views of the same underlying array rather than storing
independent arrays per channel.
"""
...
def __iter__(self) -> Iterator[str]:
"""Iterate over channel names."""
...
def __len__(self) -> int:
"""Return the number of channels."""
...
@property
def domain(self) -> DomainT:
"""Physical domain shared by all channels.
All :class:`Representation` objects must share the same domain.
"""
...
@property
def grid(self) -> GridT:
"""Grid specification shared by all channels.
All :class:`Representation` objects must share the same grid.
"""
...
@property
def kind(self) -> KindT:
"""Semantic kind shared by all channels.
All :class:`Representation` objects must share the same kind.
"""
...
@property
def channel_names(self) -> tuple[str, ...]:
"""Names of all channels and their order."""
...
def get_kernel(self, backend: str | None = None) -> Array[Any, Any]:
"""Return an array of the conventional shape ``(n_batches, n_channels, 1, n_features, *grid_like)``
for downstream processing (e.g., by noise models to compute inner products).
Arguments
---------
backend: Optional string specifying the desired array library for the returned array.
If ``None``, use the underlying array library of the entries.
.. note::
This method can be trivially implemented if the underlying data entries are already stored in the conventional shape.
Otherwise, it can be implemented by stacking the representations of individual channels along the channel dimension.
"""
...
[docs]
@runtime_checkable
class Data(_ChannelMapping[Literal["time"], Grid1D, None], Protocol):
"""API contract for data containers of gravitational wave signals.
:class:`Data` objects represent the output from (pre-processed) L1 Data, which is the source of truth
for the entire L2D pipeline. They should not be modified by L2D analysis.
Maps channel names to time-domain :class:`TDRepresentation` objects. All channels share
the same time grid.
"""
[docs]
def __getitem__(self, key: str) -> TDRepresentation[None]:
"""Get a channel representation by name.
See :meth:`_ChannelMapping.__getitem__`.
"""
...
@property
def domain(self) -> Literal["time"]:
"""Physical domain (always 'time' for `Data`).
See :attr:`_ChannelMapping.domain`.
"""
...
@property
def grid(self) -> Grid1D:
"""1D time grid specification shared by all channels.
See :attr:`_ChannelMapping.grid`.
"""
...
@property
def times(self) -> Axis:
"""Time axis."""
...