Skip to main content

gstools_core/
lib.rs

1//! GSTools-Core
2//!
3//! `gstools_core` is a Rust implementation of the core algorithms of [GSTools].
4//! At the moment, it is a drop in replacement for the Cython code included in GSTools.
5//!
6//! This crate includes
7//! - [randomization methods](field) for the random field generation
8//! - the [matrix operations](krige) of the kriging methods
9//! - the [variogram estimation](variogram)
10//!
11//! [GSTools]: https://github.com/GeoStat-Framework/GSTools
12
13#![warn(missing_docs)]
14
15use pyo3::prelude::pymodule;
16
17pub mod field;
18pub mod krige;
19mod short_vec;
20pub mod variogram;
21
22#[pymodule]
23mod gstools_core {
24    use crate::field::{summator, summator_fourier, summator_incompr};
25    use crate::krige::{calculator_field_krige, calculator_field_krige_and_variance};
26    use crate::variogram::{
27        variogram_directional, variogram_ma_structured, variogram_structured,
28        variogram_unstructured,
29    };
30    use numpy::{IntoPyArray, PyArray1, PyArray2, PyReadonlyArray1, PyReadonlyArray2};
31    use pyo3::prelude::*;
32
33    #[pymodule_init]
34    fn init(m: &Bound<'_, PyModule>) -> PyResult<()> {
35        m.add("__version__", env!("CARGO_PKG_VERSION"))?;
36        Ok(())
37    }
38
39    #[pyfunction(name = "summate")]
40    fn summate_py<'py>(
41        py: Python<'py>,
42        cov_samples: PyReadonlyArray2<f64>,
43        z1: PyReadonlyArray1<f64>,
44        z2: PyReadonlyArray1<f64>,
45        pos: PyReadonlyArray2<f64>,
46        num_threads: Option<usize>,
47    ) -> Bound<'py, PyArray1<f64>> {
48        let cov_samples = cov_samples.as_array();
49        let z1 = z1.as_array();
50        let z2 = z2.as_array();
51        let pos = pos.as_array();
52        summator(cov_samples, z1, z2, pos, num_threads).into_pyarray(py)
53    }
54
55    #[pyfunction(name = "summate_incompr")]
56    fn summate_incompr_py<'py>(
57        py: Python<'py>,
58        cov_samples: PyReadonlyArray2<f64>,
59        z1: PyReadonlyArray1<f64>,
60        z2: PyReadonlyArray1<f64>,
61        pos: PyReadonlyArray2<f64>,
62        num_threads: Option<usize>,
63    ) -> Bound<'py, PyArray2<f64>> {
64        let cov_samples = cov_samples.as_array();
65        let z1 = z1.as_array();
66        let z2 = z2.as_array();
67        let pos = pos.as_array();
68        summator_incompr(cov_samples, z1, z2, pos, num_threads).into_pyarray(py)
69    }
70
71    #[pyfunction(name = "summate_fourier")]
72    fn summate_fourier_py<'py>(
73        py: Python<'py>,
74        spectrum_factor: PyReadonlyArray1<f64>,
75        modes: PyReadonlyArray2<f64>,
76        z1: PyReadonlyArray1<f64>,
77        z2: PyReadonlyArray1<f64>,
78        pos: PyReadonlyArray2<f64>,
79        num_threads: Option<usize>,
80    ) -> Bound<'py, PyArray1<f64>> {
81        let spectrum_factor = spectrum_factor.as_array();
82        let modes = modes.as_array();
83        let z1 = z1.as_array();
84        let z2 = z2.as_array();
85        let pos = pos.as_array();
86        summator_fourier(spectrum_factor, modes, z1, z2, pos, num_threads).into_pyarray(py)
87    }
88
89    #[pyfunction(name = "calc_field_krige_and_variance")]
90    fn calc_field_krige_and_variance_py<'py>(
91        py: Python<'py>,
92        krige_mat: PyReadonlyArray2<f64>,
93        krig_vecs: PyReadonlyArray2<f64>,
94        cond: PyReadonlyArray1<f64>,
95        num_threads: Option<usize>,
96    ) -> (Bound<'py, PyArray1<f64>>, Bound<'py, PyArray1<f64>>) {
97        let krige_mat = krige_mat.as_array();
98        let krig_vecs = krig_vecs.as_array();
99        let cond = cond.as_array();
100        let (field, error) =
101            calculator_field_krige_and_variance(krige_mat, krig_vecs, cond, num_threads);
102        let field = field.into_pyarray(py);
103        let error = error.into_pyarray(py);
104        (field, error)
105    }
106
107    #[pyfunction(name = "calc_field_krige")]
108    fn calc_field_krige_py<'py>(
109        py: Python<'py>,
110        krige_mat: PyReadonlyArray2<f64>,
111        krig_vecs: PyReadonlyArray2<f64>,
112        cond: PyReadonlyArray1<f64>,
113        num_threads: Option<usize>,
114    ) -> Bound<'py, PyArray1<f64>> {
115        let krige_mat = krige_mat.as_array();
116        let krig_vecs = krig_vecs.as_array();
117        let cond = cond.as_array();
118        calculator_field_krige(krige_mat, krig_vecs, cond, num_threads).into_pyarray(py)
119    }
120
121    #[pyfunction(name = "variogram_structured")]
122    fn variogram_structured_py<'py>(
123        py: Python<'py>,
124        f: PyReadonlyArray2<f64>,
125        estimator_type: Option<char>,
126        num_threads: Option<usize>,
127    ) -> Bound<'py, PyArray1<f64>> {
128        let f = f.as_array();
129        let estimator_type = estimator_type.unwrap_or('m');
130        variogram_structured(f, estimator_type, num_threads).into_pyarray(py)
131    }
132
133    #[pyfunction(name = "variogram_ma_structured")]
134    fn variogram_ma_structured_py<'py>(
135        py: Python<'py>,
136        f: PyReadonlyArray2<f64>,
137        mask: PyReadonlyArray2<bool>,
138        estimator_type: Option<char>,
139        num_threads: Option<usize>,
140    ) -> Bound<'py, PyArray1<f64>> {
141        let f = f.as_array();
142        let mask = mask.as_array();
143        let estimator_type = estimator_type.unwrap_or('m');
144        variogram_ma_structured(f, mask, estimator_type, num_threads).into_pyarray(py)
145    }
146
147    #[pyfunction(name = "variogram_directional")]
148    #[allow(clippy::too_many_arguments)]
149    fn variogram_directional_py<'py>(
150        py: Python<'py>,
151        f: PyReadonlyArray2<f64>,
152        bin_edges: PyReadonlyArray1<f64>,
153        pos: PyReadonlyArray2<f64>,
154        direction: PyReadonlyArray2<f64>, //should be normed
155        angles_tol: Option<f64>,
156        bandwidth: Option<f64>,
157        separate_dirs: Option<bool>,
158        estimator_type: Option<char>,
159        num_threads: Option<usize>,
160    ) -> (Bound<'py, PyArray2<f64>>, Bound<'py, PyArray2<u64>>) {
161        let f = f.as_array();
162        let bin_edges = bin_edges.as_array();
163        let pos = pos.as_array();
164        let direction = direction.as_array();
165        let angles_tol = angles_tol.unwrap_or(std::f64::consts::PI / 8.0);
166        let bandwidth = bandwidth.unwrap_or(-1.0);
167        let separate_dirs = separate_dirs.unwrap_or(false);
168        let estimator_type = estimator_type.unwrap_or('m');
169        let (variogram, counts) = variogram_directional(
170            f,
171            bin_edges,
172            pos,
173            direction,
174            angles_tol,
175            bandwidth,
176            separate_dirs,
177            estimator_type,
178            num_threads,
179        );
180        let variogram = variogram.into_pyarray(py);
181        let counts = counts.into_pyarray(py);
182
183        (variogram, counts)
184    }
185
186    #[pyfunction(name = "variogram_unstructured")]
187    fn variogram_unstructured_py<'py>(
188        py: Python<'py>,
189        f: PyReadonlyArray2<f64>,
190        bin_edges: PyReadonlyArray1<f64>,
191        pos: PyReadonlyArray2<f64>,
192        estimator_type: Option<char>,
193        distance_type: Option<char>,
194        num_threads: Option<usize>,
195    ) -> (Bound<'py, PyArray1<f64>>, Bound<'py, PyArray1<u64>>) {
196        let f = f.as_array();
197        let bin_edges = bin_edges.as_array();
198        let pos = pos.as_array();
199        let estimator_type = estimator_type.unwrap_or('m');
200        let distance_type = distance_type.unwrap_or('e');
201        let (variogram, counts) = variogram_unstructured(
202            f,
203            bin_edges,
204            pos,
205            estimator_type,
206            distance_type,
207            num_threads,
208        );
209        let variogram = variogram.into_pyarray(py);
210        let counts = counts.into_pyarray(py);
211
212        (variogram, counts)
213    }
214}