Skip to content

Commit

Permalink
Add Fourier method
Browse files Browse the repository at this point in the history
  • Loading branch information
LSchueler committed Jul 2, 2024
1 parent 530ffe4 commit cc167a6
Show file tree
Hide file tree
Showing 2 changed files with 154 additions and 3 deletions.
136 changes: 134 additions & 2 deletions src/field.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand Down Expand Up @@ -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!(
Expand All @@ -250,7 +382,7 @@ mod tests {
}

#[test]
fn test_summate_incompr_3d() {
fn test_summate_incompr_2d() {
let setup = Setup::new();

assert_ulps_eq!(
Expand Down
21 changes: 20 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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>(
Expand Down

0 comments on commit cc167a6

Please sign in to comment.