Ln redshift models

.github/workflows/python-package.yml

@@ -44,7 +44,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.11", "3.12", "3.13"]
 
     steps:
     - uses: actions/checkout@v4

docs/guide/installation.rst

@@ -50,7 +50,7 @@ Installation
          Additional backends can also be specified via editable installations using
          :code:`pip install -e .[jax]`.
 
-Supported python versions: 3.10-3.13.
+Supported python versions: 3.11-3.13.
 This will install :code:`gwpopulation` and all required dependencies (e.g., :code:`bilby`)
 if they are not already available in your system.
 

gwpopulation/models/__init__.py

@@ -83,6 +83,8 @@ def ln_power_law_mass_model(dataset, alpha, minimum_mass, maximum_mass):
 A function defined elsewhere can be attached instead, e.g.,
 :code:`power_law_mass_model.ln_prob = ln_power_law_mass_model`.
 
+The models in :code:`gwpopulation.models` provide their log densities.
+
 We provide a number of pre-defined models along with standard base models that
 can be extended. See the API documentation for more details on implemented models.
 """

gwpopulation/models/redshift.py

@@ -98,7 +98,7 @@ def normalisation(self, parameters):
         return norm
 
     def probability(self, dataset, **parameters):
-        """
+        r"""
         Compute the normalized probability of a merger occurring at the
         specified redshift.
 
@@ -121,6 +121,31 @@ def probability(self, dataset, **parameters):
         )
         return differential_volume / normalisation
 
+    def ln_prob(self, dataset, **parameters):
+        r"""
+        Compute the log of the normalized probability of a merger occurring at
+        the specified redshift, see :func:`probability`.
+
+        Parameters
+        ----------
+        dataset: dict[str, array-like]
+            Dictionary of sample points that contains :code:`redshift`
+            as a key.
+        parameters: dict[str, float]
+            Dictionary of parameters :math:`\Lambda`.
+
+        Returns
+        -------
+        array-like
+            The log of the normalized probability per sample, :code:`-inf`
+            above :code:`self.z_max`
+        """
+        normalisation = self.normalisation(parameters=parameters)
+        ln_differential_volume = self.ln_differential_spacetime_volume(
+            dataset=dataset, bounds=True, **parameters
+        )
+        return ln_differential_volume - xp.log(normalisation)
+
     def psi_of_z(self, redshift, **parameters):
         r"""
         Method encoding the redshift evolution of the merger rate.
@@ -145,6 +170,26 @@ def psi_of_z(self, redshift, **parameters):
         """
         raise NotImplementedError
 
+    def ln_psi_of_z(self, redshift, **parameters):
+        r"""
+        The log of :func:`psi_of_z`. Child classes can overwrite this with a
+        version that does not underflow or overflow, e.g., in single
+        precision; by default this is the log of :func:`psi_of_z`.
+
+        Parameters
+        ----------
+        redshift: array-like
+            The redshifts at which to evaluate the model.
+        parameters: dict[str, float]
+            Dictionary of parameters :math:`\Lambda`.
+
+        Returns
+        -------
+        array-like
+            :math:`\ln \psi(z | \Lambda)` for in input redshifts.
+        """
+        return xp.log(self.psi_of_z(redshift=redshift, **parameters))
+
     def dvc_dz(self, redshift, **parameters):
         r"""
 
@@ -194,6 +239,35 @@ def differential_spacetime_volume(self, dataset, bounds=False, **parameters):
 
         return differential_volume
 
+    def ln_differential_spacetime_volume(self, dataset, bounds=False, **parameters):
+        r"""
+        Compute the log of the differential spacetime volume, see
+        :func:`differential_spacetime_volume`.
+
+        Parameters
+        ----------
+        dataset: dict
+            Dictionary containing entry "redshift"
+        bounds: bool
+            Whether to return :code:`-inf` above :code:`self.z_max`
+        parameters: dict
+            Dictionary of parameters
+
+        Returns
+        -------
+        ln_differential_volume: (float, array-like)
+            Log of the differential spacetime volume
+        """
+        redshift = dataset["redshift"]
+        ln_differential_volume = self.ln_psi_of_z(redshift=redshift, **parameters)
+        ln_differential_volume -= xp.log1p(redshift)
+        ln_differential_volume += xp.log(self.dvc_dz(redshift=redshift, **parameters))
+        if bounds:
+            ln_differential_volume = xp.where(
+                redshift <= self.z_max, ln_differential_volume, -xp.inf
+            )
+        return ln_differential_volume
+
 
 class PowerLawRedshift(_Redshift):
     base_variable_names = ["lamb"]
@@ -214,6 +288,21 @@ def psi_of_z(self, redshift, **parameters):
         """
         return (1 + redshift) ** parameters["lamb"]
 
+    def ln_psi_of_z(self, redshift, **parameters):
+        r"""
+        The log of :func:`psi_of_z`
+
+        .. math::
+
+            \ln \psi(z|\lambda) = \lambda \ln(1 + z)
+
+        Parameters
+        ----------
+        lamb: float
+            The spectral index :math:`\lambda`.
+        """
+        return parameters["lamb"] * xp.log1p(redshift)
+
 
 class MadauDickinsonRedshift(_Redshift):
 
@@ -248,6 +337,35 @@ def psi_of_z(self, redshift, **parameters):
         psi_of_z *= 1 + (1 + z_peak) ** (-kappa)
         return psi_of_z
 
+    def ln_psi_of_z(self, redshift, **parameters):
+        r"""
+        The log of :func:`psi_of_z`
+
+        .. math::
+
+            \ln \psi(z|\gamma, \kappa, z_p) = \gamma \ln(1 + z)
+            - \ln\left(1 + \left(\frac{1 + z}{1 + z_p}\right)^\kappa\right)
+            + \ln\left(1 + (1 + z_p)^{-\kappa}\right)
+
+        Parameters
+        ----------
+        gamma: float
+            Slope of the distribution at low redshift, :math:`\gamma`.
+        kappa: float
+            Slope of the distribution at high redshift, :math:`\kappa`.
+        z_peak: float
+            Redshift at which the distribution peaks, :math:`z_p`.
+        """
+        gamma = parameters["gamma"]
+        kappa = parameters["kappa"]
+        ln_1pz_peak = xp.log1p(parameters["z_peak"])
+        ln_1pz = xp.log1p(redshift)
+        return (
+            gamma * ln_1pz
+            - xp.logaddexp(0.0, kappa * (ln_1pz - ln_1pz_peak))
+            + xp.logaddexp(0.0, -kappa * ln_1pz_peak)
+        )
+
 
 def total_four_volume(lamb, analysis_time, max_redshift=2.3):
     r"""

pyproject.toml

@@ -14,14 +14,13 @@ write_to = "gwpopulation/_version.py"
 name = "gwpopulation"
 description = "Unified population inference"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 license = { file = "LICENSE.md" }
 authors = [
     { name = "Colm Talbot", email = "talbotcolm@gmail.com" }
 ]
 urls = { "Homepage" = "https://github.com/ColmTalbot/gwpopulation" }
 classifiers = [
-    "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",

test/ln_model_test.py

@@ -11,7 +11,7 @@
 
 import gwpopulation
 from gwpopulation.hyperpe import Model
-from gwpopulation.models import mass, spin
+from gwpopulation.models import mass, redshift, spin
 from gwpopulation.utils import to_numpy
 
 N_TEST = 10
@@ -113,6 +113,10 @@ def _outside_chi_eff_chi_p(data, parameters):
     return _outside_chi_eff(data, parameters) | _outside_chi_p(data, parameters)
 
 
+def _outside_redshift(data, parameters):
+    return data["redshift"] > 2.3
+
+
 def _nowhere(data, parameters):
     return np.zeros(data["mass_1"].shape, dtype=bool)
 
@@ -278,6 +282,16 @@ def _model(component):
         _spline_nodes("cos_tilt", -1, 1),
         _nowhere,
     ),
+    PowerLawRedshift=(
+        _model(redshift.PowerLawRedshift),
+        dict(lamb=(-15, 15)),
+        _outside_redshift,
+    ),
+    MadauDickinsonRedshift=(
+        _model(redshift.MadauDickinsonRedshift),
+        dict(gamma=(-15, 15), kappa=(-15, 15), z_peak=(0, 5)),
+        _outside_redshift,
+    ),
 )
 
 # the linear densities of these models do not underflow or overflow in single
@@ -287,6 +301,8 @@ def _model(component):
     "independent_spin_orientation_gaussian_isotropic",
     "SplineSpinMagnitudeIdentical",
     "SplineSpinTiltIdentical",
+    "PowerLawRedshift",
+    "MadauDickinsonRedshift",
 ]