google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Tests for the sinkhorn module."""
17

18
from __future__ import absolute_import
19
from __future__ import division
20
from __future__ import print_function
21

22
from absl.testing import absltest
23
import jax.numpy as np
24

25
from soft_sort.jax import sinkhorn
26

27

28
class SinkhornTestCase(absltest.TestCase):
29
  """Test case for Sinkhorn1D."""
30

31
  def setUp(self):
32
    super(SinkhornTestCase, self).setUp()
33
    self.no_decay = dict(
34
        epsilon_0=1e-3, epsilon_decay=1.00, max_iterations=5000)
35
    self.decay = dict(
36
        epsilon_0=1e-1, epsilon_decay=0.95, max_iterations=5000)
37
    self.x = np.array([[0.2, -0.5, 0.4, 0.44, -0.9, 0.9]])
38
    self.y = 1 / self.x.shape[1] * np.cumsum(np.ones(self.x.shape), axis=-1)
39
    self.a = 1.0 / self.x.shape[1] * np.ones(self.x.shape)
40
    self.b = 1.0 / self.y.shape[1] * np.ones(self.y.shape)
41

42
  def test_sinkhorn(self):
43
    """Tests that the __call__ methods returns transport maps."""
44
    f, g, eps, cost, _, iterations = sinkhorn.sinkhorn_iterations(
45
        self.x, self.y, self.a, self.b, **self.decay)
46
    p = sinkhorn.transport(cost, f, g, eps)
47
    self.assertTupleEqual(p.shape, self.x.shape + (self.y.shape[1],))
48
    self.assertLessEqual(iterations, self.decay['max_iterations'])
49

50
  def test_decay(self):
51
    """Tests that applying the epsilon decay scheme speeds up convergence."""
52
    decay_iterations = sinkhorn.sinkhorn_iterations(
53
        self.x, self.y, self.a, self.b, **self.decay)[-1]
54
    no_decay_iterations = sinkhorn.sinkhorn_iterations(
55
        self.x, self.y, self.a, self.b, **self.no_decay)[-1]
56
    self.assertLess(decay_iterations, no_decay_iterations)
57

58

59
if __name__ == '__main__':
60
  absltest.main()
61