Feature/init plan util

pyproject.toml

@@ -19,14 +19,15 @@ classifiers = [
 dynamic = ["version"]
 requires-python = ">=3.7"
 dependencies = [
+  "dijkstra3d>=1.12.1",
   "joblib>=1.2.0",
   "numpy>=1.20",
-  "torch>=1.13",
   "rich>=13.3.1",
+  "POT>=0.9.0",
   "scikit-learn",
   "scipy",
+  "torch>=1.13",
   "tvb-gdist>=2.1.1",
-  "dijkstra3d>=1.12.1",
 ]
 
 [project.optional-dependencies]

src/fugw/mappings/dense.py

@@ -3,7 +3,7 @@
 
 from fugw.solvers.dense import FUGWSolver
 from fugw.mappings.utils import BaseMapping, console
-from fugw.utils import _make_tensor
+from fugw.utils import _make_tensor, init_plan_dense
 
 
 class FUGW(BaseMapping):
@@ -85,6 +85,7 @@ def fit(
             will be set to 1 / m.
         init_plan: ndarray(n, m) or None
             Transport plan to use at initialisation.
+            If None, an entropic initialization will be used.
         init_duals: tuple of [ndarray(n), ndarray(m)] or None
             Dual potentials to use at initialisation.
         solver: "sinkhorn" or "mm" or "ibpp"
@@ -140,11 +141,21 @@ def fit(
         else:
             wt = _make_tensor(target_weights, device=device)
 
-        # If initial plan is provided, move it to device
+        # If initial plan is provided, move it to device.
+        # Otherwise, initialize it with entropic initialization
         pi_init = (
             _make_tensor(init_plan, device=device)
             if init_plan is not None
-            else None
+            else _make_tensor(
+                init_plan_dense(
+                    source_features.shape[1],
+                    target_features.shape[1],
+                    weights_source=ws,
+                    weights_target=wt,
+                    method="entropic",
+                ),
+                device=device,
+            )
         )
 
         # Compute distance matrix between features

src/fugw/utils.py

@@ -3,6 +3,7 @@
 import numpy as np
 import torch
 
+from ot import emd_1d
 from rich.console import Console
 from rich.progress import (
     BarColumn,
@@ -189,6 +190,68 @@ def _add_dict(d, new_d):
     return d
 
 
+def init_plan_dense(
+    n_source,
+    n_target,
+    weights_source=None,
+    weights_target=None,
+    method="entropic",
+):
+    """Initialize transport plan with dense tensor.
+
+    Generate a matrix satisfying the constraints of a transport plan.
+    In particular, marginal constraints on lines and columns are satisfied.
+
+    Parameters
+    ----------
+    n_source: int
+        Number of source points
+    n_target: int
+        Number of target points
+    weights_source: torch.Tensor of size(n_source), optional, defaults to None
+        Source weights used in entropic init
+    weights_target: torch.Tensor of size(n_target), optional, defaults to None
+        Target weights used in entropic init
+    method: str, optional, defaults to "entropic"
+        Method to use for initialization.
+        Can be "entropic", "permutation" or "identity".
+        If "entropic", weights_source and weights_target must be provided ;
+        the initial plan is then given by the product of the two arrays.
+        If "permutation", the initial plan is the solution to a 1D
+        optimal transport problem between two random arrays, which can be
+        understood as a soft permutation between source and target points.
+        If "identity", the number of source and target points must be equal ;
+        the initial plan is then the identity matrix.
+
+    Returns
+    -------
+    init_plan: torch.Tensor of size(n_source, n_target)
+    """
+
+    if method == "identity":
+        assert n_source == n_target, (
+            "Number of source and target points must be equal "
+            "when using identity initialization."
+        )
+        plan = torch.eye(n_source, dtype=torch.float32)
+        plan = plan / plan.sum()
+    elif method == "entropic":
+        if weights_source is None:
+            weights_source = torch.ones(n_source, dtype=torch.float32)
+        if weights_target is None:
+            weights_target = torch.ones(n_target, dtype=torch.float32)
+        plan = weights_source[:, None] * weights_target[None, :]
+        plan = plan / plan.sum()
+    elif method == "permutation":
+        xa = torch.rand(n_source)
+        xb = torch.rand(n_target)
+        plan = emd_1d(xa, xb).to(dtype=torch.float32)
+    else:
+        raise Exception(f"Unknown initialisation method {method}")
+
+    return plan
+
+
 def save_mapping(mapping, fname):
     """Save mapping in pickle file, separating hyperparams and weights.
 

tests/test_utils.py

@@ -11,6 +11,7 @@
 from fugw.utils import (
     _init_mock_distribution,
     _make_tensor,
+    init_plan_dense,
     load_mapping,
     save_mapping,
 )
@@ -156,3 +157,30 @@ def test_saving_and_loading(device, return_numpy, solver):
 
             weights = pickle.load(f)
             assert weights.shape == (n_voxels_source, n_voxels_target)
+
+
+@pytest.mark.parametrize(
+    "method", ["identity", "entropic", "permutation", "unknown"]
+)
+def test_init_plan(method):
+    n_source = 101
+    n_target = 99
+
+    if method == "unknown":
+        with pytest.raises(Exception, match="Unknown initialisation method.*"):
+            init_plan_dense(n_source, n_target, method=method)
+    else:
+        if method == "identity":
+            with pytest.raises(
+                AssertionError, match="Number of source and target.*"
+            ):
+                init_plan_dense(n_source, n_target, method=method)
+
+            n_source = 100
+            n_target = 100
+
+        plan = init_plan_dense(n_source, n_target, method=method)
+        assert plan.shape == (n_source, n_target)
+        # Check that plan satisfies marginal constraints
+        assert torch.allclose(plan.sum(dim=0), torch.ones(n_target) / n_target)
+        assert torch.allclose(plan.sum(dim=1), torch.ones(n_source) / n_source)