You are now the owner of this repo.

24966502451

Committed 26 Apr 2026 08:38PM UTC coverage: 84.304% (-0.1%) from 84.419%

Build # 24966502451

Build Type

Pull #114

github

Committed by

thomasckng

Commit Message

fix: refactor navigation tree generation to remove API entries more cleanly

Pull Request Pull Request #114: Merge recent updates from dev branch

Coverage Stats

16 of 24 new or added lines in 1 file covered. (66.67%)

3099 of 3676 relevant lines covered (84.3%)

0.84 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

94.04

/src/ripplegw/interfaces.py

from abc import ABC, abstractmethod

import jax.numpy as jnp
from jaxtyping import Array, Float

from .waveforms.IMRPhenomD import gen_IMRPhenomD_hphc
from .waveforms.IMRPhenomPv2 import gen_IMRPhenomPv2_hphc
from .waveforms.TaylorF2 import gen_TaylorF2_hphc
from .waveforms.IMRPhenomD_NRTidalv2 import gen_IMRPhenomD_NRTidalv2_hphc
from .waveforms.IMRPhenomXAS import gen_IMRPhenomXAS_hphc
from .waveforms.IMRPhenomXAS_NRTidalv3 import gen_IMRPhenomXAS_NRTidalv3_hphc
from .waveforms.SineGaussian import gen_SineGaussian_hphc
from .waveforms.IMRPhenomXPHM import generate_xphm
from .conversions import Mc_eta_to_ms


class Waveform(ABC):
    """Abstract base class for gravitational waveform models.

    Subclasses implement the frequency- (or time-) domain waveform and expose it
    via ``__call__``, returning a dictionary with polarization keys ``"p"`` (plus)
    and ``"c"`` (cross).
    """

    def __init__(self):
        pass

    @property
    @abstractmethod
    def parameter_names(self) -> tuple[str, ...]:
        """Ordered tuple of parameter names required by this waveform model.

        Returns:
            tuple[str, ...]: Parameter names in the order they are consumed,
                matching the keys expected in the ``params`` dict passed to
                ``__call__``.
        """
        raise NotImplementedError(
            "Waveform.parameter_names must be implemented by subclasses"
        )

    @abstractmethod
    def __call__(
        self, axis: Float[Array, " n_freq"], params: dict[str, Float]
    ) -> dict[str, Float[Array, " n_freq"]]:
        """Evaluate the waveform.

        Args:
            axis (Float[Array, " n_freq"]): Frequency (or time) grid.
            params (dict[str, Float]): Source parameter dictionary.

        Returns:
            dict[str, Float[Array, " n_freq"]]: Dictionary with keys ``"p"``
                (plus polarization) and ``"c"`` (cross polarization).
        """
        raise NotImplementedError("Waveform.__call__ must be implemented by subclasses")


class IMRPhenomD(Waveform):
    """IMRPhenomD frequency-domain waveform (non-precessing, aligned spins).

    Attributes:
        f_ref (float): Reference frequency in Hz.
    """

    f_ref: float

    def __init__(self, f_ref: float = 20.0) -> None:
        """
        Args:
            f_ref (float): Reference frequency in Hz. Defaults to 20.0.
        """
        self.f_ref = f_ref

    @property
    def parameter_names(self) -> tuple[str, ...]:
        return ("M_c", "eta", "s1_z", "s2_z", "d_L", "phase_c", "iota")

    def __call__(
        self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
    ) -> dict[str, Float[Array, " n_freq"]]:
        """Evaluate the IMRPhenomD waveform.

        Args:
            frequency (Float[Array, " n_freq"]): Frequency array in Hz.
            params (dict[str, Float]): Source parameters with keys
                ``M_c``, ``eta``, ``s1_z``, ``s2_z``, ``d_L``,
                ``phase_c``, ``iota``.

        Returns:
            dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
                polarizations.
        """
        output = {}
        theta = jnp.array(
            [
                params["M_c"],
                params["eta"],
                params["s1_z"],
                params["s2_z"],
                params["d_L"],
                0,
                params["phase_c"],
                params["iota"],
            ]
        )
        hp, hc = gen_IMRPhenomD_hphc(frequency, theta, self.f_ref)
        output["p"] = hp
        output["c"] = hc
        return output

    def __repr__(self):
        return f"IMRPhenomD(f_ref={self.f_ref})"


class IMRPhenomPv2(Waveform):
    """IMRPhenomPv2 frequency-domain waveform (precessing spins).

    Attributes:
        f_ref (float): Reference frequency in Hz.
    """

    f_ref: float

    def __init__(self, f_ref: float = 20.0) -> None:
        """
        Args:
            f_ref (float): Reference frequency in Hz. Defaults to 20.0.
        """
        self.f_ref = f_ref

    @property
    def parameter_names(self) -> tuple[str, ...]:
        return (
            "M_c",
            "eta",
            "s1_x",
            "s1_y",
            "s1_z",
            "s2_x",
            "s2_y",
            "s2_z",
            "d_L",
            "phase_c",
            "iota",
        )

    def __call__(
        self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
    ) -> dict[str, Float[Array, " n_freq"]]:
        """Evaluate the IMRPhenomPv2 waveform.

        Args:
            frequency (Float[Array, " n_freq"]): Frequency array in Hz.
            params (dict[str, Float]): Source parameters with keys
                ``M_c``, ``eta``, ``s1_x``, ``s1_y``, ``s1_z``,
                ``s2_x``, ``s2_y``, ``s2_z``, ``d_L``, ``phase_c``, ``iota``.

        Returns:
            dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
                polarizations.
        """
        output = {}
        theta = jnp.array(
            [
                params["M_c"],
                params["eta"],
                params["s1_x"],
                params["s1_y"],
                params["s1_z"],
                params["s2_x"],
                params["s2_y"],
                params["s2_z"],
                params["d_L"],
                0,
                params["phase_c"],
                params["iota"],
            ]
        )
        hp, hc = gen_IMRPhenomPv2_hphc(frequency, theta, self.f_ref)
        output["p"] = hp
        output["c"] = hc
        return output

    def __repr__(self):
        return f"IMRPhenomPv2(f_ref={self.f_ref})"


class TaylorF2(Waveform):
    """TaylorF2 post-Newtonian frequency-domain waveform including tidal effects.

    Attributes:
        f_ref (float): Reference frequency in Hz.
        use_lambda_tildes (bool): If True, expects ``lambda_tilde`` and
            ``delta_lambda_tilde``; otherwise expects ``lambda_1`` and ``lambda_2``.
    """

    f_ref: float
    use_lambda_tildes: bool

    def __init__(self, f_ref: float = 20.0, use_lambda_tildes: bool = False) -> None:
        """
        Args:
            f_ref (float): Reference frequency in Hz. Defaults to 20.0.
            use_lambda_tildes (bool): Whether to parameterise tidal deformability
                via ``lambda_tilde`` / ``delta_lambda_tilde`` (as in Eq. 5-6 of
                arXiv:1402.5156) instead of ``lambda_1`` / ``lambda_2``.
                Defaults to False.
        """
        self.f_ref = f_ref
        self.use_lambda_tildes = use_lambda_tildes

    @property
    def parameter_names(self) -> tuple[str, ...]:
        return (
            "M_c",
            "eta",
            "s1_z",
            "s2_z",
            *(
                ("lambda_tilde", "delta_lambda_tilde")
                if self.use_lambda_tildes
                else ("lambda_1", "lambda_2")
            ),
            "d_L",
            "phase_c",
            "iota",
        )

    def __call__(
        self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
    ) -> dict[str, Float[Array, " n_freq"]]:
        """Evaluate the TaylorF2 waveform.

        Args:
            frequency (Float[Array, " n_freq"]): Frequency array in Hz.
            params (dict[str, Float]): Source parameters with keys ``M_c``,
                ``eta``, ``s1_z``, ``s2_z``, ``d_L``, ``phase_c``, ``iota``,
                plus tidal keys depending on ``use_lambda_tildes``.

        Returns:
            dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
                polarizations.
        """
        output = {}

        if self.use_lambda_tildes:
            first_lambda_param = params["lambda_tilde"]
            second_lambda_param = params["delta_lambda_tilde"]
        else:
            first_lambda_param = params["lambda_1"]
            second_lambda_param = params["lambda_2"]

        theta = jnp.array(
            [
                params["M_c"],
                params["eta"],
                params["s1_z"],
                params["s2_z"],
                first_lambda_param,
                second_lambda_param,
                params["d_L"],
                0,
                params["phase_c"],
                params["iota"],
            ]
        )
        hp, hc = gen_TaylorF2_hphc(
            frequency, theta, self.f_ref, use_lambda_tildes=self.use_lambda_tildes
        )
        output["p"] = hp
        output["c"] = hc
        return output

    def __repr__(self):
        return f"TaylorF2(f_ref={self.f_ref})"


class IMRPhenomD_NRTidalv2(Waveform):
    """IMRPhenomD_NRTidalv2 frequency-domain waveform (non-precessing, NRTidalv2 tides).

    Attributes:
        f_ref (float): Reference frequency in Hz.
        use_lambda_tildes (bool): If True, expects ``lambda_tilde`` /
            ``delta_lambda_tilde``; otherwise ``lambda_1`` / ``lambda_2``.
        no_taper (bool): If True, the Planck taper in the amplitude is disabled.
    """

    f_ref: float
    use_lambda_tildes: bool

    def __init__(
        self,
        f_ref: float = 20.0,
        use_lambda_tildes: bool = False,
        no_taper: bool = False,
    ) -> None:
        """
        Args:
            f_ref (float): Reference frequency in Hz. Defaults to 20.0.
            use_lambda_tildes (bool): Whether to parameterise tidal deformability
                via ``lambda_tilde`` / ``delta_lambda_tilde`` (Eq. 5-6 of
                arXiv:1402.5156) instead of ``lambda_1`` / ``lambda_2``.
                Defaults to False.
            no_taper (bool): Whether to remove the Planck taper in the amplitude
                (useful for relative binning runs). Defaults to False.
        """
        self.f_ref = f_ref
        self.use_lambda_tildes = use_lambda_tildes
        self.no_taper = no_taper

    @property
    def parameter_names(self) -> tuple[str, ...]:
        return (
            "M_c",
            "eta",
            "s1_z",
            "s2_z",
            *(
                ("lambda_tilde", "delta_lambda_tilde")
                if self.use_lambda_tildes
                else ("lambda_1", "lambda_2")
            ),
            "d_L",
            "phase_c",
            "iota",
        )

    def __call__(
        self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
    ) -> dict[str, Float[Array, " n_freq"]]:
        """Evaluate the IMRPhenomD_NRTidalv2 waveform.

        Args:
            frequency (Float[Array, " n_freq"]): Frequency array in Hz.
            params (dict[str, Float]): Source parameters with keys ``M_c``,
                ``eta``, ``s1_z``, ``s2_z``, ``d_L``, ``phase_c``, ``iota``,
                plus tidal keys depending on ``use_lambda_tildes``.

        Returns:
            dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
                polarizations.
        """
        output = {}

        if self.use_lambda_tildes:
            first_lambda_param = params["lambda_tilde"]
            second_lambda_param = params["delta_lambda_tilde"]
        else:
            first_lambda_param = params["lambda_1"]
            second_lambda_param = params["lambda_2"]

        theta = jnp.array(
            [
                params["M_c"],
                params["eta"],
                params["s1_z"],
                params["s2_z"],
                first_lambda_param,
                second_lambda_param,
                params["d_L"],
                0,
                params["phase_c"],
                params["iota"],
            ]
        )

        hp, hc = gen_IMRPhenomD_NRTidalv2_hphc(
            frequency,
            theta,
            self.f_ref,
            use_lambda_tildes=self.use_lambda_tildes,
            no_taper=self.no_taper,
        )
        output["p"] = hp
        output["c"] = hc
        return output

    def __repr__(self):
        return f"IMRPhenomD_NRTidalv2(f_ref={self.f_ref})"


class IMRPhenomXAS(Waveform):
    """IMRPhenomXAS frequency-domain waveform (non-precessing, aligned spins, X family).

    Attributes:
        f_ref (float): Reference frequency in Hz.
    """

    f_ref: float

    def __init__(self, f_ref: float = 20.0) -> None:
        """
        Args:
            f_ref (float): Reference frequency in Hz. Defaults to 20.0.
        """
        self.f_ref = f_ref

    @property
    def parameter_names(self) -> tuple[str, ...]:
        return ("M_c", "eta", "s1_z", "s2_z", "d_L", "phase_c", "iota")

    def __call__(
        self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
    ) -> dict[str, Float[Array, " n_freq"]]:
        """Evaluate the IMRPhenomXAS waveform.

        Args:
            frequency (Float[Array, " n_freq"]): Frequency array in Hz.
            params (dict[str, Float]): Source parameters with keys
                ``M_c``, ``eta``, ``s1_z``, ``s2_z``, ``d_L``,
                ``phase_c``, ``iota``.

        Returns:
            dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
                polarizations.
        """
        output = {}
        theta = jnp.array(
            [
                params["M_c"],
                params["eta"],
                params["s1_z"],
                params["s2_z"],
                params["d_L"],
                0,
                params["phase_c"],
                params["iota"],
            ]
        )
        hp, hc = gen_IMRPhenomXAS_hphc(frequency, theta, self.f_ref)
        output["p"] = hp
        output["c"] = hc
        return output

    def __repr__(self):
        return f"IMRPhenomXAS(f_ref={self.f_ref})"


class IMRPhenomXAS_NRTidalv3(Waveform):
    """IMRPhenomXAS_NRTidalv3 frequency-domain waveform (non-precessing, NRTidalv3 tides).

    Attributes:
        f_ref (float): Reference frequency in Hz.
        use_lambda_tildes (bool): If True, expects ``lambda_tilde`` /
            ``delta_lambda_tilde``; otherwise ``lambda_1`` / ``lambda_2``.
        no_taper (bool): If True, the Planck taper in the amplitude is disabled.
    """

    f_ref: float
    use_lambda_tildes: bool
    no_taper: bool

    def __init__(
        self,
        f_ref: float = 20.0,
        use_lambda_tildes: bool = False,
        no_taper: bool = False,
    ) -> None:
        """
        Args:
            f_ref (float): Reference frequency in Hz. Defaults to 20.0.
            use_lambda_tildes (bool): Whether to parameterise tidal deformability
                via ``lambda_tilde`` / ``delta_lambda_tilde`` rather than
                ``lambda_1`` / ``lambda_2``. Defaults to False.
            no_taper (bool): Whether to disable tapering (useful for relative
                binning runs). Defaults to False.
        """
        self.f_ref = f_ref
        self.use_lambda_tildes = use_lambda_tildes
        self.no_taper = no_taper

    @property
    def parameter_names(self) -> tuple[str, ...]:
        return (
            "M_c",
            "eta",
            "s1_z",
            "s2_z",
            *(
                ("lambda_tilde", "delta_lambda_tilde")
                if self.use_lambda_tildes
                else ("lambda_1", "lambda_2")
            ),
            "d_L",
            "phase_c",
            "iota",
        )

    def __call__(
        self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
    ) -> dict[str, Float[Array, " n_freq"]]:
        """Evaluate the IMRPhenomXAS_NRTidalv3 waveform.

        Args:
            frequency (Float[Array, " n_freq"]): Frequency array in Hz.
            params (dict[str, Float]): Source parameters with keys ``M_c``,
                ``eta``, ``s1_z``, ``s2_z``, ``d_L``, ``phase_c``, ``iota``,
                plus tidal keys depending on ``use_lambda_tildes``.

        Returns:
            dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
                polarizations.
        """
        output = {}

        if self.use_lambda_tildes:
            first_lambda_param = params["lambda_tilde"]
            second_lambda_param = params["delta_lambda_tilde"]
        else:
            first_lambda_param = params["lambda_1"]
            second_lambda_param = params["lambda_2"]

        theta = jnp.array(
            [
                params["M_c"],
                params["eta"],
                params["s1_z"],
                params["s2_z"],
                first_lambda_param,
                second_lambda_param,
                params["d_L"],
                0,
                params["phase_c"],
                params["iota"],
            ]
        )
        hp, hc = gen_IMRPhenomXAS_NRTidalv3_hphc(
            frequency,
            theta,
            self.f_ref,
            use_lambda_tildes=self.use_lambda_tildes,
            no_taper=self.no_taper,
        )
        output["p"] = hp
        output["c"] = hc
        return output

    def __repr__(self):
        return f"IMRPhenomXAS_NRTidalv3(f_ref={self.f_ref})"


class IMRPhenomXPHM(Waveform):
    """IMRPhenomXPHM frequency-domain waveform (precessing spins, higher-order modes).

    Attributes:
        f_ref (float): Reference frequency in Hz.
    """

    f_ref: float

    def __init__(self, f_ref: float = 20.0) -> None:
        """
        Args:
            f_ref (float): Reference frequency in Hz. Defaults to 20.0.
        """
        self.f_ref = f_ref

    @property
    def parameter_names(self) -> tuple[str, ...]:
        return (
            "M_c",
            "eta",
            "s1_x",
            "s1_y",
            "s1_z",
            "s2_x",
            "s2_y",
            "s2_z",
            "d_L",
            "phase_c",
            "iota",
        )

    def __call__(
        self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
    ) -> dict[str, Float[Array, " n_freq"]]:
        """Evaluate the IMRPhenomXPHM waveform.

        Args:
            frequency (Float[Array, " n_freq"]): Frequency array in Hz.
            params (dict[str, Float]): Source parameters with keys
                ``M_c``, ``eta``, ``s1_x``, ``s1_y``, ``s1_z``,
                ``s2_x``, ``s2_y``, ``s2_z``, ``d_L``, ``phase_c``, ``iota``.

        Returns:
            dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
                polarizations.
        """
        output = {}
        m1, m2 = Mc_eta_to_ms(jnp.array([params["M_c"], params["eta"]]))
        hp, hc = generate_xphm(
            m1,
            m2,
            params["s1_x"],
            params["s1_y"],
            params["s1_z"],
            params["s2_x"],
            params["s2_y"],
            params["s2_z"],
            params["d_L"],
            params["iota"],
            params["phase_c"],
            frequency,
            self.f_ref,
        )
        output["p"] = hp
        output["c"] = hc
        return output

    def __repr__(self):
        return f"IMRPhenomXPHM(f_ref={self.f_ref})"


class SineGaussian(Waveform):
    """Sine-Gaussian time-domain burst waveform."""

    def __init__(self) -> None:
        pass

    @property
    def parameter_names(self) -> tuple[str, ...]:
        return ("Q", "f_0", "hrss", "phase", "e")

    def __call__(
        self, t: Float[Array, " n_time"], params: dict[str, Float]
    ) -> dict[str, Float[Array, " n_time"]]:
        """
        Args:
            t: Time grid centered at t=0. Create using
               ``jnp.arange(-duration/2, duration/2, 1/fs)``.
            params: Dictionary with keys ``Q`` (quality factor), ``f_0``
                (central frequency in Hz), ``hrss``, ``phase`` (phase),
                ``e`` (eccentricity).
        """
        output = {}
        theta = jnp.array(
            [
                params["Q"],
                params["f_0"],
                params["hrss"],
                params["phase"],
                params["e"],
            ]
        )
        hp, hc = gen_SineGaussian_hphc(t, theta)
        output["p"] = hp
        output["c"] = hc
        return output

    def __repr__(self):
        return "SineGaussian()"


#: Mapping from model name strings to :class:`Waveform` subclasses.
#: Useful for selecting waveform models by name at runtime, e.g. from a
#: configuration file.
waveform_preset: dict[str, type[Waveform]] = {
    "IMRPhenomD": IMRPhenomD,
    "IMRPhenomPv2": IMRPhenomPv2,
    "TaylorF2": TaylorF2,
    "IMRPhenomD_NRTidalv2": IMRPhenomD_NRTidalv2,
    "IMRPhenomXAS": IMRPhenomXAS,
    "IMRPhenomXAS_NRTidalv3": IMRPhenomXAS_NRTidalv3,
    "IMRPhenomXPHM": IMRPhenomXPHM,
    "SineGaussian": SineGaussian,
}

1	from abc import ABC, abstractmethod	1✔
2
3	import jax.numpy as jnp	1✔
4	from jaxtyping import Array, Float	1✔
5
6	from .waveforms.IMRPhenomD import gen_IMRPhenomD_hphc	1✔
7	from .waveforms.IMRPhenomPv2 import gen_IMRPhenomPv2_hphc	1✔
8	from .waveforms.TaylorF2 import gen_TaylorF2_hphc	1✔
9	from .waveforms.IMRPhenomD_NRTidalv2 import gen_IMRPhenomD_NRTidalv2_hphc	1✔
10	from .waveforms.IMRPhenomXAS import gen_IMRPhenomXAS_hphc	1✔
11	from .waveforms.IMRPhenomXAS_NRTidalv3 import gen_IMRPhenomXAS_NRTidalv3_hphc	1✔
12	from .waveforms.SineGaussian import gen_SineGaussian_hphc	1✔
13	from .waveforms.IMRPhenomXPHM import generate_xphm	1✔
14	from .conversions import Mc_eta_to_ms	1✔
15
16
17	class Waveform(ABC):	1✔
18	"""Abstract base class for gravitational waveform models.
19
20	Subclasses implement the frequency- (or time-) domain waveform and expose it
21	via ``__call__``, returning a dictionary with polarization keys ``"p"`` (plus)
22	and ``"c"`` (cross).
23	"""
24
25	def __init__(self):	1✔
26	pass	×
27
28	@property
29	@abstractmethod
30	def parameter_names(self) -> tuple[str, ...]:
31	"""Ordered tuple of parameter names required by this waveform model.
32
33	Returns:
34	tuple[str, ...]: Parameter names in the order they are consumed,
35	matching the keys expected in the ``params`` dict passed to
36	``__call__``.
37	"""
38	raise NotImplementedError(
39	"Waveform.parameter_names must be implemented by subclasses"
40	)
41
42	@abstractmethod
43	def __call__(
44	self, axis: Float[Array, " n_freq"], params: dict[str, Float]
45	) -> dict[str, Float[Array, " n_freq"]]:
46	"""Evaluate the waveform.
47
48	Args:
49	axis (Float[Array, " n_freq"]): Frequency (or time) grid.
50	params (dict[str, Float]): Source parameter dictionary.
51
52	Returns:
53	dict[str, Float[Array, " n_freq"]]: Dictionary with keys ``"p"``
54	(plus polarization) and ``"c"`` (cross polarization).
55	"""
56	raise NotImplementedError("Waveform.__call__ must be implemented by subclasses")
57
58
59	class IMRPhenomD(Waveform):	1✔
60	"""IMRPhenomD frequency-domain waveform (non-precessing, aligned spins).
61
62	Attributes:
63	f_ref (float): Reference frequency in Hz.
64	"""
65
66	f_ref: float	1✔
67
68	def __init__(self, f_ref: float = 20.0) -> None:	1✔
69	"""
70	Args:
71	f_ref (float): Reference frequency in Hz. Defaults to 20.0.
72	"""
73	self.f_ref = f_ref	1✔
74
75	@property	1✔
76	def parameter_names(self) -> tuple[str, ...]:	1✔
NEW 77	return ("M_c", "eta", "s1_z", "s2_z", "d_L", "phase_c", "iota")	×
78
79	def __call__(	1✔
80	self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
81	) -> dict[str, Float[Array, " n_freq"]]:
82	"""Evaluate the IMRPhenomD waveform.
83
84	Args:
85	frequency (Float[Array, " n_freq"]): Frequency array in Hz.
86	params (dict[str, Float]): Source parameters with keys
87	``M_c``, ``eta``, ``s1_z``, ``s2_z``, ``d_L``,
88	``phase_c``, ``iota``.
89
90	Returns:
91	dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
92	polarizations.
93	"""
94	output = {}	1✔
95	theta = jnp.array(	1✔
96	[
97	params["M_c"],
98	params["eta"],
99	params["s1_z"],
100	params["s2_z"],
101	params["d_L"],
102	0,
103	params["phase_c"],
104	params["iota"],
105	]
106	)
107	hp, hc = gen_IMRPhenomD_hphc(frequency, theta, self.f_ref)	1✔
108	output["p"] = hp	1✔
109	output["c"] = hc	1✔
110	return output	1✔
111
112	def __repr__(self):
113	return f"IMRPhenomD(f_ref={self.f_ref})"
114
115
116	class IMRPhenomPv2(Waveform):	1✔
117	"""IMRPhenomPv2 frequency-domain waveform (precessing spins).
118
119	Attributes:
120	f_ref (float): Reference frequency in Hz.
121	"""
122
123	f_ref: float	1✔
124
125	def __init__(self, f_ref: float = 20.0) -> None:	1✔
126	"""
127	Args:
128	f_ref (float): Reference frequency in Hz. Defaults to 20.0.
129	"""
130	self.f_ref = f_ref	1✔
131
132	@property	1✔
133	def parameter_names(self) -> tuple[str, ...]:	1✔
NEW 134	return (	×
135	"M_c",
136	"eta",
137	"s1_x",
138	"s1_y",
139	"s1_z",
140	"s2_x",
141	"s2_y",
142	"s2_z",
143	"d_L",
144	"phase_c",
145	"iota",
146	)
147
148	def __call__(	1✔
149	self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
150	) -> dict[str, Float[Array, " n_freq"]]:
151	"""Evaluate the IMRPhenomPv2 waveform.
152
153	Args:
154	frequency (Float[Array, " n_freq"]): Frequency array in Hz.
155	params (dict[str, Float]): Source parameters with keys
156	``M_c``, ``eta``, ``s1_x``, ``s1_y``, ``s1_z``,
157	``s2_x``, ``s2_y``, ``s2_z``, ``d_L``, ``phase_c``, ``iota``.
158
159	Returns:
160	dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
161	polarizations.
162	"""
163	output = {}	1✔
164	theta = jnp.array(	1✔
165	[
166	params["M_c"],
167	params["eta"],
168	params["s1_x"],
169	params["s1_y"],
170	params["s1_z"],
171	params["s2_x"],
172	params["s2_y"],
173	params["s2_z"],
174	params["d_L"],
175	0,
176	params["phase_c"],
177	params["iota"],
178	]
179	)
180	hp, hc = gen_IMRPhenomPv2_hphc(frequency, theta, self.f_ref)	1✔
181	output["p"] = hp	1✔
182	output["c"] = hc	1✔
183	return output	1✔
184
185	def __repr__(self):
186	return f"IMRPhenomPv2(f_ref={self.f_ref})"
187
188
189	class TaylorF2(Waveform):	1✔
190	"""TaylorF2 post-Newtonian frequency-domain waveform including tidal effects.
191
192	Attributes:
193	f_ref (float): Reference frequency in Hz.
194	use_lambda_tildes (bool): If True, expects ``lambda_tilde`` and
195	``delta_lambda_tilde``; otherwise expects ``lambda_1`` and ``lambda_2``.
196	"""
197
198	f_ref: float	1✔
199	use_lambda_tildes: bool	1✔
200
201	def __init__(self, f_ref: float = 20.0, use_lambda_tildes: bool = False) -> None:	1✔
202	"""
203	Args:
204	f_ref (float): Reference frequency in Hz. Defaults to 20.0.
205	use_lambda_tildes (bool): Whether to parameterise tidal deformability
206	via ``lambda_tilde`` / ``delta_lambda_tilde`` (as in Eq. 5-6 of
207	arXiv:1402.5156) instead of ``lambda_1`` / ``lambda_2``.
208	Defaults to False.
209	"""
210	self.f_ref = f_ref	1✔
211	self.use_lambda_tildes = use_lambda_tildes	1✔
212
213	@property	1✔
214	def parameter_names(self) -> tuple[str, ...]:	1✔
NEW 215	return (	×
216	"M_c",
217	"eta",
218	"s1_z",
219	"s2_z",
220	*(
221	("lambda_tilde", "delta_lambda_tilde")
222	if self.use_lambda_tildes
223	else ("lambda_1", "lambda_2")
224	),
225	"d_L",
226	"phase_c",
227	"iota",
228	)
229
230	def __call__(	1✔
231	self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
232	) -> dict[str, Float[Array, " n_freq"]]:
233	"""Evaluate the TaylorF2 waveform.
234
235	Args:
236	frequency (Float[Array, " n_freq"]): Frequency array in Hz.
237	params (dict[str, Float]): Source parameters with keys ``M_c``,
238	``eta``, ``s1_z``, ``s2_z``, ``d_L``, ``phase_c``, ``iota``,
239	plus tidal keys depending on ``use_lambda_tildes``.
240
241	Returns:
242	dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
243	polarizations.
244	"""
245	output = {}	1✔
246
247	if self.use_lambda_tildes:	1✔
248	first_lambda_param = params["lambda_tilde"]	1✔
249	second_lambda_param = params["delta_lambda_tilde"]	1✔
250	else:
251	first_lambda_param = params["lambda_1"]	1✔
252	second_lambda_param = params["lambda_2"]	1✔
253
254	theta = jnp.array(	1✔
255	[
256	params["M_c"],
257	params["eta"],
258	params["s1_z"],
259	params["s2_z"],
260	first_lambda_param,
261	second_lambda_param,
262	params["d_L"],
263	0,
264	params["phase_c"],
265	params["iota"],
266	]
267	)
268	hp, hc = gen_TaylorF2_hphc(	1✔
269	frequency, theta, self.f_ref, use_lambda_tildes=self.use_lambda_tildes
270	)
271	output["p"] = hp	1✔
272	output["c"] = hc	1✔
273	return output	1✔
274
275	def __repr__(self):
276	return f"TaylorF2(f_ref={self.f_ref})"
277
278
279	class IMRPhenomD_NRTidalv2(Waveform):	1✔
280	"""IMRPhenomD_NRTidalv2 frequency-domain waveform (non-precessing, NRTidalv2 tides).
281
282	Attributes:
283	f_ref (float): Reference frequency in Hz.
284	use_lambda_tildes (bool): If True, expects ``lambda_tilde`` /
285	``delta_lambda_tilde``; otherwise ``lambda_1`` / ``lambda_2``.
286	no_taper (bool): If True, the Planck taper in the amplitude is disabled.
287	"""
288
289	f_ref: float	1✔
290	use_lambda_tildes: bool	1✔
291
292	def __init__(	1✔
293	self,
294	f_ref: float = 20.0,
295	use_lambda_tildes: bool = False,
296	no_taper: bool = False,
297	) -> None:
298	"""
299	Args:
300	f_ref (float): Reference frequency in Hz. Defaults to 20.0.
301	use_lambda_tildes (bool): Whether to parameterise tidal deformability
302	via ``lambda_tilde`` / ``delta_lambda_tilde`` (Eq. 5-6 of
303	arXiv:1402.5156) instead of ``lambda_1`` / ``lambda_2``.
304	Defaults to False.
305	no_taper (bool): Whether to remove the Planck taper in the amplitude
306	(useful for relative binning runs). Defaults to False.
307	"""
308	self.f_ref = f_ref	1✔
309	self.use_lambda_tildes = use_lambda_tildes	1✔
310	self.no_taper = no_taper	1✔
311
312	@property	1✔
313	def parameter_names(self) -> tuple[str, ...]:	1✔
NEW 314	return (	×
315	"M_c",
316	"eta",
317	"s1_z",
318	"s2_z",
319	*(
320	("lambda_tilde", "delta_lambda_tilde")
321	if self.use_lambda_tildes
322	else ("lambda_1", "lambda_2")
323	),
324	"d_L",
325	"phase_c",
326	"iota",
327	)
328
329	def __call__(	1✔
330	self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
331	) -> dict[str, Float[Array, " n_freq"]]:
332	"""Evaluate the IMRPhenomD_NRTidalv2 waveform.
333
334	Args:
335	frequency (Float[Array, " n_freq"]): Frequency array in Hz.
336	params (dict[str, Float]): Source parameters with keys ``M_c``,
337	``eta``, ``s1_z``, ``s2_z``, ``d_L``, ``phase_c``, ``iota``,
338	plus tidal keys depending on ``use_lambda_tildes``.
339
340	Returns:
341	dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
342	polarizations.
343	"""
344	output = {}	1✔
345
346	if self.use_lambda_tildes:	1✔
347	first_lambda_param = params["lambda_tilde"]	1✔
348	second_lambda_param = params["delta_lambda_tilde"]	1✔
349	else:
350	first_lambda_param = params["lambda_1"]	1✔
351	second_lambda_param = params["lambda_2"]	1✔
352
353	theta = jnp.array(	1✔
354	[
355	params["M_c"],
356	params["eta"],
357	params["s1_z"],
358	params["s2_z"],
359	first_lambda_param,
360	second_lambda_param,
361	params["d_L"],
362	0,
363	params["phase_c"],
364	params["iota"],
365	]
366	)
367
368	hp, hc = gen_IMRPhenomD_NRTidalv2_hphc(	1✔
369	frequency,
370	theta,
371	self.f_ref,
372	use_lambda_tildes=self.use_lambda_tildes,
373	no_taper=self.no_taper,
374	)
375	output["p"] = hp	1✔
376	output["c"] = hc	1✔
377	return output	1✔
378
379	def __repr__(self):
380	return f"IMRPhenomD_NRTidalv2(f_ref={self.f_ref})"
381
382
383	class IMRPhenomXAS(Waveform):	1✔
384	"""IMRPhenomXAS frequency-domain waveform (non-precessing, aligned spins, X family).
385
386	Attributes:
387	f_ref (float): Reference frequency in Hz.
388	"""
389
390	f_ref: float	1✔
391
392	def __init__(self, f_ref: float = 20.0) -> None:	1✔
393	"""
394	Args:
395	f_ref (float): Reference frequency in Hz. Defaults to 20.0.
396	"""
397	self.f_ref = f_ref	1✔
398
399	@property	1✔
400	def parameter_names(self) -> tuple[str, ...]:	1✔
NEW 401	return ("M_c", "eta", "s1_z", "s2_z", "d_L", "phase_c", "iota")	×
402
403	def __call__(	1✔
404	self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
405	) -> dict[str, Float[Array, " n_freq"]]:
406	"""Evaluate the IMRPhenomXAS waveform.
407
408	Args:
409	frequency (Float[Array, " n_freq"]): Frequency array in Hz.
410	params (dict[str, Float]): Source parameters with keys
411	``M_c``, ``eta``, ``s1_z``, ``s2_z``, ``d_L``,
412	``phase_c``, ``iota``.
413
414	Returns:
415	dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
416	polarizations.
417	"""
418	output = {}	1✔
419	theta = jnp.array(	1✔
420	[
421	params["M_c"],
422	params["eta"],
423	params["s1_z"],
424	params["s2_z"],
425	params["d_L"],
426	0,
427	params["phase_c"],
428	params["iota"],
429	]
430	)
431	hp, hc = gen_IMRPhenomXAS_hphc(frequency, theta, self.f_ref)	1✔
432	output["p"] = hp	1✔
433	output["c"] = hc	1✔
434	return output	1✔
435
436	def __repr__(self):
437	return f"IMRPhenomXAS(f_ref={self.f_ref})"
438
439
440	class IMRPhenomXAS_NRTidalv3(Waveform):	1✔
441	"""IMRPhenomXAS_NRTidalv3 frequency-domain waveform (non-precessing, NRTidalv3 tides).
442
443	Attributes:
444	f_ref (float): Reference frequency in Hz.
445	use_lambda_tildes (bool): If True, expects ``lambda_tilde`` /
446	``delta_lambda_tilde``; otherwise ``lambda_1`` / ``lambda_2``.
447	no_taper (bool): If True, the Planck taper in the amplitude is disabled.
448	"""
449
450	f_ref: float	1✔
451	use_lambda_tildes: bool	1✔
452	no_taper: bool	1✔
453
454	def __init__(	1✔
455	self,
456	f_ref: float = 20.0,
457	use_lambda_tildes: bool = False,
458	no_taper: bool = False,
459	) -> None:
460	"""
461	Args:
462	f_ref (float): Reference frequency in Hz. Defaults to 20.0.
463	use_lambda_tildes (bool): Whether to parameterise tidal deformability
464	via ``lambda_tilde`` / ``delta_lambda_tilde`` rather than
465	``lambda_1`` / ``lambda_2``. Defaults to False.
466	no_taper (bool): Whether to disable tapering (useful for relative
467	binning runs). Defaults to False.
468	"""
469	self.f_ref = f_ref	1✔
470	self.use_lambda_tildes = use_lambda_tildes	1✔
471	self.no_taper = no_taper	1✔
472
473	@property	1✔
474	def parameter_names(self) -> tuple[str, ...]:	1✔
NEW 475	return (	×
476	"M_c",
477	"eta",
478	"s1_z",
479	"s2_z",
480	*(
481	("lambda_tilde", "delta_lambda_tilde")
482	if self.use_lambda_tildes
483	else ("lambda_1", "lambda_2")
484	),
485	"d_L",
486	"phase_c",
487	"iota",
488	)
489
490	def __call__(	1✔
491	self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
492	) -> dict[str, Float[Array, " n_freq"]]:
493	"""Evaluate the IMRPhenomXAS_NRTidalv3 waveform.
494
495	Args:
496	frequency (Float[Array, " n_freq"]): Frequency array in Hz.
497	params (dict[str, Float]): Source parameters with keys ``M_c``,
498	``eta``, ``s1_z``, ``s2_z``, ``d_L``, ``phase_c``, ``iota``,
499	plus tidal keys depending on ``use_lambda_tildes``.
500
501	Returns:
502	dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
503	polarizations.
504	"""
505	output = {}	1✔
506
507	if self.use_lambda_tildes:	1✔
508	first_lambda_param = params["lambda_tilde"]	1✔
509	second_lambda_param = params["delta_lambda_tilde"]	1✔
510	else:
511	first_lambda_param = params["lambda_1"]	1✔
512	second_lambda_param = params["lambda_2"]	1✔
513
514	theta = jnp.array(	1✔
515	[
516	params["M_c"],
517	params["eta"],
518	params["s1_z"],
519	params["s2_z"],
520	first_lambda_param,
521	second_lambda_param,
522	params["d_L"],
523	0,
524	params["phase_c"],
525	params["iota"],
526	]
527	)
528	hp, hc = gen_IMRPhenomXAS_NRTidalv3_hphc(	1✔
529	frequency,
530	theta,
531	self.f_ref,
532	use_lambda_tildes=self.use_lambda_tildes,
533	no_taper=self.no_taper,
534	)
535	output["p"] = hp	1✔
536	output["c"] = hc	1✔
537	return output	1✔
538
539	def __repr__(self):
540	return f"IMRPhenomXAS_NRTidalv3(f_ref={self.f_ref})"
541
542
543	class IMRPhenomXPHM(Waveform):	1✔
544	"""IMRPhenomXPHM frequency-domain waveform (precessing spins, higher-order modes).
545
546	Attributes:
547	f_ref (float): Reference frequency in Hz.
548	"""
549
550	f_ref: float	1✔
551
552	def __init__(self, f_ref: float = 20.0) -> None:	1✔
553	"""
554	Args:
555	f_ref (float): Reference frequency in Hz. Defaults to 20.0.
556	"""
557	self.f_ref = f_ref	1✔
558
559	@property	1✔
560	def parameter_names(self) -> tuple[str, ...]:	1✔
NEW 561	return (	×
562	"M_c",
563	"eta",
564	"s1_x",
565	"s1_y",
566	"s1_z",
567	"s2_x",
568	"s2_y",
569	"s2_z",
570	"d_L",
571	"phase_c",
572	"iota",
573	)
574
575	def __call__(	1✔
576	self, frequency: Float[Array, " n_freq"], params: dict[str, Float]
577	) -> dict[str, Float[Array, " n_freq"]]:
578	"""Evaluate the IMRPhenomXPHM waveform.
579
580	Args:
581	frequency (Float[Array, " n_freq"]): Frequency array in Hz.
582	params (dict[str, Float]): Source parameters with keys
583	``M_c``, ``eta``, ``s1_x``, ``s1_y``, ``s1_z``,
584	``s2_x``, ``s2_y``, ``s2_z``, ``d_L``, ``phase_c``, ``iota``.
585
586	Returns:
587	dict[str, Float[Array, " n_freq"]]: Plus (``"p"``) and cross (``"c"``)
588	polarizations.
589	"""
590	output = {}	1✔
591	m1, m2 = Mc_eta_to_ms(jnp.array([params["M_c"], params["eta"]]))	1✔
592	hp, hc = generate_xphm(	1✔
593	m1,
594	m2,
595	params["s1_x"],
596	params["s1_y"],
597	params["s1_z"],
598	params["s2_x"],
599	params["s2_y"],
600	params["s2_z"],
601	params["d_L"],
602	params["iota"],
603	params["phase_c"],
604	frequency,
605	self.f_ref,
606	)
607	output["p"] = hp	1✔
608	output["c"] = hc	1✔
609	return output	1✔
610
611	def __repr__(self):
612	return f"IMRPhenomXPHM(f_ref={self.f_ref})"
613
614
615	class SineGaussian(Waveform):	1✔
616	"""Sine-Gaussian time-domain burst waveform."""
617
618	def __init__(self) -> None:	1✔
619	pass	1✔
620
621	@property	1✔
622	def parameter_names(self) -> tuple[str, ...]:	1✔
NEW 623	return ("Q", "f_0", "hrss", "phase", "e")	×
624
625	def __call__(	1✔
626	self, t: Float[Array, " n_time"], params: dict[str, Float]
627	) -> dict[str, Float[Array, " n_time"]]:
628	"""
629	Args:
630	t: Time grid centered at t=0. Create using
631	``jnp.arange(-duration/2, duration/2, 1/fs)``.
632	params: Dictionary with keys ``Q`` (quality factor), ``f_0``
633	(central frequency in Hz), ``hrss``, ``phase`` (phase),
634	``e`` (eccentricity).
635	"""
636	output = {}	1✔
637	theta = jnp.array(	1✔
638	[
639	params["Q"],
640	params["f_0"],
641	params["hrss"],
642	params["phase"],
643	params["e"],
644	]
645	)
646	hp, hc = gen_SineGaussian_hphc(t, theta)	1✔
647	output["p"] = hp	1✔
648	output["c"] = hc	1✔
649	return output	1✔
650
651	def __repr__(self):
652	return "SineGaussian()"
653
654
655	#: Mapping from model name strings to :class:`Waveform` subclasses.
656	#: Useful for selecting waveform models by name at runtime, e.g. from a
657	#: configuration file.
658	waveform_preset: dict[str, type[Waveform]] = {	1✔
659	"IMRPhenomD": IMRPhenomD,
660	"IMRPhenomPv2": IMRPhenomPv2,
661	"TaylorF2": TaylorF2,
662	"IMRPhenomD_NRTidalv2": IMRPhenomD_NRTidalv2,
663	"IMRPhenomXAS": IMRPhenomXAS,
664	"IMRPhenomXAS_NRTidalv3": IMRPhenomXAS_NRTidalv3,
665	"IMRPhenomXPHM": IMRPhenomXPHM,
666	"SineGaussian": SineGaussian,
667	}

GW-JAX-Team / ripple / 24966502451

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous