From a6d5d75c6ca584ab9736b528810b3595f2571d82 Mon Sep 17 00:00:00 2001
From: tgnassou <66993815+tgnassou@users.noreply.github.com>
Date: Tue, 28 Feb 2023 08:35:12 +0100
Subject: [MRG] Add method argument to sinkhorn Transport (#440)

* add method argument to sinkhron transport'

* update release
---
 ot/da.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

(limited to 'ot')

diff --git a/ot/da.py b/ot/da.py
index 35e303b..5067a69 100644
--- a/ot/da.py
+++ b/ot/da.py
@@ -1417,12 +1417,13 @@ class SinkhornTransport(BaseTransport):
             Sciences, 7(3), 1853-1882.
     """
 
-    def __init__(self, reg_e=1., max_iter=1000,
+    def __init__(self, reg_e=1., method="sinkhorn", max_iter=1000,
                  tol=10e-9, verbose=False, log=False,
                  metric="sqeuclidean", norm=None,
                  distribution_estimation=distribution_estimation_uniform,
                  out_of_sample_map='ferradans', limit_max=np.infty):
         self.reg_e = reg_e
+        self.method = method
         self.max_iter = max_iter
         self.tol = tol
         self.verbose = verbose
@@ -1463,7 +1464,7 @@ class SinkhornTransport(BaseTransport):
         # coupling estimation
         returned_ = sinkhorn(
             a=self.mu_s, b=self.mu_t, M=self.cost_, reg=self.reg_e,
-            numItermax=self.max_iter, stopThr=self.tol,
+            method=self.method, numItermax=self.max_iter, stopThr=self.tol,
             verbose=self.verbose, log=self.log)
 
         # deal with the value of log
-- 
cgit v1.2.3