feat: Add a transfer function shader widget (#582)

* feat: transfer function widget

* refactor: fix linting warnings

* refactor: TF uses uint64 as number, not string

* chore: cleanup TODO and comments

* refactor: clean up transfer function code a little

* fix: TF control points didn't always make it to JSON

* refactor: clearer comments and update loop

* fix: no longer able to place control points on top of eachother

* fix: can grab control points in TF that are close in X by breaking ties with Y

* fix: bind remove TF point to shift+dblclick
You could accidentally remove points trying to move them before

* feat: clearer name of TF input value

* feat: Python control over transfer function shader contro

* docs: fix typo in python docs

* test (Python): shader control transfer function test

* fix: user can't specify transfer function points outside input range

* docs: transfer function UI control

* docs: code comment on transfer functions

* chore: format and lint

* chore(python): format

* refactor(in progress): store control points in abs value

* refactor(progress): transfer functino

* progress(refactor): tf refactor

* progress(refactor): tf refactor

* progress(refactor): fix all type errors in tf file

* progress(refactor): fix type errors

* refactor(progress): fix more type errors

* fix(tests): remove unused imports

* tests(fix): browser test whole TF

* fix: transfer function correct interpolation

* fix: dynamic transfer function size and correct GPU control

* feat: remove range and size as TF params
Instead they are computed based on input

* fix: parse shader directive for new TF

* fix: JSON state parsing and saving for new TF

* fix: new TF runs, but UI control is broken

* fix: remove accidental log

* fix: control points display in correct position in UI

* fix: find control point near cursor in new TF

* fix: moving control points and setting color

* fix: correct number of lines

* fix: display UI panel texture for TF

* fix: render data from TF texture

* fix: remove fixed size TF texture

* fix: link UI to JSON for control points

* fix: remove temps and TODOs

* fix: handle control point range for 0 and 1 points

* fix: test

* fix: unused code

* fix: can no longer lose control of point that you were trying to move

* fix: compute range after removing a point

* refactor: clearer range update

* fix: tf UI has correct texture indicator

* feat: default intensity for transfer functions

* fix: don't crash on window[0] === window[1] in TF UI panel

* fix: userIntensity always overwrites default

* docs: update transfer function docs

* tests: fix a test for TFs with uint64 data

* feat: use non-interpolated value in TF for consistency and efficiency

* fix: tests

* Python(fix): fix transfer function control input

* docs: fix formatting

* Python: format

* fix: remove accidental test change

* refactor: clarifications

* docs: while inverted windows are not supported, remove from docs

* fix: correctly draw lines with a points beside window, one left, one right

* feat: a little bit cleaner interaction with TF UI window

* refactor: clarify the transfer function lines drawing

python/README.md

@@ -152,8 +152,8 @@ headlessly on Firefox using `xvfb-run`.  On other platforms, tests can't be run
 
 ```shell
 # For headless using Firefox on xvfb (Linux only)
-sudo apt-get instrall xvfb # On Debian-based systems
-tox -e firefox-xvfb  # Run tests using non-headless Firefox
+sudo apt-get install xvfb # On Debian-based systems
+tox -e firefox-xvfb  # Run tests using headless Firefox
 
 # For non-headless using Chrome
 tox -e chrome

python/neuroglancer/__init__.py

@@ -89,6 +89,7 @@
     LayerDataSource,  # noqa: F401
     LayerDataSources,  # noqa: F401
     InvlerpParameters,  # noqa: F401
+    TransferFunctionParameters,  # noqa: F401
     ImageLayer,  # noqa: F401
     SkeletonRenderingOptions,  # noqa: F401
     StarredSegments,  # noqa: F401

python/neuroglancer/viewer_state.py

@@ -518,6 +518,16 @@ class InvlerpParameters(JsonObjectWrapper):
     channel = wrapped_property("channel", optional(typed_list(int)))
 
 
+@export
+class TransferFunctionParameters(JsonObjectWrapper):
+    window = wrapped_property("window", optional(array_wrapper(numbers.Number, 2)))
+    channel = wrapped_property("channel", optional(typed_list(int)))
+    controlPoints = wrapped_property(
+        "controlPoints", optional(typed_list(typed_list(number_or_string)))
+    )
+    defaultColor = wrapped_property("defaultColor", optional(str))
+
+
 _UINT64_STR_PATTERN = re.compile("[0-9]+")
 
 
@@ -530,9 +540,13 @@ def _shader_control_parameters(v, _readonly=False):
     if isinstance(v, numbers.Number):
         return v
     if isinstance(v, dict):
+        if "controlPoints" in v:
+            return TransferFunctionParameters(v, _readonly=_readonly)
         return InvlerpParameters(v, _readonly=_readonly)
     if isinstance(v, InvlerpParameters):
         return v
+    if isinstance(v, TransferFunctionParameters):
+        return v
     raise TypeError(f"Unexpected shader control parameters type: {type(v)}")
 
 

python/tests/shader_controls_test.py

@@ -62,6 +62,81 @@ def expect_color(color):
     expect_color([0, 0, 0, 255])
 
 
+def test_transfer_function(webdriver):
+    shader = """
+#uicontrol transferFunction colormap
+void main() {
+    emitRGBA(colormap());
+}
+"""
+    shaderControls = {
+        "colormap": {
+            "controlPoints": [[0, "#000000", 0.0], [84, "#ffffff", 1.0]],
+            "window": [0, 50],
+            "channel": [],
+            "defaultColor": "#ff00ff",
+        }
+    }
+    with webdriver.viewer.txn() as s:
+        s.dimensions = neuroglancer.CoordinateSpace(
+            names=["x", "y"], units="nm", scales=[1, 1]
+        )
+        s.position = [0.5, 0.5]
+        s.layers.append(
+            name="image",
+            layer=neuroglancer.ImageLayer(
+                source=neuroglancer.LocalVolume(
+                    dimensions=s.dimensions,
+                    data=np.full(shape=(1, 1), dtype=np.uint64, fill_value=63),
+                ),
+            ),
+            visible=True,
+            shader=shader,
+            shader_controls=shaderControls,
+            opacity=1.0,
+            blend="additive",
+        )
+        s.layout = "xy"
+        s.cross_section_scale = 1e-6
+        s.show_axis_lines = False
+    control = webdriver.viewer.state.layers["image"].shader_controls["colormap"]
+    assert isinstance(control, neuroglancer.TransferFunctionParameters)
+    np.testing.assert_equal(control.window, [0, 50])
+    assert control.defaultColor == "#ff00ff"
+
+    def expect_color(color):
+        webdriver.sync()
+        screenshot = webdriver.viewer.screenshot(size=[10, 10]).screenshot
+        np.testing.assert_array_equal(
+            screenshot.image_pixels,
+            np.tile(np.array(color, dtype=np.uint8), (10, 10, 1)),
+        )
+
+    # Ensure that the value 63 is mapped to the expected color.
+    # The value 63 is 3/4 of the way between 0 and 84, so the expected color
+    # is 3/4 of the way between black and white.
+    # Additionally, the opacity is 0.75, and the mode is additive, so the
+    # the final color is 0.75 * 0.75 * 255.
+    mapped_opacity = 0.75
+    mapped_color = 0.75 * 255
+    mapped_value = int(mapped_color * mapped_opacity)
+    expected_color = [mapped_value] * 3 + [255]
+    expect_color(expected_color)
+    with webdriver.viewer.txn() as s:
+        s.layers["image"].shader_controls = {
+            "colormap": neuroglancer.TransferFunctionParameters(
+                controlPoints=[[0, "#000000", 0.0], [84, "#ffffff", 1.0]],
+                window=[500, 5000],
+                channel=[],
+                defaultColor="#ff0000",
+            )
+        }
+    control = webdriver.viewer.state.layers["image"].shader_controls["colormap"]
+    np.testing.assert_equal(control.window, [500, 5000])
+    assert control.defaultColor == "#ff0000"
+    expect_color(expected_color)
+
+
 def test_slider(webdriver):
     with webdriver.viewer.txn() as s:
         s.dimensions = neuroglancer.CoordinateSpace(

src/sliceview/image_layer_rendering.md

@@ -155,6 +155,42 @@ annotation layers). The one-parameter overload simply computes the inverse linea
 the specified value within the range specified by the control. The zero-parameter overload returns
 the inverse linear interpolation of the data value for configured channel/property.
 
+### `transferFunction` controls
+
+The `transferFunction` control type allows the user to specify a function which maps
+each data value in a numerical interval to an output color and opacity. The mapping function
+is defined by a series of control points. Each control point is a color and opacity value at a
+specific data input value. In between control points, the color and opacity is linearly interpolated.
+Any input data before the first control point is mapped to a completely transparent output.
+Any data point after the last control point is mapped to the same output as the last control point.
+
+Directive syntax:
+
+```glsl
+#uicontrol transferFunction <name>(window=[lower, higher], controlPoints=[[input, hexColorString, opacity]], channel=[], defaultColor="#rrggbb")
+#uicontrol transferFunction colormap(window=[0, 100], controlPoints=[[10.0, "#000000", 0.0], [100.0, "#ffffff", 1.0]], channel=[], defaultColor="#00ffaa")
+```
+
+The following parameters are supported:
+
+- `window`: Optional. The portion of the input range to view the transfer function over.
+  Must be specified as an array. May be overridden using the UI control. Defaults to the min and max
+  of the control point input values, if control points are specified, or otherwise to the full range of the
+  data type for integer data types, and `[0, 1]` for float32. It is not valid to specify an
+  inverted interval like `[50, 20]`, or an interval where the start and end points are the same, e.g. `[20, 20]`.
+
+- `controlPoints`: Optional. The points which define the input to output mapping.
+  Must be specified as an array, with each value in the array of the form `[inputValue, hexStringColor, floatOpacity]`.
+  The default transfer function is a simple interpolation from transparent black to fully opaque white.
+
+- `channel`: Optional. The channel to perform the mapping on.
+  If the rank of the channel coordinate space is 1, may be specified as a single number,
+  e.g. `channel=2`. Otherwise, must be specified as an array, e.g. `channel=[2, 3]`. May be
+  overriden using the UI control. If not specified, defaults to all-zero channel coordinates.
+
+- `defaultColor`: Optional. The default color for new control points added via the UI control.
+  Defaults to `#ffffff`, and must be specified as a hex string if provided `#rrggbb`.
+
 ## API
 
 ### Retrieving voxel channel value

src/sliceview/volume/renderlayer.ts

@@ -582,6 +582,7 @@ void main() {
 
     const endShader = () => {
       if (shader === null) return;
+      shader.unbindTransferFunctionTextures();
       if (prevChunkFormat !== null) {
         prevChunkFormat!.endDrawing(gl, shader);
       }

src/util/array.spec.ts

@@ -22,6 +22,7 @@ import {
   spliceArray,
   tile2dArray,
   transposeArray2d,
+  findClosestMatchInSortedArray,
 } from "#src/util/array.js";
 
 describe("partitionArray", () => {
@@ -206,3 +207,24 @@ describe("getMergeSplices", () => {
     ]);
   });
 });
+
+describe("findClosestMatchInSortedArray", () => {
+  const compare = (a: number, b: number) => a - b;
+  it("works for empty array", () => {
+    expect(findClosestMatchInSortedArray([], 0, compare)).toEqual(-1);
+  });
+  it("works for simple examples", () => {
+    expect(findClosestMatchInSortedArray([0, 1, 2, 3], 0, compare)).toEqual(0);
+    expect(findClosestMatchInSortedArray([0, 1, 2, 3], 1, compare)).toEqual(1);
+    expect(findClosestMatchInSortedArray([0, 1, 2, 3], 2, compare)).toEqual(2);
+    expect(findClosestMatchInSortedArray([0, 1, 2, 3], 3, compare)).toEqual(3);
+    expect(findClosestMatchInSortedArray([0, 1, 2, 3], 4, compare)).toEqual(3);
+    expect(findClosestMatchInSortedArray([0, 1, 2, 3], -1, compare)).toEqual(0);
+    expect(findClosestMatchInSortedArray([0, 1, 2, 3], 1.5, compare)).toEqual(
+      1,
+    );
+    expect(findClosestMatchInSortedArray([0, 1, 2, 3], 1.6, compare)).toEqual(
+      2,
+    );
+  });
+});

src/util/array.ts

@@ -184,6 +184,39 @@ export function binarySearch<T>(
   return ~low;
 }
 
+/**
+ * Returns the index of the element in `haystack` that is closest to `needle`, according to
+ * `compare`.  If there are multiple elements that are equally close, the index of the first such
+ * element encountered is returned.  If `haystack` is empty, returns -1.
+ */
+export function findClosestMatchInSortedArray<T>(
+  haystack: ArrayLike<T>,
+  needle: T,
+  compare: (a: T, b: T) => number,
+  low = 0,
+  high = haystack.length,
+): number {
+  let bestIndex = -1;
+  let bestDistance = Infinity;
+  while (low < high) {
+    const mid = (low + high - 1) >> 1;
+    const compareResult = compare(needle, haystack[mid]);
+    if (compareResult > 0) {
+      low = mid + 1;
+    } else if (compareResult < 0) {
+      high = mid;
+    } else {
+      return mid;
+    }
+    const distance = Math.abs(compareResult);
+    if (distance < bestDistance) {
+      bestDistance = distance;
+      bestIndex = mid;
+    }
+  }
+  return bestIndex;
+}
+
 /**
  * Returns the first index in `[begin, end)` for which `predicate` is `true`, or returns `end` if no
  * such index exists.

src/volume_rendering/volume_render_layer.ts

@@ -570,6 +570,7 @@ void main() {
 
     const endShader = () => {
       if (shader === null) return;
+      shader.unbindTransferFunctionTextures();
       if (prevChunkFormat !== null) {
         prevChunkFormat!.endDrawing(gl, shader);
       }

src/webgl/rectangle_grid_buffer.spec.ts

@@ -0,0 +1,41 @@
+/**
+ * @license
+ * Copyright 2024 Google Inc.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { describe, it, expect } from "vitest";
+import { createGriddedRectangleArray } from "#src/webgl/rectangle_grid_buffer.js";
+
+describe("createGriddedRectangleArray", () => {
+  it("creates a set of two squares for grid size=2 and rectangle width&height=2", () => {
+    const result = createGriddedRectangleArray(2, -1, 1, 1, -1);
+    expect(result).toEqual(
+      new Float32Array([
+        -1, 1, 0, 1, 0, -1 /* triangle in top right for first grid */, -1, 1, 0,
+        -1, -1, -1 /* triangle in bottom left for first grid */, 0, 1, 1, 1, 1,
+        -1 /* triangle in top right for second grid */, 0, 1, 1, -1, 0,
+        -1 /* triangle in bottom left for second grid */,
+      ]),
+    );
+    const resultReverse = createGriddedRectangleArray(2, 1, -1, -1, 1);
+    expect(resultReverse).toEqual(
+      new Float32Array([
+        1, -1, 0, -1, 0, 1 /* triangle in top right for first grid */, 1, -1, 0,
+        1, 1, 1 /* triangle in bottom left for first grid */, 0, -1, -1, -1, -1,
+        1 /* triangle in top right for second grid */, 0, -1, -1, 1, 0,
+        1 /* triangle in bottom left for second grid */,
+      ]),
+    );
+  });
+});

src/webgl/rectangle_grid_buffer.ts

@@ -0,0 +1,80 @@
+/**
+ * @license
+ * Copyright 2024 Google Inc.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { getMemoizedBuffer } from "#src/webgl/buffer.js";
+import type { GL } from "#src/webgl/context.js";
+import { VERTICES_PER_QUAD } from "#src/webgl/quad.js";
+
+/**
+ * Create a Float32Array of vertices gridded in a rectangle
+ * Only grids along the x-axis are created, the y-axis is assumed to be the same for all grids
+ */
+export function createGriddedRectangleArray(
+  numGrids: number,
+  startX = -1,
+  endX = 1,
+  startY = 1,
+  endY = -1,
+): Float32Array {
+  const result = new Float32Array(numGrids * VERTICES_PER_QUAD * 2);
+  const step = (endX - startX) / numGrids;
+  let currentx = startX;
+  for (let i = 0; i < numGrids; ++i) {
+    const index = i * VERTICES_PER_QUAD * 2;
+
+    // Triangle 1 - top-left, top-right, bottom-right
+    result[index] = currentx; // top-left x
+    result[index + 1] = startY; // top-left y
+    result[index + 2] = currentx + step; // top-right x
+    result[index + 3] = startY; // top-right y
+    result[index + 4] = currentx + step; // bottom-right x
+    result[index + 5] = endY; // bottom-right y
+
+    // Triangle 2 - top-left, bottom-right, bottom-left
+    result[index + 6] = currentx; // top-left x
+    result[index + 7] = startY; // top-left y
+    result[index + 8] = currentx + step; // bottom-right x
+    result[index + 9] = endY; // bottom-right y
+    result[index + 10] = currentx; // bottom-left x
+    result[index + 11] = endY; // bottom-left y
+    currentx += step;
+  }
+  return result;
+}
+
+/**
+ * Get a buffer of vertices representing a rectangle that is gridded
+ * along the x dimension, useful for drawing grids, such as a lookup table / heatmap
+ */
+export function getGriddedRectangleBuffer(
+  gl: GL,
+  numGrids: number,
+  startX = -1,
+  endX = 1,
+  startY = 1,
+  endY = -1,
+) {
+  return getMemoizedBuffer(
+    gl,
+    WebGL2RenderingContext.ARRAY_BUFFER,
+    createGriddedRectangleArray,
+    numGrids,
+    startX,
+    endX,
+    startY,
+    endY,
+  ).value;
+}