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

use crate::{Distribution, DistributionError, RandomVariable, SampleableDistribution};
use rand::prelude::*;
use rayon::prelude::*;
use std::f64::consts::PI;

/// Sample `b` from posterior p(b|a) with likelihood p(a|b) and prior p(b)
/// `b` must be generated by elliptical distribution
pub struct EllipticalSliceSampler<'a, L, P, A, B>
where
    L: Distribution<Value = A, Condition = B>,
    P: SampleableDistribution<Value = B, Condition = ()>,
    A: RandomVariable,
    B: RandomVariable,
{
    value: &'a A,
    likelihood: &'a L,
    prior: &'a P,
}

impl<'a, L, P, A, B> EllipticalSliceSampler<'a, L, P, A, B>
where
    L: Distribution<Value = A, Condition = B>,
    P: SampleableDistribution<Value = B, Condition = ()>,
    A: RandomVariable,
    B: RandomVariable,
{
    pub fn new(value: &'a A, likelihood: &'a L, prior: &'a P) -> Self {
        Self {
            value,
            likelihood,
            prior,
        }
    }

    fn step(mut v: Vec<f64>, theta: f64, nu: &Vec<f64>) -> Vec<f64> {
        let cos = theta.cos();
        let sin = theta.sin();
        v.par_iter_mut()
            .zip(nu.par_iter())
            .for_each(|(bi, &nui)| *bi = *bi * cos + nui * sin);

        v
    }

    pub fn sample(&self, rng: &mut dyn RngCore) -> Result<B, DistributionError> {
        let nu = self.prior.sample(&(), rng)?;

        let mut b = self.prior.sample(&(), rng)?;

        let rho = self.likelihood.p_kernel(self.value, &b)? * rng.gen_range(0.0..1.0);
        let mut theta = rng.gen_range(0.0..2.0 * PI);

        let mut start = theta - 2.0 * PI;
        let mut end = theta;

        let nu = nu.transform_vec();

        loop {
            let mut buf = b.transform_vec();
            buf.0 = Self::step(buf.0, theta, &nu.0);

            b = B::restore(&buf.0, &buf.1)?;
            if rho < self.likelihood.p_kernel(self.value, &b)? {
                break;
            }

            if 0.0 < theta {
                end = 0.0;
            } else {
                start = 0.0;
            }
            theta = rng.gen_range(start..end);
        }

        Ok(b)
    }
}

#[cfg(test)]
mod tests {
    use crate::distribution::Distribution;
    use crate::*;
    use core::f64::consts::PI;
    use rand::prelude::*;

    #[test]
    fn it_works() {
        let distr = InstantDistribution::new(
            |x: &f64, theta: &f64| {
                let mu = *theta;
                Ok(1.0 / (2.0 * PI * 10.0f64.powi(2)).sqrt()
                    * (-(x - mu).powi(2) / (2.0 * 10.0f64.powi(2))).exp())
            },
            |theta: &f64, rng| Normal.sample(&NormalParams::new(*theta, 10.0).unwrap(), rng),
        );
        let distr_2 = InstantDistribution::new(
            |x: &f64, _theta: &()| {
                let mu: f64 = 10.0;
                Ok(1.0 / (2.0 * PI * 10.0f64.powi(2)).sqrt()
                    * (-(x - mu).powi(2) / (2.0 * 10.0f64.powi(2))).exp())
            },
            |_theta, rng| {
                Normal.sample(
                    &NormalParams::new(10.0, (10.0f64 * 0.1).abs()).unwrap(),
                    rng,
                )
            },
        );

        let value = 1.0;
        let likelihood = distr;
        let prior = distr_2;

        let start = EllipticalSliceSampler::new(&value, &likelihood, &prior);
        let mut rng = StdRng::from_seed([1; 32]);
        let x = start.sample(&mut rng).unwrap();

        println!("{}", x);
    }
}