lisaanalysistools 1.1.19__cp310-cp310-macosx_14_0_arm64.whl

lisatools/cutils/pycppdetector.pyx

@@ -0,0 +1,222 @@
+import numpy as np
+cimport numpy as np
+from libcpp.string cimport string
+from libcpp cimport bool
+from libc.stdint cimport uintptr_t
+
+from gpubackendtools import wrapper
+
+
+cdef class pycppDetector:
+
+    def __cinit__(self, 
+        *args, 
+        **kwargs
+    ):
+        (
+            dt,
+            N, 
+            n_arr,
+            L_arr, 
+            x_arr,
+            links,
+            sc_r, 
+            sc_e,
+            armlength
+        ), tkwargs = wrapper(*args, **kwargs)
+
+        self.dt = dt
+        self.N = N 
+        self.n_arr = n_arr
+        self.L_arr = L_arr 
+        self.x_arr = x_arr
+        self.links = links
+        self.sc_r = sc_r
+        self.sc_e = sc_e
+        self.armlength = armlength
+
+        cdef size_t n_arr_in = n_arr
+        cdef size_t L_arr_in = L_arr
+        cdef size_t x_arr_in = x_arr
+        cdef size_t links_in = links
+        cdef size_t sc_r_in = sc_r
+        cdef size_t sc_e_in = sc_e
+        
+        self.g = new OrbitsWrap(
+            dt,
+            N,
+            <double*> n_arr_in,
+            <double*> L_arr_in, 
+            <double*> x_arr_in, 
+            <int*> links_in, 
+            <int*> sc_r_in, 
+            <int*> sc_e_in,
+            armlength
+        )
+
+    def __dealloc__(self):
+        self.g.dealloc()
+        if self.g:
+            del self.g
+
+    def __reduce__(self):
+        return (rebuild, (self.dt, self.N, self.n_arr, self.L_arr, self.x_arr, self.links, self.sc_r, self.sc_e, self.armlength))
+
+    def get_window(self, t: float) -> int:
+        return self.g.get_window(t)
+
+    def get_normal_unit_vec_single(self, t: float, link: int) -> np.ndarray:
+        cdef VecWrap *out = new VecWrap(0.0, 0.0, 0.0)
+        self.g.get_normal_unit_vec_ptr(out, t, link)
+
+        return np.array([out.x, out.y, out.z])
+
+    def get_normal_unit_vec_arr(self, t: np.ndarray, link: int) -> np.ndarray:
+        output = np.zeros((len(t), 3), dtype=float)
+        assert t.ndim == 1
+        for i in range(len(t)):
+            output[i] = self.get_normal_unit_vec_single(t[i], link)
+
+        return output
+
+    def get_normal_unit_vec(self, t: np.ndarray | float, link: int) -> np.ndarray:
+
+        if isinstance(t, float):
+            return self.get_normal_unit_vec_single(t, link)
+        elif isinstance(t, np.ndarray):
+            return self.get_normal_unit_vec_arr(t, link)
+
+    def get_link_ind(self, link: int) -> int:
+        return self.g.get_link_ind(link)
+
+    def get_sc_ind(self, sc: int) -> int:
+        return self.g.get_sc_ind(sc)
+
+    def get_light_travel_time_single(self, t: float, link: int) -> float:
+        return self.g.get_light_travel_time(t, link)
+
+    def get_light_travel_time_arr(self, t: np.ndarray, link: int) -> np.ndarray:
+        output = np.zeros((len(t),), dtype=float)
+        assert t.ndim == 1
+        for i in range(len(t)):
+            output[i] = self.get_light_travel_time_single(t[i], link)
+
+        return output
+
+    def get_light_travel_time(self, t: np.ndarray | float, link: int) -> np.ndarray:
+
+        if isinstance(t, float):
+            print("t", t)
+            return self.get_light_travel_time_single(t, link)
+        elif isinstance(t, np.ndarray):
+            return self.get_light_travel_time_arr(t, link)
+    
+    def get_pos_single(self, t: float, sc: int) -> np.ndarray:
+  
+        cdef VecWrap *out = new VecWrap(0.0, 0.0, 0.0)
+        self.g.get_pos_ptr(out, t, sc)
+
+        return np.array([out.x, out.y, out.z])
+
+    def get_pos_arr(self, t: np.ndarray, sc: int) -> np.ndarray:
+        output = np.zeros((len(t), 3), dtype=float)
+        assert t.ndim == 1
+        for i in range(len(t)):
+            output[i] = self.get_pos_single(t[i], sc)
+
+        return output
+
+    def get_pos(self, t: np.ndarray | float, sc: int) -> np.ndarray:
+
+        if isinstance(t, float):
+            return self.get_pos_single(t, sc)
+        elif isinstance(t, np.ndarray):
+            return self.get_pos_arr(t, sc)
+
+    def get_pos_arr_wrap(self, *args, **kwargs):
+        (pos_x, pos_y, pos_z, t, sc, num), tkwargs = wrapper(*args, **kwargs)
+
+        cdef size_t pos_x_in = pos_x
+        cdef size_t pos_y_in = pos_y
+        cdef size_t pos_z_in = pos_z
+        cdef size_t t_in = t
+        cdef size_t sc_in = sc
+
+        self.g.get_pos_arr(
+            <double *>pos_x_in,
+            <double *>pos_y_in,
+            <double *>pos_z_in,
+            <double *>t_in,
+            <int *>sc_in,
+            num
+        )
+    
+    def get_normal_unit_vec_arr_wrap(self, *args, **kwargs):
+        (normal_unit_vec_x, normal_unit_vec_y, normal_unit_vec_z, t, sc, num), tkwargs = wrapper(*args, **kwargs)
+
+        cdef size_t normal_unit_vec_x_in = normal_unit_vec_x
+        cdef size_t normal_unit_vec_y_in = normal_unit_vec_y
+        cdef size_t normal_unit_vec_z_in = normal_unit_vec_z
+        cdef size_t t_in = t
+        cdef size_t sc_in = sc
+
+        self.g.get_normal_unit_vec_arr(
+            <double *>normal_unit_vec_x_in,
+            <double *>normal_unit_vec_y_in,
+            <double *>normal_unit_vec_z_in,
+            <double *>t_in,
+            <int *>sc_in,
+            num
+        )
+
+    def get_light_travel_time_arr_wrap(self, *args, **kwargs):
+        (ltt, t, link, num), tkwargs = wrapper(*args, **kwargs)
+
+        cdef size_t ltt_in = ltt
+        cdef size_t t_in = t
+        cdef size_t link_in = link
+
+        self.g.get_light_travel_time_arr(
+            <double *>ltt_in,
+            <double *>t_in,
+            <int *>link_in,
+            num
+        )
+
+    def get_sc_r_from_arr(self, i: int) -> int:
+        return self.g.get_sc_r_from_arr(i)
+
+    def get_sc_e_from_arr(self, i: int) -> int:
+        return self.g.get_sc_e_from_arr(i)
+
+    def get_link_from_arr(self, i: int) -> int:
+        return self.g.get_link_from_arr(i)
+        
+
+    @property
+    def ptr(self) -> long:
+        return <uintptr_t>self.g
+    
+        
+def rebuild(dt,
+    N, 
+    n_arr,
+    L_arr, 
+    x_arr,
+    links,
+    sc_r, 
+    sc_e,
+    armlength
+):
+    c = pycppDetector(
+        dt,
+        N, 
+        n_arr,
+        L_arr, 
+        x_arr,
+        links,
+        sc_r, 
+        sc_e,
+        armlength
+    )
+    return c

lisatools/datacontainer.py

@@ -0,0 +1,312 @@
+from __future__ import annotations
+import warnings
+from abc import ABC
+from typing import Any, Tuple, Optional, List
+
+import math
+import numpy as np
+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 Data, residual, or template information.
+
+    This class abstracts the connection with the sensitivity matrices to make this analysis
+    as generic as possible for the user frontend, while handling
+    special computations in the backend.
+
+    Args:
+        data_res_in: Data, residual, or template input information. Can be a list, numpy array
+            or another :class:`DataResidualArray`.
+        dt: Timestep in seconds.
+        f_arr: Frequency array.
+        df: Delta f in frequency domain.
+        **kwargs: For future compatibility.
+
+    """
+
+    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)
+            xp = get_array_module(self.data_res_arr)
+            self._f_arr = xp.asarray(np.fft.rfftfreq(self.data_length, dt))
+
+            # transform data
+            tmp = xp.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):
+        """Maximum frequency."""
+        return self._fmax
+
+    @property
+    def f_arr(self):
+        """Frequency array."""
+        return self._f_arr
+
+    @property
+    def dt(self):
+        """Time step in seconds."""
+        if self._dt is None:
+            raise ValueError("dt cannot be determined from this f_arr input.")
+
+        return self._dt
+
+    @property
+    def Tobs(self):
+        """Observation time in seconds"""
+        if self._Tobs is None:
+            raise ValueError("Tobs cannot be determined from this f_arr input.")
+
+        return self._Tobs
+
+    @property
+    def df(self):
+        """Delta f in the frequency domain."""
+        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:
+        """Frequency array"""
+        return self._f_arr
+
+    @property
+    def data_res_arr(self) -> np.ndarray:
+        """Actual data residual array"""
+        return self._data_res_arr
+
+    @data_res_arr.setter
+    def data_res_arr(self, data_res_arr: List[np.ndarray] | np.ndarray) -> None:
+        """Set ``data_res_arr``."""
+        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)
+        xp = get_array_module(new_out[0])
+        self._data_res_arr = xp.asarray(list(new_out), dtype=new_out[0].dtype)
+
+    def __getitem__(self, index: tuple) -> np.ndarray:
+        """Index this class directly in ``self.data_res_arr``."""
+        return self.data_res_arr[index]
+
+    def __setitem__(self, index: tuple, value: float | np.ndarray) -> np.ndarray:
+        """Index this class directly in ``self.data_res_arr``."""
+        self.data_res_arr[index] = value
+
+    @property
+    def ndim(self) -> int:
+        """Number of dimensions in the `data_res_arr`."""
+        return self.data_res_arr.ndim
+
+    def flatten(self) -> np.ndarray:
+        """Flatten the ``data_res_arr``."""
+        return self.data_res_arr.flatten()
+
+    @property
+    def shape(self) -> tuple:
+        """Shape of ``data_res_arr``."""
+        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)