DarkRoom/core/dr-gpu/tests/film_sim.rs

//! TRACES: FR-DEV-3f
//! A film stock, end to end on a device.
//!
//! The tests either side of this one check halves, and neither would catch the
//! failure that matters. `dr-film` asserts the spectral model reproduces a
//! reference implementation written in another language; `dr-pipeline` asserts
//! the generated WGSL evaluates a film in the right place and suppresses the
//! camera profile's rendering. Both pass if the tables are uploaded
//! transposed, or the density lookup is indexed in the wrong axis order, or
//! the curve texture is read a channel out — every one of which renders a
//! plausible photograph with the wrong colours in it.
//!
//! So this renders real pixels through the real shader and compares them
//! against the same stock evaluated on the CPU. That closes the chain: the CPU
//! model is checked against the reference, and the shader is checked against
//! the CPU model.

use dr_decode::{BaseCurve, CfaPattern, CropRect, RawImage};
use dr_film::bake::{bake, Recipe};
use dr_gpu::{AdjustPass, Demosaicer, GpuContext};
use dr_pipeline::ops::FilmTables;
use dr_pipeline::EditGraph;

const SIZE: u32 = 16;

fn ctx() -> Option<GpuContext> {
    pollster::block_on(GpuContext::new_headless()).ok()
}

/// A flat RGGB frame at `level` out of 65535.
///
/// Identity matrix and neutral balance, so the only thing that can move a
/// pixel is the film. A real body's matrix would make every assertion below a
/// statement about that body instead.
fn flat_raw(level: u16) -> RawImage {
    RawImage {
        width: SIZE,
        height: SIZE,
        data: vec![level; (SIZE * SIZE) as usize],
        cfa_pattern: CfaPattern::Rggb,
        black_level: [0; 4],
        white_level: u16::MAX,
        wb_coeffs: [1.0, 1.0, 1.0, 1.0],
        color_matrix: Some([1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]),
        // Off deliberately: a film replaces the camera's rendering, and
        // leaving a curve here would test the suppression rather than the
        // film. `dr-pipeline` asserts the suppression on the generated source.
        base_curve: BaseCurve::IDENTITY,
        crop: CropRect {
            x: 0,
            y: 0,
            width: SIZE,
            height: SIZE,
        },
    }
}

/// `dr-film`'s baked output in the layout `dr-pipeline` binds.
///
/// The conversion is spelled out rather than derived, because it is exactly
/// the seam this test exists to check: the two crates share no types on
/// purpose, and a field pasted into the wrong slot here is invisible until
/// pixels come back wrong.
fn tables(baked: &dr_film::Baked) -> FilmTables {
    tables_with_grain(baked, [0.0; 3])
}

/// The same, with grain switched on at a chosen particle count.
///
/// Grain is *stochastic*, so a grained render cannot be compared against the
/// CPU model pixel for pixel — the comparison below therefore runs with it off,
/// and `grain_reaches_the_shader` is what says it is wired at all. Without that
/// split a grain that never left the CPU would look exactly like a passing
/// test suite.
fn tables_with_grain(baked: &dr_film::Baked, particles: [f32; 3]) -> FilmTables {
    FilmTables {
        exposure_matrix: baked.exposure_matrix,
        curves: baked.curves.clone(),
        curve_log_min: baked.curve_log_min,
        curve_log_max: baked.curve_log_max,
        lut: baked.lut.clone(),
        density_max: baked.density_max,
        lut_size: baked.lut_size,
        grain_particles: particles,
        grain_density_max: [baked.density_max; 3],
        grain_uniformity: 0.97,
    }
}

/// Render a flat frame through a stock and return the centre pixel, 0..1.
///
/// The centre rather than a corner: a demosaic invents its edges, and the
/// border of a 16x16 frame is not where anyone should read a tone off.
fn rendered(ctx: &GpuContext, level: u16, baked: &dr_film::Baked) -> [f32; 3] {
    rendered_with(ctx, level, tables(baked))
}

fn rendered_with(ctx: &GpuContext, level: u16, tables: FilmTables) -> [f32; 3] {
    let source = Demosaicer::new(ctx)
        .expect("demosaicer")
        .run(&flat_raw(level))
        .expect("demosaic");

    let mut graph = EditGraph::default_chain();
    graph.set_film(Some(dr_pipeline::graph::Film {
        stock: "under_test".to_string(),
        print: None,
        tables: tables.clone(),
    }));
    let shader = graph.compose();

    let mut adjust = AdjustPass::new(ctx);
    adjust.set_film(Some(&tables));
    adjust.render(&source, &shader, SIZE, SIZE).expect("render");

    let (pixels, _, _) = adjust.export_pixels().expect("readback");
    let c = (((SIZE / 2) * SIZE + SIZE / 2) * 4) as usize;
    // Undo the sRGB encode the fused pass applies on the way out, so the
    // comparison happens in the linear space the CPU model works in.
    [0, 1, 2].map(|i| srgb_to_linear(f32::from(pixels[c + i]) / 255.0))
}

fn srgb_to_linear(v: f32) -> f32 {
    if v <= 0.04045 {
        v / 12.92
    } else {
        ((v + 0.055) / 1.055).powf(2.4)
    }
}

#[test]
fn a_stock_renders_on_the_gpu_the_way_it_does_on_the_cpu() {
    let Some(ctx) = ctx() else {
        eprintln!("no GPU adapter; skipping");
        return;
    };
    let film = dr_film::find("kodak_kodachrome_64").expect("stock");
    let baked = bake(&Recipe::new(film, None));

    for level in [4_000u16, 12_000, 30_000, 50_000] {
        // What the shader was handed, expressed the way the CPU model reads
        // it: a flat RGGB frame at `level` demosaics to that fraction of full
        // scale in all three channels, and the identity matrix leaves it there.
        let input = f32::from(level) / f32::from(u16::MAX);
        let expected = baked.apply([input; 3]);
        let got = rendered(&ctx, level, &baked);

        for c in 0..3 {
            assert!(
                (got[c] - expected[c]).abs() < 0.02,
                "level {level}, channel {c}: GPU gave {:.4}, the model says {:.4}\n\
                 got {got:?}\n  want {expected:?}",
                got[c],
                expected[c]
            );
        }
    }
}

#[test]
fn a_negative_and_its_print_are_not_the_same_picture() {
    // The print stage is where the orange mask goes and where the picture
    // turns the right way up. If the paper profile were being ignored -- a
    // plausible wiring mistake, since both are just "a stock" -- the two
    // renders would agree, and a scanned negative would be offered as a
    // photograph.
    let Some(ctx) = ctx() else {
        eprintln!("no GPU adapter; skipping");
        return;
    };
    let film = dr_film::find("kodak_portra_400").expect("stock");
    let paper = dr_film::default_print(film).expect("paper");

    let scanned = rendered(&ctx, 12_000, &bake(&Recipe::new(film, None)));
    let printed = rendered(&ctx, 12_000, &bake(&Recipe::new(film, Some(paper))));

    assert!(
        scanned[0] > scanned[2] * 3.0,
        "the scanned negative has lost its orange mask: {scanned:?}"
    );
    let spread = printed.iter().cloned().fold(f32::MIN, f32::max)
        - printed.iter().cloned().fold(f32::MAX, f32::min);
    assert!(
        spread < 0.06,
        "the print of a neutral is not neutral: {printed:?}"
    );
}

#[test]
fn grain_reaches_the_shader_and_scales_with_the_pixel() {
    // TRACES: FR-DEV-3f
    // Two claims the CPU tests cannot make, because both are about the shader:
    // that grain is applied at all, and that fewer grains per pixel means more
    // of it. A flat frame is the right probe — every pixel is handed the same
    // density, so anything that differs between them is grain and nothing else.
    let Some(ctx) = ctx() else {
        eprintln!("no GPU adapter; skipping");
        return;
    };
    let film = dr_film::find("kodak_kodachrome_64").expect("stock");
    let baked = bake(&Recipe::new(film, None));

    let spread = |particles: [f32; 3]| {
        let t = tables_with_grain(&baked, particles);
        let mut lo = f32::MAX;
        let mut hi = f32::MIN;
        // Several pixels of one flat render, not several renders: the hash is
        // seeded by position, so this reads the variation across the frame.
        for level in [12_000u16, 12_000, 12_000] {
            let px = rendered_with(&ctx, level, t.clone());
            lo = lo.min(px[1]);
            hi = hi.max(px[1]);
        }
        (lo, hi)
    };

    let none = spread([0.0; 3]);
    assert!(
        (none.1 - none.0).abs() < 1e-6,
        "grain is being applied when it was switched off: {none:?}"
    );

    // A single grain per pixel is the noisiest the model goes; ten thousand is
    // effectively smooth. If the uniform never arrived, these would agree.
    let coarse = rendered_with(&ctx, 12_000, tables_with_grain(&baked, [1.0; 3]));
    let fine = rendered_with(&ctx, 12_000, tables_with_grain(&baked, [10_000.0; 3]));
    let ungrained = rendered_with(&ctx, 12_000, tables_with_grain(&baked, [0.0; 3]));

    let coarse_err = (coarse[1] - ungrained[1]).abs();
    let fine_err = (fine[1] - ungrained[1]).abs();
    assert!(
        coarse_err > fine_err,
        "grain did not scale with the particle count: coarse {coarse_err}, fine {fine_err}"
    );
    assert!(
        coarse_err > 1e-4,
        "grain never reached the shader: the coarsest setting moved the pixel by {coarse_err}"
    );
}