mxbiflow 0.1.1__py3-none-any.whl

mxbiflow/tasks/GNGSiD/stages/detect_stage/detect_stage.py

@@ -0,0 +1,161 @@
+from random import choice, choices, randint
+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.detect_stage.detect_stage_models import (
+    DetectStageConfig,
+    config,
+)
+from mxbi.tasks.GNGSiD.tasks.detect.models import TrialConfig
+from mxbi.tasks.GNGSiD.tasks.detect.scene import GNGSiDDetectScene
+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 GNGSiDDetectStage:
+    STAGE_NAME: Final[str] = "GNGSiD_DETECT_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]
+
+        _is_go = choices(
+            [True, False],
+            weights=[_levels_config.go_task_prob, _levels_config.nogo_task_prob],
+        )[0]
+
+        _stimulus_duration = randint(
+            _levels_config.min_stimulus_duration,
+            _levels_config.max_stimulus_duration,
+        )
+
+        _master_amp, _digital_amp = self._prepare_stimulus_intensity(
+            self._animal_state.name, _fixed_config.stimulus_freq
+        )
+
+        _config = TrialConfig(
+            level=_levels_config.level,
+            stimulation_size=_fixed_config.stimulation_size,
+            stimulus_duration=_stimulus_duration,
+            time_out=_fixed_config.time_out,
+            inter_trial_interval=_fixed_config.inter_trial_interval,
+            reward_duration=_fixed_config.reward_duration,
+            reward_delay=_fixed_config.reward_delay,
+            go=_is_go,
+            visual_stimulus_delay=_fixed_config.visual_stimulus_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._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 = GNGSiDDetectScene(
+            theater,
+            session_state.session_config,
+            animal_state,
+            session_state.session_config.screen_type,
+            _config,
+            self._presistent_data,
+        )
+
+        self._data_logger = DataLogger(
+            self._session_state,
+            self._animal_state.name,
+            self.STAGE_NAME,
+            DataLoggerType.JSONL,
+        )
+
+    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) -> DetectStageConfig:
+        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/stages/detect_stage/detect_stage_models.py

@@ -0,0 +1,65 @@
+from pathlib import Path
+
+from mxbi.config import Configure
+from mxbi.models.animal import ScheduleCondition
+from mxbi.tasks.GNGSiD.models import LevelID, MonkeyName
+from pydantic import BaseModel, ConfigDict, RootModel
+
+CONFIG_PATH = Path(__file__).parent / "config.json"
+
+
+class DetectStageParams(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    # visual stimulation config
+    stimulation_size: int
+    visual_stimulus_delay: int
+
+    # audio stimulation config
+    stimulus_freq: int
+    stimulus_freq_duration: int
+    stimulus_interval: int
+
+    # trial lifecycle
+    time_out: int
+    inter_trial_interval: int
+
+    # reward config
+    reward_duration: int
+    reward_delay: int
+
+
+class DetectStageLeveledParams(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    level: int
+    min_stimulus_duration: int
+    max_stimulus_duration: int
+
+    go_task_prob: float
+    nogo_task_prob: float
+
+
+class DetectStageConfig(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    condition: ScheduleCondition
+
+    params: DetectStageParams
+    levels_table: dict[LevelID, DetectStageLeveledParams]
+
+
+class DetectStageConfigs(RootModel):
+    model_config = ConfigDict(frozen=True)
+
+    root: dict[MonkeyName, DetectStageConfig]
+
+
+def load_config() -> DetectStageConfigs:
+    configs = Configure(CONFIG_PATH, DetectStageConfigs).value
+    for config in configs.root.values():
+        config.condition.level_count = len(config.levels_table)
+    return configs
+
+
+config = load_config()

mxbiflow/tasks/GNGSiD/stages/discriminate_stage/config.json

@@ -0,0 +1,70 @@
+{
+    "default": {
+        "condition": {
+            "config": {
+                "evaluation_interval": 20,
+                "difficulty_increase_threshold": 0.8,
+                "difficulty_decrease_threshold": 0.45
+            }
+        },
+        "params": {
+            "level": 0,
+            "stimulation_size": 250,
+            "visual_stimulus_delay": 500,
+            "min_stimulus_duration": 2500,
+            "max_stimulus_duration": 6500,
+            "extra_response_time": 1000,
+            "stimulus_configs": [
+                {
+                    "stimulus_freq_high": 2640,
+                    "stimulus_freq_high_duration": 100,
+                    "stimulus_freq_low": 1760,
+                    "stimulus_freq_low_duration": 100
+                },
+                {
+                    "stimulus_freq_high": 5920,
+                    "stimulus_freq_high_duration": 100,
+                    "stimulus_freq_low": 3950,
+                    "stimulus_freq_low_duration": 100
+                },
+                {
+                    "stimulus_freq_high": 7000,
+                    "stimulus_freq_high_duration": 100,
+                    "stimulus_freq_low": 4700,
+                    "stimulus_freq_low_duration": 100
+                }
+            ],
+            "stimulus_interval": 200,
+            "time_out": 10000,
+            "inter_trial_interval": 2000,
+            "reward_duration": 1000,
+            "medium_reward_duration": 800,
+            "medium_reward_threshold": 2000,
+            "low_reward_duration": 500,
+            "reward_delay": 500,
+            "attention_duration": 2800
+        },
+        "levels_table": {
+            "0": {
+                "level": 0,
+                "stimulus_trial_prob": 0.2,
+                "nostimulus_trial_prob": 0.8
+            },
+            "1": {
+                "level": 1,
+                "stimulus_trial_prob": 0.3,
+                "nostimulus_trial_prob": 0.7
+            },
+            "2": {
+                "level": 2,
+                "stimulus_trial_prob": 0.4,
+                "nostimulus_trial_prob": 0.6
+            },
+            "3": {
+                "level": 3,
+                "stimulus_trial_prob": 0.5,
+                "nostimulus_trial_prob": 0.5
+            }
+        }
+    }
+}

mxbiflow/tasks/GNGSiD/stages/discriminate_stage/discriminate_stage.py

@@ -0,0 +1,173 @@
+from random import choice, choices, randint
+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.discriminate_stage.discriminate_stage_models import (
+    DiscriminateStageConfig,
+    config,
+)
+from mxbi.tasks.GNGSiD.tasks.discriminate.discriminate_models import TrialConfig
+from mxbi.tasks.GNGSiD.tasks.discriminate.discriminate_scene import (
+    GNGSiDDiscriminateScene,
+)
+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 GNGSiDDiscriminateStage:
+    STAGE_NAME: Final[str] = "GNGSiD_DISCRIMINATE_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]
+
+        _stimulus_config = choice(_fixed_config.stimulus_configs)
+        _stimulus_duration = randint(
+            _fixed_config.min_stimulus_duration, _fixed_config.max_stimulus_duration
+        )
+
+        _is_stimulus_trial = choices(
+            [True, False],
+            weights=[
+                _levels_config.stimulus_trial_prob,
+                _levels_config.nostimulus_trial_prob,
+            ],
+        )[0]
+
+        _high_master_amp, _high_digital_amp = self._prepare_stimulus_intensity(
+            animal_state.name, _stimulus_config.stimulus_freq_high
+        )
+        _low_master_amp, _low_digital_amp = self._prepare_stimulus_intensity(
+            animal_state.name, _stimulus_config.stimulus_freq_low
+        )
+
+        _config = TrialConfig(
+            level=_levels_config.level,
+            stimulation_size=_fixed_config.stimulation_size,
+            stimulus_duration=_stimulus_duration,
+            time_out=_fixed_config.time_out,
+            inter_trial_interval=_fixed_config.inter_trial_interval,
+            reward_duration=_fixed_config.reward_duration,
+            reward_delay=_fixed_config.reward_delay,
+            is_stimulus_trial=_is_stimulus_trial,
+            visual_stimulus_delay=_fixed_config.visual_stimulus_delay,
+            medium_reward_duration=_fixed_config.medium_reward_duration,
+            medium_reward_threshold=_fixed_config.medium_reward_threshold,
+            low_reward_duration=_fixed_config.low_reward_duration,
+            attention_duration=_fixed_config.attention_duration,
+            stimulus_freq_low=_stimulus_config.stimulus_freq_low,
+            stimulus_freq_low_duration=_stimulus_config.stimulus_freq_low_duration,
+            stimulus_freq_low_master_amp=_low_master_amp,
+            stimulus_freq_low_digital_amp=_low_digital_amp,
+            stimulus_freq_high=_stimulus_config.stimulus_freq_high,
+            stimulus_freq_high_duration=_stimulus_config.stimulus_freq_high_duration,
+            stimulus_freq_high_master_amp=_high_master_amp,
+            stimulus_freq_high_digital_amp=_high_digital_amp,
+            stimulus_interval=_fixed_config.stimulus_interval,
+            extra_response_time=_fixed_config.extra_response_time,
+        )
+
+        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 = GNGSiDDiscriminateScene(
+            theater,
+            session_state.session_config,
+            animal_state,
+            session_state.session_config.screen_type,
+            _config,
+            self._presistent_data,
+        )
+
+        self._data_logger = DataLogger(
+            self._session_state,
+            self._animal_state.name,
+            self.STAGE_NAME,
+            DataLoggerType.JSONL,
+        )
+
+    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) -> DiscriminateStageConfig:
+        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):
+        stimulus_intensity = choice([55, 60, 65, 70, 75])
+
+        return self._theater.acontroller.get_amp_value(frequency, stimulus_intensity)

mxbiflow/tasks/GNGSiD/stages/discriminate_stage/discriminate_stage_models.py

@@ -0,0 +1,80 @@
+from pathlib import Path
+from typing import Dict, List
+
+from mxbi.config import Configure
+from mxbi.models.animal import ScheduleCondition
+from mxbi.tasks.GNGSiD.models import LevelID, MonkeyName
+from pydantic import BaseModel, ConfigDict, RootModel
+
+CONFIG_PATH: Path = Path(__file__).parent / "config.json"
+
+
+class StimulusConfig(BaseModel):
+    stimulus_freq_high: int
+    stimulus_freq_high_duration: int
+    stimulus_freq_low: int
+    stimulus_freq_low_duration: int
+
+
+class DiscriminateStageParams(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    # visual stimulation config
+    stimulation_size: int
+    visual_stimulus_delay: int
+
+    # dynamic reward config
+    medium_reward_duration: int
+    medium_reward_threshold: int
+    low_reward_duration: int
+
+    # audio stimulation config
+    min_stimulus_duration: int
+    max_stimulus_duration: int
+    stimulus_configs: List[StimulusConfig]
+    stimulus_interval: int
+    extra_response_time: int
+
+    # trial lifecycle
+    time_out: int
+    inter_trial_interval: int
+
+    # reward config
+    reward_duration: int
+    reward_delay: int
+
+    attention_duration: int
+
+
+class DiscriminateStageLeveledParams(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    level: int
+
+    stimulus_trial_prob: float
+    nostimulus_trial_prob: float
+
+
+class DiscriminateStageConfig(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    condition: ScheduleCondition
+
+    params: DiscriminateStageParams
+    levels_table: Dict[LevelID, DiscriminateStageLeveledParams]
+
+
+class DiscriminateStageConfigs(RootModel):
+    model_config = ConfigDict(frozen=True)
+
+    root: Dict[MonkeyName, DiscriminateStageConfig]
+
+
+def load_config() -> DiscriminateStageConfigs:
+    configs = Configure(CONFIG_PATH, DiscriminateStageConfigs).value
+    for config in configs.root.values():
+        config.condition.level_count = len(config.levels_table)
+    return configs
+
+
+config = load_config()

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

@@ -0,0 +1,83 @@
+{
+    "default": {
+        "condition": {
+            "config": {
+                "evaluation_interval": 20,
+                "difficulty_increase_threshold": 0.8,
+                "difficulty_decrease_threshold": 0.45,
+                "next_task": "gngsid_detect_stage"
+            }
+        },
+        "params": {
+            "stimulus_duration": 1000,
+            "stimulus_freq": 2000,
+            "stimulus_freq_duration": 100,
+            "stimulus_interval": 130,
+            "time_out": 10000,
+            "inter_trial_interval": 2000,
+            "reward_duration": 1000
+        },
+        "levels_table": {
+            "0": {
+                "level": 0,
+                "stimulation_size": 430,
+                "reward_delay": 0
+            },
+            "1": {
+                "level": 1,
+                "stimulation_size": 420,
+                "reward_delay": 0
+            },
+            "2": {
+                "level": 2,
+                "stimulation_size": 400,
+                "reward_delay": 100
+            },
+            "3": {
+                "level": 3,
+                "stimulation_size": 380,
+                "reward_delay": 200
+            },
+            "4": {
+                "level": 4,
+                "stimulation_size": 360,
+                "reward_delay": 300
+            },
+            "5": {
+                "level": 5,
+                "stimulation_size": 345,
+                "reward_delay": 400
+            },
+            "6": {
+                "level": 6,
+                "stimulation_size": 330,
+                "reward_delay": 500
+            },
+            "7": {
+                "level": 7,
+                "stimulation_size": 315,
+                "reward_delay": 600
+            },
+            "8": {
+                "level": 8,
+                "stimulation_size": 300,
+                "reward_delay": 700
+            },
+            "9": {
+                "level": 9,
+                "stimulation_size": 285,
+                "reward_delay": 800
+            },
+            "10": {
+                "level": 10,
+                "stimulation_size": 270,
+                "reward_delay": 900
+            },
+            "11": {
+                "level": 11,
+                "stimulation_size": 260,
+                "reward_delay": 1000
+            }
+        }
+    }
+}

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

@@ -0,0 +1,58 @@
+from pathlib import Path
+
+from mxbi.config import Configure
+from mxbi.models.animal import ScheduleCondition
+from mxbi.tasks.GNGSiD.models import LevelID, MonkeyName
+from pydantic import BaseModel, ConfigDict, RootModel
+
+CONFIG_PATH = Path(__file__).parent / "config.json"
+
+
+class SizeReductionStageParameters(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    # audio stimulation config
+    stimulus_duration: int
+    stimulus_freq: int
+    stimulus_freq_duration: int
+    stimulus_interval: int
+
+    # trial lifecycle
+    time_out: int
+    inter_trial_interval: int
+
+    # reward config
+    reward_duration: int
+
+
+class SizeReductionStageLeveledParameters(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    level: int
+    stimulation_size: int
+    reward_delay: int
+
+
+class SizeReductionStageConfig(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    condition: ScheduleCondition
+
+    params: SizeReductionStageParameters
+    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()