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

//! Vignetting correction — the corner falloff a lens imposes.
//!
//! Every lens delivers less light to the corners than to the centre, by up to
//! two stops wide open. This restores it.
//!
//! # Why this one *is* an `Operation`
//!
//! Distortion and CA are [`crate::lens::Warp`]s because they change which
//! pixel is read. Vignetting does not: it applies a gain to the pixel already
//! there. That the gain happens to depend on the pixel's *position* does not
//! make it a coordinate transform — the sampling is unchanged, so it composes
//! as an ordinary colour fragment and costs no extra texture read.
//!
//! It needs the pixel's normalised radius, which a colour fragment is not
//! otherwise given. The composer publishes `radius` in the shader prologue for
//! exactly this reason: it is derived from coordinates the prologue has
//! already computed, so making it available costs nothing.
//!
//! # The model
//!
//! Lensfun's `pa` polynomial, in even powers of the radius:
//!
//! ```text
//! attenuation = 1 + k1·r² + k2·r⁴ + k3·r⁶
//! ```
//!
//! Only even powers, because vignetting is symmetric about the optical axis —
//! an odd term would describe a lens brighter on one side than the other,
//! which is a mount fault rather than a lens characteristic.
//!
//! The value is what the lens *did*, so correcting means dividing by it. That
//! division is the whole reason this operation must run before the tonal
//! stages: a corner recovered by two stops has to be recovered while the
//! highlight headroom to hold it still exists (ARCH §5.2).
use std::sync::{Arc, LazyLock};

use crate::descriptor::{Attribute, LocalizedKey, OpDescriptor, OpId, ParamDescriptor, ParamId};
use crate::operation::{Helper, Operation, Uniform};

pub const ID: OpId = OpId("vignetting");
pub const AMOUNT: ParamId = ParamId("amount");

/// The `k1` coefficient at full manual travel.
///
/// Chosen so +100 lifts the extreme corner by roughly a stop, which covers a
/// fast prime wide open — the case that actually needs correcting.
const MAX_K1: f32 = -0.5;

static DESCRIPTOR: LazyLock<Arc<OpDescriptor>> = LazyLock::new(|| {
    Arc::new(OpDescriptor {
        // Optics rather than effect: this carries lens-profile coefficients
        // and corrects what the lens did. A *creative* vignette is a different
        // operation that does not exist yet, and would be `Effect`.
        attributes: vec![Attribute::Optics],
        id: ID,
        label: LocalizedKey("op.vignetting"),
        // Bidirectional deliberately. Negative values *add* falloff, which is a
        // legitimate creative choice as well as a correction, and a control that
        // only removed vignetting would need a second one beside it to put any
        // back.
        params: vec![ParamDescriptor::amount("amount", "param.vignetting.amount")],
    })
});

/// The `pa` polynomial coefficients, as Lensfun stores them.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Pa {
    pub k1: f32,
    pub k2: f32,
    pub k3: f32,
}

#[derive(Debug, Default, Clone)]
pub struct Vignetting {
    amount: f32,
    profile: Option<Pa>,
}

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

    /// Apply a lens profile's falloff coefficients.
    pub fn set_profile(&mut self, profile: Option<Pa>) {
        self.profile = profile;
    }

    /// The effective coefficients: profile plus manual trim on `k1`.
    ///
    /// The trim lands on `k1` alone. It is the dominant term, and adjusting
    /// the higher orders by eye would change the *shape* of the falloff rather
    /// than its depth — which is not what a photographer reaching for this
    /// slider wants.
    fn coefficients(&self) -> Pa {
        let trim = self.amount / 100.0 * MAX_K1;
        match self.profile {
            Some(p) => Pa {
                k1: p.k1 + trim,
                k2: p.k2,
                k3: p.k3,
            },
            None => Pa {
                k1: trim,
                k2: 0.0,
                k3: 0.0,
            },
        }
    }
}

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

    fn set_param(&mut self, id: ParamId, value: f32) {
        match id {
            AMOUNT => self.amount = value,
            _ => log::warn!("vignetting: unknown parameter {id}"),
        }
    }

    fn param(&self, id: ParamId) -> f32 {
        match id {
            AMOUNT => self.amount,
            _ => 0.0,
        }
    }

    fn is_active(&self) -> bool {
        let c = self.coefficients();
        c.k1 != 0.0 || c.k2 != 0.0 || c.k3 != 0.0
    }

    fn set_lens_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.vignetting));
    }

    fn wgsl_body(&self) -> String {
        // `radius` comes from the prologue: the pixel's distance from the
        // optical axis, normalised so the corner is 1.
        "\
c = c / vignette_attenuation(radius, vig_k1, vig_k2, vig_k3);"
            .into()
    }

    fn uniforms(&self) -> Vec<Uniform> {
        let c = self.coefficients();
        vec![
            Uniform {
                name: "vig_k1",
                value: c.k1,
            },
            Uniform {
                name: "vig_k2",
                value: c.k2,
            },
            Uniform {
                name: "vig_k3",
                value: c.k3,
            },
        ]
    }

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

static VIGNETTE: &[Helper] = &[Helper {
    name: "vignette_attenuation",
    source: "\
// Lensfun's `pa` vignetting polynomial: 1 + k1*r^2 + k2*r^4 + k3*r^6.
//
// Returns what the lens *did* to this pixel, so correcting divides by it.
// Even powers only: vignetting is symmetric about the optical axis, and an
// odd term would describe a lens brighter on one side than the other.
//
// The result is floored well above zero. A profile evaluated slightly outside
// its calibrated range can produce a near-zero or negative attenuation, and
// dividing by that turns the extreme corners into blown or inverted pixels —
// a far more visible fault than the under-correction the floor causes.
fn vignette_attenuation(r: f32, k1: f32, k2: f32, k3: f32) -> f32 {
    let r2 = r * r;
    let a = 1.0 + r2 * (k1 + r2 * (k2 + r2 * k3));
    return max(a, 0.05);
}",
}];

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

    /// The attenuation the shader would compute, mirrored on the CPU so the
    /// maths is testable without a device (ARCH §6.5a).
    fn attenuation(c: Pa, r: f32) -> f32 {
        let r2 = r * r;
        (1.0 + r2 * (c.k1 + r2 * (c.k2 + r2 * c.k3))).max(0.05)
    }

    #[test]
    fn neutral_does_nothing() {
        let v = Vignetting::new();
        assert!(!v.is_active());
        let c = v.coefficients();
        assert_eq!((c.k1, c.k2, c.k3), (0.0, 0.0, 0.0));
    }

    #[test]
    fn a_neutral_polynomial_is_unity_everywhere() {
        // Otherwise opening an image would rescale its brightness for nothing.
        let c = Vignetting::new().coefficients();
        for r in [0.0, 0.25, 0.5, 0.75, 1.0] {
            assert!((attenuation(c, r) - 1.0).abs() < 1e-6, "r={r} was scaled");
        }
    }

    #[test]
    fn the_centre_is_never_altered() {
        // r = 0 kills every term, so the optical axis keeps its exposure
        // whatever the coefficients. If this failed, the correction would be
        // an exposure control with a gradient attached.
        for amount in [-100.0, -50.0, 50.0, 100.0] {
            let mut v = Vignetting::new();
            v.set_param(AMOUNT, amount);
            assert!(
                (attenuation(v.coefficients(), 0.0) - 1.0).abs() < 1e-6,
                "amount {amount} changed the centre"
            );
        }
    }

    #[test]
    fn a_positive_amount_brightens_the_corners() {
        // The correcting direction: the lens darkened the corners, so the
        // attenuation there must be below 1 and the division lifts them.
        let mut v = Vignetting::new();
        v.set_param(AMOUNT, 100.0);
        let a = attenuation(v.coefficients(), 1.0);
        assert!(a < 1.0, "corner attenuation was {a}, expected < 1");
    }

    #[test]
    fn a_negative_amount_darkens_them_instead() {
        // The creative direction. A control that only removed vignetting
        // would need a second one beside it to add any back.
        let mut v = Vignetting::new();
        v.set_param(AMOUNT, -100.0);
        assert!(attenuation(v.coefficients(), 1.0) > 1.0);
    }

    #[test]
    fn falloff_increases_monotonically_with_radius() {
        // Vignetting is a smooth darkening toward the corners. A polynomial
        // that reversed partway would produce a visible bright ring, which
        // reads as a rendering fault rather than as a wrong setting.
        let mut v = Vignetting::new();
        v.set_param(AMOUNT, 100.0);
        let c = v.coefficients();
        let mut prev = f32::MAX;
        for i in 0..=100 {
            let a = attenuation(c, i as f32 / 100.0);
            assert!(
                a <= prev + 1e-6,
                "attenuation rose at r={}",
                i as f32 / 100.0
            );
            prev = a;
        }
    }

    #[test]
    fn full_correction_is_worth_about_a_stop() {
        // The calibration behind MAX_K1. If this drifts, the slider either
        // cannot fix a fast prime or overshoots wildly at half travel.
        let mut v = Vignetting::new();
        v.set_param(AMOUNT, 100.0);
        let gain = 1.0 / attenuation(v.coefficients(), 1.0);
        assert!(
            (1.5..=2.5).contains(&gain),
            "corner gain was {gain}x, expected roughly a stop"
        );
    }

    #[test]
    fn the_attenuation_never_reaches_zero() {
        // Division by a near-zero attenuation blows the corners to white or
        // inverts them. The floor must hold even for coefficients well past
        // anything the sliders can reach.
        let extreme = Pa {
            k1: -5.0,
            k2: -5.0,
            k3: -5.0,
        };
        for i in 0..=100 {
            let a = attenuation(extreme, i as f32 / 100.0);
            assert!(a >= 0.05, "attenuation fell to {a}");
            assert!(a.is_finite() && a > 0.0);
        }
    }

    #[test]
    fn a_profile_corrects_with_the_slider_at_zero() {
        let mut v = Vignetting::new();
        assert!(!v.is_active());
        v.set_profile(Some(Pa {
            k1: -0.3499,
            k2: -0.914,
            k3: 0.6689,
        }));
        assert!(v.is_active());
        assert_eq!(v.param(AMOUNT), 0.0);
    }

    #[test]
    fn a_real_profile_darkens_the_corners() {
        // Coefficients taken from the Lensfun entry for the Canon EF 16-35mm
        // f/2.8L at 20mm, f/2.8 — a real lens wide open, which is the case
        // this correction exists for.
        let profile = Pa {
            k1: -0.3499,
            k2: -0.914,
            k3: 0.6689,
        };
        let mut v = Vignetting::new();
        v.set_profile(Some(profile));

        let c = v.coefficients();
        assert!(attenuation(c, 1.0) < attenuation(c, 0.0));
        assert!((attenuation(c, 0.0) - 1.0).abs() < 1e-6);
    }

    #[test]
    fn the_slider_trims_a_loaded_profile_rather_than_replacing_it() {
        let profile = Pa {
            k1: -0.35,
            k2: -0.91,
            k3: 0.67,
        };
        let mut v = Vignetting::new();
        v.set_profile(Some(profile));
        v.set_param(AMOUNT, 100.0);

        let c = v.coefficients();
        assert_eq!(c.k2, profile.k2, "the higher orders must survive a trim");
        assert_eq!(c.k3, profile.k3);
        assert!((c.k1 - (profile.k1 + MAX_K1)).abs() < 1e-6);
    }

    #[test]
    fn a_profile_can_be_cleared() {
        let mut v = Vignetting::new();
        v.set_profile(Some(Pa {
            k1: -0.3,
            k2: 0.0,
            k3: 0.0,
        }));
        assert!(v.is_active());
        v.set_profile(None);
        assert!(!v.is_active());
    }

    #[test]
    fn the_wgsl_body_reads_its_declared_uniforms_and_the_radius() {
        let mut v = Vignetting::new();
        v.set_param(AMOUNT, 50.0);
        let body = v.wgsl_body();
        for u in v.uniforms() {
            assert!(body.contains(u.name), "{} is declared but unused", u.name);
        }
        assert!(
            body.contains("radius"),
            "vignetting is radial and must read the prologue's radius"
        );
    }

    #[test]
    fn every_default_is_neutral() {
        let mut v = Vignetting::new();
        for p in &DESCRIPTOR.params {
            v.set_param(p.id, p.default);
        }
        assert!(!v.is_active());
    }
}