Zig zag and Boomerang reference implementation

zagboom.jl

# Zig zag and Boomerang reference implementation 

"""
    poisson_time(a,b,u)
obtaining waiting time for inhomogeneous Poisson Process
with rate of the form λ(t) = (a + b*t)^+, `a`,`b` ∈ R, `u` uniform random variable
"""
function poisson_time(a,b,u)
    if b > 0
        if a < 0
            τ = sqrt(-log(u)*2.0/b) - a/b
        else #a[i]>0
            τ = sqrt((a/b)^2 - log(u)*2.0/b) - a/b
        end
    elseif  b == 0
        if a > 0
            τ = -log(u)/a
        else #a[i] <= 0
            τ = Inf
        end
    else #b[i] < 0
        if a <= 0
            τ = Inf
        elseif -log(u) <= -a^2/b + a^2/(2*b)
            τ = - sqrt((a/b)^2 - log(u)*2.0/b) - a/b
        else
            τ = Inf
        end
    end
end

"""
    h_poisson_time(a, u)
obtaining waiting time for homogeneous Poisson Process
with rate of the form λ(t) = a, `a` ≥ 0, `u` uniform random variable
"""
function poisson_time(a,u)
    -log(u)/a
end
"""
    Continuous_Dynamics
Abstract type for the deterministic dynamics of PDMPs
"""
abstract type Continuous_Dynamics end
"""
    Linear <: Continuous_Dynamics
Dynamics preserving the Lebesgue measure
"""
struct Linear <: Continuous_Dynamics  end
"""
    Circular <: Continuous_Dynamics
Dynamics preserving the standard Gaussian measure (Boomerang)
"""
struct Circular <: Continuous_Dynamics  end

# Linear dynamics (time, space, velocity)
function move_forward(τ, t, x, θ, ::Linear)
    τ + t, x + θ*τ , θ
end

# Circular dynamics (time, space, velocity)
# dx = -x dt; dv = -v dt 
function move_forward(τ, t, x, θ, ::Circular)
    x_new = x*cos(τ) + θ*sin(τ)
    θ = -x*sin(τ) + θ*cos(τ)
    t + τ, x_new, θ
end


# negative log-density with respect to Lebesgue
ϕ(x, ::Linear) = (cos(2pi*x) + x^2/2)
# negative log-density with respect to standard Gaissoian
ϕ(x, ::Circular) =  cos(2pi*x)
# gradient of ϕ(x, ::Linear) (REPLACE IT WITH AUTOMATIC DIFFERENTIATION)
∇ϕ(x, ::Linear) = x - 2*pi*sin(2*π*x)
# gradient of ϕ(x, ::Circular)
∇ϕ(x, ::Circular) = - 2*pi*sin(2*π*x)

λ(x, θ, F::Continuous_Dynamics) = max(0, θ*∇ϕ(x, F))

# affine bounds for Zig-Zag
λ_bar(x, θ, ::Linear) = max(0, 2*pi + θ*x)
# constant bound for Boomerang with global bounded |∇ϕ(x)|
# suppose |∇ϕ(x, :Circular)| ≤ C. Then λ(x(t),θ(t)) ≤ C*sqrt(x(0)^2 + θ(0)^2)
λ_bar(x, θ, ::Circular) = sqrt(θ^2 + x^2)*2*pi #Global bound


# waiting times
ab(x, θ, ::Linear) = (2*pi + θ*x, 1.0)
ab(x, θ, ::Circular) = (2*pi*sqrt(x^2 + θ^2), 0.0)
waiting_time(x,θ, Flow::Continuous_Dynamics) = poisson_time(ab(x, θ, Flow)...,rand())
waiting_time_ref(λref, ::Linear) = Inf
waiting_time_ref(λref, ::Circular) = poisson_time(λref, 0.0, rand())

# Not efficient to save all this stuff. Ideally with initial condition and time events
# you can reconstruct the whole path
struct Skeleton
    t::Float64 # time
    x::Float64 #position
    θ::Float64 #velocity
end

# Algorithm for one dimensional pdmp (ZigZag or Boomerang)
function pdmp(x, θ, T, λref, Flow::Continuous_Dynamics)
    t = 0
    Ξ = Vector{Skeleton}()
    push!(Ξ, Skeleton(0.0, x, θ))
    τref = waiting_time_ref(λref, Flow)
    τ =  waiting_time(x, θ, Flow)
    while t<T
        if τref < τ
            t, x, θ = move_forward(τref, t, x, θ, Flow)
            θ = randn()
            τref = waiting_time_ref(λref, Flow)
            τ =  waiting_time(x, θ, Flow)
            push!(Ξ, Skeleton(t, x, θ))
        else
            t, x, θ = move_forward(τ, t, x, θ, Flow)
            τref -= τ
            τ = waiting_time(x, θ, Flow)
            if  λ(x, θ, Flow) >= λ_bar(x, θ, Flow)
                error("out of bounds")
            elseif rand()*λ_bar(x, θ, Flow) < λ(x, θ, Flow)
                θ = -θ #In multidimension the change of velocity is different:
                        #reflection symmetric on the normal vector of the contour
                push!(Ξ, Skeleton(t, x, θ))
            end
        end
    end
    return Ξ, Flow
end
waiting_time_ref(0.01, Circular())


# For plotting: discretization of the circular dynamics
function discretization(x::Vector{Skeleton}, Flow::Circular, dt)
    k = 1
    ξ = x[k].x
    θ = x[k].θ
    τ = x[k+1].t
    clock = 0.0
    Ω = trajectory([clock => ξ ])
    #error("STOP")
    while k < length(x)-1
        while clock + dt <= τ
            clock, ξ, θ = move_forward(dt, clock, ξ, θ, Flow)
            push!(Ω, clock => ξ)
        end
        clock, ξ, θ = move_forward(τ - clock, clock, ξ, θ, Flow)
        push!(Ω, clock => ξ)
        k += 1
        ξ = x[k].x
        θ = x[k].θ
        τ = x[k+1].t
    end
    Ω
end

using Makie
using Trajectories
# Example: ZigZag
out = pdmp(randn(), 1.0, 100.0, 0.0, Linear())
t(x) = x.t
x(x) = x.x
p1 = Makie.lines(t.(out[1]), x.(out[1]))
save("zigzag.png", p1)


# Example: Boomerang
out = pdmp(randn(), 1.0, 100.0, 0.5, Circular())
dt = 0.01
sum(t.(out[1][1:end-1]) .< t.(out[1][2:end]))
t.(out[1])

xx = discretization(out...,dt)
xx.t
p2 = Makie.lines(xx.t[3000:3500], xx.x[3000:3500])
save("boomerang.png", p2)

zigzag1d.jl

# Bare bones implementation of a 1d-ZigZag process

using ForwardDiff
using Trajectories


"""
    poisson_time(a,b,u)
obtaining waiting time for Inhomogeneous Poisson Process
with rate of the form λ(t) = (a + b*t)^+, `a`,`b` ∈ R, `u` random variable
"""
function poisson_time(a,b,u)
    if b > 0
        if a < 0
            τ = sqrt(-log(u)*2.0/b) - a/b
        else #a[i]>0
            τ = sqrt((a/b)^2 - log(u)*2.0/b) - a/b
        end
    elseif  b == 0
        if a > 0
            τ = -log(u)/a
        else #a[i] <= 0
            τ = Inf
        end
    else #b[i] < 0
        if a <= 0
            τ = Inf
        elseif -log(u) <= -a^2/b + a^2/(2*b)
            τ = - sqrt((a/b)^2 - log(u)*2.0/b) - a/b
        else
            τ = Inf
        end
    end
end

ϕ(x) = -(cos(2pi*x) - x^2/2)
∇ϕ(x) = x + 2*pi*sin(2*π*x)
bound(x, θ) = max(θ*(x - 2*pi), θ*(x + 2*pi))
ab(x, θ) = bound(x, θ), 1.0

λ(x, θ) = max(θ*∇ϕ(x), 0.0)

λmax(x, θ) =  max(bound(x, θ), 0.0)
randλmax(x, θ) = poisson_time(ab(x, θ)..., rand())

function next(t, ξ, T, θ)
    λm = 15.0
    while t < T
        Δt = randλmax(ξ, θ)
        t = t + Δt
        ξ = ξ + θ*Δt
        λm = λmax(ξ, θ)
        println(t, " ", ξ, " ", θ, " ", λ(ξ, θ), " ", λm)
        if λm < λ(ξ, θ)
            error("out of bounds: ", round(ξ, digits=2), " ", θ )
        end
        if rand()*λm < λ(ξ, θ)
            return t, ξ, -θ
        end
    end
    return t, ξ, θ
end

t, ξ, θ = 0.0, 1.0, 1

t, ξ, θ = next(t, ξ, T, θ)

function zigzag(t, ξ, T, θ=1)
    x = trajectory([t => ξ])
    while t < T
        t, ξ, θ = next(t, ξ, T, θ)
        push!(x, t=>ξ)
        println(t=>ξ)
    end
    x
end
T = 100.0
x = zigzag(t, ξ, T)

using Makie
p = lines(x.t, x.x)
save("zigzag.png", p)

How can I edit your script? It seems to me that I can only comment on the bottom.