Skip to content

Update and extend documentation #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 22, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/src/api.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# API

## Exact optimal transport (Kantorovich) problem
## Exact optimal transport

```@docs
emd
Expand Down
101 changes: 64 additions & 37 deletions src/lib.jl
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
"""
emd(μ, ν, C; kwargs...)

Compute the optimal transport map for the Monge-Kantorovich problem with source and target
Compute the optimal transport plan for the Monge-Kantorovich problem with source and target
marginals `μ` and `ν` and cost matrix `C` of size `(length(μ), length(ν))`.

The optimal transport map `γ` is of the same size as `C` and solves
The optimal transport plan `γ` is of the same size as `C` and solves
```math
\\inf_{\\gamma \\in \\Pi(\\mu, \\nu)} \\langle \\gamma, C \\rangle.
```
Expand All @@ -20,9 +20,7 @@ julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0;
2.0 0.0;
0.5 1.5];
julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> emd(μ, ν, C)
3×2 Matrix{Float64}:
Expand Down Expand Up @@ -59,9 +57,7 @@ julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0;
2.0 0.0;
0.5 1.5];
julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> emd2(μ, ν, C)
0.95
Expand All @@ -76,11 +72,11 @@ end
"""
sinkhorn(μ, ν, C, ε; kwargs...)

Compute the optimal transport map for the entropic regularization optimal transport problem
Compute the optimal transport plan for the entropic regularization optimal transport problem
with source and target marginals `μ` and `ν`, cost matrix `C` of size
`(length(μ), length(ν))`, and entropic regularization parameter `ε`.

The optimal transport map `γ` is of the same size as `C` and solves
The optimal transport plan `γ` is of the same size as `C` and solves
```math
\\inf_{\\gamma \\in \\Pi(\\mu, \\nu)} \\langle \\gamma, C \\rangle
+ \\varepsilon \\Omega(\\gamma),
Expand All @@ -94,14 +90,12 @@ Transport package. Keyword arguments are listed in the documentation of the Pyth

# Examples

```jldoctest
```jldoctest sinkhorn
julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0;
2.0 0.0;
0.5 1.5];
julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> sinkhorn(μ, ν, C, 0.01)
3×2 Matrix{Float64}:
Expand All @@ -110,6 +104,18 @@ julia> sinkhorn(μ, ν, C, 0.01)
0.0 0.3
```

It is possible to provide multiple target marginals as columns of a matrix. In this case the
optimal transport costs are returned:

```jldoctest sinkhorn
julia> ν = [0.0 0.5; 1.0 0.5];

julia> round.(sinkhorn(μ, ν, C, 0.01); sigdigits=6)
2-element Vector{Float64}:
0.95
0.45
```

See also: [`sinkhorn2`](@ref)
"""
function sinkhorn(μ, ν, C, ε; kwargs...)
Expand Down Expand Up @@ -137,20 +143,29 @@ Transport package. Keyword arguments are listed in the documentation of the Pyth

# Examples

```jldoctest
```jldoctest sinkhorn2
julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0;
2.0 0.0;
0.5 1.5];
julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> round.(sinkhorn2(μ, ν, C, 0.01); sigdigits=6)
1-element Vector{Float64}:
0.95
```

It is possible to provide multiple target marginals as columns of a matrix.

```jldoctest sinkhorn2
julia> ν = [0.0 0.5; 1.0 0.5];

julia> round.(sinkhorn2(μ, ν, C, 0.01); sigdigits=6)
2-element Vector{Float64}:
0.95
0.45
```

See also: [`sinkhorn`](@ref)
"""
function sinkhorn2(μ, ν, C, ε; kwargs...)
Expand All @@ -160,12 +175,12 @@ end
"""
sinkhorn_unbalanced(μ, ν, C, ε, λ; kwargs...)

Compute the optimal transport map for the unbalanced entropic regularization optimal
Compute the optimal transport plan for the unbalanced entropic regularization optimal
transport problem with source and target marginals `μ` and `ν`, cost matrix `C` of size
`(length(μ), length(ν))`, entropic regularization parameter `ε`, and marginal relaxation
term `λ`.

The optimal transport map `γ` is of the same size as `C` and solves
The optimal transport plan `γ` is of the same size as `C` and solves
```math
\\inf_{\\gamma} \\langle \\gamma, C \\rangle
+ \\varepsilon \\Omega(\\gamma)
Expand All @@ -182,14 +197,12 @@ Python function.

# Examples

```jldoctest
```jldoctest sinkhorn_unbalanced
julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0;
2.0 0.0;
0.5 1.5];
julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> sinkhorn_unbalanced(μ, ν, C, 0.01, 1_000)
3×2 Matrix{Float64}:
Expand All @@ -198,6 +211,18 @@ julia> sinkhorn_unbalanced(μ, ν, C, 0.01, 1_000)
0.0 0.29983
```

It is possible to provide multiple target marginals as columns of a matrix. In this case the
optimal transport costs are returned:

```jldoctest sinkhorn_unbalanced
julia> ν = [0.0 0.5; 1.0 0.5];

julia> round.(sinkhorn_unbalanced(μ, ν, C, 0.01, 1_000); sigdigits=6)
2-element Vector{Float64}:
0.949709
0.449411
```

See also: [`sinkhorn_unbalanced2`](@ref)
"""
function sinkhorn_unbalanced(μ, ν, C, ε, λ; kwargs...)
Expand Down Expand Up @@ -231,20 +256,29 @@ Python function.

# Examples

```jldoctest
```jldoctest sinkhorn_unbalanced2
julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0;
2.0 0.0;
0.5 1.5];
julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> round.(sinkhorn_unbalanced2(μ, ν, C, 0.01, 1_000); sigdigits=6)
1-element Vector{Float64}:
0.949709
```

It is possible to provide multiple target marginals as columns of a matrix:

```jldoctest sinkhorn_unbalanced2
julia> ν = [0.0 0.5; 1.0 0.5];

julia> round.(sinkhorn_unbalanced2(μ, ν, C, 0.01, 1_000); sigdigits=6)
2-element Vector{Float64}:
0.949709
0.449411
```

See also: [`sinkhorn_unbalanced`](@ref)
"""
function sinkhorn_unbalanced2(μ, ν, C, ε, λ; kwargs...)
Expand All @@ -263,7 +297,7 @@ The Wasserstein barycenter is a histogram and solves
```math
\\inf_{a} \\sum_{i} W_{\\varepsilon,C}(a, a_i),
```
where the histograms ``a_i`` are columns of matrix `A` and ``W_{\\varepsilon,C}(a, a_i)}``
where the histograms ``a_i`` are columns of matrix `A` and ``W_{\\varepsilon,C}(a, a_i)``
is the optimal transport cost for the entropically regularized optimal transport problem
with marginals ``a`` and ``a_i``, cost matrix ``C``, and entropic regularization parameter
``\\varepsilon``. Optionally, weights of the histograms ``a_i`` can be provided with the
Expand All @@ -287,11 +321,4 @@ julia> isapprox(sum(barycenter(A, C, 0.01; method="sinkhorn_stabilized")), 1; at
true
```
"""
function barycenter(A, C, ε; kwargs...)
return pot.barycenter(
PyCall.PyReverseDims(permutedims(A)),
PyCall.PyReverseDims(permutedims(C)),
ε;
kwargs...,
)
end
barycenter(A, C, ε; kwargs...) = pot.barycenter(A, C, ε; kwargs...)
4 changes: 1 addition & 3 deletions src/smooth.jl
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,7 @@ julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0;
2.0 0.0;
0.5 1.5];
julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> smooth_ot_dual(μ, ν, C, 0.01)
3×2 Matrix{Float64}:
Expand Down