Feat metadata mapping

scripts/walk_through_session.py

@@ -1,11 +1,10 @@
 import logging
 import os
 
-from aind_behavior_dynamic_foraging.task_logic.trial_generators.coupled_warmup_trial_generator import (
+from aind_behavior_dynamic_foraging.data_contract import dataset as df_foraging_dataset
+from aind_behavior_dynamic_foraging.task_logic.trial_generators.coupled_trial_generators.coupled_warmup_trial_generator import (
     CoupledWarmupTrialGeneratorSpec,
 )
-
-from aind_behavior_dynamic_foraging.data_contract import dataset as df_foraging_dataset
 from aind_behavior_dynamic_foraging.task_logic.trial_models import TrialOutcome
 
 logging.basicConfig(
@@ -20,10 +19,10 @@ def walk_through_session(data_directory: os.PathLike):
     software_events.load_all()
 
     trial_outcomes = software_events["TrialOutcome"].data["data"].iloc
-    warmup_trial_generator = CoupledWarmupTrialGeneratorSpec().create_generator()
+    trial_generator = CoupledWarmupTrialGeneratorSpec().create_generator()
     for i, outcome in enumerate(trial_outcomes):
-        warmup_trial_generator.update(TrialOutcome.model_validate(outcome))
-        trial = warmup_trial_generator.next()
+        trial_generator.update(TrialOutcome.model_validate(outcome))
+        trial = trial_generator.next()
 
         if not trial:
             print(f"Session finished at trial {i}")

src/Extensions/bonsai.py

@@ -1,8 +1,16 @@
+import logging
+import sys
 from typing import TYPE_CHECKING
 
 from pydantic import TypeAdapter
 
-from aind_behavior_dynamic_foraging.task_logic import TrialGeneratorSpec
+logging.basicConfig(
+    stream=sys.stdout,
+    level=logging.DEBUG,
+    format='{"name": "%(name)s", "level": %(levelno)d, "msg": "%(message)s"}',
+)
+
+from aind_behavior_dynamic_foraging.task_logic import TrialGeneratorSpec  # noqa
 
 if TYPE_CHECKING:
     from aind_behavior_dynamic_foraging.task_logic.trial_generators._base import ITrialGenerator

src/aind_behavior_dynamic_foraging/data_contract/_dataset.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from aind_behavior_curriculum import TrainerState
+from aind_behavior_curriculum import Metrics, TrainerState
 from aind_behavior_services.session import Session
 from contraqctor.contract import Dataset, DataStreamCollection
 from contraqctor.contract.camera import Camera
@@ -61,10 +61,18 @@ def make_dataset(
                 data_streams=[
                     Json(
                         name="PreviousMetrics",
-                        reader_params=Json.make_params(
+                        reader_params=PydanticModel.make_params(
+                            model=Metrics,
                             path=root_path / "behavior/previous_metrics.json",
                         ),
                     ),
+                    PydanticModel(
+                        name="Metrics",
+                        reader_params=PydanticModel.make_params(
+                            model=Metrics,
+                            path=root_path / "behavior/metrics.json",
+                        ),
+                    ),
                     PydanticModel(
                         name="TrainerState",
                         reader_params=PydanticModel.make_params(

src/aind_behavior_dynamic_foraging/data_contract/utils.py

@@ -0,0 +1,34 @@
+import os
+from typing import Optional
+
+from aind_behavior_dynamic_foraging.data_contract import dataset
+from aind_behavior_dynamic_foraging.task_logic import AindDynamicForagingTaskLogic
+
+
+def calculate_consumed_water(session_path: os.PathLike) -> Optional[float]:
+    """Calculate the total volume of water consumed during a session.
+
+    Args:
+        session_path (os.PathLike): Path to the session directory.
+
+    Returns:
+        Optional[float]: Total volume of water consumed in milliliters, or None if unavailable.
+    """
+
+    trial_outcomes = dataset(session_path)["Behavior"]["SoftwareEvents"]["TrialOutcome"].load().data["data"]
+    is_right_choice = [to["is_right_choice"] for to in trial_outcomes]
+    is_rewarded = [to["is_rewarded"] for to in trial_outcomes]
+
+    task_logic_data = dataset(session_path)["Behavior"]["InputSchemas"]["TaskLogic"].load().data
+    task_logic = AindDynamicForagingTaskLogic.model_validate(task_logic_data)
+    right_reward_size = task_logic.task_parameters.reward_size.right_value_volume
+    left_reward_size = task_logic.task_parameters.reward_size.left_value_volume
+
+    total = 0
+    for choice, rewarded in zip(is_right_choice, is_rewarded):
+        if rewarded:
+            if choice is True:
+                total += right_reward_size * 1e-3
+            if choice is False:
+                total += left_reward_size * 1e-3
+    return total

src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_warmup_trial_generator.py

@@ -46,7 +46,7 @@ class CoupledWarmupTrialGeneratorSpec(BaseCoupledTrialGeneratorSpec):
         default=True, description="Whether uncollected rewards carry over to the next trial."
     )
 
-    def create_generator(self) -> "CoupledWarmupTrialGeneratorSpec":
+    def create_generator(self) -> "CoupledWarmupTrialGenerator":
         return CoupledWarmupTrialGenerator(self)
 
 

workspace/aind_behavior_dynamic_foraging_curricula/pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "aind-behavior-curriculum >= 0.0.38",
     "numpy>=2.4.2",
     "pydantic-settings",
-    "aind-behavior-dynamic-foraging==0.0.2rc24"
+    "aind-behavior-dynamic-foraging==0.0.2rc30"
 ]
 
 [tool.uv.sources]

workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py

@@ -71,7 +71,7 @@ def metrics_from_dataset(
     logger.debug(f"Calculated foraging efficiency as {foraging_efficiency}")
 
     try:
-        prev_metrics = DynamicForagingMetrics(**dataset["Behavior"]["PreviousMetrics"].data)
+        prev_metrics = DynamicForagingMetrics.model_validate(dataset["Behavior"]["PreviousMetrics"].data)
         prev_stage = prev_metrics.stage_name
     except FileNotFoundError:
         logger.info("No previous metrics found.")

workspace/aind_behavior_dynamic_foraging_curricula/tests/test_metrics.py

@@ -3,8 +3,6 @@
 from typing import Optional
 from unittest.mock import MagicMock, PropertyMock, patch
 
-import numpy as np
-
 from aind_behavior_dynamic_foraging_curricula.metrics import (
     metrics_from_dataset,
 )
@@ -110,17 +108,6 @@ def test_previous_metrics_accumulate(self):
         self.assertEqual(len(result.foraging_efficiency_per_session), 2)
         self.assertEqual(len(result.unignored_trials_per_session), 2)
 
-    def test_foraging_efficiency_is_finite_and_positive(self):
-        trials = [
-            _make_trial(True, True, 0.7, 0.3),
-            _make_trial(True, False, 0.7, 0.3),
-            _make_trial(False, True, 0.7, 0.3),
-        ]
-        with _patch_dataset(trials):
-            result = metrics_from_dataset(self.tmp_path)
-        self.assertGreater(result.foraging_efficiency_per_session[-1], 0)
-        self.assertTrue(np.isfinite(result.foraging_efficiency_per_session[-1]))
-
 
 if __name__ == "__main__":
     unittest.main()

workspace/aind_behavior_dynamic_foraging_metadata_mapper/pyproject.toml

@@ -0,0 +1,78 @@
+[build-system]
+requires = ["uv_build>=0.8.22"]
+build-backend = "uv_build"
+
+[project]
+name = "aind-behavior-dynamic-foraging-metadata-mapper"
+description = "A library of mapping for the Dynamic Foraging task."
+authors = [
+    {name = "Bruno Cruz", email = "bruno.cruz@alleninstitute.org"},
+    {name = "Micah Woodard", email = "micah.woodard@alleninstitute.org"}
+    ]
+license = "MIT"
+version = "0.0.1"
+requires-python = ">=3.11"
+classifiers = [
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Operating System :: Microsoft :: Windows",
+]
+readme = {file = "README.md", content-type = "text/markdown"}
+
+dependencies = [
+    "numpy>=2.4.2",
+    "pydantic-settings",
+    "aind-behavior-dynamic-foraging==0.0.2rc30",
+    "aind-data-schema>=2.6.0",
+    "cyclopts>=4.10.0"
+]
+
+[tool.uv.sources]
+aind-behavior-dynamic-foraging = { workspace = true }
+
+[dependency-groups]
+
+dev = [
+    'ruff',
+    'pytest',
+    'pytest-cov',
+    'codespell',
+]
+
+docs = [
+    'mkdocs',
+    'mkdocs-material',
+    'mkdocstrings[python]',
+    'pymdown-extensions',
+    'ruff',
+]
+
+[tool.uv]
+default-groups = ['dev']
+
+[tool.ruff]
+line-length = 120
+target-version = 'py311'
+
+[tool.ruff.lint]
+extend-select = ['Q', 'RUF100', 'C90', 'I']
+extend-ignore = []
+mccabe = { max-complexity = 14 }
+pydocstyle = { convention = 'google' }
+
+[tool.codespell]
+skip = '.git,*.pdf,*.svg,uv.lock'
+ignore-words-list = 'nd'
+
+[tool.pytest.ini_options]
+addopts = "--strict-markers --tb=short --cov=src --cov-report=term-missing --cov-fail-under=70"
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+
+[project.scripts]
+acquisition = "aind_behavior_dynamic_foraging_metadata_mapper.acquisition:app"
+instrument = "aind_behavior_dynamic_foraging_metadata_mapper.instrument:app"
+mapper = "aind_behavior_dynamic_foraging_metadata_mapper.cli:main"

workspace/aind_behavior_dynamic_foraging_metadata_mapper/src/aind_behavior_dynamic_foraging_metadata_mapper/__init__.py

@@ -0,0 +1,7 @@
+from .acquisition import acqusition_from_dataset
+from .instrument import instrument_from_dataset
+
+__all__ = [
+    "acqusition_from_dataset",
+    "instrument_from_dataset",
+]