Implementation Details of PDMPFlux.jl

Simulating PDMPs with Automatic Differentiation

Julia
MCMC
Author

Hirofumi Shiba

Published

12/31/2024

Modified

1/01/2025

1 Directory Structure of PDMPFlux.jl/src

1.1 Composites.jl

composites / constructs such as BoundBox, PDMPState, PDMPHistory are defined.

  • BoundBox is basically a grid, together with the values on it, to perform thinning.

  • PDMPState is used to keep the position of the sampler. The whole SamplingLoop.jl Section 1.3 is implemented as functions that modify the PDMPState in place. In addition to the \((x,v,t)\) tuple, the fields include

    • state.rate and state.rate_vect are the rate function and its vectorized version. These slots are determined after checking the field pdmp.signed_bound in init_state() in AbstractPDMP.jl Section 2.1.
    • state.upper_bound_func() is a function that takes \((x,v,t)\) and returns the BoundBox object. Basically, state.upper_bound_func() is the only field that differs among samplers.
    • boolean flags to indicate whether Poisson thinning proposal is accepted or not, which is used within ac_step_with_proxy() in the SamplingLoop.jl Section 1.3.
    • proposed jump time tp and time already spent ts. They are using in SamplingLoop.jl Section 1.3 to conduct poisson thinning, where tp is the proposed jump time, and is added to ts when rejected.
    • statistics to diagnose the sampler dynamic.

1.2 sample.jl

sample.jl contains functions, called by users, to start MCMC sampling. The ProgressBar package is used to be friendly to users.

sample.jl contains 2 functions, and sample() that calls them in sequel.

  • sample_skeleton() -> PDMPHistory initializes the progress bar and PDMPState & PDMPHistory objects, using init_state() and PDMPHistory() constructor respectively. Then call get_event_state() in SamplingLoop.jl Section 1.3 for iter::Int times.
  • sample_from_skeleton() -> Matrix{Float64} is a function that generates samples, which might subsequently be called by the plotting functions in plot.jl.

1.3 SamplingLoop.jl

This module contains 10 functions. Some of them are summarized in the following figure:

The main loop for sampling is implemented here.

  • get_event_state() -> PDMPState is the entering point from sample.jl.

    • There is only one thing in get_event_state(), i.e., callling one_step_of_thinning() until state.accept becomes true.

    • get_event_state() returns a PDMPState object with the field \((x,v,t)\) indicating where an accepted event happens, which is pushed to PDMPHistory back in sample_skeleton() in sample.jl Section 1.2.

  • one_step_of_thinning() returns where simulation has completed up to, with state.accept = false if any proposal isnot accepted yet.

    • Firstly, it proposes the next jump event, by simulating exp_rv and calling next_event(upper_bound, exp_rv) -> (tp, lambda_bar).

      • upper_bound is the BoundBox object, which is calculated by state.upper_bound_func() and the current state (x,v,t).
      • tp is the proposed jump time.
      • lambda_bar is a upper bound, with error, for the rate function \(\lambda\).

    • Secondly, it checks whether tp <= state.horizon.

      • If not, it calls move_to_horizon() and continues the loop one_step_of_thinning() inside get_event_state().
      • If yes, it proceeds to moves_until_horizon(), where another loop, calling ac_step(), is performed until one of the following 2 happens
        • when state.accept becomes true, it gets out of the both loops.
        • when tp > state.horizon, it calls move_to_horizon2() to get out of ac_step() loop, and continues the loop one_step_of_thinning() inside get_event_state().

  • ac_step(), standing for acceptance-rejection step, is the core of Poisson thinning.

    • calculating the acceptance rate ar, which might exceed \(1.0\).

      In that case, erroneous_acceptance_rate() is called to shrink the horizon by half, and then continues the ac_step() loop in moves_until_horizon() function. Shrinking horizon leads to a finer grid, since n_grid is fixed.

    • Otherwise, it performs the Poisson thinning using the proxy rate lambda_bar, in ac_step_with_proxy() call.

      • Within ac_step_with_proxy(), either if_accept() or if_reject() is called depending on the result of Poisson thinning, followed by move_to_horizon2() if accepted. Whether accepted or not is informed by state.accept flag, which will lead you out of all the loops.
      • In if_accept(), the flags are updated as state.accept = true & state.accept = true, getting out of both ac_step() loop and one_step_of_thinning() loop, with the correct \((x,v,t)\) stored in appropriate state’s fields, which is pushed to PDMPHistory back in sample_skeleton() in sample.jl Section 1.2.
      • if_reject() being called, we are still in the ac_step() loop in moves_until_horizon(), until acceptance or tp > state.horizon.

1.4 UpperBound.jl

This module contains 2 functions, upper_bound_constant() and upper_bound_adaptive().

  • upper_bound_constant() computes the constant upper bound using the Brent’s algorithm.
  • upper_bound_adaptive() computes the adaptive upper bound using the Brent’s algorithm.

1.5 plot.jl

This file contains functions to plot the samples generated by sample.jl.

2 Implementation of the Samplers

2.1 Samplers/AbstractPDMP.jl

This module contains one line

abstract type AbstractPDMP end

and one function

init_state(pdmp::AbstractPDMP, xinit::Array{Float64}, vinit::Array{Float64}, seed::Int) -> PDMPState

init_state() is basically a constructor for PDMPState. This is defined in AbstractPDMP.jl, because the composite PDMPState must have different field values depending on the type of the samplers.

init_state() mainly does two things
  1. Check the field pdmp.signed_bound and initialize the 3 fields rate, rate_vect, refresh_rate accordingly.
  2. Check the field pdmp.grid_size & pdmp.vectorized_bound and initialize the field upper_bound_func as one of the functions in the UpperBound.jl Section 1.4 module accordingly.

2.2 Samplers/StickyZigZagSamplers.jl

This module implements the Sticky Zig-Zag sampler, a variant of the Zig-Zag sampler that allows for sticky behavior at boundaries or other specified conditions.

  • The sampler modifies the velocity component of the state when certain conditions are met.
  • It is particularly useful for models with reflective or sticky boundaries.