From 2a87a522c7d02cf285f346b256209449017111fd Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Fri, 23 Oct 2026 13:40:05 -0400 Subject: [PATCH] FEAT: log densities for the spin models Add ln_ versions of gaussian_chi_eff, gaussian_chi_p, the beta spin magnitude models (iid and independent), the isotropic + Gaussian orientation mixtures and iid_spin, and GaussianChiEffChiP.ln_prob. hyperpe.Model uses these, so in single precision the densities no longer underflow to zero far from narrow peaks: in a float32 survey (10k samples x 10 draws per model), gaussian_chi_eff and GaussianChiEffChiP lost 6% of points to underflow and now lose none. - utils.ln_beta_dist, with beta_dist now exp(ln_beta_dist). It uses xlogy, so a Beta(1, b) density at a = 0 (and Beta(a, 1) at a = amax) is now finite rather than 0; elsewhere values agree to rounding. - utils.ln_unnormalized_2d_gaussian, the exponent of unnormalized_2d_gaussian, which is now its exp. GaussianChiEffChiP's log density is this minus the log of the unchanged grid normalization. - The orientation mixtures use logaddexp; as in the linear model, the isotropic component is not truncated. test/ln_model_test.py checks each log density against the log of the density (numpy and jax backends) and -inf outside the support, and the jit-compiled float32 log density against float64 on the same rounded inputs. It also checks that the linear float32 density loses points for the hyperparameters used, so the comparison covers underflow. Co-Authored-By: Claude Opus 5.5 --- gwpopulation/models/spin.py | 226 +++++++++++++++++++++++++++++++- gwpopulation/utils.py | 82 ++++++++++-- test/ln_model_test.py | 252 ++++++++++++++++++++++++++++++++++++ 3 files changed, 550 insertions(+), 10 deletions(-) create mode 100644 test/ln_model_test.py diff --git a/gwpopulation/models/spin.py b/gwpopulation/models/spin.py index 373a87e2..ac143bf2 100644 --- a/gwpopulation/models/spin.py +++ b/gwpopulation/models/spin.py @@ -2,9 +2,19 @@ Implemented spin models """ +import math + import numpy as xp -from ..utils import beta_dist, trapezoid, truncnorm, unnormalized_2d_gaussian +from ..utils import ( + beta_dist, + ln_beta_dist, + ln_truncnorm, + ln_unnormalized_2d_gaussian, + trapezoid, + truncnorm, + unnormalized_2d_gaussian, +) from .interped import InterpolatedNoBaseModelIdentical __all__ = [ @@ -12,12 +22,19 @@ "SplineSpinMagnitudeIdentical", "SplineSpinTiltIdentical", "iid_spin", + "ln_iid_spin", "iid_spin_magnitude_beta", + "ln_iid_spin_magnitude_beta", "independent_spin_magnitude_beta", + "ln_independent_spin_magnitude_beta", "iid_spin_orientation_gaussian_isotropic", + "ln_iid_spin_orientation_gaussian_isotropic", "independent_spin_orientation_gaussian_isotropic", + "ln_independent_spin_orientation_gaussian_isotropic", "gaussian_chi_eff", + "ln_gaussian_chi_eff", "gaussian_chi_p", + "ln_gaussian_chi_p", ] @@ -47,6 +64,30 @@ def iid_spin(dataset, xi_spin, sigma_spin, amax, alpha_chi, beta_chi): return prior +def ln_iid_spin(dataset, xi_spin, sigma_spin, amax, alpha_chi, beta_chi): + r""" + Log of :func:`iid_spin`, computed without leaving log space so that it + does not underflow, e.g., in single precision. + + Parameters + ---------- + dataset: dict + Dictionary of numpy arrays containing 'a_1', 'a_2', 'cos_tilt_1', and + 'cos_tilt_2'. + xi_spin: float + Fraction of black holes in preferentially aligned component. + sigma_spin: float + Width of preferentially aligned component. + alpha_chi, beta_chi: float + Parameters of Beta distribution for both black holes. + amax: float + Maximum black hole spin. + """ + return ln_iid_spin_orientation_gaussian_isotropic( + dataset, xi_spin, sigma_spin + ) + ln_iid_spin_magnitude_beta(dataset, amax, alpha_chi, beta_chi) + + def iid_spin_magnitude_beta(dataset, amax=1, alpha_chi=1, beta_chi=1): """ Independent and identically distributed beta distributions for both spin magnitudes. @@ -67,6 +108,24 @@ def iid_spin_magnitude_beta(dataset, amax=1, alpha_chi=1, beta_chi=1): ) +def ln_iid_spin_magnitude_beta(dataset, amax=1, alpha_chi=1, beta_chi=1): + """ + Log of :func:`iid_spin_magnitude_beta`. + + Parameters + ---------- + dataset: dict + Dictionary of numpy arrays containing 'a_1' and 'a_2'. + alpha_chi, beta_chi: float + Parameters of Beta distribution for both black holes. + amax: float + Maximum black hole spin. + """ + return ln_independent_spin_magnitude_beta( + dataset, alpha_chi, alpha_chi, beta_chi, beta_chi, amax, amax + ) + + def independent_spin_magnitude_beta( dataset, alpha_chi_1, alpha_chi_2, beta_chi_1, beta_chi_2, amax_1, amax_2 ): @@ -92,6 +151,28 @@ def independent_spin_magnitude_beta( return prior +def ln_independent_spin_magnitude_beta( + dataset, alpha_chi_1, alpha_chi_2, beta_chi_1, beta_chi_2, amax_1, amax_2 +): + """ + Log of :func:`independent_spin_magnitude_beta`. + + Parameters + ---------- + dataset: dict + Dictionary of numpy arrays containing 'a_1' and 'a_2'. + alpha_chi_1, beta_chi_1: float + Parameters of Beta distribution for more massive black hole. + alpha_chi_2, beta_chi_2: float + Parameters of Beta distribution for less massive black hole. + amax_1, amax_2: float + Maximum spin of the more/less massive black hole. + """ + return ln_beta_dist( + dataset["a_1"], alpha_chi_1, beta_chi_1, scale=amax_1 + ) + ln_beta_dist(dataset["a_2"], alpha_chi_2, beta_chi_2, scale=amax_2) + + def iid_spin_orientation_gaussian_isotropic(dataset, xi_spin, sigma_spin, mu_spin=1): r""" A mixture model of spin orientations with isotropic and normally @@ -124,6 +205,26 @@ def iid_spin_orientation_gaussian_isotropic(dataset, xi_spin, sigma_spin, mu_spi ) +def ln_iid_spin_orientation_gaussian_isotropic(dataset, xi_spin, sigma_spin, mu_spin=1): + r""" + Log of :func:`iid_spin_orientation_gaussian_isotropic`. + + Parameters + ---------- + dataset: dict + Dictionary of numpy arrays for 'cos_tilt_1' and 'cos_tilt_2'. + xi_spin: float + Fraction of black holes in gaussian component (:math:`\xi`). + sigma_spin: float + Width of gaussian component. + mu_spin: float + Mean of the gaussian component (default: cos_tilt = 1). + """ + return ln_independent_spin_orientation_gaussian_isotropic( + dataset, xi_spin, sigma_spin, sigma_spin, mu_spin, mu_spin + ) + + def independent_spin_orientation_gaussian_isotropic( dataset, xi_spin, sigma_1, sigma_2, mu_1=1, mu_2=1 ): @@ -166,6 +267,40 @@ def independent_spin_orientation_gaussian_isotropic( return prior +def ln_independent_spin_orientation_gaussian_isotropic( + dataset, xi_spin, sigma_1, sigma_2, mu_1=1, mu_2=1 +): + r""" + Log of :func:`independent_spin_orientation_gaussian_isotropic`, computed + without leaving log space so that the Gaussian component does not + underflow far from its mean, e.g., in single precision. + + Parameters + ---------- + dataset: dict + Dictionary of numpy arrays for 'cos_tilt_1' and 'cos_tilt_2'. + xi_spin: float + Fraction of black holes in gaussian component (:math:`\xi`). + mu_1: float + Mean of the gaussian component for the more + massive black hole (:math:`\mu_1`). + mu_2: float + Mean of the gaussian component for the less + massive black hole (:math:`\mu_2`). + sigma_1: float + Width of the gaussian component for the more + massive black hole (:math:`\sigma_1`). + sigma_2: float + Width of the gaussian component for the less + massive black hole (:math:`\sigma_2`). + """ + ln_gaussian = ln_truncnorm( + dataset["cos_tilt_1"], mu=mu_1, sigma=sigma_1, high=1, low=-1 + ) + ln_truncnorm(dataset["cos_tilt_2"], mu=mu_2, sigma=sigma_2, high=1, low=-1) + # as for the linear model, the isotropic component is not truncated + return xp.logaddexp(xp.log1p(-xi_spin) - math.log(4), xp.log(xi_spin) + ln_gaussian) + + def gaussian_chi_eff(dataset, mu_chi_eff, sigma_chi_eff): r""" A Gaussian in chi effective distribution @@ -197,6 +332,29 @@ def gaussian_chi_eff(dataset, mu_chi_eff, sigma_chi_eff): ) +def ln_gaussian_chi_eff(dataset, mu_chi_eff, sigma_chi_eff): + r""" + Log of :func:`gaussian_chi_eff`, which does not underflow far from the + mean, e.g., in single precision. + + Parameters + ---------- + dataset: dict + Input data, must contain `chi_eff` (:math:`\chi_{\text{eff}}`) + mu_chi_eff: float + Mean of the distribution (:math:`\mu_\chi`) + sigma_chi_eff: float + Standard deviation of the distribution (:math:`\sigma_\chi`) + + Returns + ------- + array-like: The log probability, :code:`-inf` outside :math:`[-1, 1]` + """ + return ln_truncnorm( + dataset["chi_eff"], mu=mu_chi_eff, sigma=sigma_chi_eff, low=-1, high=1 + ) + + def gaussian_chi_p(dataset, mu_chi_p, sigma_chi_p): r""" A Gaussian distribution in precessing effective spin (chi p) @@ -224,6 +382,27 @@ def gaussian_chi_p(dataset, mu_chi_p, sigma_chi_p): return truncnorm(dataset["chi_p"], mu=mu_chi_p, sigma=sigma_chi_p, low=0, high=1) +def ln_gaussian_chi_p(dataset, mu_chi_p, sigma_chi_p): + r""" + Log of :func:`gaussian_chi_p`, which does not underflow far from the + mean, e.g., in single precision. + + Parameters + ---------- + dataset: dict + Input data, must contain `chi_p` (:math:`\chi_p`) + mu_chi_p: float + Mean of the distribution (:math:`\mu_\chi`) + sigma_chi_p: float + Standard deviation of the distribution (:math:`\sigma_\chi`) + + Returns + ------- + array-like: The log probability, :code:`-inf` outside :math:`[0, 1]` + """ + return ln_truncnorm(dataset["chi_p"], mu=mu_chi_p, sigma=sigma_chi_p, low=0, high=1) + + class GaussianChiEffChiP(object): r""" A covariant Gaussian in effective aligned and precessing spins. @@ -291,6 +470,51 @@ def __call__( prob *= (dataset["chi_p"] <= 1) * (dataset["chi_p"] >= 0) return prob + def ln_prob( + self, dataset, mu_chi_eff, sigma_chi_eff, mu_chi_p, sigma_chi_p, spin_covariance + ): + """ + The log of :func:`__call__`: the quadratic form of the Gaussian minus + the log of its normalization, which does not underflow far from the + mean, e.g., in single precision. + + Parameters + ---------- + dataset: dict + Dictionary of numpy arrays for :code:`chi_eff` and :code:`chi_p`. + mu_chi_eff, sigma_chi_eff, mu_chi_p, sigma_chi_p, spin_covariance: float + As for :func:`__call__`. + + Returns + ------- + array-like + The log probability, :code:`-inf` outside the support. + """ + ln_prob = ln_unnormalized_2d_gaussian( + dataset["chi_eff"], + dataset["chi_p"], + mu_chi_eff, + mu_chi_p, + sigma_chi_eff, + sigma_chi_p, + spin_covariance, + ) + ln_prob -= xp.log( + self._normalization( + mu_chi_eff=mu_chi_eff, + sigma_chi_eff=sigma_chi_eff, + mu_chi_p=mu_chi_p, + sigma_chi_p=sigma_chi_p, + spin_covariance=spin_covariance, + ) + ) + in_bounds = ( + (xp.abs(dataset["chi_eff"]) <= 1) + & (dataset["chi_p"] <= 1) + & (dataset["chi_p"] >= 0) + ) + return xp.where(in_bounds, ln_prob, -xp.inf) + def _normalization( self, mu_chi_eff, sigma_chi_eff, mu_chi_p, sigma_chi_p, spin_covariance ): diff --git a/gwpopulation/utils.py b/gwpopulation/utils.py index e98ff487..3b7f4d48 100644 --- a/gwpopulation/utils.py +++ b/gwpopulation/utils.py @@ -85,13 +85,43 @@ def beta_dist(xx, alpha, beta, scale=1): The distribution evaluated at `xx` """ - ln_beta = (alpha - 1) * xp.log(xx) + (beta - 1) * xp.log(scale - xx) - ln_beta -= scs.betaln(alpha, beta) - ln_beta -= (alpha + beta - 1) * xp.log(scale) - prob = xp.exp(ln_beta) - prob = xp.nan_to_num(prob) - prob *= (xx >= 0) * (xx <= scale) - return prob + # positional arguments, as the conditions have been checked + return xp.nan_to_num(xp.exp(ln_beta_dist(xx, alpha, beta, scale))) + + +@apply_conditions(dict(alpha=(gt, 0), beta=(gt, 0), scale=(gt, 0))) +def ln_beta_dist(xx, alpha, beta, scale=1): + r""" + Log beta distribution probability + + .. math:: + \ln p(x) = (\alpha - 1) \ln x + (\beta - 1) \ln(x_\max - x) + - \ln B(\alpha, \beta) - (\alpha + \beta - 1) \ln x_\max + + Parameters + ---------- + xx: float, array-like + The abscissa values (:math:`x`) + alpha: float + The Beta alpha parameter (:math:`\alpha`) + beta: float + The Beta beta parameter (:math:`\beta`) + scale: float, array-like + A scale factor for the distribution of the distribution (:math:`x_\max`) + + Returns + ------- + ln_prob: float, array-like + The log of the distribution evaluated at `xx`, :code:`-inf` outside + :math:`[0, x_\max]` + + """ + # with numpy, the logs warn where the result is masked below + with np.errstate(divide="ignore", invalid="ignore"): + # xlogy keeps x^0 = 1 at x = 0, e.g., alpha = 1 at xx = 0 + ln_prob = scs.xlogy(alpha - 1, xx) + scs.xlogy(beta - 1, scale - xx) + ln_prob -= scs.betaln(alpha, beta) + (alpha + beta - 1) * xp.log(scale) + return xp.where((xx >= 0) & (xx <= scale), ln_prob, -xp.inf) @apply_conditions(dict(low=(ge, 0), alpha=(ne, 1))) @@ -355,15 +385,49 @@ def unnormalized_2d_gaussian(xx, yy, mu_x, mu_y, sigma_x, sigma_y, covariance): prob: array-like The unnormalized probability distribution (:math:`p(x)`) """ + return xp.exp( + ln_unnormalized_2d_gaussian( + xx, yy, mu_x, mu_y, sigma_x, sigma_y, covariance=covariance + ) + ) + + +def ln_unnormalized_2d_gaussian(xx, yy, mu_x, mu_y, sigma_x, sigma_y, covariance): + r""" + The log of :func:`unnormalized_2d_gaussian`, the quadratic form in its + exponent, which does not underflow far from the mean. + + Parameters + ========== + xx: array-like + Input data in first dimension (:math:`x`) + yy: array-like + Input data in second dimension (:math:`y`) + mu_x: float + Mean in the first dimension (:math:`\mu_{x}`) + sigma_x: float + Standard deviation in the first dimension (:math:`\sigma_{x}`) + mu_y: float + Mean in the second dimension (:math:`\mu_{y}`) + sigma_y: float + Standard deviation in the second dimension (:math:`\sigma_{y}`) + covariance: float + The normalized covariance (bounded in [0, 1]) (:math:`\rho`) + + Returns + ------- + ln_prob: array-like + The log of the unnormalized probability distribution + (:math:`\ln p(x)`) + """ determinant = sigma_x**2 * sigma_y**2 * (1 - covariance) residual_x = (mu_x - xx) * sigma_y residual_y = (mu_y - yy) * sigma_x - prob = xp.exp( + return ( -(residual_x**2 + residual_y**2 - 2 * residual_x * residual_y * covariance) / 2 / determinant ) - return prob def von_mises(xx, mu, kappa): diff --git a/test/ln_model_test.py b/test/ln_model_test.py new file mode 100644 index 00000000..f931df18 --- /dev/null +++ b/test/ln_model_test.py @@ -0,0 +1,252 @@ +""" +The log densities of the population models: each against the log of the +density, in single precision against double precision, and :code:`-inf` +outside the support. +""" + +import sys + +import numpy as np +import pytest + +import gwpopulation +from gwpopulation.hyperpe import Model +from gwpopulation.models import spin +from gwpopulation.utils import to_numpy + +N_TEST = 10 +N_SAMPLES = 2000 + +# Hyperparameters are drawn uniformly from these ranges with a seeded generator, +# so that the single precision comparisons are reproducible. The ranges include +# narrow peaks and steep spin magnitude distributions, so that the densities +# underflow in single precision. +TILT = dict(xi_spin=(0, 1), sigma_spin=(0.1, 4)) +TILTS = dict( + xi_spin=(0, 1), sigma_1=(0.1, 4), sigma_2=(0.1, 4), mu_1=(-1, 1), mu_2=(-1, 1) +) +MAGNITUDE = dict(amax=(0.3, 1), alpha_chi=(1, 30), beta_chi=(1, 30)) +MAGNITUDES = dict( + amax_1=(0.3, 1), + amax_2=(0.3, 1), + alpha_chi_1=(1, 30), + alpha_chi_2=(1, 30), + beta_chi_1=(1, 30), + beta_chi_2=(1, 30), +) +CHI_EFF = dict(mu_chi_eff=(-1, 1), sigma_chi_eff=(0.02, 0.3)) +CHI_P = dict(mu_chi_p=(0, 1), sigma_chi_p=(0.02, 0.1)) + + +# points that are outside the support for any of the hyperparameters above +def _outside_magnitude(data, parameters): + return (data["a_1"] > parameters["amax"]) | (data["a_2"] > parameters["amax"]) + + +def _outside_magnitudes(data, parameters): + return (data["a_1"] > parameters["amax_1"]) | (data["a_2"] > parameters["amax_2"]) + + +def _outside_chi_eff(data, parameters): + return np.abs(data["chi_eff"]) > 1 + + +def _outside_chi_p(data, parameters): + return (data["chi_p"] < 0) | (data["chi_p"] > 1) + + +def _outside_chi_eff_chi_p(data, parameters): + return _outside_chi_eff(data, parameters) | _outside_chi_p(data, parameters) + + +def _nowhere(data, parameters): + return np.zeros(data["mass_1"].shape, dtype=bool) + + +def _model(component): + return lambda: Model([component()], cache=False) + + +# name: (builder of a Model, ranges, outside the support) +MODELS = dict( + iid_spin=(_model(lambda: spin.iid_spin), TILT | MAGNITUDE, _outside_magnitude), + iid_spin_magnitude_beta=( + _model(lambda: spin.iid_spin_magnitude_beta), + MAGNITUDE, + _outside_magnitude, + ), + independent_spin_magnitude_beta=( + _model(lambda: spin.independent_spin_magnitude_beta), + MAGNITUDES, + _outside_magnitudes, + ), + iid_spin_orientation_gaussian_isotropic=( + _model(lambda: spin.iid_spin_orientation_gaussian_isotropic), + TILT | dict(mu_spin=(-1, 1)), + _nowhere, + ), + independent_spin_orientation_gaussian_isotropic=( + _model(lambda: spin.independent_spin_orientation_gaussian_isotropic), + TILTS, + _nowhere, + ), + # without the isotropic component, the density can underflow + aligned_spin_orientation=( + _model(lambda: spin.iid_spin_orientation_gaussian_isotropic), + dict(xi_spin=1.0, sigma_spin=(0.05, 0.3), mu_spin=(0, 1)), + _nowhere, + ), + gaussian_chi_eff=( + _model(lambda: spin.gaussian_chi_eff), + CHI_EFF, + _outside_chi_eff, + ), + gaussian_chi_p=(_model(lambda: spin.gaussian_chi_p), CHI_P, _outside_chi_p), + GaussianChiEffChiP=( + _model(spin.GaussianChiEffChiP), + CHI_EFF | CHI_P | dict(spin_covariance=(-0.75, 0.75)), + _outside_chi_eff_chi_p, + ), +) + +# the linear densities of these models do not underflow or overflow in single +# precision for the ranges above +LINEAR_SINGLE_PRECISION = [ + "iid_spin_orientation_gaussian_isotropic", + "independent_spin_orientation_gaussian_isotropic", +] + + +def _never(data, parameters, ln_expected): + return np.zeros(ln_expected.shape, dtype=bool) + + +# Single precision tolerances, where points may be lost (lost), and the log +# densities above which points are compared (compare_above). +SINGLE_PRECISION = { + name: dict(rtol=2e-5, atol=2e-5, lost=_never, compare_above=-np.inf) + for name in MODELS +} + + +def _dataset(): + rng = np.random.default_rng(10) + mass_1 = np.exp(rng.uniform(np.log(2), np.log(150), N_SAMPLES)) + mass_ratio = rng.uniform(0.01, 1.1, N_SAMPLES) + return dict( + mass_1=mass_1, + mass_ratio=mass_ratio, + mass_2=mass_1 * mass_ratio, + a_1=rng.uniform(0, 1, N_SAMPLES), + a_2=rng.uniform(0, 1, N_SAMPLES), + cos_tilt_1=rng.uniform(-1, 1, N_SAMPLES), + cos_tilt_2=rng.uniform(-1, 1, N_SAMPLES), + chi_eff=rng.uniform(-1.2, 1.2, N_SAMPLES), + chi_p=rng.uniform(-0.2, 1.2, N_SAMPLES), + redshift=rng.uniform(0, 2.5, N_SAMPLES), + ) + + +def _draws(name): + rng = np.random.default_rng(sum(map(ord, name))) + ranges = MODELS[name][1] + return [ + { + key: float(rng.uniform(*value)) if isinstance(value, tuple) else value + for key, value in ranges.items() + } + for _ in range(N_TEST) + ] + + +def _single(values): + """Round to single precision, keeping double precision types""" + if isinstance(values, dict): + return {key: _single(value) for key, value in values.items()} + elif isinstance(values, np.ndarray): + return values.astype(np.float32).astype(np.float64) + return float(np.float32(values)) + + +@pytest.mark.parametrize("name", MODELS) +def test_ln_function_is_used(name): + model = MODELS[name][0]() + component = model.models[0] + ln_function = model.ln_function(component) + if hasattr(component, "ln_prob"): + assert ln_function == component.ln_prob + else: + module = sys.modules[component.__module__] + assert ln_function is getattr(module, f"ln_{component.__name__}") + + +@pytest.mark.parametrize("name", MODELS) +def test_ln_prob_matches_log_prob(backend, name): + gwpopulation.set_backend(backend) + xp = gwpopulation.utils.xp + builder, _, outside = MODELS[name] + model = builder() + data = _dataset() + dataset = {key: xp.asarray(value) for key, value in data.items()} + for parameters in _draws(name): + prob = to_numpy(model.prob(dataset, **parameters)) + ln_prob = to_numpy(model.ln_prob(dataset, **parameters)) + assert not np.any(np.isnan(ln_prob)) + positive = prob > 1e-300 + assert np.any(positive) + np.testing.assert_allclose( + ln_prob[positive], np.log(prob[positive]), rtol=1e-10, atol=1e-12 + ) + # where the density underflows, the log density is finite + assert np.all(ln_prob[prob == 0] < np.log(1e-300)) + unsupported = outside(data, parameters) + assert np.all(prob[unsupported] == 0) + assert np.all(np.isneginf(ln_prob[unsupported])) + gwpopulation.set_backend("numpy") + + +@pytest.mark.parametrize("name", MODELS) +def test_ln_prob_single_precision(name): + jax = pytest.importorskip("jax") + builder, _, outside = MODELS[name] + # round the inputs so that only the arithmetic differs + data = _single(_dataset()) + draws = [_single(parameters) for parameters in _draws(name)] + gwpopulation.set_backend("numpy") + model = builder() + expected = [to_numpy(model.ln_prob(data, **parameters)) for parameters in draws] + + gwpopulation.set_backend("jax", enable_x64=False) + try: + xp = gwpopulation.utils.xp + dataset = { + key: xp.asarray(value, dtype=xp.float32) for key, value in data.items() + } + model = builder() + jitted = jax.jit(lambda parameters: model.ln_prob(dataset, **parameters)) + settings = SINGLE_PRECISION[name] + n_linear_lost = 0 + for parameters, ln_expected in zip(draws, expected): + ln_prob = jitted(parameters) + assert ln_prob.dtype == xp.float32 + ln_prob = to_numpy(ln_prob) + assert not np.any(np.isnan(ln_prob)) + assert np.all(np.isneginf(ln_prob[outside(data, parameters)])) + finite = np.isfinite(ln_expected) + assert not np.any(np.isfinite(ln_prob) & ~finite) + lost = finite & np.isneginf(ln_prob) + assert np.all(settings["lost"](data, parameters, ln_expected)[lost]) + compare = finite & ~lost & (ln_expected > settings["compare_above"]) + np.testing.assert_allclose( + ln_prob[compare], + ln_expected[compare], + rtol=settings["rtol"], + atol=settings["atol"], + ) + prob = to_numpy(model.prob(dataset, **parameters)) + n_linear_lost += np.sum(finite & ~(prob > 0)) + # the log density is needed: the density underflows or overflows + assert (n_linear_lost > 0) != (name in LINEAR_SINGLE_PRECISION) + finally: + gwpopulation.set_backend("jax", enable_x64=True) + gwpopulation.set_backend("numpy")