diff --git a/preliminary_OT_tests/OT_chained_multilayer.jl b/preliminary_OT_tests/OT_chained_multilayer.jl deleted file mode 100644 index 607c7d9..0000000 --- a/preliminary_OT_tests/OT_chained_multilayer.jl +++ /dev/null @@ -1,126 +0,0 @@ -using SequentialMeasureTransport -using ApproxFun -import SequentialMeasureTransport as SMT -using LinearAlgebra -using Distributions -using Plots -using MosekTools - -include("helpers.jl") - - -## define a cost function -_d = 1 -c(x, y) = norm(x - y)^2 -c(x) = norm(x[1:_d] - x[_d+1:end])^2 - -## define the marginals -p = MixtureModel(Normal[ - Normal(0.3, 0.1), - Normal(0.7, 0.1) -]) -q = Normal(0.5, 0.2) - -rng = 0.0:0.001:1.0 -@time res = compute_Sinkhorn_Wasserstein_barycenter(rng, [x->pdf(p, x)[1], x->pdf(q, x)[1]], - [0.5, 0.5], c, 0.05) - -ϵ = 0.01 -plot(rng, compute_Sinkhorn_Wasserstein_barycenter(rng, [x->pdf(p, x)[1], x->pdf(q, x)[1]], - [1.0, 0.0], c, ϵ)) - -plot!(rng, compute_Sinkhorn_Wasserstein_barycenter(rng, [x->pdf(p, x)[1], x->pdf(q, x)[1]], - [0.8, 0.2], c, ϵ)) - -plot!(rng, compute_Sinkhorn_Wasserstein_barycenter(rng, [x->pdf(p, x)[1], x->pdf(q, x)[1]], - [0.6, 0.4], c, ϵ)) - -plot!(rng, compute_Sinkhorn_Wasserstein_barycenter(rng, [x->pdf(p, x)[1], x->pdf(q, x)[1]], - [0.4, 0.6], c, ϵ)) - -plot!(rng, compute_Sinkhorn_Wasserstein_barycenter(rng, [x->pdf(p, x)[1], x->pdf(q, x)[1]], - [0.2, 0.8], c, ϵ)) - -plot!(rng, compute_Sinkhorn_Wasserstein_barycenter(rng, [x->pdf(p, x)[1], x->pdf(q, x)[1]], - [0.0, 1.0], c, ϵ)) -### First, we work on the bounded domain [0, 1]^2 - -ϵ = 0.5 - -# model = SMT.PSDModel(Legendre(0.0..1.0)^2, :downward_closed, 4) -# _smp = SMT.Sampler(model) -smp_list = SMT.ConditionalMapping{2, 0, Float64}[] -# λ_marg_list = 1e12 * ones(20) - -for k=1:2 - XY = rand(2, 2000) - X = rand(1, 1000) - Y = rand(1, 1000) - - model = SMT.PSDModel(Legendre(0.0..1.0)^2, :downward_closed, 4) - - prec = if k > 1 - SMT.CondSampler(smp_list[1:k-1], nothing) - else - nothing - end - use_prec = if k > 1 - true - else - false - end - - SMT.Statistics.entropic_OT!(model, c, x->pdf(p, x)[1], x->pdf(q, x)[1], - ϵ, eachcol(XY), trace=true, - X=eachcol(X), Y=eachcol(Y), - optimizer=Mosek.Optimizer, - set_start_value=false, - preconditioner=prec, - use_preconditioner_cost=use_prec, - normalization=true) - - normalize!(model) - push!(smp_list, SMT.Sampler(model)) -end - -rng = 0.0:0.01:1.0 -_smp = SMT.CondSampler([smp_list[1:2]...], nothing) -contourf(rng, rng, (x, y) -> pdf(_smp, [x, y]), alpha=1.0, label="transported model") - -M_sink = compute_Sinkhorn(rng, x->pdf(p, x)[1], - x->pdf(q, x)[1], c, 0.05) -contourf(rng, rng, M_sink') - -contour( - rng, rng, (x, y) -> pdf(_smp, [x, y]), levels=5 -) -contour!(rng, rng, M_sink', levels=5) - -plot(rng, x->left_pdf(_smp, x)) -plot!(rng, x->pdf(p, x)[1]) - -plot(rng, x->right_pdf(_smp, x)) -plot!(rng, x->pdf(q, x)[1]) - -compute_Sinkhorn_distance(c, _smp, N=3000) -compute_Sinkhorn_distance(c, M_sink, rng) - -distance_list = Float64[] -ϵ_list = Float64[] -rng = 0:0.005:1 -for k=1:10 - ϵ = 0.5^(k) - push!(ϵ_list, ϵ) - M_sink = compute_Sinkhorn(rng, x->pdf(p, x)[1], - x->pdf(q, x)[1], c, ϵ) - push!(distance_list, compute_Sinkhorn_distance(c, M_sink, rng)) -end -plot(ϵ_list, distance_list, m=:o, label="Sinkhorn distance", xscale=:log10) - - -distance_chained_list = Float64[] -for k=1:length(smp_list) - _smp = SMT.CondSampler(smp_list[1:k], nothing) - push!(distance_chained_list, compute_Sinkhorn_distance(c, _smp, N=5000)) -end -plot(1:length(smp_list), distance_chained_list, m=:o, label="Chained Sinkhorn distance") \ No newline at end of file diff --git a/preliminary_OT_tests/OT_multilayer.jl b/preliminary_OT_tests/OT_multilayer.jl deleted file mode 100644 index 5eacfcc..0000000 --- a/preliminary_OT_tests/OT_multilayer.jl +++ /dev/null @@ -1,97 +0,0 @@ -using SequentialMeasureTransport -using ApproxFun -import SequentialMeasureTransport as SMT -using LinearAlgebra -using Distributions -using Plots -using MosekTools - -include("helpers.jl") - - -## define a cost function -c(x) = (x[1] - x[2])^2 - -## define the marginals -p = MixtureModel(Normal[ - Normal(0.3, 0.1), - Normal(0.7, 0.1) -]) -q = Normal(0.5, 0.2) - - - -# model = SMT.PSDModel(Legendre(0.0..1.0)^2, :downward_closed, 4) -# _smp = SMT.Sampler(model) -smp_list = SMT.ConditionalMapping{2, 0, Float64}[] -ϵ_list = [0.8, 0.5, 0.3, 0.2, 0.1] - -for k=1:1 - XY = rand(2, 2000) - X = rand(1, 1000) - Y = rand(1, 1000) - - model = SMT.PSDModel(Legendre(0.0..1.0)^2, :downward_closed, 5) - - prec = if k > 1 - SMT.CondSampler(smp_list[1:k-1], nothing) - else - nothing - end - - SMT.Statistics.entropic_OT!(model, c, x->pdf(p, x)[1], x->pdf(q, x)[1], - ϵ_list[k], eachcol(XY), trace=true, - X=eachcol(X), Y=eachcol(Y), - optimizer=Mosek.Optimizer, - set_start_value=false, - preconditioner=prec, - # use_preconditioner_cost=use_prec, - normalization=true) - - normalize!(model) - push!(smp_list, SMT.Sampler(model)) -end - - - -rng = 0.01:0.01:0.99 -_smp = SMT.CondSampler(smp_list[1:1], nothing) -contourf(rng, rng, (x, y) -> pdf(_smp, [x, y]), alpha=1.0, label="transported model") - - -contourf(rng, rng, (x, y)->c([x, y])) -contourf(rng, rng, (x, y)->c(SMT.pushforward(_smp, [x, y]))) - -M_sink = compute_Sinkhorn(rng, x->pdf(p, x)[1], - x->pdf(q, x)[1], c, 0.1) -contourf(rng, rng, M_sink') - -contour( - rng, rng, (x, y) -> 0.5*pdf(_smp, [x, y]), levels=5 -) -contour!(rng, rng, M_sink', levels=5) - -rng = 0.01:0.01:0.99 -plot(rng, x->left_pdf(_smp, x)) -plot!(rng, x->pdf(p, x)[1]) - -plot(rng, x->right_pdf(_smp, x)) -plot!(rng, x->pdf(q, x)[1]) - - -distance_list = [] -distance_chained_list = [] -for k=1:length(smp_list) - _smp = SMT.CondSampler(smp_list[1:k], nothing) - push!(distance_chained_list, compute_Sinkhorn_distance(c, _smp, N=5000)) - M_sink = compute_Sinkhorn(rng, x->pdf(p, x)[1], - x->pdf(q, x)[1], c, ϵ_list[k]) - push!(distance_list, compute_Sinkhorn_distance(c, M_sink, rng)) -end - -plot(ϵ_list, distance_list, label="Sinkhorn distance using Sinhorn algorithm", - xscale=:log10, linewidth=2) -plot!(ϵ_list, distance_chained_list, label="Sinkhorn distance using Sequential measure transport", linewidth=2) -scatter!(ϵ_list, distance_chained_list, label=nothing) -xlabel!("ϵ") -ylabel!("Sinkhorn distance") \ No newline at end of file diff --git a/preliminary_OT_tests/OT_single_layer.jl b/preliminary_OT_tests/OT_single_layer.jl deleted file mode 100644 index a095a00..0000000 --- a/preliminary_OT_tests/OT_single_layer.jl +++ /dev/null @@ -1,89 +0,0 @@ -using SequentialMeasureTransport -using ApproxFun -import SequentialMeasureTransport as SMT -using LinearAlgebra -using Distributions -using Plots -using MosekTools - -include("helpers.jl") - - -## define a cost function -c(x) = (x[1] - x[2])^2 - -## define the marginals -p = MixtureModel(Normal[ - Normal(0.3, 0.1), - Normal(0.7, 0.1) -]) -q = Normal(0.5, 0.2) - - -### First, we work on the bounded domain [0, 1]^2 - -model = SMT.PSDModel(Legendre(0.0..1.0)^2, :downward_closed, 5) - -# we generate some random points to estimate the cost function -XY = rand(2, 2000) - -## by default, the marginals are satisfied by the marginals of XY. -## we can also specify the marginal point constraints explicitly. -X = rand(1, 1000) -Y = rand(1, 1000) - -SMT.Statistics.entropic_OT!(model, c, x->pdf(p, x)[1], x->pdf(q, x)[1], - 0.3, eachcol(XY), trace=true, - X=eachcol(X), Y=eachcol(Y), - optimizer=Mosek.Optimizer, - set_start_value=false,) - -normalize!(model) - -rng = 0.0:0.01:1.0 -contourf(rng, rng, (x, y) -> model([x, y]), alpha=1.0, label="transported model") - - -### Let us now demonstrate to work on the unbounded domain R^2 -# first, we need a preconditioner for the unbounded domain - -ref_map = SMT.ReferenceMaps.AlgebraicReference{2, Float64}() - -model = SMT.PSDModel(Legendre(0..1)^2, :downward_closed, 5) - -## define the marginals -p = MixtureModel(Normal[ - Normal(-1.0, 0.3), - Normal(1.0, 0.3) -]) -q = Normal(0.0, 0.5) - - -## Let us generate some random points in R^2. -## Best is, we first sample in [0, 1]^2 and then apply the inverse of the reference map. -XY_R2 = SMT.pullback.(Ref(ref_map), eachcol(XY)) - -X_R_Y_R = SMT.pullback.(Ref(ref_map), eachcol([X ; Y])) - -X_R = [[x[1]] for x in X_R_Y_R] -Y_R = [[x[2]] for x in X_R_Y_R] - -rng = 0:0.01:1.0 -cost_pb = SMT.pushforward(ref_map, c) -marg_dist_pb = SMT.pushforward(ref_map, x->pdf(p, x[1]) * pdf(q, x[2])) -contourf(rng, rng, (x, y) -> cost_pb([x, y]), alpha=1.0, label="cost function") - -SMT.Statistics.entropic_OT!(model, c, x->pdf(p, x)[1], x->pdf(q, x)[1], - 0.3, eachcol(XY), trace=true, - X=X_R, Y=Y_R, - optimizer=Mosek.Optimizer, - # preconditioner=precond, - reference=ref_map, # this is a diagonal map - set_start_value=false,) - -normalize!(model) -rng = 0.0:0.01:0.0 -model_R = SMT.pushforward(ref_map, x->model(x)) -contour(rng, rng, (x, y) -> model([x, y]), alpha=1.0, label="transported model") - -plot(rng, x->model([x, 0.5]), label="x-marginal") \ No newline at end of file diff --git a/preliminary_OT_tests/data/Vancouver.jpg b/preliminary_OT_tests/data/Vancouver.jpg deleted file mode 100644 index 76b3a7f..0000000 Binary files a/preliminary_OT_tests/data/Vancouver.jpg and /dev/null differ diff --git a/preliminary_OT_tests/data/new_img2.jpg b/preliminary_OT_tests/data/new_img2.jpg deleted file mode 100644 index 3f1088f..0000000 Binary files a/preliminary_OT_tests/data/new_img2.jpg and /dev/null differ diff --git a/preliminary_OT_tests/data/pink_bananas.jpg b/preliminary_OT_tests/data/pink_bananas.jpg deleted file mode 100644 index 8bd78d3..0000000 Binary files a/preliminary_OT_tests/data/pink_bananas.jpg and /dev/null differ diff --git a/preliminary_OT_tests/helpers.jl b/preliminary_OT_tests/helpers.jl deleted file mode 100644 index 12d0b92..0000000 --- a/preliminary_OT_tests/helpers.jl +++ /dev/null @@ -1,118 +0,0 @@ -using LinearAlgebra -using FastGaussQuadrature -using ImageFiltering - -function compute_Sinkhorn(rng, left_marg_d, right_marg_d, c, ϵ; iter=200) - M_epsilon = [exp(-c([x,y])/ϵ) for x in rng, y in rng] - left_mat = [left_marg_d([x]) for x in rng] - right_mat = [right_marg_d([x]) for x in rng] - left_margin(M) = M * ones(size(M, 1)) - right_margin(M) = M' * ones(size(M, 2)) - for _=1:iter - M_epsilon = diagm(left_mat ./ left_margin(M_epsilon)) * M_epsilon - M_epsilon = M_epsilon * diagm(right_mat ./ right_margin(M_epsilon)) - end - M_epsilon = length(rng)^2 * M_epsilon / sum(M_epsilon) - # for (i, ii) in enumerate(rng), (j, jj) in enumerate(rng) - # M_epsilon[i, j] *= 1/(left_marg_d([ii]) * right_marg_d([jj])) - # end - return M_epsilon -end - -function compute_Sinkhorn(rng, left_array::AbstractVector{T}, right_array::AbstractVector{T}, - c, ϵ; iter=200, domain_size=1.0) where {T<:Number} - M_epsilon = zeros(length(rng), length(rng)) - for (i, x) in enumerate(rng), (j, y) in enumerate(rng) - M_epsilon[i, j] = exp(-c(SA[x...],SA[y...])/ϵ) - end - proj_mat = ones(size(M_epsilon, 1)) - left_margin(M) = M * proj_mat - right_margin(M) = M' * proj_mat - for _=1:iter - M_epsilon .= diagm(left_array ./ left_margin(M_epsilon)) * M_epsilon - M_epsilon .= M_epsilon * diagm(right_array ./ right_margin(M_epsilon)) - end - M_epsilon = length(rng)^2 * M_epsilon / (sum(M_epsilon) * domain_size^2) - # for (i, ii) in enumerate(rng), (j, jj) in enumerate(rng) - # M_epsilon[i, j] *= 1/(left_marg_d([ii]) * right_marg_d([jj])) - # end - return M_epsilon -end - -function Barycentric_map_from_sinkhorn(M_sink, left_marg_vec, right_marg_vec, i) - res = SA[0.0, 0.0, 0.0] - for j in 1:length(right_marg_vec) - res += M_sink[i, j] * right_marg_vec[j] - end - return (1/left_marg_vec[i]) * res -end - -""" -Implemented for 1D distributions as in: - Justin Solomon, Fernando de Goes, Gabriel Peyré, Marco Cuturi, - Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. 2015. - Convolutional wasserstein distances: efficient optimal transportation - on geometric domains. ACM Trans. Graph. 34, 4, Article 66 (August 2015), - 11 pages. https://doi.org/10.1145/2766963 -""" -function compute_Sinkhorn_Wasserstein_barycenter(rng, - list_marg::Vector{Function}, - weights::Vector{T}, - ϵ; iter=200) where {T<:Number} - v = ones(length(list_marg), length(rng)) - w = ones(length(list_marg), length(rng)) - d = ones(length(list_marg), length(rng)) - marg_vec = [list_marg[i](x) for i in 1:length(list_marg), x in rng] - a = 1 / length(rng)^2 - mu = ones(length(rng)) - for _=1:iter - mu .= 1.0 - for i=1:length(list_marg) - w[i, :] .= marg_vec[i, :] ./ imfilter(a * v[i, :], Kernel.gaussian((ϵ,))) - d[i, :] .= v[i, :] .* imfilter(a * w[i, :], Kernel.gaussian((ϵ,))) - mu .= mu .* (w[i, :].^ weights[i]) - end - - for i=1:length(list_marg) - v[i, :] .= (v[i, :] .* mu) ./ d[i, :] - end - end - return mu -end - - -function left_pdf(sampler, x) - p, w = FastGaussQuadrature.gausslegendre(50) - ## scale to [0, 1] - p .= (p .+ 1) ./ 2 - w .= w ./ 2 - return sum(w .* [pdf(sampler, [x, y]) for y in p]) -end - -function right_pdf(sampler, x) - p, w = FastGaussQuadrature.gausslegendre(50) - ## scale to [0, 1] - p .= (p .+ 1) ./ 2 - w .= w ./ 2 - return sum(w .* [pdf(sampler, [y, x]) for y in p]) -end - -function compute_Sinkhorn_distance(c, - smp::SMT.ConditionalMapping{d}; - N=10000) where {d} - _d = d ÷ 2 - X = rand(smp, N) - return mapreduce(x -> c(x[1:_d], x[_d+1:end]), +, X) / N -end - -function compute_Sinkhorn_distance(c, - M_sink::AbstractArray{T, d}, - rng) where {d, T<:Number} - _d = d ÷ 2 - density = (rng[end] - rng[1]) / length(rng) - res = zero(T) - for ((i, x), (j, y)) in Iterators.product(enumerate(rng), enumerate(rng)) - res += c(x, y) * M_sink[i, j] * density^d - end - return res -end \ No newline at end of file diff --git a/preliminary_OT_tests/img_color_transfer.jl b/preliminary_OT_tests/img_color_transfer.jl deleted file mode 100644 index 5645005..0000000 --- a/preliminary_OT_tests/img_color_transfer.jl +++ /dev/null @@ -1,178 +0,0 @@ -using Images -import SequentialMeasureTransport as SMT -using LinearAlgebra -using Transducers -using Plots -using StaticArrays - -include("helpers.jl") - - -img1 = load("preliminary_OT_tests/data/pink_bananas.jpg") -img2 = load("preliminary_OT_tests/data/Vancouver.jpg") - -# img1 = imresize(img1, ratio=0.25) - -length1 = length(img1) -length2 = length(img2) - -target_size = 30000 -iter_1_count = length1 ÷ target_size -iter_2_count = length2 ÷ target_size - -colors_img1 = [[Float64(ColorTypes.red(img1[i])), - Float64(ColorTypes.green(img1[i])), - Float64(ColorTypes.blue(img1[i]))] for i in 1:iter_1_count:length1] - -colors_img2 = [[Float64(ColorTypes.red(img2[i])), - Float64(ColorTypes.green(img2[i])), - Float64(ColorTypes.blue(img2[i]))] for i in 1:iter_2_count:length2] - - -gaussian(x::AbstractVector{T}, y::AbstractVector{T}; σ=0.05) where {T<:Number} = (1)/(2π * σ^2)^(3/2) * exp(-norm(x - y)^2 / (2 * σ^2)) - -function density_map(x::AbstractVector{T}, color_map::AbstractVector{<:Vector{T}}) where {T} - (1/length(color_map)) * foldxl(+, color_map |> Map(y -> gaussian(x, y)); init=zero(T), simd=true) -end - - -@time density_map(rand(3), colors_img1) -@time density_map(rand(3), colors_img2) - -size_x = 15 -rng = range(0, 1, length=size_x) -left_marg_vec = zeros([length(rng) for _ in 1:3]...) -@time Threads.@threads for i=1:size_x - x = rng[i] - for (j, y) in enumerate(rng) - for (k, z) in enumerate(rng) - @inbounds left_marg_vec[i, j, k] = density_map(SA[x, y, z], colors_img1) - end - end -end -## normalize Vector -left_marg_vec .*= (size_x^3 / sum(left_marg_vec)) -plot(1:size_x, i->sum(left_marg_vec[i, :, :]), color=:red, linewidth=2, label="red") -plot!(1:size_x, i->sum(left_marg_vec[:, i, :]), color=:green, linewidth=2, label="green") -plot!(1:size_x, i->sum(left_marg_vec[:, :, i]), color=:blue, linewidth=2, label="blue") - -right_marg_vec = zeros([length(rng) for _ in 1:3]...) -@time Threads.@threads for i=1:size_x - x = rng[i] - for (j, y) in enumerate(rng) - for (k, z) in enumerate(rng) - @inbounds right_marg_vec[i, j, k] = density_map(SA[x, y, z], colors_img2) - end - end -end -## normalize Vector -right_marg_vec .*= (size_x^3 / sum(right_marg_vec)) -plot(1:size_x, i->sum(right_marg_vec[i, :, :]), color=:red, linewidth=2, label="red") -plot!(1:size_x, i->sum(right_marg_vec[:, i, :]), color=:green, linewidth=2, label="green") -plot!(1:size_x, i->sum(right_marg_vec[:, :, i]), color=:blue, linewidth=2, label="blue") - - -## Compute OT -rng_Sink = Iterators.product(rng, rng, rng) -c(x, y) = norm(x - y)^2 -@time M_sink = compute_Sinkhorn(rng_Sink, vec(left_marg_vec), - vec(right_marg_vec), c, 0.01; iter=100) - -# contour(M_sink') - -rng_map_vec = rng_Sink |> Map(x -> SA[x...]) |> collect - - -cart_to_linear = reshape(collect(1:length(left_marg_vec)), size_x, size_x, size_x) -function color_to_index(color; size_x=size_x) - x_1 = round(Int, color[1] * (size_x-1) + 1) - x_2 = round(Int, color[2] * (size_x-1) + 1) - x_3 = round(Int, color[3] * (size_x-1) + 1) - return cart_to_linear[x_1, x_2, x_3] -end - -function RGB_to_SA(rgb) - return SA[ColorTypes.red(rgb) |> Float64, - ColorTypes.green(rgb) |> Float64, - ColorTypes.blue(rgb) |> Float64] -end - - -function color_transfer(img::AbstractArray{T, 2}, M_sink, marginal) where {T} - img_size = size(img) - new_img = similar(img) - Threads.@threads for i=1:img_size[1] - for j=1:img_size[2] - color_vec = RGB_to_SA(img[i, j]) - color_index = color_to_index(color_vec) - ret = Barycentric_map_from_sinkhorn(M_sink, vec(marginal), rng_map_vec, color_index) - ret = min.(ret, 1.0) - new_img[i, j] = RGB(ret...) - end - end - return new_img -end - -img1 -img2 -new_img2 = color_transfer(img2, M_sink', right_marg_vec) -new_img1 = color_transfer(img1, M_sink, left_marg_vec) - - -img1 -img2 -new_img1 -new_img2 - -## check that color spectrum is correct afterwards - -colors_new_img1 = [[Float64(ColorTypes.red(new_img1[i])), - Float64(ColorTypes.green(new_img1[i])), - Float64(ColorTypes.blue(new_img1[i]))] for i in 1:iter_1_count:length1] - -colors_new_img2 = [[Float64(ColorTypes.red(new_img2[i])), - Float64(ColorTypes.green(new_img2[i])), - Float64(ColorTypes.blue(new_img2[i]))] for i in 1:iter_2_count:length2] - - -rng = range(0, 1, length=size_x) -color_density_new1 = zeros([length(rng) for _ in 1:3]...) -@time Threads.@threads for i=1:size_x - x = rng[i] - for (j, y) in enumerate(rng) - for (k, z) in enumerate(rng) - @inbounds color_density_new1[i, j, k] = density_map(SA[x, y, z], colors_new_img1) - end - end -end -## normalize Vector -color_density_new1 .*= (size_x^3 / sum(color_density_new1)) -plot(1:size_x, i->sum(color_density_new1[i, :, :]), color=:red, linewidth=2, label="red") -plot!(1:size_x, i->sum(color_density_new1[:, i, :]), color=:green, linewidth=2, label="green") -plot!(1:size_x, i->sum(color_density_new1[:, :, i]), color=:blue, linewidth=2, label="blue") -## compare to what it should be -plot!(1:size_x, i->sum(right_marg_vec[i, :, :]), color=:red, linewidth=2, linestyle=:dot, label="reference red") -plot!(1:size_x, i->sum(right_marg_vec[:, i, :]), color=:green, linewidth=2, linestyle=:dot, label="reference green") -plot!(1:size_x, i->sum(right_marg_vec[:, :, i]), color=:blue, linewidth=2, linestyle=:dot, label="reference blue") - - -color_density_new2 = zeros([length(rng) for _ in 1:3]...) -@time Threads.@threads for i=1:size_x - x = rng[i] - for (j, y) in enumerate(rng) - for (k, z) in enumerate(rng) - @inbounds color_density_new2[i, j, k] = density_map(SA[x, y, z], colors_new_img2) - end - end -end -## normalize Vector -color_density_new2 .*= (size_x^3 / sum(color_density_new2)) -plot(1:size_x, i->sum(color_density_new2[i, :, :]), color=:red, linewidth=2, label="red") -plot!(1:size_x, i->sum(color_density_new2[:, i, :]), color=:green, linewidth=2, label="green") -plot!(1:size_x, i->sum(color_density_new2[:, :, i]), color=:blue, linewidth=2, label="blue") -## compare to what it should be -plot!(1:size_x, i->sum(left_marg_vec[i, :, :]), color=:red, linewidth=2, linestyle=:dot, label="reference red") -plot!(1:size_x, i->sum(left_marg_vec[:, i, :]), color=:green, linewidth=2, linestyle=:dot, label="reference green") -plot!(1:size_x, i->sum(left_marg_vec[:, :, i]), color=:blue, linewidth=2, linestyle=:dot, label="reference blue") - -# save("preliminary_OT_tests/data/new_img2.jpg", new_img2) \ No newline at end of file diff --git a/src/OptimalTransport.jl b/src/OptimalTransport.jl new file mode 100644 index 0000000..a821b5c --- /dev/null +++ b/src/OptimalTransport.jl @@ -0,0 +1,142 @@ +module OptimalTransport + +using ..SequentialMeasureTransport +import ..SequentialMeasureTransport as SMT +using ..SequentialMeasureTransport: PSDDataVector +using Distributions +using FastGaussQuadrature: gausslegendre + + +function entropic_OT!(model::SMT.PSDModelOrthonormal{d2, T}, + cost::Function, + p::Function, + q::Function, + ϵ::T, + XY::PSDDataVector{T}; + X=nothing, Y=nothing, + preconditioner::Union{<:SMT.ConditionalMapping{d2, 0, T}, Nothing}=nothing, + reference::Union{<:SMT.ReferenceMap{d2, 0, T}, Nothing}=nothing, + use_putinar=true, + use_preconditioner_cost=false, + λ_marg=nothing, + kwargs...) where {d2, T<:Number} + @assert d2 % 2 == 0 + d = d2 ÷ 2 + reverse_KL_cost = begin + if use_preconditioner_cost + let p=p, q=q + x->p(x[1:d]) * q(x[d+1:end]) + end + else + _rev_KL_density = let p=p, q=q + x -> p(x[1:d]) * q(x[d+1:end]) + end + if reference !== nothing + _rev_KL_density = SMT.pushforward(reference, _rev_KL_density) + end + if preconditioner === nothing + _rev_KL_density + else + SMT.pullback(preconditioner, _rev_KL_density) + end + end + end + + cost_pb = begin + _cost = let cost=cost + x -> cost(x) + end + if reference !== nothing + _cost = SMT.pushforward(reference, _cost) + else + _cost + end + end + + if preconditioner !== nothing + cost_pb = let cost_pb=cost_pb + x -> cost_pb(SMT.pushforward(preconditioner, x)) + end + end + + ξ = map(x->reverse_KL_cost(x), XY) + ξ2 = map(x->cost_pb(x), XY) + if λ_marg === nothing + ## estimate the order of the reverse KL cost to find an acceptable λ_marg + ## to do that, we calculate KL(U||reverse_KL_cost) where U is the distribution of XY + _order_rev_KL = (sum(ξ2) - ϵ * sum(log.(ξ))) / length(ξ) + λ_marg = 10.0*_order_rev_KL + @info "Estimated order of the reverse KL cost: $_order_rev_KL \n + Setting λ_marg to $λ_marg" + end + + model_for_marg = if preconditioner === nothing + model + else + SMT._add_mapping(model, preconditioner) + end + if X === nothing + X = [x[1:d] for x in XY] + end + if Y === nothing + Y = [x[d+1:end] for x in XY] + end + + + _p, _q = if reference !== nothing + _p = SMT.pushforward(reference[1:d], p) + _q = SMT.pushforward(reference[d+1:end], q) + _p, _q + else + p, q + end + + ## pushforward the samples + if reference !== nothing + _XY_marg = SMT.pushforward.(Ref(reference), [[x;y] for (x, y) in zip(X, Y)]) + X = [x[1:d] for x in _XY_marg] + Y = [x[d+1:end] for x in _XY_marg] + end + + ## evaluate the marginals on the original samples + p_X = map(_p, X) + q_Y = map(_q, Y) + e_X = collect(1:d) + e_Y = collect(d+1:d2) + if use_putinar && (typeof(model) <: SMT.PSDModelPolynomial) + D, C = SMT.get_semialgebraic_domain_constraints(model) + return SMT._OT_JuMP!(model, cost_pb, ϵ, XY, ξ; mat_list=D, coef_list=C, + model_for_marginals=model_for_marg, + marg_regularization = [(e_X, X, p_X), (e_Y, Y, q_Y)], + λ_marg_reg=λ_marg, + kwargs...) + else + return SMT._OT_JuMP!(model, cost_pb, ϵ, XY, ξ; + model_for_marginals=model_for_marg, + marg_regularization = [(e_X, X, p_X), (e_Y, Y, q_Y)], + λ_marg_reg=λ_marg, + kwargs...) + end +end + +function Wasserstein_Barycenter(model::SMT.PSDModelOrthonormal{d2, T}, + measures::AbstractVector{<:Function}, + weights::AbstractVector{T}, + ϵ::T, + XY::PSDDataVector{T}; + X=nothing, Y=nothing, + preconditioner::Union{<:SMT.ConditionalMapping{d2, 0, T}, Nothing}=nothing, + reference::Union{<:SMT.ReferenceMap{d2, 0, T}, Nothing}=nothing, + use_putinar=true, + use_preconditioner_cost=false, + λ_marg=nothing, + kwargs... + ) where {d2, T<:Number} + d = d2 * (length(measures) + 1) + + throw(error("Not implemented yet")) + +end + + +end \ No newline at end of file diff --git a/src/SequentialMeasureTransport.jl b/src/SequentialMeasureTransport.jl index c1edb3b..632454f 100644 --- a/src/SequentialMeasureTransport.jl +++ b/src/SequentialMeasureTransport.jl @@ -86,4 +86,8 @@ include("Samplers/sampler.jl") include("statistics.jl") using .Statistics +# methods to create Optimal Transport plans and more +include("OptimalTransport.jl") +using .OptimalTransport + end # module PositiveSemidefiniteModels diff --git a/src/statistics.jl b/src/statistics.jl index 247c3ee..8874e5d 100644 --- a/src/statistics.jl +++ b/src/statistics.jl @@ -225,142 +225,6 @@ function reversed_KL_fit!(model::PSDModel{T}, end end -function entropic_OT!(model::PSDModelOrthonormal{d2, T}, - cost::Function, - p::Function, - q::Function, - ϵ::T, - XY::PSDDataVector{T}; - X=nothing, Y=nothing, - preconditioner::Union{<:SMT.ConditionalMapping{d2, 0, T}, Nothing}=nothing, - reference::Union{<:SMT.ReferenceMap{d2, 0, T}, Nothing}=nothing, - use_putinar=true, - use_preconditioner_cost=false, - λ_marg=nothing, - kwargs...) where {d2, T<:Number} - @assert d2 % 2 == 0 - d = d2 ÷ 2 - reverse_KL_cost = begin - if use_preconditioner_cost - let p=p, q=q - # x -> pdf(prec, x) - x->p(x[1:d]) * q(x[d+1:end]) - end - else - _rev_KL_density = let p=p, q=q - x -> p(x[1:d]) * q(x[d+1:end]) - end - if reference !== nothing - _rev_KL_density = SMT.pushforward(reference, _rev_KL_density) - end - if preconditioner === nothing - _rev_KL_density - else - SMT.pullback(preconditioner, _rev_KL_density) - end - end - end - - cost_pb = begin - _cost = let cost=cost - x -> cost(x) - end - if reference !== nothing - _cost = SMT.pushforward(reference, _cost) - else - _cost - end - end - - if preconditioner !== nothing - cost_pb = let cost_pb=cost_pb - x -> cost_pb(SMT.pushforward(preconditioner, x)) - end - end - - ξ = map(x->reverse_KL_cost(x), XY) - ξ2 = map(x->cost_pb(x), XY) - if λ_marg === nothing - ## estimate the order of the reverse KL cost to find an acceptable λ_marg - ## to do that, we calculate KL(U||reverse_KL_cost) where U is the distribution of XY - _order_rev_KL = (sum(ξ2) - ϵ * sum(log.(ξ))) / length(ξ) - λ_marg = 10.0*_order_rev_KL - @info "Estimated order of the reverse KL cost: $_order_rev_KL \n - Setting λ_marg to $λ_marg" - end - - model_for_marg = if preconditioner === nothing - model - else - SMT._add_mapping(model, preconditioner) - end - if X === nothing - X = [x[1:d] for x in XY] - end - if Y === nothing - Y = [x[d+1:end] for x in XY] - end - - - _p, _q = if reference !== nothing - _p = SMT.pushforward(reference[1:d], p) - _q = SMT.pushforward(reference[d+1:end], q) - _p, _q - else - p, q - end - - ## if there is a reference, only now, we - ## pushforward the samples - if reference !== nothing - _XY_marg = SMT.pushforward.(Ref(reference), [[x;y] for (x, y) in zip(X, Y)]) - X = [x[1:d] for x in _XY_marg] - Y = [x[d+1:end] for x in _XY_marg] - end - - ## evaluate the marginals on the original samples - p_X = map(_p, X) - q_Y = map(_q, Y) - # e_X = diagm([ones(T, d); zeros(T, d)])[:, 1:d] - # e_Y = diagm([zeros(T, d); ones(T, d)])[:, d+1:end] - e_X = collect(1:d) - e_Y = collect(d+1:d2) - # @show e_X - # @show e_Y - if use_putinar && (typeof(model) <: PSDModelPolynomial) - D, C = SMT.get_semialgebraic_domain_constraints(model) - return SMT._OT_JuMP!(model, cost_pb, ϵ, XY, ξ; mat_list=D, coef_list=C, - model_for_marginals=model_for_marg, - marg_regularization = [(e_X, X, p_X), (e_Y, Y, q_Y)], - λ_marg_reg=λ_marg, - kwargs...) - else - return SMT._OT_JuMP!(model, cost_pb, ϵ, XY, ξ; - model_for_marginals=model_for_marg, - marg_regularization = [(e_X, X, p_X), (e_Y, Y, q_Y)], - λ_marg_reg=λ_marg, - kwargs...) - end -end - -function Wasserstein_Barycenter(model::PSDModelOrthonormal{d2, T}, - measures::AbstractVector{<:Function}, - weights::AbstractVector{T}, - ϵ::T, - XY::PSDDataVector{T}; - X=nothing, Y=nothing, - preconditioner::Union{<:SMT.ConditionalMapping{d2, 0, T}, Nothing}=nothing, - reference::Union{<:SMT.ReferenceMap{d2, 0, T}, Nothing}=nothing, - use_putinar=true, - use_preconditioner_cost=false, - λ_marg=nothing, - kwargs... - ) where {d2, T<:Number} - d = d2 * (length(measures) + 1) - - - -end function conditional_expectation(model::PSDModelOrthonormal{d, T}, dim::Int) where {d, T<:Number}