modppl 0.3.0

a experimental library for probabilistic programming in Rust.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

use crate::{DynGenFn, DynGenFnHandler, DynTrie, GenFn, ArgDiff, Trace, ParticleSystem};
use rand::rngs::ThreadRng;


/// Combinator struct for kernels that use the `DynGenFnHandler` DSL (`sample_at` and `trace_at`).
/// Supports memory-efficient extension via the `GfDiff::Extend` flag (eg. as passed during a `ParticleSystem::step`).
pub struct DynUnfold<State> {
    /// A stochastic kernel that takes in a mutable reference to a `DynGenFnHandler<A,T>` and some `State`, effectfully mutates it, and produces a new `State`.
    pub kernel: DynGenFn<(i64,State),State>
}

impl<State> DynUnfold<State> {
    /// Dynamically construct a `DynUnfold` from a kernel function at run-time.
    pub fn new(kernel: DynGenFn<(i64,State),State>) -> Self {
        DynUnfold { kernel }
    }
}

///
pub type DynParticles<State> = ParticleSystem<State,Vec<DynTrie>,Vec<State>,DynUnfold<State>>;

impl<State: Clone> GenFn<(i64,State),Vec<DynTrie>,Vec<State>> for DynUnfold<State> {
    fn simulate(&self, final_t_and_args: (i64, State)) -> Trace<(i64,State),Vec<DynTrie>,Vec<State>> {
        let (final_t, mut state) = final_t_and_args;
        assert!(final_t >= 1);
        let mut vec_trace = Trace { args: (final_t, state.clone()), data: vec![], retv: Some(vec![]), logjp: 0. };
        for t in 0..final_t {
            let mut g = DynGenFnHandler::Simulate {
                prng: &mut ThreadRng::default(),
                trace: Trace { args: (t as i64, state.clone()), data: DynTrie::new(), retv: None, logjp: 0. },
            };
            state = (self.kernel.func)(&mut g, (t as i64, state.clone()));
            let DynGenFnHandler::Simulate {prng: _, trace} = g else { unreachable!() };
            vec_trace.retv.as_mut().unwrap().push(state.clone());
            vec_trace.data.push(trace.data);
            vec_trace.logjp += trace.logjp;
        }
        vec_trace
    }

    fn generate(&self, final_t_and_args: (i64, State), vec_constraints: Vec<DynTrie>) 
        -> (Trace<(i64,State),Vec<DynTrie>,Vec<State>>, f64)
    {
        let (final_t, mut state) = final_t_and_args;
        assert!(final_t >= 1);
        let mut vec_trace = Trace { args: (final_t, state.clone()), data: vec![], retv: Some(vec![]), logjp: 0. };
        let mut gen_weight = 0.;
        for (t,constraints) in vec_constraints.into_iter().enumerate() {
            let mut g = DynGenFnHandler::Generate {
                prng: &mut ThreadRng::default(),
                trace: Trace { args: (t as i64, state.clone()), data: DynTrie::new(), retv: None, logjp: 0. },
                weight: 0.,
                constraints
            };
            state = (self.kernel.func)(&mut g, (t as i64, state.clone()));
            let DynGenFnHandler::Generate {prng: _, trace, weight, constraints} = g else { unreachable!() };
            assert!(constraints.is_empty());
            vec_trace.retv.as_mut().unwrap().push(state.clone());
            vec_trace.data.push(trace.data);
            vec_trace.logjp += trace.logjp;
            gen_weight += weight;
        }
        (vec_trace, gen_weight)
    }

    fn update(&self,
        mut vec_trace: Trace<(i64,State),Vec<DynTrie>,Vec<State>>,
        final_t_and_args: (i64, State),
        diff: ArgDiff,
        vec_constraints: Vec<DynTrie>
    ) -> (Trace<(i64,State),Vec<DynTrie>,Vec<State>>, Vec<DynTrie>, f64) {
        let (final_t, _) = final_t_and_args;
        assert!(final_t >= 1);
        let prev_t = vec_trace.args.0;
        assert!(final_t - prev_t == vec_constraints.len() as i64);
        let mut state = vec_trace.retv.as_ref().unwrap().last().unwrap().clone();
        let mut update_weight = 0.;
        match diff {
            ArgDiff::Extend => {
                for (t,constraints) in vec_constraints.into_iter().enumerate() {
                    let mut g = DynGenFnHandler::Generate {
                        prng: &mut ThreadRng::default(),
                        trace: Trace { args: (prev_t + (t as i64), state.clone()), data: DynTrie::new(), retv: None, logjp: 0. },
                        weight: 0.,
                        constraints
                    };
                    state = (self.kernel.func)(&mut g, (prev_t + (t as i64), state.clone()));
                    let DynGenFnHandler::Generate {prng: _, trace, weight, constraints} = g else { unreachable!() };
                    assert!(constraints.is_empty());
                    vec_trace.args.0 += 1;
                    vec_trace.retv.as_mut().unwrap().push(state.clone());
                    vec_trace.data.push(trace.data);
                    vec_trace.logjp += trace.logjp;
                    update_weight += weight;
                }
            },
            _ => { panic!("Can't handle GF change type: {:?}", diff) },
        }
        (vec_trace, (prev_t..final_t).map(|_| DynTrie::new()).collect::<_>(), update_weight)
    }
}