"""Representation types."""
from __future__ import annotations
import abc
import copy as _copy
import logging
from types import EllipsisType, ModuleType
from typing import (
TYPE_CHECKING,
Any,
Final,
Literal,
Protocol,
Self,
cast,
overload,
)
import array_api_compat as xpc
from .._l2d_interface import contract
from ..utils import deprecated
from .misc import (
AnyArray,
AnyAxis,
AnyGrid,
Axis,
AxLike,
Domain,
Grid1D,
Grid2D,
Grid2DCartesian,
Grid2DSparse,
Interpolator,
Linspace,
axis,
build_grid2d,
)
if TYPE_CHECKING:
class Representation[GridT: AnyGrid](
contract.Representation[Domain, GridT, str | None],
Protocol,
):
"""Protocol for any representation type."""
entries: AnyArray # type: ignore[assignment] # Necessary due to missing data array API
@property
def grid(self) -> GridT: ... # noqa: D102
def __init__(
self,
grid: AnyGrid,
entries: AnyArray,
): ...
def create_like(self, entries: Any) -> Self:
"""Create a new instance with the same grid and type but different entries.""" # noqa: E501
...
def __xp__(self, api_version: str | None = None) -> ModuleType:
"""Return the array API namespace of the representation."""
...
from .. import utils
from . import _mixins, tapering
log = logging.getLogger(__name__)
_slice = slice # Alias for slice
def _get_entry_grid_shape(entries: AnyArray):
return entries.shape[4:] # Remove batch, channels, harmonics, features dimensions
def _get_full_slice[T: (slice, AnyArray)](
grid_slices: tuple[T, ...], /
) -> tuple[EllipsisType, *tuple[T, ...]]:
"""Return the slice tuple for the canonical entries array given the grid slices."""
return (..., *grid_slices)
def _subset_grid_2d_sparse[Axis0: "AnyAxis", Axis1: "AnyAxis"](
grid: Grid2DSparse[Axis0, Axis1],
entries: AnyArray,
slices: tuple[slice, slice],
) -> tuple[Grid2DSparse[Axis0, Axis1], AnyArray]:
_sparse_idx = grid.indices
xp = xpc.get_namespace(_sparse_idx)
# Only keep indices between the subset slices
mask = xp.ones(_sparse_idx.shape[0], dtype=bool)
for i, slc in enumerate(slices):
start, stop, _ = slc.indices(len(grid[i]))
mask = mask & (_sparse_idx[:, i] >= start) & (_sparse_idx[:, i] < stop)
_new_sparse_idx = _sparse_idx[mask]
_offsets = xp.asarray(
[slc.indices(len(grid[i]))[0] for i, slc in enumerate(slices)],
dtype=_new_sparse_idx.dtype,
)
_new_sparse_idx = _new_sparse_idx - _offsets
new_grid = Grid2DSparse[Axis0, Axis1](
grid.axis0[slices[0]],
grid.axis1[slices[1]],
sparse_indices=_new_sparse_idx,
)
new_entries = entries[..., mask]
return new_grid, new_entries
def _take_subset[GridT: AnyGrid](
grid: GridT,
entries: AnyArray,
slices: tuple[slice, ...],
) -> tuple[GridT, AnyArray]:
"""Return the grid and entries subset given the slices on each grid dimension."""
if len(grid) != len(slices):
msg = (
f"Number of slices {len(slices)} "
f"does not match number of grid dimensions {len(grid)}."
)
raise ValueError(msg)
# Slice each grid dimension
_grid = tuple(g[s] for g, s in zip(grid, slices, strict=True))
_entries = entries[_get_full_slice(slices)]
return cast("GridT", _grid), _entries
def _get_subset_slice(
grid1d: AnyAxis,
/,
*,
interval: tuple[float, float] | None = None,
slice: _slice | None = None,
) -> _slice:
if interval is None:
# Note that slice(None) is not None
if slice is None:
return _slice(None)
# Otherwise we use the input slice
else:
if slice is not None:
msg = "Only one of `interval` and `slice` should be provided."
raise ValueError(msg)
slice = utils.get_subset_slice(grid1d.asarray(), interval[0], interval[1])
# slice is always a slice object at this point
return slice
def _set_value(
entries: AnyArray, slice: tuple[_slice | EllipsisType, ...], value: Any
) -> None:
try:
entries[slice] = value
except TypeError:
entries = cast("AnyArray", entries.at[slice].set(value)) # type: ignore[assignment, union-attr]
def _get_axis_onset(axis: AnyAxis) -> float:
return axis.start
def _get_axis_end(axis: AnyAxis) -> float:
return axis.stop
class _InitMixin[GridT: AnyGrid](abc.ABC):
@property
def grid(self) -> GridT:
return cast("GridT", self._grid)
def __xp__(self, api_version: str | None = None):
return xpc.get_namespace(self.entries, api_version=api_version)
@property
def xp(self):
return self.__xp__()
def __init__(
self,
grid: AnyGrid, # on purpose not GridT to allow more flexible input types
entries: AnyArray,
):
self._grid: AnyGrid = grid
self.entries: AnyArray = entries
def __repr__(self) -> str:
return (
f"{type(self).__name__}(grid={self.grid!r}, "
f"entries={self.entries!r}, {self.entries.dtype!r})"
)
@property
@abc.abstractmethod
def domain(self) -> Domain:
"""Physical domain of the representation."""
@property
@abc.abstractmethod
def kind(self) -> str | None:
"""Optional semantic kind of representation."""
@property
def n_batches(self) -> int | None:
"""Return the number of batches."""
return self.entries.shape[0]
@property
def n_channels(self) -> int | None:
"""Return the number of channels."""
return self.entries.shape[1]
@property
def n_harmonics(self) -> int | None:
"""Return the number of harmonics."""
return self.entries.shape[2]
@property
def n_features(self) -> int | None:
"""Return the number of features."""
return self.entries.shape[3]
@property
def _grid_shape(self) -> tuple[int | None, ...]:
"""Return the shape of the grid dimensions."""
return _get_entry_grid_shape(self.entries)
def get_kernel(self) -> AnyArray:
"""Return the entries of the representation."""
return self.entries
class _Subset1DMixin[GridT: "Grid1D[AnyAxis]"](_InitMixin[GridT], abc.ABC):
def get_subset(
self,
*,
interval: tuple[float, float] | None = None,
slice: _slice | None = None,
copy: bool = True,
) -> Self:
"""Return the subset as a new instance."""
_slice = _get_subset_slice(self.grid[0], interval=interval, slice=slice)
grid, entries = _take_subset(self.grid, self.entries, (_slice,))
entries = _copy.copy(entries) if copy else entries
return type(self)(grid=grid, entries=entries)
def __getitem__(self, slice: _slice) -> Self:
"""Return the view of a subset of the series."""
return self.get_subset(slice=slice, copy=False)
def __setitem__(self, slice: _slice, value: Any) -> None:
"""Set entries at slice location."""
# NOTE this method does not check the compatibility of the grids,
# and assumes that the slice is correct.
_set_value(self.entries, _get_full_slice((slice,)), value)
def _embed_entries_to_grid_2d_sparse[
Axis0: "AnyAxis",
Axis1: "AnyAxis",
](
source_grid: Grid2DSparse[AnyAxis, AnyAxis],
source_entries: AnyArray,
embedding_grid: Grid2D[Axis0, Axis1],
*,
known_slices: tuple[slice, ...] | None = None,
) -> tuple[Grid2DSparse[Axis0, Axis1], AnyArray]:
_sparse_idx = source_grid.indices
xp = xpc.get_namespace(_sparse_idx)
# The embedding amounts to compute new sparse indices
if known_slices is not None:
support_slices = known_slices
else:
support_slices = tuple(
utils.get_subset_slice(
eg.asarray(),
min=float(sg.asarray()[0]),
max=float(sg.asarray()[-1]),
)
for sg, eg in zip(source_grid, embedding_grid, strict=True)
)
_new_sparse_idx = [
_sparse_idx[:, i] + _slc.start for i, _slc in enumerate(support_slices)
]
new_sparse_idx = xp.stack(_new_sparse_idx, axis=1)
new_grid = Grid2DSparse[Axis0, Axis1](
embedding_grid[0],
embedding_grid[1],
sparse_indices=new_sparse_idx,
)
return new_grid, source_entries
class _ArithmeticReprOnGrid[GridT: "AnyGrid"](
_mixins.BinaryUnaryOpMixin[AnyArray],
_InitMixin[GridT],
abc.ABC,
):
# Provides implementations for arithmetic operations
def create_like(self, entries: AnyArray):
"""Create a new instance with the same grid as the current one."""
return type(self)(grid=self.grid, entries=entries)
def __xp__(self, api_version: str | None = None):
return xpc.get_namespace(self.entries, api_version=api_version)
@property
def xp(self):
return self.__xp__()
def _unwrap(self, other: object):
if hasattr(other, "grid") and hasattr(other, "entries"):
other = cast("_ArithmeticReprOnGrid[GridT]", other)
if _mixins.check_grid_compatibility(self.grid, other.grid):
return other.entries
msg = f"Grid mismatch: expected {self.grid}, got {other.grid}."
raise ValueError(msg)
return other
def add(
self, other: Self, slice: tuple[_slice, ...], *, inplace: bool = False
) -> Self:
"""Add another series on a sub-grid with known slice.
This method adds another series on a sub-grid of the current series
with a known slice, which is used to select the entries of the current
series to be added with.
If `inplace` is True, the current series is modified in place
and returned (equivalent to calling :meth:`.iadd`). Otherwise,
a new series is returned with the result of the addition.
Default is False.
See Also
--------
:meth:`.iadd`
"""
if inplace:
return self.iadd(other, slice)
self_copy = self.create_like(_copy.copy(self.entries))
self_copy.iadd(other, slice)
return self_copy
def iadd(self, other: Self, slice: tuple[_slice | EllipsisType, ...]) -> Self:
"""Add another series on a sub-grid with known slice in place.
See Also
--------
:meth:`.add`
"""
try:
_set_value(self.entries, slice, self.entries[slice] + other.entries)
except ValueError as e:
msg = (
"You may want to first embed the series instances to super-grids "
"before adding them, if their grids are not compatible."
)
raise ValueError(msg) from e
return self
def __iadd__(self, other: object) -> Self:
"""Add another series in place.
Note
----
Compared to :meth:`.iadd`, this method computes automatically the slice
of the subgrid to apply with a generic algorithm. This is not the most
efficient in some cases. If you have a specialised and more efficient
way of computing the slice, you should use the :meth:`.iadd` method
instead.
See Also
--------
:meth:`.iadd`
:meth:`.add`
:meth:`.__add__`
"""
if isinstance(other, type(self)):
_slices = tuple(
utils.get_subset_slice(
self.grid[idx].asarray(), _axis.start, _axis.stop
)
for idx, _axis in enumerate(other.grid)
)
if len(self.grid) < len(other.grid):
msg = (
"The current series has fewer grid dimensions "
"than the other series. "
"In-place addition is not supported in this case. "
"You may want to first embed the current series to a super-grid "
"that is compatible with the other series before adding them."
)
raise ValueError(msg)
return self.iadd(other, slice=_get_full_slice(_slices))
return super().__iadd__(other)
class _Uniform1DMixin(abc.ABC):
@property
@abc.abstractmethod
def grid(self) -> Grid1D[Axis[Linspace]]: ...
@property
def resolution(self) -> float:
return self.grid[0].ax.step
def _validate_shape(entries: AnyArray, expected_shape: tuple[int | None, ...]) -> None:
if entries.shape != expected_shape:
msg = (
"Invalid shape for `entries`. "
f"Expected {expected_shape}, got {entries.shape}."
)
raise ValueError(msg)
@overload
def frequency_series(
frequencies: Linspace,
entries: AnyArray,
) -> UniformFrequencySeries: ...
@overload
def frequency_series(
frequencies: Axis[Linspace],
entries: AnyArray,
) -> UniformFrequencySeries: ...
@overload
def frequency_series[AxisT: "AnyAxis"](
frequencies: AxisT,
entries: AnyArray,
) -> FrequencySeries[AxisT]: ...
@overload
def frequency_series(
frequencies: AnyArray,
entries: AnyArray,
) -> FrequencySeries[Axis[AnyArray]]: ...
[docs]
def frequency_series[AxisT: "AnyAxis"](
frequencies: AnyAxis | AxLike,
entries: AnyArray,
):
"""Build an :class:`~types.FrequencySeries` or a :class:`~types.UniformFrequencySeries`.
Parameters
----------
frequencies:
Either a :class:`~typed_lisa_toolkit.types.Linspace` or a 1D
:class:`array <typed_lisa_toolkit.types.misc.Array>` of positive frequencies.
entries: :class:`~typed_lisa_toolkit.types.misc.Array`
An array of shape ``(n_batch, 1, 1, 1, Nf)`` where ``Nf`` is the size of ``frequencies``.
Note
----
See the :external+l2d-interface:ref:`general description <shape_convention>`
of the shape convention for `entries`.
Example
-------
.. code-block:: python
import jax.numpy as jnp
import typed_lisa_toolkit as tlt
fs1 = tlt.frequency_series(
tlt.linspace(0, 1, 10), jnp.ones((1, 1, 1, 1, 10))
) # UniformFrequencySeries
fs2 = tlt.frequency_series(
jnp.array([0, 0.1, 0.3, 0.6, 1]), jnp.ones((1, 1, 1, 1, 5))
) # FrequencySeries[Array[Any]]
""" # noqa: E501
_validate_shape(
entries,
(entries.shape[0], 1, 1, 1, len(frequencies)),
)
if isinstance(frequencies, Linspace):
return UniformFrequencySeries((axis(frequencies),), entries)
if isinstance(frequencies, Axis):
if isinstance(frequencies.ax, Linspace):
return UniformFrequencySeries((frequencies,), entries)
return FrequencySeries[AxisT]((frequencies,), entries)
return FrequencySeries[AxisT]((axis(frequencies),), entries)
@overload
def time_series(
times: Linspace,
entries: AnyArray,
) -> UniformTimeSeries: ...
@overload
def time_series(
times: Axis[Linspace],
entries: AnyArray,
) -> UniformTimeSeries: ...
@overload
def time_series(
times: AnyArray,
entries: AnyArray,
) -> TimeSeries[Axis[AnyArray]]: ...
@overload
def time_series[AxisT: "AnyAxis"](
times: AxisT,
entries: AnyArray,
) -> TimeSeries[AxisT]: ...
[docs]
def time_series[AxisT: "AnyAxis"](
times: AnyAxis | AxLike,
entries: AnyArray,
):
"""Build a :class:`~types.TimeSeries` or a :class:`~types.UniformTimeSeries`.
Parameters
----------
times:
Either a :class:`~typed_lisa_toolkit.types.Linspace` or a 1D
:class:`array <typed_lisa_toolkit.types.misc.Array>` of time points.
entries: :class:`~typed_lisa_toolkit.types.misc.Array`
An array of shape ``(n_batch, 1, 1, 1, Nt)``
where ``Nt`` is the size of ``times``.
Note
----
See the :external+l2d-interface:ref:`general description <shape_convention>`
of the shape convention for `entries`.
Example
-------
.. code-block:: python
import jax.numpy as jnp
import typed_lisa_toolkit as tlt
ts1 = tlt.time_series(
tlt.linspace(0, 1, 10), jnp.ones((1, 1, 1, 1, 10))
) # UniformTimeSeries
ts2 = tlt.time_series(
jnp.array([0, 0.1, 0.3, 0.6, 1]), jnp.ones((1, 1, 1, 1, 5))
) # TimeSeries[Array]
"""
_validate_shape(entries, (entries.shape[0], 1, 1, 1, len(times)))
if isinstance(times, Linspace):
return UniformTimeSeries((axis(times),), entries)
if isinstance(times, Axis):
if isinstance(times.ax, Linspace):
return UniformTimeSeries((times,), entries)
return TimeSeries[AxisT]((times,), entries)
return TimeSeries[AxisT]((axis(times),), entries)
@overload
def frequency_phasor[AT: AxLike](
frequencies: AT,
amplitudes: AnyArray,
phases: AnyArray,
) -> FrequencyPhasor[Axis[AT]]: ...
@overload
def frequency_phasor[AxisT: "AnyAxis"](
frequencies: AxisT,
amplitudes: AnyArray,
phases: AnyArray,
) -> FrequencyPhasor[AxisT]: ...
[docs]
def frequency_phasor(
frequencies: AnyAxis | AxLike,
amplitudes: AnyArray,
phases: AnyArray,
):
"""Build a :class:`~types.FrequencyPhasor`.
Parameters
----------
frequencies:
Either an :class:`~typed_lisa_toolkit.types.Axis`,
a :class:`~typed_lisa_toolkit.types.Linspace`, or a 1D
:class:`array <typed_lisa_toolkit.types.misc.Array>` of positive frequencies.
In the last two cases, the axis will be automatically created with the
frequencies as values.
amplitudes: :class:`~typed_lisa_toolkit.types.misc.Array`
Either an array of shape ``(n_batch, n_channels, n_harmonics, 1, Nf)``
where ``Nf`` is the size of ``frequencies``,
or a 1D array of shape ``(Nf,)`` that will be broadcasted to
the shape ``(1, 1, 1, 1, Nf)``.
Must be of the same shape as ``phases``.
phases: :class:`~typed_lisa_toolkit.types.misc.Array`
Either an array of shape ``(n_batch, n_channels, n_harmonics, 1, Nf)``
where ``Nf`` is the size of ``frequencies``,
or a 1D array of shape ``(Nf,)`` that will be broadcasted to
the shape ``(1, 1, 1, 1, Nf)``
Must be of the same shape as ``amplitudes``.
Note
----
See the :external+l2d-interface:ref:`general description <shape_convention>`
of the shape convention.
"""
if amplitudes.shape != phases.shape:
msg = (
"Amplitudes and phases must have the same shape. "
f"Got {amplitudes.shape} and {phases.shape}."
)
raise ValueError(msg)
if amplitudes.ndim == 1:
pass
else:
_validate_shape(
amplitudes,
(
amplitudes.shape[0],
amplitudes.shape[1],
amplitudes.shape[2],
1,
len(frequencies),
),
)
if isinstance(frequencies, Linspace):
_axis = axis(frequencies)
elif isinstance(frequencies, Axis):
_axis = frequencies
else:
_axis = axis(frequencies)
return FrequencyPhasor[Any].make(
axis=_axis,
amplitudes=amplitudes,
phases=phases,
)
@overload
def time_phasor[AT: AxLike](
times: AT,
amplitudes: AnyArray,
phases: AnyArray,
) -> TimePhasor[Axis[AT]]: ...
@overload
def time_phasor[AxisT: "AnyAxis"](
times: AxisT,
amplitudes: AnyArray,
phases: AnyArray,
) -> TimePhasor[AxisT]: ...
[docs]
def time_phasor(
times: AnyAxis | AxLike,
amplitudes: AnyArray,
phases: AnyArray,
):
"""Build a :class:`~types.TimePhasor`.
Parameters
----------
times:
Either an :class:`~typed_lisa_toolkit.types.Axis`,
a :class:`~typed_lisa_toolkit.types.Linspace`, or a 1D
:class:`array <typed_lisa_toolkit.types.misc.Array>` of time points.
In the last two cases, the axis will be automatically created with the
time points as values.
amplitudes: :class:`~typed_lisa_toolkit.types.misc.Array`
Either an array of shape ``(n_batch, n_channels, n_harmonics, 1, Nt)``
where ``Nt`` is the size of ``times``,
or a 1D array of shape ``(Nt,)`` that will be broadcasted to
the shape ``(1, 1, 1, 1, Nt)``.
Must be of the same shape as ``phases``.
phases: :class:`~typed_lisa_toolkit.types.misc.Array`
Either an array of shape ``(n_batch, n_channels, n_harmonics, 1, Nt)``
where ``Nt`` is the size of ``times``,
or a 1D array of shape ``(Nt,)`` that will be broadcasted to
the shape ``(1, 1, 1, 1, Nt)``
Must be of the same shape as ``amplitudes``.
Note
----
See the :external+l2d-interface:ref:`general description <shape_convention>`
of the shape convention.
"""
if amplitudes.shape != phases.shape:
msg = (
"Amplitudes and phases must have the same shape. "
f"Got {amplitudes.shape} and {phases.shape}."
)
raise ValueError(msg)
if amplitudes.ndim == 1:
pass
else:
_validate_shape(
amplitudes,
(
amplitudes.shape[0],
amplitudes.shape[1],
amplitudes.shape[2],
1,
len(times),
),
)
if isinstance(times, Linspace):
_axis = axis(times)
elif isinstance(times, Axis):
_axis = times
else:
_axis = axis(times)
return TimePhasor[Any].make(
axis=_axis,
amplitudes=amplitudes,
phases=phases,
)
@overload
def phasor[AT: AxLike](
frequencies: AT,
amplitudes: AnyArray,
phases: AnyArray,
) -> Phasor[Axis[AT]]: ...
@overload
def phasor[AxisT: "AnyAxis"](
frequencies: AxisT,
amplitudes: AnyArray,
phases: AnyArray,
) -> Phasor[AxisT]: ...
[docs]
@deprecated("phasor", "function", "0.8.0", alternative="frequency_phasor")
def phasor(
frequencies: AnyAxis | AxLike,
amplitudes: AnyArray,
phases: AnyArray,
):
"""Alias for :func:`frequency_phasor` (*Deprecated*).
.. warning::
This function is deprecated and will be removed in 0.8.0; use
:func:`frequency_phasor` instead.
"""
return frequency_phasor(
frequencies=frequencies,
amplitudes=amplitudes,
phases=phases,
)
@overload
def stft[FreqAxisT: "AnyAxis", TimeAxisT: "AnyAxis"](
frequencies: FreqAxisT,
times: TimeAxisT,
entries: AnyArray,
*,
sparse_indices: None = None,
) -> STFT[Grid2DCartesian[FreqAxisT, TimeAxisT]]: ...
@overload
def stft[FAT: "AxLike", TAT: "AxLike"](
frequencies: FAT,
times: TAT,
entries: AnyArray,
*,
sparse_indices: None = None,
) -> STFT[Grid2DCartesian[Axis[FAT], Axis[TAT]]]: ...
@overload
def stft[FreqAxisT: "AnyAxis", TimeAxisT: "AnyAxis"](
frequencies: FreqAxisT,
times: TimeAxisT,
entries: AnyArray,
*,
sparse_indices: AnyArray,
) -> STFT[Grid2DSparse[FreqAxisT, TimeAxisT]]: ...
@overload
def stft[FAT: "AxLike", TAT: "AxLike"](
frequencies: FAT,
times: TAT,
entries: AnyArray,
*,
sparse_indices: AnyArray,
) -> STFT[Grid2DSparse[Axis[FAT], Axis[TAT]]]: ...
[docs]
def stft(
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
*,
sparse_indices: AnyArray | None = None,
):
"""Build an :class:`~types.ShortTimeFourierTransform`.
Parameters
----------
frequencies: FreqAxisT
Either a :class:`~typed_lisa_toolkit.types.Linspace` or a 1D
:class:`array <typed_lisa_toolkit.types.misc.Array>` of positive frequencies.
times: TimeAxisT
Either a :class:`~typed_lisa_toolkit.types.Linspace` or a 1D
:class:`array <typed_lisa_toolkit.types.misc.Array>` of time points.
entries: :class:`~typed_lisa_toolkit.types.misc.Array`
For dense grid: an array of shape ``(n_batch, 1, 1, 1, Nf, Nt)``
where ``Nf`` and ``Nt`` are the sizes of ``frequencies``
and ``times`` respectively.
For sparse grid: an array of shape ``(n_batch, 1, 1, 1, num_sparse_points)``
where ``num_sparse_points`` is the number of rows in ``sparse_indices``.
Note
----
See the :external+l2d-interface:ref:`general description <shape_convention>`
of the shape convention for `entries`.
"""
if sparse_indices is None:
_validate_shape(
entries,
(
entries.shape[0],
1,
1,
1,
len(frequencies),
len(times),
),
)
else:
_validate_shape(
entries,
(
entries.shape[0],
1,
1,
1,
len(sparse_indices),
),
)
_freq = axis(frequencies) if not isinstance(frequencies, Axis) else frequencies
_time = axis(times) if not isinstance(times, Axis) else times
grid = build_grid2d(_freq, _time, sparse_indices=sparse_indices)
if sparse_indices is None:
return STFT[Any](grid, entries)
return STFT[Any](grid, entries)
@overload
def wdm(
frequencies: AxLike,
times: AxLike,
entries: AnyArray,
*,
sparse_indices: None = None,
) -> WDM[Grid2DCartesian[Axis[Linspace], Axis[Linspace]]]: ...
@overload
def wdm[AxisT: "Axis[Linspace]"](
frequencies: AxisT,
times: AxisT,
entries: AnyArray,
*,
sparse_indices: None = None,
) -> WDM[Grid2DCartesian[AxisT, AxisT]]: ...
@overload
def wdm(
frequencies: AxLike,
times: AxLike,
entries: AnyArray,
*,
sparse_indices: AnyArray,
) -> WDM[Grid2DSparse[Axis[Linspace], Axis[Linspace]]]: ...
@overload
def wdm[AxisT: "Axis[Linspace]"](
frequencies: AxisT,
times: AxisT,
entries: AnyArray,
*,
sparse_indices: AnyArray,
) -> WDM[Grid2DSparse[AxisT, AxisT]]: ...
[docs]
def wdm(
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
*,
sparse_indices: AnyArray | None = None,
):
"""Build a :class:`~types.WilsonDaubechiesMeyer`.
Parameters
----------
frequencies: :class:`~typed_lisa_toolkit.types.misc.AnyAxis`
Evenly-spaced frequencies with separation ΔF and size ``Nf+1``.
times: :class:`~typed_lisa_toolkit.types.misc.AnyAxis`
Evenly-spaced times with separation ΔT and size ``Nt``.
entries: :class:`~typed_lisa_toolkit.types.misc.Array`
For dense grid: an array of shape ``(n_batch, 1, 1, 1, Nf+1, Nt)``.
For sparse grid: an array of shape ``(n_batch, 1, 1, 1, num_sparse_points)``
where ``num_sparse_points`` is the number of rows in ``sparse_indices``.
Note
----
See the :external+l2d-interface:ref:`general description <shape_convention>`
of the shape convention for `entries`.
"""
if sparse_indices is None:
_validate_shape(
entries,
(
entries.shape[0],
1,
1,
1,
len(frequencies),
len(times),
),
)
else:
_validate_shape(
entries,
(
entries.shape[0],
1,
1,
1,
len(sparse_indices),
),
)
freq_msg = "Frequencies axis must be uniformly spaced for WDM representation."
if isinstance(frequencies, Axis):
try:
freq_ax = Linspace.make(frequencies.ax)
except ValueError as e:
raise ValueError(freq_msg) from e
else:
try:
freq_ax = Linspace.make(frequencies)
except ValueError as e:
raise ValueError(freq_msg) from e
time_msg = "Times axis must be uniformly spaced for WDM representation."
if isinstance(times, Axis):
try:
time_ax = Linspace.make(times.ax)
except ValueError as e:
raise ValueError(time_msg) from e
else:
try:
time_ax = Linspace.make(times)
except ValueError as e:
raise ValueError(time_msg) from e
grid = build_grid2d(axis(freq_ax), axis(time_ax), sparse_indices=sparse_indices)
return WDM[Any](grid, entries)
class _Series1D[AxisT: "AnyAxis"]( # pyright: ignore[reportUnsafeMultipleInheritance]
_ArithmeticReprOnGrid["Grid1D[AxisT]"],
_Subset1DMixin["Grid1D[AxisT]"],
abc.ABC,
): ...
# 1D series are parameterized by axis type since there is only one type of 1D grid.
# If in the future we want to support more than one type of 1D grid, we should
# refactor a bit so that 1D series are parameterized by grid type as has been done
# for _TFRep.
[docs]
class FrequencySeries[AxisT: "AnyAxis"](_Series1D[AxisT]):
"""A series of numbers on a frequency grid.
.. note::
To construct a :class:`.FrequencySeries`, use the factory
function :func:`~typed_lisa_toolkit.frequency_series`.
"""
DOMAIN: Final = "frequency"
KIND: Final = None
@property
def domain(self) -> Literal["frequency"]:
"""The physical domain of the representation."""
return self.DOMAIN
@property
def kind(self) -> None:
"""The semantic kind of the representation."""
return self.KIND
[docs]
def angle(self, **kwargs: Any):
"""Return the angle of the series."""
return super().angle(**kwargs).unwrap(period=2 * self.xp.pi)
@property
def frequencies(self) -> AxisT:
"""The frequencies of the series."""
return self.grid[0]
@property
def f_min(self) -> float:
"""The minimum frequency of the series."""
return _get_axis_onset(self.frequencies)
@property
def f_max(self) -> float:
"""The maximum frequency of the series."""
return _get_axis_end(self.frequencies)
[docs]
def get_time_shifted(self, shift: float) -> Self:
"""Shift the series in time."""
return self * self.xp.exp(
-2j * self.xp.pi * self.xp.asarray(self.frequencies) * shift,
)
[docs]
def get_embedded[AT: "AnyAxis"](
self,
embedding_grid: Grid1D[AT],
*,
known_slices: tuple[slice, ...] | None = None,
):
"""Return the series embedded in a new 1D grid."""
grid, entries = _mixins.embed_entries_to_grid(
self.grid,
self.entries,
embedding_grid,
known_slices=known_slices,
)
return frequency_series(grid[0], entries)
[docs]
class TimeSeries[AxisT: "AnyAxis"](_Series1D[AxisT]):
"""A series of numbers on a time grid.
.. note::
To construct a :class:`.TimeSeries`, use the factory
function :func:`~typed_lisa_toolkit.time_series`.
"""
DOMAIN: Final = "time"
KIND: Final = None
@property
def domain(self) -> Literal["time"]:
"""The physical domain of the representation."""
return self.DOMAIN
@property
def kind(self) -> None:
"""The semantic kind of the representation."""
return self.KIND
@property
def times(self) -> AxisT:
"""The times of the series."""
return self.grid[0]
@property
def t_start(self) -> float:
"""The onset time of the series."""
return _get_axis_onset(self.times)
@property
def t_end(self) -> float:
"""The end time of the series."""
return _get_axis_end(self.times)
[docs]
def get_embedded[AT: "AnyAxis"](
self,
embedding_grid: Grid1D[AT],
*,
known_slices: tuple[slice, ...] | None = None,
):
"""Return the series embedded in a new 1D grid."""
grid, entries = _mixins.embed_entries_to_grid(
self.grid,
self.entries,
embedding_grid,
known_slices=known_slices,
)
return time_series(grid[0], entries)
class _BasePhasor[AxisT: "AnyAxis"](
_Subset1DMixin["Grid1D[AxisT]"],
abc.ABC,
):
KIND: Final = "phasor"
@property
def kind(self) -> Literal["phasor"]:
"""The semantic kind of the representation."""
return self.KIND
@property
@abc.abstractmethod
def domain(self) -> Domain:
"""Physical domain of the representation."""
@property
def phases(self) -> AnyArray:
"""The phases of the phasors."""
return self.xp.real(self.entries[..., slice(1, 2), :])
@property
def amplitudes(self) -> AnyArray:
"""The amplitudes of the phasors."""
return self.entries[..., slice(0, 1), :]
@property
def axis(self) -> AxisT:
return self.grid[0]
@property
def axis_onset(self) -> float:
return _get_axis_onset(self.axis)
@property
def axis_end(self) -> float:
return _get_axis_end(self.axis)
@classmethod
def make(
cls,
*,
axis: AnyAxis,
amplitudes: AnyArray,
phases: AnyArray,
):
"""Create a phasor from amplitudes and phases."""
xp = xpc.get_namespace(amplitudes, phases)
# If amplitudes and phases are 1D
if amplitudes.ndim == 1 and phases.ndim == 1:
return cls(
grid=(axis,),
entries=xp.stack((amplitudes, phases), axis=0)[None, None, None, ...],
)
# If amplitudes and phases are already in the shape of entries
full_shape_size = 5
if (
amplitudes.shape == phases.shape
and len(amplitudes.shape) == full_shape_size
and amplitudes.shape[4] == len(axis)
and amplitudes.shape[3] == 1
):
return cls(
grid=(axis,),
entries=xp.stack((amplitudes[:, :, 0], phases[:, :, 0]), axis=3),
)
_msg = (
"Amplitudes and phases must be either 1D arrays of shape (n_axis,) "
"or 5D arrays of shape "
"(n_batches, n_channels, n_harmonics, 1, n_axis), "
f"but got shapes {amplitudes.shape} and {phases.shape}."
)
raise ValueError(_msg)
def __setitem__(self, slice: _slice, value: Any) -> None:
"""Set the entries and phases of a subset of the phasor."""
_set_value(self.entries, _get_full_slice((slice,)), value)
def create_like(self, entries: AnyArray):
"""Create a new instance with the same grid as the current one."""
return type(self)(grid=self.grid, entries=entries)
def get_embedded[AT: "AnyAxis"](
self,
embedding_grid: Grid1D[AT],
*,
known_slices: tuple[slice, ...] | None = None,
) -> Self:
"""Return the phasor embedded in a new 1D grid."""
grid, entries = _mixins.embed_entries_to_grid(
self.grid,
self.entries,
embedding_grid,
known_slices=known_slices,
)
return type(self)(grid=grid, entries=entries)
def _get_interpolated_impl(
self,
axis_: AnyAxis | AnyArray,
interpolator: Interpolator,
):
"""Implement get_interpolated for subclasses."""
xp = xpc.get_namespace(self.amplitudes, self.phases)
_axis_array = axis_.asarray(xp) if isinstance(axis_, Axis) else axis_
_axis_obj = axis_ if isinstance(axis_, Axis) else axis(axis_)
self_axis = self.axis.asarray(xp)
if self.entries.shape != (1, 1, 1, 2, len(self_axis)):
_msg = (
f"Only 1D phasors with shape (1, 1, 1, 2, {len(self_axis)}) "
"are supported "
f"for interpolation, but got shape {self.entries.shape}."
)
raise ValueError(_msg)
_axes_to_squeeze = (0, 1, 2, 3)
amp_real = xp.squeeze(xp.real(self.amplitudes), axis=_axes_to_squeeze)
amp_imag = xp.squeeze(xp.imag(self.amplitudes), axis=_axes_to_squeeze)
amplitudes_real = interpolator(self_axis, amp_real)(_axis_array)
amplitudes_imag = interpolator(self_axis, amp_imag)(_axis_array)
amplitudes = cast(
"AnyArray",
amplitudes_real + cast("AnyArray", amplitudes_imag * 1j),
)
phases = interpolator(
self_axis, xp.squeeze(self.phases, axis=_axes_to_squeeze)
)(_axis_array)
return type(self).make(
axis=_axis_obj,
amplitudes=amplitudes,
phases=phases,
)
[docs]
class FrequencyPhasor[AxisT: "AnyAxis"](
_BasePhasor[AxisT],
):
"""Frequency-domain phasor representation.
.. note::
To construct a :class:`.FrequencyPhasor`, use the factory function
:func:`~typed_lisa_toolkit.frequency_phasor`.
A phasor is a couple of amplitude and phase that represent a complex number.
This class encapsulates a sequence of phasors at different frequencies, which
can be used to represent a waveform. This representation is useful for
interpolating waveforms generated on a sparse grid of frequencies to a dense
grid of frequencies.
The input phases are expected to be smooth, without zigzags, so as the real
and imaginary parts of the amplitudes. This is crucial for the interpolation
to work properly.
.. note:: The so-called amplitude is itself complex number in general.
"""
DOMAIN: Final = "frequency"
@property
def domain(self) -> Literal["frequency"]:
"""The physical domain of the representation."""
return self.DOMAIN
@property
def frequencies(self) -> AxisT:
"""The frequencies of the phasors."""
return self.axis
@property
def f_min(self) -> float:
"""The minimum frequency of the series."""
return _get_axis_onset(self.frequencies)
@property
def f_max(self) -> float:
"""The maximum frequency of the series."""
return _get_axis_end(self.frequencies)
@overload
def get_interpolated[AT: "AnyAxis"](
self,
axis_: AT,
interpolator: Interpolator,
) -> FrequencyPhasor[AT]: ...
@overload
def get_interpolated(
self,
axis_: AnyArray,
interpolator: Interpolator,
) -> FrequencyPhasor[Axis[AnyArray]]: ...
[docs]
def get_interpolated(
self,
axis_: AnyAxis | AnyArray,
interpolator: Interpolator,
) -> FrequencyPhasor[AnyAxis]:
"""Get the phasors interpolated to the given frequencies."""
return self._get_interpolated_impl(axis_, interpolator)
[docs]
def to_frequency_series(self):
r"""Convert to a :class:`.FrequencySeries` or :class:`.UniformFrequencySeries`.
This method returns a :class:`.FrequencySeries`
or :class:`.UniformFrequencySeries`
by applying the formula:
.. math::
X(f) = A(f) \cdot \exp(1j \cdot \phi(f))
"""
xp = xpc.get_namespace(self.amplitudes, self.phases)
return frequency_series(
self.frequencies,
self.amplitudes * xp.exp(self.phases * 1j),
)
[docs]
def to_series(self):
"""Alias for :meth:`.to_frequency_series`."""
return self.to_frequency_series()
[docs]
class TimePhasor[AxisT: "AnyAxis"](
_BasePhasor[AxisT],
):
"""Time-domain phasor representation.
.. note::
To construct a :class:`.TimePhasor`, use the factory function
:func:`~typed_lisa_toolkit.time_phasor`.
A phasor is a couple of amplitude and phase that represent a complex number.
This class encapsulates a sequence of phasors at different times, which
can be used to represent a waveform. This representation is useful for
interpolating waveforms generated on a sparse grid of times to a dense
grid of times.
The input phases are expected to be smooth, without zigzags, so as the real
and imaginary parts of the amplitudes. This is crucial for the interpolation
to work properly.
.. note:: The so-called amplitude is itself complex number in general.
"""
DOMAIN: Final = "time"
@property
def domain(self) -> Literal["time"]:
"""The physical domain of the representation."""
return self.DOMAIN
@property
def times(self) -> AxisT:
"""The times of the phasors."""
return self.axis
@property
def t_start(self) -> float:
"""The onset time of the series."""
return _get_axis_onset(self.times)
@property
def t_end(self) -> float:
"""The end time of the series."""
return _get_axis_end(self.times)
@overload
def get_interpolated[AT: "AnyAxis"](
self,
axis_: AT,
interpolator: Interpolator,
) -> TimePhasor[AT]: ...
@overload
def get_interpolated(
self,
axis_: AnyArray,
interpolator: Interpolator,
) -> TimePhasor[Axis[AnyArray]]: ...
[docs]
def get_interpolated(
self,
axis_: AnyAxis | AnyArray,
interpolator: Interpolator,
) -> TimePhasor[AnyAxis]:
"""Get the phasors interpolated to the given times."""
return self._get_interpolated_impl(axis_, interpolator)
[docs]
def to_time_series(self):
r"""Convert to a :class:`.TimeSeries` or :class:`.UniformTimeSeries`.
This method returns a :class:`.TimeSeries` or :class:`.UniformTimeSeries`
by applying the formula:
.. math::
X(t) = A(t) \cdot \exp(1j \cdot \phi(t))
"""
xp = xpc.get_namespace(self.amplitudes, self.phases)
return time_series(
self.times,
self.amplitudes * xp.exp(self.phases * 1j),
)
[docs]
def to_series(self):
"""Alias for :meth:`.to_time_series`."""
return self.to_time_series()
# Backward compatibility alias
Phasor = FrequencyPhasor
"""Alias for :class:`.FrequencyPhasor` (*Deprecated*)."""
class _FreqProperty2D:
def __get__[FreqAxisT: AnyAxis, TimeAxisT: AnyAxis](
self,
instance: _TFRep[Grid2D[FreqAxisT, TimeAxisT]],
owner: Any,
) -> FreqAxisT: ...
class _TimeProperty2D:
def __get__[FreqAxisT: AnyAxis, TimeAxisT: AnyAxis](
self,
instance: _TFRep[Grid2D[FreqAxisT, TimeAxisT]],
owner: Any,
) -> TimeAxisT: ...
class _TFRep[ # pyright: ignore[reportUnsafeMultipleInheritance]
GridT: Grid2D[AnyAxis, AnyAxis]
](
_ArithmeticReprOnGrid[GridT],
_InitMixin[GridT],
abc.ABC,
):
DOMAIN: Final = "time-frequency"
@property
def domain(self) -> Literal["time-frequency"]:
"""The physical domain of the representation."""
return self.DOMAIN
@property
def times(self): # pyright: ignore[reportRedeclaration]
"""The time grid of the time-frequency representation."""
return self.grid[1]
times: _TimeProperty2D # For correct type hinting
@property
def t_start(self) -> float:
"""The onset time of the series."""
return _get_axis_onset(self.times)
@property
def t_end(self) -> float:
"""The end time of the series."""
return _get_axis_end(self.times)
@property
def frequencies(self): # pyright: ignore[reportRedeclaration]
"""The frequency grid of the time-frequency representation."""
return self.grid[0]
frequencies: _FreqProperty2D # For correct type hinting
@property
def f_min(self) -> float:
"""The minimum frequency of the series."""
return _get_axis_onset(self.frequencies)
@property
def f_max(self) -> float:
"""The maximum frequency of the series."""
return _get_axis_end(self.frequencies)
def get_subset(
self,
*,
time_interval: tuple[float, float] | None = None,
freq_interval: tuple[float, float] | None = None,
slices: tuple[_slice, _slice] | None = None,
copy: bool = True,
) -> Self:
"""Return the subset as a new instance."""
_freq_slice = _get_subset_slice(
self.frequencies,
interval=freq_interval,
slice=slices[0] if slices is not None else None,
)
_time_slice = _get_subset_slice(
self.times,
interval=time_interval,
slice=slices[1] if slices is not None else None,
)
if isinstance(self.grid, Grid2DSparse):
grid, entries = _subset_grid_2d_sparse(
self.grid,
self.entries,
(_freq_slice, _time_slice),
)
return type(self)(grid=grid, entries=entries)
grid, entries = _take_subset(
self.grid,
self.entries,
(_freq_slice, _time_slice),
)
entries = _copy.copy(entries) if copy else entries
return type(self)(grid=grid, entries=entries)
STFT = ShortTimeFourierTransform
[docs]
class WilsonDaubechiesMeyer[GridT: Grid2D[Axis[Linspace], Axis[Linspace]]](
_TFRep[GridT]
):
"""
Wilson-Daubechies-Meyer time-frequency representation.
.. note::
To construct a :class:`.WilsonDaubechiesMeyer`, use the factory function
:func:`~typed_lisa_toolkit.wdm`.
This represents data using an evenly-spaced 2D grid in the time-frequency plane with
shape (Nf+1, Nt). Each "pixel" has size ΔT ΔF = 1/2. The times range approximately
from 0 to the final observation time, while the frequencies range from 0 to the
Nyquist limit (half the sampling rate).
.. attention:: The frequency grid has size Nf+1, not Nf. This is because correctly
inverting WDM transforms requires some information from the Nyquist band m=Nf.
This data layout contains redundant information but is simpler to interpret and
to work with. It follows the convention of `wdm-transform`_.
.. _wdm-transform: https://github.com/pywavelet/wdm_transform
"""
KIND: Final = "wdm"
@property
def kind(self) -> str:
"""The semantic kind of the representation."""
return self.KIND
@property
def dT(self) -> float: # noqa: N802
"""Time resolution (ΔT) of the time-frequency grid."""
return self.times.ax.step
@property
def dF(self) -> float: # noqa: N802
"""Frequency resolution (ΔF) of the time-frequency grid."""
return self.frequencies.ax.step
[docs]
def is_critically_sampled(self) -> bool:
"""Return True if :attr:`.dT` * :attr:`.dF` = 1/2."""
# I don't like how this method is implemented,
# but I don't see a better way for now.
return bool(self.xp.isclose(self.dT * self.dF, 1 / 2))
@property
def Nt(self) -> int: # noqa: N802
"""Number of time bins.
.. note::
Throughout this codebase, a grid of N points is considered to have N-1 bins,
since the first point is the start of the first bin and the last point is
the end of the last bin.
"""
return self.times.ax.num
@property
def Nf(self) -> int: # noqa: N802
"""Number of frequency points."""
return self.frequencies.ax.num - 1
@property
def ND(self) -> int: # noqa: N802
"""Total number of data points in the time-frequency plane."""
return self.Nt * self.Nf
@property
def sample_interval(self) -> float:
"""
Time resolution of a TimeSeries corresponding to this WDM.
Smaller than the wavelet time bin :attr:`.dT`.
"""
return self.dT / self.Nf
dt: property = sample_interval
"""Alias for :attr:`.sample_interval`."""
@property
def df(self) -> float:
"""
Frequency resolution of a FrequencySeries corresponding to this WDM.
Smaller than the wavelet frequency bin :attr:`.dF`.
"""
return 1 / (self.ND * self.dt)
@property
def shape(self) -> tuple[int, int]:
"""Shape (:attr:`.Nf`, :attr:`.Nt`) of the wavelet grid."""
return self.Nf, self.Nt
@property
def fs(self) -> float:
"""Sampling rate 1/Δt for the time series equivalent to this WDM."""
return 1 / self.dt
@property
def nyquist(self) -> float:
"""Nyquist frequency (half the sampling rate)."""
return self.fs / 2
[docs]
def get_embedded[AT0: "Axis[Linspace]", AT1: "Axis[Linspace]"](
self,
embedding_grid: Grid2D[AT0, AT1],
*,
known_slices: tuple[slice, ...] | None = None,
):
"""Return the representation embedded in a new 2D grid."""
if isinstance(self.grid, Grid2DSparse):
grid, entries = _embed_entries_to_grid_2d_sparse(
self.grid,
self.entries,
embedding_grid,
known_slices=known_slices,
)
return wdm(grid[0], grid[1], entries, sparse_indices=grid.indices)
grid, entries = _mixins.embed_entries_to_grid(
self.grid,
self.entries,
embedding_grid,
known_slices=known_slices,
)
return wdm(grid[0], grid[1], entries)
WDM = WilsonDaubechiesMeyer