cross corr should work with stimulus data, untested

mesmerize/plotting/widgets/cross_correlation/control_widget.ui

@@ -199,6 +199,22 @@
              </item>
             </layout>
            </item>
+           <item>
+            <widget class="QLabel" name="label_13">
+             <property name="font">
+              <font>
+               <weight>75</weight>
+               <bold>true</bold>
+              </font>
+             </property>
+             <property name="text">
+              <string>Stimulus Type:</string>
+             </property>
+            </widget>
+           </item>
+           <item>
+            <widget class="QComboBox" name="comboBoxStimulusType"/>
+           </item>
           </layout>
          </item>
         </layout>
@@ -349,6 +365,23 @@
    <header>....pyqtgraphCore</header>
   </customwidget>
  </customwidgets>
+ <tabstops>
+  <tabstop>comboBoxLabelsColumn</tabstop>
+  <tabstop>radioButtonLag</tabstop>
+  <tabstop>radioButtonMaxima</tabstop>
+  <tabstop>checkBoxAbsoluteValue</tabstop>
+  <tabstop>doubleSpinBoxLagThreshold</tabstop>
+  <tabstop>doubleSpinBoxMaximaThreshold</tabstop>
+  <tabstop>comboBoxStimulusType</tabstop>
+  <tabstop>pushButtonComputeAllData</tabstop>
+  <tabstop>pushButtonExportAllData</tabstop>
+  <tabstop>lineEditCurve1UUID</tabstop>
+  <tabstop>graphicsViewCurve1</tabstop>
+  <tabstop>lineEditCurve2UUID</tabstop>
+  <tabstop>graphicsViewCurve2</tabstop>
+  <tabstop>listWidgetSampleID</tabstop>
+  <tabstop>comboBoxDataColumn</tabstop>
+ </tabstops>
  <resources/>
  <connections/>
 </ui>

mesmerize/plotting/widgets/cross_correlation/control_widget_pytemplate.py

@@ -113,6 +113,16 @@ def setupUi(self, CrossCorrelationControls):
         self.verticalLayout_4.addWidget(self.doubleSpinBoxMaximaThreshold)
         self.horizontalLayout.addLayout(self.verticalLayout_4)
         self.verticalLayout_7.addLayout(self.horizontalLayout)
+        self.label_13 = QtWidgets.QLabel(self.layoutWidget)
+        font = QtGui.QFont()
+        font.setBold(True)
+        font.setWeight(75)
+        self.label_13.setFont(font)
+        self.label_13.setObjectName("label_13")
+        self.verticalLayout_7.addWidget(self.label_13)
+        self.comboBoxStimulusType = QtWidgets.QComboBox(self.layoutWidget)
+        self.comboBoxStimulusType.setObjectName("comboBoxStimulusType")
+        self.verticalLayout_7.addWidget(self.comboBoxStimulusType)
         self.horizontalLayout_3.addLayout(self.verticalLayout_7)
         self.verticalLayoutLeftMain.addLayout(self.horizontalLayout_3)
         self.horizontalLayout_6 = QtWidgets.QHBoxLayout()
@@ -204,6 +214,20 @@ def setupUi(self, CrossCorrelationControls):
 
         self.retranslateUi(CrossCorrelationControls)
         QtCore.QMetaObject.connectSlotsByName(CrossCorrelationControls)
+        CrossCorrelationControls.setTabOrder(self.comboBoxLabelsColumn, self.radioButtonLag)
+        CrossCorrelationControls.setTabOrder(self.radioButtonLag, self.radioButtonMaxima)
+        CrossCorrelationControls.setTabOrder(self.radioButtonMaxima, self.checkBoxAbsoluteValue)
+        CrossCorrelationControls.setTabOrder(self.checkBoxAbsoluteValue, self.doubleSpinBoxLagThreshold)
+        CrossCorrelationControls.setTabOrder(self.doubleSpinBoxLagThreshold, self.doubleSpinBoxMaximaThreshold)
+        CrossCorrelationControls.setTabOrder(self.doubleSpinBoxMaximaThreshold, self.comboBoxStimulusType)
+        CrossCorrelationControls.setTabOrder(self.comboBoxStimulusType, self.pushButtonComputeAllData)
+        CrossCorrelationControls.setTabOrder(self.pushButtonComputeAllData, self.pushButtonExportAllData)
+        CrossCorrelationControls.setTabOrder(self.pushButtonExportAllData, self.lineEditCurve1UUID)
+        CrossCorrelationControls.setTabOrder(self.lineEditCurve1UUID, self.graphicsViewCurve1)
+        CrossCorrelationControls.setTabOrder(self.graphicsViewCurve1, self.lineEditCurve2UUID)
+        CrossCorrelationControls.setTabOrder(self.lineEditCurve2UUID, self.graphicsViewCurve2)
+        CrossCorrelationControls.setTabOrder(self.graphicsViewCurve2, self.listWidgetSampleID)
+        CrossCorrelationControls.setTabOrder(self.listWidgetSampleID, self.comboBoxDataColumn)
 
     def retranslateUi(self, CrossCorrelationControls):
         _translate = QtCore.QCoreApplication.translate
@@ -218,6 +242,7 @@ def retranslateUi(self, CrossCorrelationControls):
         self.label_3.setText(_translate("CrossCorrelationControls", "lag"))
         self.checkBoxAbsoluteValue.setText(_translate("CrossCorrelationControls", "absolute value"))
         self.label_4.setText(_translate("CrossCorrelationControls", "maxima"))
+        self.label_13.setText(_translate("CrossCorrelationControls", "Stimulus Type:"))
         self.pushButtonComputeAllData.setText(_translate("CrossCorrelationControls", "Compute all data"))
         self.pushButtonExportAllData.setText(_translate("CrossCorrelationControls", "Export all data"))
         self.label_6.setText(_translate("CrossCorrelationControls", "Curve 1"))

mesmerize/plotting/widgets/cross_correlation/stims.py

@@ -0,0 +1,51 @@
+"""
+@author: kushal
+
+Chatzigeorgiou Group
+Sars International Centre for Marine Molecular Biology
+
+GNU GENERAL PUBLIC LICENSE Version 3, 29 June 2007
+"""
+
+import pandas as pd
+import numpy as np
+from typing import Tuple, List
+
+
+def get_binary_stims(
+        stim_df: pd.DataFrame,
+        index_size: int,
+        start_offset: int = 0,
+        end_offset: int = 0
+) -> Tuple[np.ndarray, List[str]]:
+    """
+
+    """
+
+    stims = stim_df['name'].unique()
+
+    if index_size < stim_df['end'].max():
+        index_size = int(stim_df['end'].max())
+
+    binarized_arrays = []
+
+    for stimulus in stims:
+        sub_df = stim_df[stim_df['name'] == stimulus]
+
+        # One binary stim array for each stimulus in the `stim_df`
+        stim_array = np.empty(index_size, dtype=np.uint8)
+        stim_array[:] = 0
+
+        # fill stim_array
+        # stim_array is used for fancy indexing
+        for ix, r in sub_df.iterrows():  # iterate through the stimulus periods
+            # apply any offsets
+            start_ix = r['start'] + start_offset
+            end_ix = r['end'] + end_offset
+
+            # set this period to 1 in the array
+            stim_array[int(start_ix):int(end_ix)] = 1
+
+        binarized_arrays.append(stim_array)
+
+    return np.vstack(binarized_arrays), stims

mesmerize/plotting/widgets/cross_correlation/widget.py

@@ -21,6 +21,7 @@
 import pandas as pd
 from ....common.qdialogs import *
 from ....common.utils import HdfTools
+from .stims import get_binary_stims
 
 
 class ControlWidget(QtWidgets.QWidget):
@@ -203,8 +204,14 @@ def set_input(self, transmission: Transmission = None):
         self.control_widget.ui.comboBoxDataColumn.clear()
         self.control_widget.ui.comboBoxDataColumn.addItems(dcols)
 
+        self.control_widget.ui.comboBoxLabelsColumn.clear()
         self.control_widget.ui.comboBoxLabelsColumn.addItems(ccols)
 
+        self.control_widget.ui.comboBoxStimulusType.clear()
+        self.control_widget.ui.comboBoxStimulusType.addItems(
+            self.transmission.STIM_DEFS
+        )
+
         self.reset_sample_id_list_widget()
         self.control_widget.ui.listWidgetSampleID.setCurrentRow(0)
 
@@ -230,14 +237,38 @@ def compute_dataframe(self):
         for sample_id in self.sample_list:
             sub_df = self.transmission.df[self.transmission.df.SampleID == sample_id]
 
+            if self.transmission.STIM_DEFS:
+                stim_type = self.control_widget.ui.comboBoxStimulusType.currentText()
+
+                stim_df = sub_df.iloc[0].stim_maps[0][0][stim_type]
+                index_size = sub_df[self.data_column].values[0].size
+
+                binary_stims_array, stim_names = get_binary_stims(
+                    stim_df=stim_df,
+                    index_size=index_size,
+                    start_offset=0,
+                    end_offset=0
+                )
+
             data = np.vstack(sub_df[self.data_column].values)
+
+            # append the stimulus array too
+            if self.transmission.STIM_DEFS:
+                data = np.vstack([data, binary_stims_array])
+
             r = get_sampling_rate(self.transmission)
             self.sampling_rate = r
 
             self.cc_data[sample_id] = compute_cc_data(data)
             self.cc_data[sample_id].lag_matrix = np.true_divide(self.cc_data[sample_id].lag_matrix, r)
             self.cc_data[sample_id].curve_uuids = np.array(list(map(str, sub_df['uuid_curve'].values))) # convert all UUIDs to str representation
-            self.cc_data[sample_id].labels = sub_df[labels_col].values.astype(np.unicode)
+
+            labels = sub_df[labels_col].values.astype(np.unicode)
+
+            if self.transmission.STIM_DEFS:
+                self.cc_data[sample_id].labels = np.concatenate([labels, stim_names]).astype(np.unicode)
+            else:
+                self.cc_data[sample_id].labels = labels
 
         self.set_current_sample()