Add vONTSS

octis/models/vONTSS_model/hyperspherical_vae/distributions/hyperspherical_uniform.py

@@ -0,0 +1,56 @@
+import math
+import torch
+
+
+class HypersphericalUniform(torch.distributions.Distribution):
+
+    support = torch.distributions.constraints.real
+    has_rsample = False
+    _mean_carrier_measure = 0
+
+    @property
+    def dim(self):
+        return self._dim
+
+    @property
+    def device(self):
+        return self._device
+
+    @device.setter
+    def device(self, val):
+        self._device = val if isinstance(val, torch.device) else torch.device(val)
+
+    def __init__(self, dim, validate_args=None, device="cpu"):
+        super(HypersphericalUniform, self).__init__(
+            torch.Size([dim]), validate_args=validate_args
+        )
+        self._dim = dim
+        self.device = device
+
+    def sample(self, shape=torch.Size()):
+        output = (
+            torch.distributions.Normal(0, 1)
+            .sample(
+                (shape if isinstance(shape, torch.Size) else torch.Size([shape]))
+                + torch.Size([self._dim + 1])
+            )
+            .to(self.device)
+        )
+
+        return output / output.norm(dim=-1, keepdim=True)
+
+    def entropy(self):
+        return self.__log_surface_area()
+
+    def log_prob(self, x):
+        return -torch.ones(x.shape[:-1], device=self.device) * self.__log_surface_area()
+
+    def __log_surface_area(self):
+        if torch.__version__ >= "1.0.0":
+            lgamma = torch.lgamma(torch.tensor([(self._dim + 1) / 2]).to(self.device))
+        else:
+            lgamma = torch.lgamma(
+                torch.Tensor([(self._dim + 1) / 2], device=self.device)
+            )
+        return math.log(2) + ((self._dim + 1) / 2) * math.log(math.pi) - lgamma
+