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")