Compute semidiscrete wasserstein according to https://papers.nips.cc/paper/6566-stochastic-optimization-for-large-scale-optimal-transport.pdf

semidiscrete.jl

using StatsFuns

function h(x::AbstractVector, v::AbstractVector, y::AbstractVector, ν::AbstractVector, ϵ::Real, c)
    dot(v, ν) - ϵ * logsumexp((v - c.(Ref(x), y)) ./ ϵ .+ log.(ν)) - ϵ
end

function h(x::AbstractVector, v::AbstractVector, y::AbstractVector, ν::AbstractVector, ϵ::Int, c)
    ϵ == 0 || error("ϵ has to be 0")
    dot(v,ν) + mininum(c.(Ref(x), y) .- v)
end

function optim_v(μ, y::AbstractVector, ν::AbstractVector, η::Real, N::Int, ϵ::Real, c)
    v = zero(ν); ṽ = zero(ν)
    for k in 1:N
        xₖ = rand(μ)
        ṽ .+= η /√(k) * gradient(ν->h(xₖ, ṽ, y, ν, ϵ, c), ν)[1]
        v = ṽ ./ k + (k - 1) / k * v
    end
    return v
end

function wasserstein_semidiscrete(μ, y, ν, ϵ, c=(x,y)->norm(x-y), η::Real = 0.1, N::Int = 100, N_MC::Int=2000)
    v = optim_v(μ, y, ν, η, N, ϵ, c)
    return mean(x->h(x, v, y, ν, ϵ, c), eachcol(rand(μ, N_MC)))
end