use crate::geom::{Index, Meters, Vec2, Vec3, Vec3u}; use crate::real::ToFloat as _; use crate::sim::{GenericSim, Sample, StaticSim}; use crate::meas::AbstractMeasurement; use crossterm::{cursor, QueueableCommand as _}; use crossterm::style::{style, Color, PrintStyledContent}; use font8x8::{BASIC_FONTS, GREEK_FONTS, UnicodeFonts as _}; use log::trace; use num::integer::Integer; use plotly; use image::{RgbImage, Rgb}; use imageproc::{pixelops, drawing}; use rayon::prelude::*; use serde::{Serialize, Deserialize}; use std::fs::File; use std::io::{BufWriter, Write as _}; use std::path::PathBuf; use std::sync::{Mutex, RwLock}; use y4m; /// Accept a value from (-\inf, \inf) and return a value in (-1, 1). /// If the input is equal to `typical`, it will be mapped to 0.5. /// If the input is equal to -`typical`, it will be mapped to -0.5. fn scale_signed(x: f32, typical: f32) -> f32 { if x >= 0.0 { scale_unsigned(x, typical) } else { -scale_unsigned(-x, typical) } } /// Accept a value from [0, \inf) and return a value in [0, 1). /// If the input is equal to `typical`, it will be mapped to 0.5. fn scale_unsigned(x: f32, typical: f32) -> f32 { // f(0) => 0 // f(1) => 0.5 // f(\inf) => 1 // f(x) = 1 - 1/(x+1) 1.0 - 1.0/(x/typical + 1.0) } fn scale_signed_to_u8(x: f32, typ: f32) -> u8 { let norm = 128.0 + 128.0*scale_signed(x, typ); norm as _ } fn scale_unsigned_to_u8(x: f32, typ: f32) -> u8 { let norm = 256.0*scale_unsigned(x, typ); norm as _ } /// Scale a vector to have magnitude between [0, 1). fn scale_vector(x: Vec2, typical_mag: f32) -> Vec2 { let new_mag = scale_unsigned(x.mag(), typical_mag); x.with_mag(new_mag) } fn im_size(state: &S, max_w: u32, max_h: u32) -> (u32, u32) { let mut width = max_w; let mut height = width * state.height() / state.width(); if height > max_h { let stretch = max_h as f32 / height as f32; width = (width as f32 * stretch).round() as _; height = max_h; } (width, height) } #[derive(Copy, Clone, Debug, Eq, PartialEq)] pub enum FieldDisplayMode { BzExy, EzBxy, BCurrent, M, } impl FieldDisplayMode { pub fn next(self) -> Self { use FieldDisplayMode::*; match self { BzExy => EzBxy, EzBxy => BCurrent, BCurrent => M, M => BzExy, } } pub fn prev(self) -> Self { use FieldDisplayMode::*; match self { BzExy => M, EzBxy => BzExy, BCurrent => EzBxy, M => BCurrent, } } } #[derive(Copy, Clone, PartialEq)] pub struct RenderConfig { mode: FieldDisplayMode, scale: f32, } impl Default for RenderConfig { fn default() -> Self { Self { mode: FieldDisplayMode::BzExy, scale: 1.0, } } } impl RenderConfig { pub fn next_mode(&mut self) { self.mode = self.mode.next(); } pub fn prev_mode(&mut self) { self.mode = self.mode.prev(); } pub fn increase_range(&mut self) { self.scale *= 2.0; } pub fn decrease_range(&mut self) { self.scale *= 0.5; } } struct RenderSteps<'a, S> { im: RgbImage, sim: &'a S, meas: &'a [Box], /// Simulation z coordinate to sample z: u32, } impl<'a, S: GenericSim> RenderSteps<'a, S> { /// Render using default configuration constants fn render(state: &'a S, measurements: &'a [Box], z: u32) -> RgbImage { Self::render_configured(state, measurements, z, (640, 480), RenderConfig::default()) } /// Render, controlling things like the size. fn render_configured( state: &'a S, measurements: &'a [Box], z: u32, max_size: (u32, u32), config: RenderConfig, ) -> RgbImage { let (width, height) = im_size(state, max_size.0, max_size.1); trace!("rendering at {}x{} with z={}", width, height, z); let mut me = Self::new(state, measurements, width, height, z); me.render_scalar_field(10.0, false, 2, |cell| { let is_vacuum = cell.conductivity() != Vec3::zero() || cell.m() != Vec3::zero(); cell.conductivity().mag().to_f32() + if is_vacuum { 0.0 } else { 5.0 } }); match config.mode { FieldDisplayMode::BzExy => { me.render_b_z_field(config.scale); me.render_e_xy_field(config.scale); }, FieldDisplayMode::EzBxy => { me.render_e_z_field(config.scale); me.render_b_xy_field(config.scale); }, FieldDisplayMode::BCurrent => { me.render_b(config.scale); me.render_current(config.scale); } FieldDisplayMode::M => { me.render_m(config.scale); } } me.render_measurements(); me.im } fn new(sim: &'a S, meas: &'a [Box], width: u32, height: u32, z: u32) -> Self { RenderSteps { im: RgbImage::new(width, height), sim, meas, z } } fn get_at_px(&self, x_px: u32, y_px: u32) -> Sample { let x_prop = x_px as f32 / self.im.width() as f32; let x_m = x_prop * (self.sim.width() as f32 * self.sim.feature_size() as f32); let y_prop = y_px as f32 / self.im.height() as f32; let y_m = y_prop * (self.sim.height() as f32 * self.sim.feature_size() as f32); let z_m = self.z as f32 * self.sim.feature_size() as f32; self.sim.sample(Meters(Vec3::new(x_m, y_m, z_m))) } ////////////// Ex/Ey/Bz configuration //////////// fn render_b_z_field(&mut self, scale: f32) { self.render_scalar_field(1.0e-4 * scale, true, 1, |cell| cell.b().z().to_f32()); } fn render_e_xy_field(&mut self, scale: f32) { self.render_vector_field(Rgb([0xff, 0xff, 0xff]), 100.0 * scale, |cell| cell.e().xy().to_f32()); // current self.render_vector_field(Rgb([0x00, 0xa0, 0x30]), 1.0e-12 * scale, |cell| { cell.e().elem_mul(cell.conductivity()).xy().to_f32() }); } ////////////// Magnitude configuration ///////////// fn render_b(&mut self, scale: f32) { self.render_scalar_field(1.0e-3 * scale, false, 1, |cell| cell.b().mag().to_f32()); } fn render_current(&mut self, scale: f32) { self.render_scalar_field(1.0e1 * scale, false, 0, |cell| { cell.e().elem_mul(cell.conductivity()).mag().to_f32() }); } ////////////// Bx/By/Ez configuration //////////// fn render_e_z_field(&mut self, scale: f32) { self.render_scalar_field(1e4 * scale, true, 1, |cell| cell.e().z().to_f32()); } fn render_b_xy_field(&mut self, scale: f32) { self.render_vector_field(Rgb([0xff, 0xff, 0xff]), 1.0e-9 * scale, |cell| cell.b().xy().to_f32()); } fn render_m(&mut self, scale: f32) { self.render_scalar_field(1.0e5 * scale, false, 1, |cell| cell.m().mag().to_f32()); self.render_vector_field(Rgb([0xff, 0xff, 0xff]), 1.0e5 * scale, |cell| cell.m().xy().to_f32()); } fn render_vector_field Vec2>(&mut self, color: Rgb, typical: f32, measure: F) { let w = self.im.width(); let h = self.im.height(); let vec_spacing = 10; for y in (0..h).into_iter().step_by(vec_spacing as _) { for x in (0..w).into_iter().step_by(vec_spacing as _) { let vec = self.field_vector(x, y, vec_spacing, &measure); let norm_vec = scale_vector(vec, typical); let alpha = 0.7*scale_unsigned(vec.mag_sq(), typical * 5.0); let vec = norm_vec * (vec_spacing as f32); let center = Vec2::new(x as f32, y as f32) + Vec2::new(vec_spacing as f32, vec_spacing as f32)*0.5; self.im.draw_field_arrow(center, vec, color, alpha as f32); } } } fn render_scalar_field f32 + Sync>(&mut self, typical: f32, signed: bool, slot: u32, measure: F) { // XXX: get_at_px borrows self, so we need to clone the image to operate on it mutably. let mut im = self.im.clone(); let w = im.width(); let h = im.height(); let samples = im.as_flat_samples_mut(); assert_eq!(samples.layout.channel_stride, 1); assert_eq!(samples.layout.width_stride, 3); assert_eq!(samples.layout.height_stride, 3*w as usize); let pixel_buf: &mut [[u8; 3]] = unsafe { std::mem::transmute(samples.samples) }; pixel_buf[..(w*h) as usize].par_iter_mut().enumerate().for_each(|(idx, px)| { let (y, x) = (idx as u32).div_rem(&w); let cell = self.get_at_px(x, y); let value = measure(&cell); let scaled = if signed { scale_signed_to_u8(value, typical) } else { scale_unsigned_to_u8(value, typical) }; px[slot as usize] = scaled; }); self.im = im; } fn render_measurements(&mut self) { for (meas_no, m) in self.meas.iter().enumerate() { let meas_string = m.eval(self.sim); for (i, c) in meas_string.chars().enumerate() { let glyph = BASIC_FONTS.get(c) .or_else(|| GREEK_FONTS.get(c)) .unwrap_or_else(|| BASIC_FONTS.get('?').unwrap()); for (y, bmp) in glyph.iter().enumerate() { for x in 0..8 { if (bmp & 1 << x) != 0 { let real_x = 2 + i as u32*8 + x; if let Some(real_y) = (y as u32 + self.im.height()).checked_sub(10 + meas_no as u32 * 8) { if real_x < self.im.width() { self.im.put_pixel(real_x, real_y, Rgb([0, 0, 0])); } } } } } } } } fn field_vector Vec2>(&self, xidx: u32, yidx: u32, size: u32, measure: &F) -> Vec2 { let mut field = Vec2::default(); let w = self.im.width(); let h = self.im.height(); let xstart = xidx.min(w); let ystart = yidx.min(h); let xend = (xstart + size).min(w); let yend = (ystart + size).min(h); for y in ystart..yend { for x in xstart..xend { field += measure(&self.get_at_px(x, y)); } } let xw = xend - xstart; let yw = yend - ystart; if xw == 0 || yw == 0 { // avoid division by zero Vec2::new(0.0, 0.0) } else { field * (1.0 / ((xw*yw) as f32)) } } } trait ImageRenderExt { fn draw_field_arrow(&mut self, center: Vec2, rel: Vec2, color: Rgb, alpha: f32); } impl ImageRenderExt for RgbImage { fn draw_field_arrow(&mut self, center: Vec2, rel: Vec2, color: Rgb, alpha: f32) { let start = (center - rel * 0.5).round(); let end = (center + rel * 0.5).round(); let i_start = (start.x().round() as _, start.y().round() as _); let i_end = (end.x().round() as _, end.y().round() as _); let interpolate_with_alpha = |left, right, left_weight| { pixelops::interpolate(left, right, left_weight*alpha) }; drawing::draw_antialiased_line_segment_mut(self, i_start, i_end, color, interpolate_with_alpha); if i_start != i_end && (0..self.width() as i32).contains(&i_end.0) && (0..self.height() as i32).contains(&i_end.1) { self.put_pixel(i_end.0 as _, i_end.1 as _, Rgb([0xff, 0, 0])); } //drawing::draw_line_segment_mut(self, i_start, i_end, color); } } pub trait Renderer: Send + Sync { fn render_z_slice(&self, state: &S, z: u32, measurements: &[Box], config: RenderConfig); // { // self.render_with_image(state, &RenderSteps::render(state, measurements, z), measurements); // } fn render(&self, state: &S, measurements: &[Box], config: RenderConfig); /// Not intended to be called directly by users; implement this if you want the image to be /// computed using default settings and you just manage where to display/save it. fn render_with_image(&self, state: &S, _im: &RgbImage, measurements: &[Box], config: RenderConfig) { self.render(state, measurements, config); } } fn default_render_z_slice>( me: &R, state: &S, z: u32, measurements: &[Box], config: RenderConfig, ) { me.render_with_image(state, &RenderSteps::render(state, measurements, z), measurements, config); } fn default_render>( me: &R, state: &S, measurements: &[Box], config: RenderConfig ) { me.render_z_slice(state, state.depth() / 2, measurements, config); } // pub struct NumericTermRenderer; // // impl Renderer for NumericTermRenderer { // fn render(&mut self, state: &SimSnapshot, _measurements: &[Box]) { // for y in 0..state.height() { // for x in 0..state.width() { // let cell = state.get((x, y).into()); // print!(" {:>10.1e}", cell.ex()); // } // print!("\n"); // for x in 0..state.width() { // let cell = state.get((x, y).into()); // print!("{:>10.1e} {:>10.1e}", cell.ey(), cell.bz()); // } // print!("\n"); // } // print!("\n"); // } // } #[derive(Default)] pub struct ColorTermRenderer; impl Renderer for ColorTermRenderer { fn render(&self, state: &S, measurements: &[Box], config: RenderConfig) { default_render(self, state, measurements, config) } fn render_z_slice( &self, state: &S, z: u32, measurements: &[Box], config: RenderConfig, ) { let (max_w, mut max_h) = crossterm::terminal::size().unwrap(); max_h = max_h.saturating_sub(2 + measurements.len() as u16); let im = RenderSteps::render_configured(state, &[], z, (max_w as _, max_h as _), config); let mut stdout = std::io::stdout(); // TODO: consider clearing line-by-line for less tearing? stdout.queue(crossterm::terminal::Clear(crossterm::terminal::ClearType::All)).unwrap(); stdout.queue(cursor::MoveTo(0, 0)).unwrap(); for row in im.rows() { for p in row { stdout.queue(PrintStyledContent(style(" ").on(Color::Rgb { r: p.0[0], g: p.0[1], b: p.0[2], }))).unwrap(); } stdout.queue(cursor::MoveDown(1)).unwrap(); stdout.queue(cursor::MoveToColumn(0)).unwrap(); } stdout.queue(PrintStyledContent(style(format!("fields: {:?} scale: {}", config.mode, config.scale)))).unwrap(); stdout.queue(cursor::MoveDown(1)).unwrap(); stdout.queue(cursor::MoveToColumn(1)).unwrap(); stdout.queue(PrintStyledContent(style(format!("z: {}", z)))).unwrap(); for m in measurements { // Measurements can be slow to compute stdout.flush().unwrap(); let meas_string = m.eval(state); stdout.queue(cursor::MoveDown(1)).unwrap(); stdout.queue(cursor::MoveToColumn(1)).unwrap(); stdout.queue(PrintStyledContent(style(meas_string))).unwrap(); } stdout.flush().unwrap(); } } pub struct Y4MRenderer { out_path: PathBuf, encoder: Mutex>>, } impl Y4MRenderer { pub fn new>(output: P) -> Self { Self { out_path: output.into(), encoder: Mutex::new(None), } } } impl Renderer for Y4MRenderer { fn render_z_slice(&self, state: &S, z: u32, measurements: &[Box], config: RenderConfig) { default_render_z_slice(self, state, z, measurements, config) } fn render(&self, state: &S, measurements: &[Box], config: RenderConfig) { default_render(self, state, measurements, config) } fn render_with_image(&self, _state: &S, im: &RgbImage, _meas: &[Box], _config: RenderConfig) { { let mut enc = self.encoder.lock().unwrap(); if enc.is_none() { let writer = File::create(&self.out_path).unwrap(); *enc = Some(y4m::encode(im.width() as usize, im.height() as usize, y4m::Ratio::new(30, 1)) .with_colorspace(y4m::Colorspace::C444) .write_header(writer) .unwrap() ); } } let mut pix_y = Vec::new(); let mut pix_u = Vec::new(); let mut pix_v = Vec::new(); for &Rgb([r, g, b]) in im.pixels() { let r = (r as f32) / (256.0); let g = (g as f32) / (256.0); let b = (b as f32) / (256.0); let y = (0.299 * r) + (0.587 * g) + (0.114 * b); let cb = 0.5 - (0.168_736 * r) - (0.331_264 * g) + (0.5 * b); let cr = 0.5 + (0.5 * r) - (0.418_688 * g) - (0.081_312 * b); pix_y.push((y * 256.0) as _); pix_u.push((cb * 256.0) as _); pix_v.push((cr * 256.0) as _); } let frame = y4m::Frame::new([&*pix_y, &*pix_u, &*pix_v], None); let mut lock = self.encoder.lock().unwrap(); let enc = lock.as_mut().unwrap(); trace!("write_frame begin"); let ret = enc.write_frame(&frame).unwrap(); trace!("write_frame end"); ret } } pub struct PlotlyRenderer { out_base: String, } fn add_scatter(plot: &mut plotly::Plot, xv: &mut Vec, yv: &mut Vec, zv: &mut Vec, colors: &mut Vec) { let xv = std::mem::replace(xv, Vec::new()); let yv = std::mem::replace(yv, Vec::new()); let zv = std::mem::replace(zv, Vec::new()); let colors = std::mem::replace(colors, Vec::new()); let scatter = plotly::Scatter::new3(xv, yv, zv) .mode(plotly::common::Mode::Markers) .marker(plotly::common::Marker::new() .opacity(0.01) //.size_array(sizes) //.opacity_array(opacities) .color_array(colors) ); plot.add_trace(scatter); } impl PlotlyRenderer { pub fn new(out_base: &str) -> Self { Self { out_base: out_base.into(), } } } impl Renderer for PlotlyRenderer { fn render_z_slice(&self, state: &S, z: u32, measurements: &[Box], config: RenderConfig) { default_render_z_slice(self, state, z, measurements, config) } fn render(&self, state: &S, _meas: &[Box], _config: RenderConfig) { use plotly::{ImageFormat, Plot, Rgba}; // use plotly::common::Marker; use plotly::layout::{AspectMode, Axis, Layout, LayoutScene}; let mut plot = Plot::new(); let scene = LayoutScene::new() .x_axis(Axis::new().range(vec![0, state.width() as i32])) .y_axis(Axis::new().range(vec![0, state.height() as i32])) .z_axis(Axis::new().range(vec![0, state.depth() as i32])) .aspect_mode(AspectMode::Cube); let layout = Layout::new() .scene(scene); plot.set_layout(layout); let mut xv = Vec::new(); let mut yv = Vec::new(); let mut zv = Vec::new(); // let mut opacities = Vec::new(); let mut colors = Vec::new(); for z in 0..state.depth() { if xv.len() >= 120000 { add_scatter(&mut plot, &mut xv, &mut yv, &mut zv, &mut colors); } for y in 0..state.height() { for x in 0..state.width() { // if x%5 == 0 || y%5 == 0 || z%5 == 0 { // continue; // } let cell = (state as &dyn GenericSim).get(Index(Vec3u::new(x, y, z))); xv.push(x); yv.push(y); zv.push(z); // opacities.push((cell.e().mag() * 0.1).min(1.0) as f64) let is_vacuum = cell.conductivity() == Vec3::zero() && cell.m() == Vec3::zero(); let mat = cell.conductivity().mag().to_f32() + if is_vacuum { 0.0 } else { 5.0 }; //let g = scale_unsigned_to_u8(mat, 10.0); //let r = scale_unsigned_to_u8(cell.m().mag(), 100.0); //let b = scale_unsigned_to_u8(cell.e().mag(), 1e2); let r = scale_unsigned_to_u8(cell.m().mag().to_f32(), 100.0); let g = scale_unsigned_to_u8(cell.e().mag().to_f32(), 1e2); let b = scale_unsigned_to_u8(mat, 10.0); let alpha = 1.0; colors.push(Rgba::new(r, g, b, alpha)); } } // let scatter = plotly::Scatter::new3(xv, yv, zv) // .mode(plotly::common::Mode::Markers) // .marker(plotly::common::Marker::new() // //.opacity(0.2) // //.size_array(sizes) // //.opacity_array(opacities) // .color_array(colors) // ); // plot.add_trace(scatter); } add_scatter(&mut plot, &mut xv, &mut yv, &mut zv, &mut colors); let name = format!("{}{}", self.out_base, state.step_no()); let (im_w, im_h) = im_size(state, 2048, 2048); plot.save(&*name, ImageFormat::PNG, im_w as _, im_h as _, 1.0); } } pub struct MultiRenderer { renderers: RwLock>>>, } impl Default for MultiRenderer { fn default() -> Self { Self { renderers: RwLock::new(Vec::new()) } } } impl MultiRenderer { pub fn new() -> Self { Default::default() } pub fn push + 'static>(&self, r: R) { self.renderers.write().unwrap().push(Box::new(r)); } pub fn with + 'static>(self, r: R) -> Self { self.push(r); self } } impl Renderer for MultiRenderer { fn render_z_slice(&self, state: &S, z: u32, measurements: &[Box], config: RenderConfig) { default_render_z_slice(self, state, z, measurements, config) } fn render(&self, state: &S, measurements: &[Box], config: RenderConfig) { if self.renderers.read().unwrap().len() != 0 { self.render_with_image(state, &RenderSteps::render(state, measurements, state.depth() / 2), measurements, config); } } fn render_with_image(&self, state: &S, im: &RgbImage, measurements: &[Box], config: RenderConfig) { for r in &*self.renderers.read().unwrap() { r.render_with_image(state, im, measurements, config); } } } #[derive(Serialize, Deserialize)] pub struct SerializedFrame { pub state: S, pub measurements: Vec>, } impl SerializedFrame { pub fn to_static(self) -> SerializedFrame { SerializedFrame { state: self.state.to_static(), measurements: self.measurements, } } } pub struct SerializerRenderer { out_base: String } impl SerializerRenderer { pub fn new(out_base: &str) -> Self { Self { out_base: out_base.into(), } } } impl Renderer for SerializerRenderer { fn render_z_slice(&self, state: &S, z: u32, measurements: &[Box], config: RenderConfig) { default_render_z_slice(self, state, z, measurements, config) } fn render(&self, state: &S, measurements: &[Box], _config: RenderConfig) { let frame = SerializedFrame { // N.B.: `to_static` isn't 100% necessary here, but it tends to result in a smaller file. state: state.to_static(), measurements: measurements.iter().cloned().collect(), }; let name = format!("{}{}.bc", self.out_base, frame.state.step_no()); let out = BufWriter::new(File::create(name).unwrap()); //serde_cbor::to_writer(out, &snap).unwrap(); bincode::serialize_into(out, &frame).unwrap(); } }