Thank you romain, this is nice.

Is it me of Cython is not particularly fast (2 sec for n=20000?)? it is probably due to the use of the dist function, you should probably implement it in cython for squared and absolute value and use the dist only for weird stuff ;)

Rémi.

If I change to the following:

        if metric == 'sqeuclidean':
            m_ij = (u[i, 0] - v[j, 0]) ** 2
        else:
            m_ij = dist(u[i].reshape((1, 1)), v[j].reshape((1, 1)),
                        metric=metric)[0, 0]

I get the same timings (in the order of 2secs)...

The slow part seems to be when we deal with G. If I remove all the stuff related to G I get:

Elapsed time : 0.0061719417572021484 s

I will check if using a sparse representation for G helps.

EDIT: OK, when I remove the overhead for G, I can see a 100x improvement in timings with this if...else, so will keep it for L1 and L2 norms and resort to dist for other distances.

Also, new timings (for larger problem than above) are:

>>> import ot
>>> import numpy as np
>>> from scipy.stats import wasserstein_distance
>>> 
>>> n = 20000
>>> m = 30000
>>> u = np.random.randn(n)
>>> v = np.random.randn(m)
>>> 
>>> ot.tic(); _ = wasserstein_distance(u, v); _ = ot.toc()
Elapsed time : 0.012831926345825195 s
>>> ot.tic(); _ = ot.emd_1d([], [], u, v, metric='euclidean', dense=False); _ = ot.toc()
Elapsed time : 0.04144096374511719 s
>>> ot.tic(); M = ot.dist(u.reshape((-1, 1)), v.reshape((-1, 1)),
...                       metric='euclidean'); _ = ot.emd([], [], M); _ = ot.toc()

RESULT MIGHT BE INACURATE
Max number of iteration reached, currently 100000. Sometimes iterations go on in cycle even though the solution has been reached, to check if it's the case here have a look at the minimal reduced cost. If it is very close to machine precision, you might actually have the correct solution, if not try setting the maximum number of iterations a bit higher
/Users/tavenard_r/Documents/costel/src/POT/ot/lp/__init__.py:106: UserWarning: numItermax reached before optimality. Try to increase numItermax.
  result_code_string = check_result(result_code)
Elapsed time : 312.5033311843872 s

We are a bit slower than scipy's implementation, not sure whether this is due to Cython or to the fact that scipy does not deal with G :/

@rtavenar did you run "cython -a" on the pyx file to see if it's white (no yellow slow python lines)?

@agramfort

1. Did not know this one, thanks for the tip
2. I've changed the np.abs, now the only yellow lines I get are the return, the np.zeros lines and the cdist, but I do not know how to remove these ones.

you cannot remove yellow lines of np.zeros or return.

For cdist either you can directly call blas functions from scipy or you need
to code the metrics in cython

I've had a look there, could not find obvious matches for distances, but maybe it's not the right place :/

Regarding coding the metrics in Cython, this is what I have done for Euclidean distance and Squared Euclidean distance up to now. The question is: should I code all of them even if they are unlikely to be used, or only a subset?

hello, I think those two are OK just be clear in the documentation that the others are slower and use cdist (such a slow function btw ;) )

Rémi

OK, and I'll also have to be clear that only strings are accepted as metrics for emd_1d

OK, so now I added proper docstrings. Let me know if something is missing or should be changed.

This is great, thank you @rtavenar for the code and optimization.

I will merge it now.