Merge pull request #466 from AllenNeuralDynamics/feat-update-aind-data-schema

Update `aind-data-schema` package version to 2.2.0

Author: bruno-f-cruz

pyproject.toml

@@ -32,7 +32,7 @@ data = ["contraqctor>=0.5.3, <0.6.0"]
 
 launcher = [
     "aind-clabe[aind-services]>=0.9",
-    "aind-data-schema>=2",
+    "aind-data-schema>=2.2",
     "aind_behavior_vr_foraging[data]",
     ]
 
@@ -42,7 +42,7 @@ dev = [
     "aind_behavior_vr_foraging[launcher]",
     "aind_behavior_vr_foraging[data]",
     "ruff",
-    "codespell"
+    "codespell",
 ]
 
 docs = [

src/aind_behavior_vr_foraging/cli.py

@@ -1,46 +1,14 @@
-import os
 import typing as t
-from pathlib import Path
 
-from pydantic import AwareDatetime, Field, RootModel
-from pydantic_settings import BaseSettings, CliApp, CliPositionalArg, CliSubCommand
+from pydantic import Field, RootModel
+from pydantic_settings import BaseSettings, CliApp, CliSubCommand
 
 from aind_behavior_vr_foraging import __semver__, regenerate
+from aind_behavior_vr_foraging.data_mappers import DataMapperCli
+from aind_behavior_vr_foraging.data_qc import DataQcCli
 from aind_behavior_vr_foraging.launcher import ClabeCli
 
 
-class DataMapperCli(BaseSettings, cli_kebab_case=True):
-    data_path: os.PathLike = Field(description="Path to the session data directory.")
-    repo_path: os.PathLike = Field(
-        default=Path("."), description="Path to the repository. By default it will use the current directory."
-    )
-    curriculum_suggestion: t.Optional[os.PathLike] = Field(
-        default=None, description="Path to curriculum suggestion file."
-    )
-    session_end_time: AwareDatetime | None = Field(
-        default=None,
-        description="End time of the session in ISO format. If not provided, will use the time the data mapping is run.",
-    )
-
-    def cli_cmd(self):
-        from aind_behavior_vr_foraging.data_mappers import cli_cmd
-
-        return cli_cmd(self)
-
-
-class DataQcCli(BaseSettings, cli_kebab_case=True):
-    data_path: CliPositionalArg[os.PathLike] = Field(description="Path to the session data directory.")
-    version: str = Field(default=__semver__, description="Version of the dataset.")
-    report_path: Path | None = Field(
-        default=None, description="Path to save the Html QC report. If not provided, report is not saved."
-    )
-
-    def cli_cmd(self):
-        from aind_behavior_vr_foraging.data_qc import cli_cmd
-
-        return cli_cmd(self)
-
-
 class VersionCli(RootModel):
     root: t.Any
 

src/aind_behavior_vr_foraging/data_mappers/__init__.py

@@ -1,31 +1,50 @@
 import logging
+import os
+import typing as t
 from pathlib import Path
-from typing import TYPE_CHECKING
 
-from ._rig import AindRigDataMapper
-from ._session import AindSessionDataMapper
+from pydantic import AwareDatetime, Field
+from pydantic_settings import BaseSettings
 
 logger = logging.getLogger(__name__)
 
-if TYPE_CHECKING:
-    from aind_behavior_vr_foraging.cli import DataMapperCli
 
-
-def cli_cmd(cli_settings: "DataMapperCli"):
-    """Generate aind-data-schema metadata for the VR Foraging dataset located at the specified path."""
-    session_mapped = AindSessionDataMapper(
-        data_path=Path(cli_settings.data_path),
-        repo_path=Path(cli_settings.repo_path),
-        session_end_time=cli_settings.session_end_time,
-    ).map()
-
-    rig_mapped = AindRigDataMapper(data_path=Path(cli_settings.data_path)).map()
-
-    assert session_mapped is not None
-
-    session_mapped.instrument_id = rig_mapped.instrument_id
-    logger.info("Writing session.json to %s", cli_settings.data_path)
-    session_mapped.write_standard_file(Path(cli_settings.data_path))
-    logger.info("Writing rig.json to %s", cli_settings.data_path)
-    rig_mapped.write_standard_file(Path(cli_settings.data_path))
-    logger.info("Mapping completed!")
+class DataMapperCli(BaseSettings, cli_kebab_case=True):
+    data_path: os.PathLike = Field(description="Path to the session data directory.")
+    repo_path: os.PathLike = Field(
+        default=Path("."), description="Path to the repository. By default it will use the current directory."
+    )
+    curriculum_suggestion: t.Optional[os.PathLike] = Field(
+        default=None, description="Path to curriculum suggestion file."
+    )
+    session_end_time: AwareDatetime | None = Field(
+        default=None,
+        description="End time of the session in ISO format. If not provided, will use the time the data mapping is run.",
+    )
+    suffix: t.Optional[str] = Field(default="vrforaging", description="Suffix to append to the output filenames.")
+
+    def cli_cmd(self):
+        """Generate aind-data-schema metadata for the VR Foraging dataset located at the specified path."""
+        from ._rig import AindInstrumentDataMapper
+        from ._session import AindAcquisitionDataMapper
+
+        session_mapper = AindAcquisitionDataMapper(
+            data_path=Path(self.data_path),
+            repo_path=Path(self.repo_path),
+            session_end_time=self.session_end_time,
+        )
+        session_mapper.map()
+
+        rig_mapper = AindInstrumentDataMapper(data_path=Path(self.data_path))
+        rig_mapper.map()
+
+        assert session_mapper.mapped is not None
+        assert rig_mapper.mapped is not None
+
+        session_mapper.mapped.instrument_id = rig_mapper.mapped.instrument_id
+        session_mapper.mapped.write_standard_file(output_directory=Path(self.data_path), filename_suffix=self.suffix)
+        rig_mapper.mapped.write_standard_file(output_directory=Path(self.data_path), filename_suffix=self.suffix)
+        logger.info(
+            "Mapping completed! Saved acquisition.json and instrument.json to %s",
+            self.data_path,
+        )

src/aind_behavior_vr_foraging/data_mappers/_rig.py

@@ -39,7 +39,7 @@ def get_spawned_device(self, name: str) -> devices.Device:
         raise ValueError(f"Device {name} not found in spawned devices of {self.device_name}")
 
 
-class AindRigDataMapper(ads.AindDataSchemaRigDataMapper):
+class AindInstrumentDataMapper(ads.AindDataSchemaRigDataMapper):
     def __init__(
         self,
         data_path: os.PathLike,
@@ -55,9 +55,6 @@ def rig_schema(self):
     def session_name(self):
         raise NotImplementedError("Method not implemented.")
 
-    def write_standard_file(self) -> None:
-        self.mapped.write_standard_file(self._data_path)
-
     def map(self) -> instrument.Instrument:
         logger.info("Mapping aind-data-schema Rig.")
         self._mapped = self._map(self._data_path)

src/aind_behavior_vr_foraging/data_mappers/_session.py

@@ -28,7 +28,7 @@
 logger = logging.getLogger(__name__)
 
 
-class AindSessionDataMapper(ads.AindDataSchemaSessionDataMapper):
+class AindAcquisitionDataMapper(ads.AindDataSchemaSessionDataMapper):
     def __init__(
         self,
         data_path: os.PathLike,
@@ -124,9 +124,6 @@ def _map(self) -> acquisition.Acquisition:
         )
         return aind_data_schema_session
 
-    def write_standard_file(self) -> None:
-        self.mapped.write_standard_file(Path(self._data_path))
-
     def _get_subject_details(self) -> acquisition.AcquisitionSubjectDetails:
         return acquisition.AcquisitionSubjectDetails(
             mouse_platform_name=TrackedDevices.WHEEL,

src/aind_behavior_vr_foraging/data_qc/__init__.py

@@ -1,20 +1,32 @@
+import logging
+import os
 from pathlib import Path
-from typing import TYPE_CHECKING
 
-if TYPE_CHECKING:
-    from ..cli import DataQcCli
+from pydantic import Field
+from pydantic_settings import BaseSettings, CliPositionalArg
 
+from aind_behavior_vr_foraging import __semver__
 
-def cli_cmd(cli_settings: "DataQcCli"):
-    """Run data quality checks on the VR Foraging dataset located at the specified path."""
-    from ..data_contract import dataset
-    from .data_qc import make_qc_runner
+logger = logging.getLogger(__name__)
 
-    vr_dataset = dataset(Path(cli_settings.data_path), cli_settings.version)
-    runner = make_qc_runner(vr_dataset)
-    results = runner.run_all_with_progress()
-    if report_path := cli_settings.report_path:
-        from contraqctor.qc.reporters import HtmlReporter
 
-        reporter = HtmlReporter(output_path=report_path)
-        reporter.report_results(results, serialize_context_exportable_obj=True)
+class DataQcCli(BaseSettings, cli_kebab_case=True):
+    data_path: CliPositionalArg[os.PathLike] = Field(description="Path to the session data directory.")
+    version: str = Field(default=__semver__, description="Version of the dataset.")
+    report_path: Path | None = Field(
+        default=None, description="Path to save the Html QC report. If not provided, report is not saved."
+    )
+
+    def cli_cmd(self):
+        """Run data quality checks on the VR Foraging dataset located at the specified path."""
+        from ..data_contract import dataset
+        from .data_qc import make_qc_runner
+
+        vr_dataset = dataset(Path(self.data_path), self.version)
+        runner = make_qc_runner(vr_dataset)
+        results = runner.run_all_with_progress()
+        if report_path := self.report_path:
+            from contraqctor.qc.reporters import HtmlReporter
+
+            reporter = HtmlReporter(output_path=report_path)
+            reporter.report_results(results, serialize_context_exportable_obj=True)

src/aind_behavior_vr_foraging/launcher.py

@@ -4,6 +4,7 @@
 
 from aind_behavior_services.calibration.aind_manipulator import ManipulatorPosition
 from aind_behavior_services.session import AindBehaviorSessionModel
+from aind_behavior_services.utils import utcnow
 from clabe import resource_monitor
 from clabe.apps import (
     AindBehaviorServicesBonsaiApp,
@@ -19,7 +20,7 @@
 from pydantic_settings import CliApp
 
 from aind_behavior_vr_foraging import data_contract
-from aind_behavior_vr_foraging.data_mappers import AindRigDataMapper, AindSessionDataMapper
+from aind_behavior_vr_foraging.data_mappers import DataMapperCli
 from aind_behavior_vr_foraging.rig import AindVrForagingRig
 from aind_behavior_vr_foraging.task_logic import AindVrForagingTaskLogic
 
@@ -77,6 +78,7 @@ async def experiment(launcher: Launcher) -> None:
 
     # Curriculum
     suggestion: CurriculumSuggestion | None = None
+    suggestion_path: Path | None = None
     if not (
         (picker.trainer_state is None)
         or (picker.trainer_state.is_on_curriculum is False)
@@ -91,20 +93,19 @@ async def experiment(launcher: Launcher) -> None:
         await trainer.run_async()
         suggestion = trainer.process_suggestion()
         # Dump suggestion for debugging (for now, but we will prob remove this later)
-        _dump_suggestion(suggestion, launcher.session_directory)
+        suggestion_path = _dump_suggestion(suggestion, launcher.session_directory)
         # Push updated trainer state back to the database
         picker.push_new_suggestion(suggestion.trainer_state)
 
     # Mappers
     assert launcher.repository.working_tree_dir is not None
-    ads_session = AindSessionDataMapper(
+
+    DataMapperCli(
         data_path=launcher.session_directory,
         repo_path=launcher.repository.working_tree_dir,  # type: ignore[arg-type]
-        curriculum_suggestion=suggestion,
-    ).map()
-    ads_session.write_standard_file(launcher.session_directory)
-    ads_rig = AindRigDataMapper(data_path=launcher.session_directory).map()
-    ads_rig.write_standard_file(launcher.session_directory)
+        curriculum_suggestion=suggestion_path,
+        session_end_time=utcnow(),
+    ).cli_cmd()
 
     # Run data qc
     if picker.ui_helper.prompt_yes_no_question("Would you like to generate a qc report?"):
@@ -142,10 +143,12 @@ async def experiment(launcher: Launcher) -> None:
     return
 
 
-def _dump_suggestion(suggestion: CurriculumSuggestion, session_directory: Path) -> None:
+def _dump_suggestion(suggestion: CurriculumSuggestion, session_directory: Path) -> Path:
     logger.info(f"Dumping curriculum suggestion to: {session_directory / 'Behavior' / 'Logs' / 'suggestion.json'}")
-    with open(session_directory / "Behavior" / "Logs" / "suggestion.json", "w", encoding="utf-8") as f:
+    suggestion_path = session_directory / "Behavior" / "Logs" / "suggestion.json"
+    with open(suggestion_path, "w", encoding="utf-8") as f:
         f.write(suggestion.model_dump_json(indent=2))
+    return suggestion_path
 
 
 class ByAnimalManipulatorModifier(ByAnimalModifier[AindVrForagingRig]):

tests/test_aind_data_mapper.py

@@ -9,7 +9,8 @@
 from aind_data_schema.core import acquisition, instrument
 from aind_data_schema.utils import compatibility_check
 
-from aind_behavior_vr_foraging.data_mappers import AindRigDataMapper, AindSessionDataMapper
+from aind_behavior_vr_foraging.data_mappers._rig import AindInstrumentDataMapper
+from aind_behavior_vr_foraging.data_mappers._session import AindAcquisitionDataMapper
 
 sys.path.append(".")
 from aind_behavior_vr_foraging.cli import DataMapperCli
@@ -41,18 +42,18 @@ def setUp(self):
         self.repo_path = Path("./")
         self.session_end_time = datetime(2023, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
 
-        self.session_mapper = AindSessionDataMapper(
+        self.session_mapper = AindAcquisitionDataMapper(
             data_path=self.data_path,
             repo_path=self.repo_path,
             session_end_time=self.session_end_time,
         )
 
-        self.rig_mapper = AindRigDataMapper(data_path=self.data_path)
+        self.rig_mapper = AindInstrumentDataMapper(data_path=self.data_path)
 
     def tearDown(self):
         self.temp_dir.cleanup()
 
-    @patch("aind_behavior_vr_foraging.data_mappers._session.AindSessionDataMapper._map")
+    @patch("aind_behavior_vr_foraging.data_mappers._session.AindAcquisitionDataMapper._map")
     def test_session_mock_map(self, mock_map):
         mock_map.return_value = MagicMock()
         result = self.session_mapper.map()
@@ -68,7 +69,7 @@ def test_session_round_trip(self):
         assert mapped is not None
         acquisition.Acquisition.model_validate_json(mapped.model_dump_json())
 
-    @patch("aind_behavior_vr_foraging.data_mappers._rig.AindRigDataMapper._map")
+    @patch("aind_behavior_vr_foraging.data_mappers._rig.AindInstrumentDataMapper._map")
     def test_rig_mock_map(self, mock_map):
         mock_map.return_value = MagicMock()
         result = self.rig_mapper.map()
@@ -91,16 +92,16 @@ def test_instrument_acquisition_compatibility(self):
         assert rig_mapped is not None
         compatibility_check.InstrumentAcquisitionCompatibility(
             instrument=rig_mapped, acquisition=session_mapped
-        ).run_compatibility_check()
+        ).run_compatibility_check(raise_for_missing_devices=True)
 
     def test_mapper_cli(self):
         DataMapperCli(
             data_path=self.data_path,
             repo_path=self.repo_path,
             session_end_time=self.session_end_time,
         ).cli_cmd()
-        instrument_path = self.data_path / "instrument.json"
-        acquisition_path = self.data_path / "acquisition.json"
+        instrument_path = self.data_path / "instrument_vrforaging.json"
+        acquisition_path = self.data_path / "acquisition_vrforaging.json"
 
         self.assertTrue(instrument_path.exists())
         self.assertTrue(acquisition_path.exists())