diff --git a/crates/resolved-cli/src/main.rs b/crates/resolved-cli/src/main.rs index 84fe347..a612d20 100644 --- a/crates/resolved-cli/src/main.rs +++ b/crates/resolved-cli/src/main.rs @@ -57,10 +57,11 @@ struct RenderArgs { /// catalog). #[arg(long, default_value = "L0")] lighting: LightingTier, - /// Camera fidelity tier: C0, or C1 (adds pixels that collect light over their area, + /// Camera fidelity tier: C0, C1 (adds pixels that collect light over their area, /// for graded edges and continuous thin parts, lens distortion, lens blur, and the /// Sun, Moon and Earth where light time and aberration make them appear, detector - /// noise, fixed pattern and blooming, and at lighting L1 stars). + /// noise, fixed pattern and blooming, and at lighting L1 stars), or C2 (adds the + /// lens blur by FFT, with diffraction or a measured PSF, whose wings reach far). #[arg(long, default_value = "C0")] camera: CameraTier, /// From camera C1: render each pixel as N x N finer ones (1 to 4; default 2), which @@ -74,6 +75,11 @@ struct RenderArgs { /// give the same image. #[arg(long, default_value_t = 0)] frame: u64, + /// From camera C2: how far in pixels the PSF must reach from a point (default 16). + /// The FFT grid grows to hold it, and the PSF then reaches as far as the grid + /// allows: 64 px at 1920 x 1080 and 2 x 2 supersampling, up to 1088 px at 1 x 1. + #[arg(long, value_name = "PIXELS")] + psf_reach: Option, /// Device: auto, gpu, cpu, or part of an adapter's name. #[arg(long, default_value = "auto")] adapter: String, @@ -162,7 +168,8 @@ fn render(args: &RenderArgs) -> Result<(), String> { } let path = args.scene.as_path(); let file = SceneFile::read(path)?; - let camera = file.camera(); + let dir = path.parent().unwrap_or(Path::new(".")); + let camera = file.camera(dir)?; // What the tiers need, gathered before the renderer exists so it can check them. let start = Instant::now(); @@ -230,12 +237,12 @@ fn render(args: &RenderArgs) -> Result<(), String> { stars, noise_seed: args.seed, fixed_pattern: None, + psf_reach: args.psf_reach, }; let start = Instant::now(); let mut renderer = Renderer::new(camera, options).map_err(|e| e.to_string())?; let setup = start.elapsed(); let start = Instant::now(); - let dir = path.parent().unwrap_or(Path::new(".")); let scene = file.scene(&mut renderer, dir, args.kernels.as_deref())?; let observer = file.observer()?; let loading = start.elapsed(); diff --git a/crates/resolved-cli/src/scene_file.rs b/crates/resolved-cli/src/scene_file.rs index 0b76411..6257575 100644 --- a/crates/resolved-cli/src/scene_file.rs +++ b/crates/resolved-cli/src/scene_file.rs @@ -12,7 +12,7 @@ use resolved::nalgebra::{Quaternion, UnitQuaternion, Vector2, Vector3}; use resolved::scene::Velocities; use resolved::{ Blooming, BrownConrady, CameraConfig, LoadOptions, Material, Mesh, ModelId, Pose, Psf, - Renderer, Scene, Spectrum, Target, + PsfImage, Renderer, Scene, Spectrum, Target, }; use serde::Deserialize; @@ -88,6 +88,13 @@ pub struct CameraFile { /// The Gaussian blur's sigma in pixels, from camera C1: one for every channel, one /// per channel, or 0 for none. pub psf_sigma: Option, + /// Diffraction by the aperture, from camera C2: the diameter of a centered + /// obstruction as a fraction of the aperture's, 0 for none. + pub psf_airy: Option, + /// A measured PSF, from camera C2: `{ image = "spot.png", pixel_scale = 0.25 }`, a + /// PNG relative to the scene file, grayscale for every channel or RGB for three, + /// whose samples are `pixel_scale` pixels apart. + pub psf_measured: Option, /// Lens distortion from camera C1, in OpenCV's order: [k1, k2, p1, p2, k3]. pub distortion: Option<[f64; 5]>, } @@ -101,6 +108,14 @@ pub enum BloomingFile { Neighbors(f64), } +/// A measured PSF in a scene file. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MeasuredFile { + pub image: PathBuf, + pub pixel_scale: f64, +} + /// One sigma, or one per channel. #[derive(Debug, Deserialize)] #[serde(untagged)] @@ -165,7 +180,7 @@ impl SceneFile { } /// The example camera with this file's overrides. - pub fn camera(&self) -> CameraConfig { + pub fn camera(&self, dir: &Path) -> Result { let c = &self.camera; let mut camera = CameraConfig::example(); let set = |field: &mut f64, value: Option| { @@ -219,17 +234,30 @@ impl SceneFile { if let Some([k1, k2, p1, p2, k3]) = c.distortion { camera.distortion = Some(BrownConrady { k1, k2, k3, p1, p2 }); } - camera.psf = match &c.psf_sigma { - None => camera.psf, - Some(Sigma::All(s)) if *s == 0.0 => Psf::None, - Some(Sigma::All(s)) => Psf::Gaussian { + let given = [ + c.psf_sigma.is_some(), + c.psf_airy.is_some(), + c.psf_measured.is_some(), + ]; + if given.iter().filter(|&&g| g).count() > 1 { + return Err("give at most one of psf_sigma, psf_airy and psf_measured".into()); + } + camera.psf = match (&c.psf_sigma, c.psf_airy, &c.psf_measured) { + (Some(Sigma::All(s)), _, _) if *s == 0.0 => Psf::None, + (Some(Sigma::All(s)), _, _) => Psf::Gaussian { sigma: vec![*s; camera.channels.len()], }, - Some(Sigma::PerChannel(sigma)) => Psf::Gaussian { + (Some(Sigma::PerChannel(sigma)), _, _) => Psf::Gaussian { sigma: sigma.clone(), }, + (_, Some(obstruction), _) => Psf::Airy { obstruction }, + (_, _, Some(m)) => Psf::Measured { + image: psf_image(&dir.join(&m.image))?, + pixel_scale: m.pixel_scale, + }, + (None, None, None) => camera.psf, }; - camera + Ok(camera) } /// Earth's and the Moon's appearance. @@ -407,6 +435,27 @@ impl SceneFile { } /// Where to look for astrokit's kernels: `explicit` alone if given; otherwise +/// A measured PSF's image from a PNG: one image from a grayscale one, three from an RGB +/// one, each sample's value as a fraction of the largest the format holds. +fn psf_image(path: &Path) -> Result { + let image = image::open(path).map_err(|e| format!("{}: {e}", path.display()))?; + let (width, height) = (image.width(), image.height()); + let channels = if image.color().has_color() { + let rgb = image.into_rgb16(); + (0..3) + .map(|k| rgb.pixels().map(|p| p[k] as f64 / 65535.0).collect()) + .collect() + } else { + let gray = image.into_luma16(); + vec![gray.pixels().map(|p| p[0] as f64 / 65535.0).collect()] + }; + Ok(PsfImage { + width, + height, + channels, + }) +} + /// `ASTROKIT_KERNELS`, then `data/kernels` under the working directory, then the /// sibling astrokit checkout this binary was built against. #[cfg(feature = "astrokit")] @@ -478,7 +527,10 @@ mod tests { #[test] fn minimal_file_uses_the_example_camera() { let file: SceneFile = toml::from_str(MINIMAL).unwrap(); - assert_eq!(file.camera(), CameraConfig::example()); + assert_eq!( + file.camera(Path::new(".")).unwrap(), + CameraConfig::example() + ); assert!(file.targets.is_empty()); let observer = file.observer().unwrap(); // Looking from +X toward the origin. @@ -493,7 +545,7 @@ mod tests { pattern_seed = 9\nblooming = {{ column = 0.25 }}\n" )) .unwrap(); - let camera = file.camera(); + let camera = file.camera(Path::new(".")).unwrap(); assert_eq!((camera.width, camera.height), (640, 1080)); assert_eq!(camera.exposure, 2e-3); assert_eq!(camera.detector.adc_bits, 10); @@ -520,18 +572,24 @@ mod tests { }, ), ("psf_sigma = 0", Psf::None), + ("psf_airy = 0.3", Psf::Airy { obstruction: 0.3 }), ] { let file: SceneFile = toml::from_str(&format!("{MINIMAL}\n[camera]\n{line}\n")).unwrap(); - assert_eq!(file.camera().psf, psf, "{line}"); + assert_eq!(file.camera(Path::new(".")).unwrap().psf, psf, "{line}"); } + let file: SceneFile = toml::from_str(&format!( + "{MINIMAL}\n[camera]\npsf_sigma = 1\npsf_airy = 0\n" + )) + .unwrap(); + assert!(file.camera(Path::new(".")).is_err()); let file: SceneFile = toml::from_str(&format!( "{MINIMAL}\n[camera]\ndistortion = [-0.1, 0.02, 1e-3, -2e-3, 0.005]\n" )) .unwrap(); assert_eq!( - file.camera().distortion, + file.camera(Path::new(".")).unwrap().distortion, Some(BrownConrady { k1: -0.1, k2: 0.02, @@ -542,6 +600,36 @@ mod tests { ); } + #[test] + fn a_measured_psf_loads_from_a_png() { + // 16-bit RGB, 3 x 2: one image per channel, each sample's value over 65,535. + let dir = std::env::temp_dir().join(format!("resolved-psf-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let values: Vec = (0..18).map(|i| i * 1000).collect(); + image::ImageBuffer::, _>::from_raw(3, 2, values) + .unwrap() + .save(dir.join("spot.png")) + .unwrap(); + let file: SceneFile = toml::from_str(&format!( + "{MINIMAL}\n[camera]\npsf_measured = {{ image = \"spot.png\", pixel_scale = 0.5 }}\n" + )) + .unwrap(); + let camera = file.camera(&dir).unwrap(); + std::fs::remove_dir_all(&dir).unwrap(); + let Psf::Measured { image, pixel_scale } = camera.psf else { + panic!("{:?}", camera.psf); + }; + assert_eq!((image.width, image.height, pixel_scale), (3, 2, 0.5)); + for (k, channel) in image.channels.iter().enumerate() { + let want: Vec = (0..6) + .map(|i| ((3 * i + k) * 1000) as f64 / 65535.0) + .collect(); + for (a, b) in channel.iter().zip(&want) { + assert!((a - b).abs() < 1e-6, "channel {k}: {channel:?}"); + } + } + } + #[test] fn earth_month_comes_from_the_setting() { let with = |extra: &str| -> SceneFile { diff --git a/crates/resolved/src/camera.rs b/crates/resolved/src/camera.rs index 1f7251d..b9d5eed 100644 --- a/crates/resolved/src/camera.rs +++ b/crates/resolved/src/camera.rs @@ -146,7 +146,8 @@ impl Detector { } /// How the optics spread light from one point over the focal plane: the point spread -/// function (PSF). Applied from camera tier C1 up. +/// function (PSF). Applied from camera tier C1 up; the Airy and measured PSFs, whose +/// wings reach far, need tier C2, which convolves by FFT ([`crate::psf`]). #[derive(Clone, Debug, Default, PartialEq)] pub enum Psf { /// None: every point lands at one point, as through a pinhole. @@ -159,6 +160,33 @@ pub enum Psf { /// red. sigma: Vec, }, + /// Diffraction by a circular aperture of the camera's f-number: the Airy pattern, + /// averaged over each channel's band as sunlight fills it. From camera C2. + Airy { + /// The diameter of a centered obstruction, such as a reflecting telescope's + /// secondary mirror, as a fraction of the aperture's: 0 for none, below 1. + obstruction: f64, + }, + /// A PSF measured in a lab or on orbit, the same across the field. From camera C2. + Measured { + /// The spot, centered on the image's center and scaled to sum to 1. + image: PsfImage, + /// The image's sample spacing in pixels: 0.25 for four samples across a pixel. + pixel_scale: f64, + }, +} + +/// A measured PSF's image: one per channel, or one for every channel. Each sample is +/// the light that falls on its square, so the image need not be normalized. +#[derive(Clone, Debug, PartialEq)] +pub struct PsfImage { + /// Samples per row. + pub width: u32, + /// Rows. + pub height: u32, + /// Per channel, or one image for all, `width` x `height` samples, rows from the + /// top. + pub channels: Vec>, } /// The largest Gaussian PSF sigma, in pixels. At 4 x 4 finer pixels it reaches 240 of @@ -167,18 +195,24 @@ pub enum Psf { pub const MAX_PSF_SIGMA: f64 = 15.0; impl Psf { - /// How far the PSF reaches in pixels `1 / scale` across: 4 sigma of the widest - /// channel, rounded up, beyond which a Gaussian holds less than 1e-4 of its light. - /// 0 for none. + /// How far camera C1's blur reaches in pixels `1 / scale` across: 4 sigma of the + /// widest channel, rounded up, beyond which a Gaussian holds less than 1e-4 of its + /// light. 0 for none, and for the Airy and measured PSFs, which C1 does not apply + /// (camera C2 sets its own reach, [`crate::psf::Layout`]). pub fn radius(&self, scale: u32) -> u32 { match self { - Psf::None => 0, Psf::Gaussian { sigma } => { (4.0 * scale as f64 * sigma.iter().fold(0.0f64, |a, &b| a.max(b))).ceil() as u32 } + _ => 0, } } + /// Whether camera C1's separable blur can apply this PSF: none or a Gaussian. + pub fn is_separable(&self) -> bool { + matches!(self, Psf::None | Psf::Gaussian { .. }) + } + /// The weights of channel `channel`'s blur along one axis over pixels `1 / scale` /// across, for offsets `-radius(scale)` to `radius(scale)` of them: the Gaussian at /// each offset, scaled to sum to 1. A 2D blur is this along rows, then along @@ -372,6 +406,7 @@ impl CameraConfig { pub fn validate(&self) -> Result<(), Error> { let fail = |what: &str| Err(Error::InvalidCamera(what.to_owned())); let positive = |x: f64| x.is_finite() && x > 0.0; + let nonnegative = |x: f64| x.is_finite() && x >= 0.0; if self.width == 0 || self.height == 0 { return fail("image size must be nonzero"); } @@ -397,6 +432,29 @@ impl CameraConfig { return fail("PSF sigma must be positive and at most 15 pixels"); } } + if let Psf::Airy { obstruction } = self.psf { + if !(0.0..1.0).contains(&obstruction) { + return fail("the Airy PSF's obstruction must be at least 0 and below 1"); + } + } + if let Psf::Measured { image, pixel_scale } = &self.psf { + let n = image.width as usize * image.height as usize; + if n == 0 || ![1, self.channels.len()].contains(&image.channels.len()) { + return fail( + "a measured PSF needs a nonempty image, one per channel or one for all", + ); + } + if !image.channels.iter().all(|c| { + c.len() == n && c.iter().all(|&x| nonnegative(x)) && c.iter().sum::() > 0.0 + }) { + return fail( + "a measured PSF's samples must be width x height, nonnegative and not all 0", + ); + } + if !positive(*pixel_scale) { + return fail("a measured PSF's pixel scale must be positive"); + } + } if let Some(distortion) = &self.distortion { let b = distortion; if ![b.k1, b.k2, b.k3, b.p1, b.p2].iter().all(|x| x.is_finite()) { @@ -413,7 +471,6 @@ impl CameraConfig { if !(1..=16).contains(&d.adc_bits) { return fail("ADC bits must be between 1 and 16"); } - let nonnegative = |x: f64| x.is_finite() && x >= 0.0; if !nonnegative(d.read_noise) || !nonnegative(d.dark_current) { return fail("read noise and dark current must be nonnegative"); } diff --git a/crates/resolved/src/fidelity.rs b/crates/resolved/src/fidelity.rs index b5145b2..7be1243 100644 --- a/crates/resolved/src/fidelity.rs +++ b/crates/resolved/src/fidelity.rs @@ -56,6 +56,13 @@ pub enum CameraTier { /// [`RendererOptions::noise_seed`](crate::RendererOptions::noise_seed) and the /// frame number. More of the camera joins this tier as it is modeled. C1, + /// C1, plus PSFs whose wings reach far: diffraction by the aperture + /// ([`Psf::Airy`](crate::Psf::Airy)) and measured spots + /// ([`Psf::Measured`](crate::Psf::Measured)). The focal plane, with a margin + /// as wide as the PSF reaches ([`psf::Layout`](crate::psf::Layout)), is convolved + /// with the PSF by FFT, Gaussian PSFs included, and stars are drawn into it before + /// the convolution, so they take the same PSF. Several times slower than C1. + C2, } /// The lighting and camera tiers a renderer draws at. @@ -81,6 +88,7 @@ impl fmt::Display for CameraTier { f.write_str(match self { Self::C0 => "C0", Self::C1 => "C1", + Self::C2 => "C2", }) } } @@ -101,12 +109,13 @@ impl FromStr for LightingTier { impl FromStr for CameraTier { type Err = String; - /// `C0` or `C1`, either case. + /// `C0`, `C1` or `C2`, either case. fn from_str(s: &str) -> Result { match s.to_ascii_uppercase().as_str() { "C0" => Ok(Self::C0), "C1" => Ok(Self::C1), - _ => Err(format!("unknown camera tier {s:?}: C0 or C1")), + "C2" => Ok(Self::C2), + _ => Err(format!("unknown camera tier {s:?}: C0, C1 or C2")), } } } @@ -121,7 +130,8 @@ mod tests { assert_eq!(LightingTier::L1.to_string(), "L1"); assert_eq!("C0".parse::().unwrap(), CameraTier::C0); assert_eq!("c1".parse::().unwrap(), CameraTier::C1); - assert!(CameraTier::C0 < CameraTier::C1); + assert_eq!("c2".parse::().unwrap(), CameraTier::C2); + assert!(CameraTier::C0 < CameraTier::C1 && CameraTier::C1 < CameraTier::C2); assert!("L9".parse::().is_err()); assert!(LightingTier::L0 < LightingTier::L1); assert_eq!(Fidelity::default().lighting, LightingTier::L0); diff --git a/crates/resolved/src/gpu/blur.rs b/crates/resolved/src/gpu/blur.rs index 85dbd0a..9497520 100644 --- a/crates/resolved/src/gpu/blur.rs +++ b/crates/resolved/src/gpu/blur.rs @@ -101,6 +101,7 @@ impl BlurPass { }; let source = [ include_str!("../shaders/distortion.wgsl"), + include_str!("../shaders/focal.wgsl"), include_str!("../shaders/blur.wgsl"), ] .concat(); diff --git a/crates/resolved/src/gpu/convolve.rs b/crates/resolved/src/gpu/convolve.rs new file mode 100644 index 0000000..9f7f602 --- /dev/null +++ b/crates/resolved/src/gpu/convolve.rs @@ -0,0 +1,378 @@ +//! Camera C2's focal plane: the rendered radiance convolved with the PSF by FFT and +//! averaged into the image's pixels, in place of the blur pass ([`super::blur`]). +//! +//! The kernels come from [`crate::psf::kernels`] and the grid from +//! [`crate::psf::Layout`]. The kernels' spectra are computed once, on the GPU, when the +//! pass is made. Each frame then takes, for each pair of channels, a forward and an +//! inverse transform of the grid ([`super::fft`]), so two for up to two channels and +//! four for three or four. Stars are added to the focal plane before the transform, so +//! they take the same PSF as everything else (`convolve.wgsl`). +//! +//! The transforms are in f32, so their round-off scales with the brightest value in +//! the grid: a pixel of the Sun's disk leaves an error of a few parts in 10^8 of its +//! value across the whole frame. + +use nalgebra::Vector2; + +use super::fft::Fft2d; +use super::profile::PassTimestamps; +use super::stars::StarDraw; +use super::targets::FrameTargets; +use super::warp::WarpUniforms; +use super::GpuContext; +use crate::camera::Intrinsics; +use crate::psf::{Kernel, Layout}; +use crate::Error; + +/// The shader's uniform block; the layout matches `Params` in `convolve.wgsl`. +#[repr(C)] +#[derive(Clone, Copy, Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable)] +struct Params { + grid: [u32; 2], + focal: [u32; 2], + scale: u32, + margin: u32, + warp: u32, + pair: u32, + stars: u32, + channels: u32, + unused: [u32; 2], +} + +/// The compiled pipelines, the transforms, the grids and the kernels' spectra. +#[derive(Debug)] +pub struct ConvolvePass { + params: Params, + fft: Fft2d, + load: wgpu::ComputePipeline, + splat: wgpu::ComputePipeline, + multiply: wgpu::ComputePipeline, + resolve: wgpu::ComputePipeline, + /// One grid per pair of channels. + grids: Vec, + scratch: wgpu::Buffer, + /// Per pair of channels, their kernels' spectra, packed as the grids are. + spectra: Vec, + /// [`Params`] per pair of channels. + uniforms: Vec, + /// [`WarpUniforms`], zero without distortion. + warp: wgpu::Buffer, + /// Stars as `Star` in `convolve.wgsl`, a buffer that grows as needed. + stars: wgpu::Buffer, +} + +impl ConvolvePass { + /// Compiles the shaders and transforms `kernels`, one per channel, for an image of + /// `size` pixels rendered `scale` times finer and placed in the grid by `layout`. + /// With `warp`, the render is undistorted and read through the distortion. Errors + /// if a grid is larger than the device can bind. + pub fn new( + ctx: &GpuContext, + kernels: &[Kernel], + layout: &Layout, + scale: u32, + size: [u32; 2], + warp: Option<&WarpUniforms>, + ) -> Result { + use wgpu::util::DeviceExt; + let fft = Fft2d::new(ctx, layout.grid[0], layout.grid[1]); + let limits = ctx.device.limits(); + let largest = limits + .max_storage_buffer_binding_size + .min(limits.max_buffer_size); + if fft.size() > largest { + return Err(Error::Unsupported(format!( + "camera C2's FFT grid of {} x {} needs {} MB per buffer, more than {} can bind ({} MB)", + layout.grid[0], + layout.grid[1], + fft.size() >> 20, + ctx.info.name, + largest >> 20 + ))); + } + let source = [ + include_str!("../shaders/distortion.wgsl"), + include_str!("../shaders/focal.wgsl"), + include_str!("../shaders/convolve.wgsl"), + ] + .concat(); + let module = ctx + .device + .create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("convolve"), + source: wgpu::ShaderSource::Wgsl(source.into()), + }); + let pipeline = |entry| { + ctx.device + .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some(entry), + layout: None, + module: &module, + entry_point: Some(entry), + compilation_options: Default::default(), + cache: None, + }) + }; + let grid = |label| { + ctx.device.create_buffer(&wgpu::BufferDescriptor { + label: Some(label), + size: fft.size(), + usage: wgpu::BufferUsages::STORAGE, + mapped_at_creation: false, + }) + }; + let channels = kernels.len(); + let pairs = channels.div_ceil(2); + let [width, height] = layout.grid; + let scratch = grid("convolve scratch"); + // Each pair's kernels, packed and wrapped so that offset (0, 0) is at the + // grid's corner, then transformed. + let spectra: Vec = (0..pairs) + .map(|pair| { + let mut values = vec![[0.0f32; 2]; (width * height) as usize]; + for (part, kernel) in kernels[2 * pair..channels.min(2 * pair + 2)] + .iter() + .enumerate() + { + let r = kernel.radius as i64; + for dy in -r..=r { + for dx in -r..=r { + let x = dx.rem_euclid(width as i64) as usize; + let y = dy.rem_euclid(height as i64) as usize; + values[y * width as usize + x][part] = kernel.at(dx, dy) as f32; + } + } + } + ctx.device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("kernel spectrum"), + contents: bytemuck::cast_slice(&values), + usage: wgpu::BufferUsages::STORAGE, + }) + }) + .collect(); + let mut encoder = ctx.device.create_command_encoder(&Default::default()); + { + let mut pass = encoder.begin_compute_pass(&Default::default()); + for spectrum in &spectra { + fft.encode(ctx, &mut pass, spectrum, &scratch, false); + } + } + ctx.queue.submit([encoder.finish()]); + let params = Params { + grid: layout.grid, + focal: size.map(|n| n * scale + 2 * layout.margin), + scale, + margin: layout.margin, + warp: warp.is_some().into(), + pair: 0, + stars: 0, + channels: channels as u32, + unused: [0; 2], + }; + let buffer = |label, contents: &[u8], usage| { + ctx.device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some(label), + contents, + usage, + }) + }; + Ok(Self { + params, + load: pipeline("load"), + splat: pipeline("splat"), + multiply: pipeline("multiply"), + resolve: pipeline("resolve"), + grids: (0..pairs).map(|_| grid("convolve grid")).collect(), + scratch, + spectra, + uniforms: (0..pairs) + .map(|pair| { + buffer( + "convolve params", + bytemuck::bytes_of(&Params { + pair: pair as u32, + ..params + }), + wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, + ) + }) + .collect(), + warp: buffer( + "convolve warp", + bytemuck::bytes_of(&warp.copied().unwrap_or_default()), + wgpu::BufferUsages::UNIFORM, + ), + stars: Self::star_buffer(ctx, 1024), + fft, + }) + } + + /// How far the PSF reaches, in image pixels, rounded up. + pub fn reach(&self) -> u32 { + self.params.margin.div_ceil(self.params.scale) + } + + fn star_buffer(ctx: &GpuContext, capacity: u64) -> wgpu::Buffer { + ctx.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("convolve stars"), + size: capacity * std::mem::size_of::() as u64, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }) + } + + /// Records convolving `targets`' focal plane, with `stars` (from [`star_light`]) + /// added, and averaging it into their blurred image. The targets need the scale and + /// margin this pass was made for. + pub fn encode( + &mut self, + ctx: &GpuContext, + encoder: &mut wgpu::CommandEncoder, + targets: &FrameTargets, + stars: &[StarDraw], + timestamps: Option, + ) { + let needed = stars.len().max(1) as u64 * std::mem::size_of::() as u64; + if needed > self.stars.size() { + self.stars = Self::star_buffer(ctx, (stars.len() as u64).next_power_of_two()); + } + if !stars.is_empty() { + ctx.queue + .write_buffer(&self.stars, 0, bytemuck::cast_slice(stars)); + } + for (pair, uniforms) in self.uniforms.iter().enumerate() { + let params = Params { + pair: pair as u32, + stars: stars.len() as u32, + ..self.params + }; + ctx.queue + .write_buffer(uniforms, 0, bytemuck::bytes_of(¶ms)); + } + let image = &targets + .blurred + .as_ref() + .expect("targets allocated with a margin")[1]; + let source = targets.radiance.create_view(&Default::default()); + let depth = targets.depth.create_view(&wgpu::TextureViewDescriptor { + aspect: wgpu::TextureAspect::DepthOnly, + ..Default::default() + }); + let image_view = image.create_view(&Default::default()); + let bind = |pipeline: &wgpu::ComputePipeline, entries: &[(u32, wgpu::BindingResource)]| { + let entries: Vec = entries + .iter() + .map(|(binding, resource)| wgpu::BindGroupEntry { + binding: *binding, + resource: resource.clone(), + }) + .collect(); + ctx.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("convolve"), + layout: &pipeline.get_bind_group_layout(0), + entries: &entries, + }) + }; + let [width, height] = self.params.grid; + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("convolve"), + timestamp_writes: timestamps.map(PassTimestamps::compute), + }); + for ((grid, spectrum), uniforms) in self.grids.iter().zip(&self.spectra).zip(&self.uniforms) + { + let load = bind( + &self.load, + &[ + (0, wgpu::BindingResource::TextureView(&source)), + (1, grid.as_entire_binding()), + (2, uniforms.as_entire_binding()), + (3, self.warp.as_entire_binding()), + ], + ); + pass.set_pipeline(&self.load); + pass.set_bind_group(0, &load, &[]); + pass.dispatch_workgroups(width.div_ceil(16), height.div_ceil(16), 1); + if !stars.is_empty() { + let splat = bind( + &self.splat, + &[ + (2, uniforms.as_entire_binding()), + (5, grid.as_entire_binding()), + (6, self.stars.as_entire_binding()), + (7, wgpu::BindingResource::TextureView(&depth)), + ], + ); + pass.set_pipeline(&self.splat); + pass.set_bind_group(0, &splat, &[]); + pass.dispatch_workgroups((stars.len() as u32).div_ceil(64), 1, 1); + } + self.fft.encode(ctx, &mut pass, grid, &self.scratch, false); + let multiply = bind( + &self.multiply, + &[ + (1, grid.as_entire_binding()), + (2, uniforms.as_entire_binding()), + (4, spectrum.as_entire_binding()), + ], + ); + pass.set_pipeline(&self.multiply); + pass.set_bind_group(0, &multiply, &[]); + pass.dispatch_workgroups(width.div_ceil(16), height.div_ceil(16), 1); + self.fft.encode(ctx, &mut pass, grid, &self.scratch, true); + } + let resolve = bind( + &self.resolve, + &[ + (2, self.uniforms[0].as_entire_binding()), + (8, self.grids[0].as_entire_binding()), + (9, self.grids[self.grids.len() - 1].as_entire_binding()), + (10, wgpu::BindingResource::TextureView(&image_view)), + ], + ); + pass.set_pipeline(&self.resolve); + pass.set_bind_group(0, &resolve, &[]); + pass.dispatch_workgroups(image.width().div_ceil(16), image.height().div_ceil(16), 1); + } +} + +/// `stars` from [`super::stars::draws`], placed in the focal plane of an image with +/// intrinsics `image` rendered `scale` times finer with a margin of `margin`, and their +/// electrons turned into the radiance, summed over rendered pixels, that the sensor +/// turns into those electrons there: it undoes the off-axis falloff and the +/// distortion's stretch at the star. +pub fn star_light( + stars: &[StarDraw], + image: &Intrinsics, + electrons_per_radiance: f64, + scale: u32, + margin: u32, +) -> Vec { + let (s, m) = (scale as f64, margin as f64); + stars + .iter() + .map(|star| { + let uv = Vector2::new(star.at[0] as f64, star.at[1] as f64); + let n = image.to_normalized(&uv); + let (sight, stretch) = match &image.distortion { + Some(d) => { + let sight = d.undistort(&n).unwrap_or(n); + (sight, d.jacobian(&sight).determinant()) + } + None => (n, 1.0), + }; + let cos2 = 1.0 / (1.0 + sight.norm_squared()); + let factor = s * s * stretch / (electrons_per_radiance * cos2 * cos2); + StarDraw { + at: [ + (uv.x * s + m) as f32, + (uv.y * s + m) as f32, + star.at[2], + star.at[3], + ], + electrons: star.electrons.map(|e| (e as f64 * factor) as f32), + } + }) + .collect() +} diff --git a/crates/resolved/src/gpu/fft.rs b/crates/resolved/src/gpu/fft.rs new file mode 100644 index 0000000..7331dc3 --- /dev/null +++ b/crates/resolved/src/gpu/fft.rs @@ -0,0 +1,178 @@ +//! 2D fast Fourier transforms of complex f32 images on the GPU (`fft.wgsl`). +//! +//! Sides are powers of two up to [`MAX_SIDE`]. A transform is two passes: the rows, +//! written transposed, then the rows of that (the columns), transposed back. Each +//! pass transforms one row per workgroup in workgroup memory, which is why a side is +//! at most 4096: 32 KB of complex f32, lavapipe's limit. + +use super::GpuContext; + +/// The longest side a transform handles. +pub const MAX_SIDE: u32 = 4096; + +/// The shader's uniform block; the layout matches `Params` in `fft.wgsl`. +#[repr(C)] +#[derive(Clone, Copy, Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable)] +struct Params { + rows: u32, + scale: f32, + unused: [u32; 2], +} + +/// One row length's pipeline and its twiddle factors, forward and inverse. +#[derive(Debug)] +struct Stage { + pipeline: wgpu::ComputePipeline, + twiddles: [wgpu::Buffer; 2], +} + +impl Stage { + fn new(ctx: &GpuContext, n: u32) -> Self { + use wgpu::util::DeviceExt; + let module = ctx + .device + .create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("fft"), + source: wgpu::ShaderSource::Wgsl( + [ + &format!( + "const N: u32 = {n}u;\nconst LOG_N: u32 = {}u;\n", + n.trailing_zeros() + ), + include_str!("../shaders/fft.wgsl"), + ] + .concat() + .into(), + ), + }); + let pipeline = ctx + .device + .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("fft"), + layout: None, + module: &module, + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + let twiddles = [-1.0, 1.0].map(|sign: f64| { + let values: Vec<[f32; 2]> = (0..(n / 2).max(1)) + .map(|k| { + let angle = sign * std::f64::consts::TAU * k as f64 / n as f64; + [angle.cos() as f32, angle.sin() as f32] + }) + .collect(); + ctx.device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("fft twiddles"), + contents: bytemuck::cast_slice(&values), + usage: wgpu::BufferUsages::STORAGE, + }) + }); + Self { pipeline, twiddles } + } +} + +/// Transforms of images `width` by `height`. +#[derive(Debug)] +pub struct Fft2d { + width: u32, + height: u32, + rows: Stage, + /// The columns' stage, when they differ in length from the rows. + columns: Option, + /// [`Params`] for the first and second pass, forward then inverse. + params: [wgpu::Buffer; 4], +} + +impl Fft2d { + /// Compiles the passes for images `width` by `height`, each a power of two from 2 + /// to [`MAX_SIDE`]. + pub fn new(ctx: &GpuContext, width: u32, height: u32) -> Self { + use wgpu::util::DeviceExt; + for side in [width, height] { + assert!( + side.is_power_of_two() && (2..=MAX_SIDE).contains(&side), + "FFT sides must be powers of two from 2 to {MAX_SIDE}, not {side}" + ); + } + let total = (width as f64 * height as f64) as f32; + let params = [ + (height, 1.0), + (width, 1.0), + (height, 1.0), + (width, 1.0 / total), + ] + .map(|(rows, scale)| { + ctx.device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("fft params"), + contents: bytemuck::bytes_of(&Params { + rows, + scale, + unused: [0; 2], + }), + usage: wgpu::BufferUsages::UNIFORM, + }) + }); + Self { + width, + height, + rows: Stage::new(ctx, width), + columns: (height != width).then(|| Stage::new(ctx, height)), + params, + } + } + + /// Bytes of one complex image. + pub fn size(&self) -> u64 { + self.width as u64 * self.height as u64 * 8 + } + + /// Records a transform of `image` (complex f32, rows of `width`), leaving the + /// result in `image` and using `scratch`, of the same size, between the passes. The + /// inverse divides by the number of values, so a forward then an inverse transform + /// gives the image back. + pub fn encode( + &self, + ctx: &GpuContext, + pass: &mut wgpu::ComputePass, + image: &wgpu::Buffer, + scratch: &wgpu::Buffer, + inverse: bool, + ) { + let direction = usize::from(inverse); + let columns = self.columns.as_ref().unwrap_or(&self.rows); + let steps = [ + (&self.rows, image, scratch, self.height, 0), + (columns, scratch, image, self.width, 1), + ]; + for (stage, src, dst, groups, step) in steps { + let bind_group = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("fft"), + layout: &stage.pipeline.get_bind_group_layout(0), + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: src.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: dst.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: stage.twiddles[direction].as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 3, + resource: self.params[2 * direction + step].as_entire_binding(), + }, + ], + }); + pass.set_pipeline(&stage.pipeline); + pass.set_bind_group(0, &bind_group, &[]); + pass.dispatch_workgroups(groups, 1, 1); + } + } +} diff --git a/crates/resolved/src/gpu/mod.rs b/crates/resolved/src/gpu/mod.rs index fb39241..044d5f3 100644 --- a/crates/resolved/src/gpu/mod.rs +++ b/crates/resolved/src/gpu/mod.rs @@ -3,7 +3,9 @@ pub mod blur; pub mod bodies; pub mod context; +pub mod convolve; pub mod earthshine; +pub mod fft; pub mod mesh; pub mod profile; pub mod sensor; diff --git a/crates/resolved/src/gpu/profile.rs b/crates/resolved/src/gpu/profile.rs index 23c272c..6cf5b7a 100644 --- a/crates/resolved/src/gpu/profile.rs +++ b/crates/resolved/src/gpu/profile.rs @@ -80,9 +80,11 @@ pub struct PassTimings { /// Target meshes. pub mesh: f64, /// Lens distortion, the point spread function and averaging finer rendering into - /// pixels (zero work below camera C1). + /// pixels (zero work below camera C1); from camera C2, the convolution by FFT, + /// stars included. pub blur: f64, - /// Stars (zero work below camera C1 or without a catalog). + /// Stars (zero work below camera C1 or without a catalog, and at C2, where the + /// convolution draws them). pub stars: f64, /// Radiance to DN. pub sensor: f64, diff --git a/crates/resolved/src/gpu/stars.rs b/crates/resolved/src/gpu/stars.rs index a7666ac..e7bce63 100644 --- a/crates/resolved/src/gpu/stars.rs +++ b/crates/resolved/src/gpu/stars.rs @@ -134,11 +134,12 @@ impl StarPass { compilation_options: Default::default(), cache: None, }); + // Camera C2 draws stars into the focal plane before its convolution instead. let sigma = match &camera.psf { - Psf::None => [0.0; 4], Psf::Gaussian { sigma } => { std::array::from_fn(|k| sigma.get(k).copied().unwrap_or(0.0) as f32) } + _ => [0.0; 4], }; Self { pipeline, diff --git a/crates/resolved/src/gpu/warp.rs b/crates/resolved/src/gpu/warp.rs index 7499507..d3c44ce 100644 --- a/crates/resolved/src/gpu/warp.rs +++ b/crates/resolved/src/gpu/warp.rs @@ -36,7 +36,7 @@ impl DistortionUniform { } /// What the focal-plane pass needs to read the render through the distortion; the -/// layout matches `Warp` in `blur.wgsl`. +/// layout matches `Warp` in `focal.wgsl`. #[repr(C)] #[derive(Clone, Copy, Debug, Default, PartialEq, bytemuck::Pod, bytemuck::Zeroable)] pub struct WarpUniforms { diff --git a/crates/resolved/src/lib.rs b/crates/resolved/src/lib.rs index f185c95..613d55b 100644 --- a/crates/resolved/src/lib.rs +++ b/crates/resolved/src/lib.rs @@ -30,6 +30,7 @@ pub mod gpu; pub mod model; pub mod moon; pub mod preview; +pub mod psf; pub mod radiometry; pub mod render; pub mod scene; @@ -39,7 +40,9 @@ pub mod spectrum; pub mod stars; pub mod sun; -pub use camera::{Blooming, BrownConrady, CameraConfig, Channel, Detector, Intrinsics, Psf}; +pub use camera::{ + Blooming, BrownConrady, CameraConfig, Channel, Detector, Intrinsics, Psf, PsfImage, +}; pub use earth::EarthTexture; pub use error::Error; pub use fidelity::{CameraTier, Fidelity, LightingTier}; diff --git a/crates/resolved/src/psf.rs b/crates/resolved/src/psf.rs new file mode 100644 index 0000000..f84436c --- /dev/null +++ b/crates/resolved/src/psf.rs @@ -0,0 +1,562 @@ +//! Camera tier C2's point spread function: the kernel it convolves the focal plane with, +//! and the FFT grid that convolution runs on ([`crate::gpu::convolve`]). +//! +//! The focal plane is rendered `scale` times finer than the pixels, and each rendered +//! pixel holds the mean radiance over its square. Treating that radiance as constant +//! over the square, the light one rendered pixel receives from another is the PSF +//! integrated over both squares. That depends only on their offset `d`, as the PSF +//! averaged over a tent two rendered pixels wide on each axis (a square's width +//! convolved with itself): +//! +//! ```text +//! K(d) = ∫∫ PSF(d + u) (1 - |u_x|) (1 - |u_y|) du, |u_x|, |u_y| < 1 +//! ``` +//! +//! with the PSF per rendered pixel's area and `d` and `u` in rendered pixels. +//! [`kernels`] computes it in f64 on the CPU once, when a renderer is made. +//! +//! The kernel reaches [`Layout::margin`] rendered pixels from its center, and the +//! focal plane is rendered with that margin, so light from just outside the frame +//! reaches into it. The convolution is circular over the grid, and the grid holds the +//! focal plane with its margin, so light carried past one edge of the grid would wrap +//! around to the other; light from the frame reaches at most the margin past its edge, +//! which is inside the grid, and nothing wraps into the frame. + +use crate::camera::{CameraConfig, Psf}; +use crate::gpu::fft::MAX_SIDE; +use crate::radiometry::{solar_spectrum, PLANCK, SPEED_OF_LIGHT}; +use crate::Error; + +/// The least reach of C2's kernel in pixels unless the renderer asks for more +/// ([`RendererOptions::psf_reach`](crate::RendererOptions::psf_reach)). The rest of the +/// FFT grid's padding is used too, so it is usually more: 64 pixels for 1920 x 1080 at +/// 2 x 2. +pub const DEFAULT_REACH: u32 = 16; + +/// Where the focal plane lies in camera C2's FFT grid. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Layout { + /// The grid's width and height, powers of two. + pub grid: [u32; 2], + /// Rendered pixels past the image on each side, and how far the kernel reaches. + pub margin: u32, +} + +impl Layout { + /// The grid for an image `width` x `height` rendered `scale` times finer: the + /// smallest that holds the image with a margin of `reach` pixels, whose margin is + /// then as wide as the narrower side of the grid allows. Errors if a side would pass + /// [`MAX_SIDE`]. + pub fn new(width: u32, height: u32, scale: u32, reach: u32) -> Result { + let side = |n: u32, s: u64| (n as u64 + 2 * reach as u64) * s; + let image = [width, height].map(|n| n as u64 * scale as u64); + let grid = [width, height].map(|n| side(n, scale as u64).next_power_of_two()); + if grid.iter().any(|&n| n > MAX_SIDE as u64) { + let hint = match (1..scale as u64).rev().find(|&s| { + [width, height] + .iter() + .all(|&n| side(n, s) <= MAX_SIDE as u64) + }) { + Some(s) => format!("; supersampling {s} fits"), + None => String::new(), + }; + return Err(Error::Unsupported(format!( + "camera C2 convolves the focal plane by FFT, at most {MAX_SIDE} rendered pixels \ + across: {width} x {height} pixels at {scale} x {scale} with a reach of \ + {reach} pixels needs {} x {}{hint}", + grid[0], grid[1] + ))); + } + let margin = (0..2).map(|i| (grid[i] - image[i]) / 2).min().unwrap_or(0) as u32; + Ok(Self { + grid: grid.map(|n| n as u32), + margin, + }) + } +} + +/// One channel's kernel over rendered pixels, `2 radius + 1` on a side, rows from the +/// top, with the offset (0, 0) at its center. It sums to 1. +#[derive(Clone, Debug, PartialEq)] +pub struct Kernel { + /// How far it reaches from its center, in rendered pixels. + pub radius: u32, + /// The values, row by row. + pub values: Vec, +} + +impl Kernel { + /// The value at offset (`dx`, `dy`) rendered pixels, 0 beyond the radius. + pub fn at(&self, dx: i64, dy: i64) -> f64 { + let r = self.radius as i64; + if dx.abs() > r || dy.abs() > r { + return 0.0; + } + self.values[((dy + r) * (2 * r + 1) + dx + r) as usize] + } + + fn normalized(radius: u32, mut values: Vec) -> Self { + let sum: f64 = values.iter().sum(); + values.iter_mut().for_each(|v| *v /= sum); + Self { radius, values } + } +} + +/// Per channel, `camera`'s PSF as camera C2's kernel over pixels rendered `scale` times +/// finer, reaching at most `reach` of them. Kernels are scaled to sum to 1, so light +/// beyond the reach is added to the rest rather than lost: for the Airy pattern, about +/// `2 λ N / (π^2 r)` of it at a reach of r. +pub fn kernels(camera: &CameraConfig, scale: u32, reach: u32) -> Vec { + let channels = camera.channels.len(); + match &camera.psf { + Psf::None => vec![ + Kernel { + radius: 0, + values: vec![1.0], + }; + channels + ], + Psf::Gaussian { sigma } => sigma + .iter() + .map(|&s| gaussian(s * scale as f64, reach)) + .collect(), + Psf::Airy { obstruction } => (0..channels) + .map(|k| { + let profile = AiryProfile::new(camera, k, *obstruction); + let pixel = camera.pixel_pitch / scale as f64; + radial(|r| profile.at(r * pixel) * pixel * pixel, reach) + }) + .collect(), + Psf::Measured { image, pixel_scale } => (0..channels) + .map(|k| { + let values = &image.channels[k.min(image.channels.len() - 1)]; + measured( + values, + image.width, + image.height, + pixel_scale * scale as f64, + reach, + ) + }) + .collect(), + } +} + +/// A channel's diffraction spot: the Airy pattern of the camera's aperture, averaged +/// over the channel's band, weighted by the photons sunlight gives it at each +/// wavelength. A channel's spot is not the spot at its center wavelength: the pattern +/// scales with the wavelength, so the rings of a wide band blur together. +#[derive(Clone, Debug, PartialEq)] +pub struct AiryProfile { + f_number: f64, + obstruction: f64, + /// Wavelength in meters and weight, the weights summing to 1. + bands: Vec<(f64, f64)>, +} + +impl AiryProfile { + /// Wavelengths are taken 4 nm apart, each weighted by the integral of sunlight + /// through the optics and the channel's QE, in photons, over its 4 nm; bands under + /// 1e-6 of the largest are left out. + pub fn new(camera: &CameraConfig, channel: usize, obstruction: f64) -> Self { + const STEP: f64 = 4e-9; + let qe = &camera.channels[channel].quantum_efficiency; + let spectra = [solar_spectrum(), &camera.transmission, qe]; + let (lo, hi) = spectra + .iter() + .filter_map(|s| s.support()) + .reduce(|(a, b), (c, d)| (a.max(c), b.min(d))) + .expect("the solar spectrum is bounded"); + let n = ((hi - lo) / STEP).ceil().max(1.0) as usize; + let mut bands: Vec<(f64, f64)> = (0..n) + .map(|i| { + let start = lo + i as f64 * STEP; + let weight: f64 = (0..4) + .map(|j| { + let l = start + (j as f64 + 0.5) * STEP / 4.0; + spectra.iter().map(|s| s.at(l)).product::() * l + / (PLANCK * SPEED_OF_LIGHT) + }) + .sum(); + (start + STEP / 2.0, weight) + }) + .collect(); + let largest = bands.iter().fold(0.0f64, |m, b| m.max(b.1)); + bands.retain(|b| b.1 >= 1e-6 * largest); + let total: f64 = bands.iter().map(|b| b.1).sum(); + bands.iter_mut().for_each(|b| b.1 /= total); + Self { + f_number: camera.f_number, + obstruction, + bands, + } + } + + /// The fraction of a point's light per square meter of focal plane at `r` meters + /// from the spot's center. It integrates to 1 over the plane. + pub fn at(&self, r: f64) -> f64 { + self.bands + .iter() + .map(|&(l, w)| w * airy(r / (l * self.f_number), self.obstruction) / (l * l)) + .sum::() + / (self.f_number * self.f_number) + } +} + +/// The Airy pattern of an aperture with a centered obstruction `e` (a fraction of its +/// diameter) at `x` = r / (λ N), per (λ N)^2 of area: it integrates to 1 over the +/// plane. Its first dark ring is at x = 1.22 without an obstruction. +pub fn airy(x: f64, e: f64) -> f64 { + let v = std::f64::consts::PI * x; + // 2 J1(v) / v, 1 at 0. + let f = |v: f64| { + if v.abs() < 1e-8 { + 1.0 + } else { + 2.0 * bessel_j1(v) / v + } + }; + let a = f(v) - e * e * f(e * v); + std::f64::consts::PI / (4.0 * (1.0 - e * e)) * a * a +} + +/// The Bessel function J1, to about 1e-8 (Numerical Recipes, `bessj1`). +fn bessel_j1(x: f64) -> f64 { + let ax = x.abs(); + if ax < 8.0 { + let y = x * x; + let num = x + * (72_362_614_232.0 + + y * (-7_895_059_235.0 + + y * (242_396_853.1 + + y * (-2_972_611.439 + y * (15_704.482_60 + y * (-30.160_366_06)))))); + let den = 144_725_228_442.0 + + y * (2_300_535_178.0 + + y * (18_583_304.74 + y * (99_447.433_94 + y * (376.999_139_7 + y)))); + num / den + } else { + let z = 8.0 / ax; + let y = z * z; + let xx = ax - 3.0 * std::f64::consts::FRAC_PI_4; + let p = 1.0 + + y * (0.183_105e-2 + + y * (-0.351_639_649_6e-4 + y * (0.245_752_017_4e-5 + y * (-0.240_337_019e-6)))); + let q = 0.046_874_999_95 + + y * (-0.200_269_087_3e-3 + + y * (0.844_919_909_6e-5 + y * (-0.882_289_87e-6 + y * 0.105_787_412e-6))); + let ans = (std::f64::consts::FRAC_2_PI / ax).sqrt() * (xx.cos() * p - z * xx.sin() * q); + if x < 0.0 { + -ans + } else { + ans + } + } +} + +/// The midpoints and weights of an `n`-point rule for the tent 1 - |u| on (-1, 1), `n` +/// even so its peak falls between points; the weights sum to 1. +fn tent_rule(n: usize) -> Vec<(f64, f64)> { + (0..n) + .map(|i| { + let u = -1.0 + (i as f64 + 0.5) * 2.0 / n as f64; + (u, (1.0 - u.abs()) * 2.0 / n as f64) + }) + .collect() +} + +/// The kernel of a PSF that depends only on the distance from its center, `psf(r)` per +/// rendered pixel's area at r rendered pixels, out to a circle of radius `reach`. The +/// tent's integral takes 16 x 16 points within 24 pixels of the center and 6 x 6 +/// beyond, where the wings change slowly across a pixel; the PSF is read from a table +/// 1/256 of a pixel apart within 16 pixels and 1/16 beyond, linearly interpolated. +fn radial(psf: impl Fn(f64) -> f64, reach: u32) -> Kernel { + const NEAR: f64 = 16.0; + let (fine, coarse) = (256.0, 16.0); + let far = reach as f64 * std::f64::consts::SQRT_2 + 2.0; + let near_table: Vec = (0..=(NEAR * fine) as usize + 1) + .map(|i| psf(i as f64 / fine)) + .collect(); + let far_table: Vec = (0..=(far * coarse) as usize + 1) + .map(|i| psf(i as f64 / coarse)) + .collect(); + let lookup = |r: f64| { + let (table, step) = if r < NEAR { + (&near_table, fine) + } else { + (&far_table, coarse) + }; + let x = r * step; + let i = (x as usize).min(table.len() - 2); + let t = x - i as f64; + table[i] * (1.0 - t) + table[i + 1] * t + }; + let (dense, sparse) = (tent_rule(16), tent_rule(6)); + let r = reach as i64; + let side = (2 * r + 1) as usize; + let mut values = vec![0.0; side * side]; + for dy in 0..=r { + for dx in dy..=r { + if dx * dx + dy * dy > r * r { + continue; + } + let rule = if ((dx * dx + dy * dy) as f64) < 24.0 * 24.0 { + &dense + } else { + &sparse + }; + let mut k = 0.0; + for &(u, wu) in rule { + for &(v, wv) in rule { + k += wu * wv * lookup((dx as f64 + u).hypot(dy as f64 + v)); + } + } + // The PSF is the same under reflections and swapping the axes. + for (x, y) in [(dx, dy), (dy, dx)] { + for (sx, sy) in [(1, 1), (-1, 1), (1, -1), (-1, -1)] { + values[((sy * y + r) * (2 * r + 1) + sx * x + r) as usize] = k; + } + } + } + } + Kernel::normalized(reach, values) +} + +/// A Gaussian of `sigma` rendered pixels, out to 8 sigma or `reach`. It is separable, +/// and so is the tent: the kernel is the product of one axis's integral with itself, +/// taken by 64 points. +fn gaussian(sigma: f64, reach: u32) -> Kernel { + let r = ((8.0 * sigma).ceil() as u32 + 1).min(reach); + let rule = tent_rule(64); + let axis: Vec = (-(r as i64)..=r as i64) + .map(|d| { + rule.iter() + .map(|&(u, w)| { + let x = (d as f64 + u) / sigma; + w * (-0.5 * x * x).exp() + }) + .sum() + }) + .collect(); + let values = axis + .iter() + .flat_map(|&y| axis.iter().map(move |&x| x * y)) + .collect(); + Kernel::normalized(r, values) +} + +/// A measured PSF: `values`, `width` x `height` samples `spacing` rendered pixels +/// apart, centered on the image's center and taken as constant over each sample's +/// square. The tent's integral over each square is exact, and separable. +fn measured(values: &[f64], width: u32, height: u32, spacing: f64, reach: u32) -> Kernel { + let r = ((width.max(height) as f64 * spacing / 2.0).ceil() as u32 + 1).min(reach); + // The integral of the tent 1 - |u| from -1 to u. + let rising = |u: f64| match u { + u if u <= -1.0 => 0.0, + u if u <= 0.0 => (u + 1.0) * (u + 1.0) / 2.0, + u if u <= 1.0 => 1.0 - (1.0 - u) * (1.0 - u) / 2.0, + _ => 1.0, + }; + // Per offset and sample, the tent at that offset integrated over the sample's + // extent along one axis. + let axis = |samples: u32| -> Vec> { + (-(r as i64)..=r as i64) + .map(|d| { + (0..samples) + .map(|i| { + let a = (i as f64 - samples as f64 / 2.0) * spacing - d as f64; + rising(a + spacing) - rising(a) + }) + .collect() + }) + .collect() + }; + let (ax, ay) = (axis(width), axis(height)); + let w = width as usize; + let side = (2 * r + 1) as usize; + let mut values_out = vec![0.0; side * side]; + for (y, row_weights) in ay.iter().enumerate() { + // The samples' rows weighted for this offset, then each column offset. + let mut row = vec![0.0; w]; + for (j, &wy) in row_weights.iter().enumerate() { + if wy != 0.0 { + for (acc, v) in row.iter_mut().zip(&values[j * w..(j + 1) * w]) { + *acc += wy * v; + } + } + } + for (x, column_weights) in ax.iter().enumerate() { + values_out[y * side + x] = column_weights.iter().zip(&row).map(|(a, b)| a * b).sum(); + } + } + Kernel::normalized(r, values_out) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::camera::{Channel, PsfImage}; + use crate::spectrum::Spectrum; + + #[test] + fn bessel_j1_matches_known_values() { + // Abramowitz and Stegun table 9.1, and J1's first zero. + for (x, want) in [ + (1.0, 0.440_050_585_7), + (5.0, -0.327_579_137_6), + (10.0, 0.043_472_746_2), + (12.0, -0.223_447_104_3), + ] { + assert!((bessel_j1(x) - want).abs() < 1e-8, "J1({x})"); + assert!((bessel_j1(-x) + want).abs() < 1e-8, "J1(-{x})"); + } + assert!(bessel_j1(3.831_705_970_2).abs() < 1e-8); + } + + #[test] + fn the_airy_pattern_holds_all_its_light() { + // Integrated out to x, the light outside is 2 / (pi^2 x) for a clear aperture + // and 1 / (1 - e) times that with an obstruction e, whose rings are brighter. + for e in [0.0, 0.3] { + let x_max: f64 = 400.0; + let n = 400_000; + let h = x_max / n as f64; + let inside: f64 = (0..n) + .map(|i| { + let x = (i as f64 + 0.5) * h; + airy(x, e) * 2.0 * std::f64::consts::PI * x * h + }) + .sum(); + let outside = 2.0 / (std::f64::consts::PI.powi(2) * x_max); + assert!( + (inside + outside / (1.0 - e) - 1.0).abs() < 1e-5, + "obstruction {e}: {inside} inside" + ); + } + } + + #[test] + fn a_narrow_channels_first_dark_ring_is_at_1_22_lambda_n() { + // A channel 2 nm wide at 550 nm, f/4: the spot's first minimum within 1 % of + // 1.22 x 550 nm x 4 = 2.684 um. + let camera = CameraConfig { + channels: vec![Channel { + name: "narrow".to_owned(), + quantum_efficiency: Spectrum::gaussian(0.6, 550e-9, 2e-9), + }], + psf: Psf::Airy { obstruction: 0.0 }, + ..CameraConfig::example() + }; + let profile = AiryProfile::new(&camera, 0, 0.0); + let want = 1.2197 * 550e-9 * 4.0; + let samples: Vec = (0..2000).map(|i| profile.at(i as f64 * 2e-9)).collect(); + let first = (1..samples.len() - 1) + .find(|&i| samples[i] < samples[i - 1] && samples[i] <= samples[i + 1]) + .unwrap() as f64 + * 2e-9; + assert!( + (first / want - 1.0).abs() < 0.01, + "first ring at {first}, expected {want}" + ); + } + + #[test] + fn kernels_sum_to_1_and_center_their_light() { + let mut camera = CameraConfig::example(); + for psf in [ + Psf::Airy { obstruction: 0.2 }, + Psf::Gaussian { + sigma: vec![0.7, 0.6, 0.5], + }, + Psf::Measured { + image: PsfImage { + width: 5, + height: 5, + channels: vec![(0..25) + .map(|i| { + [0.0, 1.0, 2.0, 1.0, 0.0][i % 5] * [0.0, 1.0, 2.0, 1.0, 0.0][i / 5] + }) + .collect()], + }, + pixel_scale: 0.5, + }, + ] { + camera.psf = psf; + for k in kernels(&camera, 2, 20) { + let r = k.radius as i64; + let sum: f64 = k.values.iter().sum(); + assert!((sum - 1.0).abs() < 1e-12, "{:?}: sum {sum}", camera.psf); + let (mut mx, mut my) = (0.0, 0.0); + for dy in -r..=r { + for dx in -r..=r { + mx += dx as f64 * k.at(dx, dy); + my += dy as f64 * k.at(dx, dy); + } + } + assert!( + mx.abs() < 1e-12 && my.abs() < 1e-12, + "{:?}: centroid {mx}, {my}", + camera.psf + ); + } + } + } + + #[test] + fn a_measured_psf_on_the_grid_is_its_samples_blurred_by_the_tent() { + // One sample a rendered pixel square, centered on a rendered pixel: the tent + // spreads its light 3/4 to its own pixel and 1/8 to each side along each axis. + let k = measured(&[1.0], 1, 1, 1.0, 4); + let along = [1.0 / 8.0, 3.0 / 4.0, 1.0 / 8.0]; + for dy in -1..=1i64 { + for dx in -1..=1i64 { + let want = along[(dx + 1) as usize] * along[(dy + 1) as usize]; + assert!((k.at(dx, dy) - want).abs() < 1e-15, "({dx}, {dy})"); + } + } + assert_eq!(k.at(2, 0), 0.0); + } + + #[test] + fn the_grid_fits_the_image_and_a_margin() { + let layout = |w, h, s, r| Layout::new(w, h, s, r).unwrap(); + let l = layout(1920, 1080, 2, DEFAULT_REACH); + assert_eq!( + l, + Layout { + grid: [4096, 4096], + margin: 128 + } + ); + let l = layout(3840, 2160, 1, DEFAULT_REACH); + assert_eq!( + l, + Layout { + grid: [4096, 4096], + margin: 128 + } + ); + let l = layout(256, 128, 2, DEFAULT_REACH); + assert_eq!( + l, + Layout { + grid: [1024, 512], + margin: 128 + } + ); + // At 1 x 1, a reach of 500 px takes the same grid as 1920 x 1080 at 2 x 2. + let l = layout(1920, 1080, 1, 500); + assert_eq!( + l, + Layout { + grid: [4096, 4096], + margin: 1088 + } + ); + let e = Layout::new(3840, 2160, 2, DEFAULT_REACH) + .unwrap_err() + .to_string(); + assert!(e.contains("supersampling 1 fits"), "{e}"); + let e = Layout::new(1920, 1080, 2, 500).unwrap_err().to_string(); + assert!(e.contains("supersampling 1 fits"), "{e}"); + } +} diff --git a/crates/resolved/src/render.rs b/crates/resolved/src/render.rs index 6bda14f..30e3fb6 100644 --- a/crates/resolved/src/render.rs +++ b/crates/resolved/src/render.rs @@ -21,6 +21,7 @@ use crate::fidelity::{CameraTier, Fidelity, LightingTier}; use crate::fixed_pattern::FixedPattern; use crate::gpu::blur::BlurPass; use crate::gpu::bodies::{BodiesPass, BodiesUniforms, BodyRadiometry}; +use crate::gpu::convolve::{star_light, ConvolvePass}; use crate::gpu::earthshine::{EarthshinePass, TargetLight}; use crate::gpu::mesh::{GpuMesh, MeshDraw, MeshPass, ShadowDraw}; use crate::gpu::profile::{PassTimings, Profiler}; @@ -32,6 +33,7 @@ use crate::gpu::warp::{WarpLayout, WarpUniforms}; use crate::gpu::{AdapterChoice, GpuContext}; use crate::model::{Mesh, ModelId}; use crate::moon::MoonTexture; +use crate::psf; use crate::radiometry::Radiometry; use crate::scene::SPEED_OF_LIGHT; use crate::scene::{Pose, Scene}; @@ -87,6 +89,13 @@ pub struct RendererOptions { /// in place of the one its fields generate. It must match the camera's size and /// channels. pub fixed_pattern: Option>, + /// From camera C2, how far in pixels the PSF must reach from a point: `None` for + /// [`psf::DEFAULT_REACH`]. The FFT grid is the smallest that holds the image with + /// this margin on each side, and the PSF then reaches as far as the grid allows + /// ([`psf::Layout`]): at 1920 x 1080, 64 pixels at 2 x 2 supersampling, and up to + /// 1088 at 1 x 1 in the same 4096 x 4096 grid. Light the PSF carries past its reach + /// is added to the rest of the spot. + pub psf_reach: Option, } /// One rendered image. @@ -136,6 +145,8 @@ pub struct Renderer { /// Whether the frame is distorted (camera C1 and a lens with distortion). distorted: bool, blur: Option, + /// From camera C2 with a PSF: the focal plane's convolution, in place of the blur. + convolve: Option, /// The pinhole projection the meshes and bodies render with: finer than the image /// by the scale, with the margin, and with distortion, undistorted and placed for /// the warp. @@ -200,7 +211,22 @@ impl Renderer { (true, _) => 4, }; let psf = if area { &camera.psf } else { &Psf::None }; - let margin = psf.radius(scale); + // From camera C2, the PSF is applied by FFT, over a margin as wide as the grid + // allows; camera C1 blurs by a separable Gaussian only. + let convolved = options.fidelity.camera >= CameraTier::C2 && *psf != Psf::None; + if area && !convolved && !psf.is_separable() { + return Err(Error::Unsupported(format!( + "camera {} blurs by a Gaussian PSF only: the Airy and measured PSFs need camera C2", + options.fidelity.camera + ))); + } + let grid = if convolved { + let reach = options.psf_reach.unwrap_or(psf::DEFAULT_REACH); + Some(psf::Layout::new(camera.width, camera.height, scale, reach)?) + } else { + None + }; + let margin = grid.map_or_else(|| psf.radius(scale), |g| g.margin); // From camera C1, lens distortion: render undistorted, and read the focal plane // through the distortion. let image = camera.intrinsics(); @@ -209,14 +235,18 @@ impl Renderer { _ => None, }; let warp = warp_layout.map(|l| WarpUniforms::new(&image, &l)); - let blur = BlurPass::new( - &ctx, - psf, - camera.channels.len(), - scale, - [camera.width, camera.height], - warp.as_ref(), - ); + let blur = if convolved { + None + } else { + BlurPass::new( + &ctx, + psf, + camera.channels.len(), + scale, + [camera.width, camera.height], + warp.as_ref(), + ) + }; // Stars are scene content (lighting L1) that only a camera collecting light over // each pixel's area (C1) can see: a single point per pixel never meets one. let stars = options @@ -230,7 +260,8 @@ impl Renderer { Blooming::None }; let layout = TargetLayout { - stars: stars.is_some(), + // Camera C2 draws the stars into its convolution instead. + stars: stars.is_some() && !convolved, charge: blooming != Blooming::None, samples, scale, @@ -333,6 +364,17 @@ impl Renderer { options.fidelity.lighting ))); } + let convolve = match grid { + Some(grid) => Some(ConvolvePass::new( + &ctx, + &psf::kernels(&camera, scale, grid.margin), + &grid, + scale, + [camera.width, camera.height], + warp.as_ref(), + )?), + None => None, + }; let profiler = if options.profile { Profiler::new(&ctx) } else { @@ -342,11 +384,12 @@ impl Renderer { bodies: BodiesPass::new(&ctx, samples, area), meshes: MeshPass::new(&ctx, samples), blur, + convolve, shadows: ShadowPass::new(&ctx, options.shadow_map_size.unwrap_or(4096)), earthshine: EarthshinePass::new(&ctx), sensor: SensorPass::new( &ctx, - stars.is_some(), + layout.stars, area && camera.detector.is_noisy(), match &options.fixed_pattern { _ if !area => SensorPattern::None, @@ -621,28 +664,53 @@ impl Renderer { timestamp_writes: stamps(index).map(|t| t.compute()), })) }; - match &self.blur { - Some(blur) => blur.encode(ctx, &mut encoder, &self.targets, stamps(4)), - None => empty(&mut encoder, "blur (none)", 4), - } - match &mut self.stars { - Some((catalog, photometry, pass)) => { - // Stars are at rest relative to the Sun, to 12 m/s. - let beta = (scene.velocities.earth + scene.velocities.observer) / SPEED_OF_LIGHT; - let draws = star_draws( - catalog, - photometry, - self.camera.exposure, - &camera, - &self.camera.intrinsics(), - &self.render_intrinsics, - &beta, - pass.reach(), - ); - let timestamps = self.profiler.as_ref().map(|p| p.pass(5)); - pass.encode(ctx, &mut encoder, &self.targets, &draws, timestamps); + // Stars are at rest relative to the Sun, to 12 m/s. + let beta = (scene.velocities.earth + scene.velocities.observer) / SPEED_OF_LIGHT; + let image = self.camera.intrinsics(); + let reach = match (&self.convolve, &self.stars) { + (Some(convolve), _) => convolve.reach(), + (None, Some((_, _, pass))) => pass.reach(), + (None, None) => 0, + }; + let draws = self.stars.as_ref().map(|(catalog, photometry, _)| { + star_draws( + catalog, + photometry, + self.camera.exposure, + &camera, + &image, + &self.render_intrinsics, + &beta, + reach, + ) + }); + match &mut self.convolve { + Some(convolve) => { + let stars = draws.map_or_else(Vec::new, |draws| { + star_light( + &draws, + &image, + self.radiometry.electrons_per_radiance, + self.targets.scale, + self.targets.margin, + ) + }); + convolve.encode(ctx, &mut encoder, &self.targets, &stars, stamps(4)); + empty(&mut encoder, "stars (in the convolution)", 5); + } + None => { + match &self.blur { + Some(blur) => blur.encode(ctx, &mut encoder, &self.targets, stamps(4)), + None => empty(&mut encoder, "blur (none)", 4), + } + match (draws, &mut self.stars) { + (Some(draws), Some((_, _, pass))) => { + let timestamps = self.profiler.as_ref().map(|p| p.pass(5)); + pass.encode(ctx, &mut encoder, &self.targets, &draws, timestamps); + } + _ => empty(&mut encoder, "stars (none)", 5), + } } - None => empty(&mut encoder, "stars (none)", 5), } self.sensor_uniforms.set_key(self.noise_seed, self.frame); self.sensor.encode( diff --git a/crates/resolved/src/shaders/blur.wgsl b/crates/resolved/src/shaders/blur.wgsl index 20f9106..f44262a 100644 --- a/crates/resolved/src/shaders/blur.wgsl +++ b/crates/resolved/src/shaders/blur.wgsl @@ -13,7 +13,7 @@ // pixel, the undistorted point the lens bends onto its center, read bilinearly. // Radiance is unchanged along a ray, so moving it is all distortion does to an // extended scene; the sensor accounts for how it stretches area. The code that -// includes this file first includes distortion.wgsl. +// includes this file first includes distortion.wgsl and focal.wgsl. struct Params { // Rendered pixels per image pixel along each axis. @@ -29,15 +29,6 @@ struct Params { unused: vec2, } -struct Warp { - distortion: Distortion, - // fx, fy, cx, cy of the image, in image pixels. - image: vec4, - // xy: whole rendered pixels from a focal-plane pixel to the render pixel that sees - // the same point without distortion. - offset: vec4, -} - // Per tap, the weight of each channel. @group(0) @binding(0) var weights: array>; @group(0) @binding(1) var source: texture_2d; @@ -45,29 +36,6 @@ struct Warp { @group(0) @binding(3) var p: Params; @group(0) @binding(4) var w: Warp; -fn texel(at: vec2) -> vec4 { - let size = vec2(textureDimensions(source)); - return textureLoad(source, clamp(at, vec2(0), size - 1), 0); -} - -// The focal plane's radiance at pixel `at`, read through the distortion from the -// undistorted render. -fn warped(at: vec2) -> vec4 { - let scale = f32(p.scale); - // This pixel's center in image pixels, then normalized and distorted. - let uv = (vec2(at) + 0.5 - f32(p.margin)) / scale; - let distorted = (uv - w.image.zw) / w.image.xy; - let shift = (undistort(w.distortion, distorted) - distorted) * w.image.xy * scale; - // In the render's texels, whose centers are at whole numbers. - let x = vec2(at) + w.offset.xy + shift; - let base = floor(x); - let f = x - base; - let i = vec2(base); - let top = mix(texel(i), texel(i + vec2(1, 0)), f.x); - let bottom = mix(texel(i + vec2(0, 1)), texel(i + vec2(1, 1)), f.x); - return mix(top, bottom, f.y); -} - // Destination pixels along the axis per workgroup, set per scale (blur.rs), and the // source texels they read: LINE x scale + 2 radius, at most 256 + 2 x 240 // (MAX_PSF_SIGMA in camera.rs). @@ -87,7 +55,7 @@ fn blur(start: vec2, i: u32, axis: vec2) { let length = LINE * p.scale + 2u * p.radius; for (var j = i; j < length; j += LINE) { if p.warp == 1u { - tile[j] = warped(min(first + j * axis, p.focal - 1u)); + tile[j] = warped(min(first + j * axis, p.focal - 1u), p.scale, p.margin); } else { tile[j] = textureLoad(source, min(first + j * axis, source_size - 1u), 0); } diff --git a/crates/resolved/src/shaders/convolve.wgsl b/crates/resolved/src/shaders/convolve.wgsl new file mode 100644 index 0000000..c6388c5 --- /dev/null +++ b/crates/resolved/src/shaders/convolve.wgsl @@ -0,0 +1,186 @@ +// Camera C2's focal plane: the rendered radiance convolved with the PSF by FFT, then +// averaged into the image's pixels (convolve.rs). +// +// Two channels share one complex transform, one as the real part and one as the +// imaginary part. For each pair of channels: `load` copies the focal plane into the +// grid, through the distortion if any, with zeros past it; `splat` adds the stars; the +// grid is transformed (fft.wgsl); `multiply` multiplies each channel's spectrum by its +// kernel's; and the grid is transformed back. `resolve` then averages each image +// pixel's rendered pixels from both pairs. The code that includes this file first +// includes distortion.wgsl and focal.wgsl. + +struct Params { + // The FFT grid's width and height. + grid: vec2, + // The focal plane's width and height, from the grid's corner. + focal: vec2, + // Rendered pixels per image pixel along each axis. + scale: u32, + // Rendered pixels past the image on each side. + margin: u32, + // 1 if the render is read through the distortion. + warp: u32, + // Which pair of channels: 0 for channels 0 and 1, 1 for 2 and 3. + pair: u32, + // Number of stars. + stars: u32, + // Number of channels. + channels: u32, + unused: vec2, +} + +struct Star { + // xy: position in the focal plane's rendered pixels. zw: position in the render's + // pixels, for the depth test. + at: vec4, + // Per channel, the radiance summed over rendered pixels that gives the star's + // electrons. + light: vec4, +} + +@group(0) @binding(0) var source: texture_2d; +@group(0) @binding(1) var data: array>; +@group(0) @binding(2) var p: Params; +@group(0) @binding(3) var w: Warp; +// The kernels' spectra, the pair's two kernels packed as `data` is. +@group(0) @binding(4) var spectrum: array>; +// `data` as words, to add the stars: WGSL has no floating-point atomics. +@group(0) @binding(5) var cells: array>; +@group(0) @binding(6) var stars: array; +@group(0) @binding(7) var depth: texture_depth_2d; +// Both pairs' convolved grids, the same one twice for two channels or fewer. +@group(0) @binding(8) var first: array>; +@group(0) @binding(9) var second: array>; +@group(0) @binding(10) var image: texture_storage_2d; + +fn index(at: vec2) -> u32 { + return at.y * p.grid.x + at.x; +} + +@compute @workgroup_size(16, 16) +fn load(@builtin(global_invocation_id) id: vec3) { + let at = id.xy; + if any(at >= p.grid) { + return; + } + var v = vec4(0.0); + if all(at < p.focal) { + if p.warp == 1u { + v = warped(at, p.scale, p.margin); + } else { + v = textureLoad(source, at, 0); + } + } + var pair = select(v.xy, v.zw, p.pair == 1u); + if 2u * p.pair + 1u >= p.channels { + pair.y = 0.0; + } + data[index(at)] = pair; +} + +// Adds `x` to word `i` of the grid, read as an f32. +fn add(i: u32, x: f32) { + var old = atomicLoad(&cells[i]); + loop { + let r = atomicCompareExchangeWeak(&cells[i], old, bitcast(bitcast(old) + x)); + if r.exchanged { + break; + } + old = r.old_value; + } +} + +// Each star's light, shared among the four rendered pixels nearest it by bilinear +// weights, which keeps its centroid where the star is. Stars lie at infinity, so +// anything drawn where a star is hides it. +@compute @workgroup_size(64) +fn splat(@builtin(global_invocation_id) id: vec3) { + if id.x >= p.stars { + return; + } + let star = stars[id.x]; + let size = textureDimensions(depth); + let render = vec2(floor(star.at.zw)); + if all(render >= vec2(0)) && all(render < vec2(size)) { + if textureLoad(depth, render, 0) > 0.0 { + return; + } + } + let light = select(star.light.xy, star.light.zw, p.pair == 1u); + // Rendered pixel centers are at half pixels. + let x = star.at.xy - 0.5; + let base = floor(x); + let f = x - base; + for (var c = 0u; c < 4u; c++) { + let corner = vec2(c & 1u, c >> 1u); + let cell = vec2(base) + vec2(corner); + if any(cell < vec2(0)) || any(vec2(cell) >= p.focal) { + continue; + } + let weights = mix(1.0 - f, f, vec2(corner)); + let i = 2u * index(vec2(cell)); + add(i, light.x * weights.x * weights.y); + add(i + 1u, light.y * weights.x * weights.y); + } +} + +fn conjugate(z: vec2) -> vec2 { + return vec2(z.x, -z.y); +} + +fn product(a: vec2, b: vec2) -> vec2 { + return vec2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x); +} + +// i z. +fn rotated(z: vec2) -> vec2 { + return vec2(-z.y, z.x); +} + +// The transform X of a + i b, for real channels a and b, holds both of theirs: +// A(k) = (X(k) + conj X(-k)) / 2 and B(k) = (X(k) - conj X(-k)) / 2i, and the +// kernels' the same way. The product A Ka + i B Kb is the transform of the two +// channels convolved, packed again. Each invocation handles k and -k together, so the +// grid is changed in place. +@compute @workgroup_size(16, 16) +fn multiply(@builtin(global_invocation_id) id: vec3) { + let at = id.xy; + if any(at >= p.grid) { + return; + } + let mirror = (p.grid - at) % p.grid; + let i = index(at); + let j = index(mirror); + if i > j { + return; + } + let x = data[i]; + let xm = conjugate(data[j]); + let k = spectrum[i]; + let km = conjugate(spectrum[j]); + let a = 0.5 * (x + xm); + let b = -rotated(0.5 * (x - xm)); + let ka = 0.5 * (k + km); + let kb = -rotated(0.5 * (k - km)); + let pa = product(a, ka); + let pb = product(b, kb); + data[j] = conjugate(pa) + rotated(conjugate(pb)); + data[i] = pa + rotated(pb); +} + +@compute @workgroup_size(16, 16) +fn resolve(@builtin(global_invocation_id) id: vec3) { + let size = textureDimensions(image); + if any(id.xy >= size) { + return; + } + var sum = vec4(0.0); + let corner = p.margin + id.xy * p.scale; + for (var y = 0u; y < p.scale; y++) { + for (var x = 0u; x < p.scale; x++) { + let i = index(corner + vec2(x, y)); + sum += vec4(first[i], second[i]); + } + } + textureStore(image, id.xy, sum / f32(p.scale * p.scale)); +} diff --git a/crates/resolved/src/shaders/fft.wgsl b/crates/resolved/src/shaders/fft.wgsl new file mode 100644 index 0000000..9f2b023 --- /dev/null +++ b/crates/resolved/src/shaders/fft.wgsl @@ -0,0 +1,56 @@ +// One pass of a 2D FFT: a 1D FFT of every row of `src`, N complex values long, written +// transposed into `dst`, so a second pass over `dst` transforms the columns and +// transposes back. N and LOG_N are declared by the code that includes this file. +// +// One workgroup per row: the row is loaded in bit-reversed order into workgroup +// memory, transformed in place by LOG_N radix-2 butterfly stages (Cooley-Tukey, +// decimation in time), and written out. + +struct Params { + // Rows in `src`, which is the row length of `dst`. + rows: u32, + // A factor on every output value: 1 forward, 1 / (width x height) at the end of an + // inverse transform. + scale: f32, + unused: vec2, +} + +@group(0) @binding(0) var src: array>; +@group(0) @binding(1) var dst: array>; +// exp(-2 pi i k / N) for k < N / 2, conjugated for an inverse transform, computed in +// f64 on the CPU. +@group(0) @binding(2) var twiddles: array>; +@group(0) @binding(3) var p: Params; + +const THREADS: u32 = 256u; + +var data: array, N>; + +fn product(a: vec2, b: vec2) -> vec2 { + return vec2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x); +} + +@compute @workgroup_size(256) +fn main(@builtin(workgroup_id) group: vec3, @builtin(local_invocation_index) t: u32) { + let row = group.x; + for (var j = t; j < N; j += THREADS) { + data[reverseBits(j) >> (32u - LOG_N)] = src[row * N + j]; + } + workgroupBarrier(); + for (var s = 1u; s <= LOG_N; s++) { + let half = 1u << (s - 1u); + let stride = N >> s; + for (var k = t; k < N / 2u; k += THREADS) { + let position = k & (half - 1u); + let i = ((k >> (s - 1u)) << s) + position; + let a = data[i]; + let b = product(data[i + half], twiddles[position * stride]); + data[i] = a + b; + data[i + half] = a - b; + } + workgroupBarrier(); + } + for (var j = t; j < N; j += THREADS) { + dst[j * p.rows + row] = data[j] * p.scale; + } +} diff --git a/crates/resolved/src/shaders/focal.wgsl b/crates/resolved/src/shaders/focal.wgsl new file mode 100644 index 0000000..0b9e49a --- /dev/null +++ b/crates/resolved/src/shaders/focal.wgsl @@ -0,0 +1,37 @@ +// Reading the focal plane from an undistorted render through the lens distortion, for +// the passes that make the focal plane's image (blur.wgsl, convolve.wgsl). The code +// that includes this file first includes distortion.wgsl, and declares the render as +// `source` and these uniforms as `w`. + +struct Warp { + distortion: Distortion, + // fx, fy, cx, cy of the image, in image pixels. + image: vec4, + // xy: whole rendered pixels from a focal-plane pixel to the render pixel that sees + // the same point without distortion. + offset: vec4, +} + +fn texel(at: vec2) -> vec4 { + let size = vec2(textureDimensions(source)); + return textureLoad(source, clamp(at, vec2(0), size - 1), 0); +} + +// The focal plane's radiance at pixel `at`, read through the distortion from the +// undistorted render, for a focal plane rendered `rendered` times finer than the image +// with a margin of `margin` of its pixels. +fn warped(at: vec2, rendered: u32, margin: u32) -> vec4 { + let scale = f32(rendered); + // This pixel's center in image pixels, then normalized and distorted. + let uv = (vec2(at) + 0.5 - f32(margin)) / scale; + let distorted = (uv - w.image.zw) / w.image.xy; + let shift = (undistort(w.distortion, distorted) - distorted) * w.image.xy * scale; + // In the render's texels, whose centers are at whole numbers. + let x = vec2(at) + w.offset.xy + shift; + let base = floor(x); + let f = x - base; + let i = vec2(base); + let top = mix(texel(i), texel(i + vec2(1, 0)), f.x); + let bottom = mix(texel(i + vec2(0, 1)), texel(i + vec2(1, 1)), f.x); + return mix(top, bottom, f.y); +} diff --git a/crates/resolved/tests/gpu/convolve.rs b/crates/resolved/tests/gpu/convolve.rs new file mode 100644 index 0000000..d2b89f6 --- /dev/null +++ b/crates/resolved/tests/gpu/convolve.rs @@ -0,0 +1,322 @@ +//! Camera tier C2's convolution of the focal plane by FFT: against the kernel and a +//! direct convolution, its flux, and its round-off near a very bright source, on the +//! GPU and on the CPU (lavapipe). + +use crate::common; + +use resolved::camera::PsfImage; +use resolved::gpu::convolve::ConvolvePass; +use resolved::gpu::stars::StarDraw; +use resolved::gpu::targets::{FrameTargets, TargetLayout}; +use resolved::gpu::GpuContext; +use resolved::psf::{kernels, Kernel, Layout, DEFAULT_REACH}; +use resolved::{CameraConfig, Psf}; + +/// The image's size in the pass tests. +const WIDTH: u32 = 48; +const HEIGHT: u32 = 40; + +/// Convolves `radiance`, the focal plane of an image of `size` pixels rendered `scale` +/// times finer with the margin `Layout` gives (rows from the top), with `kernels`, adds +/// `stars`, and averages it into the image. +fn convolve( + ctx: &GpuContext, + kernels: &[Kernel], + scale: u32, + size: (u32, u32), + radiance: &[[f32; 4]], + stars: &[StarDraw], +) -> Vec<[f32; 4]> { + let grid = Layout::new(size.0, size.1, scale, DEFAULT_REACH).unwrap(); + let layout = TargetLayout { + scale, + margin: grid.margin, + ..TargetLayout::new(size.0, size.1, kernels.len() as u32) + }; + let targets = FrameTargets::new(ctx, &layout); + let (w, h) = (targets.radiance.width(), targets.radiance.height()); + assert_eq!(radiance.len(), (w * h) as usize); + ctx.queue.write_texture( + wgpu::TexelCopyTextureInfo { + texture: &targets.radiance, + mip_level: 0, + origin: wgpu::Origin3d::ZERO, + aspect: wgpu::TextureAspect::All, + }, + bytemuck::cast_slice(radiance), + wgpu::TexelCopyBufferLayout { + offset: 0, + bytes_per_row: Some(w * 16), + rows_per_image: Some(h), + }, + targets.radiance.size(), + ); + let mut pass = ConvolvePass::new(ctx, kernels, &grid, scale, [size.0, size.1], None).unwrap(); + let mut encoder = ctx.device.create_command_encoder(&Default::default()); + pass.encode(ctx, &mut encoder, &targets, stars, None); + ctx.queue.submit([encoder.finish()]); + targets.read_radiance(ctx) +} + +/// The focal plane's size in rendered pixels, and its margin. +fn focal(scale: u32, size: (u32, u32)) -> (usize, usize, usize) { + let m = Layout::new(size.0, size.1, scale, DEFAULT_REACH) + .unwrap() + .margin as usize; + let s = scale as usize; + (size.0 as usize * s + 2 * m, size.1 as usize * s + 2 * m, m) +} + +/// Kernels from lopsided measured spots, a different one per channel, so that a kernel +/// flipped or shifted shows: `n` x `n` samples `spacing` pixels apart. +fn lopsided(n: u32, spacing: f64, scale: u32, reach: u32) -> Vec { + let spot = |k: u32| -> Vec { + (0..n * n) + .map(|i| { + let (x, y) = ((i % n) as f64, (i / n) as f64); + (1.0 + x * (k + 1) as f64 + 0.3 * y * y) + * ((x * 7.0 + y * 3.0 + k as f64).sin() + 1.5) + }) + .collect() + }; + let camera = CameraConfig { + psf: Psf::Measured { + image: PsfImage { + width: n, + height: n, + channels: (0..3).map(spot).collect(), + }, + pixel_scale: spacing, + }, + ..CameraConfig::example() + }; + kernels(&camera, scale, reach) +} + +#[test] +fn a_point_of_light_spreads_as_the_kernel() { + // One rendered pixel lit, at 1 x 1: each image pixel around it reads the kernel at + // its offset, to 1e-6 of the kernel's peak, and every other pixel 0 to the same. + // The three kernels differ and none is symmetric, so a pair's channels mixed, or a + // kernel flipped, shows. + let size = (WIDTH, HEIGHT); + let (w, h, m) = focal(1, size); + let ks = lopsided(7, 1.0, 1, 64); + let (u0, v0) = (20usize, 17usize); + let mut radiance = vec![[0.0f32; 4]; w * h]; + radiance[(v0 + m) * w + u0 + m] = [1.0; 4]; + for ctx in common::contexts() { + let name = &ctx.info.name; + let image = convolve(ctx, &ks, 1, size, &radiance, &[]); + for (k, kernel) in ks.iter().enumerate() { + let peak = kernel.values.iter().fold(0.0f64, |a, &b| a.max(b)); + let worst = image + .iter() + .enumerate() + .map(|(i, p)| { + let (u, v) = ((i % WIDTH as usize) as i64, (i / WIDTH as usize) as i64); + let want = kernel.at(u - u0 as i64, v - v0 as i64); + (p[k] as f64 - want).abs() + }) + .fold(0.0f64, f64::max); + assert!( + worst < 1e-6 * peak, + "{name}: channel {k}: off by {worst} of {peak}" + ); + } + } +} + +#[test] +fn the_convolution_is_a_direct_one() { + // Random radiance over the focal plane, margin included, at 2 x 2, convolved with + // the three lopsided kernels and averaged into pixels in f64: equal to 1e-6 of the + // brightest rendered pixel. + let size = (WIDTH, HEIGHT); + let (w, h, m) = focal(2, size); + let ks = lopsided(5, 0.75, 2, 64); + let mut state = 0x2545_f491_4f6c_dd1du64; + let mut random = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state >> 11) as f64 / (1u64 << 53) as f64 + }; + let radiance: Vec<[f32; 4]> = (0..w * h) + .map(|_| [random() as f32, random() as f32, random() as f32, 0.0]) + .collect(); + let mut expected = vec![[0.0f64; 3]; (WIDTH * HEIGHT) as usize]; + for (i, e) in expected.iter_mut().enumerate() { + let (u, v) = (i % WIDTH as usize, i / WIDTH as usize); + for (k, kernel) in ks.iter().enumerate() { + let r = kernel.radius as i64; + let mut sum = 0.0; + for (a, b) in (0..2).flat_map(|a| (0..2).map(move |b| (a, b))) { + let (x, y) = ((m + 2 * u + a) as i64, (m + 2 * v + b) as i64); + for dy in -r..=r { + for dx in -r..=r { + let at = ((y - dy) as usize) * w + (x - dx) as usize; + sum += kernel.at(dx, dy) * radiance[at][k] as f64; + } + } + } + e[k] = sum / 4.0; + } + } + for ctx in common::contexts() { + let name = &ctx.info.name; + let image = convolve(ctx, &ks, 2, size, &radiance, &[]); + let worst = image + .iter() + .zip(&expected) + .flat_map(|(p, e)| (0..3).map(move |k| (p[k] as f64 - e[k]).abs())) + .fold(0.0f64, f64::max); + assert!(worst < 1e-6, "{name}: off by {worst}"); + } +} + +#[test] +fn diffraction_keeps_the_light_that_lands_in_the_frame() { + // A point of light at the image's center, at 2 x 2, through the example camera's + // f/4 aperture: the image holds the kernel's sum over the offsets inside it, to + // 1e-6, and that is all but the far wings. + let size = (WIDTH, HEIGHT); + let (w, h, m) = focal(2, size); + let grid = Layout::new(size.0, size.1, 2, DEFAULT_REACH).unwrap(); + let camera = CameraConfig { + psf: Psf::Airy { obstruction: 0.0 }, + ..CameraConfig::example() + }; + let ks = kernels(&camera, 2, grid.margin); + let (x0, y0) = (m + WIDTH as usize, m + HEIGHT as usize); + let mut radiance = vec![[0.0f32; 4]; w * h]; + radiance[y0 * w + x0] = [4.0; 4]; + for ctx in common::contexts() { + let name = &ctx.info.name; + let image = convolve(ctx, &ks, 2, size, &radiance, &[]); + for (k, kernel) in ks.iter().enumerate() { + let total: f64 = image.iter().map(|p| p[k] as f64).sum(); + let inside: f64 = (0..2 * HEIGHT as i64) + .flat_map(|y| (0..2 * WIDTH as i64).map(move |x| (x, y))) + .map(|(x, y)| kernel.at(x + m as i64 - x0 as i64, y + m as i64 - y0 as i64)) + .sum(); + assert!( + (total - inside).abs() < 1e-6, + "{name}: channel {k}: {total} vs {inside}" + ); + assert!(inside > 0.99, "{name}: channel {k}: {inside} inside"); + } + } +} + +#[test] +fn round_off_far_from_the_sun() { + // A disk 20 pixels in radius, 1e9 per rendered pixel, as bright as the Sun's in + // electrons per pixel at a few ms, at the image's left. Beyond it, out to past the + // kernel's reach of 32 px, the image is the disk's diffraction wings computed + // directly in f64, plus round-off that scales with the disk's brightness: measured + // at most 1.15e-7 of it on both devices, the same as at 2048 x 2048 (64 on a disk + // of 5.2e8), at and past the reach as near the disk. + let size = (128, 64); + let (w, h, m) = focal(2, size); + let grid = Layout::new(size.0, size.1, 2, DEFAULT_REACH).unwrap(); + let camera = CameraConfig { + psf: Psf::Airy { obstruction: 0.0 }, + ..CameraConfig::example() + }; + let ks = kernels(&camera, 2, grid.margin); + let bright = 1e9; + let (cx, cy) = ((m + 40) as f64, (m + 64) as f64); + let lit = |x: usize, y: usize| (x as f64 + 0.5 - cx).hypot(y as f64 + 0.5 - cy) < 40.0; + let radiance: Vec<[f32; 4]> = (0..w * h) + .map(|i| { + [if lit(i % w, i / w) { + bright as f32 + } else { + 0.0 + }; 4] + }) + .collect(); + let disk: Vec<(i64, i64)> = (0..w * h) + .filter(|&i| lit(i % w, i / w)) + .map(|i| ((i % w) as i64, (i / w) as i64)) + .collect(); + for ctx in common::contexts() { + let name = &ctx.info.name; + let image = convolve(ctx, &ks, 2, size, &radiance, &[]); + let mut worst = 0.0f64; + for v in (0..64).step_by(9) { + for u in (44..128).step_by(4) { + let mut want = 0.0; + for (a, b) in [(0, 0), (1, 0), (0, 1), (1, 1)] { + let (x, y) = ((m + 2 * u + a) as i64, (m + 2 * v + b) as i64); + want += disk + .iter() + .map(|&(sx, sy)| ks[1].at(x - sx, y - sy)) + .sum::() + * bright + / 4.0; + } + let got = image[v * 128 + u][1] as f64; + worst = worst.max((got - want).abs()); + } + } + assert!( + worst < 2e-7 * bright, + "{name}: off by {worst} ({:e} of the disk)", + worst / bright + ); + } +} + +#[test] +fn stars_keep_their_light_and_centroid() { + // A star between rendered pixels, at 2 x 2, shared among the four nearest and + // spread by a Gaussian of 0.6 px: the image holds a quarter of its light (the + // average over a pixel's four rendered pixels), centered where it is to 5e-4 px. + // The bilinear share keeps the centroid exactly; the transforms' round-off, summed + // over the frame, moves it by 1.3e-4 px on both devices. + let size = (WIDTH, HEIGHT); + let (w, h, m) = focal(2, size); + let camera = CameraConfig { + psf: Psf::Gaussian { + sigma: vec![0.6; 3], + }, + ..CameraConfig::example() + }; + let ks = kernels(&camera, 2, 64); + let radiance = vec![[0.0f32; 4]; w * h]; + let (u, v) = (20.3f64, 17.62f64); + let star = StarDraw { + at: [ + (u * 2.0 + m as f64) as f32, + (v * 2.0 + m as f64) as f32, + -1.0, + -1.0, + ], + electrons: [400.0, 800.0, 1200.0, 0.0], + }; + for ctx in common::contexts() { + let name = &ctx.info.name; + let image = convolve(ctx, &ks, 2, size, &radiance, &[star]); + for k in 0..3 { + let (mut total, mut x, mut y) = (0.0, 0.0, 0.0); + for (i, p) in image.iter().enumerate() { + let l = p[k] as f64; + total += l; + x += l * ((i % WIDTH as usize) as f64 + 0.5); + y += l * ((i / WIDTH as usize) as f64 + 0.5); + } + let want = star.electrons[k] as f64 / 4.0; + assert!( + (total / want - 1.0).abs() < 1e-5, + "{name}: channel {k}: {total}" + ); + let (x, y) = (x / total, y / total); + assert!( + (x - u).abs() < 5e-4 && (y - v).abs() < 5e-4, + "{name}: channel {k}: centroid ({x}, {y})" + ); + } + } +} diff --git a/crates/resolved/tests/gpu/fft.rs b/crates/resolved/tests/gpu/fft.rs new file mode 100644 index 0000000..e8e6c60 --- /dev/null +++ b/crates/resolved/tests/gpu/fft.rs @@ -0,0 +1,120 @@ +//! The 2D FFT against a direct DFT on the CPU, on the GPU and on the CPU (lavapipe). + +use crate::common; + +use resolved::gpu::fft::Fft2d; +use resolved::gpu::GpuContext; + +/// A storage buffer holding `values`. +fn upload(ctx: &GpuContext, values: &[[f32; 2]]) -> wgpu::Buffer { + use wgpu::util::DeviceExt; + ctx.device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("fft test"), + contents: bytemuck::cast_slice(values), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + }) +} + +/// Transforms `values` (rows of `width`) on the GPU and reads the result back. +pub fn transform( + ctx: &GpuContext, + width: u32, + height: u32, + values: &[[f32; 2]], + inverse: bool, +) -> Vec<[f32; 2]> { + let fft = Fft2d::new(ctx, width, height); + let image = upload(ctx, values); + let scratch = upload(ctx, values); + let staging = ctx.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("fft staging"), + size: fft.size(), + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let mut encoder = ctx.device.create_command_encoder(&Default::default()); + { + let mut pass = encoder.begin_compute_pass(&Default::default()); + fft.encode(ctx, &mut pass, &image, &scratch, inverse); + } + encoder.copy_buffer_to_buffer(&image, 0, &staging, 0, fft.size()); + ctx.queue.submit([encoder.finish()]); + staging + .slice(..) + .map_async(wgpu::MapMode::Read, |r| r.unwrap()); + ctx.device + .poll(wgpu::PollType::wait_indefinitely()) + .unwrap(); + let view = staging.slice(..).get_mapped_range().unwrap(); + let out = bytemuck::cast_slice(&view).to_vec(); + out +} + +/// A repeatable pseudorandom value in [-1, 1). +fn noise(i: usize) -> f32 { + let x = (i as u64 + 1).wrapping_mul(0x9E37_79B9_7F4A_7C15); + ((x >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0 +} + +#[test] +fn transforms_match_the_direct_dft() { + // Random complex images, square and not, against the DFT summed in f64: within + // 1e-6 of the largest output, both ways, and forward then inverse returns the + // input to 1e-6. + for ctx in common::contexts() { + let name = &ctx.info.name; + for (width, height) in [(32u32, 16u32), (8, 64), (2, 2)] { + let n = (width * height) as usize; + let input: Vec<[f32; 2]> = (0..n).map(|i| [noise(2 * i), noise(2 * i + 1)]).collect(); + for inverse in [false, true] { + let sign = if inverse { 1.0 } else { -1.0 }; + let scale = if inverse { 1.0 / n as f64 } else { 1.0 }; + let want: Vec<[f64; 2]> = (0..n) + .map(|o| { + let (ku, kv) = (o % width as usize, o / width as usize); + let mut sum = [0.0f64; 2]; + for (i, x) in input.iter().enumerate() { + let (u, v) = (i % width as usize, i / width as usize); + let angle = sign + * std::f64::consts::TAU + * ((ku * u) as f64 / width as f64 + + (kv * v) as f64 / height as f64); + let (c, s) = (angle.cos(), angle.sin()); + sum[0] += x[0] as f64 * c - x[1] as f64 * s; + sum[1] += x[0] as f64 * s + x[1] as f64 * c; + } + sum.map(|x| x * scale) + }) + .collect(); + let got = transform(ctx, width, height, &input, inverse); + let peak = want.iter().flatten().fold(0.0f64, |m, x| m.max(x.abs())); + let worst = got + .iter() + .zip(&want) + .flat_map(|(g, w)| [(g[0] as f64 - w[0]).abs(), (g[1] as f64 - w[1]).abs()]) + .fold(0.0f64, f64::max); + assert!( + worst < 1e-6 * peak.max(1.0), + "{name}: {width} x {height}, inverse {inverse}: off by {worst} of {peak}" + ); + } + let back = transform( + ctx, + width, + height, + &transform(ctx, width, height, &input, false), + true, + ); + let worst = back + .iter() + .zip(&input) + .flat_map(|(b, x)| [(b[0] - x[0]).abs(), (b[1] - x[1]).abs()]) + .fold(0.0f32, f32::max); + assert!( + worst < 1e-6, + "{name}: {width} x {height} round trip off by {worst}" + ); + } + } +} diff --git a/crates/resolved/tests/gpu/main.rs b/crates/resolved/tests/gpu/main.rs index 3ee1796..0808125 100644 --- a/crates/resolved/tests/gpu/main.rs +++ b/crates/resolved/tests/gpu/main.rs @@ -8,10 +8,12 @@ mod analytic; mod apparent; mod blooming; mod bodies; +mod convolve; mod distortion; mod earth_texture; mod earthshine; mod eclipses; +mod fft; mod fixed_pattern; mod gpu_context; mod mesh; diff --git a/crates/resolved/tests/gpu/psf.rs b/crates/resolved/tests/gpu/psf.rs index 736e2f4..52ea216 100644 --- a/crates/resolved/tests/gpu/psf.rs +++ b/crates/resolved/tests/gpu/psf.rs @@ -402,3 +402,72 @@ fn the_renderer_blurs_finer_pixels_from_camera_c1() { assert!(renderer(adapter, small_camera(sharp()), CameraTier::C1, Some(5)).is_err()); } } + +#[test] +fn the_renderer_convolves_from_camera_c2() { + // At C2 the focal plane is convolved by FFT. A Gaussian PSF keeps the Moon's total + // and centroid to C1's bounds, and differs from C1's blur by at most 0.34 % of the + // brightest pixel: C2 integrates the Gaussian over each rendered pixel, C1 samples + // it at their centers. Diffraction loses 7e-5 more of the light, past the frame's + // edges, and moves the centroid by 2e-3 px, away from the bottom edge the Moon is + // nearer. C1 refuses the Airy PSF, and C2 a grid larger than 4096 across. + let _devices = common::device_lock(); + for adapter in &ADAPTERS { + let mut reference = + renderer(adapter, small_camera(Psf::None), CameraTier::C1, Some(1)).unwrap(); + let reference = render(&mut reference).radiance.unwrap(); + let mut c1 = renderer(adapter, small_camera(sharp()), CameraTier::C1, None).unwrap(); + let c1 = render(&mut c1).radiance.unwrap(); + for psf in [sharp(), Psf::Airy { obstruction: 0.0 }] { + let mut r = renderer(adapter, small_camera(psf.clone()), CameraTier::C2, None).unwrap(); + let name = format!("{}, {psf:?}", r.adapter_info().name); + let image = render(&mut r).radiance.unwrap(); + for k in 0..3 { + let moments = |image: &[[f32; 4]]| { + let (mut total, mut u, mut v) = (0.0, 0.0, 0.0); + for (i, p) in image.iter().enumerate() { + let l = p[k] as f64; + total += l; + u += l * (i % 160) as f64; + v += l * (i / 160) as f64; + } + (total, u / total, v / total) + }; + let ((ta, ua, va), (tb, ub, vb)) = (moments(&reference), moments(&image)); + assert!( + (tb / ta - 1.0).abs() < 3e-4, + "{name}: channel {k}: {tb} vs {ta}" + ); + assert!( + (ub - ua).abs() < 0.01 && (vb - va).abs() < 0.01, + "{name}: channel {k}: centroid ({ub}, {vb}) vs ({ua}, {va})" + ); + if psf == sharp() { + let peak = c1.iter().fold(0.0f32, |m, p| m.max(p[k])); + let worst = c1 + .iter() + .zip(&image) + .fold(0.0f32, |m, (a, b)| m.max((a[k] - b[k]).abs())); + assert!( + worst < 0.0035 * peak, + "{name}: channel {k}: {worst} of {peak}" + ); + } + } + } + let airy = small_camera(Psf::Airy { obstruction: 0.0 }); + assert!(matches!( + renderer(adapter, airy, CameraTier::C1, None), + Err(Error::Unsupported(_)) + )); + let wide = CameraConfig { + width: 3840, + height: 2160, + ..small_camera(sharp()) + }; + assert!(matches!( + renderer(adapter, wide, CameraTier::C2, None), + Err(Error::Unsupported(_)) + )); + } +} diff --git a/crates/resolved/tests/gpu/stars.rs b/crates/resolved/tests/gpu/stars.rs index 8917b3a..b71d18e 100644 --- a/crates/resolved/tests/gpu/stars.rs +++ b/crates/resolved/tests/gpu/stars.rs @@ -121,7 +121,8 @@ fn a_star_lands_where_its_apparent_position_projects() { // Earth's 30 km/s of aberration: the centroid of its electrons tracks where // Intrinsics::project puts its apparent direction to 0.01 px, and they total what // StarPhotometry gives, times the aperture, the exposure and cos(off axis), to - // 0.2 %. + // 0.2 %. The same at camera C2, which draws it into the focal plane before its + // convolution. let _devices = common::device_lock(); let config = camera(0.02); let k = config.intrinsics(); @@ -130,15 +131,12 @@ fn a_star_lands_where_its_apparent_position_projects() { let beta = Vector3::new(0.0, 29.8e3, 0.0) / SPEED_OF_LIGHT; let apparent = aberrate(&star, &beta); let position = Vector3::new(4.2e7, 0.0, 0.0); - for adapter in &ADAPTERS { - let mut r = renderer( - adapter, - config.clone(), - LightingTier::L1, - CameraTier::C1, - &[vega()], - ); - let name = r.adapter_info().name.clone(); + for (adapter, tier) in ADAPTERS + .iter() + .flat_map(|a| [(a, CameraTier::C1), (a, CameraTier::C2)]) + { + let mut r = renderer(adapter, config.clone(), LightingTier::L1, tier, &[vega()]); + let name = format!("{}, {tier}", r.adapter_info().name); for step in 0..10 { // Look 40.3 px left of and 30.1 px above the star, plus a tenth of a pixel // per step: off axis, where the distortion moves it about 0.3 px. diff --git a/docs/optics.md b/docs/optics.md index 872675c..7b690b2 100644 --- a/docs/optics.md +++ b/docs/optics.md @@ -727,7 +727,9 @@ take 80 bytes per pixel of GPU memory: 166 MB at 1920 x 1080 and 660 MB at ## Blur: the point spread function -Code: `camera.rs` (`Psf`); `gpu/blur.rs`, `shaders/blur.wgsl`. +Code: `camera.rs` (`Psf`); `gpu/blur.rs`, `shaders/blur.wgsl`; from camera +C2, `psf.rs`, `gpu/convolve.rs`, `gpu/fft.rs`, `shaders/convolve.wgsl`, +`shaders/fft.wgsl`. No camera focuses a point to a point. Light from one point of the scene lands on the detector as a small spot, and the shape of that spot is the @@ -816,9 +818,9 @@ small-pixel averages, and the same error remains at a smaller scale: (sigma 0.4 px, points at 16 and edges at 129 positions within a pixel; the worst edges lie on a pixel boundary.) C1 renders 2 x 2 by default; -`RendererOptions::supersampling` sets 1 to 4. Point sources, stars in -particular, need better than that, and will be drawn with the blurred spot -computed exactly over each pixel instead. +`RendererOptions::supersampling` sets 1 to 4. Point sources need better +than that, so stars are drawn with the blurred spot computed exactly over +each pixel instead (see Stars). A 2D Gaussian is the product of a 1D Gaussian across and one down, and averaging an n x n block is likewise an average across and one down. So the @@ -860,6 +862,83 @@ as an error. On the CPU through lavapipe, 2 x 2 adds 9 to 13 ms to 1920 x 1080 frames of 116 to 680 ms, and 3 x 3 at 3840 x 2160 runs out of memory. +### Diffraction and the PSF's wings (camera C2) + +A Gaussian falls to nothing within a few sigma. A real camera's spot does +not: diffraction's rings fade only as the cube of the distance, and dust and +rough lens surfaces scatter a little light much farther. Next to most of a +scene that faint halo, the PSF's wings, is invisible. Next to the Sun it is +not. In the sun-glare example scene each pixel of the Sun's disk collects +about 5e8 electrons in its 0.4 ms exposure, 26,000 full wells, over about +5,700 pixels. Diffraction at f/4 spreads a fraction of about +λN / (π^3 r^3) of a point's light into each pixel r pixels away (λN, the +wavelength times the f-number, is 0.39 px for green light in the example +camera), so 500 px from the Sun each pixel still gets about 300 electrons +from it, and 1000 px away about 40: a glow over the whole frame that raises +the background, can hide a faint target, and pulls a limb fit toward the +bright side. + +Camera tier C2 models this. `Psf::Airy { obstruction }` is diffraction by +the camera's aperture: the Airy pattern, set by the f-number, with an +optional central obstruction such as a reflecting telescope's secondary +mirror (which moves light from the core into the rings). The pattern scales +with the wavelength, so a channel's spot is the pattern averaged over the +channel's band, each wavelength weighted by the photons sunlight gives it +there; across an 80 nm band the outer rings blur into a smooth halo. +`Psf::Measured { image, pixel_scale }` takes a spot measured in a lab or in +flight, the same over the whole field, one image per channel or one for +all. Measured: a 2 nm wide channel's first dark ring falls within 1 % of +1.22 λN, where the textbook puts it. + +A kernel reaching hundreds of pixels is far too slow to apply pixel by +pixel. C2 instead uses the convolution theorem: transform the image into +its frequencies with a fast Fourier transform (FFT), multiply by the +kernel's frequencies, computed once at setup, and transform back. The cost +no longer depends on the kernel's size, only on the image's. The transform +runs on the GPU in compute shaders, sized to powers of two up to 4096 on a +side, with two channels sharing one complex transform, so three channels +take four transforms per frame. C2 applies every PSF this way, Gaussians +too, and draws stars into the focal plane before the transform, so they get +the same wings; each star's light is shared among its four nearest rendered +pixels, which keeps its center exact but can widen its spot by up to 0.25 px +RMS at 2 x 2. + +As at C1, the transform acts on the scene rendered finer than the pixels, +before they average it. The kernel is the PSF integrated over a rendered +pixel twice, once for the pixel the light leaves and once for the one it +lands in, so a Gaussian at C2 differs slightly from C1's, which samples the +Gaussian at the rendered pixels' centers (by up to 0.34 % of the brightest +pixel, measured on the Moon). The transform treats the image as repeating, +so the renderer draws a margin around it as wide as the kernel reaches, and +the grid holds the image and that margin. The reach is therefore set by the +grid: at least `RendererOptions::psf_reach` (16 px by default), and as far +as the smallest grid that holds it allows. Light the kernel would carry +farther is added back to the spot instead, so none is lost, but the glow +stops at the reach, and with the Sun in the frame its edge shows in a +brightened preview. + +| Image, rendered pixels per pixel | Grid | Reach | RTX 3060 | lavapipe | +|---|---|---|---|---| +| 1920 x 1080, 2 x 2 | 4096 x 4096 | 64 px | 30 ms | 560 ms | +| 1920 x 1080, 1 x 1 | 2048 x 2048 | 64 px | 5.8 ms | 213 ms | +| 1920 x 1080, 1 x 1, reach 1000 | 4096 x 4096 | 1088 px | 29 ms | 556 ms | +| 3840 x 2160, 1 x 1 | 4096 x 4096 | 128 px | 30 ms | 623 ms | + +(The convolution pass alone, in the sun-glare scene at lighting L1. The +whole frame took 35 ms at 1920 x 1080 and 2 x 2, against 5.9 ms at C1.) A +grid of 4096 x 4096 takes 640 MB on the device. 3840 x 2160 at 2 x 2 would +need a grid 8192 across and is refused; 1 x 1 fits. + +The transforms work in 32-bit floating point, which keeps about seven +significant digits, and their rounding error spreads over the whole grid in +proportion to its brightest values. With the Sun in the frame that leaves a +floor of about 10 electrons RMS, up to 165, across the frame (measured in +the sun-glare scene with a noiseless detector, where C1 reads exactly 0), +above the example detector's 5 electrons of read noise. A sunlit target at +half the full well leaves a floor near 1e-3 electrons, far below any noise. +Computing the brightest sources' wings separately, outside the transform, +would remove the Sun's floor; that is not done yet. + ## Earth, Moon and Sun Code: `bodies.rs`, `gpu/bodies.rs`, `shaders/bodies.wgsl`. @@ -1638,7 +1717,9 @@ since all of a star's light that enters lands in its spot. A star is not rendered and then blurred: a point of light rendered into a pixel loses where in the pixel it fell (see the blur section). Instead its electrons are spread over the pixels around it as the Gaussian PSF integrated over each -pixel's area, which is exact for any position within a pixel. The detector +pixel's area, which is exact for any position within a pixel. (Camera C2 +draws stars into the focal plane before its convolution instead, so they +take the Airy or measured PSF too; see the PSF's wings.) The detector adds them to the light from everything else before the full well, so a bright star saturates like any other light. A star is hidden wherever anything nearer is drawn, Earth, the Moon or a target: stars are at @@ -1672,7 +1753,8 @@ is a run of the script with `--gaia 16`, at twice the size. ## Not modeled yet - Earthshine on the Moon's night side -- The diffraction pattern's rings, and a lens's measured point spread function +- A point spread function that changes across the field, and the PSF's wings + past the reach of camera C2's grid - Non-Lambertian materials (BRDFs), glints, and the full Moon's edge as bright as its center - Clouds on Earth diff --git a/readme.md b/readme.md index 5f9f4f5..98d7f7b 100644 --- a/readme.md +++ b/readme.md @@ -170,10 +170,12 @@ cargo run --release -p resolved-showcase -- --serve 0.0.0.0:8010 3 times and rolls highlights off to white at saturation. `--out` is the same either way. - `--timings N` renders N frames and reports setup and per-frame times. -- `--lighting L0|L1` and `--camera C0|C1` pick the fidelity tiers (default +- `--lighting L0|L1` and `--camera C0|C1|C2` pick the fidelity tiers (default L0 and C0). `--supersampling N` renders each pixel as N x N finer ones from C1 - (default 2). A tier the scene or the data cannot support stops with an error - naming what is missing. + (default 2). `--psf-reach N` sets how far in pixels camera C2's PSF must reach + (default 16; it reaches as far as its FFT grid allows, 64 px at 1920 x 1080 and + 2 x 2). A tier the scene or the data cannot support stops with an error naming + what is missing. - `--seed N` and `--frame N` pick the detector noise at camera C1 (default 0 and 0): the same seed and frame give the same image. - `--data DIR` says where the tiers' data is (`earth/bmng-2004-MM.jpg`, @@ -208,9 +210,12 @@ lens blur (a Gaussian PSF) applied to the scene rendered 2 x 2 finer than the pixels, the Sun, Moon and Earth at their apparent positions (light time and aberration), detector noise (shot, dark current and read noise, seeded and reproducible), the detector's fixed pattern (PRNU, DSNU and hot pixels), -and blooming along columns or into neighbors. Planned, in order: +and blooming along columns or into neighbors. Camera C2 adds diffraction +(the Airy pattern, averaged over each channel's band) and measured PSFs, +applied with their wings by FFT. Planned, in order: -- Camera: the C1 checkpoint, then diffraction and measured PSFs (C2). +- Camera: the C1 and C2 checkpoint, their rows in the table above with + sample images. - Lighting: BRDF and spectral materials, atmospheric scattering, clouds, lunar terrain. - Portable ray tracing, path tracing, and label outputs for ML.