gwkokab.models.transformations

Provides implementation of various transformations using Transform.

Classes

BlockTransform	A transform that applies multiple sub-transforms to disjoint slices of the event
ComponentMassesAndRedshiftToDetectedMassAndRedshift	Transforms component masses and redshift to detected masses and redshift.
ComponentMassesToChirpMassAndDelta	Transforms component masses to chirp mass and delta.
ComponentMassesToChirpMassAndSymmetricMassRatio	Transforms component masses to chirp mass and symmetric mass ratio.
ComponentMassesToMassRatioAndSecondaryMass	Transforms component masses and redshift to mass ratio and secondary mass.
ComponentMassesToPrimaryMassAndMassRatio	Transforms component masses and redshift to primary mass and mass ratio.
ComponentMassesToTotalMassAndMassRatio	Transforms component masses to total mass and mass ratio.
DeltaToSymmetricMassRatio	Transforms delta to symmetric mass ratio.
PrimaryMassAndMassRatioToComponentMassesTransform	Transforms a primary mass and mass ratio to component masses.
RedshiftToLuminosityDistance	Transforms redshift to luminosity distance.
SourceMassAndRedshiftToDetectedMassAndRedshift	Transforms source mass and redshift to detected mass and redshift.

Module Contents

class gwkokab.models.transformations.BlockTransform(*transforms: numpyro.distributions.transforms.Transform, event_slices: Sequence[int | Tuple[int, int]])

Bases: numpyro.distributions.transforms.Transform

A transform that applies multiple sub-transforms to disjoint slices of the event dimension.

This class implements a block-separable transformation of the form

\[T(x) = \big( T_1(x_{S_1}), \; T_2(x_{S_2}), \; \dots, \; T_K(x_{S_K}) \big),\]

where each \(T_i\) is a Transform and \(S_i\) is a slice of the event dimension specified by event_slices. The slices must be pairwise disjoint so that no parameters or coordinates are shared between sub-transforms.

Because each sub-transform acts independently on its own coordinate block, the Jacobian matrix has block-diagonal structure:

\[\begin{split}J_T(x) = \begin{pmatrix} J_{T_1}(x_{S_1}) & 0 & \cdots & 0 \\ 0 & J_{T_2}(x_{S_2}) & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & J_{T_K}(x_{S_K}) \end{pmatrix}.\end{split}\]

Consequently, the log absolute determinant of the Jacobian factorizes:

\[\log \left| \det J_T(x) \right| = \sum_{i=1}^K \log \left| \det J_{T_i}(x_{S_i}) \right|.\]

Parameters:

• *transforms (Transform) – A sequence of sub-transforms \(T_1, \dots, T_K\). Each transform is applied independently to its corresponding slice of the input.

• event_slices (Sequence[Union[int, Tuple[int, int]]]) – A sequence specifying the slices \(S_i\) of the event dimension. Each entry is either: - an integer j (interpreted as selecting x[..., j]), or - a tuple (start, end) denoting the half-open interval \([\mathrm{start}, \mathrm{end})\).

Notes

• The overall transformation is equivalent to a product of independent transforms acting on different subspaces.

• No checks are performed to ensure that slices fully cover the event dimension or that the resulting concatenation is contiguous.

Warning

• event_slices must be non-overlapping. Overlapping slices violate the independence assumption and produce incorrect Jacobians.

• Each sub-transform must be dimensionally compatible with the slice it receives.

• If a slice misses part of the event dimension or overlaps with another, the forward and inverse mappings may not be valid.

Examples

>>> t1 = AffineTransform(loc=0.0, scale=1.0)
>>> t2 = ExpTransform()
>>> bt = BlockTransform(t1, t2, event_slices=[(0, 3), (3, 4)])
>>> x = jnp.array([1.0, 2.0, 3.0, 0.5])
>>> y = bt(x)
>>> x_recovered = bt.inv(y)

eq(value, static: bool = False)

log_abs_det_jacobian(x: jaxtyping.Array, y: jaxtyping.Array, intermediates=None)

codomain

domain

event_slices

transforms = ()

class gwkokab.models.transformations.ComponentMassesAndRedshiftToDetectedMassAndRedshift

Bases: numpyro.distributions.transforms.Transform

Transforms component masses and redshift to detected masses and redshift.

\[f: (m_1, m_2, z) \to (m_{1, \text{detected}}, m_{2, \text{detected}}, z)\]

eq(other, static: bool = False)

log_abs_det_jacobian(x, y, intermediates=None)

\[\ln\left(|\mathrm{det}(J_f)|\right) = 2\ln(1+z)\]

codomain

\(\mathcal{C}(f) = \mathbb{R}^3_+\)

domain

\(\mathcal{D}(f) = \mathbb{R}^3_+\)

class gwkokab.models.transformations.ComponentMassesToChirpMassAndDelta

Bases: numpyro.distributions.transforms.Transform

Transforms component masses to chirp mass and delta.

\[f: (m_1, m_2) \to (M_c, \delta)\]

eq(other, static: bool = False)

log_abs_det_jacobian(x, y, intermediates=None)

\[\ln\left(|\mathrm{det}(J_f)|\right) = \ln(2M_c) - 2\ln(m_1+m_2)\]

codomain

\(\mathcal{C}(f) = \mathbb{R}^2_+\times[0, 1]\)

domain

\(\mathcal{D}(f)=\{(m_1,m_2)\in\mathbb{R}^2_+\mid m_1\geq m_2>0\}\)

class gwkokab.models.transformations.ComponentMassesToChirpMassAndSymmetricMassRatio

Bases: numpyro.distributions.transforms.Transform

Transforms component masses to chirp mass and symmetric mass ratio.

\[f: (m_1, m_2)\to \left(\frac{(m_1m_2)^{3/5}}{(m_1+m_2)^{1/5}}, \frac{m_1m_2}{(m_1+m_2)^{2}}\right)\]

eq(other, static: bool = False)

log_abs_det_jacobian(x, y, intermediates=None)

\[\ln\left(|\mathrm{det}(J_f)|\right)=\frac{6}{5}\ln(\eta)+\frac{1}{2}\ln(1-4\eta)-\ln(M_c)\]

codomain

\(\mathcal{C}(f) = \mathbb{R}^2_+\times[0, 0.25]\)

domain

\(\mathcal{D}(f)=\{(m_1,m_2)\in\mathbb{R}^2_+\mid m_1\geq m_2>0\}\)

class gwkokab.models.transformations.ComponentMassesToMassRatioAndSecondaryMass

Bases: numpyro.distributions.transforms.Transform

Transforms component masses and redshift to mass ratio and secondary mass.

\[f: (m_1, m_2) \to (q, m_2)\]

\[f^{-1}: (q, m_2) \to (m_1, m_2)\]

eq(other, static: bool = False)

log_abs_det_jacobian(x, y, intermediates=None)

\[\ln\left(|\mathrm{det}(J_f)|\right) = \ln(q) - \ln(m_1)\]

codomain

\(\mathcal{C}(f) = [0, 1]\times\mathbb{R}_+\)

domain

\(\mathcal{D}(f)=\{(m_1,m_2)\in\mathbb{R}^2_+\mid m_1\geq m_2>0\}\)

class gwkokab.models.transformations.ComponentMassesToPrimaryMassAndMassRatio

Bases: numpyro.distributions.transforms.Transform

Transforms component masses and redshift to primary mass and mass ratio.

\[f: (m_1, m_2) \to (m_1, q)\]

\[f^{-1}: (m_1, q) \to (m_1, m_2)\]

eq(other, static: bool = False)

log_abs_det_jacobian(x, y, intermediates=None)

codomain

\(\mathcal{C}(f) = \mathbb{R}^2_+\times(0, 1)\)

domain

\(\mathcal{D}(f)=\{(m_1,m_2)\in\mathbb{R}^2_+\mid m_1\geq m_2>0\}\)

class gwkokab.models.transformations.ComponentMassesToTotalMassAndMassRatio

Bases: numpyro.distributions.transforms.Transform

Transforms component masses to total mass and mass ratio.

\[f: (m_1, m_2) \to (M, q)\]

eq(other, static: bool = False)

log_abs_det_jacobian(x, y, intermediates=None)

codomain

domain

class gwkokab.models.transformations.DeltaToSymmetricMassRatio

Bases: numpyro.distributions.transforms.Transform

Transforms delta to symmetric mass ratio.

\[\eta = f(\delta) = \frac{1-\delta^2}{4}\]

\[\delta = f^{-1}(\eta) = \sqrt{1-4\eta}\]

eq(other, static: bool = False)

log_abs_det_jacobian(x, y, intermediates=None)

\[\ln\left(|\mathrm{det}(J_f)|\right) = \ln(\delta) - \ln(2)\]

codomain

\(\mathcal{C}(f) = [0, 0.25]\)

domain

\(\mathcal{D}(f) = [0, 1]\)

class gwkokab.models.transformations.PrimaryMassAndMassRatioToComponentMassesTransform

Bases: numpyro.distributions.transforms.Transform

Transforms a primary mass and mass ratio to component masses.

\[f: (m_1, q)\to (m_1, m_1q)\]

\[f^{-1}: (m_1, m_2)\to (m_1, m_2/m_1)\]

eq(other, static: bool = False)

log_abs_det_jacobian(x: jaxtyping.Array, y: jaxtyping.Array, intermediates=None)

\[\ln\left(|\mathrm{det}(J_f)|\right) = \ln(|m_1|)\]

codomain

\(\mathcal{C}(f)=\{(m_1, m_2)\in\mathbb{R}^2_+\mid m_1\geq m_2>0\}\)

domain

\(\mathcal{D}(f) = \mathbb{R}^2_+\times[0, 1]\)

class gwkokab.models.transformations.RedshiftToLuminosityDistance(cosmology: gwkokab.cosmology.Cosmology)

Bases: numpyro.distributions.transforms.Transform

Transforms redshift to luminosity distance.

\[f: z \to D_L\]

\[f^{-1}: D_L \to z\]

\[\ln\left(|\mathrm{det}(J_f)|\right) = \ln\left(D_c + \frac{c(1+z)}{H_0 E(z)}\right)\]

log_abs_det_jacobian(x, y, intermediates=None)

codomain

\(\mathcal{C}(f) = \mathbb{R}_+\)

cosmology

domain

\(\mathcal{D}(f) = \mathbb{R}_+\)

class gwkokab.models.transformations.SourceMassAndRedshiftToDetectedMassAndRedshift

Bases: numpyro.distributions.transforms.Transform

Transforms source mass and redshift to detected mass and redshift.

\[f: (m_{\text{source}}, z) \to (m_{\text{detected}}, z)\]

eq(other, static: bool = False)

log_abs_det_jacobian(x, y, intermediates=None)

\[\ln\left(|\mathrm{det}(J_f)|\right) = \ln(1+z)\]

codomain

\(\mathcal{C}(f) = \mathbb{R}^2_+\)

domain

\(\mathcal{D}(f) = \mathbb{R}^2_+\)