Scatterbrain rendering in vis, including shader generation

packages/geometry/src/index.ts

@@ -37,4 +37,4 @@ export {
     PLANE_XZ,
     PLANE_YZ,
 } from './plane';
-export { type SpatialTreeInterface, visitBFS } from './spatialIndexing/tree';
+export { type SpatialTreeInterface, visitBFS, visitBFSMaybe } from './spatialIndexing/tree';

packages/geometry/src/spatialIndexing/tree.ts

@@ -14,12 +14,7 @@ export function visitBFS<Tree>(
 ): void {
     const frontier: Tree[] = [tree];
     while (frontier.length > 0) {
-        const cur = frontier.shift();
-        if (cur === undefined) {
-            // TODO: Consider logging a warning or error here, as this should never happen,
-            // but this package doesn't depend on the package where our logger lives
-            continue;
-        }
+        const cur = frontier.shift()!;
         visitor(cur);
         for (const c of children(cur)) {
             if (traversalPredicate?.(c) ?? true) {
@@ -30,3 +25,19 @@ export function visitBFS<Tree>(
         }
     }
 }
+export function visitBFSMaybe<Tree>(
+    tree: Tree,
+    children: (t: Tree) => ReadonlyArray<Tree>,
+    visitor: (tree: Tree) => boolean,
+): void {
+    const frontier: Tree[] = [tree];
+    while (frontier.length > 0) {
+        const cur = frontier.shift()!;
+
+        if (visitor(cur)) {
+            for (const c of children(cur)) {
+                frontier.push(c);
+            }
+        }
+    }
+}

packages/scatterbrain/package.json

@@ -0,0 +1,67 @@
+{
+    "name": "@alleninstitute/vis-scatterbrain",
+    "version": "0.0.1",
+    "contributors": [
+        {
+            "name": "Lane Sawyer",
+            "email": "lane.sawyer@alleninstitute.org"
+        },
+        {
+            "name": "James Gerstenberger",
+            "email": "james.gerstenberger@alleninstitute.org"
+        },
+        {
+            "name": "Noah Shepard",
+            "email": "noah.shepard@alleninstitute.org"
+        },
+        {
+            "name": "Skyler Moosman",
+            "email": "skyler.moosman@alleninstitute.org"
+        },
+        {
+            "name": "Su Li",
+            "email": "su.li@alleninstitute.org"
+        },
+        {
+            "name": "Joel Arbuckle",
+            "email": "joel.arbuckle@alleninstitute.org"
+        }
+    ],
+    "license": "BSD-3-Clause",
+    "source": "src/index.ts",
+    "main": "dist/main.js",
+    "types": "dist/types.d.ts",
+    "type": "module",
+    "files": [
+        "dist"
+    ],
+    "scripts": {
+        "typecheck": "tsc --noEmit",
+        "build": "parcel build --no-cache",
+        "dev": "parcel watch --port 1239",
+        "demo": "vite",
+        "test": "vitest --watch",
+        "test:ci": "vitest run",
+        "coverage": "vitest run --coverage",
+        "changelog": "git-cliff -o changelog.md"
+    },
+    "repository": {
+        "type": "git",
+        "url": "https://github.com/AllenInstitute/vis.git"
+    },
+    "dependencies": {
+        "@alleninstitute/vis-core": "workspace:*",
+        "@alleninstitute/vis-geometry": "workspace:*",
+        "lodash": "4.17.21",
+        "regl": "2.1.0",
+        "ts-pattern": "5.9.0",
+        "zod": "4.3.6"
+    },
+    "publishConfig": {
+        "registry": "https://npm.pkg.github.com/AllenInstitute"
+    },
+    "packageManager": "pnpm@9.14.2",
+    "devDependencies": {
+        "@types/lodash": "4.17.23"
+    }
+}

packages/scatterbrain/readme.md

@@ -0,0 +1,23 @@
+
+# Scatterbrain
+
+rendering utilities to render scatter plots, with specific support for 'scatterbrain' style (quad-tree spatial indexed) data sources as used for years now in ABC-atlas
+
+## Why? ##
+
+We wrote the ABC-atlas version of Scatterbrain rendering in a hot hot hurry a few years ago. At the time, we were preparing for a lot of variation along certain paths of development,
+and as is often the case, those guesses were a bit off the mark. As a result, the flexibility points we built into that version are not helping us much. For example, we take great pains to
+generate shaders with readable "column names" as users select different filter settings. However, for reasons, the names are just referenceIds from the backend, and we really never bother to debug the shaders in that way - they either fail up-front, or work fine, or are broken in more subtle ways that tend to have nothing to do with the names of the data.
+
+## What is the goal here? ##
+
+The goal is to modernize and simplify the WebGL powered features of ABC-atlas. We'd like to be able to understand the management of rendering resources, have more comprehensible interop
+with the rest of the React UI system, and in general reduce the confusingly generic (and needlessly generic) nature of various parts, as well as having less verbose and weird areas around
+the actual features we did build (but didn't really prepare for) like:
+
+hovering to report cell-info to the rest of the system
+multiple gene coloring / filtering
+slide-view / regular view
+
+to this end, we are gonna stick to the Renderer<> interface as given in packages/core/src/abstract/types.ts. We've seen that work with both the somewhat Ill-fated renderServer, as well as the new shared-cache
+stuff, so it seems like it might be ok.

packages/scatterbrain/src/dataset.ts

@@ -0,0 +1,202 @@
+import {
+    Box2D,
+    type box2D,
+    type box3D,
+    Box3D,
+    Vec2,
+    type vec2,
+    type vec3,
+    Vec3,
+    visitBFSMaybe,
+} from '@alleninstitute/vis-geometry';
+import type { PointAttribute, ScatterbrainDataset, SlideviewScatterbrainDataset, TreeNode, volumeBound } from './types';
+import reduce from 'lodash/reduce';
+import * as z from 'zod';
+
+type Dataset = ScatterbrainDataset | SlideviewScatterbrainDataset;
+// figure out that path through the tree, given a TreeNode name
+// these names are structured data - so it should always be possible
+export type NodeWithBounds = { node: TreeNode; bounds: box2D };
+
+// adapted from Potree createChildAABB
+// note that if you do not do indexing in precisely the same order
+// as potree octrees, this will not work correctly at all
+function getChildBoundsUsingPotreeIndexing(parentBounds: box3D, index: number) {
+    const min = parentBounds.minCorner;
+    const size = Vec3.scale(Box3D.size(parentBounds), 0.5);
+    const offset: vec3 = [
+        (index & 0b0100) > 0 ? size[0] : 0,
+        (index & 0b0010) > 0 ? size[1] : 0,
+        (index & 0b0001) > 0 ? size[2] : 0,
+    ];
+    const newMin = Vec3.add(min, offset);
+    return {
+        minCorner: newMin,
+        maxCorner: Vec3.add(newMin, size),
+    };
+}
+function children(node: TreeNode) {
+    return node.children ?? [];
+}
+function sanitizeName(fileName: string) {
+    return fileName.replace('.bin', '');
+}
+function bounds(rootBounds: box3D, path: string) {
+    // path is a name like r01373 - indicating a path through the tree, each character is a child index
+    let bounds = rootBounds;
+    for (let i = 1; i < path.length; i++) {
+        const ci = Number(path[i]);
+        bounds = getChildBoundsUsingPotreeIndexing(bounds, ci);
+    }
+    return bounds;
+}
+function dropZ(box: box3D) {
+    return {
+        minCorner: Vec3.xy(box.minCorner),
+        maxCorner: Vec3.xy(box.maxCorner),
+    };
+}
+
+function getVisibleItemsInTree(
+    dataset: { root: TreeNode; boundingBox: volumeBound },
+    camera: { view: box2D; screenResolution: vec2 },
+    limit: number,
+) {
+    const { root, boundingBox } = dataset;
+    const hits: { node: TreeNode; bounds: box2D }[] = [];
+    const rootBounds = Box3D.create(
+        [boundingBox.lx, boundingBox.ly, boundingBox.lz],
+        [boundingBox.ux, boundingBox.uy, boundingBox.uz],
+    );
+    visitBFSMaybe<TreeNode>(root, children, (t) => {
+        const B = dropZ(bounds(rootBounds, sanitizeName(t.file)));
+        if (Box2D.intersection(B, camera.view) && Box2D.size(B)[0] > limit) {
+            // this node is big enough to render - that means we should check its children as well
+            hits.push({ node: t, bounds: B });
+            return true;
+        }
+        return false;
+    });
+    return hits;
+}
+
+export function getVisibleItems(
+    dataset: Dataset,
+    camera: { view: box2D; screenResolution: vec2; layout?: Record<string, vec2> },
+    visibilitySizeThreshold: number,
+) {
+    // determine
+    if (dataset.type === 'normal') {
+        return getVisibleItemsInTree(dataset.metadata, camera, visibilitySizeThreshold);
+    }
+    // by default, if we pass NO layout info
+    const size: vec2 = [
+        dataset.metadata.spatialUnit.maxX - dataset.metadata.spatialUnit.minX,
+        dataset.metadata.spatialUnit.maxY - dataset.metadata.spatialUnit.minY,
+    ];
+    // then it means we want to draw all the slides on top of each other
+    const defaultVisibility = camera.layout === undefined;
+    const hits: NodeWithBounds[] = [];
+    for (const slide of dataset.metadata.slides) {
+        if (!defaultVisibility && camera.layout?.[slide.featureTypeValueReferenceId] === undefined) {
+            // the camera has a layout, but this slide isn't in it - dont draw it
+            continue;
+        }
+        const grid = camera.layout?.[slide.featureTypeValueReferenceId] ?? [0, 0];
+
+        const offset = Vec2.mul(grid, size);
+        // offset the camera by the opposite of the offset
+        hits.push(
+            ...getVisibleItemsInTree(
+                slide.tree,
+                { ...camera, view: Box2D.translate(camera.view, offset) },
+                visibilitySizeThreshold,
+            ),
+        );
+    }
+    return hits;
+}
+const pointAttrSchema = z.object({
+    name: z.string(),
+    size: z.number(),
+    elements: z.number(), // values per point (so a vector xy would have 2)
+    elementSize: z.number(), // size of an element, given in bytes (for example float would have 4)
+    type: z.union([
+        z.literal('uint8'),
+        z.literal('uint16'),
+        z.literal('uint32'),
+        z.literal('int8'),
+        z.literal('int16'),
+        z.literal('int32'),
+        z.literal('float'),
+    ]),
+    description: z.string(),
+});
+const commonMetadataSchema = z.object({
+    geneFileEndpoint: z.string(),
+    metadataFileEndpoint: z.string(),
+    visualizationReferenceId: z.string(),
+    spatialColumn: z.string(),
+    pointAttributes: z.array(pointAttrSchema).transform((attrs) => {
+        return reduce<PointAttribute, Record<string, PointAttribute>>(
+            attrs,
+            (acc, attr) => ({ ...acc, [attr.name]: attr }),
+            {},
+        );
+    }),
+});
+const bbSchema = z.object({
+    lx: z.number(),
+    ly: z.number(),
+    lz: z.number(),
+    ux: z.number(),
+    uy: z.number(),
+    uz: z.number(),
+});
+const treeNodeSchema = z.object({
+    file: z.string(),
+    numSpecimens: z.number(),
+    get children() {
+        return z.union([z.undefined(), z.array(treeNodeSchema)]);
+    },
+});
+const treeSchema = z.object({
+    points: z.number(),
+    boundingBox: bbSchema,
+    tightBoundingBox: bbSchema,
+    root: treeNodeSchema,
+});
+const scatterbrainMetadataSchema = z.object({
+    ...treeSchema.shape,
+    ...commonMetadataSchema.shape,
+});
+const slideSchema = z.object({
+    featureTypeValueReferenceId: z.string(),
+    tree: treeSchema,
+});
+const spatialRefFrameSchema = z.object({
+    anatomicalOrigin: z.string(),
+    direction: z.string(),
+    unit: z.string(),
+    minX: z.number(),
+    maxX: z.number(),
+    minY: z.number(),
+    maxY: z.number(),
+});
+const slideviewMetadataSchema = z.object({
+    ...commonMetadataSchema.shape,
+    slides: z.array(slideSchema),
+    spatialUnit: spatialRefFrameSchema,
+});
+
+// biome-ignore lint/suspicious/noExplicitAny: this fn is intended to accept the return value of JSON.parse() - any is appropriate here
+export function loadDataset(raw: any): Dataset | undefined {
+    if (typeof raw !== 'object' || !raw) return undefined;
+
+    if (raw.slides) {
+        const metadata = slideviewMetadataSchema.safeParse(raw);
+        return metadata.success ? { type: 'slideview', metadata: metadata.data } : undefined;
+    }
+    const metadata = scatterbrainMetadataSchema.safeParse(raw);
+    return metadata.success ? { type: 'normal', metadata: metadata.data } : undefined;
+}

packages/scatterbrain/src/index.ts

@@ -0,0 +1,8 @@
+export {
+    buildRenderFrameFn as buildScatterbrainRenderFn,
+    buildScatterbrainCacheClient,
+    setCategoricalLookupTableValues,
+    updateCategoricalValue,
+} from './renderer';
+export * from './types';
+export { getVisibleItems, loadDataset as loadScatterbrainDataset } from './dataset';