gen_rs/inference/
particle_filter.rs

1// mostly copied verbatim from: https://github.com/OpenGen/GenTL/blob/main/include/gentl/inference/particle_filter.h
2
3use rand::rngs::ThreadRng;
4use crate::{Trace,GenFn,GfDiff,Distribution,categorical,mathutils::logsumexp};
5
6
7/// Basic particle filter for generative functions with a time parameter as the first input argument.
8pub struct ParticleSystem<Args: Clone,Data: Clone,Ret: Clone,F: GenFn<(i64,Args),Data,Ret>> {
9    num_particles: usize,
10    model: Box<F>,
11
12    /// Persistent traces contained within the system
13    pub traces: Vec<Trace<(i64,Args),Data,Ret>>,
14
15    log_weights: Vec<f64>,
16    log_normalized_weights: Vec<f64>,
17    two_times_log_normalized_weights: Vec<f64>,
18    normalized_weights: Vec<f64>,
19
20    parents: Vec<usize>,
21    rng: ThreadRng,
22
23    log_ml_estimate: f64
24}
25
26impl<Args: Clone,Data: Clone,Ret: Clone,F: GenFn<(i64,Args),Data,Ret>> ParticleSystem<Args,Data,Ret,F> {
27    fn normalize_weights(&mut self) -> f64 {
28        let log_total_weight = logsumexp(&self.log_weights);
29        for i in 0..self.num_particles {
30            self.log_normalized_weights[i] = self.log_weights[i] - log_total_weight;
31            self.two_times_log_normalized_weights[i] = 2.0 * self.log_normalized_weights[i];
32            self.normalized_weights[i] = self.log_normalized_weights[i].exp();
33        }
34        log_total_weight
35    }
36
37    fn multinomial_resampling(&mut self) {
38        for i in 0..self.num_particles {
39            self.parents[i] = categorical.random(&mut self.rng, self.normalized_weights.clone());
40        }
41    }
42
43    /// Construct a new particle filter under the `model` with `num_particles` particles.
44    pub fn new(model: F, num_particles: usize, rng: ThreadRng) -> Self {
45        ParticleSystem {
46            num_particles,
47            model: Box::new(model),
48            traces: vec![],
49            log_weights: vec![0.; num_particles],
50            log_normalized_weights: vec![0.; num_particles],
51            two_times_log_normalized_weights: vec![0.; num_particles],
52            normalized_weights: vec![0.; num_particles],
53            parents: vec![0; num_particles],
54            rng: rng,
55            log_ml_estimate: 0.
56        }
57    }
58
59    /// Initialize the particle filter by generating `self.num_particles` traces from the `model` with `(1, args)`.
60    pub fn init_step(
61        &mut self,
62        args: Args,
63        constraints: Data
64    ) {
65        for i in 0..self.num_particles {
66            let (trace, log_weight) = self.model.generate((1, args.clone()), constraints.clone());
67            self.traces.push(trace);
68            self.log_weights[i] = log_weight;
69        }
70    }
71
72    /// Extend the current filter from `t` to `t+1` with new `constraints`.
73    pub fn step(self, constraints: Data) -> Self {
74        let mut tmp_traces = vec![];
75        let mut tmp_log_weights = vec![];
76        for (i, trace) in self.traces.into_iter().enumerate() {
77            let args = trace.args.clone();
78            let new_args = (args.0 + 1, args.1);
79            let (new_trace, _, log_weight) = self.model.update(trace, new_args, GfDiff::Extend, constraints.clone());
80            tmp_traces.push(new_trace);
81            tmp_log_weights.push(self.log_weights[i] + log_weight);
82        }
83        ParticleSystem {
84            num_particles: self.num_particles,
85            model: self.model,
86            traces: tmp_traces,
87            log_weights: tmp_log_weights,
88            log_normalized_weights: self.log_normalized_weights,
89            two_times_log_normalized_weights: self.two_times_log_normalized_weights,
90            normalized_weights: self.normalized_weights,
91            parents: self.parents,
92            rng: self.rng,
93            log_ml_estimate: self.log_ml_estimate
94        }
95    }
96
97    /// Calculate the effective sample size (ESS) with the current paticle weights.
98    pub fn effective_sample_size(&self) -> f64 {
99        (-logsumexp(&self.two_times_log_normalized_weights)).exp()
100    }
101
102    /// Perform multinomial resampling based on the normalized particle weights, and return the log total weight.
103    pub fn resample(&mut self) -> f64 {
104        let log_total_weight = self.normalize_weights();
105        self.log_ml_estimate += log_total_weight - (self.num_particles as f64).ln();
106
107        self.multinomial_resampling();
108
109        let mut tmp_traces = vec![];
110        for i in 0..self.num_particles {
111            tmp_traces.push(self.traces[self.parents[i]].clone());
112        }
113        self.traces = tmp_traces;
114        self.log_weights.fill(0.);
115        log_total_weight
116    }
117
118    /// Return the current log marginal likelihood estimate from the particles.
119    pub fn log_marginal_likelihood_estimate(&self) -> f64 {
120        self.log_ml_estimate + logsumexp(&self.log_weights) - (self.num_particles as f64).ln()
121    }
122}
gen_rs/inference/particle_filter.rs

gen_rs/inference/
particle_filter.rs
