Merge pull request #1 from guillaumehu/doc_otplan

doc & sigma int

torchcfm/conditional_flow_matching.py

@@ -54,7 +54,16 @@ def __init__(self, sigma: float = 0.0):
         ----------
         sigma : float
         """
-        self.sigma = sigma
+        self._sigma = sigma
+
+    @property
+    def sigma(self):
+        if isinstance(self._sigma, float):
+            return self._sigma
+        elif isinstance(self._sigma, int):
+            return float(self._sigma)
+        else:
+            raise ValueError("Sigma must be a float or int.")
 
     def compute_mu_t(self, x0, x1, t):
         """

torchcfm/optimal_transport.py

@@ -19,6 +19,21 @@ def __init__(
         normalize_cost=False,
         **kwargs,
     ):
+        r"""Initialize the OTPlanSampler class.
+
+        Parameters
+        ----------
+        method : str
+            The method used to compute the OT plan. Can be one of "exact", "sinkhorn",
+            "unbalanced", or "partial".
+        reg : float (default : 0.05)
+            Entropic regularization coefficients.
+        reg_m : float (default : 1.0)
+            Marginal relaxation term for unbalanced OT (`method='unbalanced'`).
+        normalize_cost : bool (default : False)
+            Whether to normalize the cost matrix by its maximum value.
+            It should be set to `False` when using minibatches.
+        """
         # ot_fn should take (a, b, M) as arguments where a, b are marginals and
         # M is a cost matrix
         if method == "exact":