From 13341c637cb03aa935e144f95034e9620b9a05bb Mon Sep 17 00:00:00 2001 From: aronStalmarck <76658998+aronStalmarck@users.noreply.github.com> Date: Sun, 19 Mar 2023 13:59:44 +0100 Subject: [PATCH 1/4] Create manifolds.jl --- src/manifolds.jl | 52 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) create mode 100644 src/manifolds.jl diff --git a/src/manifolds.jl b/src/manifolds.jl new file mode 100644 index 0000000..2dad53a --- /dev/null +++ b/src/manifolds.jl @@ -0,0 +1,52 @@ +""" +If you want to use a manifold not in Manifolds.jl, you can define a new type that inherits from Manifold and implement the following functions: +- project(manifold, point) +- shortest_path_interpolation(rng, process, point_0, point_t, s, t) + +# x_t = Array(euclidean_dimension, batch_dims...) +""" + +struct ManifoldBrownianDiffusion{M <: AbstractManifold, T <: Real} <: SamplingProcess + manifold::M + rate::T + getsteps::Function +end + +ManifoldBrownianDiffusion(manifold::AbstractManifold, rate::T) where T <: Real = ManifoldBrownianDiffusion(manifold, rate, (t) -> [t]) + +function project!(X, manifold) + for ix in CartesianIndices(size(X)[2]) + X[:, ix] = project(manifold, X[:, ix]) + end +end + +function sampleforward(rng::AbstractRNG, process::ManifoldBrownianDiffusion{M, T}, t::Real, x_0) where M where T + x_t = project!(x_0, process.manifold) + for step in process.getsteps(t * process.rate) + x_t .+= randn(rng, T, size(x_t)...) .* sqrt(step) + project!(x_t, process.manifold) + end + return x_t +end + +shortestpath_interpolation(rng::AbstractRNG, process::ManifoldBrownianDiffusion, p_0::AbstractVector, p_t::AbstractVector, s, t) = + shortest_geodesic(process.manifold, p_0, p_t, s / t) + +shortestpath_interpolation_all(rng::AbstractRNG, process::ManifoldBrownianDiffusion, x_0, x_t, s, t) = + mapslices(indicesofpoint -> shortestpath_interpolation(rng, process, x_0[indicesofpoint], x_t[indicesofpoint], s, t), CartesianIndices(x_0), dims = [1]) + +# Empirically, this is good - should be the same as diffusing C from B as is done in the rotational/angular cases (nice symmetries) for spheres at least. +function endpoint_conditioned_sample(rng::AbstractRNG, process::ManifoldBrownianDiffusion, s::Real, t::Real, x_0, x_t) + B = sampleforward(rng, process, s, x_0) + C = sampleforward(rng, process, t-s, x_t) + return shortestpath_interpolation_all(rng, process, B, C, s, t) +end + +# This is closer to the bridge trick for the angular diffusion but the code is not as nice, and empirically it is not as good - but should give the same result, will redo tests. +#= +function endpoint_conditioned_sample(rng::AbstractRNG, process::ManifoldBrownianDiffusion, s::Real, t::Real, x_0, x_t) + B = sampleforward(rng, process, s, x_0) + C = sampleforward(rng, process, t-s, B) + return [exp(process.manifold, b, log(process.manifold, c, shortest_geodesic(process.manifold, c, endpoint, (t-s) / t))) for (b, c, endpoint) in zip(B, C, x_t)] +end +=# From 46eec75136439bec70a9ecd5f7962b17db53e7e1 Mon Sep 17 00:00:00 2001 From: aronStalmarck <76658998+aronStalmarck@users.noreply.github.com> Date: Sun, 19 Mar 2023 14:07:43 +0100 Subject: [PATCH 2/4] Bugfix --- src/manifolds.jl | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/manifolds.jl b/src/manifolds.jl index 2dad53a..7e56d64 100644 --- a/src/manifolds.jl +++ b/src/manifolds.jl @@ -2,10 +2,10 @@ If you want to use a manifold not in Manifolds.jl, you can define a new type that inherits from Manifold and implement the following functions: - project(manifold, point) - shortest_path_interpolation(rng, process, point_0, point_t, s, t) - -# x_t = Array(euclidean_dimension, batch_dims...) """ +# x_t = Array(euclidean_dimension, batch_size) + struct ManifoldBrownianDiffusion{M <: AbstractManifold, T <: Real} <: SamplingProcess manifold::M rate::T @@ -14,17 +14,18 @@ end ManifoldBrownianDiffusion(manifold::AbstractManifold, rate::T) where T <: Real = ManifoldBrownianDiffusion(manifold, rate, (t) -> [t]) -function project!(X, manifold) - for ix in CartesianIndices(size(X)[2]) - X[:, ix] = project(manifold, X[:, ix]) +function project!(x_to, x_from, manifold) + for ix in CartesianIndices(size(x_from)[2:end]) + x_to[:, ix] = project(manifold, x_from[:, ix]) end end function sampleforward(rng::AbstractRNG, process::ManifoldBrownianDiffusion{M, T}, t::Real, x_0) where M where T - x_t = project!(x_0, process.manifold) + x_t = similar(x_0) + project!(x_t, x_0, process.manifold) for step in process.getsteps(t * process.rate) x_t .+= randn(rng, T, size(x_t)...) .* sqrt(step) - project!(x_t, process.manifold) + project!(x_t, x_t, process.manifold) end return x_t end @@ -35,14 +36,14 @@ shortestpath_interpolation(rng::AbstractRNG, process::ManifoldBrownianDiffusion, shortestpath_interpolation_all(rng::AbstractRNG, process::ManifoldBrownianDiffusion, x_0, x_t, s, t) = mapslices(indicesofpoint -> shortestpath_interpolation(rng, process, x_0[indicesofpoint], x_t[indicesofpoint], s, t), CartesianIndices(x_0), dims = [1]) -# Empirically, this is good - should be the same as diffusing C from B as is done in the rotational/angular cases (nice symmetries) for spheres at least. +# Empirically, this is really good, but I am not completely sure why. function endpoint_conditioned_sample(rng::AbstractRNG, process::ManifoldBrownianDiffusion, s::Real, t::Real, x_0, x_t) B = sampleforward(rng, process, s, x_0) C = sampleforward(rng, process, t-s, x_t) return shortestpath_interpolation_all(rng, process, B, C, s, t) end -# This is closer to the bridge trick for the angular diffusion but the code is not as nice, and empirically it is not as good - but should give the same result, will redo tests. +# This does the same as for the angular diffusion but the code is not as nice, and empirically it is not as good. #= function endpoint_conditioned_sample(rng::AbstractRNG, process::ManifoldBrownianDiffusion, s::Real, t::Real, x_0, x_t) B = sampleforward(rng, process, s, x_0) From 56bb5ee61fb069cee63a8fea4f7ea22d98e88ecb Mon Sep 17 00:00:00 2001 From: aronStalmarck <76658998+aronStalmarck@users.noreply.github.com> Date: Wed, 22 Mar 2023 12:19:01 +0100 Subject: [PATCH 3/4] Added example for manifolds --- examples/manifolds.jl | 55 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 examples/manifolds.jl diff --git a/examples/manifolds.jl b/examples/manifolds.jl new file mode 100644 index 0000000..01c9186 --- /dev/null +++ b/examples/manifolds.jl @@ -0,0 +1,55 @@ +using Pkg; +Pkg.add(url="https://github.com/JuliaManifolds/ManifoldMeasures.jl") +Pkg.add(["MeasureBase", "Plots"]) +using Manifolds, ManifoldMeasures, MeasureBase, Plots + +function expectation(x_t, t; process = process, target_dist = target_dist, n_samples = 1000) + x0samples = [sampleforward(process, t, x_t) for i in 1:n_samples] + sampleweights = [mapslices(p -> target_pdf(target_dist, p), x0sample, dims = [1]) for x0sample in x0samples] + x_0_expectation = similar(x_t) + for pointindex in CartesianIndices(size(x_t)[2:end]) + x_0_expectation[:, pointindex] = sum( (samp -> samp[:, pointindex]).(x0samples) .* (sampweights -> sampweights[pointindex]).(sampleweights) ) + x_0_expectation[:, pointindex] = project(process.manifold, x_0_expectation[:, pointindex]) + end + x_0_expectation +end + +target_pdf(target_dist, point) = sum( MeasureBase.density_def(dist, point) * weight for (dist, weight) in zip(target_dist.dists, target_dist.weights) ) + +# N-dimensional sphere - 1 is a circle, 2 is a sphere, 3 is a quaternion, etc. +N = 2 +manifold = Sphere(N) +# Distributions on the sphere +dists = [ + ManifoldMeasures.VonMisesFisher(manifold, μ = project(manifold, [1.0, 1.0, 1.0]), κ = 70), + ManifoldMeasures.VonMisesFisher(manifold, μ = project(manifold, [-1.0, -1.0, -1.0]), κ = 70), + ManifoldMeasures.VonMisesFisher(manifold, μ = project(manifold, [0.99995, 9.99934e-5, 0.0]), κ = 1.5) + ] +unnormalized_weights = [1, 1, 1] +target_dist = (dists = dists, weights = unnormalized_weights ./ sum(unnormalized_weights)) + +# Diffusion process +process = ManifoldBrownianDiffusion(manifold, 1.0) +d = (1, ) +x_T = hcat(rand(uniform_distribution(manifold, zeros(N + 1)), d...)...) +timesteps = timeschedule(exp, log, 0.001, 20, 100) + +@time diffusion_samples = samplebackward(expectation, process, timesteps, x_T) + +function target_sample(target_dist) + r = rand() + for (dist, weight) in zip(target_dist.dists, target_dist.weights) + r -= weight + r < 0 && return rand(dist) + end +end + +target_samples = hcat([target_sample(target_dist) for i in eachindex(diffusion_samples)]...) + +coordvectors(samples) = [samples[i, :] for i in 1:size(samples)[1]] + +pl_S1_diffusion_samples = plot(title = "Diffusion samples", coordvectors(diffusion_samples)..., size=(400, 400), st = :scatter, xlim = (-1.1, 1.1), ylim = (-1.1, 1.1), + alpha = 0.3, color = "blue") +pl_S1_target_samples = plot(title = "Target samples", coordvectors(target_samples)..., size=(400, 400), st = :scatter, xlim = (-1.1, 1.1), ylim = (-1.1, 1.1), + alpha = 0.3, color="red") +plot(pl_S1_diffusion_samples, pl_S1_target_samples, size = (800, 400)) From 437b49aba8bdf05c7ba4df0e970bfafe0caf0485 Mon Sep 17 00:00:00 2001 From: aronStalmarck <76658998+aronStalmarck@users.noreply.github.com> Date: Wed, 22 Mar 2023 12:22:03 +0100 Subject: [PATCH 4/4] Update manifolds.jl --- src/manifolds.jl | 55 ++++++++++++++++++++++++++++++++++-------------- 1 file changed, 39 insertions(+), 16 deletions(-) diff --git a/src/manifolds.jl b/src/manifolds.jl index 7e56d64..b7a2511 100644 --- a/src/manifolds.jl +++ b/src/manifolds.jl @@ -2,9 +2,9 @@ If you want to use a manifold not in Manifolds.jl, you can define a new type that inherits from Manifold and implement the following functions: - project(manifold, point) - shortest_path_interpolation(rng, process, point_0, point_t, s, t) -""" -# x_t = Array(euclidean_dimension, batch_size) +x_t = Array(euclidean_dimension, batch_dims...) +""" struct ManifoldBrownianDiffusion{M <: AbstractManifold, T <: Real} <: SamplingProcess manifold::M @@ -12,11 +12,13 @@ struct ManifoldBrownianDiffusion{M <: AbstractManifold, T <: Real} <: SamplingPr getsteps::Function end -ManifoldBrownianDiffusion(manifold::AbstractManifold, rate::T) where T <: Real = ManifoldBrownianDiffusion(manifold, rate, (t) -> [t]) +ManifoldBrownianDiffusion(manifold::AbstractManifold, rate::T) where T <: Real = ManifoldBrownianDiffusion(manifold, rate, (t) -> range(min(t, 0.05), t, step=0.05)) + +pointindices(X) = CartesianIndices(size(X)[2:end]) function project!(x_to, x_from, manifold) - for ix in CartesianIndices(size(x_from)[2:end]) - x_to[:, ix] = project(manifold, x_from[:, ix]) + for pointindex in pointindices(x_from) + x_to[:, pointindex] = project(manifold, x_from[:, pointindex]) end end @@ -33,21 +35,42 @@ end shortestpath_interpolation(rng::AbstractRNG, process::ManifoldBrownianDiffusion, p_0::AbstractVector, p_t::AbstractVector, s, t) = shortest_geodesic(process.manifold, p_0, p_t, s / t) -shortestpath_interpolation_all(rng::AbstractRNG, process::ManifoldBrownianDiffusion, x_0, x_t, s, t) = - mapslices(indicesofpoint -> shortestpath_interpolation(rng, process, x_0[indicesofpoint], x_t[indicesofpoint], s, t), CartesianIndices(x_0), dims = [1]) +function shortestpath_interpolation!(x_s, rng::AbstractRNG, process::ManifoldBrownianDiffusion, x_0, x_t, s, t) + for pointindex in pointindices(x_s) + x_s[:, pointindex] = shortestpath_interpolation(rng, process, x_0[:, pointindex], x_t[:, pointindex], s, t) + end +end -# Empirically, this is really good, but I am not completely sure why. +# Empirically, this is good - should be the same as diffusing C from B as is done in the rotational/angular cases (nice symmetries) for spheres at least. function endpoint_conditioned_sample(rng::AbstractRNG, process::ManifoldBrownianDiffusion, s::Real, t::Real, x_0, x_t) B = sampleforward(rng, process, s, x_0) C = sampleforward(rng, process, t-s, x_t) - return shortestpath_interpolation_all(rng, process, B, C, s, t) + shortestpath_interpolation!(C, rng, process, B, C, s, t) + return C end -# This does the same as for the angular diffusion but the code is not as nice, and empirically it is not as good. -#= -function endpoint_conditioned_sample(rng::AbstractRNG, process::ManifoldBrownianDiffusion, s::Real, t::Real, x_0, x_t) - B = sampleforward(rng, process, s, x_0) - C = sampleforward(rng, process, t-s, B) - return [exp(process.manifold, b, log(process.manifold, c, shortest_geodesic(process.manifold, c, endpoint, (t-s) / t))) for (b, c, endpoint) in zip(B, C, x_t)] +# This has not been tested yet. An optional way of a heuristic brownian bridge for general manifolds. +function endpoint_conditioned_sample_distance_proportional(rng::AbstractRNG, process::ManifoldBrownianDiffusion, s::Real, t::Real, x_0, x_t) + x_s = similar(x_0) + D = process.rate # Diffusion coefficient + dt = min(s, 0.05) # timestep + for pointindex in pointindices(x_from) + p_0 = x_0[:, pointindex] + p_t = x_t[:, pointindex] + d = distance(process.manifold, p_0, p_t) + t_remaining = t + p_cur = p_0 + while (t - t_remaining) < s + dp = randn(rng, size(p_cur)...) .* (2 * D * dt) + p_new = p_cur .+ dp + d_remaining = distance(process.manifold, p_new, p_t) + factor = 1 - (d_remaining / d) * (t_remaining / T) + p_adjusted = shortest_geodesic(p_new, p_t, factor) + + p_cur = project(process.manifold, p_adjusted) + t_remaining -= dt + end + x_s[:, pointindex] = p_cur + end + x_s end -=#