Merge pull request #5 from int-brain-lab/zapit

initial commit for zapit task

iblrig_custom_tasks/_sp_passiveVideo/task.py

@@ -8,16 +8,20 @@
 from collections import defaultdict
 import logging
 
-import vlc
 import pandas as pd
 from pybpodapi.protocol import Bpod
 
 import iblrig.misc
-from iblrig.base_tasks import BaseSession, BpodMixin
-
+from iblrig.base_tasks import BpodMixin
 
 _logger = logging.getLogger(__name__)
 
+# this allows the CI and automated tests to import the file and make sure it is valid without having vlc
+try:
+    import vlc
+except (ModuleNotFoundError, FileNotFoundError):
+    _logger.error(f'VLC not installed. Please install VLC to use this task. {__file__}')
+
 
 class Player:
     """A VLC player."""
@@ -95,7 +99,7 @@ def get_ended_time(self, repeat=-1):
             return ends[repeat]
 
 
-class Session(BaseSession, BpodMixin):
+class Session(BpodMixin):
     """Play a single video."""
 
     protocol_name = '_sp_passiveVideo'

iblrig_custom_tasks/nate_optoBiasedChoiceWorld/task.py

@@ -1,26 +1,28 @@
 """
 This task is a replica of BiasedChoiceWorldSession with the addition of optogenetic stimulation
 An `opto_stimulation` column is added to the trials_table, which is a boolean array of length NTRIALS_INIT
-The PROBABILITY_OPTO_STIMULATION parameter is used to determine the probability of optogenetic stimulation for each trial
+The PROBABILITY_OPTO_STIMULATION parameter is used to determine the probability of optogenetic stimulation
+for each trial
 
 Additionally the state machine is modified to add output TTLs for optogenetic stimulation
 """
-
-import numpy as np
-import yaml
+import logging
 from pathlib import Path
 from typing import Literal
 
-from pybpodapi.protocol import StateMachine
+import numpy as np
+import yaml
 
-from iblrig.base_choice_world import BiasedChoiceWorldSession
-from iblutil.util import setup_logger
 import iblrig
+from iblrig.base_choice_world import SOFTCODE, BiasedChoiceWorldSession
+from pybpodapi.protocol import StateMachine
 
-log = setup_logger(__name__)
+log = logging.getLogger('iblrig.task')
 
 INTERACTIVE_DELAY = 1.0
 NTRIALS_INIT = 2000
+SOFTCODE_STOP_ZAPIT = max(SOFTCODE).value + 1
+SOFTCODE_FIRE_ZAPIT = max(SOFTCODE).value + 2
 
 # read defaults from task_parameters.yaml
 with open(Path(__file__).parent.joinpath('task_parameters.yaml')) as f:
@@ -32,14 +34,27 @@ class OptoStateMachine(StateMachine):
     This class just adds output TTL on BNC2 for defined states
     """
 
-    def __init__(self, bpod, is_opto_stimulation=False, states_opto_ttls=None):
+    def __init__(
+        self,
+        bpod,
+        is_opto_stimulation=False,
+        states_opto_ttls=None,
+        states_opto_stop=None,
+    ):
         super().__init__(bpod)
         self.is_opto_stimulation = is_opto_stimulation
         self.states_opto_ttls = states_opto_ttls or []
+        self.states_opto_stop = states_opto_stop or []
 
     def add_state(self, **kwargs):
-        if self.is_opto_stimulation and kwargs['state_name'] in self.states_opto_ttls:
-            kwargs['output_actions'].append(('BNC2', 255))
+        if self.is_opto_stimulation:
+            if kwargs['state_name'] in self.states_opto_ttls:
+                kwargs['output_actions'] += [
+                    ('SoftCode', SOFTCODE_FIRE_ZAPIT),
+                    ('BNC2', 255),
+                ]
+            elif kwargs['state_name'] in self.states_opto_stop:
+                kwargs['output_actions'] += [('SoftCode', SOFTCODE_STOP_ZAPIT)]
         super().add_state(**kwargs)
 
 
@@ -51,15 +66,14 @@ def __init__(
         *args,
         probability_opto_stim: float = DEFAULTS['PROBABILITY_OPTO_STIM'],
         contrast_set_probability_type: Literal['skew_zero', 'uniform'] = DEFAULTS['CONTRAST_SET_PROBABILITY_TYPE'],
-        opto_stim_states: list[str] = DEFAULTS['OPTO_STIM_STATES'],
+        opto_ttl_states: list[str] = DEFAULTS['OPTO_TTL_STATES'],
+        opto_stop_states: list[str] = DEFAULTS['OPTO_STOP_STATES'],
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
-        print(probability_opto_stim)
-        print(contrast_set_probability_type)
-        print(opto_stim_states)
         self.task_params['CONTRAST_SET_PROBABILITY_TYPE'] = contrast_set_probability_type
-        self.task_params['OPTO_STIM_STATES'] = opto_stim_states
+        self.task_params['OPTO_TTL_STATES'] = opto_ttl_states
+        self.task_params['OPTO_STOP_STATES'] = opto_stop_states
         self.task_params['PROBABILITY_OPTO_STIM'] = probability_opto_stim
 
         # loads in the settings in order to determine the main sync and thus the pipeline extractor tasks
@@ -69,18 +83,47 @@ def __init__(
 
         # generates the opto stimulation for each trial
         self.trials_table['opto_stimulation'] = np.random.choice(
-            [0, 1], p=[1 - probability_opto_stim, probability_opto_stim], size=NTRIALS_INIT
+            [0, 1],
+            p=[1 - probability_opto_stim, probability_opto_stim],
+            size=NTRIALS_INIT,
         ).astype(bool)
 
+    def start_hardware(self):
+        super().start_hardware()
+        # add the softcodes for the zapit opto stimulation
+        soft_code_dict = self.bpod.softcodes
+        soft_code_dict.update({SOFTCODE_STOP_ZAPIT: self.zapit_stop_laser})
+        soft_code_dict.update({SOFTCODE_FIRE_ZAPIT: self.zapit_fire_laser})
+        self.bpod.register_softcodes(soft_code_dict)
+
+    def zapit_arm_laser(self):
+        log.warning('Arming laser')
+        # TODO: insert code for arming the laser here
+
+    def zapit_fire_laser(self):
+        # just logging - actual firing will be triggered by the state machine via TTL
+        log.warning('Firing laser')
+
+    def zapit_stop_laser(self):
+        log.warning('Stopping laser')
+        # TODO: insert code for stopping the laser here
+
     def _instantiate_state_machine(self, trial_number=None):
         """
         We override this using the custom class OptoStateMachine that appends TTLs for optogenetic stimulation where needed
         :param trial_number:
         :return:
         """
         is_opto_stimulation = self.trials_table.at[trial_number, 'opto_stimulation']
-        states_opto_ttls = self.task_params['OPTO_STIM_STATES']
-        return OptoStateMachine(self.bpod, is_opto_stimulation=is_opto_stimulation, states_opto_ttls=states_opto_ttls)
+        # we start the laser waiting for a TTL trigger before sending out the state machine on opto trials
+        if is_opto_stimulation:
+            self.zapit_arm_laser()
+        return OptoStateMachine(
+            self.bpod,
+            is_opto_stimulation=is_opto_stimulation,
+            states_opto_ttls=self.task_params['OPTO_TTL_STATES'],
+            states_opto_stop=self.task_params['OPTO_STOP_STATES'],
+        )
 
     @staticmethod
     def extra_parser():
@@ -104,14 +147,23 @@ def extra_parser():
             help=f'probability type for contrast set (default: {DEFAULTS["CONTRAST_SET_PROBABILITY_TYPE"]})',
         )
         parser.add_argument(
-            '--opto_stim_states',
-            option_strings=['--opto_stim_states'],
-            dest='opto_stim_states',
-            default=DEFAULTS['OPTO_STIM_STATES'],
+            '--opto_ttl_states',
+            option_strings=['--opto_ttl_states'],
+            dest='opto_ttl_states',
+            default=DEFAULTS['OPTO_TTL_STATES'],
             nargs='+',
             type=str,
             help='list of the state machine states where opto stim should be delivered',
         )
+        parser.add_argument(
+            '--opto_stop_states',
+            option_strings=['--opto_stop_states'],
+            dest='opto_stop_states',
+            default=DEFAULTS['OPTO_STOP_STATES'],
+            nargs='+',
+            type=str,
+            help='list of the state machine states where opto stim should be stopped',
+        )
         return parser
 
 

iblrig_custom_tasks/nate_optoBiasedChoiceWorld/task_parameters.yaml

@@ -1,6 +1,7 @@
 'CONTRAST_SET_PROBABILITY_TYPE': uniform  # uniform or skew_zero: uniform makes the 0 contrast as likely as the others, while skew_zero makes it half as likely as other contrasts
-'OPTO_STIM_STATES': # list of the state machine states where opto stim should be delivered
+'OPTO_TTL_STATES': # list of the state machine states where opto stim should be delivered
   - trial_start
+'OPTO_STOP_STATES':
   - no_go
   - error
   - reward

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name            = "project_extraction"
-version         = "0.2.1.post0"
+version         = "0.2.2"
 description     = "Custom extractors for satellite tasks"
 dynamic         = [ "readme" ]
 keywords        = [ "IBL", "neuro-science" ]
@@ -20,3 +20,42 @@ readme = { file = "README.md", content-type = "text/markdown" }
 
 [tool.setuptools.packages]
 find = {}
+
+[tool.ruff]
+ignore = [
+    "PLR0912", # Too many branches
+    "PLR0915", # Too many statements
+    "PLR2004", # Magic value used in comparison, consider replacing with a constant variable
+]
+exclude = [
+    ".mypy_cache",
+    "dist",
+    "docs",
+    "iblrig/gui/*_rc.py",
+    "iblrig/gui/ui_*.py",
+    "venv",
+]
+indent-width = 4
+line-length = 130
+target-version = "py310"
+
+[tool.ruff.lint]
+select = [
+    "B",   # flake8-bugbear
+    "E",   # pycodestyle
+    "F",   # Pyflakes
+    "I",   # isort
+    "N",   # pep8-naming
+    "PL",  # pylint
+    "SIM", # flake8-simplify
+    "UP",  # pyupgrade
+]
+
+[tool.ruff.format]
+quote-style = "single"
+
+[tool.ruff.lint.pydocstyle]
+convention = "numpy"
+
+[tool.ruff.lint.isort]
+known-first-party = [ "ibl*", "one*", "pybpod*" ]