[MRG] Adding greenkhorn by arakotom · Pull Request #66 · PythonOT/POT

>>> a=[.5,.5]
>>> b=[.5,.5]
>>> M=[[0.,1.],[1.,0.]]
>>> ot.sinkhorn(a,b,M,1)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ot.greenkhorn

m = b.shape[0]

# Next 3 lines equivalent to K= np.exp(-M/reg), but faster to compute
K = np.empty(M.shape, dtype=M.dtype)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

K = np.empty_like(M)

np.divide(M, -reg, out=K)
np.exp(K, out=K)

u = np.ones(n)/n

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

np.full(n, 1. / n)

np.exp(K, out=K)

u = np.ones(n)/n
v = np.ones(m)/m

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

np.full(m, 1. / m)


u = np.ones(n)/n
v = np.ones(m)/m
G = np.diag(u)@K@np.diag(v)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

use broadcasting to avoid filling diagonal matrices

G = u[:, np.newaxis] * K * v[, np.newaxis]

G[:,i_2] = u*K[:,i_2]*v[i_2]
#aviol = (G@one_m - a)
#aviol_2 = (G.T@one_n - b)
viol = viol + ( -old_v + v[i_2])*K[:,i_2]*u

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

viol += ...

np.testing.assert_allclose(G0, Ges, atol=1e-05)
np.testing.assert_allclose(G0, Gerr)

np.testing.assert_allclose(G0, G_green, atol = 1e-32)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

a pep8 checker would tell you but you should not put spaces around = in function signatures. It's to visually distinguish what is a function parameter from a variable assignment.

np.testing.assert_allclose(G0, Gerr)

np.testing.assert_allclose(G0, G_green, atol = 1e-32)
print(G0,G_green)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

and you should always put a space after a ,

one_n = np.ones(n)
one_m = np.ones(m)
viol = G@one_m - a
viol_2 = G.T@one_n - b

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

here to allocate arrays of ones to compute sum of rows and columns. I would just use np.sum(..., axis=)

log['u'] = u
log['v'] = v

while i < numItermax and stopThr_val > stopThr:

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rather than using a while you could use a for loop. For optim solvers I tend to do:

for i in range(numItermax):
      ...
      if stopping condition satisfied do:
              break
else:
     print("Solver did not converge")

so you can easily print a message when you did not converge.