hard code training phase temporarily

int-brain-lab · Jul 6, 2023 · 1e6d304 · 1e6d304
1 parent 4aab43e
commit 1e6d304
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/iblrig/base_choice_world.py b/iblrig/base_choice_world.py
@@ -672,9 +672,9 @@ def show_trial_log(self):
 class TrainingChoiceWorldSession(ActiveChoiceWorldSession):
     protocol_name = "_iblrig_tasks_trainingChoiceWorld"
 
-    def __init__(self, **kwargs):
-        super(TrainingChoiceWorldSession, self).__init__(**kwargs)
-        self.training_phase = 0
+    def __init__(self, training_phase=0, **kwargs):
+        super(TrainingChoiceWorldSession, self).__init__(training_phase=0, **kwargs)
+        self.training_phase = training_phase
         self.var = {
             "training_phase_trial_counts": np.zeros(6),
             "last_10_responses_sides": np.zeros(10),

diff --git a/iblrig_tasks/_iblrig_tasks_trainingChoiceWorld/task.py b/iblrig_tasks/_iblrig_tasks_trainingChoiceWorld/task.py
@@ -13,5 +13,10 @@ class Session(TrainingChoiceWorldSession):
 
 if __name__ == "__main__":  # pragma: no cover
     kwargs = iblrig.misc.get_task_runner_argument_parser()
-    sess = Session(**kwargs)
+    training_phase = 0
+    if kwargs['subject'] == 'ZFM-05923':
+        training_phase = 5
+    if kwargs['subject'] == 'ZFM-06440':
+        training_phase = 0
+    sess = Session(training_phase=training_phase, **kwargs)
     sess.run()