From 6f76fe2a90ada690746f4d901bda33611019d6a8 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Fri, 23 Oct 2026 14:32:42 -0400 Subject: [PATCH] BUG: finite gradients of the log densities outside the support With jax, the gradient of jnp.logaddexp(-inf, -inf) is nan, so the gradient of every log-space mixture with respect to its mixing fraction was nan whenever a sample is outside the support of every component, e.g., below mmin: ln_two_component_single and the models built on it, ln_three_component_single, ln_double_power_law_peak_primary_mass (where the nan also reached every power-law parameter) and ln_independent_spin_orientation_gaussian_isotropic. This breaks gradient-based samplers, e.g., NUTS, using the log density. - utils.logaddexp: with jax, both terms are replaced by zero where both are -inf before calling xp.logaddexp, and its result by -inf after, so the gradient there is zero. Used by every mixture. - BaseSmoothedMassDistribution.ln_p_q: the log of one, rather than of zero, where the mass ratio normalization vanishes (m1 < mmin), which gave nan gradients with respect to beta, mmin and delta_m. - ln_beta_dist: one in place of scale - xx above the support, which gave nan gradients with respect to beta for spin magnitudes above amax, also for beta_dist. The values are unchanged: bit-identical for every model with numpy and with jit-compiled jax in double precision. In single precision under jit, XLA fuses ln_two_component_single differently, changing values at the level of float32 rounding (at most 4e-6) with the same accuracy. test_ln_prob_gradient checks that the gradient of the log of the summed density with respect to every hyperparameter, with samples outside the support, is finite in single and double precision and, in double precision, matches the gradient computed from the density wherever that is finite. Still not finite: the derivative with respect to a mixing fraction of exactly zero or one in log space, and the gradients of the linear smoothed models' p_q with respect to beta, mmin and delta_m. Co-Authored-By: Claude Opus 5.5 --- gwpopulation/models/mass.py | 16 +++++----- gwpopulation/models/spin.py | 3 +- gwpopulation/utils.py | 40 ++++++++++++++++++++++-- test/ln_model_test.py | 61 +++++++++++++++++++++++++++++++++++++ 4 files changed, 110 insertions(+), 10 deletions(-) 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")