Initial commit

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Nicola De Cao
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,84 @@
+# The Power Spherical distribution
+
+## Overview
+This library contains a Pytorch implementation of the Power Spherical distribution, as presented in [[1]](#citation)(http://arxiv.org/abs/2006.TBA).
+
+## Dependencies
+
+* **python>=3.6**
+* **pytorch>=1.5**: https://pytorch.org
+
+*Notice that older version could work but they were not tested.*
+
+Optional dependency for [examples](https://github.com/nicola-decao/power_spherical/blob/master/example.ipynb) needed for plotting and numerical checks (again older version could work but they were not tested):
+* **numpy>=1.18.1**: https://numpy.org
+* **matplotlib>=3.1.1**: https://matplotlib.org
+* **quadpy>=0.14.11**: https://pypi.org/project/quadpy
+
+## Installation
+
+To install, run
+
+```bash
+$ python setup.py install
+```
+
+## Structure
+* [distributions](https://github.com/nicola-decao/power_spherical/blob/master/power_spherical/distributions.py): Pytorch implementation of the Power Spherical and hyperspherical Uniform distributions. Both inherit from `torch.distributions.Distribution`.
+* [examples](https://github.com/nicola-decao/power_spherical/blob/master/example.ipynb): Example code for using the library within a PyTorch project.
+
+## Usage
+Please have a look into the [examples](https://github.com/nicola-decao/power_spherical/blob/master/example.ipynb). We adapted our implementation to follow the structure of the [Pytorch probability distributions](https://pytorch.org/docs/stable/distributions.html).
+
+Here a minimal example that demonstrate differentiable sampling:
+```python
+>>> from power_spherical import PowerSpherical
+>>> p = PowerSpherical(
+      loc=torch.tensor([0., 1.], requires_grad=True),
+      scale=torch.tensor(4., requires_grad=True),
+    )
+>>> p.rsample()
+
+tensor([-0.1786,  0.9839], grad_fn=<SubBackward0>)
+```
+and computing KL divergence with the uniform distribution:
+```python
+>>> from power_spherical import HypersphericalUniform
+>>> q = HypersphericalUniform(dim=2)
+>>> torch.distributions.kl_divergence(p, q)
+
+tensor(1.2486, grad_fn=<AddBackward0>)
+```
+
+Examples of 2D and 3D plots are show in [examples](https://github.com/nicola-decao/power_spherical/blob/master/example.ipynb) and will generate something similar to these figures below.
+<p align="center">
+  <img class="paper_logo" src="https://i.imgur.com/4iITHS5.png" width=40%>
+  <img class="paper_logo" src="https://i.imgur.com/zXZWr9H.png" width=40%>
+</p>
+
+Please cite [[1](#citation)] in your work when using this library in your experiments.
+
+## Feedback
+For questions and comments, feel free to contact [Nicola De Cao](mailto:[email protected]).
+
+## License
+MIT
+
+## Citation
+```
+[1] De Cao, N., Aziz, W. (2020). 
+The Power Spherical distrbution.
+arXiv preprint arXiv:2006.TBA.
+```
+
+BibTeX format:
+```
+@article{decao2020power,
+  title={The Power Spherical distrbution},
+  author={
+    De Cao, Nicola and
+    Aziz, Wilker},
+  journal={arXiv preprint arXiv:2006.TBA},
+  year={2020}
+}
+```

power_spherical/__init__.py

@@ -0,0 +1,3 @@
+from .distributions import HypersphericalUniform
+from .distributions import PowerSpherical
+from .distributions import MarginalTDistribution

power_spherical/distributions.py

@@ -0,0 +1,214 @@
+import math
+import torch
+from torch.distributions.kl import register_kl
+
+
+_EPS = 1e-7
+
+
+class _TTransform(torch.distributions.Transform):
+    def _call(self, x):
+        t = x[..., 0].unsqueeze(-1)
+        v = x[..., 1:]
+        return torch.cat((t, v * torch.sqrt(torch.clamp(1 - t ** 2, _EPS))), -1)
+
+    def _inverse(self, y):
+        t = y[..., 0].unsqueeze(-1)
+        v = y[..., 1:]
+        return torch.cat((t, v / torch.sqrt(torch.clamp(1 - t ** 2, _EPS))), -1)
+
+    def log_abs_det_jacobian(self, x, y):
+        t = x[..., 0]
+        return ((x.shape[-1] - 3) / 2) * torch.log(torch.clamp(1 - t ** 2, _EPS))
+
+
+class _HouseholderRotationTransform(torch.distributions.Transform):
+    def __init__(self, loc):
+        super().__init__()
+        self.loc = loc
+        self.e1 = torch.zeros_like(self.loc)
+        self.e1[..., 0] = 1
+
+    def _call(self, x):
+        u = self.e1 - self.loc
+        u = u / (u.norm(dim=-1, keepdim=True) + _EPS)
+        return x - 2 * (x * u).sum(-1, keepdim=True) * u
+
+    def _inverse(self, y):
+        u = self.e1 - self.loc
+        u = u / (u.norm(dim=-1, keepdim=True) + _EPS)
+        return y - 2 * (y * u).sum(-1, keepdim=True) * u
+
+    def log_abs_det_jacobian(self, x, y):
+        return 0
+
+
+class HypersphericalUniform(torch.distributions.Distribution):
+
+    arg_constraints = {
+        "dim": torch.distributions.constraints.positive_integer,
+    }
+
+    def __init__(self, dim, device="cpu", dtype=torch.float32, validate_args=None):
+        super().__init__(validate_args=validate_args)
+        self.dim, self.device, self.dtype = dim, device, dtype
+
+    def rsample(self, sample_shape=()):
+        v = torch.empty(
+            sample_shape + (self.dim,), device=self.device, dtype=self.dtype
+        ).normal_()
+        return v / (v.norm(dim=-1, keepdim=True) + _EPS)
+
+    def log_prob(self, value):
+        return torch.full_like(
+            value[..., 0],
+            math.lgamma(self.dim / 2)
+            - (math.log(2) + (self.dim / 2) * math.log(math.pi)),
+            device=self.device,
+            dtype=self.dtype,
+        )
+
+    def entropy(self):
+        return -self.log_prob(torch.empty(1))
+
+    def __repr__(self):
+        return "HypersphericalUniform(dim={}, device={}, dtype={})".format(
+            self.dim, self.device, self.dtype
+        )
+
+
+class MarginalTDistribution(torch.distributions.TransformedDistribution):
+
+    arg_constraints = {
+        "dim": torch.distributions.constraints.positive_integer,
+        "scale": torch.distributions.constraints.positive,
+    }
+
+    has_rsample = True
+
+    def __init__(self, dim, scale, validate_args=None):
+        super().__init__(
+            torch.distributions.Beta(
+                (dim - 1) / 2 + scale, (dim - 1) / 2, validate_args=validate_args
+            ),
+            transforms=torch.distributions.AffineTransform(loc=-1, scale=2),
+        )
+        self.dim, self.scale = dim, scale
+
+    def entropy(self):
+        return self.base_dist.entropy() + math.log(2)
+
+    @property
+    def mean(self):
+        return 2 * self.base_dist.mean - 1
+
+    @property
+    def stddev(self):
+        return self.variance.sqrt()
+
+    @property
+    def variance(self):
+        return 4 * self.base_dist.variance
+
+
+class _JointTSDistribution(torch.distributions.Distribution):
+    def __init__(self, marginal_t, marginal_s):
+        super().__init__()
+        self.marginal_t, self.marginal_s = marginal_t, marginal_s
+
+    def rsample(self, sample_shape=()):
+        return torch.cat(
+            (
+                self.marginal_t.rsample(sample_shape).unsqueeze(-1),
+                self.marginal_s.rsample(sample_shape + self.marginal_t.scale.shape),
+            ),
+            -1,
+        )
+
+    def log_prob(self, value):
+        return self.marginal_t.log_prob(value[..., 0]) + self.marginal_s.log_prob(
+            value[..., 1:]
+        )
+
+    def entropy(self):
+        return self.marginal_t.entropy() + self.marginal_s.entropy()
+
+
+class PowerSpherical(torch.distributions.TransformedDistribution):
+
+    arg_constraints = {
+        "loc": torch.distributions.constraints.real,
+        "scale": torch.distributions.constraints.positive,
+    }
+
+    has_rsample = True
+
+    def __init__(self, loc, scale, validate_args=None):
+
+        if isinstance(scale, torch.Tensor) and len(scale.shape) > 1:
+            assert (
+                scale.shape[-1] != 1
+            ), "`scale' cannot have 1 as last dimention when `isinstance(scale, torch.Tensor)' and `len(scale.shape) > 1'."
+
+        super().__init__(
+            _JointTSDistribution(
+                MarginalTDistribution(
+                    loc.shape[-1], scale, validate_args=validate_args
+                ),
+                HypersphericalUniform(
+                    loc.shape[-1] - 1,
+                    device=loc.device,
+                    dtype=loc.dtype,
+                    validate_args=validate_args,
+                ),
+            ),
+            [_TTransform(), _HouseholderRotationTransform(loc),],
+        )
+        self.loc, self.scale, = loc, scale
+
+    def log_prob(self, value):
+        return self.log_normalizer() + self.scale * torch.log1p(
+            (self.loc * value).sum(-1)
+        )
+
+    def log_normalizer(self):
+        alpha = self.base_dist.marginal_t.base_dist.concentration1
+        beta = self.base_dist.marginal_t.base_dist.concentration0
+        return -(
+            (alpha + beta) * math.log(2)
+            + torch.lgamma(alpha)
+            - torch.lgamma(alpha + beta)
+            + beta * math.log(math.pi)
+        )
+
+    def entropy(self):
+        alpha = self.base_dist.marginal_t.base_dist.concentration1
+        beta = self.base_dist.marginal_t.base_dist.concentration0
+        return -(
+            self.log_normalizer()
+            + self.scale
+            * (math.log(2) + torch.digamma(alpha) - torch.digamma(alpha + beta))
+        )
+
+    @property
+    def mean(self):
+        return self.loc * self.base_dist.marginal_t.mean
+
+    @property
+    def stddev(self):
+        return self.variance.sqrt()
+
+    @property
+    def variance(self):
+        alpha = self.base_dist.marginal_t.base_dist.concentration1
+        beta = self.base_dist.marginal_t.base_dist.concentration0
+        ratio = (alpha + beta) / (2 * beta)
+        return self.base_dist.marginal_t.variance * (
+            (1 - ratio) * self.loc.unsqueeze(-1) @ self.loc.unsqueeze(-2)
+            + ratio * torch.eye(self.loc.shape[-1])
+        )
+
+
+@register_kl(PowerSpherical, HypersphericalUniform)
+def _kl_powerspherical_uniform(p, q):
+    return -p.entropy() + q.entropy()

setup.py

@@ -0,0 +1,18 @@
+import os
+from setuptools import setup
+from setuptools import find_packages
+
+setup(
+    name="power_spherical",
+    version="0.1.0",
+    author="Nicola De Cao",
+    author_email="[email protected]",
+    description="Pytorch implementation of the Power Spherical distribution",
+    license="MIT",
+    keywords="pytorch machine-learning deep-learning manifold-learning",
+    url="https://nicola-decao.github.io",
+    download_url="https://github.com/nicola-decao/power_spherical",
+    long_description=open(os.path.join(os.path.dirname(__file__), "README.md")).read(),
+    install_requires=["torch>=1.5.0"],
+    packages=find_packages(),
+)