Skip to content

Commit

Permalink
Noah/layers pt2 (#23)
Browse files Browse the repository at this point in the history
* delete ye olde build script

* add some react for UI - think about if I like it or it will make my life nice

* fix a perf bug in which we consider rendering every slice in a grid regardless of zoom level.
some (better?) smoke and mirrors trickery to try to obscure redraw flickering when panning/zooming

* working but messy impl of a worker pool that handles decoding zarr chunks for us

# Conflicts:
#	apps/layers/package.json

* play around with WW pool size, and fix a bug in which the optional render-working budget parameters were being ignored (for the slice renderer specifically)

* move a shocking amount of code in to render slide-view style annotations

# Conflicts:
#	apps/layers/package.json

* good enough for now

# Conflicts:
#	apps/layers/src/demo.ts

* respect queue params

* enable screenshots, with a default output resolution of roughly 85MP

# Conflicts:
#	apps/layers/package.json
#	apps/layers/src/demo.ts
#	pnpm-lock.yaml

* start thinking about upside down data...

* its all upside down now  great

* minor tidy, plus a rough attempt at a less flickery stand-in algorithm

* tools to add layers during demotime

# Conflicts:
#	apps/layers/src/demo.ts

* add a versa layer

* add scatterplot-slideview real quick

# Conflicts:
#	apps/layers/src/demo.ts

* start some cleanup so I can merge this...

# Conflicts:
#	apps/layers/src/demo.ts

* Merge branch 'noah/layered-demo' into noah/layered-with-react-whynot

# Conflicts:
#	apps/layers/src/demo.ts

* quickly change the underlying cache type for scatterplots for much better perf (gpu buffer not client-buffer)

* try out sds components for quick hacky ui fun - delete old ui code

* add a bunch of per-layer ui elements

* prev/next layer buttons

* take a snapshot button

* quickly re-enable drawing layers

* a bit hacky, but non-flickering drawings are worth it for a demo

* change moduleResolution in the apps tsconfig to make the zarr library that we use extensively get resolved correctly. this is an issue on their end: gzuidhof/zarr.js#152

* cleanup some increasingly scary cherrypicks, and finally tidy up those last little demo ts errors.

* clean up a bunch of low hanging fruit

* fix up the scatterplot (standalone) demo

* readme and demo script

* a little more

* copy in the latest and greatest annotation stuff in

* minor cleanups

* fix wrongly named example
  • Loading branch information
froyo-np authored May 13, 2024
1 parent 46a326c commit 031149a
Show file tree
Hide file tree
Showing 44 changed files with 2,604 additions and 1,009 deletions.
File renamed without changes.
30 changes: 30 additions & 0 deletions apps/common/src/loaders/ome-zarr/fetchSlice.worker.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@

// a web-worker which fetches slices of data, decodes them, and returns the result as a flat float32 array, using transferables
import type { NestedArray, TypedArray } from 'zarr'
import { getSlice, type ZarrDataset, type ZarrRequest } from "./zarr-data";

const ctx = self;
type ZarrSliceRequest = {
id: string;
type: 'ZarrSliceRequest'
metadata: ZarrDataset
req: ZarrRequest,
layerIndex: number
}
function isSliceRequest(payload: any): payload is ZarrSliceRequest {
return typeof payload === 'object' && payload['type'] === 'ZarrSliceRequest';
}
ctx.onmessage = (msg: MessageEvent<unknown>) => {
const { data } = msg;
if (isSliceRequest(data)) {
const { metadata, req, layerIndex, id } = data;
getSlice(metadata, req, layerIndex).then((result: {
shape: number[],
buffer: NestedArray<TypedArray>
}) => {
const { shape, buffer } = result;
const R = new Float32Array(buffer.flatten());
ctx.postMessage({ type: 'slice', id, shape, data: R }, { transfer: [R.buffer] })
})
}
}
75 changes: 75 additions & 0 deletions apps/common/src/loaders/ome-zarr/sliceWorkerPool.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
import { uniqueId } from "lodash";
import type { ZarrDataset, ZarrRequest } from "./zarr-data";

type PromisifiedMessage = {
requestCacheKey: string;
resolve: (t: Slice) => void;
reject: (reason: unknown) => void;
promise?: Promise<Slice> | undefined;
};
type ExpectedResultSlice = {
type: 'slice',
id: string;
} & Slice;
type Slice = {
data: Float32Array;
shape: number[]
}
function isExpectedResult(obj: any): obj is ExpectedResultSlice {
return (typeof obj === 'object' && 'type' in obj && obj.type === 'slice')
}
export class SliceWorkerPool {
private workers: Worker[];
private promises: Record<string, PromisifiedMessage>;
private which: number;
constructor(size: number) {
this.workers = new Array(size);
for (let i = 0; i < size; i++) {
this.workers[i] = new Worker(new URL('./fetchSlice.worker.ts', import.meta.url), { type: 'module' });
this.workers[i].onmessage = (msg) => this.handleResponse(msg)
}
this.promises = {};
this.which = 0;
}

handleResponse(msg: MessageEvent<unknown>) {
const { data: payload } = msg;
if (isExpectedResult(payload)) {
const prom = this.promises[payload.id];
if (prom) {
const { data, shape } = payload
prom.resolve({ data, shape });
delete this.promises[payload.id]
}
}
}
private roundRobin() {
this.which = (this.which + 1) % this.workers.length
}
requestSlice(dataset: ZarrDataset, req: ZarrRequest, layerIndex: number) {
const reqId = uniqueId('rq');
const cacheKey = JSON.stringify({ url: dataset.url, req, layerIndex });
// TODO caching I guess...
const eventually = new Promise<Slice>((resolve, reject) => {
this.promises[reqId] = {
requestCacheKey: cacheKey,
resolve,
reject,
promise: undefined, // ill get added to the map once I am fully defined!
};
this.workers[this.which].postMessage({ id: reqId, type: 'ZarrSliceRequest', metadata: dataset, req, layerIndex });
this.roundRobin();
});
this.promises[reqId].promise = eventually;
return eventually;
}
}

// a singleton...
let slicePool: SliceWorkerPool;
export function getSlicePool() {
if (!slicePool) {
slicePool = new SliceWorkerPool(16);
}
return slicePool;
}
73 changes: 0 additions & 73 deletions apps/common/src/loaders/ome-zarr/zarr-data.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,47 +42,6 @@ type ZarrAttrs = {
multiscales: ReadonlyArray<ZarrAttr>;
};

// function getSpatialDimensionShape(dataset: DatasetWithShape, axes: readonly AxisDesc[]) {
// const dims = axes.reduce(
// (shape, ax, i) => (ax.type === "spatial" ? { ...shape, [ax.name]: dataset.shape[i] } : shape),
// {} as Record<string, number>
// );
// return dims;
// }
// function getSpatialOrdering
// function getBoundsInMillimeters(data: ZarrDataset) {
// if (data.multiscales.length !== 1) {
// throw new Error("cant support multi-scene zarr file...");
// }
// const scene = data.multiscales[0];
// const { axes, datasets } = scene;
// if (datasets.length < 1) {
// throw new Error("malformed dataset - no voxels!");
// }
// const dataset = datasets[0];
// const spatialResolution = getSpatialDimensionShape(dataset, axes);
// // apply transforms
// dataset.coordinateTransformations.forEach((trn) => {});
// const dimensions = getNumVoxelsInXYZ(getXYZIndexing(axes), dataset.shape);

// let bounds: box3D = Box3D.create([0, 0, 0], dimensions);
// dataset.coordinateTransformations.forEach((trn) => {
// // specification for coordinate transforms given here: https://ngff.openmicroscopy.org/latest/#trafo-md
// // from the above doc, its not super clear if the given transformation is in the order of the axes metadata (https://ngff.openmicroscopy.org/latest/#axes-md)
// // or some other order
// // all files I've seen so far have both in xyz order, so its a bit ambiguous.
// if (isScaleTransform(trn) && trn.scale.length >= 3) {
// bounds = applyScaleToXYZBounds(bounds, trn, axes);
// } else {
// throw new Error(`unsupported coordinate transformation type - please implement`);
// }
// });
// // finally - convert whatever the axes units are to millimeters, or risk crashing into mars
// // get the units of each axis in xyz order...

// return Box3D.map(bounds, (corner) => unitsToMillimeters(corner, axes));
// }

async function getRawInfo(store: HTTPStore) {
const group = await openGroup(store);
// TODO HACK ALERT: I am once again doing the thing that I hate, in which I promise to my friend Typescript that
Expand Down Expand Up @@ -125,18 +84,6 @@ export function uvForPlane(plane: AxisAlignedPlane) {
export function sliceDimensionForPlane(plane: AxisAlignedPlane) {
return sliceDimension[plane];
}
// function sizeOnScreen(full: box2D, relativeView: box2D, screen: vec2) {
// const pxView = Box2D.scale(relativeView, Box2D.size(full));
// const onScreen = Box2D.intersection(pxView, full);
// if (!onScreen) return [0, 0];

// const effective = Box2D.size(onScreen);
// // as a parameter, how much is on screen?
// const p = Vec2.div(effective, Box2D.size(full));
// const limit = Vec2.mul(p, screen);

// return limit[0] * limit[1] < effective[0] * effective[1] ? limit : effective
// }
export type ZarrRequest = Record<OmeDimension, number | Interval | null>;
export function pickBestScale(
dataset: ZarrDataset,
Expand Down Expand Up @@ -173,8 +120,6 @@ export function pickBestScale(
: bestSoFar,
datasets[0]
);
// console.log('choose layer: ', choice.path);
// console.log('---->', planeSizeInVoxels(plane, axes, choice))
return choice ?? datasets[datasets.length - 1];
}
function indexFor(dim: OmeDimension, axes: readonly AxisDesc[]) {
Expand Down Expand Up @@ -287,24 +232,6 @@ export async function getSlice(metadata: ZarrDataset, r: ZarrRequest, layerIndex
buffer: result,
};
}
// export async function getRGBSlice(metadata: ZarrDataset, r: ZarrRequest, layerIndex: number) {
// dieIfMalformed(r);
// // put the request in native order
// const store = new HTTPStore(metadata.url);
// const scene = metadata.multiscales[0];
// const { axes } = scene;
// const level = scene.datasets[layerIndex] ?? scene.datasets[scene.datasets.length - 1];
// const arr = await openArray({ store, path: level.path, mode: "r" });
// const result = await arr.get(buildQuery(r, axes, level.shape));
// if (typeof result == "number" || result.shape.length !== 2) {
// throw new Error("oh noes, slice came back all weird");
// }
// return {
// shape: result.shape as unknown as vec2,
// buffer: result.flatten(),
// };
// }

export async function load(url: string) {
const store = new HTTPStore(url);
return loadMetadata(store, await getRawInfo(store));
Expand Down
38 changes: 23 additions & 15 deletions apps/common/src/loaders/scatterplot/data.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,18 @@
// todo rename this file

import { Box2D, visitBFS, type box2D, type vec2 } from "@alleninstitute/vis-geometry";
import { fetchColumn, type ColumnarTree, type SlideViewDataset, type loadDataset } from "./scatterbrain-loader";
import { fetchColumn, type ColumnData, type ColumnRequest, type ColumnarNode, type ColumnarTree, type SlideViewDataset, type loadDataset } from "./scatterbrain-loader";
import REGL from 'regl'
export type Dataset = ReturnType<typeof loadDataset>
export type RenderSettings = {
dataset: Dataset;
view: box2D;
colorBy: ColumnRequest;
pointSize: number;
target: REGL.Framebuffer2D | null;
regl: REGL.Regl;
}

function isVisible(view: box2D, sizeLimit: number, tree: ColumnarTree<vec2>) {
const { bounds } = tree.content;
return Box2D.size(bounds)[0] > sizeLimit && !!Box2D.intersection(view, tree.content.bounds);
Expand All @@ -23,12 +27,12 @@ export function getVisibleItems(dataset: Dataset, view: box2D, sizeLimit: number
(tree) => isVisible(view, sizeLimit, tree));
return hits;
}
export function getVisibleItemsInSlide(dataset: SlideViewDataset, slide: string, view: box2D, sizeLimit: number){
const theSlide = dataset.slides[slide];
if(!theSlide){
export function getVisibleItemsInSlide(dataset: SlideViewDataset, slide: string, view: box2D, sizeLimit: number) {
const theSlide = dataset.slides[slide];
if (!theSlide) {
console.log('nope', Object.keys(dataset.slides))
return []
}
}

const hits: ColumnarTree<vec2>[] = []
const tree = theSlide.tree;
Expand All @@ -38,19 +42,23 @@ export function getVisibleItemsInSlide(dataset: SlideViewDataset, slide: string,
(tree) => isVisible(view, sizeLimit, tree));
return hits;
}
function toReglBuffer(c: ColumnData, regl: REGL.Regl) {
return {
type: 'vbo',
data: regl.buffer(c)
} as const;
}
function fetchAndUpload(settings: { dataset: Dataset, regl: REGL.Regl }, node: ColumnarNode<vec2>, req: ColumnRequest, signal?: AbortSignal | undefined) {
const { dataset, regl } = settings;
return fetchColumn(node, dataset, req, signal).then(cd => toReglBuffer(cd, regl))
}
export function fetchItem(item: ColumnarTree<vec2>, settings: RenderSettings, signal?: AbortSignal) {
const { dataset } = settings;
const position = () => fetchColumn(item.content, settings.dataset, {
name: dataset.spatialColumn,
type: 'METADATA',
}, signal);
const color = () => fetchColumn(item.content, settings.dataset, {
type: 'QUANTITATIVE',
name: '442'
}, signal);
const { dataset, colorBy } = settings;
const position = () => fetchAndUpload(settings, item.content, { type: 'METADATA', name: dataset.spatialColumn }, signal);
const color = () => fetchAndUpload(settings, item.content, colorBy, signal)
return {
position,
color
}
} as const
}

9 changes: 9 additions & 0 deletions apps/common/src/loaders/scatterplot/scatterbrain-loader.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { Box2D, Box3D, Vec3, type box, type box3D, type vec2, type vec3 } from "@alleninstitute/vis-geometry";
import { MakeTaggedBufferView, type TaggedTypedArray, type WebGLSafeBasicType } from "../../typed-array";
import type REGL from "regl";

type volumeBound = {
lx: number;
Expand Down Expand Up @@ -244,9 +245,17 @@ type QuantitativeColumn = {
name: string;
};
export type ColumnRequest = MetadataColumn | QuantitativeColumn;
export type ColumnBuffer = {
type: 'vbo',
data: REGL.Buffer
}
export type ColumnData = TaggedTypedArray & {
elements: number; // per vector entry - for example 'xy' would have elements: 2
};
export async function loadScatterbrainJson(url: string) {
// obviously, we should check or something
return fetch(url).then(stuff => stuff.json() as unknown as ColumnarMetadata)
}

export async function fetchColumn(
node: ColumnarNode<ReadonlyArray<number>>,
Expand Down
26 changes: 18 additions & 8 deletions apps/layers/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -30,28 +30,38 @@
"private": true,
"scripts": {
"preinstall": "npx only-allow pnpm",
"demo": "esbuild src/demo.ts --bundle --outfile=dst/demo.js && cp public/* dst/",
"typecheck": "tsc --noEmit",
"start": "parcel public/demo.html"
},
"devDependencies": {
"esbuild": "^0.19.12",
"@parcel/packager-ts": "^2.12.0",
"@parcel/transformer-typescript-types": "^2.12.0",
"@types/file-saver": "^2.0.7",
"@types/react": "^18.3.0",
"@types/react-dom": "^18.3.0",
"esbuild": "^0.19.12",
"parcel": "2.12.0",
"typescript": "^5.3.3"
},
"dependencies": {
"@alleninstitute/vis-geometry": "workspace:*",
"@alleninstitute/vis-scatterbrain": "workspace:*",
"@thi.ng/imgui": "^2.2.54",
"@thi.ng/layout": "^3.0.36",
"@thi.ng/rdom-canvas": "^0.5.83",
"@thi.ng/rstream": "^8.4.0",
"@thi.ng/rstream-gestures": "^5.0.70",
"@czi-sds/components": "^20.0.1",
"@emotion/css": "^11.11.2",
"@emotion/react": "^11.11.4",
"@emotion/styled": "^11.11.5",
"@mui/base": "5.0.0-beta.40",
"@mui/icons-material": "^5.15.15",
"@mui/lab": "5.0.0-alpha.170",
"@mui/material": "^5.15.15",
"@types/lodash": "^4.14.202",
"file-saver": "^2.0.5",
"json5": "^2.2.3",
"kiwi-schema": "^0.5.0",
"lodash": "^4.17.21",
"regl": "^2.1.0"
"regl": "^2.1.0",
"react": "^18.3.0",
"react-dom": "^18.3.0",
"zarr": "^0.6.2"
}
}
44 changes: 44 additions & 0 deletions apps/layers/readme.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
# Layered Rendering Example App
## Run it
1. download/clone this repository.
2. in the root directory (vis/) run `pnpm install`
3. run `pnpm build`
4. `cd apps/layers`
5. `pnpm run start`
6. navigate to the running app (default `localhost://1234`)

## Why?
the goal of this (rather complicated) example app is not to show off a cool app - rather its goal is to show that we can build complexity by composing simple, focused modules. As we (the AIBS Apps team) have developed ABC-Atlas, we've tried to make sure our visualization code stays general, and that each part does as little as possible. The result is that it was fairly easy to combine those components into this new app, which mixes a (terrible) UI, scatter-plot rendering, polygon-mesh rendering (for annotations) and multi-channel volumetric rendering into independent layers. Although each of these data types appear different, our caching, fetching, visibility determination, and render-scheduling code is the same regardless of the datatype to be rendered. All that is required is to fill in a simple interface, and provide a low-level "renderer" plugin in order to add a new "layer" type.
## Demo Script
### Programmatic Configuration
After starting the app in a browser, you'll be greeted by a blank screen. We're going to demonstrate programmatic access to the features of this demo. The goal here is not to make users invoke command-line arguments, but rather just an easy way for interested parties to "peak under the hood". All the visualizations are configured here via simple json objects - it would not be a stretch to read these configuration options at initialization-time via URL parameters for example.

Open your browser's developer tools via ctrl+shift+i, ctrl+opt+i, or F12, and go to the developer console tab. This console is a running REPL environment with direct access to the demo. You can explore what features are available by typing "demo." at which point you should be given some auto-complete options, corresponding to the methods available on the demo object. The first thing we should take a look at is `demo.addLayer`, which accepts a configuration object for some data to be viewed as a layer in the app. A list of pre-constructed objects is provided in the `examples` object. lets run `examples.reconstructed` to take a peek:<BR>
```{
"type": "ScatterPlotGridConfig",
"colorBy": {
"name": "88",
"type": "QUANTITATIVE"
},
"url": "https://.../4STCSZBXHYOI0JUUA3M/ScatterBrain.json"
}
```
note - the above config has been slightly altered here for readability!

Some things to observe about this config object: <BR>
1. its simple json
2. the "type" tag tells the app what fields to expect next.
3. this data is a grid of scatter-plots (slide view!), its colored by a quantitative feature (gene #88).
4. lastly - do not take this structure too seriously! I made up all the layer-config types in a hurry, and we can easily make new ones, compose existing ones, or change anything at all!

### Explore a few layers
Lets continue by adding one layer:
`demo.addLayer(examples.reconstructed)`<BR>
now, we should be able to navigate the (previously blank) view with the mouse, and we can watch the data load as it becomes visible (and thus, worth fetching & caching). Note also that the UI populates, with very simple (ugly?) controls to change settings, in this case only the gene by which to color can be edited. This ui took about 60 seconds to write, and I would not hesitate to throw it away.
Lets add a second layer: `demo.addLayer(examples.structureAnnotation)`<BR>
This data is the CCF annotations (the regions to which the scatterplot data is registered), and should line up with the scatterplot grid. This layer type has only a few simple options for altering the opacity of the polygons.<BR>
Lets add another layer: `demo.addLayer(examples.tissuecyte396)`<BR>
This should load up a contact-sheet style view, with 142 slices. You should note that the slices are not expected to line up with those in the other two layers (the slide-view scatterplot) - the data in each are in different coordinate systems. Perhaps obviously, the whole point of having multiple layers in one view would be to observe interesting relationships between them, and in a (non-demo) app, it would be important to allow configurations to specify mappings / transforms / etc. to get things to "line up". You should also see some double-post, "invlerp" sliders that control how the raw channels of information get mapped to visible-light color. You should be able to move the min & max for each channel, although do note that there seems to be very little data in the 3rd channel, mapped to the color blue.

### Take a picture
At the top left of the UI, you might notice a little 📸 button. This will take a low-resolution snapshot of the current view. If you'd like to be a little patient, enter the dev console again, and run `demo.requestSnapshot(10000)` which, depending on the aspect ratio of your browser window, will (slowly!) capture an ~80 million pixel image of the current screen - the main reason its slow is because, to properly serve that higher resolution, the rendering system must fetch many, many chunks of high-resolution data to render. Note that making changes to the UI during the snapshot process may cause those changes to take partial effect in the final snapshot output. This is a well understood class of bugs that are easy to address when planning to "do software correctly" - however this is a quick demo so I took a bunch of shortcuts!
Loading

0 comments on commit 031149a

Please sign in to comment.