Virtual Brownian Tree

[frankschae]

This PR adds a first version of a Virtual Brownian Tree (VBT) based on counter-based PRNGs.

The counter-based PRNGs should be good for our purpose:
http://www.thesalmons.org/john/random123/papers/random123sc11.pdf
https://github.com/sunoru/RandomNumbers.jl/issues/37

Though using some other RNGs with a split function (see, e.g., https://dl.acm.org/doi/pdf/10.1145/2578854.2503784) might be better. However, I didn't found any deterministic split(seed) function in Base.Random, RandomNumbers or Random123 .. but I found quite a few discussions that simply using a randjump in the sequence of a MarsenneTwister is dangerous.

https://discourse.julialang.org/t/best-practices-for-parallel-generation-of-pseudo-random-numbers/26504
https://discourse.julialang.org/t/parallel-random-number-generation/1950/10

The VBT is build up very similar to NoiseGrid in general.

The VBT has three unique arguments:

• The VBT has a finite search_depth keyword which restricts the depth in the tree that should be searched. Setting this argument is necessary at the moment due to the choice of the PRNG (see below).
• The search_depth will be chosen automatically based on an atol or can be set manually.
• tree_depth stores a small cache of time steps, noise values and seeds such that one can trade memory for speed if desired.

Since there was no pre-defined splitting function (or at least I haven't fount one). I implemented our own splitting function split_VBT_seed(rng::Random123.AbstractR123, parent_seed, current_depth, Nt)

This splitting function makes use of the following properties:

• The total number of accessible timesteps / noise values is Nt = 2^D + 1 where D is the maximum search_depth
• If the starting point is labeled with the counter_start = seeded_counter + 1 and the end point is labeled with counter_end = 2^D+1, then one can define unique counters by t_mid = (tstart - tend)/2 , (Nt+1)/2, and so on. Ultimately, the node with index j at level l is given by seeded_counter + 1 + (2*j+1)*Nt/2^l .
• We can thus iteratively construct two new counters (for the left and right process of the bisection) given the parent counter and the current depth of the tree by: parent_seed -+ (Nt-1)/2^(current_depth+1).

[ChrisRackauckas]

Though using some other RNGs with a split function (see, e.g., https://dl.acm.org/doi/pdf/10.1145/2578854.2503784) might be better. However, I didn't found any deterministic split(seed) function in Base.Random, RandomNumbers or Random123 .. but I found quite a few discussions that simply using a randjump in the sequence of a MarsenneTwister is dangerous.

Yes, avoid MarsenneTwister here.

Since there was no pre-defined splitting function (or at least I haven't fount one). I implemented our own splitting function split_VBT_seed(rng::Random123.AbstractR123, parent_seed, current_depth, Nt)

Is there a good way to do this for the CUDA.jl RNG as well? We can add that to the @require block. Pinging @maleadt in case he knows about this portion.

[ChrisRackauckas]

I emailed someone for comments on RNGs. But as is, I think the PR is really great and I don't see major issues with it. I would just test for the cache size and get a fully non-allocating version before merging.

Handling RNG splitting for the GPU RNG can get its own issue. Handling other CPU-based RNGs, such as the "standard ones" could get another low priority issue, low priority since using Random123 isn't a bad idea anyways (but it's good to have an issue track the current limitation). We'll need to get this into the DiffEq docs as well, and get a full-scale test setup with Neural SDEs which will likely be the main use case.

[frankschae]

Is there a good way to do this for the CUDA.jl RNG as well? We can add that to the @require block. Pinging @maleadt in case he knows about this portion.

On the CURAND site https://github.com/JuliaAttic/CURAND.jl (which seems to be called for the RNGs in CUDA.jl https://github.com/JuliaGPU/CUDA.jl/blob/fd59518cf6080f74fc8e2ff309c31955009bb39a/src/random.jl#L13). There is

CURAND_RNG_PSEUDO_PHILOX4_32_10

This should be a counter-based PRNG (https://sunoru.github.io/RandomNumbers.jl/stable/man/random123/#Random123-Family-1) for which we might be able to do the same.

[mschauer]

You should perhaps test the distributional properties of the resulting bridge. Can you also plot a trajectory for me to eyeball?

[ChrisRackauckas]

Yes, put it to the test in the bridge test which tests the mean and variance. You'll see this is done for the NoiseProcess, so add a VBT noise process to that test.

[frankschae]

It looks like there is a problem with the distributional properties. For this test:

using DiffEqNoiseProcess, DiffEqBase, Test, Random, DiffEqBase.EnsembleAnalysis
W = VirtualBrownianTree(0.0,0.0; Wend=1.0,tree_depth=3)
prob = NoiseProblem(W,(0.0,1.0))
function prob_func(prob,i,repeat)
  # to re-instantiate PRNG
  Wtmp = VirtualBrownianTree(0.0,0.0; Wend=1.0,tree_depth=3)
  remake(prob, noise=Wtmp)
end
ensemble_prob = EnsembleProblem(prob, prob_func=prob_func)
@time sol = solve(ensemble_prob,dt=0.125,trajectories=10000)

# Spot check the mean and the variance
qs = 0:0.125:1
for i in 2:8
  q = qs[i]
  @test ≈(timestep_mean(sol,i),q,atol=1e-2) ##fails
  @test ≈(timestep_meanvar(sol,i)[2],(1-q)*q,atol=1e-2) ##fails
end
@test ≈(timestep_mean(sol,1)[1],0.0,atol=1e-16)
@test ≈(timestep_meanvar(sol,1)[2],0.0,atol=1e-16)
@test ≈(timestep_mean(sol,Int(2^(W.tree_depth)+1))[1],W.W[end],atol=1e-16)
@test ≈(timestep_meanvar(sol,Int(2^(W.tree_depth)+1))[2],0.0,atol=1e-16)


using Plots
plt1 = plot(timeseries_steps_mean(sol),label=false)
plot!(qs,x->x, label="mean")

plt2 = plot(timeseries_steps_meanvar(sol)[2],label=false)
plot!(qs,x->(1-x)*x, label="var")

plt = plot(plt1,plt2, layout=(2,1))

The plot has some weird kinks:

This originates from the stage where the cache is built. From a print within create_VBT_cache, I obtain for a tree with a small search_depth=5 W = VirtualBrownianTree(0.0,0.0; Wend=1.0,tree_depth=3, search_depth=5)

level: 1
Int(seed_v) = 17
(t0, t, t1) = (0.0, 0.5, 1.0)
(W0, w, W1) = (0.0, 0.8071189451673877, 1.0)
(q, h) = (1//2, 1.0)
-------
level: 2
Int(seed_v) = 9
(t0, t, t1) = (0.0, 0.25, 0.5)
(W0, w, W1) = (0.0, -0.015543175417636345, 0.8071189451673877)
(q, h) = (1//2, 0.5)
-------
Int(seed_v) = 25
(t0, t, t1) = (0.5, 0.75, 1.0)
(W0, w, W1) = (0.8071189451673877, 0.40242391761082197, 1.0)
(q, h) = (1//2, 0.5)
-------
level: 3
Int(seed_v) = 5
(t0, t, t1) = (0.0, 0.125, 0.25)
(W0, w, W1) = (0.0, 0.10451036399995671, -0.015543175417636345)
(q, h) = (1//2, 0.25)
-------
Int(seed_v) = 13
(t0, t, t1) = (0.25, 0.375, 0.5)
(W0, w, W1) = (-0.015543175417636345, 0.5858275647585797, 0.8071189451673877)
(q, h) = (1//2, 0.25)
-------
Int(seed_v) = 21
(t0, t, t1) = (0.5, 0.625, 0.75)
(W0, w, W1) = (0.8071189451673877, 0.19220819331483543, 0.40242391761082197)
(q, h) = (1//2, 0.25)
-------
Int(seed_v) = 29
(t0, t, t1) = (0.75, 0.875, 1.0)
(W0, w, W1) = (0.40242391761082197, 0.6737623750611852, 1.0)
(q, h) = (1//2, 0.25)
———

(the small search_depth is used to check if the seeds are assigned properly). However, I don't see the error in these assignments..

[mschauer]

Looks like the mean is computed interpolating 0 and the right endpoint instead of the left and the right endpoint.

[ChrisRackauckas]

It was mentioned to me that we might want to make use of http://www.sprng.org/, so that should go into a follow up issue as well.

[frankschae]

@mschauer
That was it! (I assumed that WHITE_NOISE_BRIDGE actually takes W0 into account) but it assumes to start from W0=0.
So I added a VBT_BRIDGE where the only change is that instead of q*Wh, the mean value q*(W0+Wh) is added.

[ChrisRackauckas]

That was it! (I assumed that WHITE_NOISE_BRIDGE actually takes W0 into account) but it assumes to start from W0=0.
So I added a VBT_BRIDGE where the only change is that instead of qWh, the mean value q(W0+Wh) is added.

Yeah, it's because of how the extra terms are stored on the stack. It's a little odd but it works out in the end. Your change looks like it's the correct thing to do here.

[frankschae]

Here are two plots from the trajectories:

using Plots
W = VirtualBrownianTree(0.0,0.0; tree_depth=10)

# step to cached value
dt = W.t[2]-W.t[1]
calculate_step!(W,dt,nothing,nothing)

Ws = []
for i in 1:length(W.t)-1
  accept_step!(W,dt,nothing,nothing)
  push!(Ws,W.curW)
end

plt = plot(Ws)
savefig(plt, "trajectory1.png")

W = VirtualBrownianTree(0.0,0.0; tree_depth=0)
# step to some interpolated values
dt = 0.001
calculate_step!(W,dt,nothing,nothing)

Ws = []
for i in 1:1000
  accept_step!(W,dt,nothing,nothing)
  push!(Ws,W.curW)
end

plt = plot(Ws)
savefig(plt, "trajectory2.png")

[andreasnoack]

@frankschae This PR introduces a dependency on a C++ compiler because Random123.jl builds a few C++ files from source, see https://github.com/sunoru/Random123.jl/tree/234208c4cbbc1fb961dcd1087fd86bfe5b282684/deps. It would probably be worthwhile to figure out if there is a way to avoid this. Either by using BinaryBuilder, converting the C++ pieces to Julia, or considering an alternative dependency.

[chriselrod]

I should be able to translate this to Julia (using llvmcall through SIMD.jl or VectorizationBase.jl).