Maintain dimension mapping when possible

packages/app/src/App.tsx

@@ -8,6 +8,7 @@ import { ReflexContainer, ReflexElement, ReflexSplitter } from 'react-reflex';
 import styles from './App.module.css';
 import BreadcrumbsBar from './breadcrumbs/BreadcrumbsBar';
 import type { FeedbackContext } from './breadcrumbs/models';
+import { DimMappingProvider } from './dimension-mapper/store';
 import EntityLoader from './EntityLoader';
 import ErrorFallback from './ErrorFallback';
 import MetadataViewer from './metadata-viewer/MetadataViewer';
@@ -84,21 +85,25 @@ function App(props: Props) {
             getFeedbackURL={getFeedbackURL}
           />
           <VisConfigProvider>
-            <ErrorBoundary
-              resetKeys={[selectedPath, isInspecting]}
-              FallbackComponent={ErrorFallback}
-            >
-              <Suspense fallback={<EntityLoader isInspecting={isInspecting} />}>
-                {isInspecting ? (
-                  <MetadataViewer
-                    path={selectedPath}
-                    onSelectPath={onSelectPath}
-                  />
-                ) : (
-                  <Visualizer path={selectedPath} />
-                )}
-              </Suspense>
-            </ErrorBoundary>
+            <DimMappingProvider>
+              <ErrorBoundary
+                resetKeys={[selectedPath, isInspecting]}
+                FallbackComponent={ErrorFallback}
+              >
+                <Suspense
+                  fallback={<EntityLoader isInspecting={isInspecting} />}
+                >
+                  {isInspecting ? (
+                    <MetadataViewer
+                      path={selectedPath}
+                      onSelectPath={onSelectPath}
+                    />
+                  ) : (
+                    <Visualizer path={selectedPath} />
+                  )}
+                </Suspense>
+              </ErrorBoundary>
+            </DimMappingProvider>
           </VisConfigProvider>
         </ReflexElement>
       </ReflexContainer>

packages/app/src/__tests__/DimensionMapper.test.tsx

@@ -1,7 +1,7 @@
 import { screen, within } from '@testing-library/react';
 import { expect, test } from 'vitest';
 
-import { renderApp, waitForAllLoaders } from '../test-utils';
+import { getDimMappingBtn, renderApp } from '../test-utils';
 import { Vis } from '../vis-packs/core/visualizations';
 
 test('control mapping for X axis when visualizing 2D dataset as Line', async () => {
@@ -100,14 +100,92 @@ test('slice through 2D dataset', async () => {
   const { user } = await renderApp({
     initialPath: '/nD_datasets/twoD',
     preferredVis: Vis.Line,
-    withFakeTimers: true, // required since React 18 upgrade (along with `waitForAllLoaders` below)
   });
 
-  await waitForAllLoaders();
-
   // Move to next slice with keyboard
   const d0Slider = screen.getByRole('slider', { name: 'D0' });
   await user.type(d0Slider, '{ArrowUp}');
 
   expect(d0Slider).toHaveAttribute('aria-valuenow', '1');
 });
+
+test('maintain mapping when switching to inspect mode and back', async () => {
+  const { user } = await renderApp({
+    initialPath: '/nD_datasets/twoD',
+    preferredVis: Vis.Heatmap,
+  });
+
+  // Swap axes for D0 and D1
+  await user.click(getDimMappingBtn('x', 0));
+
+  // Toggle inspect mode
+  await user.click(screen.getByRole('tab', { name: 'Inspect' }));
+  await user.click(screen.getByRole('tab', { name: 'Display' }));
+
+  expect(getDimMappingBtn('x', 0)).toBeChecked();
+  expect(getDimMappingBtn('x', 1)).not.toBeChecked();
+});
+
+test('maintain mapping when switching to visualization with same axes count', async () => {
+  const { user, selectVisTab } = await renderApp({
+    initialPath: '/nD_datasets/twoD',
+    preferredVis: Vis.Heatmap,
+  });
+
+  // Swap axes for D0 and D1
+  await user.click(getDimMappingBtn('x', 0));
+
+  // Switch to Matrix visualization
+  await selectVisTab(Vis.Matrix);
+
+  expect(getDimMappingBtn('x', 0)).toBeChecked();
+  expect(getDimMappingBtn('x', 1)).not.toBeChecked();
+});
+
+test('maintain mapping when switching to dataset with same dimensions', async () => {
+  const { user, selectExplorerNode } = await renderApp({
+    initialPath: '/nD_datasets/twoD_bool',
+    preferredVis: Vis.Line,
+  });
+
+  // Swap axes for D0 and D1
+  await user.click(getDimMappingBtn('x', 0));
+
+  // Switch to dataset with same dimensions
+  await selectExplorerNode('twoD_enum');
+
+  expect(getDimMappingBtn('x', 0)).toBeChecked();
+  expect(getDimMappingBtn('x', 1)).not.toBeChecked();
+});
+
+test('reset mapping when switching to visualization with different axes count', async () => {
+  const { user, selectVisTab } = await renderApp({
+    initialPath: '/nD_datasets/twoD',
+    preferredVis: Vis.Heatmap,
+  });
+
+  // Swap axes for D0 and D1
+  await user.click(getDimMappingBtn('x', 0));
+
+  // Switch to Line visualization
+  await selectVisTab(Vis.Line);
+
+  expect(getDimMappingBtn('x', 0)).not.toBeChecked();
+  expect(getDimMappingBtn('x', 1)).toBeChecked();
+});
+
+test('reset mapping when switching to dataset with different dimensions', async () => {
+  const { user, selectExplorerNode } = await renderApp({
+    initialPath: '/nD_datasets/twoD',
+    preferredVis: Vis.Heatmap,
+  });
+
+  // Swap axes for D0 and D1
+  await user.click(getDimMappingBtn('x', 0));
+
+  // Switch to dataset with different dimensions
+  await selectExplorerNode('twoD_cplx');
+
+  expect(getDimMappingBtn('x', 0)).not.toBeChecked();
+  expect(getDimMappingBtn('x', 1)).toBeChecked();
+});

packages/app/src/dimension-mapper/store.tsx

@@ -0,0 +1,72 @@
+import type { ArrayShape } from '@h5web/shared/hdf5-models';
+import type { PropsWithChildren } from 'react';
+import { createContext, useContext, useEffect, useState } from 'react';
+import type { StoreApi } from 'zustand';
+import { createStore, useStore } from 'zustand';
+
+import { areSameDims } from '../vis-packs/nexus/utils';
+import type { DimensionMapping } from './models';
+
+interface DimMappingState {
+  dims: ArrayShape;
+  axesCount: number;
+  mapping: DimensionMapping;
+  setMapping: (mapping: DimensionMapping) => void;
+  reset: (
+    dims: ArrayShape,
+    axesCount: number,
+    mapping: DimensionMapping,
+  ) => void;
+}
+
+function createLineConfigStore() {
+  return createStore<DimMappingState>((set) => ({
+    dims: [],
+    axesCount: 0,
+    mapping: [],
+    setMapping: (mapping) => set({ mapping }),
+    reset: (dims, axesCount, mapping) => {
+      set({ dims, axesCount, mapping });
+    },
+  }));
+}
+
+const StoreContext = createContext({} as StoreApi<DimMappingState>);
+
+interface Props {}
+export function DimMappingProvider(props: PropsWithChildren<Props>) {
+  const { children } = props;
+
+  const [store] = useState(createLineConfigStore);
+
+  return (
+    <StoreContext.Provider value={store}>{children}</StoreContext.Provider>
+  );
+}
+
+export function useDimMappingState(
+  dims: number[],
+  axesCount: number,
+): [DimensionMapping, (mapping: DimensionMapping) => void] {
+  const state = useStore(useContext(StoreContext));
+
+  /* If current mapping was initialised with different axes count and dimensions,
+   * need to compute new mapping and reset state. */
+  const isStale =
+    axesCount !== state.axesCount || !areSameDims(dims, state.dims);
+
+  const mapping = isStale
+    ? [
+        ...Array.from({ length: dims.length - axesCount }, () => 0),
+        ...(dims.length > 0
+          ? ['y' as const, 'x' as const].slice(-axesCount)
+          : []),
+      ]
+    : state.mapping;
+
+  useEffect(() => {
+    state.reset(dims, axesCount, mapping);
+  }, [isStale]); // eslint-disable-line react-hooks/exhaustive-deps
+
+  return [mapping, state.setMapping];
+}

packages/app/src/test-utils.tsx

@@ -118,6 +118,11 @@ export function getSelectedVisTab(): string {
   return selectedTab.textContent;
 }
 
+export function getDimMappingBtn(axis: 'x' | 'y', dim: number): HTMLElement {
+  const radioGroup = screen.getByLabelText(`Dimension as ${axis} axis`);
+  return within(radioGroup).getByRole('radio', { name: `D${dim}` });
+}
+
 /**
  * Mock a console method.
  * Mocks are automatically cleared and restored after every test but you

packages/app/src/vis-packs/core/complex/ComplexLineVisContainer.tsx

@@ -5,7 +5,7 @@ import {
 } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import type { VisContainerProps } from '../../models';
 import VisBoundary from '../../VisBoundary';
 import { useValuesInCache } from '../hooks';

packages/app/src/vis-packs/core/complex/ComplexVisContainer.tsx

@@ -6,7 +6,7 @@ import {
 } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import type { VisContainerProps } from '../../models';
 import VisBoundary from '../../VisBoundary';
 import { useHeatmapConfig } from '../heatmap/config';

packages/app/src/vis-packs/core/compound/CompoundVisContainer.tsx

@@ -6,7 +6,7 @@ import {
 } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import type { VisContainerProps } from '../../models';
 import VisBoundary from '../../VisBoundary';
 import { useValuesInCache } from '../hooks';

packages/app/src/vis-packs/core/heatmap/HeatmapVisContainer.tsx

@@ -6,7 +6,7 @@ import {
 } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import type { VisContainerProps } from '../../models';
 import VisBoundary from '../../VisBoundary';
 import { useIgnoreFillValue, useValuesInCache } from '../hooks';

packages/app/src/vis-packs/core/line/LineVisContainer.tsx

@@ -5,7 +5,7 @@ import {
 } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import type { VisContainerProps } from '../../models';
 import VisBoundary from '../../VisBoundary';
 import { useIgnoreFillValue, useValuesInCache } from '../hooks';

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

@@ -5,7 +5,7 @@ import {
 } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import type { VisContainerProps } from '../../models';
 import VisBoundary from '../../VisBoundary';
 import { useValuesInCache } from '../hooks';

packages/app/src/vis-packs/core/rgb/RgbVisContainer.tsx

@@ -6,7 +6,7 @@ import {
 } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import { useDataContext } from '../../../providers/DataProvider';
 import type { VisContainerProps } from '../../models';
 import VisBoundary from '../../VisBoundary';

packages/app/src/vis-packs/core/surface/SurfaceVisContainer.tsx

@@ -6,7 +6,7 @@ import {
 } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import type { VisContainerProps } from '../../models';
 import VisBoundary from '../../VisBoundary';
 import { useValuesInCache } from '../hooks';

packages/app/src/vis-packs/nexus/containers/NxComplexImageContainer.tsx

@@ -2,7 +2,7 @@ import { assertGroup, assertMinDims } from '@h5web/shared/guards';
 import { useState } from 'react';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import { useComplexConfig } from '../../core/complex/config';
 import MappedComplexVis from '../../core/complex/MappedComplexVis';
 import { useHeatmapConfig } from '../../core/heatmap/config';

packages/app/src/vis-packs/nexus/containers/NxComplexSpectrumContainer.tsx

@@ -2,7 +2,7 @@ import { ScaleType } from '@h5web/lib';
 import { assertGroup, isAxisScaleType } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import { useComplexLineConfig } from '../../core/complex/lineConfig';
 import MappedComplexLineVis from '../../core/complex/MappedComplexLineVis';
 import { useLineConfig } from '../../core/line/config';

packages/app/src/vis-packs/nexus/containers/NxImageContainer.tsx

@@ -2,7 +2,7 @@ import { assertGroup, assertMinDims } from '@h5web/shared/guards';
 import { useState } from 'react';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import { useHeatmapConfig } from '../../core/heatmap/config';
 import MappedHeatmapVis from '../../core/heatmap/MappedHeatmapVis';
 import { getSliceSelection } from '../../core/utils';

packages/app/src/vis-packs/nexus/containers/NxRgbContainer.tsx

@@ -1,7 +1,7 @@
 import { assertGroup, assertMinDims } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import { useRgbConfig } from '../../core/rgb/config';
 import MappedRgbVis from '../../core/rgb/MappedRgbVis';
 import { getSliceSelection } from '../../core/utils';

packages/app/src/vis-packs/nexus/containers/NxSpectrumContainer.tsx

@@ -2,7 +2,7 @@ import { ScaleType } from '@h5web/lib';
 import { assertGroup, isAxisScaleType } from '@h5web/shared/guards';
 
 import DimensionMapper from '../../../dimension-mapper/DimensionMapper';
-import { useDimMappingState } from '../../../dimension-mapper/hooks';
+import { useDimMappingState } from '../../../dimension-mapper/store';
 import { useLineConfig } from '../../core/line/config';
 import MappedLineVis from '../../core/line/MappedLineVis';
 import { getSliceSelection } from '../../core/utils';