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

//! Geometric distortion correction.
//!
//! Straightens the lines a lens bends: barrel distortion on wide angles,
//! pincushion on telephotos. A [`crate::lens::Warp`] rather than an
//! [`crate::operation::Operation`], because it changes *where* a pixel is read
//! from rather than what its value becomes.
//!
//! # The model
//!
//! Lensfun's `ptlens` model, matched deliberately so a lens profile from the
//! Lensfun database applies with no conversion:
//!
//! ```text
//! r_d = r_u · (a·r_u³ + b·r_u² + c·r_u + 1 − a − b − c)
//! ```
//!
//! The `1 − a − b − c` term is not decoration: it forces the polynomial to
//! equal 1 at `r_u = 1`, pinning the image corner in place. Without it every
//! coefficient change would also rescale the frame, so the distortion slider
//! would double as a zoom and no setting would leave the framing alone.
//!
//! `a` and `b` are the higher-order terms that describe a lens's real,
//! slightly wavy profile; `c` alone gives the simple barrel/pincushion shape.
//! The manual control drives `c` only — a single slider cannot meaningfully
//! set three correlated coefficients, and hand-correcting a lens with no
//! profile is a "make the horizon straight" task, which one term does well.
//! The full triple is reachable by loading a profile.
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("distortion");
pub const AMOUNT: ParamId = ParamId("amount");

static DESCRIPTOR: LazyLock<Arc<OpDescriptor>> = LazyLock::new(|| {
    Arc::new(OpDescriptor {
        attributes: vec![Attribute::Optics],
        id: ID,
        label: LocalizedKey("op.distortion"),
        // ±100 maps to a ±0.25 cubic coefficient. That covers an uncorrected
        // fisheye at one end and strong pincushion at the other; beyond it the
        // inverse mapping stops being single-valued near the corners and the
        // correction folds the image over itself.
        params: vec![ParamDescriptor::scalar(
            "amount",
            "param.distortion.amount",
            -100.0,
            100.0,
            0.0,
            Unit::None,
            Scale::Linear,
            0,
        )],
    })
});

/// The cubic coefficient at full slider travel.
const MAX_COEFF: f32 = 0.25;

#[derive(Debug, Default, Clone)]
pub struct Distortion {
    amount: f32,
    /// Profile coefficients, when a lens profile is loaded. `None` means the
    /// manual slider drives `c` alone.
    profile: Option<PtLens>,
}

/// The three `ptlens` coefficients, as Lensfun stores them.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PtLens {
    pub a: f32,
    pub b: f32,
    pub c: f32,
}

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

    /// Apply a lens profile's coefficients.
    ///
    /// The manual slider then acts as a *trim* on top: photographers routinely
    /// find a profile slightly over- or under-corrects on their copy of a
    /// lens, and having to choose between "profile" and "manual" would make
    /// that untunable.
    pub fn set_profile(&mut self, profile: Option<PtLens>) {
        self.profile = profile;
    }

    /// The effective coefficients: profile plus manual trim.
    fn coefficients(&self) -> PtLens {
        let trim = self.amount / 100.0 * MAX_COEFF;
        match self.profile {
            Some(p) => PtLens {
                a: p.a,
                b: p.b,
                c: p.c + trim,
            },
            None => PtLens {
                a: 0.0,
                b: 0.0,
                c: trim,
            },
        }
    }
}

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

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

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

    fn is_active(&self) -> bool {
        // A loaded profile corrects even with the slider at zero — that is
        // the whole point of a profile.
        let c = self.coefficients();
        c.a != 0.0 || c.b != 0.0 || c.c != 0.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.distortion));
    }

    fn wgsl_body(&self) -> String {
        // Written against `p`, which is already normalised and centred.
        "\
let r = length(p);
p = p * ptlens_scale(r, dist_a, dist_b, dist_c);"
            .into()
    }

    fn uniforms(&self) -> Vec<Uniform> {
        let c = self.coefficients();
        vec![
            Uniform {
                name: "dist_a",
                value: c.a,
            },
            Uniform {
                name: "dist_b",
                value: c.b,
            },
            Uniform {
                name: "dist_c",
                value: c.c,
            },
        ]
    }

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

static PTLENS: &[Helper] = &[Helper {
    name: "ptlens_scale",
    source: "\
// The `ptlens` radial polynomial (Lensfun's model).
//
// Returns the factor mapping an undistorted radius to the distorted radius
// it should be sampled from. The trailing `1 - a - b - c` normalises the
// polynomial to 1 at r = 1, which pins the corner and stops a coefficient
// change from also rescaling the frame.
fn ptlens_scale(r: f32, a: f32, b: f32, c: f32) -> f32 {
    let d = 1.0 - a - b - c;
    return ((a * r + b) * r + c) * r + d;
}",
}];

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

    /// The scale factor the shader would compute, mirrored on the CPU so the
    /// maths is testable without a device (ARCH §6.5a).
    fn scale(c: PtLens, r: f32) -> f32 {
        let d = 1.0 - c.a - c.b - c.c;
        ((c.a * r + c.b) * r + c.c) * r + d
    }

    #[test]
    fn neutral_does_nothing() {
        let d = Distortion::new();
        assert!(!d.is_active());
        let c = d.coefficients();
        assert_eq!((c.a, c.b, c.c), (0.0, 0.0, 0.0));
    }

    #[test]
    fn a_neutral_polynomial_is_the_identity() {
        // Every radius must map to itself when no correction is set,
        // otherwise opening an image would resample it for nothing.
        let c = Distortion::new().coefficients();
        for r in [0.0, 0.25, 0.5, 0.75, 1.0] {
            assert!((scale(c, r) - 1.0).abs() < 1e-6, "r={r} was rescaled");
        }
    }

    #[test]
    fn the_corner_is_pinned_whatever_the_coefficients() {
        // The property the `1 - a - b - c` term exists for: correction must
        // not silently zoom the frame. If this fails, the distortion slider
        // doubles as a crop and no setting leaves framing untouched.
        for amount in [-100.0, -50.0, -1.0, 1.0, 50.0, 100.0] {
            let mut d = Distortion::new();
            d.set_param(AMOUNT, amount);
            let s = scale(d.coefficients(), 1.0);
            assert!(
                (s - 1.0).abs() < 1e-5,
                "amount {amount} moved the corner by {}",
                s - 1.0
            );
        }
    }

    #[test]
    fn the_centre_never_moves() {
        // r = 0 is the optical axis; a radial model must leave it fixed, and
        // `p * scale` does so for any finite scale.
        let mut d = Distortion::new();
        d.set_param(AMOUNT, 100.0);
        assert!(scale(d.coefficients(), 0.0).is_finite());
    }

    #[test]
    fn positive_amounts_correct_barrel_distortion() {
        // Barrel distortion pushes detail outward, so correcting it must
        // sample from further out at mid radii — an inverse map (see the
        // `warp` module docs), which is why "correct barrel" magnifies.
        let mut d = Distortion::new();
        d.set_param(AMOUNT, 100.0);
        let s = scale(d.coefficients(), 0.5);
        assert!(s < 1.0, "mid-radius scale was {s}, expected < 1");
    }

    #[test]
    fn negative_amounts_go_the_other_way() {
        let mut pin = Distortion::new();
        pin.set_param(AMOUNT, -100.0);
        let mut bar = Distortion::new();
        bar.set_param(AMOUNT, 100.0);
        assert!(scale(pin.coefficients(), 0.5) > scale(bar.coefficients(), 0.5));
    }

    #[test]
    fn the_mapping_stays_monotonic_across_the_whole_range() {
        // If radius stops increasing with radius, the correction folds the
        // image over itself and produces a mirrored ring. This is what bounds
        // the slider at ±100, so it is worth asserting rather than trusting.
        for amount in [-100.0, -50.0, 0.0, 50.0, 100.0] {
            let mut d = Distortion::new();
            d.set_param(AMOUNT, amount);
            let c = d.coefficients();
            let mut prev = 0.0;
            for i in 1..=100 {
                let r = i as f32 / 100.0;
                let mapped = r * scale(c, r);
                assert!(
                    mapped > prev,
                    "amount {amount}: mapping folded at r={r} ({mapped} <= {prev})"
                );
                prev = mapped;
            }
        }
    }

    #[test]
    fn a_profile_corrects_with_the_slider_at_zero() {
        // Loading a lens profile must do something on its own; requiring the
        // user to also move a slider would make profiles pointless.
        let mut d = Distortion::new();
        assert!(!d.is_active());
        d.set_profile(Some(PtLens {
            a: 0.0168,
            b: -0.0320,
            c: -0.0287,
        }));
        assert!(d.is_active());
        assert_eq!(d.param(AMOUNT), 0.0);
    }

    #[test]
    fn the_slider_trims_a_loaded_profile_rather_than_replacing_it() {
        // A profile that over-corrects on this copy of the lens must stay
        // tunable, so the manual control adds to `c` and leaves a and b.
        let profile = PtLens {
            a: 0.01,
            b: -0.02,
            c: 0.03,
        };
        let mut d = Distortion::new();
        d.set_profile(Some(profile));
        d.set_param(AMOUNT, 100.0);

        let c = d.coefficients();
        assert_eq!(c.a, profile.a, "the profile's a must survive a trim");
        assert_eq!(c.b, profile.b);
        assert!((c.c - (profile.c + MAX_COEFF)).abs() < 1e-6);
    }

    #[test]
    fn a_profile_can_be_cleared() {
        let mut d = Distortion::new();
        d.set_profile(Some(PtLens {
            a: 0.01,
            b: 0.0,
            c: 0.0,
        }));
        assert!(d.is_active());
        d.set_profile(None);
        assert!(!d.is_active(), "clearing a profile must return to neutral");
    }

    #[test]
    fn the_wgsl_body_reads_its_declared_uniforms() {
        // The composer rewrites bare names; a body naming something it did
        // not declare would compile to a reference to a nonexistent field.
        let mut d = Distortion::new();
        d.set_param(AMOUNT, 50.0);
        let body = d.wgsl_body();
        for u in d.uniforms() {
            assert!(body.contains(u.name), "{} is declared but unused", u.name);
        }
    }
}