move operators to separate file

This commit is contained in:
Magnus Ulimoen 2019-08-09 16:49:19 +02:00
parent 46d71b469d
commit 8dc26ec8d8
2 changed files with 332 additions and 230 deletions

View File

@ -1,6 +1,9 @@
use ndarray::{s, Array2, ArrayView1, ArrayView2, ArrayViewMut1, ArrayViewMut2, Zip}; use ndarray::{Array2, Zip};
use wasm_bindgen::prelude::*; use wasm_bindgen::prelude::*;
mod operators;
use operators::{diffx, diffx_periodic, diffy, diffy_periodic, dissx, dissy};
#[cfg(feature = "wee_alloc")] #[cfg(feature = "wee_alloc")]
#[global_allocator] #[global_allocator]
static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT; static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;
@ -80,36 +83,36 @@ impl Universe {
// y.assign(&self.e_x) // y.assign(&self.e_x)
let mut k0 = &mut buffers.k1; // Only need k0 for first step in RK let mut k0 = &mut buffers.k1; // Only need k0 for first step in RK
k0.fill(0.0); k0.fill(0.0);
diffx(y.view(), k0.view_mut()); diffx_periodic(y.view(), k0.view_mut());
diffy(y.view(), k0.view_mut()); diffy_periodic(y.view(), k0.view_mut());
// y.assign(&self.e_x); // y.assign(&self.e_x);
y.scaled_add(-1.0 / 2.0 * dt, &k0); y.scaled_add(-1.0 / 2.0 * dt, &k0);
let mut k1 = buffers.k1; let mut k1 = buffers.k1;
k1.fill(0.0); k1.fill(0.0);
diffx(y.view(), k1.view_mut()); diffx_periodic(y.view(), k1.view_mut());
diffy(y.view(), k1.view_mut()); diffy_periodic(y.view(), k1.view_mut());
y.assign(&self.e_x); y.assign(&self.e_x);
y.scaled_add(-1.0 / 2.0 * dt, &k1); y.scaled_add(-1.0 / 2.0 * dt, &k1);
let mut k2 = buffers.k2; let mut k2 = buffers.k2;
k2.fill(0.0); k2.fill(0.0);
diffx(y.view(), k2.view_mut()); diffx_periodic(y.view(), k2.view_mut());
diffy(y.view(), k2.view_mut()); diffy_periodic(y.view(), k2.view_mut());
y.assign(&self.e_x); y.assign(&self.e_x);
y.scaled_add(-1.0 / 2.0 * dt, &k2); y.scaled_add(-1.0 / 2.0 * dt, &k2);
let mut k3 = buffers.k3; let mut k3 = buffers.k3;
k3.fill(0.0); k3.fill(0.0);
diffx(y.view(), k3.view_mut()); diffx_periodic(y.view(), k3.view_mut());
diffy(y.view(), k3.view_mut()); diffy_periodic(y.view(), k3.view_mut());
y.assign(&self.e_x); y.assign(&self.e_x);
y.scaled_add(-dt, &k2); y.scaled_add(-dt, &k2);
let mut k4 = buffers.k4; let mut k4 = buffers.k4;
k4.fill(0.0); k4.fill(0.0);
diffx(y.view(), k4.view_mut()); diffx_periodic(y.view(), k4.view_mut());
diffy(y.view(), k4.view_mut()); diffy_periodic(y.view(), k4.view_mut());
Zip::from(&mut fut.e_x) Zip::from(&mut fut.e_x)
.and(&self.e_x) .and(&self.e_x)
@ -193,225 +196,6 @@ impl Universe {
} }
} }
fn diffx(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>) {
for j in 0..prev.shape()[0] {
upwind4_periodic(prev.slice(s!(j, ..)), fut.slice_mut(s!(j, ..)));
}
}
fn diffy(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>) {
for i in 0..prev.shape()[1] {
upwind4_periodic(prev.slice(s!(.., i)), fut.slice_mut(s!(.., i)));
}
}
fn dissx(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>) {
for j in 0..prev.shape()[0] {
upwind4_diss(prev.slice(s!(j, ..)), fut.slice_mut(s!(j, ..)));
}
}
fn dissy(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>) {
for i in 0..prev.shape()[1] {
upwind4_diss(prev.slice(s!(.., i)), fut.slice_mut(s!(.., i)));
}
}
fn trad4_periodic(prev: ArrayView1<f32>, mut fut: ArrayViewMut1<f32>) {
assert_eq!(prev.shape(), fut.shape());
let nx = prev.shape()[0];
let dx = 1.0 / (nx - 1) as f32;
let diag = [1.0 / 12.0, -2.0 / 3.0, 0.0, 2.0 / 3.0, -1.0 / 12.0];
let diff = diag[0] * prev[(nx - 2)]
+ diag[1] * prev[(nx - 1)]
+ diag[2] * prev[(0)]
+ diag[3] * prev[(1)]
+ diag[4] * prev[(2)];
fut[(0)] += diff / dx;
let diff = diag[0] * prev[(nx - 1)]
+ diag[1] * prev[(0)]
+ diag[2] * prev[(1)]
+ diag[3] * prev[(2)]
+ diag[4] * prev[(3)];
fut[(1)] += diff / dx;
for i in 2..nx - 2 {
let diff = diag[0] * prev[(i - 2)]
+ diag[1] * prev[(i - 1)]
+ diag[2] * prev[(i)]
+ diag[3] * prev[(i + 1)]
+ diag[4] * prev[(i + 2)];
fut[(i)] += diff / dx;
}
let diff = diag[0] * prev[(nx - 4)]
+ diag[1] * prev[(nx - 3)]
+ diag[2] * prev[(nx - 2)]
+ diag[3] * prev[(nx - 1)]
+ diag[4] * prev[(0)];
fut[(nx - 2)] += diff / dx;
let diff = diag[0] * prev[(nx - 3)]
+ diag[1] * prev[(nx - 2)]
+ diag[2] * prev[(nx - 1)]
+ diag[3] * prev[(0)]
+ diag[4] * prev[(1)];
fut[(nx - 1)] += diff / dx;
}
fn upwind4_periodic(prev: ArrayView1<f32>, mut fut: ArrayViewMut1<f32>) {
assert_eq!(prev.shape(), fut.shape());
let nx = prev.shape()[0];
let dx = 1.0 / (nx - 1) as f32;
let diag = [
-1.0 / 24.0,
1.0 / 4.0,
-7.0 / 8.0,
0.0,
7.0 / 8.0,
-1.0 / 4.0,
1.0 / 24.0,
];
let diff = diag[0] * prev[(nx - 3)]
+ diag[1] * prev[(nx - 2)]
+ diag[2] * prev[(nx - 1)]
+ diag[3] * prev[(0)]
+ diag[4] * prev[(1)]
+ diag[5] * prev[(2)]
+ diag[6] * prev[(3)];
fut[0] += diff / dx;
let diff = diag[0] * prev[(nx - 2)]
+ diag[1] * prev[(nx - 1)]
+ diag[2] * prev[(0)]
+ diag[3] * prev[(1)]
+ diag[4] * prev[(2)]
+ diag[5] * prev[(3)]
+ diag[6] * prev[(4)];
fut[1] += diff / dx;
let diff = diag[0] * prev[(nx - 1)]
+ diag[1] * prev[(0)]
+ diag[2] * prev[(1)]
+ diag[3] * prev[(2)]
+ diag[4] * prev[(3)]
+ diag[5] * prev[(4)]
+ diag[6] * prev[(5)];
fut[2] += diff / dx;
for i in 3..nx - 3 {
let diff = diag[0] * prev[(i - 3)]
+ diag[1] * prev[(i - 2)]
+ diag[2] * prev[(i - 1)]
+ diag[3] * prev[(i)]
+ diag[4] * prev[(i + 1)]
+ diag[5] * prev[(i + 2)]
+ diag[6] * prev[(i + 3)];
fut[(i)] += diff / dx;
}
let diff = diag[0] * prev[(nx - 6)]
+ diag[1] * prev[(nx - 5)]
+ diag[2] * prev[(nx - 4)]
+ diag[3] * prev[(nx - 3)]
+ diag[4] * prev[(nx - 2)]
+ diag[5] * prev[(nx - 1)]
+ diag[6] * prev[(0)];
fut[(nx - 3)] += diff / dx;
let diff = diag[0] * prev[(nx - 5)]
+ diag[1] * prev[(nx - 4)]
+ diag[2] * prev[(nx - 3)]
+ diag[3] * prev[(nx - 2)]
+ diag[4] * prev[(nx - 1)]
+ diag[5] * prev[(0)]
+ diag[6] * prev[(1)];
fut[(nx - 2)] += diff / dx;
let diff = diag[0] * prev[(nx - 4)]
+ diag[1] * prev[(nx - 3)]
+ diag[2] * prev[(nx - 2)]
+ diag[3] * prev[(nx - 1)]
+ diag[4] * prev[(0)]
+ diag[5] * prev[(1)]
+ diag[6] * prev[(2)];
fut[(nx - 1)] += diff / dx;
}
fn upwind4_diss(prev: ArrayView1<f32>, mut fut: ArrayViewMut1<f32>) {
assert_eq!(prev.shape(), fut.shape());
let nx = prev.shape()[0];
let dx = 1.0 / (nx - 1) as f32;
let diag = [
1.0 / 24.0,
-1.0 / 4.0,
5.0 / 8.0,
-5.0 / 6.0,
5.0 / 8.0,
-1.0 / 4.0,
1.0 / 24.0,
];
let diff = diag[0] * prev[(nx - 3)]
+ diag[1] * prev[(nx - 2)]
+ diag[2] * prev[(nx - 1)]
+ diag[3] * prev[(0)]
+ diag[4] * prev[(1)]
+ diag[5] * prev[(2)]
+ diag[6] * prev[(3)];
fut[0] += diff / dx;
let diff = diag[0] * prev[(nx - 2)]
+ diag[1] * prev[(nx - 1)]
+ diag[2] * prev[(0)]
+ diag[3] * prev[(1)]
+ diag[4] * prev[(2)]
+ diag[5] * prev[(3)]
+ diag[6] * prev[(4)];
fut[1] += diff / dx;
let diff = diag[0] * prev[(nx - 1)]
+ diag[1] * prev[(0)]
+ diag[2] * prev[(1)]
+ diag[3] * prev[(2)]
+ diag[4] * prev[(3)]
+ diag[5] * prev[(4)]
+ diag[6] * prev[(5)];
fut[2] += diff / dx;
for i in 3..nx - 3 {
let diff = diag[0] * prev[(i - 3)]
+ diag[1] * prev[(i - 2)]
+ diag[2] * prev[(i - 1)]
+ diag[3] * prev[(i)]
+ diag[4] * prev[(i + 1)]
+ diag[5] * prev[(i + 2)]
+ diag[6] * prev[(i + 3)];
fut[(i)] += diff / dx;
}
let diff = diag[0] * prev[(nx - 6)]
+ diag[1] * prev[(nx - 5)]
+ diag[2] * prev[(nx - 4)]
+ diag[3] * prev[(nx - 3)]
+ diag[4] * prev[(nx - 2)]
+ diag[5] * prev[(nx - 1)]
+ diag[6] * prev[(0)];
fut[(nx - 3)] += diff / dx;
let diff = diag[0] * prev[(nx - 5)]
+ diag[1] * prev[(nx - 4)]
+ diag[2] * prev[(nx - 3)]
+ diag[3] * prev[(nx - 2)]
+ diag[4] * prev[(nx - 1)]
+ diag[5] * prev[(0)]
+ diag[6] * prev[(1)];
fut[(nx - 2)] += diff / dx;
let diff = diag[0] * prev[(nx - 4)]
+ diag[1] * prev[(nx - 3)]
+ diag[2] * prev[(nx - 2)]
+ diag[3] * prev[(nx - 1)]
+ diag[4] * prev[(0)]
+ diag[5] * prev[(1)]
+ diag[6] * prev[(2)];
fut[(nx - 1)] += diff / dx;
}
#[wasm_bindgen] #[wasm_bindgen]
pub struct WorkBuffers { pub struct WorkBuffers {
k1: Array2<f32>, k1: Array2<f32>,

318
src/operators.rs Normal file
View File

@ -0,0 +1,318 @@
use ndarray::{arr1, arr2, s, ArrayView1, ArrayView2, ArrayViewMut1, ArrayViewMut2};
pub(crate) fn diffx(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>) {
for j in 0..prev.shape()[0] {
upwind4(prev.slice(s!(j, ..)), fut.slice_mut(s!(j, ..)));
}
}
pub(crate) fn diffx_periodic(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>) {
for j in 0..prev.shape()[0] {
upwind4_periodic(prev.slice(s!(j, ..)), fut.slice_mut(s!(j, ..)));
}
}
pub(crate) fn diffy(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>) {
for i in 0..prev.shape()[1] {
upwind4(prev.slice(s!(.., i)), fut.slice_mut(s!(.., i)));
}
}
pub(crate) fn diffy_periodic(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>) {
for i in 0..prev.shape()[1] {
upwind4_periodic(prev.slice(s!(.., i)), fut.slice_mut(s!(.., i)));
}
}
pub(crate) fn dissx(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>) {
for j in 0..prev.shape()[0] {
upwind4_diss(prev.slice(s!(j, ..)), fut.slice_mut(s!(j, ..)));
}
}
pub(crate) fn dissy(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>) {
for i in 0..prev.shape()[1] {
upwind4_diss(prev.slice(s!(.., i)), fut.slice_mut(s!(.., i)));
}
}
enum Border {
N,
S,
E,
W,
}
fn apply_periodic_SAT(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>, border: Border) {}
fn trad4_periodic(prev: ArrayView1<f32>, mut fut: ArrayViewMut1<f32>) {
assert_eq!(prev.shape(), fut.shape());
let nx = prev.shape()[0];
let dx = 1.0 / (nx - 1) as f32;
let diag = [1.0 / 12.0, -2.0 / 3.0, 0.0, 2.0 / 3.0, -1.0 / 12.0];
let diff = diag[0] * prev[(nx - 2)]
+ diag[1] * prev[(nx - 1)]
+ diag[2] * prev[(0)]
+ diag[3] * prev[(1)]
+ diag[4] * prev[(2)];
fut[(0)] += diff / dx;
let diff = diag[0] * prev[(nx - 1)]
+ diag[1] * prev[(0)]
+ diag[2] * prev[(1)]
+ diag[3] * prev[(2)]
+ diag[4] * prev[(3)];
fut[(1)] += diff / dx;
for i in 2..nx - 2 {
let diff = diag[0] * prev[(i - 2)]
+ diag[1] * prev[(i - 1)]
+ diag[2] * prev[(i)]
+ diag[3] * prev[(i + 1)]
+ diag[4] * prev[(i + 2)];
fut[(i)] += diff / dx;
}
let diff = diag[0] * prev[(nx - 4)]
+ diag[1] * prev[(nx - 3)]
+ diag[2] * prev[(nx - 2)]
+ diag[3] * prev[(nx - 1)]
+ diag[4] * prev[(0)];
fut[(nx - 2)] += diff / dx;
let diff = diag[0] * prev[(nx - 3)]
+ diag[1] * prev[(nx - 2)]
+ diag[2] * prev[(nx - 1)]
+ diag[3] * prev[(0)]
+ diag[4] * prev[(1)];
fut[(nx - 1)] += diff / dx;
}
fn upwind4(prev: ArrayView1<f32>, mut fut: ArrayViewMut1<f32>) {
assert_eq!(prev.shape(), fut.shape());
let nx = prev.shape()[0];
let dx = 1.0 / (nx - 1) as f32;
let diag = arr1(&[
-1.0 / 24.0,
1.0 / 4.0,
-7.0 / 8.0,
0.0,
7.0 / 8.0,
-1.0 / 4.0,
1.0 / 24.0,
]);
let block = arr2(&[
[
-72.0 / 49.0f32,
187.0 / 98.0,
-20.0 / 49.0,
-3.0 / 98.0,
0.0,
0.0,
0.0,
],
[
-187.0 / 366.0,
0.0,
69.0 / 122.0,
-16.0 / 183.0,
2.0 / 61.0,
0.0,
0.0,
],
[
20.0 / 123.0,
-69.0 / 82.0,
0.0,
227.0 / 246.0,
-12.0 / 41.0,
2.0 / 41.0,
0.0,
],
[
3.0 / 298.0,
16.0 / 149.0,
-227.0 / 298.0,
0.0,
126.0 / 149.0,
-36.0 / 149.0,
6.0 / 149.0,
],
]);
let h_block = [49.0 / 144.0, 61.0 / 48.0, 41.0 / 48.0, 149.0 / 144.0];
let first_elems = prev.slice(s!(..7));
for i in 0..4 {
let diff = first_elems.dot(&block.slice(s!(i, ..)));
fut[i] += diff / dx * h_block[i];
}
for i in 4..nx - 4 {
let diff = diag.dot(&prev.slice(s!(i - 3..i + 3 + 1)));
fut[(i)] += diff / dx;
}
let last_elems = prev.slice(s!(nx - 7..));
for i in 0..4 {
let ii = nx - 4 + i;
let block = block.slice(s!(3 - i, ..;-1));
let diff = last_elems.dot(&block);
fut[ii] += diff / dx * h_block[3 - i];
}
}
fn upwind4_periodic(prev: ArrayView1<f32>, mut fut: ArrayViewMut1<f32>) {
assert_eq!(prev.shape(), fut.shape());
let nx = prev.shape()[0];
let dx = 1.0 / (nx - 1) as f32;
let diag = [
-1.0 / 24.0,
1.0 / 4.0,
-7.0 / 8.0,
0.0,
7.0 / 8.0,
-1.0 / 4.0,
1.0 / 24.0,
];
let diff = diag[0] * prev[(nx - 3)]
+ diag[1] * prev[(nx - 2)]
+ diag[2] * prev[(nx - 1)]
+ diag[3] * prev[(0)]
+ diag[4] * prev[(1)]
+ diag[5] * prev[(2)]
+ diag[6] * prev[(3)];
fut[0] += diff / dx;
let diff = diag[0] * prev[(nx - 2)]
+ diag[1] * prev[(nx - 1)]
+ diag[2] * prev[(0)]
+ diag[3] * prev[(1)]
+ diag[4] * prev[(2)]
+ diag[5] * prev[(3)]
+ diag[6] * prev[(4)];
fut[1] += diff / dx;
let diff = diag[0] * prev[(nx - 1)]
+ diag[1] * prev[(0)]
+ diag[2] * prev[(1)]
+ diag[3] * prev[(2)]
+ diag[4] * prev[(3)]
+ diag[5] * prev[(4)]
+ diag[6] * prev[(5)];
fut[2] += diff / dx;
for i in 3..nx - 3 {
let diff = diag[0] * prev[(i - 3)]
+ diag[1] * prev[(i - 2)]
+ diag[2] * prev[(i - 1)]
+ diag[3] * prev[(i)]
+ diag[4] * prev[(i + 1)]
+ diag[5] * prev[(i + 2)]
+ diag[6] * prev[(i + 3)];
fut[(i)] += diff / dx;
}
let diff = diag[0] * prev[(nx - 6)]
+ diag[1] * prev[(nx - 5)]
+ diag[2] * prev[(nx - 4)]
+ diag[3] * prev[(nx - 3)]
+ diag[4] * prev[(nx - 2)]
+ diag[5] * prev[(nx - 1)]
+ diag[6] * prev[(0)];
fut[(nx - 3)] += diff / dx;
let diff = diag[0] * prev[(nx - 5)]
+ diag[1] * prev[(nx - 4)]
+ diag[2] * prev[(nx - 3)]
+ diag[3] * prev[(nx - 2)]
+ diag[4] * prev[(nx - 1)]
+ diag[5] * prev[(0)]
+ diag[6] * prev[(1)];
fut[(nx - 2)] += diff / dx;
let diff = diag[0] * prev[(nx - 4)]
+ diag[1] * prev[(nx - 3)]
+ diag[2] * prev[(nx - 2)]
+ diag[3] * prev[(nx - 1)]
+ diag[4] * prev[(0)]
+ diag[5] * prev[(1)]
+ diag[6] * prev[(2)];
fut[(nx - 1)] += diff / dx;
}
fn upwind4_diss(prev: ArrayView1<f32>, mut fut: ArrayViewMut1<f32>) {
assert_eq!(prev.shape(), fut.shape());
let nx = prev.shape()[0];
let dx = 1.0 / (nx - 1) as f32;
let diag = [
1.0 / 24.0,
-1.0 / 4.0,
5.0 / 8.0,
-5.0 / 6.0,
5.0 / 8.0,
-1.0 / 4.0,
1.0 / 24.0,
];
let diff = diag[0] * prev[(nx - 3)]
+ diag[1] * prev[(nx - 2)]
+ diag[2] * prev[(nx - 1)]
+ diag[3] * prev[(0)]
+ diag[4] * prev[(1)]
+ diag[5] * prev[(2)]
+ diag[6] * prev[(3)];
fut[0] += diff / dx;
let diff = diag[0] * prev[(nx - 2)]
+ diag[1] * prev[(nx - 1)]
+ diag[2] * prev[(0)]
+ diag[3] * prev[(1)]
+ diag[4] * prev[(2)]
+ diag[5] * prev[(3)]
+ diag[6] * prev[(4)];
fut[1] += diff / dx;
let diff = diag[0] * prev[(nx - 1)]
+ diag[1] * prev[(0)]
+ diag[2] * prev[(1)]
+ diag[3] * prev[(2)]
+ diag[4] * prev[(3)]
+ diag[5] * prev[(4)]
+ diag[6] * prev[(5)];
fut[2] += diff / dx;
for i in 3..nx - 3 {
let diff = diag[0] * prev[(i - 3)]
+ diag[1] * prev[(i - 2)]
+ diag[2] * prev[(i - 1)]
+ diag[3] * prev[(i)]
+ diag[4] * prev[(i + 1)]
+ diag[5] * prev[(i + 2)]
+ diag[6] * prev[(i + 3)];
fut[(i)] += diff / dx;
}
let diff = diag[0] * prev[(nx - 6)]
+ diag[1] * prev[(nx - 5)]
+ diag[2] * prev[(nx - 4)]
+ diag[3] * prev[(nx - 3)]
+ diag[4] * prev[(nx - 2)]
+ diag[5] * prev[(nx - 1)]
+ diag[6] * prev[(0)];
fut[(nx - 3)] += diff / dx;
let diff = diag[0] * prev[(nx - 5)]
+ diag[1] * prev[(nx - 4)]
+ diag[2] * prev[(nx - 3)]
+ diag[3] * prev[(nx - 2)]
+ diag[4] * prev[(nx - 1)]
+ diag[5] * prev[(0)]
+ diag[6] * prev[(1)];
fut[(nx - 2)] += diff / dx;
let diff = diag[0] * prev[(nx - 4)]
+ diag[1] * prev[(nx - 3)]
+ diag[2] * prev[(nx - 2)]
+ diag[3] * prev[(nx - 1)]
+ diag[4] * prev[(0)]
+ diag[5] * prev[(1)]
+ diag[6] * prev[(2)];
fut[(nx - 1)] += diff / dx;
}