Finite Discrete Measure

Project.toml

@@ -32,6 +32,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tulip = "6dd1b50a-3aae-11e9-10b5-ef983d2400fa"
+PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 
 [targets]
 test = ["ForwardDiff", "Pkg", "PythonOT", "Random", "SafeTestsets", "Test", "Tulip"]

docs/src/index.md

@@ -36,3 +36,8 @@ sinkhorn_unbalanced2
 ```@docs
 quadreg
 ```
+
+## Auxiliary functions
+```@docs
+FiniteDiscreteMeasure
+```

src/OptimalTransport.jl

@@ -4,6 +4,7 @@
 
 module OptimalTransport
 
+using LinearAlgebra: AbstractMatrix
 using Distances
 using LinearAlgebra
 using IterativeSolvers, SparseArrays
@@ -17,8 +18,11 @@ export sinkhorn, sinkhorn2
 export emd, emd2
 export sinkhorn_stabilized, sinkhorn_stabilized_epsscaling, sinkhorn_barycenter
 export sinkhorn_unbalanced, sinkhorn_unbalanced2
+export sinkhorn_divergence
 export quadreg
 export ot_cost, ot_plan, wasserstein, squared2wasserstein
+export FiniteDiscreteMeasure
+export cost_matrix
 
 const MOI = MathOptInterface
 

src/utils.jl

@@ -54,3 +54,53 @@ function checkbalanced(x::AbstractVecOrMat, y::AbstractVecOrMat)
         throw(ArgumentError("source and target marginals are not balanced"))
     return nothing
 end
+
+struct FiniteDiscreteMeasure{X<:AbstractArray,P<:AbstractVector}
+    support::X
+    p::P
+
+    function FiniteDiscreteMeasure{X,P}(support::X, p::P) where {X,P}
+        size(support, 1) == length(p) ||
+            error("number of rows of `support` and `p` must be equal")
+        sum(p) ≈ 1 || error("`p` must sum to 1")
+        all(p .>= 0) || error("values of `p` must be greater of equal than 0")
+        return new{X,P}(support, p)
+    end
+end
+
+"""
+    FiniteDiscreteMeasure(support::AbstractArray, p::AbstractVector)
+Construct a finite discrete probability measure with support `support` and corresponding weights `p`.
+"""
+function FiniteDiscreteMeasure(support::AbstractArray, p::AbstractVector)
+    P = sum(p)
+    if size(support, 2) == 1
+        return if P ≈ 1
+            DiscreteNonParametric(vec(support), p)
+        else
+            DiscreteNonParametric(vec(support), p ./ P)
+        end
+    else
+        return if P ≈ 1
+            FiniteDiscreteMeasure{typeof(support),typeof(p)}(support, p)
+        else
+            FiniteDiscreteMeasure{typeof(support),typeof(p)}(support, p ./ P)
+        end
+    end
+end
+
+"""
+    FiniteDiscreteMeasure(support::AbstractArray)
+Construct a finite discrete probability measure with support `support` and equal probability for each point.
+"""
+function FiniteDiscreteMeasure(support::AbstractArray)
+    p = ones(size(support)[1]) ./ size(support)[1]
+    if size(support, 2) == 1
+        return DiscreteNonParametric(vec(support), p)
+    else
+        return FiniteDiscreteMeasure{typeof(support),typeof(p)}(support, p)
+    end
+end
+
+Distributions.support(d::FiniteDiscreteMeasure) = d.support
+Distributions.probs(d::FiniteDiscreteMeasure) = d.p

test/entropic.jl

@@ -4,6 +4,7 @@ using Distances
 using ForwardDiff
 using LogExpFunctions
 using PythonOT: PythonOT
+using Distributions
 
 using Random
 using Test

test/runtests.jl

@@ -1,3 +1,4 @@
+using LinearAlgebra: symmetric
 using OptimalTransport
 using Pkg: Pkg
 using SafeTestsets

test/utils.jl

@@ -3,6 +3,7 @@ using OptimalTransport
 using LinearAlgebra
 using Random
 using Test
+using Distributions
 
 Random.seed!(100)
 
@@ -95,4 +96,53 @@ Random.seed!(100)
             x2, y2 .* hcat(rand(), ones(1, size(y2, 2) - 1))
         )
     end
+
+    @testset "finitediscretemeasure.jl" begin
+        @testset "Univariate Finite Discrete Measure" begin
+            n = 100
+            μsupp = rand(n)
+            νsupp = rand(n, 1)
+            μ = FiniteDiscreteMeasure(μsupp)
+            ν = FiniteDiscreteMeasure(νsupp, rand(n))
+            # check if it assigns equal probabilities to all entries
+            @test μ.p ≈ ones(n) ./ n
+            @test probs(μ) ≈ ones(n) ./ n
+            # check if it probabilities sum to 1
+            @test sum(ν.p) ≈ 1
+            @test sum(probs(ν)) ≈ 1
+            # check if probabilities are all positive (non-negative)
+            @test all(ν.p .>= 0)
+            @test all(probs(ν) .>= 0)
+            # check if it assigns to DiscreteNonParametric when Vector/Matrix is 1D
+            @test typeof(μ) <: DiscreteNonParametric
+            @test typeof(ν) <: DiscreteNonParametric
+            # check if support is correctly assinged
+            @test sort(μsupp) == μ.support
+            @test sort(μsupp) == support(μ)
+            @test sort(vec(νsupp)) == ν.support
+            @test sort(vec(νsupp)) == support(ν)
+        end
+        @testset "Multivariate Finite Discrete Measure" begin
+            n = 10
+            m = 3
+            μsupp = rand(n, m)
+            νsupp = rand(n, m)
+            μ = FiniteDiscreteMeasure(μsupp)
+            ν = FiniteDiscreteMeasure(νsupp, rand(n))
+            # check if it assigns equal probabilities to all entries
+            @test μ.p ≈ ones(n) ./ n
+            @test probs(μ) ≈ ones(n) ./ n
+            # check if it probabilities sum to 1
+            @test sum(ν.p) ≈ 1
+            @test sum(probs(ν)) ≈ 1
+            # check if probabilities are all positive (non-negative)
+            @test all(ν.p .>= 0)
+            @test all(probs(ν) .>= 0)
+            # check if support is correctly assinged
+            @test μsupp == μ.support
+            @test μsupp == support(μ)
+            @test νsupp == ν.support
+            @test νsupp == support(ν)
+        end
+    end
 end