diff --git a/Project.toml b/Project.toml index e8283793..3bc4bad0 100644 --- a/Project.toml +++ b/Project.toml @@ -4,7 +4,9 @@ authors = ["Sebastiano Grazzi and Moritz Schauer"] version = "0.8.0" [deps] +ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" +FunctionWranglers = "b554fa5d-8411-41ee-8688-d60d0f644aa9" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -12,10 +14,12 @@ RandomNumbers = "e6cf234a-135c-5ec9-84dd-332b85af5143" Requires = "ae029012-a4dd-5104-9daa-d747884805df" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" Trajectories = "2c80a279-213e-54d7-a557-e9a14725db56" [compat] DataStructures = "0.17, 0.18" +FunctionWranglers = "0.1.3" ProgressMeter = "1" Requires = "1.1" Trajectories = "0.2" diff --git a/casestudies/precision/precision.jl b/casestudies/precision/precision.jl index 7129806e..dac163d5 100644 --- a/casestudies/precision/precision.jl +++ b/casestudies/precision/precision.jl @@ -79,8 +79,9 @@ Z = ZigZag(Γ̂, μ̂) κ = 0.01ones(d) -trc, _ = @time ZigZagBoomerang.sspdmp(∇ϕ, t0, x0, θ0, T, c, G, Z, κ, YY, (𝕀, 𝕁), N; structured=true, adapt=true, progress=true) +trc__, _ = @time ZigZagBoomerang.sspdmp(∇ϕ, t0, x0, θ0, T, c, G, Z, κ, YY, (𝕀, 𝕁), N; adapt=true, progress=true) +trc = trc__ J = [1,2,5] subtrc = subtrace(trc, J) diff --git a/casestudies/precision/precision2.jl b/casestudies/precision/precision2.jl new file mode 100644 index 00000000..2851e486 --- /dev/null +++ b/casestudies/precision/precision2.jl @@ -0,0 +1,185 @@ +using Random +Random.seed!(5) +using Statistics, ZigZagBoomerang, LinearAlgebra, Test, SparseArrays +using ZigZagBoomerang +const Zig = ZigZagBoomerang +using ZigZagBoomerang: sλ, sλ̄, reflect!, Rng, ab, smove_forward!, neighbours +using StructArrays +using StructArrays: components +using LinearAlgebra + +function countids(f, s) + res = Dict{Int, Int}() + for c in s; + i = f(c) + res[i] = get(res, i, 0) + 1 + end + return res +end +seed = (UInt(1),UInt(1)) + +n = 400 +d = n*(n+1)÷2 +N = 1000 + +T = 600.0 + +outer(x) = x*x' +function backform(u, 𝕀) + L = zeros(𝕀[end][1], 𝕀[end][2]) + for (x, i) in zip(u, 𝕀) + L[i] = x + end + L +end +transform(L, 𝕀) = L[𝕀] + +#const σ2 = 1.0 +#const γ0 = 0.1 +dia = -0.3ones(n-1) +Γtrue = sparse(SymTridiagonal(1.0ones(n), dia)) +Γtrue[1,1] = Γtrue[end,end] = 1/2 + +Ltrue_ = cholesky(Γtrue).L +Y = Ltrue_\randn(n, N) +YY = Y*Y' +Ltrue = L = sparse(Ltrue_) +# Compute an unbiased estimate of the i'th partial derivative of the negative loglikelihood in a smart way +function ϕ(L, Y) + # L = reshape(x, d, d) + sum(2*(diag(L) .- 1.0).^2)/2 - N*sum(log.(diag(L).^2)) + tr(Y'*(L*(L'*Y)))/2 +end +using ForwardDiff +ϕ(Ltrue, Y) +utrue_ = Vector(vec(Ltrue)) +@test tr(YY*L*L') ≈ sum(Y[:, i]'*L*L'*Y[:,i] for i in 1:N) + +𝕀 = [c for c in CartesianIndices((n,n)) if c[1] >= c[2]] +𝕁 = [[(i,CartesianIndex(c[1], c2[1])) for (i,c2) in enumerate(𝕀) if c2[1] >= c[2] && c[2] == c2[2]] for c in 𝕀] +utrue = Vector(L[𝕀]) +function ∇ϕ(u, i, YY, (𝕀, 𝕁), N) + c = 0.0 + for (j, c2) in 𝕁[i] + #c += YY[ii[1], c2[1]] * u[j] #L[j,ii[2]] + c += YY[c2] * u[j] #L[j,ii[2]] + + end + if 𝕀[i][1] == 𝕀[i][2] + c += -N/u[i] #+ 2(u[i]-1) + end + c +end +#∇ϕ(utrue, 2, YY, (𝕀, 𝕁), N) + + +function next_rand_reflect(j, i, t′, u, P, YY, (𝕀, 𝕁), N) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + if m[j] == 1 + return 0, Inf + end + if !(𝕀[j][1] == 𝕀[j][2]) + b[j] = ab(G1, j, x, θ, c, F) + else + b[j] = ab(G1, j, x, θ, c, F) .+ (N/(x[j]), N*2/(x[j]^2)) + end + t_old[j] = t′ + 0, t[j] + poisson_time(b[j], rand(P.rng)) +end +function next_reset(j, _, t′, u, P, YY, (𝕀, 𝕁), N) + 0, !(𝕀[j][1] == 𝕀[j][2]) ? Inf : t′ + 0.5*u.x[j] +end + + +function freeze!(args...) + Zig.freeze!(0.0, args...) +end + +function next_freezeunfreeze(args...) + Zig.next_freezeunfreeze(0.0, 0.04, args...) +end + + +t0 = 0.0 +t = zeros(d) + +x0 = utrue #+ 0.01*randn(d) # jiggle the starting point to see convergence + +#te = reshape(ForwardDiff.gradient(u -> ϕ(reshape(u, n, n), Y), backform(x0, 𝕀)[:]), n, n) +#@test norm(Vector(te[𝕀]) - [∇ϕ(x0, i, YY, (𝕀, 𝕁), N) for i in 1:d]) < 10d^2*eps() + +θ0 = ones(d) + +Γ̂ = inv(cov(Y')) + +# precision bounds +c = 0.01ones(d) +dt = T/500 +Γ̂Z = sparse(1.0I(d))*N +#L̂ = cholesky(sparse(SymTridiagonal(cov(Y')))).L +#μ̂ = transform(Matrix(sparse(L̂)), 𝕀) +L̂ = cholesky(Symmetric(Γ̂)).L +μ̂ = transform(L̂, 𝕀) + +F = Z = ZigZag(Γ̂Z, μ̂) + +# Graphical structure of posterior +G = [i => first.(j) for (i,j) in enumerate(𝕁)] +# Graphical structure of bounds +G1 = [i => rowvals(F.Γ)[nzrange(F.Γ, i)] for i in eachindex(θ0)] +# What is needed to update clocks +G2 = [i => setdiff(union((G1[j].second for j in G1[i].second)...), G[i].second) for i in eachindex(G1)] + + +b = [ab(G1, i, x0, θ0, c, F) for i in eachindex(θ0)] + +u0 = StructArray(t=t, x=x0, θ=θ0, θ_old = zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b) +rng = Rng(seed) +t_old = copy(t) +adapt = true +factor = 1.7 +P = SPDMP(G, G1, G2, ∇ϕ, F, rng, adapt, factor) +action! = (Zig.reset!, Zig.rand_reflect!, freeze!) +next_action = FunctionWrangler((next_reset, next_rand_reflect, next_freezeunfreeze)) +h = Schedule(action!, next_action, u0, T, (P, YY, (𝕀, 𝕁), N)) + +trc_ = @time simulate(h); +trc = Zig.FactTrace(F, t0, x0, θ0, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_]) + + +#trc, _ = @time ZigZagBoomerang.sspdmp(∇ϕ, t0, x0, θ0, T, c, G, Z, κ, YY, (𝕀, 𝕁), N; structured=true, adapt=true, progress=true) + +J = [1,2,5] +J, C = Zig.sep([(i,c) for (i,c) in enumerate(𝕀) if abs(c[1] - c[2]) <= 1]) + +subtrc = subtrace(trc, J) + +ts, xs = ZigZagBoomerang.sep(collect(discretize(subtrc, dt))) +#ts, xs = ZigZagBoomerang.sep(subtrc) + +# posterior mean +u = mean(trc) +Lhat = backform(u, 𝕀) +utrue - u + +using Makie +ina(i) = "$(𝕀[J[i]][1]),$(𝕀[J[i]][2])" +fig = Figure(resolution=(1800,1000)) +ax = fig[1,1:3] = Axis(fig, title="Error Gamma") +ax1 = fig[2,1] = Axis(fig, title="x$(ina(1))") +ax2 = fig[2,2] = Axis(fig, title="x$(ina(2))") +ax3 = fig[2,3] = Axis(fig, title="x$(ina(3))") +linkaxes!(ax1, ax2, ax3) + +heatmap!(ax, [Matrix(Γ̂); Matrix(Γtrue); outer(Lhat); Matrix(Γtrue) - outer(Lhat)], colormap=:vik, colorrange=[-1/4,1/4]) +lines!(fig[2,1], ts, getindex.(xs, 1)) +lines!(fig[2,1], ts, fill(utrue[J[1]], length(ts)), color=:green) +lines!(fig[2,2], ts, getindex.(xs, 2)) +lines!(fig[2,2], ts, fill(utrue[J[2]], length(ts)), color=:green) +lines!(fig[2,3], ts, getindex.(xs, 3)) +lines!(fig[2,3], ts, fill(utrue[J[3]], length(ts)), color=:green) +display(fig) + +using FileIO +save("precision.png", fig) \ No newline at end of file diff --git a/research/juliacon/animations.jl b/research/juliacon/animations.jl new file mode 100644 index 00000000..79d2e12f --- /dev/null +++ b/research/juliacon/animations.jl @@ -0,0 +1,185 @@ +using Base: Float64 +using ZigZagBoomerang +const ZZB = ZigZagBoomerang +### animation 1: no reflection (lebsgue measure) +x0, v0 = -2.0, 1.0 +dt = 0.01 +# fake trace from -2.0 to +2.0 +trace1 = [(t, x0 + t) for t in 0.0:dt:4.0] +# animate: todo + +### animation 1b: no reflection (lebsgue measure) with stickyness at 0 +x0, v0 = -2.0, 1.0 +dt = 0.01 +dtinv = Int(1/dt) +κ = 1.0 +T = 10.0 +wait = floor(-log(rand())/κ*dtinv)/dtinv # 10^-2 +# fake trace from -2.0 to +2.0 +trace1b = [[(t, x0 + t) for t in 0.0:dt:2.0]..., [(2.0 + t, 0.0) for t in dt:dt:wait]...,[(2.0 + wait + t, 0.0+ t) for t in dt:dt:T]... ] +# animate: todo + + + +### animation 2: Boomerang sampler with rereshments but no reflections +B = Boomerang1d(0.0) +t = 0.0 +T = 10.0 +x = x0 +v = v0 +out2 = [(0.0, x, v)] +while t < T + global x, v, t, out + ref = t -log(rand())/1.0 + τ = min(ref, T) + t, x, v = ZZB.move_forward(τ, t, x, v, B) + push!(out2, (t,x,v)) +end +trace2 = ZigZagBoomerang.discretize(out2, B, dt) +# animate: todo + +### animation 3: Zig-Zag with reflection for a multimodal density +∇ϕ(x) = (x-1.0)/1.0 + (x + 1.0)/1.0 +c = 3.0 +out3, acc = ZZB.pdmp(∇ϕ, x0, v0, T, 1.2π, ZigZag1d()) +trace3 = ZigZagBoomerang.discretize(out3, B, dt) + + +### animation 4: same as animation 3 + stickiness at 0 +# include("C:\\Users\\sebas\\.julia\\dev\\ZigZagBoomerang\\scripts\\sticky\\ss_1d_zigzag.jl") +const ZZB = ZigZagBoomerang +function freezing_time(x, θ) + if θ*x > 0 + return Inf + else + return -x/θ + end +end + +# k determines the weight on the Dirac measure. The smaller k, the higher the wegiht +function ss_pdmp(∇ϕ, x, θ, T, c, k, Flow::ZZB.ContinuousDynamics; adapt=false, factor=2.0) + t = zero(T) + Ξ = [(t, x, θ)] + num = acc = 0 + a, b = ZZB.ab(x, θ, c, Flow) + t_ref = t + ZZB.waiting_time_ref(Flow) + t′ = t + poisson_time(a, b, rand()) + tˣ = t + freezing_time(x, θ) + while t < T + if tˣ < min(t_ref, t′) + t, x, θ = ZZB.move_forward(tˣ - t, t, x, θ, Flow) # go to 0 + @assert -0.0001 < x < 0.0001 #check + #t, x , θ = tˣ, 0.0, θ #go to 0 + push!(Ξ, (t, x, 0.0)) + t = t - log(rand())/k #wait exponential time + push!(Ξ, (t, x, θ)) + tˣ = Inf + elseif t_ref < t′ + t, x, θ = ZZB.move_forward(t_ref - t, t, x, θ, Flow) + θ = sqrt(Flow.Σ)*randn() + t_ref = t + ZZB.waiting_time_ref(Flow) + tˣ = t + freezing_time(x, θ) + push!(Ξ, (t, x, θ)) + else + τ = t′ - t + t, x, θ = ZZB.move_forward(τ, t, x, θ, Flow) + l, lb = ZZB.λ(∇ϕ, x, θ, Flow), ZZB.λ_bar(τ, a, b) + num += 1 + if rand()*lb < l + acc += 1 + if l >= lb + !adapt && error("Tuning parameter `c` too small.") + c *= factor + end + θ = -θ # In multi dimensions the change of velocity is different: + # reflection symmetric on the normal vector of the contour + tˣ = t + freezing_time(x, θ) + push!(Ξ, (t, x, θ)) + end + end + a, b = ZZB.ab(x, θ, c, Flow) + t′ = t + poisson_time(a, b, rand()) + end + return Ξ, acc/num +end + +∇ϕ(x) = (x-1.0)/1.0 + (x + 1.0)/1.0 +k = 1.0 # rate of stickying time +c = 3.0 +k = 1.0 +out4, acc = ss_pdmp(∇ϕ, x0, v0, T, c, k, ZigZag1d()) +trace4 = ZigZagBoomerang.discretize(out4, B, dt) + +### animation 5: as 4) + boundaries +function hitting_bounds(x, θ, b) + b1, b2 = b + if θ*(x-b1) > 0 + return -(x-b2)/θ + else + return -(x-b1)/θ + end +end + +function ss_pdmp_b(b, ∇ϕ, x, θ, T, c, k, Flow::ZZB.ContinuousDynamics; adapt=false, factor=2.0) + b1, b2 = b + t = zero(T) + Ξ = [(t, x, θ)] + num = acc = 0 + a, b = ZZB.ab(x, θ, c, Flow) + t_ref = t + ZZB.waiting_time_ref(Flow) + t′ = t + poisson_time(a, b, rand()) + tˣ = t + freezing_time(x, θ) + tb = t + hitting_bounds(x, θ, b) + while t < T + if tˣ < min(t_ref, t′, tb) + t, x, θ = ZZB.move_forward(tˣ - t, t, x, θ, Flow) # go to 0 + @assert -0.0001 < x < 0.0001 #check + #t, x , θ = tˣ, 0.0, θ #go to 0 + push!(Ξ, (t, x, 0.0)) + t = t - log(rand())/k #wait exponential time + push!(Ξ, (t, x, θ)) + tˣ = Inf + elseif tb < min(t_ref, t′) + t, x, θ = ZZB.move_forward(tb - t, t, x, θ, Flow) # go to 0 + #t, x , θ = tˣ, 0.0, θ #go to 0 + @assert (b1-0.0001 < x < b1 + 0.0001 || b2 - 0.0001 < x < b2 + 0.0001 ) #check + θ > 0.0 ? x = b2 : x = b1 + θ = -θ + push!(Ξ, (t, x, θ)) + tˣ = t + hitting_bounds(x, θ, b) + elseif t_ref < t′ + t, x, θ = ZZB.move_forward(t_ref - t, t, x, θ, Flow) + θ = sqrt(Flow.Σ)*randn() + t_ref = t + ZZB.waiting_time_ref(Flow) + tˣ = t + freezing_time(x, θ) + push!(Ξ, (t, x, θ)) + else + τ = t′ - t + t, x, θ = ZZB.move_forward(τ, t, x, θ, Flow) + l, lb = ZZB.λ(∇ϕ, x, θ, Flow), ZZB.λ_bar(τ, a, b) + num += 1 + if rand()*lb < l + acc += 1 + if l >= lb + !adapt && error("Tuning parameter `c` too small.") + c *= factor + end + θ = -θ # In multi dimensions the change of velocity is different: + # reflection symmetric on the normal vector of the contour + tˣ = t + freezing_time(x, θ) + push!(Ξ, (t, x, θ)) + end + end + a, b = ZZB.ab(x, θ, c, Flow) + t′ = t + poisson_time(a, b, rand()) + end + return Ξ, acc/num +end + +∇ϕ(x) = (x-1.0)/1.0 + (x + 1.0)/1.0 +k = 1.0 # rate of stickying time +c = 3.0 +k = 1.0 +b = (-2.0, 2.0) #buondaries +out5, acc = ss_pdmp_b(∇ϕ, x0, v0, T, c, k, ZigZag1d()) +trace5 = ZigZagBoomerang.discretize(out4, B, dt) \ No newline at end of file diff --git a/research/sticky/comparisons/comparison_rj_vs_sticky.png b/research/sticky/comparisons/comparison_rj_vs_sticky.png new file mode 100644 index 00000000..ecebe639 Binary files /dev/null and b/research/sticky/comparisons/comparison_rj_vs_sticky.png differ diff --git a/research/sticky/comparisons/gibbs_linear.jl b/research/sticky/comparisons/gibbs_linear.jl new file mode 100644 index 00000000..33cf40a5 --- /dev/null +++ b/research/sticky/comparisons/gibbs_linear.jl @@ -0,0 +1,75 @@ +using LinearAlgebra +using Distributions + +function gibbs_linear(Y, X, w, iter, β, Z, σ2, c) + ββ = Vector{Float64}[] + ZZ = Vector{Float64}[] + push!(ββ, deepcopy(β)) + push!(ZZ, deepcopy(Z)) + for k in 2:iter + # Update active coordinates + Z = update_Z!(Y, X, w, Z, σ2, c) + # Update β + β = update_β!(Y, X, β, Z, σ2, c) + push!(ββ, deepcopy(β)) + push!(ZZ, deepcopy(Z)) + end + return ββ, ZZ +end + +function update_Z!(Y, X, w, Z, σ2, c) + for i in eachindex(Z) + Z[i] = 0 + c0 = compute_terms(Y, X, Z, σ2, c) + Z[i] = 1 + c1 = compute_terms(Y, X, Z, σ2, c) + p = w[i]/(1 + (1-w[i])*c0*sqrt(c)/(w[i]*c1)) # c comes from det(V0| Z_i=1)/det(V0|Z_i=0) = c^γ /c^(γ-1) + @assert 0<=p<=1 + rand() < p ? Z[i] = 1 : Z[i] = 0 + end + Z +end + +function compute_terms(Y, X, Z, σ2, c) #ok + γ = sum(Z) + if γ == 0 + # error("should not be here") + return 0.0 + end + Xz = view(X, :, Z) + βhatz = (Xz'*Xz) \ Xz'*Y + Vz = inv(Symmetric(I(γ)./c + Xz'*Xz)) + Snz = dot(Y - Xz*βhatz, Y - Xz*βhatz) + dot(Xz*βhatz, Xz*βhatz) - (Xz'*Y)'*Vz*(Xz'*Y) + return sqrt(det(Vz))*exp(-1/(2*σ2)*Snz) +end + +function update_β!(Y, X, β, Z, σ2, c) # OK + γ = sum(Z) + if γ == 0 + return β + end + Xz = view(X, :, Z) + Vz = inv(Symmetric(I(γ)./c + Xz'*Xz)) + μ_z = Vz*Xz'*Y + Σ_z = σ2*Vz + βz = rand(MvNormal(μ_z, Σ_z)) + β[Z] .= βz + β +end + +# Y = X β + ϵ +N = 100 +p = 50 +X = randn(N, p) +β_full = randn(p) +β_sparse = β_full.*abs.(β_full) .> 0.3 # make β spase +σ2 = 1.0 +Y = X*β_sparse .+ randn(N)*sqrt(σ2) +@assert N == length(Y) + +w = fill(0.5, p) +Z = w .< 0 +c = 1.0 +iter = 100 +β0 = randn(p) +gibbs_linear(Y, X, w, iter, β0, Z, σ2, c) \ No newline at end of file diff --git a/research/sticky/comparisons/gibbs_logistic.jl b/research/sticky/comparisons/gibbs_logistic.jl new file mode 100644 index 00000000..4603030f --- /dev/null +++ b/research/sticky/comparisons/gibbs_logistic.jl @@ -0,0 +1,99 @@ +using LinearAlgebra +using Distributions +include("./poyla.jl") +# pgdist = PolyaGamma(1,3.0) +# rand(pgdist) +function gibbs_logistic(Y, X, w, iter, β, Z, ω, σ2) + ββ = Vector{Float64}[] + ZZ = Vector{Float64}[] + ωω = Vector{Float64}[] + push!(ββ, deepcopy(β)) + push!(ZZ, deepcopy(Z)) + push!(ωω, deepcopy(ω)) + for k in 2:iter + # Update active coordinates + Z = update_Z!(Y, X, w, Z, ω, σ2) + # Update β + β = update_β!(Y, X, β, Z, ω, σ2) + # Update ω + ω = update_ω!(Y, X, w, β, Z, ω, σ2) + push!(ββ, deepcopy(β)) + push!(ZZ, deepcopy(Z)) + end + return ββ, ZZ +end + +function update_ω!(Y, X, w, β, Z, ω, σ2) + Xz = view(X, :, Z) + βz = view(β, Z) + ω .= rand.(PolyaGamma.(1, Xz*βz)) + ω +end + +function update_Z!(Y, X, w, Z, ω, σ2) + for i in eachindex(Z) + Z[i] = 0 + c0 = compute_terms_logistic(Y, X, Z, ω, σ2) + Z[i] = 1 + c1 = compute_terms_logistic(Y, X, Z, ω, σ2) + # println("c0 is $(c0)") + # println("c1 is $(c1)") + p = w[i]/(1 + (1-w[i])*sqrt(2π*σ2)*c0/(w[i]*c1)) # c comes from det(V0| Z_i=1)/det(V0|Z_i=0) = c^γ /c^(γ-1) + # p = w[i]/(1 + (1-w[i])*c0*sqrt(c)/(w[i]*c1)) # c comes from det(V0| Z_i=1)/det(V0|Z_i=0) = c^γ /c^(γ-1) + if !(0<=p<=1) + error("p equal to $p") + end + rand() < p ? Z[i] = 1 : Z[i] = 0 + end + Z +end + +function update_β!(Y, X, β, Z, ω, σ2) + γ = sum(Z) + if γ == 0 + return β + end + Ω = diagm(ω) + Xz = view(X, :, Z) + Vz = inv(Symmetric(I(γ)./σ2 + Xz'*Ω*Xz)) + μz = Vz*Xz'*(Y .- 0.5) + βz = rand(MvNormal(μz, Vz)) + β[Z] .= βz + β +end + +function compute_terms_logistic(Y, X, Z, ω, σ2) + γ = sum(Z) + if γ == 0 + # error("should not be here") + return 0.0 + end + Xz = view(X, :, Z) + Ω = diagm(ω) + Vz = inv(Symmetric(I(γ)./σ2 + Xz'*Ω*Xz)) + Snz = exp(0.5*(Y .- 0.5)'*Xz*Vz*Xz'*(Y .- 0.5)) + return sqrt(det(2*π*Vz))*Snz +end + +#logistic +sigmoid(x) = exp(x)/(1 + exp(x)) + +# Y = X β + ϵ +N = 50 +p = 5 +X = randn(N, p) +β_full = randn(p) +β_sparse = β_full.*(abs.(β_full) .> 0.3) # make β spase +P = sigmoid.(X*β_sparse) +Y = rand(N) .<= P +@assert N == length(Y) + +w = fill(0.5, p) +Z = w .> 0 +σ2 = 1.0 # prior on β +iter = 100 +β0 = randn(p) +ω = rand.(PolyaGamma.(1, X*β0)) +gibbs_logistic(Y, X, w, iter, β0, Z, ω, σ2) + +X*β0 \ No newline at end of file diff --git a/research/sticky/comparisons/inclusion.jl b/research/sticky/comparisons/inclusion.jl new file mode 100644 index 00000000..eb52c191 --- /dev/null +++ b/research/sticky/comparisons/inclusion.jl @@ -0,0 +1,59 @@ +function augment(t, x) + z = abs.(x) .>= 2*eps() + return z +end +z = augment.(t, x) + +function inclusion_probability(ts0, z) + T = ts0[end] + dict = Dict(z[1] => 0.0) + for (δt, i) in zip(diff(ts0), eachindex(ts0)) + if haskey(dict, z[i]) + dict[z[i]] += δt/T + else + push!(dict, z[i] => δt/T) + end + end +return dict +end + +function inclusion_probability(trace) + ts, xs = splitpairs(trace0) + zs = augment.(ts, xs) + inclusion_probability(ts, zs) +end + +Random.seed!(10) +function ϕ(x, i, μ) + x[i] - μ[i] +end +κ = .5 +n = 2 +μ = zeros(n) +x0 = randn(2) +θ0 = [1.0, -1.0] + +T = 10.0 +c = 0.1*ones(n) +@time trace0, _ = ZigZagBoomerang.sspdmp(ϕ, 0.0, x0, θ0, T, c, ZigZag(sparse(1.0I,n,n), zeros(n)), [κ,κ], μ) +ts0, xs0 = splitpairs(trace0) +zs0 = augment.(ts0, xs0) +zz = inclusion_probability(ts0, zs0) +@assert sum(values(zz)) == 1.0 + + +### Inclusion discrete time sampler +function inclusion_probability_discrete(z, N) + dict = Dict(z[1] => 0.0) + for i in eachindex(z) + if haskey(dict, z[i]) + dict[z[i]] += 1/N + else + push!(dict, z[i] => 1/N) + end + end +return dict +end + +error("") + diff --git a/research/sticky/comparisons/poyla.jl b/research/sticky/comparisons/poyla.jl new file mode 100644 index 00000000..f56b4d5f --- /dev/null +++ b/research/sticky/comparisons/poyla.jl @@ -0,0 +1,166 @@ +# copy and paste from https://github.com/currymj/PolyaGammaDistribution.jl +# https://github.com/currymj/PolyaGammaDistribution.jl +using Distributions, Random +using Statistics +using SpecialFunctions +const __TRUNC = 0.64; +const __TRUNC_RECIP = 1.0 / __TRUNC; +const twoπ = 6.2831853071795864769 +""" + PolyaGamma(b::Int, c::Real) +## Arguments +- `b::Int` +- `c::Real` exponential tilting +## Keyword Arguments +- `nmax::Int` Sampling parameter +- `trunc::Int` Sampling parameter +Create a PolyaGamma sampler with parameters `b` and `c` +""" +struct PolyaGamma{Tc,A} <: Distributions.ContinuousUnivariateDistribution + # For sum of Gammas. + b::Int + c::Tc + trunc::Int + nmax::Int + bvec::A + #Constructor + function PolyaGamma{T}(b::Int, c::T, trunc::Int, nmax::Int) where {T<:Real} + if trunc < 1 + @warn "trunc < 1. Setting trunc=1." + trunc = 1 + end + bvec = [convert(T, (twoπ * (k - 0.5))^2) for k in 1:trunc] + return new{typeof(c),typeof(bvec)}(b, c, trunc, nmax, bvec) + end +end + +Statistics.mean(d::PolyaGamma) = d.b / (2 * d.c) * tanh(0.5 * d.c) + +function PolyaGamma(b::Int, c::T; nmax::Int=10, trunc::Int=200) where {T<:Real} + return PolyaGamma{T}(b, c, trunc, nmax) +end + +function Distributions.pdf(d::PolyaGamma, x) + return cosh(d.c / 2)^d.b * 2.0^(d.b - 1) / gamma(d.b) * sum( + ((-1)^n) * gamma(n + d.b) / gamma(n + 1) * (2 * n + b) / (sqrt(2 * π * x^3)) * + exp(-(2 * n + b)^2 / (8 * x) - c^2 / 2 * x) for n in 0:(d.nmax) + ) +end + +## Sampling +function Distributions.rand(rng::AbstractRNG, d::PolyaGamma{T}) where {T<:Real} + if iszero(d.b) + return zero(T) + end + return sum(Base.Fix1(draw_like_devroye, rng), d.c * ones(d.b)) +end + +## Utility functions +function a(n::Int, x::Real) + k = (n + 0.5) * π + if x > __TRUNC + return k * exp(-0.5 * k^2 * x) + elseif x > 0 + expnt = -1.5 * (log(0.5 * π) + log(x)) + log(k) - 2.0 * (n + 0.5)^2 / x + return exp(expnt) + end +end + +function mass_texpon(z::Real) + t = __TRUNC + + fz = 0.125 * π^2 + 0.5 * z^2 + b = sqrt(1.0 / t) * (t * z - 1) + a = sqrt(1.0 / t) * (t * z + 1) * -1.0 + + x0 = log(fz) + fz * t + xb = x0 - z + logcdf(Distributions.Normal(), b) + xa = x0 + z + logcdf(Distributions.Normal(), a) + + qdivp = 4 / π * (exp(xb) + exp(xa)) + + return 1.0 / (1.0 + qdivp) +end + +function rtigauss(rng::AbstractRNG, z::Real) + z = abs(z) + t = __TRUNC + x = t + 1.0 + if __TRUNC_RECIP > z + alpha = 0.0 + rate = 1.0 + d_exp = Exponential(1.0 / rate) + while (rand(rng) > alpha) + e1 = rand(rng, d_exp) + e2 = rand(rng, d_exp) + while e1^2 > 2 * e2 / t + e1 = rand(rng, d_exp) + e2 = rand(rng, d_exp) + end + x = 1 + e1 * t + x = t / x^2 + alpha = exp(-0.5 * z^2 * x) + end + else + mu = 1.0 / z + while (x > t) + y = randn(rng)^2 + half_mu = 0.5 * mu + mu_Y = mu * y + x = mu + half_mu * mu_Y - half_mu * sqrt(4 * mu_Y + mu_Y^2) + if rand(rng) > mu / (mu + x) + x = mu^2 / x + end + end + end + return x +end + +# //////////////////////////////////////////////////////////////////////////////// +# // Sample // +# //////////////////////////////////////////////////////////////////////////////// + +function draw_like_devroye(rng::AbstractRNG, c::Real) + # Change the parameter. + c = 0.5 * abs(c) + + # Now sample 0.25 * J^*(1, Z := Z/2). + fz = 0.125 * π^2 + 0.5 * c^2 + # ... Problems with large Z? Try using q_over_p. + # double p = 0.5 * __PI * exp(-1.0 * fz * __TRUNC) / fz; + # double q = 2 * exp(-1.0 * Z) * pigauss(__TRUNC, Z); + + x = 0.0 + s = 1.0 + y = 0.0 + # int iter = 0; If you want to keep track of iterations. + d_exp = Exponential() + while true + if rand(rng) < mass_texpon(c) + x = __TRUNC + rand(rng, d_exp) / fz + else + x = rtigauss(rng, c) + end + s = a(0, x) + y = rand(rng) * s + n = 0 + go = true + + # Cap the number of iterations? + while (go) + n = n + 1 + if isodd(n) + s = s - a(n, x) + if y <= s + return 0.25 * x + end + else + s = s + a(n, x) + if y > s + go = false + end + end + end + # Need Y <= S in event that Y = S, e.g. when x = 0. + end +end # draw_like_devroye \ No newline at end of file diff --git a/research/sticky/comparisons/rev_jumps_vs_sticky.jl b/research/sticky/comparisons/rev_jumps_vs_sticky.jl new file mode 100644 index 00000000..4a4b6259 --- /dev/null +++ b/research/sticky/comparisons/rev_jumps_vs_sticky.jl @@ -0,0 +1,204 @@ +using ZigZagBoomerang +const Zig = ZigZagBoomerang +using ZigZagBoomerang: sλ, sλ̄, reflect!, Rng, ab, smove_forward!, neighbours +using Random +using StructArrays +using SparseArrays +using StructArrays: components +using LinearAlgebra +using ZigZagBoomerang: SPriorityQueue, enqueue!, lastiterate + +function next_rand_reflect(j, i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + if m[j] == 1 + return 0, Inf + end + b[j] = ab(G1, j, x, θ, c, F) + t_old[j] = t′ + return 0, t[j] + poisson_time(b[j], rand(P.rng)) +end + +function multiple_freeze!(ξξ, rev, i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + # find integer + fake = (t′ - t[i]) == 0.0 + ev, jj = findmin([norm((x[i] + θ[i]*(t′ - t[i]) - ξ)) for ξ in ξξ]) + @assert ev < 1e-7 + smove_forward!(G, i, t, x, θ, m, t′, F) + smove_forward!(G2, i, t, x, θ, m, t′, F) + + if m[i] == 0 # to freeze + if rev == true # tuning parameter of Chevallier + pp = 0.6 + else + pp = 1.0 + end + if rand() < pp + x[i] = ξξ[jj] + t[i] = t′ + m[i] = 1 + θ_old[i], θ[i] = θ[i], 0.0 + else + x[i] = ξξ[jj] + θ[i]*2*eps() + t[i] = t′ + end + else # to unfreeze + m[i] = 0 + t[i] = t′ + if rev == true && rand() < 0.5 #reversibility + θ[i] = -θ_old[i] + else + θ[i] = θ_old[i] + end + end + return true, neighbours(G1, i) +end + +function next_freezeunfreeze(ξξ, rev, κ, j, i, t′, u, P::SPDMP, args...) + t, x, θ, θ_old, m = components(u) + if m[j] == 0 + ev, kk = findmin([θ[j]*(x[j] - ξ) >= 0 ? Inf : t[j] - (x[j] - ξ)/θ[j] for ξ in ξξ]) + return 0, ev + else + if rev == true # tuning parameter of Chevallier + pp = 0.6 + else + pp = 1.0 + end + new_time = poisson_time(κ*pp, rand(P.rng)) + # println("rev new time equal to $(new_time)") + return 0, t[j] + new_time + end +end +# poisson_time(0.1, rand(P.rng)) +sticky_floors = collect(-5.0:0.01:5.0) +# sticky_floors = [0.0] +function multiple_rev_freeze!(args...) + multiple_freeze!(sticky_floors, true, args...) +end + +function multiple_nonrev_freeze!(args...) + multiple_freeze!(sticky_floors, false, args...) +end + +function next_freezeunfreeze_rev(args...) + next_freezeunfreeze(sticky_floors, true, 10.0, args...) +end + +function next_freezeunfreeze_nonrev(args...) + next_freezeunfreeze(sticky_floors, false, 10.0, args...) +end + + + + +T = 10000.0 +d = 20 +seed = (UInt(1),UInt(1)) +Random.seed!(1) + + + +S = I(d)*0.01 .+ 0.99 +Γ = sparse(inv(S)) + +∇ϕ(x, i, Γ) = Zig.idot(Γ, i, x) # sparse computation + +############# TEST 2 + +# t, x, θ, θ_old, m, c, t_old, b +t0 = 0.0 +t = fill(t0, d) +x = rand(d) +θ = θ0 = rand([-1.0, 1.0], d) +F = ZigZag(Γ, x*0) + +c = .01*[norm(Γ[:, i], 2) for i in 1:d] +G = G1 = [[i] => rowvals(F.Γ)[nzrange(F.Γ, i)] for i in eachindex(θ0)] +G2 = [i => setdiff(union((G1[j].second for j in G1[i].second)...), G[i].second) for i in eachindex(G1)] + +b = [ab(G1, i, x, θ, c, F) for i in eachindex(θ)] + +u0 = StructArray(t=t, x=x, θ=θ, θ_old = zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b) + + +rng = Rng(seed) +t_old = copy(t) +adapt = false +factor = 1.7 +P = SPDMP(G, G1, G2, ∇ϕ, F, rng, adapt, factor) + +action2! = (Zig.rand_reflect!, multiple_nonrev_freeze!,) +next_action2 = FunctionWrangler((next_rand_reflect, next_freezeunfreeze_nonrev,)) + +h2 = Schedule(action2!, next_action2, u0, T, (P, Γ)) +trc_2 = @time collect(h2); +trc2 = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_2]) +ts2, xs2 = Zig.sep(Zig.subtrace(trc2, [1,2])) + + +######################## Test 1 +# t, x, θ, θ_old, m, c, t_old, b +t0 = 0.0 +t = fill(t0, d) +x = rand(d) +θ = θ0 = rand([-1.0, 1.0], d) +F = ZigZag(Γ, x*0) + +c = .01*[norm(Γ[:, i], 2) for i in 1:d] +G = G1 = [[i] => collect(1:d) for i in eachindex(θ0)] +G2 = [i => setdiff(union((G1[j].second for j in G1[i].second)...), G[i].second) for i in eachindex(G1)] + +b = [ab(G1, i, x, θ, c, F) for i in eachindex(θ)] + +u0 = StructArray(t=t, x=x, θ=θ, θ_old = zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b) + + +rng = Rng(seed) +t_old = copy(t) +adapt = false +factor = 1.7 +P = SPDMP(G, G1, G2, ∇ϕ, F, rng, adapt, factor) + + +action1! = (Zig.rand_reflect!, multiple_rev_freeze!,) +next_action1 = FunctionWrangler((next_rand_reflect, next_freezeunfreeze_rev,)) + +h1 = Schedule(action1!, next_action1, u0, T, (P, Γ)) +trc_1 = @time collect(h1); +trc1 = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_1]) +ts1, xs1 = Zig.sep(Zig.subtrace(trc1, [1,2])) + + + + + +# error("") +using GLMakie + +# f1 = lines(ts1, getindex.(xs1, 2)) +# fig1 = lines(getindex.(xs1, 1), getindex.(xs1, 2)) +# save("rev_.png",fig1) +# save("rev_trace.png",f1) + +# f2 = lines(ts2, getindex.(xs2, 2)) +# fig2 = lines(getindex.(xs2, 1), getindex.(xs2, 2)) +# save("nonrev_.png",fig2) +# save("nonrev_trace.png",f2) + +fig = Figure() +ax1 = fig[1,1] = Axis(fig) +lines!(ax1, ts1, getindex.(xs1, 2)) +ax2 = fig[1,2] = Axis(fig) +fig1 = lines!(ax2, getindex.(xs1, 1), getindex.(xs1, 2)) + +ax3 = fig[2,1] = Axis(fig) +lines!(ax3, ts2, getindex.(xs2, 2)) +ax4 = fig[2,2] = Axis(fig) +fig1 = lines!(ax4, getindex.(xs2, 1), getindex.(xs2, 2)) +fig +save("trace_comparison.png", fig) \ No newline at end of file diff --git a/research/sticky/comparisons/sticky_varying_p.jl b/research/sticky/comparisons/sticky_varying_p.jl new file mode 100644 index 00000000..f588f6d3 --- /dev/null +++ b/research/sticky/comparisons/sticky_varying_p.jl @@ -0,0 +1,224 @@ +using ZigZagBoomerang +const Zig = ZigZagBoomerang +using ZigZagBoomerang: sλ, sλ̄, reflect!, Rng, ab, smove_forward!, neighbours +using Random +using StructArrays +using SparseArrays +using StructArrays: components +using LinearAlgebra +using ZigZagBoomerang: SPriorityQueue, enqueue!, lastiterate + +function next_rand_reflect(j, i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + if m[j] == 1 + return 0, Inf + end + b[j] = ab(G1, j, x, θ, c, F) + t_old[j] = t′ + return 0, t[j] + poisson_time(b[j], rand(P.rng)) +end + +function multiple_freeze!(ξξ, pp, i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + # find integer + ev, jj = findmin([norm((x[i] + θ[i]*(t′ - t[i]) - ξ)) for ξ in ξξ]) + @assert ev < 1e-7 + smove_forward!(G, i, t, x, θ, m, t′, F) + smove_forward!(G2, i, t, x, θ, m, t′, F) + if m[i] == 0 # to freeze + # if rev == true # tuning parameter of Chevallier + # pp = 0.6 + # else + # pp = 1.0 + # end + if rand() < pp + x[i] = ξξ[jj] + t[i] = t′ + m[i] = 1 + θ_old[i], θ[i] = θ[i], 0.0 + else + x[i] = ξξ[jj] + θ[i]*2*eps() + t[i] = t′ + end + else # to unfreeze + m[i] = 0 + t[i] = t′ + θ[i] = θ_old[i] + end + return true, neighbours(G1, i) +end + +function next_freezeunfreeze(ξξ, pp, κ, j, i, t′, u, P::SPDMP, args...) + t, x, θ, θ_old, m = components(u) + if m[j] == 0 + ev, kk = findmin([θ[j]*(x[j] - ξ) >= 0 ? Inf : t[j] - (x[j] - ξ)/θ[j] for ξ in ξξ]) + return 0, ev + else + # if rev == true # tuning parameter of Chevallier + # pp = 0.6 + # else + # pp = 1.0 + # end + new_time = poisson_time(κ*pp, rand(P.rng)) + # println("rev new time equal to $(new_time)") + return 0, t[j] + new_time + end +end + + +#### test 1 α = 0.1 +T = 5000.0 +d = 20 +seed = (UInt(1),UInt(1)) +Random.seed!(1) +S = I(d)*0.01 .+ 0.99 +Γ = sparse(inv(S)) +∇ϕ(x, i, Γ) = Zig.idot(Γ, i, x) # sparse computation +# poisson_time(0.1, rand(P.rng)) +sticky_floors = collect(-5.0:0.01:5.0) +# sticky_floors = [0.0] +pp = 0.1 +function multiple_freeze!(args...) + multiple_freeze!(sticky_floors, pp, args...) +end +function next_freezeunfreeze(args...) + next_freezeunfreeze(sticky_floors, pp, 10.0, args...) +end +t0 = 0.0 +t = fill(t0, d) +x = rand(d) +θ = θ0 = rand([-1.0, 1.0], d) +F = ZigZag(Γ, x*0) +c = .01*[norm(Γ[:, i], 2) for i in 1:d] +G = G1 = [[i] => rowvals(F.Γ)[nzrange(F.Γ, i)] for i in eachindex(θ0)] +G2 = [i => setdiff(union((G1[j].second for j in G1[i].second)...), G[i].second) for i in eachindex(G1)] +b = [ab(G1, i, x, θ, c, F) for i in eachindex(θ)] +u0 = StructArray(t=t, x=x, θ=θ, θ_old = zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b) +rng = Rng(seed) +adapt = false +factor = 1.7 +P = SPDMP(G, G1, G2, ∇ϕ, F, rng, adapt, factor) +action2! = (Zig.rand_reflect!, multiple_freeze!,) +next_action2 = FunctionWrangler((next_rand_reflect, next_freezeunfreeze,)) +h2 = Schedule(action2!, next_action2, u0, T, (P, Γ)) +trc_2 = @time collect(h2); +trc2 = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_2]) +ts1, xs1 = Zig.sep(Zig.subtrace(trc2, [1,2])) +using GLMakie +fig = Figure() +ax1 = fig[1,1] = Axis(fig) +lines!(ax1, ts1, getindex.(xs1, 2)) +ax2 = fig[1,2] = Axis(fig) +lines!(ax2, getindex.(xs1, 1), getindex.(xs1, 2)) + + + +#### test 2 α = 0.5 +T = 5000.0 +d = 20 +seed = (UInt(1),UInt(1)) +Random.seed!(1) +S = I(d)*0.01 .+ 0.99 +Γ = sparse(inv(S)) +∇ϕ(x, i, Γ) = Zig.idot(Γ, i, x) # sparse computation +# poisson_time(0.1, rand(P.rng)) +sticky_floors = collect(-5.0:0.01:5.0) +# sticky_floors = [0.0] +pp = 0.5 +function multiple_freeze!(args...) + multiple_freeze!(sticky_floors, pp, args...) +end +function next_freezeunfreeze(args...) + next_freezeunfreeze(sticky_floors, pp, 10.0, args...) +end +t0 = 0.0 +t = fill(t0, d) +x = rand(d) +θ = θ0 = rand([-1.0, 1.0], d) +F = ZigZag(Γ, x*0) +c = .01*[norm(Γ[:, i], 2) for i in 1:d] +G = G1 = [[i] => rowvals(F.Γ)[nzrange(F.Γ, i)] for i in eachindex(θ0)] +G2 = [i => setdiff(union((G1[j].second for j in G1[i].second)...), G[i].second) for i in eachindex(G1)] +b = [ab(G1, i, x, θ, c, F) for i in eachindex(θ)] +u0 = StructArray(t=t, x=x, θ=θ, θ_old = zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b) +rng = Rng(seed) +adapt = false +factor = 1.7 +P = SPDMP(G, G1, G2, ∇ϕ, F, rng, adapt, factor) +action2! = (Zig.rand_reflect!, multiple_freeze!,) +next_action2 = FunctionWrangler((next_rand_reflect, next_freezeunfreeze,)) +h2 = Schedule(action2!, next_action2, u0, T, (P, Γ)) +trc_2 = @time collect(h2); +trc2 = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_2]) +ts1, xs1 = Zig.sep(Zig.subtrace(trc2, [1,2])) +ax2 = fig[2,1] = Axis(fig) +lines!(ax2, ts1, getindex.(xs1, 2)) +ax3 = fig[2,2] = Axis(fig) +lines!(ax3, getindex.(xs1, 1), getindex.(xs1, 2)) + + + +#### test 3 α = 1.0 +T = 5000.0 +d = 20 +seed = (UInt(1),UInt(1)) +Random.seed!(1) +S = I(d)*0.01 .+ 0.99 +Γ = sparse(inv(S)) +∇ϕ(x, i, Γ) = Zig.idot(Γ, i, x) # sparse computation +# poisson_time(0.1, rand(P.rng)) +sticky_floors = collect(-5.0:0.01:5.0) +# sticky_floors = [0.0] +pp = 1.0 +function multiple_freeze!(args...) + multiple_freeze!(sticky_floors, pp, args...) +end +function next_freezeunfreeze(args...) + next_freezeunfreeze(sticky_floors, pp, 10.0, args...) +end +t0 = 0.0 +t = fill(t0, d) +x = rand(d) +θ = θ0 = rand([-1.0, 1.0], d) +F = ZigZag(Γ, x*0) +c = .01*[norm(Γ[:, i], 2) for i in 1:d] +G = G1 = [[i] => rowvals(F.Γ)[nzrange(F.Γ, i)] for i in eachindex(θ0)] +G2 = [i => setdiff(union((G1[j].second for j in G1[i].second)...), G[i].second) for i in eachindex(G1)] +b = [ab(G1, i, x, θ, c, F) for i in eachindex(θ)] +u0 = StructArray(t=t, x=x, θ=θ, θ_old = zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b) +rng = Rng(seed) +adapt = false +factor = 1.7 +P = SPDMP(G, G1, G2, ∇ϕ, F, rng, adapt, factor) +action2! = (Zig.rand_reflect!, multiple_freeze!,) +next_action2 = FunctionWrangler((next_rand_reflect, next_freezeunfreeze,)) +h2 = Schedule(action2!, next_action2, u0, T, (P, Γ)) +trc_2 = @time collect(h2); +trc2 = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_2]) +ts1, xs1 = Zig.sep(Zig.subtrace(trc2, [1,2])) +ax5= fig[3,1] = Axis(fig) +lines!(ax5, ts1, getindex.(xs1, 2)) +ax6 = fig[3,2] = Axis(fig) +lines!(ax6, getindex.(xs1, 1), getindex.(xs1, 2)) + +save("varying_p.png", fig) + + + + + + + + + + + + + + + + diff --git a/research/sticky/comparisons/varying_p.png b/research/sticky/comparisons/varying_p.png new file mode 100644 index 00000000..e06ff068 Binary files /dev/null and b/research/sticky/comparisons/varying_p.png differ diff --git a/research/sticky/trajectories/boomyphase.png b/research/sticky/trajectories/boomyphase.png index a8d1e24c..6806427d 100644 Binary files a/research/sticky/trajectories/boomyphase.png and b/research/sticky/trajectories/boomyphase.png differ diff --git a/research/sticky/trajectories/boomytrace.png b/research/sticky/trajectories/boomytrace.png index 079a93d2..dc46582a 100644 Binary files a/research/sticky/trajectories/boomytrace.png and b/research/sticky/trajectories/boomytrace.png differ diff --git a/research/sticky/trajectories/bouncyphase.png b/research/sticky/trajectories/bouncyphase.png index 6c440490..a715a8d8 100644 Binary files a/research/sticky/trajectories/bouncyphase.png and b/research/sticky/trajectories/bouncyphase.png differ diff --git a/research/sticky/trajectories/bouncytrace.png b/research/sticky/trajectories/bouncytrace.png index 2e4074b7..913b3214 100644 Binary files a/research/sticky/trajectories/bouncytrace.png and b/research/sticky/trajectories/bouncytrace.png differ diff --git a/research/sticky/trajectories/sticky_traj.jl b/research/sticky/trajectories/sticky_traj.jl index ea7ba1c9..d6d1ec72 100644 --- a/research/sticky/trajectories/sticky_traj.jl +++ b/research/sticky/trajectories/sticky_traj.jl @@ -1,12 +1,13 @@ using Pkg -Pkg.activate(joinpath(@__DIR__, "..", "..")) -using ZigZagBoomerang +#Pkg.activate(joinpath(@__DIR__, "..", "..")) +using ZigZagBoomerang +using Makie Pkg.activate(@__DIR__) cd(@__DIR__) #using Revise -using DataStructures +#using DataStructures using LinearAlgebra using Random using SparseArrays @@ -15,7 +16,7 @@ using FileIO using Statistics using Makie, AbstractPlotting using ForwardDiff -using StructArrays +#using StructArrays const ρ0 = 0.0 Σ = [1 ρ0; ρ0 1] const d = 2 @@ -51,15 +52,15 @@ function gradϕ!(y, x) end -κ = [0.1, 0.1] +κ = [Inf, Inf] t0 = 0.0 -x0 = randn(d) +x0 = rand(d) θ0 = [1.0, 1.0] c = 20.0 T = 1000.0 #sspdmp(∇ϕi, t0, x0, θ0, T, c*ones(d), ZigZag(sparse(I(d)), 0*x0), κ; adapt=false) -trace, (tT, xT, θT), (acc, num) = sspdmp(gradϕ, t0, x0, θ0, T, c*ones(d), ZigZag(sparse(I(d)), 0*x0), κ; adapt=false) +trace, (tT, xT, θT), (acc, num) = sspdmp(gradϕ, t0, x0, θ0, T, c*ones(d), ZigZag(sparse(I(d)), 0*x0), κ; reversible=true, adapt=false) #ts, xs = ZigZagBoomerang.sep(collect(discretize(trace, 0.05))) ts, xs = ZigZagBoomerang.sep(collect(trace)) @@ -68,15 +69,20 @@ p1 = lines(first.(xs), last.(xs), color=ts) scene, layout = layoutscene(resolution = (1200, 900)) -layout[1, 1] = ax1 = Axis(scene) -layout[2, 1] = ax2 = Axis(scene) +layout[1:2, 1] = ax0 = Axis(scene) +layout[1, 2] = ax1 = Axis(scene) +layout[2, 2] = ax2 = Axis(scene) +lines!(ax0, first.(xs), last.(xs), color=ts) linkyaxes!(ax1, ax2) linkxaxes!(ax1, ax2) lines!(ax1, ts, first.(xs)) lines!(ax2, ts, last.(xs)) p1a = scene + +error("here") + trace, (tT, xT, θT), (acc, num) = sspdmp(gradϕ!, t0, x0, θ0, T, c, BouncyParticle(sparse(I(d)), 0*x0, 0.1), κ; adapt=false) #ts, xs = ZigZagBoomerang.sep(collect(discretize(trace, 0.05))) diff --git a/research/sticky/trajectories/zigzagphase.png b/research/sticky/trajectories/zigzagphase.png index 13368a56..e6ecdf14 100644 Binary files a/research/sticky/trajectories/zigzagphase.png and b/research/sticky/trajectories/zigzagphase.png differ diff --git a/research/sticky/trajectories/zigzagtrace.png b/research/sticky/trajectories/zigzagtrace.png index 429b3d73..f7095c23 100644 Binary files a/research/sticky/trajectories/zigzagtrace.png and b/research/sticky/trajectories/zigzagtrace.png differ diff --git a/src/ZigZagBoomerang.jl b/src/ZigZagBoomerang.jl index 9270c191..85df9821 100644 --- a/src/ZigZagBoomerang.jl +++ b/src/ZigZagBoomerang.jl @@ -6,6 +6,12 @@ using ProgressMeter using RandomNumbers.Xorshifts using RandomNumbers: gen_seed const Rng = Xoroshiro128Plus +using FunctionWranglers +export FunctionWrangler +using ConcreteStructs + + + Seed() = gen_seed(UInt64, 2) # ZigZag1d and Boomerang1d reference implementation @@ -34,6 +40,12 @@ export pdmp, spdmp, eventtime, eventposition include("ss_fact.jl") export sspdmp +include("engine.jl") +export simulate, Schedule + +include("zigzag.jl") +export SPDMP + include("ss_not_fact.jl") export sticky_pdmp diff --git a/src/common.jl b/src/common.jl index bbb8cf9b..0cc53208 100644 --- a/src/common.jl +++ b/src/common.jl @@ -59,3 +59,24 @@ Splits a vector of pairs into a pair of vectors. """ splitpairs(x) = first.(x) => last.(x) export splitpairs + + +""" + lastiterate(itr) + +Last iterate of `itr`. +""" +function lastiterate(itr) + ϕ = iterate(itr) + if ϕ === nothing + error("empty") + end + x, state = ϕ + while true + ϕ = iterate(itr, state) + if ϕ === nothing + return x + end + x, state = ϕ + end +end diff --git a/src/discrete-continuous_zigzag.jl b/src/discrete-continuous_zigzag.jl new file mode 100644 index 00000000..7d8f2b34 --- /dev/null +++ b/src/discrete-continuous_zigzag.jl @@ -0,0 +1,108 @@ +using ZigZagBoomerang +const Zig = ZigZagBoomerang +using ZigZagBoomerang: sλ, sλ̄, reflect!, Rng, ab, smove_forward!, neighbours +using ZigZagBoomerang: next_rand_reflect, reflect! +using Random +using StructArrays +using StructArrays: components +using LinearAlgebra +using ZigZagBoomerang: SPriorityQueue, enqueue!, lastiterate + +import ZigZagBoomerang: next_rand_reflect, rand_reflect!, reflect!, reset! + + +function rand_reflect!(i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + smove_forward!(G, i, t, x, θ, m, t′, F) + ∇ϕi = P.∇ϕ[1](x, i, args...) #change just here from the original function + l, lb = sλ(∇ϕi, i, x, θ, F), sλ̄(b[i], t[i] - t_old[i]) + if rand(P.rng)*lb < l + if l >= lb + !P.adapt && error("Tuning parameter `c` too small.") + adapt!(c, i, P.factor) + end + smove_forward!(G2, i, t, x, θ, m, t′, F) + ZigZagBoomerang.reflect!(i, ∇ϕi, x, θ, F) + return true, neighbours(G1, i) + else + return false, G1[i].first + end + +end + +# random reflection of the Zig-Zag +function next_jump_move(j, i, t′, u, P::SPDMP, args...) + t, x, θ, θ_old, m, c, t_old, b = components(u) + 0, t[j] + randexp() # new exponential time +end + + +function jump_move!(i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + if d[i] == 0 # make sure that jumps are allowed + error("Discrete moves not available") + end + t, x, θ = smove_forward!(G, i, t, x, θ, m, t′, F) + x′ = x + θ # propose a jump + ϕx′ = P.∇ϕ[2](x′) + ϕx = P.∇ϕ[2](x) + if ϕx′ < ϕx || randexp() > ϕx′ - ϕx + x .= x′ # jump + else + θ[i] *= -1 + end + return true, neighbours(G1, i) +end + +# Discrete Gaussian random variable +# t, x, θ, c, t_old, b, d +d = 1 +t0 = 0.0 +t = fill(t0, d) +x = [1.0] +θ = θ0 = rand([-1.0, 1.0], d) +Γ = I(1) +F = ZigZag(Γ, x*0) +c = (zero(x) .+ 0.1) + +G = G1 = [[1] =>[1]] +G2 = [[1] => []] +b = [Zig.ab(G1, i, x, θ, c, F) for i in eachindex(θ)] # does not really matter + +u0 = StructArray(t=t, x=x, θ=θ, θ_old=zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b, θ2 = [1]) + +Random.seed!(1) +seed = (UInt(1),UInt(1)) +rng = Rng(seed) +T = 50.0 +t_old = copy(t) +adapt = false +factor = 1.7 +σ2 = 5.0 +∇ϕ(x,i) = x[i]/σ2 +ϕ(x) = sum(x.^2)/σ2 +P = SPDMP(G, G1, G2, (∇ϕ, ϕ), F, rng, adapt, factor) + + +action! = (rand_reflect!, jump_move!) +next_action = FunctionWrangler((next_rand_reflect, next_jump_move)) + +# action! = (reset!, rand_reflect!, circle_hit1!) +# next_action = FunctionWrangler((Zig.never_reset, next_rand_reflect, n +h = Schedule(action!, next_action, u0, T, (P, )) +trc_ = Zig.simulate(h, progress=true) +trc = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_]) + + + + + + +ts, xs = getindex.(trc.events,1), getindex.(trc.events,3) +# ts, xs = Zig.sep(Zig.discretize(trc, 0.001)) +using Makie +fig = scatter(ts, xs) \ No newline at end of file diff --git a/src/discrete_zigzag.jl b/src/discrete_zigzag.jl new file mode 100644 index 00000000..624d7473 --- /dev/null +++ b/src/discrete_zigzag.jl @@ -0,0 +1,91 @@ +using ZigZagBoomerang +const Zig = ZigZagBoomerang +using ZigZagBoomerang: sλ, sλ̄, reflect!, Rng, ab, smove_forward!, neighbours +using Random +using StructArrays +using SparseArrays +using StructArrays: components +using LinearAlgebra +using ZigZagBoomerang: SPriorityQueue, enqueue!, lastiterate + +# import standard reflection of zigzag + +# t, x, θ, θ_old, m, c, t_old, b, θ2 = components(u) +#introducing θ2 which is used for giving the direction when jumping + +# random reflection of the Zig-Zag +function next_discrete_move(j, i, t′, u, P::SPDMP, args...) + t, x, θ, θ_old, m, c, t_old, b, θ2 = components(u) + @assert m[j] == 1 + d[j] == 1 || return false, Inf # if it s not a discrete random variable, then return Inf + 0, t[j] + randexp() # new exponential time +end + + +function discrete_move!(i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b, θ2 = components(u) + if d[i] == 0 #make sure it s a discrete random variable + error("Discrete moves not available for continuous components") + end + @assert m[i] == 1 + t, x, θ = smove_forward!(G, i, t, x, θ, m, t′, F) + x′ = x + θ2 #propose a jump + ϕx′ = P.∇ϕ[2](x′) + ϕx = P.∇ϕ[2](x) + if ϕx′ < ϕx || randexp() > ϕx′ - ϕx + x .= x′ + else + θ2[i] *= -1 + end + return true, neighbours(G1, i) +end +# Discrete Gaussian random variable +# t, x, θ, c, t_old, b, d +d = 1 +t0 = 0.0 +t = fill(t0, d) +x = [1.0] +θ = θ0 = rand([-1.0, 1.0], d) +Γ = I(1) +F = ZigZag(Γ, x*0) +c = (zero(x) .+ 0.1) + +G = G1 = [[1] =>[1]] +G2 = [[1] => []] +b = [Zig.ab(G1, i, x, θ, c, F) for i in eachindex(θ)] # does not really matter + +u0 = StructArray(t=t, x=x, θ=θ, θ_old=zeros(d), m=ones(Int,d), c=c, t_old=copy(t), b=b, θd = [1]) + +Random.seed!(1) +seed = (UInt(1),UInt(1)) +rng = Rng(seed) +T = 10.0 +t_old = copy(t) +adapt = false +factor = 1.7 +σ2 = 5.0 +∇ϕ(x,i) = x[i]/σ2 +ϕ(x) = sum(x.^2)/σ2 +P = SPDMP(G, G1, G2, (∇ϕ, ϕ), F, rng, adapt, factor) + + +action! = (discrete_move!,) +next_action = FunctionWrangler((next_discrete_move,)) + +# action! = (reset!, rand_reflect!, circle_hit1!) +# next_action = FunctionWrangler((Zig.never_reset, next_rand_reflect, n +h = Schedule(action!, next_action, u0, T, (P, )) +trc_ = Zig.simulate(h, progress=true) +trc = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_]) + + + + + + +ts, xs = getindex.(trc.events,1), getindex.(trc.events,3) +# ts, xs = Zig.sep(Zig.discretize(trc, 0.001)) +using Makie +fig = lines(ts, xs) \ No newline at end of file diff --git a/src/engine.jl b/src/engine.jl new file mode 100644 index 00000000..1a6e1e55 --- /dev/null +++ b/src/engine.jl @@ -0,0 +1,128 @@ +using FunctionWranglers +using StructArrays +using StructArrays: components +using Random, LinearAlgebra +include("wranglers.jl") +include("switch.jl") + + +struct Schedule{T1,T2,T3,T4,T5} + action!::T1 + next_action::T2 + state::T3 + T::T4 + args::T5 +end +Base.eltype(::Type{Schedule{T1,T2,T3,T4,T5}}) where {T1,T2,T3,T4,T5} = Tuple{Float64, Int, eltype(T3), Int, Int} +function Base.iterate(handler::Schedule) + d = length(handler.state) + action = ones(Int, d) # reset! events to trigger next_action event computations + Q = SPriorityQueue(zeros(d)) + u = deepcopy(handler.state) + iterate(handler, (u, action, Q)) +end +Base.IteratorSize(::Type{<:Schedule}) = Base.SizeUnknown() +Base.IteratorEltype(::Type{<:Schedule}) = Base.HasEltype() + +function Base.iterate(handler, (u, action, Q)) + action! = handler.action! + next_action = handler.next_action + ev = handle!(u, action!, next_action, action, Q, handler.args...) + ev[1] > handler.T && return nothing + ev, (u, action, Q) +end + +@inline function queue_next!(J, next_action, Q, action, i, t′, u, args...) + for j in J # make typestable? + τ, e = _sfindmin(next_action, Inf, 0, 1, j, i, t′, u, args...) + Q[j] = τ + action[j] = e + end +end + +function simulate(handler; progress=true, progress_stops = 20) + T = handler.T + if progress + prg = Progress(progress_stops, 1) + else + prg = missing + end + stops = ismissing(prg) ? 0 : max(prg.n - 1, 0) # allow one stop for cleanup + tstop = T/stops + + u = deepcopy(handler.state) + action! = handler.action! + next_action = handler.next_action + d = length(handler.state) + action = ones(Int, d) # resets + + Q = SPriorityQueue(zeros(d)) + + ev = handle!(u, action!, next_action, action, Q, handler.args...) + evs = [ev] + while true + ev = handle!(u, action!, next_action, action, Q, handler.args...) + t′ = ev[1] + t′ > T && break + push!(evs, ev) + if t′ > tstop + tstop += T/stops + next!(prg) + end + end + ismissing(prg) || ProgressMeter.finish!(prg) + return evs +end + +#= +function handle!(u, action!, next_action, action, Q, args...) + # Who is (i) next_action, when (t′) and what (j) happens? + done = false + local e, t′, i + while !done + i, t′ = peek(Q) + e = action[i] + + if e == 1 + done, affected1 = action![1](i, t′, u, args...) + queue_next!(affected1, next_action, Q, action, i, t′, u, args...) + elseif e == 2 + done, affected2 = action![2](i, t′, u, args...) + queue_next!(affected2, next_action, Q, action, i, t′, u, args...) + elseif e == 3 + done, affected3 = action![3](i, t′, u, args...) + queue_next!(affected3, next_action, Q, action, i, t′, u, args...) + elseif e == 4 + done, affected4 = action![4](i, t′, u, args...) + queue_next!(affected4, next_action, Q, action, i, t′, u, args...) + elseif e == 5 + done, affected5 = action![5](i, t′, u, args...) + queue_next!(affected5, next_action, Q, action, i, t′, u, args...) + elseif e == 6 + done, affected6 = action![6](i, t′, u, args...) + queue_next!(affected6, next_action, Q, action, i, t′, u, args...) + end + end + + (t′, i, u[i], action[i], num) +end +=# + +function handle!(u, action!, next_action, action, Q, args::Vararg{Any, N}) where {N} + # Who is (i) next_action, when (t′) and what (j) happens? + done = false + local e, t′, i + num = 0 + while !done + num += 1 + i, t′ = peek(Q) + e = action[i] + + #done = action_nextaction(action!, next_action, Q, action, e, i, t′, u, args...) + done = switch(e, action!, next_action, (Q, action), i, t′, u, args...) + end + (t′, i, u[i], action[i], num) +end + + + diff --git a/src/engineex.jl b/src/engineex.jl new file mode 100644 index 00000000..c4fcd26d --- /dev/null +++ b/src/engineex.jl @@ -0,0 +1,98 @@ +using ZigZagBoomerang +using StructArrays +using StructArrays: components +using Random +using LinearAlgebra + +function lastiterate(itr) + ϕ = iterate(itr) + if ϕ === nothing + error("empty") + end + x, state = ϕ + while true + ϕ = iterate(itr, state) + if ϕ === nothing + return x + end + x, state = ϕ + end +end +using Random +#include("engine.jl") +T = 100.0 +d = 10000 +function reset!(i, t′, u, args...) + false, i +end + +function next_reset(j, i, t′, u, args...) + 0, Inf +end + +function action1!(i, t′, u, args...) + t, x, θ, m = components(u) + x[i] += θ[i]*(t′ - t[i]) + # if rand() < 0.1 + θ[i] = -abs(θ[i]) + # end + t[i] = t′ + x[i] = 1.0 + true, i + +end + +function action2!(i, t′, u, args...) + t, x, θ, m = components(u) + x[i] += θ[i]*(t′ - t[i]) + @assert norm(x[i]) < 1e-7 + θ[i] = abs(θ[i]) + x[i] = 0.0 + t[i] = t′ + true, i +end + +function next_action1(j, i, t′, u, args...) + t, x, θ, m = components(u) + # t[j] + randexp() + 0, θ[j]*(x[j]-1) >= 0 ? Inf : t[j] - (x[j]-1)/θ[j] +end +function next_action2(j, i, t′, u, args...) + t, x, θ, m = components(u) + 0, θ[j]*x[j] >= 0 ? Inf : t[j] - x[j]/θ[j] +end + + + +action! = (reset!, action1!, action2!) +#action! = FunctionWrangler(action!) +next_action = FunctionWrangler((next_reset, next_action1, next_action2)) + +u0 = StructArray(t=zeros(d), x=zeros(d), θ=ones(d), m=zeros(Int,d)) +u0.θ[1] = 1/Base.MathConstants.golden +Random.seed!(1) +h = Schedule(action!, next_action, u0, T, ()) +l_ = lastiterate(h) +l_ = @time lastiterate(h) + +Random.seed!(1) +l = simulate(h) +@assert l[3].x == l_[3].x +l = @time simulate(h) +trc_ = @time collect(h) +#@code_warntype handler(zeros(d), T, (f1!, f2!)); + +#using ProfileView + +#ProfileView.@profview handler(zeros(d), 10T); + +subtrace = [t for t in trc_ if t[2] == 1] +lines(getindex.(subtrace, 1), getfield.(getindex.(subtrace, 3), :x)) + +trc = Zig.FactTrace(ZigZag(sparse(1.0I(d)), zeros(d)), t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_]) +ts, xs = Zig.sep(discretize(Zig.subtrace(trc, [1,4]), 0.01)) + +cummean(x) = cumsum(x) ./ eachindex(x) +fig = lines(ts, getindex.(xs, 2)) +lines!(ts, cummean(getindex.(xs, 2)), color=:green, linewidth=2.0) +fig = lines(getindex.(xs, 1), getindex.(xs, 2), color=:red) \ No newline at end of file diff --git a/src/priorityqueue.jl b/src/priorityqueue.jl index 0d32f6f8..ea7693d1 100644 --- a/src/priorityqueue.jl +++ b/src/priorityqueue.jl @@ -20,24 +20,14 @@ struct SPriorityQueue{K,V,O<:Ordering} <: AbstractDict{K,V} SPriorityQueue{K, V, O}(xs::Array{Pair{K,V}, 1}, o::O, index) where {K,V,O<:Ordering} = new(xs, o, index) - function SPriorityQueue{K,V,O}(o::O, itr) where {K,V,O<:Ordering} - xs = Vector{Pair{K,V}}(undef, length(itr)) - index = Dict{K, Int}() - for (i, (k, v)) in enumerate(itr) - xs[i] = Pair{K,V}(k, v) - push!(index, i) - end - pq = new{K,V,O}(xs, o, index) - - # heapify - for i in heapparent(length(pq.xs)):-1:1 - percolate_down!(pq, i) - end - - return pq +end +function SPriorityQueue(τ) + Q = SPriorityQueue{Int,eltype(τ)}() + for i in eachindex(τ) + enqueue!(Q, i => τ[i]) end + Q end - Base.length(pq::SPriorityQueue) = length(pq.xs) Base.isempty(pq::SPriorityQueue) = isempty(pq.xs) @@ -45,9 +35,10 @@ Base.peek(pq::SPriorityQueue) = pq.xs[1] function percolate_down!(pq::SPriorityQueue, i::Integer) x = pq.xs[i] - @inbounds while (l = heapleft(i)) <= length(pq) + L = length(pq) + @inbounds while (l = heapleft(i)) <= L r = heapright(i) - j = r > length(pq) || lt(pq.o, pq.xs[l].second, pq.xs[r].second) ? l : r + j = r > L || lt(pq.o, pq.xs[l].second, pq.xs[r].second) ? l : r if lt(pq.o, pq.xs[j].second, x.second) pq.index[pq.xs[j].first] = i pq.xs[i] = pq.xs[j] diff --git a/src/rev_jumps_vs_sticky.jl b/src/rev_jumps_vs_sticky.jl new file mode 100644 index 00000000..57a89a30 --- /dev/null +++ b/src/rev_jumps_vs_sticky.jl @@ -0,0 +1,176 @@ +using ZigZagBoomerang +const Zig = ZigZagBoomerang +using ZigZagBoomerang: sλ, sλ̄, reflect!, Rng, ab, smove_forward!, neighbours +using Random +using StructArrays +using SparseArrays +using StructArrays: components +using LinearAlgebra +using ZigZagBoomerang: SPriorityQueue, enqueue!, lastiterate + +function next_rand_reflect(j, i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + if m[j] == 1 + return 0, Inf + end + b[j] = ab(G1, j, x, θ, c, F) + t_old[j] = t′ + return 0, t[j] + poisson_time(b[j], rand(P.rng)) +end + +function multiple_freeze!(ξξ, rev, i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + # find integer + fake = (t′ - t[i]) == 0.0 + ev, jj = findmin([norm((x[i] + θ[i]*(t′ - t[i]) - ξ)) for ξ in ξξ]) + @assert ev < 1e-7 + smove_forward!(G, i, t, x, θ, m, t′, F) + smove_forward!(G2, i, t, x, θ, m, t′, F) + + if m[i] == 0 # to freeze + x[i] = ξξ[jj] + t[i] = t′ + m[i] = 1 + θ_old[i], θ[i] = θ[i], 0.0 + else # to unfreeze + m[i] = 0 + t[i] = t′ + if rev == true && rand() < 0.5 && fake == false #reversibility + θ[i] = -θ_old[i] + else + θ[i] = θ_old[i] + end + end + return true, neighbours(G1, i) +end + +function next_freezeunfreeze(ξξ, rev, κ, j, i, t′, u, P::SPDMP, args...) + t, x, θ, θ_old, m = components(u) + if m[j] == 0 + ev, kk = findmin([θ[j]*(x[j] - ξ) >= 0 ? Inf : t[j] - (x[j] - ξ)/θ[j] for ξ in ξξ]) + return 0, ev + else + if rev == true # tuning parameter of Chevallier + if rand() > 0.6 + return 0, t[j] #don't stick + else + return 0, t[j] + poisson_time(κ*0.6, rand(P.rng)) + end + else # rev = false + return 0, t[j] + poisson_time(κ, rand(P.rng)) + end + end +end + +sticky_floors = collect(-5.0:0.5:5.0) + +function multiple_rev_freeze!(args...) + multiple_freeze!(sticky_floors, true, args...) +end + +function multiple_nonrev_freeze!(args...) + multiple_freeze!(sticky_floors, false, args...) +end + +function next_freezeunfreeze_rev(args...) + next_freezeunfreeze(sticky_floors, true, 1.0, args...) +end + +function next_freezeunfreeze_nonrev(args...) + next_freezeunfreeze(sticky_floors, false, 1.0, args...) +end + + +T = 100.0 +d = 50 +seed = (UInt(1),UInt(1)) +Random.seed!(1) + + + +S = I(d)*0.1 .+ 0.9 +Γ = inv(S) + +∇ϕ(x, i, Γ) = Zig.idot(Γ, i, x) # sparse computation + +# t, x, θ, θ_old, m, c, t_old, b +t0 = 0.0 +t = fill(t0, d) +x = rand(d) +θ = θ0 = rand([-1.0, 1.0], d) +F = ZigZag(Γ, x*0) + + + +c = .01*[norm(Γ[:, i], 2) for i in 1:d] +G = G1 = [[i] => collect(1:d) for i in eachindex(θ0)] +G2 = [i => setdiff(union((G1[j].second for j in G1[i].second)...), G[i].second) for i in eachindex(G1)] + +b = [ab(G1, i, x, θ, c, F) for i in eachindex(θ)] + +u0 = StructArray(t=t, x=x, θ=θ, θ_old = zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b) + + +rng = Rng(seed) +t_old = copy(t) +adapt = false +factor = 1.7 +P = SPDMP(G, G1, G2, ∇ϕ, F, rng, adapt, factor) + +action1! = (Zig.rand_reflect!, multiple_rev_freeze!,) +next_action1 = FunctionWrangler((next_rand_reflect, next_freezeunfreeze_rev,)) + + +# h = Schedule(action!, next_action, u0, T, (P, Γ)) + +# l_ = lastiterate(h) + +# using ProfileView, Profile +# Profile.init(10000000, 0.00001) +# ProfileView.@profview lastiterate(h) +# l_ = @time lastiterate(h) + +h1 = Schedule(action1!, next_action1, u0, T, (P, Γ)) +trc_1 = @time collect(h1); +trc1 = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_1]) + + +action2! = (Zig.rand_reflect!, multiple_nonrev_freeze!,) +next_action2 = FunctionWrangler((next_rand_reflect, next_freezeunfreeze_nonrev,)) + +h2 = Schedule(action2!, next_action2, u0, T, (P, Γ)) +trc_2 = @time collect(h2); +trc2 = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_2]) + +# h = Schedule(action!, next_action, u0, T, (P, Γ)) + +# l_ = lastiterate(h) + +# using ProfileView, Profile +# Profile.init(10000000, 0.00001) +# ProfileView.@profview lastiterate(h) +# l_ = @time lastiterate(h) + + + +#trace, _, acc = @time spdmp(∇ϕ, t0, x, θ, T, c, G, F, Γ); +#@code_warntype handler(zeros(d), T, (f1!, f2!)); + +#using ProfileView +error() +#ProfileView.@profview handler(zeros(d), 10T); +using GLMakie +#subtrace1 = [t for t in trc_ if t[2] == 1] +#lines(getindex.(subtrace1, 1), getfield.(getindex.(subtrace1, 3), :x)) + +ts1, xs1 = Zig.sep(Zig.subtrace(trc1, [1,4])) +lines(ts1, getindex.(xs1, 2)) +fig = lines(getindex.(xs1, 1), getindex.(xs1, 2)) + +ts2, xs2 = Zig.sep(Zig.subtrace(trc2, [1,4])) +lines(ts2, getindex.(xs2, 2)) +fig = lines(getindex.(xs2, 1), getindex.(xs2, 2)) \ No newline at end of file diff --git a/src/switch.jl b/src/switch.jl new file mode 100644 index 00000000..970e7334 --- /dev/null +++ b/src/switch.jl @@ -0,0 +1,59 @@ + + +@inline @generated function _switch(i, f, g, q, args...) + argnames = [:(args[$i]) for i = 1:length(args)] + n = length(f.parameters) + qu = quote + if i == $n + uJ = f[$n]($(argnames...)) + for j in uJ[2] + τ, e = _sfindmin(g, Inf, 0, 1, j, $(argnames...)) + q[1][j] = τ + q[2][j] = e + end + uJ[1] + else + throw(BoundsError()) + end + end + for k in n-1:-1:1 + qu = quote + if i == $k + uJ = f[$k]($(argnames...)) + for j in uJ[2] + τ, e = _sfindmin(g, Inf, 0, 1, j, $(argnames...)) + q[1][j] = τ + q[2][j] = e + end + uJ[1] + else $qu + end + end + end + return qu +end +switch(i, f, g, q, args...) = _switch(i, f, g, q, args...) +#time switch(2, (+, -), nothing, nothing, 1, 2) + + +@inline @generated function _switch1(i, f, args...) + argnames = [:(args[$i]) for i = 1:length(args)] + n = length(f.parameters) + qu = quote + if i == $n + f[$n]($(argnames...)) + else + BoundsError() + end + end + for k in n-1:-1:1 + qu = quote + if i == $k + f[$k]($(argnames...)) + else $qu + end + end + end + return qu +end +switch1(i, f, args...) = _switch1(i, f, args...) diff --git a/src/trace.jl b/src/trace.jl index 89133b86..d740c6e5 100644 --- a/src/trace.jl +++ b/src/trace.jl @@ -13,7 +13,7 @@ struct FactTrace{FT,T,S,S2,R} <: Trace end """ - FactTrace + PDMPTrace See [`Trace`](@ref). """ diff --git a/src/wranglers.jl b/src/wranglers.jl new file mode 100644 index 00000000..718d951e --- /dev/null +++ b/src/wranglers.jl @@ -0,0 +1,46 @@ +@inline @generated function _action_nextaction(wrangler::FunctionWrangler{TOp, TNext}, next_action, Q, action, idx, args...) where {TNext, TOp} + argnames = [:(args[$i]) for i = 1:length(args)] + return quote + if idx == 1 + v, affected = wrangler.op($(argnames...)) + for j in affected + #τ, e = sfindmin(next_action, j, $(argnames...)) + τ, e = _sfindmin(next_action, Inf, 0, 1, j, $(argnames...)) + Q[j] = τ + action[j] = e + end + v + else + return _action_nextaction(wrangler.next, next_action, Q, action, idx - 1, $(argnames...)) + end + end +end + +""" + sindex(wrangler::FunctionWrangler, idx, args...) + +Call the `idx`-th function with args. +Note that this call iterates the wrangler from 1 to `idx`. Try to +put the frequently called functions at the beginning to minimize overhead. +""" +function action_nextaction(wrangler::FunctionWrangler, next_action, Q, action, idx, args...) + @assert idx <= length(wrangler) + _action_nextaction(wrangler, next_action, Q, action, idx, args...) +end + +@inline @generated function _sfindmin(wrangler::FunctionWrangler{TOp, TNext}, tmin, j, n, args...) where {TNext, TOp} + TOp === Nothing && return :(tmin, j) + argnames = [:(args[$i]) for i = 1:length(args)] + return quote + k, t = wrangler.op($(argnames...)) + tmin, j = t ≤ tmin ? (t, k > 0 ? k : n) : (tmin, j) + return _sfindmin(wrangler.next, tmin, j, n + 1, $(argnames...)) + end +end + +""" + sfindmin(wrangler::FunctionWrangler, args...) + +Look for the function which returns smallest value for the given arguments, and returns its index. +""" +sfindmin(wrangler::FunctionWrangler, args...) = _sfindmin(wrangler, Inf, 0, 1, args...) diff --git a/src/zigzag.jl b/src/zigzag.jl new file mode 100644 index 00000000..015c7e09 --- /dev/null +++ b/src/zigzag.jl @@ -0,0 +1,172 @@ +using Random +using ConcreteStructs +@concrete struct SPDMP + G + G1 + G2 + ∇ϕ + F + rng + adapt + factor +end + +function ZigZagBoomerang.smove_forward!(G, i, t, x, θ, m, t′, Z::Union{BouncyParticle, ZigZag}) + nhd = neighbours(G, i) + for i in nhd + if m[i] != 1 # not frozen + t[i], x[i] = t′, x[i] + θ[i]*(t′ - t[i]) + else # frozen + t[i] = t′ + end + end + t, x, θ +end +function ZigZagBoomerang.smove_forward!(G, i, t, x, θ, m, t′, B::Union{Boomerang, FactBoomerang}) + nhd = neighbours(G, i) + for i in nhd + τ = t′ - t[i] + s, c = sincos(τ) + if m[i] != 1 # not frozen + t[i], x[i], θ[i] = t′, (x[i] - B.μ[i])*c + θ[i]*s + B.μ[i], + -(x[i] - B.μ[i])*s + θ[i]*c + else # frozen + t[i] = t′ + end + end + t, x, θ +end + +function reset!(i, t′, u, P::SPDMP, args...) + t, x, θ, θ_old, m = components(u) + smove_forward!(P.G, i, t, x, θ, m, t′, P.F) + smove_forward!(P.G2, i, t, x, θ, m, t′, P.F) + + false, P.G1[i].first +end + +function never_reset(j, _, t′, u, P, args...) + 0, Inf +end + + +function rand_reflect!(i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + smove_forward!(G, i, t, x, θ, m, t′, F) + ∇ϕi = P.∇ϕ(x, i, args...) + l, lb = sλ(∇ϕi, i, x, θ, F), sλ̄(b[i], t[i] - t_old[i]) + if rand(P.rng)*lb < l + if l >= lb + !P.adapt && error("Tuning parameter `c` too small.") + adapt!(c, i, P.factor) + end + smove_forward!(G2, i, t, x, θ, m, t′, F) + ZigZagBoomerang.reflect!(i, ∇ϕi, x, θ, F) + return true, neighbours(G1, i) + else + return false, G1[i].first + end + +end + +function freeze!(ξ, i, t′, u, P::SPDMP, args...) + + + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + + @assert norm(x[i] + θ[i]*(t′ - t[i]) - ξ) < 1e-7 + smove_forward!(G, i, t, x, θ, m, t′, F) + smove_forward!(G2, i, t, x, θ, m, t′, F) + + if m[i] == 0 # to freeze + x[i] = ξ + t[i] = t′ + m[i] = 1 + θ_old[i], θ[i] = θ[i], 0.0 + else # to unfreeze + m[i] = 0 + t[i] = t′ + θ[i] = θ_old[i] + end + + return true, neighbours(G1, i) + +end + +function discontinuity_at!(ξ, a, dir, i, t′, u, P::SPDMP, args...) + + + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + + @assert norm(x[i] + θ[i]*(t′ - t[i]) - ξ) < 1e-7 + smove_forward!(G, i, t, x, θ, m, t′, F) + + x[i] = ξ + t[i] = t′ + if dir*θ[i] > 0 && rand(P.rng) < a + θ[i] *= -1 + smove_forward!(G2, i, t, x, θ, m, t′, F) + return true, neighbours(G1, i) + else + return false, G1[i].first + end +end + +function reflect!(ξ, dir, i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + + @assert dir*(x[i] + θ[i]*(t′ - t[i]) - ξ) < 1e-7 + smove_forward!(G, i, t, x, θ, m, t′, F) + smove_forward!(G2, i, t, x, θ, m, t′, F) + + x[i] = ξ + θ[i] = dir*abs(θ[i]) + t[i] = t′ + return true, neighbours(G1, i) +end + + +function next_rand_reflect(j, i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + if m[j] == 1 + return 0, Inf + end + b[j] = ab(G1, j, x, θ, c, F) + t_old[j] = t′ + 0, t[j] + poisson_time(b[j], rand(P.rng)) +end + +function next_reflect(ξ, dir, j, i, t′, u, P::SPDMP, args...) + t, x, θ, θ_old, m = components(u) + + if dir*(x[j] - ξ) < 0 + return 0, t[j] + end + 0, θ[j]*(x[j] - ξ) >= 0 ? Inf : t[j] - (x[j] - ξ)/θ[j] +end + + +function next_hit(ξ, j, i, t′, u, P::SPDMP, args...) + t, x, θ, θ_old, m = components(u) + 0, θ[j]*(x[j] - ξ) >= 0 ? Inf : t[j] - (x[j] - ξ)/θ[j] +end + + +function next_freezeunfreeze(ξ, κ, j, i, t′, u, P::SPDMP, args...) + t, x, θ, θ_old, m = components(u) + if m[j] == 0 + 0, θ[j]*(x[j] - ξ) >= 0 ? Inf : t[j] - (x[j] - ξ)/θ[j] + else + 0, t[j] + poisson_time(κ, rand(P.rng)) + end +end diff --git a/src/zigzag_circ.jl b/src/zigzag_circ.jl new file mode 100644 index 00000000..e43c01ac --- /dev/null +++ b/src/zigzag_circ.jl @@ -0,0 +1,99 @@ +# flexible functions defining boundaries of the form of a $d$-dimensional ball +using ZigZagBoomerang +const Zig = ZigZagBoomerang +using ZigZagBoomerang: sλ, sλ̄, reflect!, Rng, ab, smove_forward!, neighbours +using ZigZagBoomerang: next_rand_reflect, reflect! +using Random +using StructArrays +using StructArrays: components +using LinearAlgebra +using ZigZagBoomerang: SPriorityQueue, enqueue!, lastiterate + +# coefficients of the quadratic equation coming for the condition \|x + θ*t - μ\|^2 = rsq +function abc_eq2d(μ, rsq, x, θ,) + a = θ[1]^2 + θ[2]^2 + b = 2*((x[1] - μ[1]) *θ[1] + (x[2] - μ[2])*θ[2]) + c = (x[1]-μ[1])^2 + (x[2]-μ[2])^2 - rsq + a, b, c +end + + +# joint reflection at the boundary for the Zig-Zag sampler +# dir = 1 you want to be outside the circle +# dir = -1 you want to be inside the circle +function circle_boundary_reflection!(μ, rsq, dir, x, θ) + for i in eachindex(x) + if dir*θ[i]*(x[i]-μ[i]) < 0.0 + θ[i]*=-1 + end + end + θ +end + + +# dir = 1 you want to be outside the circle +# dir = -1 you want to be inside the circle +function next_circle_hit(μ, rsq, dir, j, i, t′, u, P::SPDMP, args...) + if j <= length(u) # not applicable + return 0, Inf + else #hitting time to the ball with radius `radius` + t, x, θ, θ_old, m, c, t_old, b = components(u) + a1, a2, a3 = abc_eq2d(μ, rsq, x, θ) #solving quadradic equation + dis = a2^2 - 4a1*a3 #discriminant + # no solutions or TABU region + if dis <= 1e-7 || dir*((x[1] - μ[1])^2 + (x[2] - μ[2])^2 - rsq - 1e-7) < 0.0 + return 0, Inf + else #pick the first positive event time + if dir == 1 + hitting_time = min((-a2 - sqrt(dis))/(2*a1),(-a2 + sqrt(dis))/(2*a1)) + if hitting_time <= 0.0 + return 0, Inf + end + else #dir == -1 + hitting_time = max((-a2 - sqrt(dis))/(2*a1),(-a2 + sqrt(dis))/(2*a1)) + @assert hitting_time > 0 + end + return 0, t′ + hitting_time #hitting time + end + end +end + +# either standard reflection, or bounce at the boundary or traverse the boundary +# α(μ, rsq, x, v) teleportation. Default value is no teleportation +function circle_hit!(α, μ, rsq, dir, i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + + if i > length(u) #hitting boundary + smove_forward!(G, i, t, x, θ, m, t′, F) + if abs((x[1] - μ[1])^2 + (x[2] - μ[2])^2 - rsq) > 1e-7 # make sure to hit be on the circle + dump(u) + error("not on the circle") + end + if α == nothing + θ .= circle_boundary_reflection!(μ, rsq, dir, x, θ) + else + xnew, θnew = α(μ, rsq, x, θ) # α(x) = -x + 2*μ + disc = ϕ(x) - ϕ(xnew) # magnitude of the discontinuity + if disc < 0.0 || rand() > 1 - exp(-disc) # teleport + # if false # never teleport + # jump on the other side drawing a line passing through the center of the ball + x .= xnew # improve by looking at G[3] + θ .= θnew # improve by looking at G[3] + else # bounce off + θ .= circle_boundary_reflection!(μ, rsq, dir, x, θ) + end + end + return true, neighbours(G1, i) + else + error("action not available for clock $i") + end +end + +# draw circle +function draw_circ(μ, rsq) + r = sqrt(rsq) + θ = LinRange(0,2π, 1000) + μ[1] .+ r*sin.(θ), μ[2] .+ r*cos.(θ) +end diff --git a/src/zigzag_circ_ex.jl b/src/zigzag_circ_ex.jl new file mode 100644 index 00000000..c65f0fd7 --- /dev/null +++ b/src/zigzag_circ_ex.jl @@ -0,0 +1,195 @@ +using ZigZagBoomerang +const Zig = ZigZagBoomerang +using ZigZagBoomerang: sλ, sλ̄, reflect!, Rng, ab, smove_forward!, neighbours +using ZigZagBoomerang: next_rand_reflect, reflect! +using Random +using StructArrays +using SparseArrays +using StructArrays: components +using LinearAlgebra +using ZigZagBoomerang: SPriorityQueue, enqueue!, lastiterate +# import standard reflection of zigzag +import ZigZagBoomerang: next_rand_reflect, rand_reflect!, reflect!, reset! + + +# overload simulate and handle! to allow having 2 more clocks than coordinates +# from `d` to `d+2` +function Zig.simulate(handler; progress=true, progress_stops = 20) + T = handler.T + if progress + prg = Zig.Progress(progress_stops, 1) + else + prg = missing + end + stops = ismissing(prg) ? 0 : max(prg.n - 1, 0) # allow one stop for cleanup + tstop = T/stops + + u = deepcopy(handler.state) + action! = handler.action! + next_action = handler.next_action + d = length(handler.state) + action = ones(Int, d+2) # resets only + + Q = Zig.SPriorityQueue(zeros(d+2)) + + evs = Zig.handle!(u, action!, next_action, action, Q, handler.args...) + # evs = [ev] + while true + ev = Zig.handle!(u, action!, next_action, action, Q, handler.args...) + t′ = ev[end][1] + t′ > T && break + append!(evs, ev) + if t′ > tstop + tstop += T/stops + Zig.next!(prg) + end + end + ismissing(prg) || Zig.ProgressMeter.finish!(prg) + return evs +end + +function Zig.handle!(u, action!, next_action, action, Q, args::Vararg{Any, N}) where {N} + # Who is (i) next_action, when (t′) and what (j) happens? + done = false + local e, t′, i + num = 0 + while !done + num += 1 + i, t′ = Zig.peek(Q) + e = action[i] + #done = action_nextaction(action!, next_action, Q, action, e, i, t′, u, args...) + done = Zig.switch(e, action!, next_action, (Q, action), i, t′, u, args...) + end + traceevent(t′, i, u, action, num) +end + +function traceevent(t′, i, u, action, num) + if 1 <= i <= length(u) + return [(t′, i, u[i], action[i], num)] + else + return [(t′, 1, u[1], action[1], num), (t′, 2, u[2], action[2], num)] + end +end + +# random reflection of the Zig-Zag +function next_rand_reflect(j, i, t′, u, P::SPDMP, args...) + 1 <= j <= length(u) || return 0, Inf + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + if m[j] == 1 + return 0, Inf + end + b[j] = ab(G1, j, x, θ, c, F) + t_old[j] = t′ + 0, t[j] + poisson_time(b[j], rand(P.rng)) +end + +function rand_reflect!(i, t′, u, P::SPDMP, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + @assert 1 <= i <= length(u) + smove_forward!(G, i, t, x, θ, m, t′, F) + ∇ϕi = P.∇ϕ(x, i, args...) + l, lb = sλ(∇ϕi, i, x, θ, F), sλ̄(b[i], t[i] - t_old[i]) + if rand(P.rng)*lb < l + if l >= lb + !P.adapt && error("Tuning parameter `c` too small.") + adapt!(c, i, P.factor) + end + smove_forward!(G2, i, t, x, θ, m, t′, F) + ZigZagBoomerang.reflect!(i, ∇ϕi, x, θ, F) + return true, neighbours(G1, i) + else + return false, G1[i].first + end + +end + +# reset +function reset!(i, t′, u, P::SPDMP, args...) # should move the coordinates if called later in the game + false, P.G1[i].first +end + + + +include("./zigzag_circ.jl") +# stay outside ball centered at μ1 with radius rsq1, if hit then teleport or reflect +# α(μ, rsq, x, v) = (-x + 2*μ, v) # teleportation +α = nothing # no teleportation +next_circle_hit1(j, i, t′, u, P::SPDMP, args...) = next_circle_hit([-0.7, -0.7], 0.5, 1, j, i, t′, u, P::SPDMP, args...) +circle_hit1!(i, t′, u, P::SPDMP, args...) = circle_hit!(α, [-0.7, -0.7], 0.5, 1, i, t′, u, P::SPDMP, args...) + +# stay outside ball centered at μ1 with radius rsq1, if hit then teleport or reflect +next_circle_hit2(j, i, t′, u, P::SPDMP, args...) = next_circle_hit([0.0, 0.0], 4.0, -1, j, i, t′, u, P::SPDMP, args...) +circle_hit2!(i, t′, u, P::SPDMP, args...) = circle_hit!(α, [0.0, 0.0], 4.0, -1, i, t′, u, P::SPDMP, args...) + +Sigma = Matrix([1.0 0.0; + 0.0 1.0]) + +Γ = sparse(Sigma^(-1)) + +ϕ(x) = -0.5*x'*Γ*x # negated log-density +∇ϕ(x, i) = Zig.idot(Γ, i, x) # sparse computation + +# t, x, θ, c, t_old, b +d = 2 +t0 = 0.0 +t = fill(t0, d) +x = [1.0, 1.0] + rand(d) +θ = θ0 = rand([-1.0, 1.0], d) +F = ZigZag(Γ, x*0) + + + +# G, G2: Clocks to coordinates +# G1: Clocks to clocks + +c = (zero(x) .+ 0.1) +G = [[i] => rowvals(F.Γ)[nzrange(F.Γ, i)] for i in eachindex(θ0)] +push!(G, [3] => [1,2]) # move all coordinates +push!(G, [4] => [1,2]) # move all coordinates +G1 = [[i] => [rowvals(F.Γ)[nzrange(F.Γ, i)]; [3, 4]] for i in eachindex(θ0)] +push!(G1, [3] => [1, 2, 3, 4]) # invalidate ALL the clocks yeah +push!(G1, [4] => [1, 2, 3, 4]) # invalidate ALL the clocks yeah +# push!(G1, [3] => [1, 2, 3]) # invalidate ALL the clocks yeah + +# G2 = [1 => [2], 2 => [1], 3=>Int[]] # +G2 = [1 => [2], 2 => [1], 3=>Int[], 4=>Int[]] # +#G2 = [i => setdiff(union((G1[j].second for j in G1[i].second)...), G[i].second) for i in eachindex(G)] + +b = [ab(G1, i, x, θ, c, F) for i in eachindex(θ)] + +u0 = StructArray(t=t, x=x, θ=θ, θ_old=zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b) + +Random.seed!(1) +seed = (UInt(1),UInt(1)) +rng = Rng(seed) +T = 500.0 +t_old = copy(t) +adapt = false +factor = 1.7 +P = SPDMP(G, G1, G2, ∇ϕ, F, rng, adapt, factor) + +action! = (reset!, rand_reflect!, circle_hit1!, circle_hit2!) +next_action = FunctionWrangler((Zig.never_reset, next_rand_reflect, next_circle_hit1, next_circle_hit2)) + +# action! = (reset!, rand_reflect!, circle_hit1!) +# next_action = FunctionWrangler((Zig.never_reset, next_rand_reflect, next_circle_hit1)) + +h = Schedule(action!, next_action, u0, T, (P, )) +trc_ = Zig.simulate(h, progress=true) +trc = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_]) +error("") + +#subtrace1 = [t for t in trc_ if t[2] == 1] +#lines(getindex.(subtrace1, 1), getfield.(getindex.(subtrace1, 3), :x)) +ts, xs = Zig.sep(Zig.discretize(trc, 0.001)) + +using Makie +# scatter(ts, getindex.(xs, 2)) +fig = lines(getindex.(xs, 1), getindex.(xs, 2), ) +lines!(draw_circ([0.0, 0.0], 4.0), color = :red) +lines!(draw_circ( [-0.7, -0.7], 0.5), color = :red) +save("two_balls_one_inside_the_other.png", fig) \ No newline at end of file diff --git a/src/zigzagex.jl b/src/zigzagex.jl new file mode 100644 index 00000000..b4f1654d --- /dev/null +++ b/src/zigzagex.jl @@ -0,0 +1,159 @@ +using ZigZagBoomerang +const Zig = ZigZagBoomerang +using ZigZagBoomerang: sλ, sλ̄, reflect!, Rng, ab, smove_forward!, neighbours +using Random +using StructArrays +using StructArrays: components +using LinearAlgebra +using ZigZagBoomerang: SPriorityQueue, enqueue!, lastiterate + +T = 500.0 +d = 80 +seed = (UInt(1),UInt(1)) + + +𝕁(j) = mod(j,2) == 0 +function next_reset(j, _, t′, u, P, args...) + 0, (!𝕁(j)) ? Inf : t′ + 0.5*u.x[j] +end + + + +function reflect0!(i, t′, u, P::SPDMP, args...) + Zig.reflect!(0.0, 1, i, t′, u, P::SPDMP, args...) +end + +function reflect1!(i, t′, u, P::SPDMP, args...) + Zig.reflect!(1.0, -1, i, t′, u, P::SPDMP, args...) +end + +function next_rand_reflect(j, i, t′, u, P, args...) + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + if m[j] == 1 + return 0, Inf + end + if !𝕁(j) + b[j] = ab(G1, j, x, θ, c, F) + else + b[j] = ab(G1, j, x, θ, c, F) .+ (1/(x[j]), 2/(x[j]^2)) + end + t_old[j] = t′ + 0, t[j] + poisson_time(b[j], rand(P.rng)) +end + +function next_reflect0(j, i, t′, u, args...) + Zig.next_reflect(0.0, 1, j, i, t′, u, args...) +#= + t, x, θ, θ_old, m = components(u) + + if x[j] < 0 + return t[j] + end + θ[j]*x[j] >= 0 ? Inf : t[j] - x[j]/θ[j]=# +end + +function next_reflect1(j, i, t′, u, args...) + Zig.next_reflect(1.0, -1, j, i, t′, u, args...) + #= + t, x, θ, θ_old, m = components(u) + if x[j] > 1 + return t[j] + end + θ[j]*(x[j]-1) >= 0 ? Inf : t[j] - (x[j]-1)/θ[j] + =# +end + +function freeze!(args...) + Zig.freeze!(0.25, args...) +end +#if !@isdefined(discontinuity!) +function discontinuity!(args...) + Zig.discontinuity_at!(0.5, 0.0, 1, args...) +end +function next_discontinuity(args...) + Zig.next_hit(0.5, args...) +end + +function next_freezeunfreeze(args...) + Zig.next_freezeunfreeze(0.25, 1.0, args...) +end + +Random.seed!(1) + +using SparseArrays + +S = 1.3I + 0.5sprandn(d, d, 0.1) +Γ = S*S' + +∇ϕ(x, i, Γ) = -(𝕁(i))/x[i] + Zig.idot(Γ, i, x) # sparse computation + +# t, x, θ, θ_old, m, c, t_old, b +t0 = 0.0 +t = fill(t0, d) +x = rand(d) +θ = θ0 = rand([-1.0, 1.0], d) +F = ZigZag(0.9Γ, x*0) + + + +c = .6*[norm(Γ[:, i], 2) for i in 1:d] +G = G1 = [[i] => rowvals(F.Γ)[nzrange(F.Γ, i)] for i in eachindex(θ0)] +G2 = [i => setdiff(union((G1[j].second for j in G1[i].second)...), G[i].second) for i in eachindex(G1)] + +b = [ab(G1, i, x, θ, c, F) for i in eachindex(θ)] + +u0 = StructArray(t=t, x=x, θ=θ, θ_old = zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b) + + +rng = Rng(seed) +t_old = copy(t) +adapt = false +factor = 1.7 +P = SPDMP(G, G1, G2, ∇ϕ, F, rng, adapt, factor) + +#action! = FunctionWrangler((reset!, rand_reflect!, reflect0!, reflect1!)) +#next_action = FunctionWrangler((next_reset, next_rand_reflect, next_reflect0, next_reflect1)) + +action! = (Zig.reset!, Zig.rand_reflect!, discontinuity!, freeze!, reflect0!, reflect1!) +next_action = FunctionWrangler((next_reset, next_rand_reflect, next_discontinuity, next_freezeunfreeze, next_reflect0, next_reflect1)) + +#action! = (reset!, rand_reflect!, discontinuity!, reflect0!, reflect1!) +#next_action = FunctionWrangler((next_reset, next_rand_reflect, next_discontinuity, next_reflect0, next_reflect1)) + +#action! = ((reset!, rand_reflect!)) +#next_action = FunctionWrangler((next_reset, next_rand_reflect)) + + +#h = Schedule(FunctionWrangler(action!), next_action, u0, T, (P, Γ)) + +h = Schedule(action!, next_action, u0, T, (P, Γ)) + +l_ = lastiterate(h) + +using ProfileView, Profile +Profile.init(10000000, 0.00001) +ProfileView.@profview lastiterate(h) +l_ = @time lastiterate(h) + +total, l = simulate(h) +_, l = @time simulate(h) +trc_ = @time collect(h); +trc = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_]) + + +#trace, _, acc = @time spdmp(∇ϕ, t0, x, θ, T, c, G, F, Γ); +#@code_warntype handler(zeros(d), T, (f1!, f2!)); + +#using ProfileView +error() +#ProfileView.@profview handler(zeros(d), 10T); +using GLMakie +#subtrace1 = [t for t in trc_ if t[2] == 1] +#lines(getindex.(subtrace1, 1), getfield.(getindex.(subtrace1, 3), :x)) + +ts, xs = Zig.sep(Zig.subtrace(trc, [1,4])) + +lines(ts, getindex.(xs, 2)) +fig = lines(getindex.(xs, 1), getindex.(xs, 2)) \ No newline at end of file diff --git a/src/zigzagex2.jl b/src/zigzagex2.jl new file mode 100644 index 00000000..fc5441bd --- /dev/null +++ b/src/zigzagex2.jl @@ -0,0 +1,266 @@ +##################################################################### +# 2d Zig-Zag ouside ball: |x| > c # +# with glued boundaries for x in Γ (exit-non-entrance): x -> -x # +# fake 3rd coordinate for reflections and hitting times # +##################################################################### + +using ZigZagBoomerang +const Zig = ZigZagBoomerang +using ZigZagBoomerang: sλ, sλ̄, reflect!, Rng, ab, smove_forward!, neighbours +using ZigZagBoomerang: next_rand_reflect, reflect! +using Random +using StructArrays +using StructArrays: components +using LinearAlgebra +using ZigZagBoomerang: SPriorityQueue, enqueue!, lastiterate +# import standard reflection of zigzag +import ZigZagBoomerang: next_rand_reflect, rand_reflect!, reflect!, reset! + + +# overload simulate and handle! to allow having more clocks than coordinates +function Zig.simulate(handler; progress=true, progress_stops = 20) + T = handler.T + if progress + prg = Zig.Progress(progress_stops, 1) + else + prg = missing + end + stops = ismissing(prg) ? 0 : max(prg.n - 1, 0) # allow one stop for cleanup + tstop = T/stops + + u = deepcopy(handler.state) + action! = handler.action! + next_action = handler.next_action + d = length(handler.state) + action = ones(Int, d+1) # resets + + Q = Zig.SPriorityQueue(zeros(d+1)) + + evs = Zig.handle!(u, action!, next_action, action, Q, handler.args...) + # evs = [ev] + while true + ev = Zig.handle!(u, action!, next_action, action, Q, handler.args...) + if length(ev) == 1 + println("reflection") + else + println("hitting the circle") + end + t′ = ev[end][1] + println("at time current time: $(t′)") + println("") + t′ > T && break + append!(evs, ev) + if t′ > tstop + tstop += T/stops + Zig.next!(prg) + end + end + ismissing(prg) || Zig.ProgressMeter.finish!(prg) + return evs +end + +function Zig.handle!(u, action!, next_action, action, Q, args::Vararg{Any, N}) where {N} + # Who is (i) next_action, when (t′) and what (j) happens? + done = false + local e, t′, i + num = 0 + while !done + num += 1 + i, t′ = Zig.peek(Q) + e = action[i] + #done = action_nextaction(action!, next_action, Q, action, e, i, t′, u, args...) + done = Zig.switch(e, action!, next_action, (Q, action), i, t′, u, args...) + end + traceevent(t′, i, u, action, num) +end + +function traceevent(t′, i, u, action, num) + if 1 <= i <= length(u) + return [(t′, i, u[i], action[i], num)] + else + return [(t′, 1, u[1], action[1], num), (t′, 2, u[2], action[2], num)] + end +end + +function next_rand_reflect(j, i, t′, u, P::SPDMP, args...) + 1 <= j <= length(u) || return 0, Inf + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + if m[j] == 1 + return 0, Inf + end + b[j] = ab(G1, j, x, θ, c, F) + t_old[j] = t′ + 0, t[j] + poisson_time(b[j], rand(P.rng)) +end + +function rand_reflect!(i, t′, u, P::SPDMP, nt, args...) + μ, rsq = nt.μ, nt.rsq + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + @assert 1 <= i <= length(u) + smove_forward!(G, i, t, x, θ, m, t′, F) + ∇ϕi = P.∇ϕ(x, i, nt, args...) + l, lb = sλ(∇ϕi, i, x, θ, F), sλ̄(b[i], t[i] - t_old[i]) + if rand(P.rng)*lb < l + if l >= lb + !P.adapt && error("Tuning parameter `c` too small.") + adapt!(c, i, P.factor) + end + smove_forward!(G2, i, t, x, θ, m, t′, F) + ZigZagBoomerang.reflect!(i, ∇ϕi, x, θ, F) + return true, neighbours(G1, i) + else + return false, G1[i].first + end + +end + + + +Random.seed!(1) + +using SparseArrays +T = 500.0 +d = 2 +seed = (UInt(1),UInt(1)) + + +function reset!(i, t′, u, P::SPDMP, args...) # should move the coordinates if called later in the game + false, P.G1[i].first +end + +# coefficients of the quadratic equation coming for the condition \|x + θ*t - μ\|^2 = rsq +function abc_eq2d(x, θ, μ, rsq) + a = θ[1]^2 + θ[2]^2 + b = 2*((x[1] - μ[1]) *θ[1] + (x[2] - μ[2])*θ[2]) + c = (x[1]-μ[1])^2 + (x[2]-μ[2])^2 - rsq + a, b, c +end + +# joint reflection at the boundary +function circle_boundary_reflection!(x, θ, μ) + for i in eachindex(x) + if θ[i]*(x[i]-μ[i]) < 0.0 + θ[i]*=-1 + end + end + θ +end + +function next_circle_hit(j, i, t′, u, P::SPDMP, nt, args...) + μ, rsq = nt.μ, nt.rsq + if j <= length(u) # not applicable + return 0, Inf + else #hitting time to the ball with radius `radius` + t, x, θ, θ_old, m, c, t_old, b = components(u) + a1, a2, a3 = abc_eq2d(x, θ, μ, rsq) #solving quadradic equation + dis = a2^2 - 4a1*a3 #discriminant + # no solutions or inside + if dis <= 1e-7 || (x[1] - μ[1])^2 + (x[2] - μ[2])^2 - rsq < 1e-7 + return 0, Inf + else #pick the first positive event time + hitting_time = min((-a2 - sqrt(dis))/(2*a1),(-a2 + sqrt(dis))/(2*a1)) + if hitting_time <= 0.0 + return 0, Inf + end + + return 0, t′ + hitting_time #hitting time + end + end +end + +# either standard reflection, or bounce at the boundary or traversing the boundary +function circle_hit!(i, t′, u, P::SPDMP, nt, args...) + μ, rsq = nt.μ, nt.rsq + G, G1, G2 = P.G, P.G1, P.G2 + F = P.F + t, x, θ, θ_old, m, c, t_old, b = components(u) + if i == 3 #hitting boundary + smove_forward!(G, i, t, x, θ, m, t′, F) + + if abs((x[1] - μ[1])^2 + (x[2] - μ[2])^2 - rsq) > 1e-7 # make sure to hit be on the circle + dump(u) + error("not on the circle") + end + disc = ϕ(x) - ϕ(-x + 2*μ) # magnitude of the discontinuity + if disc < 0.0 || rand() > 1 - exp(-disc) # teleport + # if false # never teleport + # jump on the other side drawing a line passing through the center of the ball + x .= -x + 2 .*μ # improve by looking at G[3] + else # bounce off + θ .= circle_boundary_reflection!(x, θ, μ) + end + return true, neighbours(G1, i) + else + error("action not available for clock $i") + end +end + + + +Sigma = Matrix([1.0 0.9; + 0.9 1.0]) + +Γ = sparse(Sigma^(-1)) + +ϕ(x) = -0.5*x'*Γ*x # negated log-density +∇ϕ(x, i, nt) = Zig.idot(nt.Γ, i, x) # sparse computation + +# t, x, θ, c, t_old, b +d = 2 +t0 = 0.0 +t = fill(t0, d) +x = [+1.0, -2.0] + rand(d) +θ = θ0 = rand([-1.0, 1.0], d) +F = ZigZag(Γ, x*0) + +# G, G2: Clocks to coordinates +# G1: Clocks to clocks + +c = (zero(x) .+ 0.1) +G = [[i] => rowvals(F.Γ)[nzrange(F.Γ, i)] for i in eachindex(θ0)] +push!(G, [3] => [1,2]) # move all coordinates +G1 = [[i] => [rowvals(F.Γ)[nzrange(F.Γ, i)];3] for i in eachindex(θ0)] +push!(G1, [3] => [1, 2, 3]) # invalidate ALL the clocks yeah +G2 = [1 => [2], 2 => [1], 3=>Int[]] # +#G2 = [i => setdiff(union((G1[j].second for j in G1[i].second)...), G[i].second) for i in eachindex(G)] + +b = [ab(G1, i, x, θ, c, F) for i in eachindex(θ)] # specific for subsampling + +u0 = StructArray(t=t, x=x, θ=θ, θ_old=zeros(d), m=zeros(Int,d), c=c, t_old=copy(t), b=b) + + +rng = Rng(seed) +t_old = copy(t) +adapt = false +factor = 1.7 +P = SPDMP(G, G1, G2, ∇ϕ, F, rng, adapt, factor) # ? + +action! = (reset!, rand_reflect!, circle_hit!) +next_action = FunctionWrangler((Zig.never_reset, next_rand_reflect, next_circle_hit)) +rsq = 0.5 +μ = [2.0, -2.0] +h = Schedule(action!, next_action, u0, T, (P, (Γ=Γ, μ=μ, rsq=rsq))) +trc_ = Zig.simulate(h, progress=true) +trc = Zig.FactTrace(F, t0, x, θ, [(ev[1], ev[2], ev[3].x, ev[3].θ) for ev in trc_]) + +ts, xs = Zig.sep(Zig.discretize(trc, 0.001)) +error("") + +using Makie +# scatter(ts, getindex.(xs, 2)) +fig = lines(getindex.(xs, 1), getindex.(xs, 2), ) +axes = Axis(fig) +#draw the circle with radius r, centered in μ +r = sqrt(rsq) +x1 = Float64.(-r:0.0001:r) +x2 = zero(x1) +for i in eachindex(x1) + x2[i] = sqrt(r^2 - x1[i]^2) +end +lines!(x1 .+ μ[1], x2 .+ μ[2], color = "red") +lines!(x1 .+ μ[1], -x2 .+ μ[2], color = "red") +# save("bounce_off_the_ball_with_teleportation.png", fig) \ No newline at end of file diff --git a/test/maintest.jl b/test/maintest.jl index 5c123479..2c6cf1f2 100644 --- a/test/maintest.jl +++ b/test/maintest.jl @@ -192,6 +192,7 @@ end end @testset "FactBoomerang1" begin + Random.seed!(1) ϕ(x) = [cos(π*x[1]) + x[1]^2/2] # not needed # gradient of ϕ(x) ∇ϕ(x) = [-π*sin(π*x[1]) + x[1]]