lisaanalysistools-cuda11x 1.1.0__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.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,54 @@
+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
+
+######################
+# Physical constants #
+######################
+
+# Mass of Jupiter
+Mjup = 1.898e27
+
+#################
+# LISA constant #
+#################
+
+# Transfer frequency
+fstar = C_SI / (lisaL * 2 * PI)

lisatools/utils/parallelbase.py

@@ -0,0 +1,11 @@
+from typing import Optional, Sequence, TypeVar, Union
+import types
+
+
+from gpubackendtools import ParallelModuleBase
+
+
+class LISAToolsParallelModule(ParallelModuleBase):
+    def __init__(self, force_backend=None):
+        force_backend_in = ('lisatools', force_backend) if isinstance(force_backend, str) else force_backend
+        super().__init__(force_backend_in)

lisatools/utils/utility.py

@@ -0,0 +1,245 @@
+from __future__ import annotations
+from typing import Tuple
+import numpy as np
+
+# from ..sensitivity import get_sensitivity
+
+try:
+    import cupy as cp
+
+except (ModuleNotFoundError, ImportError):
+    import numpy as cp
+
+    pass
+
+
+def get_array_module(arr: np.ndarray | cp.ndarray) -> object:
+    """Return array library of an array (np/cp).
+
+    Args:
+        arr: Numpy or Cupy array.
+
+    """
+    if isinstance(arr, np.ndarray):
+        return np
+    elif isinstance(arr, cp.ndarray):
+        return cp
+    else:
+        raise ValueError("arr must be a numpy or cupy array.")
+
+
+def generate_noise_fd(freqs, df, *sensitivity_args, func=None, **sensitivity_kwargs):
+    if func is None:
+        func = get_sensitivity
+
+    norm = 0.5 * (1.0 / df) ** 0.5
+    psd = func(freqs, *sensitivity_args, **sensitivity_kwargs)
+    noise_to_add = psd ** (1 / 2) * (
+        np.random.normal(0, norm, len(freqs))
+        + 1j * np.random.normal(0, norm, len(freqs))
+    )
+    return noise_to_add
+
+
+def AET(
+    X: float | np.ndarray, Y: float | np.ndarray, Z: float | np.ndarray
+) -> Tuple[float | np.ndarray, float | np.ndarray, float | np.ndarray]:
+    """Transform to AET from XYZ
+
+    .. math::
+
+        A = (Z - X) / \\sqrt(2)
+
+    .. math::
+
+        E = (X - 2Y + Z) / \\sqrt(6)
+
+    .. math::
+
+        T = (X + Y + Z) / \\sqrt(3)
+
+    Args:
+        X: X-channel information.
+        Y: Y-channel information.
+        Z: Z-channel information.
+
+    Returns:
+        A, E, T Channels.
+
+    """
+    return (
+        (Z - X) / np.sqrt(2.0),
+        (X - 2.0 * Y + Z) / np.sqrt(6.0),
+        (X + Y + Z) / np.sqrt(3.0),
+    )
+
+
+def searchsorted2d_vec(a, b, xp=None, gpu=None, **kwargs):
+    if xp is None:
+        xp = np
+    else:
+        try:
+            xp.cuda.runtime.setDevice(gpu)
+        except AttributeError:
+            pass
+
+    m, n = a.shape
+    max_num = xp.maximum(a.max() - a.min(), b.max() - b.min()) + 1
+    r = max_num * xp.arange(a.shape[0])[:, None]
+    p = xp.searchsorted((a + r).ravel(), (b + r).ravel(), **kwargs).reshape(m, -1)
+
+    out = p - n * (xp.arange(m)[:, None])
+    try:
+        xp.cuda.runtime.deviceSynchronize()
+    except AttributeError:
+        pass
+
+    return out
+
+
+def get_groups_from_band_structure(
+    f0, band_edges, f0_2=None, xp=None, num_groups_base=3, fix_f_test=None, inds=None
+):
+    if num_groups_base not in [2, 3, 4]:
+        raise ValueError("num_groups_base must be 2 or 3 or 4.")
+    if xp is None:
+        xp = np
+
+    else:
+        try:
+            xp.cuda.runtime.setDevice(xp.cuda.runtime.getDevice())
+
+        except AttributeError:
+            # it is numpy
+            pass
+
+    if not isinstance(f0, xp.ndarray) or not isinstance(band_edges, xp.ndarray):
+        raise TypeError(
+            "f0 and band_edges must be xp.ndarray with xp as numpy or cupy as given by the xp kwarg."
+        )
+
+    shape = f0.shape
+
+    # remove any above or below bands
+    bad = (f0 < band_edges.min()) | (f0 > band_edges.max())
+
+    band_indices = xp.searchsorted(band_edges, f0.flatten()).reshape(shape) - 1
+
+    # sort the bands in, but keep places with inds_band_indices
+    band_indices_sorted = xp.sort(band_indices, axis=-1)
+    inds_band_indices = xp.argsort(band_indices, axis=-1)
+
+    if f0_2 is not None:
+        assert f0_2.shape == f0.shape
+        band_indices_2 = xp.searchsorted(band_edges, f0_2.flatten()).reshape(shape) - 1
+        band_indices_2_sorted = xp.take_along_axis(
+            band_indices_2, inds_band_indices, axis=-1
+        )
+
+        # very important: ensures the proposed new point is not further than 1 band away.
+        diff = 1 if num_groups_base > 2 else 0
+        keep = (
+            np.abs(band_indices_2_sorted.flatten() - band_indices_sorted.flatten())
+            <= diff
+        )
+        if fix_f_test is not None:
+            keep[fix_f_test.flatten()] = False
+        remove = ~keep
+
+    else:
+        keep = np.ones(np.prod(band_indices_sorted.shape), dtype=bool)
+
+    # temperature index associated with each band
+    temp_inds = xp.repeat(
+        xp.arange(band_indices_sorted.shape[0]), np.prod(band_indices_sorted.shape[1:])
+    )[
+        keep
+    ]  # .reshape(shape)
+
+    # walker index associated with each band
+    walker_inds = (
+        xp.tile(
+            xp.arange(band_indices_sorted.shape[1]),
+            (band_indices_sorted.shape[0], band_indices_sorted.shape[2], 1),
+        )
+        .transpose((0, 2, 1))
+        .flatten()[keep]
+    )
+
+    if f0_2 is not None:
+        temp_inds_remove = xp.repeat(
+            xp.arange(band_indices_sorted.shape[0]),
+            np.prod(band_indices_sorted.shape[1:]),
+        )[
+            remove
+        ]  # .reshape(shape)
+
+        # walker index associated with each band
+        walker_inds_remove = (
+            xp.tile(
+                xp.arange(band_indices_sorted.shape[1]),
+                (band_indices_sorted.shape[0], band_indices_sorted.shape[2], 1),
+            )
+            .transpose((0, 2, 1))
+            .flatten()[remove]
+        )
+        inds_band_indices_remove = inds_band_indices.flatten()[remove]
+
+    # special indexing method
+    band_indices_sorted_special = (
+        band_indices_sorted.flatten()[keep]
+        + int(1e12) * temp_inds
+        + int(1e6) * walker_inds
+    )
+
+    # get the unique special indicators
+    (
+        unique_special,
+        unique_special_start_inds,
+        unique_special_reverse,
+        unique_special_counts,
+    ) = np.unique(
+        band_indices_sorted_special,
+        return_index=True,
+        return_inverse=True,
+        return_counts=True,
+    )
+
+    # this basically makes mini arange setups for each band
+    # the added_contribution for the first unique band index is removed
+    added_contribution = xp.arange(band_indices_sorted_special.shape[0])
+
+    # gets the groups
+    combined = band_indices_sorted_special + added_contribution
+    groups = (
+        combined - (combined[unique_special_start_inds])[unique_special_reverse]
+    )  # .reshape(shape)
+
+    groups_even_odd_tmp = xp.asarray(
+        [
+            (num_groups_base * groups + i)
+            * (band_indices_sorted.flatten()[keep] % num_groups_base == i)
+            for i in range(num_groups_base)
+        ]
+    )
+    groups_even_odd = xp.sum(groups_even_odd_tmp, axis=0)
+
+    groups_out = -2 * xp.ones_like(f0, dtype=int)
+    groups_out[(temp_inds, walker_inds, inds_band_indices.flatten()[keep])] = (
+        groups_even_odd
+    )
+
+    groups_out[bad] = -1
+
+    """if f0_2 is not None and not np.all(keep):
+        fix = (temp_inds_remove, walker_inds_remove, inds_band_indices_remove)
+        fix_2 = band_indices_2[fix]
+        fix_1 = band_indices[fix]"""
+
+    return groups_out
+
+
+autodoc_type_aliases = {
+    "Iterable": "Iterable",
+    "ArrayLike": "ArrayLike",
+}

lisatools_backend_cuda11x/git_version.py

@@ -0,0 +1,7 @@
+"""Metadata deduced from git at build time."""
+
+id: str
+short_id: str
+
+id = ""
+short_id = ""