Implement PQ transfer functions

jxl/src/color/tf.rs

@@ -3,6 +3,8 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+use crate::util::eval_rational_poly;
+
 const SRGB_POWTABLE_UPPER: [u8; 16] = [
     0x00, 0x0a, 0x19, 0x26, 0x32, 0x41, 0x4d, 0x5c, 0x68, 0x75, 0x83, 0x8f, 0xa0, 0xaa, 0xb9, 0xc6,
 ];
@@ -11,6 +13,39 @@ const SRGB_POWTABLE_LOWER: [u8; 16] = [
     0x00, 0xb7, 0x04, 0x0d, 0xcb, 0xe7, 0x41, 0x68, 0x51, 0xd1, 0xeb, 0xf2, 0x00, 0xb7, 0x04, 0x0d,
 ];
 
+const PQ_M1: f64 = 2610.0 / 16384.0;
+const PQ_M2: f64 = (2523.0 / 4096.0) * 128.0;
+const PQ_C1: f64 = 3424.0 / 4096.0;
+const PQ_C2: f64 = (2413.0 / 4096.0) * 32.0;
+const PQ_C3: f64 = (2392.0 / 4096.0) * 32.0;
+
+const PQ_EOTF_P: [f32; 5] = [
+    2.6297566e-4,
+    -6.235531e-3,
+    7.386023e-1,
+    2.6455317,
+    5.500349e-1,
+];
+const PQ_EOTF_Q: [f32; 5] = [
+    4.213501e2,
+    -4.2873682e2,
+    1.7436467e2,
+    -3.3907887e1,
+    2.6771877,
+];
+
+const PQ_INV_EOTF_P: [f32; 5] = [1.351392e-2, -1.095778, 5.522776e1, 1.492516e2, 4.838434e1];
+const PQ_INV_EOTF_Q: [f32; 5] = [1.012416, 2.016708e1, 9.26371e1, 1.120607e2, 2.590418e1];
+const PQ_INV_EOTF_P_SMALL: [f32; 5] = [
+    9.863406e-6,
+    3.881234e-1,
+    1.352821e2,
+    6.889862e4,
+    -2.864824e5,
+];
+const PQ_INV_EOTF_Q_SMALL: [f32; 5] =
+    [3.371868e1, 1.477719e3, 1.608477e4, -4.389884e4, -2.072546e5];
+
 /// Converts the linear samples with the sRGB transfer curve.
 // Fast linear to sRGB conversion, ported from libjxl. Max error ~1.7e-4
 pub fn linear_to_srgb_fast(samples: &mut [f32]) {
@@ -62,7 +97,7 @@ pub fn linear_to_srgb(samples: &mut [f32]) {
         *x = if a <= 0.0031308 {
             a * 12.92
         } else {
-            crate::util::eval_rational_poly(a.sqrt(), P, Q)
+            eval_rational_poly(a.sqrt(), P, Q)
         }
         .copysign(*x);
     }
@@ -93,7 +128,7 @@ pub fn srgb_to_linear(samples: &mut [f32]) {
         *x = if a <= 0.04045 {
             a / 12.92
         } else {
-            crate::util::eval_rational_poly(a, P, Q)
+            eval_rational_poly(a, P, Q)
         }
         .copysign(*x);
     }
@@ -125,6 +160,88 @@ pub fn bt709_to_linear(samples: &mut [f32]) {
     }
 }
 
+/// Converts linear sample to PQ signal using PQ inverse EOTF, where linear sample value of 1.0
+/// represents `intensity_target` display nits.
+///
+/// This version uses original EOTF using double precision arithmetic internally.
+pub fn linear_to_pq_precise(intensity_target: f32, samples: &mut [f32]) {
+    let mult = intensity_target as f64 * 10000f64.recip();
+
+    for s in samples {
+        if *s == 0.0 {
+            continue;
+        }
+
+        let a = s.abs() as f64;
+        let xp = (a * mult).powf(PQ_M1);
+        let num = PQ_C1 + xp * PQ_C2;
+        let den = 1.0 + xp * PQ_C3;
+        let e = (num / den).powf(PQ_M2);
+        *s = (e as f32).copysign(*s);
+    }
+}
+
+/// Converts PQ signal to linear sample using PQ EOTF, where linear sample value of 1.0 represents
+/// `intensity_target` display nits.
+///
+/// This version uses original EOTF using double precision arithmetic internally.
+pub fn pq_to_linear_precise(intensity_target: f32, samples: &mut [f32]) {
+    let mult = 10000.0 / intensity_target as f64;
+
+    for s in samples {
+        if *s == 0.0 {
+            continue;
+        }
+
+        let a = s.abs() as f64;
+        let xp = a.powf(PQ_M2.recip());
+        let num = (xp - PQ_C1).max(0.0);
+        let den = PQ_C2 - PQ_C3 * xp;
+        let y = (num / den).powf(PQ_M1.recip());
+        *s = ((y * mult) as f32).copysign(*s);
+    }
+}
+
+/// Converts linear sample to PQ signal using PQ inverse EOTF, where linear sample value of 1.0
+/// represents `intensity_target` display nits.
+///
+/// This version uses approximate curve using rational polynomial.
+// Max error: ~7e-7 at intensity_target = 10000
+pub fn linear_to_pq(intensity_target: f32, samples: &mut [f32]) {
+    let y_mult = intensity_target * 10000f32.recip();
+
+    for s in samples {
+        let a = s.abs();
+        let a_scaled = a * y_mult;
+        let a_1_4 = a_scaled.sqrt().sqrt();
+
+        let y = if a < 1e-4 {
+            eval_rational_poly(a_1_4, PQ_INV_EOTF_P_SMALL, PQ_INV_EOTF_Q_SMALL)
+        } else {
+            eval_rational_poly(a_1_4, PQ_INV_EOTF_P, PQ_INV_EOTF_Q)
+        };
+
+        *s = y.copysign(*s);
+    }
+}
+
+/// Converts PQ signal to linear sample using PQ EOTF, where linear sample value of 1.0 represents
+/// `intensity_target` display nits.
+///
+/// This version uses approximate curve using rational polynomial.
+// Max error: ~3e-6 at intensity_target = 10000
+pub fn pq_to_linear(intensity_target: f32, samples: &mut [f32]) {
+    let y_mult = 10000.0 / intensity_target;
+
+    for s in samples {
+        let a = s.abs();
+        // a + a * a
+        let x = a.mul_add(a, a);
+        let y = eval_rational_poly(x, PQ_EOTF_P, PQ_EOTF_Q);
+        *s = (y * y_mult).copysign(*s);
+    }
+}
+
 #[cfg(test)]
 mod test {
     use test_log::test;
@@ -137,11 +254,10 @@ mod test {
     ) -> arbtest::arbitrary::Result<Vec<f32>> {
         const DENOM: u32 = 1 << 24;
 
-        let len = u.arbitrary_len::<u32>()?;
-        let mut samples = Vec::with_capacity(len);
+        let mut samples = Vec::new();
 
         // uniform distribution in [-1.0, 1.0]
-        for _ in 0..len {
+        while !u.is_empty() {
             let a: u32 = u.int_in_range(0..=DENOM)?;
             let signed: bool = u.arbitrary()?;
             let x = a as f32 / DENOM as f32;
@@ -189,4 +305,33 @@ mod test {
             Ok(())
         });
     }
+
+    #[test]
+    fn linear_to_pq_arb() {
+        arbtest::arbtest(|u| {
+            let intensity_target = u.int_in_range(9900..=10100)? as f32;
+            let mut samples: Vec<f32> = arb_samples(u)?;
+            let mut precise = samples.clone();
+
+            linear_to_pq(intensity_target, &mut samples);
+            linear_to_pq_precise(intensity_target, &mut precise);
+            // Error seems to increase at intensity_target < 10000
+            assert_all_almost_eq!(&samples, &precise, 8e-7);
+            Ok(())
+        });
+    }
+
+    #[test]
+    fn pq_to_linear_arb() {
+        arbtest::arbtest(|u| {
+            let intensity_target = u.int_in_range(9900..=10100)? as f32;
+            let mut samples: Vec<f32> = arb_samples(u)?;
+            let mut precise = samples.clone();
+
+            pq_to_linear(intensity_target, &mut samples);
+            pq_to_linear_precise(intensity_target, &mut precise);
+            assert_all_almost_eq!(&samples, &precise, 3e-6);
+            Ok(())
+        });
+    }
 }