np-services 0.1.59__py3-none-any.whl → 0.1.73__py3-none-any.whl

np_services/scripts/pretest.py

@@ -1,50 +1,93 @@
+from __future__ import annotations
+
+import abc
 import argparse
+import copy
 import contextlib
+import dataclasses
 import functools
+import json
 import os
 import pathlib
+import sys
 import tempfile
 import time
 import logging
-from typing import Iterable, Literal
+from typing import Iterable, Literal, Type
 
 import upath
 
 
-os.environ["USE_TEST_RIG"] = "0"
-os.environ["AIBS_RIG_ID"] = "NP.3"
-
 import np_session
 import np_services
 import np_config
 import npc_sync
-import npc_ephys
 import npc_mvr
 import npc_stim
 
 logger = logging.getLogger()
 
-DEFAULT_SERVICES = (np_services.MouseDirector, )
-DEFAULT_STIM = np_services.ScriptCamstim
-DEFAULT_RECORDERS = (np_services.Sync, np_services.OpenEphys, np_services.VideoMVR, )
+DEFAULT_SERVICES: tuple[np_services.Testable, ...] = (np_services.MouseDirector, )
+DEFAULT_RECORDERS: tuple[np_services.Startable, ...] = (np_services.Sync, np_services.VideoMVR, )
+
+@dataclasses.dataclass
+class PretestConfig:
+    check_sync_barcodes: bool = False
+    check_ephys_barcodes: bool = False # this is error prone with short recordings, so is separated from checking sync alone
+    check_licks: bool = False
+    check_opto: bool = False
+    check_audio: bool = False
+    check_running: bool = False
+    
+    @property
+    def check_barcodes(self) -> bool:
+        return self.check_sync_barcodes or self.check_ephys_barcodes
+
+
+class PretestSession(abc.ABC):
+
+    def __init__(self, pretest_config: PretestConfig) -> None:
+        self.pretest_config = pretest_config
+    
+    @property
+    def services(self) -> tuple[np_services.Testable | np_services.Startable, ...]: 
+        return DEFAULT_SERVICES + self.recorders + (self.stim, )
+    
+    @property
+    def recorders(self) -> tuple[np_services.Startable, ...]:
+        if self.pretest_config.check_ephys_barcodes or self.pretest_config.check_sync_barcodes:
+            return DEFAULT_RECORDERS + (np_services.OpenEphys, )
+        return DEFAULT_RECORDERS
+    
+    @property
+    @abc.abstractmethod     
+    def stim(self) -> Type[np_services.Camstim]: ...
 
-class DynamicRoutingPretest:
+    @abc.abstractmethod     
+    def configure_services(self) -> None: ...
+    
+    @abc.abstractmethod
+    def run_pretest_stim(self) -> None: ...
+    
+class DynamicRoutingPretest(PretestSession):
     """Modified version of class in np_workflows."""
+    
     use_github: bool = True
+    task_name: str = "" # unused in pretest
     
+    @property
+    def stim(self) -> Type[np_services.ScriptCamstim]:
+        return np_services.ScriptCamstim
+        
     @property
     def rig(self) -> np_config.Rig:
         return np_config.Rig()
-
-    @property
-    def config(self) -> dict:
-        return self.rig.config
     
     @property
     def commit_hash(self) -> str:
         if hasattr(self, '_commit_hash'):
             return self._commit_hash
-        self._commit_hash = self.config['dynamicrouting_task_script']['commit_hash']
+        self._commit_hash = self.rig.config['dynamicrouting_task_script']['commit_hash']
         return self.commit_hash
     
     @commit_hash.setter
@@ -55,7 +98,7 @@ class DynamicRoutingPretest:
     def github_url(self) -> str:
         if hasattr(self, '_github_url'):
             return self._github_url
-        self._github_url = self.config['dynamicrouting_task_script']['url']
+        self._github_url = self.rig.config['dynamicrouting_task_script']['url']
         return self.github_url
     
     @github_url.setter
@@ -70,21 +113,6 @@ class DynamicRoutingPretest:
     def base_path(self) -> pathlib.Path:
         return pathlib.Path('//allen/programs/mindscope/workgroups/dynamicrouting/DynamicRoutingTask/')
 
-    @property
-    def task_name(self) -> str:
-        """For sending to runTask.py and controlling implementation details of the task."""
-        if hasattr(self, '_task_name'): 
-            return self._task_name 
-        return ""
-
-    @task_name.setter
-    def task_name(self, task_name: str) -> None:
-        self._task_name = task_name
-        if task_name not in self.preset_task_names:
-            print(f"{task_name = !r} doesn't correspond to a preset value, but the attribute is updated anyway!")
-        else:
-            print(f"Updated {self.__class__.__name__}.{task_name = !r}")
-
     @property
     def mouse(self) -> np_session.Mouse:
         return np_session.Mouse(366122)
@@ -135,9 +163,11 @@ class DynamicRoutingPretest:
         
         rig = str(self.rig).replace('.', '')
         locs_root = self.base_path / 'OptoGui' / f'{dirname}'
-        available_locs = sorted(tuple(locs_root.glob(f"{file_prefix}_{self.mouse.id}_{rig}_*")), reverse=True)
+        # use any available locs file - as long as the light switches on the
+        # values don't matter
+        available_locs = sorted(tuple(locs_root.glob(f"{file_prefix}*")), reverse=True)
         if not available_locs:
-            raise FileNotFoundError(f"No optotagging locs found for {self.mouse}/{rig} - have you run OptoGui?")
+            raise FileNotFoundError(f"No optotagging locs found - have you run OptoGui?")
         return available_locs[0]
         
         
@@ -197,16 +227,16 @@ class DynamicRoutingPretest:
     
     def run_script(self, stim: Literal['sound_test', 'mapping', 'task', 'opto', 'optotagging', 'spontaneous', 'spontaneous_rewards']) -> None:
         
-        params = getattr(self, f'{stim.replace(" ", "_")}_params')
+        params = copy.deepcopy(getattr(self, f'{stim.replace(" ", "_")}_params'))
         
         # add mouse and user info for MPE
         params['mouse_id'] = str(self.mouse.id)
         params['user_id'] = 'ben.hardcastle'
         
-        script: str = params['taskScript']
-        params['taskScript'] = (self.task_script_base / script).as_posix()
+        if self.task_script_base.as_posix() not in params['taskScript']:
+            params['taskScript'] = (self.task_script_base / params['taskScript']).as_posix()
         
-        params['maxTrials'] = 30
+        params['maxTrials'] = 30    
         
         if self.use_github:
         
@@ -217,23 +247,53 @@ class DynamicRoutingPretest:
                 }
             params['task_script_commit_hash'] = self.commit_hash
 
-            np_services.ScriptCamstim.script = self.camstim_script.read_text()
+            self.stim.script = self.camstim_script.read_text()
         else:
-            np_services.ScriptCamstim.script = self.camstim_script.as_posix()
+            self.stim.script = self.camstim_script.as_posix()
         
-        np_services.ScriptCamstim.params = params
+        self.stim.params = params
         
 
-        np_services.ScriptCamstim.start()
+        self.stim.start()
         with contextlib.suppress(np_services.resources.zro.ZroError):
-            while not np_services.ScriptCamstim.is_ready_to_start():
+            while not self.stim.is_ready_to_start():
                 time.sleep(1)
 
 
         with contextlib.suppress(np_services.resources.zro.ZroError):
-            np_services.ScriptCamstim.finalize()
-     
-def configure_services(services: Iterable[np_services.Testable]) -> None:
+            self.stim.finalize()
+            
+    def run_pretest_stim(self) -> None:
+        if self.pretest_config.check_audio:
+            print("Starting audio test - check for sound...", end="", flush=True)
+            self.run_script('sound_test')
+            print(" done")
+        print("Starting stim - check for opto, spin wheel and tap lick spout...", end="", flush=True)
+        self.run_script('optotagging') # vis stim with opto (for checking vsyncs, running,  )
+        print(" done")
+        
+    def configure_services(self) -> None:
+        self.stim.script = '//allen/programs/mindscope/workgroups/dynamicrouting/DynamicRoutingTask/runTask.py'
+        self.stim.data_root = pathlib.Path('//allen/programs/mindscope/workgroups/dynamicrouting/DynamicRoutingTask/Data/366122')
+
+                
+class LegacyNP0Pretest(PretestSession):
+    def __init__(self, pretest_config: PretestConfig) -> None:
+        self.pretest_config = pretest_config
+        
+    @property
+    def stim(self) -> Type[np_services.SessionCamstim]:
+        return np_services.SessionCamstim
+    
+    def run_pretest_stim(self) -> None:
+        self.stim.start()
+                                
+    def configure_services(self) -> None:
+        self.stim.lims_user_id = "ben.hardcastle"
+        self.stim.labtracks_mouse_id = 598796
+        self.stim.override_params = json.loads(pathlib.Path("//allen/programs/mindscope/workgroups/dynamicrouting/ben/np0_pretest/params.json").read_bytes())
+
+def configure_services(session: PretestSession) -> None:
     """For each service, apply every key in self.config['service'] as an attribute."""
 
     def apply_config(service) -> None:
@@ -244,73 +304,110 @@ def configure_services(services: Iterable[np_services.Testable]) -> None:
                     f"{service.__name__} | Configuring {service.__name__}.{key} = {getattr(service, key)}"
                 )
 
-    for service in services:
+    for service in session.services:
         for base in service.__class__.__bases__:
             apply_config(base)
         apply_config(service)
-        
-    np_services.ScriptCamstim.script = '//allen/programs/mindscope/workgroups/dynamicrouting/DynamicRoutingTask/runTask.py'
-    np_services.ScriptCamstim.data_root = pathlib.Path('//allen/programs/mindscope/workgroups/dynamicrouting/DynamicRoutingTask/Data/366122')
-
+    
     np_services.MouseDirector.user = 'ben.hardcastle'
     np_services.MouseDirector.mouse = 366122
+    
+    session.configure_services()
 
-    np_services.OpenEphys.folder = '_test_'
+    if session.pretest_config.check_barcodes:
+        np_services.OpenEphys.folder = '_test_'
 
 
 @functools.cache
 def get_temp_dir() -> pathlib.Path:
     return pathlib.Path(tempfile.mkdtemp())
 
-def main(
-    recorders: Iterable[np_services.Testable] = DEFAULT_RECORDERS,
-    stim: np_services.Startable = DEFAULT_STIM,
-    other: Iterable[np_services.Testable] = DEFAULT_SERVICES,
-    check_licks: bool = False,
-    check_opto: bool = False,
-    check_audio: bool = False,
+def run_pretest(
+    config: PretestConfig = PretestConfig(),
     ) -> None:
     print("Starting pretest")
-    configure_services((*recorders, stim, *other))
-    for service in (*recorders, stim, *other):
+    session: PretestSession
+    if np_config.Rig().idx == 0:
+        session = LegacyNP0Pretest(config)
+    else:
+        session = DynamicRoutingPretest(config)
+    configure_services(session)
+    
+    for service in session.services:
         if isinstance(service, np_services.Initializable):
             service.initialize()
             
-    stoppables = tuple(_ for _ in recorders if isinstance(_, np_services.Stoppable))
+    stoppables = tuple(_ for _ in session.recorders if isinstance(_, np_services.Stoppable))
     with np_services.stop_on_error(*stoppables):
         for service in stoppables:
             if isinstance(service, np_services.Startable):
                 service.start()
         t0 = time.time()
-        DynamicRoutingPretest().run_script('optotagging')
+        session.run_pretest_stim()
         t1 = time.time()
-        time.sleep(max(0, 70 - (t1 - t0))) # long enough to capture 2 sets of barcodes on sync/openephys (cannot scale time with 1 set)
+        if not config.check_barcodes:
+            min_wait_time = 0
+        elif config.check_sync_barcodes and not config.check_ephys_barcodes:
+            min_wait_time = 35 # long enough to capture 1 sets of barcodes on sync
+        else:
+            min_wait_time = 70 # long enough to capture 2 sets of barcodes on sync/openephys (cannot scale time with 1 set)
+        time.sleep(max(0, min_wait_time - (t1 - t0))) 
         for service in reversed(stoppables):
             if isinstance(service, np_services.Stoppable):
                 service.stop()
 
-    for service in (*recorders, stim, *other):
+    for service in session.services:
         if isinstance(service, np_services.Finalizable):
             service.finalize()
 
-    np_services.VideoMVR.sync_path = np_services.OpenEphys.sync_path = stim.sync_path = np_services.Sync.data_files[0]
+    np_services.VideoMVR.sync_path = np_services.OpenEphys.sync_path = session.stim.sync_path = np_services.Sync.data_files[0]
     
-    for service in (*recorders, stim, *other):
+    # validate
+    for service in session.services:
         if isinstance(service, np_services.Validatable):
-            service.validate()
-    if any((check_licks, check_opto, check_audio)):
+            if service is not np_services.OpenEphys:
+                service.validate()
+            elif config.check_ephys_barcodes:
+                # try validating ephys without sync (currently error prone with short pretest-like recordings)
+                import npc_ephys
+                npc_ephys.validate_ephys(
+                    root_paths=service.data_files,
+                    sync_path_or_dataset=False,
+                    ignore_small_folders=False,
+                )
+            else:
+                # barcodes on sync will be validated, open ephys will be assumed to work correctly
+                continue
+    assert np_services.Sync.data_files is not None, "No sync file found"
+    assert session.stim.data_files is not None, "No stim file found"
+    
+    if any((config.check_licks, config.check_opto, config.check_audio)):
         npc_sync.SyncDataset(np_services.Sync.data_files[0]).validate(
-            licks=check_licks, opto=check_opto, audio=check_audio,
+            licks=config.check_licks, opto=config.check_opto, audio=config.check_audio,
         )
+    if config.check_running:
+        speed, timestamps  = npc_stim.get_running_speed_from_stim_files(*session.stim.data_files, sync=np_services.Sync.data_files[0])
+        if not speed.size or not timestamps.size:
+            raise AssertionError("No running data found")
 
-def parse_args() -> dict:
+def parse_args() -> PretestConfig:
     parser = argparse.ArgumentParser(description="Run pretest")
+    parser.add_argument("--check_ephys_barcodes", action="store_true", help="Check barcodes from Arduino are being received on open ephys and time-alignment is possible (currently error prone with short pretest-like recordings)", default=False)
+    parser.add_argument("--check_sync_barcodes", action="store_true", help="Check barcodes from Arduino are being received on sync", default=False)
     parser.add_argument("--check_licks", action="store_true", help="Check lick sensor line on sync", default=False)
     parser.add_argument("--check_opto", action="store_true", help="Check opto-running line on sync", default=False)
     parser.add_argument("--check_audio", action="store_true", help="Check audio-running line on sync", default=False)
-    return vars(parser.parse_args())
+    parser.add_argument("--check_running", action="store_true", help="Check running-wheel encoder data in stim files", default=False)
+    return PretestConfig(**vars(parser.parse_args()))
+
+def main() -> None:
+    logging.basicConfig(
+        level="INFO",
+        format="%(name)s | %(levelname)s | %(message)s",
+        datefmt="%H:%M:%S",
+        stream=sys.stdout,
+    )
+    run_pretest(parse_args())
 
 if __name__ == '__main__':
-    main(
-        **parse_args()
-    )
+    main()

np_services/stim_computer_theme_changer.py

@@ -1,42 +1,42 @@
-import pathlib
-from typing import ClassVar
-
-import fabric
-import np_config
-
-from np_services.proxies import ScriptCamstim
-
-class DesktopThemeChanger(ScriptCamstim):
-    """Base class for setting the background wallpaper on a remote machine,
-    then hiding all desktop icons, hiding the taskbar, minimizing all windows"""
-    local_file: ClassVar[str | pathlib.Path] 
-    remote_file: ClassVar[str | pathlib.Path] 
-    
-    extra_args: ClassVar[list[str]] = []
-    ssh: ClassVar[fabric.Connection]
-    user: ClassVar[str] = 'svc_neuropix'
-    password: ClassVar[str] = np_config.fetch('logins')['svc_neuropix']['password']
-    
-    # @classmethod
-    # def initialize(cls):
-    #     super().initialize()
-    #     cls.get_ssh().put(cls.local_file, cls.remote_file)
-    
-    @classmethod
-    def start(cls):
-        cls.get_ssh().run(f'powershell.exe -ExecutionPolicy bypass {cls.remote_file}')
-        
-class DarkDesktopChanger(DesktopThemeChanger):
-    local_file = pathlib.Path(__file__).parent / 'resources' / 'black_wallpaper.ps1'
-    # remote_file: ClassVar[str | pathlib.Path] = 'c:/users/svc_neuropix/desktop/black_wallpaper.ps1'
-    remote_file: ClassVar[str | pathlib.Path] = R'\\allen\programs\mindscope\workgroups\dynamicrouting\ben\black_wallpaper.ps1'
-    
-class GreyDesktopChanger(DesktopThemeChanger):
-    local_file = pathlib.Path(__file__).parent / 'resources' / 'grey_wallpaper.ps1'
-    # remote_file: ClassVar[str | pathlib.Path] = 'c:/users/svc_neuropix/desktop/grey_wallpaper.ps1'
-    remote_file: ClassVar[str | pathlib.Path] = R'\\allen\programs\mindscope\workgroups\dynamicrouting\ben\grey_wallpaper.ps1'
-    
-class DesktopResetter(DesktopThemeChanger):
-    local_file = pathlib.Path(__file__).parent / 'resources' / 'reset_wallpaper.ps1'
-    # remote_file: ClassVar[str | pathlib.Path] = 'c:/users/svc_neuropix/desktop/reset_wallpaper.ps1'
+import pathlib
+from typing import ClassVar
+
+import fabric
+import np_config
+
+from np_services.proxies import ScriptCamstim
+
+class DesktopThemeChanger(ScriptCamstim):
+    """Base class for setting the background wallpaper on a remote machine,
+    then hiding all desktop icons, hiding the taskbar, minimizing all windows"""
+    local_file: ClassVar[str | pathlib.Path] 
+    remote_file: ClassVar[str | pathlib.Path] 
+    
+    extra_args: ClassVar[list[str]] = []
+    ssh: ClassVar[fabric.Connection]
+    user: ClassVar[str] = 'svc_neuropix'
+    password: ClassVar[str] = np_config.fetch('logins')['svc_neuropix']['password']
+    
+    # @classmethod
+    # def initialize(cls):
+    #     super().initialize()
+    #     cls.get_ssh().put(cls.local_file, cls.remote_file)
+    
+    @classmethod
+    def start(cls):
+        cls.get_ssh().run(f'powershell.exe -ExecutionPolicy bypass {cls.remote_file}')
+        
+class DarkDesktopChanger(DesktopThemeChanger):
+    local_file = pathlib.Path(__file__).parent / 'resources' / 'black_wallpaper.ps1'
+    # remote_file: ClassVar[str | pathlib.Path] = 'c:/users/svc_neuropix/desktop/black_wallpaper.ps1'
+    remote_file: ClassVar[str | pathlib.Path] = R'\\allen\programs\mindscope\workgroups\dynamicrouting\ben\black_wallpaper.ps1'
+    
+class GreyDesktopChanger(DesktopThemeChanger):
+    local_file = pathlib.Path(__file__).parent / 'resources' / 'grey_wallpaper.ps1'
+    # remote_file: ClassVar[str | pathlib.Path] = 'c:/users/svc_neuropix/desktop/grey_wallpaper.ps1'
+    remote_file: ClassVar[str | pathlib.Path] = R'\\allen\programs\mindscope\workgroups\dynamicrouting\ben\grey_wallpaper.ps1'
+    
+class DesktopResetter(DesktopThemeChanger):
+    local_file = pathlib.Path(__file__).parent / 'resources' / 'reset_wallpaper.ps1'
+    # remote_file: ClassVar[str | pathlib.Path] = 'c:/users/svc_neuropix/desktop/reset_wallpaper.ps1'
     remote_file: ClassVar[str | pathlib.Path] = R'\\allen\programs\mindscope\workgroups\dynamicrouting\ben\reset_wallpaper.ps1'