Improve tracklet visualization with highlights (#42)

* Initial loading of track lineage

* Quiet noisy logs

* Properly compute track lineage

* Update convert_tracks, in particular the comments on time/size

* Render a section of the track in a different color; colors follow points

rebase to get sub-track coloring, colors follow points

* Lint

* Set line width in world coords

* Remove updating camera target on array change

not relevant to this branch, and wasn't working well anyway

* Update points geometry allocation to use maxPointsPerTimepoint

* Refactor track highlight to use a separate Line2

* Add track length slider

* Fix bug where highlights would disappear if tracklets were shorter than the highlight length

* API tidying, adding some comments

* Refactor so Track is an Object3D

* Refactor to clean up and slightly improve performance

* Apply suggestions from code review

Co-authored-by: Andy Sweet <[email protected]>

* Add TrackLine with custom shader, based on Line2

* Use TrackLine instead of combined Line2 visuals

* Add some comments, remove unused features from TrackLine

* Update src/lib/three/TrackGeometry.ts

Co-authored-by: Andy Sweet <[email protected]>

* Update src/lib/three/TrackGeometry.ts

Co-authored-by: Andy Sweet <[email protected]>

* Apply suggestions from code review

Co-authored-by: Andy Sweet <[email protected]>

* Simplify track fetching/adding based on code review

Co-authored-by: Andy Sweet <[email protected]>

* Move Track into its own file

* Add more comments to modified Line2 code

* Don't store LUT texture on Track

* Update src/Track.ts

Co-authored-by: Andy Sweet <[email protected]>

* Refactor to combine Track + TrackLine -> Track

---------

Co-authored-by: Andy Sweet <[email protected]>

src/PointCanvas.ts

@@ -20,10 +20,10 @@ import { EffectComposer } from "three/addons/postprocessing/EffectComposer.js";
 import { RenderPass } from "three/addons/postprocessing/RenderPass.js";
 import { OutputPass } from "three/addons/postprocessing/OutputPass.js";
 import { UnrealBloomPass } from "three/addons/postprocessing/UnrealBloomPass.js";
-import { Lut } from "three/addons/math/Lut.js";
-import { Line2, LineGeometry, LineMaterial } from "three/examples/jsm/Addons.js";
 
-type Tracks = Map<number, Line2>;
+import { Track } from "./lib/three/Track";
+
+type Tracks = Map<number, Track>;
 
 export class PointCanvas {
     scene: Scene;
@@ -34,6 +34,10 @@ export class PointCanvas {
     controls: OrbitControls;
     bloomPass: UnrealBloomPass;
     tracks: Tracks = new Map();
+    // this is used to initialize the points geometry, and kept to initialize the
+    // tracks but could be pulled from the points geometry when adding tracks
+    // private here to consolidate external access via `TrackManager` instead
+    private maxPointsPerTimepoint = 0;
 
     constructor(width: number, height: number) {
         this.scene = new Scene();
@@ -125,15 +129,19 @@ export class PointCanvas {
         this.composer.setSize(width, height);
     }
 
-    initPointsGeometry(numPoints: number) {
+    initPointsGeometry(maxPointsPerTimepoint: number) {
+        this.maxPointsPerTimepoint = maxPointsPerTimepoint;
         const geometry = this.points.geometry;
-        if (!geometry.hasAttribute("position") || geometry.getAttribute("position").count !== numPoints) {
-            geometry.setAttribute("position", new Float32BufferAttribute(new Float32Array(3 * numPoints), 3));
+        if (!geometry.hasAttribute("position") || geometry.getAttribute("position").count !== maxPointsPerTimepoint) {
+            geometry.setAttribute(
+                "position",
+                new Float32BufferAttribute(new Float32Array(3 * maxPointsPerTimepoint), 3),
+            );
             // prevent drawing uninitialized points at the origin
             geometry.setDrawRange(0, 0);
         }
-        if (!geometry.hasAttribute("color") || geometry.getAttribute("color").count !== numPoints) {
-            geometry.setAttribute("color", new Float32BufferAttribute(new Float32Array(3 * numPoints), 3));
+        if (!geometry.hasAttribute("color") || geometry.getAttribute("color").count !== maxPointsPerTimepoint) {
+            geometry.setAttribute("color", new Float32BufferAttribute(new Float32Array(3 * maxPointsPerTimepoint), 3));
         }
         // Initialize all the colors immediately.
         this.resetPointColors();
@@ -151,45 +159,29 @@ export class PointCanvas {
         this.points.geometry.computeBoundingSphere();
     }
 
-    addTrack(trackID: number, positions: Float32Array) {
+    addTrack(trackID: number, positions: Float32Array, ids: Int32Array): Track | null {
         if (this.tracks.has(trackID)) {
+            // this is a warning because it should alert us to duplicate fetching
             console.warn("Track with ID %d already exists", trackID);
-            return;
-        }
-        const pos = [];
-        const colors = [];
-        const lut = new Lut("rainbow", 256);
-
-        for (let i = 0; i < positions.length; i += 3) {
-            pos.push(positions[i], positions[i + 1], positions[i + 2]);
-            const color = lut.getColor(i / positions.length);
-            colors.push(color.r, color.g, color.b);
+            return null;
         }
-
-        const geometry = new LineGeometry();
-        geometry.setPositions(positions);
-        geometry.setColors(colors);
-        const material = new LineMaterial({
-            linewidth: 0.003,
-            vertexColors: true,
-        });
-        const track = new Line2(geometry, material);
-        this.scene.add(track);
+        const track = Track.new(positions, ids, this.maxPointsPerTimepoint);
         this.tracks.set(trackID, track);
+        this.scene.add(track);
+        return track;
+    }
+
+    updateAllTrackHighlights(minTime: number, maxTime: number) {
+        for (const track of this.tracks.values()) {
+            track.updateHighlightLine(minTime, maxTime);
+        }
     }
 
     removeTrack(trackID: number) {
         const track = this.tracks.get(trackID);
         if (track) {
             this.scene.remove(track);
-            track.geometry.dispose();
-            if (Array.isArray(track.material)) {
-                for (const material of track.material) {
-                    material.dispose();
-                }
-            } else {
-                track.material.dispose();
-            }
+            track.dispose();
             this.tracks.delete(trackID);
         } else {
             console.warn("No track with ID %d to remove", trackID);

src/PointSelectionBox.ts

@@ -24,10 +24,7 @@ const _vectemp1 = new Vector3();
 const _vectemp2 = new Vector3();
 const _vectemp3 = new Vector3();
 
-interface PointsCollection {
-    // object_id : [point_index, point_index, ...]
-    [key: number]: number[];
-}
+type PointsCollection = Map<number, number[]>;
 
 class PointSelectionBox {
     camera: OrthographicCamera | PerspectiveCamera;
@@ -42,14 +39,14 @@ class PointSelectionBox {
         this.scene = scene;
         this.startPoint = new Vector3();
         this.endPoint = new Vector3();
-        this.collection = {};
+        this.collection = new Map();
         this.deep = deep;
     }
 
     select(startPoint?: Vector3, endPoint?: Vector3) {
         this.startPoint = startPoint ?? this.startPoint;
         this.endPoint = endPoint ?? this.endPoint;
-        this.collection = {};
+        this.collection = new Map();
 
         this.updateFrustum(this.startPoint, this.endPoint);
         this.searchChildInFrustum(_frustum, this.scene);
@@ -164,10 +161,11 @@ class PointSelectionBox {
             for (let i = start; i < end; i++) {
                 _vec3.set(positionAttribute.getX(i), positionAttribute.getY(i), positionAttribute.getZ(i));
                 if (frustum.containsPoint(_vec3)) {
-                    if (!this.collection[object.id]) {
-                        this.collection[object.id] = [i];
+                    const objectCollection = this.collection.get(object.id);
+                    if (!objectCollection) {
+                        this.collection.set(object.id, [i]);
                     } else {
-                        this.collection[object.id].push(i);
+                        objectCollection.push(i);
                     }
                 }
             }

src/TrackManager.ts

@@ -67,6 +67,7 @@ export class TrackManager {
     pointsToTracks: SparseZarrArray;
     tracksToPoints: SparseZarrArray;
     tracksToTracks: SparseZarrArray;
+    maxPointsPerTimepoint: number;
 
     constructor(
         store: string,
@@ -80,6 +81,7 @@ export class TrackManager {
         this.pointsToTracks = pointsToTracks;
         this.tracksToPoints = tracksToPoints;
         this.tracksToTracks = tracksToTracks;
+        this.maxPointsPerTimepoint = points.shape[1] / 3;
     }
 
     async fetchPointsAtTime(timeIndex: number): Promise<Float32Array> {
@@ -106,15 +108,23 @@ export class TrackManager {
         return trackIDs.data;
     }
 
-    async fetchPointsForTrack(trackID: number): Promise<Float32Array> {
+    async fetchPointsForTrack(trackID: number): Promise<[Float32Array, Int32Array]> {
         const rowStartEnd = await this.tracksToPoints.getIndPtr(slice(trackID, trackID + 2));
-        const points = await this.tracksToPoints.data.get([slice(rowStartEnd[0], rowStartEnd[1]), slice(null)]);
+        const points = (await this.tracksToPoints.data.get([slice(rowStartEnd[0], rowStartEnd[1]), slice(null)])).data;
+        // TODO: can bake this into the data array
+        const pointIDs = (await this.tracksToPoints.indices.get([slice(rowStartEnd[0], rowStartEnd[1])])).data;
+
+        if (points.length !== pointIDs.length) {
+            console.error("points and pointIDs are different lengths: %d, %d", points.length, pointIDs.length);
+        }
+
         // flatten the resulting n x 3 array in to a 1D [xyzxyzxyz...] array
-        const flatPoints = new Float32Array(points.data.length * 3);
-        for (let i = 0; i < points.data.length; i++) {
-            flatPoints.set(points.data[i], i * 3);
+        const flatPoints = new Float32Array(points.length * 3);
+        for (let i = 0; i < points.length; i++) {
+            flatPoints.set(points[i], i * 3);
         }
-        return flatPoints;
+
+        return [flatPoints, pointIDs];
     }
 
     async fetchLineageForTrack(trackID: number): Promise<Int32Array> {

src/hooks/useSelectionBox.ts

@@ -100,15 +100,5 @@ export default function useSelectionBox(canvas: PointCanvas | undefined) {
         }
     }, [selecting]);
 
-    useEffect(() => {
-        if (selectedPoints && canvas && canvas.points.id in selectedPoints) {
-            console.debug("highlighting points:", selectedPoints[canvas.points.id]);
-            canvas.highlightPoints(selectedPoints[canvas.points.id]);
-        } else if (canvas) {
-            console.log("resetting point colors");
-            canvas.resetPointColors();
-        }
-    }, [selectedPoints]);
-
     return { selectedPoints, setSelectedPoints, selecting, setSelecting };
 }

src/lib/three/Track.ts

@@ -0,0 +1,59 @@
+/**
+ * This class gives a specific type for the Track object, but it's mostly a
+ * Mesh with a custom initializer.
+ *
+ * see:
+ *  https://github.com/mrdoob/three.js/blob/5ed5417d63e4eeba5087437cc27ab1e3d0813aea/examples/jsm/lines/Line2.js
+ *  https://github.com/mrdoob/three.js/blob/5ed5417d63e4eeba5087437cc27ab1e3d0813aea/examples/jsm/lines/LineSegments2.js
+ */
+import { Mesh } from "three";
+
+import { TrackGeometry } from "./TrackGeometry.js";
+import { TrackMaterial } from "./TrackMaterial.js";
+
+export class Track extends Mesh {
+    isTrack = true;
+    type = "Track";
+    declare geometry: TrackGeometry;
+    declare material: TrackMaterial;
+
+    static new(positions: Float32Array, pointIDs: Int32Array, maxPointsPerTimepoint: number) {
+        const geometry = new TrackGeometry();
+        const material = new TrackMaterial({
+            vertexColors: true,
+            trackwidth: 0.3,
+            highlightwidth: 2.0,
+            showtrack: true,
+            transparent: true,
+            opacity: 0.5,
+        });
+        const track = new Track(geometry, material);
+
+        const time: number[] = [];
+        const pos = Array.from(positions);
+        const colors: number[] = [];
+        const n = pos.length / 3;
+        for (const [i, id] of pointIDs.entries()) {
+            const t = Math.floor(id / maxPointsPerTimepoint);
+            time.push(t);
+            // TODO: use a LUT for the main track, too
+            colors.push(((0.9 * (n - i)) / n) ** 3, ((0.9 * (n - i)) / n) ** 3, (0.9 * (n - i)) / n);
+        }
+        track.geometry.setPositions(pos);
+        track.geometry.setColors(colors);
+        track.geometry.setTime(time);
+        track.geometry.computeBoundingSphere();
+        return track;
+    }
+
+    updateHighlightLine(minTime: number, maxTime: number) {
+        this.material.minTime = minTime;
+        this.material.maxTime = maxTime;
+        this.material.needsUpdate = true;
+    }
+
+    dispose() {
+        this.geometry.dispose();
+        this.material.dispose();
+    }
+}

src/lib/three/TrackGeometry.ts

@@ -0,0 +1,79 @@
+/**
+ * This class maintains the geometry of the track, adding instanced attributes for time.
+ * see:
+ *  https://github.com/mrdoob/three.js/blob/dev/examples/jsm/lines/LineGeometry.js
+ */
+import { InstancedInterleavedBuffer, InterleavedBufferAttribute } from "three";
+import { LineSegmentsGeometry } from "three/examples/jsm/Addons.js";
+
+class TrackGeometry extends LineSegmentsGeometry {
+    isTrackGeometry = true;
+    type = "TrackGeometry";
+
+    setPositions(array: number[] | Float32Array) {
+        // converts [ x1, y1, z1, x2, y2, z2, ... ] to pairs format
+
+        const length = array.length - 3;
+        const points = new Float32Array(2 * length);
+
+        for (let i = 0; i < length; i += 3) {
+            points[2 * i] = array[i];
+            points[2 * i + 1] = array[i + 1];
+            points[2 * i + 2] = array[i + 2];
+
+            points[2 * i + 3] = array[i + 3];
+            points[2 * i + 4] = array[i + 4];
+            points[2 * i + 5] = array[i + 5];
+        }
+
+        super.setPositions(points);
+
+        return this;
+    }
+
+    setColors(array: number[] | Float32Array) {
+        // converts [ r1, g1, b1, r2, g2, b2, ... ] to pairs format
+
+        const length = array.length - 3;
+        const colors = new Float32Array(2 * length);
+
+        for (let i = 0; i < length; i += 3) {
+            colors[2 * i] = array[i];
+            colors[2 * i + 1] = array[i + 1];
+            colors[2 * i + 2] = array[i + 2];
+
+            colors[2 * i + 3] = array[i + 3];
+            colors[2 * i + 4] = array[i + 4];
+            colors[2 * i + 5] = array[i + 5];
+        }
+
+        super.setColors(colors);
+
+        return this;
+    }
+
+    setTime(array: number[]) {
+        // TRACK SPECIFIC CODE ADDED
+        // converts [ t1, t2, ... ] to pairs format
+        // this ecodes the timepoint of each point in the track
+        // this can be used in shaders to highlight specific segments in the
+        // track based on time
+
+        // float32 should be sufficient given we're expecting ~1000 timepoints
+        const length = array.length - 1;
+        const times = new Float32Array(2 * length);
+
+        for (let i = 0; i < length; i++) {
+            times[2 * i] = array[i];
+            times[2 * i + 1] = array[i + 1];
+        }
+
+        const time = new InstancedInterleavedBuffer(times, 2, 1);
+        this.setAttribute("instanceTimeStart", new InterleavedBufferAttribute(time, 1, 0));
+        this.setAttribute("instanceTimeEnd", new InterleavedBufferAttribute(time, 1, 1));
+
+        return this;
+    }
+}
+
+export { TrackGeometry };