separate maxwell solver

main.js

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

src/lib.rs

@@ -1,175 +1,52 @@
-use ndarray::{Array2, Zip};
 use wasm_bindgen::prelude::*;
 
+mod maxwell;
 mod operators;
-use operators::{diffx_periodic, diffy_periodic};
+use maxwell::{System, WorkBuffers};
 
 #[cfg(feature = "wee_alloc")]
 #[global_allocator]
 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]
 pub fn set_panic_hook() {
     #[cfg(feature = "console_error_panic_hook")]
     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]
 impl Universe {
     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 init(&mut self, x0: f32, y0: f32) {
+        self.sys.0.set_gaussian(x0, y0);
     }
 
-    pub fn advance(&self, fut: &mut Universe, dt: f32, work_buffers: Option<WorkBuffers>) {
+    pub fn advance(&mut self, dt: f32) {
-        assert_eq!(self.ex.shape(), fut.ex.shape());
+        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 buffers = work_buffers
-            .unwrap_or_else(|| WorkBuffers::new(self.ex.shape()[1], self.ex.shape()[0]));
-
-        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 {
-        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 {
-        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 {
-        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),
-            ],
-        }
     }
 }

src/maxwell.rs

@@ -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),
+            ],
+        }
+    }
+}