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::{CfaPattern, CropRect, RawImage};
use dr_film::bake::{bake, Recipe, Settings};
use dr_gpu::{AdjustPass, Demosaicer, GpuContext, LabelField, MaskPass};
use dr_pipeline::mask::{MaskLayer, MaskSource};
use dr_pipeline::ops::film_sim;
use dr_pipeline::ops::film_sim::FORMAT_COUNT;
use dr_pipeline::ops::{FilmTables, PaperTables};
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.
        samples_per_pixel: 1,
        profile: None,
        profile_tables: None,
        baseline_exposure: 0.0,
        make: String::new(),
        model: String::new(),
        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 {
    // The paper, when there is one, rides behind the film: its curve as one
    // more row, its cube stacked after the film's.
    let mut curves = baked.curves.clone();
    let mut lut = baked.lut.clone();
    let paper = baked.paper.as_ref().map(|p| {
        curves.extend_from_slice(&p.curves);
        lut.extend_from_slice(&p.lut);
        PaperTables {
            balance: p.balance,
            log_min: p.log_min,
            log_max: p.log_max,
            density_max: p.density_max,
        }
    });
    FilmTables {
        exposure_matrix: baked.exposure_matrix,
        curves,
        push_stations: baked.push_stations.clone(),
        curve_log_min: baked.curve_log_min,
        curve_log_max: baked.curve_log_max,
        lut,
        density_max: baked.density_max,
        lut_size: baked.lut_size,
        paper,
        grain_particles: [particles; FORMAT_COUNT],
        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}"
    );
}

/// The same render, with the film's sliders set and mask layers laid over it.
///
/// Every layer is a `Regions` mask over a field splitting the frame down the
/// middle: region 0, the left half, at full weight, and the right half
/// untouched. Returns the left and right centre pixels, linear.
fn rendered_split(
    ctx: &GpuContext,
    level: u16,
    tables: FilmTables,
    global: Settings,
    layers: Vec<MaskLayer>,
) -> ([f32; 3], [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(),
    }));
    graph.set_param(film_sim::ID, film_sim::EXPOSURE, global.exposure_ev);
    graph.set_param(film_sim::ID, film_sim::PUSH, global.push_stops);
    graph.set_param(
        film_sim::ID,
        film_sim::PRINT_EXPOSURE,
        global.print_exposure_ev,
    );
    for layer in layers {
        graph.masks_mut().push(layer);
    }
    let shader = graph.compose();

    let labels: Vec<u32> = (0..SIZE * SIZE)
        .map(|i| u32::from(i % SIZE >= SIZE / 2))
        .collect();
    let field = LabelField::upload(ctx, &labels, SIZE, SIZE, 2).expect("label upload");
    let mut masks = MaskPass::new(ctx).expect("mask pass");
    let array = masks
        .render(graph.masks(), Some(&field), None, None, SIZE, SIZE)
        .expect("rasterise");

    let mut adjust = AdjustPass::new(ctx);
    adjust.set_film(Some(&tables));
    adjust
        .render_masked(&source, &shader, SIZE, SIZE, Some(array))
        .expect("render");
    let (pixels, _, _) = adjust.export_pixels().expect("readback");
    let at = |x: u32| {
        let c = (((SIZE / 2) * SIZE + x) * 4) as usize;
        [0, 1, 2].map(|i| srgb_to_linear(f32::from(pixels[c + i]) / 255.0))
    };
    (at(SIZE / 4), at(3 * SIZE / 4))
}

/// A layer over the left half holding these film offsets.
fn left_half(id: &str, offsets: &[(dr_pipeline::descriptor::ParamId, f32)]) -> MaskLayer {
    let mut layer = MaskLayer::new(
        id,
        MaskSource::Regions {
            signature: 1,
            level: 2,
            ids: vec![0],
        },
    );
    for (param, v) in offsets {
        layer.set_param(film_sim::ID.0, *param, *v);
    }
    layer
}

fn assert_close(got: [f32; 3], want: [f32; 3], what: &str) {
    for c in 0..3 {
        assert!(
            (got[c] - want[c]).abs() < 0.02,
            "{what}, channel {c}: GPU gave {got:?}, the model says {want:?}"
        );
    }
}

#[test]
fn a_print_renders_on_the_gpu_the_way_it_does_on_the_cpu_at_any_setting() {
    // TRACES: FR-DEV-3f
    // The print path — the film's lookup into the paper's log exposure, the
    // enlarger added between, the paper's curve and its own lookup — is read
    // from the same two textures as the film, at offsets. Every one of those
    // offsets is a way to render a plausible print of the wrong thing.
    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 baked = bake(&Recipe::new(film, Some(paper)));

    for settings in [
        Settings::default(),
        Settings {
            print_exposure_ev: -1.3,
            ..Settings::default()
        },
        Settings {
            exposure_ev: 0.4,
            print_exposure_ev: 0.8,
            ..Settings::default()
        },
    ] {
        for level in [6_000u16, 20_000] {
            let input = f32::from(level) / f32::from(u16::MAX);
            let (got, _) = rendered_split(&ctx, level, tables(&baked), settings, Vec::new());
            assert_close(
                got,
                baked.apply_at([input; 3], &settings),
                &format!("{settings:?} at {level}"),
            );
        }
    }
}

#[test]
fn a_push_between_two_measured_processes_renders_as_the_model_does() {
    // TRACES: FR-DEV-3f
    // Double-X measures five processes; a push between two is a mix of two
    // rows of the curve texture, found by searching the stations uniform.
    let Some(ctx) = ctx() else {
        eprintln!("no GPU adapter; skipping");
        return;
    };
    let film = dr_film::find("kodak_doublex").expect("stock");
    let baked = bake(&Recipe::new(film, None));
    assert!(baked.curve_rows > 2, "Double-X has a development series");

    for push in [-0.8f32, 0.4, 1.3, 2.9] {
        let settings = Settings {
            push_stops: push,
            ..Settings::default()
        };
        let level = 12_000u16;
        let input = f32::from(level) / f32::from(u16::MAX);
        let (got, _) = rendered_split(&ctx, level, tables(&baked), settings, Vec::new());
        assert_close(
            got,
            baked.apply_at([input; 3], &settings),
            &format!("push {push}"),
        );
    }
}

#[test]
fn a_layer_develops_its_region_on_its_own_settings() {
    // TRACES: FR-DEV-3f
    // Offsets to the photograph's: print exposure +1 on a photograph at +0.5
    // is +1.5 under the layer, and the rest of the print is untouched. Before
    // film was blended as settings the layer's sliders moved and nothing
    // happened, because the layer's copy of the node had no stock.
    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 baked = bake(&Recipe::new(film, Some(paper)));
    let level = 12_000u16;
    let input = f32::from(level) / f32::from(u16::MAX);

    let global = Settings {
        print_exposure_ev: 0.5,
        ..Settings::default()
    };
    let layer = left_half(
        "burn",
        &[(film_sim::PRINT_EXPOSURE, 1.0), (film_sim::EXPOSURE, -0.5)],
    );
    let (left, right) = rendered_split(&ctx, level, tables(&baked), global, vec![layer]);

    let under = Settings {
        exposure_ev: -0.5,
        print_exposure_ev: 1.5,
        ..Settings::default()
    };
    assert_close(left, baked.apply_at([input; 3], &under), "under the layer");
    assert_close(right, baked.apply_at([input; 3], &global), "outside it");
    assert!(
        left[1] < right[1] - 0.01,
        "the burn did not darken: {left:?} vs {right:?}"
    );
}

#[test]
fn overlapping_layers_take_the_average_of_their_settings() {
    // TRACES: FR-DEV-3f
    // Three layers over the same pixels at full weight: the plain mean of what
    // each asks for, and the global setting has no weight left. Summed, the
    // offsets would be -2 stops of print exposure and +1 of exposure; the mean
    // is (-1, -1, 0) / 3 and (0, 0, +1) / 3.
    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 baked = bake(&Recipe::new(film, Some(paper)));
    let level = 12_000u16;
    let input = f32::from(level) / f32::from(u16::MAX);

    let layers = vec![
        left_half("a", &[(film_sim::PRINT_EXPOSURE, -1.0)]),
        left_half("b", &[(film_sim::PRINT_EXPOSURE, -1.0)]),
        left_half("c", &[(film_sim::EXPOSURE, 1.0)]),
    ];
    let (left, right) = rendered_split(&ctx, level, tables(&baked), Settings::default(), layers);

    let mean = Settings {
        exposure_ev: 1.0 / 3.0,
        print_exposure_ev: -2.0 / 3.0,
        ..Settings::default()
    };
    let summed = Settings {
        exposure_ev: 1.0,
        print_exposure_ev: -2.0,
        ..Settings::default()
    };
    let (m, s) = (
        baked.apply_at([input; 3], &mean)[1],
        baked.apply_at([input; 3], &summed)[1],
    );
    // Three times the tolerance the GPU is held to below, or a sum could pass
    // for a mean.
    assert!(
        (m - s).abs() > 0.06,
        "the mean and the sum render alike ({m} vs {s}), so this proves nothing"
    );
    assert_close(left, baked.apply_at([input; 3], &mean), "under all three");
    assert_close(right, baked.apply([input; 3]), "outside them");
}