feat: add optional calibration fit to Calibration (#1494)

* feat: working on CalibrationFit

* chore: lint

* chore: docstring int he wrong place

Author: dbirman

src/aind_data_schema/components/measurements.py

@@ -1,14 +1,52 @@
 """Calibration data models"""
 
+from enum import Enum
 from typing import List, Literal, Optional
 
-from aind_data_schema_models.units import UNITS, PowerUnit, TimeUnit, VolumeUnit
+from aind_data_schema_models.units import UNITS, PowerUnit, TimeUnit, VolumeUnit, VoltageUnit
+from pydantic import model_validator
 
-from aind_data_schema.base import AwareDatetimeWithDefault, Discriminated, Field
+from aind_data_schema.base import AwareDatetimeWithDefault, DataModel, Discriminated, Field, GenericModel
 from aind_data_schema.components.configs import DeviceConfig
 from aind_data_schema.components.reagent import Reagent
 
 
+class FitType(Enum):
+    """Type of fit for calibration data"""
+
+    LINEAR_INTERPOLATION = "linear_interpolation"
+    LINEAR = "linear"
+    OTHER = "other"
+
+
+class CalibrationFit(DataModel):
+    """Fit equation for calibration data"""
+
+    fit_type: FitType = Field(
+        ...,
+        title="Fit type",
+    )
+    fit_parameters: Optional[GenericModel] = Field(
+        default=None,
+        title="Fit parameters",
+        description="Parameters of the fit equation, e.g. slope and intercept for linear fit",
+    )
+
+    @model_validator(mode="before")
+    def validate_fit_type(cls, values):
+        """Ensure that parameters are provided for linear and other fits"""
+        fit_type = values.get("fit_type")
+        fit_parameters = values.get("fit_parameters")
+
+        if fit_type in {FitType.LINEAR, FitType.OTHER} and not fit_parameters:
+            raise ValueError(f"Fit parameters must be provided for {fit_type.value} fit type")
+
+        if fit_type == FitType.LINEAR_INTERPOLATION and fit_parameters is not None:
+            raise ValueError("Fit parameters should not be provided for linear interpolation fit type")
+
+        return values
+
+
 class Calibration(DeviceConfig):
     """Generic calibration class"""
 
@@ -25,17 +63,21 @@ class Calibration(DeviceConfig):
         ..., description="Calibration output (provide the average if repeated)", title="Outputs"
     )
     output_unit: UNITS = Field(..., title="Output unit")
+    fit: Optional[CalibrationFit] = Field(
+        default=None,
+        title="Fit",
+        description="Fit equation for the calibration data used during data acquisition",
+    )
     notes: Optional[str] = Field(
         default=None,
         title="Notes",
-        description="Fit equation, etc",
     )
 
 
 class VolumeCalibration(Calibration):
-    """Calibration of a liquid delivery device"""
+    """Calibration of a liquid delivery device based on solenoid/valve opening times"""
 
-    input: List[float] = Field(..., title="Input times", description="Length of time solenoid is open")
+    input: List[float] = Field(..., title="Input times", description="Length of time solenoid/valve is open")
     input_unit: TimeUnit = Field(..., title="Input unit")
     repeats: Optional[int] = Field(
         default=None,
@@ -51,10 +93,10 @@ class VolumeCalibration(Calibration):
 
 
 class PowerCalibration(Calibration):
-    """Calibration of a laser device"""
+    """Calibration of a device that outputs power based on input strength"""
 
-    input: List[float] = Field(..., title="Input", description="Power or percentage input strength")
-    input_unit: PowerUnit = Field(..., title="Input unit")
+    input: List[float] = Field(..., title="Input", description="Power, voltage, or percentage input strength")
+    input_unit: PowerUnit | VoltageUnit = Field(..., title="Input unit")
     output: List[float] = Field(..., title="Output", description="Power output (provide the average if repeated)")
     output_unit: PowerUnit = Field(..., title="Output unit")
 

tests/test_measurements.py

@@ -0,0 +1,67 @@
+"""Tests for CalibrationFit from measurements module"""
+
+import unittest
+from pydantic import ValidationError
+from aind_data_schema.components.measurements import CalibrationFit, FitType
+
+
+class TestCalibrationFit(unittest.TestCase):
+    """Tests for CalibrationFit class"""
+
+    def test_linear_interpolation_without_parameters(self):
+        """Test that linear interpolation fit type works without parameters"""
+        fit = CalibrationFit(fit_type=FitType.LINEAR_INTERPOLATION)
+        self.assertEqual(fit.fit_type, FitType.LINEAR_INTERPOLATION.value)
+        self.assertIsNone(fit.fit_parameters)
+
+    def test_linear_interpolation_with_parameters_raises_error(self):
+        """Test that linear interpolation fit type raises error with parameters"""
+        with self.assertRaises(ValidationError) as context:
+            CalibrationFit(fit_type=FitType.LINEAR_INTERPOLATION, fit_parameters={"slope": 1.0, "intercept": 0.0})
+        self.assertIn("Fit parameters should not be provided for linear interpolation fit type", str(context.exception))
+
+    def test_linear_fit_with_parameters(self):
+        """Test that linear fit type works with parameters"""
+        parameters = {"slope": 1.5, "intercept": 2.0}
+        fit = CalibrationFit(fit_type=FitType.LINEAR, fit_parameters=parameters)
+        self.assertEqual(fit.fit_type, FitType.LINEAR.value)
+        # Compare the fit_parameters as a dict using model_dump()
+        self.assertIsNotNone(fit.fit_parameters)
+        self.assertEqual(fit.fit_parameters.model_dump(), parameters)
+
+    def test_linear_fit_without_parameters_raises_error(self):
+        """Test that linear fit type raises error without parameters"""
+        with self.assertRaises(ValidationError) as context:
+            CalibrationFit(fit_type=FitType.LINEAR)
+        self.assertIn("Fit parameters must be provided for linear fit type", str(context.exception))
+
+    def test_other_fit_with_parameters(self):
+        """Test that other fit type works with parameters"""
+        parameters = {"a": 1.0, "b": 2.0, "c": 3.0}
+        fit = CalibrationFit(fit_type=FitType.OTHER, fit_parameters=parameters)
+        self.assertEqual(fit.fit_type, FitType.OTHER.value)
+        # Compare the fit_parameters as a dict using model_dump()
+        self.assertIsNotNone(fit.fit_parameters)
+        self.assertEqual(fit.fit_parameters.model_dump(), parameters)
+
+    def test_other_fit_without_parameters_raises_error(self):
+        """Test that other fit type raises error without parameters"""
+        with self.assertRaises(ValidationError) as context:
+            CalibrationFit(fit_type=FitType.OTHER)
+        self.assertIn("Fit parameters must be provided for other fit type", str(context.exception))
+
+    def test_linear_fit_with_none_parameters_raises_error(self):
+        """Test that linear fit type raises error with None parameters"""
+        with self.assertRaises(ValidationError) as context:
+            CalibrationFit(fit_type=FitType.LINEAR, fit_parameters=None)
+        self.assertIn("Fit parameters must be provided for linear fit type", str(context.exception))
+
+    def test_other_fit_with_none_parameters_raises_error(self):
+        """Test that other fit type raises error with None parameters"""
+        with self.assertRaises(ValidationError) as context:
+            CalibrationFit(fit_type=FitType.OTHER, fit_parameters=None)
+        self.assertIn("Fit parameters must be provided for other fit type", str(context.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()