NB loss and poisson loss

Author: romankouz

spatial/models/monet_ae.py

@@ -98,13 +98,22 @@ def __init__(
 
     @staticmethod
     def calc_loss(pred, val, losstype, celltype_data=None, celltype=None):
-        # standard losses
         if losstype == "mse_against_log1pdata":
-            return torch.mean((pred - torch.log(1 + val)) ** 2)
+            return torch.mean((pred - torch.log1p(val)) ** 2)
         elif losstype == "mse":
             return torch.mean((pred - val) ** 2)
         elif losstype == "mae":
             return torch.mean(torch.abs(pred - val))
+        elif losstype == "poisson_nll":
+            mu = F.softplus(pred)  # ensure positive mean
+            dist = Poisson(mu)
+            return -dist.log_prob(val).mean()
+        elif losstype == "nb_nll":
+            mu = F.softplus(pred)  # mean
+            theta = torch.tensor(1.0, device=pred.device)  # dispersion
+            probs = mu / (mu + theta)
+            dist = NegativeBinomial(total_count=theta, probs=probs)
+            return -dist.log_prob(val).mean()
         else:
             raise NotImplementedError