From 0e8e260a25314e3ff1d4e107fa5ccaaf5181cba5 Mon Sep 17 00:00:00 2001
From: Rudy Celekli <rudy@gradiahq.com>
Date: Wed, 4 Nov 2026 10:41:53 -0500
Subject: [PATCH] Keep Gamma log-score gradients exact for small positive
 targets

---
 ngboost/distns/gamma.py      |  3 +--
 tests/test_gamma_logscore.py | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_gamma_logscore.py

diff --git a/ngboost/distns/gamma.py b/ngboost/distns/gamma.py
index c2ea71a3..7d3c4a23 100644
--- a/ngboost/distns/gamma.py
+++ b/ngboost/distns/gamma.py
@@ -16,7 +16,7 @@ def d_score(self, Y):
         D = np.zeros((len(Y), 2))
         # d(-log(PDF))/dalpha
         D[:, 0] = self.alpha * (
-            sp.special.digamma(self.alpha) - np.log(self.eps + self.beta * Y)
+            sp.special.digamma(self.alpha) - np.log(self.beta) - np.log(Y)
         )
         # d(-log(PDF))/dbeta
         D[:, 1] = (self.beta * Y) - self.alpha
@@ -42,7 +42,6 @@ def __init__(self, params):
         self.dist = dist(
             a=self.alpha, loc=np.zeros_like(self.alpha), scale=1 / self.beta
         )
-        self.eps = 1e-10
 
     def fit(Y):
         a, _, scale = dist.fit(Y, floc=0)
diff --git a/tests/test_gamma_logscore.py b/tests/test_gamma_logscore.py
new file mode 100644
index 00000000..70ff4b5e
--- /dev/null
+++ b/tests/test_gamma_logscore.py
@@ -0,0 +1,34 @@
+"""The Gamma score gradient must remain exact on its positive support."""
+
+import numpy as np
+import pytest
+from scipy.stats import gamma
+
+from ngboost.distns import Gamma
+from ngboost.manifold import manifold
+from ngboost.scores import LogScore
+
+
+@pytest.mark.parametrize("alpha", [0.5, 2.0])
+def test_gamma_logscore_positive_tail_gradient(alpha):
+    beta = np.array([1.0, 1.0, 1e-100, 1e100, 0.5, 2.0])
+    y = np.array([1e-20, 1e-200, 1e-100, 1e-100, 0.8, 3.0])
+    params = np.array([np.full(len(y), np.log(alpha)), np.log(beta)])
+    score_dist = manifold(LogScore, Gamma)
+    obj = score_dist(params)
+    # Closed-form digamma values avoid using the implementation's special function.
+    psi = 1 - np.euler_gamma if alpha == 2 else -np.euler_gamma - 2 * np.log(2)
+    expected = np.column_stack(
+        (alpha * (psi - np.log(beta) - np.log(y)), beta * y - alpha)
+    )
+    np.testing.assert_allclose(obj.score(y), -gamma.logpdf(y, alpha, scale=1 / beta))
+    np.testing.assert_allclose(obj.d_score(y), expected, rtol=1e-12, atol=1e-12)
+    for i in range(2):
+        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
+        )