Feat fast retract

pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
     "aind-data-schema==1.1.0",
     "aind-data-schema-models==0.5.6",
     "pydantic==2.10.6",
-    "stagewidget==1.0.4.dev11",
+    "stagewidget==1.0.5.dev0",
     "python-logging-loki >=0.3.1, <2",
     "pykeepass >=4.0.7, <5",
     "pyyaml >=6, <7",

src/foraging_gui/Foraging.py

@@ -125,14 +125,6 @@
     SessionParametersWidget,
 )
 from foraging_gui.settings_model import BonsaiSettingsModel, DFTSettingsModel
-from foraging_gui.stage import Stage
-from foraging_gui.Visualization import (
-    PlotLickDistribution,
-    PlotTimeDistribution,
-    PlotV,
-)
-from foraging_gui.warning_widget import WarningWidget
-from foraging_gui.settings_model import BonsaiSettingsModel, DFTSettingsModel
 from foraging_gui.sound_button import SoundButton
 from foraging_gui.stage import Stage
 from foraging_gui.Visualization import (
@@ -357,6 +349,11 @@ def __init__(self, parent=None, box_number=1, start_bonsai_ide=True):
 
         # Stage Widget
         self.stage_widget = None
+        # initialize empty timers
+        self.left_retract_timer = QTimer(timeout=lambda: None)
+        self.left_retract_timer.setSingleShot(True)
+        self.right_retract_timer = QTimer(timeout=lambda: None)
+        self.right_retract_timer.setSingleShot(True)
         try:
             self._load_stage()
         except IOError as e:
@@ -624,6 +621,103 @@ def _insert_stage_widget(self, widget_to_replace: str) -> None:
             self.stage_widget = get_stage_widget()
             layout.addWidget(self.stage_widget)
 
+    def retract_lick_spout(self, lick_spout_licked: Literal["Left", "Right"], delta: float = -10) -> None:
+        """
+        Fast retract lick spout based on lick spout licked
+
+        :param lick_spout_licked: lick spout that was licked. Opposite lickspout will be retracted
+        :param delta: change in pos to move lick spout to. Default is -10
+
+        """
+        # disconnect so it's only triggered once
+        try:
+            self.Channel2.mouseLicked.disconnect(self.retract_lick_spout)
+        except TypeError:
+            pass
+
+        lick_spout_retract = "right" if lick_spout_licked == "Left" else "left"
+        timer = getattr(self, f"{lick_spout_retract}_retract_timer")
+        tp = self.task_logic.task_parameters
+        motor = 1 if lick_spout_licked == "Left" else 2
+        at_origin = list(self._GetPositions().values())[motor] == 0.0
+
+        if tp.lick_spout_retraction and self.stage_widget is not None and not at_origin:
+            logger.info(f"Retracting {lick_spout_retract} lick spout.")
+            curr_pos = self.stage_widget.stage_model.get_current_positions_mm(motor, rel_to_monument=True)
+            self.stage_widget.stage_model.quick_move(motor=motor, distance=delta, skip_if_busy=False)
+            logger.info(f"Retracting {lick_spout_retract} lick spout to pos {delta} from current pos {curr_pos}.")
+            # configure timer to un-retract lick spout
+            timer.timeout.disconnect()
+            timer.timeout.connect(lambda: self.un_retract_lick_spout(lick_spout_retract.title(), curr_pos))
+            timer.setInterval(self.operation_control_model.lick_spout_retraction_specs.wait_time*1000)
+            timer.setSingleShot(True)
+            timer.start()
+
+        elif self.stage_widget is None:
+            logger.info("Can't fast retract stage because AIND stage is not being used.",
+                        extra={"tags": [self.warning_log_tag]})
+
+        elif tp.lick_spout_retraction or at_origin:
+            try:
+                self.Channel2.mouseLicked.connect(self.retract_lick_spout, type=Qt.UniqueConnection)
+            except TypeError:  # signal already connected
+                logger.debug("Mouse lick signal already connected.")
+            logger.debug("Cannot retract stage because " + "lickspouts at origin." if at_origin
+                        else "retraction turned off.")
+
+    def un_retract_lick_spout(self, lick_spout_retracted: Literal["Left", "Right"], pos: float) -> None:
+        """
+        Un-retract specified lick spout
+
+        :param lick_spout_retracted: lick spout that was retracted.
+        :param pos: pos to move lick spout to.
+
+        """
+
+        if self.stage_widget is not None:
+            logger.info(f"Un-retracting {lick_spout_retracted.lower()} lick spout. Moving back to position {pos}.")
+            speed = self.operation_control_model.lick_spout_retraction_specs.un_retract_speed.value
+            motor = 2 if lick_spout_retracted == "Left" else 1
+            # set un-retract speed. Set back to normal when user presses stop
+            self.stage_widget.stage_model.update_speed(value=speed)
+            self.stage_widget.stage_model.update_position(positions={motor: pos})
+
+            # create thread to track when stage has made it back to previous position and reconnect fast retract
+            Thread(target=self.reset_fast_retract, args=(motor, pos)).start()
+
+        else:
+            logger.info("Can't un retract lick spout because no AIND stage connected")
+
+    def reset_fast_retract(self, motor: int, pos: float):
+        """"
+            Resets AIND stage fast retract when spout has returned to original pos
+
+            :param motor: motor index to query
+            :param pos: original pos to wait for
+        """
+
+        if self.stage_widget is not None:
+
+            while (curr := self.stage_widget.stage_model.get_current_positions_mm(motor, rel_to_monument=True)) < pos:
+                logger.info(f"Stage at {curr}")
+                time.sleep(.1)
+
+            logger.info("Setting stage to normal speed.")
+            try:
+                self.stage_widget.stage_model.move_worker.finished.disconnect(self.reset_fast_retract)
+            except TypeError:  # signal isn't connected
+                pass
+            self.stage_widget.stage_model.update_speed(value=1)
+
+            logger.info("Reconnecting fast retract at pos.")
+            try:
+                self.Channel2.mouseLicked.connect(self.retract_lick_spout, type=Qt.UniqueConnection)
+            except TypeError:  # signal already connected
+                logger.debug("Mouse lick signal already connected.")
+
+        else:
+            logger.info("Can't set stage speed because no AIND stage connected")
+
     def _LoadUI(self):
         """
         Determine which user interface to use
@@ -845,7 +939,7 @@ def update_loaded_mouse_offset(self):
             return
 
         elif list(current_positions.keys()) == ["x", "y", "z"]:
-            logging.info(
+            logging.debug(
                 "Can't update loaded mouse offset with non AIND stage coordinates."
             )
         else:
@@ -885,7 +979,7 @@ def update_operational_control_stage_positions(self, *args) -> None:
             self.operation_control_model.stage_specs.step_size = step
 
 
-        
+
 
     def update_stage_positions_from_operational_control(
         self, oc: OperationalControl = None
@@ -2759,6 +2853,7 @@ def _ConnectOSC(self):
         self.client2 = OSCStreamingClient()
         self.client2.connect((self.ip, self.request_port2))
         self.Channel2 = rigcontrol.RigClient(self.client2)
+
         # manually give water
         self.client3 = OSCStreamingClient()  # Create client
         self.client3.connect((self.ip, self.request_port3))
@@ -3576,7 +3671,7 @@ def _Save(self, ForceSave=0, SaveAs=0, SaveContinue=0, BackupSave=0):
             and self.InitializeBonsaiSuccessfully == 1
             and BackupSave == 0
         ):
-            self.GeneratedTrials._get_irregular_timestamp(self.Channel2)
+            self.GeneratedTrials._get_irregular_timestamp(self.Channel2, self.data_lock)
 
         # Create new folders.
         if self.CreateNewFolder == 1:
@@ -4209,7 +4304,11 @@ def _LoadVisualization(self):
             return
 
         self.PlotM = PlotV(
-            win=self, GeneratedTrials=self.GeneratedTrials, width=5, height=4
+            win=self,
+            data_lock=self.data_lock,
+            GeneratedTrials=self.GeneratedTrials,
+            width=5,
+            height=4
         )
         self.PlotM.setSizePolicy(
             QSizePolicy.MinimumExpanding, QSizePolicy.MinimumExpanding
@@ -5007,6 +5106,11 @@ def _Start(self):
             # set flag to perform habituation period
             self.behavior_baseline_period.set()
 
+            try:    # connect signal for fast retraction
+                self.Channel2.mouseLicked.connect(self.retract_lick_spout, type=Qt.UniqueConnection)
+            except TypeError:  # signal already connected
+                logger.debug("Mouse lick signal already connected.")
+
             self.session_run = True  # session has been started
         else:
             # Prompt user to confirm stopping trials
@@ -5047,6 +5151,12 @@ def _Start(self):
             self.sound_button.setEnabled(True)
             self.behavior_baseline_period.clear()   # set flag to break out of habituation period
 
+            # disconnect fast retract signals if connected
+            try:
+                self.Channel2.mouseLicked.disconnect(self.retract_lick_spout)
+            except TypeError:
+                pass
+
         if (self.StartANewSession == 1) and (self.ANewTrial == 0):
             # If we are starting a new session, we should wait for the last trial to finish
             self._StopCurrentSession()
@@ -5102,7 +5212,11 @@ def _Start(self):
             self.GeneratedTrials = GeneratedTrials
             self.StartANewSession = 0
             PlotM = PlotV(
-                win=self, GeneratedTrials=GeneratedTrials, width=5, height=4
+                win=self,
+                data_lock=self.data_lock,
+                GeneratedTrials=GeneratedTrials,
+                width=5,
+                height=4
             )
             # PlotM.finish=1
             self.PlotM = PlotM
@@ -5156,6 +5270,7 @@ def _Start(self):
                 GeneratedTrials._get_irregular_timestamp, self.Channel2
             )
             workerLick.signals.finished.connect(self._thread_complete2)
+
             workerPlot = Worker(
                 PlotM._Update,
                 GeneratedTrials=GeneratedTrials,
@@ -5230,7 +5345,6 @@ def _Start(self):
                 "Running photometry baseline",
                 extra={"tags": [self.warning_log_tag]},
             )
-
         self._StartTrialLoop(GeneratedTrials, worker1, worker_save)
 
         if self.actionDrawing_after_stopping.isChecked() == True:

src/foraging_gui/MyFunctions.py

@@ -7,6 +7,7 @@
 from datetime import datetime
 from itertools import accumulate
 from sys import platform as PLATFORM
+from threading import Lock
 
 import numpy as np
 import requests
@@ -2883,7 +2884,7 @@ def _add_one_trial(self):
                         self.BlockLenHistory[i][-1] + 1
                     )
 
-    def _GetAnimalResponse(self, Channel1, Channel3, data_lock):
+    def _GetAnimalResponse(self, Channel1, Channel3, data_lock: Lock):
         """Get the animal's response"""
         self._CheckSimulationSession()
         if self.CurrentSimulation:
@@ -3114,67 +3115,88 @@ def _GiveRight(self, channel3):
         channel3.ManualWater_Right(int(1))
         channel3.RightValue1(float(self.win.right_valve_open_time * 1000))
 
-    def _get_irregular_timestamp(self, Channel2):
+    def _get_irregular_timestamp(self, Channel2, data_lock: Lock):
         """Get timestamps occurred irregularly (e.g. licks and reward delivery time)"""
+
         while not Channel2.msgs.empty():
             Rec = Channel2.receive()
+
             if Rec[0].address == "/LeftLickTime":
-                self.B_LeftLickTime = np.append(
-                    self.B_LeftLickTime, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_LeftLickTime = np.append(
+                        self.B_LeftLickTime, Rec[1][1][0]
+                    )
+
             elif Rec[0].address == "/RightLickTime":
-                self.B_RightLickTime = np.append(
-                    self.B_RightLickTime, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_RightLickTime = np.append(
+                        self.B_RightLickTime, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/LeftRewardDeliveryTime":
-                self.B_LeftRewardDeliveryTime = np.append(
-                    self.B_LeftRewardDeliveryTime, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_LeftRewardDeliveryTime = np.append(
+                        self.B_LeftRewardDeliveryTime, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/RightRewardDeliveryTime":
-                self.B_RightRewardDeliveryTime = np.append(
-                    self.B_RightRewardDeliveryTime, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_RightRewardDeliveryTime = np.append(
+                        self.B_RightRewardDeliveryTime, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/LeftRewardDeliveryTimeHarp":
-                self.B_LeftRewardDeliveryTimeHarp = np.append(
-                    self.B_LeftRewardDeliveryTimeHarp, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_LeftRewardDeliveryTimeHarp = np.append(
+                        self.B_LeftRewardDeliveryTimeHarp, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/RightRewardDeliveryTimeHarp":
-                self.B_RightRewardDeliveryTimeHarp = np.append(
-                    self.B_RightRewardDeliveryTimeHarp, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_RightRewardDeliveryTimeHarp = np.append(
+                        self.B_RightRewardDeliveryTimeHarp, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/PhotometryRising":
-                self.B_PhotometryRisingTimeHarp = np.append(
-                    self.B_PhotometryRisingTimeHarp, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_PhotometryRisingTimeHarp = np.append(
+                        self.B_PhotometryRisingTimeHarp, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/PhotometryFalling":
-                self.B_PhotometryFallingTimeHarp = np.append(
-                    self.B_PhotometryFallingTimeHarp, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_PhotometryFallingTimeHarp = np.append(
+                        self.B_PhotometryFallingTimeHarp, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/OptogeneticsTimeHarp":
-                self.B_OptogeneticsTimeHarp = np.append(
-                    self.B_OptogeneticsTimeHarp, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_OptogeneticsTimeHarp = np.append(
+                        self.B_OptogeneticsTimeHarp, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/ManualLeftWaterStartTime":
-                self.B_ManualLeftWaterStartTime = np.append(
-                    self.B_ManualLeftWaterStartTime, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_ManualLeftWaterStartTime = np.append(
+                        self.B_ManualLeftWaterStartTime, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/ManualRightWaterStartTime":
-                self.B_ManualRightWaterStartTime = np.append(
-                    self.B_ManualRightWaterStartTime, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_ManualRightWaterStartTime = np.append(
+                        self.B_ManualRightWaterStartTime, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/EarnedLeftWaterStartTime":
-                self.B_EarnedLeftWaterStartTime = np.append(
-                    self.B_EarnedLeftWaterStartTime, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_EarnedLeftWaterStartTime = np.append(
+                        self.B_EarnedLeftWaterStartTime, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/EarnedRightWaterStartTime":
-                self.B_EarnedRightWaterStartTime = np.append(
-                    self.B_EarnedRightWaterStartTime, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_EarnedRightWaterStartTime = np.append(
+                        self.B_EarnedRightWaterStartTime, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/AutoLeftWaterStartTime":
-                self.B_AutoLeftWaterStartTime = np.append(
-                    self.B_AutoLeftWaterStartTime, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_AutoLeftWaterStartTime = np.append(
+                        self.B_AutoLeftWaterStartTime, Rec[1][1][0]
+                    )
             elif Rec[0].address == "/AutoRightWaterStartTime":
+                with data_lock:
+                    self.B_AutoRightWaterStartTime = np.append(
+                        self.B_AutoRightWaterStartTime, Rec[1][1][0]
+                    )
                 self.B_AutoRightWaterStartTime = np.append(
                     self.B_AutoRightWaterStartTime, Rec[1][1][0]
                 )