Refactor optimal transport utilities and add log-Sinkhorn

Removed unimplemented Hungarian algorithm and streamline Sinkhorn implementation. Introduced log-Sinkhorn for numerical stability and modularized cost computation (e.g., Euclidean distance). Updated the interface for improved usability and scalability.

Author: LarsKue

bayesflow/utils/optimal_transport/__init__.py

@@ -1,2 +1 @@
 from .optimal_transport import optimal_transport
-from .sinkhorn import sinkhorn, sinkhorn_indices, sinkhorn_plan

bayesflow/utils/optimal_transport/euclidean.py

@@ -0,0 +1,11 @@
+import keras
+
+
+def euclidean(x1, x2):
+    # TODO: rename and move this function
+    result = x1[:, None] - x2[None, :]
+    shape = list(keras.ops.shape(result))
+    shape[2:] = [-1]
+    result = keras.ops.reshape(result, shape)
+    result = keras.ops.norm(result, ord=2, axis=-1)
+    return result

bayesflow/utils/optimal_transport/hungarian.py

@@ -0,0 +1,80 @@
+import keras
+
+from .. import logging
+from ..tensor_utils import is_symbolic_tensor
+
+from .euclidean import euclidean
+
+
+def log_sinkhorn(x1, x2, seed: int = None, **kwargs):
+    """
+    Log-stabilized version of :py:func:`~bayesflow.utils.optimal_transport.sinkhorn.sinkhorn`.
+    Significantly slower than the unstabilized version, so use only when you need numerical stability.
+    """
+    log_plan = log_sinkhorn_plan(x1, x2, **kwargs)
+    assignments = keras.random.categorical(keras.ops.exp(log_plan), num_samples=1, seed=seed)
+    assignments = keras.ops.squeeze(assignments, axis=1)
+
+    return assignments
+
+
+def log_sinkhorn_plan(x1, x2, regularization: float = 1.0, rtol=1e-5, atol=1e-8, max_steps=None):
+    """
+    Log-stabilized version of :py:func:`~bayesflow.utils.optimal_transport.sinkhorn.sinkhorn_plan`.
+    Significantly slower than the unstabilized version, so use only when you need numerical stability.
+    """
+    cost = euclidean(x1, x2)
+
+    log_plan = cost / -(regularization * keras.ops.mean(cost) + 1e-16)
+
+    if is_symbolic_tensor(log_plan):
+        return log_plan
+
+    def contains_nans(plan):
+        return keras.ops.any(keras.ops.isnan(plan))
+
+    def is_converged(plan):
+        # for convergence, the plan should be doubly stochastic
+        conv0 = keras.ops.all(keras.ops.isclose(keras.ops.logsumexp(plan, axis=0), 0.0, rtol=rtol, atol=atol))
+        conv1 = keras.ops.all(keras.ops.isclose(keras.ops.logsumexp(plan, axis=1), 0.0, rtol=rtol, atol=atol))
+        return conv0 & conv1
+
+    def cond(_, plan):
+        # break the while loop if the plan contains nans or is converged
+        return ~(contains_nans(plan) | is_converged(plan))
+
+    def body(steps, plan):
+        # Sinkhorn-Knopp: repeatedly normalize the transport plan along each dimension
+        plan = keras.ops.log_softmax(plan, axis=0)
+        plan = keras.ops.log_softmax(plan, axis=1)
+
+        return steps + 1, plan
+
+    steps = 0
+    steps, log_plan = keras.ops.while_loop(cond, body, (steps, log_plan), maximum_iterations=max_steps)
+
+    def do_nothing():
+        pass
+
+    def log_steps():
+        msg = "Log-Sinkhorn-Knopp converged after {:d} steps."
+
+        logging.info(msg, steps)
+
+    def warn_convergence():
+        marginals = keras.ops.logsumexp(log_plan, axis=0)
+        deviations = keras.ops.abs(marginals)
+        badness = 100.0 * keras.ops.exp(keras.ops.max(deviations))
+
+        msg = "Log-Sinkhorn-Knopp did not converge after {:d} steps (badness: {:.1f}%)."
+
+        logging.warning(msg, max_steps, badness)
+
+    def warn_nans():
+        msg = "Log-Sinkhorn-Knopp produced NaNs."
+        logging.warning(msg)
+
+    keras.ops.cond(contains_nans(log_plan), warn_nans, do_nothing)
+    keras.ops.cond(is_converged(log_plan), log_steps, warn_convergence)
+
+    return log_plan

bayesflow/utils/optimal_transport/log_sinkhorn.py

@@ -1,12 +1,15 @@
-from bayesflow.types import Tensor
-
-from .hungarian import hungarian
+from .log_sinkhorn import log_sinkhorn
 from .sinkhorn import sinkhorn
 
+methods = {
+    "sinkhorn": sinkhorn,
+    "sinkhorn_knopp": sinkhorn,
+    "log_sinkhorn": log_sinkhorn,
+    "log_sinkhorn_knopp": log_sinkhorn,
+}
+
 
-def optimal_transport(
-    x1: Tensor, x2: Tensor, *aux: Tensor, method: str = "sinkhorn_knopp", **kwargs
-) -> (Tensor, Tensor):
+def optimal_transport(x1, x2, method="log_sinkhorn", return_assignments=False, **kwargs):
     """Matches elements from x2 onto x1, such that the transport cost between them is minimized, according to the method
     and cost matrix used.
 
@@ -22,27 +25,21 @@ def optimal_transport(
     :param x2: Tensor of shape (m, ...)
         Samples from the second distribution.
 
-    :param aux: Tensors of shape (n, ...)
-        Auxiliary tensors to be permuted along with x1.
-        Note that x2 is never permuted for all currently available methods.
-
     :param method: Method used to compute the transport cost.
-        Default: 'sinkhorn_knopp'
+        Default: 'log_sinkhorn'
 
-    :param kwargs: Additional keyword arguments passed to the optimization method.
+    :param return_assignments: Whether to return the assignment indices.
+        Default: False
+
+    :param kwargs: Additional keyword arguments that are passed to the optimization method.
 
     :return: Tensors of shapes (n, ...) and (m, ...)
         x1 and x2 in optimal transport permutation order.
     """
-    methods = {
-        "hungarian": hungarian,
-        "sinkhorn": sinkhorn,
-        "sinkhorn_knopp": sinkhorn,
-    }
-
-    method = method.lower()
+    assignments = methods[method.lower()](x1, x2, **kwargs)
+    x2 = x2[assignments]
 
-    if method not in methods:
-        raise ValueError(f"Unsupported method name: '{method}'.")
+    if return_assignments:
+        return x1, x2, assignments
 
-    return methods[method](x1, x2, *aux, **kwargs)
+    return x1, x2