mxbiflow 0.1.1__py3-none-any.whl

mxbiflow/tasks/default/initial_habituation_training/tasks/stay_to_reward/stay_to_reward.py

@@ -0,0 +1,244 @@
+from datetime import datetime
+from math import ceil
+from random import uniform
+from tkinter import Frame
+from typing import TYPE_CHECKING, Final
+
+from mxbi.tasks.default.initial_habituation_training.tasks.stay_to_reward.stay_to_reward_models import (
+    DataToShow,
+    Result,
+    TrialData,
+)
+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
+    from mxbi.tasks.default.initial_habituation_training.stages.models import (
+        StageContext,
+    )
+    from mxbi.tasks.default.initial_habituation_training.tasks.stay_to_reward.stay_to_reward_models import (
+        TrialConfig,
+    )
+    from mxbi.theater import Theater
+    from numpy import int16
+    from numpy.typing import NDArray
+
+
+class DefaultStayToRewardScene:
+    def __init__(
+        self,
+        theater: "Theater",
+        animal_state: "AnimalState",
+        screen_type: "ScreenConfig",
+        trial_config: "TrialConfig",
+        context: "StageContext",
+        background: "CanvasWithInnerBorder",
+    ):
+        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._context: "Final[StageContext]" = context
+        self._background = background
+
+        self._play_future: Future[bool] | None = None
+
+        self._tone = 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._cleanup()
+
+    # endregion
+
+    # region lifecycle
+    def _on_trial_start(self) -> None:
+        self._create_view()
+        self._init_data()
+        self._bind_events()
+        self._start_tracking_data()
+
+        if (
+            self._trial_config.entry_reward
+            and self._animal_state.current_animal_session_trial_id == 1
+        ):
+            self._direct_stimulus()
+        else:
+            self._start_timing()
+            self._stimulus_loop()
+
+    def _on_trial_end(self) -> None:
+        if self._play_future is not None:
+            self._play_future.add_done_callback(lambda _: self._cleanup())
+        else:
+            self._cleanup()
+
+    def _cleanup(self) -> None:
+        self._data.trial_end_time = datetime.now().timestamp()
+        self._data.stay_duration = (
+            self._data.trial_end_time - self._data.trial_start_time
+        )
+
+        self._theater.aplayer.stop()
+        self._trigger.destroy()
+        self._show_data_widget.destroy()
+        self._theater.root.quit()
+
+    # endregion
+
+    # region views
+
+    def _create_view(self) -> None:
+        self._create_background()
+        self._create_trigger()
+        self._create_show_data_widget()
+
+    def _create_background(self) -> None:
+        self._background.place(relx=0.5, rely=0.5, anchor="center")
+
+    def _create_trigger(self) -> None:
+        self._trigger = Frame(bg="red", width=0, height=0)
+        self._trigger.place(relx=0.5, rely=0.5, anchor="center")
+        self._trigger.lower()
+
+    def _create_show_data_widget(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,
+            ald=self._animal_state.current_animal_session_trial_id,
+            level=self._animal_state.level,
+            id=self._animal_state.trial_id,
+            lid=self._animal_state.current_level_trial_id,
+            t_dur=f"{self._context.duration} s",
+            dur="0 s",
+            rewards=0,
+        )
+        self._show_data_widget.show_data(_data.model_dump())
+
+    # endregion
+
+    # region event binding
+    def _bind_events(self) -> None:
+        self._trigger.focus_set()
+
+    def _start_timing(self) -> None:
+        target_ms = int(self._trial_config.target * 1000)
+        self._trigger.after(target_ms, self._on_correct)
+
+    def _stimulus_loop(self) -> None:
+        stimulus_interval = uniform(
+            self._trial_config.min_stimulus_interval,
+            self._trial_config.max_stimulus_interval,
+        )
+        stimulus_interval_ms = int(stimulus_interval * 1000)
+
+        self._trigger.after(stimulus_interval_ms, self._give_stimulus)
+
+    # endregion
+
+    # region data
+    def _start_tracking_data(self) -> None:
+        data = DataToShow(
+            name=self._animal_state.name,
+            ald=self._animal_state.current_animal_session_trial_id,
+            level=self._animal_state.level,
+            id=self._animal_state.trial_id,
+            lid=self._animal_state.current_level_trial_id,
+            t_dur=f"{self._context.duration} s",
+            dur=f"{int(self._data.stay_duration)} s",
+            rewards=self._context.rewards,
+        )
+
+        self._data.stay_duration += 1
+        self._context.duration += 1
+
+        self._show_data_widget.update_data(data.model_dump())
+        self._trigger.after(1000, self._start_tracking_data)
+
+    def _init_data(self) -> None:
+        self._data = TrialData(
+            level=self._trial_config.level,
+            level_trial_id=self._animal_state.current_level_trial_id,
+            animal_session_trial_id=self._animal_state.current_animal_session_trial_id,
+            animal=self._animal_state.name,
+            trial_id=self._animal_state.trial_id,
+            trial_start_time=datetime.now().timestamp(),
+            trial_end_time=0,
+            stay_duration=0,
+            result=Result.CORRECT,
+        )
+
+    # endregion
+
+    # region stimulus and reward
+    def _prepare_stimulus(self) -> "NDArray[int16]":
+        unit_duration = (
+            self._trial_config.stimulus_freq_duration
+            + self._trial_config.stimulus_freq_interval
+        )
+        times = ceil(self._trial_config.stimulus_duration / unit_duration)
+        times = max(times, 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_freq_interval,
+        )
+
+        return self._theater.aplayer.generate_stimulus([freq_1, freq_2], times)
+
+    def _set_stimulus_intensity(self) -> None:
+        self._theater.acontroller.set_master_volume(self._trial_config.stimulus_density)
+
+    def _give_stimulus(self) -> None:
+        self._play_future = self._theater.aplayer.play_stimulus(self._tone)
+        self._play_future.add_done_callback(self._on_stimulus_complete)
+
+    def _direct_stimulus(self) -> None:
+        self._play_future = self._theater.aplayer.play_stimulus(self._tone)
+        self._play_future.add_done_callback(self._on_direct_stimulus_complete)
+
+    def _on_direct_stimulus_complete(self, future: "Future[bool]") -> None:
+        if future.result():
+            self._trigger.after(0, self._give_reward)
+
+            self._trigger.after(
+                self._trial_config.reward_duration,
+                lambda: (self._start_timing(), self._stimulus_loop()),
+            )
+
+    def _on_stimulus_complete(self, future: "Future[bool]") -> None:
+        if future.result():
+            self._trigger.after(0, self._give_reward)
+
+            self._trigger.after(self._trial_config.reward_duration, self._stimulus_loop)
+
+    def _on_correct(self) -> None:
+        self._data.result = Result.CORRECT
+        self._on_trial_end()
+
+    def _give_reward(self) -> None:
+        self._context.rewards += 1
+        self._theater.reward.give_reward(self._trial_config.reward_duration)
+
+    # endregion

mxbiflow/tasks/default/initial_habituation_training/tasks/stay_to_reward/stay_to_reward_models.py

@@ -0,0 +1,50 @@
+from enum import StrEnum, auto
+
+from pydantic import BaseModel
+
+
+class Result(StrEnum):
+    CORRECT = auto()
+    INCORRECT = auto()
+
+
+class TrialConfig(BaseModel):
+    level: int
+    entry_reward: bool
+
+    min_stimulus_interval: float  # seconds
+    max_stimulus_interval: float  # seconds
+    target: float  # seconds
+
+    reward_duration: int = 1000  # milliseconds
+    stimulus_duration: int = 1000  # milliseconds
+
+    stimulus_density: int  # volume 0-100
+
+    stimulus_freq: int = 2000
+    stimulus_freq_duration: int = 100
+    stimulus_freq_interval: int = 100
+
+
+class TrialData(BaseModel):
+    level: int
+    animal: str
+    trial_id: int
+    level_trial_id: int
+    animal_session_trial_id: int
+    trial_start_time: float
+    trial_end_time: float
+    result: Result
+
+    stay_duration: float
+
+
+class DataToShow(BaseModel):
+    level: int
+    name: str
+    id: int
+    lid: int
+    ald: int
+    t_dur: str
+    dur: str
+    rewards: int

mxbiflow/tasks/task_protocol.py

@@ -0,0 +1,26 @@
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from mxbi.models.animal import AnimalState, ScheduleCondition
+    from mxbi.models.session import SessionState
+    from mxbi.models.task import Feedback
+    from mxbi.theater import Theater
+
+
+class Task(Protocol):
+    def __init__(
+        self,
+        theater: "Theater",
+        session_state: "SessionState",
+        animal_state: "AnimalState",
+    ) -> None: ...
+    def start(self) -> "Feedback": ...
+
+    def quit(self) -> None: ...
+
+    def on_idle(self) -> None: ...
+
+    def on_return(self) -> None: ...
+
+    @property
+    def condition(self) -> "ScheduleCondition | None": ...

mxbiflow/tasks/task_table.py

@@ -0,0 +1,29 @@
+from mxbi.models.task import TaskEnum
+from mxbi.tasks.cross_modal.stage import CrossModalTask
+from mxbi.tasks.default.error_task.error_scene import ErrorScene
+from mxbi.tasks.default.idle_task.idle_scene import IDLEScene
+from mxbi.tasks.default.initial_habituation_training.stages.initial_habituation_training_stage import (
+    InitialHabituationTrainingStage,
+)
+from mxbi.tasks.GNGSiD.stages.detect_stage.detect_stage import GNGSiDDetectStage
+from mxbi.tasks.GNGSiD.stages.discriminate_stage.discriminate_stage import (
+    GNGSiDDiscriminateStage,
+)
+from mxbi.tasks.GNGSiD.stages.size_reduction_stage.size_reduction_stage import (
+    SizeReductionStage,
+)
+from mxbi.tasks.task_protocol import Task
+from mxbi.tasks.two_alternative_choice.stages.size_reduction_stage.size_reduction_stage import (
+    TWOACSizeReductionStage,
+)
+
+task_table: dict[TaskEnum, type[Task]] = {
+    TaskEnum.IDEL: IDLEScene,
+    TaskEnum.ERROR: ErrorScene,
+    TaskEnum.HABITUATION: InitialHabituationTrainingStage,
+    TaskEnum.GNGSiD_SIZE_REDUCTION_STAGE: SizeReductionStage,
+    TaskEnum.GNGSiD_DETECT_STAGE: GNGSiDDetectStage,
+    TaskEnum.GNGSiD_DISCRIMINATE_STAGE: GNGSiDDiscriminateStage,
+    TaskEnum.TWOAC_SIZE_REDUCTION_STAGE: TWOACSizeReductionStage,
+    TaskEnum.CROSS_MODAL: CrossModalTask,
+}

mxbiflow/tasks/two_alternative_choice/assets/starter.py

@@ -0,0 +1,27 @@
+from tkinter import Canvas
+
+from mxbi.utils.tkinter.create_circle import create_circle
+
+
+class Starter(Canvas):
+    def __init__(self, master, size) -> None:
+        super().__init__(
+            master, width=size, height=size, bg="blue", highlightthickness=0
+        )
+
+        create_circle(
+            size * 0.5,
+            size * 0.5,
+            (size - 1) / 2,
+            self,
+            "white",
+        )
+
+
+if __name__ == "__main__":
+    from tkinter import Tk
+
+    root = Tk()
+    root.geometry("300x300")
+    Starter(root, 300).pack()
+    root.mainloop()

mxbiflow/tasks/two_alternative_choice/models.py

@@ -0,0 +1,68 @@
+from enum import StrEnum, auto
+from typing import TypeAlias
+
+from pydantic import BaseModel
+
+LevelID: TypeAlias = int
+MonkeyName: TypeAlias = str
+
+
+class Result(StrEnum):
+    CORRECT = auto()
+    INCORRECT = auto()
+    TIMEOUT = auto()
+    CANCEL = auto()
+
+
+class TouchEvent(BaseModel):
+    time: float
+    x: int
+    y: int
+
+
+class BaseTrialConfig(BaseModel):
+    # basic config
+    level: int = 0
+
+    # visual stimulation config
+    stimulation_size: int = 0
+
+    # audio stimulation config
+    stimulus_duration: int = 0
+
+    # trial lifecycle
+    time_out: int = 0
+    inter_trial_interval: int = 0
+
+    # reward config
+    reward_duration: int = 0
+    reward_delay: int = 0
+
+
+class BaseTrialData(BaseModel):
+    animal: str
+    trial_id: int
+    current_level_trial_id: int
+    trial_start_time: float
+    trial_end_time: float
+    result: Result
+    correct_rate: float
+    touch_events: list[TouchEvent]
+
+
+class BaseDataToShow(BaseModel):
+    name: str
+    id: int
+    level_id: int
+    level: int
+    rewards: int
+    correct: int
+    incorrect: int
+    timeout: int
+
+
+class PersistentData(BaseModel):
+    rewards: int
+    correct: int
+    incorrect: int
+    timeout: int

mxbiflow/tasks/two_alternative_choice/stages/size_reduction_stage/config.json

@@ -0,0 +1,118 @@
+{
+    "default": {
+        "condition": {
+            "config": {
+                "evaluation_interval": 20,
+                "difficulty_increase_threshold": 0.8,
+                "difficulty_decrease_threshold": 0.45,
+                "next_task": "gngsid_detect_stage"
+            }
+        },
+        "trial_config": {
+            "stimulus_duration": 1000,
+            "time_out": 10000,
+            "inter_trial_interval": 1000,
+            "reward_duration": 1000,
+            "reward_delay": 0,
+            "stimulus_freq": 2000,
+            "stimulus_freq_duration": 100,
+            "stimulus_interval": 100
+        },
+        "levels_table": {
+            "1": {
+                "level": 1,
+                "stimulation_size": 300
+            },
+            "2": {
+                "level": 2,
+                "stimulation_size": 295
+            },
+            "3": {
+                "level": 3,
+                "stimulation_size": 390
+            },
+            "4": {
+                "level": 4,
+                "stimulation_size": 285
+            },
+            "5": {
+                "level": 5,
+                "stimulation_size": 280
+            },
+            "6": {
+                "level": 6,
+                "stimulation_size": 275
+            },
+            "7": {
+                "level": 7,
+                "stimulation_size": 270
+            },
+            "8": {
+                "level": 8,
+                "stimulation_size": 260
+            },
+            "9": {
+                "level": 9,
+                "stimulation_size": 255
+            },
+            "10": {
+                "level": 10,
+                "stimulation_size": 240
+            },
+            "11": {
+                "level": 11,
+                "stimulation_size": 230
+            },
+            "12": {
+                "level": 12,
+                "stimulation_size": 225
+            },
+            "13": {
+                "level": 13,
+                "stimulation_size": 220
+            },
+            "14": {
+                "level": 14,
+                "stimulation_size": 210
+            },
+            "15": {
+                "level": 15,
+                "stimulation_size": 200
+            }
+        }
+    },
+    "debug": {
+        "condition": {
+            "config": {
+                "evaluation_interval": 2,
+                "difficulty_increase_threshold": 0.8,
+                "difficulty_decrease_threshold": 0.45,
+                "next_task": "gngsid_detect_stage"
+            }
+        },
+        "trial_config": {
+            "stimulus_duration": 1000,
+            "time_out": 10000,
+            "inter_trial_interval": 1000,
+            "reward_duration": 1000,
+            "reward_delay": 0,
+            "stimulus_freq": 2000,
+            "stimulus_freq_duration": 100,
+            "stimulus_interval": 100
+        },
+        "levels_table": {
+            "1": {
+                "level": 1,
+                "stimulation_size": 300
+            },
+            "2": {
+                "level": 2,
+                "stimulation_size": 295
+            },
+            "3": {
+                "level": 3,
+                "stimulation_size": 390
+            }
+        }
+    }
+}

mxbiflow/tasks/two_alternative_choice/stages/size_reduction_stage/size_reduction_models.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+
+from mxbi.config import Configure
+from mxbi.models.animal import ScheduleCondition
+from mxbi.tasks.two_alternative_choice.models import LevelID, MonkeyName
+from mxbi.tasks.two_alternative_choice.tasks.touch.touch_models import TrialConfig
+from pydantic import BaseModel, ConfigDict, RootModel
+
+CONFIG_PATH = Path(__file__).parent / "config.json"
+
+
+class SizeReductionStageLeveledParameters(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    level: int
+    stimulation_size: int
+
+
+class SizeReductionStageConfig(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    condition: ScheduleCondition
+    trial_config: TrialConfig
+
+    levels_table: dict[LevelID, SizeReductionStageLeveledParameters]
+
+
+class SizeReductionStageConfigs(RootModel):
+    model_config = ConfigDict(frozen=True)
+
+    root: dict[MonkeyName, SizeReductionStageConfig]
+
+
+def load_config() -> SizeReductionStageConfigs:
+    configs = Configure(CONFIG_PATH, SizeReductionStageConfigs).value
+    for config in configs.root.values():
+        config.condition.level_count = len(config.levels_table)
+    return configs
+
+
+config = load_config()