Merge pull request #18 from mochic/enh/aind-data-schema-upgrade

Enh/aind data schema upgrade
AllenNeuralDynamics · Dec 19, 2023 · 4f4fe01 · 4f4fe01
2 parents 11609ab + ef412a5
commit 4f4fe01
Show file tree

Hide file tree

Showing 10 changed files with 202 additions and 103 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -17,9 +17,10 @@ readme = "README.md"
 dynamic = ["version"]
 
 dependencies = [
-    "aind-data-schema==0.15.9",
+    "aind-data-schema==0.19.0",
     "aind-metadata-service[client]",
-    "scanimage-tiff-reader==1.4.1.4"
+    "scanimage-tiff-reader==1.4.1.4",
+    "pydantic-settings",
 ]
 
 [project.optional-dependencies]

diff --git a/src/aind_metadata_mapper/bergamo/session.py b/src/aind_metadata_mapper/bergamo/session.py
@@ -7,27 +7,29 @@
 import re
 import sys
 from dataclasses import dataclass
-from datetime import datetime, time
+from datetime import datetime
 from os import PathLike
 from pathlib import Path
 from typing import Dict, List, Tuple, Union
 
 import numpy as np
-from aind_data_schema.session import (
-    Detector,
+
+# from aind_data_schema.models.devices import Detector, Laser
+from aind_data_schema.core.session import (
+    DetectorConfig,
     FieldOfView,
-    Laser,
+    LaserConfig,
     Modality,
     Session,
     Stream,
 )
-from aind_data_schema.stimulus import (
+from aind_data_schema.models.stimulus import (
     PhotoStimulation,
     PhotoStimulationGroup,
     StimulusEpoch,
 )
-from aind_data_schema.utils.units import PowerUnit, SizeUnit
-from pydantic import BaseSettings, Extra
+from aind_data_schema.models.units import PowerUnit, SizeUnit
+from pydantic_settings import BaseSettings
 from ScanImageTiffReader import ScanImageTiffReader
 
 from aind_metadata_mapper.core import BaseEtl
@@ -44,8 +46,12 @@ class UserSettings(BaseSettings):
     session_end_time: datetime
     stream_start_time: datetime
     stream_end_time: datetime
-    stimulus_start_time: time
-    stimulus_end_time: time
+    stimulus_start_time: datetime
+    stimulus_end_time: datetime
+
+    # TODO: Look into whether defaults can be set for these fields
+    mouse_platform_name: str
+    active_mouse_platform: bool
 
     # Data that might change but can have default values
     session_type: str = "BCI"
@@ -79,7 +85,6 @@ def num_of_photo_stim_groups(self):
     class Config:
         """Config to set env var prefix to BERGAMO"""
 
-        extra = Extra.forbid
         env_prefix = "BERGAMO_"
 
 
@@ -390,12 +395,14 @@ def _transform(self, extracted_source: RawImageInfo) -> Session:
         ]
 
         data_stream = Stream(
+            mouse_platform_name=self.user_settings.mouse_platform_name,
+            active_mouse_platform=self.user_settings.active_mouse_platform,
             stream_start_time=self.user_settings.stream_start_time,
             stream_end_time=self.user_settings.stream_end_time,
             stream_modalities=[Modality.POPHYS],
             camera_names=list(self.user_settings.camera_names),
             light_sources=[
-                Laser(
+                LaserConfig(
                     name=self.user_settings.laser_a_name,
                     wavelength=self.user_settings.laser_a_wavelength,
                     wavelength_unit=self.user_settings.laser_a_wavelength_unit,
@@ -406,7 +413,7 @@ def _transform(self, extracted_source: RawImageInfo) -> Session:
                 ),
             ],
             detectors=[
-                Detector(
+                DetectorConfig(
                     name=self.user_settings.detector_a_name,
                     exposure_time=self.user_settings.detector_a_exposure_time,
                     trigger_type=self.user_settings.detector_a_trigger_type,

diff --git a/src/aind_metadata_mapper/core.py b/src/aind_metadata_mapper/core.py
@@ -8,7 +8,7 @@
 from typing import Any, Union
 
 from aind_data_schema.base import AindCoreModel
-from pydantic import validate_model
+from pydantic import ValidationError
 
 
 class BaseEtl(ABC):
@@ -84,22 +84,14 @@ def _run_validation_check(model_instance: AindCoreModel) -> None:
           Contents from the service response.
         """
         try:
-            *_, validation_error = validate_model(
-                model_instance.__class__, model_instance.__dict__
-            )
-            if validation_error:
-                logging.warning(
-                    f"Validation errors were found. This may be due to "
-                    f"mismatched versions or data not found in the "
-                    f"databases. Error: {validation_error}"
-                )
-            else:
-                logging.debug("No validation errors detected.")
-        except AttributeError as e:
+            model_instance.model_validate(model_instance.__dict__)
+            logging.debug("No validation errors detected.")
+        except ValidationError:
             logging.warning(
-                f"An error was detected during the validation check. "
-                f"This may be due to mismatched versions of aind-data-schema "
-                f"from the source and this local package. {repr(e)}"
+                "Validation errors were found. This may be due to "
+                "mismatched versions or data not found in the "
+                "databases.",
+                exc_info=True,
             )
 
     def run_job(self) -> None:

diff --git a/src/aind_metadata_mapper/procedures.py b/src/aind_metadata_mapper/procedures.py
@@ -6,7 +6,7 @@
 from os import PathLike
 from pathlib import Path
 
-from aind_data_schema.procedures import Procedures
+from aind_data_schema.core.procedures import Procedures
 from aind_metadata_service.client import AindMetadataServiceClient, StatusCodes
 from requests import Response
 
@@ -74,7 +74,7 @@ def _transform(self, extracted_source: Response) -> Procedures:
             case _:
                 raise Exception("An error occurred connecting to the server.")
 
-        procedures = Procedures.construct(**contents)
+        procedures = Procedures.model_construct(**contents)
         return procedures
 
     @classmethod

diff --git a/src/aind_metadata_mapper/subject.py b/src/aind_metadata_mapper/subject.py
@@ -6,7 +6,7 @@
 from os import PathLike
 from pathlib import Path
 
-from aind_data_schema.subject import Subject
+from aind_data_schema.core.subject import Subject
 from aind_metadata_service.client import AindMetadataServiceClient, StatusCodes
 from requests import Response
 
@@ -75,7 +75,7 @@ def _transform(self, extracted_source: Response) -> Subject:
             case _:
                 raise Exception("An error occurred connecting to the server.")
 
-        subject = Subject.construct(**contents)
+        subject = Subject.model_construct(**contents)
         return subject
 
     @classmethod

diff --git a/tests/resources/bergamo/expected_session.json b/tests/resources/bergamo/expected_session.json
@@ -1,6 +1,6 @@
 {
-   "describedBy": "https://raw.githubusercontent.com/AllenNeuralDynamics/aind-data-schema/main/src/aind_data_schema/session.py",
-   "schema_version": "0.0.1",
+   "describedBy": "https://raw.githubusercontent.com/AllenNeuralDynamics/aind-data-schema/main/src/aind_data_schema/core/session.py",
+   "schema_version": "0.1.0",
    "experimenter_full_name": [
       "John Smith",
       "Jane Smith"
@@ -10,23 +10,17 @@
    "session_type": "BCI",
    "iacuc_protocol": "2115",
    "rig_id": "Bergamo photostim.",
-   "calibrations": null,
-   "maintenance": null,
-   "subject_id": 12345,
+   "calibrations": [],
+   "maintenance": [],
+   "subject_id": "12345",
    "animal_weight_prior": null,
    "animal_weight_post": null,
    "weight_unit": "gram",
    "data_streams": [
       {
          "stream_start_time": "2023-10-10T15:00:00",
          "stream_end_time": "2023-10-10T16:00:00",
-         "stream_modalities": [
-            {
-               "name": "Planar optical physiology",
-               "abbreviation": "ophys"
-            }
-         ],
-         "daq_names": null,
+         "daq_names": [],
          "camera_names": [
             "Side Camera"
          ],
@@ -36,43 +30,55 @@
                "name": "Laser A",
                "wavelength": 920,
                "wavelength_unit": "nanometer",
-               "excitation_power": 15,
+               "excitation_power": "15",
                "excitation_power_unit": "percent"
             }
          ],
-         "ephys_modules": null,
-         "manipulator_modules": null,
+         "ephys_modules": [],
+         "stick_microscopes": [],
+         "manipulator_modules": [],
          "detectors": [
             {
                "name": "PMT A",
-               "exposure_time": 0.1,
+               "exposure_time": "0.1",
                "exposure_time_unit": "millisecond",
                "trigger_type": "Internal"
             }
          ],
-         "fiber_photometry_assemblies": null,
+         "fiber_connections": [],
+         "fiber_modules": [],
          "ophys_fovs": [
             {
                "index": 0,
                "imaging_depth": 150,
                "imaging_depth_unit": "micrometer",
                "targeted_structure": "M1",
-               "fov_coordinate_ml": 1.5,
-               "fov_coordinate_ap": 1.5,
+               "fov_coordinate_ml": "1.5",
+               "fov_coordinate_ap": "1.5",
                "fov_coordinate_unit": "micrometer",
                "fov_reference": "Bregma",
                "fov_width": 512,
                "fov_height": 512,
                "fov_size_unit": "pixel",
                "magnification": "16x",
-               "fov_scale_factor": 1.2,
+               "fov_scale_factor": "1.2",
                "fov_scale_factor_unit": "um/pixel",
-               "frame_rate": 30.0119,
-               "frame_rate_unit": "hertz"
+               "frame_rate": "30.0119",
+               "frame_rate_unit": "hertz",
+               "coupled_fov_index": null
             }
          ],
+         "slap_fovs": null,
          "stack_parameters": null,
-         "stimulus_device_names": null,
+         "stimulus_device_names": [],
+         "mouse_platform_name": "some_mouse_platform_name",
+         "active_mouse_platform": true,
+         "stream_modalities": [
+            {
+               "name": "Planar optical physiology",
+               "abbreviation": "ophys"
+            }
+         ],
          "notes": null
       }
    ],
@@ -86,42 +92,43 @@
                {
                   "group_index": 0,
                   "number_of_neurons": 10,
-                  "stimulation_laser_power": 20,
+                  "stimulation_laser_power": "20",
                   "stimulation_laser_power_unit": "milliwatt",
                   "number_trials": 5,
                   "number_spirals": 10,
-                  "spiral_duration": 0.01,
+                  "spiral_duration": "0.01",
                   "spiral_duration_unit": "second",
-                  "inter_spiral_interval": 0.001,
+                  "inter_spiral_interval": "0.001",
                   "inter_spiral_interval_unit": "second",
-                  "other_parameters": null,
+                  "other_parameters": {},
                   "notes": null
                },
                {
                   "group_index": 0,
                   "number_of_neurons": 10,
-                  "stimulation_laser_power": 20,
+                  "stimulation_laser_power": "20",
                   "stimulation_laser_power_unit": "milliwatt",
                   "number_trials": 5,
                   "number_spirals": 10,
-                  "spiral_duration": 0.01,
+                  "spiral_duration": "0.01",
                   "spiral_duration_unit": "second",
-                  "inter_spiral_interval": 0.001,
+                  "inter_spiral_interval": "0.001",
                   "inter_spiral_interval_unit": "second",
-                  "other_parameters": null,
+                  "other_parameters": {},
                   "notes": null
                }
             ],
-            "inter_trial_interval": 10,
+            "inter_trial_interval": "10",
             "inter_trial_interval_unit": "second",
-            "other_parameters": null,
+            "other_parameters": {},
             "notes": null
          },
-         "stimulus_start_time": "15:15:00",
-         "stimulus_end_time": "15:45:00"
+         "stimulus_start_time": "2023-10-10T15:15:00",
+         "stimulus_end_time": "2023-10-10T15:45:00"
       }
    ],
    "reward_delivery": null,
-   "stick_microscopes": null,
+   "reward_consumed_total": null,
+   "reward_consumed_unit": "microliter",
    "notes": null
 }