Sine Skewed Toridial distribution

[OlaRonning]

This PR introduces the Sine Skewed Toridial distribution described by Jose Ameijeiras-Alonso and Christophe Ley.

The distribution enables skewing a base distribution on a d-dimensional torus, which is useful, for example, with dihedral angles (1-torus) of peptides as can be seen on a Ramachandran plot.

Missing prior suggestion in docstring; inference method and inferable params as suggested in #2821.

[OlaRonning]

@fehiepsi, thanks for the review! 👍

[fehiepsi]

Could you clarify the input shapes for base density and skewness? It is not clear to me what's the requirements so I couldn't check if the expand method and the reshape ops in sample and log_prob are valid.

[OlaRonning]

Hi @fehiepsi,

Thanks for the review! Sorry for the input dimensions being vague. Let me clarify; the requirements are base_dist.shape() == skewness.shape skewness.shape broadcastable to base_dist.shape() and base_dist.event_shape[-1] == 2. So,

ss = SineSkewed(base_dist=BaseDist(Size([3,3,2])).to_event(1), skewness=Size([3,3,2]))
assert ss.event_shape == (2,)
assert ss.batch_shape == (3,3)

ss1 = SineSkewed(base_dist=BaseDist(Size([3,3,2])).to_event(1), skewness=Size([3,3,2])).to_event(1)
assert ss1.event_shape == (3,2)
assert ss1.batch_shape == (3,)

ss2 = SineSkewed(base_dist=BaseDist(Size([3,3,2])).to_event(2), skewness=Size([3,3,2]))
assert ss2.event_shape == (3,2)
assert ss2.batch_shape == (3,)

try:
  SineSkewed(base_dist=BaseDist(Size([3,3,2])).to_event(2), skewness=Size([3,3,2])).to_event(1)
  assert 1==0
except AssertionError:
  pass

try:
  SineSkewed(base_dist=BaseDist(Size([3,3,2])).to_event(1), skewness=Size([3,3,2])).to_event(2)  
  assert 1==0
except AssertionError:
  pass

I would like for the construction ss=SineSkewed(base_dist=BaseDist(Size([3,3,2])), skewness=Size([3,3,2])).to_event(1) also to work; however, I cannot figure out how to ensure the to_event has been called on the ss object before it's used. Edit1: typo. Edit2: changed to reflect restriction to 0<event_dim<=2.

[fehiepsi]

Looks much cleaner than before. 👍

[fritzo]

LGTM after a couple tiny nits

[fehiepsi]

It would be nice to have a constraint+transform for this (in a follow-up PR). I believe we can use signed stick-breaking transform here. This way users can define distributions over skewness (or just simply use pyro.param with correct constraint). Without that, it is a bit cumbersome for users to define correct skewness over general d-torus.

[fehiepsi]

With that, we can have correct constraints in the distribution definition. :D

[OlaRonning]

That would be neat; I'll add it to the backlog.