Merge pull request #64 from pbarbarant/feat/add_callbacks_barycenter

Feat/add callbacks barycenter

src/fugw/mappings/barycenter.py

@@ -121,6 +121,7 @@ def compute_all_ot_plans(
         barycenter_geometry,
         solver,
         solver_params,
+        callback_barycenter,
         device,
         verbose,
     ):
@@ -182,6 +183,7 @@ def fit(
         solver_params={},
         nits_barycenter=5,
         device="auto",
+        callback_barycenter=None,
         verbose=False,
     ):
         """Compute barycentric features and geometry
@@ -209,6 +211,11 @@ def fit(
         device: "auto" or torch.device
             if "auto": use first available gpu if it's available,
             cpu otherwise.
+        callback_barycenter: callable or None
+            Callback function called at the end of each barycenter step.
+            It will be called with the following arguments:
+
+                - locals (dictionary containing all local variables)
 
         Returns
         -------
@@ -270,7 +277,7 @@ def fit(
         duals = None
         losses_each_bar_step = []
 
-        for _ in range(nits_barycenter):
+        for idx in range(nits_barycenter):
             # Transport all elements
             plans, losses = self.compute_all_ot_plans(
                 plans,
@@ -283,6 +290,7 @@ def fit(
                 barycenter_geometry,
                 solver,
                 solver_params,
+                callback_barycenter,
                 device,
                 verbose,
             )
@@ -298,6 +306,9 @@ def fit(
                     plans, weights_list, geometry_list, self.force_psd, device
                 )
 
+            if callback_barycenter is not None:
+                callback_barycenter(locals())
+
         return (
             barycenter_weights,
             barycenter_features,

src/fugw/mappings/sparse_barycenter.py

@@ -156,6 +156,7 @@ def fit(
         mesh_sample=None,
         nits_barycenter=5,
         device="auto",
+        callback_barycenter=None,
         verbose=False,
     ):
         """Compute barycentric features and geometry
@@ -192,6 +193,11 @@ def fit(
         device: "auto" or torch.device
             if "auto": use first available gpu if it's available,
             cpu otherwise.
+        callback_barycenter: callable or None
+            Callback function called at the end of each barycenter step.
+            It will be called with the following arguments:
+
+                - locals (dictionary containing all local variables)
 
         Returns
         -------
@@ -248,7 +254,7 @@ def fit(
         mask = None
         losses_each_bar_step = []
 
-        for _ in range(nits_barycenter):
+        for idx in range(nits_barycenter):
             # Transport all elements
             plans, losses = self.compute_all_ot_plans(
                 plans,
@@ -274,6 +280,9 @@ def fit(
                 plans, weights_list, features_list, device
             )
 
+            if callback_barycenter is not None:
+                callback_barycenter(locals())
+
         return (
             barycenter_weights,
             barycenter_features,

tests/mappings/test_barycenter.py

@@ -1,3 +1,5 @@
+from itertools import product
+
 import numpy as np
 import pytest
 import torch
@@ -9,9 +11,14 @@
 if torch.cuda.is_available():
     devices.append(torch.device("cuda:0"))
 
+callbacks = [None, lambda x: x["plans"]]
+
 
-@pytest.mark.parametrize("device", devices)
-def test_fugw_barycenter(device):
+@pytest.mark.parametrize(
+    "device, callback",
+    product(devices, callbacks),
+)
+def test_fugw_barycenter(device, callback):
     np.random.seed(0)
     n_subjects = 4
     n_voxels = 100
@@ -32,5 +39,9 @@ def test_fugw_barycenter(device):
 
     fugw_barycenter = FUGWBarycenter()
     fugw_barycenter.fit(
-        weights_list, features_list, geometry_list, device=device
+        weights_list,
+        features_list,
+        geometry_list,
+        device=device,
+        callback_barycenter=callback,
     )

tests/mappings/test_sparse_barycenter.py

@@ -1,3 +1,5 @@
+from itertools import product
+
 import numpy as np
 import pytest
 import torch
@@ -9,10 +11,15 @@
 if torch.cuda.is_available():
     devices.append(torch.device("cuda:0"))
 
+callbacks = [None, lambda x: x["plans"]]
+
 
 @pytest.mark.skip_if_no_mkl
-@pytest.mark.parametrize("device", devices)
-def test_fugw_barycenter(device):
+@pytest.mark.parametrize(
+    "device, callback",
+    product(devices, callbacks),
+)
+def test_fugw_barycenter(device, callback):
     np.random.seed(0)
     n_subjects = 4
     n_voxels = 100