Merge pull request #478 from int-brain-lab/training_plots

Training plots

Author: mayofaulkner

brainbox/behavior/training.py

@@ -349,7 +349,7 @@ def compute_training_info(trials, trials_all):
     perf_easy = np.array([compute_performance_easy(trials[k]) for k in trials.keys()])
     n_trials = np.array([compute_n_trials(trials[k]) for k in trials.keys()])
     psych = compute_psychometric(trials_all, signed_contrast=signed_contrast)
-    rt = compute_median_reaction_time(trials_all, signed_contrast=signed_contrast)
+    rt = compute_median_reaction_time(trials_all, contrast=0, signed_contrast=signed_contrast)
 
     return perf_easy, n_trials, psych, rt
 
@@ -376,7 +376,7 @@ def compute_bias_info(trials, trials_all):
     n_trials = np.array([compute_n_trials(trials[k]) for k in trials.keys()])
     psych_20 = compute_psychometric(trials_all, signed_contrast=signed_contrast, block=0.2)
     psych_80 = compute_psychometric(trials_all, signed_contrast=signed_contrast, block=0.8)
-    rt = compute_median_reaction_time(trials_all, signed_contrast=signed_contrast)
+    rt = compute_median_reaction_time(trials_all, contrast=0, signed_contrast=signed_contrast)
 
     return perf_easy, n_trials, psych_20, psych_80, rt
 
@@ -452,7 +452,7 @@ def compute_n_trials(trials):
     return trials['choice'].shape[0]
 
 
-def compute_psychometric(trials, signed_contrast=None, block=None):
+def compute_psychometric(trials, signed_contrast=None, block=None, plotting=False):
     """
     Compute psychometric fit parameters for trials object
 
@@ -479,17 +479,27 @@ def compute_psychometric(trials, signed_contrast=None, block=None):
     prob_choose_right, contrasts, n_contrasts = compute_performance(trials, signed_contrast=signed_contrast, block=block,
                                                                     prob_right=True)
 
-    psych, _ = psy.mle_fit_psycho(
-        np.vstack([contrasts, n_contrasts, prob_choose_right]),
-        P_model='erf_psycho_2gammas',
-        parstart=np.array([np.mean(contrasts), 20., 0.05, 0.05]),
-        parmin=np.array([np.min(contrasts), 0., 0., 0.]),
-        parmax=np.array([np.max(contrasts), 100., 1, 1]))
+    if plotting:
+        psych, _ = psy.mle_fit_psycho(
+            np.vstack([contrasts, n_contrasts, prob_choose_right]),
+            P_model='erf_psycho_2gammas',
+            parstart=np.array([0., 40., 0.1, 0.1]),
+            parmin=np.array([-50., 10., 0., 0.]),
+            parmax=np.array([50., 50., 0.2, 0.2]),
+            nfits=10)
+    else:
+
+        psych, _ = psy.mle_fit_psycho(
+            np.vstack([contrasts, n_contrasts, prob_choose_right]),
+            P_model='erf_psycho_2gammas',
+            parstart=np.array([np.mean(contrasts), 20., 0.05, 0.05]),
+            parmin=np.array([np.min(contrasts), 0., 0., 0.]),
+            parmax=np.array([np.max(contrasts), 100., 1, 1]))
 
     return psych
 
 
-def compute_median_reaction_time(trials, stim_on_type='stimOn_times', signed_contrast=None):
+def compute_median_reaction_time(trials, stim_on_type='stimOn_times', contrast=None, signed_contrast=None):
     """
     Compute median reaction time on zero contrast trials from trials object
 
@@ -505,10 +515,15 @@ def compute_median_reaction_time(trials, stim_on_type='stimOn_times', signed_con
     """
     if signed_contrast is None:
         signed_contrast = get_signed_contrast(trials)
-    zero_trials = (trials.response_times - trials[stim_on_type])[signed_contrast == 0]
-    if np.any(zero_trials):
+
+    if contrast is None:
+        contrast_idx = np.full(trials.probabilityLeft.shape, True, dtype=bool)
+    else:
+        contrast_idx = signed_contrast == contrast
+
+    if np.any(contrast_idx):
         reaction_time = np.nanmedian((trials.response_times - trials[stim_on_type])
-                                     [signed_contrast == 0])
+                                     [contrast_idx])
     else:
         reaction_time = np.nan
 
@@ -590,15 +605,15 @@ def plot_psychometric(trials, ax=None, title=None, **kwargs):
     contrasts_fit = np.arange(-100, 100)
 
     prob_right_50, contrasts_50, _ = compute_performance(trials, signed_contrast=signed_contrast, block=0.5, prob_right=True)
-    pars_50 = compute_psychometric(trials, signed_contrast=signed_contrast, block=0.5)
+    pars_50 = compute_psychometric(trials, signed_contrast=signed_contrast, block=0.5, plotting=True)
     prob_right_fit_50 = psy.erf_psycho_2gammas(pars_50, contrasts_fit)
 
     prob_right_20, contrasts_20, _ = compute_performance(trials, signed_contrast=signed_contrast, block=0.2, prob_right=True)
-    pars_20 = compute_psychometric(trials, signed_contrast=signed_contrast, block=0.2)
+    pars_20 = compute_psychometric(trials, signed_contrast=signed_contrast, block=0.2, plotting=True)
     prob_right_fit_20 = psy.erf_psycho_2gammas(pars_20, contrasts_fit)
 
     prob_right_80, contrasts_80, _ = compute_performance(trials, signed_contrast=signed_contrast, block=0.8, prob_right=True)
-    pars_80 = compute_psychometric(trials, signed_contrast=signed_contrast, block=0.8)
+    pars_80 = compute_psychometric(trials, signed_contrast=signed_contrast, block=0.8, plotting=True)
     prob_right_fit_80 = psy.erf_psycho_2gammas(pars_80, contrasts_fit)
 
     cmap = sns.diverging_palette(20, 220, n=3, center="dark")

brainbox/tests/test_behavior.py

@@ -152,7 +152,7 @@ def test_concatenate_and_computations(self):
         perf_easy = np.array([train.compute_performance_easy(trials[k]) for k in trials.keys()])
         n_trials = np.array([train.compute_n_trials(trials[k]) for k in trials.keys()])
         psych = train.compute_psychometric(trials_all)
-        rt = train.compute_median_reaction_time(trials_all)
+        rt = train.compute_median_reaction_time(trials_all, contrast=0)
         np.testing.assert_allclose(perf_easy, [0.91489362, 0.9, 0.90853659])
         np.testing.assert_array_equal(n_trials, [617, 532, 719])
         np.testing.assert_allclose(psych, [4.04487042, 21.6293942, 1.91451396e-02, 1.72669957e-01],

ibllib/pipes/ephys_preprocessing.py

@@ -22,6 +22,7 @@
 from ibllib.io.extractors import ephys_fpga, ephys_passive, camera
 from ibllib.pipes import tasks
 from ibllib.pipes.training_preprocessing import TrainingRegisterRaw as EphysRegisterRaw
+from ibllib.pipes.training_preprocessing import TrainingStatus as EphysTrainingStatus
 from ibllib.pipes.misc import create_alyx_probe_insertions
 from ibllib.qc.alignment_qc import get_aligned_channels
 from ibllib.qc.task_extractors import TaskQCExtractor
@@ -1287,6 +1288,7 @@ def __init__(self, session_path=None, **kwargs):
             self.session_path, parents=[tasks["EphysVideoCompress"], tasks["EphysPulses"], tasks["EphysTrials"]])
         tasks["EphysCellsQc"] = EphysCellsQc(self.session_path, parents=[tasks["SpikeSorting"]])
         tasks["EphysDLC"] = EphysDLC(self.session_path, parents=[tasks["EphysVideoCompress"]])
+        tasks['EphysTrainingStatus'] = EphysTrainingStatus(self.session_path, parents=[tasks["EphysTrials"]])
         # level 3
         tasks["EphysPostDLC"] = EphysPostDLC(self.session_path, parents=[tasks["EphysDLC"], tasks["EphysTrials"],
                                                                          tasks["EphysVideoSyncQc"]])

ibllib/pipes/training_preprocessing.py

@@ -1,7 +1,7 @@
 import logging
 from collections import OrderedDict
 
-from ibllib.pipes import tasks
+from ibllib.pipes import tasks, training_status
 from ibllib.io import ffmpeg
 from ibllib.io.extractors.base import get_session_extractor_type
 from ibllib.io.extractors import training_audio, bpod_trials, camera
@@ -107,6 +107,27 @@ def _run(self):
         pass
 
 
+class TrainingStatus(tasks.Task):
+    priority = 90
+    level = 1
+    force = False
+    signature = {
+        'input_files': [('_iblrig_taskData.raw.*', 'raw_behavior_data', True),
+                        ('_iblrig_taskSettings.raw.*', 'raw_behavior_data', True),
+                        ('*trials.table.pqt', 'alf', True)],
+        'output_files': []
+    }
+
+    def _run(self, upload=True):
+        """
+        Extracts training status for subject
+        """
+        df = training_status.get_latest_training_information(self.session_path, self.one)
+        training_status.make_plots(self.session_path, self.one, df=df, save=True, upload=upload)
+        output_files = []
+        return output_files
+
+
 class TrainingExtractionPipeline(tasks.Pipeline):
     label = __name__
 
@@ -120,6 +141,7 @@ def __init__(self, session_path, **kwargs):
         tasks['TrainingVideoCompress'] = TrainingVideoCompress(self.session_path)
         tasks['TrainingAudio'] = TrainingAudio(self.session_path)
         # level 1
+        tasks['TrainingStatus'] = TrainingStatus(self.session_path, parents=[tasks['TrainingTrials']])
         tasks['TrainingDLC'] = TrainingDLC(
             self.session_path, parents=[tasks['TrainingVideoCompress']])
         self.tasks = tasks