Optimal Transport for Multivariate Gaussians (#85)

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

.gitignore

@@ -32,3 +32,4 @@ docs/src/examples/
 
 # Files generated by Jupyter Notebooks
 *.ipynb_checkpoints
+*.ipynb

Project.toml

@@ -1,7 +1,7 @@
 name = "OptimalTransport"
 uuid = "7e02d93a-ae51-4f58-b602-d97af76e3b33"
 authors = ["zsteve <stephenz@student.unimelb.edu.au>"]
-version = "0.3.8"
+version = "0.3.9"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
@@ -10,6 +10,7 @@ IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
@@ -20,6 +21,7 @@ Distributions = "0.25"
 IterativeSolvers = "0.8.4, 0.9"
 LogExpFunctions = "0.2"
 MathOptInterface = "0.9"
+PDMats = "0.11"
 QuadGK = "2"
 StatsBase = "0.33.8"
 julia = "1"
@@ -32,6 +34,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"
+HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
 
 [targets]
-test = ["ForwardDiff", "Pkg", "PythonOT", "Random", "SafeTestsets", "Test", "Tulip"]
+test = ["ForwardDiff", "Pkg", "PythonOT", "Random", "SafeTestsets", "Test", "Tulip", "HCubature"]

src/OptimalTransport.jl

@@ -10,6 +10,7 @@ using IterativeSolvers, SparseArrays
 using LogExpFunctions: LogExpFunctions
 using MathOptInterface
 using Distributions
+using PDMats
 using QuadGK
 using StatsBase: StatsBase
 
@@ -22,6 +23,7 @@ export ot_cost, ot_plan, wasserstein, squared2wasserstein
 
 const MOI = MathOptInterface
 
+include("distances/bures.jl")
 include("utils.jl")
 include("exact.jl")
 include("wasserstein.jl")

src/distances/bures.jl

@@ -0,0 +1,74 @@
+# Code from @devmotion
+# https://github.com/devmotion/\
+# CalibrationErrorsDistributions.jl/blob/main/src/distances/bures.jl
+
+"""
+ tr_sqrt(A::AbstractMatrix)
+
+Compute ``\\operatorname{tr}\\big(A^{1/2}\\big)``.
+"""
+tr_sqrt(A::AbstractMatrix) = LinearAlgebra.tr(sqrt(A))
+tr_sqrt(A::PDMats.PDMat) = tr_sqrt(A.mat)
+tr_sqrt(A::PDMats.PDiagMat) = sum(sqrt, A.diag)
+tr_sqrt(A::PDMats.ScalMat) = A.dim * sqrt(A.value)
+
+"""
+ _gaussian_ot_A(A::AbstractMatrix, B::AbstractMatrix)
+
+Compute
+```math
+A^{1/2} B A^{1/2}.
+```
+"""
+function _gaussian_ot_A(A::AbstractMatrix, B::AbstractMatrix)
+ sqrt_A = sqrt(A)
+ return sqrt_A * B * sqrt_A
+end
+function _gaussian_ot_A(A::PDMats.PDiagMat, B::AbstractMatrix)
+ return sqrt.(A.diag) .* B .* sqrt.(A.diag')
+end
+function _gaussian_ot_A(A::StridedMatrix, B::PDMats.PDMat)
+ return PDMats.X_A_Xt(B, sqrt(A))
+end
+_gaussian_ot_A(A::PDMats.PDMat, B::PDMats.PDMat) = _gaussian_ot_A(A.mat, B)
+_gaussian_ot_A(A::AbstractMatrix, B::PDMats.PDiagMat) = _gaussian_ot_A(B, A)
+_gaussian_ot_A(A::PDMats.PDMat, B::StridedMatrix) = _gaussian_ot_A(B, A)
+
+"""
+ sqbures(A::AbstractMatrix, B::AbstractMatrix)
+
+Compute the squared Bures metric
+```math
+\\operatorname{tr}(A) + \\operatorname{tr}(B)
+- \\operatorname{tr}\\Big({\\big(A^{1/2} B A^{1/2}\\big)}^{1/2}\\Big).
+```
+"""
+function sqbures(A::AbstractMatrix, B::AbstractMatrix)
+ return LinearAlgebra.tr(A) + LinearAlgebra.tr(B) - 2 * tr_sqrt(_gaussian_ot_A(A, B))
+end
+
+# diagonal matrix
+function sqbures(A::PDMats.PDiagMat, B::PDMats.PDiagMat)
+ if !(A.dim == B.dim)
+ throw(ArgumentError("matrices must have the same dimensions."))
+ end
+ return sum(zip(A.diag, B.diag)) do (x, y)
+ abs2(sqrt(x) - sqrt(y))
+ end
+end
+
+# scaled identity matrix
+function sqbures(A::PDMats.ScalMat, B::AbstractMatrix)
+ return LinearAlgebra.tr(A) + LinearAlgebra.tr(B) - 2 * sqrt(A.value) * tr_sqrt(B)
+end
+sqbures(A::AbstractMatrix, B::PDMats.ScalMat) = sqbures(B, A)
+sqbures(A::PDMats.ScalMat, B::PDMats.ScalMat) = A.dim * abs2(sqrt(A.value) - sqrt(B.value))
+
+# combinations
+function sqbures(A::PDMats.PDiagMat, B::PDMats.ScalMat)
+ sqrt_B = sqrt(B.value)
+ return sum(A.diag) do x
+ abs2(sqrt(x) - sqrt_B)
+ end
+end
+sqbures(A::PDMats.ScalMat, B::PDMats.PDiagMat) = sqbures(B, A)

src/exact.jl

@@ -345,3 +345,82 @@ end
 function _ot_cost(c, μ::DiscreteNonParametric, ν::DiscreteNonParametric, plan)
  return dot(plan, StatsBase.pairwise(c, support(μ), support(ν)))
 end
+
+################
+# OT Gaussians
+################
+
+"""
+ ot_cost(::SqEuclidean, μ::MvNormal, ν::MvNormal)
+
+Compute the squared 2-Wasserstein distance between normal distributions `μ` and `ν` as
+source and target marginals.
+
+In this setting, the optimal transport cost can be computed as
+```math
+W_2^2(\\mu, \\nu) = \\|m_\\mu - m_\\nu \\|^2 + \\mathcal{B}(\\Sigma_\\mu, \\Sigma_\\nu)^2,
+```
+where ``\\mu = \\mathcal{N}(m_\\mu, \\Sigma_\\mu)``,
+``\\nu = \\mathcal{N}(m_\\nu, \\Sigma_\\nu)``, and ``\\mathcal{B}`` is the Bures metric.
+
+See also: [`ot_plan`](@ref), [`emd2`](@ref)
+"""
+function ot_cost(::SqEuclidean, μ::MvNormal, ν::MvNormal)
+ return sqeuclidean(μ.μ, ν.μ) + sqbures(μ.Σ, ν.Σ)
+end
+
+"""
+ ot_cost(::SqEuclidean, μ::Normal, ν::Normal)
+
+Compute the squared 2-Wasserstein distance between univariate normal distributions `μ` and
+`ν` as source and target marginals.
+
+See also: [`ot_plan`](@ref), [`emd2`](@ref)
+"""
+function ot_cost(::SqEuclidean, μ::Normal, ν::Normal)
+ return (μ.μ - ν.μ)^2 + (μ.σ - ν.σ)^2
+end
+
+"""
+ ot_plan(::SqEuclidean, μ::MvNormal, ν::MvNormal)
+
+Compute the optimal transport plan for the Monge-Kantorovich problem with multivariate
+normal distributions `μ` and `ν` as source and target marginals and cost function
+``c(x, y) = \\|x - y\\|_2^2``.
+
+In this setting, for ``\\mu = \\mathcal{N}(m_\\mu, \\Sigma_\\mu)`` and
+``\\nu = \\mathcal{N}(m_\\nu, \\Sigma_\\nu)``, the optimal transport plan is the Monge
+map
+```math
+T \\colon x \\mapsto m_\\nu
++ \\Sigma_\\mu^{-1/2}
+{\\big(\\Sigma_\\mu^{1/2} \\Sigma_\\nu \\Sigma_\\mu^{1/2}\\big)}^{1/2}\\Sigma_\\mu^{-1/2}
+(x - m_\\mu).
+
+See also: [`ot_cost`](@ref), [`emd`](@ref)
+"""
+function ot_plan(::SqEuclidean, μ::MvNormal, ν::MvNormal)
+ Σμsqrt = μ.Σ^(-1 / 2)
+ A = Σμsqrt * sqrt(_gaussian_ot_A(μ.Σ, ν.Σ)) * Σμsqrt
+ mμ = μ.μ
+ mν = ν.μ
+ T(x) = mν + A * (x - mμ)
+ return T
+end
+
+"""
+ ot_plan(::SqEuclidean, μ::Normal, ν::Normal)
+
+Compute the optimal transport plan for the Monge-Kantorovich problem with
+normal distributions `μ` and `ν` as source and target marginals and cost function
+``c(x, y) = \\|x - y\\|_2^2``.
+
+See also: [`ot_cost`](@ref), [`emd`](@ref)
+"""
+function ot_plan(::SqEuclidean, μ::Normal, ν::Normal)
+ mμ = μ.μ
+ mν = ν.μ
+ a = ν.σ / μ.σ
+ T(x) = mν + a * (x - mμ)
+ return T
+end

test/bures.jl

@@ -0,0 +1,27 @@
+# Code from @devmotion
+# https://github.com/devmotion/\
+# CalibrationErrorsDistributions.jl/blob/main/src/distances/bures.jl
+using OptimalTransport
+
+using LinearAlgebra
+using Random
+using PDMats
+
+@testset "bures.jl" begin
+ function _sqbures(A, B)
+ sqrt_A = sqrt(A)
+ return tr(A) + tr(B) - 2 * tr(sqrt(sqrt_A * B * sqrt_A'))
+ end
+
+ function rand_matrices(n)
+ A = randn(n, n)
+ B = A' * A + I
+ return B, PDMat(B), PDiagMat(diag(B)), ScalMat(n, B[1])
+ end
+
+ for (x, y) in Iterators.product(rand_matrices(10), rand_matrices(10))
+ xfull = Matrix(x)
+ yfull = Matrix(y)
+ @test OptimalTransport.sqbures(x, y) ≈ _sqbures(xfull, yfull)
+ end
+end

test/exact.jl

@@ -5,6 +5,7 @@ using PythonOT: PythonOT
 using Tulip
 using MathOptInterface
 using Distributions
+using HCubature
 
 using LinearAlgebra
 using Random
@@ -164,4 +165,55 @@ Random.seed!(100)
  @test c2 ≈ c
  end
  end
+
+ @testset "Multivariate Gaussians" begin
+ @testset "translation with constant covariance" begin
+ m = randn(100)
+ τ = rand(100)
+ Σ = Matrix(Hermitian(rand(100, 100) + 100I))
+ μ = MvNormal(m, Σ)
+ ν = MvNormal(m .+ τ, Σ)
+ @test ot_cost(SqEuclidean(), μ, ν) ≈ norm(τ)^2
+
+ x = rand(100, 10)
+ T = ot_plan(SqEuclidean(), μ, ν)
+ @test pdf(ν, mapslices(T, x; dims=1)) ≈ pdf(μ, x)
+ end
+
+ @testset "comparison to grid approximation" begin
+ μ = MvNormal([0, 0], [1 0; 0 2])
+ ν = MvNormal([10, 10], [2 0; 0 1])
+ # Constructing circular grid approximation
+ # Angular grid step
+ θ = collect(0:0.2:(2π))
+ θx = cos.(θ)
+ θy = sin.(θ)
+ # Radius grid step
+ δ = collect(0:0.2:1)
+ μsupp = [0.0 0.0]
+ νsupp = [10.0 10.0]
+ for i in δ[2:end]
+ a = [θx .* i θy .* i * 2]
+ b = [θx .* i * 2 θy .* i] .+ [10 10]
+ μsupp = vcat(μsupp, a)
+ νsupp = vcat(νsupp, b)
+ end
+
+ # Create discretized distribution
+ μprobs = pdf(μ, μsupp')
+ μprobs = μprobs ./ sum(μprobs)
+ νprobs = pdf(ν, νsupp')
+ νprobs = νprobs ./ sum(νprobs)
+ C = pairwise(SqEuclidean(), μsupp', νsupp')
+ @test emd2(μprobs, νprobs, C, Tulip.Optimizer()) ≈ ot_cost(SqEuclidean(), μ, ν) rtol =
+ 1e-3
+
+ # Use hcubature integration to perform ``\\int c(x,T(x)) d\\mu``
+ T = ot_plan(SqEuclidean(), μ, ν)
+ c_hcubature, _ = hcubature([-10, -10], [10, 10]) do x
+ return sqeuclidean(x, T(x)) * pdf(μ, x)
+ end
+ @test ot_cost(SqEuclidean(), μ, ν) ≈ c_hcubature rtol = 1e-3
+ end
+ end
 end

test/runtests.jl

@@ -26,6 +26,9 @@ const GROUP = get(ENV, "GROUP", "All")
  @safetestset "Wasserstein distance" begin
  include("wasserstein.jl")
  end
+ @safetestset "Bures distance" begin
+ include("bures.jl")
+ end
  end
 
  # CUDA requires Julia >= 1.6