Refactor optimal transport logic in flow matching.

Updated the optimal transport implementation to prioritize resampling x1 due to potential data noise and outliers, ensuring better stability across batches. Adjusted conditions to align properly with assignments returned by the optimal transport function. These changes improve robustness and clarity in the resampling process.

Author: LarsKue

bayesflow/networks/flow_matching/flow_matching.py

@@ -256,9 +256,20 @@ def compute_metrics(
             x0 = self.base_distribution.sample(keras.ops.shape(x1)[:-1])
 
             if self.use_optimal_transport:
-                x1, x0, conditions = optimal_transport(
-                    x1, x0, conditions, seed=self.seed_generator, **self.optimal_transport_kwargs
+                # we must choose between resampling x0 or x1
+                # since the data is possibly noisy and may contain outliers, it is better
+                # to possibly drop some samples from x1 than from x0
+                # in the marginal over multiple batches, this is not a problem
+                x0, x1, assignments = optimal_transport(
+                    x0,
+                    x1,
+                    seed=self.seed_generator,
+                    **self.optimal_transport_kwargs,
+                    return_assignments=True,
                 )
+                if conditions is not None:
+                    # conditions must be resampled along with x1
+                    conditions = conditions[assignments]
 
             t = keras.random.uniform((keras.ops.shape(x0)[0],), seed=self.seed_generator)
             t = expand_right_as(t, x0)