mxbiflow 0.1.1__py3-none-any.whl

mxbiflow/models/reward.py

@@ -0,0 +1,7 @@
+from enum import StrEnum, auto
+
+
+class RewardEnum(StrEnum):
+    AGUM_ONE_FIFTH = auto()
+    AGUM_TWO_FIFTHS = auto()
+    M_JUICE = auto()

mxbiflow/models/session.py

@@ -0,0 +1,145 @@
+import os
+import tempfile
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+
+from pydantic import BaseModel, Field, PrivateAttr, ValidationError
+
+from .animal import Animal, AnimalConfig
+from .reward import RewardEnum
+
+
+class SessionConfig(BaseModel):
+    experimenter: str = Field(default="auto", frozen=True)
+    reward_type: RewardEnum = Field(default=RewardEnum.AGUM_ONE_FIFTH, frozen=True)
+    send_email: bool = Field(default=False, frozen=True)
+    sync_data: bool = Field(default=False, frozen=True)
+    note: str = Field(default="", max_length=1000)
+
+    animals: list[AnimalConfig] = Field(default_factory=list)
+
+
+class DailySessionCounter(BaseModel):
+    day: str = Field(
+        default_factory=lambda: datetime.now(timezone.utc).date().isoformat()
+    )
+    last_session_id: int = Field(default=0, ge=0)
+
+
+@dataclass
+class DailySessionIdStore:
+    path: Path
+
+    def _today_utc(self):
+        return datetime.now(timezone.utc).date().isoformat()
+
+    def _load(self) -> DailySessionCounter:
+        if not self.path.exists():
+            return DailySessionCounter()
+
+        text = self.path.read_text()
+
+        try:
+            return DailySessionCounter.model_validate_json(text)
+        except ValueError, ValidationError:
+            return DailySessionCounter()
+
+    def _atomic_write(self, data: DailySessionCounter) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        fd, tmp = tempfile.mkstemp(
+            dir=str(self.path.parent),
+            prefix=self.path.name,
+            suffix=".tmp",
+        )
+
+        try:
+            with os.fdopen(fd, "w", encoding="utf-8") as f:
+                f.write(data.model_dump_json())
+                f.flush()
+                os.fsync(f.fileno())
+
+            os.replace(tmp, self.path)
+        finally:
+            try:
+                os.remove(tmp)
+            except FileNotFoundError:
+                pass
+
+    @property
+    def session_id(self) -> int:
+        today = self._today_utc()
+        data = self._load()
+
+        if data.day != today:
+            data.day = today
+            data.last_session_id = 0
+
+        data.last_session_id += 1
+        self._atomic_write(data)
+
+        return data.last_session_id
+
+
+class Session(BaseModel):
+    experimenter: str = Field(default="auto", frozen=True)
+    reward_type: RewardEnum = Field(default=RewardEnum.AGUM_ONE_FIFTH, frozen=True)
+    send_email: bool = Field(default=False, frozen=True)
+    sync_data: bool = Field(default=False, frozen=True)
+
+    session_id: int = Field(default=0, ge=0)
+    start_at: float = Field(default=0, ge=0)
+    end_at: float = Field(default=0, ge=0)
+    note: str = Field(frozen=True)
+
+    _current_animal: str | None = PrivateAttr(default=None)
+    animals: dict[str, Animal] = Field(default_factory=dict, frozen=True)
+
+    def start(self):
+        self.start_at = datetime.now(timezone.utc).timestamp()
+
+    def end(self):
+        self.end_at = datetime.now(timezone.utc).timestamp()
+
+    @property
+    def current_animal(self) -> Animal | None:
+        key = self._current_animal
+        if key is None:
+            return None
+
+        return self.animals[key]
+
+    @property
+    def require_current_animal(self) -> Animal:
+        key = self._current_animal
+        if key is None:
+            raise RuntimeError(
+                f"Animal: {key} is not set. Please call set_current_animal() first"
+            )
+
+        return self.animals[key]
+
+    def clear_current_animal(self):
+        a = self.current_animal
+        if a is None:
+            return
+
+        a.end_animal_session()
+        self._current_animal = None
+
+    def set_current_animal(self, animal: str):
+        self.clear_current_animal()
+
+        try:
+            a = self.animals[animal]
+        except KeyError:
+            raise ValueError(f"animal {animal} not found")
+
+        self._current_animal = animal
+        a.start_animal_session()
+
+
+class Options(BaseModel):
+    mxbis: list[str] = Field(default_factory=list, frozen=True)
+    experimenter: list[str] = Field(default_factory=list, frozen=True)
+    animals: dict[str, str] = Field(default_factory=dict, frozen=True)

mxbiflow/mxbiflow.py

@@ -0,0 +1,43 @@
+from pymxbi import MXBI
+
+from .models.session import Session
+from .timer import FrameTimer
+
+
+class MXBIFlow:
+    def __init__(self, session: Session, mxbi: MXBI) -> None:
+        self._session = session
+        self._timer = FrameTimer()
+        self._mxbi = mxbi
+
+    @property
+    def session(self) -> Session:
+        return self._session
+
+    @property
+    def timer(self) -> FrameTimer:
+        return self._timer
+
+    @property
+    def mxbi(self) -> MXBI:
+        return self._mxbi
+
+    def update(self) -> None:
+        self._timer.update()
+
+
+_current_mxbiflow: MXBIFlow | None = None
+
+
+def get_mxbiflow() -> MXBIFlow:
+    global _current_mxbiflow
+    if _current_mxbiflow is None:
+        raise RuntimeError("MXBIFlow not initialized")
+    return _current_mxbiflow
+
+
+def set_mxbiflow(mxbiflow: MXBIFlow) -> None:
+    global _current_mxbiflow
+    if _current_mxbiflow is not None:
+        raise RuntimeError("MXBIFlow already initialized")
+    _current_mxbiflow = mxbiflow

mxbiflow/path.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+
+# TODO: XDG Base Directory Specification
+ROOT_DIR_PATH = Path(__file__).parents[2]
+
+CONFIG_DIR_PATH = ROOT_DIR_PATH / "config"
+CONFIG_SESSION_FILENAME = "config_session.json"
+CONFIG_SESSION_PATH = CONFIG_DIR_PATH / CONFIG_SESSION_FILENAME
+
+OPTIONS_SESSION_FILENAME = "options_session.json"
+OPTIONS_SESSION_PATH = CONFIG_DIR_PATH / OPTIONS_SESSION_FILENAME
+
+ANIMAL_DB_FILENAME = "animal_db.json"
+ANIMAL_DB_PATH = CONFIG_DIR_PATH / ANIMAL_DB_FILENAME
+
+CROSS_MODAL_CONFIG_FILENAME = "config_cross_modal.json"
+CROSS_MODAL_CONFIG_PATH = CONFIG_DIR_PATH / CROSS_MODAL_CONFIG_FILENAME
+
+DATA_DIR_PATH = ROOT_DIR_PATH / "data"
+
+LOG_PATH = ROOT_DIR_PATH / "log"
+
+SAMBA_MOUNT_PATH = ROOT_DIR_PATH / "samba_mount"
+SAMBA_BACKUP_DIR_NAME = "backup"
+SAMBA_BACKUP_DIR_PATH = SAMBA_MOUNT_PATH / SAMBA_BACKUP_DIR_NAME
+
+SERVICE_DIR_NAME = "services"
+SERVICE_DIR_PATH = ROOT_DIR_PATH / SERVICE_DIR_NAME
+MOUNT_SERVICE_NAME = "mount.service"
+MOUNT_SERVICE_PATH = SERVICE_DIR_PATH / MOUNT_SERVICE_NAME
+SYNC_SERVICE_NAME = "sync.service"
+SYNC_SERVICE_PATH = SERVICE_DIR_PATH / SYNC_SERVICE_NAME
+
+MXBI_CONFIG_PATH = CONFIG_DIR_PATH / "mxbi.json"
+SESSION_CONFIG_PATH = CONFIG_DIR_PATH / "session.json"
+OPTIONS_PATH = CONFIG_DIR_PATH / "options.json"
+STAGE_PATH = CONFIG_DIR_PATH / "stage.json"
+SESSION_COUNTER_PATH = CONFIG_DIR_PATH / "session_counter.json"
+
+ASSETS_DIR_PATH = ROOT_DIR_PATH / "assets"
+ASSET_CLICKER_PATH = ASSETS_DIR_PATH / "clicker.wav"

mxbiflow/scene/__init__.py

@@ -0,0 +1,8 @@
+from .scene_manager import SceneManager, Scenes
+from .scene_protocol import SceneProtocol
+
+__all__ = [
+    "SceneManager",
+    "Scenes",
+    "SceneProtocol",
+]

mxbiflow/scene/scene_manager.py

@@ -0,0 +1,64 @@
+from pathlib import Path
+
+from pydantic import RootModel
+from pygame import Event, Surface
+
+from .scene_protocol import SceneProtocol
+
+
+class Scenes(RootModel):
+    root: list[str]
+
+
+class SceneManager:
+    scenes: dict[str, type[SceneProtocol]] = {}
+
+    def __init__(self) -> None:
+        self.current: SceneProtocol | None = None
+        self._pending: SceneProtocol | None = None
+
+    def persist(self, path: Path) -> None:
+        scenes = Scenes(root=list(self.scenes.keys()))
+        json_data = scenes.model_dump_json()
+
+        with path.open("w") as f:
+            f.write(json_data)
+
+    def register(self, scene: type[SceneProtocol], name: str | None = None) -> None:
+        if name is None:
+            name = scene.__name__.lower()
+
+        self.scenes[name] = scene
+
+    def switch(self, scene: type[SceneProtocol], defer: bool = True) -> None:
+        if defer:
+            self._pending = scene()
+        else:
+            self._switch(scene())
+
+    def _switch(self, next_scene: SceneProtocol) -> None:
+        prev = self.current
+        if prev is not None and prev.running:
+            prev.quit()
+
+        self.current = next_scene
+        next_scene.start()
+
+    def apply_pending(self) -> None:
+        if self._pending is None:
+            return
+        next_scene = self._pending
+        self._pending = None
+        self._switch(next_scene)
+
+    def handle_event(self, event: Event) -> None:
+        if self.current:
+            self.current.handle_event(event)
+
+    def update(self, dt_s: float) -> None:
+        if self.current:
+            self.current.update(dt_s)
+
+    def draw(self, screen: Surface) -> None:
+        if self.current:
+            self.current.draw(screen)

mxbiflow/scene/scene_protocol.py

@@ -0,0 +1,22 @@
+from typing import Protocol, runtime_checkable
+
+from pygame.event import Event
+from pygame.surface import Surface
+
+
+@runtime_checkable
+class SceneProtocol(Protocol):
+    _running: bool
+
+    def start(self) -> None: ...
+
+    def quit(self) -> None: ...
+
+    @property
+    def running(self) -> bool: ...
+
+    def handle_event(self, event: Event) -> None: ...
+
+    def update(self, dt_s: float) -> None: ...
+
+    def draw(self, screen: Surface) -> None: ...

mxbiflow/scheduler.py

@@ -0,0 +1,90 @@
+from pygame import Event
+from pymxbi.detector.detector import DetectorEvent
+
+from .default import IDLE
+from .detector_bridge import EVT_DETECTOR, DetectorMsg
+from .models.session import Session
+from .scene import SceneManager
+
+
+class Scheduler:
+    def __init__(
+        self,
+        session: Session,
+        scene_manager: SceneManager,
+    ) -> None:
+        self._session = session
+        self._scene_manager = scene_manager
+        self._scenes = self._scene_manager.scenes
+
+        self._need_refresh = False
+
+    def _mark_refresh(self) -> None:
+        self._need_refresh = True
+
+    def _set_current_animal(self, animal: str) -> None:
+        self._session.set_current_animal(animal)
+        self._need_refresh = True
+
+    def _clear_current_animal(self) -> None:
+        if self._session.current_animal is None:
+            return
+
+        self._session.clear_current_animal()
+        self._mark_refresh()
+
+    def _set_next_stage(self, stage: str) -> None:
+        if self._session.current_animal is None:
+            return
+
+        self._session.current_animal.set_current_stage(stage)
+        self._need_refresh = True
+
+    def _handle_fault_event(self) -> None: ...
+
+    def handle_event(self, event: Event) -> None:
+        if event.type != EVT_DETECTOR:
+            return
+
+        msg: DetectorMsg = event.msg
+
+        match msg.kind:
+            case DetectorEvent.FAULT_DETECTED:
+                self._handle_fault_event()
+
+            case DetectorEvent.ANIMAL_ENTERED:
+                if msg.animal is None:
+                    return
+                self._set_current_animal(msg.animal)
+
+            case DetectorEvent.ANIMAL_LEFT:
+                self._clear_current_animal()
+
+    def _shoud_refresh(self) -> bool:
+        current = self._scene_manager.current
+        return (
+            current is not None
+            and not isinstance(current, IDLE)
+            and not current.running
+        )
+
+    def update(self) -> None:
+        if self._shoud_refresh():
+            self._mark_refresh()
+
+        if not self._need_refresh:
+            return
+
+        self._need_refresh = False
+        self._refresh_by_state()
+
+    def _refresh_by_state(self) -> None:
+        animal = self._session.current_animal
+
+        if animal is None:
+            if not isinstance(self._scene_manager.current, IDLE):
+                self._scene_manager.switch(self._scenes[IDLE.__name__.lower()])
+            return
+
+        stage = animal.current_stage.stage_name
+        self._scene_manager.switch(self._scenes[stage])

mxbiflow/tasks/GNGSiD/models.py

@@ -0,0 +1,70 @@
+from enum import StrEnum, auto
+from typing import TypeAlias
+
+from pydantic import BaseModel, ConfigDict
+
+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):
+    model_config = ConfigDict(frozen=True)
+
+    # basic config
+    level: int
+
+    # visual stimulation config
+    stimulation_size: int
+
+    # audio stimulation config
+    stimulus_duration: int
+
+    # trial lifecycle
+    time_out: int
+    inter_trial_interval: int
+
+    # reward config
+    reward_duration: int
+    reward_delay: int
+
+
+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/GNGSiD/stages/detect_stage/config.json

@@ -0,0 +1,116 @@
+{
+    "default": {
+        "condition": {
+            "config": {
+                "evaluation_interval": 2,
+                "difficulty_increase_threshold": 0.8,
+                "difficulty_decrease_threshold": 0.45,
+                "next_task": "gngsid_discriminate_stage"
+            }
+        },
+        "params": {
+            "stimulation_size": 250,
+            "visual_stimulus_delay": 500,
+            "stimulus_freq": 2000,
+            "stimulus_freq_duration": 100,
+            "stimulus_interval": 130,
+            "time_out": 10000,
+            "inter_trial_interval": 2000,
+            "reward_duration": 1000,
+            "reward_delay": 500
+        },
+        "levels_table": {
+            "0": {
+                "level": 0,
+                "min_stimulus_duration": 2500,
+                "max_stimulus_duration": 6500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "1": {
+                "level": 1,
+                "min_stimulus_duration": 1000,
+                "max_stimulus_duration": 2500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "2": {
+                "level": 2,
+                "min_stimulus_duration": 3500,
+                "max_stimulus_duration": 6500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "3": {
+                "level": 3,
+                "min_stimulus_duration": 1000,
+                "max_stimulus_duration": 2500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "4": {
+                "level": 4,
+                "min_stimulus_duration": 3500,
+                "max_stimulus_duration": 6500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "5": {
+                "level": 5,
+                "min_stimulus_duration": 1000,
+                "max_stimulus_duration": 2500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "6": {
+                "level": 6,
+                "min_stimulus_duration": 3500,
+                "max_stimulus_duration": 6500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "7": {
+                "level": 7,
+                "min_stimulus_duration": 1000,
+                "max_stimulus_duration": 3500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "8": {
+                "level": 8,
+                "min_stimulus_duration": 3500,
+                "max_stimulus_duration": 6500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "9": {
+                "level": 9,
+                "min_stimulus_duration": 1000,
+                "max_stimulus_duration": 2500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "10": {
+                "level": 10,
+                "min_stimulus_duration": 3500,
+                "max_stimulus_duration": 6500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "11": {
+                "level": 11,
+                "min_stimulus_duration": 1000,
+                "max_stimulus_duration": 2500,
+                "go_task_prob": 0.8,
+                "nogo_task_prob": 0.2
+            },
+            "12": {
+                "level": 12,
+                "min_stimulus_duration": 3000,
+                "max_stimulus_duration": 5000,
+                "go_task_prob": 0.5,
+                "nogo_task_prob": 0.5
+            }
+        }
+    }
+}