DarkRoom/core/dr-pipeline/src/ops/aberration.rs

//! Lateral chromatic aberration correction.
//!
//! The purple-and-green fringing on high-contrast edges toward the frame
//! corners. A lens focuses short wavelengths and long wavelengths at slightly
//! different magnifications, so the red, green and blue images it projects are
//! very slightly different sizes. Rescaling two of them about the optical axis
//! puts them back on top of each other.
//!
//! # Why this cannot be an `Operation`
//!
//! This is the correction that forced [`crate::lens::Warp`] to exist. An
//! [`crate::operation::Operation`] receives `c` — a colour already sampled,
//! with all three channels fetched from *one* coordinate. Lateral CA needs
//! three *different* coordinates, and by the time an operation runs, the
//! information needed to pick them is gone. So it declares
//! [`Warp::splits_channels`] and the composer emits the three-sample path.
//!
//! # The model
//!
//! Lensfun's `poly3`, reduced to its linear term: a per-channel radial scale
//! with green as the fixed reference.
//!
//! ```text
//! r_red  = r · v_red
//! r_blue = r · v_blue
//! ```
//!
//! Green is never moved, and that is a deliberate asymmetry rather than an
//! arbitrary choice of reference. Green carries most of the luminance a Bayer
//! sensor records — twice the photosites of red or blue — so leaving it
//! untouched means a mis-set correction shifts the channels that contribute
//! least to perceived sharpness. Scaling all three about a virtual reference
//! would soften the image even when the correction is right.
use std::sync::{Arc, LazyLock};

use crate::descriptor::{
    Attribute, LocalizedKey, OpDescriptor, OpId, ParamDescriptor, ParamId, Scale, Unit,
};
use crate::lens::Warp;
use crate::operation::{Helper, Uniform};

pub const ID: OpId = OpId("aberration");
pub const RED: ParamId = ParamId("red");
pub const BLUE: ParamId = ParamId("blue");

/// The radial scale at full slider travel, as a fraction.
///
/// Lateral CA is a tiny effect — the database's own coefficients sit within
/// ±0.1% — so a slider spanning ±0.5% covers every real lens with enough
/// resolution left to tune by eye at 100%.
const MAX_SCALE: f32 = 0.005;

static DESCRIPTOR: LazyLock<Arc<OpDescriptor>> = LazyLock::new(|| {
    Arc::new(OpDescriptor {
        attributes: vec![Attribute::Optics],
        id: ID,
        label: LocalizedKey("op.aberration"),
        params: vec![
            // Two independent controls rather than one: the red and blue
            // displacements are caused by different ends of the spectrum and are
            // not symmetric, so a single "fringing" slider could not remove both.
            ParamDescriptor::scalar(
                "red",
                "param.aberration.red",
                -100.0,
                100.0,
                0.0,
                Unit::None,
                Scale::Linear,
                0,
            ),
            ParamDescriptor::scalar(
                "blue",
                "param.aberration.blue",
                -100.0,
                100.0,
                0.0,
                Unit::None,
                Scale::Linear,
                0,
            ),
        ],
    })
});

#[derive(Debug, Default, Clone)]
pub struct Aberration {
    red: f32,
    blue: f32,
    /// Per-channel scales from a lens profile, when one is loaded.
    profile: Option<(f32, f32)>,
}

impl Aberration {
    pub fn new() -> Self {
        Self::default()
    }

    /// Apply a profile's red and blue radial scales.
    ///
    /// As with distortion, the sliders then trim rather than replace: CA
    /// varies between copies of a lens and with focus distance, so a profile
    /// gets close and the user finishes the job.
    pub fn set_profile(&mut self, scales: Option<(f32, f32)>) {
        self.profile = scales;
    }

    /// The effective per-channel scales, profile plus manual trim.
    fn scales(&self) -> (f32, f32) {
        let (base_r, base_b) = self.profile.unwrap_or((1.0, 1.0));
        (
            base_r + self.red / 100.0 * MAX_SCALE,
            base_b + self.blue / 100.0 * MAX_SCALE,
        )
    }
}

impl Warp for Aberration {
    fn descriptor(&self) -> Arc<OpDescriptor> {
        DESCRIPTOR.clone()
    }

    fn set_param(&mut self, id: ParamId, value: f32) {
        match id {
            RED => self.red = value,
            BLUE => self.blue = value,
            _ => log::warn!("aberration: unknown parameter {id}"),
        }
    }

    fn param(&self, id: ParamId) -> f32 {
        match id {
            RED => self.red,
            BLUE => self.blue,
            _ => 0.0,
        }
    }

    fn is_active(&self) -> bool {
        let (r, b) = self.scales();
        r != 1.0 || b != 1.0
    }

    fn set_profile(&mut self, profile: Option<&crate::lens::LensProfile>) {
        // The inherent `set_profile` taking just this correction's own
        // coefficients, not this trait method: an inherent method wins over a
        // trait one of the same name, so this is a narrowing and not a loop.
        self.set_profile(
            profile
                .and_then(|p| p.tca)
                .map(|t| (t.red_scale, t.blue_scale)),
        );
    }

    fn wgsl_body(&self) -> String {
        // `p_r` and `p_b` enter equal to `p` and are carried out of the block.
        // Green is deliberately absent: it is the reference and never moves.
        "\
p_r = p * ca_red;
p_b = p * ca_blue;"
            .into()
    }

    fn uniforms(&self) -> Vec<Uniform> {
        let (red, blue) = self.scales();
        vec![
            Uniform {
                name: "ca_red",
                value: red,
            },
            Uniform {
                name: "ca_blue",
                value: blue,
            },
        ]
    }

    fn splits_channels(&self) -> bool {
        true
    }

    fn helpers(&self) -> &'static [Helper] {
        &[]
    }
}

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

    #[test]
    fn neutral_does_nothing() {
        let a = Aberration::new();
        assert!(!a.is_active());
        assert_eq!(a.scales(), (1.0, 1.0));
    }

    #[test]
    fn a_neutral_correction_leaves_every_channel_coincident() {
        // If the channels diverge at neutral, an image with no CA correction
        // is resampled into colour fringing that was not there.
        let (r, b) = Aberration::new().scales();
        assert_eq!(r, 1.0);
        assert_eq!(b, 1.0);
    }

    #[test]
    fn the_channels_are_controlled_independently() {
        // Red and blue displacement have different causes and are not
        // symmetric; one slider could not remove both.
        let mut a = Aberration::new();
        a.set_param(RED, 100.0);
        let (r, b) = a.scales();
        assert!(r > 1.0, "red should be scaled");
        assert_eq!(b, 1.0, "blue must be untouched by the red control");
    }

    #[test]
    fn green_is_never_scaled() {
        // The reference channel. Asserted through the shader body, since that
        // is where a stray green term would actually do damage.
        let mut a = Aberration::new();
        a.set_param(RED, 50.0);
        a.set_param(BLUE, -50.0);
        let body = a.wgsl_body();
        assert!(body.contains("p_r"), "red must be displaced");
        assert!(body.contains("p_b"), "blue must be displaced");
        assert!(
            !body.contains("p_g"),
            "green is the reference and must never be displaced"
        );
    }

    #[test]
    fn the_centre_never_moves() {
        // A radial scale about the optical axis leaves r = 0 fixed whatever
        // the coefficients, which is why CA correction cannot shift a frame.
        let mut a = Aberration::new();
        a.set_param(RED, 100.0);
        a.set_param(BLUE, -100.0);
        let (r, b) = a.scales();
        for scale in [r, b] {
            assert_eq!(0.0 * scale, 0.0);
        }
    }

    #[test]
    fn the_correction_stays_subpixel_at_the_extremes() {
        // Lateral CA is a fraction of a percent. If full travel displaced a
        // corner by more than a pixel or two on a 6000px frame, the slider
        // would be a smear control rather than a correction.
        let mut a = Aberration::new();
        a.set_param(RED, 100.0);
        a.set_param(BLUE, -100.0);
        let (r, b) = a.scales();
        // Half-diagonal of a 6000x4000 frame, the worst case.
        let half_diag = ((6000.0f32 / 2.0).powi(2) + (4000.0f32 / 2.0).powi(2)).sqrt();
        for scale in [r, b] {
            let px = (scale - 1.0).abs() * half_diag;
            assert!(px < 25.0, "full travel displaces the corner by {px} px");
        }
    }

    #[test]
    fn it_always_requests_the_per_channel_path() {
        // The declaration that makes the composer emit three samples. Without
        // it the fragment would write `p_r`/`p_b` that nothing reads.
        assert!(Aberration::new().splits_channels());
    }

    #[test]
    fn a_profile_corrects_with_both_sliders_at_zero() {
        let mut a = Aberration::new();
        assert!(!a.is_active());
        a.set_profile(Some((1.0003211, 1.0000667)));
        assert!(a.is_active());
        assert_eq!(a.param(RED), 0.0);
        assert_eq!(a.param(BLUE), 0.0);
    }

    #[test]
    fn the_sliders_trim_a_loaded_profile() {
        // CA varies between copies of a lens and with focus distance, so a
        // profile must remain tunable rather than being all-or-nothing.
        let profile = (1.0003, 1.0001);
        let mut a = Aberration::new();
        a.set_profile(Some(profile));
        a.set_param(RED, 100.0);

        let (r, b) = a.scales();
        assert!((r - (profile.0 + MAX_SCALE)).abs() < 1e-9);
        assert_eq!(b, profile.1, "the red trim must not disturb blue");
    }

    #[test]
    fn a_profile_can_be_cleared() {
        let mut a = Aberration::new();
        a.set_profile(Some((1.0003, 1.0001)));
        assert!(a.is_active());
        a.set_profile(None);
        assert!(!a.is_active());
    }

    #[test]
    fn the_wgsl_body_reads_its_declared_uniforms() {
        let mut a = Aberration::new();
        a.set_param(RED, 50.0);
        let body = a.wgsl_body();
        for u in a.uniforms() {
            assert!(body.contains(u.name), "{} is declared but unused", u.name);
        }
    }

    #[test]
    fn every_default_is_neutral() {
        let mut a = Aberration::new();
        for p in &DESCRIPTOR.params {
            a.set_param(p.id, p.default);
        }
        assert!(!a.is_active(), "descriptor defaults must be neutral");
    }
}