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_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] + )