def schedule_task(module, task, task_type, **kwargs):
if module not in _tasks:
raise Exception(f"Unknown task module: {module}")
if task not in _tasks[module]:
raise Exception(f"Unknown task {task} in module {module}.")
if task_type == "datetime":
if type(kwargs["dt"]) is str:
dt = parser.parse(kwargs["dt"]).astimezone()
else:
dt = kwargs["dt"]
cancel = sched.on_datetime(lambda: run(module, task),
dt,
pool=_task_pool)
_log.info(f"Scheduled task {module}.{task} at {dt}")
elif task_type == "interval":
cancel = sched.repeat(
lambda: run(module, task),
kwargs["interval"],
kwargs["repeats"],
pool=_task_pool,
)
_log.info(
f"Scheduled task {module}.{task}. interval={kwargs['interval']} seconds, repeats={kwargs['repeats']}"
)
elif task_type == "timeofday":
if type(kwargs["dt"]) is str:
dt = parser.parse(kwargs["dt"]).astimezone()
else:
dt = kwargs["dt"]
time = dt.time()
cancel = sched.timeofday(lambda: run(module, task),
time,
kwargs["repeats"],
pool=_task_pool)
if kwargs["repeats"] is True:
_log.info(
f"Scheduled task {module}.{task} at {dt}. Repeats every day")
else:
_log.info(
f"Scheduled task {module}.{task} at {time}. Repeats for {kwargs['repeats']} day(s)"
)
else:
raise ValueError(f"Unknown task type: {task_type}")
_scheduled_tasks.append({
"task": f"{module}.{task}",
"task_type": task_type,
"params": kwargs,
"created": datetime.now(),
"cancel_fn": cancel,
"task_id": get_new_scheduled_task_id(),
})
def led_stimulus(self):
params = exp.get_merged_params()
if exp.state["arena", "signal_led"]:
arena.signal_led(False)
self.stim_cancel = schedule.repeat(
lambda: arena.signal_led(not exp.state["arena", "signal_led"]),
params["led_duration"],
2 * params.get("led_blinks", 1),
)
def led_blink(num_blinks, led_duration):
interval = led_duration / num_blinks / 2
def toggle_led():
arena.run_command("toggle", "Signal LED")
arena.run_command("set", "Signal LED", [0])
return schedule.repeat(toggle_led, interval, repeats=num_blinks * 2)
def run(self, params):
# resetting init values
self.in_trial = False
self.got_detection = False
self.time = 0.0
self.cancel_trials = None
self.cancel_logic_trial = None
self.prev_det = None
self.prev_trial_detection = False
self.stim_cancel = None
self.consq_end = False
self.data_dir = session_state["data_dir"]
self.cur_trial = params["$num_trials"]
self.consecutive = params["consecutive"]
# determining the experiment type
self.ex_type = ("consecutive" if self.consecutive else
("continuous" if params["continuous"] else "regular"))
# logging yolo detections
self.yolo_log = data_log.QueuedDataLogger(
columns=[
("time", "timestamptz not null"),
("x1", "double precision"),
("y1", "double precision"),
("x2", "double precision"),
("y2", "double precision"),
("confidence", "double precision"),
],
csv_path=session_state["data_dir"] / "head_bbox.csv",
table_name="bbox_position",
)
self.yolo_log.start()
# detecting Aruco squares within arena
self.detectAruco()
if params["record_all"]: # record start at init
video_system.start_record()
if not self.consecutive: # no need to schedule next trials if consecutive
self.cancel_trials = schedule.repeat(
self.period_call,
params.get("exp_interval", 100),
params.get("num_of_exp", 1) - 1,
) # schedule the next trials
# starting image observers
yolo = exp.image_observers["head_bbox"]
yolo.on_detection = self.on_yolo_detection
yolo.start_observing()
self.log.info("Start Observing")
monitor.change_color("white")
def end_logic_trial(self):
params = exp.get_merged_params()
self.stim_cancel() # canceling stimulus, if active.
if params["stimulus"] == "monitor":
monitor.chnage_color("black")
timestap = time.time()
# logging trial data
if self.in_trial and not self.got_detection:
self.log.info("Logic trial ended, failure")
exp.event_logger.log(
"learn_exp/logical_trial_ended",
{
"type": self.ex_type,
"success": False
},
)
elif self.in_trial and self.got_detection:
self.log.info("Logic trial ended, success")
exp.event_logger.log("learn_exp/logical_trial_ended", {
"type": self.ex_type,
"success": True
})
else:
self.log.info("Logic trial ended")
# continuous trial: schedule the next.
self.in_trial = False
self.got_detection = False
if params.get("continuous", False):
if params["record_exp"] and not params["record_all"]:
video_system.stop_record()
self.cancel_trials()
self.cancel_trials = schedule.repeat(self.period_call, interval,
self.cur_trial - 1)
elif self.consecutive:
if params["record_exp"] and not params["record_all"]:
schedule.once(lambda: video_system.stop_record(),
params.get("record_overhead", 0))
if self.cur_trial > 0:
exp.next_trial()
else:
if params["record_exp"] and not params["record_all"]:
schedule.once(lambda: video_system.stop_record(),
params.get("record_overhead", 0))
def run_block(self):
interval = exp.get_params()["interval"]
self.cancel_timer = schedule.repeat(
self.timer_fn, interval,
exp.get_params().get("$num_trials", True))
def init(logger, config):
"""
Initialize the arena module.
Connects to MQTT, inits sensor data logger, sends arena defaults, and subscribes for
sensor updates.
- arena_defaults: A dict with signal_led and day_lights keys with default values.
"""
global _log, _arena_log, _config, _arena_state, _interfaces_config, _trigger_interface
_log = logger
_config = config
_arena_state = state.get_cursor("arena")
_arena_state.set_self({
"values": {},
"timestamp":
None,
"displays":
dict([(d, False) for d in config.arena["displays"].keys()]),
})
for display in config.arena["displays"].keys():
switch_display(False, display)
try:
with open(config.arena_config_path, "r") as f:
arena_config = json.load(f)
interfaces_config = []
for interfaces in arena_config.values():
interfaces_config += interfaces
except json.JSONDecodeError:
_log.exception("Exception while parsing arena config:")
return
_interfaces_config = interfaces_config
for ifs in interfaces_config:
if ifs["type"] == "trigger":
_trigger_interface = ifs["name"]
break
if "data_log" in config.arena:
log_conf = config.arena["data_log"]
columns = [("time", "timestamptz not null")] + log_conf["columns"]
_arena_log = data_log.QueuedDataLogger(
table_name=log_conf["table_name"],
log_to_db=True,
columns=columns,
)
_arena_log.start()
while not mqtt.client.is_connected:
time.sleep(0.01)
topic = config.arena["receive_topic"]
mqtt.client.subscribe_callback(f"{topic}/all_values",
mqtt.mqtt_json_callback(_on_all_values))
mqtt.client.subscribe_callback(f"{topic}/value",
mqtt.mqtt_json_callback(_on_value))
mqtt.client.subscribe_callback(f"{topic}/info/#",
mqtt.mqtt_json_callback(_on_info))
mqtt.client.subscribe_callback(f"{topic}/error/#",
mqtt.mqtt_json_callback(_on_error))
poll()
schedule.repeat(poll, config.arena["poll_interval"], pool="arena")
def led_stimulus(self):
self.stim_cancel = schedule.repeat(
lambda: arena.run_command("toggle", "Signal LED"),
exp.get_params()["led_duration"],
2 * exp.get_params().get("led_blinks", 1),
)