BUG: finite gradients of the log densities outside the support

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

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

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,

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