1057 brain atlas model for targeted structures (#1149)

* feat: replacing str structures with CCFStructure.ONE_OF

* feat: add validator to coerce str -> CCFStructure

* chore: docstrings

* tests: fixing tests that had fake CCF areas

* chore: lint

* refactor: add ValueError warning when area doesn't exist

* tests: adding coverage on validator

* chore: docstring

examples/bergamo_ophys_session.json

@@ -53,7 +53,12 @@
                "index": 0,
                "imaging_depth": 150,
                "imaging_depth_unit": "micrometer",
-               "targeted_structure": "M1",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Primary motor area",
+                  "acronym": "MOp",
+                  "id": "985"
+               },
                "fov_coordinate_ml": "1.5",
                "fov_coordinate_ap": "1.5",
                "fov_coordinate_unit": "micrometer",

examples/bergamo_ophys_session.py

@@ -55,7 +55,7 @@
                 FieldOfView(
                     index=0,
                     imaging_depth=150,
-                    targeted_structure="M1",
+                    targeted_structure="MOp",
                     fov_coordinate_ml=1.5,
                     fov_coordinate_ap=1.5,
                     fov_reference="Bregma",

examples/ephys_session.json

@@ -37,7 +37,12 @@
                "coordinate_transform": "behavior/calibration_info_np2_2023_04_24.npy",
                "calibration_date": "2023-04-25T00:00:00Z",
                "notes": "Moved Y to avoid blood vessel, X to avoid edge. Mouse made some noise during the recording with a sudden shift in signals. Lots of motion. Maybe some implant motion.",
-               "primary_targeted_structure": "LGd",
+               "primary_targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Dorsal part of the lateral geniculate complex",
+                  "acronym": "LGd",
+                  "id": "170"
+               },
                "other_targeted_structure": null,
                "targeted_ccf_coordinates": [
                   {
@@ -70,7 +75,12 @@
                "coordinate_transform": "behavior/calibration_info_np2_2023_04_24.py",
                "calibration_date": "2023-04-25T00:00:00Z",
                "notes": "Trouble penetrating. Lots of compression, needed to move probe. Small amount of surface bleeding/bruising. Initial Target: X;10070.3\tY:7476.6",
-               "primary_targeted_structure": "LC",
+               "primary_targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Locus ceruleus",
+                  "acronym": "LC",
+                  "id": "147"
+               },
                "other_targeted_structure": null,
                "targeted_ccf_coordinates": [
                   {
@@ -172,7 +182,12 @@
                "coordinate_transform": "behavior/calibration_info_np2_2023_04_24.npy",
                "calibration_date": "2023-04-25T00:00:00Z",
                "notes": "Moved Y to avoid blood vessel, X to avoid edge. Mouse made some noise during the recording with a sudden shift in signals. Lots of motion. Maybe some implant motion.",
-               "primary_targeted_structure": "LGd",
+               "primary_targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Dorsal part of the lateral geniculate complex",
+                  "acronym": "LGd",
+                  "id": "170"
+               },
                "other_targeted_structure": null,
                "targeted_ccf_coordinates": [
                   {
@@ -205,7 +220,12 @@
                "coordinate_transform": "behavior/calibration_info_np2_2023_04_24.py",
                "calibration_date": "2023-04-25T00:00:00Z",
                "notes": "Trouble penetrating. Lots of compression, needed to move probe. Small amount of surface bleeding/bruising. Initial Target: X;10070.3\tY:7476.6",
-               "primary_targeted_structure": "LC",
+               "primary_targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Locus ceruleus",
+                  "acronym": "LC",
+                  "id": "147"
+               },
                "other_targeted_structure": null,
                "targeted_ccf_coordinates": [
                   {

examples/multiplane_ophys_session.json

@@ -50,7 +50,12 @@
                "index": 0,
                "imaging_depth": 190,
                "imaging_depth_unit": "micrometer",
-               "targeted_structure": "VISp",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Primary visual area",
+                  "acronym": "VISp",
+                  "id": "385"
+               },
                "fov_coordinate_ml": "1.5",
                "fov_coordinate_ap": "1.5",
                "fov_coordinate_unit": "micrometer",
@@ -76,7 +81,12 @@
                "index": 1,
                "imaging_depth": 232,
                "imaging_depth_unit": "micrometer",
-               "targeted_structure": "VISp",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Primary visual area",
+                  "acronym": "VISp",
+                  "id": "385"
+               },
                "fov_coordinate_ml": "1.5",
                "fov_coordinate_ap": "1.5",
                "fov_coordinate_unit": "micrometer",
@@ -102,7 +112,12 @@
                "index": 2,
                "imaging_depth": 136,
                "imaging_depth_unit": "micrometer",
-               "targeted_structure": "VISp",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Primary visual area",
+                  "acronym": "VISp",
+                  "id": "385"
+               },
                "fov_coordinate_ml": "1.5",
                "fov_coordinate_ap": "1.5",
                "fov_coordinate_unit": "micrometer",
@@ -128,7 +143,12 @@
                "index": 3,
                "imaging_depth": 282,
                "imaging_depth_unit": "micrometer",
-               "targeted_structure": "VISp",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Primary visual area",
+                  "acronym": "VISp",
+                  "id": "385"
+               },
                "fov_coordinate_ml": "1.5",
                "fov_coordinate_ap": "1.5",
                "fov_coordinate_unit": "micrometer",
@@ -154,7 +174,12 @@
                "index": 4,
                "imaging_depth": 72,
                "imaging_depth_unit": "micrometer",
-               "targeted_structure": "VISp",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Primary visual area",
+                  "acronym": "VISp",
+                  "id": "385"
+               },
                "fov_coordinate_ml": "1.5",
                "fov_coordinate_ap": "1.5",
                "fov_coordinate_unit": "micrometer",
@@ -180,7 +205,12 @@
                "index": 5,
                "imaging_depth": 326,
                "imaging_depth_unit": "micrometer",
-               "targeted_structure": "VISp",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Primary visual area",
+                  "acronym": "VISp",
+                  "id": "385"
+               },
                "fov_coordinate_ml": "1.5",
                "fov_coordinate_ap": "1.5",
                "fov_coordinate_unit": "micrometer",
@@ -206,7 +236,12 @@
                "index": 6,
                "imaging_depth": 30,
                "imaging_depth_unit": "micrometer",
-               "targeted_structure": "VISp",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Primary visual area",
+                  "acronym": "VISp",
+                  "id": "385"
+               },
                "fov_coordinate_ml": "1.5",
                "fov_coordinate_ap": "1.5",
                "fov_coordinate_unit": "micrometer",
@@ -232,7 +267,12 @@
                "index": 7,
                "imaging_depth": 364,
                "imaging_depth_unit": "micrometer",
-               "targeted_structure": "VISp",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Primary visual area",
+                  "acronym": "VISp",
+                  "id": "385"
+               },
                "fov_coordinate_ml": "1.5",
                "fov_coordinate_ap": "1.5",
                "fov_coordinate_unit": "micrometer",

examples/ophys_procedures.json

@@ -65,7 +65,12 @@
                "bregma_to_lambda_unit": "millimeter",
                "injection_angle": "0",
                "injection_angle_unit": "degrees",
-               "targeted_structure": "VTA",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Ventral tegmental area",
+                  "acronym": "VTA",
+                  "id": "749"
+               },
                "injection_hemisphere": "Left",
                "procedure_type": "Nanoject injection",
                "injection_volume": [
@@ -96,7 +101,12 @@
                         "total_length": "0.5",
                         "length_unit": "millimeter"
                      },
-                     "targeted_structure": "VTA",
+                     "targeted_structure": {
+                        "atlas": "CCFv3",
+                        "name": "Ventral tegmental area",
+                        "acronym": "VTA",
+                        "id": "749"
+                     },
                      "stereotactic_coordinate_ap": "-3.05",
                      "stereotactic_coordinate_ml": "-0.6",
                      "stereotactic_coordinate_dv": "-4",

examples/procedures.json

@@ -68,7 +68,12 @@
                "bregma_to_lambda_unit": "millimeter",
                "injection_angle": "10",
                "injection_angle_unit": "degrees",
-               "targeted_structure": "VISp",
+               "targeted_structure": {
+                  "atlas": "CCFv3",
+                  "name": "Primary visual area",
+                  "acronym": "VISp",
+                  "id": "385"
+               },
                "injection_hemisphere": "Left",
                "procedure_type": "Nanoject injection",
                "injection_volume": [

src/aind_data_schema/base.py

@@ -19,6 +19,7 @@
 )
 from pydantic.functional_validators import WrapValidator
 from typing_extensions import Annotated
+from aind_data_schema_models.brain_atlas import CCFStructure
 
 
 def _coerce_naive_datetime(v: Any, handler: ValidatorFunctionWrapHandler) -> AwareDatetime:
@@ -90,10 +91,23 @@ def validate_fieldnames(self):
 
 
 class AindModel(BaseModel, Generic[AindGenericType]):
-    """BaseModel that disallows extra fields"""
+    """BaseModel that disallows extra fields
+
+    Also performs validation checks / coercion / upgrades where necessary
+    """
 
     model_config = ConfigDict(extra="forbid", use_enum_values=True)
 
+    @model_validator(mode="before")
+    def coerce_targeted_structures(cls, values):
+        """If a user passes a targeted_structure as a str, convert to CCFStructure"""
+        for field_name, value in values.items():
+            if "targeted_structure" in field_name and isinstance(value, str):
+                if not hasattr(CCFStructure, value.upper()):
+                    raise ValueError(f"{value} is not a valid CCF structure")
+                values[field_name] = getattr(CCFStructure, value.upper())
+        return values
+
 
 class AindCoreModel(AindModel):
     """Generic base class to hold common fields/validators/etc for all basic AIND schema"""

src/aind_data_schema/core/procedures.py

@@ -20,6 +20,7 @@
     VolumeUnit,
     create_unit_with_value,
 )
+from aind_data_schema_models.brain_atlas import CCFStructure
 from pydantic import Field, SkipValidation, field_serializer, field_validator, model_validator
 from pydantic_core.core_schema import ValidationInfo
 from typing_extensions import Annotated
@@ -222,7 +223,7 @@ class Sectioning(AindModel):
     section_distance_unit: SizeUnit = Field(default=SizeUnit.MM, title="Distance unit")
     reference_location: CoordinateReferenceLocation = Field(..., title="Reference location for distance measurement")
     section_strategy: SectionStrategy = Field(..., title="Slice strategy")
-    targeted_structure: str = Field(..., title="Targeted structure", description="Use Allen Brain Atlas Ontology")
+    targeted_structure: CCFStructure.ONE_OF = Field(..., title="Targeted structure")
 
     @field_validator("output_specimen_ids")
     def check_output_id_length(cls, v, info: ValidationInfo):
@@ -434,7 +435,7 @@ class BrainInjection(Injection):
     bregma_to_lambda_unit: SizeUnit = Field(default=SizeUnit.MM, title="Bregma to lambda unit")
     injection_angle: Decimal = Field(..., title="Injection angle (deg)")
     injection_angle_unit: AngleUnit = Field(default=AngleUnit.DEG, title="Injection angle unit")
-    targeted_structure: Optional[str] = Field(default=None, title="Injection targeted brain structure")
+    targeted_structure: Optional[CCFStructure.ONE_OF] = Field(default=None, title="Injection targeted brain structure")
     injection_hemisphere: Optional[Side] = Field(default=None, title="Injection hemisphere")
 
 
@@ -509,7 +510,7 @@ class OphysProbe(AindModel):
     """Description of an implanted ophys probe"""
 
     ophys_probe: FiberProbe = Field(..., title="Fiber probe")
-    targeted_structure: str = Field(..., title="Targeted structure")
+    targeted_structure: CCFStructure.ONE_OF = Field(..., title="Targeted structure")
     stereotactic_coordinate_ap: Decimal = Field(..., title="Stereotactic coordinate A/P (mm)")
     stereotactic_coordinate_ml: Decimal = Field(..., title="Stereotactic coordinate M/L (mm)")
     stereotactic_coordinate_dv: Decimal = Field(

src/aind_data_schema/core/session.py

@@ -17,6 +17,7 @@
     TimeUnit,
     VolumeUnit,
 )
+from aind_data_schema_models.brain_atlas import CCFStructure
 from pydantic import Field, SkipValidation, field_validator, model_validator
 from pydantic_core.core_schema import ValidationInfo
 from typing_extensions import Annotated
@@ -95,7 +96,7 @@ class FieldOfView(AindModel):
     index: int = Field(..., title="Index")
     imaging_depth: int = Field(..., title="Imaging depth (um)")
     imaging_depth_unit: SizeUnit = Field(default=SizeUnit.UM, title="Imaging depth unit")
-    targeted_structure: str = Field(..., title="Targeted structure")
+    targeted_structure: CCFStructure.ONE_OF = Field(..., title="Targeted structure")
     fov_coordinate_ml: Decimal = Field(..., title="FOV coordinate ML")
     fov_coordinate_ap: Decimal = Field(..., title="FOV coordinate AP")
     fov_coordinate_unit: SizeUnit = Field(default=SizeUnit.UM, title="FOV coordinate unit")
@@ -162,7 +163,7 @@ class Stack(AindModel):
     fov_scale_factor_unit: str = Field(default="um/pixel", title="FOV scale factor unit")
     frame_rate: Decimal = Field(..., title="Frame rate (Hz)")
     frame_rate_unit: FrequencyUnit = Field(default=FrequencyUnit.HZ, title="Frame rate unit")
-    targeted_structure: Optional[str] = Field(default=None, title="Targeted structure")
+    targeted_structure: Optional[CCFStructure.ONE_OF] = Field(default=None, title="Targeted structure")
 
 
 class SlapSessionType(str, Enum):
@@ -207,8 +208,10 @@ class DomeModule(AindModel):
 class ManipulatorModule(DomeModule):
     """A dome module connected to a 3-axis manipulator"""
 
-    primary_targeted_structure: str = Field(..., title="Targeted structure")
-    other_targeted_structure: Optional[List[str]] = Field(default=None, title="Other targeted structure")
+    primary_targeted_structure: CCFStructure.ONE_OF = Field(..., title="Targeted structure")
+    other_targeted_structure: Optional[List[CCFStructure.ONE_OF]] = Field(
+        default=None, title="Other targeted structure"
+    )
     targeted_ccf_coordinates: List[CcfCoords] = Field(
         default=[],
         title="Targeted CCF coordinates",

tests/test_base.py

@@ -8,8 +8,9 @@
 
 from pydantic import ValidationError, create_model
 
-from aind_data_schema.base import AindGeneric, AwareDatetimeWithDefault, is_dict_corrupt
+from aind_data_schema.base import AindGeneric, AwareDatetimeWithDefault, is_dict_corrupt, AindModel
 from aind_data_schema.core.subject import Subject
+from aind_data_schema_models.brain_atlas import CCFStructure
 
 
 class BaseTests(unittest.TestCase):
@@ -112,6 +113,15 @@ def test_aind_generic_validate_fieldnames(self):
                 AindGeneric.model_validate(params)
             self.assertIn(expected_error, repr(e.exception))
 
+    def test_ccf_validator(self):
+        """Tests that CCFStructure validator works"""
+
+        class StructureModel(AindModel):
+            """Test model with a targeted_structure"""
+            targeted_structure: CCFStructure.ONE_OF
+
+        self.assertRaises(ValueError, StructureModel, targeted_structure="invalid")
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_procedures.py

@@ -194,7 +194,7 @@ def test_injection_material_check(self):
                             injection_angle=1,
                             injection_volume=[1],
                             recovery_time=10,
-                            targeted_structure="VISp6",
+                            targeted_structure="VISpl6a",
                         ),
                         FiberImplant(
                             protocol_id="dx.doi.org/120.123/fkjd",

tests/test_rig_session_compatibility.py

@@ -799,7 +799,7 @@ def read_json(filepath: Path) -> dict:
                             assembly_name="Fiber Module A",
                             arc_angle=30,
                             module_angle=180,
-                            primary_targeted_structure="Structure A",
+                            primary_targeted_structure="VISp",
                             manipulator_coordinates=Coordinates3d(x=30.5, y=70, z=180),
                         )
                     ],

tests/test_session.py

@@ -64,7 +64,7 @@ def test_constructors(self):
                             assembly_name="Ephys_assemblyA",
                             arc_angle=0,
                             module_angle=10,
-                            primary_targeted_structure="VISlm",
+                            primary_targeted_structure="VISl",
                             targeted_ccf_coordinates=[CcfCoords(ml="1", ap="1", dv="1")],
                             manipulator_coordinates=Coordinates3d(x="1", y="1", z="1"),
                         ),