legume-numeric 0.8.11

Numeric and ML foundation for the legume ecosystem (matrix, Leiden, candle, 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
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
141
142
143
144
145
146
147
148
149
150
151

use rand::rngs::SmallRng;

use crate::mcmc::engine::traits::EssParamSummary;
use crate::mcmc::engine::{elliptical_slice_step, McmcConfig, McmcModel};

use super::prior::ComponentPrior;
use super::{compute_pip, compute_posterior_mean_beta, SparseResult, SparseSample};

/// Blackbox sparse regression model for the MCMC engine.
///
/// Uses ESS for both inclusion weights and effect sizes. The log-likelihood
/// is a blackbox function of the combined effect theta.
pub(crate) struct SparseModel<P: ComponentPrior, F: Fn(&P::Theta) -> f32> {
    lnpdf: F,
    prior: P,
    p: usize,
    num_components: usize,
}

/// Internal Gibbs state for the blackbox sampler.
pub(crate) struct SparseState<P: ComponentPrior> {
    pub inclusions: Vec<P::Inclusion>,
    pub alphas: Vec<Vec<f32>>,
    pub effects: Vec<P::Effect>,
    pub theta: P::Theta,
}

impl<P: ComponentPrior, F: Fn(&P::Theta) -> f32> McmcModel for SparseModel<P, F> {
    type State = SparseState<P>;
    type Sample = SparseSample;
    type Result = SparseResult;

    fn init(&self, rng: &mut SmallRng) -> SparseState<P> {
        let inclusions: Vec<P::Inclusion> = (0..self.num_components)
            .map(|_| self.prior.draw_inclusion(self.p, rng))
            .collect();
        let alphas: Vec<Vec<f32>> = inclusions
            .iter()
            .map(|inc| self.prior.to_alpha(inc))
            .collect();
        let effects: Vec<P::Effect> = (0..self.num_components)
            .map(|_| self.prior.draw_effect(rng))
            .collect();
        let theta = self.prior.combine(&alphas, &effects);
        SparseState {
            inclusions,
            alphas,
            effects,
            theta,
        }
    }

    fn sweep(&self, state: &mut SparseState<P>, rng: &mut SmallRng) {
        ////////////////////////////////////
        // ESS update of inclusion params //
        ////////////////////////////////////
        for l in 0..self.num_components {
            self.prior.remove_component_inplace(
                &mut state.theta,
                &state.alphas[l],
                &state.effects[l],
            );

            let lnpdf_inc = |inc: &P::Inclusion| -> f32 {
                let alpha_new = self.prior.to_alpha(inc);
                let mut theta_new = state.theta.clone();
                self.prior
                    .add_component_inplace(&mut theta_new, &alpha_new, &state.effects[l]);
                (self.lnpdf)(&theta_new)
            };

            let cur_ll = lnpdf_inc(&state.inclusions[l]);
            let prior_sample = self.prior.draw_inclusion(self.p, rng);
            let (new_inc, _) =
                elliptical_slice_step(&state.inclusions[l], &prior_sample, &lnpdf_inc, cur_ll, rng);

            state.inclusions[l] = new_inc;
            state.alphas[l] = self.prior.to_alpha(&state.inclusions[l]);
            self.prior
                .add_component_inplace(&mut state.theta, &state.alphas[l], &state.effects[l]);
        }

        ////////////////////////////////
        // ESS update of effect sizes //
        ////////////////////////////////
        for l in 0..self.num_components {
            self.prior.remove_component_inplace(
                &mut state.theta,
                &state.alphas[l],
                &state.effects[l],
            );

            let alpha_l = &state.alphas[l];
            let lnpdf_eff = |eff: &P::Effect| -> f32 {
                let mut theta_new = state.theta.clone();
                self.prior
                    .add_component_inplace(&mut theta_new, alpha_l, eff);
                (self.lnpdf)(&theta_new)
            };

            let cur_ll = lnpdf_eff(&state.effects[l]);
            let prior_eff = self.prior.draw_effect(rng);
            let (new_eff, _) =
                elliptical_slice_step(&state.effects[l], &prior_eff, &lnpdf_eff, cur_ll, rng);

            state.effects[l] = new_eff;
            self.prior
                .add_component_inplace(&mut state.theta, &state.alphas[l], &state.effects[l]);
        }
    }

    fn collect(&self, state: &SparseState<P>) -> SparseSample {
        SparseSample {
            alphas: state.alphas.clone(),
            betas: state.effects.iter().map(|e| e.as_slice()[0]).collect(),
            sigma2_eps: f32::NAN,
            effect_var: f32::NAN,
        }
    }

    fn summarize(&self, samples: Vec<SparseSample>) -> SparseResult {
        let pip = compute_pip(&samples, self.p);
        let posterior_mean_beta = compute_posterior_mean_beta(&samples, self.p);
        SparseResult {
            samples,
            pip,
            posterior_mean_beta,
            sigma2_eps_mean: f32::NAN,
            effect_var_mean: f32::NAN,
        }
    }
}

/// Blackbox MCMC sparse regression.
///
/// Convenience wrapper around [`SparseModel`] + [`run_mcmc`].
pub fn mcmc_sparse<P: ComponentPrior>(
    lnpdf: &impl Fn(&P::Theta) -> f32,
    p: usize,
    prior: &P,
    config: &McmcConfig,
    num_components: usize,
) -> SparseResult {
    let model = SparseModel {
        lnpdf,
        prior: prior.clone(),
        p,
        num_components,
    };
    crate::mcmc::engine::run_mcmc(&model, config)
}