[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

@@ -7,6 +7,11 @@ import userEvent from '@testing-library/user-event';
 import { act } from 'react-dom/test-utils';
 import { Provider } from 'react-redux';
 import { loadBackendStatus } from 'redux/actions/backendStatus';
+
+import { seekFromS3 } from 'utils/work/seekWorkResponse';
+import expressionDataFAKEGENE from '__test__/data/gene_expression_FAKEGENE.json';
+
+import { loadGeneExpression } from 'redux/actions/genes';
 import { loadPlotConfig, updatePlotConfig } from 'redux/actions/componentConfig';
 import { makeStore } from 'redux/store';
 
@@ -45,6 +50,32 @@ 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';
@@ -93,7 +124,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 +145,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 +156,22 @@ 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();
   });
 
-  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 +182,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 are no custom cell sets 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 +197,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);
 

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

@@ -12,7 +12,9 @@ const populateHeatmapData = (
   selectedGenes, downsampling = false, vitessce = false,
 ) => {
   const { hierarchy, properties, hidden } = cellSets;
-  const { selectedTracks, groupedTracks } = heatmapSettings;
+  const {
+    selectedTracks, groupedTracks, selectedCellSet, selectedPoints,
+  } = heatmapSettings;
 
   const maxCells = 1000;
   const getCellsInSet = (cellSetName) => properties[cellSetName].cellIds;
@@ -80,18 +82,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;
@@ -144,8 +149,21 @@ const populateHeatmapData = (
   // For now, this is statically defined. In the future, these values are
   // controlled from the settings panel in the heatmap.
 
+  // Get cellIds of the base set that is chosen in the first option
+  let baseCellIds = hierarchy.find(
+    (cellSet) => cellSet.key === selectedCellSet,
+  ).children.reduce(
+    (acc, cellSet) => [
+      ...acc,
+      ...properties[cellSet.key].cellIds,
+    ], [],
+  );
+
+  baseCellIds = new Set(baseCellIds);
+
   // Do downsampling and return cellIds with their order by groupings.
-  const cellOrder = generateCellOrder(groupedTracks);
+  const cellOrder = generateCellOrder(groupedTracks)
+    .filter((cellId) => baseCellIds.has(cellId));
   const geneOrder = selectedGenes;
 
   if (!vitessce) {

src/components/plots/helpers/heatmap/vega/generateVegaData.js

@@ -14,14 +14,12 @@ const generateVegaData = (
     trackGroupData: [],
   };
 
-  const cells = new Set(cellOrder);
-
   data.geneExpressionsData = generateVegaGeneExpressionsData(
     cellOrder, geneOrder, expression, heatmapSettings,
   );
 
   const trackData = trackOrder.map((rootNode) => generateVegaHeatmapTracksData(
-    cells,
+    cellOrder,
     rootNode,
     cellSets,
     heatmapSettings,

src/components/plots/helpers/heatmap/vega/utils/generateVegaHeatmapTracksData.js

@@ -1,4 +1,4 @@
-const generateVegaHeatmapTracksData = (cells, track, cellSets, heatmapSettings) => {
+const generateVegaHeatmapTracksData = (cellOrder, track, cellSets, heatmapSettings) => {
   const { hierarchy, properties } = cellSets;
 
   const getCellClusterFromCellId = (clusters, cellId) => {
@@ -27,15 +27,18 @@ const generateVegaHeatmapTracksData = (cells, track, cellSets, heatmapSettings)
 
   const clusterSeparationLines = [];
   if (heatmapSettings.guardLines) {
-    let currentCluster = getCellClusterFromCellId(childrenCellSets, cells[0]);
-    cells.forEach((cell) => {
+    let currentCluster = getCellClusterFromCellId(childrenCellSets, cellOrder[0]);
+    cellOrder.forEach((cell) => {
       const isTheSameCluster = properties[currentCluster]?.cellIds?.has(cell);
       if (!isTheSameCluster) {
         currentCluster = getCellClusterFromCellId(childrenCellSets, cell);
         clusterSeparationLines.push(cell);
       }
     });
   }
+
+  const cellOrderSet = new Set(cellOrder);
+
   childrenCellSets.forEach(({ key }) => {
     const { cellIds, name, color } = properties[key];
 
@@ -47,9 +50,7 @@ const generateVegaHeatmapTracksData = (cells, track, cellSets, heatmapSettings)
       trackName: properties[track].name,
     });
 
-    const intersectionSet = [cellIds, cells].reduce(
-      (acc, curr) => new Set([...acc].filter((x) => curr.has(x))),
-    );
+    const intersectionSet = [...cellIds].filter((x) => cellOrderSet.has(x));
 
     intersectionSet.forEach((cellId) => trackColorData.push({
       cellId,