Add Fourier method

Author: LSchueler

src/field.rs

@@ -181,6 +181,73 @@ pub fn summator_incompr(
     }
 }
 
+/// The Fourier method for scalar fields.
+///
+/// Computes the periodic, isotropic spatial random field $u(x)$ by the Fourier
+/// method according to
+/// $$
+/// u(x) = \sum_{i=1}^N \sqrt{2S(k_i\Delta k}\left(
+/// z_{1,i} \cdot \cos(\langle k_i, x \rangle) +
+/// z_{2,i} \cdot \sin(\langle k_i, x \rangle)\right)
+/// $$
+/// with
+/// * $S$ being the spectrum of the covariance model
+/// * $N$ being the number of Fourier modes
+/// * $z_1, z_2$ being independent samples from a standard normal distribution
+/// * $k$ being the equidistant Fourier grid
+/// and are given by the argument `modes`.
+/// * $\Delta k$ being the cell size of the Fourier grid
+///
+/// # Arguments
+///
+/// * `spectrum_factor` - the pre-calculated factor $\sqrt{2S(k_i\Delta k)}
+/// <br>&nbsp;&nbsp;&nbsp;&nbsp; dim = Fourier modes $N$
+/// * `modes` - equidistant Fourier grid, $k$ in Eq. above
+/// <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = (spatial dim. of field $d$, Fourier modes $N$)
+/// * `z1` - independent samples from a standard normal distribution
+/// <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = Fourier modes $N$
+/// * `z2` - independent samples from a standard normal distribution
+/// <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = Fourier modes $N$
+/// * `pos` - the position $x$ where the spatial random field is calculated
+/// <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = (spatial dim. of field $d$, no. of spatial points where field is calculated $j$)
+/// * `num_threads` - the number of parallel threads used, if None, use rayon's default
+///
+/// # Returns
+///
+/// * `summed_modes` - the isotropic spatial field
+/// <br>&nbsp;&nbsp;&nbsp;&nbsp;dim = no. of spatial points where field is calculated $j$
+pub fn summator_fourier(
+    spectrum_factor: ArrayView1<'_, f64>,
+    modes: ArrayView2<'_, f64>,
+    z1: ArrayView1<'_, f64>,
+    z2: ArrayView1<'_, f64>,
+    pos: ArrayView2<'_, f64>,
+    num_threads: Option<usize>,
+) -> Array1<f64> {
+    assert_eq!(modes.dim().0, pos.dim().0);
+    assert_eq!(modes.dim().1, z1.dim());
+    assert_eq!(modes.dim().1, z2.dim());
+
+    rayon::ThreadPoolBuilder::new()
+        .num_threads(num_threads.unwrap_or(rayon::current_num_threads()))
+        .build()
+        .unwrap()
+        .install(|| {
+            Zip::from(pos.columns()).par_map_collect(|pos| {
+                Zip::from(modes.columns())
+                    .and(spectrum_factor)
+                    .and(z1)
+                    .and(z2)
+                    .fold(0.0, |sum, k, &spectrum_factor, &z1, &z2| {
+                        let phase = k.dot(&pos);
+                        let z12 = spectrum_factor * (z1 * phase.cos() + z2 * phase.sin());
+
+                        sum + z12
+                    })
+            })
+        })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -224,8 +291,73 @@ mod tests {
         }
     }
 
+    struct SetupFourier {
+        spectrum_factor: Array1<f64>,
+        modes: Array2<f64>,
+        z_1: Array1<f64>,
+        z_2: Array1<f64>,
+        pos: Array2<f64>,
+    }
+
+    impl SetupFourier {
+        fn new() -> Self {
+            Self {
+                spectrum_factor: arr1(&[
+                    -2.15, 1.04, 0.69, -1.09, -1.54, -2.32, -1.81, -2.78, 1.57, -3.44,
+                ]),
+                modes: arr2(&[
+                    [
+                        -2.15, 1.04, 0.69, -1.09, -1.54, -2.32, -1.81, -2.78, 1.57, -3.44,
+                    ],
+                    [
+                        0.19, -1.24, -2.10, -2.86, -0.63, -0.51, -1.68, -0.07, 0.29, -0.007,
+                    ],
+                    [
+                        0.98, -2.83, -0.10, 3.23, 0.51, 0.13, -1.03, 1.53, -0.51, 2.82,
+                    ],
+                ]),
+                z_1: arr1(&[
+                    -1.93, 0.46, 0.66, 0.02, -0.10, 1.29, 0.93, -1.14, 1.81, 1.47,
+                ]),
+                z_2: arr1(&[
+                    -0.26, 0.98, -1.30, 0.66, 0.57, -0.25, -0.31, -0.29, 0.69, 1.14,
+                ]),
+                pos: arr2(&[
+                    [0.00, 1.43, 2.86, 4.29, 5.71, 7.14, 9.57, 10.00],
+                    [-5.00, -3.57, -2.14, -0.71, 0.71, 2.14, 3.57, 5.00],
+                    [-6.00, -4.00, -2.00, 0.00, 2.00, 4.00, 6.00, 8.00],
+                ]),
+            }
+        }
+    }
+
+    #[test]
+    fn test_summate_fourier_2d() {
+        let setup = SetupFourier::new();
+        assert_eq!(
+            summator_fourier(
+                setup.spectrum_factor.view(),
+                setup.modes.view(),
+                setup.z_1.view(),
+                setup.z_2.view(),
+                setup.pos.view(),
+                None,
+            ),
+            arr1(&[
+                1.0666558330143816,
+                -3.5855143411414883,
+                -2.70208228699285,
+                9.808554698975039,
+                0.01634921830347258,
+                -2.2356422006860663,
+                14.730786907708966,
+                -2.851408419726332,
+            ])
+        );
+    }
+
     #[test]
-    fn test_summate_3d() {
+    fn test_summate_2d() {
         let setup = Setup::new();
 
         assert_eq!(
@@ -250,7 +382,7 @@ mod tests {
     }
 
     #[test]
-    fn test_summate_incompr_3d() {
+    fn test_summate_incompr_2d() {
         let setup = Setup::new();
 
         assert_ulps_eq!(

src/lib.rs

@@ -15,7 +15,7 @@
 use numpy::{IntoPyArray, PyArray1, PyArray2, PyReadonlyArray1, PyReadonlyArray2};
 use pyo3::prelude::{pymodule, PyModule, PyResult, Python};
 
-use field::{summator, summator_incompr};
+use field::{summator, summator_incompr, summator_fourier};
 use krige::{calculator_field_krige, calculator_field_krige_and_variance};
 use variogram::{
     variogram_directional, variogram_ma_structured, variogram_structured, variogram_unstructured,
@@ -64,6 +64,25 @@ fn gstools_core(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
         summator_incompr(cov_samples, z1, z2, pos, num_threads).into_pyarray(py)
     }
 
+    #[pyfn(m)]
+    #[pyo3(name = "summate_fourier")]
+    fn summate_fourier_py<'py>(
+        py: Python<'py>,
+        spectrum_factor: PyReadonlyArray1<f64>,
+        modes: PyReadonlyArray2<f64>,
+        z1: PyReadonlyArray1<f64>,
+        z2: PyReadonlyArray1<f64>,
+        pos: PyReadonlyArray2<f64>,
+        num_threads: Option<usize>,
+    ) -> &'py PyArray1<f64> {
+        let spectrum_factor = spectrum_factor.as_array();
+        let modes = modes.as_array();
+        let z1 = z1.as_array();
+        let z2 = z2.as_array();
+        let pos = pos.as_array();
+        summator_fourier(spectrum_factor, modes, z1, z2, pos, num_threads).into_pyarray(py)
+    }
+
     #[pyfn(m)]
     #[pyo3(name = "calc_field_krige_and_variance")]
     fn calc_field_krige_and_variance_py<'py>(