diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 36275bb1..84530c19 100644 --- a/.github/workflows/python-package.yml +++ b/.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 diff --git a/docs/guide/installation.rst b/docs/guide/installation.rst index add8c952..34254b5f 100644 --- a/docs/guide/installation.rst +++ b/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. diff --git a/gwpopulation/models/__init__.py b/gwpopulation/models/__init__.py index aab7eb3b..6b849a06 100644 --- a/gwpopulation/models/__init__.py +++ b/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. """ diff --git a/gwpopulation/models/redshift.py b/gwpopulation/models/redshift.py index 6444b7d4..d1d63df4 100644 --- a/gwpopulation/models/redshift.py +++ b/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""" diff --git a/pyproject.toml b/pyproject.toml index 4334fd33..eb0a163a 100644 --- a/pyproject.toml +++ b/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", diff --git a/test/ln_model_test.py b/test/ln_model_test.py index 9fbdcd69..3507eea3 100644 --- a/test/ln_model_test.py +++ b/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", ]