"""Data container types."""
from __future__ import annotations
import abc
import logging
import pathlib
from collections.abc import Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Literal, Protocol, Self, cast, overload
import array_api_compat as xpc
import h5py
import numpy as np
import numpy.typing as npt
from ..utils import deprecated, extend_to, warn_external
from . import _mixins, tapering
from . import representations as reps
from .misc import (
AnyArray,
AnyAxis,
AnyGrid,
Axis,
AxLike,
Domain,
Grid1D,
Grid2D,
Grid2DCartesian,
Grid2DSparse,
Linspace,
axis,
build_grid2d,
linspace_from_step,
)
if TYPE_CHECKING:
from MojitoProcessor import SignalProcessor
# from .waveforms import ProjectedWaveform
from .representations import Representation
AnyReps = Representation[AnyGrid]
class _SubsettableRep1D(AnyReps, Protocol):
def get_subset(
self,
*,
interval: tuple[float, float] | None = None,
slice: slice | None = None,
copy: bool = True,
) -> Self: ...
class _SubsettableRep2D(AnyReps, Protocol):
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: ...
log = logging.getLogger(__name__)
def _save_axis(grp: h5py.Group, name: str, axis: AnyAxis) -> None:
"""Serialize one grid axis, preserving Linspace compactness."""
axis_grp = grp.create_group(name)
if isinstance(axis.ax, Linspace):
axis_grp.attrs["linspace"] = True
axis_grp.attrs["start"] = axis.start
axis_grp.attrs["step"] = axis.ax.step
axis_grp.attrs["num"] = axis.ax.num
return
axis_grp.create_dataset("values", data=np.asarray(axis.ax))
def _attr_float(attrs: Any, key: str) -> float:
return float(attrs[key])
def _attr_int(attrs: Any, key: str) -> int:
return int(attrs[key])
def _attr_bool(attrs: Any, key: str, *, default: bool = False) -> bool:
return bool(attrs.get(key, default))
def _load_axis(node: h5py.Group | h5py.Dataset) -> AnyAxis:
"""Deserialize one grid axis.
Supports both new grouped axis format and old raw-dataset axis format.
"""
if isinstance(node, h5py.Dataset):
return cast("AnyAxis", node[()])
linspace = _attr_bool(node.attrs, "linspace", default=False)
if linspace:
return axis(
Linspace(
start=_attr_float(node.attrs, "start"),
step=_attr_float(node.attrs, "step"),
num=_attr_int(node.attrs, "num"),
)
)
return cast("AnyAxis", axis(cast("h5py.Dataset", node["values"])[()]))
def _save_grid(grp: h5py.Group, grid: AnyGrid) -> None:
"""Serialize a representation grid into an HDF5 group."""
if isinstance(grid, Grid2DSparse):
grp.attrs["sparse"] = True
grp.attrs["dim"] = 2
_save_axis(grp, "axis0", grid[0])
_save_axis(grp, "axis1", grid[1])
grp.create_dataset("sparse_indices", data=np.asarray(grid.indices))
return
if len(grid) == 1:
grp.attrs["dim"] = 1
_save_axis(grp, "axis0", grid[0])
return
if len(grid) == 2: # noqa: PLR2004
grp.attrs["sparse"] = False
grp.attrs["dim"] = 2
_save_axis(grp, "axis0", grid[0])
_save_axis(grp, "axis1", grid[1])
return
def _load_grid(node: h5py.Group | h5py.Dataset) -> AnyGrid:
"""Deserialize a representation grid from HDF5."""
# Backward compatibility: previous format stored `grid` as a single dataset.
if isinstance(node, h5py.Dataset):
return cast("AnyGrid", node[()])
sparse = _attr_bool(node.attrs, "sparse", default=False)
dim = _attr_int(node.attrs, "dim") if "dim" in node.attrs else -1
if not sparse and dim == 1:
axis0 = _load_axis(cast("h5py.Group | h5py.Dataset", node["axis0"]))
return (axis0,)
if not sparse and dim == 2: # noqa: PLR2004
axis0 = _load_axis(cast("h5py.Group | h5py.Dataset", node["axis0"]))
axis1 = _load_axis(cast("h5py.Group | h5py.Dataset", node["axis1"]))
return (axis0, axis1)
if sparse and dim == 2: # noqa: PLR2004
axis0 = _load_axis(cast("h5py.Group | h5py.Dataset", node["axis0"]))
axis1 = _load_axis(cast("h5py.Group | h5py.Dataset", node["axis1"]))
sparse_indices = cast("h5py.Dataset", node["sparse_indices"])[()]
return Grid2DSparse(axis0, axis1, sparse_indices=sparse_indices)
msg = f"Unknown grid serialization format with attributes: {node.attrs}"
raise ValueError(msg)
def _load_data[DataT: "Data[AnyReps]"](
cls: type[DataT],
file_path: str | pathlib.Path,
*,
legacy: bool = False,
) -> DataT:
"""Load data from an HDF5 file."""
if legacy:
return cls._load_legacy(file_path) # pyright: ignore[reportPrivateUsage]
with h5py.File(str(file_path), "r") as f:
channels_attr = cast("Iterable[object]", f.attrs["channels"])
channels = tuple(str(ch) for ch in channels_attr)
additions = cls._additional_load(f) # pyright: ignore[reportPrivateUsage]
data_group = cast("h5py.Group", f["data"])
grid_node = cast("h5py.Group | h5py.Dataset", data_group["grid"])
grid_data = _load_grid(grid_node)
entries_data = cast("h5py.Dataset", data_group["entries"])[()]
return cls(grid_data, entries_data, channels=channels, **additions)
class _SubsetMixin1D[RepT: _SubsettableRep1D](_mixins.ChannelMapping[RepT], 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."""
subset_dict = {
chnname: chn.get_subset(interval=interval, slice=slice, copy=copy)
for chnname, chn in self.items()
}
return type(self).from_dict(subset_dict).set_name(self.name)
class _SubsetMixin2D[RepT: _SubsettableRep2D](_mixins.ChannelMapping[RepT], abc.ABC):
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."""
subset_dict = {
chnname: chn.get_subset(
time_interval=time_interval,
freq_interval=freq_interval,
slices=slices,
copy=copy,
)
for chnname, chn in self.items()
}
return type(self).from_dict(subset_dict).set_name(self.name)
class Data[RepT: "AnyReps"](_mixins.ChannelMapping[RepT], abc.ABC):
"""Channel-indexed data containers.
Stores a single homogeneous _representation with channels as the first dimension,
providing per-channel access via views and the Mapping protocol.
"""
_REP_TYPE: type[RepT]
def _additional_save(self, f: h5py.File):
del f
def save(self, file_path: str | pathlib.Path):
"""Save the data to an HDF5 file.
The data are saved in the following structure:
- The root level contains the attribute ``type`` with the class name.
- Each channel is saved as a group.
- The group contains two datasets:
- ``grid`` for the grid.
- ``entries`` for the entries.
- For :class:`.TimedFSData`, there will be a dataset ``times`` at
the root level containing the time grid.
"""
with h5py.File(str(file_path), "a") as f:
f.attrs["domain"] = self.domain
f.attrs["kind"] = str(self.kind)
f.attrs["channels"] = self.channel_names
self._additional_save(f)
grp = f.create_group("data")
_save_grid(grp.create_group("grid"), self.grid)
grp.create_dataset("entries", data=cast("Any", self.get_kernel()))
@classmethod
def _additional_load(cls, f: h5py.File) -> dict[str, Any]:
del f
return {}
@classmethod
def _load_legacy(cls, file_path: str | pathlib.Path) -> Self:
"""Load the data from an HDF5 file."""
with h5py.File(str(file_path), "r") as f:
additions = cls._additional_load(f)
dict_ = {
chnname: cls._REP_TYPE(
grid=(f[chnname]["grid"][...],), # pyright: ignore[reportArgumentType, reportIndexIssue]
entries=f[chnname]["entries"][...][None, None, None, None, ...], # pyright: ignore[reportArgumentType, reportIndexIssue]
)
for chnname in f
if isinstance(f[chnname], h5py.Group)
}
return cls.from_dict(dict_, **additions)
@classmethod
@deprecated("load", "method", "0.8.0", alternative="load_data")
def load(cls, file_path: str | pathlib.Path, *, legacy: bool = False):
"""Load the data from an HDF5 file (*Deprecated*).
Warning
-------
This method is deprecated and will be removed in 0.8.0;
use :func:`~typed_lisa_toolkit.load_data` instead.
"""
return _load_data(cls, file_path, legacy=legacy)
class _SeriesData[RepT: reps.UniformTimeSeries | reps.UniformFrequencySeries]( # pyright: ignore[reportUnsafeMultipleInheritance]
Data[RepT],
_SubsetMixin1D[RepT],
abc.ABC,
):
def get_embedded(
self,
embedding_grid: Grid1D[AnyAxis],
*,
known_slices: tuple[slice, ...] | None = None,
):
"""Return data embedded on 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,
channels=self.channel_names,
name=self.name,
)
[docs]
class TSData(_SeriesData[reps.UniformTimeSeries]):
"""Multi-channel time series data container.
.. note::
To construct a :class:`.TSData`,
use the factory function :func:`~typed_lisa_toolkit.tsdata`.
"""
_REP_TYPE: type[reps.UniformTimeSeries] = reps.UniformTimeSeries
[docs]
@classmethod
def from_entries(
cls,
*,
times: AnyAxis,
entries: AnyArray,
channels: tuple[str, ...],
name: str | None = None,
) -> Self:
"""Construct from raw time-domain entries and explicit channel names.
Warning
-------
This method is considered an expert-level API; for most users,
prefer to construct a :class:`.TSData` with the factory function
:func:`~typed_lisa_toolkit.tsdata`.
"""
return cls((times,), entries, channels=channels, name=name)
@property
def kind(self) -> None:
"""Semantic kind of the data."""
return None
@property
def times(self):
"""Return the times."""
return self.grid[0]
@property
def dt(self) -> float:
"""Return the time step."""
return self.times.ax.step
@property
def t_start(self) -> float:
"""Return the start time."""
return self.times.start
@property
def t_end(self) -> float:
"""Return the end time."""
return self.times.stop
[docs]
def get_frequencies(self):
"""Return the frequencies grid matching the time grid."""
return self.xp.fft.rfftfreq(len(self.times), d=self.dt)
@overload
def to_fsdata(
self,
*,
keep_times: Literal[False],
tapering: tapering.Tapering | None = None,
) -> FSData: ...
@overload
def to_fsdata(
self,
*,
keep_times: Literal[True] = True,
tapering: tapering.Tapering | None = None,
) -> TimedFSData: ...
[docs]
@deprecated("to_fsdata", "method", "0.8.0", alternative="shop.time2freq")
def to_fsdata(
self,
*,
keep_times: bool = True,
tapering: tapering.Tapering | None = None,
):
"""Return the frequency series data (*Deprecated*).
.. warning::
This method is deprecated and will be removed in 0.8.0; use
:func:`~typed_lisa_toolkit.shop.time2freq` instead.
Returns
-------
:class:`.FSData` | :class:`.TimedFSData`
The frequency series data. If `keep_times` is `True`, the time grid is kept.
"""
from ..shop import time2freq
_window = tapering(self.xp.asarray(self.times)) if tapering is not None else 1
return time2freq(self * _window, keep_time=keep_times)
[docs]
def get_zero_padded(
self,
pad_time: tuple[float, float],
tapering: tapering.Tapering | None = None,
) -> TSData:
"""Return the zero-padded data."""
xp = xpc.get_namespace(self.get_kernel())
pad_width = tuple(
int(xp.round(xp.asarray(time / self.dt))) for time in pad_time
)
known_slices = (slice(pad_width[0], pad_width[0] + len(self.times)),)
target = axis(
linspace_from_step(
self.times.start - pad_width[0] * self.dt,
self.dt,
len(self.times) + sum(pad_width),
)
)
taper = tapering(xp.asarray(self.times)) if tapering is not None else 1
padded_entries = extend_to(target, known_slices=known_slices)(
self.times, self.get_kernel() * taper
)
return tsdata(
times=target,
entries=padded_entries,
channels=self.channel_names,
name=self.name,
)
[docs]
class FSData(_SeriesData[reps.UniformFrequencySeries]):
"""Multi-channel frequency series data container.
.. note::
To construct a :class:`.FSData`, use the factory
function :func:`~typed_lisa_toolkit.fsdata`.
"""
_REP_TYPE: type[reps.UniformFrequencySeries] = reps.UniformFrequencySeries
[docs]
@classmethod
def from_entries(
cls,
*,
frequencies: AnyAxis,
entries: AnyArray,
channels: tuple[str, ...],
name: str | None = None,
) -> Self:
"""Construct from raw frequency-domain entries and explicit channel names.
Warning
-------
This method is considered an expert-level API; for most users,
prefer to construct a :class:`.FSData` with the factory function
:func:`~typed_lisa_toolkit.fsdata`.
"""
return cls((frequencies,), entries, channels=channels, name=name)
@property
def kind(self) -> None | str:
"""Semantic kind of the data."""
return None
@property
def frequencies(self):
"""Return the frequencies."""
return self.grid[0]
@property
def df(self):
"""Return the frequency step."""
return self.frequencies.ax.step
@property
def f_min(self):
"""Return the minimum frequency."""
return self.frequencies.start
@property
def f_max(self):
"""Return the maximum frequency."""
return self.frequencies.stop
@overload
def set_times(self, times: AxLike) -> TimedFSData: ...
@overload
def set_times(self, times: AnyAxis) -> TimedFSData: ...
[docs]
def set_times(self, times: AxLike | AnyAxis) -> TimedFSData:
"""Return a :class:`.TimedFSData` with the time grid set."""
return TimedFSData(
self.grid,
self.entries,
channels=self.channel_names,
name=self.name,
).set_times(times)
[docs]
@deprecated("to_tsdata", "method", "0.8.0", alternative="shop.freq2time")
def to_tsdata(
self,
times: npt.NDArray[np.floating[Any]] | Linspace,
*,
tapering: tapering.Tapering | None = None,
):
"""Return the time series data (*Deprecated*).
.. warning::
This method is deprecated and will be removed in 0.8.0; use
:func:`~typed_lisa_toolkit.shop.freq2time` instead.
"""
from ..shop import freq2time
_window = tapering(self.xp.asarray(times)) if tapering is not None else 1
return freq2time(self * _window, times=times)
# def _get_plotter(self):
# from ..viz import plotters
# return plotters.FSDataPlotter
[docs]
class TimedFSData(FSData):
"""Multi-channel frequency series data with time information.
.. note::
To construct a :class:`.TimedFSData`, use the factory
function :func:`~typed_lisa_toolkit.fsdata`.
"""
@property
def kind(self):
"""Semantic kind of the data."""
return "timed"
def create_like(self, entries: AnyArray) -> Self: # noqa: D102
return super().create_like(entries).set_times(self.times)
[docs]
def get_subset( # noqa: D102
self,
*,
interval: tuple[float, float] | None = None,
slice: slice | None = None,
copy: bool = True,
) -> Self:
return (
super()
.get_subset(interval=interval, slice=slice, copy=copy)
.set_times(self.times)
)
[docs]
def get_embedded( # noqa: D102
self,
embedding_grid: Grid1D[AnyAxis],
*,
known_slices: tuple[slice, ...] | None = None,
) -> Self:
return (
super()
.get_embedded(
embedding_grid,
known_slices=known_slices,
)
.set_times(self.times)
)
def _additional_save(self, f: h5py.File):
f.create_dataset("times", data=np.asarray(self.times))
@classmethod
def _additional_load(cls, f: h5py.File):
times_data = cast("h5py.Dataset", f["times"])[()]
return {"times": times_data}
@overload
def set_times(self, times: AxLike) -> Self: ...
@overload
def set_times(self, times: AnyAxis) -> Self: ...
[docs]
def set_times(self, times: AxLike | AnyAxis) -> Self:
"""Set the time grid.
.. note::
This method returns ``self`` to allow for fluent method chaining.
"""
_times = axis(times) if not isinstance(times, Axis) else times
self._times: Axis[Linspace] = axis(Linspace.make(_times.ax))
self._dt: float = self._times.ax.step
return self
@property
def times(self):
"""Associated time grid."""
return self._times
@property
def dt(self) -> float:
"""Step size of the associated time grid."""
return self._dt
[docs]
def drop_times(self):
"""Drop the time grid."""
return FSData(
self.grid,
self.get_kernel(),
channels=self.channel_names,
name=self.name,
)
[docs]
@classmethod
def from_dict(cls, data_dict: Mapping[str, AnyReps], /, **kwargs: Any) -> Self: # noqa: D102
times = kwargs.pop("times", None)
obj = super().from_dict(data_dict, **kwargs)
return obj.set_times(times) if times is not None else obj
[docs]
def to_tsdata(
self,
times: npt.NDArray[np.floating[Any]] | Linspace | None = None,
*,
tapering: tapering.Tapering | None = None,
):
"""Return the time series data with times grid (*Deprecated*).
.. warning::
This method is deprecated and will be removed in 0.8.0; use
:func:`~typed_lisa_toolkit.shop.freq2time` instead.
"""
del times
return super().to_tsdata(times=self.times, tapering=tapering)
class _Grid2DData[ # pyright: ignore[reportUnsafeMultipleInheritance]
RepT: reps.STFT[Grid2D[Axis[Linspace], Axis[Linspace]]]
| reps.WDM[Grid2D[Axis[Linspace], Axis[Linspace]]],
](
Data[RepT],
_SubsetMixin2D[RepT],
abc.ABC,
):
@classmethod
def from_entries(
cls,
*,
frequencies: AnyAxis,
times: AnyAxis,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: AnyArray | None = None,
name: str | None = None,
) -> Self:
"""Construct from raw time-frequency entries and explicit channel names."""
grid2d = build_grid2d(frequencies, times, sparse_indices=sparse_indices)
return cls(grid2d, entries, channels=channels, name=name)
[docs]
class STFTData[GridT: Grid2D[Axis[Linspace], Axis[Linspace]]](
_Grid2DData[reps.STFT[GridT]]
):
"""Multi-channel short-time Fourier transform data container.
.. note::
To construct a :class:`.STFTData`, use the factory
function :func:`~typed_lisa_toolkit.stftdata`.
"""
_REP_TYPE: type[reps.STFT[GridT]] = reps.STFT[GridT]
@property
def kind(self) -> str:
"""Semantic kind of the data."""
return "stft"
[docs]
class WDMData[GridT: Grid2D[Axis[Linspace], Axis[Linspace]]](
_Grid2DData[reps.WDM[GridT]]
):
"""Multi-channel wavelet domain model data container.
.. note::
To construct a :class:`.WDMData`, use the factory
function :func:`~typed_lisa_toolkit.wdmdata`.
"""
_REP_TYPE: type[reps.WDM[GridT]] = reps.WDM[GridT]
@property
def kind(self) -> str:
"""Semantic kind of the data."""
return "wdm"
def _enforce_uniform(ary: AnyAxis, /) -> Axis[Linspace]:
"""Enforce that the given array is uniform and return it as a Linspace."""
try:
return axis(Linspace.make(ary.ax))
except ValueError as e:
msg = "To construct data objects, the grid axes must be uniform"
raise ValueError(msg) from e
def _validate_rep_shape(mapping: Mapping[Any, AnyReps], /):
_msg = (
"For data objects, the representation entries must have shape "
"(n_batches, 1, 1, n_features, ...)."
)
for rep in mapping.values():
if rep.entries.shape[1:3] != (1, 1):
raise ValueError(_msg)
def _enforce_uniform_mapping[RepT: AnyReps](
mapping: Mapping[str, RepT], /
) -> Mapping[str, RepT]:
def axis_gen(grid: AnyGrid):
for _axis in grid:
if isinstance(_axis.ax, Linspace):
yield _axis
else:
msg = (
f"Linspace axes expected; found array axis {_axis}"
"Convert the axes to Linspace with `tlt.linsapce_from_array`"
"before constructing the data object."
)
warn_external(msg, category=UserWarning)
yield _enforce_uniform(_axis)
def new_grid(grid: AnyGrid) -> AnyGrid:
_grid = tuple(axis_gen(grid))
if isinstance(grid, Grid2DSparse):
_grid = Grid2DSparse[Axis[Linspace], Axis[Linspace]](
_grid[0],
_grid[1],
sparse_indices=grid.indices,
)
return _grid # pyright: ignore[reportReturnType]
return {
ch: type(rep)(new_grid(rep.grid), rep.entries) for ch, rep in mapping.items()
}
_func_deprecation_msg = (
"The factory function `{deprecated_func_name}` is deprecated "
"and will be removed in 0.8.0;"
" use the factory function {func_name} instead."
)
@overload
def tsdata(
mapping: Mapping[str, reps.TimeSeries[Axis[Linspace]]],
/,
*,
name: str | None = None,
) -> TSData: ...
@overload
def tsdata(
*,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
name: str | None = None,
) -> TSData: ...
[docs]
def tsdata(
mapping: Mapping[str, reps.TimeSeries[Axis[Linspace]]] | None = None,
/,
*,
times: AnyAxis | AxLike | None = None,
entries: AnyArray | None = None,
channels: tuple[str, ...] | None = None,
name: str | None = None,
) -> TSData:
"""Construct a :class:`~types.TSData`.
This function provides **two mutually exclusive** construction ways:
**First**, from a positional-only `mapping` argument:
Parameters
----------
mapping:
A mapping from channel names to :class:`~types.TimeSeries`
with :class:`~types.Linspace` axes.
name: str, optional
Name of the data.
**Second**, from several keyword arguments:
Parameters
----------
times: :class:`~types.misc.Axis`
A uniform array of shape ``(n_times,)`` or a :class:`~types.Linspace`
object representing the time grid of the data.
entries: :class:`~types.misc.AnyArray`
An array of shape ``(n_batch, n_channels, n_harmonics, n_features, Nt)``
where ``Nt`` is the size of ``times``.
See the :external+l2d-interface:attr:`convention <l2d_interface.contract.Representation.entries>`
for more details.
channels:
A tuple of channel names in the data.
name: str, optional
Name of the data.
""" # noqa: E501
if mapping is not None:
_msg_from_mapping = (
"Cannot specify `times`, `entries`, or `channels` "
"when using `mapping` to construct TSData."
)
if not all(arg is None for arg in (times, entries, channels)):
raise ValueError(_msg_from_mapping)
_ = _mixins.validate_maps_to_reps(mapping)
_ = _validate_rep_shape(mapping)
_mapping = _enforce_uniform_mapping(mapping)
return TSData.from_dict(_mapping, name=name)
_msg = (
"Must specify `times`, `entries`, and `channels`"
"when not using `mapping` to construct TSData."
)
if not (times is not None and entries is not None and channels is not None):
raise ValueError(_msg)
_times = axis(times) if not isinstance(times, Axis) else times
_times = _enforce_uniform(_times)
return TSData.from_entries(
times=_times,
entries=entries,
channels=channels,
name=name,
)
@overload
def fsdata(
mapping: Mapping[str, reps.FrequencySeries[Axis[Linspace]]],
/,
*,
times: AnyAxis | AxLike,
name: str | None = None,
) -> TimedFSData: ...
@overload
def fsdata(
mapping: Mapping[str, reps.FrequencySeries[Axis[Linspace]]],
/,
*,
name: str | None = None,
) -> FSData: ...
@overload
def fsdata(
*,
frequencies: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
times: None = None,
name: str | None = None,
) -> FSData: ...
@overload
def fsdata(
*,
frequencies: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
times: AnyAxis | AxLike,
name: str | None = None,
) -> TimedFSData: ...
[docs]
def fsdata(
mapping: Mapping[str, reps.FrequencySeries[Axis[Linspace]]] | None = None,
/,
*,
frequencies: AnyAxis | AxLike | None = None,
entries: AnyArray | None = None,
channels: tuple[str, ...] | None = None,
times: AnyAxis | AxLike | None = None,
name: str | None = None,
):
"""Construct :class:`~types.FSData` or :class:`.TimedFSData`.
This function provides **two mutually exclusive** construction ways:
**First**, from a positional-only `mapping` argument:
Parameters
----------
mapping:
A mapping from channel names to :class:`~types.FrequencySeries`
with :class:`~types.Linspace` axes.
times: :class:`~types.misc.Axis`, optional
A uniform array of shape ``(n_times,)`` or a :class:`~types.Linspace`
object representing the associated time grid of the data.
Returns a :class:`~types.TimedFSData` if provided, otherwise returns
a :class:`~types.FSData`.
name: str, optional
Name of the data.
**Second**, from several keyword arguments:
Parameters
----------
frequencies: :class:`~types.misc.Axis`
A uniform array of shape ``(n_freqs,)`` or a :class:`~types.Linspace`
object representing the frequency grid of the data.
entries: :class:`~types.misc.AnyArray`
A array of shape ``(n_batch, n_channels, n_harmonics, n_features, Nf)``
where ``Nf`` is the size of ``frequencies``.
See the :external+l2d-interface:attr:`convention <l2d_interface.contract.Representation.entries>`
for more details.
channels:
A tuple of channel names in the data.
times: :class:`~types.misc.Axis`, optional
A uniform array of shape ``(n_times,)`` or a :class:`~types.Linspace`
object representing the associated time grid of the data.
Returns a :class:`~types.TimedFSData` if provided, otherwise returns
a :class:`~types.FSData`.
name: str, optional
Name of the data.
Note
----
The associated time grid does not count as a grid dimension of the data,
hence the shape of `entries` does include the time dimension.
""" # noqa: E501
if mapping is not None:
_msg_from_mapping = (
"Cannot specify `frequencies`, `entries`, or `channels` "
"when using `mapping` to construct FSData."
)
if not all(arg is None for arg in (frequencies, entries, channels)):
raise ValueError(_msg_from_mapping)
_ = _mixins.validate_maps_to_reps(mapping)
_ = _validate_rep_shape(mapping)
_mapping = _enforce_uniform_mapping(mapping)
_fsdata = FSData.from_dict(_mapping, name=name)
else:
_msg = (
"Must specify `frequencies`, `entries`, and `channels` when "
"not using `mapping` to construct FSData."
)
if not (
frequencies is not None and entries is not None and channels is not None
):
raise ValueError(_msg)
_freqs = axis(frequencies) if not isinstance(frequencies, Axis) else frequencies
_freqs = _enforce_uniform(_freqs)
_fsdata = FSData.from_entries(
frequencies=_freqs,
entries=entries,
channels=channels,
name=name,
)
if times is not None:
_times = axis(times) if not isinstance(times, Axis) else times
_times = _enforce_uniform(_times)
return _fsdata.set_times(_times)
return _fsdata
[docs]
@deprecated("timed_fsdata", "function", "0.8.0", alternative="fsdata")
def timed_fsdata(
mapping: Mapping[str, reps.FrequencySeries[Axis[Linspace]]],
times: AnyAxis | AxLike,
name: str | None = None,
) -> TimedFSData:
"""Construct :class:`~types.TimedFSData` (*Deprecated*).
Warning
-------
This function is deprecated and will be removed in 0.8.0;
use the :func:`.fsdata` function with the `times` argument instead.
"""
return fsdata(mapping, times=times, name=name)
@overload
def stftdata[GridT: Grid2D[Axis[Linspace], Axis[Linspace]]](
mapping: Mapping[str, reps.STFT[GridT]],
/,
*,
name: str | None = None,
) -> STFTData[GridT]: ...
@overload
def stftdata(
*,
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: None = None,
name: str | None = None,
) -> STFTData[Grid2DCartesian[Axis[Linspace], Axis[Linspace]]]: ...
@overload
def stftdata(
*,
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: AnyArray,
name: str | None = None,
) -> STFTData[Grid2DSparse[Axis[Linspace], Axis[Linspace]]]: ...
[docs]
def stftdata[GridT: Grid2D[Axis[Linspace], Axis[Linspace]]](
mapping: Mapping[str, reps.STFT[GridT]] | None = None,
/,
*,
frequencies: AnyAxis | AxLike | None = None,
times: AnyAxis | AxLike | None = None,
entries: AnyArray | None = None,
channels: tuple[str, ...] | None = None,
sparse_indices: AnyArray | None = None,
name: str | None = None,
):
"""Construct :class:`~types.STFTData`.
This function provides **two mutually exclusive** construction ways:
**First**, from a positional-only `mapping` argument:
Parameters
----------
mapping:
A mapping from channel names to :class:`~types.STFT`
with :class:`~types.misc.Grid2D`
or :class:`~types.Linspace` axes.
name: str, optional
Name of the data.
**Second**, from several keyword arguments:
Parameters
----------
frequencies: :class:`~types.misc.Axis`
A uniform array of shape ``(n_freqs,)`` or a :class:`~types.Linspace`
object representing the frequency grid of the data.
times: :class:`~types.misc.Axis`
A uniform array of shape ``(n_times,)`` or a :class:`~types.Linspace`
object representing the time grid of the data.
entries: :class:`~types.misc.AnyArray`
An array of shape ``(n_batch, n_channels, n_harmonics, n_features, n_freqs, n_times)`` or
``(n_batch, n_channels, n_harmonics, n_features, n_sparse)``. See the
:external+l2d-interface:attr:`convention <l2d_interface.contract.Representation.entries>`
for more details.
channels:
A tuple of channel names in the data.
sparse_indices: :class:`~types.misc.AnyArray`, optional
An array of shape ``(n_sparse, 2)`` containing the indices of the
non-zero entries in the sparse grid.
See :attr:`~types.misc.Grid2DSparse.indices` for more details.
If not provided, the data is treated as dense.
name: str, optional
Name of the data.
""" # noqa: E501
if mapping is not None:
_msg_from_mapping = (
"Cannot specify `frequencies`, `times`, `entries`, `channels`, "
"or `sparse_indices` when using `mapping` to construct STFTData."
)
if not all(
arg is None
for arg in (frequencies, times, entries, channels, sparse_indices)
):
raise ValueError(_msg_from_mapping)
_ = _mixins.validate_maps_to_reps(mapping)
_ = _validate_rep_shape(mapping)
_mapping = _enforce_uniform_mapping(mapping)
return STFTData[GridT].from_dict(_mapping, name=name)
_msg = (
"Must specify `frequencies`, `times`, `entries`, and `channels` "
"when not using `mapping` to construct STFTData."
)
if not (
frequencies is not None
and times is not None
and entries is not None
and channels is not None
):
raise ValueError(_msg)
_freqs = axis(frequencies) if not isinstance(frequencies, Axis) else frequencies
_times = axis(times) if not isinstance(times, Axis) else times
_freqs = _enforce_uniform(_freqs)
_times = _enforce_uniform(_times)
return STFTData[Any].from_entries(
frequencies=_freqs,
times=_times,
entries=entries,
channels=channels,
sparse_indices=sparse_indices,
name=name,
)
@overload
def wdmdata[GridT: Grid2D[Axis[Linspace], Axis[Linspace]]](
mapping: Mapping[str, reps.WDM[GridT]],
/,
*,
name: str | None = None,
) -> WDMData[GridT]: ...
@overload
def wdmdata(
*,
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: None = None,
name: str | None = None,
) -> WDMData[Grid2DCartesian[Axis[Linspace], Axis[Linspace]]]: ...
@overload
def wdmdata(
*,
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: AnyArray,
name: str | None = None,
) -> WDMData[Grid2DSparse[Axis[Linspace], Axis[Linspace]]]: ...
[docs]
def wdmdata[GridT: Grid2D[Axis[Linspace], Axis[Linspace]]](
mapping: Mapping[str, reps.WDM[GridT]] | None = None,
/,
*,
frequencies: AnyAxis | AxLike | None = None,
times: AnyAxis | AxLike | None = None,
entries: AnyArray | None = None,
channels: tuple[str, ...] | None = None,
sparse_indices: AnyArray | None = None,
name: str | None = None,
):
"""Construct :class:`~types.WDMData`.
This function provides **two mutually exclusive** construction ways:
**First**, from a positional-only `mapping` argument:
Parameters
----------
mapping:
A mapping from channel names to :class:`~types.WDM`
with :class:`~types.misc.Grid2D`
or :class:`~types.Linspace` axes.
name: str, optional
Name of the data.
**Second**, from several keyword arguments:
Parameters
----------
frequencies: :class:`~types.misc.Axis`
A uniform array of shape ``(n_freqs,)`` or a :class:`~types.Linspace`
object representing the frequency grid of the data.
times: :class:`~types.misc.Axis`
A uniform array of shape ``(n_times,)`` or a :class:`~types.Linspace`
object representing the time grid of the data.
entries: :class:`~types.misc.AnyArray`
An array of shape ``(n_batch, n_channels, n_harmonics, n_features, n_freqs, n_times)`` or
``(n_batch, n_channels, n_harmonics, n_features, n_sparse)``. See the
:external+l2d-interface:attr:`convention <l2d_interface.contract.Representation.entries>`
for more details.
channels:
A tuple of channel names in the data.
sparse_indices: :class:`~types.misc.AnyArray`, optional
An array of shape ``(n_sparse, 2)`` containing the indices of
the non-zero entries in the sparse grid.
See :attr:`~types.misc.Grid2DSparse.indices` for more details.
If not provided, the data is treated as dense.
name: str, optional
Name of the data.
""" # noqa: E501
if mapping is not None:
_msg_from_mapping = (
"Cannot specify `frequencies`, `times`, `entries`, `channels`, "
"or `sparse_indices` when using `mapping` to construct WDMData."
)
if not all(
arg is None
for arg in (frequencies, times, entries, channels, sparse_indices)
):
raise ValueError(_msg_from_mapping)
_ = _mixins.validate_maps_to_reps(mapping)
_ = _validate_rep_shape(mapping)
_mapping = _enforce_uniform_mapping(mapping)
return WDMData[GridT].from_dict(_mapping, name=name)
_msg = (
"Must specify `frequencies`, `times`, `entries`, and `channels` "
"when not using `mapping` to construct WDMData."
)
if not (
frequencies is not None
and times is not None
and entries is not None
and channels is not None
):
raise ValueError(_msg)
_freqs = axis(frequencies) if not isinstance(frequencies, Axis) else frequencies
_times = axis(times) if not isinstance(times, Axis) else times
_freqs = _enforce_uniform(_freqs)
_times = _enforce_uniform(_times)
return WDMData[Any].from_entries(
frequencies=_freqs,
times=_times,
entries=entries,
channels=channels,
sparse_indices=sparse_indices,
name=name,
)
[docs]
@deprecated("construct_tsdata", "function", "0.8.0", alternative="tsdata")
def construct_tsdata(
*,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
name: str | None = None,
) -> TSData:
"""Construct a :class:`~types.TSData` (*Deprecated*).
Parameters
----------
times: :class:`~types.misc.Axis`
A uniform array of shape ``(n_times,)`` or a :class:`~types.Linspace`
object representing the time grid of the data.
entries: :class:`~types.misc.AnyArray`
An array of shape ``(n_batch, n_channels, n_harmonics, n_features, Nt)``
where ``Nt`` is the size of ``times``.
See the :external+l2d-interface:attr:`convention <l2d_interface.contract.Representation.entries>`
for more details.
channels:
Names of the channels in the data.
name: str, optional
Name of the data.
Warning
-------
This function is deprecated and will be removed in 0.8.0; use the :func:`.tsdata`
function with the keyword arguments instead.
""" # noqa: E501
return tsdata(times=times, entries=entries, channels=channels, name=name)
@overload
def construct_fsdata(
*,
frequencies: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
times: AnyAxis | AxLike,
name: str | None = None,
) -> TimedFSData: ...
@overload
def construct_fsdata(
*,
frequencies: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
name: str | None = None,
) -> FSData: ...
[docs]
@deprecated("construct_fsdata", "function", "0.8.0", alternative="fsdata")
def construct_fsdata(
*,
frequencies: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
name: str | None = None,
times: AnyAxis | AxLike | None = None,
) -> FSData:
"""Construct an :class:`~types.data.FSData` (*Deprecated*).
Parameters
----------
frequencies: :class:`~types.misc.Axis`
A uniform array of shape ``(n_freqs,)`` or a :class:`~types.Linspace`
object representing the frequency grid of the data.
entries: :class:`~types.misc.AnyArray`
An array of shape ``(n_batch, n_channels, n_harmonics, n_features, Nf)``
where ``Nf`` is the size of ``frequencies``.
See the :external+l2d-interface:attr:`convention <l2d_interface.contract.Representation.entries>`
for more details.
channels:
Names of the channels in the data.
name: str, optional
Name of the data.
Warning
-------
This function is deprecated and will be removed in 0.8.0; use the
:func:`.fsdata` function with the keyword arguments instead.
""" # noqa: E501
return fsdata(
frequencies=frequencies,
entries=entries,
channels=channels,
name=name,
times=times,
)
[docs]
@deprecated("construct_timed_fsdata", "function", "0.8.0", alternative="fsdata")
def construct_timed_fsdata(
*,
frequencies: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
times: AnyAxis | AxLike,
name: str | None = None,
) -> TimedFSData:
"""Construct a :class:`~types.data.TimedFSData` (*Deprecated*).
Parameters
----------
frequencies: :class:`~types.misc.Axis`
A uniform array of shape ``(n_freqs,)`` or a :class:`~types.Linspace`
object representing the frequency grid of the data.
entries: :class:`~types.misc.AnyArray`
An array of shape ``(n_batch, n_channels, n_harmonics, n_features, Nf)``
where ``Nf`` is the size of ``frequencies``.
See the :external+l2d-interface:attr:`convention <l2d_interface.contract.Representation.entries>`
for more details.
channels: tuple[str, ...]
Names of the channels in the data.
times: :class:`~types.misc.Axis`
An array of shape ``(n_times,)`` or a :class:`~types.Linspace`
object representing the associated time grid of the data.
name: str | None
Name of the data. Default is ``None``.
Note
----
The associated time grid does not count as a grid dimension of the data,
hence the shape of `entries` does include the time dimension.
Warning
-------
This function is deprecated and will be removed in 0.8.0; use the
:func:`.fsdata` function with the `times` argument instead.
""" # noqa: E501
return fsdata(
frequencies=frequencies,
entries=entries,
channels=channels,
times=times,
name=name,
)
@overload
def construct_stftdata(
*,
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: None = None,
name: str | None = None,
) -> STFTData[Grid2DCartesian[Axis[Linspace], Axis[Linspace]]]: ...
@overload
def construct_stftdata(
*,
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: AnyArray,
name: str | None = None,
) -> STFTData[Grid2DSparse[Axis[Linspace], Axis[Linspace]]]: ...
[docs]
@deprecated("construct_stftdata", "function", "0.8.0", alternative="stftdata")
def construct_stftdata(
*,
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: AnyArray | None = None,
name: str | None = None,
):
"""Construct an :class:`~types.data.STFTData` (*Deprecated*).
Parameters
----------
frequencies: :class:`~types.misc.Axis`
A uniform array of shape ``(n_freqs,)`` or a :class:`~types.Linspace`
object representing the frequency grid of the data.
times: :class:`~types.misc.Axis`
A uniform array of shape ``(n_times,)`` or a :class:`~types.Linspace`
object representing the time grid of the data.
entries: :class:`~types.misc.AnyArray`
An array of shape ``(n_batch, n_channels, n_harmonics, n_features, n_freqs, n_times)`` or
``(n_batch, n_channels, n_harmonics, n_features, n_sparse)``. See the
:external+l2d-interface:attr:`convention <l2d_interface.contract.Representation.entries>`
for more details.
channels: tuple[str, ...]
Names of the channels in the data.
sparse_indices: :class:`~types.misc.AnyArray`, optional
An array of shape ``(n_sparse, 2)`` containing the indices of
the non-zero entries in the sparse grid.
See :attr:`~types.misc.Grid2DSparse.indices` for more details.
If not provided, the data is treated as dense.
name: str, optional
Name of the data.
Warning
-------
This function is deprecated and will be removed in 0.8.0; use the
:func:`.stftdata` function with the keyword arguments instead.
""" # noqa: E501
return stftdata(
frequencies=frequencies,
times=times,
entries=entries,
channels=channels,
sparse_indices=sparse_indices,
name=name,
)
@overload
def construct_wdmdata(
*,
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: None = None,
name: str | None = None,
) -> WDMData[Grid2DCartesian[Axis[Linspace], Axis[Linspace]]]: ...
@overload
def construct_wdmdata(
*,
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: AnyArray,
name: str | None = None,
) -> WDMData[Grid2DSparse[Axis[Linspace], Axis[Linspace]]]: ...
[docs]
@deprecated("construct_wdmdata", "function", "0.8.0", alternative="wdmdata")
def construct_wdmdata(
*,
frequencies: AnyAxis | AxLike,
times: AnyAxis | AxLike,
entries: AnyArray,
channels: tuple[str, ...],
sparse_indices: AnyArray | None = None,
name: str | None = None,
):
"""Construct a :class:`~types.data.WDMData` (*Deprecated*).
Parameters
----------
frequencies: :class:`~types.misc.Axis`
A uniform array of shape ``(n_freqs,)`` or a :class:`~types.Linspace`
object representing the frequency grid of the data.
times: :class:`~types.misc.Axis`
A uniform array of shape ``(n_times,)`` or a :class:`~types.Linspace`
object representing the time grid of the data.
entries: :class:`~types.misc.AnyArray`
An array of shape ``(n_batch, n_channels, n_harmonics, n_features, n_freqs, n_times)`` or
``(n_batch, n_channels, n_harmonics, n_features, n_sparse)``. See the
:external+l2d-interface:attr:`convention <l2d_interface.contract.Representation.entries>`
for more details.
channels:
Names of the channels in the data.
sparse_indices: :class:`~types.misc.AnyArray`, optional
An array of shape ``(n_sparse, 2)`` containing the indices of
the non-zero entries in the sparse grid.
See :attr:`~types.misc.Grid2DSparse.indices` for more details.
If not provided, the data is treated as dense.
name: str, optional
Name of the data. Default is ``None``.
Warning
-------
This function is deprecated and will be removed in 0.8.0; use the
:func:`.wdmdata` function with the keyword arguments instead.
""" # noqa: E501
return wdmdata(
frequencies=frequencies,
times=times,
entries=entries,
channels=channels,
sparse_indices=sparse_indices,
name=name,
)
@overload
def load_data(
file_path: str | pathlib.Path,
*,
domain: Literal["time"],
kind: str | None = None,
sparse: bool = False,
legacy: bool = False,
) -> TSData: ...
@overload
def load_data(
file_path: str | pathlib.Path,
*,
domain: Literal["frequency"],
kind: None = None,
sparse: bool = False,
legacy: bool = False,
) -> FSData: ...
@overload
def load_data(
file_path: str | pathlib.Path,
*,
domain: Literal["frequency"],
kind: Literal["timed"],
sparse: bool = False,
legacy: bool = False,
) -> TimedFSData: ...
@overload
def load_data(
file_path: str | pathlib.Path,
*,
domain: Literal["time-frequency"],
kind: Literal["stft"],
sparse: Literal[False] = False,
legacy: bool = False,
) -> STFTData[Grid2DCartesian[Axis[Linspace], Axis[Linspace]]]: ...
@overload
def load_data(
file_path: str | pathlib.Path,
*,
domain: Literal["time-frequency"],
kind: Literal["stft"],
sparse: Literal[True],
legacy: bool = False,
) -> STFTData[Grid2DSparse[Axis[Linspace], Axis[Linspace]]]: ...
@overload
def load_data(
file_path: str | pathlib.Path,
*,
domain: Literal["time-frequency"],
kind: Literal["wdm"],
sparse: Literal[False] = False,
legacy: bool = False,
) -> WDMData[Grid2DCartesian[Axis[Linspace], Axis[Linspace]]]: ...
@overload
def load_data(
file_path: str | pathlib.Path,
*,
domain: Literal["time-frequency"],
kind: Literal["wdm"],
sparse: Literal[True],
legacy: bool = False,
) -> WDMData[Grid2DSparse[Axis[Linspace], Axis[Linspace]]]: ...
[docs]
def load_data(
file_path: str | pathlib.Path,
*,
domain: Domain | None = None,
kind: str | None = None,
sparse: bool = False,
legacy: bool = False,
):
"""Load the data from a saved HDF5 file."""
if legacy:
msg = (
"The `legacy` mode of `load_data` is deprecated "
"and will be removed in 0.9.0;"
)
warn_external(
msg,
DeprecationWarning,
)
with h5py.File(str(file_path), "r") as f:
data_type = str(f.attrs["type"])
if data_type == "TSData":
return _load_data(TSData, file_path, legacy=legacy)
if data_type == "FSData":
return _load_data(FSData, file_path, legacy=legacy)
if data_type == "TimedFSData":
return _load_data(TimedFSData, file_path, legacy=legacy)
_msg = (
f"Unsupported data type in legacy mode: {data_type}. "
"Only TSData, FSData, and TimedFSData are supported in legacy mode."
)
raise ValueError(_msg)
with h5py.File(str(file_path), "r") as f:
domain_attr = str(f.attrs["domain"])
kind_attr = str(f.attrs["kind"])
expected_kind = str(kind)
if domain is None:
_msg = (
"The `domain` argument of `load_data` is not provided. "
"From 0.8.0, this argument will be required. "
"Currently, it is inferred from the data, "
"but this behavior is deprecated."
)
warn_external(
_msg,
FutureWarning,
)
elif domain_attr != domain:
_msg = (
"The domain of the data does not match the provided `domain` argument. "
)
f"The data has domain {domain_attr}, but `domain` is {domain}."
raise ValueError(_msg)
if kind_attr != expected_kind:
_msg = (
"The kind of the data does not match the provided `kind` argument. "
f"The data has kind {kind_attr}, but `kind` is {kind}."
)
raise ValueError(_msg)
if domain_attr in {"time", "frequency"} and sparse:
_msg = f"Sparse grid is not supported for domain {domain_attr}."
raise ValueError(_msg)
classes = {
("time", "None", False): TSData,
("frequency", "None", False): FSData,
("frequency", "timed", False): TimedFSData,
("time-frequency", "stft", False): STFTData[
Grid2DCartesian[Axis[Linspace], Axis[Linspace]]
],
("time-frequency", "stft", True): STFTData[
Grid2DSparse[Axis[Linspace], Axis[Linspace]]
],
("time-frequency", "wdm", False): WDMData[
Grid2DCartesian[Axis[Linspace], Axis[Linspace]]
],
("time-frequency", "wdm", True): WDMData[
Grid2DSparse[Axis[Linspace], Axis[Linspace]]
],
}
cls = classes.get((domain_attr, kind_attr, sparse))
if cls is None:
_msg = (
"Unsupported combination of domain, kind, "
f"and sparse: domain={domain_attr}, "
f"kind={kind_attr}, sparse={sparse}. "
"Supported combinations are: "
f"{', '.join(str(key) for key in classes)}."
)
raise ValueError(_msg)
return _load_data(cls, file_path)
[docs]
def load_sangria(
file_path: str | pathlib.Path,
name: Literal[
"obs/tdi",
"sky/mbhb/tdi",
"sky/igb/tdi",
"sky/vgb/tdi",
"sky/dgb/tdi",
] = "obs/tdi",
channels: Literal["AE", "AET", "XYZ"] = "AE",
) -> TSData:
"""Load Sangria dataset or Sangria HM dataset."""
from ..shop import aet2xyz, xyz2aet
def transform_channels(data: TSData) -> TSData:
channel_names = tuple(name for name in channels)
if set(channel_names).issubset(data.channel_names):
return data.pick(channel_names)
if set(channel_names).issubset(("X", "Y", "Z")):
return aet2xyz(data).pick(channel_names)
if set(channel_names).issubset(("A", "E", "T")):
return xyz2aet(data).pick(channel_names)
msg = (
"Cannot transform the channels to the requested ones. "
f"The data has channels {data.channel_names}, "
f"but the requested channels are {channel_names}."
)
raise ValueError(msg)
with h5py.File(str(file_path), "r") as fid:
dataset = np.array(fid[name])
if dtype_names := dataset.dtype.names:
domain = dtype_names[0]
channel_names = dtype_names[1:]
# I assume we only need to read LDC data in the time domain.
# When this is not true anymore, we can enhance this function.
if domain == "t":
tsdata = TSData.from_dict(
{
chnname: reps.time_series(
_enforce_uniform(axis(dataset[domain].squeeze())),
entries=dataset[chnname].squeeze()[None, None, None, None, :],
)
for chnname in channel_names
},
)
return transform_channels(tsdata)
msg = "The dataset does not have the expected structure."
raise ValueError(msg)
[docs]
def load_ldc_data(file_path: str | pathlib.Path, *args: Any, **kwargs: Any) -> TSData:
"""Load the LDC dataset."""
# Currently among LDC datasets only Sangria and Sangria HM are supported
return load_sangria(file_path, *args, **kwargs)
[docs]
def load_mojito(
file_path: str | pathlib.Path | Sequence[str | pathlib.Path],
time_interval: tuple[float, float]
| tuple[float, None]
| tuple[None, float]
| tuple[None, None] = (None, None),
*,
relative_time: bool = True,
):
"""Load the Mojito data.
The downloading and cache management of Mojito is handled
by `mojito <https://mojito-e66317.io.esa.int/>`_. Especially,
we recommend using `mojito` API to feed the file path
to this loader. If requested data is missing on the machine,
`mojito` will automatically download it from the server and cache it locally.
See `this page <https://mojito-e66317.io.esa.int/content/pouring/download.html>`_
for more details.
Parameters
----------
file_path: str | pathlib.Path | Sequence[str | pathlib.Path]
The file path(s) to the Mojito data file(s). If a sequence of file
paths is provided, they are combined in a meaningful way (see
`this section <https://mojito-e66317.io.esa.int/content/drinking/reader.html#reading-multiple-files>`_
for more details).
time_interval: tuple[float|None, float|None], optional
A tuple specifying the start and end time (in seconds) of the data segment to
load.
If the start time is ``None``,
it defaults to the start of the data. If the end time is ``None``, it defaults
to the end of the data. If both are ``None``, or if `time_interval`
is not provided, the entire data will be loaded.
We recommend using :func:`~shop.year2second`,
:func:`~shop.month2second`, etc. to convert time units to seconds.
relative_time: bool, optional
Whether the time in the returned data is relative to the start of the data.
Default is ``True``. If ``False``, the time will be absolute, i.e., directly
corresponding to the time in the Mojito data file.
Returns
-------
:class:`~types.TSData`
The loaded data segment, formatted as a :class:`~types.TSData` object.
There are four channels in the returned data: "X", "Y", "Z", and "flag",
where "X", "Y", and "Z" are the TDI observables in Doppler units,
and "flag" is the quality flag of the TDI observables.
The shape of the kernel is ``(1, 4, 1, 1, len(times))``.
Example
-------
>>> import mojito.download
>>> import typed_lisa_tools as tlt
>>>
>>> combined = tlt.load_mojito(mojito.download.download-brick("combined"))
"""
import mojito.reader
_fp = file_path if isinstance(file_path, (str, pathlib.Path)) else list(file_path)
with mojito.reader.MojitoL1File(_fp) as f:
sampling = f.tdis.time_sampling
delta = sampling.t0 if relative_time else 0
tmin = time_interval[0] + delta if time_interval[0] is not None else None
tmax = time_interval[1] + delta if time_interval[1] is not None else None
_slice = sampling.slice_between(tmin, tmax)
times = axis(
linspace_from_step(sampling.t0, sampling.dt, sampling.size)[_slice]
)
xyz = f.tdis.xyz_doppler[_slice, :]
xyz_flags = f.tdis.xyz_flags[_slice]
_xyz_data = tsdata(
times=times, entries=xyz.T[None, :, None, None, :], channels=("X", "Y", "Z")
)
_flag = tsdata(
times=times,
entries=xyz_flags.T[None, None, None, None, :],
channels=("flag",),
)
return type(_xyz_data).from_dict(dict(**_xyz_data, **_flag))
# NOTE Not a good idea to use GW signal container to store orbits,
# considering designing OrbitData
# def load_mojito_orbits(
# file_path: str | pathlib.Path | Sequence[str | pathlib.Path],
# time_interval: tuple[float, float]
# | tuple[float, None]
# | tuple[None, float]
# | tuple[None, None] = (None, None),
# ):
# """Load the orbits from the Mojito data.
# Note
# ----
# To use this loader, install TLT with the `mojito` extra.
# The downloading and cache management of Mojito is handled
# by `mojito <https://mojito-e66317.io.esa.int/>`_. Especially,
# we recommend using `mojito` API to feed the file path
# to this loader. If requested data is missing on the machine,
# `mojito` will automatically download it from the server and cache it locally.
# See `this page <https://mojito-e66317.io.esa.int/content/pouring/download.html>`_
# for more details.
# Parameters
# ----------
# file_path: str | pathlib.Path | Sequence[str | pathlib.Path]
# The file path(s) to the Mojito data file(s). If a sequence of file
# paths is provided, they are combined in a meaningful way (see
# `this section <https://mojito-e66317.io.esa.int/content/drinking/reader.html#reading-multiple-files>`_
# for more details).
# time_interval: tuple[float|None, float|None], optional
# A tuple specifying the start and end time (relative to the start of the data,
# in seconds) of the data segment to load.
# If the start time is ``None``,
# it defaults to the start of the data. If the end time is ``None``, it defaults
# to the end of the data. If both are ``None``, or if `time_interval`
# is not provided, the entire data will be loaded.
# We recommend using :func:`~shop.year2second`,
# :func:`~shop.month2second`, etc. to convert time units to seconds.
# Returns
# -------
# :class:`~types.TSData`
# The loaded orbits, formatted as a :class:`~types.TSData` object.
# There are 18 channels in the returned data, in the format of
# "SAT{satellite_id}_POS{component_id}"and"SAT{satellite_id}_VEL{component_id}",
# where {satellite_id} is 1, 2, or 3, and {component_id} is 1, 2, or 3,
# representing the position and velocity components of the three satellites.
# The shape of the kernel is ``(1, 18, 1, 1, len(times))``.
# Example
# -------
# >>> import mojito.download
# >>> import typed_lisa_tools as tlt
# >>>
# >>> orbits = tlt.load_mojito(mojito.download.download-brick("combined"))
# """
# import mojito.reader
# _fp = file_path if isinstance(file_path, (str, pathlib.Path)) else list(file_path)
# with mojito.reader.MojitoL1File(_fp) as f:
# sampling = f.orbits.time_sampling
# tmin = time_interval[0]+sampling.t0 if time_interval[0] is not None else None
# tmax = time_interval[1]+sampling.t0 if time_interval[1] is not None else None
# _slice = sampling.slice_between(tmin, tmax)
# times = axis(
# linspace_from_step(sampling.t0, sampling.dt, sampling.size)[_slice]
# )
# positions = f.orbits.positions[_slice, :, :] # time, satellite, component
# velocities = f.orbits.velocities[_slice, :, :]
# # Channels will be satellite+component, e.g., "SAT1_POS1", "SAT1_POS2",
# # "SAT1_POS3", "SAT2_POS1", etc.
# def _get_entry(grp: AnyArray, sat: tuple[int, str], comp: tuple[int, str]):
# sat_id, sat_name = sat
# comp_id, comp_name = comp
# key = f"{sat_name}_{comp_name}"
# return {
# key: reps.time_series(
# _enforce_uniform(times),
# entries=grp[:, sat_id, comp_id][None, None, None, None, :],
# )
# }
# registry = [
# _get_entry(positions, sat, comp)
# for sat in enumerate(["SAT1", "SAT2", "SAT3"])
# for comp in enumerate(["POS1", "POS2", "POS3"])
# ] + [
# _get_entry(velocities, sat, comp)
# for sat in enumerate(["SAT1", "SAT2", "SAT3"])
# for comp in enumerate(["VEL1", "VEL2", "VEL3"])
# ]
# mapping = {key: entry for d in registry for key, entry in d.items()}
# return tsdata(mapping)
[docs]
def load_preprocessed_mojito(processed_data: SignalProcessor):
"""Load the data from a preprocessed Mojito data segment.
This loader is intended for use with
`mojito-processor <https://github.com/OllieBurke/MojitoProcessor>`_.
It loads a preprocessed Mojito data segment and formats it into
a :class:`~types.TSData` object.
"""
channel_names = tuple(processed_data.channels)
_data = cast("dict[str, AnyArray]", processed_data.data)
_mapping = {
chnname: reps.time_series(
_enforce_uniform(axis(processed_data.t)),
_data[chnname][None, None, None, None, :],
)
for chnname in channel_names
}
return tsdata(cast("Mapping[str, reps.UniformTimeSeries]", _mapping))