Sine Skewed Toridial distribution

docs/source/distributions.rst

@@ -316,6 +316,13 @@ Rejector
     :undoc-members:
     :show-inheritance:
 
+SineSkewed
+----------
+.. autoclass:: pyro.distributions.SineSkewed
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 SoftLaplace
 -------------
 .. autoclass:: pyro.distributions.SoftLaplace

pyro/distributions/__init__.py

@@ -58,9 +58,11 @@
     RelaxedBernoulliStraightThrough,
     RelaxedOneHotCategoricalStraightThrough,
 )
+from pyro.distributions.sine_skewed import SineSkewed
 from pyro.distributions.softlaplace import SoftLaplace
 from pyro.distributions.spanning_tree import SpanningTree
 from pyro.distributions.stable import Stable
+from pyro.distributions.torch import *  # noqa F403
 from pyro.distributions.torch import __all__ as torch_dists
 from pyro.distributions.torch_distribution import (
     ExpandedDistribution,
@@ -128,6 +130,7 @@
     "Rejector",
     "RelaxedBernoulliStraightThrough",
     "RelaxedOneHotCategoricalStraightThrough",
+    "SineSkewed",
     "SoftLaplace",
     "SpanningTree",
     "Stable",

pyro/distributions/sine_skewed.py

@@ -0,0 +1,96 @@
+import warnings
+from math import pi
+
+import torch
+from torch import broadcast_shapes
+from torch.distributions import Uniform
+
+from pyro.distributions import constraints
+
+from .torch_distribution import TorchDistribution
+
+
+class SineSkewed(TorchDistribution):
+    """The Sine Skewed distribution [1] is a distribution for breaking pointwise-symmetry on a base-distribution over
+    the d-dimensional torus defined as ⨂^d S^1 where S^1 is the circle. So for example the 0-torus is a point, the
+    1-torus is a circle and the 2-tours is commonly associated with the donut shape (some may object to this simile).
+
+    The skewness parameter can be inferred using :class:`~pyro.infer.HMC` or :class:`~pyro.infer.NUTS`.
+    For example, the following will produce a uniform prior over skewness for the 2-torus,::
+
+        def model(...):
+            ...
+            skew_phi = pyro.sample(f'skew_phi', Uniform(-1., 1.))
+            psi_bound = 1 - skewness_phi.abs()
+            skew_psi = pyro.sample(f'skew_psi', Uniform(-1, 1.))
+            skewness = torch.stack((skew_phi, psi_bound * skew_psi), dim=0)
+            ...
+
+    In the context of :class:`~pyro.infer.SVI`, this distribution can be freely used as a likelihood, but use as a
+    latent variables will lead to slow inference for 2 and higher order toruses. This is because the base_dist
+    cannot be reparameterized.
+
+    .. note:: An event in the base distribution must be on a d-torus, so the event_shape must be (d,).
+
+    .. note:: For the skewness parameter, it must hold that the sum of the absolute value of its weights for an event
+        must be less than or equal to one. See eq. 2.1 in [1].
+
+    ** References: **
+      1. Sine-skewed toroidal distributions and their application in protein bioinformatics
+         Ameijeiras-Alonso, J., Ley, C. (2019)
+
+    :param base_dist: base density on a d-dimensional torus.
+    :param skewness: skewness of the distribution.
+    """
+    arg_constraints = {'skewness': constraints.independent(constraints.interval(-1., 1.), 1)}
+
+    support = constraints.independent(constraints.real, 1)
+
+    def __init__(self, base_dist: TorchDistribution, skewness, validate_args=None):
+        if (skewness.abs().sum(-1) > 1.).any():
+            warnings.warn("Total skewness weight shouldn't exceed one.", UserWarning)
+
+        batch_shape = broadcast_shapes(base_dist.batch_shape, skewness.shape[:-1])
+        event_shape = skewness.shape[-1:]
+        self.skewness = skewness.broadcast_to(batch_shape + event_shape)
+        self.base_dist = base_dist.expand(batch_shape)
+        super().__init__(batch_shape, event_shape, validate_args=validate_args)
+
+        if self._validate_args and base_dist.mean.device != skewness.device:
+            raise ValueError(f"base_density: {base_dist.__class__.__name__} and SineSkewed "
+                             f"must be on same device.")
+
+    def __repr__(self):
+        args_string = ', '.join(['{}: {}'.format(p, getattr(self, p)
+                                if getattr(self, p).numel() == 1
+                                else getattr(self, p).size()) for p in self.arg_constraints.keys()])
+        return self.__class__.__name__ + '(' + f'base_density: {str(self.base_dist)}, ' + args_string + ')'
+
+    def sample(self, sample_shape=torch.Size()):
+        bd = self.base_dist
+        ys = bd.sample(sample_shape)
+        u = Uniform(0., torch.ones(torch.Size([]), device=self.skewness.device)).sample(sample_shape + self.batch_shape)
+
+        # Section 2.3 step 3 in [1]
+        mask = u <= .5 + .5 * (self.skewness * torch.sin((ys - bd.mean) % (2 * pi))).sum(-1)
+        mask = mask[..., None]
+        samples = (torch.where(mask, ys, -ys + 2 * bd.mean) + pi) % (2 * pi) - pi
+        return samples
+
+    def log_prob(self, value):
+        if self._validate_args:
+            self._validate_sample(value)
+
+        # Eq. 2.1 in [1]
+        skew_prob = torch.log(1 + (self.skewness * torch.sin((value - self.base_dist.mean) % (2 * pi))).sum(-1))
+        return self.base_dist.log_prob(value) + skew_prob
+
+    def expand(self, batch_shape, _instance=None):
+        batch_shape = torch.Size(batch_shape)
+        new = self._get_checked_instance(SineSkewed, _instance)
+        base_dist = self.base_dist.expand(batch_shape)
+        new.base_dist = base_dist
+        new.skewness = self.skewness.expand(batch_shape + (-1,))
+        super(SineSkewed, new).__init__(batch_shape, self.event_shape, validate_args=False)
+        new._validate_args = self._validate_args
+        return new

tests/distributions/conftest.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import math
+from math import pi
 
 import numpy as np
 import pytest
@@ -15,7 +16,7 @@
     ShapeAugmentedDirichlet,
     ShapeAugmentedGamma,
 )
-from tests.distributions.dist_fixture import Fixture
+from tests.distributions.dist_fixture import Fixture, tensor_wrap
 
 
 class FoldedNormal(dist.FoldedDistribution):
@@ -187,7 +188,7 @@ def __init__(self, rate, *, validate_args=None):
             ],
             # This hack seems to be the best option right now, as 'scale' is not handled well by get_scipy_batch_logpdf
             scipy_arg_fn=lambda loc, covariance_matrix=None:
-                ((), {"mean": np.array(loc), "cov": np.array([[1.0, 0.5], [0.5, 1.0]])}),
+            ((), {"mean": np.array(loc), "cov": np.array([[1.0, 0.5], [0.5, 1.0]])}),
             prec=0.01,
             min_samples=500000),
     Fixture(pyro_dist=dist.LowRankMultivariateNormal,
@@ -197,7 +198,7 @@ def __init__(self, rate, *, validate_args=None):
                  'test_data': [[2.0, 1.0], [9.0, 3.4]]},
             ],
             scipy_arg_fn=lambda loc, cov_diag=None, cov_factor=None:
-                ((), {"mean": np.array(loc), "cov": np.array([[1.5, 0.5], [0.5, 0.75]])}),
+            ((), {"mean": np.array(loc), "cov": np.array([[1.5, 0.5], [0.5, 0.75]])}),
             prec=0.01,
             min_samples=500000),
     Fixture(pyro_dist=FoldedNormal,
@@ -280,12 +281,12 @@ def __init__(self, rate, *, validate_args=None):
     Fixture(pyro_dist=dist.LKJ,
             examples=[
                 {'dim': 3, 'concentration': 1., 'test_data':
-                    [[[1.0000, -0.8221,  0.7655], [-0.8221,  1.0000, -0.5293], [0.7655,  -0.5293,  1.0000]],
-                     [[1.0000, -0.5345, -0.5459], [-0.5345,  1.0000, -0.0333], [-0.5459, -0.0333,  1.0000]],
-                     [[1.0000, -0.3758, -0.2409], [-0.3758,  1.0000,  0.4653], [-0.2409,  0.4653,  1.0000]],
-                     [[1.0000, -0.8800, -0.9493], [-0.8800,  1.0000,  0.9088], [-0.9493,  0.9088,  1.0000]],
-                     [[1.0000,  0.2284, -0.1283], [0.2284,   1.0000,  0.0146], [-0.1283,  0.0146,  1.0000]]]},
-                ]),
+                    [[[1.0000, -0.8221, 0.7655], [-0.8221, 1.0000, -0.5293], [0.7655, -0.5293, 1.0000]],
+                     [[1.0000, -0.5345, -0.5459], [-0.5345, 1.0000, -0.0333], [-0.5459, -0.0333, 1.0000]],
+                     [[1.0000, -0.3758, -0.2409], [-0.3758, 1.0000, 0.4653], [-0.2409, 0.4653, 1.0000]],
+                     [[1.0000, -0.8800, -0.9493], [-0.8800, 1.0000, 0.9088], [-0.9493, 0.9088, 1.0000]],
+                     [[1.0000, 0.2284, -0.1283], [0.2284, 1.0000, 0.0146], [-0.1283, 0.0146, 1.0000]]]},
+            ]),
     Fixture(pyro_dist=dist.LKJCholesky,
             examples=[
                 {
@@ -305,19 +306,19 @@ def __init__(self, rate, *, validate_args=None):
             examples=[
                 {'stability': [1.5], 'skew': 0.1, 'test_data': [-10.]},
                 {'stability': [1.5], 'skew': 0.1, 'scale': 2.0, 'loc': -2.0, 'test_data': [10.]},
-                ]),
+            ]),
     Fixture(pyro_dist=dist.MultivariateStudentT,
             examples=[
                 {'df': 1.5, 'loc': [0.2, 0.3], 'scale_tril': [[0.8, 0.0], [1.3, 0.4]],
                  'test_data': [-3., 2]},
-                ]),
+            ]),
     Fixture(pyro_dist=dist.ProjectedNormal,
             examples=[
                 {'concentration': [0., 0.], 'test_data': [1., 0.]},
                 {'concentration': [2., 3.], 'test_data': [0., 1.]},
                 {'concentration': [0., 0., 0.], 'test_data': [1., 0., 0.]},
                 {'concentration': [-1., 2., 3.], 'test_data': [0., 0., 1.]},
-                ]),
+            ]),
     Fixture(pyro_dist=dist.SoftLaplace,
             examples=[
                 {'loc': [2.0], 'scale': [4.0],
@@ -328,7 +329,16 @@ def __init__(self, rate, *, validate_args=None):
                  'test_data': [[[2.0]]]},
                 {'loc': [2.0, 50.0], 'scale': [4.0, 100.0],
                  'test_data': [[2.0, 50.0], [2.0, 50.0]]},
-                ]),
+            ]),
+    Fixture(pyro_dist=dist.SineSkewed,
+            examples=[
+                {'base_dist': dist.VonMises(*tensor_wrap([0.], [1.])).to_event(1),
+                 'skewness': [.342355], 'test_data': [.1]},
+                {'base_dist': dist.Uniform(*tensor_wrap([-pi, -pi], [pi, pi])).to_event(1),
+                 'skewness': [-pi / 4, .1], 'test_data': [pi / 2, -2 * pi / 3]},
+                {'base_dist': dist.VonMises(*tensor_wrap([0., -1.234], [1., 10.])).to_event(1),
+                 'skewness': [[.342355, -.0001], [.91, 0.09]], 'test_data': [[.1, -3.2], [-2., 0.]]},
+            ])
 ]
 
 discrete_dists = [

tests/distributions/test_cuda.py

@@ -15,6 +15,8 @@
 
 @requires_cuda
 def test_sample(dist):
+    if dist.pyro_dist.__name__ == 'SineSkewed':
+        pytest.xfail(reason="Fixture with distribution param not handled.")
     for idx in range(len(dist.dist_params)):
 
         # Compute CPU value.
@@ -77,6 +79,8 @@ def test_rsample(dist):
 
 @requires_cuda
 def test_log_prob(dist):
+    if dist.pyro_dist.__name__ == 'SineSkewed':
+        pytest.xfail(reason="Fixture with distribution param not handled.")
     for idx in range(len(dist.dist_params)):
 
         # Compute CPU value.

tests/distributions/test_distributions.py

@@ -40,7 +40,7 @@ def test_support_shape(dist):
 
 
 def test_infer_shapes(dist):
-    if "LKJ" in dist.pyro_dist.__name__:
+    if "LKJ" in dist.pyro_dist.__name__ or "SineSkewed" == dist.pyro_dist.__name__:
         pytest.xfail(reason="cannot statically compute shape")
     for idx in range(dist.get_num_test_data()):
         dist_params = dist.get_dist_params(idx)

tests/distributions/test_sine_skewed.py

@@ -0,0 +1,79 @@
+from math import pi
+
+import pytest
+import torch
+
+import pyro
+from pyro.distributions import Normal, SineSkewed, Uniform, VonMises, constraints
+from pyro.infer import SVI, Trace_ELBO
+from pyro.optim import Adam
+from tests.common import assert_equal
+
+BASE_DISTS = [(Uniform, [-pi, pi]), (VonMises, (0., 1.))]
+
+
+def _skewness(event_shape):
+    skewness = torch.zeros(event_shape.numel())
+    done = False
+    while not done:
+        for i in range(event_shape.numel()):
+            max_ = 1. - skewness.abs().sum(-1)
+            if torch.any(max_ < 1e-15):
+                break
+            skewness[i] = Uniform(-max_, max_).sample()
+        done = not torch.any(max_ < 1e-15)
+
+    if event_shape == tuple():
+        skewness = skewness.reshape(event_shape)
+    else:
+        skewness = skewness.view(event_shape)
+    return skewness
+
+
+@pytest.mark.parametrize('expand_shape',
+                         [(1,), (2,), (4,), (1, 1), (1, 2), (10, 10), (1, 3, 1), (10, 1, 5), (1, 1, 1), (3, 2, 3)])
+@pytest.mark.parametrize('dist', BASE_DISTS)
+def test_ss_multidim_log_prob(expand_shape, dist):
+    base_dist = dist[0](*(torch.tensor(param).expand(expand_shape) for param in dist[1])).to_event(1)
+
+    loc = base_dist.sample((10,)) + Normal(0., 1e-3).sample()
+
+    base_prob = base_dist.log_prob(loc)
+    skewness = _skewness(base_dist.event_shape)
+
+    ss = SineSkewed(base_dist, skewness)
+    assert_equal(base_prob.shape, ss.log_prob(loc).shape)
+    assert_equal(ss.sample().shape, torch.Size(expand_shape))
+
+
+@pytest.mark.parametrize('dist', BASE_DISTS)
+@pytest.mark.parametrize('dim', [1, 2])
+def test_ss_mle(dim, dist):
+    base_dist = dist[0](*(torch.tensor(param).expand((dim,)) for param in dist[1])).to_event(1)
+
+    skewness_tar = _skewness(base_dist.event_shape)
+    data = SineSkewed(base_dist, skewness_tar).sample((1000,))
+
+    def model(data, batch_shape):
+        skews = []
+        for i in range(dim):
+            skews.append(pyro.param(f'skew{i}', .5 * torch.ones(batch_shape), constraint=constraints.interval(-1, 1)))
+
+        skewness = torch.stack(skews, dim=-1)
+        with pyro.plate("data", data.size(-len(data.size()))):
+            pyro.sample('obs', SineSkewed(base_dist, skewness), obs=data)
+
+    def guide(data, batch_shape):
+        pass
+
+    pyro.clear_param_store()
+    adam = Adam({"lr": .1})
+    svi = SVI(model, guide, adam, loss=Trace_ELBO())
+
+    losses = []
+    steps = 80
+    for step in range(steps):
+        losses.append(svi.step(data, base_dist.batch_shape))
+
+    act_skewness = torch.stack([v for k, v in pyro.get_param_store().items() if 'skew' in k], dim=-1)
+    assert_equal(act_skewness, skewness_tar, 1e-1)