gwkokab.models.transformations¶

Provides implementation of various transformations using Transform.

Most of these are changes of mass coordinates: a population model is naturally written in one parameterisation – \((m_1, q)\), say – while the data or the selection function live in another, such as \((m_1, m_2)\) or \((M_c, \eta)\). Each transform supplies the forward and inverse maps together with the log Jacobian determinant that keeps the density normalised across the change. RedshiftToLuminosityDistance does the same for the redshift/distance pair, and BlockTransform composes several transforms that act on disjoint slices of one event vector.

The module also registers biject_to rules for the custom constraints defined in gwkokab.models.constraints, so that samplers working in an unconstrained space know how to reach an ordered vector or a mass sandwich.

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]])[source]¶

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)[source]¶

Compare two block transforms structurally.

Parameters:

value (object) – The object to compare against.
static (bool, optional) – Forwarded to each sub-transform’s eq. Defaults to False.

Returns:

True if value is a BlockTransform with the same event slices and pairwise equal sub-transforms.

Return type:

bool

log_abs_det_jacobian(x: jaxtyping.Array, y: jaxtyping.Array, intermediates=None)[source]¶

Log absolute determinant of the block-diagonal Jacobian.

Because the blocks are independent, the determinant factorises and the log determinant is the sum of the sub-transforms’ contributions:

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

An elementwise sub-transform reports one term per coordinate while one with an event dimension has already reduced its block, so any axes beyond the batch shape are summed away before the terms are added.

Parameters:

x (Array) – Values in the domain.
y (Array) – The corresponding values in the codomain.
intermediates (Any, optional) – Passed through to each sub-transform. Defaults to None.

Returns:

The log absolute determinant, of the batch shape of x.

Return type:

Array

property codomain: numpyro.distributions.constraints.Constraint¶

Codomain of the block transform.

Returns:: An all_constraint() requiring each slice to lie in the codomain of its own sub-transform.
Return type:: constraints.Constraint

property domain: numpyro.distributions.constraints.Constraint¶

Domain of the block transform.

Returns:: An all_constraint() requiring each slice to lie in the domain of its own sub-transform.
Return type:: constraints.Constraint

event_slices¶

transforms = ()¶

class gwkokab.models.transformations.ComponentMassesAndRedshiftToDetectedMassAndRedshift[source]¶

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)[source]¶

Compare two transforms.

Parameters:

other (object) – The object to compare against.
static (bool, optional) – Accepted for API compatibility with numpyro.distributions.transforms.Transform.eq() and ignored. Defaults to False.

Returns:

True if other is also a ComponentMassesAndRedshiftToDetectedMassAndRedshift. The transform is stateless, so type identity is enough.

Return type:

bool

log_abs_det_jacobian(x, y, intermediates=None)[source]¶

Log absolute determinant of the Jacobian of the forward map.

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

Parameters:

x (Array) – Values in the domain.
y (Array) – The corresponding values in the codomain.
intermediates (Any, optional) – Accepted for API compatibility and ignored. Defaults to None.

Returns:

The log absolute determinant of the Jacobian.

Return type:

Array

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

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

class gwkokab.models.transformations.ComponentMassesToChirpMassAndDelta[source]¶

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)[source]¶

Compare two transforms.

Parameters:

other (object) – The object to compare against.
static (bool, optional) – Accepted for API compatibility with numpyro.distributions.transforms.Transform.eq() and ignored. Defaults to False.

Returns:

True if other is also a ComponentMassesToChirpMassAndDelta. The transform is stateless, so type identity is enough.

Return type:

bool

log_abs_det_jacobian(x, y, intermediates=None)[source]¶

Log absolute determinant of the Jacobian of the forward map.

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

Parameters:

x (Array) – Values in the domain.
y (Array) – The corresponding values in the codomain.
intermediates (Any, optional) – Accepted for API compatibility and ignored. Defaults to None.

Returns:

The log absolute determinant of the Jacobian.

Return type:

Array

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[source]¶

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)[source]¶

Compare two transforms.

Parameters:

other (object) – The object to compare against.
static (bool, optional) – Accepted for API compatibility with numpyro.distributions.transforms.Transform.eq() and ignored. Defaults to False.

Returns:

True if other is also a ComponentMassesToChirpMassAndSymmetricMassRatio. The transform is stateless, so type identity is enough.

Return type:

bool

log_abs_det_jacobian(x, y, intermediates=None)[source]¶

Log absolute determinant of the Jacobian of the forward map.

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

Parameters:

x (Array) – Values in the domain.
y (Array) – The corresponding values in the codomain.
intermediates (Any, optional) – Accepted for API compatibility and ignored. Defaults to None.

Returns:

The log absolute determinant of the Jacobian.

Return type:

Array

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[source]¶

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)[source]¶

Compare two transforms.

Parameters:

other (object) – The object to compare against.
static (bool, optional) – Accepted for API compatibility with numpyro.distributions.transforms.Transform.eq() and ignored. Defaults to False.

Returns:

True if other is also a ComponentMassesToMassRatioAndSecondaryMass. The transform is stateless, so type identity is enough.

Return type:

bool

log_abs_det_jacobian(x, y, intermediates=None)[source]¶

Log absolute determinant of the Jacobian of the forward map.

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

Parameters:

x (Array) – Values in the domain.
y (Array) – The corresponding values in the codomain.
intermediates (Any, optional) – Accepted for API compatibility and ignored. Defaults to None.

Returns:

The log absolute determinant of the Jacobian.

Return type:

Array

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[source]¶

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)[source]¶

Compare two transforms.

Parameters:

other (object) – The object to compare against.
static (bool, optional) – Accepted for API compatibility with numpyro.distributions.transforms.Transform.eq() and ignored. Defaults to False.

Returns:

True if other is also a ComponentMassesToPrimaryMassAndMassRatio. The transform is stateless, so type identity is enough.

Return type:

bool

log_abs_det_jacobian(x, y, intermediates=None)[source]¶

Log absolute determinant of the Jacobian of the forward map.

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

Parameters:

x (Array) – Values in the domain.
y (Array) – The corresponding values in the codomain.
intermediates (Any, optional) – Accepted for API compatibility and ignored. Defaults to None.

Returns:

The log absolute determinant of the Jacobian.

Return type:

Array

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[source]¶

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)[source]¶

Compare two transforms.

Parameters:

other (object) – The object to compare against.
static (bool, optional) – Accepted for API compatibility with numpyro.distributions.transforms.Transform.eq() and ignored. Defaults to False.

Returns:

True if other is also a ComponentMassesToTotalMassAndMassRatio. The transform is stateless, so type identity is enough.

Return type:

bool

log_abs_det_jacobian(x, y, intermediates=None)[source]¶

Log absolute determinant of the Jacobian of the forward map.

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

Parameters:

x (Array) – Values in the domain.
y (Array) – The corresponding values in the codomain.
intermediates (Any, optional) – Accepted for API compatibility and ignored. Defaults to None.

Returns:

The log absolute determinant of the Jacobian.

Return type:

Array

codomain¶

domain¶

class gwkokab.models.transformations.DeltaToSymmetricMassRatio[source]¶

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)[source]¶

Compare two transforms.

Parameters:

other (object) – The object to compare against.
static (bool, optional) – Accepted for API compatibility with numpyro.distributions.transforms.Transform.eq() and ignored. Defaults to False.

Returns:

True if other is also a DeltaToSymmetricMassRatio. The transform is stateless, so type identity is enough.

Return type:

bool

log_abs_det_jacobian(x, y, intermediates=None)[source]¶

Log absolute determinant of the Jacobian of the forward map.

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

Parameters:

x (Array) – Values in the domain.
y (Array) – The corresponding values in the codomain.
intermediates (Any, optional) – Accepted for API compatibility and ignored. Defaults to None.

Returns:

The log absolute determinant of the Jacobian.

Return type:

Array

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

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

class gwkokab.models.transformations.PrimaryMassAndMassRatioToComponentMassesTransform[source]¶

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)[source]¶

Compare two transforms.

Parameters:

other (object) – The object to compare against.
static (bool, optional) – Accepted for API compatibility with numpyro.distributions.transforms.Transform.eq() and ignored. Defaults to False.

Returns:

True if other is also a PrimaryMassAndMassRatioToComponentMassesTransform. The transform is stateless, so type identity is enough.

Return type:

bool

log_abs_det_jacobian(x: jaxtyping.Array, y: jaxtyping.Array, intermediates=None)[source]¶

Log absolute determinant of the Jacobian of the forward map.

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

Parameters:

x (Array) – Values in the domain.
y (Array) – The corresponding values in the codomain.
intermediates (Any, optional) – Accepted for API compatibility and ignored. Defaults to None.

Returns:

The log absolute determinant of the Jacobian.

Return type:

Array

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)[source]¶

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)\]

Both directions go through the cosmology’s precomputed grid, so the transform is differentiable but only as accurate as that grid.

Parameters:: cosmology (Cosmology) – The cosmology supplying the distance/redshift relation.

log_abs_det_jacobian(x, y, intermediates=None)[source]¶

Log absolute determinant of the Jacobian of the forward map.

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

Parameters:

x (Array) – Redshift.
y (Array) – The corresponding luminosity distance. Unused; kept for API compatibility.
intermediates (Any, optional) – Accepted for API compatibility and ignored. Defaults to None.

Returns:

The log absolute determinant of the Jacobian.

Return type:

Array

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

cosmology¶

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

class gwkokab.models.transformations.SourceMassAndRedshiftToDetectedMassAndRedshift[source]¶

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)[source]¶

Compare two transforms.

Parameters:

other (object) – The object to compare against.
static (bool, optional) – Accepted for API compatibility with numpyro.distributions.transforms.Transform.eq() and ignored. Defaults to False.

Returns:

True if other is also a SourceMassAndRedshiftToDetectedMassAndRedshift. The transform is stateless, so type identity is enough.

Return type:

bool

log_abs_det_jacobian(x, y, intermediates=None)[source]¶

Log absolute determinant of the Jacobian of the forward map.

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

Parameters:

x (Array) – Values in the domain.
y (Array) – The corresponding values in the codomain.
intermediates (Any, optional) – Accepted for API compatibility and ignored. Defaults to None.

Returns:

The log absolute determinant of the Jacobian.

Return type:

Array

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

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