1#![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>, 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}