diff --git a/packages/geometry/src/index.ts b/packages/geometry/src/index.ts index 12994db1..7d7c54fd 100644 --- a/packages/geometry/src/index.ts +++ b/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'; diff --git a/packages/geometry/src/spatialIndexing/tree.ts b/packages/geometry/src/spatialIndexing/tree.ts index 0cf516c2..1a2edac8 100644 --- a/packages/geometry/src/spatialIndexing/tree.ts +++ b/packages/geometry/src/spatialIndexing/tree.ts @@ -14,12 +14,7 @@ export function visitBFS( ): 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( } } } +export function visitBFSMaybe( + tree: Tree, + children: (t: Tree) => ReadonlyArray, + 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); + } + } + } +} diff --git a/packages/scatterbrain/package.json b/packages/scatterbrain/package.json new file mode 100644 index 00000000..7568f60c --- /dev/null +++ b/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" + } +} diff --git a/packages/scatterbrain/readme.md b/packages/scatterbrain/readme.md new file mode 100644 index 00000000..7d033ea6 --- /dev/null +++ b/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. diff --git a/packages/scatterbrain/src/dataset.ts b/packages/scatterbrain/src/dataset.ts new file mode 100644 index 00000000..98d00044 --- /dev/null +++ b/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(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 }, + 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>( + 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; +} diff --git a/packages/scatterbrain/src/index.ts b/packages/scatterbrain/src/index.ts new file mode 100644 index 00000000..bcc39ea3 --- /dev/null +++ b/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'; diff --git a/packages/scatterbrain/src/renderer.ts b/packages/scatterbrain/src/renderer.ts new file mode 100644 index 00000000..10a9e983 --- /dev/null +++ b/packages/scatterbrain/src/renderer.ts @@ -0,0 +1,222 @@ +import type { SharedPriorityCache } from '@alleninstitute/vis-core'; +import type REGL from 'regl'; +import type { ColumnRequest, ScatterbrainDataset, SlideviewScatterbrainDataset, TreeNode } from './types'; +import { Box2D, type box2D, type vec4 } from '@alleninstitute/vis-geometry'; +import { MakeTaggedBufferView } from './typed-array'; +import keys from 'lodash/keys'; +import reduce from 'lodash/reduce'; +import { getVisibleItems, type NodeWithBounds } from './dataset'; +import { buildScatterbrainRenderCommand, type Config, configureShader, type ShaderSettings, VBO } from './shader'; +export type Item = Readonly<{ + dataset: SlideviewScatterbrainDataset | ScatterbrainDataset; + node: TreeNode; + bounds: box2D; + columns: Record; +}>; +type Content = Record; + +export function buildScatterbrainCacheClient( + allNeededColumns: readonly string[], + regl: REGL.Regl, + cache: SharedPriorityCache, + onDataArrived: () => void, +) { + const client = cache.registerClient({ + cacheKeys: (item) => { + const { dataset, node, columns } = item; + return reduce, Record>( + 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 Promise>>( + columns, + (getters, col, key) => { + const { url, type } = getColumnInfo(col); + return { + ...getters, + [key]: (signal) => + fetch(url, { signal }).then((b) => + b.arrayBuffer().then((buff) => { + const typed = MakeTaggedBufferView(type, buff); + return new VBO({ + buffer: regl.buffer({ type: type, data: typed.data }), + bytes: buff.byteLength, + type: 'buffer', + }); + }), + ), + }; + }, + {}, + ); + return proms; + }, + isValue: (v): v is Content => { + for (const column of allNeededColumns) { + if (!(column in v)) { + return false; + } + } + return true; + }, + onDataArrived, + }); + return client; +} + +function columnsForItem( + config: Config, + col2shader: Record, + dataset: ScatterbrainDataset | SlideviewScatterbrainDataset, +) { + const columns: Record = {}; + const s2c = reduce( + keys(col2shader), + (acc, col) => ({ ...acc, [col2shader[col]]: col }), + {} as Record, + ); + + for (const c of config.categoricalColumns) { + columns[c] = { type: 'METADATA', name: s2c[c] }; + } + for (const m of config.quantitativeColumns) { + columns[m] = { type: 'QUANTITATIVE', name: s2c[m] }; + } + columns[config.positionColumn] = { type: 'METADATA', name: dataset.metadata.spatialColumn }; + return (item: T) => { + return { ...item, dataset, columns }; + }; +} + +/** + * a helper function that MUTATES ALL the values in the given @param texture + * to set them to the color and filter status as given in the categories record + * note that the texture's maping to categories is based on a lexical sorting of the names of the + * categories + * @param categories + * @param regl + * @param texture + */ +export function setCategoricalLookupTableValues( + categories: Record>, + texture: REGL.Texture2D, +) { + const categoryKeys = keys(categories).toSorted(); + const columns = categoryKeys.length; + const rows = reduce(categoryKeys, (highest, category) => Math.max(highest, keys(categories[category]).length), 1); + const data = new Uint8Array(columns * rows * 4); + const rgbf = [0, 0, 0, 0]; + const empty = [0, 0, 0, 0] as const; + // write the rgb of the color, and encode the filter boolean into the alpha channel + for (let columnIndex = 0; columnIndex < columns; columnIndex += 1) { + const category = categories[categoryKeys[columnIndex]]; + const nRows = keys(category).length; + for (let rowIndex = 0; rowIndex < nRows; rowIndex += 1) { + const color = category[rowIndex]?.color ?? empty; + const filtered = category[rowIndex]?.filteredIn ?? false; + rgbf[0] = color[0] * 255; + rgbf[1] = color[1] * 255; + rgbf[2] = color[2] * 255; + rgbf[3] = filtered ? 255 : 0; + data.set(rgbf, rowIndex * columns * 4 + columnIndex * 4); + } + } + // calling a texture as a function is REGL shorthand for total re-init of this texture, capable of resizing if needed + // warning - this is not likely to be fast + texture({ data, width: columns, height: rows }); +} +/** + * same as setCategoricalLookupTableValues, except it only writes a single value update to the texture. + * note that the list of categories given must match those used to construct the texture, and are needed here + * due to the lexical sorting order determining the column order of the @param texture + * @param categories + * @param update + * @param regl + * @param texture + */ +export function updateCategoricalValue( + categories: readonly string[], + update: { category: string; row: number; color: vec4; filteredIn: boolean }, + texture: REGL.Texture2D, +) { + const { category, row, color, filteredIn } = update; + const col = categories.toSorted().indexOf(category); + if (texture.width <= col || texture.height <= row || row < 0 || col < 0) { + // todo - it might be better to let regl throw the same error... think about it + throw new Error( + `attempted to update metadata lookup table with invalid coordinates: row=${row},col=${col} is not within ${texture.width}, ${texture.height}`, + ); + } + const data = new Uint8Array(4); + data[0] = color[0] * 255; + data[1] = color[1] * 255; + data[2] = color[2] * 255; + data[3] = filteredIn ? 255 : 0; + texture.subimage(data, col, row); +} + +type ScatterbrainRenderProps = Omit>[0], 'item'> & { + visibilityThresholdPx: number; + dataset: ScatterbrainDataset | SlideviewScatterbrainDataset; + client: ReturnType; +}; +/** + * + * @param regl a regl context + * @param settings settings describing the data and how it should be rendered + * @returns a pair of functions: + * render: when called with renderable data, will determine the set of visible data, request that data from the client, and then draw all currently available data + * connectToCache - called to produce a cacheClient, which must be passed to the render function + */ +export function buildRenderFrameFn(regl: REGL.Regl, settings: ShaderSettings) { + const { dataset } = settings; + const { config, columnNameToShaderName } = configureShader(settings); + + const prepareQtCell = columnsForItem(config, columnNameToShaderName, dataset); + const drawQtCell = buildScatterbrainRenderCommand(config, regl); + + const render = (props: ScatterbrainRenderProps) => { + const { camera, dataset, client, visibilityThresholdPx } = props; + // compute the size of a screen pixel in data-space units + const visibilityThreshold = (visibilityThresholdPx * Box2D.size(camera.view)[0]) / camera.screenResolution[0]; // (units*pixel)/pixel ==> units + const visibleQtNodes = getVisibleItems(dataset, camera, visibilityThreshold).map(prepareQtCell); + client.setPriorities(visibleQtNodes, []); + for (const node of visibleQtNodes) { + if (client.has(node)) { + const drawable = client.get(node); + if (drawable) { + drawQtCell({ + ...props, + item: { + columnData: drawable, + count: node.node.numSpecimens, + }, + }); + } + } + } + }; + const connectToCache = (cache: SharedPriorityCache, onDataArrived: () => void) => { + const allColumns = [...config.categoricalColumns, ...config.quantitativeColumns, config.positionColumn]; + const client = buildScatterbrainCacheClient(allColumns, regl, cache, onDataArrived); + return client; + }; + return { render, connectToCache }; +} diff --git a/packages/scatterbrain/src/shader.test.ts b/packages/scatterbrain/src/shader.test.ts new file mode 100644 index 00000000..948f9777 --- /dev/null +++ b/packages/scatterbrain/src/shader.test.ts @@ -0,0 +1,520 @@ +import { describe, expect, test } from 'vitest'; +import { buildShaders, type Config, configureShader } from './shader'; +import type { ScatterbrainDataset } from './types'; + +const tenx: ScatterbrainDataset = { + type: 'normal', + metadata: { + points: 1494801, + boundingBox: { + lx: -11.541528999999999, + ly: -13.448518, + lz: 0, + ux: 24.993372, + uy: 23.086382999999998, + uz: 36.534901, + }, + tightBoundingBox: { + lx: -11.541528999999999, + ly: -13.448518, + lz: 0, + ux: 24.993372, + uy: 22.9617, + uz: 0, + }, + root: { + file: 'r.bin', + numSpecimens: 100496, + children: [ + { + file: 'r0.bin', + numSpecimens: 33084, + children: [ + { + file: 'r00.bin', + numSpecimens: 1004, + children: [], + }, + { + file: 'r02.bin', + numSpecimens: 8959, + children: [], + }, + { + file: 'r04.bin', + numSpecimens: 18257, + children: [], + }, + { + file: 'r06.bin', + numSpecimens: 21375, + children: [ + { + file: 'r060.bin', + numSpecimens: 1150, + children: [], + }, + { + file: 'r062.bin', + numSpecimens: 818, + children: [], + }, + { + file: 'r064.bin', + numSpecimens: 65649, + children: [], + }, + { + file: 'r066.bin', + numSpecimens: 5242, + children: [], + }, + ], + }, + ], + }, + { + file: 'r2.bin', + numSpecimens: 83987, + children: [ + { + file: 'r20.bin', + numSpecimens: 41058, + children: [ + { + file: 'r200.bin', + numSpecimens: 10658, + children: [], + }, + { + file: 'r202.bin', + numSpecimens: 37462, + children: [], + }, + { + file: 'r204.bin', + numSpecimens: 4385, + children: [], + }, + { + file: 'r206.bin', + numSpecimens: 23558, + children: [], + }, + ], + }, + { + file: 'r22.bin', + numSpecimens: 19404, + children: [ + { + file: 'r220.bin', + numSpecimens: 33089, + children: [], + }, + { + file: 'r224.bin', + numSpecimens: 33541, + children: [], + }, + { + file: 'r226.bin', + numSpecimens: 303, + children: [], + }, + ], + }, + { + file: 'r24.bin', + numSpecimens: 69249, + children: [ + { + file: 'r240.bin', + numSpecimens: 25571, + children: [], + }, + { + file: 'r242.bin', + numSpecimens: 32062, + children: [], + }, + { + file: 'r244.bin', + numSpecimens: 40682, + children: [], + }, + { + file: 'r246.bin', + numSpecimens: 60049, + children: [ + { + file: 'r2460.bin', + numSpecimens: 49708, + children: [], + }, + { + file: 'r2462.bin', + numSpecimens: 12312, + children: [], + }, + { + file: 'r2464.bin', + numSpecimens: 8059, + children: [], + }, + { + file: 'r2466.bin', + numSpecimens: 3863, + children: [], + }, + ], + }, + ], + }, + { + file: 'r26.bin', + numSpecimens: 40202, + children: [ + { + file: 'r260.bin', + numSpecimens: 18073, + children: [], + }, + { + file: 'r262.bin', + numSpecimens: 6172, + children: [], + }, + { + file: 'r264.bin', + numSpecimens: 25204, + children: [], + }, + { + file: 'r266.bin', + numSpecimens: 5508, + children: [], + }, + ], + }, + ], + }, + { + file: 'r4.bin', + numSpecimens: 41387, + children: [ + { + file: 'r40.bin', + numSpecimens: 9924, + children: [], + }, + { + file: 'r42.bin', + numSpecimens: 40161, + children: [], + }, + { + file: 'r44.bin', + numSpecimens: 1883, + children: [], + }, + { + file: 'r46.bin', + numSpecimens: 33158, + children: [], + }, + ], + }, + { + file: 'r6.bin', + numSpecimens: 74145, + children: [ + { + file: 'r60.bin', + numSpecimens: 53224, + children: [ + { + file: 'r600.bin', + numSpecimens: 41323, + children: [ + { + file: 'r6000.bin', + numSpecimens: 1895, + children: [], + }, + { + file: 'r6002.bin', + numSpecimens: 24075, + children: [], + }, + { + file: 'r6004.bin', + numSpecimens: 1609, + children: [], + }, + { + file: 'r6006.bin', + numSpecimens: 18091, + children: [], + }, + ], + }, + { + file: 'r602.bin', + numSpecimens: 30526, + children: [], + }, + { + file: 'r604.bin', + numSpecimens: 4394, + children: [], + }, + { + file: 'r606.bin', + numSpecimens: 21398, + children: [], + }, + ], + }, + { + file: 'r62.bin', + numSpecimens: 46391, + children: [ + { + file: 'r620.bin', + numSpecimens: 9355, + children: [], + }, + { + file: 'r622.bin', + numSpecimens: 20763, + children: [], + }, + { + file: 'r624.bin', + numSpecimens: 7436, + children: [], + }, + { + file: 'r626.bin', + numSpecimens: 11283, + children: [], + }, + ], + }, + { + file: 'r64.bin', + numSpecimens: 42149, + children: [], + }, + { + file: 'r66.bin', + numSpecimens: 20038, + children: [], + }, + ], + }, + ], + }, + spatialColumn: '488I12FURRB8ZY5KJ8TCoordinates', + visualizationReferenceId: '488I12FURRB8ZY5KJ8T', + geneFileEndpoint: + 'https://bkp-2d-visualizations-stage.s3.amazonaws.com/wmb_tenx_01172024_stage-20240128193624/488I12FURRB8ZY5KJ8T/data/', + metadataFileEndpoint: + 'https://bkp-2d-visualizations-stage.s3.amazonaws.com/wmb_tenx_01172024_stage-20240128193624/488I12FURRB8ZY5KJ8T/metadata/', + pointAttributes: { + FS00DXV0T9R1X9FJ4QE: { + name: 'FS00DXV0T9R1X9FJ4QE', + size: 2, + elements: 1, + elementSize: 2, + type: 'uint16', + description: 'Class', + }, + QY5S8KMO5HLJUF0P00K: { + name: 'QY5S8KMO5HLJUF0P00K', + size: 2, + elements: 1, + elementSize: 2, + type: 'uint16', + description: 'Subclass', + }, + '15BK47DCIOF1SLLUW9P': { + name: '15BK47DCIOF1SLLUW9P', + size: 2, + elements: 1, + elementSize: 2, + type: 'uint16', + description: 'Supertype', + }, + CBGC0U30VV9JPR60TJU: { + name: 'CBGC0U30VV9JPR60TJU', + size: 2, + elements: 1, + elementSize: 2, + type: 'uint16', + description: 'Cluster', + }, + '4MV7HA5DG2XJZ3UD8G9': { + name: '4MV7HA5DG2XJZ3UD8G9', + size: 2, + elements: 1, + elementSize: 2, + type: 'uint16', + description: 'Neurotransmitter Type', + }, + Y937CVUSVZC7KYOHWVO: { + name: 'Y937CVUSVZC7KYOHWVO', + size: 2, + elements: 1, + elementSize: 2, + type: 'uint16', + description: 'Dissection Region', + }, + KRP9GYF002I5OPM7JSR: { + name: 'KRP9GYF002I5OPM7JSR', + size: 2, + elements: 1, + elementSize: 2, + type: 'uint16', + description: 'Donor ID', + }, + N3YEG845JSIPMS3C0MJ: { + name: 'N3YEG845JSIPMS3C0MJ', + size: 2, + elements: 1, + elementSize: 2, + type: 'uint16', + description: 'Platform', + }, + O95N6FNAK13WZWEIU5N: { + name: 'O95N6FNAK13WZWEIU5N', + size: 2, + elements: 1, + elementSize: 2, + type: 'uint16', + description: 'Sex', + }, + Q0LG0S1W23HUAKA2SW3: { + name: 'Q0LG0S1W23HUAKA2SW3', + size: 2, + elements: 1, + elementSize: 2, + type: 'uint16', + description: 'Genotype', + }, + '488I12FURRB8ZY5KJ8TCoordinates': { + name: '488I12FURRB8ZY5KJ8TCoordinates', + size: 8, + elements: 2, + elementSize: 4, + type: 'float', + description: 'Pallium-Glut', + }, + }, + }, +}; +describe('configure', () => { + test('can we generate a sensible shader from settings', () => { + const { config, columnNameToShaderName } = configureShader({ + dataset: tenx, + categoricalFilters: {}, + mode: 'color', + colorBy: { kind: 'quantitative', column: '123', gradient: 'viridis', range: { min: 0, max: 10 } }, + quantitativeFilters: [], + }); + const shaders = buildShaders(config); + + const expectedConfig: Config = { + categoricalColumns: [], + categoricalTable: 'lookup', + colorByColumn: 'COLOR_BY_MEASURE', + gradientTable: 'gradient', + mode: 'color', + tableSize: [1, 1], + quantitativeColumns: ['COLOR_BY_MEASURE'], + positionColumn: 'position', + }; + const expectedShader = /*glsl*/ ` + precision highp float; + // attribs // + + attribute vec2 position; + + attribute float COLOR_BY_MEASURE; + + // uniforms // + + uniform vec4 view; + uniform vec2 screenSize; + uniform vec2 offset; + uniform vec4 spatialFilterBox; + uniform vec4 filteredOutColor; + uniform float hoveredValue; + + uniform sampler2D gradient; + uniform sampler2D lookup; + // quantitative columns each need a range value - its the min,max in a vec2 + uniform vec2 COLOR_BY_MEASURE_range; + + + // utility functions // + + vec4 applyCamera(vec3 dataPos){ + vec2 size = view.zw-view.xy; + vec2 unit = (dataPos.xy-view.xy)/size; + return vec4((unit*2.0)-1.0,0.0,1.0); + } + float rangeParameter(float v, vec2 range){ + return (v-range.x)/(range.y-range.x); + } + float within(float v, vec2 range){ + return step(range.x,v)*step(v,range.y); + } + + + // per-point interface functions // + + float isHovered(){ + + return 0.0; + } + vec3 getDataPosition(){ + return vec3(position+offset,0.0); + } + float isFilteredIn(){ + + vec3 p = getDataPosition(); + return within(p.x,spatialFilterBox.xz)*within(p.y,spatialFilterBox.yw) + * 1.0 + * within(COLOR_BY_MEASURE,COLOR_BY_MEASURE_range); + + } + // the primary per-point functions, called directly // + vec4 getClipPosition(){ + return applyCamera(getDataPosition()); + } + float getPointSize(){ + return mix(2.0,6.0,isHovered()); + } + vec4 getColor(){ + + return mix(filteredOutColor, + texture2D(gradient,vec2(rangeParameter(COLOR_BY_MEASURE,COLOR_BY_MEASURE_range),0.5)) + ,isFilteredIn()); + + } + varying vec4 color; + void main(){ + color = getColor(); + gl_PointSize = getPointSize(); + gl_Position = getClipPosition(); + }`; + expect(config).toEqual(expectedConfig); + expect(shaders.vs.replace(/\s/g, '')).toEqual(expectedShader.replace(/\s/g, '')); + expect(columnNameToShaderName).toEqual({ + '123': 'COLOR_BY_MEASURE', + '488I12FURRB8ZY5KJ8TCoordinates': 'position', + }); + }); +}); diff --git a/packages/scatterbrain/src/shader.ts b/packages/scatterbrain/src/shader.ts new file mode 100644 index 00000000..f27a2482 --- /dev/null +++ b/packages/scatterbrain/src/shader.ts @@ -0,0 +1,388 @@ +/** biome-ignore-all lint/style/noUnusedTemplateLiteral: not at all helpful*/ + +import type REGL from 'regl'; +import type { ScatterbrainDataset, SlideviewScatterbrainDataset } from './types'; +import type { Cacheable, CachedVertexBuffer } from '@alleninstitute/vis-core'; +import { Box2D, type box2D, type Interval, type vec2, type vec4 } from '@alleninstitute/vis-geometry'; +import * as lodash from 'lodash'; +const { keys, mapValues, reduce } = lodash; + +// the set of columns and what to do with them can vary +// there might be 3 categorical columns and 2 range columns +// each range column (a vertex attrib) uses a uniform vec2 as its filter range +// due to a variety of limitations in WebGL / GLSL 1 - this is about as general as we can +// get without a great deal of extra performance cost + +// scatterbrain does scatterplot rendering +// its main claim to fame is handling complex filtering +// when generating shaders, most of the variable parts flow through these +// utilities - this type's fields are the names of GLSL functions - the contents +// of the string must be the body of that function. All functions should take NO arguments +// and simply refer to global uniforms / attribs directly. +// as a quick reminder - recursion of any kind is forbidden in GLSL - so be careful +// as it will not be possible to detect this until runtime. +// its not unreasonable for some of these utils to call each other - just be sure to +// avoid recursion! +// note - you dont have to use these! these are just kinda like guide-rails for +// patterns we've seen in our shaders so far! you could easily generate your own +// totally custom shaders! + +type ScatterbrainShaderUtils = { + uniforms: string; // the GLSL declarations of the uniforms for this shader + attributes: string; // the GLSL declarations of the vertex attributes for this shader + commonUtilsGLSL: string; // prepend any GLSL to the final vertex shader + isFilteredIn: string; // ()->float + isHovered: string; // ()->float + getColor: string; // ()-> vec4 + getDataPosition: string; // ()-> vec3 // the position of the point in data-space + getClipPosition: string; // ()-> vec4 // the position of the point in clip space - (hint - apply the camera to data-space) + getPointSize: string; // ()->float +}; +export class VBO implements Cacheable { + buffer: CachedVertexBuffer; + constructor(buffer: CachedVertexBuffer) { + this.buffer = buffer; + } + destroy() { + this.buffer.buffer.destroy(); + } + sizeInBytes() { + return this.buffer.bytes; + } +} + +function buildVertexShader(utils: ScatterbrainShaderUtils) { + return /*glsl*/ ` + precision highp float; + // attribs // + ${utils.attributes} + // uniforms // + ${utils.uniforms} + + // utility functions // + ${utils.commonUtilsGLSL} + + // per-point interface functions // + + float isHovered(){ + ${utils.isHovered} + } + vec3 getDataPosition(){ + ${utils.getDataPosition} + } + float isFilteredIn(){ + ${utils.isFilteredIn} + } + // the primary per-point functions, called directly // + vec4 getClipPosition(){ + ${utils.getClipPosition} + } + float getPointSize(){ + ${utils.getPointSize} + } + vec4 getColor(){ + ${utils.getColor} + } + varying vec4 color; + void main(){ + color = getColor(); + gl_PointSize = getPointSize(); + gl_Position = getClipPosition(); + } + `; + // note that only getColor, getPointSize, and getPosition are called + // that should make clear that the other fns are indended to simply be useful + // concepts that the other main fns can call +} +export function buildShaders(config: Config) { + return { + vs: buildVertexShader(generate(config)), + fs: /*glsl*/ ` + precision highp float; + varying vec4 color; + void main(){ + gl_FragColor = color; + } + `, + }; +} +export type Config = { + mode: 'color' | 'info'; + quantitativeColumns: string[]; + categoricalColumns: string[]; + categoricalTable: string; + tableSize: vec2; + gradientTable: string; + positionColumn: string; + colorByColumn: string; +}; +function rangeFor(col: string) { + return `${col}_range`; +} +export type RenderProps = { + target: REGL.Framebuffer2D | null; + categoricalLookupTable: REGL.Texture2D; + gradient: REGL.Texture2D; + camera: { view: box2D; screenResolution: vec2 }; + offset: vec2; + filteredOutColor: vec4; + spatialFilterBox: box2D; + quantitativeRangeFilters: Record; + hoveredValue: number; + item: { + count: number; + columnData: Record; + }; +}; +export function buildScatterbrainRenderCommand(config: Config, regl: REGL.Regl) { + const prop = (p: string) => regl.prop(p); + const { quantitativeColumns, categoricalColumns, categoricalTable, gradientTable, positionColumn } = config; + const ranges = reduce( + quantitativeColumns, + (unis, col) => ({ ...unis, [rangeFor(col)]: prop(rangeFor(col)) }), + {} as Record>, + ); + const { vs, fs } = buildShaders(config); + const uniforms = { + [categoricalTable]: prop('categoricalLookupTable'), + [gradientTable]: prop('gradient'), + ...ranges, + spatialFilterBox: prop('spatialFilterBox'), + filteredOutColor: prop('filteredOutColor'), + view: prop('view'), + hoveredValue: prop('hoveredValue'), + screenSize: prop('screenSize'), + offset: prop('offset'), + }; + const cmd = regl({ + vert: vs, + frag: fs, + attributes: [positionColumn, ...categoricalColumns, ...quantitativeColumns].reduce( + (attribs, col) => ({ ...attribs, [col]: regl.prop(col) }), + {}, + ), + uniforms, + blend: { + enable: false, + }, + primitive: 'points', + framebuffer: prop('target'), + count: prop('count'), + }); + + return (props: RenderProps) => { + const { + target, + hoveredValue, + spatialFilterBox, + filteredOutColor, + gradient, + camera, + offset, + quantitativeRangeFilters, + categoricalLookupTable, + item, + } = props; + const filterRanges = reduce( + keys(quantitativeRangeFilters), + (acc, cur) => ({ ...acc, [rangeFor(cur)]: quantitativeRangeFilters[cur] }), + {}, + ); + const { view, screenResolution } = camera; + const { count, columnData } = item; + const rawBuffers = mapValues(columnData, (vbo) => vbo.buffer.buffer); + cmd({ + target, + gradient, + hoveredValue, + filteredOutColor, + spatialFilterBox: Box2D.toFlatArray(spatialFilterBox), + categoricalLookupTable, + offset, + count, + view: Box2D.toFlatArray(view), + screenSize: screenResolution, + ...filterRanges, + ...rawBuffers, + }); + }; +} + +function rangeFilterExpression(quantitativeColumns: readonly string[]) { + return quantitativeColumns.map((attrib) => /*glsl*/ `within(${attrib},${rangeFor(attrib)})`).join(' * '); +} +function categoricalFilterExpression(categoricalColumns: readonly string[], tableSize: vec2, tableName: string) { + // categorical columns are in order - this array will have the same order as the col in the texture + const [w, h] = tableSize; + return categoricalColumns + .map( + (attrib, i) => + /*glsl*/ `step(0.01,texture2D(${tableName},vec2(${i.toFixed(0)}.5,${attrib}+0.5)/vec2(${w.toFixed(1)},${h.toFixed(1)})).a)`, + ) + .join(' * '); +} + +export function generate(config: Config): ScatterbrainShaderUtils { + const { + mode, + quantitativeColumns, + categoricalColumns, + categoricalTable, + tableSize, + gradientTable, + positionColumn, + colorByColumn, + } = config; + const catFilter = categoricalFilterExpression(categoricalColumns, tableSize, categoricalTable); + const rangeFilter = rangeFilterExpression(quantitativeColumns); + const uniforms = /*glsl*/ ` + uniform vec4 view; + uniform vec2 screenSize; + uniform vec2 offset; + uniform vec4 spatialFilterBox; + uniform vec4 filteredOutColor; + uniform float hoveredValue; + + uniform sampler2D ${gradientTable}; + uniform sampler2D ${categoricalTable}; + // quantitative columns each need a range value - its the min,max in a vec2 + ${quantitativeColumns.map((col) => /*glsl*/ `uniform vec2 ${rangeFor(col)};`).join('\n')} + `; + + const attributes = /*glsl*/ ` + attribute vec2 ${positionColumn}; + ${categoricalColumns.map((col) => /*glsl*/ `attribute float ${col};`).join('\n')} + ${quantitativeColumns.map((col) => /*glsl*/ `attribute float ${col};`).join('\n')} + `; + + const commonUtilsGLSL = /*glsl*/ ` + vec4 applyCamera(vec3 dataPos){ + vec2 size = view.zw-view.xy; + vec2 unit = (dataPos.xy-view.xy)/size; + return vec4((unit*2.0)-1.0,0.0,1.0); + } + float rangeParameter(float v, vec2 range){ + return (v-range.x)/(range.y-range.x); + } + float within(float v, vec2 range){ + return step(range.x,v)*step(v,range.y); + } + `; + const categoryColumnIndex = categoricalColumns.indexOf(colorByColumn); + const isCategoricalColor = categoryColumnIndex > -1; + const hoverCategoryExpr = /*glsl*/ `1.0-step(0.1,abs(${colorByColumn}-hoveredValue))`; + const isHovered = /*glsl*/ ` + return ${isCategoricalColor ? hoverCategoryExpr : '0.0'};`; + const isFilteredIn = /*glsl*/ ` + vec3 p = getDataPosition(); + return within(p.x,spatialFilterBox.xz)*within(p.y,spatialFilterBox.yw) + * ${catFilter.length > 0 ? catFilter : '1.0'} + * ${rangeFilter.length > 0 ? rangeFilter : '1.0'}; + `; + + const getDataPosition = /*glsl*/ `return vec3(${positionColumn}+offset,0.0);`; + const getClipPosition = /*glsl*/ `return applyCamera(getDataPosition());`; + const getPointSize = /*glsl*/ `return mix(2.0,6.0,isHovered());`; // todo! + // todo - use config options! + // if the colorByColumn is a categorical column, generate that + // else, use a range-colorby + const [w, h] = tableSize; + const colorByCategorical = /*glsl*/ ` + vec4(texture2D(${categoricalTable},vec2(${categoryColumnIndex.toFixed(0)}.5,${colorByColumn}+0.5)/vec2(${w.toFixed(1)},${h.toFixed(1)})).rgb,1.0)`; + + const colorByQuantitative = /*glsl*/ ` + texture2D(${gradientTable},vec2(rangeParameter(${colorByColumn},${rangeFor(colorByColumn)}),0.5)) + `; + const colorize = categoryColumnIndex !== -1 ? colorByCategorical : colorByQuantitative; + + const colorByCategoricalId = /*glsl*/ ` + float G = mod(${colorByColumn},256.0); + float R = mod(${colorByColumn}/256.0,256.0); + return vec4(R/255.0,G/255.0,0,1); + `; + const colorByQuantitativeValue = /*glsl*/ ` + return vec4(0,rangeParameter(${colorByColumn},${rangeFor(colorByColumn)}),0,1); + `; + const getColor = + mode === 'color' + ? /*glsl*/ ` + return mix(filteredOutColor,${colorize},isFilteredIn()); + ` + : categoryColumnIndex === -1 + ? colorByQuantitativeValue + : colorByCategoricalId; + return { + attributes, + uniforms, + commonUtilsGLSL, + getClipPosition, + getColor, + getDataPosition, + getPointSize, + isFilteredIn, + isHovered, + }; +} + +// these settings impact how the shader is generated - +// that means changing them may require re-building the renderer (and the shader beneath it) +export type ShaderSettings = { + dataset: ScatterbrainDataset | SlideviewScatterbrainDataset; + categoricalFilters: Record; // category name -> maximum # of distinct values in that category + quantitativeFilters: readonly string[]; // the names of quantitative variables + mode: 'color' | 'info'; + colorBy: + | { kind: 'metadata'; column: string } + | { kind: 'quantitative'; column: string; gradient: 'viridis' | 'inferno'; range: Interval }; +}; + +export function configureShader(settings: ShaderSettings): { + config: Config; + columnNameToShaderName: Record; +} { + // given settings that make sense to a caller (stuff about the data we want to visualize) + // produce an object that can be used to set up some internal config of the shader that would + // do the visualization + const { dataset, categoricalFilters, quantitativeFilters, colorBy, mode } = settings; + // figure out the columns we care about + // assign them names that are safe to use in the shader (A,B,C, whatever) + const categories = keys(categoricalFilters).toSorted(); + const numCategories = categories.length; + const longestCategory = reduce( + keys(categoricalFilters), + (highest, cur) => Math.max(highest, categoricalFilters[cur]), + 0, + ); + + // the goal here is to associate column names with shader-safe names + const initialQuantitativeAttrs: Record = + colorBy.kind === 'metadata' ? {} : { [colorBy.column]: 'COLOR_BY_MEASURE' }; + const initialCategoricalAttrs: Record = + colorBy.kind === 'metadata' ? { [colorBy.column]: 'COLOR_BY_CATEGORY' } : {}; + // we map each quantitative filter name to the shader-safe attribute name: MEASURE_{i} + const qAttrs = reduce( + quantitativeFilters.toSorted(), + (quantAttrs, quantFilter, i) => ({ ...quantAttrs, [quantFilter]: `MEASURE_${i.toFixed(0)}` }), + initialQuantitativeAttrs, + ); + // we map each categorical filter's name to the shader-safe attribute name: CATEGORY_{i} + const cAttrs = reduce( + categories, + (catAttrs, categoricalFilter, i) => ({ ...catAttrs, [categoricalFilter]: `CATEGORY_${i.toFixed(0)}` }), + initialCategoricalAttrs, + ); + + const colToAttribute = { ...qAttrs, ...cAttrs, [dataset.metadata.spatialColumn]: 'position' }; + + const config: Config = { + categoricalColumns: keys(cAttrs).map((columnName) => colToAttribute[columnName]), + quantitativeColumns: keys(qAttrs).map((columnName) => colToAttribute[columnName]), + categoricalTable: 'lookup', + gradientTable: 'gradient', + colorByColumn: colToAttribute[colorBy.column], + mode, + positionColumn: 'position', + tableSize: [Math.max(numCategories, 1), Math.max(1, longestCategory)], + }; + return { config, columnNameToShaderName: colToAttribute }; +} diff --git a/packages/scatterbrain/src/typed-array.ts b/packages/scatterbrain/src/typed-array.ts new file mode 100644 index 00000000..9edb675b --- /dev/null +++ b/packages/scatterbrain/src/typed-array.ts @@ -0,0 +1,89 @@ +// lets help the compiler to know that these two types are related: +export type TaggedFloat32Array = { + type: 'float'; + data: Float32Array; +}; + +export type TaggedUint32Array = { + type: 'uint32'; + data: Uint32Array; +}; +export type TaggedInt32Array = { + type: 'int32'; + data: Int32Array; +}; + +export type TaggedUint16Array = { + type: 'uint16'; + data: Uint16Array; +}; +export type TaggedInt16Array = { + type: 'int16'; + data: Int16Array; +}; + +export type TaggedUint8Array = { + type: 'uint8'; + data: Uint8Array; +}; +export type TaggedInt8Array = { + type: 'int8'; + data: Int8Array; +}; +export type TaggedTypedArray = + | TaggedFloat32Array + | TaggedUint32Array + | TaggedInt32Array + | TaggedUint16Array + | TaggedInt16Array + | TaggedUint8Array + | TaggedInt8Array; +export type WebGLSafeBasicType = TaggedTypedArray['type']; + +export const BufferConstructors = { + uint8: Uint8Array, + uint16: Uint16Array, + uint32: Uint32Array, + int8: Int8Array, + int16: Int16Array, + int32: Int32Array, + float: Float32Array, +} as const; + +const SizeInBytes = { + uint8: 1, + uint16: 2, + uint32: 4, + int8: 1, + int16: 2, + int32: 4, + float: 4, +} as const; +export function sizeInBytes(type: WebGLSafeBasicType) { + return SizeInBytes[type]; +} +export function MakeTaggedBufferView(type: WebGLSafeBasicType, buffer: ArrayBuffer): TaggedTypedArray { + // note that TS is not smart enough to realize the mapping here, so we have 7 identical, spoonfed + // cases.... + switch (type) { + case 'uint8': + return { type, data: new BufferConstructors[type](buffer) }; + case 'int8': + return { type, data: new BufferConstructors[type](buffer) }; + case 'uint16': + return { type, data: new BufferConstructors[type](buffer) }; + case 'int16': + return { type, data: new BufferConstructors[type](buffer) }; + case 'uint32': + return { type, data: new BufferConstructors[type](buffer) }; + case 'int32': + return { type, data: new BufferConstructors[type](buffer) }; + case 'float': + return { type, data: new BufferConstructors[type](buffer) }; + default: { + // will be a compile error if we ever add any basic types + const unreachable: never = type; + throw new Error(`unsupported type requested: ${unreachable}`); + } + } +} diff --git a/packages/scatterbrain/src/types.ts b/packages/scatterbrain/src/types.ts new file mode 100644 index 00000000..50a442ec --- /dev/null +++ b/packages/scatterbrain/src/types.ts @@ -0,0 +1,84 @@ +/// Types describing the metadata that gets loaded from scatterbrain.json files /// +// there are 2 variants, slideview and regular - they are distinguished at runtime +// by checking the parsed metadata for the 'slides' field +export type WebGLSafeBasicType = 'uint8' | 'uint16' | 'int8' | 'int16' | 'uint32' | 'int32' | 'float'; + +export type volumeBound = { + lx: number; + ly: number; + lz: number; + ux: number; + uy: number; + uz: number; +}; +export type PointAttribute = { + name: string; + size: number; // elements * elementSize - todo ask Peter to remove + elements: number; // values per point (so a vector xy would have 2) + elementSize: number; // size of an element, given in bytes (for example float would have 4) + type: WebGLSafeBasicType; + description: string; +}; +export type TreeNode = { + file: string; + numSpecimens: number; + children: undefined | TreeNode[]; +}; + +type MetadataColumn = { + type: 'METADATA'; + name: string; +}; +type QuantitativeColumn = { + type: 'QUANTITATIVE'; + name: string; +}; +export type ColumnRequest = MetadataColumn | QuantitativeColumn; +type CommonMetadata = { + geneFileEndpoint: string; + metadataFileEndpoint: string; + visualizationReferenceId: string; + spatialColumn: string; + pointAttributes: Record; +}; +// scatterbrain distinguishes 2 kinds of datasets - those arranged at the topmost level into slides +// and 'regular' - which is just a simple, 2D point cloud +export type ScatterbrainMetadata = CommonMetadata & { + points: number; + boundingBox: volumeBound; + tightBoundingBox: volumeBound; + root: TreeNode; +}; + +// slideview variant: +type Slide = { + featureTypeValueReferenceId: string; + tree: { + root: TreeNode; + points: number; + boundingBox: volumeBound; + tightBoundingBox: volumeBound; + }; +}; +type SpatialReferenceFrame = { + anatomicalOrigin: string; + direction: string; + unit: string; + minX: number; + maxX: number; + minY: number; + maxY: number; +}; + +export type SlideviewMetadata = CommonMetadata & { + slides: Slide[]; + spatialUnit: SpatialReferenceFrame; +}; + +/// Types related to the rendering of scatterbrain datasets /// + +// a Dataset is the top level entity +// an Item is a chunk of that dataset - esentially a singular, loadable, thing + +export type SlideviewScatterbrainDataset = { type: 'slideview'; metadata: SlideviewMetadata }; +export type ScatterbrainDataset = { type: 'normal'; metadata: ScatterbrainMetadata }; diff --git a/packages/scatterbrain/tsconfig.json b/packages/scatterbrain/tsconfig.json new file mode 100644 index 00000000..d8a6412f --- /dev/null +++ b/packages/scatterbrain/tsconfig.json @@ -0,0 +1,13 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "paths": { + "~/*": ["./*"] + }, + "moduleResolution": "Bundler", + "module": "es6", + "target": "es2024", + "lib": ["es2024", "DOM"] + }, + "include": ["./src/index.ts"] +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 5b1e73ab..65ab413f 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -90,6 +90,31 @@ importers: specifier: 4.3.6 version: 4.3.6 + packages/scatterbrain: + dependencies: + '@alleninstitute/vis-core': + specifier: workspace:* + version: link:../core + '@alleninstitute/vis-geometry': + specifier: workspace:* + version: link:../geometry + lodash: + specifier: 4.17.21 + version: 4.17.21 + regl: + specifier: 2.1.0 + version: 2.1.0 + ts-pattern: + specifier: 5.9.0 + version: 5.9.0 + zod: + specifier: 4.3.6 + version: 4.3.6 + devDependencies: + '@types/lodash': + specifier: 4.17.23 + version: 4.17.23 + site: dependencies: '@alleninstitute/vis-core': @@ -104,6 +129,9 @@ importers: '@alleninstitute/vis-omezarr': specifier: workspace:* version: link:../packages/omezarr + '@alleninstitute/vis-scatterbrain': + specifier: workspace:* + version: link:../packages/scatterbrain '@astrojs/check': specifier: 0.9.6 version: 0.9.6(prettier@3.8.1)(typescript@5.9.3) @@ -3045,6 +3073,9 @@ packages: regex@6.0.1: resolution: {integrity: sha512-uorlqlzAKjKQZ5P+kTJr3eeJGSVroLKoHmquUj4zHWuR+hEyNqlXsSKlYYF5F4NI6nl7tWCs0apKJ0lmfsXAPA==} + regl@2.1.0: + resolution: {integrity: sha512-oWUce/aVoEvW5l2V0LK7O5KJMzUSKeiOwFuJehzpSFd43dO5spP9r+sSUfhKtsky4u6MCqWJaRL+abzExynfTg==} + regl@2.1.1: resolution: {integrity: sha512-+IOGrxl3FZ8ZM9ixCWQZzFRiRn7Rzn9bu3iFHwg/yz4tlOUQgbO4PHLgG+1ZT60zcIV8tief6Qrmyl8qcoJP0g==} @@ -3265,6 +3296,9 @@ packages: trough@2.2.0: resolution: {integrity: sha512-tmMpK00BjZiUyVyvrBK7knerNgmgvcV/KLVyuma/SC+TQN167GrMRciANTz09+k3zW8L8t60jWO1GpfkZdjTaw==} + ts-pattern@5.9.0: + resolution: {integrity: sha512-6s5V71mX8qBUmlgbrfL33xDUwO0fq48rxAu2LBE11WBeGdpCPOsXksQbZJHvHwhrd3QjUusd3mAOM5Gg0mFBLg==} + tsconfck@3.1.6: resolution: {integrity: sha512-ks6Vjr/jEw0P1gmOVwutM3B7fWxoWBL2KRDb1JfqGVawBmO5UsvmWOQFGHBPl5yxYz4eERr19E6L7NMv+Fej4w==} engines: {node: ^18 || >=20} @@ -7330,6 +7364,8 @@ snapshots: dependencies: regex-utilities: 2.3.0 + regl@2.1.0: {} + regl@2.1.1: {} rehype-expressive-code@0.41.3: @@ -7652,6 +7688,8 @@ snapshots: trough@2.2.0: {} + ts-pattern@5.9.0: {} + tsconfck@3.1.6(typescript@5.9.3): optionalDependencies: typescript: 5.9.3 diff --git a/site/package.json b/site/package.json index aeafdd56..308577e8 100644 --- a/site/package.json +++ b/site/package.json @@ -49,6 +49,7 @@ "@alleninstitute/vis-dzi": "workspace:*", "@alleninstitute/vis-geometry": "workspace:*", "@alleninstitute/vis-omezarr": "workspace:*", + "@alleninstitute/vis-scatterbrain": "workspace:*", "@astrojs/check": "0.9.6", "@astrojs/mdx": "4.3.13", "@astrojs/react": "4.4.2", diff --git a/site/src/content/docs/examples/scatterbrain.mdx b/site/src/content/docs/examples/scatterbrain.mdx new file mode 100644 index 00000000..c786aba9 --- /dev/null +++ b/site/src/content/docs/examples/scatterbrain.mdx @@ -0,0 +1,8 @@ +--- +title: Scatterbrain +tableOfContents: false +--- + +import { scatterBrainDemo } from '../../../examples/scatterbrain/demo.tsx'; + + diff --git a/site/src/examples/scatterbrain/demo.tsx b/site/src/examples/scatterbrain/demo.tsx new file mode 100644 index 00000000..f154221c --- /dev/null +++ b/site/src/examples/scatterbrain/demo.tsx @@ -0,0 +1,109 @@ +import type { vec2, vec4 } from '@alleninstitute/vis-geometry'; +import { SharedCacheContext, SharedCacheProvider } from '../common/react/priority-cache-provider'; +import { useContext, useEffect, useRef, useState } from 'react'; +import { + buildScatterbrainRenderFn, + loadScatterbrainDataset, + setCategoricalLookupTableValues, + type Dataset, + type ShaderSettings, +} from '@alleninstitute/vis-scatterbrain'; + +const screenSize: vec2 = [800, 800]; +const tenx = + 'https://bkp-2d-visualizations-stage.s3.amazonaws.com/wmb_tenx_01172024_stage-20240128193624/G4I4GFJXJB9ATZ3PTX1/ScatterBrain.json'; +export function scatterBrainDemo() { + return ( + + + + ); +} + +const makeFakeColors = (n: number) => { + const stuff: Record = {}; + 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; +}; +// fake color and filter tables, as a demo: +const categories = { + '4MV7HA5DG2XJZ3UD8G9': makeFakeColors(40), // nt type + FS00DXV0T9R1X9FJ4QE: makeFakeColors(40), // class +}; +const settings: Omit = { + 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: [], +}; +async function loadRawJson() { + return await (await fetch(tenx)).json(); +} +type Props = { screenSize: vec2 }; +function Demo(props: Props) { + const { screenSize } = props; + const cnvs = useRef(null); + const server = useContext(SharedCacheContext); + const [dataset, setDataset] = useState(undefined); + useEffect(() => { + loadRawJson().then((raw) => setDataset(loadScatterbrainDataset(raw))); + }, []); + // todo handlers, etc + useEffect(() => { + // build the renderer + if (server && dataset && cnvs.current) { + const ctx = cnvs.current.getContext('2d'); + const { cache, regl } = server; + const lookup = regl.texture({ width: 10, height: 10, format: 'rgba' }); + 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 gradient = regl.texture({ width: 256, height: 1, format: 'rgba', data: gradientData }); + const tgt = regl.framebuffer(screenSize[0], screenSize[1]); + // make up random colors for the coloring, and add random filtering + + setCategoricalLookupTableValues(categories, lookup); + + const { render, connectToCache } = buildScatterbrainRenderFn( + // @ts-expect-error we'll deal with this later + regl, + { ...settings, dataset }, + ); + // this ts error is bogus, dont know why + const renderOneFrame = () => { + render({ + client, + visibilityThresholdPx: 10, + camera: { view: { minCorner: [-17, -17], maxCorner: [26, 26] }, screenResolution: [800, 800] }, + categoricalLookupTable: lookup, + dataset, + filteredOutColor: [0, 0, 0, 1], + gradient, + hoveredValue: 22, + offset: [0, 0], + quantitativeRangeFilters: {}, + spatialFilterBox: { minCorner: [-17, -17], maxCorner: [30, 30] }, + target: tgt, + }); + const bytes = regl.read({ framebuffer: tgt }); + const img = new ImageData(new Uint8ClampedArray(bytes), screenSize[0], screenSize[1]); + ctx!.putImageData(img, 0, 0); + }; + const client = connectToCache(cache, renderOneFrame); + renderOneFrame(); + } + }, [dataset, server, screenSize]); + return ; +}