factor out diffx/diffy

This commit is contained in:
Magnus Ulimoen 2019-07-21 18:16:39 +02:00
parent 87ca45cbb8
commit 79941a5b09
1 changed files with 48 additions and 38 deletions

View File

@ -1,4 +1,4 @@
use ndarray::{Array2, ArrayView2, ArrayViewMut2};
use ndarray::{s, Array2, ArrayView1, ArrayView2, ArrayViewMut1, ArrayViewMut2};
use wasm_bindgen::prelude::*;
#[cfg(feature = "wee_alloc")]
@ -62,7 +62,8 @@ impl Universe {
fut.e_x.assign(&self.e_x);
traditional4(self.e_x.view(), fut.e_x.view_mut(), dt);
diffx(self.e_x.view(), fut.e_x.view_mut(), dt);
diffy(self.e_x.view(), fut.e_x.view_mut(), dt);
}
pub fn get_ptr(&mut self) -> *mut u8 {
@ -70,47 +71,56 @@ impl Universe {
}
}
fn traditional4(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>, dt: f32) {
fn diffx(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>, dt: f32) {
for j in 0..prev.shape()[0] {
trad4(prev.slice(s!(j, ..)), fut.slice_mut(s!(j, ..)), dt);
}
}
fn diffy(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>, dt: f32) {
for i in 0..prev.shape()[1] {
trad4(prev.slice(s!(.., i)), fut.slice_mut(s!(.., i)), dt);
}
}
fn trad4(prev: ArrayView1<f32>, mut fut: ArrayViewMut1<f32>, dt: f32) {
assert_eq!(prev.shape(), fut.shape());
let nx = prev.shape()[1];
let ny = prev.shape()[0];
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];
for j in 0..ny {
let diff = diag[0] * prev[(j, nx - 2)]
+ diag[1] * prev[(j, nx - 1)]
+ diag[2] * prev[(j, 0)]
+ diag[3] * prev[(j, 1)]
+ diag[4] * prev[(j, 2)];
fut[(j, 0)] += dt / dx * diff;
let diff = diag[0] * prev[(j, nx - 1)]
+ diag[1] * prev[(j, 0)]
+ diag[2] * prev[(j, 1)]
+ diag[3] * prev[(j, 2)]
+ diag[4] * prev[(j, 3)];
fut[(j, 1)] += dt / dx * diff;
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)] += dt / dx * diff;
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)] += dt / dx * diff;
for i in 2..nx - 2 {
let diff = diag[0] * prev[(j, i - 2)]
+ diag[1] * prev[(j, i - 1)]
+ diag[2] * prev[(j, i)]
+ diag[3] * prev[(j, i + 1)]
+ diag[4] * prev[(j, i + 2)];
fut[(j, i)] += dt / dx * diff;
}
let diff = diag[0] * prev[(j, nx - 4)]
+ diag[1] * prev[(j, nx - 3)]
+ diag[2] * prev[(j, nx - 2)]
+ diag[3] * prev[(j, nx - 1)]
+ diag[4] * prev[(j, 0)];
fut[(j, nx - 2)] += dt / dx * diff;
let diff = diag[0] * prev[(j, nx - 3)]
+ diag[1] * prev[(j, nx - 2)]
+ diag[2] * prev[(j, nx - 1)]
+ diag[3] * prev[(j, 0)]
+ diag[4] * prev[(j, 1)];
fut[(j, nx - 1)] += dt / dx * diff;
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)] += dt / dx * diff;
}
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)] += dt / dx * diff;
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)] += dt / dx * diff;
}