modppl 0.3.1

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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

use super::types_2d::{uniform_2d, Bounds, Point};
use modppl::{mvnormal, ArgDiff, Distribution, GenFn, Real, Trace};
use nalgebra::DMatrix;
use rand::rngs::ThreadRng;

pub struct PointedModel {
    pub obs_cov: DMatrix<Real>,
}

pub type PointedBuffer = (Option<Point>, Option<Point>);
pub type PointedTrace = Trace<Bounds, PointedBuffer, Point>;

impl GenFn<Bounds, PointedBuffer, Point> for PointedModel {
    fn simulate(&self, bounds: Bounds) -> PointedTrace {
        let mut rng = ThreadRng::default();
        let mut logjp = 0.;
        let latent = uniform_2d.random(&mut rng, bounds);
        logjp += uniform_2d.logpdf(&latent, bounds);
        let obs = mvnormal.random(&mut rng, (latent.clone(), self.obs_cov.clone()));
        logjp += mvnormal.logpdf(&obs, (obs.clone(), self.obs_cov.clone()));
        PointedTrace::new(bounds, (Some(latent), Some(obs.clone())), obs, logjp)
    }

    fn generate(&self, bounds: Bounds, constraints: PointedBuffer) -> (PointedTrace, Real) {
        let mut rng = ThreadRng::default();
        let mut logjp = 0.;
        let mut weight = 0.;
        let mut choices = (None, None);

        // manual latent branch
        let latent_choice = match constraints.0 {
            Some(constrained_latent) => {
                let new_weight = uniform_2d.logpdf(&constrained_latent, bounds);
                weight += new_weight;
                logjp += new_weight;
                constrained_latent
            }
            None => {
                let latent_choice = uniform_2d.random(&mut rng, bounds);
                let new_weight = uniform_2d.logpdf(&latent_choice, bounds);
                logjp += new_weight;
                latent_choice
            }
        };
        choices.0 = Some(latent_choice.clone());

        // manual obs branch
        let obs_choice = match constraints.1 {
            Some(constrained_obs) => {
                let new_weight =
                    mvnormal.logpdf(&constrained_obs, (latent_choice, self.obs_cov.clone()));
                weight += new_weight;
                logjp += new_weight;
                constrained_obs
            }
            None => {
                let obs_choice =
                    mvnormal.random(&mut rng, (latent_choice.clone(), self.obs_cov.clone()));
                let new_weight =
                    mvnormal.logpdf(&obs_choice, (latent_choice, self.obs_cov.clone()));
                logjp += new_weight;
                obs_choice
            }
        };
        choices.1 = Some(obs_choice.clone());

        (
            PointedTrace::new(bounds, choices, obs_choice, logjp),
            weight,
        )
    }

    fn update(
        &self,
        trace: PointedTrace,
        args: Bounds,
        diff: ArgDiff,
        constraints: PointedBuffer,
    ) -> (PointedTrace, PointedBuffer, Real) {
        match diff {
            ArgDiff::NoChange => {
                let prev_choices = trace.data;
                let bounds = trace.args;
                let mut discard = (None, None);

                let mut new_logjp = trace.logjp;
                let mut visited_obs = false;

                let mut latent_choice = prev_choices.0.clone();
                if let Some(latent_constraint) = constraints.0 {
                    discard.0 = latent_choice;
                    latent_choice = Some(latent_constraint.clone());
                    new_logjp -= uniform_2d.logpdf(&prev_choices.0.clone().unwrap(), bounds);
                    new_logjp += uniform_2d.logpdf(&latent_constraint, bounds);

                    visited_obs = true;
                    new_logjp -= mvnormal.logpdf(
                        &prev_choices.1.clone().unwrap(),
                        (prev_choices.0.clone().unwrap(), self.obs_cov.clone()),
                    );
                }

                let mut obs_choice = prev_choices.1.clone();
                if let Some(obs_constraint) = constraints.1 {
                    discard.1 = obs_choice;
                    obs_choice = Some(obs_constraint);
                    if !visited_obs {
                        new_logjp -= mvnormal.logpdf(
                            &prev_choices.1.unwrap(),
                            (prev_choices.0.clone().unwrap(), self.obs_cov.clone()),
                        );
                    }
                    new_logjp += mvnormal.logpdf(
                        &obs_choice.clone().unwrap(),
                        (latent_choice.clone().unwrap(), self.obs_cov.clone()),
                    );
                } else if visited_obs {
                    new_logjp += mvnormal.logpdf(
                        &obs_choice.clone().unwrap(),
                        (latent_choice.clone().unwrap(), self.obs_cov.clone()),
                    );
                }

                (
                    PointedTrace::new(
                        args,
                        (latent_choice, obs_choice.clone()),
                        obs_choice.unwrap(),
                        new_logjp,
                    ),
                    discard,
                    new_logjp - trace.logjp,
                )
            }
            _ => {
                panic!("Can't handle GF change type: {:?}", diff)
            }
        }
    }
}