resolved/crates/resolved/tests/sensor.rs

//! The sensor pass against the CPU reference (`Radiometry::electrons` then
//! `Detector::dn`), on the GPU and on the CPU (lavapipe).

mod common;

use resolved::bodies::Appearance;
use resolved::gpu::bodies::{BodiesPass, BodiesUniforms, BodyRadiometry};
use resolved::gpu::sensor::{SensorPass, SensorUniforms};
use resolved::gpu::targets::FrameTargets;
use resolved::gpu::GpuContext;
use resolved::nalgebra::{UnitQuaternion, Vector2, Vector3};
use resolved::radiometry::{Radiometry, AU};
use resolved::{CameraConfig, Pose, Scene};

/// Runs the sensor pass over `radiance` (example camera, rows from the top) and
/// returns the DN image.
fn sense(ctx: &GpuContext, name: &str, radiance: &[[f32; 4]]) -> Vec<u16> {
    let config = CameraConfig::example();
    let targets = FrameTargets::new(ctx, config.width, config.height, 3);
    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(config.width * 16),
            rows_per_image: Some(config.height),
        },
        targets.radiance.size(),
    );
    let pass = SensorPass::new(ctx);
    let uniforms = SensorUniforms::new(&config, &Radiometry::new(&config));
    let mut encoder = ctx.device.create_command_encoder(&Default::default());
    pass.encode(ctx, &mut encoder, &targets, &uniforms);
    ctx.queue.submit([encoder.finish()]);
    let dn = targets.read_dn(ctx);
    common::save_dn(name, ctx, &dn);
    dn
}

/// The DN the CPU reference gives for `radiance` in channel `c` at pixel (u, v).
fn reference(config: &CameraConfig, r: &Radiometry, radiance: f32, u: usize, v: usize) -> u16 {
    let ray = config
        .intrinsics()
        .ray(&Vector2::new(u as f64 + 0.5, v as f64 + 0.5));
    config.detector.dn(r.electrons(radiance as f64, ray.z))
}

/// Channel radiance of a white Lambertian surface facing the Sun at 1 AU.
fn white() -> [f32; 4] {
    let r = Radiometry::new(&CameraConfig::example());
    let mut out = [0.0; 4];
    for (k, o) in out.iter_mut().take(3).enumerate() {
        *o = r.lambert_radiance(k, 1.0, 1.0, AU) as f32;
    }
    out
}

#[test]
fn every_pixel_matches_the_cpu_reference() {
    // A pattern from black to three times the white level (past saturation), different
    // in each channel and pixel.
    let config = CameraConfig::example();
    let r = Radiometry::new(&config);
    let w = white();
    let level = |u: usize, v: usize, c: usize| ((u * 7 + v * 13 + c * 5) % 1000) as f32 / 333.0;
    let radiance: Vec<[f32; 4]> = (0..1080)
        .flat_map(|v| (0..1920).map(move |u| (u, v)))
        .map(|(u, v)| [0, 1, 2, 3].map(|c| if c < 3 { w[c] * level(u, v, c) } else { 0.0 }))
        .collect();
    for ctx in common::contexts() {
        let dn = sense(ctx, "sensor-pattern", &radiance);
        let mut off_by_one = 0;
        // Every 8th row, plus the last: enough to catch any indexing mistake.
        for v in (0..1080).step_by(8).chain([1079]) {
            for u in 0..1920 {
                for c in 0..3 {
                    let got = dn[(v * 1920 + u) * 3 + c];
                    let want = reference(&config, &r, radiance[v * 1920 + u][c], u, v);
                    assert!(
                        got.abs_diff(want) <= 1,
                        "{} ({u},{v},{c}): {got} vs {want}",
                        ctx.info.name
                    );
                    off_by_one += usize::from(got != want);
                }
            }
        }
        // f32 on the GPU and f64 on the CPU disagree only where a value lands within
        // rounding of a half-DN boundary: measured 557 of all 6.2 million values on the
        // RTX 3060 and 699 on lavapipe, about 1 in 10,000.
        assert!(
            off_by_one < 1920 * 136 * 3 / 5000,
            "{}: {off_by_one} off by one",
            ctx.info.name
        );
    }
}

#[test]
fn worked_example_white_surface_saturation_and_bias() {
    let w = white();
    for ctx in common::contexts() {
        // A white surface facing the Sun fills the frame: docs/optics.md's worked example.
        let dn = sense(ctx, "sensor-white", &vec![w; 1920 * 1080]);
        let green = |u: usize, v: usize| dn[(v * 1920 + u) * 3 + 1];
        assert_eq!(green(960, 540), 2292);
        // cos^4 falloff: 6.9 degrees off axis at the corner, 10,819 electrons instead of
        // 11,139, so 64 + 10,819 / 5 = 2227.8.
        assert_eq!(green(0, 0), 2228);

        // Looking straight at the Sun: the full well caps the signal at 20,000 / 5 + 64.
        let sun = w.map(|x| x * 5e4);
        assert!(sense(ctx, "sensor-sun", &vec![sun; 1920 * 1080])
            .iter()
            .all(|&d| d == 4064));
        // Darkness reads as the bias.
        assert!(sense(ctx, "sensor-dark", &vec![[0.0; 4]; 1920 * 1080])
            .iter()
            .all(|&d| d == 64));
    }
}

#[test]
fn sunlit_earth_through_the_whole_chain() {
    // Earth seen from 200,000 km with the Sun behind the camera: the sub-camera point
    // has albedo 0.3 facing the Sun, so 0.3 x 11,138.8 electrons in green: 732 DN.
    let config = CameraConfig::example();
    let radiometry = Radiometry::new(&config);
    let scene = Scene {
        earth_attitude: UnitQuaternion::identity(),
        sun: Vector3::new(AU, 0.0, 0.0),
        moon: Vector3::new(0.0, -3.844e8, 0.0),
        targets: vec![],
    };
    let camera = Pose::looking_at(
        Vector3::new(2.0e8, 0.0, 0.0),
        Vector3::zeros(),
        Vector3::z(),
    )
    .unwrap();
    let body_radiometry = BodyRadiometry::new(&radiometry, &Appearance::default()).unwrap();
    let bodies = BodiesUniforms::new(
        &scene,
        &camera,
        &config.intrinsics(),
        config.near_clip,
        &body_radiometry,
    );
    let sensor = SensorUniforms::new(&config, &radiometry);
    for ctx in common::contexts() {
        let targets = FrameTargets::new(ctx, config.width, config.height, 3);
        let mut encoder = ctx.device.create_command_encoder(&Default::default());
        BodiesPass::new(ctx).encode(ctx, &mut encoder, &targets, &bodies);
        SensorPass::new(ctx).encode(ctx, &mut encoder, &targets, &sensor);
        ctx.queue.submit([encoder.finish()]);
        let dn = targets.read_dn(ctx);
        common::save_dn("sensor-earth", ctx, &dn);
        assert_eq!(dn[(540 * 1920 + 960) * 3 + 1], 732);
        assert_eq!(dn[0], 64, "space reads as the bias");
    }
}