@playcanvas/splat-transform 1.9.0 → 1.9.1

package/dist/index.mjs

@@ -742,6 +742,16 @@ class Progress {
         if (!quiet)
             impl.onProgress(this.currentNode);
     }
+    /**
+     * Cancel the current progress node, popping it from the stack without
+     * completing remaining steps. Use this before early exits (e.g. break)
+     * to keep the progress stack balanced.
+     */
+    cancel() {
+        if (!this.currentNode)
+            return;
+        this.currentNode = this.currentNode.parent;
+    }
     /**
      * Advance to the next step. Auto-increments the step counter.
      * Auto-ends when all steps are complete.
@@ -929,33 +939,74 @@ const sortMortonOrder = (dataTable, indices) => {
     generate(indices);
 };
 
+const nthElement = (arr, lo, hi, k, values) => {
+    while (lo < hi) {
+        const mid = (lo + hi) >> 1;
+        const va = values[arr[lo]], vb = values[arr[mid]], vc = values[arr[hi]];
+        let pivotIdx;
+        if ((vb - va) * (vc - vb) >= 0)
+            pivotIdx = mid;
+        else if ((va - vb) * (vc - va) >= 0)
+            pivotIdx = lo;
+        else
+            pivotIdx = hi;
+        const pivotVal = values[arr[pivotIdx]];
+        let tmp = arr[pivotIdx];
+        arr[pivotIdx] = arr[hi];
+        arr[hi] = tmp;
+        let store = lo;
+        for (let i = lo; i < hi; i++) {
+            if (values[arr[i]] < pivotVal) {
+                tmp = arr[i];
+                arr[i] = arr[store];
+                arr[store] = tmp;
+                store++;
+            }
+        }
+        tmp = arr[store];
+        arr[store] = arr[hi];
+        arr[hi] = tmp;
+        if (store === k)
+            return;
+        else if (store < k)
+            lo = store + 1;
+        else
+            hi = store - 1;
+    }
+};
 class KdTree {
     centroids;
     root;
+    colData;
     constructor(centroids) {
-        const build = (indices, depth) => {
-            const { centroids } = this;
-            const values = centroids.columns[depth % centroids.numColumns].data;
-            indices.sort((a, b) => values[a] - values[b]);
-            if (indices.length === 1) {
-                return {
-                    index: indices[0],
-                    count: 1
-                };
-            }
-            else if (indices.length === 2) {
+        const numCols = centroids.numColumns;
+        const colData = centroids.columns.map(c => c.data);
+        const indices = new Uint32Array(centroids.numRows);
+        for (let i = 0; i < indices.length; ++i) {
+            indices[i] = i;
+        }
+        const build = (lo, hi, depth) => {
+            const count = hi - lo + 1;
+            if (count === 1) {
+                return { index: indices[lo], count: 1 };
+            }
+            const values = colData[depth % numCols];
+            if (count === 2) {
+                if (values[indices[lo]] > values[indices[hi]]) {
+                    const tmp = indices[lo];
+                    indices[lo] = indices[hi];
+                    indices[hi] = tmp;
+                }
                 return {
-                    index: indices[0],
+                    index: indices[lo],
                     count: 2,
-                    right: {
-                        index: indices[1],
-                        count: 1
-                    }
+                    right: { index: indices[hi], count: 1 }
                 };
             }
-            const mid = indices.length >> 1;
-            const left = build(indices.subarray(0, mid), depth + 1);
-            const right = build(indices.subarray(mid + 1), depth + 1);
+            const mid = lo + (count >> 1);
+            nthElement(indices, lo, hi, mid, values);
+            const left = build(lo, mid - 1, depth + 1);
+            const right = build(mid + 1, hi, depth + 1);
             return {
                 index: indices[mid],
                 count: 1 + left.count + right.count,
@@ -963,48 +1014,39 @@ class KdTree {
                 right
             };
         };
-        const indices = new Uint32Array(centroids.numRows);
-        for (let i = 0; i < indices.length; ++i) {
-            indices[i] = i;
-        }
         this.centroids = centroids;
-        this.root = build(indices, 0);
+        this.colData = colData;
+        this.root = build(0, indices.length - 1, 0);
     }
     findNearest(point, filterFunc) {
-        const { centroids } = this;
-        const { numColumns } = centroids;
-        const calcDistance = (index) => {
-            let l = 0;
-            for (let i = 0; i < numColumns; ++i) {
-                const v = centroids.columns[i].data[index] - point[i];
-                l += v * v;
-            }
-            return l;
-        };
+        const colData = this.colData;
+        const numCols = colData.length;
         let mind = Infinity;
         let mini = -1;
         let cnt = 0;
-        const recurse = (node, depth) => {
-            const axis = depth % numColumns;
-            const distance = point[axis] - centroids.columns[axis].data[node.index];
+        const recurse = (node, axis) => {
+            const distance = point[axis] - colData[axis][node.index];
             const next = (distance > 0) ? node.right : node.left;
+            const nextAxis = axis + 1 < numCols ? axis + 1 : 0;
             cnt++;
             if (next) {
-                recurse(next, depth + 1);
+                recurse(next, nextAxis);
             }
-            // check index
             if (!filterFunc || filterFunc(node.index)) {
-                const thisd = calcDistance(node.index);
+                let thisd = 0;
+                for (let c = 0; c < numCols; c++) {
+                    const v = colData[c][node.index] - point[c];
+                    thisd += v * v;
+                }
                 if (thisd < mind) {
                     mind = thisd;
                     mini = node.index;
                 }
             }
-            // check the other side
             if (distance * distance < mind) {
                 const other = next === node.right ? node.left : node.right;
                 if (other) {
-                    recurse(other, depth + 1);
+                    recurse(other, nextAxis);
                 }
             }
         };
@@ -1016,16 +1058,8 @@ class KdTree {
             return { indices: new Int32Array(0), distances: new Float32Array(0) };
         }
         k = Math.min(k, this.centroids.numRows);
-        const { centroids } = this;
-        const { numColumns } = centroids;
-        const calcDistance = (index) => {
-            let l = 0;
-            for (let i = 0; i < numColumns; ++i) {
-                const v = centroids.columns[i].data[index] - point[i];
-                l += v * v;
-            }
-            return l;
-        };
+        const colData = this.colData;
+        const numCols = colData.length;
         // Bounded max-heap: stores (distance, index) pairs sorted so the
         // farthest element is at position 0, enabling O(1) pruning bound.
         const heapDist = new Float32Array(k).fill(Infinity);
@@ -1033,14 +1067,12 @@ class KdTree {
         let heapSize = 0;
         const heapPush = (dist, idx) => {
             if (heapSize < k) {
-                // Heap not full yet -- insert via sift-up
                 let pos = heapSize++;
                 heapDist[pos] = dist;
                 heapIdx[pos] = idx;
                 while (pos > 0) {
                     const parent = (pos - 1) >> 1;
                     if (heapDist[parent] < heapDist[pos]) {
-                        // swap
                         const td = heapDist[parent];
                         heapDist[parent] = heapDist[pos];
                         heapDist[pos] = td;
@@ -1055,7 +1087,6 @@ class KdTree {
                 }
             }
             else if (dist < heapDist[0]) {
-                // Replace root (farthest) and sift-down
                 heapDist[0] = dist;
                 heapIdx[0] = idx;
                 let pos = 0;
@@ -1079,22 +1110,26 @@ class KdTree {
                 }
             }
         };
-        const recurse = (node, depth) => {
-            const axis = depth % numColumns;
-            const distance = point[axis] - centroids.columns[axis].data[node.index];
+        const recurse = (node, axis) => {
+            const distance = point[axis] - colData[axis][node.index];
             const next = (distance > 0) ? node.right : node.left;
+            const nextAxis = axis + 1 < numCols ? axis + 1 : 0;
             if (next) {
-                recurse(next, depth + 1);
+                recurse(next, nextAxis);
             }
             if (!filterFunc || filterFunc(node.index)) {
-                const thisd = calcDistance(node.index);
+                let thisd = 0;
+                for (let c = 0; c < numCols; c++) {
+                    const v = colData[c][node.index] - point[c];
+                    thisd += v * v;
+                }
                 heapPush(thisd, node.index);
             }
             const bound = heapSize < k ? Infinity : heapDist[0];
             if (distance * distance < bound) {
                 const other = next === node.right ? node.left : node.right;
                 if (other) {
-                    recurse(other, depth + 1);
+                    recurse(other, nextAxis);
                 }
             }
         };
@@ -1129,6 +1164,54 @@ const OPACITY_PRUNE_THRESHOLD = 0.1;
 const KNN_K = 16;
 const MC_SAMPLES = 1;
 const EPS_COV = 1e-8;
+const PROGRESS_TICKS = 100;
+// Radix sort edge indices by their Float32 costs.
+// Converts floats to sortable uint32 keys (preserving order), then does
+// 4-pass LSD radix sort with 8-bit radix. Returns the number of valid
+// (finite-cost) edges written to `out`.
+const radixSortIndicesByFloat = (out, count, keys) => {
+    const keyBits = new Uint32Array(keys.buffer, keys.byteOffset, keys.length);
+    const sortKeys = new Uint32Array(count);
+    let validCount = 0;
+    for (let i = 0; i < count; i++) {
+        const bits = keyBits[i];
+        if ((bits & 0x7F800000) === 0x7F800000)
+            continue;
+        sortKeys[validCount] = (bits & 0x80000000) ? ~bits >>> 0 : (bits | 0x80000000) >>> 0;
+        out[validCount] = i;
+        validCount++;
+    }
+    if (validCount <= 1)
+        return validCount;
+    const n = validCount;
+    const temp = new Uint32Array(n);
+    const tempKeys = new Uint32Array(n);
+    const counts = new Uint32Array(256);
+    for (let pass = 0; pass < 4; pass++) {
+        const shift = pass << 3;
+        const srcIdx = (pass & 1) ? temp : out;
+        const dstIdx = (pass & 1) ? out : temp;
+        const srcK = (pass & 1) ? tempKeys : sortKeys;
+        const dstK = (pass & 1) ? sortKeys : tempKeys;
+        counts.fill(0);
+        for (let i = 0; i < n; i++) {
+            counts[(srcK[i] >>> shift) & 0xFF]++;
+        }
+        let sum = 0;
+        for (let b = 0; b < 256; b++) {
+            const c = counts[b];
+            counts[b] = sum;
+            sum += c;
+        }
+        for (let i = 0; i < n; i++) {
+            const bucket = (srcK[i] >>> shift) & 0xFF;
+            const pos = counts[bucket]++;
+            dstIdx[pos] = srcIdx[i];
+            dstK[pos] = srcK[i];
+        }
+    }
+    return validCount;
+};
 // ---------- sigmoid / logit ----------
 const sigmoid$1 = (x) => 1 / (1 + Math.exp(-x));
 const logit = (p) => {
@@ -1616,7 +1699,6 @@ const simplifyGaussians = (dataTable, targetCount) => {
     let current;
     if (keptIndices.length < N && keptIndices.length > targetCount) {
         current = dataTable.permuteRows(keptIndices);
-        logger.debug(`opacity pruning: ${N} -> ${current.numRows} splats (threshold=${pruneThreshold.toFixed(4)})`);
     }
     else {
         current = dataTable;
@@ -1627,7 +1709,8 @@ const simplifyGaussians = (dataTable, targetCount) => {
     while (current.numRows > targetCount) {
         const n = current.numRows;
         const kEff = Math.min(Math.max(1, KNN_K), Math.max(1, n - 1));
-        logger.debug(`merging iteration: ${n} -> ${targetCount} splats`);
+        logger.progress.begin(5);
+        logger.progress.step('Building KD-tree');
         const cx = current.getColumnByName('x').data;
         const cy = current.getColumnByName('y').data;
         const cz = current.getColumnByName('z').data;
@@ -1640,16 +1723,21 @@ const simplifyGaussians = (dataTable, targetCount) => {
         const cr2 = current.getColumnByName('rot_2').data;
         const cr3 = current.getColumnByName('rot_3').data;
         const cache = buildPerSplatCache(n, cx, cy, cz, cop, cs0, cs1, cs2, cr0, cr1, cr2, cr3);
-        // Build KNN graph
         const posTable = new DataTable([
             new Column('x', cx instanceof Float32Array ? cx : new Float32Array(cx)),
             new Column('y', cy instanceof Float32Array ? cy : new Float32Array(cy)),
             new Column('z', cz instanceof Float32Array ? cz : new Float32Array(cz))
         ]);
         const kdTree = new KdTree(posTable);
-        const edgeSet = new Set();
-        const edges = [];
+        logger.progress.step('Finding nearest neighbors');
+        let edgeCapacity = Math.ceil(n * kEff / 2);
+        let edgeU = new Uint32Array(edgeCapacity);
+        let edgeV = new Uint32Array(edgeCapacity);
+        let edgeCount = 0;
         const queryPoint = new Float32Array(3);
+        const knnInterval = Math.max(1, Math.ceil(n / PROGRESS_TICKS));
+        const knnTicks = Math.ceil(n / knnInterval);
+        logger.progress.begin(knnTicks);
         for (let i = 0; i < n; i++) {
             queryPoint[0] = cx[i];
             queryPoint[1] = cy[i];
@@ -1657,46 +1745,58 @@ const simplifyGaussians = (dataTable, targetCount) => {
             const knn = kdTree.findKNearest(queryPoint, kEff + 1);
             for (let ki = 0; ki < knn.indices.length; ki++) {
                 const j = knn.indices[ki];
-                if (j === i || j < 0)
+                if (j <= i)
                     continue;
-                const u = Math.min(i, j);
-                const v = Math.max(i, j);
-                const key = `${u},${v}`;
-                if (!edgeSet.has(key)) {
-                    edgeSet.add(key);
-                    edges.push([u, v]);
+                if (edgeCount === edgeCapacity) {
+                    edgeCapacity *= 2;
+                    const newU = new Uint32Array(edgeCapacity);
+                    const newV = new Uint32Array(edgeCapacity);
+                    newU.set(edgeU);
+                    newV.set(edgeV);
+                    edgeU = newU;
+                    edgeV = newV;
                 }
+                edgeU[edgeCount] = i;
+                edgeV[edgeCount] = j;
+                edgeCount++;
             }
+            if ((i + 1) % knnInterval === 0)
+                logger.progress.step();
         }
-        if (edges.length === 0)
+        if (n % knnInterval !== 0)
+            logger.progress.step();
+        if (edgeCount === 0) {
+            logger.progress.cancel();
             break;
-        // Compute edge costs
+        }
+        logger.progress.step('Computing edge costs');
         const appData = [];
         for (let ai = 0; ai < allAppearanceCols.length; ai++) {
             const col = current.getColumnByName(allAppearanceCols[ai]);
             if (col)
                 appData.push(col.data);
         }
-        const costs = new Float32Array(edges.length);
-        for (let e = 0; e < edges.length; e++) {
-            costs[e] = computeEdgeCost(edges[e][0], edges[e][1], cx, cy, cz, cache, Z, appData, appData.length);
-        }
-        // Greedy disjoint pair selection
-        const valid = [];
-        for (let i = 0; i < edges.length; i++) {
-            if (Number.isFinite(costs[i]))
-                valid.push(i);
-        }
-        valid.sort((a, b) => {
-            const d = costs[a] - costs[b];
-            return d !== 0 ? d : a - b;
-        });
         const mergesNeeded = n - targetCount;
+        const costs = new Float32Array(edgeCount);
+        const costInterval = Math.max(1, Math.ceil(edgeCount / PROGRESS_TICKS));
+        const costTicks = Math.ceil(edgeCount / costInterval);
+        logger.progress.begin(costTicks);
+        for (let e = 0; e < edgeCount; e++) {
+            costs[e] = computeEdgeCost(edgeU[e], edgeV[e], cx, cy, cz, cache, Z, appData, appData.length);
+            if ((e + 1) % costInterval === 0)
+                logger.progress.step();
+        }
+        if (edgeCount % costInterval !== 0)
+            logger.progress.step();
+        logger.progress.step('Merging splats');
+        // Sort and greedy disjoint pair selection
+        const sorted = new Uint32Array(edgeCount);
+        const validCount = radixSortIndicesByFloat(sorted, edgeCount, costs);
         const used = new Uint8Array(n);
         const pairs = [];
-        for (let t = 0; t < valid.length; t++) {
-            const e = valid[t];
-            const u = edges[e][0], v = edges[e][1];
+        for (let t = 0; t < validCount; t++) {
+            const e = sorted[t];
+            const u = edgeU[e], v = edgeV[e];
             if (used[u] || used[v])
                 continue;
             used[u] = 1;
@@ -1705,9 +1805,10 @@ const simplifyGaussians = (dataTable, targetCount) => {
             if (pairs.length >= mergesNeeded)
                 break;
         }
-        if (pairs.length === 0)
+        if (pairs.length === 0) {
+            logger.progress.cancel();
             break;
-        logger.debug(`selected ${pairs.length} merge pairs from ${edges.length} edges`);
+        }
         // Mark which indices are consumed by merging
         const usedSet = new Uint8Array(n);
         for (let p = 0; p < pairs.length; p++) {
@@ -1735,7 +1836,7 @@ const simplifyGaussians = (dataTable, targetCount) => {
                 newTable.columns[c].data[dst] = cols[c].data[src];
             }
         }
-        // Merge pairs -- cache column refs and handled set once
+        // Merge pairs
         const mergeOut = {
             mu: new Float64Array(3),
             sc: new Float64Array(3),
@@ -1763,6 +1864,9 @@ const simplifyGaussians = (dataTable, targetCount) => {
             .filter(col => !handledCols.has(col.name))
             .map(col => ({ src: col, dst: newTable.getColumnByName(col.name) }))
             .filter(pair => pair.dst);
+        const mergeInterval = Math.max(1, Math.ceil(pairs.length / PROGRESS_TICKS));
+        const mergeTicks = Math.ceil(pairs.length / mergeInterval);
+        logger.progress.begin(mergeTicks);
         for (let p = 0; p < pairs.length; p++, dst++) {
             const pi = pairs[p][0], pj = pairs[p][1];
             momentMatch(pi, pj, cx, cy, cz, cop, cs0, cs1, cs2, cr0, cr1, cr2, cr3, mergeOut, appData, appData.length);
@@ -1785,7 +1889,12 @@ const simplifyGaussians = (dataTable, targetCount) => {
             for (let u = 0; u < unhandledColPairs.length; u++) {
                 unhandledColPairs[u].dst.data[dst] = unhandledColPairs[u].src.data[dominant];
             }
+            if ((p + 1) % mergeInterval === 0)
+                logger.progress.step();
         }
+        if (pairs.length % mergeInterval !== 0)
+            logger.progress.step();
+        logger.progress.step('Finalizing');
         current = newTable;
     }
     return current;
@@ -5613,7 +5722,7 @@ class CompressedChunk {
     }
 }
 
-var version = "1.9.0";
+var version = "1.9.1";
 
 const generatedByString = `Generated by splat-transform ${version}`;
 const chunkProps = [