Improve matrix vis performance through virtualisation of index cells

packages/app/src/vis-packs/core/configs.ts

@@ -1,3 +1,4 @@
+export { MatrixConfigProvider } from './matrix/config';
 export { LineConfigProvider } from './line/config';
 export { HeatmapConfigProvider } from './heatmap/config';
 export { ComplexConfigProvider } from './complex/config';

packages/app/src/vis-packs/core/matrix/MappedMatrixVis.tsx

@@ -5,6 +5,7 @@ import { createPortal } from 'react-dom';
 import type { DimensionMapping } from '../../../dimension-mapper/models';
 import { useMappedArray, useSlicedDimsAndMapping } from '../hooks';
 import MatrixToolbar from './MatrixToolbar';
+import { useMatrixVisConfig } from './config';
 
 interface Props {
   value: Primitive<PrintableType>[];
@@ -19,6 +20,8 @@ function MappedMatrixVis(props: Props) {
   const { value, dims, dimMapping, formatter, cellWidth, toolbarContainer } =
     props;
 
+  const sticky = useMatrixVisConfig((state) => state.sticky);
+
   const [slicedDims, slicedMapping] = useSlicedDimsAndMapping(dims, dimMapping);
   const [mappedArray] = useMappedArray(value, slicedDims, slicedMapping);
 
@@ -34,6 +37,7 @@ function MappedMatrixVis(props: Props) {
         dataArray={mappedArray}
         formatter={formatter}
         cellWidth={cellWidth}
+        sticky={sticky}
       />
     </>
   );

packages/app/src/vis-packs/core/matrix/MatrixToolbar.tsx

@@ -1,8 +1,9 @@
-import { DownloadBtn, Toolbar } from '@h5web/lib';
+import { DownloadBtn, Separator, ToggleBtn, Toolbar } from '@h5web/lib';
 import type { Primitive, PrintableType } from '@h5web/shared';
 import type { NdArray } from 'ndarray';
-import { FiDownload } from 'react-icons/fi';
+import { FiAnchor, FiDownload } from 'react-icons/fi';
 
+import { useMatrixVisConfig } from './config';
 import { sliceToCsv } from './utils';
 
 interface Props {
@@ -11,13 +12,23 @@ interface Props {
 
 function MatrixToolbar(props: Props) {
   const { currentSlice } = props;
+  const { sticky, toggleSticky } = useMatrixVisConfig();
 
   if (currentSlice && currentSlice.shape.length > 2) {
     throw new Error('Expected current slice to have at most two dimensions');
   }
 
   return (
     <Toolbar>
+      <ToggleBtn
+        label="Freeze indices"
+        icon={FiAnchor}
+        value={sticky}
+        onToggle={toggleSticky}
+      />
+
+      <Separator />
+
       {currentSlice && (
         <DownloadBtn
           icon={FiDownload}

packages/app/src/vis-packs/core/matrix/config.tsx

@@ -0,0 +1,33 @@
+import create from 'zustand';
+import createContext from 'zustand/context';
+import { persist } from 'zustand/middleware';
+
+import type { ConfigProviderProps } from '../../models';
+
+interface MatrixVisConfig {
+  sticky: boolean;
+  toggleSticky: () => void;
+}
+
+function createStore() {
+  return create<MatrixVisConfig>(
+    persist(
+      (set) => ({
+        sticky: false,
+        toggleSticky: () => set((state) => ({ sticky: !state.sticky })),
+      }),
+      {
+        name: 'h5web:matrix',
+        version: 1,
+      }
+    )
+  );
+}
+
+const { Provider, useStore } = createContext<MatrixVisConfig>();
+export const useMatrixVisConfig = useStore;
+
+export function MatrixConfigProvider(props: ConfigProviderProps) {
+  const { children } = props;
+  return <Provider createStore={createStore}>{children}</Provider>;
+}

packages/app/src/vis-packs/core/visualizations.ts

@@ -25,6 +25,7 @@ import {
   ComplexConfigProvider,
   ComplexLineConfigProvider,
   RgbConfigProvider,
+  MatrixConfigProvider,
 } from './configs';
 import {
   RawVisContainer,
@@ -76,6 +77,7 @@ export const CORE_VIS: Record<Vis, CoreVisDef> = {
     name: Vis.Matrix,
     Icon: FiGrid,
     Container: MatrixVisContainer,
+    ConfigProvider: MatrixConfigProvider,
     supportsDataset: (dataset) => {
       return hasPrintableType(dataset) && hasArrayShape(dataset);
     },

packages/lib/src/vis/matrix/Cell.tsx

@@ -1,12 +1,12 @@
-import { useContext } from 'react';
+import { memo, useContext } from 'react';
 import type { GridChildComponentProps } from 'react-window';
 
-import { GridSettingsContext } from './GridSettingsContext';
 import styles from './MatrixVis.module.css';
+import { SettingsContext } from './context';
 
 function Cell(props: GridChildComponentProps) {
   const { rowIndex, columnIndex, style } = props;
-  const { cellFormatter } = useContext(GridSettingsContext);
+  const { cellFormatter } = useContext(SettingsContext);
 
   // Disable index columns (rendering done by the innerElementType)
   if (rowIndex * columnIndex === 0) {
@@ -20,7 +20,7 @@ function Cell(props: GridChildComponentProps) {
       role="cell"
       aria-rowindex={rowIndex}
       aria-colindex={columnIndex}
-      data-bg={(rowIndex + columnIndex) % 2 === 1 ? '' : undefined}
+      data-bg={(rowIndex + columnIndex) % 2 === 1 || undefined}
     >
       {
         // -1 to account for the index row and column
@@ -30,4 +30,4 @@ function Cell(props: GridChildComponentProps) {
   );
 }
 
-export default Cell;
+export default memo(Cell);

packages/lib/src/vis/matrix/Grid.tsx

@@ -0,0 +1,38 @@
+import { useMeasure } from '@react-hookz/web';
+import { useContext } from 'react';
+import { FixedSizeGrid as IndexedGrid } from 'react-window';
+
+import Cell from './Cell';
+import styles from './MatrixVis.module.css';
+import StickyGrid from './StickyGrid';
+import { SettingsContext } from './context';
+
+function Grid() {
+  const { rowCount, columnCount, cellSize, setRenderedItems } =
+    useContext(SettingsContext);
+
+  const [wrapperSize, wrapperRef] = useMeasure<HTMLDivElement>();
+
+  return (
+    <div ref={wrapperRef} className={styles.wrapper}>
+      {wrapperSize && (
+        <IndexedGrid
+          className={styles.grid}
+          innerElementType={StickyGrid}
+          width={wrapperSize.width}
+          height={wrapperSize.height}
+          rowHeight={cellSize.height}
+          rowCount={rowCount}
+          columnWidth={cellSize.width}
+          columnCount={columnCount}
+          overscanCount={8}
+          onItemsRendered={setRenderedItems}
+        >
+          {Cell}
+        </IndexedGrid>
+      )}
+    </div>
+  );
+}
+
+export default Grid;

packages/lib/src/vis/matrix/HeaderCells.tsx

@@ -0,0 +1,33 @@
+import { range } from 'lodash';
+import { useContext } from 'react';
+
+import styles from './MatrixVis.module.css';
+import { SettingsContext } from './context';
+
+interface Props {
+  indexMin: number;
+  indexMax: number;
+  transform: string /* to compensate for header cells not rendered in the range [0 indexMin] */;
+}
+
+function HeaderCells(props: Props) {
+  const { indexMin, indexMax, transform } = props;
+  const { cellSize } = useContext(SettingsContext);
+
+  return (
+    <>
+      {range(indexMin, indexMax).map((index) => (
+        <div
+          key={index.toString()}
+          className={styles.indexCell}
+          style={{ ...cellSize, transform }}
+          data-bg={index % 2 === 1 ? '' : undefined}
+        >
+          {index >= 0 && index}
+        </div>
+      ))}
+    </>
+  );
+}
+
+export default HeaderCells;