legume-numeric 0.9.3

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

#[cfg(feature = "tensor")]
use candle_core::Tensor;
use nalgebra::{DMatrix, DVector};

/// Minimal trait for types usable as ESS parameters.
/// Only requires linear combination — prior sampling is handled
/// externally via the `prior_draw` closure in `EssSampler::run`.
pub trait EssParam: Clone {
    /// Elliptical combination: `a * self + b * other`
    fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self;
}

/// Optional trait enabling summary statistics (mean, variance, quantile).
/// Requires flattening to a slice.
pub trait EssParamSummary: EssParam {
    fn as_slice(&self) -> &[f32];
    fn dim(&self) -> usize;
}

///////////////////////////
// nalgebra DVector<f32> //
///////////////////////////

impl EssParam for DVector<f32> {
    fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self {
        self * a + other * b
    }
}

impl EssParamSummary for DVector<f32> {
    fn as_slice(&self) -> &[f32] {
        self.as_slice()
    }

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

///////////////////////////
// nalgebra DMatrix<f32> //
///////////////////////////

impl EssParam for DMatrix<f32> {
    fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self {
        self * a + other * b
    }
}

impl EssParamSummary for DMatrix<f32> {
    fn as_slice(&self) -> &[f32] {
        self.as_slice()
    }

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

////////////////
// scalar f32 //
////////////////

/// A single scalar parameter, without the heap allocation per proposal that a
/// 1-element `DVector` would cost in the sampler's inner loop.
impl EssParam for f32 {
    fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self {
        a * self + b * other
    }
}

impl EssParamSummary for f32 {
    fn as_slice(&self) -> &[f32] {
        std::slice::from_ref(self)
    }

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

/////////////////////////////////////
// Vec<P> for composite parameters //
/////////////////////////////////////

impl<P: EssParam> EssParam for Vec<P> {
    fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self {
        self.iter()
            .zip(other.iter())
            .map(|(s, o)| s.linear_combine(a, o, b))
            .collect()
    }
}

///////////////////
// candle Tensor //
///////////////////

#[cfg(feature = "tensor")]
impl EssParam for Tensor {
    fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self {
        (self * a as f64)
            .unwrap()
            .add(&(other * b as f64).unwrap())
            .unwrap()
    }
}