clipping for poisson

CDonnerer · Mar 11, 2023 · d0aa713 · d0aa713
1 parent 07c411c
commit d0aa713
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/xgboost_distribution/distributions/poisson.py b/src/xgboost_distribution/distributions/poisson.py
@@ -11,6 +11,10 @@
     check_all_integer,
 )
 
+MIN_LOG_MU = np.log(np.finfo("float32").tiny) + 1
+MAX_LOG_MU = np.log(np.finfo("float32").max) - 1
+
+
 Params = namedtuple("Params", ("mu"))
 
 
@@ -66,7 +70,7 @@ def loss(self, y, params):
         return "Poisson-NLL", -poisson.logpmf(y, mu=mu)
 
     def predict(self, params):
-        log_mu = params  # params are shape (n,)
+        log_mu = np.clip(params, a_min=MIN_LOG_MU, a_max=MAX_LOG_MU)
         mu = np.exp(log_mu)
         return Params(mu=mu)