Improve unit details state management

packages/column-views/CHANGELOG.md

@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file. The format
 is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this
 project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.2.1] - 2025-11-29
+
+- Start unifying state management components
+- Create a hoistable store for column state
+- Begin using `jotai` for some aspects of state management
+
 ## [2.2.0] - 2025-11-28
 
 - Update SGP and PBDB facets

packages/column-views/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@macrostrat/column-views",
-  "version": "2.2.0",
+  "version": "2.2.1",
   "description": "Data views for Macrostrat stratigraphic columns",
   "type": "module",
   "source": "src/index.ts",
@@ -63,6 +63,8 @@
     "d3-path": "^3.1.0",
     "d3-scale": "^4.0.2",
     "d3-shape": "^3.2.0",
+    "jotai": "^2.15.1",
+    "jotai-scope": "^0.9.7",
     "mapbox-gl": "^2.15.0||^3.0.0",
     "react-spring": "^9.7.5",
     "topojson-client": "^3.1.0",

packages/column-views/src/column.ts

@@ -16,12 +16,12 @@ import {
   ComponentType,
 } from "react";
 import styles from "./column.module.sass";
+import { UnitComponent } from "./units";
 import {
   UnitSelectionProvider,
-  UnitComponent,
   UnitKeyboardNavigation,
   useUnitSelectionDispatch,
-} from "./units";
+} from "./data-provider";
 
 import {
   Identifier,
@@ -95,7 +95,6 @@ export function Column(props: ColumnProps) {
     selectedUnit,
     children,
     units: rawUnits,
-    axisType = ColumnAxisType.AGE,
     t_age,
     b_age,
     t_pos,
@@ -109,6 +108,7 @@ export function Column(props: ColumnProps) {
     allowUnitSelection,
     hybridScale,
     scale,
+    axisType,
     ...rest
   } = props;
   const ref = useRef<HTMLElement>();
@@ -125,8 +125,20 @@ export function Column(props: ColumnProps) {
     _minPixelScale = pixelScale;
   }
 
+  // Handle special cases for hybrid scales (WIP, we need to regularize this)
+  let _axisType = axisType ?? ColumnAxisType.AGE;
+  let ageAxisComponent = CompositeAgeAxis;
+  if (
+    hybridScale?.type === HybridScaleType.ApproximateHeight &&
+    _axisType != ColumnAxisType.AGE
+  ) {
+    // Use approximate height axis for non-age columns if a non-age axis type is requested
+    ageAxisComponent = ApproximateHeightAxis;
+    _axisType = ColumnAxisType.AGE;
+  }
+
   const { sections, units, totalHeight } = usePreparedColumnUnits(rawUnits, {
-    axisType,
+    axisType: _axisType,
     t_age,
     b_age,
     t_pos,
@@ -154,15 +166,6 @@ export function Column(props: ColumnProps) {
     );
   }
 
-  let ageAxisComponent = CompositeAgeAxis;
-  if (
-    hybridScale?.type === HybridScaleType.ApproximateHeight &&
-    axisType != ColumnAxisType.AGE
-  ) {
-    // Use approximate height axis for non-age columns if a non-age axis type is requested
-    ageAxisComponent = ApproximateHeightAxis;
-  }
-
   let main: any = h(
     ColumnInner,
     { columnRef: ref, ageAxisComponent, ...rest },
@@ -196,7 +199,7 @@ export function Column(props: ColumnProps) {
 
   return h(
     MacrostratColumnDataProvider,
-    { units, sections, totalHeight, axisType },
+    { units, sections, totalHeight, axisType: _axisType },
     main,
   );
 }

packages/column-views/src/correlation-chart/main.ts

@@ -3,7 +3,7 @@ import {
   UnitSelectionProvider,
   UnitKeyboardNavigation,
   useUnitSelectionStore,
-} from "../units";
+} from "../data-provider";
 import { UnitSelectionPopover } from "../unit-details";
 import hyper from "@macrostrat/hyper";
 import styles from "./main.module.sass";

packages/column-views/src/data-provider/index.ts

@@ -2,3 +2,4 @@ export * from "./fetch";
 export * from "./base";
 export * from "./column-nav";
 export * from "./store";
+export * from "./unit-selection";

packages/column-views/src/data-provider/store.ts

@@ -6,6 +6,61 @@ import {
 } from "../prepare-units/composite-scale";
 import { ColumnAxisType } from "@macrostrat/column-components";
 import type { ExtUnit, PackageLayoutData } from "../prepare-units";
+// An isolated jotai store for Macrostrat column usage
+// TODO: there might be a better way to do this using the MacrostratDataProvider or similar
+import { createIsolation } from "jotai-scope";
+import { atom, type WritableAtom } from "jotai";
+
+const { Provider, useAtom, useAtomValue, useStore } = createIsolation();
+
+type ProviderProps = {
+  children: ReactNode;
+  initialValues?: Iterable<[WritableAtom<any, any, any>, any]>;
+};
+
+const columnUnitsAtom = atom<ExtUnit[]>();
+
+const columnUnitsMapAtom = atom<Map<number, ExtUnit> | null>((get) => {
+  const units = get(columnUnitsAtom);
+  if (!units) return null;
+  const unitMap = new Map<number, ExtUnit>();
+  units.forEach((unit) => {
+    unitMap.set(unit.unit_id, unit);
+  });
+  return unitMap;
+});
+
+function ScopedProvider({ children, ...rest }: ProviderProps) {
+  // Always use the same store instance in this tree
+  let val = null;
+  try {
+    val = useStore();
+  } catch {
+    // No store found, create a new one
+    val = null;
+  }
+  return h(Provider, { store: val, ...rest }, children);
+}
+
+export function MacrostratColumnStateProvider({
+  children,
+  units,
+}: {
+  children: ReactNode;
+  units: ExtUnit[];
+}) {
+  /** Top-level provider for Macrostrat column data.
+   * It is either provided by the Column component itself, or
+   * can be hoisted higher in the tree to provide a common data context
+   */
+  return h(
+    ScopedProvider,
+    {
+      initialValues: [[columnUnitsAtom, units]],
+    },
+    children,
+  );
+}
 
 export interface MacrostratColumnDataContext {
   units: ExtUnit[];
@@ -24,7 +79,7 @@ export function MacrostratColumnDataProvider({
   totalHeight,
   axisType,
 }: MacrostratColumnDataContext & { children: ReactNode }) {
-  /** Provider for Macrostrat column data.
+  /** Internal provider for Macrostrat column data.
    * As a general rule, we want to provide data and column-axis
    * height calculations through the context, since these need to
    * be accessed by any component that lays out information on the
@@ -45,7 +100,11 @@ export function MacrostratColumnDataProvider({
     };
   }, [units, sections, totalHeight, axisType]);
 
-  return h(MacrostratColumnDataContext.Provider, { value }, children);
+  return h(
+    MacrostratColumnStateProvider,
+    { units },
+    h(MacrostratColumnDataContext.Provider, { value }, children),
+  );
 }
 
 export function useMacrostratColumnData() {
@@ -59,19 +118,15 @@ export function useMacrostratColumnData() {
 }
 
 export function useMacrostratUnits() {
-  return useMacrostratColumnData().units;
+  return useAtomValue(columnUnitsAtom);
 }
 
 export function useColumnUnitsMap(): Map<number, ExtUnit> | null {
-  const ctx = useContext(MacrostratColumnDataContext);
-  return useMemo(() => {
-    if (ctx == null) return null;
-    const unitMap = new Map<number, ExtUnit>();
-    ctx.units.forEach((unit) => {
-      unitMap.set(unit.unit_id, unit);
-    });
-    return unitMap;
-  }, [ctx?.units]);
+  try {
+    return useAtomValue(columnUnitsMapAtom);
+  } catch {
+    return null;
+  }
 }
 
 export function useCompositeScale(): CompositeColumnScale {

packages/column-views/src/units/selection.ts → packages/column-views/src/data-provider/unit-selection.ts

@@ -12,7 +12,7 @@ import {
   useCallback,
 } from "react";
 import { createStore, StoreApi, useStore } from "zustand";
-import type { RectBounds, IUnit } from "./types";
+import type { RectBounds, IUnit } from "../units/types";
 
 type UnitSelectDispatch = (
   unit: number | BaseUnit | null,

packages/column-views/src/maps/column-correlation/state.ts

@@ -6,7 +6,6 @@ import { lineIntersect } from "@turf/line-intersect";
 import distance from "@turf/distance";
 import { nearestPointOnLine } from "@turf/nearest-point-on-line";
 import { centroid } from "@turf/centroid";
-import mapboxgl from "mapbox-gl";
 import {
   createContext,
   useState,

packages/column-views/src/unit-details/panel.ts

@@ -593,7 +593,7 @@ function UnitIDList({ units, onSelectUnit, showNames = false }) {
     const u1 = units.filter((d) => d != 0);
     if (showNames) {
       return u1.map((unitID) => {
-        const unitData = unitsMap.get(unitID);
+        const unitData = unitsMap?.get(unitID);
         let name: string = undefined;
         if (unitData != null) {
           name = defaultNameFunction(unitData);

packages/column-views/src/unit-details/popover.ts

@@ -1,7 +1,7 @@
 import hyper from "@macrostrat/hyper";
 import { Popover } from "@blueprintjs/core";
 import styles from "./popover.module.sass";
-import { useSelectedUnit, useUnitSelectionStore } from "../units";
+import { useSelectedUnit, useUnitSelectionStore } from "../data-provider";
 import { UnitDetailsPanel } from "./panel";
 
 const h = hyper.styled(styles);

packages/column-views/src/units/boxes.ts

@@ -12,12 +12,11 @@ import { SizeAwareLabel, Clickable } from "@macrostrat/ui-components";
 import hyper from "@macrostrat/hyper";
 import { forwardRef, ReactNode, useContext, useMemo } from "react";
 import { resolveID, scalePattern } from "./resolvers";
-import { useUnitSelectionTarget } from "./selection";
+import { useUnitSelectionTarget, useLithologies } from "../data-provider";
 import { IUnit } from "./types";
 import styles from "./boxes.module.sass";
 import classNames from "classnames";
 import { getUnitHeightRange } from "../prepare-units/utils";
-import { useLithologies } from "../data-provider";
 import { getMixedUnitColor } from "./colors";
 import type { RectBounds } from "./types";
 

packages/column-views/src/units/composite.ts

@@ -16,8 +16,11 @@ import {
 import { BaseUnit } from "@macrostrat/api-types";
 import { LabeledUnit, UnitBoxes } from "./boxes";
 import styles from "./composite.module.sass";
-import { useCompositeScale, useMacrostratColumnData } from "../data-provider";
-import { useUnitSelectionDispatch } from "../units";
+import {
+  useCompositeScale,
+  useMacrostratColumnData,
+  useUnitSelectionDispatch,
+} from "../data-provider";
 
 const h = hyperStyled(styles);
 

packages/column-views/src/units/index.ts

@@ -17,7 +17,6 @@ import { Lithology } from "@macrostrat/api-types";
 
 export * from "./composite";
 export * from "./types";
-export * from "./selection";
 export * from "./colors";
 
 export function UnitsColumn({ width = 100 }) {

packages/column-views/stories/nonlinear-scale.stories.ts

@@ -167,7 +167,7 @@ export const WithApproximateHeightScaleOnly: Story = {
     id: 448,
     // Ordered time bins
     pixelScale: 0.5,
-    axisType: ColumnAxisType.AGE,
+    axisType: ColumnAxisType.HEIGHT,
     hybridScale: {
       type: HybridScaleType.ApproximateHeight,
     },
@@ -177,9 +177,5 @@ export const WithApproximateHeightScaleOnly: Story = {
     showTimescale: false,
     unconformityHeight: 20,
     unconformityLabels: false,
-    children: h(ApproximateHeightAxis, {
-      // Move to the left side
-      style: { order: -1, marginRight: "8px" },
-    }),
   },
 };