added negative binomial and beta output

pts/dataset/artificial.py

@@ -67,4 +67,4 @@ def constant_dataset() -> Tuple[DatasetInfo, Dataset, Dataset]:
         test_statistics=calculate_dataset_statistics(test_ds),
     )
 
-    return info, train_ds, test_ds
+    return info, train_ds, test_ds

pts/modules/__init__.py

@@ -1,4 +1,4 @@
-from .distribution_output import ArgProj, Output, DistributionOutput, StudentTOutput
+from .distribution_output import ArgProj, Output, DistributionOutput, StudentTOutput, BetaOutput, NegativeBinomialOutput
 from .lambda_layer import LambdaLayer
 from .feature import FeatureEmbedder, FeatureAssembler
 from .scaler import MeanScaler, NOPScaler

pts/modules/distribution_output.py

@@ -5,7 +5,14 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.distributions import Distribution, StudentT, TransformedDistribution, AffineTransform
+from torch.distributions import (
+    Distribution,
+    Beta,
+    NegativeBinomial,
+    StudentT,
+    TransformedDistribution,
+    AffineTransform,
+)
 
 from .lambda_layer import LambdaLayer
 
@@ -75,6 +82,50 @@ class DistributionOutput(Output):
             return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)])
 
 
+class BetaOutput(DistributionOutput):
+    args_dim: Dict[str, int] = {"concentration1": 1, "concentration0": 1}
+    distr_cls: type = Beta
+
+    @classmethod
+    def domain_map(cls, concentration1, concentration0):
+        concentration1 = F.softplus(concentration1) + 1e-8
+        concentration0 = F.softplus(concentration0) + 1e-8
+        return concentration1.squeeze(-1), concentration0.squeeze(-1)
+    
+    @property
+    def event_shape(self) -> Tuple:
+        return ()
+
+
+class NegativeBinomialOutput(DistributionOutput):
+    args_dim: Dict[str, int] = {"mu": 1, "alpha": 1}
+    distr_cls: Distribution = NegativeBinomial
+
+    @classmethod
+    def domain_map(cls, mu, alpha):
+        mu = F.softplus(mu) + 1e-8
+        alpha = F.softplus(alpha) + 1e-8
+        return mu.squeeze(-1), alpha.squeeze(-1)
+
+    def distribution(
+        self, distr_args, scale: Optional[torch.Tensor] = None
+    ) -> Distribution:
+        mu, alpha = distr_args
+
+        if scale is not None:
+            mu *= scale
+            alpha *= torch.sqrt(scale + 1.0)
+
+        n = 1.0 / alpha
+        p = mu * alpha / (1.0 + mu * alpha)
+
+        return NegativeBinomial(total_count=n, probs=p)
+    
+    @property
+    def event_shape(self) -> Tuple:
+        return ()
+
+
 class StudentTOutput(DistributionOutput):
     args_dim: Dict[str, int] = {"df": 1, "loc": 1, "scale": 1}
     distr_cls: type = StudentT
@@ -87,4 +138,4 @@ class StudentTOutput(DistributionOutput):
 
     @property
     def event_shape(self) -> Tuple:
-        return ()
+        return ()