new validator and tests

Author: Sun-flow

src/aind_data_schema/core/mri_session.py

@@ -5,7 +5,7 @@
 from enum import Enum
 from typing import List, Literal, Optional
 
-from pydantic import Field, ValidationInfo, field_validator
+from pydantic import Field, ValidationInfo, field_validator, model_validator
 
 from aind_data_schema.base import AindCoreModel, AindGenericType, AindModel
 from aind_data_schema.core.procedures import Anaesthetic
@@ -47,13 +47,13 @@ class MRIScan(AindModel):
     scan_sequence_type: MriScanSequence = Field(..., title="Scan sequence")
     rare_factor: Optional[int] = Field(None, title="RARE factor")
     echo_time: Decimal = Field(..., title="Echo time (ms)")
-    effective_echo_time: Decimal = Field(..., title="Effective echo time (ms)")
+    effective_echo_time: Optional[Decimal] = Field(None, title="Effective echo time (ms)")
     echo_time_unit: TimeUnit = Field(TimeUnit.MS, title="Echo time unit")
     repetition_time: Decimal = Field(..., title="Repetition time (ms)")
     repetition_time_unit: TimeUnit = Field(TimeUnit.MS, title="Repetition time unit")
     # fields required to get correct orientation
-    vc_orientation: Rotation3dTransform = Field(..., title="Scan orientation")
-    vc_position: Translation3dTransform = Field(..., title="Scan position")
+    vc_orientation: Optional[Rotation3dTransform] = Field(None, title="Scan orientation")
+    vc_position: Optional[Translation3dTransform] = Field(None, title="Scan position")
     subject_position: SubjectPosition = Field(..., title="Subject position")
     # other fields
     voxel_sizes: Scale3dTransform = Field(..., title="Voxel sizes", description="Resolution")
@@ -77,6 +77,14 @@ def validate_other(cls, value: Optional[str], info: ValidationInfo) -> Optional[
                 " Describe the scan_sequence_type in the notes field."
             )
         return value
+    
+    @model_validator(mode="after")
+    def validate_primary_scan(self):
+        if self.primary_scan:
+            if not self.vc_orientation or not self.vc_position:
+                raise ValueError(
+                    "Primary scan must have vc_orientation and vc_position"
+                )
 
 
 class MriSession(AindCoreModel):

tests/test_imaging.py

@@ -319,5 +319,29 @@ def test_validators(self):
         self.assertEqual(expected_exception, repr(e.exception))
 
 
+        with self.assertRaises(ValueError) as e:
+            ms.MRIScan(
+                scan_index=1,
+                scan_type="3D Scan",
+                scan_sequence_type="RARE",
+                rare_factor=4,
+                primary_scan=True,
+                subject_position="Supine",
+                voxel_sizes=Scale3dTransform(scale=[0.1, 0.1, 0.1]),
+                echo_time=2.2,
+                effective_echo_time=2.0,
+                repetition_time=1.2,
+                additional_scan_parameters={"number_averages": 3},
+            )
+
+        expected_exception = (
+            "1 validation error for MRIScan\n"
+            "  Value error, Primary scan must have vc_orientation and vc_position [type=value_error, "
+            "input_value={'scan_index': 1, 'scan_t... {'number_averages': 3}}, input_type=dict]\n"
+            f"    For further information visit https://errors.pydantic.dev/{PYD_VERSION}/v/value_error"
+        )
+        self.assertEqual(expected_exception, repr(e.exception))
+
+
 if __name__ == "__main__":
     unittest.main()