lisaanalysistools 1.0.0__cp312-cp312-macosx_10_9_x86_64.whl

lisatools/datacontainer.py

@@ -0,0 +1,292 @@
+import warnings
+from abc import ABC
+from typing import Any, Tuple, Optional, List
+
+import math
+import numpy as np
+from numpy.typing import ArrayLike
+from scipy import interpolate
+import matplotlib.pyplot as plt
+
+try:
+    import cupy as cp
+
+except (ModuleNotFoundError, ImportError):
+    import numpy as cp
+
+from . import detector as lisa_models
+from .utils.utility import AET, get_array_module
+from .utils.constants import *
+from .stochastic import (
+    StochasticContribution,
+    FittedHyperbolicTangentGalacticForeground,
+)
+from .sensitivity import SensitivityMatrix
+
+
+class DataResidualArray:
+    pass
+
+
+class DataResidualArray:
+    """Container to hold sensitivity information.
+
+    Args:
+        f: Frequency array.
+        sens_mat: Input sensitivity list. The shape of the nested lists should represent the shape of the
+            desired matrix. Each entry in the list must be an array, :class:`Sensitivity`-derived object,
+            or a string corresponding to the :class:`Sensitivity` object.
+        **sens_kwargs: Keyword arguments to pass to :method:`Sensitivity.get_Sn`.
+
+    """
+
+    def __init__(
+        self,
+        data_res_in: List[np.ndarray] | np.ndarray | DataResidualArray,
+        dt: Optional[float] = None,
+        f_arr: Optional[np.ndarray] = None,
+        df: Optional[float] = None,
+        **kwargs: dict,
+    ) -> None:
+        if isinstance(data_res_in, DataResidualArray):
+            for key, item in data_res_in.__dict__.items():
+                setattr(self, key, item)
+
+        else:
+            self._check_inputs(dt=dt, f_arr=f_arr, df=df)
+            self.data_res_arr = data_res_in
+            self._store_time_and_frequency_information(dt=dt, f_arr=f_arr, df=df)
+
+    @property
+    def init_kwargs(self) -> dict:
+        """Initial dt, df, f_arr"""
+        return self._init_kwargs
+
+    @init_kwargs.setter
+    def init_kwargs(self, init_kwargs: dict) -> None:
+        """Set initial kwargs."""
+        self._init_kwargs = init_kwargs
+
+    def _check_inputs(
+        self,
+        dt: Optional[float] = None,
+        f_arr: Optional[np.ndarray] = None,
+        df: Optional[float] = None,
+    ):
+        number_of_none = 0
+
+        number_of_none += 1 if dt is None else 0
+        number_of_none += 1 if f_arr is None else 0
+        number_of_none += 1 if df is None else 0
+
+        if number_of_none == 3:
+            raise ValueError("Must provide either df, dt, or f_arr.")
+
+        elif number_of_none == 1:
+            raise ValueError(
+                "Can only provide one of dt, f_arr, or df. Not more than one."
+            )
+        self.init_kwargs = dict(dt=dt, f_arr=f_arr, df=df)
+
+    def _store_time_and_frequency_information(
+        self,
+        dt: Optional[float] = None,
+        f_arr: Optional[np.ndarray] = None,
+        df: Optional[float] = None,
+    ):
+        if dt is not None:
+            self._dt = dt
+            self._Tobs = self.data_length * dt
+            self._df = 1 / self._Tobs
+            self._fmax = 1 / (2 * dt)
+            self._f_arr = np.fft.rfftfreq(self.data_length, dt)
+
+            # transform data
+            tmp = (
+                get_array_module(self.data_res_arr).fft.rfft(self.data_res_arr, axis=-1)
+                * self._dt
+            )
+            del self._data_res_arr
+            self._data_res_arr = tmp
+            self.data_length = self._data_res_arr.shape[-1]
+
+        elif df is not None:
+            self._df = df
+            self._Tobs = 1 / self._df
+            self._fmax = (self.data_length - 1) * df
+            self._dt = 1 / (2 * self._fmax)
+            self._f_arr = np.arange(0.0, self._fmax, self._df)
+
+        elif f_arr is not None:
+            self._f_arr = f_arr
+            self._fmax = f_arr.max()
+            # constant spacing
+            if np.all(np.diff(f_arr) == np.diff(f_arr)[0]):
+                self._df = np.diff(f_arr)[0].item()
+
+                if f_arr[0] == 0.0:
+                    # could be fft because of constant spacing and f_arr[0] == 0.0
+                    self._Tobs = 1 / self._df
+                    self._dt = 1 / (2 * self._fmax)
+
+                else:
+                    # cannot be fft basis
+                    self._Tobs = None
+                    self._dt = None
+
+            else:
+                self._df = None
+                self._Tobs = None
+                self._dt = None
+
+        if len(self.f_arr) != self.data_length:
+            raise ValueError(
+                "Entered or determined f_arr does not have the same length as the data channel inputs."
+            )
+
+    @property
+    def fmax(self):
+        return self._fmax
+
+    @property
+    def f_arr(self):
+        return self._f_arr
+
+    @property
+    def dt(self):
+        if self._dt is None:
+            raise ValueError("dt cannot be determined from this f_arr input.")
+
+        return self._dt
+
+    @property
+    def Tobs(self):
+        if self._Tobs is None:
+            raise ValueError("Tobs cannot be determined from this f_arr input.")
+
+        return self._Tobs
+
+    @property
+    def df(self):
+        if self._df is None:
+            raise ValueError("df cannot be determined from this f_arr input.")
+
+        return self._df
+
+    @property
+    def frequency_arr(self) -> np.ndarray:
+        return self._f_arr
+
+    @property
+    def data_res_arr(self) -> np.ndarray:
+        return self._data_res_arr
+
+    @data_res_arr.setter
+    def data_res_arr(self, data_res_arr: List[np.ndarray] | np.ndarray) -> None:
+        self._data_res_arr_input = data_res_arr
+
+        if (
+            isinstance(data_res_arr, np.ndarray) or isinstance(data_res_arr, cp.ndarray)
+        ) and data_res_arr.ndim == 1:
+            data_res_arr = [data_res_arr]
+
+        elif (
+            isinstance(data_res_arr, np.ndarray) or isinstance(data_res_arr, cp.ndarray)
+        ) and data_res_arr.ndim == 2:
+            data_res_arr = list(data_res_arr)
+
+        new_out = np.full(len(data_res_arr), None, dtype=object)
+        self.data_length = None
+        for i in range(len(data_res_arr)):
+            current_data = data_res_arr[i]
+            if isinstance(current_data, np.ndarray) or isinstance(
+                current_data, cp.ndarray
+            ):
+                if self.data_length is None:
+                    self.data_length = len(current_data)
+                else:
+                    assert len(current_data) == self.data_length
+
+                new_out[i] = current_data
+            else:
+                raise ValueError
+
+        self.nchannels = len(new_out)
+        self._data_res_arr = np.asarray(list(new_out), dtype=complex)
+
+    def __getitem__(self, index: tuple) -> np.ndarray:
+        return self.data_res_arr[index]
+
+    @property
+    def ndim(self) -> int:
+        return self.data_res_arr.ndim
+
+    def flatten(self) -> np.ndarray:
+        return self.data_res_arr.flatten()
+
+    @property
+    def shape(self) -> tuple:
+        return self.data_res_arr.shape
+
+    def loglog(
+        self,
+        ax: Optional[List[plt.Axes] | plt.Axes] = None,
+        fig: Optional[plt.Figure] = None,
+        inds: Optional[List[int] | int] = None,
+        char_strain: Optional[bool] = False,
+        **kwargs: dict,
+    ) -> Tuple[plt.Figure, plt.Axes]:
+        """Produce a log-log plot of the data.
+
+        Args:
+            ax: Matplotlib Axes objects to add plots. Either a list of Axes objects or a single Axes object.
+            fig: Matplotlib figure object.
+            inds: Integer index to select out which data to add to a single access.
+                A list can be provided if ax is a list. They must be the same length.
+            char_strain: If ``True`` return plot in characteristic strain representation.
+            **kwargs: Keyword arguments to be passed to ``loglog`` function in matplotlib.
+
+        Returns:
+            Matplotlib figure and axes objects in a 2-tuple.
+
+
+        """
+        if ax is None and fig is None:
+            nrows = 1
+            ncols = self.shape[0]
+
+            fig, ax = plt.subplots(nrows, ncols, sharex=True, sharey=True)
+            ax = ax.ravel()
+            inds_list = range(len(ax))
+
+        elif ax is not None:
+            if isinstance(ax, list):
+                assert len(ax) == np.prod(self.shape[:-1])
+                if inds is None:
+                    inds_list = list(np.arange(np.prod(self.shape[:-1])))
+                else:
+                    assert isinstance(inds, list) and len(inds) == len(ax)
+                    inds_list = inds
+
+            elif isinstance(ax, plt.Axes):
+                assert inds is not None and (
+                    isinstance(inds, tuple) or isinstance(inds, int)
+                )
+                ax = [ax]
+                inds_list = [inds]
+
+        elif fig is not None:
+            raise NotImplementedError
+
+        for i, ax_tmp in zip(inds_list, ax):
+            plot_in = np.abs(self.data_res_arr[i])
+            if char_strain:
+                plot_in *= self.frequency_arr
+            ax_tmp.loglog(self.frequency_arr, plot_in, **kwargs)
+
+        return (fig, ax)
+
+    @property
+    def char_strain(self) -> np.ndarray:
+        """Characteristic strain representation of the data."""
+        return np.sqrt(self.f_arr) * np.abs(self.data_res_arr)

lisatools/detector.py

@@ -0,0 +1,410 @@
+import os
+from abc import ABC, abstractmethod
+from typing import Any, List, Tuple, Optional
+from dataclasses import dataclass
+
+import numpy as np
+import h5py
+from scipy import interpolate
+
+from lisatools.cutils.detector import pycppDetector
+
+SC = [1, 2, 3]
+LINKS = [12, 23, 31, 13, 32, 21]
+
+LINEAR_INTERP_TIMESTEP = 600.00  # sec (0.25 hr)
+
+
+class Orbits(ABC):
+    """LISA Orbit Base Class
+
+    Args:
+        filename: File name. File should be in the style of LISAOrbits
+
+
+    """
+
+    def __init__(self, filename: str) -> None:
+        self.filename = filename
+        self._setup()
+        self.configured = False
+
+    @property
+    def LINKS(self) -> List[int]:
+        """Link order."""
+        return LINKS
+
+    @property
+    def SC(self) -> List[int]:
+        """Spacecraft order."""
+        return SC
+
+    @property
+    def link_space_craft_r(self) -> List[int]:
+        """Receiver (first) spacecraft"""
+        return [int(str(link_i)[0]) for link_i in self.LINKS]
+
+    @property
+    def link_space_craft_e(self) -> List[int]:
+        """Sender (second) spacecraft"""
+        return [int(str(link_i)[1]) for link_i in self.LINKS]
+
+    def _setup(self) -> None:
+        with self.open() as f:
+            for key in f.attrs.keys():
+                setattr(self, key + "_base", f.attrs[key])
+
+    @property
+    def filename(self) -> str:
+        """Orbit file name."""
+        return self._filename
+
+    @filename.setter
+    def filename(self, filename: str) -> None:
+        """Set file name."""
+        assert isinstance(filename, str)
+        assert os.path.exists(filename)
+        self._filename = filename
+
+    def open(self) -> h5py.File:
+        """Opens the h5 file in the proper mode.
+
+        Returns:
+            H5 file object: Opened file.
+
+        Raises:
+            RuntimeError: If backend is opened for writing when it is read-only.
+
+        """
+        f = h5py.File(self.filename, "r")
+        return f
+
+    @property
+    def t_base(self) -> np.ndarray:
+        """Light travel times along links from file."""
+        with self.open() as f:
+            t_base = np.arange(self.size_base) * self.dt_base
+        return t_base
+
+    @property
+    def ltt_base(self) -> np.ndarray:
+        """Light travel times along links from file."""
+        with self.open() as f:
+            ltt = f["tcb"]["ltt"][:]
+        return ltt
+
+    @property
+    def n_base(self) -> np.ndarray:
+        """Normal unit vectors towards receiver along links from file."""
+        with self.open() as f:
+            n = f["tcb"]["n"][:]
+        return n
+
+    @property
+    def x_base(self) -> np.ndarray:
+        """Light travel times along links from file."""
+        with self.open() as f:
+            x = f["tcb"]["x"][:]
+        return x
+
+    @property
+    def v_base(self) -> np.ndarray:
+        """Light travel times along links from file."""
+        with self.open() as f:
+            v = f["tcb"]["v"][:]
+        return v
+
+    @property
+    def t(self) -> np.ndarray:
+        """Configured time array."""
+        self._check_configured()
+        return self._t
+
+    @t.setter
+    def t(self, t: np.ndarray):
+        """Set configured time array."""
+        assert isinstance(t, np.ndarray) and t.ndim == 1
+        self._t = t
+
+    @property
+    def ltt(self) -> np.ndarray:
+        """Light travel time."""
+        self._check_configured()
+        return self._ltt
+
+    @ltt.setter
+    def ltt(self, ltt: np.ndarray) -> np.ndarray:
+        """Set light travel time."""
+        assert ltt.shape[0] == len(self.t)
+
+    @property
+    def n(self) -> np.ndarray:
+        """Light travel time."""
+        self._check_configured()
+        return self._n
+
+    @n.setter
+    def n(self, n: np.ndarray) -> np.ndarray:
+        """Set light travel time."""
+        return self._n
+
+    @property
+    def x(self) -> np.ndarray:
+        """Light travel time."""
+        self._check_configured()
+        return self._x
+
+    @x.setter
+    def x(self, x: np.ndarray) -> np.ndarray:
+        """Set light travel time."""
+        return self._x
+
+    @property
+    def v(self) -> np.ndarray:
+        """Light travel time."""
+        self._check_configured()
+        return self._v
+
+    @v.setter
+    def v(self, v: np.ndarray) -> np.ndarray:
+        """Set light travel time."""
+        return self._v
+
+    def configure(
+        self,
+        t_arr: Optional[np.ndarray] = None,
+        dt: Optional[float] = None,
+        linear_interp_setup: Optional[bool] = False,
+    ) -> None:
+        """Configure the orbits to match the signal response generator time basis.
+
+        The base orbits will be scaled up or down as needed using Cubic Spline interpolation.
+        The higherarchy of consideration to each keyword argument if multiple are given:
+        ``linear_interp_setup``, ``t_arr``, ``dt``.
+
+        If nothing is provided, the base points are used.
+
+        Args:
+            t_arr: New time array.
+            dt: New time step. Will take the time duration to be that of the input data.
+            linear_interp_setup: If ``True``, it will create a dense grid designed for linear interpolation with a constant time step.
+
+        """
+
+        x_orig = self.t_base
+
+        if linear_interp_setup:
+            make_cpp = True
+            dt = LINEAR_INTERP_TIMESTEP
+            Tobs = self.t_base[-1]
+            Nobs = int(Tobs / dt)
+            t_arr = np.arange(Nobs) * dt
+            if t_arr[-1] < self.t_base[-1]:
+                t_arr = np.concatenate([t_arr, self.t_base[-1:]])
+        elif t_arr is not None:
+            assert np.all(t_arr >= self.t_base[0]) and np.all(t_arr <= self.t_base[-1])
+
+        elif dt is not None:
+            make_cpp = True
+            Tobs = self.t_base[-1]
+            Nobs = int(Tobs / dt)
+            t_arr = np.arange(Nobs) * dt
+            if t_arr[-1] < self.t_base[-1]:
+                t_arr = np.concatenate([t_arr, self.t_base[-1:]])
+
+        else:
+            make_cpp = False
+            t_arr = self.t_base
+
+        x_new = t_arr.copy()
+        self.t = t_arr.copy()
+
+        for which in ["ltt", "x", "n", "v"]:
+            arr = getattr(self, which + "_base")
+            arr_tmp = arr.reshape(self.size_base, -1)
+            arr_out_tmp = np.zeros((len(x_new), arr_tmp.shape[-1]))
+            for i in range(arr_tmp.shape[-1]):
+                arr_out_tmp[:, i] = interpolate.CubicSpline(x_orig, arr_tmp[:, i])(
+                    x_new
+                )
+            arr_out = arr_out_tmp.reshape((len(x_new),) + arr.shape[1:])
+            setattr(self, "_" + which, arr_out)
+
+        self.configured = True
+
+        lsr = np.asarray(self.link_space_craft_r).copy().astype(np.int32)
+        lse = np.asarray(self.link_space_craft_e).copy().astype(np.int32)
+        ll = np.asarray(self.LINKS).copy().astype(np.int32)
+
+        if make_cpp:
+            self.pycppdetector_args = (
+                dt,
+                len(self.t),
+                self.n.flatten().copy(),
+                self.ltt.flatten().copy(),
+                self.x.flatten().copy(),
+                ll,
+                lsr,
+                lse,
+            )
+            self.dt = dt
+        else:
+            self.pycppdetector_args = None
+            self.dt = dt
+
+    @property
+    def dt(self) -> float:
+        """new time step if it exists"""
+        if self._dt is None:
+            raise ValueError("dt not available for t_arr only.")
+        return self._dt
+
+    @dt.setter
+    def dt(self, dt: float) -> None:
+        self._dt = dt
+
+    @property
+    def pycppdetector(self) -> pycppDetector:
+        """C++ class"""
+        if self._pycppdetector_args is None:
+            raise ValueError(
+                "Asking for c++ class. Need to set linear_interp_setup = True when configuring."
+            )
+        self._pycppdetector = pycppDetector(*self._pycppdetector_args)
+        return self._pycppdetector
+
+    @property
+    def pycppdetector_args(self) -> tuple:
+        return self._pycppdetector_args
+
+    @pycppdetector_args.setter
+    def pycppdetector_args(self, pycppdetector_args: tuple) -> None:
+        self._pycppdetector_args = pycppdetector_args
+
+    @property
+    def size(self) -> int:
+        """Number of time points."""
+        self._check_configured()
+        return len(self.t)
+
+    def _check_configured(self) -> None:
+        if not self.configured:
+            raise ValueError(
+                "Cannot request property. Need to use configure() method first."
+            )
+
+    def get_light_travel_times(
+        self, t: float | np.ndarray, link: int
+    ) -> float | np.ndarray:
+        return self.pycppdetector.get_light_travel_time(t, link)
+
+    def get_normal_unit_vec(self, t: float | np.ndarray, link: int) -> np.ndarray:
+        return self.pycppdetector.get_normal_unit_vec(t, link)
+
+    def get_pos(self, t: float | np.ndarray, sc: int) -> np.ndarray:
+        return self.pycppdetector.get_pos(t, sc)
+
+    @property
+    def ptr(self) -> int:
+        """pointer to c-class"""
+        return self.pycppdetector.ptr
+
+
+class EqualArmlengthOrbits(Orbits):
+    """Equal Armlength Orbits"""
+
+    def __init__(self):
+        # TODO: fix this up
+        super().__init__(
+            "/Users/mlkatz1/Research/LISAanalysistools/examples/equalarmlength-orbits.h5"
+        )
+
+
+class DefaultOrbits(EqualArmlengthOrbits):
+    """Set default orbit class to Equal Arm Length orbits for now."""
+
+    pass
+
+
+@dataclass
+class LISAModelSettings:
+    """Required LISA model settings:
+
+    TODO: rename these
+
+    Args:
+        Soms_d: OMS displacement noise.
+        Sa_a: Acceleration noise.
+        name: Name of model.
+
+    """
+
+    Soms_d: float
+    Sa_a: float
+    orbits: Orbits
+    name: str
+
+
+class LISAModel(LISAModelSettings, ABC):
+    """Model for the LISA Constellation"""
+
+    def __str__(self) -> str:
+        out = "LISA Constellation Configurations Settings:\n"
+        for key, item in self.__dict__.items():
+            out += f"{key}: {item}\n"
+        return out
+
+
+scirdv1 = LISAModel((15.0e-12) ** 2, (3.0e-15) ** 2, DefaultOrbits(), "scirdv1")
+proposal = LISAModel((10.0e-12) ** 2, (3.0e-15) ** 2, DefaultOrbits(), "proposal")
+mrdv1 = LISAModel((10.0e-12) ** 2, (2.4e-15) ** 2, DefaultOrbits(), "mrdv1")
+sangria = LISAModel((10.0e-12) ** 2, (2.4e-15) ** 2, DefaultOrbits(), "sangria")
+
+__stock_list_models__ = [scirdv1, proposal, mrdv1, sangria]
+__stock_list_models_name__ = [tmp.name for tmp in __stock_list_models__]
+
+
+def get_available_default_lisa_models() -> List[LISAModel]:
+    """Get list of default LISA models
+
+    Returns:
+        List of LISA models.
+
+    """
+    return __stock_list_models__
+
+
+def get_default_lisa_model_from_str(model: str) -> LISAModel:
+    """Return a LISA model from a ``str`` input.
+
+    Args:
+        model: Model indicated with a ``str``.
+
+    Returns:
+        LISA model associated to that ``str``.
+
+    """
+    if model not in __stock_list_models_name__:
+        raise ValueError(
+            "Requested string model is not available. See lisatools.detector documentation."
+        )
+    return globals()[model]
+
+
+def check_lisa_model(model: Any) -> LISAModel:
+    """Check input LISA model.
+
+    Args:
+        model: LISA model to check.
+
+    Returns:
+        LISA Model checked. Adjusted from ``str`` if ``str`` input.
+
+    """
+    if isinstance(model, str):
+        model = get_default_lisa_model_from_str(model)
+
+    if not isinstance(model, LISAModel):
+        raise ValueError("model argument not given correctly.")
+
+    return model