nuts-rs 0.18.1

Sample from unnormalized densities using Hamiltonian MCMC
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

use std::collections::HashMap;

use nuts_storable::HasDims;
use thiserror::Error;

use super::{CpuLogpFunc, CpuMathError, LogpError};

#[derive(Clone, Debug)]
pub struct NormalLogp {
    pub dim: usize,
    pub mu: f64,
}

impl NormalLogp {
    pub(crate) fn new(dim: usize, mu: f64) -> NormalLogp {
        NormalLogp { dim, mu }
    }
}

#[derive(Error, Debug)]
pub enum NormalLogpError {}

impl LogpError for NormalLogpError {
    fn is_recoverable(&self) -> bool {
        false
    }
}

impl HasDims for NormalLogp {
    fn dim_sizes(&self) -> HashMap<String, u64> {
        vec![
            ("unconstrained_parameter".to_string(), self.dim as u64),
            ("dim".to_string(), self.dim as u64),
        ]
        .into_iter()
        .collect()
    }
}

impl CpuLogpFunc for NormalLogp {
    type LogpError = NormalLogpError;
    type FlowParameters = ();
    type ExpandedVector = Vec<f64>;

    fn dim(&self) -> usize {
        self.dim
    }

    fn logp(&mut self, position: &[f64], gradient: &mut [f64]) -> Result<f64, NormalLogpError> {
        assert!(gradient.len() == position.len());
        let mut logp = 0f64;
        for (p, g) in position.iter().zip(gradient.iter_mut()) {
            let val = *p - self.mu;
            logp -= val * val / 2.;
            *g = -val;
        }
        Ok(logp)
    }

    fn expand_vector<R: rand::Rng + ?Sized>(
        &mut self,
        _rng: &mut R,
        array: &[f64],
    ) -> Result<Self::ExpandedVector, CpuMathError> {
        Ok(array.to_vec())
    }
}

impl HasDims for &NormalLogp {
    fn dim_sizes(&self) -> HashMap<String, u64> {
        (*self).dim_sizes()
    }
}

impl CpuLogpFunc for &NormalLogp {
    type LogpError = NormalLogpError;
    type FlowParameters = ();
    type ExpandedVector = Vec<f64>;

    fn dim(&self) -> usize {
        self.dim
    }

    fn logp(&mut self, position: &[f64], gradient: &mut [f64]) -> Result<f64, NormalLogpError> {
        assert!(gradient.len() == position.len());
        let mut logp = 0f64;
        for (p, g) in position.iter().zip(gradient.iter_mut()) {
            let val = *p - self.mu;
            logp -= val * val / 2.;
            *g = -val;
        }
        Ok(logp)
    }

    fn expand_vector<R: rand::Rng + ?Sized>(
        &mut self,
        _rng: &mut R,
        array: &[f64],
    ) -> Result<Self::ExpandedVector, CpuMathError> {
        Ok(array.to_vec())
    }
}