diff --git a/gwpopulation/models/mass.py b/gwpopulation/models/mass.py index e23a7ddf..5a389e07 100644 --- a/gwpopulation/models/mass.py +++ b/gwpopulation/models/mass.py @@ -12,6 +12,7 @@ ln_powerlaw, ln_powerlaw_from_logs, ln_truncnorm, + logaddexp, powerlaw, trapezoid, truncnorm, @@ -266,7 +267,7 @@ def ln_double_power_law_peak_primary_mass( ln_p_norm = ln_truncnorm( mass, mu=mpp, sigma=sigpp, high=gaussian_mass_maximum, low=mmin ) - return xp.logaddexp(xp.log1p(-lam) + ln_p_pow, xp.log(lam) + ln_p_norm) + return logaddexp(xp.log1p(-lam) + ln_p_pow, xp.log(lam) + ln_p_norm) def double_power_law_primary_power_law_mass_ratio( @@ -655,7 +656,7 @@ def _ln_two_component_single( ln_p_norm = ln_truncnorm( mass, mu=mpp, sigma=sigpp, high=gaussian_mass_maximum, low=mmin ) - return xp.logaddexp(xp.log1p(-lam) + ln_p_pow, xp.log(lam) + ln_p_norm) + return logaddexp(xp.log1p(-lam) + ln_p_pow, xp.log(lam) + ln_p_norm) def _dataset_log(dataset, key): @@ -779,8 +780,8 @@ def ln_three_component_single( ln_p_norm2 = ln_truncnorm( mass, mu=mpp_2, sigma=sigpp_2, high=gaussian_mass_maximum, low=mmin ) - ln_p_peaks = xp.logaddexp(xp.log(lam_1) + ln_p_norm1, xp.log1p(-lam_1) + ln_p_norm2) - return xp.logaddexp(xp.log1p(-lam) + ln_p_pow, xp.log(lam) + ln_p_peaks) + ln_p_peaks = logaddexp(xp.log(lam_1) + ln_p_norm1, xp.log1p(-lam_1) + ln_p_norm2) + return logaddexp(xp.log1p(-lam) + ln_p_pow, xp.log(lam) + ln_p_peaks) def two_component_primary_mass_ratio( @@ -1234,9 +1235,10 @@ def ln_p_q(self, dataset, beta, mmin, delta_m): norm = self._norm_p_q_at( dataset["mass_1"], beta=beta, mmin=mmin, delta_m=delta_m ) - # the normalization, and so p_q, is zero for primary masses below mmin - with np.errstate(divide="ignore", invalid="ignore"): - return xp.where(norm > 0, ln_p_q - xp.log(norm), -xp.inf) + # the normalization, and so p_q, is zero for primary masses below mmin; + # the log of one there, rather than of zero, keeps the gradient finite + positive = norm > 0 + return xp.where(positive, ln_p_q - xp.log(xp.where(positive, norm, 1)), -xp.inf) def _norm_p_q_at(self, mass_1, beta, mmin, delta_m): """ diff --git a/gwpopulation/models/spin.py b/gwpopulation/models/spin.py index 5135a53c..c922aaf1 100644 --- a/gwpopulation/models/spin.py +++ b/gwpopulation/models/spin.py @@ -11,6 +11,7 @@ ln_beta_dist, ln_truncnorm, ln_unnormalized_2d_gaussian, + logaddexp, trapezoid, truncnorm, unnormalized_2d_gaussian, @@ -298,7 +299,7 @@ def ln_independent_spin_orientation_gaussian_isotropic( 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) + return 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): diff --git a/gwpopulation/utils.py b/gwpopulation/utils.py index 50df7efa..059978d2 100644 --- a/gwpopulation/utils.py +++ b/gwpopulation/utils.py @@ -116,12 +116,16 @@ def ln_beta_dist(xx, alpha, beta, scale=1): :math:`[0, x_\max]` """ + inside = (xx >= 0) & (xx <= scale) + # one, rather than a negative number, above the support keeps the gradient + # finite there, e.g., with respect to beta for spin magnitudes above amax + scale_minus_xx = xp.where(inside, scale - xx, 1) # 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.xlogy(alpha - 1, xx) + scs.xlogy(beta - 1, scale_minus_xx) ln_prob -= scs.betaln(alpha, beta) + (alpha + beta - 1) * xp.log(scale) - return xp.where((xx >= 0) & (xx <= scale), ln_prob, -xp.inf) + return xp.where(inside, ln_prob, -xp.inf) @apply_conditions(dict(low=(ge, 0), alpha=(ne, 1))) @@ -565,6 +569,38 @@ def trapezoid(y, x=None, dx=1.0, axis=-1): raise NotImplementedError(f"trapz/trapezoid not implemented for {xp.__name__}") +def logaddexp(ln_a, ln_b): + """ + :math:`\\ln(e^a + e^b)`, as :code:`xp.logaddexp`, with a finite gradient + where both terms are :code:`-inf`. + + The gradient of :code:`jax.numpy.logaddexp(-inf, -inf)` is :code:`nan`, + so the gradient of a log-space mixture, e.g., with respect to its mixing + fraction, is :code:`nan` wherever a sample is outside the support of + every component. With :code:`jax`, both terms are replaced by zero there + before calling :code:`xp.logaddexp`, and its result by :code:`-inf` + after, so the values are unchanged and the gradient there is zero. + + Parameters + ========== + ln_a: array-like + The log of the first term (:math:`a`). + ln_b: array-like + The log of the second term (:math:`b`). + + Returns + ======= + array-like + :math:`\\ln(e^a + e^b)`. + """ + if "jax" not in xp.__name__: + return xp.logaddexp(ln_a, ln_b) + both_zero = xp.maximum(ln_a, ln_b) == -xp.inf + ln_a = xp.where(both_zero, 0.0, ln_a) + ln_b = xp.where(both_zero, 0.0, ln_b) + return xp.where(both_zero, -xp.inf, xp.logaddexp(ln_a, ln_b)) + + def logsumexp(values, axis=-1): """ The log of the sum of the exponentials of the input along an axis, diff --git a/test/ln_model_test.py b/test/ln_model_test.py index 3507eea3..720827a3 100644 --- a/test/ln_model_test.py +++ b/test/ln_model_test.py @@ -15,6 +15,8 @@ from gwpopulation.utils import to_numpy N_TEST = 10 +# gradients are slower to evaluate, so they are tested at fewer draws +N_GRADIENT_TEST = 3 N_SAMPLES = 2000 # Hyperparameters are drawn uniformly from these ranges with a seeded generator, @@ -459,3 +461,62 @@ def test_ln_prob_single_precision(name): finally: gwpopulation.set_backend("jax", enable_x64=True) gwpopulation.set_backend("numpy") + + +@pytest.mark.parametrize("enable_x64", [True, False]) +@pytest.mark.parametrize("name", MODELS) +def test_ln_prob_gradient(name, enable_x64): + """ + The gradient of the log of the summed density with respect to each + hyperparameter is finite, although some samples are outside the support, + e.g., for mixtures, where the gradient of :code:`logaddexp(-inf, -inf)` + is :code:`nan`. In double precision, it matches the gradient computed + from the density wherever that is finite. + """ + jax = pytest.importorskip("jax") + from gwpopulation.utils import logsumexp + + builder, ranges, _ = MODELS[name] + fixed = { + key: value for key, value in ranges.items() if not isinstance(value, tuple) + } + gwpopulation.set_backend("jax", enable_x64=enable_x64) + try: + xp = gwpopulation.utils.xp + dtype = xp.float64 if enable_x64 else xp.float32 + dataset = { + key: xp.asarray(value, dtype=dtype) for key, value in _dataset().items() + } + model = builder() + draws = [ + {key: value for key, value in parameters.items() if key not in fixed} + for parameters in _draws(name)[:N_GRADIENT_TEST] + ] + # the first call caches the interpolants of the spline and smoothed models + model.ln_prob(dataset, **draws[0], **fixed) + + def ln_total(parameters): + return logsumexp(model.ln_prob(dataset, **parameters, **fixed)) + + def total(parameters): + return xp.log(xp.sum(model.prob(dataset, **parameters, **fixed))) + + ln_gradient = jax.jit(jax.grad(ln_total)) + gradient = jax.jit(jax.grad(total)) + for parameters in draws: + ln_result = { + key: float(value) for key, value in ln_gradient(parameters).items() + } + for key, value in ln_result.items(): + assert np.isfinite(value), f"d/d{key} = {value} at {parameters}" + if not enable_x64: + continue + result = {key: float(value) for key, value in gradient(parameters).items()} + for key, value in result.items(): + if np.isfinite(value): + np.testing.assert_allclose( + ln_result[key], value, rtol=1e-8, atol=1e-10, err_msg=key + ) + finally: + gwpopulation.set_backend("jax", enable_x64=True) + gwpopulation.set_backend("numpy")