Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 9 additions & 8 deletions src/aind_behavior_vr_foraging/data_mappers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import pydantic_settings
from aind_behavior_services.session import AindBehaviorSessionModel
from aind_behavior_services.utils import model_from_json_file
from clabe.apps import CurriculumSuggestion
from clabe.apps import BonsaiAppSettings, CurriculumSuggestion
from git import Repo

from aind_behavior_vr_foraging.rig import AindVrForagingRig
Expand All @@ -21,10 +21,6 @@

class DataMapperCli(pydantic_settings.BaseSettings, cli_kebab_case=True):
data_path: os.PathLike = pydantic.Field(description="Path to the session data directory.")
db_root: os.PathLike = pydantic.Field(
default=Path(r"\\allen\aind\scratch\AindBehavior.db\AindVrForaging"),
description="Root directory for the database for additional metadata.",
)
repo_path: os.PathLike = pydantic.Field(
default=Path("."), description="Path to the repository. By default it will use the current directory."
)
Expand All @@ -45,15 +41,20 @@ def cli_cmd(self):
curriculum_suggestion = None

repo = Repo(self.repo_path)
settings = BonsaiAppSettings(
workflow=Path(repo.working_dir) / "src" / "main.bonsai",
executable=Path(repo.working_dir) / "bonsai/bonsai.exe",
)

session_mapped = AindSessionDataMapper(
session=session,
rig=rig,
task_logic_model=task_logic,
task_logic=task_logic,
repository=repo,
script_path=Path("./src/main.bonsai"),
bonsai_app_settings=settings,
curriculum_suggestion=curriculum_suggestion,
).map()
rig_mapped = AindRigDataMapper(rig_schema_filename=f"{rig.rig_name}.json", db_root=Path(self.db_root)).map()
rig_mapped = AindRigDataMapper(rig=rig).map()

assert session.session_name is not None
assert session_mapped is not None
Expand Down
11 changes: 6 additions & 5 deletions src/aind_behavior_vr_foraging/data_mappers/_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from aind_data_schema.core import acquisition
from aind_data_schema_models import units
from aind_data_schema_models.modalities import Modality
from clabe.apps import CurriculumSuggestion
from clabe.apps import BonsaiAppSettings, CurriculumSuggestion
from clabe.data_mapper import aind_data_schema as ads
from clabe.data_mapper import helpers as data_mapper_helpers

Expand All @@ -32,8 +32,8 @@ def __init__(
session: AindBehaviorSessionModel,
rig: AindVrForagingRig,
task_logic: AindVrForagingTaskLogic,
bonsai_app_settings: BonsaiAppSettings = BonsaiAppSettings(workflow=Path("./src/main.bonsai")),
repository: Union[os.PathLike, git.Repo] = Path("."),
script_path: os.PathLike = Path("./src/main.bonsai"),
session_end_time: Optional[datetime.datetime] = None,
curriculum_suggestion: Optional[CurriculumSuggestion] = None,
):
Expand All @@ -43,7 +43,7 @@ def __init__(
self.repository = repository
if isinstance(self.repository, os.PathLike | str):
self.repository = git.Repo(Path(self.repository))
self.script_path = script_path
self.bonsai_app = bonsai_app_settings
self._session_end_time = session_end_time
self._mapped: Optional[acquisition.Acquisition] = None
self.curriculum = curriculum_suggestion
Expand Down Expand Up @@ -260,7 +260,8 @@ def _map_compression(compression: AbsRig.cameras.VideoWriter) -> acquisition.Cod
return list(map(_map_camera, cameras.keys(), cameras.values()))

def _get_bonsai_as_code(self) -> acquisition.Code:
bonsai_env = data_mapper_helpers.snapshot_bonsai_environment(Path("./bonsai/bonsai.config"))
bonsai_folder = Path(self.bonsai_app.executable).parent
bonsai_env = data_mapper_helpers.snapshot_bonsai_environment(bonsai_folder / "bonsai.config")
bonsai_version = bonsai_env.get("Bonsai", "unknown")
assert isinstance(self.repository, git.Repo)

Expand All @@ -270,7 +271,7 @@ def _get_bonsai_as_code(self) -> acquisition.Code:
version=self.repository.head.commit.hexsha,
language="Bonsai",
language_version=bonsai_version,
run_script=Path("./src/main.bonsai"),
run_script=Path(self.bonsai_app.workflow),
)

def _get_python_as_code(self) -> acquisition.Code:
Expand Down
6 changes: 5 additions & 1 deletion src/aind_behavior_vr_foraging/launcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,11 @@ def experiment(launcher: Launcher) -> None:

# Mappers
ads_session = AindSessionDataMapper(
rig=rig, session=session, task_logic=task_logic, curriculum_suggestion=suggestion
rig=rig,
session=session,
task_logic=task_logic,
curriculum_suggestion=suggestion,
bonsai_app_settings=bonsai_app.settings,
).map()
ads_session.write_standard_file(launcher.session_directory)
ads_rig = AindRigDataMapper(rig=rig).map()
Expand Down
9 changes: 6 additions & 3 deletions tests/test_aind_data_mapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

from aind_data_schema.core import acquisition, instrument
from aind_data_schema.utils import compatibility_check
from clabe.apps import BonsaiAppSettings
from git import Repo

from aind_behavior_vr_foraging.data_mappers import AindRigDataMapper, AindSessionDataMapper
Expand All @@ -22,7 +23,7 @@ def setUp(self):
self.rig = rig
self.task_logic = task_logic
self.repository = Repo(Path("./"))
self.script_path = Path("./src/main.bonsai")
self.bonsai_app_settings = BonsaiAppSettings(workflow=Path("./src/main.bonsai"))
self.session_end_time = datetime.now()
self.session_directory = None

Expand All @@ -31,7 +32,7 @@ def setUp(self):
rig=self.rig,
task_logic=self.task_logic,
repository=self.repository,
script_path=self.script_path,
bonsai_app_settings=self.bonsai_app_settings,
session_end_time=self.session_end_time,
)

Expand Down Expand Up @@ -84,12 +85,14 @@ def setUp(self):
self.rig_mapper = AindRigDataMapper(
rig=self.rig,
)
self.bonsai_app_settings = BonsaiAppSettings(workflow=Path("./src/main.bonsai"))

self.session_mapper = AindSessionDataMapper(
session=self.session,
rig=self.rig,
task_logic=self.task_logic,
repository=Repo(Path("./")),
script_path=Path("./src/main.bonsai"),
bonsai_app_settings=self.bonsai_app_settings,
session_end_time=datetime.now(),
)

Expand Down