JPEG XL: faster weighted predictor, sRGB encoding and Gaborish

The weighted predictor keeps its four subpredictors in fields instead of
allocating arrays per pixel (lossless -e 7: 24% faster, identical output).
Linear light becomes sRGB like libjxl's TF_SRGB, a rational polynomial of the
square root within 5e-7 instead of a power (the conformance results are
unchanged). Gaborish mirrors only at the edges (identical output). Lossy
2048x1536: about 1.2 s -> 0.95 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Lfv41PqHEGbU7SLg5onpG1

Author: matmen

src/codecs/jxl/color.ts

@@ -7,7 +7,15 @@ import { ImageHeader, TF_PQ, TF_HLG, PRI_BT2100, PRI_P3, EXTRA_ALPHA, EXTRA_SPOT
 import { toSRGB, hdrToLinear, ToneMapping, REFERENCE_WHITE } from '../hdr.ts';
 import type { Planes } from './frame.ts';
 
-const srgbEncode = (v: number): number => v <= 0.0031308 ? 12.92 * v : 1.055 * v ** (1 / 2.4) - 0.055;
+// linear light to sRGB like libjxl (TF_SRGB::EncodedFromDisplay, lib/jxl/cms/transfer_functions-inl.h): a rational
+// polynomial of the square root instead of the power, within 5e-7
+const [P0, P1, P2, P3, P4] = [-5.135152395e-04, 5.287254571e-03, 3.903842876e-01, 1.474205315e+00, 7.352629620e-01].map(Math.fround);
+const [Q0, Q1, Q2, Q3, Q4] = [1.004519624e-02, 3.036675394e-01, 1.340816930e+00, 9.258482155e-01, 2.424867759e-02].map(Math.fround);
+const srgbEncode = (v: number): number => {
+    if (v <= 0.0031308) return 12.92 * v;
+    const x = Math.sqrt(v);
+    return ((((P4 * x + P3) * x + P2) * x + P1) * x + P0) / ((((Q4 * x + Q3) * x + Q2) * x + Q1) * x + Q0);
+};
 
 /**
  * How frames become output colors: XYB to sRGB (or to linear light, 1 being SDR white, for HDR images), other images

src/codecs/jxl/filters.ts

@@ -13,10 +13,15 @@ export function gaborish(channels: Float32Array[], w: number, h: number, weights
         const out = new Float32Array(w * h);
         for (let y = 0; y < h; y++) {
             const t = mirror(y - 1, h) * w, m = y * w, b = mirror(y + 1, h) * w;
-            for (let x = 0; x < w; x++) {
-                const l = mirror(x - 1, w), r = mirror(x + 1, w);
-                out[m + x] = (src[m + x] + (src[t + x] + src[b + x] + src[m + l] + src[m + r]) * w0 + (src[t + l] + src[t + r] + src[b + l] + src[b + r]) * w1) * norm;
+            const at = (x: number, l: number, r: number): number =>
+                (src[m + x] + (src[t + x] + src[b + x] + src[m + l] + src[m + r]) * w0 + (src[t + l] + src[t + r] + src[b + l] + src[b + r]) * w1) * norm;
+            out[m] = at(0, mirror(-1, w), mirror(1, w));
+            // (the columns between the edges: no mirroring)
+            for (let x = 1; x < w - 1; x++) {
+                out[m + x] = (src[m + x] + (src[t + x] + src[b + x] + src[m + x - 1] + src[m + x + 1]) * w0 +
+                    (src[t + x - 1] + src[t + x + 1] + src[b + x - 1] + src[b + x + 1]) * w1) * norm;
             }
+            if (w > 1) out[m + w - 1] = at(w - 1, w - 2, mirror(w, w));
         }
         return out;
     });

src/codecs/jxl/modular.ts

@@ -44,33 +44,47 @@ function mulShift24(a: number, b: number): number {
     return Number((BigInt(a) * BigInt(b)) >> 24n);
 }
 
+// a weight of the subpredictors from their errors around (x, y): the errors' sum (wrapping, as unsigned 32 bits)
+function weight(w: number, e: Uint32Array, posN: number, posNE: number, posNW: number): number {
+    const sum = (e[posN] + e[posNE] + e[posNW]) >>> 0;
+    const shift = sum === 0xffffffff ? 27 : Math.max(floorLog2(sum + 1) - 5, 0);
+    return 4 + ((w * DIV_LOOKUP[sum >>> shift]) >>> shift);
+}
+
 /** The self-correcting (weighted) predictor's state over a channel */
 class WeightedPredictor {
-    private readonly predErrors: Uint32Array[];
+    // (the subpredictors' errors and the error of the prediction, of two rows)
+    private readonly e0: Uint32Array;
+    private readonly e1: Uint32Array;
+    private readonly e2: Uint32Array;
+    private readonly e3: Uint32Array;
     private readonly error: Int32Array;
-    private readonly prediction = new Float64Array(4);
+    private readonly w0: number;
+    private readonly w1: number;
+    private readonly w2: number;
+    private readonly w3: number;
+    private p0 = 0;
+    private p1 = 0;
+    private p2 = 0;
+    private p3 = 0;
     private pred = 0;
     /** the property (the largest error of the neighbors) of the last prediction */
     maxError = 0;
 
     constructor(private readonly header: WPHeader, private readonly xsize: number) {
-        this.predErrors = [0, 1, 2, 3].map(() => new Uint32Array((xsize + 2) * 2));
-        this.error = new Int32Array((xsize + 2) * 2);
+        const n = (xsize + 2) * 2;
+        [this.e0, this.e1, this.e2, this.e3] = [0, 1, 2, 3].map(() => new Uint32Array(n));
+        this.error = new Int32Array(n);
+        [this.w0, this.w1, this.w2, this.w3] = header.w;
     }
 
     /** the prediction (rounded) of (x, y) from its neighbors */
     predict(x: number, y: number, N: number, W: number, NE: number, NW: number, NN: number): number {
         const h = this.header, xsize = this.xsize;
         const cur = y & 1 ? 0 : xsize + 2, prev = y & 1 ? xsize + 2 : 0;
         const posN = prev + x, posNE = x < xsize - 1 ? posN + 1 : posN, posNW = x > 0 ? posN - 1 : posN;
-        const weights = [0, 0, 0, 0];
-        for (let i = 0; i < 4; i++) {
-            const e = this.predErrors[i];
-            // (wrapping, as unsigned 32 bits)
-            const sum = (e[posN] + e[posNE] + e[posNW]) >>> 0;
-            const shift = Math.max((sum === 0xffffffff ? 32 : floorLog2(sum + 1)) - 5, 0);
-            weights[i] = 4 + ((h.w[i] * DIV_LOOKUP[Math.floor(sum / 2 ** shift)]) >>> shift);
-        }
+        let w0 = weight(this.w0, this.e0, posN, posNE, posNW), w1 = weight(this.w1, this.e1, posN, posNE, posNW);
+        let w2 = weight(this.w2, this.e2, posN, posNE, posNW), w3 = weight(this.w3, this.e3, posN, posNE, posNW);
         N *= 8;
         W *= 8;
         NE *= 8;
@@ -83,35 +97,44 @@ class WeightedPredictor {
         if (Math.abs(teNW) > Math.abs(p)) p = teNW;
         if (Math.abs(teNE) > Math.abs(p)) p = teNE;
         this.maxError = p;
-        const sumWN = teN + teW, pr = this.prediction;
-        pr[0] = W + NE - N;
-        pr[1] = N - Math.floor((sumWN + teNE) * h.p1 / 32);
-        pr[2] = W - Math.floor((sumWN + teNW) * h.p2 / 32);
-        pr[3] = N - Math.floor((teNW * h.p3a + teN * h.p3b + teNE * h.p3c + (NN - N) * h.p3d + (NW - W) * h.p3e) / 32);
+        const sumWN = teN + teW;
+        const p0 = this.p0 = W + NE - N;
+        const p1 = this.p1 = N - Math.floor((sumWN + teNE) * h.p1 / 32);
+        const p2 = this.p2 = W - Math.floor((sumWN + teNW) * h.p2 / 32);
+        const p3 = this.p3 = N - Math.floor((teNW * h.p3a + teN * h.p3b + teNE * h.p3c + (NN - N) * h.p3d + (NW - W) * h.p3e) / 32);
         // the weighted average
-        let weightSum = weights[0] + weights[1] + weights[2] + weights[3];
-        const logWeight = floorLog2(weightSum) - 4;
-        weightSum = 0;
-        for (let i = 0; i < 4; i++) weightSum += weights[i] >>= logWeight;
-        let sum = (weightSum >> 1) - 1;
-        for (let i = 0; i < 4; i++) sum += pr[i] * weights[i];
+        const logWeight = floorLog2(w0 + w1 + w2 + w3) - 4;
+        w0 >>= logWeight;
+        w1 >>= logWeight;
+        w2 >>= logWeight;
+        w3 >>= logWeight;
+        const weightSum = w0 + w1 + w2 + w3;
+        const sum = (weightSum >> 1) - 1 + p0 * w0 + p1 * w1 + p2 * w2 + p3 * w3;
         let pred = mulShift24(sum, DIV_LOOKUP[weightSum - 1]);
         // (clamped unless the errors have the same sign)
-        if (((teN ^ teW) | (teN ^ teNW)) <= 0) pred = Math.min(Math.max(pred, Math.min(W, NE, N)), Math.max(W, NE, N));
+        if (((teN ^ teW) | (teN ^ teNW)) <= 0) {
+            const lo = W < NE ? (W < N ? W : N) : (NE < N ? NE : N), hi = W > NE ? (W > N ? W : N) : (NE > N ? NE : N);
+            pred = pred < lo ? lo : pred > hi ? hi : pred;
+        }
         this.pred = pred;
         return Math.floor((pred + 3) / 8);
     }
 
     update(value: number, x: number, y: number): void {
-        const xsize = this.xsize, cur = y & 1 ? 0 : xsize + 2, prev = y & 1 ? xsize + 2 : 0;
+        const xsize = this.xsize, cur = y & 1 ? 0 : xsize + 2, i = cur + x, j = (y & 1 ? xsize + 2 : 0) + x + 1;
         const val = value * 8;
-        this.error[cur + x] = this.pred - val;
-        for (let i = 0; i < 4; i++) {
-            const e = Math.floor((Math.abs(this.prediction[i] - val) + 3) / 8) >>> 0;
-            const errors = this.predErrors[i];
-            errors[cur + x] = e;
-            errors[prev + x + 1] = (errors[prev + x + 1] + e) >>> 0;
-        }
+        this.error[i] = this.pred - val;
+        const e0 = Math.floor((Math.abs(this.p0 - val) + 3) / 8) >>> 0, e1 = Math.floor((Math.abs(this.p1 - val) + 3) / 8) >>> 0;
+        const e2 = Math.floor((Math.abs(this.p2 - val) + 3) / 8) >>> 0, e3 = Math.floor((Math.abs(this.p3 - val) + 3) / 8) >>> 0;
+        const E0 = this.e0, E1 = this.e1, E2 = this.e2, E3 = this.e3;
+        E0[i] = e0;
+        E0[j] = (E0[j] + e0) >>> 0;
+        E1[i] = e1;
+        E1[j] = (E1[j] + e1) >>> 0;
+        E2[i] = e2;
+        E2[j] = (E2[j] + e2) >>> 0;
+        E3[i] = e3;
+        E3[j] = (E3[j] + e3) >>> 0;
     }
 }