DarkRoom/core/dr-pipeline/src/neutral.rs

//! TRACES: FR-DEV-3 | FR-DEV-3a
//! Turning a colour that ought to be grey into white balance parameters.
//!
//! # Why this is in the core and not in the interface
//!
//! FR-DEV-3 asks for a white balance picker: point at something neutral and
//! the correction follows. The pointing is the interface's — it owns the
//! canvas, the pixel and the modality — but the *arithmetic* is not, and the
//! reason is ARCH §4.3a rather than tidiness. How far a hundred units of
//! temperature move red against blue is declared in
//! `core/dr-pipeline/ops/white_balance.yaml`, in one expression, and an
//! interface that inverted it would be holding a second copy of a number it
//! is not allowed to know. The copy would then be wrong the first time
//! somebody adjusted the range, silently, in the direction of "the picker is
//! slightly off".
//!
//! So the interface hands over a colour and this hands back a moved graph.
//!
//! # What it is allowed to assume, which is a widget kind and not an operation
//!
//! Nothing here names white balance. It looks for the operation that asks for
//! [`WidgetKind::WhitePoint`], and everything else is that widget kind's own
//! contract — the same kind of contract [`WidgetKind::ToneCurve`] carries when
//! it says its parameters are point coordinates interleaved:
//!
//! * The **first two parameters** the presentation names are the axes. The
//!   first trades red against blue, the second green against magenta.
//! * The operation publishes **exactly three uniforms: the linear per-channel
//!   gains, in red, green, blue order.** That is what an eyedropper needs in
//!   order to be invertible at all, and it is the honest way to say so —
//!   a node whose effect on a grey is not three gains is not a node an
//!   eyedropper can drive, and declaring the widget on one is the mistake
//!   this refuses rather than approximates.
//!
//! Checked rather than trusted: a node that declares the widget and publishes
//! four uniforms gets no picker, and its sliders go on working.
//!
//! # Why a search rather than an inversion
//!
//! The declared expression is arbitrary — `exp2` today, a table lookup or a
//! polynomial tomorrow — and only its *monotonicity* is part of the bargain:
//! warmer is warmer all the way along, or the slider is not a slider. A
//! bisection needs exactly that and nothing more, so it stays correct across
//! every expression the declaration language can grow, where a closed-form
//! inverse would be a second definition of the node maintained here.
//!
//! It is also free. Twenty halvings of two ranges, three times over, is a few
//! hundred evaluations of two small arithmetic expressions — on a click, not
//! on a frame.

use crate::descriptor::{OpId, ParamId, ParamKind, WidgetKind};
use crate::graph::{EditGraph, OpCapability};

/// Halvings per axis.
///
/// Twenty resolves a ±100 control to about a five-thousandth of a unit, which
/// is far under the precision any of them is displayed at. There is no reason
/// to stop earlier: each step is two multiplications.
const STEPS: u32 = 20;

/// How many times the two axes are solved in turn.
///
/// One pass suffices for a node whose axes are independent, which is what
/// "trades red against blue" and "trades green against magenta" describe. The
/// repeats are for a node whose are not — a green correction that also lifts
/// red would leave the first axis a little out — and they cost nothing worth
/// counting.
const ROUNDS: u32 = 3;

/// Below this a channel carries no ratio worth balancing.
///
/// A sample in the deep shadows has no white balance in it: the logarithms
/// below would run away and the picker would slam a slider to its stop.
/// Refusing is the honest answer, and the caller reports that the point was
/// not usable rather than moving the photograph. The other end — a blown
/// highlight, where every channel has stopped counting — is refused by the
/// caller before the sample is taken, because only the caller can see the
/// sensor value; see [`crate::operation::CLIP_ONSET`].
const FLOOR: f32 = 1e-4;

/// TRACES: FR-DEV-3
/// Move the graph so that `sample` renders neutral.
///
/// `sample` is the linear triple the operation's own gains multiply — camera
/// RGB with the camera's as-shot balance on, *before* the camera matrix and
/// the view transform, and with the sampling operation at its defaults. Not the
/// pixel on the screen: the matrix mixes the channels on the way there, so
/// a colour read after it does not answer to these gains, and a solve over
/// one lands somewhere no sample asked for. Returns whether the graph was
/// moved: `false` where the chain offers no white point widget, or where the
/// colour has no balance in it to correct.
///
/// **Absolute, not relative.** The values written depend on the colour and not
/// on where the sliders happened to be, so sampling the same wall twice lands
/// in the same place — and sampling a second, better neutral corrects the
/// photograph rather than correcting the correction.
pub fn neutralise(graph: &mut EditGraph, sample: [f32; 3]) -> bool {
    if sample.iter().any(|c| !c.is_finite() || *c < FLOOR) {
        return false;
    }
    let Some(axes) = sampler(graph) else {
        return false;
    };
    // The uniform contract, checked before anything moves. A node that
    // declares the widget and publishes something other than three gains
    // cannot be driven from a pixel, and half-solving it would leave the
    // photograph somewhere nobody asked for.
    if gains(graph, axes.op).is_none() {
        return false;
    }

    for _ in 0..ROUNDS {
        solve(
            graph,
            axes.op,
            axes.warm,
            axes.warm_axis,
            sample,
            warm_error,
        );
        solve(
            graph,
            axes.op,
            axes.green,
            axes.green_axis,
            sample,
            green_error,
        );
    }
    true
}

/// The operation an eyedropper writes to, and the two axes it moves.
struct Axes {
    op: OpId,
    /// Red against blue, with its declared travel and display precision.
    warm: ParamId,
    warm_axis: (f32, f32, u8),
    /// Green against magenta, on the same terms.
    green: ParamId,
    green_axis: (f32, f32, u8),
}

/// Find the operation that asked to be driven from a pixel.
///
/// The *first* one, if a chain ever carries two. Two white points in one chain
/// is a chain that has already decided something odd, and picking the first is
/// at least the same answer every time.
fn sampler(graph: &EditGraph) -> Option<Axes> {
    graph.capabilities().into_iter().find_map(|cap| {
        let presentation = cap.presentation.as_ref()?;
        if !presentation.widgets.contains(&WidgetKind::WhitePoint) {
            return None;
        }
        let mut owned = presentation.params.iter().copied();
        let warm = owned.next()?;
        let green = owned.next()?;
        Some(Axes {
            op: cap.id,
            warm,
            warm_axis: axis(&cap, warm)?,
            green,
            green_axis: axis(&cap, green)?,
        })
    })
}

/// How far a parameter may be moved and how finely, from its own declaration.
///
/// A non-scalar axis is refused rather than coerced: a bisection over a list
/// of named alternatives is meaningless, and an operation offering one has not
/// declared what this widget kind needs.
fn axis(cap: &OpCapability, param: ParamId) -> Option<(f32, f32, u8)> {
    cap.params
        .iter()
        .find(|p| p.id == param)
        .and_then(|p| match &p.kind {
            ParamKind::Scalar {
                min,
                max,
                precision,
                ..
            } => Some((*min, *max, *precision)),
            _ => None,
        })
}

/// Round to the precision the parameter is displayed at.
///
/// **The search has to stop somewhere, and this is the honest place.** Twenty
/// halvings land on something like -63.0117, which the panel would draw as
/// -63 and a photographer could never reproduce by hand — so the two would
/// disagree about what the control says, and a slider nudged one unit either
/// way would silently discard a fraction nobody could see.
///
/// It also fixes the case that matters most: sampling something that is
/// *already* neutral. Unrounded, that lands a ten-thousandth off zero, which
/// is a photograph marked modified and a step on the undo stack for a
/// correction of nothing.
fn snap(value: f32, precision: u8) -> f32 {
    let scale = 10f32.powi(i32::from(precision));
    (value * scale).round() / scale
}

/// The three linear gains this operation currently applies, in RGB order.
///
/// `None` unless there are exactly three, which is the widget kind's contract
/// stated as a check. See the module note for why it is a contract and not a
/// guess.
fn gains(graph: &EditGraph, op: OpId) -> Option<[f32; 3]> {
    let uniforms = graph.uniforms_of(op)?;
    match uniforms.as_slice() {
        [r, g, b] => Some([r.value, g.value, b.value]),
        _ => None,
    }
}

/// How far red sits from blue in the corrected sample, in stops.
///
/// Logarithms because the correction is multiplicative and the controls are
/// symmetric: warming by n and cooling by n are exact reciprocals, so an error
/// measured as a ratio is linear in the thing being searched and a difference
/// of products is not.
fn warm_error(gains: [f32; 3], sample: [f32; 3]) -> f32 {
    (gains[0] * sample[0]).ln() - (gains[2] * sample[2]).ln()
}

/// How far green sits from the red-blue midpoint, in stops.
///
/// Against the midpoint rather than against red alone, so that the two axes do
/// not fight: with red and blue already balanced the two are the same number,
/// and while they are not, green is being pulled toward where the other axis
/// is heading rather than toward one end of it.
fn green_error(gains: [f32; 3], sample: [f32; 3]) -> f32 {
    let r = (gains[0] * sample[0]).ln();
    let b = (gains[2] * sample[2]).ln();
    (gains[1] * sample[1]).ln() - 0.5 * (r + b)
}

/// Drive one axis until `error` crosses zero, and leave it there.
///
/// The direction is read off the ends rather than assumed: which way "warmer"
/// runs is the node's business, and a picker that had guessed would move the
/// slider the wrong way on the first node that disagreed with it.
fn solve(
    graph: &mut EditGraph,
    op: OpId,
    param: ParamId,
    (min, max, precision): (f32, f32, u8),
    sample: [f32; 3],
    error: fn([f32; 3], [f32; 3]) -> f32,
) {
    let at = |graph: &mut EditGraph, value: f32| -> f32 {
        graph.set_param(op, param, value);
        gains(graph, op).map_or(0.0, |g| error(g, sample))
    };

    let mut low = min;
    let mut high = max;
    let low_error = at(graph, low);
    let high_error = at(graph, high);

    // No crossing inside the declared travel: the correction this colour needs
    // is more than the control can give. Taken as far as it goes, which is
    // what a photographer would do by hand — refusing would leave a badly cast
    // frame uncorrected because it could not be corrected *entirely*.
    if low_error.is_sign_positive() == high_error.is_sign_positive() {
        let end = if low_error.abs() <= high_error.abs() {
            low
        } else {
            high
        };
        graph.set_param(op, param, snap(end, precision));
        return;
    }

    for _ in 0..STEPS {
        let middle = 0.5 * (low + high);
        if at(graph, middle).is_sign_positive() == low_error.is_sign_positive() {
            low = middle;
        } else {
            high = middle;
        }
    }
    graph.set_param(op, param, snap(0.5 * (low + high), precision));
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The spread of a triple, relative to its middle channel.
    fn cast(rendered: [f32; 3]) -> f32 {
        let high = rendered.iter().copied().fold(f32::MIN, f32::max);
        let low = rendered.iter().copied().fold(f32::MAX, f32::min);
        (high - low) / rendered[1].max(f32::EPSILON)
    }

    /// What the sampled colour becomes once the graph has been moved.
    fn corrected(graph: &EditGraph, sample: [f32; 3]) -> [f32; 3] {
        let axes = sampler(graph).expect("the default chain declares a white point");
        let g = gains(graph, axes.op).expect("three gains");
        [g[0] * sample[0], g[1] * sample[1], g[2] * sample[2]]
    }

    /// TRACES: FR-DEV-3
    /// The whole point: a warm grey comes back grey.
    #[test]
    fn sampling_a_warm_grey_makes_it_grey() {
        let mut graph = EditGraph::default_chain();
        let sample = [0.62, 0.50, 0.40];
        assert!(cast(sample) > 0.3, "the premise: this is a strong cast");

        assert!(neutralise(&mut graph, sample));
        assert!(
            cast(corrected(&graph, sample)) < 0.02,
            "the sample should render neutral: {:?}",
            corrected(&graph, sample)
        );
    }

    /// TRACES: FR-DEV-3
    /// And so does a cold one, which is the other half of the same claim: the
    /// search finds its own direction rather than being told one.
    #[test]
    fn sampling_a_cold_grey_makes_it_grey_too() {
        let mut graph = EditGraph::default_chain();
        let sample = [0.40, 0.50, 0.62];

        assert!(neutralise(&mut graph, sample));
        assert!(
            cast(corrected(&graph, sample)) < 0.02,
            "{:?}",
            corrected(&graph, sample)
        );
    }

    /// TRACES: FR-DEV-3
    /// Sampling is absolute: the second sample corrects the photograph, not
    /// the previous correction.
    ///
    /// The failure this guards is the one every relative picker has — sample a
    /// wall, decide a cloud was better, sample the cloud, and land somewhere
    /// neither of them describes.
    #[test]
    fn a_second_sample_replaces_the_first_rather_than_compounding_it() {
        let mut graph = EditGraph::default_chain();

        assert!(neutralise(&mut graph, [0.62, 0.50, 0.40]));
        assert!(neutralise(&mut graph, [0.44, 0.50, 0.56]));

        let second = [0.44, 0.50, 0.56];
        assert!(
            cast(corrected(&graph, second)) < 0.02,
            "{:?}",
            corrected(&graph, second)
        );
    }

    /// TRACES: FR-DEV-3
    /// A sample with nothing in it is refused rather than acted on.
    ///
    /// Black has no white balance. A picker that treated it as one would take
    /// the ratio of two numbers that are both noise and slam the controls to
    /// their stops, which reads as the feature being broken rather than as the
    /// point having been a poor choice.
    #[test]
    fn a_sample_with_no_balance_in_it_moves_nothing() {
        let mut graph = EditGraph::default_chain();
        let before = crate::Preset::capture(&graph);

        assert!(!neutralise(&mut graph, [0.0, 0.0, 0.0]));
        assert!(!neutralise(&mut graph, [f32::NAN, 0.5, 0.5]));

        assert_eq!(crate::Preset::capture(&graph), before);
    }

    /// TRACES: FR-DEV-3a
    /// The picker is found through the widget kind, not through a name.
    ///
    /// If this ever fails because the declaration moved, the right fix is in
    /// the declaration: nothing here is entitled to know which node it is.
    #[test]
    fn the_chain_offers_exactly_one_white_point() {
        let graph = EditGraph::default_chain();
        let found = graph
            .capabilities()
            .into_iter()
            .filter(|c| {
                c.presentation
                    .as_ref()
                    .is_some_and(|p| p.widgets.contains(&WidgetKind::WhitePoint))
            })
            .count();
        assert_eq!(found, 1, "one operation should ask to be driven by a pixel");
    }
}