refactor: replace all str structures with CCFStructure

AllenNeuralDynamics · Jan 10, 2025 · 546eafc · 546eafc
1 parent 4e5b7c4
commit 546eafc
Show file tree

Hide file tree

Showing 11 changed files with 38 additions and 45 deletions.
diff --git a/examples/bergamo_ophys_session.py b/examples/bergamo_ophys_session.py
@@ -15,6 +15,7 @@
     StimulusModality,
     Stream,
 )
+from aind_data_schema_models.brain_atlas import CCFStructure
 
 # If a timezone isn't specified, the timezone of the computer running this
 # script will be used as default
@@ -56,7 +57,7 @@
                 FieldOfView(
                     index=0,
                     imaging_depth=150,
-                    targeted_structure="MOp",
+                    targeted_structure=CCFStructure.MOP,
                     fov_coordinate_ml=1.5,
                     fov_coordinate_ap=1.5,
                     fov_reference="Bregma",

diff --git a/examples/ephys_session.py b/examples/ephys_session.py
@@ -16,6 +16,7 @@
     Stream,
     VisualStimulation,
 )
+from aind_data_schema_models.brain_atlas import CCFStructure
 
 session = Session(
     experimenter_full_name=["Max Quibble", "Finn Tickle"],
@@ -122,7 +123,7 @@
                     arc_angle=5.2,
                     module_angle=8,
                     coordinate_transform="behavior/calibration_info_np2_2023_04_24.npy",
-                    primary_targeted_structure="LGd",
+                    primary_targeted_structure=CCFStructure.LGD,
                     manipulator_coordinates=Coordinates3d(x=8422, y=4205, z=11087.5),
                     calibration_date=datetime(year=2023, month=4, day=25, tzinfo=timezone.utc),
                     notes=(
@@ -137,7 +138,7 @@
                     targeted_ccf_coordinates=[CcfCoords(ml=6637.28, ap=4265.02, dv=10707.35)],
                     assembly_name="Ephys_assemblyB",
                     coordinate_transform="behavior/calibration_info_np2_2023_04_24.py",
-                    primary_targeted_structure="LC",
+                    primary_targeted_structure=CCFStructure.LC,
                     manipulator_coordinates=Coordinates3d(x=9015, y=7144, z=13262),
                     calibration_date=datetime(year=2023, month=4, day=25, tzinfo=timezone.utc),
                     notes=(
@@ -170,7 +171,7 @@
                     targeted_ccf_coordinates=[CcfCoords(ml=8150, ap=3250, dv=7800)],
                     assembly_name="Ephys_assemblyA",
                     coordinate_transform="behavior/calibration_info_np2_2023_04_24.npy",
-                    primary_targeted_structure="LGd",
+                    primary_targeted_structure=CCFStructure.LGD,
                     manipulator_coordinates=Coordinates3d(x=8422, y=4205, z=11087.5),
                     calibration_date=datetime(year=2023, month=4, day=25, tzinfo=timezone.utc),
                     notes=(
@@ -185,7 +186,7 @@
                     targeted_ccf_coordinates=[CcfCoords(ml=6637.28, ap=4265.02, dv=10707.35)],
                     assembly_name="Ephys_assemblyB",
                     coordinate_transform="behavior/calibration_info_np2_2023_04_24.py",
-                    primary_targeted_structure="LC",
+                    primary_targeted_structure=CCFStructure.LC,
                     manipulator_coordinates=Coordinates3d(x=9015, y=7144, z=13262),
                     calibration_date=datetime(year=2023, month=4, day=25, tzinfo=timezone.utc),
                     notes=(

diff --git a/examples/exaspim_acquisition.py b/examples/exaspim_acquisition.py
@@ -44,10 +44,7 @@
             calibration_date=t,
             device_name="Laser_1",
             description="Laser power calibration",
-            input={
-                "power_setting": 100.0,
-                "power_unit": PowerUnit.PERCENT
-            },
+            input={"power_setting": 100.0, "power_unit": PowerUnit.PERCENT},
             output={
                 "power_measurement": 50.0,
                 "power_unit": PowerUnit.MW,

diff --git a/examples/multiplane_ophys_session.py b/examples/multiplane_ophys_session.py
@@ -6,6 +6,7 @@
 from aind_data_schema_models.units import PowerUnit, SizeUnit, FrequencyUnit
 
 from aind_data_schema.core.session import FieldOfView, LaserConfig, Session, Stream
+from aind_data_schema_models.brain_atlas import CCFStructure
 
 # If a timezone isn't specified, the timezone of the computer running this
 # script will be used as default
@@ -60,7 +61,7 @@
                     power_unit=PowerUnit.PERCENT,
                     scanimage_roi_index=0,
                     imaging_depth=190,
-                    targeted_structure="VISp",
+                    targeted_structure=CCFStructure.VISP,
                     scanfield_z=230,
                     scanfield_z_unit=SizeUnit.UM,
                     coupled_fov_index=1,
@@ -82,7 +83,7 @@
                     power_unit=PowerUnit.PERCENT,
                     scanimage_roi_index=0,
                     imaging_depth=232,
-                    targeted_structure="VISp",
+                    targeted_structure=CCFStructure.VISP,
                     scanfield_z=257,
                     scanfield_z_unit=SizeUnit.UM,
                     coupled_fov_index=0,
@@ -104,7 +105,7 @@
                     power_unit=PowerUnit.PERCENT,
                     scanimage_roi_index=0,
                     imaging_depth=136,
-                    targeted_structure="VISp",
+                    targeted_structure=CCFStructure.VISP,
                     scanfield_z=176,
                     scanfield_z_unit=SizeUnit.UM,
                     coupled_fov_index=3,
@@ -126,7 +127,7 @@
                     power_unit=PowerUnit.PERCENT,
                     scanimage_roi_index=0,
                     imaging_depth=282,
-                    targeted_structure="VISp",
+                    targeted_structure=CCFStructure.VISP,
                     scanfield_z=307,
                     scanfield_z_unit=SizeUnit.UM,
                     coupled_fov_index=2,
@@ -148,7 +149,7 @@
                     power_unit=PowerUnit.PERCENT,
                     scanimage_roi_index=0,
                     imaging_depth=72,
-                    targeted_structure="VISp",
+                    targeted_structure=CCFStructure.VISP,
                     scanfield_z=112,
                     scanfield_z_unit=SizeUnit.UM,
                     coupled_fov_index=5,
@@ -170,7 +171,7 @@
                     power_unit=PowerUnit.PERCENT,
                     scanimage_roi_index=0,
                     imaging_depth=326,
-                    targeted_structure="VISp",
+                    targeted_structure=CCFStructure.VISP,
                     scanfield_z=351,
                     scanfield_z_unit=SizeUnit.UM,
                     coupled_fov_index=4,
@@ -192,7 +193,7 @@
                     power_unit=PowerUnit.PERCENT,
                     scanimage_roi_index=0,
                     imaging_depth=30,
-                    targeted_structure="VISp",
+                    targeted_structure=CCFStructure.VISP,
                     scanfield_z=70,
                     scanfield_z_unit=SizeUnit.UM,
                     coupled_fov_index=7,
@@ -214,7 +215,7 @@
                     power_unit=PowerUnit.PERCENT,
                     scanimage_roi_index=0,
                     imaging_depth=364,
-                    targeted_structure="VISp",
+                    targeted_structure=CCFStructure.VISP,
                     scanfield_z=389,
                     scanfield_z_unit=SizeUnit.UM,
                     coupled_fov_index=6,

diff --git a/examples/ophys_procedures.py b/examples/ophys_procedures.py
@@ -21,6 +21,7 @@
     ViralMaterial,
     WaterRestriction,
 )
+from aind_data_schema_models.brain_atlas import CCFStructure
 
 t = datetime.datetime(2022, 7, 12, 7, 00, 00)
 t2 = datetime.datetime(2022, 9, 23, 10, 22, 00)
@@ -67,7 +68,7 @@
                     injection_coordinate_reference="Bregma",
                     injection_angle=0,
                     injection_volume=[400],
-                    targeted_structure="VTA",
+                    targeted_structure=CCFStructure.VTA,
                 ),
                 FiberImplant(
                     protocol_id="TO ENTER",
@@ -80,7 +81,7 @@
                                 ferrule_material="Ceramic",
                                 total_length=0.5,
                             ),
-                            targeted_structure="VTA",
+                            targeted_structure=CCFStructure.VTA,
                             angle=0,
                             stereotactic_coordinate_ap=-3.05,
                             stereotactic_coordinate_ml=-0.6,

diff --git a/examples/procedures.py b/examples/procedures.py
@@ -12,6 +12,7 @@
     TarsVirusIdentifiers,
     ViralMaterial,
 )
+from aind_data_schema_models.brain_atlas import CCFStructure
 
 # If a timezone isn't specified, the timezone of the computer running this
 # script will be used as default
@@ -61,7 +62,7 @@
                     bregma_to_lambda_distance=4.1,
                     injection_angle=10,
                     injection_volume=[200],
-                    targeted_structure="VISp",
+                    targeted_structure=CCFStructure.VISP,
                 ),
             ],
         ),

diff --git a/src/aind_data_schema/core/procedures.py b/src/aind_data_schema/core/procedures.py
@@ -716,12 +716,7 @@ class Procedures(AindCoreModel):
     )
     subject_procedures: List[
         Annotated[
-            Union[
-                Surgery,
-                TrainingProtocol,
-                WaterRestriction,
-                OtherSubjectProcedure
-            ],
+            Union[Surgery, TrainingProtocol, WaterRestriction, OtherSubjectProcedure],
             Field(discriminator="procedure_type"),
         ]
     ] = Field(default=[], title="Subject Procedures")

diff --git a/tests/test_imaging.py b/tests/test_imaging.py
@@ -43,10 +43,7 @@ def test_constructors(self):
                     calibration_date=datetime.now(tz=timezone.utc),
                     description="Laser power calibration",
                     device_name="Laser 1",
-                    input={
-                        "power_setting": 100.0,
-                        "power_unit": PowerUnit.PERCENT
-                    },
+                    input={"power_setting": 100.0, "power_unit": PowerUnit.PERCENT},
                     output={
                         "power_measurement": 50.0,
                         "power_unit": PowerUnit.MW,
@@ -150,10 +147,7 @@ def test_axis(self):
                         calibration_date=datetime.now(tz=timezone.utc),
                         description="Laser power calibration",
                         device_name="Laser 1",
-                        input={
-                            "power_setting": 100.0,
-                            "power_unit": PowerUnit.PERCENT
-                        },
+                        input={"power_setting": 100.0, "power_unit": PowerUnit.PERCENT},
                         output={
                             "power_measurement": 50.0,
                             "power_unit": PowerUnit.MW,

diff --git a/tests/test_procedures.py b/tests/test_procedures.py
@@ -24,6 +24,7 @@
     TarsVirusIdentifiers,
     ViralMaterial,
 )
+from aind_data_schema_models.brain_atlas import CCFStructure
 
 PYD_VERSION = re.match(r"(\d+.\d+).\d+", pyd_version).group(1)
 
@@ -206,7 +207,7 @@ def test_injection_material_check(self):
                             injection_volume=[1],
                             recovery_time=10,
                             recovery_time_unit=TimeUnit.M,
-                            targeted_structure="VISp6",
+                            targeted_structure=CCFStructure.VISP6,
                         ),
                         FiberImplant(
                             protocol_id="dx.doi.org/120.123/fkjd",
@@ -222,7 +223,7 @@ def test_injection_material_check(self):
                                         ferrule_material="Ceramic",
                                         total_length=10,
                                     ),
-                                    targeted_structure="MOp",
+                                    targeted_structure=CCFStructure.MOP,
                                     stereotactic_coordinate_ap=1,
                                     stereotactic_coordinate_dv=2,
                                     stereotactic_coordinate_ml=3,
@@ -367,7 +368,7 @@ def test_sectioning(self):
             section_distance_from_reference=0.3,
             reference_location="Bregma",
             section_strategy="Whole Brain",
-            targeted_structure="MOp",
+            targeted_structure=CCFStructure.MOP,
         )
         self.assertEqual(section.number_of_slices, len(section.output_specimen_ids))
 
@@ -381,7 +382,7 @@ def test_sectioning(self):
                 section_distance_from_reference=0.3,
                 reference_location="Bregma",
                 section_strategy="Whole Brain",
-                targeted_structure="MOp",
+                targeted_structure=CCFStructure.MOP,
             )
 
 

diff --git a/tests/test_rig_session_compatibility.py b/tests/test_rig_session_compatibility.py
@@ -50,6 +50,7 @@
     VisualStimulation,
 )
 from aind_data_schema.utils.compatibility_check import RigSessionCompatibility
+from aind_data_schema_models.brain_atlas import CCFStructure
 
 EXAMPLES_DIR = Path(__file__).parents[1] / "examples"
 EPHYS_RIG_JSON = EXAMPLES_DIR / "ephys_rig.json"
@@ -351,7 +352,7 @@
                     arc_angle=5.2,
                     module_angle=8,
                     coordinate_transform="behavior/calibration_info_np2_2023_04_24.npy",
-                    primary_targeted_structure="LGd",
+                    primary_targeted_structure=CCFStructure.LGD,
                     manipulator_coordinates=Coordinates3d(x=8422, y=4205, z=11087.5),
                     calibration_date=datetime(year=2023, month=4, day=25, tzinfo=timezone.utc),
                     notes=(
@@ -366,7 +367,7 @@
                     targeted_ccf_coordinates=[CcfCoords(ml=6637.28, ap=4265.02, dv=10707.35)],
                     assembly_name="ephys module 2",
                     coordinate_transform="behavior/calibration_info_np2_2023_04_24.py",
-                    primary_targeted_structure="LC",
+                    primary_targeted_structure=CCFStructure.LC,
                     manipulator_coordinates=Coordinates3d(x=9015, y=7144, z=13262),
                     calibration_date=datetime(year=2023, month=4, day=25, tzinfo=timezone.utc),
                     notes=(
@@ -420,7 +421,7 @@
                     targeted_ccf_coordinates=[CcfCoords(ml=8150, ap=3250, dv=7800)],
                     assembly_name="ephys module 1",
                     coordinate_transform="behavior/calibration_info_np2_2023_04_24.npy",
-                    primary_targeted_structure="LGd",
+                    primary_targeted_structure=CCFStructure.LGD,
                     manipulator_coordinates=Coordinates3d(x=8422, y=4205, z=11087.5),
                     calibration_date=datetime(year=2023, month=4, day=25, tzinfo=timezone.utc),
                     notes=(
@@ -437,7 +438,7 @@
                     targeted_ccf_coordinates=[CcfCoords(ml=8150, ap=3250, dv=7800)],
                     assembly_name="ephys module 1",
                     coordinate_transform="behavior/calibration_info_np2_2023_04_24.npy",
-                    primary_targeted_structure="LGd",
+                    primary_targeted_structure=CCFStructure.LGD,
                     manipulator_coordinates=Coordinates3d(x=8422, y=4205, z=11087.5),
                     calibration_date=datetime(year=2023, month=4, day=25, tzinfo=timezone.utc),
                     notes=(
@@ -452,7 +453,7 @@
                     targeted_ccf_coordinates=[CcfCoords(ml=6637.28, ap=4265.02, dv=10707.35)],
                     assembly_name="ephys module 2",
                     coordinate_transform="behavior/calibration_info_np2_2023_04_24.py",
-                    primary_targeted_structure="LC",
+                    primary_targeted_structure=CCFStructure.LC,
                     manipulator_coordinates=Coordinates3d(x=9015, y=7144, z=13262),
                     calibration_date=datetime(year=2023, month=4, day=25, tzinfo=timezone.utc),
                     notes=(
@@ -812,7 +813,7 @@ def read_json(filepath: Path) -> dict:
                             assembly_name="Fiber Module A",
                             arc_angle=30,
                             module_angle=180,
-                            primary_targeted_structure="VISp",
+                            primary_targeted_structure=CCFStructure.VISP,
                             manipulator_coordinates=Coordinates3d(x=30.5, y=70, z=180),
                         )
                     ],

diff --git a/tests/test_session.py b/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="VISl",
+                            primary_targeted_structure=CCFStructure.VISL,
                             targeted_ccf_coordinates=[CcfCoords(ml="1", ap="1", dv="1")],
                             manipulator_coordinates=Coordinates3d(x="1", y="1", z="1"),
                         ),