Add IndexedStackColorMapper for ImageStack glyph

bokehjs/src/lib/models/annotations/color_bar.ts

@@ -4,7 +4,7 @@ import {CategoricalAxis, LinearAxis, LogAxis} from "../axes"
 import type {TickFormatter} from "../formatters/tick_formatter"
 import {BasicTickFormatter, LogTickFormatter, CategoricalTickFormatter} from "../formatters"
 import {ColorMapper} from "../mappers/color_mapper"
-import {LinearColorMapper, LogColorMapper, ScanningColorMapper, CategoricalColorMapper, ContinuousColorMapper, WeightedStackColorMapper} from "../mappers"
+import {IndexedStackColorMapper, LinearColorMapper, LogColorMapper, ScanningColorMapper, CategoricalColorMapper, ContinuousColorMapper, WeightedStackColorMapper} from "../mappers"
 import type {Range} from "../ranges"
 import {Range1d, FactorRange} from "../ranges"
 import type {Scale} from "../scales"
@@ -41,13 +41,19 @@ export class ColorBarView extends BaseColorBarView {
     this.connect(this.model.color_mapper.metrics_change, () => this._metrics_changed())
     this.connect(this.model.properties.display_low.change, () => this._metrics_changed())
     this.connect(this.model.properties.display_high.change, () => this._metrics_changed())
+
+    if (this.model.color_mapper instanceof IndexedStackColorMapper) {
+      this.connect(this.model.color_mapper.single_color_mapper.metrics_change, () => this._metrics_changed())
+    }
   }
 
   get color_mapper(): ColorMapper {
     // Color mapper that is used to render this colorbar.
     let mapper = this.model.color_mapper
     if (mapper instanceof WeightedStackColorMapper)
       mapper = mapper.alpha_mapper
+    else if (mapper instanceof IndexedStackColorMapper)
+      mapper = mapper.single_color_mapper
     return mapper
   }
 

bokehjs/src/lib/models/mappers/continuous_color_mapper.ts

@@ -112,9 +112,9 @@ export abstract class ContinuousColorMapper extends ColorMapper {
     }
   }
 
-  protected _scan_data: {min: number, max: number} | null = null
+  public _scan_data: {min: number, max: number} | null = null
 
-  protected abstract scan(data: Arrayable<number>, n: number): {min: number, max: number}
+  public abstract scan(data: Arrayable<number>, n: number): {min: number, max: number}
 
   public _v_compute<T>(data: Arrayable<number>, values: Arrayable<T>,
       palette: Arrayable<T>, colors: {nan_color: T, low_color?: T, high_color?: T}): void {
@@ -149,7 +149,7 @@ export abstract class ContinuousColorMapper extends ColorMapper {
     }
   }
 
-  protected cmap<T>(value: number, palette: Arrayable<T>, low_color: T, high_color: T): T {
+  public cmap<T>(value: number, palette: Arrayable<T>, low_color: T, high_color: T): T {
     const index = this.value_to_index(value, palette.length)
     if (index < 0)
       return low_color

bokehjs/src/lib/models/mappers/index.ts

@@ -3,6 +3,7 @@ export {CategoricalMarkerMapper}  from "./categorical_marker_mapper"
 export {CategoricalPatternMapper} from "./categorical_pattern_mapper"
 export {ContinuousColorMapper}    from "./continuous_color_mapper"
 export {ColorMapper}              from "./color_mapper"
+export {IndexedStackColorMapper}  from "./indexed_stack_color_mapper"
 export {LinearColorMapper}        from "./linear_color_mapper"
 export {LogColorMapper}           from "./log_color_mapper"
 export {ScanningColorMapper}      from "./scanning_color_mapper"

bokehjs/src/lib/models/mappers/indexed_stack_color_mapper.ts

@@ -0,0 +1,72 @@
+import {_convert_palette, _convert_color} from "./color_mapper"
+import {ContinuousColorMapper} from "./continuous_color_mapper"
+import {StackColorMapper} from "./stack_color_mapper"
+import type * as p from "core/properties"
+import type {Arrayable} from "core/types"
+
+export namespace IndexedStackColorMapper {
+  export type Attrs = p.AttrsOf<Props>
+
+  export type Props = StackColorMapper.Props & {
+    single_color_mapper: p.Property<ContinuousColorMapper>
+    index: p.Property<number>
+    map_all: p.Property<boolean>
+  }
+}
+
+export interface IndexedStackColorMapper extends IndexedStackColorMapper.Attrs {}
+
+export class IndexedStackColorMapper extends StackColorMapper {
+  declare properties: IndexedStackColorMapper.Props
+
+  constructor(attrs?: Partial<IndexedStackColorMapper.Attrs>) {
+    super(attrs)
+  }
+
+  static {
+    this.define<IndexedStackColorMapper.Props>(({Boolean, Number, Ref}) => ({
+      single_color_mapper: [ Ref(ContinuousColorMapper) ],
+      index:               [ Number, 0 ],
+      map_all:             [ Boolean, false ],
+    }))
+  }
+
+  protected override _v_compute<T>(data: Arrayable<number>, values: Arrayable<T>,
+      palette: Arrayable<T>, colors: {nan_color: T, low_color?: T, high_color?: T}): void {
+    const {index} = this
+    const n = values.length
+    const max_index = data.length / n
+    // TODO: check index is within bounds of 0 .. max_index-1
+
+    const sub_data = new Array<number>(n)
+    for (let i = 0; i < n; i++)
+      sub_data[i] = data[i*max_index + index]  // do a slice or similar
+
+    if (this.map_all) {
+      // Need to reimplement single_color_mapper._v_compute as need to use our own scan_data
+      const {nan_color} = colors
+      let {low_color, high_color} = colors
+      if (low_color == null)
+        low_color = palette[0]
+      if (high_color == null)
+        high_color = palette[palette.length-1]
+
+      // Calling scan using all of the 3D data.
+      this.single_color_mapper._scan_data = this.single_color_mapper.scan(data, palette.length)
+      this.single_color_mapper.metrics_change.emit()
+
+      for (let i = 0, end = sub_data.length; i < end; i++) {
+        const d = sub_data[i]
+
+        if (isNaN(d))
+          values[i] = nan_color
+        else
+          values[i] = this.single_color_mapper.cmap(d, palette, low_color, high_color)
+      }
+    } else {
+      this.single_color_mapper._v_compute(sub_data, values, palette, colors)
+    }
+
+    this.metrics_change.emit()
+  }
+}

bokehjs/src/lib/models/mappers/linear_color_mapper.ts

@@ -26,7 +26,7 @@ export class LinearColorMapper extends ContinuousColorMapper {
     super(attrs)
   }
 
-  protected scan(data: Arrayable<number>, n: number): LinearScanData {
+  public scan(data: Arrayable<number>, n: number): LinearScanData {
     const low = this.low != null ? this.low : min(data)
     const high = this.high != null ? this.high : max(data)
 

bokehjs/src/lib/models/mappers/log_color_mapper.ts

@@ -25,7 +25,7 @@ export class LogColorMapper extends ContinuousColorMapper {
     super(attrs)
   }
 
-  protected scan(data: Arrayable<number>, n: number): LogScanData {
+  public scan(data: Arrayable<number>, n: number): LogScanData {
     const low = this.low != null ? this.low : min(data)
     const high = this.high != null ? this.high : max(data)
     const scale = n / Math.log(high / low)  // subtract the low offset

src/bokeh/models/mappers.py

@@ -59,6 +59,7 @@
     'CategoricalMarkerMapper',
     'CategoricalPatternMapper',
     'ContinuousColorMapper',
+    'IndexedStackColorMapper',
     'LinearColorMapper',
     'LogColorMapper',
     'EqHistColorMapper',
@@ -332,6 +333,37 @@ class StackColorMapper(ColorMapper):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
 
+class IndexedStackColorMapper(StackColorMapper):
+    ''' Color maps a single 2D array from a 3D data array of shape
+    ``(ny, nx, nstack)`` in an ``ImageStack`` glyph to an RGBA image.
+
+    The mapped 2D array is determined by the ``index`` attribute and in terms
+    of the ``ImageStack`` data is ``data[:, :, index]``.
+
+    If ``map_all`` is ``False`` the color map limits and properties are
+    determined from just the 2D array displayed. If ``map_all`` is ``True``
+    they are determined from the whole 3D array.
+
+    '''
+
+    # explicit __init__ to support Init signatures
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    single_color_mapper = Instance(ContinuousColorMapper, help="""
+    Color mapper used to calculate the colors of the 2D array displayed.
+    """)
+
+    index = Int(default=0, help="""
+    Final dimension of the ``ImageStack`` 3D data array that is color mapped.
+    """)
+
+    map_all = Bool(default=False, help="""
+    Whether to calculate the color mapping limits and properties based on all
+    of the 3D data array or just the 2D array that is displayed.
+    """)
+
+
 class WeightedStackColorMapper(StackColorMapper):
     ''' Maps 3D data arrays of shape ``(ny, nx, nstack)`` to 2D RGBA images
     of shape ``(ny, nx)`` using a palette of length ``nstack``.

tests/baselines/defaults.json5

@@ -457,6 +457,15 @@
   StackColorMapper: {
     __extends__: "ColorMapper",
   },
+  IndexedStackColorMapper: {
+    __extends__: "StackColorMapper",
+    single_color_mapper: {
+      type: "symbol",
+      name: "unset",
+    },
+    index: 0,
+    map_all: false,
+  },
   WeightedStackColorMapper: {
     __extends__: "StackColorMapper",
     alpha_mapper: {