Extend Delta distribution to multiple dimensions (#76)

* adding vmf distribution

* introduce DeltaNormal and DeltaVMF distributions

* small changes

* add dependencies + fix things

* fix docs

* remove vmf distribution

* fix docstrings

* fix docstring

* make delta/delta_normal backwards compatible

* add test

* extend docs

* add one more assert to test

* add failing test

* fix docstrings

* change spherical to isotropic

* Adapt tests for delta distribution

* fix indent

* Update test_distributions.py

* Increase allowed tolerance in num. test

Author: gonlairo

cebra/data/single_session.py

@@ -17,6 +17,7 @@
 
 import abc
 import collections
+import warnings
 from typing import List
 
 import literate_dataclasses as dataclasses
@@ -164,9 +165,9 @@ class ContinuousDataLoader(cebra_data.Loader):
     * auxiliary variables, using the empirical distribution of how behavior various across
       ``time_offset`` timesteps (``time_delta``). Sampling for this setting is implemented
       in :py:class:`cebra.distributions.continuous.TimedeltaDistribution`.
-    * alternatively, the distribution can be selected to be a Gaussian distribution parametrized
-      by a fixed ``delta`` around the reference sample, using the implementation in
-      :py:class:`cebra.distributions.continuous.DeltaDistribution`.
+    * alternatively, the distribution can be selected to be a Gaussian distribution
+      parametrized by a fixed ``delta`` around the reference sample, using the implementation in
+      :py:class:`cebra.distributions.continuous.DeltaNormalDistribution`.
 
     Args:
         See dataclass fields.
@@ -208,8 +209,14 @@ def _init_distribution(self):
                     self.dataset.continuous_index,
                     self.time_offset,
                     device=self.device)
-            elif self.conditional == "delta":
-                self.distribution = cebra.distributions.DeltaDistribution(
+
+            elif self.conditional in ("delta", "delta_normal"):
+                if self.conditional == "delta":
+                    warnings.warn(
+                        '"delta" distribution will be deprecated in an upcoming release. Please use "delta_normal" instead.',
+                        DeprecationWarning)
+
+                self.distribution = cebra.distributions.DeltaNormalDistribution(
                     self.dataset.continuous_index,
                     self.delta,
                     device=self.device)

cebra/distributions/__init__.py

@@ -51,6 +51,7 @@
     "Discrete",
     "DiscreteUniform",
     "DiscreteEmpirical",
+    "DeltaNormalDistribution",
     "MultivariateDiscrete",
     "MultisessionSampler",
 ]

cebra/distributions/continuous.py

@@ -242,14 +242,16 @@ def sample_conditional(self, reference_idx: torch.Tensor) -> torch.Tensor:
         return self.index.search(query)
 
 
-class DeltaDistribution(abc_.JointDistribution, abc_.HasGenerator):
+class DeltaNormalDistribution(abc_.JointDistribution, abc_.HasGenerator):
     """Define a conditional distribution based on behavioral changes over time.
 
-    Takes a continuous index, and uses sample from Gaussian distribution to sample positive
+    Takes a continuous index, and uses sample from Gaussian distribution to sample positive pairs.
+    Note that if the continuous index is multidimensional, the Gaussian distribution will have
+    isotropic covariance matrix i.e. Σ = sigma^2 * I.
 
     Args:
-        continuous: The multidimensional, continuous index
-        delta: Standard deviation of Gaussian distribution to sample positive pair
+        continuous: The multidimensional, continuous index.
+        delta: Standard deviation of Gaussian distribution to sample positive pair.
 
     """
 
@@ -277,12 +279,14 @@ def sample_conditional(self, reference_idx: torch.Tensor) -> torch.Tensor:
                 "Pass a 1D array of indices of reference samples.")
 
         # TODO(stes): Set seed
+        mean = self.data[reference_idx]
         query = torch.distributions.Normal(
-            self.data[reference_idx].squeeze(),
-            torch.ones_like(reference_idx, device=self.device) * self.std,
+            loc=mean,
+            scale=torch.ones_like(mean, device=self.device) * self.std,
         ).sample()
 
-        return self.index.search(query.unsqueeze(-1))
+        query = query.unsqueeze(-1) if query.dim() == 1 else query
+        return self.index.search(query)
 
 
 class CEBRADistribution(abc_.JointDistribution):

tests/test_distributions.py

@@ -10,13 +10,15 @@
 # https://github.com/AdaptiveMotorControlLab/CEBRA/LICENSE.md
 #
 import functools
+from typing import Literal, Optional
 
 import numpy as np
 import pytest
 import torch
 
 import cebra.datasets as cebra_datasets
 import cebra.distributions as cebra_distr
+import cebra.distributions.base as cebra_distr_base
 
 
 def assert_is_tensor(T, device=None):
@@ -284,3 +286,93 @@ def test_multi_session_time_contrastive(time_offset):
     # NOTE(celia): test the private function ``_inverse_idx()``, with idx arrays flat
     assert (idx.flatten()[rev_idx.flatten()].all() == np.arange(
         len(rev_idx.flatten())).all())
+
+
+class OldDeltaDistribution(cebra_distr_base.JointDistribution,
+                           cebra_distr_base.HasGenerator):
+    """
+    Old version of the Delta Distribution where it only works for 1d
+    behavior variable.
+
+    """
+
+    def __init__(self,
+                 continuous: torch.Tensor,
+                 delta: float = 0.1,
+                 device: Literal["cpu", "cuda"] = "cpu",
+                 seed: Optional[int] = 1812):
+        cebra_distr_base.HasGenerator.__init__(self, device=device, seed=seed)
+        torch.manual_seed(seed)
+        self.data = continuous
+        self.std = delta
+        self.index = cebra_distr.ContinuousIndex(self.data)
+        self.prior = cebra_distr.Prior(self.data, device=device, seed=seed)
+
+    def sample_prior(self, num_samples: int) -> torch.Tensor:
+        """See :py:meth:`.Prior.sample_prior`."""
+        return self.prior.sample_prior(num_samples)
+
+    def sample_conditional(self, reference_idx: torch.Tensor) -> torch.Tensor:
+        """Return indices from the conditional distribution."""
+
+        if reference_idx.dim() != 1:
+            raise ValueError(
+                f"Reference indices have wrong shape: {reference_idx.shape}. "
+                "Pass a 1D array of indices of reference samples.")
+
+        # TODO(stes): Set seed
+        query = torch.distributions.Normal(
+            self.data[reference_idx].squeeze(),
+            torch.ones_like(reference_idx, device=self.device) * self.std,
+        ).sample()
+
+        return self.index.search(query.unsqueeze(-1))
+
+
+def test_old_vs_new_delta_normal_with_1Dindex():
+    _, continuous = prepare()
+    assert continuous.dim() == 2
+    num_samples = len(continuous)
+    reference_idx = torch.randint(0, num_samples, (num_samples,))
+
+    new_distribution = cebra_distr.DeltaNormalDistribution(
+        continuous=continuous[:, 0].unsqueeze(-1), delta=0.1)
+
+    old_distribution = OldDeltaDistribution(
+        continuous=continuous[:, 0].unsqueeze(-1), delta=0.1)
+
+    torch.manual_seed(1812)
+    old_positives = old_distribution.sample_conditional(reference_idx)
+    torch.manual_seed(1812)
+    new_positives = new_distribution.sample_conditional(reference_idx)
+
+    assert not torch.equal(old_positives, reference_idx)
+    assert not torch.equal(new_positives, reference_idx)
+    assert torch.equal(old_positives, new_positives)
+
+
+@pytest.mark.parametrize("delta,numerical_check", [(0.01, True), (0.025, True), (1., False), (5., False)])
+def test_new_delta_normal_with_multidimensional_index(delta, numerical_check):
+    continuous = torch.rand(100_000, 3).to("cpu")
+    num_samples = 1000
+    delta_normal_multidim = cebra_distr.DeltaNormalDistribution(
+        delta=delta, continuous=continuous)
+    reference_idx = delta_normal_multidim.sample_prior(num_samples)
+    positive_idx = delta_normal_multidim.sample_conditional(reference_idx)
+
+    assert positive_idx.dim() == 1
+    assert len(positive_idx) == num_samples
+    assert not torch.equal(positive_idx, reference_idx)
+
+    if numerical_check:
+        reference_samples = continuous[reference_idx]
+        positive_samples = continuous[positive_idx]
+        diff = positive_samples - reference_samples
+        #TODO(stes): Improve test, use lower error margin here
+        assert torch.isclose(diff.std(), torch.tensor(delta), rtol=0.1)
+    else:
+        #TODO(stes): Add a warning message to the delta distribution.
+        pytest.skip(
+          "multivariate delta distribution can not accurately sample with the "
+          "given parameters. TODO: Add a warning message for these cases."
+        )