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 (
@@ -347,6 +339,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:
@@ -596,6 +593,7 @@ def _load_stage(self) -> None:
                 else "widget_2"
             )
             self._insert_stage_widget(widget_to_replace)
+
         else:
             self._GetPositions()
 
@@ -619,6 +617,88 @@ 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"], pos: float = 0) -> 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 pos: pos to move lick spout to. Default is 0
+
+        """
+        # 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
+        at_origin = list(self._GetPositions().values())[1:3] == [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.")
+            motor = 1 if lick_spout_licked == "Left" else 2                             # TODO: is this the correct mapping
+            curr_pos = self.stage_widget.stage_model.get_current_positions_mm(motor)    # TODO: Do I need to set rel_to_monument to True?
+            self.stage_widget.stage_model.quick_move(motor=motor, distance=pos-curr_pos, skip_if_busy=True)
+
+            # configure timer to un-retract lick spout
+            timer.timeout.disconnect()
+            timer.timeout.connect(lambda: self.un_retract_lick_spout(lick_spout_licked, 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_licked: Literal["Left", "Right"], pos: float = 0) -> None:
+        """
+        Un-retract specified lick spout
+
+        :param lick_spout_licked: lick spout that was licked. Opposite licks pout will be un-retracted
+        :param pos: pos to move lick spout to. Default is 0
+
+        """
+        if self.stage_widget is not None:
+            logger.info("Un-retracting lick spout.")
+            speed = self.operation_control_model.lick_spout_retraction_specs.un_retract_speed.value
+            motor = 1 if lick_spout_licked == "Left" else 2
+            self.stage_widget.stage_model.update_speed(value=speed)
+            self.stage_widget.stage_model.update_position(positions={motor:pos})
+            self.stage_widget.stage_model.move_worker.finished.connect(self.set_stage_speed_to_normal,
+                                                                       type=Qt.UniqueConnection)
+        else:
+            logger.info("Can't un retract lick spout because no AIND stage connected")
+        try:
+            self.Channel2.mouseLicked.connect(self.retract_lick_spout, type=Qt.UniqueConnection)
+        except TypeError:   # signal already connected
+            logger.debug("Mouse lick signal already connected.")
+
+    def set_stage_speed_to_normal(self):
+        """"
+        Sets AIND stage to normal speed
+        """
+
+        if self.stage_widget is not None:
+            logger.info("Setting stage to normal speed.")
+            try:
+                self.stage_widget.stage_model.move_worker.finished.disconnect(self.set_stage_speed_to_normal)
+            except TypeError:   # signal isn't connected
+                pass
+            self.stage_widget.stage_model.update_speed(value=1)
+
+        else:
+            logger.info("Can't set stage speed because no AIND stage connected")
+
     def _LoadUI(self):
         """
         Determine which user interface to use
@@ -844,7 +924,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:
@@ -2767,6 +2847,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))
@@ -3584,7 +3665,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:
@@ -4217,7 +4298,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
@@ -5024,6 +5109,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
@@ -5080,6 +5170,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()
@@ -5135,7 +5231,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
@@ -5187,10 +5287,12 @@ def _Start(self):
                 self.data_lock,
             )
             worker1.signals.finished.connect(self._thread_complete)
+
             workerLick = Worker(
                 GeneratedTrials._get_irregular_timestamp, self.Channel2
             )
             workerLick.signals.finished.connect(self._thread_complete2)
+
             workerPlot = Worker(
                 PlotM._Update,
                 GeneratedTrials=GeneratedTrials,
@@ -5277,7 +5379,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
@@ -2848,7 +2849,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:
@@ -3079,70 +3080,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":
-                self.B_AutoRightWaterStartTime = np.append(
-                    self.B_AutoRightWaterStartTime, Rec[1][1][0]
-                )
+                with data_lock:
+                    self.B_AutoRightWaterStartTime = np.append(
+                        self.B_AutoRightWaterStartTime, Rec[1][1][0]
+                    )
 
     def _DeletePreviousLicks(self, Channel2):
         """Delete licks from the previous session"""