From 5e8b239daee3acb9d71efe4372c0ee506d058677 Mon Sep 17 00:00:00 2001
From: mschauer <moritzschauer@web.de>
Date: Mon, 20 Jun 2022 10:30:13 +0200
Subject: [PATCH 1/8] Use findmin

---
 src/not_fact_samplers.jl | 96 +++++++++++++++++++++-------------------
 src/notfactiter.jl       | 17 ++++---
 2 files changed, 59 insertions(+), 54 deletions(-)

diff --git a/src/not_fact_samplers.jl b/src/not_fact_samplers.jl
index 02124a6..d3e3cdb 100644
--- a/src/not_fact_samplers.jl
+++ b/src/not_fact_samplers.jl
@@ -29,15 +29,15 @@ end
 function ab(x, θ, C::LocalBound, ∇ϕx::AbstractVector, v, B::BouncyParticle)
     (C.c + dot(θ, ∇ϕx), v, 2sqrt(length(θ))/C.c/norm(θ, 2))
 end
-function ab(x, θ, C::LocalBound, vdϕ::Number, v, B::BouncyParticle)
+function ab(t, x, θ, C::LocalBound, vdϕ::Number, v, B::BouncyParticle)
     @assert vdϕ isa Number
-    (C.c + vdϕ, v, 2sqrt(length(θ))/C.c/norm(θ, 2))
+    (C.c + vdϕ, v, t + 2sqrt(length(θ))/C.c/norm(θ, 2))
 end
 
 function ab(x, θ, C::GlobalBound, ∇ϕx, v, B::Boomerang)
     (sqrt(normsq(θ) + normsq((x - B.μ)))*C.c, 0.0, Inf)
 end
-ab(x, θ, c, Flow) =  ab(x, θ, GlobalBound(c), nothing, nothing, Flow)
+ab(x, θ, c, flow) =  ab(x, θ, GlobalBound(c), nothing, nothing, flow)
 
 function event(t, x, θ, Z::Union{BouncyParticle,Boomerang})
     t, copy(x), copy(θ), nothing
@@ -150,34 +150,41 @@ function pdmp(∇ϕ!, t0, x0, θ0, T, c::Bound, Flow::Union{BouncyParticle, Boom
 end
 
 
-##################################
+##################################    ##################################
+function next_event1(rng, u::Tuple, abc, flow)
+    t, x, v = u
+    a, b, Δ = abc
+    τ = t + poisson_time(a, b, rand(rng))
+    τrefresh = t + waiting_time_ref(rng, flow)
+    when, what = findmin((τ, Δ, τrefresh))
+    return when, (:bounce, :expire, :refresh)[what]
+end
 
-function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, c::Bound, abc, (t′, renew), τref, (acc, num),
-    Flow::BouncyParticle, args...; subsample=false, oscn=false, factor=1.5, adapt=false) where {F1, F2}
+function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, c::Bound, abc, (t′, action), τref, (acc, num),
+    flow::BouncyParticle, args...; subsample=false, oscn=false, factor=1.5, adapt=false) where {F1, F2}
     while true
-        if τref < t′
-            t, _ = move_forward!(τref - t, t, x, θ, Flow)
-            refresh!(rng, θ, Flow)
+        if action == :refresh
+            t, _ = move_forward!(t′ - t, t, x, θ, flow)
+            refresh!(rng, θ, flow)
             θdϕ, v = dϕ(t, x, θ, args...) 
-            #∇ϕx = grad_correct!(∇ϕx, x, Flow)
-            l = λ(θdϕ, Flow) 
-            τref = t + waiting_time_ref(rng, Flow)
-            abc = ab(x, θ, c, θdϕ, v, Flow)
-            t′, renew = next_time(t, abc, rand(rng))
-            return t, (acc, num), c, abc, (t′, renew), τref
-        elseif renew
+            #∇ϕx = grad_correct!(∇ϕx, x, flow)
+            l = λ(θdϕ, flow) 
+            abc = ab(t, x, θ, c, θdϕ, v, flow)
+            t′, action = next_event1(rng, (t, x, θ), abc, flow)
+            return t, (acc, num), c, abc, (t′, action), τref
+        elseif action == :expire
             τ = t′ - t
-            t, _ = move_forward!(τ, t, x, θ, Flow) 
+            t, _ = move_forward!(τ, t, x, θ, flow) 
             θdϕ, v = dϕ(t, x, θ, args...) 
-            #∇ϕx = grad_correct!(∇ϕx, x, Flow)
-            abc = ab(x, θ, c, θdϕ, v, Flow)
-            t′, renew = next_time(t, abc, rand(rng))
-        else
+            #∇ϕx = grad_correct!(∇ϕx, x, flow)
+            abc = ab(t, x, θ, c, θdϕ, v, flow)
+            t′, action = next_event1(rng, (t, x, θ), abc, flow)
+        else # action == :reflect
             τ = t′ - t
-            t, _ = move_forward!(τ, t, x, θ, Flow)
+            t, _ = move_forward!(τ, t, x, θ, flow)
             θdϕ, v = dϕ(t, x, θ, args...) 
-            #∇ϕx = grad_correct!(∇ϕx, x, Flow)
-            l, lb = λ(θdϕ, Flow), pos(abc[1] + abc[2]*τ)
+            #∇ϕx = grad_correct!(∇ϕx, x, flow)
+            l, lb = λ(θdϕ, flow), pos(abc[1] + abc[2]*τ)
             num += 1
             if rand(rng)*lb <= l
                 acc += 1
@@ -191,25 +198,25 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, c::Bound, abc,
                     #error("subsampling needs to be seeded by time")
                 end
                 if oscn
-                    @assert Flow.L == I
-                    oscn!(rng, θ, ∇ϕx, Flow.ρ; normalize=false)
+                    @assert flow.L == I
+                    oscn!(rng, θ, ∇ϕx, flow.ρ; normalize=false)
                 else
-                    reflect!(∇ϕx, x, θ, Flow)
+                    reflect!(∇ϕx, x, θ, flow)
                 end
                 θdϕ, v = dϕ(t, x, θ, args...) 
-                #∇ϕx = grad_correct!(∇ϕx, x, Flow)
-                abc = ab(x, θ, c, θdϕ, v, Flow)
-                t′, renew = next_time(t, abc, rand(rng))
-                !subsample && return t, (acc, num), c, abc, (t′, renew), τref
+                #∇ϕx = grad_correct!(∇ϕx, x, flow)
+                abc = ab(t, x, θ, c, θdϕ, v, flow)
+                t′, action = next_event1(rng, (t, x, θ), abc, flow)
+                !subsample && return t, (acc, num), c, abc, (t′, action), τref
             else
-                abc = ab(x, θ, c, θdϕ, v, Flow)
-                t′, renew = next_time(t, abc, rand(rng))
+                abc = ab(t, x, θ, c, θdϕ, v, flow)
+                t′, action = next_event1(rng, (t, x, θ), abc, flow)
             end
         end
     end
 end
 """
-    pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, Flow::BouncyParticle, args...; oscn=false, adapt=false, subsample=false, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
+    pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; oscn=false, adapt=false, subsample=false, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
 
 The first directional derivative `dϕ[1]` tells me if I move up or down the potential. The second directional derivative `dϕ[2]` tells me how fast the first changes. The gradient `∇ϕ!` tells me the surface I want to reflect on.
 
@@ -256,19 +263,18 @@ The remaining arguments:
     t, x = ZigZagBoomerang.sep(trace)
 
 """
-function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, Flow::BouncyParticle, args...; oscn=false, adapt=false, subsample=false, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
+function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; oscn=false, adapt=false, subsample=false, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
     t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
     rng = Rng(seed)
-    Ξ = Trace(t0, x0, θ0, Flow)
-    τref = waiting_time_ref(rng, Flow)
+    Ξ = Trace(t0, x0, θ0, flow)
     θdϕ, v = dϕ(t, x, θ, args...) 
     #@assert v2 ≈ v
     #@assert θdϕ ≈ dot(∇ϕx, θ)
     
-    #∇ϕx = grad_correct!(∇ϕx, x, Flow)
+    #∇ϕx = grad_correct!(∇ϕx, x, flow)
     num = acc = 0
     #l = 0.0
-    abc = ab(x, θ, c, θdϕ, v, Flow)
+    abc = ab(t, x, θ, c, θdϕ, v, flow)
     if progress
         prg = Progress(progress_stops, 1)
     else
@@ -276,11 +282,11 @@ function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, Flow::BouncyParticle, args.
     end
     stops = ismissing(prg) ? 0 : max(prg.n - 1, 0) # allow one stop for cleanup
     tstop = T/stops
-
-    t′, renew = next_time(t, abc, rand(rng))
+    tref = nothing
+    t′, action = next_event1(rng, (t, x, θ), abc, flow)
     while t < T
-        t, (acc, num), c, abc, (t′, renew), τref = pdmp_inner!(rng, dϕ, ∇ϕ!, ∇ϕx, t, x, θ, c, abc, (t′, renew), τref, (acc, num), Flow, args...; oscn=oscn, subsample=subsample, factor=factor, adapt=adapt)
-        push!(Ξ, event(t, x, θ, Flow))
+        t, (acc, num), c, abc, (t′, action), τref = pdmp_inner!(rng, dϕ, ∇ϕ!, ∇ϕx, t, x, θ, c, abc, (t′, action), τref, (acc, num), flow, args...; oscn=oscn, subsample=subsample, factor=factor, adapt=adapt)
+        push!(Ξ, event(t, x, θ, flow))
 
         if t > tstop
             tstop += T/stops
@@ -299,5 +305,5 @@ struct Wrapper{F}
 end
 (F::Wrapper)(y, t, x, θ, args...) = F.f(y, x, args...), nothing
 
-pdmp(∇ϕ!, t0, x0, θ0, T, c, Flow::Union{BouncyParticle, Boomerang}, args...; nargs...) = 
-pdmp(Wrapper(∇ϕ!), t0, x0, θ0, T, GlobalBound(c), Flow, args...; nargs...)
+pdmp(∇ϕ!, t0, x0, θ0, T, c, flow::Union{BouncyParticle, Boomerang}, args...; nargs...) = 
+pdmp(Wrapper(∇ϕ!), t0, x0, θ0, T, GlobalBound(c), flow, args...; nargs...)
diff --git a/src/notfactiter.jl b/src/notfactiter.jl
index fe45707..619684f 100644
--- a/src/notfactiter.jl
+++ b/src/notfactiter.jl
@@ -58,31 +58,30 @@ end
 
 function iterate(FS::NotFactSampler{<:Any, <:Tuple})
     t0, (x0, θ0) = FS.u0
-    Flow = FS.F
+    flow = FS.F
     n = length(x0)
     t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
     c = FS.c
     rng = FS.rng
-    τref = waiting_time_ref(rng, Flow)
+    τref = NaN
 
 
     dϕ, ∇ϕ! = FS.∇ϕ![1], FS.∇ϕ![2]  
     θdϕ, v = dϕ(t, x, θ, FS.args...) 
     num = acc = 0
-    abc = ab(x, θ, c, θdϕ, v, Flow)
-
-    t′, renew = next_time(t, abc, rand(rng))
-    iterate(FS, ((t => (x, θ)), ∇ϕx, (acc, num), c, abc, (t′, renew), τref))
+    abc = ab(t, x, θ, c, θdϕ, v, flow)
+    t′, action = next_event1(rng, (t, x, θ), abc, flow)
+    iterate(FS, ((t => (x, θ)), ∇ϕx, (acc, num), c, abc, (t′, action), τref))
 end
 using Test
 
 
-function iterate(FS::NotFactSampler{<:Any, <:Tuple},  (u, ∇ϕx, (acc, num), c, abc, (t′, renew), τref))
+function iterate(FS::NotFactSampler{<:Any, <:Tuple},  (u, ∇ϕx, (acc, num), c, abc, (t′, action), τref))
     t, (x, θ) = u
     dϕ, ∇ϕ! = FS.∇ϕ![1], FS.∇ϕ![2]  
-    t, (acc, num), c, abc, (t′, renew), τref = pdmp_inner!(FS.rng, dϕ, ∇ϕ!, ∇ϕx, t, x, θ, c, abc, (t′, renew), τref, (acc, num), FS.F, FS.args...; FS.kargs...)
+    t, (acc, num), c, abc, (t′, action), τref = pdmp_inner!(FS.rng, dϕ, ∇ϕ!, ∇ϕx, t, x, θ, c, abc, (t′, action), τref, (acc, num), FS.F, FS.args...; FS.kargs...)
    
     ev = rawevent(t, x, θ, FS.F)
     u = t => (x, θ)
-    return ev, (u, ∇ϕx, (acc, num), c, abc, (t′, renew), τref)
+    return ev, (u, ∇ϕx, (acc, num), c, abc, (t′, action), τref)
 end

From c79d701a150d4827bd1d4e1e897cafaaf619bf2b Mon Sep 17 00:00:00 2001
From: mschauer <moritzschauer@web.de>
Date: Mon, 20 Jun 2022 16:04:08 +0200
Subject: [PATCH 2/8] Move function

---
 src/not_fact_samplers.jl | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/not_fact_samplers.jl b/src/not_fact_samplers.jl
index d3e3cdb..6d5eb33 100644
--- a/src/not_fact_samplers.jl
+++ b/src/not_fact_samplers.jl
@@ -29,10 +29,7 @@ end
 function ab(x, θ, C::LocalBound, ∇ϕx::AbstractVector, v, B::BouncyParticle)
     (C.c + dot(θ, ∇ϕx), v, 2sqrt(length(θ))/C.c/norm(θ, 2))
 end
-function ab(t, x, θ, C::LocalBound, vdϕ::Number, v, B::BouncyParticle)
-    @assert vdϕ isa Number
-    (C.c + vdϕ, v, t + 2sqrt(length(θ))/C.c/norm(θ, 2))
-end
+
 
 function ab(x, θ, C::GlobalBound, ∇ϕx, v, B::Boomerang)
     (sqrt(normsq(θ) + normsq((x - B.μ)))*C.c, 0.0, Inf)
@@ -151,6 +148,11 @@ end
 
 
 ##################################    ##################################
+function ab(t, x, θ, C::LocalBound, vdϕ::Number, v, B::BouncyParticle)
+    @assert vdϕ isa Number
+    (C.c + vdϕ, v, t + 2sqrt(length(θ))/C.c/norm(θ, 2))
+end
+
 function next_event1(rng, u::Tuple, abc, flow)
     t, x, v = u
     a, b, Δ = abc

From f61abfd5f288c08da30722e2ad1d1675a0c88902 Mon Sep 17 00:00:00 2001
From: mschauer <moritzschauer@web.de>
Date: Mon, 20 Jun 2022 19:21:52 +0200
Subject: [PATCH 3/8] Use non-equidistant sample scheme

---
 src/not_fact_samplers.jl | 34 ++++++++++++++++++++++------------
 src/notfactiter.jl       | 18 +++++++++---------
 2 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/src/not_fact_samplers.jl b/src/not_fact_samplers.jl
index 6d5eb33..bda7180 100644
--- a/src/not_fact_samplers.jl
+++ b/src/not_fact_samplers.jl
@@ -154,33 +154,42 @@ function ab(t, x, θ, C::LocalBound, vdϕ::Number, v, B::BouncyParticle)
 end
 
 function next_event1(rng, u::Tuple, abc, flow)
-    t, x, v = u
+    t, x, v, V = u
     a, b, Δ = abc
     τ = t + poisson_time(a, b, rand(rng))
-    τrefresh = t + waiting_time_ref(rng, flow)
+    τrefresh = t + waiting_time_ref(rng, flow)/V
     when, what = findmin((τ, Δ, τrefresh))
     return when, (:bounce, :expire, :refresh)[what]
 end
 
-function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, c::Bound, abc, (t′, action), τref, (acc, num),
+function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, abc, (t′, action), Δrec, (acc, num),
     flow::BouncyParticle, args...; subsample=false, oscn=false, factor=1.5, adapt=false) where {F1, F2}
     while true
+        if t + Δrec/V <= t′ # record! (large speed, more records)
+            t, _ = move_forward!(V\Δrec, t, x, θ, flow) 
+            Δrec = 1/flow.λref
+            return t, V, (acc, num), c, abc, (t′, action), Δrec
+        end
+        Δrec = Δrec - (t′ - t)*V # coming closer to rec
+        @assert Δrec > 0.0
         if action == :refresh
+            @assert Δrec >= 0
             t, _ = move_forward!(t′ - t, t, x, θ, flow)
             refresh!(rng, θ, flow)
+            V = norm(θ, 2)
             θdϕ, v = dϕ(t, x, θ, args...) 
             #∇ϕx = grad_correct!(∇ϕx, x, flow)
             l = λ(θdϕ, flow) 
             abc = ab(t, x, θ, c, θdϕ, v, flow)
-            t′, action = next_event1(rng, (t, x, θ), abc, flow)
-            return t, (acc, num), c, abc, (t′, action), τref
+            t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
+            #return t, V, (acc, num), c, abc, (t′, action), Δrec
         elseif action == :expire
             τ = t′ - t
             t, _ = move_forward!(τ, t, x, θ, flow) 
             θdϕ, v = dϕ(t, x, θ, args...) 
             #∇ϕx = grad_correct!(∇ϕx, x, flow)
             abc = ab(t, x, θ, c, θdϕ, v, flow)
-            t′, action = next_event1(rng, (t, x, θ), abc, flow)
+            t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
         else # action == :reflect
             τ = t′ - t
             t, _ = move_forward!(τ, t, x, θ, flow)
@@ -208,11 +217,11 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, c::Bound, abc,
                 θdϕ, v = dϕ(t, x, θ, args...) 
                 #∇ϕx = grad_correct!(∇ϕx, x, flow)
                 abc = ab(t, x, θ, c, θdϕ, v, flow)
-                t′, action = next_event1(rng, (t, x, θ), abc, flow)
-                !subsample && return t, (acc, num), c, abc, (t′, action), τref
+                t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
+                #!subsample && return t, V, (acc, num), c, abc, (t′, action), Δrec
             else
                 abc = ab(t, x, θ, c, θdϕ, v, flow)
-                t′, action = next_event1(rng, (t, x, θ), abc, flow)
+                t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
             end
         end
     end
@@ -267,6 +276,7 @@ The remaining arguments:
 """
 function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; oscn=false, adapt=false, subsample=false, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
     t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
+    V = norm(θ, 2)
     rng = Rng(seed)
     Ξ = Trace(t0, x0, θ0, flow)
     θdϕ, v = dϕ(t, x, θ, args...) 
@@ -284,10 +294,10 @@ function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args.
     end
     stops = ismissing(prg) ? 0 : max(prg.n - 1, 0) # allow one stop for cleanup
     tstop = T/stops
-    tref = nothing
-    t′, action = next_event1(rng, (t, x, θ), abc, flow)
+    Δrec = 1/flow.λref
+    t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
     while t < T
-        t, (acc, num), c, abc, (t′, action), τref = pdmp_inner!(rng, dϕ, ∇ϕ!, ∇ϕx, t, x, θ, c, abc, (t′, action), τref, (acc, num), flow, args...; oscn=oscn, subsample=subsample, factor=factor, adapt=adapt)
+        t, V, (acc, num), c, abc, (t′, action), Δrec = pdmp_inner!(rng, dϕ, ∇ϕ!, ∇ϕx, t, x, θ, V, c, abc, (t′, action), Δrec, (acc, num), flow, args...; oscn=oscn, subsample=subsample, factor=factor, adapt=adapt)
         push!(Ξ, event(t, x, θ, flow))
 
         if t > tstop
diff --git a/src/notfactiter.jl b/src/notfactiter.jl
index 619684f..d26331e 100644
--- a/src/notfactiter.jl
+++ b/src/notfactiter.jl
@@ -61,27 +61,27 @@ function iterate(FS::NotFactSampler{<:Any, <:Tuple})
     flow = FS.F
     n = length(x0)
     t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
+    V = norm(θ, 2)
     c = FS.c
     rng = FS.rng
-    τref = NaN
-
+    Δrec = 1/flow.λref
 
     dϕ, ∇ϕ! = FS.∇ϕ![1], FS.∇ϕ![2]  
     θdϕ, v = dϕ(t, x, θ, FS.args...) 
     num = acc = 0
     abc = ab(t, x, θ, c, θdϕ, v, flow)
-    t′, action = next_event1(rng, (t, x, θ), abc, flow)
-    iterate(FS, ((t => (x, θ)), ∇ϕx, (acc, num), c, abc, (t′, action), τref))
+    t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
+    iterate(FS, ((t => (x, θ, V)), ∇ϕx, (acc, num), c, abc, (t′, action), Δrec))
 end
 using Test
 
 
-function iterate(FS::NotFactSampler{<:Any, <:Tuple},  (u, ∇ϕx, (acc, num), c, abc, (t′, action), τref))
-    t, (x, θ) = u
+function iterate(FS::NotFactSampler{<:Any, <:Tuple},  (u, ∇ϕx, (acc, num), c, abc, (t′, action), Δrec))
+    t, (x, θ, V) = u
     dϕ, ∇ϕ! = FS.∇ϕ![1], FS.∇ϕ![2]  
-    t, (acc, num), c, abc, (t′, action), τref = pdmp_inner!(FS.rng, dϕ, ∇ϕ!, ∇ϕx, t, x, θ, c, abc, (t′, action), τref, (acc, num), FS.F, FS.args...; FS.kargs...)
+    t, V, (acc, num), c, abc, (t′, action), Δrec = pdmp_inner!(FS.rng, dϕ, ∇ϕ!, ∇ϕx, t, x, θ, V, c, abc, (t′, action), Δrec, (acc, num), FS.F, FS.args...; FS.kargs...)
    
     ev = rawevent(t, x, θ, FS.F)
-    u = t => (x, θ)
-    return ev, (u, ∇ϕx, (acc, num), c, abc, (t′, action), τref)
+    u = t => (x, θ, V)
+    return ev, (u, ∇ϕx, (acc, num), c, abc, (t′, action), Δrec)
 end

From a1b0400cbd7a6860ca393ae67aa34a47e41d4373 Mon Sep 17 00:00:00 2001
From: mschauer <moritzschauer@web.de>
Date: Wed, 22 Jun 2022 10:39:03 +0200
Subject: [PATCH 4/8] Adapt mass matrix to Fisher information

---
 src/not_fact_samplers.jl | 52 +++++++++++++++++++++++++++------
 src/notfactiter.jl       |  5 ++++
 test/maintest.jl         | 63 +++++++++++++++++++++++++++++++---------
 3 files changed, 97 insertions(+), 23 deletions(-)

diff --git a/src/not_fact_samplers.jl b/src/not_fact_samplers.jl
index bda7180..7627ead 100644
--- a/src/not_fact_samplers.jl
+++ b/src/not_fact_samplers.jl
@@ -147,7 +147,15 @@ function pdmp(∇ϕ!, t0, x0, θ0, T, c::Bound, Flow::Union{BouncyParticle, Boom
 end
 
 
-##################################    ##################################
+################################## ModernBPS ##################################
+function mass_adapt_init(M)
+    M.diag
+end
+
+function mass_adapt!(M, m)
+     @. M.diag = 1/m
+end
+
 function ab(t, x, θ, C::LocalBound, vdϕ::Number, v, B::BouncyParticle)
     @assert vdϕ isa Number
     (C.c + vdϕ, v, t + 2sqrt(length(θ))/C.c/norm(θ, 2))
@@ -164,10 +172,19 @@ end
 
 function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, abc, (t′, action), Δrec, (acc, num),
     flow::BouncyParticle, args...; subsample=false, oscn=false, factor=1.5, adapt=false) where {F1, F2}
+    if action == :invalid # invalidated event
+        θdϕ, v = dϕ(t, x, θ, args...) 
+        #∇ϕx = grad_correct!(∇ϕx, x, flow)
+        abc = ab(t, x, θ, c, θdϕ, v, flow)
+        t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
+    end
+
     while true
         if t + Δrec/V <= t′ # record! (large speed, more records)
             t, _ = move_forward!(V\Δrec, t, x, θ, flow) 
             Δrec = 1/flow.λref
+            ∇ϕ!(∇ϕx, t, x, args...)
+
             return t, V, (acc, num), c, abc, (t′, action), Δrec
         end
         Δrec = Δrec - (t′ - t)*V # coming closer to rec
@@ -179,7 +196,6 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
             V = norm(θ, 2)
             θdϕ, v = dϕ(t, x, θ, args...) 
             #∇ϕx = grad_correct!(∇ϕx, x, flow)
-            l = λ(θdϕ, flow) 
             abc = ab(t, x, θ, c, θdϕ, v, flow)
             t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
             #return t, V, (acc, num), c, abc, (t′, action), Δrec
@@ -227,9 +243,12 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
     end
 end
 """
-    pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; oscn=false, adapt=false, subsample=false, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
+    pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; oscn=false, adapt=false, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
+
+The first directional derivative `dϕ[1]` tells me if I move up or down the potential. The second directional derivative `dϕ[2]` tells me how fast the first changes.
+The gradient `∇ϕ!` tells me the surface I want to reflect on. Refreshes proportional to speed. 
+Keeps only samples at intervals proportional to those times.
 
-The first directional derivative `dϕ[1]` tells me if I move up or down the potential. The second directional derivative `dϕ[2]` tells me how fast the first changes. The gradient `∇ϕ!` tells me the surface I want to reflect on.
 
      dϕ = function (t, x, v, args...) # two directional derivatives
          u = ForwardDiff.derivative(t -> -ℓ(x + t*v), Dual{:hSrkahPmmC}(0.0, 1.0))
@@ -261,28 +280,30 @@ The remaining arguments:
     trace, final, (acc, num), cs = @time pdmp(
             dneglogp, # return first two directional derivatives of negative target log-likelihood in direction v
             ∇neglogp!, # return gradient of negative target log-likelihood
-            t0, x0, θ0, T, # initial state and duration
+            t0, x0, θ0, #initial state 
+            T, # duration (Real) or number of samples (Int)
             ZZB.LocalBound(c), # use Hessian information 
             Z; # sampler
             oscn=false, # no orthogonal subspace pCR
             adapt=true, # adapt bound c
+            adapt_mass=false # adapt PDiag U matrix to fisher information estimate
             progress=true, # show progress bar
-            subsample=true # keep only samples at refreshment times
     )
 
     # to obtain direction change times and points of piecewise linear trace
     t, x = ZigZagBoomerang.sep(trace)
 
 """
-function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; oscn=false, adapt=false, subsample=false, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
+function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; adapt_mass=false, oscn=false, adapt=false, subsample=true, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
     t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
+    subsample==true || throw(ArgumentError("`subsample=true` required."))
     V = norm(θ, 2)
     rng = Rng(seed)
     Ξ = Trace(t0, x0, θ0, flow)
     θdϕ, v = dϕ(t, x, θ, args...) 
     #@assert v2 ≈ v
     #@assert θdϕ ≈ dot(∇ϕx, θ)
-    
+   
     #∇ϕx = grad_correct!(∇ϕx, x, flow)
     num = acc = 0
     #l = 0.0
@@ -296,9 +317,22 @@ function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args.
     tstop = T/stops
     Δrec = 1/flow.λref
     t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
-    while t < T
+    m = ones(length(θ))
+    if adapt_mass 
+        m .= mass_adapt_init(flow.U)
+    end
+    iter = 1
+    while T isa Int ? iter < T : t < T
         t, V, (acc, num), c, abc, (t′, action), Δrec = pdmp_inner!(rng, dϕ, ∇ϕ!, ∇ϕx, t, x, θ, V, c, abc, (t′, action), Δrec, (acc, num), flow, args...; oscn=oscn, subsample=subsample, factor=factor, adapt=adapt)
         push!(Ξ, event(t, x, θ, flow))
+        iter += 1
+        if adapt_mass
+            @. m =  m + (∇ϕx^2 - m)/iter # running average
+            action = :invalid
+            PDMats.whiten!(flow.U, θ)
+            mass_adapt!(flow.U, m)
+            PDMats.unwhiten!(flow.U, θ)
+        end
 
         if t > tstop
             tstop += T/stops
diff --git a/src/notfactiter.jl b/src/notfactiter.jl
index d26331e..1cfcb47 100644
--- a/src/notfactiter.jl
+++ b/src/notfactiter.jl
@@ -55,6 +55,11 @@ function rawevent(t, x, θ, Z::Union{BouncyParticle,Boomerang})
 end
 
 ######
+function set_action(state, action)
+    u, ∇ϕx, (acc, num), c, abc, (t′, _action), Δrec = state
+    state = u, ∇ϕx, (acc, num), c, abc, (t′, action), Δrec 
+    state
+end
 
 function iterate(FS::NotFactSampler{<:Any, <:Tuple})
     t0, (x0, θ0) = FS.u0
diff --git a/test/maintest.jl b/test/maintest.jl
index 147d549..ba2e2c5 100644
--- a/test/maintest.jl
+++ b/test/maintest.jl
@@ -171,40 +171,75 @@ end
     @test mean(abs.(cov(xs) - inv(Matrix(Γ0)))) < 2/sqrt(T)
 end
 
-@testset "Bouncy Particle Sampler (mass matrix)" begin
+@testset "Bouncy Particle Sampler (arbitrary mass matrix)" begin
     Random.seed!(2)
     t0 = 0.0
     θ0 = randn(d)
     x0 = randn(d)
-    M = LowerTriangular(I + randn(d,d))
-    c = 1.1
-    B = BouncyParticle(nothing, nothing, # ignored
+    M = LowerTriangular(I + 0.4randn(d,d))
+    c = 20.0
+    B = BouncyParticle(missing, missing, # ignored
         1.0, # momentum refreshment rate 
-        0.0, # momentum correlation / only gradually change momentum in refreshment/momentum update
-        0.0, # ignored
+        0.9, # momentum correlation / only gradually change momentum in refreshment/momentum update
+        missing, # metric
         M # cholesky of momentum precision
     ) 
 
     ∇ϕ!(y, t, x, args...) = mul!(y, Γ, x)
     dϕ(t, x, v, args...) =  dot(v, Γ, x), dot(v, Γ, v)
-    T = 500.0
+    n = 800
     trace, _, acc, more = @time pdmp(
         dϕ, # return first two directional derivatives of negative target log-likelihood in direction v
         ∇ϕ!, # return gradient of negative target log-likelihood
-        t0, x0, θ0, T, # initial state and duration
+        t0, x0, θ0, # initial state and duration
+        n, # number of samples
         ZigZagBoomerang.LocalBound(c), # use Hessian information 
         B; # sampler
         adapt=false, # adapt bound c
         progress=true, # show progress bar
-        subsample=false # keep only samples at refreshment times
     )
-    #pdmp(∇ϕ!, t0, x0, θ0, T, c, B, progress=true)
     @show more
     @show acc[1]/acc[2]
-    dt = 0.1
-    ts, xs = sep(collect(discretize(trace, dt)))
-    @test mean(abs.(mean(xs))) < 2/sqrt(T)
-    @test mean(abs.(cov(xs) - inv(Matrix(Γ)))) < 2/sqrt(T)
+    ts, xs = sep(trace)
+    @show length(ts)
+    @test mean(abs.(mean(xs))) < 3/sqrt(length(ts))
+    @test mean(abs.(cov(xs) - inv(Matrix(Γ)))) < 3/sqrt(length(ts))
+end
+
+@testset "Bouncy Particle Sampler (adapted mass matrix)" begin
+    Random.seed!(2)
+    t0 = 0.0
+    θ0 = randn(d)
+    x0 = randn(d)
+    M = ZigZagBoomerang.PDMats.PDiagMat(ones(d))
+    c = 1.1
+    B = BouncyParticle(missing, missing, # ignored
+        1.0, # momentum refreshment rate 
+        0.9, # momentum correlation / only gradually change momentum in refreshment/momentum update
+        M, # metric
+        missing # cholesky of momentum precision
+    ) 
+
+    ∇ϕ!(y, t, x, args...) = mul!(y, Γ, x)
+    dϕ(t, x, v, args...) =  dot(v, Γ, x), dot(v, Γ, v)
+    n = 800
+    trace, _, acc, more = @time pdmp(
+        dϕ, # return first two directional derivatives of negative target log-likelihood in direction v
+        ∇ϕ!, # return gradient of negative target log-likelihood
+        t0, x0, θ0, # initial state and duration
+        n, # number of samples
+        ZigZagBoomerang.LocalBound(c), # use Hessian information 
+        B; # sampler
+        adapt=false, # adapt bound c
+        adapt_mass=true,
+        progress=true, # show progress bar
+    )
+    @show more
+    @show acc[1]/acc[2]
+    ts, xs = sep(trace)
+    @show length(ts)
+    @test mean(abs.(mean(xs))) < 2/sqrt(length(ts))
+    @test mean(abs.(cov(xs) - inv(Matrix(Γ)))) < 2/sqrt(length(ts))
 end
 
 @testset "ZigZag (independent)" begin

From 21f30950fa263d9ee9289723a07cc3986bd0499b Mon Sep 17 00:00:00 2001
From: mschauer <moritzschauer@web.de>
Date: Wed, 22 Jun 2022 14:42:47 +0200
Subject: [PATCH 5/8] Fixes

---
 src/dynamics.jl          | 13 -----------
 src/not_fact_samplers.jl | 49 ++++++++++++++++++++++++++++------------
 src/notfactiter.jl       |  4 ++--
 3 files changed, 36 insertions(+), 30 deletions(-)

diff --git a/src/dynamics.jl b/src/dynamics.jl
index 1fcf06a..0b9c183 100644
--- a/src/dynamics.jl
+++ b/src/dynamics.jl
@@ -125,16 +125,3 @@ function refresh!(rng, θ, F::Boomerang)
     θ
 end
 
-# Use Geometry and pdmats if L is not provided
-function ZigZagBoomerang.reflect!(∇ϕx, x, v, F::BouncyParticle{<:Any, <:Any, <:Any, <:AbstractPDMat}) # Seth's version
-    z = F.U * ∇ϕx
-    v .-= (2*dot(∇ϕx, v)/dot(∇ϕx, z)) * z
-    v
-end
-function ZigZagBoomerang.refresh!(rng, v, F::BouncyParticle{<:Any, <:Any, <:Any, <:AbstractPDMat})
-    ρ̄ = sqrt(1-F.ρ^2)
-    v .*= F.ρ
-    u = ρ̄*PDMats.unwhiten(F.U, randn(rng, length(v)))
-    v .+= u
-    v
-end
\ No newline at end of file
diff --git a/src/not_fact_samplers.jl b/src/not_fact_samplers.jl
index 7627ead..17c65e8 100644
--- a/src/not_fact_samplers.jl
+++ b/src/not_fact_samplers.jl
@@ -148,6 +148,20 @@ end
 
 
 ################################## ModernBPS ##################################
+# Use Geometry and pdmats if L is not provided
+function ZigZagBoomerang.reflect!(∇ϕx, x, v, F::BouncyParticle{<:Any, <:Any, <:Any, <:AbstractPDMat}) # Seth's version
+    z = F.U * ∇ϕx
+    v .-= (2*dot(∇ϕx, v)/dot(∇ϕx, z)) * z
+    v
+end
+function ZigZagBoomerang.refresh!(rng, v, F::BouncyParticle{<:Any, <:Any, <:Any, <:AbstractPDMat})
+    ρ̄ = sqrt(1-F.ρ^2)
+    v .*= F.ρ
+    u = ρ̄*PDMats.unwhiten(F.U, randn(rng, length(v)))
+    v .+= u
+    speed(v, F)
+end
+
 function mass_adapt_init(M)
     M.diag
 end
@@ -155,10 +169,11 @@ end
 function mass_adapt!(M, m)
      @. M.diag = 1/m
 end
+_working = 4
+speed(θ, F) = norm(whiten(F.U, θ), 2)    
 
-function ab(t, x, θ, C::LocalBound, vdϕ::Number, v, B::BouncyParticle)
-    @assert vdϕ isa Number
-    (C.c + vdϕ, v, t + 2sqrt(length(θ))/C.c/norm(θ, 2))
+function ab(t, x, θ, V, C::LocalBound, vdϕ::Number, v, B::BouncyParticle)
+    (C.c + vdϕ, v, t + 2sqrt(length(θ))/C.c/V)
 end
 
 function next_event1(rng, u::Tuple, abc, flow)
@@ -175,7 +190,7 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
     if action == :invalid # invalidated event
         θdϕ, v = dϕ(t, x, θ, args...) 
         #∇ϕx = grad_correct!(∇ϕx, x, flow)
-        abc = ab(t, x, θ, c, θdϕ, v, flow)
+        abc = ab(t, x, θ, V, c, θdϕ, v, flow)
         t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
     end
 
@@ -183,8 +198,10 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
         if t + Δrec/V <= t′ # record! (large speed, more records)
             t, _ = move_forward!(V\Δrec, t, x, θ, flow) 
             Δrec = 1/flow.λref
+            θdϕ, v = dϕ(t, x, θ, args...) 
             ∇ϕ!(∇ϕx, t, x, args...)
-
+            abc = ab(t, x, θ, V, c, θdϕ, v, flow)
+            t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
             return t, V, (acc, num), c, abc, (t′, action), Δrec
         end
         Δrec = Δrec - (t′ - t)*V # coming closer to rec
@@ -192,11 +209,10 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
         if action == :refresh
             @assert Δrec >= 0
             t, _ = move_forward!(t′ - t, t, x, θ, flow)
-            refresh!(rng, θ, flow)
-            V = norm(θ, 2)
+            V = refresh!(rng, θ, flow)
             θdϕ, v = dϕ(t, x, θ, args...) 
             #∇ϕx = grad_correct!(∇ϕx, x, flow)
-            abc = ab(t, x, θ, c, θdϕ, v, flow)
+            abc = ab(t, x, θ, V, c, θdϕ, v, flow)
             t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
             #return t, V, (acc, num), c, abc, (t′, action), Δrec
         elseif action == :expire
@@ -204,14 +220,14 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
             t, _ = move_forward!(τ, t, x, θ, flow) 
             θdϕ, v = dϕ(t, x, θ, args...) 
             #∇ϕx = grad_correct!(∇ϕx, x, flow)
-            abc = ab(t, x, θ, c, θdϕ, v, flow)
+            abc = ab(t, x, θ, V, c, θdϕ, v, flow)
             t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
         else # action == :reflect
             τ = t′ - t
             t, _ = move_forward!(τ, t, x, θ, flow)
             θdϕ, v = dϕ(t, x, θ, args...) 
             #∇ϕx = grad_correct!(∇ϕx, x, flow)
-            l, lb = λ(θdϕ, flow), pos(abc[1] + abc[2]*τ)
+            l, lb = θdϕ, pos(abc[1] + abc[2]*τ)
             num += 1
             if rand(rng)*lb <= l
                 acc += 1
@@ -228,15 +244,17 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
                     @assert flow.L == I
                     oscn!(rng, θ, ∇ϕx, flow.ρ; normalize=false)
                 else
+                    Vold = V
                     reflect!(∇ϕx, x, θ, flow)
+                    V = speed(θ, flow)
+                    @assert Vold ≈ V
                 end
                 θdϕ, v = dϕ(t, x, θ, args...) 
                 #∇ϕx = grad_correct!(∇ϕx, x, flow)
-                abc = ab(t, x, θ, c, θdϕ, v, flow)
+                abc = ab(t, x, θ, V, c, θdϕ, v, flow)
                 t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
-                #!subsample && return t, V, (acc, num), c, abc, (t′, action), Δrec
             else
-                abc = ab(t, x, θ, c, θdϕ, v, flow)
+                abc = ab(t, x, θ, V, c, θdϕ, v, flow)
                 t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
             end
         end
@@ -297,7 +315,7 @@ The remaining arguments:
 function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; adapt_mass=false, oscn=false, adapt=false, subsample=true, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
     t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
     subsample==true || throw(ArgumentError("`subsample=true` required."))
-    V = norm(θ, 2)
+    V = speed(θ, flow)
     rng = Rng(seed)
     Ξ = Trace(t0, x0, θ0, flow)
     θdϕ, v = dϕ(t, x, θ, args...) 
@@ -307,7 +325,7 @@ function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args.
     #∇ϕx = grad_correct!(∇ϕx, x, flow)
     num = acc = 0
     #l = 0.0
-    abc = ab(t, x, θ, c, θdϕ, v, flow)
+    abc = ab(t, x, θ, V, c, θdϕ, v, flow)
     if progress
         prg = Progress(progress_stops, 1)
     else
@@ -343,6 +361,7 @@ function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args.
     return Ξ, (t, x, θ), (acc, num), c
 end
 
+##########
 wrap(f) = wrap_(f,  methods(f)...)
 wrap_(f, args...) = f
 @inline wrap_(f, m) = m.nargs <= 4 ? Wrapper(f) : f
diff --git a/src/notfactiter.jl b/src/notfactiter.jl
index 1cfcb47..d42da4e 100644
--- a/src/notfactiter.jl
+++ b/src/notfactiter.jl
@@ -66,7 +66,7 @@ function iterate(FS::NotFactSampler{<:Any, <:Tuple})
     flow = FS.F
     n = length(x0)
     t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
-    V = norm(θ, 2)
+    V = speed(θ, flow)
     c = FS.c
     rng = FS.rng
     Δrec = 1/flow.λref
@@ -74,7 +74,7 @@ function iterate(FS::NotFactSampler{<:Any, <:Tuple})
     dϕ, ∇ϕ! = FS.∇ϕ![1], FS.∇ϕ![2]  
     θdϕ, v = dϕ(t, x, θ, FS.args...) 
     num = acc = 0
-    abc = ab(t, x, θ, c, θdϕ, v, flow)
+    abc = ab(t, x, θ, V, c, θdϕ, v, flow)
     t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
     iterate(FS, ((t => (x, θ, V)), ∇ϕx, (acc, num), c, abc, (t′, action), Δrec))
 end

From e1ea285ff538da714e9023b5751a3a3e5641ae71 Mon Sep 17 00:00:00 2001
From: mschauer <moritzschauer@web.de>
Date: Mon, 27 Jun 2022 15:39:58 +0200
Subject: [PATCH 6/8] Add local

---
 src/ZigZagBoomerang.jl   |  1 +
 src/invchol.jl           | 28 ++++++++++++++++
 src/not_fact_samplers.jl | 72 +++++++++++++++++++++-------------------
 src/notfactiter.jl       |  4 ++-
 4 files changed, 70 insertions(+), 35 deletions(-)
 create mode 100644 src/invchol.jl

diff --git a/src/ZigZagBoomerang.jl b/src/ZigZagBoomerang.jl
index db4ae8e..5b4a9bb 100644
--- a/src/ZigZagBoomerang.jl
+++ b/src/ZigZagBoomerang.jl
@@ -15,6 +15,7 @@ Seed() = gen_seed(UInt64, 2)
 include("types.jl")
 include("common.jl")
 include("oscn.jl")
+include("invchol.jl")
 include("dynamics.jl")
 export ZigZag1d, Boomerang1d, ZigZag, FactBoomerang
 const LocalZigZag = ZigZag
diff --git a/src/invchol.jl b/src/invchol.jl
new file mode 100644
index 0000000..6d138e7
--- /dev/null
+++ b/src/invchol.jl
@@ -0,0 +1,28 @@
+using LinearAlgebra
+struct InvChol{T} <: AbstractPDMat{Float64}
+    R::T
+    InvChol(R::S) where {S<:UpperTriangular} = new{S}(R)
+end
+PDMats.dim(M::InvChol) = size(M.R,1)
+function PDMats.unwhiten!(r::Vector{Float64}, M::InvChol{UpperTriangular{Float64, Matrix{Float64}}}, z::Vector{Float64})
+    r .= z
+    LinearAlgebra.naivesub!(M.R, r)
+    r
+end    
+function LinearAlgebra.mul!(r::AbstractVector, M::InvChol, x::AbstractVector, alpha::Number, beta::Number)
+    @assert beta==0
+    @. r = alpha*x
+    LinearAlgebra.naivesub!(M.R', r)
+    LinearAlgebra.naivesub!(M.R, r)
+    r
+end
+
+function PDMats.whiten!(r::Vector{Float64}, M::InvChol{UpperTriangular{Float64, Matrix{Float64}}}, x::Vector{Float64})
+    mul!(r, M.R, x)
+    r
+end    
+function Base.show(io::IOContext, m::MIME{Symbol("text/plain")}, M::InvChol) 
+    print(io, "CholInv: ")
+    show(io, m, M.R)
+end
+Base.Matrix(M::InvChol) = M.R'*M.R
\ No newline at end of file
diff --git a/src/not_fact_samplers.jl b/src/not_fact_samplers.jl
index 17c65e8..15bbe60 100644
--- a/src/not_fact_samplers.jl
+++ b/src/not_fact_samplers.jl
@@ -148,19 +148,26 @@ end
 
 
 ################################## ModernBPS ##################################
+function localσ(t, x, v, F)
+    1.0
+end
 # Use Geometry and pdmats if L is not provided
-function ZigZagBoomerang.reflect!(∇ϕx, x, v, F::BouncyParticle{<:Any, <:Any, <:Any, <:AbstractPDMat}) # Seth's version
+function ZigZagBoomerang.reflect!(∇ϕx, t, x, v, F::BouncyParticle{<:Any, <:Any, <:Any, <:AbstractPDMat}) # Seth's version
     z = F.U * ∇ϕx
     v .-= (2*dot(∇ϕx, v)/dot(∇ϕx, z)) * z
     v
 end
-function ZigZagBoomerang.refresh!(rng, v, F::BouncyParticle{<:Any, <:Any, <:Any, <:AbstractPDMat})
+function ZigZagBoomerang.refresh!(rng, t, x, v, F::BouncyParticle{<:Any, <:Any, <:Any, <:AbstractPDMat})
     ρ̄ = sqrt(1-F.ρ^2)
     v .*= F.ρ
-    u = ρ̄*PDMats.unwhiten(F.U, randn(rng, length(v)))
+    s = localσ(t, x, v, F)
+    u = ρ̄*PDMats.unwhiten(F.U, s*randn(rng, length(v)))
     v .+= u
     speed(v, F)
 end
+function mass_adapt_init(M::InvChol)
+    Cholesky(M.R)
+end
 
 function mass_adapt_init(M)
     M.diag
@@ -169,7 +176,10 @@ end
 function mass_adapt!(M, m)
      @. M.diag = 1/m
 end
-_working = 4
+function mass_adapt!(M::InvChol, m)
+    M.R.data .= m.U.data 
+end
+
 speed(θ, F) = norm(whiten(F.U, θ), 2)    
 
 function ab(t, x, θ, V, C::LocalBound, vdϕ::Number, v, B::BouncyParticle)
@@ -188,7 +198,8 @@ end
 function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, abc, (t′, action), Δrec, (acc, num),
     flow::BouncyParticle, args...; subsample=false, oscn=false, factor=1.5, adapt=false) where {F1, F2}
     if action == :invalid # invalidated event
-        θdϕ, v = dϕ(t, x, θ, args...) 
+        θdϕ, v = dϕ(t, x, θ, flow, args...) 
+        ∇ϕ!(∇ϕx, t, x, θ, flow, args...)
         #∇ϕx = grad_correct!(∇ϕx, x, flow)
         abc = ab(t, x, θ, V, c, θdϕ, v, flow)
         t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
@@ -198,8 +209,8 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
         if t + Δrec/V <= t′ # record! (large speed, more records)
             t, _ = move_forward!(V\Δrec, t, x, θ, flow) 
             Δrec = 1/flow.λref
-            θdϕ, v = dϕ(t, x, θ, args...) 
-            ∇ϕ!(∇ϕx, t, x, args...)
+            θdϕ, v = dϕ(t, x, θ, flow, args...) 
+            ∇ϕ!(∇ϕx, t, x, θ, flow, args...)
             abc = ab(t, x, θ, V, c, θdϕ, v, flow)
             t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
             return t, V, (acc, num), c, abc, (t′, action), Δrec
@@ -209,8 +220,8 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
         if action == :refresh
             @assert Δrec >= 0
             t, _ = move_forward!(t′ - t, t, x, θ, flow)
-            V = refresh!(rng, θ, flow)
-            θdϕ, v = dϕ(t, x, θ, args...) 
+            V = refresh!(rng, t, x, θ, flow)
+            θdϕ, v = dϕ(t, x, θ, flow, args...) 
             #∇ϕx = grad_correct!(∇ϕx, x, flow)
             abc = ab(t, x, θ, V, c, θdϕ, v, flow)
             t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
@@ -218,15 +229,13 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
         elseif action == :expire
             τ = t′ - t
             t, _ = move_forward!(τ, t, x, θ, flow) 
-            θdϕ, v = dϕ(t, x, θ, args...) 
-            #∇ϕx = grad_correct!(∇ϕx, x, flow)
+            θdϕ, v = dϕ(t, x, θ, flow, args...) 
             abc = ab(t, x, θ, V, c, θdϕ, v, flow)
             t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
         else # action == :reflect
             τ = t′ - t
             t, _ = move_forward!(τ, t, x, θ, flow)
-            θdϕ, v = dϕ(t, x, θ, args...) 
-            #∇ϕx = grad_correct!(∇ϕx, x, flow)
+            θdϕ, v = dϕ(t,x, θ, flow, args...) 
             l, lb = θdϕ, pos(abc[1] + abc[2]*τ)
             num += 1
             if rand(rng)*lb <= l
@@ -235,22 +244,15 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
                     !adapt && error("Tuning parameter `c` too small.")
                     c *= factor
                 end
-                ∇ϕ!(∇ϕx, t, x, args...)
-                if !(dot(θ, ∇ϕx) ≈ θdϕ)
-                    #@show  dot(θ, ∇ϕx) θdϕ
-                    #error("subsampling needs to be seeded by time")
-                end
+                ∇ϕ!(∇ϕx, t, x, θ, flow, args...)
                 if oscn
                     @assert flow.L == I
                     oscn!(rng, θ, ∇ϕx, flow.ρ; normalize=false)
                 else
-                    Vold = V
-                    reflect!(∇ϕx, x, θ, flow)
-                    V = speed(θ, flow)
-                    @assert Vold ≈ V
+                    reflect!(∇ϕx, t, x, θ, flow)
+                    #V = speed(θ, flow)
                 end
                 θdϕ, v = dϕ(t, x, θ, args...) 
-                #∇ϕx = grad_correct!(∇ϕx, x, flow)
                 abc = ab(t, x, θ, V, c, θdϕ, v, flow)
                 t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
             else
@@ -312,17 +314,15 @@ The remaining arguments:
     t, x = ZigZagBoomerang.sep(trace)
 
 """
-function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; adapt_mass=false, oscn=false, adapt=false, subsample=true, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
+function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; iter_offset=0, adapt_mass=false, oscn=false, adapt=false, subsample=true, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
     t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
     subsample==true || throw(ArgumentError("`subsample=true` required."))
     V = speed(θ, flow)
     rng = Rng(seed)
     Ξ = Trace(t0, x0, θ0, flow)
-    θdϕ, v = dϕ(t, x, θ, args...) 
-    #@assert v2 ≈ v
-    #@assert θdϕ ≈ dot(∇ϕx, θ)
-   
-    #∇ϕx = grad_correct!(∇ϕx, x, flow)
+    θdϕ, v = dϕ(t, x, θ, flow, args...) 
+    ∇ϕ!(∇ϕx, t, x, θ, flow, args...)
+
     num = acc = 0
     #l = 0.0
     abc = ab(t, x, θ, V, c, θdϕ, v, flow)
@@ -335,17 +335,21 @@ function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args.
     tstop = T/stops
     Δrec = 1/flow.λref
     t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
-    m = ones(length(θ))
     if adapt_mass 
-        m .= mass_adapt_init(flow.U)
+        m = mass_adapt_init(flow.U)
     end
-    iter = 1
+    iter = iter_offset
     while T isa Int ? iter < T : t < T
         t, V, (acc, num), c, abc, (t′, action), Δrec = pdmp_inner!(rng, dϕ, ∇ϕ!, ∇ϕx, t, x, θ, V, c, abc, (t′, action), Δrec, (acc, num), flow, args...; oscn=oscn, subsample=subsample, factor=factor, adapt=adapt)
         push!(Ξ, event(t, x, θ, flow))
         iter += 1
-        if adapt_mass
-            @. m =  m + (∇ϕx^2 - m)/iter # running average
+        if adapt_mass # todo: make function
+            if m isa Cholesky
+                m.factors .*= sqrt(1-1/iter)
+                LinearAlgebra.lowrankupdate!(HC, ∇ϕx/sqrt(iter))
+            else
+                @. m =  m + (∇ϕx^2 - m)/iter # running average
+            end
             action = :invalid
             PDMats.whiten!(flow.U, θ)
             mass_adapt!(flow.U, m)
diff --git a/src/notfactiter.jl b/src/notfactiter.jl
index d42da4e..4284289 100644
--- a/src/notfactiter.jl
+++ b/src/notfactiter.jl
@@ -72,7 +72,9 @@ function iterate(FS::NotFactSampler{<:Any, <:Tuple})
     Δrec = 1/flow.λref
 
     dϕ, ∇ϕ! = FS.∇ϕ![1], FS.∇ϕ![2]  
-    θdϕ, v = dϕ(t, x, θ, FS.args...) 
+    θdϕ, v = dϕ(t, x, θ, flow, FS.args...) 
+    ∇ϕ!(∇ϕx, t, x, θ, flow, FS.args...)
+   
     num = acc = 0
     abc = ab(t, x, θ, V, c, θdϕ, v, flow)
     t′, action = next_event1(rng, (t, x, θ, V), abc, flow)

From e5c7a5f05d0a4293149176446f5b5dbf917a20ea Mon Sep 17 00:00:00 2001
From: mschauer <moritzschauer@web.de>
Date: Tue, 28 Jun 2022 14:11:01 +0200
Subject: [PATCH 7/8] Allow local

---
 src/not_fact_samplers.jl | 32 +++++++++++++++-----------------
 src/notfactiter.jl       |  2 +-
 2 files changed, 16 insertions(+), 18 deletions(-)

diff --git a/src/not_fact_samplers.jl b/src/not_fact_samplers.jl
index 15bbe60..ad393cd 100644
--- a/src/not_fact_samplers.jl
+++ b/src/not_fact_samplers.jl
@@ -148,22 +148,23 @@ end
 
 
 ################################## ModernBPS ##################################
-function localσ(t, x, v, F)
+function local_speed(t, x, v, F)
     1.0
 end
+
 # Use Geometry and pdmats if L is not provided
 function ZigZagBoomerang.reflect!(∇ϕx, t, x, v, F::BouncyParticle{<:Any, <:Any, <:Any, <:AbstractPDMat}) # Seth's version
-    z = F.U * ∇ϕx
+    z = F.U * ∇ϕx # constant factor cancels
     v .-= (2*dot(∇ϕx, v)/dot(∇ϕx, z)) * z
     v
 end
 function ZigZagBoomerang.refresh!(rng, t, x, v, F::BouncyParticle{<:Any, <:Any, <:Any, <:AbstractPDMat})
     ρ̄ = sqrt(1-F.ρ^2)
     v .*= F.ρ
-    s = localσ(t, x, v, F)
-    u = ρ̄*PDMats.unwhiten(F.U, s*randn(rng, length(v)))
+    s = local_speed(t, x, v, F)
+    u = (s*ρ̄)*PDMats.unwhiten(F.U, randn(rng, length(v)))
     v .+= u
-    speed(v, F)
+    record_rate(v, F)
 end
 function mass_adapt_init(M::InvChol)
     Cholesky(M.R)
@@ -180,7 +181,7 @@ function mass_adapt!(M::InvChol, m)
     M.R.data .= m.U.data 
 end
 
-speed(θ, F) = norm(whiten(F.U, θ), 2)    
+record_rate(θ, F) = norm(whiten(F.U, θ))    
 
 function ab(t, x, θ, V, C::LocalBound, vdϕ::Number, v, B::BouncyParticle)
     (C.c + vdϕ, v, t + 2sqrt(length(θ))/C.c/V)
@@ -200,13 +201,12 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
     if action == :invalid # invalidated event
         θdϕ, v = dϕ(t, x, θ, flow, args...) 
         ∇ϕ!(∇ϕx, t, x, θ, flow, args...)
-        #∇ϕx = grad_correct!(∇ϕx, x, flow)
         abc = ab(t, x, θ, V, c, θdϕ, v, flow)
         t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
     end
 
     while true
-        if t + Δrec/V <= t′ # record! (large speed, more records)
+        if t + Δrec/V <= t′ # record! (large V, more records)
             t, _ = move_forward!(V\Δrec, t, x, θ, flow) 
             Δrec = 1/flow.λref
             θdϕ, v = dϕ(t, x, θ, flow, args...) 
@@ -222,10 +222,8 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
             t, _ = move_forward!(t′ - t, t, x, θ, flow)
             V = refresh!(rng, t, x, θ, flow)
             θdϕ, v = dϕ(t, x, θ, flow, args...) 
-            #∇ϕx = grad_correct!(∇ϕx, x, flow)
             abc = ab(t, x, θ, V, c, θdϕ, v, flow)
             t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
-            #return t, V, (acc, num), c, abc, (t′, action), Δrec
         elseif action == :expire
             τ = t′ - t
             t, _ = move_forward!(τ, t, x, θ, flow) 
@@ -235,7 +233,7 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
         else # action == :reflect
             τ = t′ - t
             t, _ = move_forward!(τ, t, x, θ, flow)
-            θdϕ, v = dϕ(t,x, θ, flow, args...) 
+            θdϕ, v = dϕ(t, x, θ, flow, args...) 
             l, lb = θdϕ, pos(abc[1] + abc[2]*τ)
             num += 1
             if rand(rng)*lb <= l
@@ -248,11 +246,12 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
                 if oscn
                     @assert flow.L == I
                     oscn!(rng, θ, ∇ϕx, flow.ρ; normalize=false)
+                    V = record_rate(θ, flow)
                 else
                     reflect!(∇ϕx, t, x, θ, flow)
-                    #V = speed(θ, flow)
+                    V = record_rate(θ, flow)
                 end
-                θdϕ, v = dϕ(t, x, θ, args...) 
+                θdϕ, v = dϕ(t, x, θ, flow, args...) 
                 abc = ab(t, x, θ, V, c, θdϕ, v, flow)
                 t′, action = next_event1(rng, (t, x, θ, V), abc, flow)
             else
@@ -290,10 +289,10 @@ The remaining arguments:
     c = 50.0 # initial guess for the bound
 
     # define BouncyParticle sampler (has two relevant parameters) 
-    Z = BouncyParticle(∅, ∅, # ignored
+    Z = BouncyParticle(∅, ∅, # information about graphical structure
         10.0, # momentum refreshment rate 
         0.95, # momentum correlation / only gradually change momentum in refreshment/momentum update
-        0.0, # ignored
+        nothing, # PDMats compatible inverse mass OR
         I # left cholesky factor of momentum precision
     ) 
 
@@ -317,14 +316,13 @@ The remaining arguments:
 function pdmp(dϕ, ∇ϕ!, t0, x0, θ0, T, c::Bound, flow::BouncyParticle, args...; iter_offset=0, adapt_mass=false, oscn=false, adapt=false, subsample=true, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
     t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
     subsample==true || throw(ArgumentError("`subsample=true` required."))
-    V = speed(θ, flow)
+    V = record_rate(θ, flow)
     rng = Rng(seed)
     Ξ = Trace(t0, x0, θ0, flow)
     θdϕ, v = dϕ(t, x, θ, flow, args...) 
     ∇ϕ!(∇ϕx, t, x, θ, flow, args...)
 
     num = acc = 0
-    #l = 0.0
     abc = ab(t, x, θ, V, c, θdϕ, v, flow)
     if progress
         prg = Progress(progress_stops, 1)
diff --git a/src/notfactiter.jl b/src/notfactiter.jl
index 4284289..7ae644e 100644
--- a/src/notfactiter.jl
+++ b/src/notfactiter.jl
@@ -66,7 +66,7 @@ function iterate(FS::NotFactSampler{<:Any, <:Tuple})
     flow = FS.F
     n = length(x0)
     t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
-    V = speed(θ, flow)
+    V = record_rate(θ, flow)
     c = FS.c
     rng = FS.rng
     Δrec = 1/flow.λref

From 34c29a7d5610fdd71a815dddbaf22a0a14430cca Mon Sep 17 00:00:00 2001
From: mschauer <moritzschauer@web.de>
Date: Fri, 1 Jul 2022 15:49:55 +0200
Subject: [PATCH 8/8] Update bounds on recordings

---
 src/not_fact_samplers.jl | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/not_fact_samplers.jl b/src/not_fact_samplers.jl
index ad393cd..f7e26ab 100644
--- a/src/not_fact_samplers.jl
+++ b/src/not_fact_samplers.jl
@@ -207,9 +207,15 @@ function pdmp_inner!(rng, dϕ::F1, ∇ϕ!::F2, ∇ϕx, t, x, θ, V, c::Bound, ab
 
     while true
         if t + Δrec/V <= t′ # record! (large V, more records)
-            t, _ = move_forward!(V\Δrec, t, x, θ, flow) 
+            τ = V\Δrec
+            t, _ = move_forward!(τ, t, x, θ, flow) 
             Δrec = 1/flow.λref
             θdϕ, v = dϕ(t, x, θ, flow, args...) 
+            l, lb = θdϕ, pos(abc[1] + abc[2]*τ)
+            if l > lb # check bounds on recordings
+                !adapt && error("Tuning parameter `c` too small.")
+                c *= factor
+            end
             ∇ϕ!(∇ϕx, t, x, θ, flow, args...)
             abc = ab(t, x, θ, V, c, θdϕ, v, flow)
             t′, action = next_event1(rng, (t, x, θ, V), abc, flow)