Merge pull request #771 from hms-dbmi-cellenics/add-granular-select-data

[BIOMAGE-603] - Add capability to select granular cell sets in heatmap

src/__test__/components/plots/styling/SelectData.test.jsx

@@ -38,11 +38,11 @@ describe('Select Data', () => {
     });
 
     // It has the first dropdown
-    expect(screen.getByText(/Select the Cell sets or Metadata that cells are grouped by \(determines the y-axis\)/)).toBeInTheDocument();
+    expect(screen.getByText(/Select the cell sets or metadata that cells are grouped by/i)).toBeInTheDocument();
     expect(screen.getByRole('combobox', { name: 'selectCellSets' }));
 
     // It has the second dropdown
-    expect(screen.getByText(/Select the Cell sets or Metadata to be shown as data:/)).toBeInTheDocument();
+    expect(screen.getByText(/Select the cell sets or metadata to be shown as data/i)).toBeInTheDocument();
     expect(screen.getByRole('combobox', { name: 'selectPoints' }));
   });
 
@@ -79,22 +79,19 @@ describe('Select Data', () => {
     expect(screen.getByText(/Error loading cell set/i)).toBeInTheDocument();
   });
 
-  it('Shows x or y axis properly according to config ', async () => {
-    await act(async () => {
-      render(selectDataFactory());
-    });
-
-    // It shows y-axis by default
-    expect(screen.getByText(/determines the y-axis/)).toBeInTheDocument();
+  it('Renders custom texts properly', async () => {
+    const firstSelectionText = 'First selection text';
+    const secondSelectionText = 'Second selection text';
 
     await act(async () => {
-      render(
-        selectDataFactory({ axisName: 'x' }),
-      );
+      render(selectDataFactory({
+        firstSelectionText,
+        secondSelectionText,
+      }));
     });
 
-    // It shows x-axis if configured
-    expect(screen.getByText(/determines the x-axis/)).toBeInTheDocument();
+    expect(screen.getByText(firstSelectionText)).toBeInTheDocument();
+    expect(screen.getByText(secondSelectionText)).toBeInTheDocument();
   });
 
   it('Changing the first option triggers onUpdate ', async () => {

src/__test__/data/cell_sets_with_scratchpad.json

@@ -84,6 +84,18 @@
                         4,
                         5
                     ]
+                },
+                {
+                    "key": "d273c823-ffa0-4f10-823d-a9ee6edb793f",
+                    "name": "Copied KO",
+                    "color": "#8c564b",
+                    "cellIds": [
+                        6,
+                        7,
+                        8,
+                        9,
+                        10
+                    ]
                 }
             ],
             "type": "cellSets"

src/__test__/pages/experiments/[experimentId]/plots-and-tables/heatmap/index.test.jsx

@@ -6,12 +6,15 @@ import React from 'react';
 import userEvent from '@testing-library/user-event';
 import { act } from 'react-dom/test-utils';
 import { Provider } from 'react-redux';
+import { seekFromS3 } from 'utils/work/seekWorkResponse';
+import expressionDataFAKEGENE from '__test__/data/gene_expression_FAKEGENE.json';
+import cellSetsWithScratchpad from '__test__/data/cell_sets_with_scratchpad.json';
+
 import { loadBackendStatus } from 'redux/actions/backendStatus';
-import { loadPlotConfig, updatePlotConfig } from 'redux/actions/componentConfig';
+import { loadGeneExpression } from 'redux/actions/genes';
+import { updatePlotConfig } from 'redux/actions/componentConfig';
 import { makeStore } from 'redux/store';
 
-import cellSetsWithScratchpad from '__test__/data/cell_sets_with_scratchpad.json';
-
 import preloadAll from 'jest-next-dynamic';
 
 import fake from '__test__/test-utils/constants';
@@ -45,9 +48,34 @@ jest.mock('redux/actions/componentConfig', () => {
   };
 });
 
+// Mock hash so we can control the ETag that is produced by hash.MD5 when fetching work requests
+// EtagParams is the object that's passed to the function which generates ETag in fetchWork
+jest.mock('object-hash', () => {
+  const objectHash = jest.requireActual('object-hash');
+  const mockWorkResultETag = jest.requireActual('__test__/test-utils/mockWorkResultETag').default;
+
+  const mockWorkRequestETag = (ETagParams) => {
+    if (ETagParams.body.name === 'ListGenes') return 'ListGenes';
+    return `${ETagParams.body.nGenes}-marker-genes`;
+  };
+
+  const mockGeneExpressionETag = (ETagParams) => `${ETagParams.missingGenesBody.genes.join('-')}-expression`;
+
+  return mockWorkResultETag(objectHash, mockWorkRequestETag, mockGeneExpressionETag);
+});
+
+jest.mock('utils/work/seekWorkResponse', () => ({
+  __esModule: true,
+  dispatchWorkRequest: jest.fn(() => true),
+  seekFromS3: jest.fn(),
+}));
+
+const mockWorkerResponses = {
+  'FAKEGENE-expression': expressionDataFAKEGENE,
+};
+
 const experimentId = fake.EXPERIMENT_ID;
 const plotUuid = 'heatmapPlotMain';
-const plotType = 'heatmap';
 let storeState = null;
 
 const customAPIResponses = {
@@ -93,7 +121,6 @@ describe('Heatmap plot', () => {
 
     // Set up state for backend status
     await storeState.dispatch(loadBackendStatus(experimentId));
-    await storeState.dispatch(loadPlotConfig(experimentId, plotUuid, plotType));
   });
 
   it('Loads controls and elements', async () => {
@@ -115,7 +142,7 @@ describe('Heatmap plot', () => {
     expect(screen.getByText(/Add some genes to this heatmap to get started/i)).toBeInTheDocument();
   });
 
-  it('Changing clusters updates the plot data', async () => {
+  it('It shows an informative text if there are cell sets to show', async () => {
     await renderHeatmapPage(storeState);
 
     // Open the Select Data panel
@@ -126,9 +153,23 @@ describe('Heatmap plot', () => {
     userEvent.click(screen.getByText(/Custom cell sets/i), null, { skipPointerEventsCheck: true });
 
     expect(updatePlotConfig).toHaveBeenCalled();
+    expect(screen.getByText(/There is no data to show/i)).toBeInTheDocument();
+    expect(screen.getByText(/Select another option from the 'Select data' menu/i)).toBeInTheDocument();
   });
 
-  it('It shows an informative text if there are cell sets to show', async () => {
+  it('Shows the plot if there are custom clusters to show', async () => {
+    const withScratchpadResponse = _.merge(
+      generateDefaultMockAPIResponses(experimentId),
+      customAPIResponses,
+      {
+        [`experiments/${experimentId}/cellSets`]: () => promiseResponse(
+          JSON.stringify(cellSetsWithScratchpad),
+        ),
+      },
+    );
+
+    fetchMock.mockIf(/.*/, mockAPI(withScratchpadResponse));
+
     await renderHeatmapPage(storeState);
 
     // Open the Select Data panel
@@ -139,10 +180,10 @@ describe('Heatmap plot', () => {
     userEvent.click(screen.getByText(/Custom cell sets/i), null, { skipPointerEventsCheck: true });
 
     expect(updatePlotConfig).toHaveBeenCalled();
-    expect(screen.getByText(/There are no custom cell sets to show/i)).toBeInTheDocument();
+    expect(screen.queryByText(/There is no data to show/i)).toBeNull();
   });
 
-  it('Shows the plot if there are custom clusters to show', async () => {
+  it('Changing chosen cluster updates the plot data', async () => {
     const withScratchpadResponse = _.merge(
       generateDefaultMockAPIResponses(experimentId),
       customAPIResponses,
@@ -154,18 +195,32 @@ describe('Heatmap plot', () => {
     );
 
     fetchMock.mockIf(/.*/, mockAPI(withScratchpadResponse));
-    await storeState.dispatch(loadPlotConfig(experimentId, plotUuid, plotType));
+
+    seekFromS3
+      .mockReset()
+      .mockImplementationOnce(() => null)
+      .mockImplementationOnce((Etag) => mockWorkerResponses[Etag]);
 
     await renderHeatmapPage(storeState);
 
+    const genesToLoad = ['FAKEGENE'];
+
+    await act(async () => {
+      await storeState.dispatch(loadGeneExpression(experimentId, genesToLoad, plotUuid));
+    });
+
     // Open the Select Data panel
     userEvent.click(screen.getByText(/Select data/i));
 
-    // Change from Louvain to Custom cell sets
-    userEvent.click(screen.getByText(/Louvain/i));
-    userEvent.click(screen.getByText(/Custom cell sets/i), null, { skipPointerEventsCheck: true });
+    // Change to display another cell set
+    userEvent.click(screen.getByText(/All/i));
+    await act(async () => {
+      userEvent.click(screen.getByText(/Copied KO/i), null, { skipPointerEventsCheck: true });
+    });
 
-    expect(updatePlotConfig).toHaveBeenCalled();
-    expect(screen.queryByText(/There are no custom cell sets to show/i)).toBeNull();
+    // 2 calls:
+    // 1 for the selected gene (loadGeneExpression)
+    // 1 for changing the cell set
+    expect(updatePlotConfig).toHaveBeenCalledTimes(2);
   });
 });

src/__test__/pages/experiments/[experimentId]/plots-and-tables/marker-heatmap/index.test.jsx

@@ -1,6 +1,6 @@
 import { render, screen } from '@testing-library/react';
 import { mount } from 'enzyme';
-import { within, fireEvent } from '@testing-library/dom';
+import { within } from '@testing-library/dom';
 import userEvent from '@testing-library/user-event';
 import fetchMock, { enableFetchMocks } from 'jest-fetch-mock';
 import _ from 'lodash';
@@ -106,7 +106,7 @@ const getCurrentGeneOrder = (component) => {
   });
   newOrder.splice(0, 1);
   return newOrder;
-}
+};
 
 const renderHeatmapPage = async (store) => {
   await act(async () => (
@@ -167,7 +167,7 @@ describe('Marker heatmap plot', () => {
     expect(screen.getByText(/Colours/i)).toBeInTheDocument();
     expect(screen.getByText(/Legend/i)).toBeInTheDocument();
   });
-  
+
   it('Loads the plot', async () => {
     await renderHeatmapPage(storeState);
 
@@ -261,6 +261,24 @@ describe('Marker heatmap plot', () => {
     expect(_.isEqual(displayedGenesList, genesToLoad)).toEqual(false);
   });
 
+  it('Shows an information text if a selected cell set does not contain enough number of samples', async () => {
+    await renderHeatmapPage(storeState);
+
+    // Open the toggle
+    userEvent.click(screen.getByText(/Select data/i));
+
+    // Click custom cell sets
+    userEvent.click(screen.getByText(/louvain clusters/i));
+    userEvent.click(screen.getByText(/Custom cell sets/i), null, { skipPointerEventsCheck: true });
+
+    // It shouldn't show the plot
+    expect(screen.queryByRole('graphics-document', { name: 'Marker heatmap' })).toBeNull();
+
+    // There is an error message
+    expect(screen.getByText(/There is no data to show/i)).toBeInTheDocument();
+    expect(screen.getByText(/Select another option from the 'Select data' menu/i)).toBeInTheDocument();
+  });
+
   it('Shows an error message if gene expression fails to load', async () => {
     seekFromS3
       .mockReset()

src/components/data-exploration/heatmap/HeatmapPlot.jsx

@@ -24,6 +24,7 @@ import HeatmapTracksCellInfo from 'components/data-exploration/heatmap/HeatmapTr
 
 import getContainingCellSetsProperties from 'utils/cellSets/getContainingCellSetsProperties';
 import useConditionalEffect from 'utils/customHooks/useConditionalEffect';
+import generateVitessceData from 'components/plots/helpers/heatmap/vitessce/generateVitessceData';
 
 const COMPONENT_TYPE = 'interactiveHeatmap';
 
@@ -134,10 +135,10 @@ const HeatmapPlot = (props) => {
       return;
     }
 
-    const data = populateHeatmapData(
-      cellSets, heatmapSettings, expressionData, selectedGenes, true, true,
-    );
+    const cellOrder = populateHeatmapData(cellSets, heatmapSettings, true);
 
+    heatmapSettings.selectedGenes = selectedGenes;
+    const data = generateVitessceData(cellOrder, heatmapSettings, expressionData, cellSets);
     setHeatmapData(data);
   }, [
     selectedGenes,

src/components/plots/helpers/heatmap/populateHeatmapData.js

@@ -1,18 +1,16 @@
 import _ from 'lodash';
 
-import generateVegaData from 'components/plots/helpers/heatmap/vega/generateVegaData';
-import generateVitessceData from 'components/plots/helpers/heatmap/vitessce/generateVitessceData';
-
 import SetOperations from 'utils/setOperations';
 import { union } from 'utils/cellSetOperations';
-import { reversed } from 'utils/arrayUtils';
 
 const populateHeatmapData = (
-  cellSets, heatmapSettings, expression,
-  selectedGenes, downsampling = false, vitessce = false,
+  cellSets, heatmapSettings,
+  downsampling = false,
 ) => {
   const { hierarchy, properties, hidden } = cellSets;
-  const { selectedTracks, groupedTracks } = heatmapSettings;
+  const {
+    groupedTracks, selectedCellSet, selectedPoints,
+  } = heatmapSettings;
 
   const maxCells = 1000;
   const getCellsInSet = (cellSetName) => properties[cellSetName].cellIds;
@@ -80,18 +78,21 @@ const populateHeatmapData = (
   };
 
   const getAllEnabledCellIds = () => {
-    // we want to avoid displaying elements which are not in a louvain cluster
-    // so initially consider as enabled only cells in louvain clusters
-    // See: https://biomage.atlassian.net/browse/BIOMAGE-809
-    const selectedCellSet = heatmapSettings?.selectedCellSet ? heatmapSettings.selectedCellSet : 'louvain';
-    const selectedClusters = hierarchy.find(
-      (clusters) => clusters.key === selectedCellSet,
-    );
-    const cellIsInLouvainCluster = getCellsSetInGroup(selectedClusters);
+    let cellIds;
+
+    if (selectedPoints === 'All') {
+      const selectedClusters = hierarchy.find(
+        (clusters) => clusters.key === selectedCellSet,
+      );
+      cellIds = getCellsSetInGroup(selectedClusters);
+    } else {
+      const cellSetKey = selectedPoints.split('/')[1];
+      cellIds = getCellsInSet(cellSetKey);
+    }
 
     // Remove cells from groups marked as hidden by the user in the UI.
     const hiddenCellIds = union(Array.from(hidden), properties);
-    const enabledCellIds = new Set([...cellIsInLouvainCluster]
+    const enabledCellIds = new Set([...cellIds]
       .filter((cellId) => !hiddenCellIds.has(cellId)));
 
     return enabledCellIds;
@@ -141,23 +142,22 @@ const populateHeatmapData = (
 
     return cellIds;
   };
-  // For now, this is statically defined. In the future, these values are
-  // controlled from the settings panel in the heatmap.
-
-  // Do downsampling and return cellIds with their order by groupings.
-  const cellOrder = generateCellOrder(groupedTracks);
-  const geneOrder = selectedGenes;
-
-  if (!vitessce) {
-    return generateVegaData(
-      cellOrder, geneOrder, reversed(selectedTracks),
-      expression, heatmapSettings, cellSets,
-    );
-  }
-
-  return generateVitessceData(
-    cellOrder, geneOrder, selectedTracks,
-    expression, heatmapSettings, cellSets,
+
+  let firstOptionCellIds = hierarchy.find(
+    (cellSet) => cellSet.key === selectedCellSet,
+  ).children.reduce(
+    (acc, cellSet) => [
+      ...acc,
+      ...properties[cellSet.key].cellIds,
+    ], [],
   );
+
+  firstOptionCellIds = new Set(firstOptionCellIds);
+
+  const filteredCellOrder = generateCellOrder(groupedTracks)
+    .filter((cellId) => firstOptionCellIds.has(cellId));
+
+  return filteredCellOrder;
 };
+
 export default populateHeatmapData;