Skip to content

Feat fast retract #1543

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 49 commits into
base: curriculum-dev-branch
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
49 commits
Select commit Hold shift + click to select a range
fb2a5a8
outlines fast retract for stage
micahwoodard May 15, 2025
405b033
initializes timers
micahwoodard May 15, 2025
fa1bea9
fixes qtimer
micahwoodard May 16, 2025
ffc2f8b
debugs
micahwoodard May 16, 2025
ec255a3
connects signals'
micahwoodard May 16, 2025
13df210
debugs
micahwoodard May 16, 2025
6e08783
removes double signal
micahwoodard May 16, 2025
9899e0a
fixes timer bug
micahwoodard May 16, 2025
00f501c
debugs
micahwoodard May 16, 2025
0ec73ad
fixes stage widget
micahwoodard May 16, 2025
6e15161
degugs
micahwoodard May 19, 2025
b1ea235
debugs
micahwoodard May 19, 2025
5837066
inits timer with dummy timeout
micahwoodard May 19, 2025
53e0053
debugs
micahwoodard May 19, 2025
70dcf1c
debugs
micahwoodard May 19, 2025
b1385a4
debugs
micahwoodard May 19, 2025
d7a711a
debugs
micahwoodard May 19, 2025
7f47da6
debugs
micahwoodard May 19, 2025
0ebca0f
debugs
micahwoodard May 19, 2025
92ede10
debugs
micahwoodard May 19, 2025
252b818
changes stage widget release
micahwoodard May 21, 2025
4d540a5
disconnects signals
micahwoodard May 22, 2025
316869e
reconnects if retraction off
micahwoodard May 22, 2025
2466011
sets stage to normal speed after unretract
micahwoodard May 22, 2025
3ed7d75
comments out setting stage to normal speed
micahwoodard May 22, 2025
985a4e6
disconnects finished signal
micahwoodard May 22, 2025
ef31ffe
handles disconnects better
micahwoodard May 22, 2025
b2058fb
formats gui better
micahwoodard May 22, 2025
748d449
moves lick emit
micahwoodard May 22, 2025
eb51cd6
debugging
micahwoodard May 23, 2025
946ac29
debugging
micahwoodard May 23, 2025
9bc2f48
puts irregular timestamp in thread
micahwoodard May 23, 2025
43acfd2
removes simulation
micahwoodard May 23, 2025
a813165
create lick bonsai signal
micahwoodard May 27, 2025
5a614bb
catch lick at RigClient
micahwoodard May 28, 2025
f2fbcc8
reverting layout file
micahwoodard May 28, 2025
d31f473
removes extra folder
micahwoodard May 28, 2025
7522660
adds connection upon channel2 initialization
micahwoodard May 30, 2025
d1905ec
updates bonsai workflow
micahwoodard May 30, 2025
7be1775
try except to disconnect signals
micahwoodard May 30, 2025
e38f90c
adds conditional to bonsai workflow
micahwoodard May 30, 2025
e561638
checks if stage is at origin
micahwoodard May 30, 2025
999c12e
adds checks for connecting and disconnect mouse licked signals
micahwoodard May 30, 2025
d21b189
check y pos of lickspout
micahwoodard May 30, 2025
f1f8ac2
moves connection in start function
micahwoodard May 30, 2025
dea892f
adds logic for lickety split lickspout
micahwoodard May 30, 2025
568106c
reverts layout file
micahwoodard May 30, 2025
545ee42
Merge branch 'curriculum-dev-branch' into feat-fast-retract
micahwoodard Jun 10, 2025
75b824d
removes super init
micahwoodard Jun 10, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
127 changes: 114 additions & 13 deletions src/foraging_gui/Foraging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these were duplicated imports so that's why they were removed

from foraging_gui.stage import Stage
from foraging_gui.Visualization import (
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -596,6 +593,7 @@ def _load_stage(self) -> None:
else "widget_2"
)
self._insert_stage_widget(widget_to_replace)

else:
self._GetPositions()

Expand All @@ -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))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shouldn't this be lick_spout_retract, not lick_spout_licked ?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh nvm, I see that function wants lick_spout_licked, a little confusing to me that you don't just pass that function which lick spout to un-retract, but this works too

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah that does seems a little unintuitive. That is easy to rework

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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down
113 changes: 66 additions & 47 deletions src/foraging_gui/MyFunctions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)"""
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

addd data_lock here since these variables are saved during the experiment


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"""
Expand Down
Loading