mxbiflow 0.1.1__py3-none-any.whl

mxbiflow/tasks/GNGSiD/stages/size_reduction_stage/size_reduction_stage.py

@@ -0,0 +1,149 @@
+from random import choice
+from typing import TYPE_CHECKING, Final
+
+from mxbi.data_logger import DataLogger, DataLoggerType
+from mxbi.models.animal import ScheduleCondition
+from mxbi.tasks.GNGSiD.models import PersistentData, Result
+from mxbi.tasks.GNGSiD.stages.size_reduction_stage.size_reduction_models import (
+    SizeReductionStageConfig,
+    config,
+)
+from mxbi.tasks.GNGSiD.tasks.touch.touch_models import TrialConfig
+from mxbi.tasks.GNGSiD.tasks.touch.touch_scene import GNGSiDTouchScene
+from mxbi.utils.logger import logger
+
+if TYPE_CHECKING:
+    from mxbi.models.animal import AnimalState
+    from mxbi.models.session import SessionState
+    from mxbi.models.task import Feedback
+    from mxbi.theater import Theater
+
+
+_presistent_data: dict[str, PersistentData] = {}
+
+
+class SizeReductionStage:
+    STAGE_NAME: Final[str] = "GNGSiD_SIZE_REDUCTION_STAGE"
+
+    def __init__(
+        self,
+        theater: "Theater",
+        session_state: "SessionState",
+        animal_state: "AnimalState",
+    ) -> None:
+        self._theater = theater
+        self._session_state = session_state
+        self._animal_state = animal_state
+
+        self._stage_config = self._load_stage_config(animal_state.name)
+
+        _fixed_config = self._stage_config.params
+        _levels_config = self._stage_config.levels_table[animal_state.level]
+
+        master_amp, digital_amp = self._prepare_stimulus_intensity(
+            self._animal_state.name, _fixed_config.stimulus_freq
+        )
+
+        _config = TrialConfig(
+            level=_levels_config.level,
+            stimulation_size=_levels_config.stimulation_size,
+            stimulus_duration=_fixed_config.stimulus_duration,
+            time_out=_fixed_config.time_out,
+            inter_trial_interval=_fixed_config.inter_trial_interval,
+            reward_duration=_fixed_config.reward_duration,
+            reward_delay=_levels_config.reward_delay,
+            stimulus_freq=_fixed_config.stimulus_freq,
+            stimulus_freq_duration=_fixed_config.stimulus_freq_duration,
+            stimulus_freq_master_amp=master_amp,
+            stimulus_freq_digital_amp=digital_amp,
+            stimulus_interval=_fixed_config.stimulus_interval,
+        )
+
+        self._data_logger = DataLogger(
+            self._session_state,
+            self._animal_state.name,
+            self.STAGE_NAME,
+            DataLoggerType.JSONL,
+        )
+
+        self._presistent_data = _presistent_data.get(self._animal_state.name)
+
+        if self._presistent_data is None:
+            self._presistent_data = PersistentData(
+                rewards=0,
+                correct=0,
+                incorrect=0,
+                timeout=0,
+            )
+            _presistent_data[self._animal_state.name] = self._presistent_data
+
+        self._task = GNGSiDTouchScene(
+            theater,
+            session_state.session_config,
+            animal_state,
+            session_state.session_config.screen_type,
+            _config,
+            self._presistent_data,
+        )
+
+    def start(self) -> "Feedback":
+        trial_data = self._task.start()
+        self._data_logger.save(trial_data.model_dump())
+
+        feedback = self._handle_result(trial_data.result)
+        logger.debug(
+            f"{self.STAGE_NAME}: "
+            f"session_id={self._session_state.session_id}, "
+            f"animal_name={self._animal_state.name}, "
+            f"animal_level={self._animal_state.level}, "
+            f"state_name={self.STAGE_NAME}, "
+            f"result={trial_data}, "
+            f"feedback={feedback}"
+        )
+
+        return feedback
+
+    def _load_stage_config(self, monkey: str) -> SizeReductionStageConfig:
+        stage_config = config.root.get(monkey) or config.root.get("default")
+        if stage_config is None:
+            raise ValueError("No default stage config found")
+        return stage_config
+
+    def _handle_result(self, result: "Result") -> "Feedback":
+        feedback = False
+        match result:
+            case Result.CORRECT:
+                _presistent_data[self._animal_state.name].correct += 1
+                feedback = True
+            case Result.INCORRECT:
+                _presistent_data[self._animal_state.name].incorrect += 1
+                feedback = False
+            case Result.TIMEOUT:
+                _presistent_data[self._animal_state.name].timeout += 1
+                feedback = False
+            case Result.CANCEL:
+                feedback = False
+
+        return feedback
+
+    def quit(self) -> None:
+        self._task.cancle()
+
+    def on_idle(self) -> None:
+        self._task.cancle()
+
+    def on_return(self) -> None:
+        self._task.cancle()
+
+    @property
+    def condition(self) -> "ScheduleCondition | None":
+        return self._stage_config.condition
+
+    def _prepare_stimulus_intensity(self, monkey: str, frequency: int):
+        bt = choice(([[10, 30], [50, 70]])) if monkey == "wolfgang" else []
+        at = [80, 80, 80] if monkey == "wolfgang" else [80, 80, 80]
+        intensity_options = at * 10 + bt
+
+        stimulus_intensity = choice(intensity_options)
+
+        return self._theater.acontroller.get_amp_value(frequency, stimulus_intensity)

mxbiflow/tasks/GNGSiD/tasks/artifacts.py

@@ -0,0 +1,13 @@
+from tkinter import Canvas
+
+
+class BaseViews:
+    def _create_background(self, master, width, height) -> None:
+        self._backgroud = Canvas(
+            master,
+            bg="black",
+            width=width,
+            height=height,
+            highlightthickness=0,
+        )
+        self._backgroud.place(relx=0.5, rely=0.5, anchor="center")

mxbiflow/tasks/GNGSiD/tasks/detect/models.py

@@ -0,0 +1,21 @@
+from mxbi.tasks.GNGSiD.models import BaseDataToShow, BaseTrialConfig, BaseTrialData
+
+
+class TrialConfig(BaseTrialConfig):
+    go: bool
+    visual_stimulus_delay: int
+
+    stimulus_freq: int
+    stimulus_freq_duration: int
+    stimulus_freq_master_amp: int
+    stimulus_freq_digital_amp: int
+
+    stimulus_interval: int
+
+
+class TrialData(BaseTrialData):
+    trial_config: TrialConfig
+
+
+class DataToShow(BaseDataToShow):
+    stimulus: bool

mxbiflow/tasks/GNGSiD/tasks/detect/scene.py

@@ -0,0 +1,271 @@
+from datetime import datetime
+from math import ceil
+from tkinter import CENTER, Canvas, Event
+from typing import TYPE_CHECKING, Final
+
+from mxbi.tasks.GNGSiD.models import Result, TouchEvent
+from mxbi.tasks.GNGSiD.tasks.detect.models import DataToShow, TrialConfig, TrialData
+from mxbi.tasks.GNGSiD.tasks.utils.targets import DetectTarget
+from mxbi.utils.aplayer import ToneConfig
+from mxbi.utils.tkinter.components.canvas_with_border import CanvasWithInnerBorder
+from mxbi.utils.tkinter.components.showdata_widget import ShowDataWidget
+
+if TYPE_CHECKING:
+    from concurrent.futures import Future
+
+    from mxbi.models.animal import AnimalState
+    from mxbi.models.session import ScreenConfig, SessionConfig
+    from mxbi.tasks.GNGSiD.models import PersistentData
+    from mxbi.theater import Theater
+    from numpy import int16
+    from numpy.typing import NDArray
+
+
+class GNGSiDDetectScene:
+    def __init__(
+        self,
+        theater: "Theater",
+        session_config: "SessionConfig",
+        animal_state: "AnimalState",
+        screen_type: "ScreenConfig",
+        trial_config: "TrialConfig",
+        persistent_data: "PersistentData",
+    ) -> None:
+        self._theater: "Final[Theater]" = theater
+        self._animal_state: "Final[AnimalState]" = animal_state
+        self._screen_type: "Final[ScreenConfig]" = screen_type
+        self._trial_config: "Final[TrialConfig]" = trial_config
+        self._persistent_data: Final["PersistentData"] = persistent_data
+
+        self._tone: Final[NDArray[int16]] = self._prepare_stimulus()
+        self._standard_reward_stimulus = self._theater.new_standard_reward_stimulus(
+            self._trial_config.stimulus_duration
+        )
+
+        self._set_stimulus_intensity()
+
+        self._on_trial_start()
+
+    # region public api
+    def start(self) -> TrialData:
+        self._theater.mainloop()
+        return self._data
+
+    def cancle(self) -> None:
+        self._data.result = Result.CANCEL
+        self._theater.aplayer.stop()
+        self._on_trial_end()
+
+    # endregion
+
+    # region Lifecycle
+    def _on_trial_start(self) -> None:
+        self._create_view()
+        self._init_data()
+        self._bind_first_stage()
+
+    def _on_inter_trial(self) -> None:
+        self._background.after(
+            self._trial_config.inter_trial_interval, self._on_trial_end
+        )
+
+    def _on_trial_end(self) -> None:
+        self._background.destroy()
+        self._theater.root.quit()
+
+    # endregion
+
+    # region Views
+    def _create_view(self) -> None:
+        self._create_background()
+        self._create_show_data_view()
+        self._create_target()
+
+    def _create_background(self) -> None:
+        self._background = CanvasWithInnerBorder(
+            master=self._theater.root,
+            bg="black",
+            width=self._screen_type.width,
+            height=self._screen_type.height,
+            border_width=40,
+        )
+        self._background.place(relx=0.5, rely=0.5, anchor="center")
+
+    def _create_show_data_view(self) -> None:
+        self._show_data_widget = ShowDataWidget(self._background)
+        self._show_data_widget.place(relx=0, rely=1, anchor="sw")
+        data = DataToShow(
+            name=self._animal_state.name,
+            id=self._animal_state.trial_id,
+            level_id=self._animal_state.current_level_trial_id,
+            level=self._trial_config.level,
+            rewards=self._persistent_data.rewards,
+            correct=self._animal_state.correct_trial,
+            incorrect=self._persistent_data.incorrect,
+            timeout=self._persistent_data.timeout,
+            stimulus=self._trial_config.go,
+        )
+        self._show_data_widget.show_data(data.model_dump())
+
+    def _create_target(self):
+        xshift = 240
+        xcenter = self._screen_type.width * 0.5 + xshift
+        ycenter = self._screen_type.height * 0.5
+
+        self._trigger_canvas = DetectTarget(
+            self._background, self._trial_config.stimulation_size
+        )
+        self._trigger_canvas.place(x=xcenter, y=ycenter, anchor="center")
+
+    def _create_wrong_view(self) -> None:
+        self._trigger_canvas = Canvas(
+            self._background,
+            bg="grey",
+            width=self._screen_type.width,
+            height=self._screen_type.height,
+        )
+        self._trigger_canvas.place(relx=0.5, rely=0.5, anchor=CENTER)
+
+    # endregion
+
+    # region event binding
+    def _bind_first_stage(self) -> None:
+        self._background.focus_set()
+        self._background.bind("<r>", lambda e: self._give_standard_stimulus())
+        self._trigger_canvas.bind("<ButtonPress>", self._on_first_touched)
+        self._trigger_canvas.after(self._trial_config.time_out, self._on_timeout)
+
+    def _bind_second_stage(self) -> None:
+        self._trigger_canvas.bind("<ButtonPress>", self._on_second_touched)
+
+        # TODO: Confirm the waiting time
+        if not self._trial_config.go:
+            self._trigger_canvas.after(
+                self._trial_config.stimulus_duration, self._on_incorrect
+            )
+
+    # endregion
+
+    # region event handlers
+    def _on_first_touched(self, event: Event) -> None:
+        self._trigger_canvas.destroy()
+        self._record_touch(event)
+
+        if self._trial_config.go:
+            future = self._give_stimulus(self._tone)
+            future.add_done_callback(self._on_stimulus_complete)
+
+        self._background.after(
+            self._trial_config.visual_stimulus_delay,
+            lambda: (self._create_target(), self._bind_second_stage()),
+        )
+
+    def _on_second_touched(self, event: Event) -> None:
+        self._trigger_canvas.destroy()
+        self._record_touch(event)
+
+        if self._trial_config.go:
+            self._on_incorrect()
+        else:
+            future = self._give_stimulus(self._tone)
+            future.add_done_callback(self._on_stimulus_complete)
+
+        self._background.after(2000, self._create_target)
+
+    def _record_touch(self, event: Event) -> None:
+        self._data.touch_events.append(
+            TouchEvent(time=datetime.now().timestamp(), x=event.x, y=event.y)
+        )
+
+    # endregion
+
+    # region result handlers
+    def _on_correct(self) -> None:
+        self._trigger_canvas.destroy()
+
+        self._give_reward()
+        self._data.result = Result.CORRECT
+        self._data.correct_rate = (self._animal_state.correct_trial + 1) / (
+            self._animal_state.current_level_trial_id + 1
+        )
+        self._on_inter_trial()
+
+    def _on_incorrect(self) -> None:
+        self._theater._aplayer.stop()
+        self._trigger_canvas.destroy()
+
+        self._data.result = Result.INCORRECT
+        self._data.correct_rate = self._animal_state.correct_trial / (
+            self._animal_state.current_level_trial_id + 1
+        )
+
+        self._create_wrong_view()
+        self._on_inter_trial()
+
+    def _on_timeout(self) -> None:
+        self._data.result = Result.TIMEOUT
+        self._data.correct_rate = self._animal_state.correct_trial / (
+            self._animal_state.current_level_trial_id + 1
+        )
+
+        self._create_wrong_view()
+        self._on_inter_trial()
+
+    # endregion
+
+    # region stimulus and reward
+    def _prepare_stimulus(self) -> "NDArray[int16]":
+        cycle = (
+            self._trial_config.stimulus_freq_duration
+            + self._trial_config.stimulus_interval
+        )
+        repeat = ceil(self._trial_config.stimulus_duration / cycle)
+        repeat = max(repeat, 1)
+
+        freq_1 = ToneConfig(
+            frequency=self._trial_config.stimulus_freq,
+            duration=self._trial_config.stimulus_freq_duration,
+        )
+        freq_2 = ToneConfig(frequency=0, duration=self._trial_config.stimulus_interval)
+
+        return self._theater.aplayer.generate_stimulus([freq_1, freq_2], repeat)
+
+    def _give_stimulus(self, tone: "NDArray[int16]") -> "Future[bool]":
+        return self._theater.aplayer.play_stimulus(tone)
+
+    def _on_stimulus_complete(self, future: "Future[bool]") -> None:
+        if future.result():
+            self._background.after(self._trial_config.reward_delay, self._on_correct)
+
+    def _give_reward(self, _=None) -> None:
+        self._persistent_data.rewards += 1
+        self._theater.reward.give_reward(self._trial_config.reward_duration)
+
+    def _set_stimulus_intensity(self) -> None:
+        self._theater.acontroller.set_master_volume(
+            self._trial_config.stimulus_freq_master_amp
+        )
+        self._theater.acontroller.set_digital_volume(
+            self._trial_config.stimulus_freq_digital_amp
+        )
+
+    def _give_standard_stimulus(self) -> None:
+        self._standard_reward_stimulus.play(self._trial_config.reward_duration)
+
+    # endregion
+
+    # region data
+    def _init_data(self):
+        self._data = TrialData(
+            animal=self._animal_state.name,
+            trial_id=self._animal_state.trial_id,
+            current_level_trial_id=self._animal_state.current_level_trial_id,
+            trial_config=self._trial_config,
+            trial_start_time=datetime.now().timestamp(),
+            trial_end_time=0,
+            result=Result.TIMEOUT,
+            correct_rate=0,
+            touch_events=[],
+        )
+
+    # endregion

mxbiflow/tasks/GNGSiD/tasks/discriminate/discriminate_models.py

@@ -0,0 +1,31 @@
+from mxbi.tasks.GNGSiD.models import BaseDataToShow, BaseTrialConfig, BaseTrialData
+
+
+class TrialConfig(BaseTrialConfig):
+    is_stimulus_trial: bool
+    visual_stimulus_delay: int
+
+    medium_reward_duration: int
+    medium_reward_threshold: int
+    low_reward_duration: int
+
+    attention_duration: int
+    extra_response_time: int
+
+    stimulus_freq_high: int
+    stimulus_freq_high_duration: int
+    stimulus_freq_high_master_amp: int
+    stimulus_freq_high_digital_amp: int
+    stimulus_freq_low: int
+    stimulus_freq_low_duration: int
+    stimulus_freq_low_master_amp: int
+    stimulus_freq_low_digital_amp: int
+    stimulus_interval: int
+
+
+class TrialData(BaseTrialData):
+    trial_config: TrialConfig
+
+
+class DataToShow(BaseDataToShow):
+    stimulus: bool