1use ndarray::{Array1, ArrayView1, ArrayView2, Zip};
4use rayon::prelude::*;
5
6pub 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
76pub 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}