mxbiflow 0.1.1__py3-none-any.whl

mxbiflow/tasks/default/error_task/error_scene.py

@@ -0,0 +1,53 @@
+from tkinter import Canvas
+from typing import TYPE_CHECKING
+
+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 ErrorScene:
+    def __init__(
+        self,
+        theater: "Theater",
+        session_state: "SessionState",
+        animal_state: "AnimalState",
+    ) -> None:
+        self._theater = theater
+        self._session_config = session_state
+        self._screen_type = self._session_config.session_config.screen_type
+
+        self._on_trial_start()
+
+    def _on_trial_start(self) -> None:
+        self._create_view()
+
+    def _create_view(self) -> None:
+        self._background = Canvas(
+            self._theater.root,
+            bg="red",
+            width=self._screen_type.width,
+            height=self._screen_type.height,
+            highlightthickness=0,
+        )
+        self._background.place(relx=0.5, rely=0.5, anchor="center")
+
+    def start(self) -> "Feedback":
+        self._theater.root.mainloop()
+
+        return True
+
+    def quit(self) -> None:
+        self._background.destroy()
+        self._theater.root.quit()
+
+    def on_idle(self) -> None:
+        self.quit()
+
+    def on_return(self) -> None: ...
+
+    @property
+    def condition(self) -> "ScheduleCondition | None":
+        return None

mxbiflow/tasks/default/idle_task/idle_scene.py

@@ -0,0 +1,85 @@
+from pathlib import Path
+from tkinter.ttk import Label
+from typing import TYPE_CHECKING
+
+from mxbi.utils.tkinter.components.canvas_with_border import CanvasWithInnerBorder
+from PIL import Image, ImageTk
+
+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
+
+ASSETS_PATH = Path(__file__).parent / "assets"
+
+
+class IDLEScene:
+    def __init__(
+        self,
+        theater: "Theater",
+        session_state: "SessionState",
+        animal_state: "AnimalState",
+    ) -> None:
+        self._theater = theater
+        self._session_config = session_state
+        self._screen_type = self._session_config.session_config.screen_type
+        self._standard_reward_stimulus = self._theater.new_standard_reward_stimulus(
+            1000
+        )
+
+        self._on_trial_start()
+
+    def start(self) -> "Feedback":
+        self._theater.root.mainloop()
+
+        return True
+
+    def _on_trial_start(self) -> None:
+        self._create_view()
+        self._bind_events()
+
+    def _on_trial_end(self) -> None:
+        self._theater.aplayer.stop()
+        self._background.destroy()
+        self._theater.root.quit()
+
+    def _create_view(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")
+
+        xshift = 240
+        xcenter = self._screen_type.width / 2 + xshift
+        ycenter = self._screen_type.height / 2
+
+        self._img = Image.open(ASSETS_PATH / "apple_v1.png")
+        self._img = self._img.resize((400, 400)).rotate(-90, expand=True)
+        self._img = ImageTk.PhotoImage(self._img)
+        self.label_apple = Label(self._background, image=self._img)
+        self.label_apple.place(x=xcenter, y=ycenter, anchor="center")
+
+    def _bind_events(self) -> None:
+        self._background.focus_set()
+        self._background.bind("<r>", lambda e: self._give_stimulus(1000))
+
+    def _give_stimulus(self, duration: int) -> None:
+        self._standard_reward_stimulus.play(1000)
+
+    def quit(self) -> None:
+        self._on_trial_end()
+
+    def on_idle(self) -> None:
+        self._on_trial_end()
+
+    def on_return(self) -> None:
+        self._on_trial_end()
+
+    @property
+    def condition(self) -> "ScheduleCondition | None":
+        return None

mxbiflow/tasks/default/initial_habituation_training/README.md

@@ -0,0 +1,188 @@
+## Automated Habituation Training (Strong Auditory Cue)
+
+**🎯 Overview**
+
+- Familiarize the animal with the MXBI apparatus and establish the fixed relationship “auditory stimulus → reward.”
+- Present a consistent routine to gradually extend stay time while keeping the reward experience stable.
+- The animal must remain in MXBI for a period to obtain a reward.
+- An auditory stimulus is presented before each reward to strengthen the association.
+
+**🔁 Procedure**
+The experimental flow is:
+
+```mermaid
+flowchart TD
+    mxbi -->|start| idle_state
+    idle_state -->|animal enter as trigger| habituation_state
+    habituation_state -->|animal exit| idle_state
+    habituation_state -->|animal stay some time| stimulus
+    stimulus -->|after stimulus| reward
+    reward -->|after reward| habituation_state
+```
+
+**🧭 Levels**
+This stage consists of a series of levels. The parameters that change as the level increases include:
+
+```json
+"0": {
+    "level": 0,
+    "evaluation_interval": 20,
+    "reward_interval": 5,
+    "reward_duration": 1000,
+    "stimulus_duration": 1000,
+    "stimulus_density": 5
+}
+```
+
+At level 0:
+
+- Stay in MXBI for 5 seconds to receive a reward.
+- Reward duration: 1 second; auditory stimulus duration: 1 second.
+- Stimulus density: 5.
+- Complete at least 20 trials to advance.
+
+As the level increases:
+
+- Required stay time increases; the auditory stimulus becomes longer and louder.
+- More trials are required to progress, but more rewards are also available.
+
+Outcome goal:
+
+- Single stay in MXBI exceeds 1 minute.
+- Strong understanding of “auditory stimulus → reward”.
+
+> **Note: We currently lack detection methods; because the animal need fully controls the experiment, traditional training approaches cannot be used**
+
+**📊 Current Level Table**
+
+| level | evaluation_interval | reward_interval_s | reward_duration_ms | stimulus_duration_ms | stimulus_density |
+| :---: | :-----------------: | :---------------: | :----------------: | :------------------: | :--------------: |
+|   0   |         20          |         2         |        1000        |         1000         |        5         |
+|   1   |         20          |         5         |        1200        |         1500         |        7         |
+|   2   |         40          |        10         |        1400        |         2000         |        9         |
+|   3   |         50          |        15         |        1600        |         2500         |        11        |
+|   4   |         50          |        20         |        1800        |         3000         |        13        |
+|   5   |         100         |        25         |        2000        |         3000         |        15        |
+|   6   |         100         |        30         |        2200        |         3000         |        17        |
+|   7   |         100         |        35         |        2400        |         3000         |        19        |
+|   8   |         100         |        40         |        2600        |         3000         |        21        |
+|   9   |         100         |        45         |        2800        |         3000         |        23        |
+|  10   |         100         |        50         |        3000        |         3000         |        25        |
+|  11   |         100         |        55         |        3000        |         3000         |        25        |
+|  12   |         100         |        60         |        3000        |         3000         |        25        |
+
+**🧪 Next Step**
+Currently, the time the animal must stay in MXBI to get a reward is controlled by `reward_interval`. As shown above, it is a fixed value. Consider changing it to a range and randomly choosing the actual `reward_interval` within that range.
+
+The complete configuration file is as follows:
+
+```json
+{
+  "default": {
+    "condition": {
+      "config": {
+        "difficulty_increase_threshold": 0.8,
+        "allow_decrease": false
+      }
+    },
+    "levels_table": {
+      "0": {
+        "level": 0,
+        "evaluation_interval": 20,
+        "reward_interval": 5,
+        "reward_duration": 1000,
+        "stimulus_duration": 1000,
+        "stimulus_density": 5
+      },
+      "1": {
+        "level": 1,
+        "evaluation_interval": 30,
+        "reward_interval": 10,
+        "reward_duration": 1200,
+        "stimulus_duration": 1500,
+        "stimulus_density": 7
+      },
+      "2 ": {
+        "level": 2,
+        "evaluation_interval": 40,
+        "reward_interval": 15,
+        "reward_duration": 1400,
+        "stimulus_duration": 2000,
+        "stimulus_density": 9
+      },
+      "3": {
+        "level": 3,
+        "evaluation_interval": 50,
+        "reward_interval": 20,
+        "reward_duration": 1600,
+        "stimulus_duration": 2500,
+        "stimulus_density": 11
+      },
+      "4": {
+        "level": 4,
+        "evaluation_interval": 50,
+        "reward_interval": 25,
+        "reward_duration": 1800,
+        "stimulus_duration": 3000,
+        "stimulus_density": 13
+      },
+      "5": {
+        "level": 5,
+        "evaluation_interval": 100,
+        "reward_interval": 30,
+        "reward_duration": 2000,
+        "stimulus_duration": 3000,
+        "stimulus_density": 15
+      },
+      "6": {
+        "level": 6,
+        "evaluation_interval": 100,
+        "reward_interval": 35,
+        "reward_duration": 2200,
+        "stimulus_duration": 3000,
+        "stimulus_density": 17
+      },
+      "7": {
+        "level": 7,
+        "evaluation_interval": 100,
+        "reward_interval": 40,
+        "reward_duration": 2400,
+        "stimulus_duration": 3000,
+        "stimulus_density": 19
+      },
+      "8": {
+        "level": 8,
+        "evaluation_interval": 100,
+        "reward_interval": 45,
+        "reward_duration": 2600,
+        "stimulus_duration": 3000,
+        "stimulus_density": 21
+      },
+      "9": {
+        "level": 9,
+        "evaluation_interval": 100,
+        "reward_interval": 50,
+        "reward_duration": 2800,
+        "stimulus_duration": 3000,
+        "stimulus_density": 23
+      },
+      "10": {
+        "level": 10,
+        "evaluation_interval": 100,
+        "reward_interval": 55,
+        "reward_duration": 3000,
+        "stimulus_duration": 3000,
+        "stimulus_density": 25
+      },
+      "11": {
+        "level": 11,
+        "evaluation_interval": 100,
+        "reward_interval": 60,
+        "reward_duration": 3000,
+        "stimulus_duration": 3000,
+        "stimulus_density": 25
+      }
+    }
+  }
+}
+```

mxbiflow/tasks/default/initial_habituation_training/stages/config.csv

@@ -0,0 +1,7 @@
+_key,level,entry_reward,evaluation_interval,min_stimulus_interval,max_stimulus_interval,target,stimulus_density
+0,0,1,10,2,2,5,5
+1,1,0.8,50,2,2,10,5
+2,2,0.6,100,3,5,15,10
+3,3,0.4,100,4,7,20,15
+4,4,0.2,100,5,9,25,20
+5,5,0,200,5,10,30,25

mxbiflow/tasks/default/initial_habituation_training/stages/config.json

@@ -0,0 +1,67 @@
+{
+    "default": {
+        "condition": {
+            "config": {
+                "difficulty_increase_threshold": 0,
+                "allow_decrease": false,
+                "present_level_trial_id": true
+            }
+        },
+        "levels_table": {
+            "0": {
+                "level": 0,
+                "entry_reward": 1,
+                "evaluation_interval": 10,
+                "min_stimulus_interval": 2,
+                "max_stimulus_interval": 2,
+                "target": 5,
+                "stimulus_density": 5
+            },
+            "1": {
+                "level": 1,
+                "entry_reward": 0.8,
+                "evaluation_interval": 50,
+                "min_stimulus_interval": 2,
+                "max_stimulus_interval": 2,
+                "target": 10,
+                "stimulus_density": 5
+            },
+            "2": {
+                "level": 2,
+                "entry_reward": 0.6,
+                "evaluation_interval": 100,
+                "min_stimulus_interval": 3,
+                "max_stimulus_interval": 5,
+                "target": 15,
+                "stimulus_density": 10
+            },
+            "3": {
+                "level": 3,
+                "entry_reward": 0.4,
+                "evaluation_interval": 100,
+                "min_stimulus_interval": 4,
+                "max_stimulus_interval": 7,
+                "target": 20,
+                "stimulus_density": 15
+            },
+            "4": {
+                "level": 4,
+                "entry_reward": 0.2,
+                "evaluation_interval": 100,
+                "min_stimulus_interval": 5,
+                "max_stimulus_interval": 9,
+                "target": 25,
+                "stimulus_density": 20
+            },
+            "5": {
+                "level": 5,
+                "entry_reward": 0,
+                "evaluation_interval": 200,
+                "min_stimulus_interval": 5,
+                "max_stimulus_interval": 10,
+                "target": 30,
+                "stimulus_density": 25
+            }
+        }
+    }
+}

mxbiflow/tasks/default/initial_habituation_training/stages/initial_habituation_training_stage.py

@@ -0,0 +1,172 @@
+from datetime import datetime
+from random import choices
+from typing import TYPE_CHECKING, Final
+
+from mxbi.data_logger import DataLogger, DataLoggerType
+from mxbi.models.animal import ScheduleCondition
+from mxbi.tasks.default.initial_habituation_training.stages.models import (
+    InitialHabituationTrainingStageConfig,
+    StageContext,
+    StageContexts,
+    config,
+)
+from mxbi.tasks.default.initial_habituation_training.tasks.stay_to_reward.stay_to_reward import (
+    DefaultStayToRewardScene,
+)
+from mxbi.tasks.default.initial_habituation_training.tasks.stay_to_reward.stay_to_reward_models import (
+    Result,
+    TrialConfig,
+)
+from mxbi.utils.logger import logger
+from mxbi.utils.tkinter.components.canvas_with_border import CanvasWithInnerBorder
+
+if TYPE_CHECKING:
+    from mxbi.models.animal import AnimalState
+    from mxbi.models.session import SessionConfig, SessionState
+    from mxbi.models.task import Feedback
+    from mxbi.theater import Theater
+
+
+contexts = StageContexts()
+background: CanvasWithInnerBorder | None = None
+
+
+def _initialize_contexts(session_config: "SessionConfig") -> None:
+    global contexts
+    for monkey in session_config.animals.keys():
+        if contexts.root.get(monkey) is None:
+            contexts.root[monkey] = StageContext()
+
+
+def _initialize_background(theater: "Theater", session_config: "SessionConfig") -> None:
+    global background
+    background = CanvasWithInnerBorder(
+        master=theater.root,
+        bg="black",
+        width=session_config.screen_type.width,
+        height=session_config.screen_type.height,
+        border_width=40,
+    )
+
+
+class InitialHabituationTrainingStage:
+    STAGE_NAME: Final[str] = "DEFAULT_INITIAL_HABITUATION_TRAINING_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
+
+        if not contexts.root:
+            _initialize_contexts(session_state.session_config)
+
+        if background is None:
+            _initialize_background(theater, session_state.session_config)
+
+        context = contexts.root[animal_state.name]
+
+        self._stage_config = self._load_stage_config(self._animal_state.name)
+
+        _levels_config = self._stage_config.levels_table[animal_state.level]
+        self._stage_config.condition.config.evaluation_interval = (
+            _levels_config.evaluation_interval
+        )
+
+        entry_reward = choices(
+            [True, False],
+            weights=[_levels_config.entry_reward, 1 - _levels_config.entry_reward],
+        )[0]
+
+        _config = TrialConfig(
+            level=_levels_config.level,
+            entry_reward=entry_reward,
+            min_stimulus_interval=_levels_config.min_stimulus_interval,
+            max_stimulus_interval=_levels_config.max_stimulus_interval,
+            target=_levels_config.target,
+            stimulus_density=_levels_config.stimulus_density,
+        )
+
+        self._data_logger = DataLogger(
+            self._session_state,
+            self._animal_state.name,
+            self.STAGE_NAME,
+            DataLoggerType.JSONL,
+        )
+
+        if self._animal_state.data_path is None:
+            self._animal_state.data_path = self._data_logger.path
+
+        assert background is not None
+
+        self._task = DefaultStayToRewardScene(
+            theater=theater,
+            animal_state=animal_state,
+            screen_type=session_state.session_config.screen_type,
+            trial_config=_config,
+            context=context,
+            background=background,
+        )
+
+    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}"
+        )
+
+        now = datetime.now().timestamp()
+        if self._animal_state.animal_session_start_time != 0.0:
+            session_duration = now - self._animal_state.animal_session_start_time
+            if session_duration >= 30:
+                print(session_duration)
+                print(self._animal_state.animal_session_start_time)
+                self._animal_state.animal_session_start_time = (
+                    datetime.now().timestamp()
+                )
+
+                self._theater._scheduler._increase_difficulty(self._animal_state)
+
+        return feedback
+
+    def _load_stage_config(self, monkey: str) -> InitialHabituationTrainingStageConfig:
+        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:
+                feedback = True
+            case Result.INCORRECT:
+                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

mxbiflow/tasks/default/initial_habituation_training/stages/models.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+
+from mxbi.config import Configure
+from mxbi.models.animal import ScheduleCondition
+from mxbi.tasks.GNGSiD.models import MonkeyName
+from pydantic import BaseModel, ConfigDict, Field, RootModel
+
+CONFIG_PATH = Path(__file__).parent / "config.json"
+
+
+class InitialHabituationTraingStageLeveledParams(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    level: int
+    entry_reward: float
+
+    evaluation_interval: int
+
+    min_stimulus_interval: float
+    max_stimulus_interval: float
+
+    target: float
+
+    stimulus_density: int
+
+
+class InitialHabituationTrainingStageConfig(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    condition: ScheduleCondition
+    levels_table: dict[int, InitialHabituationTraingStageLeveledParams]
+
+
+class InitialHabituationTrainingStageConfigs(RootModel):
+    model_config = ConfigDict(frozen=True)
+
+    root: dict[MonkeyName, InitialHabituationTrainingStageConfig]
+
+
+class StageContext(BaseModel):
+    duration: int = 0
+    rewards: int = 0
+
+
+class StageContexts(RootModel):
+    root: dict[MonkeyName, StageContext] = Field(default_factory=dict)
+
+
+def load_config() -> InitialHabituationTrainingStageConfigs:
+    configs = Configure(CONFIG_PATH, InitialHabituationTrainingStageConfigs).value
+    for config in configs.root.values():
+        config.condition.level_count = len(config.levels_table)
+    return configs
+
+
+config = load_config()