gwkokab.models.transformations ============================== .. py:module:: gwkokab.models.transformations .. autoapi-nested-parse:: Provides implementation of various transformations using :class:`~numpyro.distributions.transforms.Transform`. Classes ------- .. autoapisummary:: gwkokab.models.transformations.BlockTransform gwkokab.models.transformations.ComponentMassesAndRedshiftToDetectedMassAndRedshift gwkokab.models.transformations.ComponentMassesToChirpMassAndDelta gwkokab.models.transformations.ComponentMassesToChirpMassAndSymmetricMassRatio gwkokab.models.transformations.ComponentMassesToMassRatioAndSecondaryMass gwkokab.models.transformations.ComponentMassesToPrimaryMassAndMassRatio gwkokab.models.transformations.ComponentMassesToTotalMassAndMassRatio gwkokab.models.transformations.DeltaToSymmetricMassRatio gwkokab.models.transformations.PrimaryMassAndMassRatioToComponentMassesTransform gwkokab.models.transformations.RedshiftToLuminosityDistance gwkokab.models.transformations.SourceMassAndRedshiftToDetectedMassAndRedshift Module Contents --------------- .. py:class:: BlockTransform(*transforms: numpyro.distributions.transforms.Transform, event_slices: Sequence[Union[int, Tuple[int, int]]]) Bases: :py:obj:`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 .. math:: T(x) = \big( T_1(x_{S_1}), \; T_2(x_{S_2}), \; \dots, \; T_K(x_{S_K}) \big), where each :math:`T_i` is a ``Transform`` and :math:`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: .. math:: 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}. Consequently, the log absolute determinant of the Jacobian factorizes: .. math:: \log \left| \det J_T(x) \right| = \sum_{i=1}^K \log \left| \det J_{T_i}(x_{S_i}) \right|. :param \*transforms: A sequence of sub-transforms :math:`T_1, \dots, T_K`. Each transform is applied independently to its corresponding slice of the input. :type \*transforms: Transform :param event_slices: A sequence specifying the slices :math:`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 :math:`[\mathrm{start}, \mathrm{end})`. :type event_slices: Sequence[Union[int, Tuple[int, int]]] .. rubric:: 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. .. rubric:: 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) .. py:method:: eq(value, static: bool = False) .. py:method:: log_abs_det_jacobian(x: jaxtyping.Array, y: jaxtyping.Array, intermediates=None) .. py:attribute:: codomain .. py:attribute:: domain .. py:attribute:: event_slices .. py:attribute:: transforms :value: () .. py:class:: ComponentMassesAndRedshiftToDetectedMassAndRedshift Bases: :py:obj:`numpyro.distributions.transforms.Transform` Transforms component masses and redshift to detected masses and redshift. .. math:: f: (m_1, m_2, z) \to (m_{1, \text{detected}}, m_{2, \text{detected}}, z) .. py:method:: eq(other, static: bool = False) .. py:method:: log_abs_det_jacobian(x, y, intermediates=None) .. math:: \ln\left(|\mathrm{det}(J_f)|\right) = 2\ln(1+z) .. py:attribute:: codomain :math:`\mathcal{C}(f) = \mathbb{R}^3_+` .. py:attribute:: domain :math:`\mathcal{D}(f) = \mathbb{R}^3_+` .. py:class:: ComponentMassesToChirpMassAndDelta Bases: :py:obj:`numpyro.distributions.transforms.Transform` Transforms component masses to chirp mass and delta. .. math:: f: (m_1, m_2) \to (M_c, \delta) .. py:method:: eq(other, static: bool = False) .. py:method:: log_abs_det_jacobian(x, y, intermediates=None) .. math:: \ln\left(|\mathrm{det}(J_f)|\right) = \ln(2M_c) - 2\ln(m_1+m_2) .. py:attribute:: codomain :math:`\mathcal{C}(f) = \mathbb{R}^2_+\times[0, 1]` .. py:attribute:: domain :math:`\mathcal{D}(f)=\{(m_1,m_2)\in\mathbb{R}^2_+\mid m_1\geq m_2>0\}` .. py:class:: ComponentMassesToChirpMassAndSymmetricMassRatio Bases: :py:obj:`numpyro.distributions.transforms.Transform` Transforms component masses to chirp mass and symmetric mass ratio. .. math:: 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) .. py:method:: eq(other, static: bool = False) .. py:method:: log_abs_det_jacobian(x, y, intermediates=None) .. math:: \ln\left(|\mathrm{det}(J_f)|\right)=\frac{6}{5}\ln(\eta)+\frac{1}{2}\ln(1-4\eta)-\ln(M_c) .. py:attribute:: codomain :math:`\mathcal{C}(f) = \mathbb{R}^2_+\times[0, 0.25]` .. py:attribute:: domain :math:`\mathcal{D}(f)=\{(m_1,m_2)\in\mathbb{R}^2_+\mid m_1\geq m_2>0\}` .. py:class:: ComponentMassesToMassRatioAndSecondaryMass Bases: :py:obj:`numpyro.distributions.transforms.Transform` Transforms component masses and redshift to mass ratio and secondary mass. .. math:: f: (m_1, m_2) \to (q, m_2) .. math:: f^{-1}: (q, m_2) \to (m_1, m_2) .. py:method:: eq(other, static: bool = False) .. py:method:: log_abs_det_jacobian(x, y, intermediates=None) .. math:: \ln\left(|\mathrm{det}(J_f)|\right) = \ln(q) - \ln(m_1) .. py:attribute:: codomain :math:`\mathcal{C}(f) = [0, 1]\times\mathbb{R}_+` .. py:attribute:: domain :math:`\mathcal{D}(f)=\{(m_1,m_2)\in\mathbb{R}^2_+\mid m_1\geq m_2>0\}` .. py:class:: ComponentMassesToPrimaryMassAndMassRatio Bases: :py:obj:`numpyro.distributions.transforms.Transform` Transforms component masses and redshift to primary mass and mass ratio. .. math:: f: (m_1, m_2) \to (m_1, q) .. math:: f^{-1}: (m_1, q) \to (m_1, m_2) .. py:method:: eq(other, static: bool = False) .. py:method:: log_abs_det_jacobian(x, y, intermediates=None) .. py:attribute:: codomain :math:`\mathcal{C}(f) = \mathbb{R}^2_+\times(0, 1)` .. py:attribute:: domain :math:`\mathcal{D}(f)=\{(m_1,m_2)\in\mathbb{R}^2_+\mid m_1\geq m_2>0\}` .. py:class:: ComponentMassesToTotalMassAndMassRatio Bases: :py:obj:`numpyro.distributions.transforms.Transform` Transforms component masses to total mass and mass ratio. .. math:: f: (m_1, m_2) \to (M, q) .. py:method:: eq(other, static: bool = False) .. py:method:: log_abs_det_jacobian(x, y, intermediates=None) .. py:attribute:: codomain .. py:attribute:: domain .. py:class:: DeltaToSymmetricMassRatio Bases: :py:obj:`numpyro.distributions.transforms.Transform` Transforms delta to symmetric mass ratio. .. math:: \eta = f(\delta) = \frac{1-\delta^2}{4} .. math:: \delta = f^{-1}(\eta) = \sqrt{1-4\eta} .. py:method:: eq(other, static: bool = False) .. py:method:: log_abs_det_jacobian(x, y, intermediates=None) .. math:: \ln\left(|\mathrm{det}(J_f)|\right) = \ln(\delta) - \ln(2) .. py:attribute:: codomain :math:`\mathcal{C}(f) = [0, 0.25]` .. py:attribute:: domain :math:`\mathcal{D}(f) = [0, 1]` .. py:class:: PrimaryMassAndMassRatioToComponentMassesTransform Bases: :py:obj:`numpyro.distributions.transforms.Transform` Transforms a primary mass and mass ratio to component masses. .. math:: f: (m_1, q)\to (m_1, m_1q) .. math:: f^{-1}: (m_1, m_2)\to (m_1, m_2/m_1) .. py:method:: eq(other, static: bool = False) .. py:method:: log_abs_det_jacobian(x: jaxtyping.Array, y: jaxtyping.Array, intermediates=None) .. math:: \ln\left(|\mathrm{det}(J_f)|\right) = \ln(|m_1|) .. py:attribute:: codomain :math:`\mathcal{C}(f)=\{(m_1, m_2)\in\mathbb{R}^2_+\mid m_1\geq m_2>0\}` .. py:attribute:: domain :math:`\mathcal{D}(f) = \mathbb{R}^2_+\times[0, 1]` .. py:class:: RedshiftToLuminosityDistance(cosmology: gwkokab.cosmology.Cosmology) Bases: :py:obj:`numpyro.distributions.transforms.Transform` Transforms redshift to luminosity distance. .. math:: f: z \to D_L .. math:: f^{-1}: D_L \to z .. math:: \ln\left(|\mathrm{det}(J_f)|\right) = \ln\left(D_c + \frac{c(1+z)}{H_0 E(z)}\right) .. py:method:: log_abs_det_jacobian(x, y, intermediates=None) .. py:attribute:: codomain :math:`\mathcal{C}(f) = \mathbb{R}_+` .. py:attribute:: cosmology .. py:attribute:: domain :math:`\mathcal{D}(f) = \mathbb{R}_+` .. py:class:: SourceMassAndRedshiftToDetectedMassAndRedshift Bases: :py:obj:`numpyro.distributions.transforms.Transform` Transforms source mass and redshift to detected mass and redshift. .. math:: f: (m_{\text{source}}, z) \to (m_{\text{detected}}, z) .. py:method:: eq(other, static: bool = False) .. py:method:: log_abs_det_jacobian(x, y, intermediates=None) .. math:: \ln\left(|\mathrm{det}(J_f)|\right) = \ln(1+z) .. py:attribute:: codomain :math:`\mathcal{C}(f) = \mathbb{R}^2_+` .. py:attribute:: domain :math:`\mathcal{D}(f) = \mathbb{R}^2_+`