Merge pull request #122 from AllenNeuralDynamics/refactor-add-slope-c…

…alibration-to-load-cell

Add slope calibration to load cells schema

src/Extensions/LoadCellsCalibrationRig.cs

@@ -166,6 +166,8 @@ public partial class LoadCellCalibrationOutput
 
         private double? _baseline;
 
+        private double? _slope;
+
         private System.Collections.Generic.List<MeasuredWeight> _weightLookup = new System.Collections.Generic.List<MeasuredWeight>();
 
         public LoadCellCalibrationOutput()
@@ -177,6 +179,7 @@ protected LoadCellCalibrationOutput(LoadCellCalibrationOutput other)
             _channel = other._channel;
             _offset = other._offset;
             _baseline = other._baseline;
+            _slope = other._slope;
             _weightLookup = other._weightLookup;
         }
 
@@ -225,6 +228,20 @@ public double? Baseline
             }
         }
 
+        [System.Xml.Serialization.XmlIgnoreAttribute()]
+        [Newtonsoft.Json.JsonPropertyAttribute("slope")]
+        public double? Slope
+        {
+            get
+            {
+                return _slope;
+            }
+            set
+            {
+                _slope = value;
+            }
+        }
+
         [System.Xml.Serialization.XmlIgnoreAttribute()]
         [Newtonsoft.Json.JsonPropertyAttribute("weight_lookup")]
         public System.Collections.Generic.List<MeasuredWeight> WeightLookup
@@ -254,6 +271,7 @@ protected virtual bool PrintMembers(System.Text.StringBuilder stringBuilder)
             stringBuilder.Append("channel = " + _channel + ", ");
             stringBuilder.Append("offset = " + _offset + ", ");
             stringBuilder.Append("baseline = " + _baseline + ", ");
+            stringBuilder.Append("slope = " + _slope + ", ");
             stringBuilder.Append("weight_lookup = " + _weightLookup);
             return true;
         }

src/aind_behavior_services/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.8.6"
+__version__ = "0.8.7"
 
 from .rig import AindBehaviorRigModel  # noqa: F401
 from .session import AindBehaviorSessionModel  # noqa: F401

src/aind_behavior_services/calibration/load_cells.py

@@ -1,6 +1,10 @@
+from __future__ import annotations
+
 from typing import Annotated, List, Literal, Optional
 
+import numpy as np
 from pydantic import BaseModel, Field, field_validator
+from sklearn.linear_model import LinearRegression
 
 from aind_behavior_services.calibration import Calibration
 from aind_behavior_services.rig import AindBehaviorRigModel, HarpLoadCells
@@ -38,6 +42,18 @@ class LoadCellCalibrationInput(BaseModel):
     )
 
 
+class LoadCellCalibrationOutput(BaseModel):
+    channel: LoadCellChannel
+    offset: Optional[LoadCellOffset] = Field(
+        default=None, title="Load cell offset applied to the wheatstone bridge circuit"
+    )
+    baseline: Optional[float] = Field(default=None, title="Load cell baseline that will be DSP subtracted")
+    slope: Optional[float] = Field(
+        default=None, title="Load cell slope that will be used to convert adc units to weight (g)."
+    )
+    weight_lookup: List[MeasuredWeight] = Field(default=[], title="Load cell weight lookup calibration table")
+
+
 class LoadCellsCalibrationInput(BaseModel):
     channels: List[LoadCellCalibrationInput] = Field(
         default=[], title="Load cells calibration data", validate_default=True
@@ -51,12 +67,32 @@ def ensure_unique_channels(cls, values: List[LoadCellCalibrationInput]) -> List[
             raise ValueError("Channels must be unique.")
         return values
 
-
-class LoadCellCalibrationOutput(BaseModel):
-    channel: LoadCellChannel
-    offset: Optional[LoadCellOffset] = Field(default=None, title="Load cell offset")
-    baseline: Optional[float] = Field(default=None, title="Load cell baseline")
-    weight_lookup: List[MeasuredWeight] = Field(default=[], title="Load cell weight lookup calibration table")
+    @classmethod
+    def calibrate_loadcell_output(cls, value: LoadCellCalibrationInput) -> "LoadCellCalibrationOutput":
+        x = np.array([m.weight for m in value.weight_measurement])
+        y = np.array([m.baseline for m in value.weight_measurement])
+
+        # Calculate the linear regression
+        model = LinearRegression()
+        model.fit(x.reshape(-1, 1), y)
+        return LoadCellCalibrationOutput(
+            channel=value.channel,
+            offset=cls.get_optimum_offset(value.offset_measurement),
+            baseline=model.intercept_,
+            slope=model.coef_[0],
+            weight_lookup=value.weight_measurement,
+        )
+
+    @staticmethod
+    def get_optimum_offset(value: Optional[List[MeasuredOffset]]) -> Optional[LoadCellOffset]:
+        if not value:
+            return None
+        if len(value) == 0:
+            return None
+        return value[np.argmin([m.baseline for m in value])].offset
+
+    def calibrate_output(self) -> LoadCellsCalibrationOutput:
+        return LoadCellsCalibrationOutput(channels=[self.calibrate_loadcell_output(c) for c in self.channels])
 
 
 class LoadCellsCalibrationOutput(BaseModel):

src/schemas/load_cells_calibration_rig.json

@@ -59,7 +59,7 @@
               "type": "null"
             }
           ],
-          "title": "Load cell offset"
+          "title": "Load cell offset applied to the wheatstone bridge circuit"
         },
         "baseline": {
           "default": null,
@@ -71,7 +71,19 @@
               "type": "null"
             }
           ],
-          "title": "Load cell baseline"
+          "title": "Load cell baseline that will be DSP subtracted"
+        },
+        "slope": {
+          "default": null,
+          "oneOf": [
+            {
+              "type": "number"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "title": "Load cell slope that will be used to convert adc units to weight (g)."
         },
         "weight_lookup": {
           "default": [],