feat: data stratification (#361)

- Replace the tabbed interface for data uploads with a stepper.
- Add utilities to group subjects by protocol.
- Add stratification to the stepper.
- Stratify by dose protocols for the time being.
- Display each protocol group as a MUI data grid.
- Display protocol groups as tabs in the trial design view.
- Add the dosing compartment qname to the protocol model.
- Add a 'cohort' field to the CSV.

Author: eatyourgreens

frontend-v2/package.json

@@ -8,6 +8,7 @@
     "@fontsource/comfortaa": "^5.0.18",
     "@mui/icons-material": "^5.11.16",
     "@mui/material": "^5.12.2",
+    "@mui/x-data-grid": "^6.19.6",
     "@reduxjs/toolkit": "^1.8.1",
     "@testing-library/jest-dom": "^5.16.4",
     "@testing-library/react": "^13.0.1",

frontend-v2/src/app/backendApi.ts

@@ -2004,6 +2004,7 @@ export type Protocol = {
   datetime?: string | null;
   name: string;
   dose_type?: DoseTypeEnum;
+  mapped_qname?: string;
   project?: number | null;
   compound?: number | null;
   time_unit?: number | null;
@@ -2019,6 +2020,7 @@ export type ProtocolRead = {
   datetime?: string | null;
   name: string;
   dose_type?: DoseTypeEnum;
+  mapped_qname?: string;
   project?: number | null;
   compound?: number | null;
   time_unit?: number | null;
@@ -2381,6 +2383,7 @@ export type PatchedProtocol = {
   datetime?: string | null;
   name?: string;
   dose_type?: DoseTypeEnum;
+  mapped_qname?: string;
   project?: number | null;
   compound?: number | null;
   time_unit?: number | null;
@@ -2396,6 +2399,7 @@ export type PatchedProtocolRead = {
   datetime?: string | null;
   name?: string;
   dose_type?: DoseTypeEnum;
+  mapped_qname?: string;
   project?: number | null;
   compound?: number | null;
   time_unit?: number | null;

frontend-v2/src/features/data/Data.tsx

@@ -1,26 +1,21 @@
-import { DynamicTabs, TabPanel } from "../../components/DynamicTabs";
-import LoadDataTab from "./LoadDataTab";
-import { SubPageName } from "../main/mainSlice";
+import { FC, useState } from "react";
+import { Button } from "@mui/material";
+import LoadDataStepper from "./LoadDataStepper";
 
-const Data: React.FC = () => {
-  const tabKeys = [
-    SubPageName.LOAD_DATA,
-    SubPageName.STRATIFICATION,
-    SubPageName.VISUALISATION
-  ];
-  return (
-    <DynamicTabs tabNames={tabKeys}>
-      <TabPanel>
-        <LoadDataTab  />
-      </TabPanel>
-      <TabPanel>
-        <LoadDataTab  />
-      </TabPanel>
-      <TabPanel>
-        <LoadDataTab  />
-      </TabPanel>
-    </DynamicTabs>
-  );
+const Data:FC = () => {
+  const [isLoading, setIsLoading] = useState(false);
+  function handleNewUpload() {
+    setIsLoading(true);
+  }
+  function onUploadComplete() {
+    setIsLoading(false);
+  }
+
+  return isLoading ?
+    <LoadDataStepper onFinish={onUploadComplete}  /> :
+    <Button variant="outlined" onClick={handleNewUpload}>
+      Upload new dataset
+    </Button>;
 }
 
 export default Data;

frontend-v2/src/features/data/LoadDataStepper.tsx

@@ -13,15 +13,16 @@ import MapObservations from './MapObservations';
 import MapDosing from './MapDosing';
 import PreviewData from './PreviewData';
 import { RootState } from "../../app/store";
-import {
-  DatasetRead,
-  useDatasetListQuery,
-  useDatasetCreateMutation,
-  useDatasetCsvUpdateMutation,
-} from '../../app/backendApi';
+import { DatasetRead, useDatasetCsvUpdateMutation } from '../../app/backendApi';
+import Stratification from './Stratification';
+import useDataset from '../../hooks/useDataset';
 
-const stepLabels = ['Upload Data', 'Map Dosing', 'Map Observations', 'Preview Dataset'];
-const stepComponents = [LoadData, MapDosing, MapObservations, PreviewData];
+interface IStepper {
+  onFinish: () => void
+}
+
+const stepLabels = ['Upload Data', 'Map Dosing', 'Map Observations', 'Stratification', 'Preview Dataset'];
+const stepComponents = [LoadData, MapDosing, MapObservations, Stratification, PreviewData];
 
 type Row = {[key: string]: string};
 type Data = Row[];
@@ -40,8 +41,7 @@ export type StepperState = {
   setAmountUnit: (amountUnit: string) => void;
 }
 
-const LoadDataStepper: FC = () => {
-  const [dataset, setDataset] = useState<null | DatasetRead>(null);
+const LoadDataStepper: FC<IStepper> = ({ onFinish }) => {
   const [data, setData] = useState<Data>([]);
   const [fields, setFields] = useState<string[]>([]);
   const [normalisedFields, setNormalisedFields] = useState<string[]>([]);
@@ -50,18 +50,10 @@ const LoadDataStepper: FC = () => {
   const selectedProject = useSelector(
     (state: RootState) => state.main.selectedProject,
   );
-  const selectedProjectOrZero = selectedProject || 0;
-  const { data: datasets = [], isLoading: isDatasetLoading } = useDatasetListQuery(
-    { projectId: selectedProjectOrZero },
-    { skip: !selectedProject },
-  );
-  const [
-    createDataset
-  ] = useDatasetCreateMutation();
   const [
-    updateDataset
+    updateDatasetCsv
   ] = useDatasetCsvUpdateMutation();
-
+  const { dataset, updateDataset } = useDataset(selectedProject);
 
   const state = {
     fields,
@@ -80,43 +72,26 @@ const LoadDataStepper: FC = () => {
   const StepComponent = stepComponents[stepState.activeStep];
   const isFinished = stepState.activeStep === stepLabels.length;
 
-  useEffect(function onDataLoad() {
-    async function addDataset() {
-      let [dataset] = datasets;
-      if (!dataset) {
-        const response = await createDataset({ 
-          dataset: { 
-            name: 'New Dataset',
-            project: selectedProjectOrZero,
-          }
-        });
-        if ('data' in response && response.data) {
-          dataset = response.data;
-        }
-      }
-      console.log({dataset});
-      setDataset(dataset);
-    }
-    if (!isDatasetLoading) {
-      addDataset();
-    }
-  }, [datasets, createDataset, isDatasetLoading]);
-
   useEffect(function onFinished() {
     if (isFinished && dataset?.id) {
       try {
         const csv = Papa.unparse(data);
-        updateDataset({
+        updateDatasetCsv({
           id: dataset.id,
           datasetCsv: {
             csv
           }
         })
+        .unwrap()
+        .then(data => {
+          updateDataset(data as unknown as DatasetRead);
+          onFinish();
+        });
       } catch (e) {
         console.error(e);
       }
     }
-  }, [isFinished, updateDataset, dataset?.id, data])
+  }, [isFinished, onFinish, updateDatasetCsv, updateDataset, dataset?.id, data])
 
   const handleNext = () => {
     setStepState((prevActiveStep) => ({ 

frontend-v2/src/features/data/LoadDataTab.tsx

@@ -84,7 +84,7 @@ const MapDosing: FC<IMapDosing> = ({ state, firstTime }: IMapDosing) => {
     const { value } = event.target;
     nextData.filter(row => administrationIdField ? row[administrationIdField] === id : true)
       .forEach(row => {
-        row['Dosing Variable'] = value;
+        row['Amount Variable'] = value;
       })
     state.setData(nextData);
   }
@@ -134,7 +134,7 @@ const MapDosing: FC<IMapDosing> = ({ state, firstTime }: IMapDosing) => {
           <TableBody>
             {uniqueAdministrationIds.map((adminId, index) => {
               const currentRow = dosingRows.find(row => administrationIdField ? row[administrationIdField] === adminId : true);
-              const selectedVariable = variables?.find(variable => variable.qname === currentRow?.['Dosing Variable']);
+              const selectedVariable = variables?.find(variable => variable.qname === currentRow?.['Amount Variable']);
               const compatibleUnits = units?.find(unit => unit.id === selectedVariable?.unit)?.compatible_units;
               const adminUnit = amountUnitField && currentRow && currentRow[amountUnitField];
               const amount = amountField && currentRow?.[amountField];

frontend-v2/src/features/data/MapDosing.tsx

@@ -11,7 +11,8 @@ const PreviewData: FC<IPreviewData> = ({ state, firstTime }: IPreviewData) => {
   const { data } = state;
   const fields = [
     ...state.fields,
-    'Dosing Variable',
+    'cohort',
+    'Amount Variable',
     'Observation Variable'
   ];
   if (!state.normalisedFields.find(field => field === 'Amount Unit')) {

frontend-v2/src/features/data/PreviewData.tsx

@@ -0,0 +1,33 @@
+import { FC } from 'react';
+import { DataGrid } from '@mui/x-data-grid';
+import { IProtocol } from './protocolUtils';
+import { StepperState } from './LoadDataStepper';
+
+interface IProtocolDataGrid {
+  protocol: IProtocol
+  state: StepperState
+}
+
+const ProtocolDataGrid: FC<IProtocolDataGrid> = ({ protocol, state }) => {
+  const idField = state.fields.find((field, index) => state.normalisedFields[index] === 'ID');
+  const amountField = state.fields.find((field, index) => state.normalisedFields[index] === 'Amount');
+  const { subjects } = protocol;
+  const protocolRows = state.data.filter(row => {
+    const subjectId = idField && row[idField];
+    const amount = amountField && +row[amountField];
+    return subjects.includes(subjectId || '') && amount;
+  }).map(row => {
+    const subjectId = (idField && +row[idField]) || 0;
+    return { id: +subjectId, ...row };
+  });
+  const protocolColumns = state.fields.map((field) => ({ field, headerName: field }));
+  return (
+    <DataGrid
+      rows={protocolRows}
+      columns={protocolColumns}
+      checkboxSelection
+    />
+  );
+}
+
+export default ProtocolDataGrid;