From 279cf7db511977a5e8e24e2170c388f3c0e7c0f1 Mon Sep 17 00:00:00 2001 From: Rudy Celekli Date: Wed, 4 Nov 2026 10:39:06 -0500 Subject: [PATCH 1/2] Preserve censored log-likelihoods and gradients in the right tail --- ngboost/distns/exponential.py | 6 +-- ngboost/distns/lognormal.py | 11 ++--- tests/test_survival_logscore.py | 73 +++++++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 12 deletions(-) create mode 100644 tests/test_survival_logscore.py diff --git a/ngboost/distns/exponential.py b/ngboost/distns/exponential.py index 8577422f..972df3e0 100644 --- a/ngboost/distns/exponential.py +++ b/ngboost/distns/exponential.py @@ -7,15 +7,11 @@ from ngboost.distns.distn import RegressionDistn from ngboost.scores import CRPScore, LogScore -eps = 1e-10 - class ExponentialLogScore(LogScore): def score(self, Y): E, T = Y["Event"], Y["Time"] - cens = (1 - E) * np.log(1 - self.dist.cdf(T) + eps) - uncens = E * self.dist.logpdf(T) - return -(cens + uncens) + return -np.where(E, self.dist.logpdf(T), self.dist.logsf(T)) def d_score(self, Y): E, T = Y["Event"], Y["Time"] diff --git a/ngboost/distns/lognormal.py b/ngboost/distns/lognormal.py index d0ec755a..585eb79f 100644 --- a/ngboost/distns/lognormal.py +++ b/ngboost/distns/lognormal.py @@ -12,9 +12,7 @@ class LogNormalLogScoreCensored(LogScore): def score(self, Y): E = Y["Event"] T = Y["Time"] - cens = (1 - E) * np.log(1 - self.dist.cdf(T) + self.eps) - uncens = E * self.dist.logpdf(T) - return -(cens + uncens) + return -np.where(E, self.dist.logpdf(T), self.dist.logsf(T)) def d_score(self, Y): E = Y["Event"][:, np.newaxis] @@ -27,10 +25,9 @@ def d_score(self, Y): D_uncens[:, 1] = 1 - ((self.loc - lT) ** 2) / (self.scale**2) D_cens = np.zeros((self.loc.shape[0], 2)) - D_cens[:, 0] = -sp.stats.norm.pdf(lT, loc=self.loc, scale=self.scale) / ( - 1 - self.dist.cdf(T) + self.eps - ) - D_cens[:, 1] = -Z * sp.stats.norm.pdf(Z) / (1 - self.dist.cdf(T) + self.eps) + hazard = np.exp(sp.stats.norm.logpdf(Z) - sp.stats.norm.logsf(Z)) + D_cens[:, 0] = -hazard / self.scale + D_cens[:, 1] = -Z * hazard return (1 - E) * D_cens + E * D_uncens diff --git a/tests/test_survival_logscore.py b/tests/test_survival_logscore.py new file mode 100644 index 00000000..86644573 --- /dev/null +++ b/tests/test_survival_logscore.py @@ -0,0 +1,73 @@ +"""Survival log scores must retain information in the right tail.""" + +import numpy as np +import pytest +import scipy.special +from scipy.stats import expon, lognorm + +from ngboost.distns import Exponential, LogNormal +from ngboost.distns.utils import SurvivalDistnClass +from ngboost.helpers import Y_from_censored +from ngboost.manifold import manifold +from ngboost.scores import LogScore + + +@pytest.mark.parametrize("dist", [Exponential, LogNormal]) +@pytest.mark.parametrize("event", [False, True]) +def test_survival_logscore_tails(dist, event): # pylint: disable=too-many-locals + if dist is Exponential: + params = np.log([[1.7, 1.7, 1.7, 1.7]]) + time = 1.7 * np.array([0.2, 5, 30, 1000]) + scipy_dist = expon(scale=1.7) + expected_grad = (event - time / 1.7)[:, None] + else: + loc, scale = 0.3, 1.7 + z = np.array([0.2, 5, 10, 40]) + time = np.exp(loc + scale * z) + params = np.repeat([[loc], [np.log(scale)]], len(z), axis=1) + scipy_dist = lognorm(s=scale, scale=np.exp(loc)) + # erfcx gives the normal inverse Mills ratio without subtracting a CDF. + mills = np.sqrt(2 / np.pi) / scipy.special.erfcx( # pylint: disable=no-member + z / np.sqrt(2) + ) + if event: + expected_grad = np.column_stack((-z / scale, 1 - z**2)) + else: + expected_grad = np.column_stack((-mills / scale, -z * mills)) + + y = Y_from_censored(time, np.full(len(time), event)) + score_dist = manifold(LogScore, SurvivalDistnClass(dist)) + obj = score_dist(params) + expected_score = -(scipy_dist.logpdf(time) if event else scipy_dist.logsf(time)) + np.testing.assert_allclose(obj.score(y), expected_score, rtol=1e-12) + np.testing.assert_allclose(obj.d_score(y), expected_grad, rtol=1e-11) + + for i in range(dist.n_params): + step = np.zeros_like(params) + step[i] = 1e-5 + numerical = ( + score_dist(params + step).score(y) - score_dist(params - step).score(y) + ) / 2e-5 + np.testing.assert_allclose( + obj.d_score(y)[:, i], numerical, rtol=1e-7, atol=1e-7 + ) + + +def test_survival_logscore_mixed_events(): + time = np.array([1.0, 2.0, 1000.0, 1000.0]) + events = np.array([True, False, True, False]) + y = Y_from_censored(time, events) + for dist, params in [ + (Exponential, np.zeros((1, 4))), + (LogNormal, np.zeros((2, 4))), + ]: + score_dist = manifold(LogScore, SurvivalDistnClass(dist)) + combined = score_dist(params) + for i in range(len(time)): + single = score_dist(params[:, i : i + 1]) + np.testing.assert_allclose( + combined.score(y)[i], single.score(y[i : i + 1])[0] + ) + np.testing.assert_allclose( + combined.d_score(y)[i], single.d_score(y[i : i + 1])[0] + ) From 4e19f3b92c16256ece49a17ff691b7cee6226df8 Mon Sep 17 00:00:00 2001 From: Rudy Celekli Date: Wed, 4 Nov 2026 10:56:37 -0500 Subject: [PATCH 2/2] Keep survival score regressions in the existing score test module --- tests/test_score.py | 63 ++++++++++++++++++++++++++++ tests/test_survival_logscore.py | 73 --------------------------------- 2 files changed, 63 insertions(+), 73 deletions(-) delete mode 100644 tests/test_survival_logscore.py diff --git a/tests/test_score.py b/tests/test_score.py index 956459af..2071ec1c 100644 --- a/tests/test_score.py +++ b/tests/test_score.py @@ -3,8 +3,10 @@ import numpy as np import pytest +import scipy.special from scipy.integrate import quad from scipy.optimize import approx_fprime +from scipy.stats import expon, lognorm from ngboost.distns import ( Beta, @@ -257,3 +259,64 @@ def test_dists_metric(dist_score_pair: DistScore, seed: int): manifold_test = manifold(LogScore, dist) FI_err = estimate_metric_err(params, manifold_test) assert FI_err < 1e-1 + + +@pytest.mark.parametrize("dist", [Exponential, LogNormal]) +@pytest.mark.parametrize("event", [False, True]) +def test_survival_logscore_tails(dist, event): # pylint: disable=too-many-locals + if dist is Exponential: + params = np.log([[1.7, 1.7, 1.7, 1.7]]) + time = 1.7 * np.array([0.2, 5, 30, 1000]) + scipy_dist = expon(scale=1.7) + expected_grad = (event - time / 1.7)[:, None] + else: + loc, scale = 0.3, 1.7 + z = np.array([0.2, 5, 10, 40]) + time = np.exp(loc + scale * z) + params = np.repeat([[loc], [np.log(scale)]], len(z), axis=1) + scipy_dist = lognorm(s=scale, scale=np.exp(loc)) + # erfcx gives the normal inverse Mills ratio without subtracting a CDF. + mills = np.sqrt(2 / np.pi) / scipy.special.erfcx( # pylint: disable=no-member + z / np.sqrt(2) + ) + if event: + expected_grad = np.column_stack((-z / scale, 1 - z**2)) + else: + expected_grad = np.column_stack((-mills / scale, -z * mills)) + + y = Y_from_censored(time, np.full(len(time), event)) + score_dist = manifold(LogScore, SurvivalDistnClass(dist)) + obj = score_dist(params) + expected_score = -(scipy_dist.logpdf(time) if event else scipy_dist.logsf(time)) + np.testing.assert_allclose(obj.score(y), expected_score, rtol=1e-12) + np.testing.assert_allclose(obj.d_score(y), expected_grad, rtol=1e-11) + + for i in range(dist.n_params): + step = np.zeros_like(params) + step[i] = 1e-5 + numerical = ( + score_dist(params + step).score(y) - score_dist(params - step).score(y) + ) / 2e-5 + np.testing.assert_allclose( + obj.d_score(y)[:, i], numerical, rtol=1e-7, atol=1e-7 + ) + + +def test_survival_logscore_mixed_events(): + time = np.array([1.0, 2.0, 1000.0, 1000.0]) + events = np.array([True, False, True, False]) + y = Y_from_censored(time, events) + for dist, params in [ + (Exponential, np.zeros((1, 4))), + (LogNormal, np.zeros((2, 4))), + ]: + score_dist = manifold(LogScore, SurvivalDistnClass(dist)) + combined = score_dist(params) + for i in range(len(time)): + single = score_dist(params[:, i : i + 1]) + np.testing.assert_allclose( + combined.score(y)[i], single.score(y[i : i + 1])[0] + ) + np.testing.assert_allclose( + combined.d_score(y)[i], single.d_score(y[i : i + 1])[0] + ) diff --git a/tests/test_survival_logscore.py b/tests/test_survival_logscore.py deleted file mode 100644 index 86644573..00000000 --- a/tests/test_survival_logscore.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Survival log scores must retain information in the right tail.""" - -import numpy as np -import pytest -import scipy.special -from scipy.stats import expon, lognorm - -from ngboost.distns import Exponential, LogNormal -from ngboost.distns.utils import SurvivalDistnClass -from ngboost.helpers import Y_from_censored -from ngboost.manifold import manifold -from ngboost.scores import LogScore - - -@pytest.mark.parametrize("dist", [Exponential, LogNormal]) -@pytest.mark.parametrize("event", [False, True]) -def test_survival_logscore_tails(dist, event): # pylint: disable=too-many-locals - if dist is Exponential: - params = np.log([[1.7, 1.7, 1.7, 1.7]]) - time = 1.7 * np.array([0.2, 5, 30, 1000]) - scipy_dist = expon(scale=1.7) - expected_grad = (event - time / 1.7)[:, None] - else: - loc, scale = 0.3, 1.7 - z = np.array([0.2, 5, 10, 40]) - time = np.exp(loc + scale * z) - params = np.repeat([[loc], [np.log(scale)]], len(z), axis=1) - scipy_dist = lognorm(s=scale, scale=np.exp(loc)) - # erfcx gives the normal inverse Mills ratio without subtracting a CDF. - mills = np.sqrt(2 / np.pi) / scipy.special.erfcx( # pylint: disable=no-member - z / np.sqrt(2) - ) - if event: - expected_grad = np.column_stack((-z / scale, 1 - z**2)) - else: - expected_grad = np.column_stack((-mills / scale, -z * mills)) - - y = Y_from_censored(time, np.full(len(time), event)) - score_dist = manifold(LogScore, SurvivalDistnClass(dist)) - obj = score_dist(params) - expected_score = -(scipy_dist.logpdf(time) if event else scipy_dist.logsf(time)) - np.testing.assert_allclose(obj.score(y), expected_score, rtol=1e-12) - np.testing.assert_allclose(obj.d_score(y), expected_grad, rtol=1e-11) - - for i in range(dist.n_params): - step = np.zeros_like(params) - step[i] = 1e-5 - numerical = ( - score_dist(params + step).score(y) - score_dist(params - step).score(y) - ) / 2e-5 - np.testing.assert_allclose( - obj.d_score(y)[:, i], numerical, rtol=1e-7, atol=1e-7 - ) - - -def test_survival_logscore_mixed_events(): - time = np.array([1.0, 2.0, 1000.0, 1000.0]) - events = np.array([True, False, True, False]) - y = Y_from_censored(time, events) - for dist, params in [ - (Exponential, np.zeros((1, 4))), - (LogNormal, np.zeros((2, 4))), - ]: - score_dist = manifold(LogScore, SurvivalDistnClass(dist)) - combined = score_dist(params) - for i in range(len(time)): - single = score_dist(params[:, i : i + 1]) - np.testing.assert_allclose( - combined.score(y)[i], single.score(y[i : i + 1])[0] - ) - np.testing.assert_allclose( - combined.d_score(y)[i], single.d_score(y[i : i + 1])[0] - )