lisaanalysistools 1.0.0__cp312-cp312-macosx_10_9_x86_64.whl

lisatools/sources/emri/__init__.py

@@ -0,0 +1 @@
+from .tdiwaveform import EMRITDIWaveform

lisatools/sources/emri/tdiwaveform.py

@@ -0,0 +1,72 @@
+from lisatools.detector import EqualArmlengthOrbits
+import numpy as np
+from typing import Optional, Any
+from copy import deepcopy
+
+from few.waveform import GenerateEMRIWaveform
+
+# imports
+from fastlisaresponse import ResponseWrapper
+from lisatools.detector import EqualArmlengthOrbits
+
+default_response_kwargs = dict(
+    t0=30000.0,
+    order=25,
+    tdi="1st generation",
+    tdi_chan="AET",
+    orbits=EqualArmlengthOrbits(),
+)
+
+
+class EMRITDIWaveform(ResponseWrapper):
+    """Generate EMRI waveforms with the TDI LISA Response
+
+    Args:
+        T: Observation time in years.
+        dt: Time cadence in seconds.
+        emri_waveform_args: Arguments for :class:`GenerateEMRIWaveforms`.
+        emri_waveform_kwargs: Keyword arguments for :class:`GenerateEMRIWaveforms`.
+        response_kwargs: Keyword arguments for :class:`ResponseWrapper`.
+
+    """
+
+    def __init__(
+        self,
+        T: Optional[float] = 1.0,
+        dt: Optional[float] = 10.0,
+        emri_waveform_args: Optional[tuple] = ("FastSchwarzschildEccentricFlux",),
+        emri_waveform_kwargs: Optional[dict] = {},
+        response_kwargs: Optional[dict] = default_response_kwargs,
+    ):
+        # sky parameters in GenerateEMRIWaveform
+        index_lambda = 8
+        index_beta = 7
+
+        for key in default_response_kwargs:
+            response_kwargs[key] = response_kwargs.get(
+                key, default_response_kwargs[key]
+            )
+        gen_wave = GenerateEMRIWaveform(
+            *emri_waveform_args,
+            sum_kwargs=dict(pad_output=True),
+            **emri_waveform_kwargs
+        )
+
+        response_kwargs_in = deepcopy(response_kwargs)
+        # parameters
+        super(EMRITDIWaveform, self).__init__(
+            gen_wave,
+            T,
+            dt,
+            index_lambda,
+            index_beta,
+            flip_hx=True,  # set to True if waveform is h+ - ihx
+            remove_sky_coords=False,
+            is_ecliptic_latitude=False,
+            remove_garbage=True,  # removes the beginning of the signal that has bad information
+            **response_kwargs_in
+        )
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        __doc__ = super(EMRITDIWaveform, self).__call__.__doc__
+        return super(EMRITDIWaveform, self).__call__(*args, **kwargs)

lisatools/stochastic.py

@@ -0,0 +1,291 @@
+import warnings
+from abc import ABC
+from typing import Any, Tuple, Optional, List, Dict
+
+import math
+import numpy as np
+from numpy.typing import ArrayLike
+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: ArrayLike) -> 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: ArrayLike, **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 years.
+
+        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__

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