diff options
| author | Mokhtar Z. Alaya <mokhtarzahdi.alaya@gmail.com> | 2020-01-07 11:18:43 +0100 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2020-01-07 11:18:43 +0100 |
| commit | d4b403ef3bda06db101d8a7c3dba1e0be36e9c80 (patch) | |
| tree | a6b739d3b5f0202b7146cc6ed47e7d1fc0efd36f /ot/bregman.py | |
| parent | 3582b6ffe571cb96b69dc6c356ef8d946df57fe7 (diff) | |
close "bottleneck" tag
Diffstat (limited to 'ot/bregman.py')
| -rw-r--r-- | ot/bregman.py | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/ot/bregman.py b/ot/bregman.py index 61b8605..58c76d0 100644 --- a/ot/bregman.py +++ b/ot/bregman.py @@ -15,7 +15,6 @@ Bregman projections for regularized OT import numpy as np import warnings from .utils import unif, dist -import bottleneck from scipy.optimize import fmin_l_bfgs_b def sinkhorn(a, b, M, reg, method='sinkhorn', numItermax=1000, @@ -1893,6 +1892,7 @@ def screenkhorn(a, b, M, reg, ns_budget, nt_budget, uniform=True, restricted=Tru aK_sort = np.sort(K_sum_cols) epsilon_u_square = a[0] / aK_sort[ns_budget - 1] else: + import bottleneck aK_sort = bottleneck.partition(K_sum_cols, ns_budget - 1)[ns_budget - 1] epsilon_u_square = a[0] / aK_sort @@ -1900,6 +1900,7 @@ def screenkhorn(a, b, M, reg, ns_budget, nt_budget, uniform=True, restricted=Tru bK_sort = np.sort(K_sum_rows) epsilon_v_square = b[0] / bK_sort[ns_budget - 1] else: + import bottleneck bK_sort = bottleneck.partition(K_sum_rows, nt_budget - 1)[nt_budget - 1] epsilon_v_square = b[0] / bK_sort else: |