diff --git a/Cargo.lock b/Cargo.lock index efe9006..39abf15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,12 +2,24 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "autocfg" version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + [[package]] name = "bulirsch" version = "0.1.12" @@ -18,6 +30,27 @@ dependencies = [ "num-traits", ] +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "color_quant" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b" + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -49,6 +82,35 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "fdeflate" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c" +dependencies = [ + "simd-adler32", +] + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + +[[package]] +name = "gif" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ae047235e33e2829703574b54fdec96bfbad892062d97fed2f76022287de61b" +dependencies = [ + "color_quant", + "weezl", +] + [[package]] name = "heck" version = "0.5.0" @@ -66,8 +128,10 @@ name = "lcurve" version = "0.4.4" dependencies = [ "bulirsch", + "gif", "ndarray", "numpy", + "png", "pyo3", "rayon", "roche-rs", @@ -104,6 +168,16 @@ version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "ndarray" version = "0.16.1" @@ -168,6 +242,19 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "png" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61" +dependencies = [ + "bitflags", + "crc32fast", + "fdeflate", + "flate2", + "miniz_oxide", +] + [[package]] name = "portable-atomic" version = "1.13.1" @@ -357,6 +444,12 @@ dependencies = [ "zmij", ] +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + [[package]] name = "syn" version = "2.0.117" @@ -380,6 +473,12 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "weezl" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a28ac98ddc8b9274cb41bb4d9d4d5c425b6020c50c46f25559911905610b4a88" + [[package]] name = "zmij" version = "1.0.21" diff --git a/Cargo.toml b/Cargo.toml index a218ca1..b4ac71a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,8 @@ ndarray = "0.16.1" bulirsch = "0.1.12" numpy = "0.28.0" rayon = "1.10" +png = "0.18" +gif = "0.13" serde_json = "1.0.149" serde-pyobject = "0.8.0" roche-rs = "0.5.3" @@ -25,4 +27,4 @@ version = "0.28.2" features = ["abi3-py38"] [profile.release] -debug=true +debug=true \ No newline at end of file diff --git a/src/binary_model.rs b/src/binary_model.rs index cb48ffe..f2f6244 100644 --- a/src/binary_model.rs +++ b/src/binary_model.rs @@ -9,7 +9,9 @@ use crate::set_disc_continuum::{set_disc_continuum, set_edge_continuum}; use crate::set_disc_grid::{set_disc_edge_grid, set_disc_grid}; use crate::set_star_continuum::set_star_continuum; use crate::set_star_grid::set_star_grid; -use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1}; +use crate::visualize::{self, RenderOptions}; +use ndarray::{Array2, Array3}; +use numpy::{IntoPyArray, PyArray1, PyArray2, PyArray3, PyReadonlyArray1}; use pyo3::prelude::*; use pyo3::types::{PyDict, PyDictMethods}; use rayon::prelude::*; @@ -19,6 +21,7 @@ use roche::{self, Etype, Point, Star, disc_eclipse}; use serde_pyobject::from_pyobject; use std::collections::HashMap; use std::f64::consts::TAU; +use std::path::PathBuf; #[pyclass] pub struct LightCurve { @@ -395,6 +398,196 @@ impl BinaryModel { log_prob, }) } + + /// + /// Render the system at a single orbital phase, projected onto the + /// plane of the sky, in the style of the original LCURVE `visualise` + /// program. Returns the image as a (height, width, 3) uint8 numpy + /// array; if `path` is given the image is also written there as a PNG. + /// + /// Keyword arguments: + /// width: image width in pixels (default 800). + /// height: image height in pixels; derived from the aspect ratio + /// of the plot limits when omitted. + /// x1, x2, y1, y2: sky-plane plot limits in units of the binary + /// separation (defaults -2, 2, -2, 2). + /// reverse: white background with dark colours if True (default), + /// black background with bright colours if False. + /// sdob: swap the colours of the two stars (default False). + /// path: optional filename for a PNG copy of the frame. + /// + #[pyo3(signature = ( + phase, + *, + width=800, + height=None, + x1=-2.0, + x2=2.0, + y1=-2.0, + y2=2.0, + reverse=true, + sdob=false, + path=None, + ))] + #[allow(clippy::too_many_arguments)] + pub fn render_frame( + &self, + py: Python, + phase: f64, + width: usize, + height: Option, + x1: f64, + x2: f64, + y1: f64, + y2: f64, + reverse: bool, + sdob: bool, + path: Option, + ) -> PyResult>> { + let options = RenderOptions::new(width, height, x1, x2, y1, y2, reverse, sdob, false)?; + let rgb = py.detach(|| -> Result, RocheError> { + let scene = visualize::build_scene( + &self.model, + &self.star1_fine_grid, + &self.star2_fine_grid, + &self.disc_grid, + )?; + let buffer = visualize::render_frame_indexed(&scene, &options, phase); + if let Some(ref path) = path { + visualize::write_png(path, &buffer, &options)?; + } + Ok(visualize::indexed_to_rgb(&buffer, &options.palette())) + })?; + let array = Array3::from_shape_vec((options.height, options.width, 3), rgb) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + Ok(array.into_pyarray(py).unbind()) + } + + /// + /// Project the visible points of every model component onto the plane + /// of the sky at the given orbital phase, in units of the binary + /// separation relative to the centre of mass. Returns a dict mapping + /// "star1", "star2", "disc", "disc_edge_outer", "disc_edge_inner" and + /// "stream" to (n, 2) float64 numpy arrays of (x, y) sky coordinates + /// (empty for absent components), plus "bright_spot", which is either + /// None or a ((x, y), cosbs) tuple where cosbs sets the marker size in + /// the rendered frames. Useful for custom plotting, e.g. matplotlib. + /// + pub fn scene_points(&self, py: Python, phase: f64) -> PyResult> { + let scene = visualize::build_scene( + &self.model, + &self.star1_fine_grid, + &self.star2_fine_grid, + &self.disc_grid, + )?; + let (star1, star2, disc, outer_edge, inner_edge, stream, spot) = + visualize::scene_points(&scene, phase); + + let to_array = |points: Vec<[f64; 2]>| -> PyResult>> { + let n = points.len(); + let flat: Vec = points.into_iter().flatten().collect(); + let array = Array2::from_shape_vec((n, 2), flat) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + Ok(array.into_pyarray(py).unbind()) + }; + + let dict = PyDict::new(py); + dict.set_item("star1", to_array(star1)?)?; + dict.set_item("star2", to_array(star2)?)?; + dict.set_item("disc", to_array(disc)?)?; + dict.set_item("disc_edge_outer", to_array(outer_edge)?)?; + dict.set_item("disc_edge_inner", to_array(inner_edge)?)?; + dict.set_item("stream", to_array(stream)?)?; + dict.set_item("bright_spot", spot.map(|(xy, cosbs)| ((xy[0], xy[1]), cosbs)))?; + Ok(dict.unbind()) + } + + /// + /// Render an animation of the orbit between two phases and write it to + /// `path`. The format is chosen from the file extension: `.gif` is + /// encoded natively, `.mp4` requires the `ffmpeg` binary on the PATH. + /// Frames are rendered in parallel. + /// + /// Keyword arguments: + /// phase1, phase2: first and last orbital phase (defaults 0 and 1). + /// nphase: number of frames (default 300). + /// fps: frames per second, 1 to 240 (default 25). + /// width, height, x1, x2, y1, y2, reverse, sdob: as render_frame. + /// + #[pyo3(signature = ( + path, + *, + phase1=0.0, + phase2=1.0, + nphase=300, + fps=25, + width=800, + height=None, + x1=-2.0, + x2=2.0, + y1=-2.0, + y2=2.0, + reverse=true, + sdob=false, + ))] + #[allow(clippy::too_many_arguments)] + pub fn render_animation( + &self, + py: Python, + path: PathBuf, + phase1: f64, + phase2: f64, + nphase: usize, + fps: usize, + width: usize, + height: Option, + x1: f64, + x2: f64, + y1: f64, + y2: f64, + reverse: bool, + sdob: bool, + ) -> PyResult<()> { + let extension = path + .extension() + .and_then(|e| e.to_str()) + .map(str::to_ascii_lowercase); + let is_mp4 = match extension.as_deref() { + Some("gif") => false, + Some("mp4") => true, + _ => { + return Err(pyo3::exceptions::PyValueError::new_err( + "path must end in .gif or .mp4", + )); + } + }; + if nphase < 1 { + return Err(pyo3::exceptions::PyValueError::new_err( + "nphase must be at least 1", + )); + } + if !(1..=240).contains(&fps) { + return Err(pyo3::exceptions::PyValueError::new_err( + "fps must be between 1 and 240", + )); + } + let options = RenderOptions::new(width, height, x1, x2, y1, y2, reverse, sdob, is_mp4)?; + py.detach(|| -> Result<(), RocheError> { + let scene = visualize::build_scene( + &self.model, + &self.star1_fine_grid, + &self.star2_fine_grid, + &self.disc_grid, + )?; + let phases = visualize::phase_grid(phase1, phase2, nphase); + if is_mp4 { + visualize::write_mp4(&path, &scene, &options, &phases, fps) + } else { + visualize::write_gif(&path, &scene, &options, &phases, fps) + } + })?; + Ok(()) + } } impl BinaryModel { diff --git a/src/lib.rs b/src/lib.rs index 1a0bd6c..09d3e4f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,6 +16,7 @@ pub mod set_disc_continuum; pub mod set_disc_grid; pub mod set_star_continuum; pub mod set_star_grid; +pub mod visualize; #[pymodule] fn lcurve(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/src/visualize.rs b/src/visualize.rs new file mode 100644 index 0000000..9854da8 --- /dev/null +++ b/src/visualize.rs @@ -0,0 +1,786 @@ +//! 3D visualization of the binary system, rendered as sky-plane projections. +//! +//! This is a port of the `visualise` program from the original cpp-lcurve: +//! grids of points representing the two stars, the disc, its edges and the +//! gas stream are projected onto the plane of the sky at a given orbital +//! phase and rasterized as coloured dots. Still frames are written as PNG, +//! animations as GIF (pure Rust) or MP4 (via an external `ffmpeg` binary). + +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use rayon::prelude::*; +use roche::errors::RocheError; +use roche::{self, Etype, Point, RocheContext, Star, Vec3}; + +use crate::model::Model; +use crate::set_star_grid::star_eclipse; + +/// Palette indices used by the rasterizer. These mirror the PGPLOT colour +/// indices of the original `visualise` program. +const BACKGROUND: u8 = 0; +const EDGES: u8 = 1; +const STAR2: u8 = 2; +const DISC: u8 = 3; +const STAR1: u8 = 4; +const NCOLOURS: usize = 5; + +/// Rendering options shared by still frames and animations. +#[derive(Clone, Copy, Debug)] +pub struct RenderOptions { + pub width: usize, + pub height: usize, + pub x1: f64, + pub x2: f64, + pub y1: f64, + pub y2: f64, + pub reverse: bool, + pub sdob: bool, +} + +impl RenderOptions { + /// Build validated options. When `height` is `None` it is derived from + /// the aspect ratio of the plot limits; `force_even` rounds both + /// dimensions down to even numbers (required by yuv420p MP4 encoding). + #[allow(clippy::too_many_arguments)] + pub fn new( + width: usize, + height: Option, + x1: f64, + x2: f64, + y1: f64, + y2: f64, + reverse: bool, + sdob: bool, + force_even: bool, + ) -> Result { + if !(x2 > x1 && y2 > y1) { + return Err(RocheError::ParameterError( + "require x2 > x1 and y2 > y1.".to_string(), + )); + } + let mut width = width; + let mut height = + height.unwrap_or_else(|| (width as f64 * (y2 - y1) / (x2 - x1)).round() as usize); + if force_even { + width &= !1; + height &= !1; + } + if !(2..=16384).contains(&width) || !(2..=16384).contains(&height) { + return Err(RocheError::ParameterError( + "width and height must be between 2 and 16384 pixels.".to_string(), + )); + } + Ok(Self { + width, + height, + x1, + x2, + y1, + y2, + reverse, + sdob, + }) + } + + /// RGB palette matching the original PGPLOT colour scheme. `reverse` + /// selects a white background with dark colours (the original default), + /// otherwise a black background with bright colours. `sdob` swaps the + /// colours of the two stars. + pub fn palette(&self) -> [[u8; 3]; NCOLOURS] { + let (bg, fg) = if self.reverse { + ([255, 255, 255], [0, 0, 0]) + } else { + ([0, 0, 0], [255, 255, 255]) + }; + let (red, green, blue) = if self.reverse { + ([102, 0, 0], [0, 77, 0], [0, 0, 128]) + } else { + ([255, 0, 0], [0, 255, 0], [0, 0, 255]) + }; + let (star2, star1) = if self.sdob { (blue, red) } else { (red, blue) }; + [bg, fg, star2, green, star1] + } +} + +/// The geometry of every plotted component, built once per render call and +/// shared (immutably) between frame-rendering threads. +pub struct Scene<'a> { + star1: &'a [Point], + star2: &'a [Point], + disc: &'a [Point], + outer_edge: Vec, + inner_edge: Vec, + stream: Vec, + add_disc: bool, + add_spot: bool, + cofm: Vec3, + iangle: f64, +} + +/// Build the scene from a model and its pre-computed star and disc grids. +/// The disc edges are rebuilt densely for display (the flux grids are too +/// coarse to read as clean rim lines) and the gas stream is integrated with +/// `strinit`/`stradv`, exactly as the original `visualise` program does. +pub fn build_scene<'a>( + model: &Model, + star1: &'a [Point], + star2: &'a [Point], + disc: &'a [Point], +) -> Result, RocheError> { + let (outer_edge, inner_edge) = if model.add_disc { + ( + build_visual_disc_edge(model, true)?, + build_visual_disc_edge(model, false)?, + ) + } else { + (Vec::new(), Vec::new()) + }; + + let stream = if model.add_spot { + build_stream(model)? + } else { + Vec::new() + }; + + Ok(Scene { + star1, + star2, + disc, + outer_edge, + inner_edge, + stream, + add_disc: model.add_disc, + add_spot: model.add_spot, + cofm: Vec3::new(model.q.value / (1.0 + model.q.value), 0.0, 0.0), + iangle: model.iangle.value, + }) +} + +/// Build a dense ring of points tracing the outer or inner disc edge for +/// display. The geometry and eclipse handling mirror `set_disc_edge_grid` +/// (top rim always visible unless eclipsed, wall visible from its own +/// side), but the sampling is chosen for image quality rather than flux +/// integration so the rim renders as a clean line. +fn build_visual_disc_edge(model: &Model, outer: bool) -> Result, RocheError> { + const EFAC: f64 = 1.0000001; + const NZ: usize = 5; + + let q: f64 = model.q.value; + let iangle: f64 = model.iangle.value; + + let roche_context1 = RocheContext::new(q, Star::Primary, model.spin1.value)?; + let roche_context2 = RocheContext::new(q, Star::Secondary, model.spin2.value)?; + + let (mut r1, mut r2) = model.get_r1r2(); + let rl1: f64 = roche_context1.x_l1; + let rl2: f64 = 1.0 - roche_context2.x_l1; + if r1 < 0.0 { + r1 = rl1; + } + if r2 < 0.0 { + r2 = rl2; + } + let ffac1: f64 = r1 / rl1; + let ffac2: f64 = r2 / rl2; + + let rdisc1: f64 = if model.rdisc1.value > 0.0 { + model.rdisc1.value + } else { + r1 + }; + let rdisc2: f64 = if model.rdisc2.value > 0.0 { + model.rdisc2.value + } else { + model.radius_spot.value + }; + + let rad: f64 = if outer { rdisc2 } else { rdisc1 }; + let h: f64 = model.height_disc.value * rad.powf(model.beta_disc.value); + + // Sample the larger circle at 1/4-degree steps and scale the smaller + // one with its circumference, so both rims render equally dense. + let ntheta: usize = ((1440.0 * rad / rdisc2.max(rdisc1)) as usize).max(180); + + let add_eclipses = |posn: &Vec3| -> Result { + let mut eclipses = Etype::new(); + if model.opaque { + for eclipse_pair in roche::disc_eclipse( + iangle, + rdisc1, + rdisc2, + model.beta_disc.value, + model.height_disc.value, + posn, + )? { + eclipses.push(eclipse_pair); + } + } + if model.eclipse1 { + star_eclipse( + &roche_context1, + r1, + ffac1, + iangle, + posn, + model.delta_phase, + model.roche1, + Star::Primary, + &mut eclipses, + )?; + } + if model.eclipse2 { + star_eclipse( + &roche_context2, + r2, + ffac2, + iangle, + posn, + model.delta_phase, + model.roche2, + Star::Secondary, + &mut eclipses, + )?; + } + Ok(eclipses) + }; + + let always_visible = Vec3::new(0.0, 0.0, 0.0); + let columns: Vec> = (0..ntheta) + .into_par_iter() + .map(|i| -> Result, RocheError> { + let theta: f64 = std::f64::consts::TAU * i as f64 / ntheta as f64; + let (sint, cost) = theta.sin_cos(); + let (r_xy, wall_dirn) = if outer { + (EFAC * rad, Vec3::new(cost, sint, 0.0)) + } else { + (rad / EFAC, Vec3::new(-cost, -sint, 0.0)) + }; + + let mut column = Vec::with_capacity(NZ + 1); + + // Top rim, visible from every direction unless eclipsed + let posn = Vec3::new(r_xy * cost, r_xy * sint, EFAC * h); + column.push(Point::new( + posn, + always_visible, + 0.0, + 1.0, + add_eclipses(&posn)?, + )); + + // Wall, visible from its own side + for j in 0..NZ { + let z: f64 = -h + 2.0 * h * j as f64 / (NZ - 1) as f64; + let posn = Vec3::new(r_xy * cost, r_xy * sint, z); + column.push(Point::new(posn, wall_dirn, 0.0, 1.0, add_eclipses(&posn)?)); + } + Ok(column) + }) + .collect::, _>>()?; + + Ok(columns.into_iter().flatten().collect()) +} + +/// Integrate the gas stream from the L1 point down to the bright-spot +/// radius, recording the phases at which each point is eclipsed by either +/// star (and the disc, if present). +fn build_stream(model: &Model) -> Result, RocheError> { + let q: f64 = model.q.value; + let iangle: f64 = model.iangle.value; + + let roche_context1 = RocheContext::new(q, Star::Primary, model.spin1.value)?; + let roche_context2 = RocheContext::new(q, Star::Secondary, model.spin2.value)?; + + let (mut r1, mut r2) = model.get_r1r2(); + let rl1: f64 = roche_context1.x_l1; + let rl2: f64 = 1.0 - roche_context2.x_l1; + if r1 <= 0.0 { + r1 = 0.99999999999 * rl1; + } + if r2 <= 0.0 { + r2 = 0.99999999999 * rl2; + } + let ffac1: f64 = r1 / rl1; + let ffac2: f64 = r2 / rl2; + + let rdisc1: f64 = if model.rdisc1.value > 0.0 { + model.rdisc1.value + } else { + r1 + }; + let rdisc2: f64 = if model.rdisc2.value > 0.0 { + model.rdisc2.value + } else { + model.radius_spot.value + }; + + // Direction chosen so the stream is visible at all inclinations < 90. + let dirn = Vec3::new(0.0, 0.0, 1.0); + + let add_eclipses = |posn: &Vec3| -> Result { + let mut eclipses = Etype::new(); + star_eclipse( + &roche_context1, + r1, + ffac1, + iangle, + posn, + model.delta_phase, + model.roche1, + Star::Primary, + &mut eclipses, + )?; + star_eclipse( + &roche_context2, + r2, + ffac2, + iangle, + posn, + model.delta_phase, + model.roche2, + Star::Secondary, + &mut eclipses, + )?; + if model.add_disc { + for eclipse_pair in roche::disc_eclipse( + iangle, + rdisc1, + rdisc2, + model.beta_disc.value, + model.height_disc.value, + posn, + )? { + eclipses.push(eclipse_pair); + } + } + Ok(eclipses) + }; + + let (mut posn, mut v) = roche::strinit(q)?; + + let mut stream: Vec = Vec::new(); + stream.push(Point::new(posn, dirn, 0.0, 1.0, add_eclipses(&posn)?)); + + let rl1_point: f64 = roche::x_l1(q)?; + let radius_spot: f64 = model.radius_spot.value; + let nstream: usize = (((rl1_point - radius_spot) / 0.001) as usize).max(1); + for i in 0..nstream { + let radius: f64 = rl1_point + (radius_spot - rl1_point) * (i as f64 + 1.0) / nstream as f64; + roche::stradv(q, &mut posn, &mut v, radius, 1.0e-10, 1.0e-3)?; + stream.push(Point::new(posn, dirn, 0.0, 1.0, add_eclipses(&posn)?)); + } + Ok(stream) +} + +/// Fill a circle of the given radius (in pixels) into the indexed buffer, +/// clipped to the image bounds. +fn fill_circle( + buffer: &mut [u8], + width: usize, + height: usize, + cx: f64, + cy: f64, + radius: f64, + colour: u8, +) { + let r = radius.max(0.5); + let x_min = ((cx - r).floor().max(0.0)) as usize; + let x_max = ((cx + r).ceil().min(width as f64 - 1.0)) as usize; + let y_min = ((cy - r).floor().max(0.0)) as usize; + let y_max = ((cy + r).ceil().min(height as f64 - 1.0)) as usize; + if cx + r < 0.0 || cy + r < 0.0 || cx - r > width as f64 - 1.0 || cy - r > height as f64 - 1.0 { + return; + } + let r2 = r * r; + for y in y_min..=y_max { + let dy = y as f64 - cy; + let row = y * width; + for x in x_min..=x_max { + let dx = x as f64 - cx; + if dx * dx + dy * dy <= r2 { + buffer[row + x] = colour; + } + } + } +} + +struct FrameGeometry { + earth: Vec3, + xsky: Vec3, + ysky: Vec3, + phase_normed: f64, + x_scale: f64, + y_scale: f64, +} + +impl FrameGeometry { + fn new(scene: &Scene, options: &RenderOptions, phase: f64) -> Self { + let earth = roche::set_earth_iangle(scene.iangle, phase); + let (sinp, cosp) = (std::f64::consts::TAU * phase).sin_cos(); + let xsky = Vec3::new(sinp, cosp, 0.0); + let ysky = earth.cross(&xsky); + Self { + earth, + xsky, + ysky, + phase_normed: phase - phase.floor(), + x_scale: (options.width - 1) as f64 / (options.x2 - options.x1), + y_scale: (options.height - 1) as f64 / (options.y2 - options.y1), + } + } + + fn project(&self, options: &RenderOptions, scene: &Scene, position: &Vec3) -> (f64, f64) { + let r: Vec3 = *position - scene.cofm; + let sx: f64 = r.dot(&self.xsky); + let sy: f64 = r.dot(&self.ysky); + let px = (sx - options.x1) * self.x_scale; + let py = (options.height - 1) as f64 - (sy - options.y1) * self.y_scale; + (px, py) + } +} + +/// Draw every point of one component that faces the observer and is not +/// eclipsed at this phase. Points with a zero direction vector (the disc +/// edges in "visual" mode) are treated as always facing the observer. +fn draw_component( + buffer: &mut [u8], + scene: &Scene, + options: &RenderOptions, + geometry: &FrameGeometry, + points: &[Point], + colour: u8, + radius: f64, +) { + for point in points { + let d = point.direction; + let faces_observer = + (d.x == 0.0 && d.y == 0.0 && d.z == 0.0) || geometry.earth.dot(&d) > 0.0; + if faces_observer && point.is_visible_phase_normed(geometry.phase_normed) { + let (px, py) = geometry.project(options, scene, &point.position); + fill_circle( + buffer, + options.width, + options.height, + px, + py, + radius, + colour, + ); + } + } +} + +/// Render one frame into a palette-indexed buffer of `width * height` bytes. +pub fn render_frame_indexed(scene: &Scene, options: &RenderOptions, phase: f64) -> Vec { + let width = options.width; + let height = options.height; + let mut buffer = vec![BACKGROUND; width * height]; + let geometry = FrameGeometry::new(scene, options, phase); + + let dot_radius: f64 = (width.min(height) as f64 / 400.0).max(1.0); + + // Depth-sort the two stars so the nearer one is drawn on top. + let star1_in_front = geometry.earth.x < 0.0; + let (back_star, back_colour, front_star, front_colour) = if star1_in_front { + (scene.star2, STAR2, scene.star1, STAR1) + } else { + (scene.star1, STAR1, scene.star2, STAR2) + }; + + draw_component( + &mut buffer, + scene, + options, + &geometry, + back_star, + back_colour, + dot_radius, + ); + if scene.add_disc { + draw_component( + &mut buffer, + scene, + options, + &geometry, + scene.disc, + DISC, + dot_radius, + ); + draw_component( + &mut buffer, + scene, + options, + &geometry, + &scene.outer_edge, + EDGES, + dot_radius, + ); + draw_component( + &mut buffer, + scene, + options, + &geometry, + &scene.inner_edge, + EDGES, + dot_radius, + ); + } + if scene.add_spot { + draw_component( + &mut buffer, + scene, + options, + &geometry, + &scene.stream, + STAR2, + dot_radius, + ); + } + draw_component( + &mut buffer, + scene, + options, + &geometry, + front_star, + front_colour, + dot_radius, + ); + + // Emphasize the bright spot at the end of the stream with a larger + // marker whose size grows as it turns towards the observer, as in the + // original program. + if let Some(spot) = scene.stream.last().filter(|_| scene.add_spot) { + let cosbs: f64 = geometry.earth.dot(&spot.position); + if cosbs > 0.0 && spot.is_visible_phase_normed(geometry.phase_normed) { + let (px, py) = geometry.project(options, scene, &spot.position); + let radius = (0.5 + 3.5 * cosbs) * width.min(height) as f64 / 200.0; + fill_circle(&mut buffer, width, height, px, py, radius, STAR2); + } + } + + buffer +} + +/// Sky-plane coordinates of every visible point of one component at the +/// given phase geometry, as (x, y) pairs relative to the centre of mass. +fn project_component(scene: &Scene, geometry: &FrameGeometry, points: &[Point]) -> Vec<[f64; 2]> { + let mut projected = Vec::new(); + for point in points { + let d = point.direction; + let faces_observer = + (d.x == 0.0 && d.y == 0.0 && d.z == 0.0) || geometry.earth.dot(&d) > 0.0; + if faces_observer && point.is_visible_phase_normed(geometry.phase_normed) { + let r: Vec3 = point.position - scene.cofm; + projected.push([r.dot(&geometry.xsky), r.dot(&geometry.ysky)]); + } + } + projected +} + +/// The visible points of every component at one phase, projected onto the +/// sky plane. Components are returned in draw order; absent components +/// (no disc / no spot) come back as empty vectors. The bright spot is the +/// last stream point paired with its `cosbs` size factor when visible. +pub type ScenePoints = ( + Vec<[f64; 2]>, + Vec<[f64; 2]>, + Vec<[f64; 2]>, + Vec<[f64; 2]>, + Vec<[f64; 2]>, + Vec<[f64; 2]>, + Option<([f64; 2], f64)>, +); + +pub fn scene_points(scene: &Scene, phase: f64) -> ScenePoints { + // Any options with valid limits produce the same geometry; the sky + // projection does not depend on the raster settings. + let options = RenderOptions { + width: 2, + height: 2, + x1: -1.0, + x2: 1.0, + y1: -1.0, + y2: 1.0, + reverse: true, + sdob: false, + }; + let geometry = FrameGeometry::new(scene, &options, phase); + + let spot = scene + .stream + .last() + .filter(|_| scene.add_spot) + .and_then(|spot| { + let cosbs: f64 = geometry.earth.dot(&spot.position); + if cosbs > 0.0 && spot.is_visible_phase_normed(geometry.phase_normed) { + let r: Vec3 = spot.position - scene.cofm; + Some(([r.dot(&geometry.xsky), r.dot(&geometry.ysky)], cosbs)) + } else { + None + } + }); + + ( + project_component(scene, &geometry, scene.star1), + project_component(scene, &geometry, scene.star2), + project_component(scene, &geometry, scene.disc), + project_component(scene, &geometry, &scene.outer_edge), + project_component(scene, &geometry, &scene.inner_edge), + project_component(scene, &geometry, &scene.stream), + spot, + ) +} + +/// Expand a palette-indexed buffer into packed RGB bytes. +pub fn indexed_to_rgb(buffer: &[u8], palette: &[[u8; 3]; NCOLOURS]) -> Vec { + let mut rgb = Vec::with_capacity(buffer.len() * 3); + for &index in buffer { + rgb.extend_from_slice(&palette[index as usize]); + } + rgb +} + +/// The uniformly spaced phases of an animation running from `phase1` to +/// `phase2` inclusive; a single frame sits at `phase1`. +pub fn phase_grid(phase1: f64, phase2: f64, nphase: usize) -> Vec { + (0..nphase) + .map(|i| { + if nphase == 1 { + phase1 + } else { + phase1 + (phase2 - phase1) * i as f64 / (nphase - 1) as f64 + } + }) + .collect() +} + +fn io_error(err: impl std::fmt::Display, context: &str) -> RocheError { + RocheError::ParameterError(format!("{context}: {err}")) +} + +/// Write one palette-indexed frame as an indexed-colour PNG. +pub fn write_png(path: &Path, buffer: &[u8], options: &RenderOptions) -> Result<(), RocheError> { + let file = std::fs::File::create(path).map_err(|e| io_error(e, "cannot create PNG"))?; + let writer = std::io::BufWriter::new(file); + let mut encoder = png::Encoder::new(writer, options.width as u32, options.height as u32); + encoder.set_color(png::ColorType::Indexed); + encoder.set_depth(png::BitDepth::Eight); + encoder.set_palette(options.palette().concat()); + let mut png_writer = encoder + .write_header() + .map_err(|e| io_error(e, "cannot write PNG header"))?; + png_writer + .write_image_data(buffer) + .map_err(|e| io_error(e, "cannot write PNG data"))?; + Ok(()) +} + +/// Render an animated GIF. Frames are rendered in parallel in chunks so the +/// peak memory use stays bounded while the (sequential) encoder keeps up. +pub fn write_gif( + path: &Path, + scene: &Scene, + options: &RenderOptions, + phases: &[f64], + fps: usize, +) -> Result<(), RocheError> { + let file = std::fs::File::create(path).map_err(|e| io_error(e, "cannot create GIF"))?; + let writer = std::io::BufWriter::new(file); + let palette = options.palette().concat(); + let mut encoder = gif::Encoder::new( + writer, + options.width as u16, + options.height as u16, + &palette, + ) + .map_err(|e| io_error(e, "cannot write GIF header"))?; + encoder + .set_repeat(gif::Repeat::Infinite) + .map_err(|e| io_error(e, "cannot set GIF repeat"))?; + + let delay_cs = (100 / fps.max(1)).max(1) as u16; + let chunk_size = rayon::current_num_threads().max(1) * 4; + for chunk in phases.chunks(chunk_size) { + let frames: Vec> = chunk + .par_iter() + .map(|&phase| render_frame_indexed(scene, options, phase)) + .collect(); + for buffer in frames { + let mut frame = gif::Frame { + width: options.width as u16, + height: options.height as u16, + delay: delay_cs, + buffer: std::borrow::Cow::Owned(buffer), + ..gif::Frame::default() + }; + frame.make_lzw_pre_encoded(); + encoder + .write_lzw_pre_encoded_frame(&frame) + .map_err(|e| io_error(e, "cannot write GIF frame"))?; + } + } + Ok(()) +} + +static MP4_RUN_COUNTER: AtomicUsize = AtomicUsize::new(0); + +/// Render an MP4 by writing PNG frames to a temporary directory (in +/// parallel) and invoking the system `ffmpeg` binary, exactly like the +/// original C++ tool. Fails with a clear message if `ffmpeg` is missing. +pub fn write_mp4( + path: &Path, + scene: &Scene, + options: &RenderOptions, + phases: &[f64], + fps: usize, +) -> Result<(), RocheError> { + let run = MP4_RUN_COUNTER.fetch_add(1, Ordering::Relaxed); + let tmpdir: PathBuf = + std::env::temp_dir().join(format!("lroche_frames_{}_{}", std::process::id(), run)); + std::fs::create_dir_all(&tmpdir).map_err(|e| io_error(e, "cannot create frame dir"))?; + + let result = (|| -> Result<(), RocheError> { + phases + .par_iter() + .enumerate() + .try_for_each(|(i, &phase)| -> Result<(), RocheError> { + let buffer = render_frame_indexed(scene, options, phase); + write_png(&tmpdir.join(format!("frame_{i:05}.png")), &buffer, options) + })?; + + let pattern = tmpdir.join("frame_%05d.png"); + let status = Command::new("ffmpeg") + .args(["-y", "-hide_banner", "-v", "error", "-framerate"]) + .arg(fps.to_string()) + .arg("-i") + .arg(&pattern) + .args([ + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + "-movflags", + "+faststart", + ]) + .arg(path) + .status(); + + match status { + Ok(status) if status.success() => Ok(()), + Ok(status) => Err(RocheError::ParameterError(format!( + "ffmpeg failed with {status}." + ))), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => { + Err(RocheError::ParameterError( + "ffmpeg not found on PATH; install it or render a .gif instead.".to_string(), + )) + } + Err(err) => Err(io_error(err, "cannot run ffmpeg")), + } + })(); + + let _ = std::fs::remove_dir_all(&tmpdir); + result +}