[WIP] Different radius per point

CONFIG.ts

@@ -10,15 +10,17 @@ const config = {
     // When opening the viewer, or refreshing the page, the viewer will revert to the following default dataset
     data:{
         // Default dataset URL (must be publically accessible)
-        default_dataset: "https://sci-imaging-vis-public-demo-data.s3.us-west-2.amazonaws.com/points-web-viewer/sparse-zarr-v2/ZSNS001_tracks_bundle.zarr"
+        default_dataset: "https://public.czbiohub.org/royerlab/zoo/Ascidian/tracks_withSize_bundle.zarr/"
     },
 
     // Default settings for certain parameters
     settings:{
         // Maximum number of cells a user can select without getting a warning
         max_num_selected_cells: 100,
         // Choose colormap for the tracks, options: viridis-inferno, magma-inferno, inferno-inferno
-        colormap_tracks: "viridis-inferno"
+        colormap_tracks: "viridis-inferno",
+        // Point size (arbitrary units)
+        point_size: 30
     }
 }
 

src/components/App.tsx

@@ -174,6 +174,7 @@ export default function App() {
                         // adding the track *in* the dispatcher creates issues with duplicate fetching
                         // but we refresh so the selected/loaded count is updated
                         canvas.addTrack(relatedTrackId, pos, ids, trackData[index]);
+                        canvas.clearPointIndicesCache;
                         dispatchCanvas({ type: ActionType.REFRESH });
                     });
                 });

src/components/DataControls.tsx

@@ -94,13 +94,15 @@ export default function DataControls(props: DataControlsProps) {
                 />
             </Tooltip>
             <Tooltip title="Copy a shareable URL for this view to your clipboard">
-                <ButtonIcon
-                    icon="Share"
-                    sdsSize="large"
-                    sdsType="secondary"
-                    disabled={!props.trackManager}
-                    onClick={copyShareableUrlToClipBoard}
-                />
+                <span>
+                    <ButtonIcon
+                        icon="Share"
+                        sdsSize="large"
+                        sdsType="secondary"
+                        disabled={!props.trackManager}
+                        onClick={copyShareableUrlToClipBoard}
+                    />
+                </span>
             </Tooltip>
             <Snackbar
                 open={copyUrlSnackBarOpen}

src/hooks/usePointCanvas.ts

@@ -135,13 +135,16 @@ function reducer(canvas: PointCanvas, action: PointCanvasAction): PointCanvas {
         case ActionType.POINT_BRIGHTNESS:
             newCanvas.pointBrightness = action.brightness;
             newCanvas.resetPointColors();
+            newCanvas.updateSelectedPointIndices();
             break;
         case ActionType.POINTS_POSITIONS:
             newCanvas.setPointsPositions(action.positions);
             newCanvas.resetPointColors();
+            newCanvas.updateSelectedPointIndices();
             break;
         case ActionType.REMOVE_ALL_TRACKS:
             newCanvas.removeAllTracks();
+            newCanvas.clearPointIndicesCache();
             newCanvas.pointBrightness = 1.0;
             newCanvas.resetPointColors();
             break;
@@ -167,14 +170,13 @@ function reducer(canvas: PointCanvas, action: PointCanvasAction): PointCanvas {
         case ActionType.ADD_SELECTED_POINT_IDS: {
             newCanvas.pointBrightness = 0.8;
             newCanvas.resetPointColors();
-            // TODO: only highlight the indices if the canvas is at the same time
-            // point as when it was selected.
-            newCanvas.highlightPoints(action.selectedPointIndices);
+            // newCanvas.highlightPoints(action.selectedPointIndices);
             const newSelectedPointIds = new Set(canvas.selectedPointIds);
             for (const trackId of action.selectedPointIds) {
                 newSelectedPointIds.add(trackId);
             }
             newCanvas.selectedPointIds = newSelectedPointIds;
+            newCanvas.highlightPoints(action.selectedPointIndices);
             break;
         }
         case ActionType.UPDATE_WITH_STATE:

src/lib/PointCanvas.ts

@@ -1,14 +1,14 @@
 import {
-    AdditiveBlending,
     AxesHelper,
     BufferGeometry,
     Color,
     Float32BufferAttribute,
     FogExp2,
+    NormalBlending,
     PerspectiveCamera,
     Points,
-    PointsMaterial,
     Scene,
+    ShaderMaterial,
     SRGBColorSpace,
     TextureLoader,
     Vector2,
@@ -23,6 +23,10 @@ import { UnrealBloomPass } from "three/addons/postprocessing/UnrealBloomPass.js"
 import { Track } from "@/lib/three/Track";
 import { PointSelector, PointSelectionMode } from "@/lib/PointSelector";
 import { ViewerState } from "./ViewerState";
+import { numberOfValuesPerPoint } from "./TrackManager";
+
+import config from "../../CONFIG.ts";
+const pointSize = config.settings.point_size;
 
 // TrackType is a place to store the visual information about a track and any track-specific attributes
 type TrackType = {
@@ -51,6 +55,7 @@ export class PointCanvas {
     readonly fetchedRootTrackIds = new Set<number>();
     // Needed to skip fetches for point IDs that been selected.
     readonly fetchedPointIds = new Set<number>();
+    selectedPointIndices: number[] = [];
 
     // All the point IDs that have been selected.
     // PointCanvas.selector.selection is the transient array of selected
@@ -67,6 +72,7 @@ export class PointCanvas {
     // 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
     maxPointsPerTimepoint = 0;
+    private pointIndicesCache: Map<number, number[]> = new Map();
 
     constructor(width: number, height: number) {
         this.scene = new Scene();
@@ -80,16 +86,44 @@ export class PointCanvas {
         );
 
         const pointsGeometry = new BufferGeometry();
-        const pointsMaterial = new PointsMaterial({
-            size: 16.0,
-            map: new TextureLoader().load("/spark1.png"),
-            vertexColors: true,
-            blending: AdditiveBlending,
-            depthTest: true,
-            alphaTest: 0.1,
-            transparent: true,
+        const pointVertexShader = `
+            attribute float size;
+            attribute vec3 color; //Declare the color attribute
+            varying vec3 vColor;
+
+            void main() {
+            vColor = color;
+            vec4 mvPosition = modelViewMatrix * vec4(position, 1.0);
+            gl_PointSize = size * (300.0 / -mvPosition.z); // Adjust scaling factor
+            gl_Position = projectionMatrix * mvPosition;
+            }
+        `;
+        const pointFragmentShader = `
+            varying vec3 vColor;
+            uniform sampler2D pointTexture;
+
+            void main() {
+            gl_FragColor = vec4(vColor, 1.0);
+            gl_FragColor = gl_FragColor * texture2D(pointTexture, gl_PointCoord);
+            if (gl_FragColor.a < .5) discard;
+            }
+        `;
+
+        const shaderMaterial = new ShaderMaterial({
+            uniforms: {
+                color: { value: new Color(0xffffff) },
+                pointTexture: { value: new TextureLoader().load("/spark1.png") },
+            },
+            vertexShader: pointVertexShader,
+            fragmentShader: pointFragmentShader,
+
+            blending: NormalBlending,
+            depthTest: true, // true
+            // alphaTest: 0.1, //no effect
+            depthWrite: true, // true
+            transparent: false, // false
         });
-        this.points = new Points(pointsGeometry, pointsMaterial);
+        this.points = new Points(pointsGeometry, shaderMaterial);
 
         this.scene.add(new AxesHelper(128));
         this.scene.add(this.points);
@@ -101,7 +135,7 @@ export class PointCanvas {
             new Vector2(width, height), // resolution
             0.4, // strength
             0, // radius
-            0, // threshold
+            0.2, // threshold
         );
         const outputPass = new OutputPass();
         this.composer = new EffectComposer(this.renderer);
@@ -165,6 +199,47 @@ export class PointCanvas {
         this.controls.update();
     };
 
+    updateSelectedPointIndices() {
+        const cacheKey = this.createCacheKey();
+
+        // Check if the result is already cached
+        if (this.pointIndicesCache.has(cacheKey)) {
+            this.selectedPointIndices = this.pointIndicesCache.get(cacheKey)!;
+            this.highlightPoints(this.selectedPointIndices);
+            return;
+        }
+
+        // If not cached: find selectedPointIndices
+        const idOffset = this.curTime * this.maxPointsPerTimepoint;
+        this.selectedPointIndices = [];
+        for (const track of this.tracks.values()) {
+            if (this.curTime < track.threeTrack.startTime || this.curTime > track.threeTrack.endTime) continue;
+            const timeIndex = this.curTime - track.threeTrack.startTime;
+            const pointId = track.threeTrack.pointIds[timeIndex];
+            this.selectedPointIndices.push(pointId - idOffset);
+        }
+
+        this.pointIndicesCache.set(cacheKey, this.selectedPointIndices);
+        this.highlightPoints(this.selectedPointIndices);
+    }
+
+    private createCacheKey(): number {
+        let hash = 0;
+        const trackIds = Array.from(this.tracks.keys()).join(",");
+        const keyString = `${this.curTime}:${trackIds}`;
+
+        for (let i = 0; i < keyString.length; i++) {
+            const char = keyString.charCodeAt(i);
+            hash = (hash << 5) - hash + char;
+            hash = hash & hash; // Convert to 32-bit integer
+        }
+        return hash;
+    }
+
+    clearPointIndicesCache() {
+        this.pointIndicesCache.clear();
+    }
+
     highlightPoints(points: number[]) {
         const colorAttribute = this.points.geometry.getAttribute("color");
         const color = new Color();
@@ -211,18 +286,29 @@ export class PointCanvas {
         if (!geometry.hasAttribute("color") || geometry.getAttribute("color").count !== maxPointsPerTimepoint) {
             geometry.setAttribute("color", new Float32BufferAttribute(new Float32Array(3 * maxPointsPerTimepoint), 3));
         }
+        if (!geometry.hasAttribute("size") || geometry.getAttribute("size").count !== maxPointsPerTimepoint) {
+            geometry.setAttribute("size", new Float32BufferAttribute(new Float32Array(maxPointsPerTimepoint), 1));
+        }
         // Initialize all the colors immediately.
         this.resetPointColors();
     }
 
     setPointsPositions(data: Float32Array) {
-        const numPoints = data.length / 3;
+        const numPoints = data.length / numberOfValuesPerPoint;
         const geometry = this.points.geometry;
         const positions = geometry.getAttribute("position");
+        const sizes = geometry.getAttribute("size");
+        const num = numberOfValuesPerPoint;
         for (let i = 0; i < numPoints; i++) {
-            positions.setXYZ(i, data[3 * i], data[3 * i + 1], data[3 * i + 2]);
+            positions.setXYZ(i, data[num * i], data[num * i + 1], data[num * i + 2]);
+            if (num == 4) {
+                sizes.setX(i, 25 * data[num * i + 3]); // factor of 21 used to match the desired size of the points
+            } else {
+                sizes.setX(i, pointSize);
+            }
         }
         positions.needsUpdate = true;
+        sizes.needsUpdate = true;
         geometry.setDrawRange(0, numPoints);
         this.points.geometry.computeBoundingSphere();
     }
@@ -258,6 +344,7 @@ export class PointCanvas {
     }
 
     removeAllTracks() {
+        console.log("removeAllTracks!");
         this.selectedPointIds = new Set();
         this.fetchedRootTrackIds.clear();
         this.fetchedPointIds.clear();

src/lib/TrackManager.ts

@@ -1,5 +1,6 @@
 // @ts-expect-error - types for zarr are not working right now, but a PR is open https://github.com/gzuidhof/zarr.js/pull/149
-import { ZarrArray, slice, Slice, openArray, NestedArray } from "zarr";
+import { ZarrArray, slice, Slice, openArray, NestedArray, HTTPStore } from "zarr";
+export let numberOfValuesPerPoint = 0; // 3 if points=[x,y,z], 4 if points=[x,y,z,size]
 
 class SparseZarrArray {
     store: string;
@@ -83,7 +84,7 @@ export class TrackManager {
         this.tracksToPoints = tracksToPoints;
         this.tracksToTracks = tracksToTracks;
         this.numTimes = points.shape[0];
-        this.maxPointsPerTimepoint = points.shape[1] / 3;
+        this.maxPointsPerTimepoint = points.shape[1] / numberOfValuesPerPoint; // default is /3
     }
 
     async fetchPointsAtTime(timeIndex: number): Promise<Float32Array> {
@@ -97,9 +98,9 @@ export class TrackManager {
         let endIndex = points.findIndex((value) => value <= -127);
         if (endIndex === -1) {
             endIndex = points.length;
-        } else if (endIndex % 3 !== 0) {
-            console.error("invalid points - %d not divisible by 3", endIndex);
-            endIndex -= endIndex % 3;
+        } else if (endIndex % numberOfValuesPerPoint !== 0) {
+            console.error("invalid points - %d not divisible by %d", endIndex, numberOfValuesPerPoint);
+            endIndex -= endIndex % numberOfValuesPerPoint;
         }
         return points.subarray(0, endIndex);
     }
@@ -149,10 +150,26 @@ export async function loadTrackManager(url: string) {
             path: "points",
             mode: "r",
         });
+
+        // load the zarr metadata (to know is radius is included)
+        try {
+            const store = new HTTPStore(url);
+            const zattrsResponse = await store.getItem("points/.zattrs");
+            const zattrs = JSON.parse(new TextDecoder().decode(zattrsResponse));
+            numberOfValuesPerPoint = zattrs["values_per_point"];
+        } catch (error) {
+            numberOfValuesPerPoint = 3;
+        }
+
         const pointsToTracks = await openSparseZarrArray(url, "points_to_tracks", false);
         const tracksToPoints = await openSparseZarrArray(url, "tracks_to_points", true);
         const tracksToTracks = await openSparseZarrArray(url, "tracks_to_tracks", true);
+
+        // make trackManager, and reset "maxPointsPerTimepoint", because tm constructor does points/3
         trackManager = new TrackManager(url, points, pointsToTracks, tracksToPoints, tracksToTracks);
+        if (numberOfValuesPerPoint == 4) {
+            trackManager.maxPointsPerTimepoint = trackManager.points.shape[1] / numberOfValuesPerPoint;
+        }
     } catch (err) {
         console.error("Error opening TrackManager: %s", err);
         trackManager = null;

src/lib/three/Track.ts

@@ -16,6 +16,9 @@ export class Track extends Mesh {
     type = "Track";
     declare geometry: TrackGeometry;
     declare material: TrackMaterial;
+    pointIds: Int32Array = new Int32Array(0);
+    startTime: number = -1;
+    endTime: number = -1;
 
     static new(positions: Float32Array, pointIDs: Int32Array, maxPointsPerTimepoint: number) {
         const geometry = new TrackGeometry();
@@ -53,6 +56,12 @@ export class Track extends Mesh {
         track.geometry.setColors(colors);
         track.geometry.setTime(time);
         track.geometry.computeBoundingSphere();
+
+        track.pointIds = pointIDs;
+        if (time.length > 0) {
+            track.startTime = time[0];
+            track.endTime = time[time.length - 1];
+        }
         return track;
     }