separate maxwell solver

This commit is contained in:
Magnus Ulimoen 2019-08-13 20:43:31 +02:00
parent 407877c2e6
commit a2185d8900
3 changed files with 175 additions and 151 deletions

18
main.js
View File

@ -1,4 +1,4 @@
import { Universe, WorkBuffers, set_panic_hook, default as init } from "./webgl.js"; import { Universe, set_panic_hook, default as init } from "./webgl.js";
async function run() { async function run() {
let wasm = await init("./webgl_bg.wasm"); let wasm = await init("./webgl_bg.wasm");
@ -148,12 +148,11 @@ async function run() {
const width = 40; const width = 40;
const height = 50; const height = 50;
let universes = [Universe.new(width, height), Universe.new(width, height)]; const universe = Universe.new(width, height);
const workbuffer = WorkBuffers.new(width, height);
const TIMEFACTOR = 1.0/7000; const TIMEFACTOR = 1.0/7000;
let t = performance.now()*TIMEFACTOR; let t = performance.now()*TIMEFACTOR;
universes[0].set_initial(0.5, 0.5); universe.init(0.5, 0.5);
function drawMe(t_draw) { function drawMe(t_draw) {
gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT); gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);
@ -166,16 +165,16 @@ async function run() {
} else { } else {
t += dt; t += dt;
} }
universes[0].advance(universes[1], dt, workbuffer); universe.advance(dt);
const field_ex = new Float32Array(wasm.memory.buffer, const field_ex = new Float32Array(wasm.memory.buffer,
universes[0].get_ex_ptr(), universe.get_ex_ptr(),
width*height); width*height);
const field_ey = new Float32Array(wasm.memory.buffer, const field_ey = new Float32Array(wasm.memory.buffer,
universes[0].get_ey_ptr(), universe.get_ey_ptr(),
width*height); width*height);
const field_hz = new Float32Array(wasm.memory.buffer, const field_hz = new Float32Array(wasm.memory.buffer,
universes[0].get_hz_ptr(), universe.get_hz_ptr(),
width*height); width*height);
{ {
const level = 0; const level = 0;
@ -197,7 +196,6 @@ async function run() {
gl.drawArrays(gl.TRIANGLE_STRIP, offset, vertexCount); gl.drawArrays(gl.TRIANGLE_STRIP, offset, vertexCount);
} }
universes = [universes[1], universes[0]];
window.requestAnimationFrame(drawMe); window.requestAnimationFrame(drawMe);
} }
@ -218,7 +216,7 @@ async function run() {
window.addEventListener('click', event => { window.addEventListener('click', event => {
const mousex = event.clientX / window.innerWidth; const mousex = event.clientX / window.innerWidth;
const mousey = event.clientY / window.innerHeight; const mousey = event.clientY / window.innerHeight;
universes[0].set_initial(mousex, 1.0 - mousey); universe.init(mousex, 1.0 - mousey);
}, {passive: true}); }, {passive: true});
resizeCanvas(); resizeCanvas();

View File

@ -1,175 +1,52 @@
use ndarray::{Array2, Zip};
use wasm_bindgen::prelude::*; use wasm_bindgen::prelude::*;
mod maxwell;
mod operators; mod operators;
use operators::{diffx_periodic, diffy_periodic}; use maxwell::{System, WorkBuffers};
#[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;
#[wasm_bindgen]
pub struct Universe {
ex: Array2<f32>,
ey: Array2<f32>,
hz: Array2<f32>,
}
#[wasm_bindgen] #[wasm_bindgen]
pub fn set_panic_hook() { pub fn set_panic_hook() {
#[cfg(feature = "console_error_panic_hook")] #[cfg(feature = "console_error_panic_hook")]
console_error_panic_hook::set_once(); console_error_panic_hook::set_once();
} }
fn gaussian(x: f32, x0: f32, y: f32, y0: f32) -> f32 { #[wasm_bindgen]
use std::f32; pub struct Universe {
let x = x - x0; sys: (System, System),
let y = y - y0; wb: WorkBuffers,
let sigma = 0.05;
1.0 / (2.0 * f32::consts::PI * sigma * sigma) * (-(x * x + y * y) / (2.0 * sigma * sigma)).exp()
} }
#[wasm_bindgen] #[wasm_bindgen]
impl Universe { impl Universe {
pub fn new(width: u32, height: u32) -> Self { pub fn new(width: u32, height: u32) -> Self {
let field = Array2::zeros((height as usize, width as usize)); Self {
let ex = field.clone(); sys: (System::new(width, height), System::new(width, height)),
let ey = field.clone(); wb: WorkBuffers::new(width as usize, height as usize),
let hz = field;
Universe { ex, ey, hz }
}
pub fn set_initial(&mut self, x0: f32, y0: f32) {
let nx = self.ex.shape()[1];
let ny = self.ex.shape()[0];
for j in 0..ny {
for i in 0..nx {
// Must divice interval on nx/ny instead of nx - 1/ny-1
// due to periodic conditions [0, 1)
let x = i as f32 / nx as f32;
let y = j as f32 / ny as f32;
self.ex[(j, i)] = 0.0;
self.ey[(j, i)] = 0.0;
self.hz[(j, i)] = gaussian(x, x0, y, y0) / 32.0;
}
} }
} }
pub fn advance(&self, fut: &mut Universe, dt: f32, work_buffers: Option<WorkBuffers>) { pub fn init(&mut self, x0: f32, y0: f32) {
assert_eq!(self.ex.shape(), fut.ex.shape()); self.sys.0.set_gaussian(x0, y0);
}
let mut buffers = work_buffers pub fn advance(&mut self, dt: f32) {
.unwrap_or_else(|| WorkBuffers::new(self.ex.shape()[1], self.ex.shape()[0])); self.sys.0.advance(&mut self.sys.1, dt, Some(&mut self.wb));
std::mem::swap(&mut self.sys.0, &mut self.sys.1);
let mut y = buffers.y;
let mut k = buffers.buf;
for i in 0..4 {
// y = y0 + c*kn
y.0.assign(&self.ex);
y.1.assign(&self.hz);
y.2.assign(&self.ey);
match i {
0 => {}
1 => {
y.0.scaled_add(1.0 / 2.0 * dt, &k[i - 1].0);
y.1.scaled_add(1.0 / 2.0 * dt, &k[i - 1].1);
y.2.scaled_add(1.0 / 2.0 * dt, &k[i - 1].2);
}
2 => {
y.0.scaled_add(1.0 / 2.0 * dt, &k[i - 1].0);
y.1.scaled_add(1.0 / 2.0 * dt, &k[i - 1].1);
y.2.scaled_add(1.0 / 2.0 * dt, &k[i - 1].2);
}
3 => {
y.0.scaled_add(dt, &k[i - 1].0);
y.1.scaled_add(dt, &k[i - 1].1);
y.2.scaled_add(dt, &k[i - 1].2);
}
_ => {
unreachable!();
}
};
k[i].0.fill(0.0);
// ex = hz_y
diffy_periodic(y.1.view(), k[i].0.view_mut());
// ey = -hz_x
k[i].2.fill(0.0);
diffx_periodic(y.1.view(), k[i].2.view_mut());
k[i].2.mapv_inplace(|v| -v);
// hz = -ey_x + ex_y
k[i].1.fill(0.0);
diffx_periodic(y.2.view(), k[i].1.view_mut());
k[i].1.mapv_inplace(|v| -v);
diffy_periodic(y.0.view(), k[i].1.view_mut());
}
Zip::from(&mut fut.ex)
.and(&self.ex)
.and(&k[0].0)
.and(&k[1].0)
.and(&k[2].0)
.and(&k[3].0)
.apply(|y1, &y0, &k1, &k2, &k3, &k4| {
*y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
});
Zip::from(&mut fut.hz)
.and(&self.hz)
.and(&k[0].1)
.and(&k[1].1)
.and(&k[2].1)
.and(&k[3].1)
.apply(|y1, &y0, &k1, &k2, &k3, &k4| {
*y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
});
Zip::from(&mut fut.ey)
.and(&self.ey)
.and(&k[0].2)
.and(&k[1].2)
.and(&k[2].2)
.and(&k[3].2)
.apply(|y1, &y0, &k1, &k2, &k3, &k4| {
*y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
});
} }
pub fn get_ex_ptr(&mut self) -> *mut u8 { pub fn get_ex_ptr(&mut self) -> *mut u8 {
self.ex.as_mut_ptr() as *mut u8 self.sys.0.ex.as_mut_ptr() as *mut u8
} }
pub fn get_ey_ptr(&mut self) -> *mut u8 { pub fn get_ey_ptr(&mut self) -> *mut u8 {
self.ey.as_mut_ptr() as *mut u8 self.sys.0.ey.as_mut_ptr() as *mut u8
} }
pub fn get_hz_ptr(&mut self) -> *mut u8 { pub fn get_hz_ptr(&mut self) -> *mut u8 {
self.hz.as_mut_ptr() as *mut u8 self.sys.0.hz.as_mut_ptr() as *mut u8
}
}
#[wasm_bindgen]
pub struct WorkBuffers {
y: (Array2<f32>, Array2<f32>, Array2<f32>),
buf: [(Array2<f32>, Array2<f32>, Array2<f32>); 4],
}
#[wasm_bindgen]
impl WorkBuffers {
pub fn new(nx: usize, ny: usize) -> Self {
let arr = Array2::zeros((ny, nx));
Self {
y: (arr.clone(), arr.clone(), arr.clone()),
buf: [
(arr.clone(), arr.clone(), arr.clone()),
(arr.clone(), arr.clone(), arr.clone()),
(arr.clone(), arr.clone(), arr.clone()),
(arr.clone(), arr.clone(), arr),
],
}
} }
} }

149
src/maxwell.rs Normal file
View File

@ -0,0 +1,149 @@
use super::operators::{diffx_periodic, diffy_periodic};
use ndarray::{Array2, Zip};
pub struct System {
pub(crate) ex: Array2<f32>,
pub(crate) ey: Array2<f32>,
pub(crate) hz: Array2<f32>,
}
fn gaussian(x: f32, x0: f32, y: f32, y0: f32) -> f32 {
use std::f32;
let x = x - x0;
let y = y - y0;
let sigma = 0.05;
1.0 / (2.0 * f32::consts::PI * sigma * sigma) * (-(x * x + y * y) / (2.0 * sigma * sigma)).exp()
}
impl System {
pub fn new(width: u32, height: u32) -> Self {
let field = Array2::zeros((height as usize, width as usize));
let ex = field.clone();
let ey = field.clone();
let hz = field;
Self { ex, ey, hz }
}
pub fn set_gaussian(&mut self, x0: f32, y0: f32) {
let nx = self.ex.shape()[1];
let ny = self.ex.shape()[0];
for j in 0..ny {
for i in 0..nx {
// Must divice interval on nx/ny instead of nx - 1/ny-1
// due to periodic conditions [0, 1)
let x = i as f32 / nx as f32;
let y = j as f32 / ny as f32;
self.ex[(j, i)] = 0.0;
self.ey[(j, i)] = 0.0;
self.hz[(j, i)] = gaussian(x, x0, y, y0) / 32.0;
}
}
}
pub fn advance(&self, fut: &mut System, dt: f32, work_buffers: Option<&mut WorkBuffers>) {
assert_eq!(self.ex.shape(), fut.ex.shape());
let mut wb: WorkBuffers;
let (y, k) = match work_buffers {
Some(x) => (&mut x.y, &mut x.buf),
None => {
wb = WorkBuffers::new(self.ex.shape()[1], self.ex.shape()[0]);
(&mut wb.y, &mut wb.buf)
}
};
for i in 0..4 {
// y = y0 + c*kn
y.0.assign(&self.ex);
y.1.assign(&self.hz);
y.2.assign(&self.ey);
match i {
0 => {}
1 => {
y.0.scaled_add(1.0 / 2.0 * dt, &k[i - 1].0);
y.1.scaled_add(1.0 / 2.0 * dt, &k[i - 1].1);
y.2.scaled_add(1.0 / 2.0 * dt, &k[i - 1].2);
}
2 => {
y.0.scaled_add(1.0 / 2.0 * dt, &k[i - 1].0);
y.1.scaled_add(1.0 / 2.0 * dt, &k[i - 1].1);
y.2.scaled_add(1.0 / 2.0 * dt, &k[i - 1].2);
}
3 => {
y.0.scaled_add(dt, &k[i - 1].0);
y.1.scaled_add(dt, &k[i - 1].1);
y.2.scaled_add(dt, &k[i - 1].2);
}
_ => {
unreachable!();
}
};
k[i].0.fill(0.0);
// ex = hz_y
diffy_periodic(y.1.view(), k[i].0.view_mut());
// ey = -hz_x
k[i].2.fill(0.0);
diffx_periodic(y.1.view(), k[i].2.view_mut());
k[i].2.mapv_inplace(|v| -v);
// hz = -ey_x + ex_y
k[i].1.fill(0.0);
diffx_periodic(y.2.view(), k[i].1.view_mut());
k[i].1.mapv_inplace(|v| -v);
diffy_periodic(y.0.view(), k[i].1.view_mut());
}
Zip::from(&mut fut.ex)
.and(&self.ex)
.and(&k[0].0)
.and(&k[1].0)
.and(&k[2].0)
.and(&k[3].0)
.apply(|y1, &y0, &k1, &k2, &k3, &k4| {
*y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
});
Zip::from(&mut fut.hz)
.and(&self.hz)
.and(&k[0].1)
.and(&k[1].1)
.and(&k[2].1)
.and(&k[3].1)
.apply(|y1, &y0, &k1, &k2, &k3, &k4| {
*y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
});
Zip::from(&mut fut.ey)
.and(&self.ey)
.and(&k[0].2)
.and(&k[1].2)
.and(&k[2].2)
.and(&k[3].2)
.apply(|y1, &y0, &k1, &k2, &k3, &k4| {
*y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
});
}
}
pub struct WorkBuffers {
y: (Array2<f32>, Array2<f32>, Array2<f32>),
buf: [(Array2<f32>, Array2<f32>, Array2<f32>); 4],
}
impl WorkBuffers {
pub fn new(nx: usize, ny: usize) -> Self {
let arr = Array2::zeros((ny, nx));
Self {
y: (arr.clone(), arr.clone(), arr.clone()),
buf: [
(arr.clone(), arr.clone(), arr.clone()),
(arr.clone(), arr.clone(), arr.clone()),
(arr.clone(), arr.clone(), arr.clone()),
(arr.clone(), arr.clone(), arr),
],
}
}
}