introduce multivariate cdf / quantiles

src/ott/tools/soft_sort.py

@@ -18,15 +18,18 @@
 import jax.numpy as jnp
 import numpy as np
 
-from ott.geometry import pointcloud
+from ott.geometry import costs, pointcloud
 from ott.problems.linear import linear_problem
+from ott.solvers import linear
 from ott.solvers.linear import sinkhorn
 
 __all__ = [
     "sort", "ranks", "sort_with", "quantile", "quantile_normalization",
     "quantize", "topk_mask"
 ]
 
+Func_t = Callable[[jnp.ndarray], jnp.ndarray]
+
 
 def transport_for_sort(
     inputs: jnp.ndarray,
@@ -450,6 +453,80 @@ def _quantile(
   return apply_on_axis(_quantile, inputs, axis, q, weight, **kwargs)
 
 
+def multiv_cdf_quantile_maps(
+    inputs: jnp.ndarray,
+    target_sampler: Optional[Callable[[jax.random.PRNGKey, Tuple[int, int]],
+                                      jnp.ndarray]] = None,
+    key: Optional[jax.random.PRNGKey] = None,
+    num_target_samples: Optional[int] = None,
+    cost_fn: Optional[costs.CostFn] = None,
+    epsilon: Optional[float] = None,
+    input_weights: Optional[jnp.ndarray] = None,
+    target_weights: Optional[jnp.ndarray] = None,
+    **kwargs: Any
+) -> [Func_t, Func_t]:
+  r"""Returns multivariate CDF and quantile maps, given input samples.
+
+  Implements the multivariate generalizations for CDF and quantiles proposed in
+  :cite:`chernozhukov:17`. The reference measure is assumed to be the uniform
+  measure by default, but can be modified. For consistency, the reference
+  measure should be symmetrically centered around
+  :math:`(\tfrac{1}{2},\cdots,\tfrac{1}{2})` and suppported on :math:`[0,1]^d`.
+
+  The implementation return two entropic map estimators, one for the CDF map,
+  the other for the quantiles map.
+
+  Args:
+    inputs: 2D array of :math:`[n, d]` vectors.
+    target_sampler: Callable that takes a ``key`` and ``[m,d]`` shape Tuple.
+      ``m`` is passed on as ``target_num_samples``, dimension ``d`` is inferred
+      directly from the shape passed in ``inputs``. This is assumed by default
+      to be :func:`jax.random.uniform`, and could be any other random sampler
+      properly wrapped to have the signature above.
+    key: rng key used by ``target_sampler``
+    num_target_samples: number ``m`` of points generated in the target
+      distribution.
+    cost_fn: :class:`~ott.geometry.costs.CostFn` object, used to compare
+      ``inputs`` and ``targets``. Passed on to instantiate a
+      :class:`~ott.geometry.pointcloud.PointCloud` object. This
+      defaults to the squared-Euclidean distance.
+    epsilon: entropic regularization parameter used to instantiate the
+      :class:`~ott.geometry.pointcloud.PointCloud` object.
+    input_weights: :math:`[n,]` vector of weights for input measure. Assumed to
+      be uniform by default.
+    target_weights: :math:`[m,]` vector of weights for target measure. Assumed
+      to be uniform by default.
+    kwargs: keyword arguments passed on to the :func:`~ott.solvers.linear.solve`
+      function, which solves the OT problem between ``inputs`` and ``targets``
+      using the Sinkhorn algorithm.
+
+  Returns:
+    Two callables, vector-to-vector mappings:
+      - multivariate CDF map, taking values in the range of the reference
+        measure.
+      - quantile map, going by default from :math:`[0, 1]^d` to the range of the
+        input measure.
+
+  Raises:
+    A ValueError in case the input and target have not the same dimension.
+  """
+  n, d = inputs.shape
+  key = jax.random.PRNGKey(0) if key is None else key
+  num_target_samples = n if num_target_samples is None else num_target_samples
+  if target_sampler is None:
+    target_sampler = jax.random.uniform
+  targets = target_sampler(key, (num_target_samples, d))
+
+  geom = pointcloud.PointCloud(
+      inputs, targets, cost_fn=cost_fn, epsilon=epsilon
+  )
+  out = linear.solve(geom, a=input_weights, b=target_weights, **kwargs)
+  potentials = out.to_dual_potentials()
+  return potentials.transport, functools.partial(
+      potentials.transport, forward=False
+  )
+
+
 def _quantile_normalization(
     inputs: jnp.ndarray, targets: jnp.ndarray, weights: float, **kwargs: Any
 ) -> jnp.ndarray:

tests/solvers/linear/sinkhorn_test.py

@@ -474,7 +474,7 @@ def test_restart(self, lse_mode: bool):
     assert num_iter_restarted == 1
 
   @pytest.mark.cpu()
-  @pytest.mark.limit_memory("35 MB")
+  @pytest.mark.limit_memory("36 MB")
   @pytest.mark.fast()
   def test_sinkhorn_online_memory_jit(self):
     # test that full matrix is not materialized.

tests/tools/soft_sort_test.py

@@ -86,6 +86,44 @@ def test_sort_batch(self, rng: jax.random.PRNGKeyArray, topk: int):
     np.testing.assert_array_equal(xs.shape, expected_shape)
     np.testing.assert_array_equal(jnp.diff(xs, axis=axis) >= 0.0, True)
 
+  def test_multiv_cdf_quantiles(self):
+    n, d = 512, 3
+    keys = jax.random.split(jax.random.PRNGKey(0), 3)
+
+    # Set central point in sampled input measure
+    z = jax.random.uniform(keys[0], (1, d))
+
+    # Sample inputs symmetrically centered on z
+    inputs = 0.34 * jax.random.normal(keys[0], (n, d)) + z
+
+    # Set central point in target distribution.
+    q = 0.5 * jnp.ones((1, d))
+
+    # Set tolerance for quantile / cdf comparisons to ground truth.
+    atol = 0.1
+
+    # Check approximate correctness of naked call to API
+    cdf, qua = soft_sort.multiv_cdf_quantile_maps(inputs)
+    np.testing.assert_allclose(cdf(z), q, atol=atol)
+    np.testing.assert_allclose(z, qua(q), atol=atol)
+
+    # Check passing custom sampler, must be still symmetric / centered on {.5}^d
+    # Check passing custom epsilon also works.
+    def ball_sampler(k: jax.random.PRNGKey, s: Tuple[int, int]) -> jnp.ndarray:
+      return 0.5 * (jax.random.ball(k, d=s[1], p=4, shape=(s[0],)) + 1.)
+
+    num_target_samples = 473
+
+    cdf, qua = soft_sort.multiv_cdf_quantile_maps(
+        inputs,
+        target_sampler=ball_sampler,
+        num_target_samples=num_target_samples,
+        key=keys[2],
+        epsilon=0.05
+    )
+    np.testing.assert_allclose(cdf(z), q, atol=atol)
+    np.testing.assert_allclose(z, qua(q), atol=atol)
+
   @pytest.mark.fast.with_args("axis,jit", [(0, False), (1, True)], only_fast=0)
   def test_ranks(self, axis, rng: jax.random.PRNGKeyArray, jit: bool):
     rng1, rng2 = jax.random.split(rng, 2)