lisaanalysistools 1.0.14__cp312-cp312-macosx_11_0_arm64.whl

lisatools/stochastic.py

@@ -0,0 +1,327 @@
+from __future__ import annotations
+import warnings
+from abc import ABC
+from typing import Any, Tuple, Optional, List, Dict
+
+import math
+import numpy as np
+from scipy import interpolate
+
+try:
+    import cupy as cp
+
+except (ModuleNotFoundError, ImportError):
+    import numpy as cp
+
+from . import detector as lisa_models
+from .utils.utility import AET
+from .utils.constants import *
+
+
+class StochasticContribution(ABC):
+    """Base Class for Stochastic Contributions to the PSD."""
+
+    ndim = None
+    added_stochastic_list = []
+
+    @classmethod
+    def _check_ndim(cls, params: np.ndarray | list) -> None:
+        """Check the dimensionality of the parameters matches the model.
+
+        Args:
+            params: Parameters for stochastic model.
+
+        """
+        if cls.ndim is None:
+            raise ValueError(
+                "When subclassing the StochasticContribution class, must set `ndim` as a static attribute."
+            )
+
+        if len(params) != cls.ndim:
+            raise ValueError("length of parameters is not equivalent to class ndim.")
+
+    @classmethod
+    def get_Sh(
+        cls, f: float | np.ndarray, *params: np.ndarray | list, **kwargs: Any
+    ) -> float | np.ndarray:
+        """Calculate the power spectral density of the stochastic contribution.
+
+        Args:
+            f: Frequency array.
+            *params: Parameters for the stochastic model.
+            **kwargs: Keyword arguments for the stochastic model.
+
+        """
+        if len(cls.added_stochastic_list) > 0:
+            cls._check_ndim(params[0])
+        return cls.specific_Sh_function(f, *params, **kwargs)
+
+    @staticmethod
+    def specific_Sh_function(
+        f: float | np.ndarray, *args: Any, **kwargs: Any
+    ) -> float | np.ndarray:
+        """Calculate the power spectral density contained in a stochastic signal contribution.
+
+        Args:
+            f: Frequency array.
+            *args: Any arguments for the function.
+            **kwargs: Any keyword arguments for the function.
+
+        Returns:
+            Power spectral density contained in stochastic signal.
+
+        """
+        raise NotImplementedError
+
+
+class StochasticContributionContainer:
+    """Container for multiple Stochastic Contributions
+
+    Args:
+        stochastic_contribution_dict: Dictionary with multiple Stochastic entries.
+            Keys are the names and values are of type :class:`StochasticContribution`.
+
+    """
+
+    def __init__(
+        self, stochastic_contribution_dict: dict[StochasticContribution]
+    ) -> None:
+        self.stochastic_contribution_dict = stochastic_contribution_dict
+
+    @property
+    def stochastic_contribution_dict(self) -> dict[StochasticContribution]:
+        """Stochastic contribution storage."""
+        return self._stochastic_contribution_dict
+
+    @stochastic_contribution_dict.setter
+    def stochastic_contribution_dict(
+        self, stochastic_contribution_dict: dict[StochasticContribution]
+    ) -> None:
+        """Set stochastic_contribution_dict."""
+        assert isinstance(stochastic_contribution_dict, dict)
+        for key, value in stochastic_contribution_dict.items():
+            if not isinstance(value, StochasticContribution):
+                raise ValueError(
+                    f"Stochastic model {key} is not of type StochasticContribution."
+                )
+        self._stochastic_contribution_dict = stochastic_contribution_dict
+
+    def get_Sh(
+        self, f: float | np.ndarray, params_dict: dict[tuple], kwargs_dict: dict[dict]
+    ) -> np.ndarray:
+        """Calculate Sh for stochastic contribution.
+
+        Args:
+            f: Frequency array.
+            params_dict: Dictionary with keys equivalent to ``self.stochastic_contribution_dict.keys()``.
+                Values are the parameters for each associated model.
+            kwargs_dict: Dictionary with keys equivalent to ``self.stochastic_contribution_dict.keys()``.
+                Values are the keyword argument dicts for each associated model.
+
+        Returns:
+            Stochastic contribution.
+
+        """
+        Sh_out = np.zeros_like(f)
+        for key in params_dict:
+            stochastic_contrib = self.stochastic_contribution_dict[key]
+            Sh_out += stochastic_contrib.get_Sh(
+                f, params_dict[key], **(kwargs_dict.get(key, {}))
+            )
+        return Sh_out
+
+    def __setitem__(self, key: str | int | tuple, val: StochasticContribution) -> None:
+        """Set an item by directly indexing the class object."""
+        self.stochastic_contribution_dict[key] = val
+
+    def __getitem__(self, key: str | int | tuple) -> StochasticContribution:
+        """Get item directly from dictionary."""
+        return self.stochastic_contribution_dict[key]
+
+
+class HyperbolicTangentGalacticForeground(StochasticContribution):
+    """Hyperbolic Tangent-based foreground fitting function."""
+
+    ndim = 5
+
+    @staticmethod
+    def specific_Sh_function(
+        f: float | np.ndarray, amp: float, fk: float, alpha: float, s1: float, s2: float
+    ) -> float | np.ndarray:
+        """Hyperbolic tangent model 1 for the Galaxy foreground noise
+
+        This model for the PSD contribution from the Galactic foreground noise is given by
+
+        .. math::
+
+            S_\\text{gal} = \\frac{A_\\text{gal}}{2}e^{-s_1f^\\alpha}f^{-7/3}\\left[ 1 + \\tanh{\\left(-s_2 (f - f_k)\\right)} \\right],
+
+        where :math:`A_\\text{gal}` is the amplitude of the stochastic signal, :math:`f_k` is the knee frequency at which a bend occurs,
+        math:`\\alpha` is a power law parameter, :math:`s_1` is a slope parameter below the knee,
+        and :math:`s_2` is a slope parameter after the knee.:
+
+        Args:
+            f: Frequency array.
+            amp: Amplitude parameter for the Galaxy.
+            fk: Knee frequency in Hz.
+            alpha: Power law parameter.
+            s1: Slope parameter below knee.
+            s2: Slope parameter above knee.
+
+        Returns:
+            PSD of the Galaxy foreground noise
+
+        """
+        Sgal = (
+            amp
+            * np.exp(-(f**alpha) * s1)
+            * (f ** (-7.0 / 3.0))
+            * 0.5
+            * (1.0 + np.tanh(-(f - fk) * s2))
+        )
+
+        return Sgal
+
+
+class FittedHyperbolicTangentGalacticForeground(HyperbolicTangentGalacticForeground):
+    # TODO: need to verify this is still working
+    ndim = 1
+    amp = 3.26651613e-44
+    alpha = 1.18300266e00
+    # Tobs should be in sec.
+    day = 86400.0
+    month = day * 30.5
+    year = 365.25 * 24.0 * 3600.0  # hard coded for initial fits
+
+    Xobs = [
+        1.0 * day,
+        3.0 * month,
+        6.0 * month,
+        1.0 * year,
+        2.0 * year,
+        4.0 * year,
+        10.0 * year,
+    ]
+    knee = [
+        1.15120924e-02,
+        4.01884128e-03,
+        3.47302482e-03,
+        2.77606177e-03,
+        2.41178384e-03,
+        2.09278117e-03,
+        1.57362626e-03,
+    ]
+    Slope1 = [
+        9.41315118e02,
+        1.36887568e03,
+        1.68729474e03,
+        1.76327234e03,
+        2.32678814e03,
+        3.01430978e03,
+        3.74970124e03,
+    ]
+
+    Slope2 = [
+        1.03239773e02,
+        1.03351646e03,
+        1.62204855e03,
+        1.68631844e03,
+        2.06821665e03,
+        2.95774596e03,
+        3.15199454e03,
+    ]
+    Tmax = 10 * YRSID_SI
+
+    @classmethod
+    def specific_Sh_function(
+        cls, f: float | np.ndarray, Tobs: float
+    ) -> float | np.ndarray:
+        """Fitted hyperbolic tangent model 1 for the Galaxy foreground noise.
+
+        This class fits the parameters for :class:`HyperbolicTangentGalacticForeground`
+        using analytic estimates from (# TODO). The fit is a function of time, so the user
+        inputs ``Tobs``.
+
+        # Sgal_1d = 2.2e-44*np.exp(-(fr**1.2)*0.9e3)*(fr**(-7./3.))*0.5*(1.0 + np.tanh(-(fr-1.4e-2)*0.7e2))
+        # Sgal_3m = 2.2e-44*np.exp(-(fr**1.2)*1.7e3)*(fr**(-7./3.))*0.5*(1.0 + np.tanh(-(fr-4.8e-3)*5.4e2))
+        # Sgal_1y = 2.2e-44*np.exp(-(fr**1.2)*2.2e3)*(fr**(-7./3.))*0.5*(1.0 + np.tanh(-(fr-3.1e-3)*1.3e3))
+        # Sgal_2y = 2.2e-44*np.exp(-(fr**1.2)*2.2e3)*(fr**(-7./3.))*0.5*(1.0 + np.tanh(-(fr-2.3e-3)*1.8e3))
+        # Sgal_4y = 2.2e-44*np.exp(-(fr**1.2)*2.9e3)*(fr**(-7./3.))*0.5*(1.0 + np.tanh(-(fr-2.0e-3)*1.9e3))
+
+        Args:
+            f: Frequency array.
+            Tobs: Observation time in seconds.
+
+        Returns:
+            PSD of the Galaxy foreground noise
+
+        """
+
+        if Tobs > cls.Tmax:
+            raise ValueError(
+                "Tobs is greater than the maximum allowable fit which is 10 years."
+            )
+
+        # Interpolate
+        tck1 = interpolate.splrep(cls.Xobs, cls.Slope1, s=0, k=1)
+        tck2 = interpolate.splrep(cls.Xobs, cls.knee, s=0, k=1)
+        tck3 = interpolate.splrep(cls.Xobs, cls.Slope2, s=0, k=1)
+        s1 = interpolate.splev(Tobs, tck1, der=0).item()
+        fk = interpolate.splev(Tobs, tck2, der=0).item()
+        s2 = interpolate.splev(Tobs, tck3, der=0).item()
+
+        return HyperbolicTangentGalacticForeground.specific_Sh_function(
+            f, cls.amp, fk, cls.alpha, s1, s2
+        )
+
+
+__stock_gb_stochastic_options__ = [
+    "HyperbolicTangentGalacticForeground",
+    "FittedHyperbolicTangentGalacticForeground",
+]
+
+
+def get_stock_gb_stochastic_options() -> List[StochasticContribution]:
+    """Get stock options for stochastic contributions.
+
+    Returns:
+        List of stock stochastic options.
+
+    """
+    return __stock_gb_stochastic_options__
+
+
+def get_default_stochastic_from_str(stochastic: str) -> StochasticContribution:
+    """Return a LISA stochastic from a ``str`` input.
+
+    Args:
+        stochastic: Stochastic contribution indicated with a ``str``.
+
+    Returns:
+        Stochastic contribution associated to that ``str``.
+
+    """
+    if stochastic not in __stock_gb_stochastic_options__:
+        raise ValueError(
+            "Requested string stochastic is not available. See lisatools.stochastic documentation."
+        )
+    return globals()[stochastic]
+
+
+def check_stochastic(stochastic: Any) -> StochasticContribution:
+    """Check input stochastic contribution.
+
+    Args:
+        stochastic: Stochastic contribution to check.
+
+    Returns:
+        Stochastic contribution checked. Adjusted from ``str`` if ``str`` input.
+
+    """
+    if isinstance(stochastic, str):
+        stochastic = get_default_stochastic_from_str(stochastic)
+
+    if not issubclass(stochastic, StochasticContribution):
+        raise ValueError("stochastic argument not given correctly.")
+
+    return stochastic

lisatools/utils/constants.py

@@ -0,0 +1,40 @@
+MSUN_SI = 1.98848e30
+YRSID_SI = 31558149.763545603
+AU_SI = 149597870700.0
+C_SI = 299792458.0
+G_SI = 6.674080e-11
+GMSUN = 1.3271244210789466e20
+MTSUN_SI = 4.925491025873693e-06
+MRSUN_SI = 1476.6250615036158
+PC_SI = 3.0856775814913674e16
+PI = 3.141592653589793238462643383279502884
+PI_2 = 1.570796326794896619231321691639751442
+PI_3 = 1.047197551196597746154214461093167628
+PI_4 = 0.785398163397448309615660845819875721
+SQRTPI = 1.772453850905516027298167483341145183
+SQRTTWOPI = 2.506628274631000502415765284811045253
+INVSQRTPI = 0.564189583547756286948079451560772585
+INVSQRTTWOPI = 0.398942280401432677939946059934381868
+GAMMA = 0.577215664901532860606512090082402431
+SQRT2 = 1.414213562373095048801688724209698079
+SQRT3 = 1.732050807568877293527446341505872367
+SQRT6 = 2.449489742783178098197284074705891392
+INVSQRT2 = 0.707106781186547524400844362104849039
+INVSQRT3 = 0.577350269189625764509148780501957455
+INVSQRT6 = 0.408248290463863016366214012450981898
+F0 = 3.168753578687779e-08
+Omega0 = 1.9909865927683788e-07
+L_SI = 2.5e9
+eorbit = 0.004824185218078991
+ConstOmega = 1.99098659277e-7
+
+#### Armlength
+lisaL = 2.5e9  # LISA's arm meters
+lisaLT = lisaL / C_SI  # LISA's armn in sec
+
+#### Noise levels
+### Optical Metrology System noise
+## Decomposition
+Sloc = (1.7e-12) ** 2  # m^2/Hz
+Ssci = (8.9e-12) ** 2  # m^2/Hz
+Soth = (2.0e-12) ** 2  # m^2/Hz

lisatools/utils/pointeradjust.py

@@ -0,0 +1,106 @@
+# check to see if cupy is available for gpus
+import numpy as np
+
+try:
+    import cupy as xp
+
+    gpu = True
+
+except (ImportError, ModuleNotFoundError) as e:
+    import numpy as np
+
+    gpu = False
+
+
+def wrapper(*args, **kwargs):
+    """Function to convert array and C/C++ class arguments to ptrs
+
+    This function checks the object type. If it is a cupy or numpy array,
+    it will determine its pointer by calling the proper attributes. If you design
+    a Cython class to be passed through python, it must have a :code:`ptr`
+    attribute.
+
+    If you use this function, you must convert input arrays to size_t data type in Cython and
+    then properly cast the pointer as it enters the c++ function. See the
+    Cython codes
+    `here <https://github.com/BlackHolePerturbationToolkit/FastEMRIWaveforms/tree/master/src>`_
+    for examples.
+
+    args:
+        *args (list): list of the arguments for a function.
+        **kwargs (dict): dictionary of keyword arguments to be converted.
+
+    returns:
+        Tuple: (targs, tkwargs) where t indicates target (with pointer values
+            rather than python objects).
+
+    """
+    # declare target containers
+    targs = []
+    tkwargs = {}
+
+    # args first
+    for arg in args:
+        if gpu:
+            # cupy arrays
+            if isinstance(arg, xp.ndarray):
+                targs.append(arg.data.mem.ptr)
+                continue
+
+        # numpy arrays
+        if isinstance(arg, np.ndarray):
+            targs.append(arg.__array_interface__["data"][0])
+            continue
+
+        try:
+            # cython classes
+            targs.append(arg.ptr)
+            continue
+        except AttributeError:
+            # regular argument
+            targs.append(arg)
+
+    # kwargs next
+    for key, arg in kwargs.items():
+        if gpu:
+            # cupy arrays
+            if isinstance(arg, xp.ndarray):
+                tkwargs[key] = arg.data.mem.ptr
+                continue
+
+        if isinstance(arg, np.ndarray):
+            # numpy arrays
+            tkwargs[key] = arg.__array_interface__["data"][0]
+            continue
+
+        try:
+            # cython classes
+            tkwargs[key] = arg.ptr
+            continue
+        except AttributeError:
+            # other arguments
+            tkwargs[key] = arg
+
+    return (targs, tkwargs)
+
+
+def pointer_adjust(func):
+    """Decorator function for cupy/numpy agnostic cython
+
+    This decorator applies :func:`few.utils.utility.wrapper` to functions
+    via the decorator construction.
+
+    If you use this decorator, you must convert input arrays to size_t data type in Cython and
+    then properly cast the pointer as it enters the c++ function. See the
+    Cython codes
+    `here <https://github.com/BlackHolePerturbationToolkit/FastEMRIWaveforms/tree/master/src>`_
+    for examples.
+
+    """
+
+    def func_wrapper(*args, **kwargs):
+        # get pointers
+        targs, tkwargs = wrapper(*args, **kwargs)
+        return func(*targs, **tkwargs)
+
+    return func_wrapper