gstools_core/
krige.rs

1//! Perform the kriging matrix operations.
2
3use ndarray::{Array1, ArrayView1, ArrayView2, Zip};
4use rayon::prelude::*;
5
6/// Calculate the interpolated field and also return the variance.
7///
8/// # Arguments
9///
10/// * `krige_mat` - the kriging matrix
11///   <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = (no. of kriging conditions `cond.dim()`, returned field `field.dim().0`)
12/// * `krige_vecs` - the right hand side of the kriging system
13///   <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = (`krig_mat.dim().0`, returned field `field.dim().0`)
14/// * `cond` - the kriging conditions
15///   <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = `krig_mat.dim().0`
16/// * `num_threads` - the number of parallel threads used, if None, use rayon's default
17///
18/// # Returns
19///
20/// * `field` - the kriging field
21///   <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = `krig_mat.dim().1`
22/// * `error` - the error variance
23///   <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = `krig_mat.dim().1`
24pub fn calculator_field_krige_and_variance(
25    krig_mat: ArrayView2<'_, f64>,
26    krig_vecs: ArrayView2<'_, f64>,
27    cond: ArrayView1<'_, f64>,
28    num_threads: Option<usize>,
29) -> (Array1<f64>, Array1<f64>) {
30    assert_eq!(krig_mat.dim().0, krig_mat.dim().1);
31    assert_eq!(krig_mat.dim().0, krig_vecs.dim().0);
32    assert_eq!(krig_mat.dim().0, cond.dim());
33
34    let mut field = Array1::<f64>::zeros(krig_vecs.shape()[1]);
35    let mut error = Array1::<f64>::zeros(krig_vecs.shape()[1]);
36
37    rayon::ThreadPoolBuilder::new()
38        .num_threads(num_threads.unwrap_or(rayon::current_num_threads()))
39        .build()
40        .unwrap()
41        .install(|| {
42            Zip::from(field.view_mut())
43                .and(error.view_mut())
44                .and(krig_vecs.columns())
45                .par_for_each(|f, e, v_col| {
46                    let acc = Zip::from(cond)
47                        .and(v_col)
48                        .and(krig_mat.columns())
49                        .into_par_iter()
50                        .fold(
51                            || (0.0, 0.0),
52                            |mut acc, (c, v, m_row)| {
53                                let krig_fac = m_row.dot(&v_col);
54                                acc.0 += c * krig_fac;
55                                acc.1 += v * krig_fac;
56                                acc
57                            },
58                        )
59                        .reduce(
60                            || (0.0, 0.0),
61                            |mut lhs, rhs| {
62                                lhs.0 += rhs.0;
63                                lhs.1 += rhs.1;
64                                lhs
65                            },
66                        );
67
68                    *f = acc.0;
69                    *e = acc.1;
70                })
71        });
72
73    (field, error)
74}
75
76/// Calculate the interpolated field.
77///
78/// # Arguments
79///
80/// * `krige_mat` - the kriging matrix
81///   <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = (no. of kriging conditions `cond.dim()`, returned field `field.dim().0`)
82/// * `krige_vecs` - the right hand side of the kriging system
83///   <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = (`krig_mat.dim().0`, returned field `field.dim().0`)
84/// * `cond` - the kriging conditions
85///   <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = `krig_mat.dim().0`
86/// * `num_threads` - the number of parallel threads used, if None, use rayon's default
87///
88/// # Returns
89///
90/// * `field` - the kriging field
91///   <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = `krig_mat.dim().1`
92pub fn calculator_field_krige(
93    krig_mat: ArrayView2<'_, f64>,
94    krig_vecs: ArrayView2<'_, f64>,
95    cond: ArrayView1<'_, f64>,
96    num_threads: Option<usize>,
97) -> Array1<f64> {
98    assert_eq!(krig_mat.dim().0, krig_mat.dim().1);
99    assert_eq!(krig_mat.dim().0, krig_vecs.dim().0);
100    assert_eq!(krig_mat.dim().0, cond.dim());
101
102    rayon::ThreadPoolBuilder::new()
103        .num_threads(num_threads.unwrap_or(rayon::current_num_threads()))
104        .build()
105        .unwrap()
106        .install(|| {
107            Zip::from(krig_vecs.columns()).par_map_collect(|v_col| {
108                Zip::from(cond)
109                    .and(krig_mat.columns())
110                    .into_par_iter()
111                    .map(|(c, m_row)| {
112                        let krig_fac = m_row.dot(&v_col);
113                        c * krig_fac
114                    })
115                    .sum()
116            })
117        })
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    use approx::assert_ulps_eq;
125    use ndarray::{arr1, arr2, Array2};
126
127    struct Setup {
128        krig_mat: Array2<f64>,
129        krig_vecs: Array2<f64>,
130        cond: Array1<f64>,
131    }
132
133    impl Setup {
134        fn new() -> Self {
135            Self {
136                krig_mat: arr2(&[
137                    [
138                        5.00000000068981e-01,
139                        -5.87287095364834e-06,
140                        7.82325812566282e-12,
141                    ],
142                    [
143                        -5.87287095378827e-06,
144                        5.00000000070158e-01,
145                        -7.67370103394336e-07,
146                    ],
147                    [
148                        7.82331319334681e-12,
149                        -7.67370103410243e-07,
150                        5.00000000001178e-01,
151                    ],
152                ]),
153                krig_vecs: arr2(&[
154                    [
155                        3.00650970845165e-01,
156                        7.92958674144233e-11,
157                        7.34102993092809e-02,
158                        1.10371060304999e-08,
159                        2.00114256042442e-01,
160                        7.23018134159345e-03,
161                    ],
162                    [
163                        5.51416575736629e-09,
164                        4.79656668238205e-09,
165                        3.91247964853073e-03,
166                        3.59846942149471e-11,
167                        2.10720573114332e-10,
168                        4.83625846265317e-04,
169                    ],
170                    [
171                        7.08796598544206e-13,
172                        1.09700007286403e-01,
173                        2.46322359027701e-05,
174                        1.75889992745405e-07,
175                        3.05671083940413e-17,
176                        2.38513785599550e-11,
177                    ],
178                ]),
179                cond: arr1(&[
180                    -1.27755407195723e+00,
181                    1.15554040655641e+00,
182                    8.47374235895458e-01,
183                ]),
184            }
185        }
186    }
187
188    #[test]
189    fn test_calculator_field_krige_and_variance() {
190        let setup = Setup::new();
191
192        let (kr_field, kr_error) = calculator_field_krige_and_variance(
193            setup.krig_mat.view(),
194            setup.krig_vecs.view(),
195            setup.cond.view(),
196            None,
197        );
198        assert_ulps_eq!(
199            kr_field,
200            arr1(&[
201                -0.19205097317842723,
202                0.04647838537175125,
203                -0.04462233428403452,
204                0.0000000674926344864219,
205                -0.12782974926973434,
206                -0.0043390949462510245
207            ]),
208            max_ulps = 6,
209        );
210        assert_ulps_eq!(
211            kr_error,
212            arr1(&[
213                0.04519550314128594,
214                0.006017045799331816,
215                0.0027021867008690937,
216                0.000000000000015529554261898964,
217                0.020022857738471924,
218                0.00002625466702800745
219            ]),
220            max_ulps = 6,
221        );
222    }
223
224    #[test]
225    fn test_calculator_field_krige() {
226        let setup = Setup::new();
227
228        assert_ulps_eq!(
229            calculator_field_krige(
230                setup.krig_mat.view(),
231                setup.krig_vecs.view(),
232                setup.cond.view(),
233                None
234            ),
235            arr1(&[
236                -0.19205097317842723,
237                0.04647838537175125,
238                -0.04462233428403452,
239                0.0000000674926344864219,
240                -0.12782974926973434,
241                -0.0043390949462510245
242            ]),
243            max_ulps = 6,
244        );
245    }
246}