np-workflows 1.6.89__py3-none-any.whl

np_workflows/experiments/openscope_loop/camstim_scripts/barcode_mapping_script.py

@@ -0,0 +1,138 @@
+"""
+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
+
+
+# get params ------------------------------------------------------------------
+# stored in json file -
+# path to json supplied by camstim via command line arg when this script is called
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "params_path",
+    nargs="?",
+    type=str,
+    default="",
+)
+args, _ = parser.parse_known_args()
+
+with open(args.params_path, "r") as f:
+    json_params = json.load(f)
+
+# Create display window
+# ----------------------------------------------------------------------------
+window = Window(
+    fullscr=True,
+    monitor=json_params["monitor"],
+    screen=0,
+    warp=Warp.Spherical,
+)
+
+# patch the Stimulus_v2 class to allow for serializing without large arrays
+# ----------------------------------------------------------------------------
+class Stimulus_v2_MinusFrameArrays(Stimulus_v2):
+
+    def __init__(self, *args, **kwargs):
+        super(Stimulus_v2_MinusFrameArrays, self).__init__(*args, **kwargs)
+
+    def package(self):
+        """
+        Package for serializing - minus large arrays of frame timing/order.
+        """
+        if not self.save_sweep_table:
+            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'])
+        return wecanpicklethat(self_dict)
+    
+# ----------------------------------------------------------------------------
+# setup mapping stim
+"""from mapping_script_v2.py"""
+
+mapping_stimuli = []
+
+# load common stimuli
+gabor_path = json_params["gabor_path"]
+flash_path = json_params["flash_path"]
+gabor = Stimulus_v2_MinusFrameArrays.from_file(gabor_path, window)
+flash = Stimulus_v2_MinusFrameArrays.from_file(flash_path, window)
+
+gabor_duration_sec = json_params["default_gabor_duration_seconds"]
+flash_duration_sec = json_params["default_flash_duration_seconds"]
+
+original_duration_sec = gabor_duration_sec + flash_duration_sec
+
+# if max total duration is set, and less than original movie length, cut down display sequence:
+max_mapping_duation_minutes = json_params[
+    "max_total_duration_minutes"
+]  # can be zero, in which case we use the full movie length
+max_mapping_duration_sec = max_mapping_duation_minutes * 60
+if 0 < max_mapping_duration_sec < original_duration_sec:
+    logging.info("Mapping duration capped at %s minutes", max_mapping_duation_minutes)
+
+    logging.info("original gabor duration: %s sec", gabor_duration_sec)
+    logging.info("original flash duration: %s sec", flash_duration_sec)
+    logging.info("max mapping duration: %s sec", max_mapping_duration_sec)
+
+    gabor_duration_sec = (
+        max_mapping_duration_sec * gabor_duration_sec
+    ) / original_duration_sec
+    flash_duration_sec = (
+        max_mapping_duration_sec * flash_duration_sec
+    ) / original_duration_sec
+
+    logging.info("modified gabor duration: %s sec", gabor_duration_sec)
+    logging.info("modified flash duration: %s sec", flash_duration_sec)
+
+# setup timing
+mapping_sequence_start_sec = 0  # if stims are daisy-chained within one script, this should be the end of the prev stim
+gabor.set_display_sequence([(mapping_sequence_start_sec, gabor_duration_sec)])
+flash.set_display_sequence(
+    [(gabor_duration_sec, gabor_duration_sec + flash_duration_sec)]
+)
+
+mapping_stimuli = [gabor, flash]
+
+mapping_sequence_end_sec = (
+    gabor_duration_sec + flash_duration_sec
+)  # if daisy-chained, the next stim in this script should start at this time
+
+# create SweepStim_v2 instance for main stimulus
+ss = SweepStim_v2(
+    window,
+    stimuli=mapping_stimuli,
+    pre_blank_sec=json_params["pre_blank_screen_sec"],
+    post_blank_sec=json_params["post_blank_screen_sec"],
+    params=json_params["sweepstim"],
+)
+
+# add in foraging so we can track wheel, potentially give rewards, etc
+f = Foraging(
+    window=window,
+    auto_update=False,
+    params=json_params["sweepstim"],
+    nidaq_tasks={
+        "digital_input": ss.di,
+        "digital_output": ss.do,
+    },
+)  # share di and do with SS
+
+ss.add_item(f, "foraging")
+
+ss.run()

np_workflows/experiments/openscope_loop/camstim_scripts/barcode_opto_script.py

@@ -0,0 +1,219 @@
+# -*- coding: utf-8 -*-
+"""
+optotagging.py
+
+runs optotagging code for ecephys pipeline experiments
+
+by joshs@alleninstitute.org
+
+(c) 2018 Allen Institute for Brain Science
+
+"""
+import camstim  # ensures "magic" gets setup properly by importing first
+import logging  # must occur after camstim import for "magic"
+from camstim.zro import agent
+
+import numpy as np
+from toolbox.IO.nidaq import AnalogOutput
+from toolbox.IO.nidaq import DigitalOutput
+
+import datetime
+import numpy as np
+import time
+import pickle as pkl
+
+
+# %%
+
+
+def run_optotagging(levels, conditions, waveforms, isis, sampleRate=10000.):
+
+    from toolbox.IO.nidaq import AnalogOutput
+    from toolbox.IO.nidaq import DigitalOutput
+
+    sweep_on = np.array([0, 0, 1, 0, 0, 0, 0, 0], dtype=np.uint8)
+    stim_on = np.array([0, 0, 1, 1, 0, 0, 0, 0], dtype=np.uint8)
+    stim_off = np.array([0, 0, 1, 0, 0, 0, 0, 0], dtype=np.uint8)
+    sweep_off = np.array([0, 0, 0, 0, 0, 0, 0, 0], dtype=np.uint8)
+
+    ao = AnalogOutput('Dev1', channels=[1])
+    ao.cfg_sample_clock(sampleRate)
+
+    do = DigitalOutput('Dev1', 2)
+
+    do.start()
+    ao.start()
+
+    do.write(sweep_on)
+    time.sleep(5)
+
+    for i, level in enumerate(levels):
+
+        print(level)
+
+        data = waveforms[conditions[i]]
+
+        do.write(stim_on)
+        ao.write(data * level)
+        do.write(stim_off)
+        time.sleep(isis[i])
+
+    do.write(sweep_off)
+    do.clear()
+    ao.clear()
+
+# %%
+
+
+def generatePulseTrain(pulseWidth, pulseInterval, numRepeats, riseTime, sampleRate=10000.):
+
+    data = np.zeros((int(sampleRate),), dtype=np.float64)
+   # rise_samples =
+
+    rise_and_fall = (
+        ((1 - np.cos(np.arange(sampleRate*riseTime/1000., dtype=np.float64)*2*np.pi/10))+1)-1)/2
+    half_length = rise_and_fall.size / 2
+    rise = rise_and_fall[:half_length]
+    fall = rise_and_fall[half_length:]
+
+    peak_samples = int(sampleRate*(pulseWidth-riseTime*2)/1000)
+    peak = np.ones((peak_samples,))
+
+    pulse = np.concatenate((rise,
+                            peak,
+                            fall))
+
+    interval = int(pulseInterval*sampleRate/1000.)
+
+    for i in range(0, numRepeats):
+        data[i*interval:i*interval+pulse.size] = pulse
+
+    return data
+
+
+# %% create waveforms
+
+def optotagging(mouseID, operation_mode='experiment', level_list=[1.15, 1.28, 1.345], genotype=None):
+
+    sampleRate = 10000
+
+    # 1 s cosine ramp:
+    data_cosine = (((1 - np.cos(np.arange(sampleRate, dtype=np.float64)
+                                * 2*np.pi/sampleRate)) + 1) - 1)/2  # create raised cosine waveform
+
+    # 1 ms cosine ramp:
+    rise_and_fall = (
+        ((1 - np.cos(np.arange(sampleRate*0.001, dtype=np.float64)*2*np.pi/10))+1)-1)/2
+    half_length = rise_and_fall.size / 2
+
+    # pulses with cosine ramp:
+    pulse_2ms = np.concatenate((rise_and_fall[:half_length], np.ones(
+        (int(sampleRate*0.001),)), rise_and_fall[half_length:]))
+    pulse_5ms = np.concatenate((rise_and_fall[:half_length], np.ones(
+        (int(sampleRate*0.004),)), rise_and_fall[half_length:]))
+    pulse_10ms = np.concatenate((rise_and_fall[:half_length], np.ones(
+        (int(sampleRate*0.009),)), rise_and_fall[half_length:]))
+
+    data_2ms_10Hz = np.zeros((sampleRate,), dtype=np.float64)
+
+    for i in range(0, 10):
+        interval = sampleRate / 10
+        data_2ms_10Hz[i*interval:i*interval+pulse_2ms.size] = pulse_2ms
+
+    data_5ms = np.zeros((sampleRate,), dtype=np.float64)
+    data_5ms[:pulse_5ms.size] = pulse_5ms
+
+    data_10ms = np.zeros((sampleRate,), dtype=np.float64)
+    data_10ms[:pulse_10ms.size] = pulse_10ms
+
+    data_10s = np.zeros((sampleRate*10,), dtype=np.float64)
+    data_10s[:-2] = 1
+
+    # %% for experiment
+
+    isi = 1.5
+    isi_rand = 0.5
+    numRepeats = 50
+
+    condition_list = [2, 3]
+    waveforms = [data_2ms_10Hz, data_5ms, data_10ms, data_cosine]
+
+    opto_levels = np.array(level_list*numRepeats*len(condition_list))  # BLUE
+    opto_conditions = condition_list*numRepeats*len(level_list)
+    opto_conditions = np.sort(opto_conditions)
+    opto_isis = np.random.random(opto_levels.shape) * isi_rand + isi
+
+    p = np.random.permutation(len(opto_levels))
+
+    # implement shuffle?
+    opto_levels = opto_levels[p]
+    opto_conditions = opto_conditions[p]
+
+    # %% for testing
+
+    if operation_mode == 'test_levels':
+        isi = 2.0
+        isi_rand = 0.0
+
+        numRepeats = 2
+
+        condition_list = [0]
+        waveforms = [data_10s, data_10s]
+
+        opto_levels = np.array(level_list*numRepeats *
+                               len(condition_list))  # BLUE
+        opto_conditions = condition_list*numRepeats*len(level_list)
+        opto_conditions = np.sort(opto_conditions)
+        opto_isis = np.random.random(opto_levels.shape) * isi_rand + isi
+
+    elif operation_mode == 'pretest':
+        numRepeats = 1
+
+        condition_list = [0]
+        data_2s = data_10s[-sampleRate*2:]
+        waveforms = [data_2s]
+
+        opto_levels = np.array(level_list*numRepeats *
+                               len(condition_list))  # BLUE
+        opto_conditions = condition_list*numRepeats*len(level_list)
+        opto_conditions = np.sort(opto_conditions)
+        opto_isis = [1]*len(opto_conditions)
+    # %%
+
+    outputDirectory = agent.OUTPUT_DIR
+    fileDate = str(datetime.datetime.now()).replace(':', '').replace(
+        '.', '').replace('-', '').replace(' ', '')[2:14]
+    fileName = outputDirectory  + "/" + fileDate + '_'+mouseID + '.opto.pkl'
+
+    print('saving info to: ' + fileName)
+    fl = open(fileName, 'wb')
+    output = {}
+
+    output['opto_levels'] = opto_levels
+    output['opto_conditions'] = opto_conditions
+    output['opto_ISIs'] = opto_isis
+    output['opto_waveforms'] = waveforms
+
+    pkl.dump(output, fl)
+    fl.close()
+    print('saved.')
+
+    # %%
+    run_optotagging(opto_levels, opto_conditions,
+                    waveforms, opto_isis, float(sampleRate))
+
+
+# %%
+if __name__ == "__main__":
+    import json
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('json_params', type=str, )
+    args, _ = parser.parse_known_args()
+
+    with open(args.json_params, 'r', ) as f:
+        json_params = json.load(f)
+
+    logging.info('Optotagging with params: %s' % json_params)
+    optotagging(**json_params)

np_workflows/experiments/openscope_loop/loop_workflow_widget.py

@@ -0,0 +1,83 @@
+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_loop.main_loop_pilot import LoopSession
+
+global_state = {}
+"""Global variable for persisting widget states."""
+
+# for widget, before creating a experiment --------------------------------------------- #
+
+class SelectedSession:
+    def __init__(self, session: str | LoopSession, mouse: str | int | np_session.Mouse):
+        if isinstance(session, str):
+            session = LoopSession(session)
+        self.session = session
+        self.mouse = str(mouse)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.session}, {self.mouse})"
+
+
+def loop_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(LoopSession.PRETEST, mouse)
+
+    session_dropdown = ipw.Select(
+        options=tuple(_.value for _ in LoopSession),
+        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_loop_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_loop/main_loop_pilot.py

@@ -0,0 +1,217 @@
+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, 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 PIL.Image
+import pydantic
+from pyparsing import Any
+from np_services import (
+    Service,
+    Finalizable,
+    ScriptCamstim, SessionCamstim, 
+    SessionCamstim,
+    OpenEphys,
+    Sync,
+    VideoMVR,
+    NewScaleCoordinateRecorder,
+    MouseDirector,
+)
+
+logger = np_logging.getLogger(__name__)
+
+
+class LoopSession(enum.Enum):
+    """Enum for the different sessions available, each with different param sets."""
+
+    PRETEST = "pretest"
+    HAB = "hab"
+    EPHYS = "ephys"
+
+
+class LoopMixin:
+    """Provides project-specific methods and attributes, mainly related to camstim scripts."""
+    
+    workflow: LoopSession
+    """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 LoopSession.PRETEST | LoopSession.EPHYS:
+                return (Sync, VideoMVR, OpenEphys)
+            case LoopSession.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_Loop_session'] = str(self.workflow)
+        if self.mouse == 366122:
+            return
+        match self.workflow:
+            case LoopSession.PRETEST:
+                return
+            case LoopSession.HAB:
+                self.session.project.state['latest_hab'] = self.session.id
+            case LoopSession.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'Loop_{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'Loop_{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: LoopSession, 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(LoopMixin, np_workflows.PipelineHab):
+    def __init__(self, *args, **kwargs):
+        self.services = (
+            MouseDirector,
+            Sync,
+            VideoMVR,
+            self.imager,
+            NewScaleCoordinateRecorder,
+            SessionCamstim,
+        )
+        super().__init__(*args, **kwargs)
+
+
+class Ephys(LoopMixin, 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: LoopSession,
+) -> Ephys | Hab:
+    """Create a new experiment for the given mouse and user."""
+    match workflow:
+        case LoopSession.PRETEST | LoopSession.EPHYS:
+            experiment = Ephys(mouse, user)
+        case LoopSession.HAB:
+            experiment = Hab(mouse, user)
+        case _:
+            raise ValueError(f"Invalid workflow type: {workflow}")
+    experiment.workflow = workflow
+    
+    with contextlib.suppress(Exception):
+        np_logging.web(f'Loop_{experiment.workflow.name.lower()}').info(f"{experiment} created")
+            
+    return experiment
+

np_workflows/experiments/openscope_psycode/__init__.py

@@ -0,0 +1,2 @@
+from .main_psycode_pilot import new_experiment, Hab, Ephys, validate_selected_workflow
+from .psycode_workflow_widget import PsyCode_workflow_widget