np-workflows 1.6.87__py3-none-any.whl → 1.6.91__py3-none-any.whl

np_workflows/__init__.py

@@ -1,7 +1,5 @@
+import np_workflows.experiments as experiments
+import np_workflows.shared.base_experiments as base_experiments
 import np_workflows.shared.npxc as npxc
 import np_workflows.shared.widgets as widgets
-import np_workflows.shared.base_experiments as base_experiments
-
-import np_workflows.experiments as experiments
-
-from np_workflows.shared import *
+from np_workflows.shared import *

np_workflows/experiments/dynamic_routing/main.py

@@ -1,48 +1,21 @@
 from __future__ import annotations
 
-import configparser
-import contextlib
-import copy
-import dataclasses
-import datetime
-import enum
-import functools
-import pathlib
-import platform
-import shutil
-import threading
-import time
-import zlib
-from typing import ClassVar, Literal, NamedTuple, NoReturn, Optional, Type, TypedDict
-
-import IPython
-import IPython.display
-import ipywidgets as ipw
-import np_config
 import np_logging
-import np_services
 import np_session
-import np_workflows
-import fabric
-import PIL.Image
-import pydantic
-from pyparsing import Any
 from np_services import (
-    Service,
-    Finalizable,
-    ScriptCamstim,
+    MouseDirector,
+    NewScaleCoordinateRecorder,
     OpenEphys,
+    ScriptCamstim,
     Sync,
-    ImageMVR,
     VideoMVR,
-    NewScaleCoordinateRecorder,
-    MouseDirector,
 )
 
 from np_workflows.shared import base_experiments
 
 logger = np_logging.getLogger(__name__)
 
+
 def new_experiment(
     mouse: int | str | np_session.Mouse,
     user: str | np_session.User,
@@ -50,29 +23,33 @@ def new_experiment(
 ) -> DRTask:
     """Create a new experiment for the given mouse and user."""
     experiment: DRTask
-    if any(tag in workflow.name for tag in ('EPHYS', 'PRETEST')):
+    if any(tag in workflow.name for tag in ("EPHYS", "PRETEST")):
         experiment = Ephys(mouse, user)
-    elif 'HAB' in workflow.name:
+    elif "HAB" in workflow.name:
         experiment = Hab(mouse, user)
-    elif 'OPTO' in workflow.name:
+    elif "OPTO" in workflow.name:
         experiment = Opto(mouse, user)
-    elif 'TRAINING' in workflow.name:
+    elif "TRAINING" in workflow.name:
         experiment = Training(mouse, user)
     else:
-        raise ValueError(f"Unknown {workflow = }. Create an experiment with e.g.\n\n\texperiment = Ephys(mouse, user)\nexperiment.session.npexp_path.mkdir()")
+        raise ValueError(
+            f"Unknown {workflow = }. Create an experiment with e.g.\n\n\texperiment = Ephys(mouse, user)\nexperiment.session.npexp_path.mkdir()"
+        )
     experiment.workflow = workflow
     experiment.log(f"{experiment} created")
-    experiment.session.npexp_path.mkdir(parents=True, exist_ok=True) 
+    experiment.session.npexp_path.mkdir(parents=True, exist_ok=True)
     return experiment
 
+
 class DRTask(base_experiments.DynamicRoutingExperiment):
     """Provides project-specific methods and attributes, mainly related to camstim scripts."""
-    
+
     default_session_subclass = np_session.DRPilotSession
-    
+
     workflow: base_experiments.DynamicRoutingExperiment.Workflow
     """Enum for workflow type, e.g. PRETEST, HAB_AUD, HAB_VIS, EPHYS_ etc."""
 
+
 class Hab(DRTask):
     def __init__(self, *args, **kwargs):
         self.services = (
@@ -82,7 +59,8 @@ class Hab(DRTask):
             ScriptCamstim,
         )
         super().__init__(*args, **kwargs)
-        
+
+
 class Opto(DRTask):
     def __init__(self, *args, **kwargs):
         self.services = (
@@ -106,6 +84,7 @@ class Ephys(DRTask):
         )
         super().__init__(*args, **kwargs)
 
+
 class Training(DRTask):
     def __init__(self, *args, **kwargs):
         self.services = (
@@ -114,4 +93,4 @@ class Training(DRTask):
             VideoMVR,
             ScriptCamstim,
         )
-        super().__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)

np_workflows/experiments/dynamic_routing/widgets.py

@@ -1,29 +1,22 @@
-import configparser
-import contextlib
-import copy
-import enum
-import functools
-from typing import ClassVar, Literal, NamedTuple, NoReturn, Optional, TypedDict
-
 import IPython.display
 import ipywidgets as ipw
-import np_config
-import np_logging
 import np_session
-import np_workflows
-from pyparsing import Any
 
-from np_workflows.experiments.dynamic_routing.main import Ephys, Hab, DRTask
 from np_workflows.shared.base_experiments import DynamicRoutingExperiment
 
-
 # for widget, before creating a experiment --------------------------------------------- #
 
 
 class SelectedWorkflow:
-    def __init__(self, workflow: str | DynamicRoutingExperiment.Workflow, mouse: str | int | np_session.Mouse):
+    def __init__(
+        self,
+        workflow: str | DynamicRoutingExperiment.Workflow,
+        mouse: str | int | np_session.Mouse,
+    ):
         if isinstance(workflow, str):
-            workflow = DynamicRoutingExperiment.Workflow[workflow] # uses enum name (not value)
+            workflow = DynamicRoutingExperiment.Workflow[
+                workflow
+            ]  # uses enum name (not value)
         self.workflow = workflow
         self.mouse = str(mouse)
 
@@ -55,8 +48,12 @@ def workflow_select_widget(
     )
     console = ipw.Output()
     with console:
-        if last_workflow := np_session.Mouse(selection.mouse).state.get('last_workflow'):
-            print(f"{mouse} last workflow: {last_workflow}\t({np_session.Mouse(selection.mouse).state.get('last_session')})")
+        if last_workflow := np_session.Mouse(selection.mouse).state.get(
+            "last_workflow"
+        ):
+            print(
+                f"{mouse} last workflow: {last_workflow}\t({np_session.Mouse(selection.mouse).state.get('last_session')})"
+            )
         print(f"Selected: {selection.workflow.name}")
 
     def update(change):
@@ -68,15 +65,23 @@ def workflow_select_widget(
             return
         if change["new"] == change["old"]:
             return
-        selection.__init__(str(workflow_dropdown.value), mouse.id if isinstance(mouse, np_session.Mouse) else str(mouse))
+        selection.__init__(
+            str(workflow_dropdown.value),
+            mouse.id if isinstance(mouse, np_session.Mouse) else str(mouse),
+        )
         with console:
             print(f"Selected: {selection.workflow}")
-    workflow_dropdown.observe(update, names='value')
 
-    IPython.display.display(ipw.VBox([ipw.HBox([workflow_dropdown, workflow_descriptions]), console]))
+    workflow_dropdown.observe(update, names="value")
+
+    IPython.display.display(
+        ipw.VBox([ipw.HBox([workflow_dropdown, workflow_descriptions]), console])
+    )
 
     return selection
-    
-    
-def photodoc_widget(session: np_session.Session, reminder: str) -> None: 
-    print(f'Take an image in Vimba Viewer and save as:\n{session.npexp_path.name}_{reminder}.png')
+
+
+def photodoc_widget(session: np_session.Session, reminder: str) -> None:
+    print(
+        f"Take an image in Vimba Viewer and save as:\n{session.npexp_path.name}_{reminder}.png"
+    )

np_workflows/experiments/openscope_P3/P3_workflow_widget.py

@@ -0,0 +1,75 @@
+import IPython.display
+import ipywidgets as ipw
+import np_session
+
+from np_workflows.experiments.openscope_P3.main_P3_pilot import P3Session
+
+global_state = {}
+"""Global variable for persisting widget states."""
+
+# for widget, before creating a experiment --------------------------------------------- #
+
+
+class SelectedSession:
+    def __init__(self, session: str | P3Session, mouse: str | int | np_session.Mouse):
+        if isinstance(session, str):
+            session = P3Session(session)
+        self.session = session
+        self.mouse = str(mouse)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.session}, {self.mouse})"
+
+
+def P3_workflow_widget(
+    mouse: str | int | np_session.Mouse,
+) -> SelectedSession:
+    """Select a stimulus session (hab, pretest, ephys) to run.
+
+    An object with mutable attributes is returned, so the selected session can be
+    updated along with the GUI selection. (Preference would be to return an enum
+    directly, and change it's value, but that doesn't seem possible.)
+
+    """
+
+    selection = SelectedSession(P3Session.PRETEST, mouse)
+
+    session_dropdown = ipw.Select(
+        options=tuple(_.value for _ in P3Session),
+        description="Session",
+    )
+
+    def update_selection():
+        selection.__init__(str(session_dropdown.value), str(mouse))
+
+    if previously_selected_value := global_state.get("selected_session"):
+        session_dropdown.value = previously_selected_value
+        update_selection()
+
+    console = ipw.Output()
+    with console:
+        if last_session := np_session.Mouse(selection.mouse).state.get(
+            "last_P3_session"
+        ):
+            print(f"{mouse} last session: {last_session}")
+        print(f"Selected: {selection.session}")
+
+    def update(change):
+        if change["name"] != "value":
+            return
+        if (options := getattr(change["owner"], "options", None)) and change[
+            "new"
+        ] not in options:
+            return
+        if change["new"] == change["old"]:
+            return
+        update_selection()
+        with console:
+            print(f"Selected: {selection.session}")
+        global_state["selected_session"] = selection.session.value
+
+    session_dropdown.observe(update, names="value")
+
+    IPython.display.display(ipw.VBox([session_dropdown, console]))
+
+    return selection

np_workflows/experiments/openscope_P3/__init__.py

@@ -0,0 +1,2 @@
+from .main_P3_pilot import Ephys, Hab, new_experiment, validate_selected_workflow
+from .P3_workflow_widget import P3_workflow_widget

np_workflows/experiments/openscope_P3/main_P3_pilot.py

@@ -0,0 +1,215 @@
+import contextlib
+import enum
+import time
+
+import np_logging
+import np_session
+from np_services import (
+    Finalizable,
+    MouseDirector,
+    NewScaleCoordinateRecorder,
+    OpenEphys,
+    Service,
+    SessionCamstim,
+    Sync,
+    VideoMVR,
+)
+
+import np_workflows
+
+logger = np_logging.getLogger(__name__)
+
+
+class P3Session(enum.Enum):
+    """Enum for the different sessions available, each with different param sets."""
+
+    PRETEST = "pretest"
+    HAB = "hab"
+    EPHYS = "ephys"
+
+
+class P3Mixin:
+    """Provides project-specific methods and attributes, mainly related to camstim scripts."""
+
+    workflow: P3Session
+    """Enum for particular workflow/session, e.g. PRETEST, HAB_60, HAB_90,
+    EPHYS."""
+
+    session: np_session.PipelineSession
+    mouse: np_session.Mouse
+    user: np_session.User
+    platform_json: np_session.PlatformJson
+
+    @property
+    def recorders(self) -> tuple[Service, ...]:
+        """Services to be started before stimuli run, and stopped after. Session-dependent."""
+        match self.workflow:
+            case P3Session.PRETEST | P3Session.EPHYS:
+                return (Sync, VideoMVR, OpenEphys)
+            case P3Session.HAB:
+                return (Sync, VideoMVR)
+
+    @property
+    def stims(self) -> tuple[Service, ...]:
+        return (SessionCamstim,)
+
+    def initialize_and_test_services(self) -> None:
+        """Configure, initialize (ie. reset), then test all services."""
+
+        MouseDirector.user = self.user.id
+        MouseDirector.mouse = self.mouse.id
+
+        OpenEphys.folder = self.session.folder
+
+        NewScaleCoordinateRecorder.log_root = self.session.npexp_path
+        NewScaleCoordinateRecorder.log_name = self.platform_json.path.name
+
+        SessionCamstim.labtracks_mouse_id = self.mouse.id
+        SessionCamstim.lims_user_id = self.user.id
+
+        self.configure_services()
+
+        super().initialize_and_test_services()
+
+    def update_state(self) -> None:
+        "Store useful but non-essential info."
+        self.mouse.state["last_session"] = self.session.id
+        self.mouse.state["last_P3_session"] = str(self.workflow)
+        if self.mouse == 366122:
+            return
+        match self.workflow:
+            case P3Session.PRETEST:
+                return
+            case P3Session.HAB:
+                self.session.project.state["latest_hab"] = self.session.id
+            case P3Session.EPHYS:
+                self.session.project.state["latest_ephys"] = self.session.id
+                self.session.project.state["sessions"] = self.session.project.state.get(
+                    "sessions", []
+                ) + [self.session.id]
+
+    def run_stim(self) -> None:
+
+        self.update_state()
+
+        if not SessionCamstim.is_ready_to_start():
+            raise RuntimeError("SessionCamstim is not ready to start.")
+
+        np_logging.web(f"P3_{self.workflow.name.lower()}").info(
+            f"Started session {self.mouse.mtrain.stage['name']}"
+        )
+        SessionCamstim.start()
+
+        with contextlib.suppress(Exception):
+            while not SessionCamstim.is_ready_to_start():
+                time.sleep(2.5)
+
+        if isinstance(SessionCamstim, Finalizable):
+            SessionCamstim.finalize()
+
+        with contextlib.suppress(Exception):
+            np_logging.web(f"P3_{self.workflow.name.lower()}").info(
+                f"Finished session {self.mouse.mtrain.stage['name']}"
+            )
+
+    def copy_data_files(self) -> None:
+        super().copy_data_files()
+
+        # When all processing completes, camstim Agent class passes data and uuid to
+        # /camstim/lims BehaviorSession class, and write_behavior_data() writes a
+        # final .pkl with default name YYYYMMDDSSSS_mouseID_foragingID.pkl
+        # - if we have a foraging ID, we can search for that
+        if None == (
+            stim_pkl := next(
+                self.session.npexp_path.glob(
+                    f"{self.session.date:%y%m%d}*_{self.session.mouse}_*.pkl"
+                ),
+                None,
+            )
+        ):
+            logger.warning(
+                "Did not find stim file on npexp matching the format `YYYYMMDDSSSS_mouseID_foragingID.pkl`"
+            )
+            return
+        assert stim_pkl
+        if not self.session.platform_json.foraging_id:
+            self.session.platform_json.foraging_id = stim_pkl.stem.split("_")[-1]
+        new_stem = f"{self.session.folder}.stim"
+        logger.debug(f"Renaming stim file copied to npexp: {stim_pkl} -> {new_stem}")
+        stim_pkl = stim_pkl.rename(stim_pkl.with_stem(new_stem))
+
+        # remove other stim pkl, which is nearly identical, if it was also copied
+        for pkl in self.session.npexp_path.glob("*.pkl"):
+            if (
+                self.session.folder not in pkl.stem
+                and abs(pkl.stat().st_size - stim_pkl.stat().st_size) < 1e6
+            ):
+                logger.debug(f"Deleting extra stim pkl copied to npexp: {pkl.stem}")
+                pkl.unlink()
+
+
+def validate_selected_workflow(session: P3Session, mouse: np_session.Mouse) -> None:
+    for workflow in ("hab", "ephys"):
+        if (
+            workflow in session.value.lower()
+            and workflow not in mouse.mtrain.stage["name"].lower()
+        ) or (
+            session.value.lower() == "ephys"
+            and "hab" in mouse.mtrain.stage["name"].lower()
+        ):
+            raise ValueError(
+                f"Workflow selected ({session.value}) does not match MTrain stage ({mouse.mtrain.stage['name']}): please check cells above."
+            )
+
+
+class Hab(P3Mixin, np_workflows.PipelineHab):
+    def __init__(self, *args, **kwargs):
+        self.services = (
+            MouseDirector,
+            Sync,
+            VideoMVR,
+            self.imager,
+            NewScaleCoordinateRecorder,
+            SessionCamstim,
+        )
+        super().__init__(*args, **kwargs)
+
+
+class Ephys(P3Mixin, np_workflows.PipelineEphys):
+    def __init__(self, *args, **kwargs):
+        self.services = (
+            MouseDirector,
+            Sync,
+            VideoMVR,
+            self.imager,
+            NewScaleCoordinateRecorder,
+            SessionCamstim,
+            OpenEphys,
+        )
+        super().__init__(*args, **kwargs)
+
+
+# --------------------------------------------------------------------------------------
+
+
+def new_experiment(
+    mouse: int | str | np_session.Mouse,
+    user: str | np_session.User,
+    workflow: P3Session,
+) -> Ephys | Hab:
+    """Create a new experiment for the given mouse and user."""
+    match workflow:
+        case P3Session.PRETEST | P3Session.EPHYS:
+            experiment = Ephys(mouse, user)
+        case P3Session.HAB:
+            experiment = Hab(mouse, user)
+        case _:
+            raise ValueError(f"Invalid workflow type: {workflow}")
+    experiment.workflow = workflow
+
+    with contextlib.suppress(Exception):
+        np_logging.web(f"P3_{experiment.workflow.name.lower()}").info(
+            f"{experiment} created"
+        )
+
+    return experiment

np_workflows/experiments/openscope_barcode/__init__.py

@@ -1,2 +1,2 @@
-from .main_barcode_pilot import new_experiment, Hab, Ephys, validate_selected_workflow
 from .barcode_workflow_widget import barcode_workflow_widget
+from .main_barcode_pilot import Ephys, Hab, new_experiment, validate_selected_workflow

np_workflows/experiments/openscope_barcode/barcode_workflow_widget.py

@@ -1,17 +1,6 @@
-import configparser
-import contextlib
-import copy
-import enum
-import functools
-from typing import ClassVar, Literal, NamedTuple, NoReturn, Optional, TypedDict
-
 import IPython.display
 import ipywidgets as ipw
-import np_config
-import np_logging
 import np_session
-import np_workflows
-from pyparsing import Any
 
 from np_workflows.experiments.openscope_barcode.main_barcode_pilot import BarcodeSession
 
@@ -20,8 +9,11 @@ global_state = {}
 
 # for widget, before creating a experiment --------------------------------------------- #
 
+
 class SelectedSession:
-    def __init__(self, session: str | BarcodeSession, mouse: str | int | np_session.Mouse):
+    def __init__(
+        self, session: str | BarcodeSession, mouse: str | int | np_session.Mouse
+    ):
         if isinstance(session, str):
             session = BarcodeSession(session)
         self.session = session
@@ -48,17 +40,19 @@ def barcode_workflow_widget(
         options=tuple(_.value for _ in BarcodeSession),
         description="Session",
     )
-    
+
     def update_selection():
         selection.__init__(str(session_dropdown.value), str(mouse))
-        
-    if (previously_selected_value := global_state.get('selected_session')):
+
+    if previously_selected_value := global_state.get("selected_session"):
         session_dropdown.value = previously_selected_value
         update_selection()
-        
+
     console = ipw.Output()
     with console:
-        if last_session := np_session.Mouse(selection.mouse).state.get('last_barcode_session'):
+        if last_session := np_session.Mouse(selection.mouse).state.get(
+            "last_barcode_session"
+        ):
             print(f"{mouse} last session: {last_session}")
         print(f"Selected: {selection.session}")
 
@@ -74,9 +68,9 @@ def barcode_workflow_widget(
         update_selection()
         with console:
             print(f"Selected: {selection.session}")
-        global_state['selected_session'] = selection.session.value
-        
-    session_dropdown.observe(update, names='value')
+        global_state["selected_session"] = selection.session.value
+
+    session_dropdown.observe(update, names="value")
 
     IPython.display.display(ipw.VBox([session_dropdown, console]))
 

np_workflows/experiments/openscope_barcode/camstim_scripts/barcode_mapping_script.py

@@ -5,16 +5,8 @@ May '23 OpenScope: Barcode stimuli
 import argparse
 import json
 import logging
-import os
-import time
-
-import numpy as np
-from psychopy import visual
-from camstim import Foraging
-from camstim import Stimulus_v2
-from camstim import SweepStim_v2
-from camstim import Warp, Window
 
+from camstim import Foraging, Stimulus_v2, SweepStim_v2, Warp, Window
 
 # get params ------------------------------------------------------------------
 # stored in json file -
@@ -29,7 +21,7 @@ parser.add_argument(
 )
 args, _ = parser.parse_known_args()
 
-with open(args.params_path, "r") as f:
+with open(args.params_path) as f:
     json_params = json.load(f)
 
 # Create display window
@@ -41,6 +33,7 @@ window = Window(
     warp=Warp.Spherical,
 )
 
+
 # patch the Stimulus_v2 class to allow for serializing without large arrays
 # ----------------------------------------------------------------------------
 class Stimulus_v2_MinusFrameArrays(Stimulus_v2):
@@ -56,11 +49,12 @@ class Stimulus_v2_MinusFrameArrays(Stimulus_v2):
             self.sweep_table = None
             self.sweep_params = self.sweep_params.keys()
         self_dict = self.__dict__
-        del self_dict['sweep_frames']
-        del self_dict['sweep_order']
-        self_dict['stim'] = str(self_dict['stim'])
+        del self_dict["sweep_frames"]
+        del self_dict["sweep_order"]
+        self_dict["stim"] = str(self_dict["stim"])
         return wecanpicklethat(self_dict)
-    
+
+
 # ----------------------------------------------------------------------------
 # setup mapping stim
 """from mapping_script_v2.py"""