Expose spherical harmonics per splat

docs/docs/ext-splats.md

@@ -81,7 +81,7 @@ Each `ExtSplat` occupies 8 × `uint32` total = 32 bytes.
 import { utils } from "@sparkjsdev/spark";
 
 utils.encodeExtSplat(extSplats.extArrays, index, x, y, z, sx, sy, sz, qx, qy, qz, qw, opacity, r, g, b);
-const { center, scales, quaternion, color, opacity } = utils.decodeExtSplat(extSplats.extArrays, index);
+const { center, scales, quaternion, color, opacity, sphericalHarmonics } = utils.decodeExtSplat(extSplats.extArrays, index, extSplats.extra);
 ```
 
 You can also use instance helpers: `setSplat`, `pushSplat`, `getSplat`, and `forEachSplat`.

docs/docs/packed-splats.md

@@ -62,7 +62,7 @@ Utility functions are provided in Javascript to pack/unpack these encodings:
 ```javascript
 
 // Set via packedSplats interface
-packedSplats.setSplat(index, center, scales, quaternion, opacity, color);
+packedSplats.setSplat(index, center, scales, quaternion, opacity, color, sphericalHarmonics);
 
 // Set underlying Uint32 array directly
 import { utils } from "@sparkjsdev/spark";
@@ -72,15 +72,15 @@ utils.setPackedSplat(packedSplats.packedArray, index, x, y, z, scaleX, scaleY, .
 utils.setPackedSplatQuat(packedSplats.packedArray, index, quatX, quatY, quatZ, quatW);
 
 // Unpack all splat components from the Uint32 array
-const { center, scales, quaternion, color, opacity } = utils.unpackSplat(packedSplats.packedArray, index);
+const { center, scales, quaternion, color, opacity, sphericalHarmonics } = utils.unpackSplat(packedSplats.packedArray, index, packedSplats.splatEncoding, packedSplats.extra);
 
 // Unpack all splats with callback
-packedSplats.forEachSplat((index, center, scales, quaternion, opacity, color) => {
+packedSplats.forEachSplat((index, center, scales, quaternion, opacity, color, sphericalHarmonics) => {
     // Use unpacked splat data. Changing the inputs directly has no effect.
     // Update just the scales component
     utils.setPackedSplatScales(packedSplat.packedArray, index, 0.005, 0.01, 0.015);
     // Update the entire splat
-    packedSplat.setSplat(index, center, scales, quaternion, opacity, color);
+    packedSplat.setSplat(index, center, scales, quaternion, opacity, color, sphericalHarmonics);
 });
 ```
 
@@ -175,19 +175,19 @@ Ensures that `this.packedArray` can fit `numSplats` splats. If it's too small, r
 
 Typically you don't need to call this, because calling `this.setSplat(index, ...)` and `this.pushSplat(...)` will automatically call `ensureSplats()` so we have enough splats.
 
-### `getSplat(index): { center, scales, quaternion, opacity, color }`
+### `getSplat(index): { center, scales, quaternion, opacity, color, sphericalHarmonics }`
 
-Unpack the 16-byte splat data at `index` into the THREE.js components `center: THREE.Vector3`, `scales: THREE.Vector3`, `quaternion: THREE.Quaternion`, `opacity: number 0..1`, `color: THREE.Color 0..1`.
+Unpack the splat data at `index`, including decoded `sphericalHarmonics.sh1`, `.sh2`, and `.sh3` bands when present.
 
-### `setSplat(index, center, scales, quaternion, opacity, color)`
+### `setSplat(index, center, scales, quaternion, opacity, color, sphericalHarmonics?)`
 
 Set all PackedSplat components at `index` with the provided splat attributes (can be the same objects returned by `getSplat`). Ensures there is capacity for at least `index+1` splats.
 
-### `pushSplat(center, scales, quaternion, opacity, color)`
+### `pushSplat(center, scales, quaternion, opacity, color, sphericalHarmonics?)`
 
 Effectively calls `this.setSplat(this.numSplats++, center, ...)`, useful on construction where you just want to iterate and create a collection of splats.
 
-### `forEachSplat(callback: (index, center, scales, quaternion, opacity, color) => void)`
+### `forEachSplat(callback: (index, center, scales, quaternion, opacity, color, sphericalHarmonics) => void)`
 
 Iterate over splats index `0..=(this.numSplats-1)`, unpack each splat and invoke the callback function with the splat attributes.
 

docs/docs/splat-mesh.md

@@ -141,13 +141,13 @@ Additional properties on a `SplatMesh` instance:
 
 Call this when you are finished with the `SplatMesh` and want to free any buffers it holds (via `packedSplats`).
 
-## `pushSplat(center, scales, quaternion, opacity, color)`
+## `pushSplat(center, scales, quaternion, opacity, color, sphericalHarmonics?)`
 
 Creates a new splat with the provided parameters (all values in "float" space, i.e. 0-1 for opacity and color) and adds it to the end of the `packedSplats`, increasing `numSplats` by 1. If necessary, reallocates the buffer with an exponential doubling strategy to fit the new data, so it's fairly efficient to `pushSplat(...)` each splat you want to create in a loop.
 
-## `forEachSplat(callback: (index, center, scales, quaternion, opacity, color) => void)`
+## `forEachSplat(callback: (index, center, scales, quaternion, opacity, color, sphericalHarmonics) => void)`
 
-This method iterates over all splats in this instance's `packedSplats`, invoking the provided callback with `index: number` in `0..=(this.numSplats-1)`, `center: THREE.Vector3`, `scales: THREE.Vector3`, `quaternion: THREE.Quaternion`, `opacity: number` (0..1), and `color: THREE.Color` (rgb values in 0..1). Note that the objects passed in as `center` etc. are the same for every callback invocation: they are reused for efficiency. *Changing these values has no effect* as they are decoded/unpacked copies of the underlying data. To update the `packedSplats`, call `.packedSplats.setSplat(index, center, scales, quaternion, opacity, color)`.
+This method iterates over all splats in this instance's `packedSplats`, invoking the provided callback with `index: number` in `0..=(this.numSplats-1)`, `center: THREE.Vector3`, `scales: THREE.Vector3`, `quaternion: THREE.Quaternion`, `opacity: number` (0..1), `color: THREE.Color` (rgb values in 0..1), and decoded `sphericalHarmonics`. Note that the objects passed in as `center` etc. are the same for every callback invocation: these objects are reused for efficiency. *Changing these values has no effect* as they are decoded/unpacked copies of the underlying data. To update the `packedSplats`, call `.packedSplats.setSplat(index, center, scales, quaternion, opacity, color, sphericalHarmonics)`.
 
 
 ## `getBoundingBox(centers_only=true)`

package.json

@@ -38,7 +38,7 @@
     "site:deploy": "npm run site:build && node scripts/deploy-site.js",
     "site:serve": "node scripts/serve-site.js site",
     "start": "npm run dev",
-    "test": "node --no-warnings --loader ts-node/esm --test test/**/*.test.ts"
+    "test": "node --no-warnings --loader ts-node/esm --loader ./test/wasm-loader.mjs --test test/**/*.test.ts"
   },
   "repository": "sparkjs-dev/spark",
   "files": ["dist"],

src/ExtSplats.ts

@@ -1,5 +1,10 @@
 import * as THREE from "three";
 import type { RgbaArray } from "./RgbaArray";
+import type {
+  ForEachSplatCallback,
+  SplatSphericalHarmonics,
+  UnpackedSplat,
+} from "./SplatData";
 import { SplatLoader } from "./SplatLoader";
 import type { SplatSource } from "./SplatMesh";
 import { workerPool } from "./SplatWorker";
@@ -24,7 +29,14 @@ import {
   unindent,
   unindentLines,
 } from "./dyno";
-import { decodeExtSplat, encodeExtSplat, getTextureSize } from "./utils";
+import {
+  decodeExtSplat,
+  encodeExt3Rgb,
+  encodeExtRgb,
+  encodeExtSh1Rgb,
+  encodeExtSplat,
+  getTextureSize,
+} from "./utils";
 
 export type ExtSplatsOptions = {
   // URL to fetch a Gaussian splat file from (supports .ply, .splat, .ksplat,
@@ -472,17 +484,11 @@ export class ExtSplats implements SplatSource {
   // Unpack the 16-byte Gsplat data at index into the Three.js components
   // center: THREE.Vector3, scales: THREE.Vector3, quaternion: THREE.Quaternion,
   // opacity: number 0..1, color: THREE.Color 0..1.
-  getSplat(index: number): {
-    center: THREE.Vector3;
-    scales: THREE.Vector3;
-    quaternion: THREE.Quaternion;
-    opacity: number;
-    color: THREE.Color;
-  } {
+  getSplat(index: number): UnpackedSplat {
     if (index >= this.numSplats) {
       throw new Error("Invalid index");
     }
-    return decodeExtSplat(this.extArrays, index);
+    return decodeExtSplat(this.extArrays, index, this.extra);
   }
 
   // Set all ExtSplat components at index with the provided Gsplat attributes
@@ -495,6 +501,7 @@ export class ExtSplats implements SplatSource {
     quaternion: THREE.Quaternion,
     opacity: number,
     color: THREE.Color,
+    sphericalHarmonics?: SplatSphericalHarmonics,
   ) {
     const extArrays = this.ensureSplats(index + 1);
     encodeExtSplat(
@@ -515,9 +522,71 @@ export class ExtSplats implements SplatSource {
       color.g,
       color.b,
     );
+    this.setSplatSphericalHarmonics(index, sphericalHarmonics);
     this.numSplats = Math.max(this.numSplats, index + 1);
   }
 
+  private ensureSplatsSh(key: string, numSplats: number): Uint32Array {
+    const current = this.extra[key] as Uint32Array | undefined;
+    const currentSplats = current?.length ? current.length / 4 : 0;
+    if (current && currentSplats >= numSplats) {
+      return current;
+    }
+    const targetSplats = getTextureSize(
+      Math.max(numSplats, 2 * currentSplats),
+    ).maxSplats;
+    const result = new Uint32Array(targetSplats * 4);
+    if (current) {
+      result.set(current);
+    }
+    this.extra[key] = result;
+    return result;
+  }
+
+  private setSplatSphericalHarmonics(
+    index: number,
+    sphericalHarmonics?: SplatSphericalHarmonics,
+  ) {
+    if (!sphericalHarmonics) {
+      return;
+    }
+    if (sphericalHarmonics.sh1) {
+      encodeExtSh1Rgb(
+        this.ensureSplatsSh("sh1", index + 1),
+        index,
+        sphericalHarmonics.sh1,
+      );
+    }
+    if (sphericalHarmonics.sh2) {
+      const sh1 = this.ensureSplatsSh("sh1", index + 1);
+      const sh2 = this.ensureSplatsSh("sh2", index + 1);
+      const base = index * 4;
+      sh1[base + 3] = encodeExtRgb(
+        sphericalHarmonics.sh2[0],
+        sphericalHarmonics.sh2[1],
+        sphericalHarmonics.sh2[2],
+      );
+      for (let coefficient = 1; coefficient < 5; ++coefficient) {
+        const offset = coefficient * 3;
+        sh2[base + coefficient - 1] = encodeExtRgb(
+          sphericalHarmonics.sh2[offset],
+          sphericalHarmonics.sh2[offset + 1],
+          sphericalHarmonics.sh2[offset + 2],
+        );
+      }
+    }
+    if (sphericalHarmonics.sh3) {
+      this.ensureSplatsSh("sh1", index + 1);
+      this.ensureSplatsSh("sh2", index + 1);
+      encodeExt3Rgb(
+        this.ensureSplatsSh("sh3a", index + 1),
+        this.ensureSplatsSh("sh3b", index + 1),
+        index,
+        sphericalHarmonics.sh3,
+      );
+    }
+  }
+
   // Effectively calls this.setSplat(this.numSplats++, center, ...), useful on
   // construction where you just want to iterate and create a collection of Gsplats.
   pushSplat(
@@ -526,6 +595,7 @@ export class ExtSplats implements SplatSource {
     quaternion: THREE.Quaternion,
     opacity: number,
     color: THREE.Color,
+    sphericalHarmonics?: SplatSphericalHarmonics,
   ) {
     const extArrays = this.ensureSplats(this.numSplats + 1);
     encodeExtSplat(
@@ -546,33 +616,26 @@ export class ExtSplats implements SplatSource {
       color.g,
       color.b,
     );
+    this.setSplatSphericalHarmonics(this.numSplats, sphericalHarmonics);
     ++this.numSplats;
   }
 
   // Iterate over Gsplats index 0..=(this.numSplats-1), unpack each Gsplat
   // and invoke the callback function with the Gsplat attributes.
-  forEachSplat(
-    callback: (
-      index: number,
-      center: THREE.Vector3,
-      scales: THREE.Vector3,
-      quaternion: THREE.Quaternion,
-      opacity: number,
-      color: THREE.Color,
-    ) => void,
-  ) {
+  forEachSplat(callback: ForEachSplatCallback) {
     if (!this.numSplats) {
       return;
     }
     for (let i = 0; i < this.numSplats; ++i) {
-      const unpacked = decodeExtSplat(this.extArrays, i);
+      const unpacked = decodeExtSplat(this.extArrays, i, this.extra);
       callback(
         i,
         unpacked.center,
         unpacked.scales,
         unpacked.quaternion,
         unpacked.opacity,
         unpacked.color,
+        unpacked.sphericalHarmonics,
       );
     }
   }