lisaanalysistools 1.1.20__cp39-cp39-macosx_15_0_arm64.whl

lisatools/sources/__init__.py

@@ -0,0 +1,6 @@
+from .utils import (
+    CalculationController,
+    BBHCalculationController,
+    GBCalculationController,
+    EMRICalculationController,
+)

lisatools/sources/bbh/__init__.py

@@ -0,0 +1 @@
+from .waveform import BBHSNRWaveform

lisatools/sources/bbh/waveform.py

@@ -0,0 +1,106 @@
+import numpy as np
+from typing import Optional, Any, Tuple
+from copy import deepcopy
+
+# imports
+from fastlisaresponse import ResponseWrapper
+from bbhx.waveformbuild import BBHWaveformFD
+from ...utils.constants import *
+
+from ..waveformbase import SNRWaveform
+
+
+class BBHSNRWaveform(SNRWaveform):
+    """Wrapper class for straightforward BBH SNR calculations.
+
+    Calculates it for A and E channels in TDI2.
+
+    Args:
+        bbh_waveform_kwargs: ``amp_phase_kwargs`` for :class:`BBHWaveformFD`.
+        response_kwargs: ``response_kwargs`` for :class:`BBHWaveformFD`.
+
+    """
+
+    def __init__(
+        self,
+        bbh_waveform_kwargs: Optional[dict] = {"run_phenomd": False},
+        response_kwargs: Optional[dict] = {"TDItag": "AET", "tdi2": True},
+    ) -> None:
+
+        if "TDItag" not in response_kwargs:
+            response_kwargs["TDItag"] = "AET"
+
+        if "tdi2" not in response_kwargs:
+            response_kwargs["tdi2"] = True
+
+        # wave generating class
+        self.wave_gen = BBHWaveformFD(
+            amp_phase_kwargs=bbh_waveform_kwargs,
+            response_kwargs=response_kwargs,
+        )
+
+    @property
+    def f_arr(self) -> np.ndarray:
+        """Frequency array."""
+        return self._f_arr
+
+    @f_arr.setter
+    def f_arr(self, f_arr: np.ndarray) -> None:
+        """Set frequency array."""
+        self._f_arr = f_arr
+
+    def __call__(
+        self,
+        *params: Any,
+        return_array: Optional[bool] = False,
+        mf_min: Optional[float] = 1e-4,
+        mf_max: Optional[float] = 0.6,
+        freqs: Optional[np.ndarray] = None,
+        **kwargs: Any
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray] | np.ndarray:
+        """Generate waveforms for SNR calculations.
+
+        Args:
+            *params: Parameters for the ``__call__`` function
+                for :class:`BBHWaveformFD`.
+            return_array: If ``True``, return ``array([A, E, T]).
+                If ``False``, return (A, E, T).
+            mf_min: Minimum dimensionless frequency to evaluate.
+            mf_max: Maximum dimensionless frequency to evaluate.
+            freqs: If ``None``, then default will be ``np.logspace(mf_min / M, mf_max / M, 1024)``.
+                Otherwise, it will calulate frequencies based on this exact array.
+            **kwargs: ``kwargs`` for the ``__call__`` function
+                for :class:`BBHWaveformFD`.
+
+        Returns:
+            Output waveform.
+
+        """
+
+        # determine frequency array (sparse, log-spaced)
+        m1 = params[0]
+        m2 = params[1]
+
+        if freqs is None:
+            min_f = mf_min / (MTSUN_SI * (m1 + m2))
+            max_f = mf_max / (MTSUN_SI * (m1 + m2))
+            self.f_arr = np.logspace(np.log10(min_f), np.log10(max_f), 1024)
+
+        else:
+            assert isinstance(freqs, np.ndarray)
+            self.f_arr = freqs
+
+        # generate waveform with proper settings
+        AET = self.wave_gen(
+            *params,
+            direct=True,
+            combine=True,
+            freqs=self.f_arr,
+            **kwargs,
+        )[0]
+
+        # prepare output
+        if return_array:
+            return AET
+        else:
+            return (AET[0], AET[1], AET[2])

lisatools/sources/defaultresponse.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+
+from ..detector import EqualArmlengthOrbits, Orbits
+
+
+@dataclass
+class DefaultResponseKwargs:
+    """Default response kwargs
+
+    Default response kwargs for
+    
+    ``t0=30000.0``
+    ``order=25``
+    ``tdi="1st generation"``
+    ``tdi_chan="AET"``
+    ``orbits=EqualArmlengthOrbits()``
+
+    """
+    # `fastlisaresponse.ResponseWrapper <https://mikekatz04.github.io/lisa-on-gpu/user/main.html#response-function-wrapper>`_.
+
+
+    t0 = 30000.0
+    order = 25
+    tdi = "1st generation"
+    tdi_chan = "AET"
+    orbits = EqualArmlengthOrbits()
+
+    @classmethod
+    def get_dict(cls) -> dict:
+        """Return default dictionary"""
+        return dict(
+            t0=cls.t0,
+            order=cls.order,
+            tdi=cls.tdi,
+            tdi_chan=cls.tdi_chan,
+            orbits=cls.orbits,
+        )

lisatools/sources/emri/__init__.py

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

lisatools/sources/emri/waveform.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+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 ..waveformbase import AETTDIWaveform
+from fastlisaresponse import ResponseWrapper
+from ...detector import EqualArmlengthOrbits
+
+default_response_kwargs = dict(
+    t0=30000.0,
+    order=25,
+    tdi="1st generation",
+    tdi_chan="AET",
+    orbits=EqualArmlengthOrbits(),
+)
+
+
+class EMRITDIWaveform(AETTDIWaveform):
+    """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] = ("FastKerrEccentricEquatorialFlux",),
+        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
+        self.response = ResponseWrapper(
+            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,
+        )
+
+    @property
+    def dt(self) -> float:
+        """timestep"""
+        return self.response.dt
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        __doc__ = ResponseWrapper.__call__.__doc__
+        return self.response(*args, **kwargs)

lisatools/sources/gb/__init__.py

@@ -0,0 +1 @@
+from .waveform import GBAETWaveform

lisatools/sources/gb/waveform.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import numpy as np
+from typing import Optional, Any, Tuple
+from copy import deepcopy
+
+from few.waveform import GenerateEMRIWaveform
+
+# imports
+from fastlisaresponse import ResponseWrapper
+from lisatools.detector import EqualArmlengthOrbits
+from ..waveformbase import AETTDIWaveform
+
+from gbgpu.gbgpu import GBGPU
+
+
+class GBAETWaveform(AETTDIWaveform):
+    """Stock AET waveform for Galactic Binaries.
+
+    Args:
+        *args: Arguments for :class:`GBGPU`.
+        **kwargs: Arguments for :class:`GBGPU`.
+
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        # wave generating class
+        self.wave_gen = GBGPU(*args, **kwargs)
+
+    @property
+    def f_arr(self) -> np.ndarray:
+        """Frequency array."""
+        return self._f_arr
+
+    @f_arr.setter
+    def f_arr(self, f_arr: np.ndarray) -> None:
+        """Set the frequency array."""
+        self._f_arr = f_arr
+
+    def __call__(
+        self, *params: Any, return_array: Optional[bool] = False, **kwargs: Any
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray] | np.ndarray:
+        """Generate waveform.
+
+        Args:
+            *params: Parameters going into :meth:`GBGPU.run_wave`.
+            return_array: If ``True``, return ``array([A, E, T]).
+                If ``False``, return (A, E, T).
+            **kwargs: Keyword arguments going into :meth:`GBGPU.run_wave`.
+
+        Returns:
+            Output waveform.
+
+        """
+        self.wave_gen.run_wave(
+            *params,
+            **kwargs,
+        )
+
+        # prepare outputs
+        self.f_arr = self.wave_gen.freqs[0]
+        A = self.wave_gen.A[0]
+        E = self.wave_gen.E[0]
+        T = self.wave_gen.X[0]
+
+        if return_array:
+            return np.array([A, E, T])
+        else:
+            return (A, E, T)