@vitessce/statistical-plots 3.5.5 → 3.5.7

package/src/Treemap.js

@@ -0,0 +1,217 @@
+/* eslint-disable indent */
+/* eslint-disable camelcase */
+import React, { useMemo, useEffect, useRef } from 'react';
+import { scaleOrdinal } from 'd3-scale';
+import { select } from 'd3-selection';
+import { treemap, treemapBinary, hierarchy as d3_hierarchy } from 'd3-hierarchy';
+import { rollup as d3_rollup } from 'd3-array';
+import { isEqual } from 'lodash-es';
+import { colorArrayToString } from '@vitessce/sets-utils';
+import { getDefaultColor, pluralize as plur } from '@vitessce/utils';
+
+// Based on Observable's built-in DOM.uid function.
+// This is intended to be used with SVG clipPaths
+// which require a unique href value to reference
+// other elements contained in the DOM.
+function uidGenerator(prefix) {
+  let i = 0;
+  return () => {
+    i += 1;
+    return { id: `${prefix}-${i}`, href: `#${prefix}-${i}` };
+  };
+}
+
+// Create a d3-scale ordinal scale mapping set paths to color strings.
+function getColorScale(setSelectionArr, setColorArr, theme) {
+  return scaleOrdinal()
+    .domain(setSelectionArr || [])
+    .range(
+      setSelectionArr
+        ?.map(setNamePath => (
+          setColorArr?.find(d => isEqual(d.path, setNamePath))?.color
+          || getDefaultColor(theme)
+        ))
+        ?.map(colorArrayToString) || [],
+    );
+}
+
+/**
+ * Renders a treemap plot using D3.
+ * References:
+ * - https://observablehq.com/@d3/treemap-component
+ * - https://observablehq.com/@d3/treemap-stratify
+ * - https://observablehq.com/@d3/json-treemap
+ * - https://observablehq.com/@d3/nested-treemap
+ * @returns
+ */
+export default function Treemap(props) {
+  const {
+    obsCounts,
+    obsColorEncoding,
+    hierarchyLevels,
+    theme,
+    width,
+    height,
+    obsType,
+    sampleType,
+    obsSetColor,
+    sampleSetColor,
+    obsSetSelection,
+    sampleSetSelection,
+    marginTop = 5,
+    marginRight = 5,
+    marginLeft = 80,
+    marginBottom,
+  } = props;
+
+  const hierarchyData = useMemo(() => {
+    // Support both sampleSet->obsSet and
+    // obsSet->sampleSet hierarchy modes
+    if (!obsCounts) {
+      return null;
+    }
+    let map;
+    if (isEqual(hierarchyLevels, ['sampleSet', 'obsSet'])) {
+      map = d3_rollup(
+        obsCounts,
+        D => D[0].value,
+        d => d.sampleSetPath,
+        d => d.obsSetPath,
+      );
+    } else if (isEqual(hierarchyLevels, ['obsSet', 'sampleSet'])) {
+      map = d3_rollup(
+        obsCounts,
+        D => D[0].value,
+        d => d.obsSetPath,
+        d => d.sampleSetPath,
+      );
+    } else {
+      throw new Error('Unexpected levels value.');
+    }
+    return d3_hierarchy(map);
+  }, [obsCounts, hierarchyLevels]);
+
+  const [obsSetColorScale, sampleSetColorScale] = useMemo(() => [
+      getColorScale(obsSetSelection, obsSetColor, theme),
+      getColorScale(sampleSetSelection, sampleSetColor, theme),
+    ], [obsSetSelection, sampleSetSelection, sampleSetColor, obsSetColor, theme]);
+
+  const treemapLeaves = useMemo(() => {
+    const treemapFunc = treemap()
+      .tile(treemapBinary)
+      .size([width, height])
+      .padding(1)
+      .round(true);
+
+    // When d3.hierarchy is passed a Map object,
+    // the nodes are represented like [key, value] tuples.
+    // So in `.sum` and `.sort` below,
+    // `d[1]` accesses the value (i.e., cell count).
+    // Reference: https://d3js.org/d3-hierarchy/hierarchy#hierarchy
+    const treemapLayout = treemapFunc(hierarchyData
+      .sum(d => d[1])
+      .sort((a, b) => b[1] - a[1]));
+    return treemapLayout.leaves();
+  }, [hierarchyData, width, height]);
+
+  const svgRef = useRef();
+
+  useEffect(() => {
+    const domElement = svgRef.current;
+
+    const svg = select(domElement);
+    svg.selectAll('g').remove();
+    svg
+      .attr('width', width)
+      .attr('height', height)
+      .attr('viewBox', [0, 0, width, height])
+      .attr('style', 'font: 10px sans-serif');
+
+    if (!treemapLeaves || !obsSetSelection || !sampleSetSelection) {
+      return;
+    }
+
+    // Add a group for each leaf of the hierarchy.
+    const leaf = svg.selectAll('g')
+      .data(treemapLeaves)
+      .join('g')
+        .attr('transform', d => `translate(${d.x0},${d.y0})`);
+
+    // Append a tooltip.
+    leaf.append('title')
+        .text((d) => {
+          const cellCount = d.data?.[1];
+          const primaryPathString = JSON.stringify(d.data[0]);
+          const secondaryPathString = JSON.stringify(d.parent.data[0]);
+          return `${cellCount.toLocaleString()} ${plur(obsType, cellCount)} in ${primaryPathString} and ${secondaryPathString}`;
+        });
+
+    const getLeafUid = uidGenerator('leaf');
+    const getClipUid = uidGenerator('clip');
+
+    const colorScale = obsColorEncoding === 'sampleSetSelection'
+      ? sampleSetColorScale
+      : obsSetColorScale;
+    const getPathForColoring = d => (
+      // eslint-disable-next-line no-nested-ternary
+      obsColorEncoding === 'sampleSetSelection'
+        ? (hierarchyLevels[0] === 'obsSet' ? d.data?.[0] : d.parent?.data?.[0])
+        : (hierarchyLevels[0] === 'sampleSet' ? d.data?.[0] : d.parent?.data?.[0])
+    );
+
+    // Append a color rectangle for each leaf.
+    leaf.append('rect')
+        .attr('id', (d) => {
+          // eslint-disable-next-line no-param-reassign
+          d.leafUid = getLeafUid();
+          return d.leafUid.id;
+        })
+        .attr('fill', d => colorScale(getPathForColoring(d)))
+        .attr('fill-opacity', 0.8)
+        .attr('width', d => d.x1 - d.x0)
+        .attr('height', d => d.y1 - d.y0);
+
+    // Append a clipPath to ensure text does not overflow.
+    leaf.append('clipPath')
+        .attr('id', (d) => {
+          // eslint-disable-next-line no-param-reassign
+          d.clipUid = getClipUid();
+          return d.clipUid.id;
+        })
+      .append('use')
+        .attr('xlink:href', d => d.leafUid.href);
+
+    // Append multiline text.
+    leaf.append('text')
+        .attr('clip-path', d => `url(${d.clipUid.href})`)
+      .selectAll('tspan')
+      .data(d => ([
+        // Each element in this array corresponds to a line of text.
+        d.data?.[0]?.at(-1),
+        d.parent?.data?.[0]?.at(-1),
+        `${d.data?.[1].toLocaleString()} ${plur(obsType, d.data?.[1])}`,
+      ]))
+      .join('tspan')
+        .attr('x', 3)
+        // eslint-disable-next-line no-unused-vars
+        .attr('y', (d, i, nodes) => `${(i === nodes.length - 1) * 0.3 + 1.1 + i * 0.9}em`)
+        .text(d => d);
+  }, [width, height, marginLeft, marginBottom, theme, marginTop, marginRight,
+    obsType, sampleType, treemapLeaves, sampleSetColor, sampleSetSelection,
+    obsSetSelection, obsSetColor, obsSetColorScale, sampleSetColorScale,
+    obsColorEncoding, hierarchyLevels,
+  ]);
+
+  return (
+    <svg
+      ref={svgRef}
+      style={{
+        top: 0,
+        left: 0,
+        width: `${width}px`,
+        height: `${height}px`,
+        position: 'relative',
+      }}
+    />
+  );
+}

package/src/TreemapOptions.js

@@ -0,0 +1,90 @@
+import React from 'react';
+import { useId } from 'react-aria';
+import { isEqual } from 'lodash-es';
+import { TableCell, TableRow } from '@material-ui/core';
+import { capitalize } from '@vitessce/utils';
+import {
+  usePlotOptionsStyles, OptionSelect, OptionsContainer,
+} from '@vitessce/vit-s';
+
+export default function TreemapOptions(props) {
+  const {
+    children,
+    obsType,
+    sampleType,
+
+    hierarchyLevels,
+    setHierarchyLevels,
+
+    obsColorEncoding,
+    setObsColorEncoding,
+
+  } = props;
+
+  const treemapOptionsId = useId();
+  const classes = usePlotOptionsStyles();
+
+  function handleColorEncodingChange(event) {
+    setObsColorEncoding(event.target.value);
+  }
+
+  function handleHierarchyLevelsOrderingChange(event) {
+    if (event.target.value === 'sampleSet') {
+      setHierarchyLevels(['sampleSet', 'obsSet']);
+    } else {
+      setHierarchyLevels(['obsSet', 'sampleSet']);
+    }
+  }
+
+  const primaryHierarchyLevel = isEqual(hierarchyLevels, ['sampleSet', 'obsSet']) ? 'sampleSet' : 'obsSet';
+
+  return (
+    <OptionsContainer>
+      {children}
+      <TableRow>
+        <TableCell className={classes.labelCell} variant="head" scope="row">
+          <label
+            htmlFor={`cell-color-encoding-select-${treemapOptionsId}`}
+          >
+            Color Encoding
+          </label>
+        </TableCell>
+        <TableCell className={classes.inputCell} variant="body">
+          <OptionSelect
+            className={classes.select}
+            value={obsColorEncoding}
+            onChange={handleColorEncodingChange}
+            inputProps={{
+              id: `cell-color-encoding-select-${treemapOptionsId}`,
+            }}
+          >
+            <option value="cellSetSelection">{capitalize(obsType)} Sets</option>
+            <option value="sampleSetSelection">{capitalize(sampleType)} Sets</option>
+          </OptionSelect>
+        </TableCell>
+      </TableRow>
+      <TableRow>
+        <TableCell className={classes.labelCell} variant="head" scope="row">
+          <label
+            htmlFor={`treemap-set-hierarchy-levels-${treemapOptionsId}`}
+          >
+            Primary Hierarchy Level
+          </label>
+        </TableCell>
+        <TableCell className={classes.inputCell} variant="body">
+          <OptionSelect
+            className={classes.select}
+            value={primaryHierarchyLevel}
+            onChange={handleHierarchyLevelsOrderingChange}
+            inputProps={{
+              id: `hierarchy-level-select-${treemapOptionsId}`,
+            }}
+          >
+            <option value="obsSet">{capitalize(obsType)} Sets</option>
+            <option value="sampleSet">{capitalize(sampleType)} Sets</option>
+          </OptionSelect>
+        </TableCell>
+      </TableRow>
+    </OptionsContainer>
+  );
+}

package/src/TreemapSubscriber.js

@@ -0,0 +1,261 @@
+/* eslint-disable no-unused-vars */
+import React, { useMemo } from 'react';
+import {
+  TitleInfo,
+  useCoordination,
+  useLoaders,
+  useUrls,
+  useReady,
+  useGridItemSize,
+  useObsFeatureMatrixIndices,
+  useObsSetsData,
+  useSampleEdgesData,
+  useSampleSetsData,
+} from '@vitessce/vit-s';
+import { ViewType, COMPONENT_COORDINATION_TYPES, ViewHelpMapping } from '@vitessce/constants-internal';
+import { treeToSelectedSetMap, treeToSetSizesBySetNames, mergeObsSets } from '@vitessce/sets-utils';
+import { pluralize as plur, commaNumber, unnestMap, capitalize } from '@vitessce/utils';
+import { InternMap } from 'internmap';
+import { isEqual } from 'lodash-es';
+import Treemap from './Treemap.js';
+import { useStyles } from './styles.js';
+import TreemapOptions from './TreemapOptions.js';
+
+const DEFAULT_HIERARCHY_LEVELS = ['obsSet', 'sampleSet'];
+
+export function TreemapSubscriber(props) {
+  const {
+    coordinationScopes,
+    removeGridComponent,
+    theme,
+    helpText = ViewHelpMapping.TREEMAP,
+  } = props;
+
+  const classes = useStyles();
+  const loaders = useLoaders();
+
+  // Get "props" from the coordination space.
+  const [{
+    dataset,
+    obsType,
+    featureType,
+    featureValueType,
+    obsFilter,
+    obsHighlight,
+    obsSetSelection,
+    obsSetFilter,
+    obsSelection,
+    obsSelectionMode,
+    obsSetHighlight,
+    obsSetColor,
+    obsColorEncoding,
+    additionalObsSets,
+    sampleType,
+    sampleSetSelection,
+    sampleSetFilter,
+    sampleSetColor,
+    sampleSelection,
+    sampleSelectionMode,
+    sampleFilter,
+    sampleFilterMode,
+    sampleHighlight,
+    hierarchyLevels,
+  }, {
+    setObsFilter,
+    setObsSelection,
+    setObsSetFilter,
+    setObsSetSelection,
+    setObsSelectionMode,
+    setObsFilterMode,
+    setObsHighlight,
+    setObsSetColor,
+    setObsColorEncoding,
+    setAdditionalObsSets,
+    setSampleFilter,
+    setSampleSetFilter,
+    setSampleFilterMode,
+    setSampleSelection,
+    setSampleSetSelection,
+    setSampleSelectionMode,
+    setSampleHighlight,
+    setSampleSetColor,
+    setHierarchyLevels,
+  }] = useCoordination(
+    COMPONENT_COORDINATION_TYPES[ViewType.TREEMAP],
+    coordinationScopes,
+  );
+
+  const [width, height, containerRef] = useGridItemSize();
+
+  // TODO: how to deal with multimodal cases (multiple obsIndex, one per modality)?
+  const [{ obsIndex }, matrixIndicesStatus, matrixIndicesUrls] = useObsFeatureMatrixIndices(
+    loaders, dataset, false,
+    { obsType, featureType, featureValueType },
+  );
+  const [{ obsSets }, obsSetsStatus, obsSetsUrls] = useObsSetsData(
+    loaders, dataset, true, {}, {},
+    { obsType },
+  );
+
+  const [{ sampleIndex, sampleSets }, sampleSetsStatus, sampleSetsUrls] = useSampleSetsData(
+    loaders,
+    dataset,
+    // TODO: support `false`, i.e., configurations in which
+    // there are no sampleSets
+    true,
+    { setSampleSetColor },
+    { sampleSetColor },
+    { sampleType },
+  );
+
+  const [{ sampleEdges }, sampleEdgesStatus, sampleEdgesUrls] = useSampleEdgesData(
+    loaders,
+    dataset,
+    // TODO: support `false`, i.e., configurations in which
+    // there are no sampleEdges
+    true,
+    {},
+    {},
+    { obsType, sampleType },
+  );
+
+  const isReady = useReady([
+    matrixIndicesStatus,
+    obsSetsStatus,
+    sampleSetsStatus,
+    sampleEdgesStatus,
+  ]);
+  const urls = useUrls([
+    matrixIndicesUrls,
+    obsSetsUrls,
+    sampleSetsUrls,
+    sampleEdgesUrls,
+  ]);
+
+  const mergedObsSets = useMemo(
+    () => mergeObsSets(obsSets, additionalObsSets),
+    [obsSets, additionalObsSets],
+  );
+  const mergedSampleSets = useMemo(
+    () => mergeObsSets(sampleSets, null),
+    [sampleSets],
+  );
+
+  const obsCount = obsIndex?.length || 0;
+  const sampleCount = sampleIndex?.length || 0;
+
+  // TODO: use obsFilter / sampleFilter to display
+  // _all_ cells/samples in gray / transparent in background,
+  // and use obsSetSelection/sampleSetSelection to display
+  // the _selected_ samples in color in the foreground.
+  const [obsCounts, sampleCounts] = useMemo(() => {
+    const obsResult = new InternMap([], JSON.stringify);
+    const sampleResult = new InternMap([], JSON.stringify);
+
+    const hasSampleSetSelection = (
+      Array.isArray(sampleSetSelection)
+      && sampleSetSelection.length > 0
+    );
+    const hasCellSetSelection = (
+      Array.isArray(obsSetSelection)
+      && obsSetSelection.length > 0
+    );
+
+    const sampleSetKeys = hasSampleSetSelection ? sampleSetSelection : [null];
+    const cellSetKeys = hasCellSetSelection ? obsSetSelection : [null];
+
+    // First level: cell set
+    cellSetKeys.forEach((cellSetKey) => {
+      obsResult.set(cellSetKey, new InternMap([], JSON.stringify));
+      // Second level: sample set
+      sampleSetKeys.forEach((sampleSetKey) => {
+        obsResult.get(cellSetKey).set(sampleSetKey, 0);
+      });
+    });
+
+    const sampleSetSizes = treeToSetSizesBySetNames(
+      mergedSampleSets, sampleSetSelection, sampleSetSelection, sampleSetColor, theme,
+    );
+
+    sampleSetKeys.forEach((sampleSetKey) => {
+      const sampleSetSize = sampleSetSizes.find(d => isEqual(d.setNamePath, sampleSetKey))?.size;
+      sampleResult.set(sampleSetKey, sampleSetSize || 0);
+    });
+
+    if (mergedObsSets && obsSetSelection) {
+      const sampleIdToSetMap = sampleSets && sampleSetSelection
+        ? treeToSelectedSetMap(sampleSets, sampleSetSelection)
+        : null;
+      const cellIdToSetMap = treeToSelectedSetMap(mergedObsSets, obsSetSelection);
+
+      for (let i = 0; i < obsIndex.length; i += 1) {
+        const obsId = obsIndex[i];
+
+        const cellSet = cellIdToSetMap?.get(obsId);
+        const sampleId = sampleEdges?.get(obsId);
+        const sampleSet = sampleId ? sampleIdToSetMap?.get(sampleId) : null;
+
+        if (hasSampleSetSelection && !sampleSet) {
+          // Skip this sample if it is not in the selected sample set.
+          // eslint-disable-next-line no-continue
+          continue;
+        }
+        const prevObsCount = obsResult.get(cellSet)?.get(sampleSet);
+        obsResult.get(cellSet)?.set(sampleSet, prevObsCount + 1);
+      }
+    }
+
+    return [
+      unnestMap(obsResult, ['obsSetPath', 'sampleSetPath', 'value']),
+      unnestMap(sampleResult, ['sampleSetPath', 'value']),
+    ];
+  }, [obsIndex, sampleEdges, sampleSets, obsSetColor,
+    sampleSetColor, mergedObsSets, obsSetSelection, mergedSampleSets,
+    // TODO: consider filtering-related coordination values
+  ]);
+
+  return (
+    <TitleInfo
+      title={`Treemap of ${capitalize(plur(obsType, 2))}`}
+      info={`${commaNumber(obsCount)} ${plur(obsType, obsCount)} from ${commaNumber(sampleCount)} ${plur(sampleType, sampleCount)}`}
+      removeGridComponent={removeGridComponent}
+      urls={urls}
+      theme={theme}
+      isReady={isReady}
+      helpText={helpText}
+      options={(
+        <TreemapOptions
+          obsType={obsType}
+          sampleType={sampleType}
+          obsColorEncoding={obsColorEncoding}
+          setObsColorEncoding={setObsColorEncoding}
+          hierarchyLevels={hierarchyLevels || DEFAULT_HIERARCHY_LEVELS}
+          setHierarchyLevels={setHierarchyLevels}
+          // TODO:
+          // - Add option to only include cells in treemap which express selected gene
+          //   above some threshold (kind of like a dot plot)
+          // - Add option to _only_ consider sampleSets or obsSets
+          //   (not both sampleSets and obsSets)
+        />
+      )}
+    >
+      <div ref={containerRef} className={classes.vegaContainer}>
+        <Treemap
+          obsCounts={obsCounts}
+          sampleCounts={sampleCounts}
+          obsColorEncoding={obsColorEncoding}
+          hierarchyLevels={hierarchyLevels || DEFAULT_HIERARCHY_LEVELS}
+          theme={theme}
+          width={width}
+          height={height}
+          obsType={obsType}
+          sampleType={sampleType}
+          obsSetColor={obsSetColor}
+          sampleSetColor={sampleSetColor}
+          obsSetSelection={obsSetSelection}
+          sampleSetSelection={sampleSetSelection}
+        />
+      </div>
+    </TitleInfo>
+  );
+}

package/src/index.js

@@ -3,6 +3,7 @@ export { CellSetSizesPlotSubscriber } from './CellSetSizesPlotSubscriber.js';
 export { ExpressionHistogramSubscriber } from './ExpressionHistogramSubscriber.js';
 export { DotPlotSubscriber } from './DotPlotSubscriber.js';
 export { FeatureBarPlotSubscriber } from './FeatureBarPlotSubscriber.js';
+export { TreemapSubscriber } from './TreemapSubscriber.js';
 export { default as CellSetSizesPlot } from './CellSetSizesPlot.js';
 export { default as CellSetExpressionPlot } from './CellSetExpressionPlot.js';
 export { default as ExpressionHistogram } from './ExpressionHistogram.js';