Noah/webgpu scatterbrain pt1

packages/scatterbrain/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@alleninstitute/vis-scatterbrain",
-    "version": "0.0.2",
+    "version": "0.0.3",
     "contributors": [
         {
             "name": "Lane Sawyer",
@@ -37,7 +37,8 @@
     ],
     "scripts": {
         "typecheck": "tsc --noEmit",
-        "build": "parcel build --no-cache",
+        "build": "vite build",
+        "oldbuild": "parcel build --no-cache",
         "dev": "parcel watch --port 1239",
         "demo": "vite",
         "test": "vitest --watch",
@@ -55,6 +56,8 @@
         "lodash": "4.17.23",
         "regl": "2.1.0",
         "ts-pattern": "5.9.0",
+        "typegpu": "0.11.2",
+        "webgpu-utils": "2.0.2",
         "zod": "4.3.6"
     },
     "publishConfig": {
@@ -63,6 +66,10 @@
     },
     "packageManager": "pnpm@9.14.2",
     "devDependencies": {
-        "@types/lodash": "4.17.24"
+        "@types/lodash": "4.17.24",
+        "@types/node": "22.19.15",
+        "@webgpu/types": "0.1.69",
+        "vite": "8.0.8",
+        "vite-plugin-dts": "4.5.4"
     }
-}
+}

packages/scatterbrain/src/cache-client.ts

@@ -0,0 +1,67 @@
+import type { Cacheable, SharedPriorityCache } from '@alleninstitute/vis-core';
+import type { ColumnRequest, Item, } from './types';
+import reduce from 'lodash/reduce';
+import type { WebGLSafeBasicType } from './typed-array';
+import { keys } from 'lodash';
+
+
+type Content<V extends Cacheable> = Record<string, V>
+
+export function buildScatterbrainCacheClient<V extends Cacheable>(
+    allNeededColumns: readonly string[],
+    cache: SharedPriorityCache,
+    toCacheValue: (buffer: ArrayBuffer, type: WebGLSafeBasicType) => V,
+    onDataArrived: () => void,
+) {
+    const client = cache.registerClient<Item, Content<V>>({
+        cacheKeys: (item) => {
+            const { dataset, node, columns } = item;
+            return reduce<Record<string, ColumnRequest>, Record<string, string>>(
+                columns,
+                (acc, col, key) => ({
+                    ...acc,
+                    [key]: `${dataset.metadata.metadataFileEndpoint}/${node.file}/${col.name}`,
+                }),
+                {},
+            );
+        },
+        fetch: (item) => {
+            const { dataset, node, columns } = item;
+            const attrs = dataset.metadata.pointAttributes;
+            const getColumnUrl = (columnName: string) =>
+                `${dataset.metadata.metadataFileEndpoint}${columnName}/${dataset.metadata.visualizationReferenceId}/${node.file}`;
+            const getGeneUrl = (columnName: string) =>
+                `${dataset.metadata.geneFileEndpoint}${columnName}/${dataset.metadata.visualizationReferenceId}/${node.file}`;
+            const getColumnInfo = (col: ColumnRequest) =>
+                col.type === 'QUANTITATIVE'
+                    ? ({ url: getGeneUrl(col.name), elements: 1, type: 'float' } as const)
+                    : { url: getColumnUrl(col.name), elements: attrs[col.name].elements, type: attrs[col.name].type };
+
+            const proms = reduce<Record<string, ColumnRequest>, Record<string, (signal: AbortSignal) => Promise<V>>>(
+                columns,
+                (getters, col, key) => {
+                    const { url, type } = getColumnInfo(col);
+                    return {
+                        ...getters,
+                        [key]: (signal) =>
+                            fetch(url, { signal }).then((b) =>
+                                b.arrayBuffer().then((buff) => toCacheValue(buff, type))
+                            ),
+                    };
+                },
+                {},
+            );
+            return proms;
+        },
+        isValue: (v): v is Content<V> => {
+            for (const column of allNeededColumns) {
+                if (!(column in v)) {
+                    return false;
+                }
+            }
+            return true;
+        },
+        onDataArrived,
+    });
+    return client;
+}

packages/scatterbrain/src/demo.html

@@ -0,0 +1,32 @@
+<!DOCTYPE html>
+<html>
+
+<head>
+  <meta charset="utf-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1, shrink-to-fit=yes">
+  <title>hello webgpu</title>
+  <style>
+    :root {
+      color-scheme: light dark;
+    }
+
+    html,
+    body {
+      margin: 0;
+      /* remove default margin */
+      height: 100%;
+      overflow: hidden;
+      /* make body fill the browser window */
+      display: flex;
+      touch-action: none;
+      place-content: center center;
+    }
+  </style>
+  <script defer src="../dist/main.js" type="module"></script>
+</head>
+
+<body>
+  <canvas id="canvas"></canvas>
+</body>
+
+</html>

packages/scatterbrain/src/demo.ts

@@ -0,0 +1,121 @@
+
+// lets try and make not a full-fledged scatterbrain shader,
+// with all its fancy filtering, hovering, dot sizes, etc
+// but instead, some subplot shaders - so we render the dots,
+// but we have no fancy filtering, just a simple highlight value,
+// and a color-by attribute
+
+
+// and lets try it with typeGPU generating our shaders for us... which I must admit seems pretty good...
+
+import type { ScatterbrainDataset, SlideviewScatterbrainDataset } from './types';
+import { SharedPriorityCache } from '@alleninstitute/vis-core';
+import { loadDataset } from './dataset';
+import { Box2D, type vec4 } from '@alleninstitute/vis-geometry';
+import { buildRenderFrameFn, type ShaderSettings } from './render/webgpu/renderer';
+
+
+const tenx =
+    'https://bkp-2d-visualizations-stage.s3.amazonaws.com/wmb_tenx_01172024_stage-20240128193624/G4I4GFJXJB9ATZ3PTX1/ScatterBrain.json';
+
+async function loadRawJson() {
+    return await (await fetch(tenx)).json();
+}
+const makeFakeColors = (n: number) => {
+    const stuff: Record<number, { color: vec4; filteredIn: boolean }> = {};
+    for (let i = 0; i < n; i++) {
+        stuff[i] = {
+            color: [Math.random(), Math.random(), Math.random(), 1],
+            // 80% of either category are filtered in, at random:
+            filteredIn: Math.random() > 0.2,
+        };
+    }
+    return stuff;
+};
+
+
+export async function whatever() {
+
+    const gradientData = new Uint8Array(256 * 4);
+    for (let i = 0; i < 256; i += 4) {
+        gradientData[i * 4 + 0] = i;
+        gradientData[i * 4 + 1] = i;
+        gradientData[i * 4 + 2] = i;
+        gradientData[i * 4 + 3] = 255;
+    }
+    const adapter = await navigator.gpu.requestAdapter()
+    const device = await adapter?.requestDevice()!;
+    // buildTest(root.device)
+
+    const categories = {
+        '4MV7HA5DG2XJZ3UD8G9': makeFakeColors(40), // nt type
+        FS00DXV0T9R1X9FJ4QE: makeFakeColors(40), // class
+    };
+
+    const settings: Omit<ShaderSettings, 'dataset'> = {
+        categoricalFilters: { '4MV7HA5DG2XJZ3UD8G9': 40, FS00DXV0T9R1X9FJ4QE: 40 },
+        colorBy: { kind: 'metadata', column: 'FS00DXV0T9R1X9FJ4QE' },
+        // an alternative color-by setting, swap it to see quantitative coloring
+        // colorBy: { kind: 'quantitative', column: '27683', gradient: 'viridis', range: { min: 0, max: 10 } },
+        mode: 'color',
+        quantitativeFilters: [],
+        highlightByColumn: { kind: 'metadata', column: 'FS00DXV0T9R1X9FJ4QE' }
+    };
+
+    const dataset = await loadDataset(await loadRawJson())
+    if (!dataset) {
+        throw new Error('blerg this data is toast')
+    }
+    const cache = new SharedPriorityCache(new Map(), 1024 * 1024 * 2000);
+    const { render, connectToCache } = buildRenderFrameFn(device, { ...settings, dataset })
+
+    const cnvs: HTMLCanvasElement = document.getElementById('canvas') as HTMLCanvasElement;
+    cnvs.width = 1500;
+    cnvs.height = 1500;
+    const ctx = cnvs.getContext('webgpu')
+    ctx?.configure({
+        device: device,
+        format: navigator.gpu.getPreferredCanvasFormat(),
+        alphaMode: 'premultiplied'
+    })
+
+    const bound = (dataset as ScatterbrainDataset).metadata.tightBoundingBox;
+    const view = Box2D.create([bound.lx, bound.ly], [bound.ux, bound.uy]);
+    const client = connectToCache(cache, () => {
+        // redraw?
+        // console.log('new data arrived...')
+        requestAnimationFrame(() => {
+            console.log('re render!')
+
+            render({
+                categories,
+                client,
+                gradient: gradientData,
+                target: ctx!.getCurrentTexture().createView(),
+                uniforms: {
+                    camera: { view, screenResolution: [1500, 1500] },
+                    filteredOutColor: [0.5, 0.5, 0.5, 1.0],
+                    highlightedValue: 22,
+                    offset: [0, 0],
+                    quantitativeRangeFilters: {},
+                    spatialFilterBox: view,
+                }
+            })
+        });
+    });
+    render({
+        categories,
+        client,
+        gradient: gradientData,
+        target: ctx!.getCurrentTexture().createView(),
+        uniforms: {
+            camera: { view, screenResolution: [1500, 1500] },
+            filteredOutColor: [0.5, 0.5, 0.5, 1.0],
+            highlightedValue: 22,
+            offset: [0, 0],
+            quantitativeRangeFilters: {},
+            spatialFilterBox: view,
+        }
+    })
+}
+whatever();

packages/scatterbrain/src/index.ts

@@ -1,8 +1,11 @@
 export {
     buildRenderFrameFn as buildScatterbrainRenderFn,
-    buildScatterbrainCacheClient,
     setCategoricalLookupTableValues,
     updateCategoricalValue,
-} from './renderer';
+} from './render/webgl/renderer';
+export {
+    buildRenderFrameFn as buildWebGPUScatterbrainRenderFn,
+} from './render/webgpu/renderer';
+export { buildScatterbrainCacheClient } from './cache-client'
 export * from './types';
-export { getVisibleItems, loadDataset as loadScatterbrainDataset } from './dataset';
+export { getVisibleItems, loadDataset as loadScatterbrainDataset } from './dataset';