DEV Community

Toshiki Teramura
Toshiki Teramura

Posted on

stencil; abstract stencil calculation

#rust #simulation #stencil

I am developing a library for stencil calculation in Rust.

https://github.com/termoshtt/stencil
Abstract stencil calculation

stencil v0.2 is published on crates.io, and you can use it through cargo:

[dependencies] stencil = "0.2"
Enter fullscreen mode Exit fullscreen mode

The full document is available on docs.rs.

StencilArray trait

This library abstract a stencil calculation through StencilArray trait:

/// Array with stencil calculations pub trait StencilArray<St>: NdArray where St: Stencil<Elem = Self::Elem, Dim = Self::Dim>, { /// Execute a stencil calculation fn stencil_map<Output, Func>(&self, out: &mut Output, Func) where Output: NdArray<Dim = Self::Dim>, Func: Fn(St) -> Output::Elem; }
Enter fullscreen mode Exit fullscreen mode

stencil_map function is the core of this crate. Basic example is here:

let mut a = torus::Torus::<f64, Ix1>::zeros(n); let mut b = a.clone(); a.coordinate_fill(|x| x.sin()); let dx = a.dx(); a.stencil_map(&mut b, |n: N1D1<f64>| (n.r - n.l) / (2.0 * dx));
Enter fullscreen mode Exit fullscreen mode

This crate defines Torus<A, D>, which represents N-dimensional torus. a denotes one-dimensional torus, we can regard it as a region with periodic boundary condition. The second argument of stencil_map, a lambda function |n: N1D1<f64>| (n.r - n.l) / (2.0 * dx) denotes a central difference. N1D1 means 1-neighbor on 1-dimensional space. So, this example describes a central difference of sin(x) function, i.e. cos(x) is set b.

N1D1<A> is a small struct:

/// one-neighbor, one-dimensional stencil #[derive(Clone, Copy)] pub struct N1D1<A: LinalgScalar> { /// left pub l: A, /// right pub r: A, /// center pub c: A, }
Enter fullscreen mode Exit fullscreen mode

and implements the trait Stencil. It has values around "what we focused on"; when we update b[i] using a, n has a[i] as n.c and its neighbors a[i-1] as n.l and a[i+1] as n.r. This crate abstracts the stencil calculation in this way.
The implementation of StencilArray for the Torus<A, D> describes how to get N1D1 from the array a. Currently, there are three Stencils: N1D1<A>, N2D1<A>, and N1D2<A>, and two StencilArrays: Torus<A, D> and Line<A, P: Padding, E: Edge>. The Line has a padding at the each edge of line, and it can be regarded as a fixed boundary condition.

Manifold trait

stencil crate introduces one more useful trait, Manifold:

/// Uniformly coordinated array pub trait Manifold: NdArray { /// Type of coordinate type Coordinate; /// Increment of coordinate fn dx(&self) -> Self::Coordinate; /// Fill manifold by a function fn coordinate_fill<F>(&mut self, F) where F: Fn(Self::Coordinate) -> Self::Elem; /// Map values on manifold using a function fn coordinate_map<F>(&mut self, F) where F: Fn(Self::Coordinate, Self::Elem) -> Self::Elem; }
Enter fullscreen mode Exit fullscreen mode

This trait gives array a coordinate. This is useful to set value of array by mathematical function as you see above example

a.coordinate_fill(|x| x.sin());
Enter fullscreen mode Exit fullscreen mode

x is set by the implementation of Manifold for each objects.

Example1: Diffusion equation

Let us solve the diffusion equation for both periodic and fixed boundary conditions:

fn periodic(r: Sender<Vec<f64>>) { let dt = 1e-3; let mut t = torus::Torus::<f64, Ix1>::zeros(32); t.coordinate_fill(|x| x.sin()); let dx = t.dx(); let mut s = t.clone(); for _step in 0..3000 { t.stencil_map(&mut s, |n: N1D1<f64>| { let d2 = (n.l + n.r - 2.0 * n.c) / (dx.powi(2)); n.c + dt * d2 }); r.send(s.as_view().to_vec()).unwrap(); swap(&mut t, &mut s); } }
Enter fullscreen mode Exit fullscreen mode

heat-periodic

fn fixed(r: Sender<Vec<f64>>) { let dt = 1e-3; let mut t = Line::<f64, P1, Closed>::new(64, 0.0, 1.0, 2.5 * PI); t.coordinate_fill(|x| x.sin()); let dx = t.dx(); let mut s = t.clone(); for _step in 0..10000 { t.stencil_map(&mut s, |n: N1D1<f64>| { let d2 = (n.l + n.r - 2.0 * n.c) / (dx.powi(2)); n.c + dt * d2 }); r.send(s.as_view().to_vec()).unwrap(); swap(&mut t, &mut s); } }
Enter fullscreen mode Exit fullscreen mode

heat-fixed

This example uses asink for storing data. Please see the last post.

Example2: KdV equation

At last, we solve KdV equation using the ZK-scheme:

fn kdv(r: Sender<Vec<f64>>) { let dt = 1e-5; let mut t = torus::Torus::<f64, Ix1>::zeros(128); t.coordinate_fill(|x| x.sin()); let dx = t.dx(); let mut b = t.clone(); let mut n = t.clone(); for step in 0..1_000_000 { t.stencil_map(&mut n, |n: N2D1<f64>| { // ZK scheme let uux = (n.l + n.c + n.r) * (n.r - n.l) / (3.0 * dx); let u3 = (n.rr - 2.0 * n.r + 2.0 * n.l - n.ll) / dx.powi(3); uux + 0.01 * u3 }); azip!(mut n, b in { *n = b - dt * *n }); if step % 1000 == 0 { r.send(n.as_view().to_vec()).unwrap(); } swap(&mut t, &mut b); swap(&mut t, &mut n); } }
Enter fullscreen mode Exit fullscreen mode

kdv

Leap-frog algorithm can also be written easily.

Please see the examples/ directory for full examples.
I'd like to get your feedback. Thanks.

Top comments (0)