Earth Mover Distance

emd.py

import numpy
import scipy
import scipy.spatial
import scipy.optimize

def constraint(alen, blen):
    return numpy.eye(alen)[..., numpy.newaxis] * numpy.ones(blen)[numpy.newaxis, numpy.newaxis]

def emd(a, b, a_weights=None, b_weights=None):
    if a_weights is None:
        a_weights = numpy.ones(len(a))
    if b_weights is None:
        b_weights = numpy.ones(len(b))
    a_weights /= a_weights.sum()
    b_weights /= b_weights.sum()
    dist = scipy.spatial.distance.cdist(a, b).reshape(-1)
    alen = len(a)
    blen = len(b)
    A_eq = numpy.vstack((constraint(alen, blen), constraint(blen, alen).transpose((0, 2, 1))))
    A_eq = A_eq.reshape((alen + blen, alen * blen))
    b_eq = numpy.concatenate((a_weights, b_weights))
    res = scipy.optimize.linprog(dist, A_eq=A_eq, b_eq=b_eq)
    return res.x.dot(dist)