Compute the LogitNormal mean by quadrature

RELEASE_NOTES.md

@@ -20,6 +20,7 @@
 
 ### Bug fixes
 
+* Fix `LogitNormal.mean()`, and so `predict()`, which averaged 1000 random draws and gave a different value on each call. It is now computed by quadrature, is the same on every call, and is accurate to about 1e-14. The `make_distribution` docstring now warns that a regression distribution with neither `scipy_dist_cls` nor `mean_fn` gets the same random mean (PR #449)
 * Fix the `CRPScore` metric of `Normal`, `NormalFixedMean`, and `LogNormal`, which was off by a factor of the predicted scale. Fits using `CRPScore` with these distributions now produce correct natural gradients (issue #421, PR #424)
 * Fix the scale gradient of the `Laplace` `CRPScore`, which dropped a `0.75 * scale` term (PR #419)
 * Fix censored scores and gradients for `Exponential` and `LogNormal`, which swapped censored and observed rows and had sign and extra-term errors (PR #420)

docs/sympy_distributions.md

@@ -108,6 +108,11 @@ The factory auto-derives score, gradients, and Fisher Information
 from the SymPy distribution, and auto-generates fit/sample/mean (regression)
 or class_probs/sample (classification).
 
+Without `scipy_dist_cls` or `mean_fn`, a regression distribution's `mean()`
+is the average of 1000 draws from `sample()`. That mean, and so `predict()`,
+changes from call to call. Pass `mean_fn` to get a deterministic mean, as
+`LogitNormal` does.
+
 ## When You Need a Manual Score Expression
 
 For distributions without a `sympy.stats` equivalent (e.g., LogitNormal), or

ngboost/distns/logitnormal.py

@@ -9,10 +9,13 @@
 
 Uses a **manual score expression** (no ``sympy.stats`` equivalent) with the
 ``make_distribution`` factory.  Fisher Information is computed via Monte Carlo.
+The mean has no closed form and is computed by deterministic quadrature.
 """
 
 import numpy as np
 import sympy as sp
+from numpy.polynomial.hermite_e import hermegauss
+from scipy.special import expit, ndtr  # pylint: disable=no-name-in-module
 
 from ngboost.distns.sympy_utils import make_distribution
 
@@ -36,10 +39,46 @@ def _sample_logitnormal(self, m):
     )
 
 
+# The mean E[expit(mu + sigma * Z)] with Z ~ N(0, 1) has no closed form.  It is
+# computed with one of two quadrature rules, chosen by sigma:
+#
+# - sigma <= 1: 64-node Gauss-Hermite quadrature over Z.
+# - sigma > 1: expit(mu + sigma * z) is close to a step in z, and Gauss-Hermite
+#   loses accuracy (error 3e-3 at sigma = 10).  Instead write the mean as
+#   P(L < mu + sigma * Z) = E[Phi((mu - L) / sigma)] with L standard logistic,
+#   and integrate over L with the trapezoidal rule, which converges
+#   geometrically for this smooth integrand.
+#
+# On its own side of sigma = 1, each rule agrees with a high-precision
+# reference to about 1e-14.
+_GH_NODES, _GH_WEIGHTS = hermegauss(64)
+_GH_WEIGHTS = _GH_WEIGHTS / np.sqrt(2 * np.pi)
+# step 0.5 on [-40, 40]; the logistic mass outside is 2 * expit(-40), about 8e-18
+_LOGIS_NODES = np.linspace(-40, 40, 161)
+_LOGIS_WEIGHTS = 0.5 * expit(_LOGIS_NODES) * expit(-_LOGIS_NODES)
+
+
+def _mean_logitnormal(self):
+    """Mean of each row by quadrature (see the note above), so it is deterministic."""
+    mu, sigma = np.broadcast_arrays(self.mu, self.sigma)
+    mean = np.empty(mu.shape)
+    narrow = sigma <= 1
+    mu_n, sigma_n = mu[narrow], sigma[narrow]
+    mean[narrow] = sum(
+        w * expit(mu_n + sigma_n * z) for z, w in zip(_GH_NODES, _GH_WEIGHTS)
+    )
+    mu_w, sigma_w = mu[~narrow], sigma[~narrow]
+    mean[~narrow] = sum(
+        w * ndtr((mu_w - t) / sigma_w) for t, w in zip(_LOGIS_NODES, _LOGIS_WEIGHTS)
+    )
+    return mean
+
+
 LogitNormal = make_distribution(
     params=[(_mu, False), (_sigma, True)],  # mu: identity, sigma: log link
     y=_y,
     score_expr=_score_expr,
     sample_fn=_sample_logitnormal,
+    mean_fn=_mean_logitnormal,
     name="LogitNormal",
 )

ngboost/distns/sympy_utils.py

@@ -549,6 +549,11 @@ def make_distribution(  # pylint: disable=R0912,R0913,R0914,R0915,R0917
     mean_fn : callable or None
         Custom ``mean(self) -> np.ndarray`` override.  When provided,
         takes precedence over scipy delegation and MC fallback.
+        A regression distribution with neither ``scipy_dist_cls`` nor
+        ``mean_fn`` gets the MC fallback: ``mean()`` averages 1000 draws
+        from ``sample()``.  That mean, and so ``predict()``, is random and
+        changes from call to call.  Pass ``mean_fn`` to make it
+        deterministic.
     class_prob_exprs : list[sympy.Expr] or None
         SymPy expressions for class probabilities ``[P(Y=0), P(Y=1), ...]``.
         When provided, the result is a ``ClassificationDistn`` subclass
@@ -817,6 +822,7 @@ def _class_probs(self):
     if mean_fn is not None:
         _mean = mean_fn
     elif not is_classification and scipy_dist_cls is None:
+        # Monte Carlo mean: random from call to call (see ``mean_fn`` above).
 
         def _mean(self):
             return np.mean(self.sample(1000), axis=0)

tests/test_logitnormal.py

@@ -0,0 +1,53 @@
+# pylint: disable=no-member,no-name-in-module
+import numpy as np
+import pytest
+from scipy.integrate import quad
+from scipy.special import expit
+from scipy.stats import norm
+
+from ngboost import NGBRegressor
+from ngboost.distns import LogitNormal
+
+SIGMAS = [0.01, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0]
+MUS = np.linspace(-10, 10, 9)
+
+
+def _logitnormal(mu, sigma):
+    return LogitNormal(np.array([mu, np.full_like(mu, np.log(sigma))]))
+
+
+def test_predict_is_deterministic():
+    rng = np.random.default_rng(0)
+    X = rng.normal(size=(100, 3))
+    Y = expit(X[:, 0] + 0.5 * rng.normal(size=100))
+    ngb = NGBRegressor(Dist=LogitNormal, n_estimators=5, verbose=False).fit(X, Y)
+    np.testing.assert_array_equal(ngb.predict(X), ngb.predict(X))
+
+
+@pytest.mark.parametrize("sigma", SIGMAS)
+def test_mean_matches_monte_carlo(sigma):
+    rng = np.random.default_rng(1)
+    draws = expit(MUS + sigma * rng.standard_normal((1_000_000, len(MUS))))
+    mc_mean = draws.mean(axis=0)
+    mc_se = draws.std(axis=0) / np.sqrt(len(draws))
+    mean = _logitnormal(MUS, sigma).mean()
+    np.testing.assert_array_less(np.abs(mean - mc_mean), 5 * mc_se + 1e-12)
+
+
+@pytest.mark.parametrize("sigma", SIGMAS)
+def test_mean_matches_adaptive_quadrature(sigma):
+    def integrand(z, mu):
+        return expit(mu + sigma * z) * norm.pdf(z)
+
+    expected = []
+    for mu in MUS:
+        # split where expit(mu + sigma * z) changes fastest
+        z0 = np.clip(-mu / sigma, -12, 12)
+        expected.append(
+            quad(integrand, -12, z0, args=(mu,), epsabs=1e-13)[0]
+            + quad(integrand, z0, 12, args=(mu,), epsabs=1e-13)[0]
+        )
+    dist = _logitnormal(MUS, sigma)
+    np.testing.assert_allclose(dist.mean(), expected, rtol=0, atol=1e-10)
+    # a single row gives the same mean as the full batch
+    assert dist[4].mean() == dist.mean()[4]