fdtd-coremem/crates/coremem/src/driver.rs

use crate::geom::{Coord, Index, Meters, Region, Vec3};
use crate::mat::{self, Pml};
use crate::meas::{self, AbstractMeasurement};
use crate::real::{self, Real};
use crate::render::{self, MultiRenderer, Renderer};
use crate::sim::{GenericSim, MaterialSim, SampleableSim, SimState};
use crate::sim::units::{Frame, Time};
use crate::sim::spirv::{self, SpirvSim};
use crate::stim::AbstractStimulus;

use log::{info, trace};
use serde::{Deserialize, Serialize};
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<S=SimState> {
    pub state: S,
    renderer: Arc<MultiRenderer<S>>,
    // TODO: use Rayon's thread pool?
    render_pool: ThreadPool,
    render_channel: (SyncSender<()>, Receiver<()>),
    time_spent_stepping: Duration,
    time_spent_on_stimuli: Duration,
    time_spent_prepping_render: Duration,
    time_spent_blocked_on_render: Duration,
    time_spent_rendering: Arc<Mutex<Duration>>,
    measurements: Vec<Box<dyn AbstractMeasurement>>,
    stimuli: StimuliAdapter,
    start_time: Instant,
    last_diag_time: Instant,
    /// simulation end time
    sim_end_time: Option<Frame>,
}

pub type CpuDriver<R=real::R32, M=mat::GenericMaterial<R>> = Driver<SimState<R, M>>;
pub type SpirvDriver<M=spirv::FullyGenericMaterial> = Driver<SpirvSim<M>>;

impl<R: Real, M: Default> Driver<SimState<R, M>> {
    pub fn new<C: Coord>(size: C, feature_size: f32) -> Self {
        Self::new_with_state(SimState::new(size.to_index(feature_size), feature_size))
    }
}

impl<M: spirv::IntoFfi> SpirvDriver<M>
    where M::Ffi: Default + 'static
{
    pub fn new_spirv<C: Coord>(size: C, feature_size: f32) -> Self {
        Self::new_with_state(SpirvSim::new(size.to_index(feature_size), feature_size))
    }
}

impl<S> Driver<S> {
    pub fn new_with_state(state: S) -> Self {
        Self {
            state,
            renderer: Arc::new(MultiRenderer::new()),
            render_pool: ThreadPool::new(3),
            render_channel: sync_channel(0),
            time_spent_stepping: Default::default(),
            time_spent_on_stimuli: Default::default(),
            time_spent_prepping_render: 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::world()),
                Box::new(meas::Power::world()),
            ],
            stimuli: StimuliAdapter::new(),
            start_time: Instant::now(),
            last_diag_time: Instant::now(),
            sim_end_time: None,
        }
    }
}

impl<S> Driver<S> {
    pub fn add_stimulus<Stim: AbstractStimulus + 'static>(&mut self, s: Stim) {
        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_stim(&mut self, steps_per_stim: u64) {
        self.stimuli.frame_interval = steps_per_stim;
    }
}

impl<S: MaterialSim> Driver<S> {
    pub fn fill_region<Reg: Region, M: Into<S::Material> + Clone>(&mut self, region: &Reg, mat: M) {
        self.state.fill_region(region, mat);
    }
    pub fn test_region_filled<Reg: Region, M: Into<S::Material> + Clone>(&mut self, region: &Reg, mat: M) -> bool {
        self.state.test_region_filled(region, mat)
    }
}

impl<S: SampleableSim> Driver<S> {
    pub fn size(&self) -> Index {
        self.state.size()
    }
    pub fn timestep(&self) -> f32 {
        self.state.timestep()
    }
    pub fn time(&self) -> f32 {
        self.state.time()
    }
}

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

    fn add_renderer<Rend: Renderer<S> + 'static>(
        &mut self, renderer: Rend, name: &str, step_frequency: u64, frame_limit: Option<u64>
    ) {
        info!("render to {} at f={} until f={:?}", name, step_frequency, frame_limit);
        self.renderer.push(renderer, step_frequency, frame_limit);
    }

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

    pub fn add_term_renderer(&mut self, step_frequency: u64, frame_limit: Option<u64>) {
        self.add_renderer(render::ColorTermRenderer, "terminal", step_frequency, frame_limit);
    }

    pub fn add_csv_renderer(&mut self, path: &str, step_frequency: u64, frame_limit: Option<u64>) {
        self.add_renderer(render::CsvRenderer::new(path), path, step_frequency, frame_limit);
    }
}

impl<S: SampleableSim + Send + Sync + Serialize + 'static> Driver<S> {
    pub fn add_serializer_renderer(&mut self, out_base: &str, step_frequency: u64, frame_limit: Option<u64>) {
        let fmt_str = format!("{out_base}{{step_no}}.bc", out_base=out_base);
        self.add_renderer(render::SerializerRenderer::new_static(&*fmt_str), &*fmt_str, step_frequency, frame_limit);
    }
}

impl<S: SampleableSim + Send + Sync + Serialize + for<'a> Deserialize<'a> + 'static> Driver<S> {
    /// instruct the driver to periodically save the simulation state to the provided path.
    /// also attempts to load an existing state file, returning `true` on success.
    pub fn add_state_file(&mut self, state_file: &str, snapshot_frequency: u64) -> bool {
        let ser = render::SerializerRenderer::new(state_file);
        let loaded = match ser.try_load() {
            Some(state) => {
                self.state = state.state;
                true
            },
            None => false,
        };
        self.add_renderer(ser, state_file, snapshot_frequency, None);
        loaded
    }
}

impl<S: GenericSim + Clone + Default + Send + Sync + 'static> Driver<S> {
    fn render(&mut self) {
        let prep_start = Instant::now();
        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 || {
            // unblock the main thread (this limits the number of renders in-flight at any time
            sender.send(()).unwrap();
            trace!("render begin");
            let start_time = Instant::now();
            renderer.render(&their_state, &*their_measurements, Default::default());
            *time_spent_rendering.lock().unwrap() += start_time.elapsed();
            trace!("render end");
        });
        self.time_spent_prepping_render += prep_start.elapsed();
        let block_start = Instant::now();
        self.render_channel.1.recv().unwrap();
        self.time_spent_blocked_on_render += block_start.elapsed();
    }
    /// Return the number of steps actually stepped
    fn step_at_most(&mut self, at_most: u32) -> u32 {
        assert!(at_most != 0);
        let start_step = self.state.step_no();
        if self.stimuli.should_apply(start_step) {
            self.stimuli.real_time = self.state.time();
            self.stimuli.time_step = self.state.timestep();
            trace!("updating stimuli");
        }

        if self.renderer.any_work_for_frame(start_step) {
            self.render();
        }

        let mut can_step = 1;
        while can_step < at_most && !self.renderer.any_work_for_frame(start_step + can_step as u64) {
            can_step += 1;
        }
        trace!("step begin");
        let start_time = Instant::now();
        self.state.step_multiple(can_step, &self.stimuli);
        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 render_prep_time = self.time_spent_prepping_render.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;
            let percent_complete = match self.sim_end_time {
                Some(t) => format!("[{:.1}%] ", 100.0 * self.state.time() / *t.to_seconds(self.timestep())),
                None => "".to_owned(),
            };
            info!(
                "{}t={:.2e} frame {:06} fps: {:6.2} (sim: {:.1}s, stim: {:.1}s, [render: {:.1}s], blocked: {:.1}s, render_prep: {:.1}s, other: {:.1}s)",
                percent_complete,
                sim_time,
                step,
                fps,
                step_time,
                stim_time,
                render_time,
                block_time,
                render_prep_time,
                overall_time - step_time - stim_time - block_time - render_prep_time,
            );
        }
        can_step as u32
    }
    pub fn step_multiple(&mut self, num_steps: u32) {
        let mut steps_remaining = num_steps;
        while steps_remaining != 0 {
            steps_remaining -= self.step_at_most(steps_remaining);
        }
    }
    pub fn step(&mut self) {
        self.step_multiple(1);
    }

    /// Returns the number of timesteps needed to reach the end time
    pub fn steps_until<T: Time>(&mut self, sim_end_time: T) -> u64 {
        let sim_end_step = sim_end_time.to_frame(self.state.timestep());
        let start_step = self.state.step_no();
        sim_end_step.saturating_sub(start_step)
    }

    pub fn step_until<T: Time>(&mut self, sim_end_time: T) {
        let sim_end_time = sim_end_time.to_frame(self.state.timestep());
        self.sim_end_time = Some(sim_end_time);
        let mut stepped = false;
        while self.dyn_state().step_no() < *sim_end_time {
            self.step_multiple(100);
            stepped = true;
        }
        if stepped {
            // render the final frame -- unless we already *have*
            self.render();
        }
        self.render_pool.join();
        self.sim_end_time = None;
    }
}

impl<S: MaterialSim> Driver<S> {
    pub fn add_pml_boundary<C: Coord, R: Real>(&mut self, thickness: C)
        where S::Material: From<Pml<R>>
    {
        let timestep = self.state.timestep();
        self.state.fill_boundary_using(thickness, |boundary_ness| {
            let b = boundary_ness.elem_pow(3.0);
            let conductivity = b * (0.5 / timestep);
            Pml::new(conductivity)
        });
    }

    pub fn add_classical_boundary<C: Coord>(&mut self, thickness: C)
        where S::Material: From<mat::IsomorphicConductor<f32>>
    {
        self.add_classical_boundary_explicit::<f32, _>(thickness)
    }

    /// the CPU code is parameterized over `Real`: you'll need to use this interface to get access
    /// to that, if using a CPU driver. otherwise, use `add_classical_boundary`
    pub fn add_classical_boundary_explicit<R: Real, C: Coord>(&mut self, thickness: C)
        where S::Material: From<mat::IsomorphicConductor<R>>
    {
        let timestep = self.state.timestep();
        self.state.fill_boundary_using(thickness, |boundary_ness| {
            let b = boundary_ness.elem_pow(3.0);
            let cond = b * (0.5 / timestep);

            let iso_cond = cond.x() + cond.y() + cond.z();
            let iso_conductor = mat::IsomorphicConductor::new(iso_cond.cast());
            iso_conductor
        });
    }
}

/// Adapts the stimuli to be applied only every so often, to improve perf
struct StimuliAdapter {
    stim: Vec<Box<dyn AbstractStimulus>>,
    /// How many frames to go between applications of the stimulus.
    frame_interval: u64,
    real_time: f32,
    time_step: f32,
}

impl AbstractStimulus for StimuliAdapter {
    fn at(&self, t_sec: f32, pos: Meters) -> (Vec3<f32>, Vec3<f32>) {
        self.stim.at(t_sec, pos)
        // TODO: remove this stuff  (here only for testing)
        /*
        if true {
            // interpolation unaware (i.e. let the Sim backend do it)
        } else if false {
            // delta-fn "interpolation"
            self.stim.at(t_sec, pos) * (self.frame_interval as f32)
        } else if false {
            // step-fn "interpolation"
            self.stim.at(self.real_time, pos)
        } else {
            // linear interpolation
            let interp_width = self.frame_interval as f32 * self.time_step;
            let prev = self.stim.at(self.real_time, pos);
            let next = self.stim.at(self.real_time + interp_width, pos);
            let interp = (t_sec - self.real_time) / interp_width;
            prev * (1.0 - interp) + next * interp
        }
        */
    }
}

impl StimuliAdapter {
    fn new() -> Self {
        Self {
            stim: Default::default(),
            frame_interval: 1,
            real_time: 0.0,
            time_step: 0.0,
        }
    }
    fn should_apply(&self, frame: u64) -> bool {
        (frame % self.frame_interval == 0) && self.stim.len() != 0
    }
    fn push(&mut self, s: Box<dyn AbstractStimulus>) {
        self.stim.push(s)
    }
}