models_utils (gluon's conflicted copy 2024-04-05)

Module providing physical constants and useful cosmological functions.

All the constants are expressed in GeV, unless differently stated. The values are taken from the Particle Data Group (PDG).

ATTRIBUTE	DESCRIPTION
`G`	Newton constant (GeV\(^{-2}\)) TYPE: `float64`
`M_pl`	Reduced Planck mass (GeV) TYPE: `float64`
`T_0`	Present day temperature of the Universe (GeV) TYPE: `float64`
`z_eq`	Redshift of matter-radiation equality TYPE: `int`
`T_eq`	Temperature of matter-radiation equality (GeV) TYPE: `float64`
`h`	Reduced Hubble parameter TYPE: `float`
`H_0`	Hubble constant (GeV) TYPE: `float64`
`H_0_Hz`	Hubble constant (Hz) TYPE: `float64`
`omega_v`	Dark energy density today (Planck 2018) TYPE: `float`
`omega_m`	Matter density today (Planck 2018) TYPE: `float`
`omega_r`	Radiation density today (Planck 2018) TYPE: `float`
`A_s`	Planck 2018 amplitude of primordial scalar power spectrum TYPE: `float64`
`f_cmb`	CMB pivot scale (Hz) TYPE: `float`
`gev_to_hz`	Conversion from GeV to Hz TYPE: `float64`
`g_rho_0`	Relativistic degrees of freedom today TYPE: `float64`
`g_s_0`	Entropic relativistic degrees of freedom today TYPE: `float64`
`priors_type`	Type for parameter priors. TYPE: `Literal['Uniform', 'Normal', 'TruncNormal', 'LinearExp', 'Constant', 'Gamma']`

Classes¶

ParamDict ¶

Bases: UserDict

UserDict child class that instantiates common or uncommon parameter priors.

Examples:

>>> import numpy as np
>>> from ptarcade.models_utils import ParamDict, prior
>>> from pprint import pprint
>>> parameters = ParamDict(
...     {
...     "log10_A_dm" : prior("Uniform", -9, -4),
...     "log10_f_dm" : prior("Uniform", pmin=-10, pmax=-5.5), # kwargs also work
...     "gamma_p" : prior("Uniform", 0, 2 * np.pi, common=False),
...     "gamma_e" : prior("Uniform", 0, 2 * np.pi),
...     "phi_hat_sq_e" : prior("Gamma", a=1, loc=0, scale=1, common=False),
...     }
... )
>>> pprint(parameters, sort_dicts=False)
{'log10_A_dm': log10_A_dm:Uniform(pmin=-9, pmax=-4),
 'log10_f_dm': log10_f_dm:Uniform(pmin=-10, pmax=-5.5),
 'gamma_p': <class 'enterprise.signals.parameter.Uniform.<locals>.Uniform'>,
 'gamma_e': gamma_e:Uniform(pmin=0, pmax=6.283185307179586),
 'phi_hat_sq_e': <class 'ptarcade.models_utils.Gamma.<locals>.Gamma'>}

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

class ParamDict(UserDict):
    """UserDict child class that instantiates common or uncommon parameter priors.

    Examples
    --------
    >>> import numpy as np
    >>> from ptarcade.models_utils import ParamDict, prior
    >>> from pprint import pprint
    >>> parameters = ParamDict(
    ...     {
    ...     "log10_A_dm" : prior("Uniform", -9, -4),
    ...     "log10_f_dm" : prior("Uniform", pmin=-10, pmax=-5.5), # kwargs also work
    ...     "gamma_p" : prior("Uniform", 0, 2 * np.pi, common=False),
    ...     "gamma_e" : prior("Uniform", 0, 2 * np.pi),
    ...     "phi_hat_sq_e" : prior("Gamma", a=1, loc=0, scale=1, common=False),
    ...     }
    ... )
    >>> pprint(parameters, sort_dicts=False)
    {'log10_A_dm': log10_A_dm:Uniform(pmin=-9, pmax=-4),
     'log10_f_dm': log10_f_dm:Uniform(pmin=-10, pmax=-5.5),
     'gamma_p': <class 'enterprise.signals.parameter.Uniform.<locals>.Uniform'>,
     'gamma_e': gamma_e:Uniform(pmin=0, pmax=6.283185307179586),
     'phi_hat_sq_e': <class 'ptarcade.models_utils.Gamma.<locals>.Gamma'>}

    """

    def __setitem__(self, key: str, prior: parameter.Parameter):
        # The "or" here makes it backwards compatible with our old syntax
        # for the parameter dictionaries
        if isinstance(prior, parameter.Parameter ) or getattr(prior, "common", False):
            super().__setitem__(key, prior(key))
        else:
            super().__setitem__(key, prior)

Functions¶

g_rho ¶

g_rho(x: array_like, is_freq: bool = False) -> array_like

Return the number of relativistic degrees of freedom as a function of T/GeV or f/Hz.

PARAMETER	DESCRIPTION
`x`	The temperature(s) [GeV] or frequency/frequencies [Hz]. TYPE: `_ArrayLikeFloat_co`
`is_freq`	True if `x` is a frequency/frequencies, False if temperature(s). Defaults to False. TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`dof`	The relativistic degrees of freedom at `x`. TYPE: `_ArrayLikeFloat_co`

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def g_rho(x: array_like, is_freq: bool = False) -> array_like:  # noqa: FBT001, FBT002
    """Return the number of relativistic degrees of freedom as a function of T/GeV or f/Hz.

    Parameters
    ----------
    x : array_like
        The temperature(s) [GeV] or frequency/frequencies [Hz].
    is_freq : bool, optional
        True if `x` is a frequency/frequencies, False if temperature(s).
        Defaults to False.

    Returns
    -------
    dof : array_like
        The relativistic degrees of freedom at `x`.

    """
    if is_freq:
        dof = np.interp(x, gs[:, 1], gs[:, 3])

    else:
        dof = np.interp(x, gs[:, 0], gs[:, 3])

    return dof

g_rho_jax ¶

g_rho_jax(
    x: jax_array_like, is_freq: bool = False
) -> jax_array_like

Return the number of relativistic degrees of freedom as a function of T/GeV or f/Hz.

PARAMETER	DESCRIPTION
`x`	The temperature(s) [GeV] or frequency/frequencies [Hz]. TYPE: `_ArrayLikeFloat_co`
`is_freq`	True if `x` is a frequency/frequencies, False if temperature(s). Defaults to False. TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`dof`	The relativistic degrees of freedom at `x`. TYPE: `_ArrayLikeFloat_co`

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def g_rho_jax(x: jax_array_like, is_freq: bool = False) -> jax_array_like:  # noqa: FBT001, FBT002
    """Return the number of relativistic degrees of freedom as a function of T/GeV or f/Hz.

    Parameters
    ----------
    x : array_like
        The temperature(s) [GeV] or frequency/frequencies [Hz].
    is_freq : bool, optional
        True if `x` is a frequency/frequencies, False if temperature(s).
        Defaults to False.

    Returns
    -------
    dof : array_like
        The relativistic degrees of freedom at `x`.

    """
    if is_freq:
        dof = jnp.interp(x, gs[:, 1], gs[:, 3])

    else:
        dof = jnp.interp(x, gs[:, 0], gs[:, 3])

    return dof

g_s ¶

g_s(x: array_like, is_freq: bool = False) -> array_like

Return the number of entropic relativistic degrees of freedom as a function of T/GeV or f/Hz.

PARAMETER	DESCRIPTION
`x`	The temperature(s) [GeV] or frequency/frequencies [Hz]. TYPE: `_ArrayLikeFloat_co`
`is_freq`	True if `x` is a frequency/frequencies, False if temperature(s). Defaults to False. TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`dof`	The entropic relativistic degrees of freedom at `x`. TYPE: `_ArrayLikeFloat_co`

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def g_s(x: array_like, is_freq: bool = False) -> array_like:  # noqa: FBT001, FBT002
    """Return the number of entropic relativistic degrees of freedom as a function of T/GeV or f/Hz.

    Parameters
    ----------
    x : array_like
        The temperature(s) [GeV] or frequency/frequencies [Hz].
    is_freq : bool, optional
        True if `x` is a frequency/frequencies, False if temperature(s).
        Defaults to False.

    Returns
    -------
    dof : array_like
        The entropic relativistic degrees of freedom at `x`.

    """
    if is_freq:
        dof = np.interp(x, gs[:, 1], gs[:, 2])

    else:
        dof = np.interp(x, gs[:, 0], gs[:, 2])

    return dof

g_s_jax ¶

g_s_jax(
    x: jax_array_like, is_freq: bool = False
) -> jax_array_like

Return the number of entropic relativistic degrees of freedom as a function of T/GeV or f/Hz.

PARAMETER	DESCRIPTION
`x`	The temperature(s) [GeV] or frequency/frequencies [Hz]. TYPE: `_ArrayLikeFloat_co`
`is_freq`	True if `x` is a frequency/frequencies, False if temperature(s). Defaults to False. TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`dof`	The entropic relativistic degrees of freedom at `x`. TYPE: `_ArrayLikeFloat_co`

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def g_s_jax(x: jax_array_like, is_freq: bool = False) -> jax_array_like:  # noqa: FBT001, FBT002
    """Return the number of entropic relativistic degrees of freedom as a function of T/GeV or f/Hz.

    Parameters
    ----------
    x : array_like
        The temperature(s) [GeV] or frequency/frequencies [Hz].
    is_freq : bool, optional
        True if `x` is a frequency/frequencies, False if temperature(s).
        Defaults to False.

    Returns
    -------
    dof : array_like
        The entropic relativistic degrees of freedom at `x`.

    """
    if is_freq:
        dof = jnp.interp(x, gs[:, 1], gs[:, 2])

    else:
        dof = jnp.interp(x, gs[:, 0], gs[:, 2])

    return dof

__g_s_0 `cached` ¶

__g_s_0(T_0: float) -> np.float64

Calculate the entropic relativistic degrees of freedom today.

This function is cached because it only needs to be ran once. It is a function instead of a constant so that if the g_star.dat file changes, this value will update as well.

PARAMETER	DESCRIPTION
`T_0`	The universe's temperature today [GeV]. TYPE: `float`

RETURNS	DESCRIPTION
`float`	The entropic relativistic degrees of freedom today.

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

@cache
def __g_s_0(T_0: float) -> np.float64:
    """Calculate the entropic relativistic degrees of freedom today.

    This function is cached because it only needs to be ran once. It is a function instead of a constant so that
    if the `g_star.dat` file changes, this value will update as well.

    Parameters
    ----------
    T_0 : float
        The universe's temperature today [GeV].

    Returns
    -------
    float
        The entropic relativistic degrees of freedom today.

    """
    return g_s(T_0) # type: ignore

__g_rho_0 `cached` ¶

__g_rho_0(T_0: float) -> np.float64

Calculate the relativistic degrees of freedom today.

This function is cached because it only needs to be ran once. It is a function instead of a constant so that if the g_star.dat file changes, this value will update as well.

PARAMETER	DESCRIPTION
`T_0`	The universe's temperature today [GeV]. TYPE: `float`

RETURNS	DESCRIPTION
`float`	The relativistic degrees of freedom today.

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

@cache
def __g_rho_0(T_0: float) -> np.float64:
    """Calculate the relativistic degrees of freedom today.

    This function is cached because it only needs to be ran once. It is a function instead of a constant so that
    if the `g_star.dat` file changes, this value will update as well.

    Parameters
    ----------
    T_0 : float
        The universe's temperature today [GeV].

    Returns
    -------
    float
        The relativistic degrees of freedom today.

    """
    return g_rho(T_0) # type: ignore

GammaPrior ¶

GammaPrior(
    value: float, a: float, loc: float, scale: float
) -> float

Prior function for Gamma parameters.

PARAMETER	DESCRIPTION
`value`	The value of the parameter. TYPE: `float`
`a`	The shape parameter of the Gamma distribution. TYPE: `float`
`loc`	The location parameter of the Gamma distribution. TYPE: `float`
`scale`	The scale parameter of the Gamma distribution. TYPE: `float`

RETURNS	DESCRIPTION
`float`	The probability density of the Gamma distribution at `value`.

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def GammaPrior(value: float, a: float, loc: float, scale: float) -> float:
    """Prior function for Gamma parameters.

    Parameters
    ----------
    value : float
        The value of the parameter.
    a : float
        The shape parameter of the Gamma distribution.
    loc : float
        The location parameter of the Gamma distribution.
    scale : float
        The scale parameter of the Gamma distribution.

    Returns
    -------
    float
        The probability density of the Gamma distribution at `value`.

    """
    return ss.gamma.pdf(value, a, loc, scale)

GammaSampler ¶

GammaSampler(
    a: float,
    loc: float,
    scale: float,
    size: int | None = None,
) -> NDArray

Sampling function for Gamma parameters.

PARAMETER	DESCRIPTION
`a`	The shape parameter of the Gamma distribution. TYPE: `float`
`loc`	The location parameter of the Gamma distribution. TYPE: `float`
`scale`	The scale parameter of the Gamma distribution. TYPE: `float`
`size`	The number of samples to draw. TYPE: `int` DEFAULT: `None`

RETURNS	DESCRIPTION
`NDArray`	A NumPy array of size `size` containing samples from the Gamma distribution.

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def GammaSampler(a: float, loc: float, scale: float, size: int | None  = None) -> NDArray:
    """Sampling function for Gamma parameters.

    Parameters
    ----------
    a : float
        The shape parameter of the Gamma distribution.
    loc : float
        The location parameter of the Gamma distribution.
    scale : float
        The scale parameter of the Gamma distribution.
    size : int, optional
        The number of samples to draw.

    Returns
    -------
    NDArray
        A NumPy array of size `size` containing samples from the Gamma distribution.

    """
    return ss.gamma.rvs(a, loc, scale, size=size)

Gamma ¶

Gamma(
    a: float,
    loc: float,
    scale: float,
    size: int | None = None,
)

Class factory for Gamma parameters.

PARAMETER	DESCRIPTION
`a`	The shape parameter of the Gamma distribution. TYPE: `float`
`loc`	The location parameter of the Gamma distribution. TYPE: `float`
`scale`	The scale parameter of the Gamma distribution. TYPE: `float`
`size`	The number of samples to draw. TYPE: `int` DEFAULT: `None`

RETURNS	DESCRIPTION
`Gamma`	Child class of [enterprise.signals.parameter.Parameter][] for Gamma distribution TYPE: `Parameter`

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def Gamma(a: float, loc: float, scale: float, size: int | None = None):
    """Class factory for Gamma parameters.

    Parameters
    ----------
    a : float
        The shape parameter of the Gamma distribution.
    loc : float
        The location parameter of the Gamma distribution.
    scale : float
        The scale parameter of the Gamma distribution.
    size : int, optional
        The number of samples to draw.

    Returns
    -------
    Gamma : parameter.Parameter
        Child class of [enterprise.signals.parameter.Parameter][] for Gamma distribution

    """
    class Gamma(parameter.Parameter):
        """Child class of enterprise.signals.parameter.Parameter."""

        _size = size
        _prior = parameter.Function(GammaPrior, a=a, loc=loc, scale=scale)
        _sampler = staticmethod(GammaSampler)
        _typename = parameter._argrepr("Gamma", a=a, loc=loc, scale=scale)

    return Gamma

omega2cross ¶

omega2cross(
    omega_hh: Callable[..., NDArray], ceffyl: bool = False
) -> Callable[..., NDArray]

Convert GW energy density.

Converts the GW energy density as a fraction of the closure density into the cross-power spectral density as a function of the frequency in Hz. This is intended to be used as a decorator on a function that returns the GW energy density as a fraction of the closure density.

PARAMETER	DESCRIPTION
`omega_hh`	The function that returns the GW energy density as a fraction of the closure density. TYPE: `Callable[..., NDArray]`

ceffyl: bool If set to tru use a version compatible with ceffyl, if set to false a version compatible with ENTERPRISE

RETURNS	DESCRIPTION
`Callable[..., NDArray]`	A function that returns the cross-power spectral density as a function of the frequency in Hz.

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def omega2cross(omega_hh: Callable[..., NDArray], ceffyl : bool = False) -> Callable[..., NDArray]:
    """Convert GW energy density.

    Converts the GW energy density as a fraction of the closure density into the cross-power spectral density
    as a function of the frequency in Hz. This is intended to be used as a decorator on a function that returns
    the GW energy density as a fraction of the closure density.

    Parameters
    ----------
    omega_hh : Callable[..., NDArray]
        The function that returns the GW energy density as a fraction of the closure density.

    ceffyl: bool
        If set to tru use a version compatible with ceffyl, if set to false a version compatible with
        ENTERPRISE

    Returns
    -------
    Callable[..., NDArray]
        A function that returns the cross-power spectral density as a function of the frequency in Hz.

    """
    if ceffyl:
        @function
        def cross(f: NDArray, Tspan: float, **kwargs):

            # fraction of the critical density in GWs
            h2_omega = omega_hh(f, **kwargs)

            # characteristic strain spectrum h_c(f)
            hcf = H_0_Hz / h * np.sqrt(3 * h2_omega / 2) / (np.pi * f)

            # cross-power spectral density S(f) (s^3)
            sf = (hcf**2 / (12 * np.pi**2 * f**3)) / Tspan

            return sf

    else:
        @function
        def cross(f: NDArray, components: int = 2, **kwargs):

            df = np.diff(np.concatenate((np.array([0]), f[::components])))

            # fraction of the critical density in GWs
            h2_omega = omega_hh(f, **kwargs)

            # characteristic strain spectrum h_c(f)
            hcf = H_0_Hz / h * np.sqrt(3 * h2_omega / 2) / (np.pi * f)

            # cross-power spectral density S(f) (s^3)
            sf = (hcf**2 / (12 * np.pi**2 * f**3)) * np.repeat(df, components)

            return sf

    return cross

prep_data ¶

prep_data(
    path: str,
) -> tuple[list[NDArray], NDArray, NDArray]

Shape tabulated data into a form that can be handled by interpn.

PARAMETER	DESCRIPTION
`path`	Path to the tabulated data file. TYPE: `str`

RETURNS	DESCRIPTION
`grids`	The grids of the tabulated data. TYPE: `list[NDArray]`
`omega_grid`	The omega grid of the tabulated data. TYPE: `NDArray`
`par_names`	The names of the parameters in the tabulated data. TYPE: `NDArray`

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def prep_data(path: str) -> tuple[list[NDArray], NDArray, NDArray]:
    """Shape tabulated data into a form that can be handled by `interpn`.

    Parameters
    ----------
    path : str
        Path to the tabulated data file.

    Returns
    -------
    grids : list[NDArray]
        The grids of the tabulated data.
    omega_grid : NDArray
        The omega grid of the tabulated data.
    par_names : NDArray
        The names of the parameters in the tabulated data.
    """
    par_names = np.loadtxt(path, max_rows=1, dtype='str')
    data = np.loadtxt(path, skiprows=1)

    spec_col = np.where(par_names=='spectrum')
    omega_grid = data.T[spec_col][0]

    data = np.delete(data, spec_col, axis=1)
    par_names = np.delete(par_names, spec_col)

    grids = [np.unique(row) for row in data.T]

    for idx, par in enumerate(par_names):
        omega_grid = omega_grid[data[:, -idx -1].argsort(kind="mergesort")]
        data = data[data[:, -idx -1].argsort(kind="mergesort")]

    grid_size = [len(x) for x in grids]
    omega_grid = omega_grid.reshape(tuple(grid_size))

    return grids, omega_grid, par_names

spec_importer ¶

spec_importer(
    path: str,
) -> Callable[[NDArray, Any], NDArray]

Import data and create a fast interpolation function.

Interpolate the GWB power spectrum from tabulated data. Return a function that interpolates over frequency.

PARAMETER	DESCRIPTION
`path`	Path to the tabulated data file. TYPE: `str`

RETURNS	DESCRIPTION
`Callable[[NDArray, P], NDArray]`	A callable object that interpolates the GWB power spectrum at a given frequency `f` and with given parameters `kwargs`.

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def spec_importer(path: str) -> Callable[[NDArray, Any],  NDArray]:
    """Import data and create a fast interpolation function.

    Interpolate the GWB power spectrum from tabulated data. Return a function that interpolates
    over frequency.

    Parameters
    ----------
    path : str
        Path to the tabulated data file.

    Returns
    -------
    Callable[[NDArray, P], NDArray]
        A callable object that interpolates the GWB power spectrum at a given frequency `f` and with given
        parameters `kwargs`.
    """
    info, data = fast_interpolate.load_data(path)
    # info is a list of (name, start, step)

    def spectrum(f: NDArray, **kwargs: Any) -> NDArray:

        # Construct right information format for interpolation
        return fast_interpolate.interp([(start, step, f if name == 'f' else kwargs[name])
                                         for (name, start, step) in info],
                                        data)
    return spectrum # type: ignore

freq_at_temp ¶

freq_at_temp(T: array_like) -> array_like

Find frequency today as function of temperature when GW was horizon size.

Calculates the GW frequency [Hz] today as a function of the universe temperature [GeV] when the GW was of horizon size.

PARAMETER	DESCRIPTION
`T`	The universe temperature [GeV] at the time when the GW was of horizon size. TYPE: `_ArrayLikeFloat_co`

RETURNS	DESCRIPTION
`NDArray`	The GW frequency [Hz] today that was of horizon size when the universe was at temperature `T` [GeV].

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def freq_at_temp(T: array_like) -> array_like:
    """Find frequency today as function of temperature when GW was horizon size.

    Calculates the GW frequency [Hz] today as a function of the universe temperature [GeV]
    when the GW was of horizon size.

    Parameters
    ----------
    T : array_like
        The universe temperature [GeV] at the time when the GW was of horizon size.

    Returns
    -------
    NDArray
        The GW frequency [Hz] today that was of horizon size when the universe was at temperature `T` [GeV].
    """
    f_0 = H_0_Hz / (2 * np.pi)

    T_ratio = T_0 / T # type: ignore
    g_ratio = g_rho_0 / g_rho(T) # type: ignore
    gs_ratio = g_s_0 / g_s(T) # type: ignore

    prefactor = f_0 * (gs_ratio) ** (1 / 3) * T_ratio
    sqr_term = np.sqrt(
        omega_v
        + (gs_ratio**-1 * T_ratio**-3 * omega_m)
        + (g_ratio**-1 * T_ratio**-4 * omega_r)
    )

    return prefactor * sqr_term

freq_at_temp_jax ¶

freq_at_temp_jax(T: array_like) -> array_like

Find frequency today as function of temperature when GW was horizon size.

Calculates the GW frequency [Hz] today as a function of the universe temperature [GeV] when the GW was of horizon size.

PARAMETER	DESCRIPTION
`T`	The universe temperature [GeV] at the time when the GW was of horizon size. TYPE: `_ArrayLikeFloat_co`

RETURNS	DESCRIPTION
`NDArray`	The GW frequency [Hz] today that was of horizon size when the universe was at temperature `T` [GeV].

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def freq_at_temp_jax(T: array_like) -> array_like:
    """Find frequency today as function of temperature when GW was horizon size.

    Calculates the GW frequency [Hz] today as a function of the universe temperature [GeV]
    when the GW was of horizon size.

    Parameters
    ----------
    T : array_like
        The universe temperature [GeV] at the time when the GW was of horizon size.

    Returns
    -------
    NDArray
        The GW frequency [Hz] today that was of horizon size when the universe was at temperature `T` [GeV].
    """
    f_0 = H_0_Hz / (2 * np.pi)

    T_ratio = T_0 / T # type: ignore
    g_ratio = g_rho_0 / g_rho_jax(T) # type: ignore
    gs_ratio = g_s_0 / g_s_jax(T) # type: ignore

    prefactor = f_0 * (gs_ratio) ** (1 / 3) * T_ratio
    sqr_term = jnp.sqrt(
        omega_v
        + (gs_ratio**-1 * T_ratio**-3 * omega_m)
        + (g_ratio**-1 * T_ratio**-4 * omega_r)
    )

    return prefactor * sqr_term

temp_at_freq ¶

temp_at_freq(f: array_like) -> NDArray

Get the temperature [GeV] of the universe when a gravitational wave of a certain frequency [Hz] today was of horizon size.

PARAMETER	DESCRIPTION
`f`	The frequency in Hz today. TYPE: `_ArrayLikeFloat_co`

RETURNS	DESCRIPTION
`NDArray`	The temperature [GeV] when the GW at frequency `f` [Hz] was of horizon size.

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def temp_at_freq(f: array_like) -> NDArray:
    """Get the temperature [GeV] of the universe when a gravitational wave of a
    certain frequency [Hz] today was of horizon size.

    Parameters
    ----------
    f : array_like
        The frequency in Hz today.

    Returns
    -------
    NDArray
        The temperature [GeV] when the GW at frequency `f` [Hz] was of horizon size.

    """
    return np.interp(f, gs[:, 1], gs[:, 0], left=np.nan, right=np.nan)

temp_at_freq_jax ¶

temp_at_freq_jax(f: array_like) -> NDArray

Get the temperature [GeV] of the universe when a gravitational wave of a certain frequency [Hz] today was of horizon size.

PARAMETER	DESCRIPTION
`f`	The frequency in Hz today. TYPE: `_ArrayLikeFloat_co`

RETURNS	DESCRIPTION
`NDArray`	The temperature [GeV] when the GW at frequency `f` [Hz] was of horizon size.

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def temp_at_freq_jax(f: array_like) -> NDArray:
    """Get the temperature [GeV] of the universe when a gravitational wave of a
    certain frequency [Hz] today was of horizon size.

    Parameters
    ----------
    f : array_like
        The frequency in Hz today.

    Returns
    -------
    NDArray
        The temperature [GeV] when the GW at frequency `f` [Hz] was of horizon size.

    """
    return jnp.interp(f, gs[:, 1], gs[:, 0], left=jnp.nan, right=jnp.nan)

prior ¶

prior(
    name: priors_type, *args: Any, **kwargs: Any
) -> parameter.Parameter

Wrap enterprise prior creation.

This function wraps the class factories in [enterprise.signals.parameter][]. It functions exactly the same as the original, except that it accepts additional kwargs and sets an additional attribute common. This attribute refers to whether the parameter this prior corresponds to is common to all pulsars. In the original implementation in enterprise, the way that this works is ambiguous. With this function, it is explicit. If the user does not pass common=False as a kwarg, then common defaults to True. This attribute will be used by ptarcade.models_utils.ParamDict objects in the model files.

PARAMETER	DESCRIPTION
`name`	The prior to use. TYPE: `priors_type`
`*args`	Positional arguments passed to the prior factory. TYPE: `Any` DEFAULT: `()`
`**kwargs`	kwargs passed to the prior factory. TYPE: `Any` DEFAULT: `{}`

RETURNS	DESCRIPTION
`prior`	The configured prior. TYPE: `Parameter`

RAISES	DESCRIPTION
`SystemExit`	If the prior name passed does not exist within [enterprise.signals.parameter]

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def prior(name: priors_type, *args: Any, **kwargs: Any) -> parameter.Parameter:
    """Wrap enterprise prior creation.

    This function wraps the class factories in [enterprise.signals.parameter][].
    It functions exactly the same as the original, except that it accepts additional
    `kwargs` and sets an additional attribute `common`. This attribute refers to
    whether the parameter this prior corresponds to is common to all pulsars. In the
    original implementation in enterprise, the way that this works is ambiguous. With
    this function, it is explicit. If the user does not pass `common=False` as a `kwarg`,
    then `common` defaults to `True`. This attribute will be used by
    [ptarcade.models_utils.ParamDict][] objects in the model files.

    Parameters
    ----------
    name : priors_type
        The prior to use.
    *args
        Positional arguments passed to the prior factory.
    **kwargs
        kwargs passed to the prior factory.

    Returns
    -------
    prior : parameter.Parameter
        The configured prior.

    Raises
    ------
    SystemExit
        If the prior name passed does not exist within [enterprise.signals.parameter]

    """
    # If "common" is passed as kwarg, remove it from the kwarg dictionary and store it.
    # If it wasn't passed, set it to True
    common = kwargs.pop("common", True)

    # Check if the user passed a correct prior name.
    # If they didn't, print an informative message
    try:
        prior_factory = getattr(parameter, name)
    except AttributeError:
        try:
            prior_factory = globals()[name]
        except KeyError:
            err = (f"The 'name' must be a string from the following list {priors_type=}.\n"
            f"You supplied {name=}.")

            log.error(err)
            raise SystemExit from None

    # Use enterprise's class factory
    prior_obj = prior_factory(*args, **kwargs)

    # Store the `common` arg for later use
    prior_obj.common = common

    return prior_obj

signal_to_noise ¶

signal_to_noise(
    signal: Callable[[NDArray, ...], NDArray],
    signal_args: dict,
    sens: NDArray,
    Tspan: float,
    interp: bool = False,
) -> array_like

Calculate the SNR of a signal given a sensitivity curve.

PARAMETER	DESCRIPTION
`signal`	Callable signal function with frequency as the first argument TYPE: `Callable[[NDArray, ...], NDArray]`
`signal_args`	Dictionary of kwargs to pass to `signal` TYPE: `dict`
`sens`	Sensitivity curve stored as a 2D NumPy array with dimensions (Hz, h*2Omega_GW) TYPE: `NDArray`
`Tspan`	Timespan of the observation TYPE: `float`

RETURNS	DESCRIPTION
`float64`	SNR ratio

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

def signal_to_noise(
    signal: Callable[[NDArray, ...], NDArray],
    signal_args: dict,
    sens: NDArray,
    Tspan: float,
    interp:bool = False
) -> array_like:
    """Calculate the SNR of a signal given a sensitivity curve.

    Parameters
    ----------
    signal : Callable[[NDArray, ...], NDArray]
        Callable signal function with frequency as the first argument
    signal_args : dict
        Dictionary of kwargs to pass to `signal`
    sens : NDArray
        Sensitivity curve stored as a 2D NumPy array with dimensions (Hz, h**2*Omega_GW)
    Tspan : float
        Timespan of the observation

    Returns
    -------
    np.float64
        SNR ratio

    """
    if interp:
        integrand = signal(sens[:, 0], **signal_args)[0,0].squeeze()
        # print("orig shape", integrand.shape)
        integrand = np.moveaxis(integrand, 1, 0)[...,0].squeeze()
        # print(integrand)
    else:
        integrand = signal(sens[:, 0], **signal_args)   # S_h / S_eff

    # print("integrand", integrand.shape)
    # print("sens", sens[:,-1].shape)
    integrand = integrand /  sens[:, -1]
    integrand**=2

    return np.sqrt(
        2 * Tspan *  # Prefactor
        np.trapz(integrand,
        sens[:, 0],  # Frequency bounds
        axis=0))  # Integrate over frequency

signal_to_noise_jax ¶

signal_to_noise_jax(
    signal: Callable[[jax_array, ...], jax_array],
    signal_args: dict,
    sens: jax_array,
    Tspan: float,
) -> jax_array_like

Calculate the signal-to-noise ratio (SNR) of a signal given a sensitivity curve.

PARAMETER	DESCRIPTION
`signal`	Callable function that takes frequency as the first argument and returns the signal waveform as a NumPy array. TYPE: `Callable[[...], ndarray]`
`signal_args`	Additional keyword arguments to be passed to the `signal` function. TYPE: `dict`
`f`	Frequency array in Hz. TYPE: `ndarray`
`sensitivity`	Sensitivity curve as a 2D NumPy array with dimensions (Hz, h*2Omega_GW). TYPE: `ndarray`
`Tspan`	Observation timespan in seconds. TYPE: `float`

RETURNS	DESCRIPTION
`ArrayLike`	The calculated signal-to-noise ratio.

Notes¶

This function uses JAX for numerical integration.

Source code in src/ptarcade/models_utils (gluon's conflicted copy 2024-04-05).py

@partial(jax.jit, static_argnums=(0, ))
def signal_to_noise_jax(
    signal: Callable[[jax_array, ...], jax_array],
    signal_args: dict,
    sens: jax_array,
    Tspan: float,
) -> jax_array_like:
    """
    Calculate the signal-to-noise ratio (SNR) of a signal given a sensitivity curve.

    Parameters
    ----------
    signal : Callable[[ ...], jnp.ndarray]
        Callable function that takes frequency as the first argument and returns
        the signal waveform as a NumPy array.
    signal_args : dict
        Additional keyword arguments to be passed to the `signal` function.
    f : jnp.ndarray
        Frequency array in Hz.
    sensitivity : jnp.ndarray
        Sensitivity curve as a 2D NumPy array with dimensions (Hz, h**2*Omega_GW).
    Tspan : float
        Observation timespan in seconds.

    Returns
    -------
    jax_array_like
        The calculated signal-to-noise ratio.

    Notes
    -----
    This function uses JAX for numerical integration.
    """
    #print(signal_args)
    integrand = (signal(sens[...,0], **signal_args) / sens[...,-1]) ** 2  # S_h / S_eff
    print(integrand)
    return jnp.sqrt(
        2 * Tspan  # Prefactor
        * jax.scipy.integrate.trapezoid(
            integrand,
            sens[...,0],  # Frequency bounds
            axis=0,  # Integrate over frequency
        )
    )

models_utils (gluon's conflicted copy 2024-04-05)

Classes¶

ParamDict ¶

Functions¶

g_rho ¶

g_rho_jax ¶

g_s ¶

g_s_jax ¶

__g_s_0 cached ¶

__g_rho_0 cached ¶

GammaPrior ¶

GammaSampler ¶

Gamma ¶

omega2cross ¶

prep_data ¶

spec_importer ¶

freq_at_temp ¶

freq_at_temp_jax ¶

temp_at_freq ¶

temp_at_freq_jax ¶

prior ¶

signal_to_noise ¶

signal_to_noise_jax ¶

Notes¶

__g_s_0 `cached` ¶

__g_rho_0 `cached` ¶