fdtd-coremem/src/driver.rs

use crate::{flt::Flt, mat};
use crate::consts;
use crate::geom::{Coord, Index, Region, Vec3, Vec3u};
use crate::mat::{Material, GenericMaterial};
use crate::meas::{self, AbstractMeasurement};
use crate::render::{self, MultiRenderer, Renderer};
use crate::sim::{GenericSim, SimState};
use crate::stim::AbstractStimulus;

use log::{info, debug, trace};
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{sync_channel, SyncSender, Receiver};
use std::time::{Duration, Instant};
use threadpool::ThreadPool;

pub struct Driver<M=GenericMaterial> {
    pub state: SimState<M>,
    renderer: Arc<MultiRenderer>,
    render_pool: ThreadPool,
    render_channel: (SyncSender<()>, Receiver<()>),
    steps_per_frame: u64,
    time_spent_stepping: Duration,
    time_spent_on_stimuli: Duration,
    time_spent_blocked_on_render: Duration,
    time_spent_rendering: Arc<Mutex<Duration>>,
    measurements: Vec<Box<dyn AbstractMeasurement>>,
    stimuli: Vec<Box<dyn AbstractStimulus>>,
    start_time: Instant,
    last_diag_time: Instant,
}

impl<M: Default> Driver<M> {
    pub fn new<C: Coord>(size: C, feature_size: Flt) -> Self {
        Driver {
            state: SimState::new(size.to_index(feature_size), feature_size),
            renderer: Arc::new(MultiRenderer::new()),
            render_pool: ThreadPool::new(3),
            render_channel: sync_channel(0),
            steps_per_frame: 1,
            time_spent_stepping: Default::default(),
            time_spent_on_stimuli: Default::default(),
            time_spent_blocked_on_render: Default::default(),
            time_spent_rendering: Default::default(),
            measurements: vec![Box::new(meas::Time), Box::new(meas::Meta), Box::new(meas::Energy)],
            stimuli: vec![],
            start_time: Instant::now(),
            last_diag_time: Instant::now(),
        }
    }
}

impl<M: Material + Clone + Send + Sync + 'static> Driver<M> {
    pub fn dyn_state(&mut self) -> &mut dyn GenericSim {
        &mut self.state
    }
}

impl<M: Material> Driver<M> {
    pub fn add_stimulus<S: AbstractStimulus + 'static>(&mut self, s: S) {
        self.stimuli.push(Box::new(s))
    }

    pub fn add_measurement<Meas: AbstractMeasurement + 'static>(&mut self, m: Meas) {
        self.measurements.push(Box::new(m));
    }

    pub fn set_steps_per_frame(&mut self, steps_per_frame: u64) {
        self.steps_per_frame = steps_per_frame;
    }

    fn add_renderer<R: Renderer + 'static>(&mut self, renderer: R, name: &str) {
        info!("render to {}", name);
        self.renderer.push(renderer);
    }

    pub fn add_y4m_renderer<S: Into<PathBuf>>(&mut self, output: S) {
        let output = output.into();
        let name = output.to_string_lossy().into_owned();
        self.add_renderer(render::Y4MRenderer::new(output), &*name);
    }

    pub fn add_plotly_renderer(&mut self, out_base: &str) {
        self.add_renderer(render::PlotlyRenderer::new(out_base), out_base);
    }

    pub fn add_term_renderer(&mut self) {
        self.add_renderer(render::ColorTermRenderer, "terminal");
    }

    pub fn add_serializer_renderer(&mut self, out_base: &str) {
        self.add_renderer(render::SerializerRenderer::new(out_base), out_base);
    }
}

impl<M: Material + Clone> Driver<M> {
    pub fn fill_region<R: Region>(&mut self, region: &R, mat: M) {
        for z in 0..self.state.depth() {
            for y in 0..self.state.height() {
                for x in 0..self.state.width() { 
                    let loc = Index((x, y, z).into());
                    let meters = loc.to_meters(self.state.feature_size());
                    if region.contains(meters) {
                        *self.state.get_mut(loc).mat_mut() = mat.clone();
                    }
                }
            }
        }
    }
}

impl<M: Material + Clone + Default + Send + Sync + 'static> Driver<M> {
    fn render(&mut self) {
        let their_state = self.state.clone();
        let their_measurements = self.measurements.clone();
        let renderer = self.renderer.clone();
        let time_spent_rendering = self.time_spent_rendering.clone();
        let sender = self.render_channel.0.clone();
        self.render_pool.execute(move || {
            sender.send(()).unwrap();
            trace!("render begin");
            let start_time = Instant::now();
            renderer.render(&their_state, &*their_measurements);
            *time_spent_rendering.lock().unwrap() += start_time.elapsed();
            trace!("render end");
        });
        let block_start = Instant::now();
        self.render_channel.1.recv().unwrap();
        self.time_spent_blocked_on_render += block_start.elapsed();
    }
    pub fn step(&mut self) {
        if self.state.step_no() % self.steps_per_frame == 0 {
            self.render();
        }
        {
            trace!("stimuli begin");
            let start_time = Instant::now();
            self.state.apply_stimulus(&self.stimuli);
            self.time_spent_on_stimuli += start_time.elapsed();
        }

        trace!("step begin");
        let start_time = Instant::now();
        self.state.step();
        self.time_spent_stepping += start_time.elapsed();
        trace!("step end");
        if self.last_diag_time.elapsed().as_secs_f64() >= 5.0 {
            self.last_diag_time = Instant::now();
            let step = self.state.step_no();
            let step_time = self.time_spent_stepping.as_secs_f64();
            let stim_time = self.time_spent_on_stimuli.as_secs_f64();
            let render_time = self.time_spent_rendering.lock().unwrap().as_secs_f64();
            let block_time = self.time_spent_blocked_on_render.as_secs_f64();
            let overall_time = self.start_time.elapsed().as_secs_f64();
            let fps = (self.state.step_no() as f64) / overall_time;
            let sim_time = self.state.time() as f64;
            info!(
                "t={:.2e} frame {:06} fps: {:6.2} (sim: {:.1}s, stim: {:.1}s, render: {:.1}s, blocked: {:.1}s, other: {:.1}s)",
                sim_time,
                step,
                fps,
                step_time,
                stim_time,
                render_time,
                block_time,
                overall_time - step_time - stim_time - render_time
            );
        }
    }

    pub fn step_until(&mut self, deadline: Flt) {
        while self.dyn_state().time() < deadline {
            self.step();
        }
    }
}

impl<M: Material + From<mat::Static>> Driver<M> {
    // pub fn add_boundary(&mut self, thickness: u32, base_conductivity: Flt) {
    //     for inset in 0..thickness {
    //         let depth = thickness - inset;
    //         let conductivity = base_conductivity * (depth*depth) as Flt;
    //         for x in inset..self.state.width() - inset {
    //             // left
    //             *self.state.get_mut((x, inset).into()).mat_mut() = mat::Static::conductor(conductivity).into();
    //             // right
    //             *self.state.get_mut((x, self.state.height() - 1 - inset).into()).mat_mut() = mat::Static::conductor(conductivity).into();
    //         }
    //         for y in inset..self.state.height() - inset {
    //             // top
    //             *self.state.get_mut((inset, y).into()).mat_mut() = mat::Static::conductor(conductivity).into();
    //             // bottom
    //             *self.state.get_mut((self.state.width() - 1 - inset, y).into()).mat_mut() = mat::Static::conductor(conductivity).into();
    //         }
    //     }
    // }

    pub fn add_upml_boundary<C: Coord>(&mut self, thickness: C) {
        // Based on explanation here (slide 63): https://empossible.net/wp-content/uploads/2020/01/Lecture-The-Perfectly-Matched-Layer.pdf
        let thickness = thickness.to_index(self.state.feature_size());
        let d = self.state.depth();
        let h = self.state.height();
        let w = self.state.width();
        for z in 0..d {
            let depth_z = if z < thickness.z() {
                thickness.z() - z
            } else if z >= d - thickness.z() {
                1 + z - (d - thickness.z())
            } else {
                0
            };
            for y in 0..h {
                let depth_y = if y < thickness.y() {
                    thickness.y() - y
                } else if y >= h - thickness.y() {
                    1 + y - (h - thickness.y())
                } else {
                    0
                };
                for x in 0..w {
                    let depth_x = if x < thickness.x() {
                        thickness.x() - x
                    } else if x >= w - thickness.x() {
                        1 + x - (w - thickness.x())
                    } else {
                        0
                    };

                    if depth_x > 0 || depth_y > 0 || depth_z > 0 {
                        let scale = 0.5 * consts::EPS0 / self.state.timestep();
                        // let scale = 1e3;
                        let cond_x = if thickness.x() != 0 {
                            scale * (depth_x as Flt/thickness.x() as Flt).powf(3.0)
                        } else {
                            0.0
                        };
                        let cond_y = if thickness.y() != 0 {
                            scale * (depth_y as Flt/thickness.y() as Flt).powf(3.0)
                        } else {
                            0.0
                        };
                        let cond_z = if thickness.z() != 0 {
                            scale * (depth_z as Flt/thickness.z() as Flt).powf(3.0)
                        } else {
                            0.0
                        };
                        let cond = Vec3::new(cond_x, cond_y, cond_z);
                        // trace!("cond: {:?}", cond);
                        let conductor = mat::Static::anisotropic_conductor(cond);
                        *self.state.get_mut(Index(Vec3u::new(x, y, z))).mat_mut() = conductor.into();
                    }
                }
            }
        }
    }
}