diff --git a/Project.toml b/Project.toml index d8894ac..edfce28 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "PythonOT" uuid = "3c485715-4278-42b2-9b5f-8f00e43c12ef" authors = ["David Widmann"] -version = "0.1.6" +version = "0.1.7" [deps] PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0" diff --git a/docs/src/api.md b/docs/src/api.md index 00409b1..f9d7320 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -37,3 +37,9 @@ sinkhorn_unbalanced2 barycenter_unbalanced mm_unbalanced ``` + +## Partial optimal transport + +```@docs +entropic_partial_wasserstein +``` diff --git a/src/PythonOT.jl b/src/PythonOT.jl index d801f3d..4a1c71b 100644 --- a/src/PythonOT.jl +++ b/src/PythonOT.jl @@ -13,7 +13,8 @@ export emd, sinkhorn_unbalanced, sinkhorn_unbalanced2, empirical_sinkhorn_divergence, - mm_unbalanced + mm_unbalanced, + entropic_partial_wasserstein const pot = PyCall.PyNULL() diff --git a/src/lib.jl b/src/lib.jl index 07bb9af..65460cb 100644 --- a/src/lib.jl +++ b/src/lib.jl @@ -306,29 +306,18 @@ Python function. # Examples ```jldoctest sinkhorn_unbalanced -julia> μ = [0.5, 0.2, 0.3]; +julia> μ = [0.5, 0.5]; -julia> ν = [0.0, 1.0]; +julia> ν = [0.5, 0.5]; -julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5]; +julia> C = [0.0 1.0; 1.0 0.0]; -julia> round.(sinkhorn_unbalanced(μ, ν, C, 0.01, 1_000); sigdigits=4) -3×2 Matrix{Float64}: - 0.0 0.5 - 0.0 0.2002 - 0.0 0.2998 +julia> round.(sinkhorn_unbalanced(μ, ν, C, 1, 1); sigdigits=7) +2×2 Matrix{Float64}: + 0.322054 0.118477 + 0.118477 0.322054 ``` -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=4) -2-element Vector{Float64}: - 0.9497 - 0.4494 ``` See also: [`sinkhorn_unbalanced2`](@ref) @@ -365,25 +354,14 @@ Python function. # Examples ```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> μ = [0.5, 0.1]; -julia> round.(sinkhorn_unbalanced2(μ, ν, C, 0.01, 1_000); sigdigits=4) -0.9497 -``` - -It is possible to provide multiple target marginals as columns of a matrix: +julia> ν = [0.5, 0.5]; -```jldoctest sinkhorn_unbalanced2 -julia> ν = [0.0 0.5; 1.0 0.5]; +julia> C = [0.0 1.0; 1.0 0.0]; -julia> round.(sinkhorn_unbalanced2(μ, ν, C, 0.01, 1_000); sigdigits=4) -2-element Vector{Float64}: - 0.9497 - 0.4494 +julia> round.(sinkhorn_unbalanced2(μ, ν, C, 1., 1.); sigdigits=8) +0.19600125 ``` See also: [`sinkhorn_unbalanced`](@ref) @@ -566,3 +544,45 @@ julia> round.(mm_unbalanced(a, b, M, 5, div="l2"), digits=2) function mm_unbalanced(a, b, M, reg_m; kwargs...) return pot.unbalanced.mm_unbalanced(a, b, M, reg_m; kwargs...) end + + +""" + entropic_partial_wasserstein(a, b, M, reg; kwargs...) + +Solves the partial optimal transport problem and returns the OT plan +The function considers the following problem: + +```math +\\gamma = \\mathop{\\arg \\min}_\\gamma \\quad \\langle \\gamma, +\\mathbf{M} \\rangle_F + \\mathrm{reg} \\cdot\\Omega(\\gamma) +``` + +- `a` and `b` are the sample weights +- `M` is the metric cost matrix +- `reg` is a regularization term > 0 + +This function is a wrapper of the function +[`entropic_partial_wasserstein`](https://pythonot.github.io/gen_modules/ot.partial.html#ot.partial.entropic_partial_wasserstein) in the +Python Optimal Transport package. Keyword arguments are listed in the documentation of the +Python function. + + +# Examples + +```jldoctest +julia> a = [.1, .2]; + +julia> b = [.1, .1]; + +julia> M = [0. 1.; 2. 3.]; + +julia> round.(entropic_partial_wasserstein(a, b, M, 1, m=0.1), digits=2) +2×2 Matrix{Float64}: + 0.06 0.02 + 0.01 0.0 +``` + +""" +function entropic_partial_wasserstein(a, b, M, reg; kwargs...) + return pot.partial.entropic_partial_wasserstein(a, b, M, reg; kwargs...) +end