Refactor olfactometer calibration to reflect channel constraints

pyproject.toml

@@ -16,7 +16,7 @@ classifiers = [
     "Programming Language :: Python :: 3.13",
     "Operating System :: Microsoft :: Windows",
 ]
-version = "0.13.7"
+version = "0.14.0rc0"
 
 dependencies = [
     'pydantic',

schema/session.json

@@ -1,14 +1,14 @@
 {
   "properties": {
     "aind_behavior_services_pkg_version": {
-      "default": "0.13.7",
+      "default": "0.14.0-rc0",
       "pattern": "^(0|[1-9]\\d*)\\.(0|[1-9]\\d*)\\.(0|[1-9]\\d*)(?:-((?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\\.(?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\\+([0-9a-zA-Z-]+(?:\\.[0-9a-zA-Z-]+)*))?$",
       "title": "aind_behavior_services package version",
       "type": "string"
     },
     "version": {
-      "const": "0.13.7",
-      "default": "0.13.7",
+      "const": "0.14.0-rc0",
+      "default": "0.14.0-rc0",
       "title": "Version",
       "type": "string"
     },

src/aind_behavior_services/rig/_harp_gen.py

@@ -276,6 +276,11 @@ class HarpCuttlefishfip(_HarpDeviceBase):
     who_am_i: Literal[1407] = 1407
 
 
+class HarpDeluxdriver(_HarpDeviceBase):
+    device_type: Literal["deLuxDriver"] = "deLuxDriver"
+    who_am_i: Literal[1410] = 1410
+
+
 _HarpDevice = Union[
     HarpDeviceGeneric,
     HarpHobgoblin,
@@ -322,6 +327,7 @@ class HarpCuttlefishfip(_HarpDeviceBase):
     HarpWhiteRabbit,
     HarpEnvironmentSensor,
     HarpCuttlefishfip,
+    HarpDeluxdriver,
 ]
 
 if TYPE_CHECKING:
@@ -379,5 +385,6 @@ class HarpCuttlefishfip(_HarpDeviceBase):
     "HarpWhiteRabbit",
     "HarpEnvironmentSensor",
     "HarpCuttlefishfip",
+    "HarpDeluxdriver",
     "HarpDevice",
 ]

src/aind_behavior_services/rig/olfactometer.py

@@ -1,7 +1,7 @@
 from enum import Enum, IntEnum
-from typing import Dict, Literal, Optional
+from typing import Any, Dict, Literal, Mapping, Optional
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 from ._base import DatedCalibration
 from ._harp_gen import (
@@ -28,24 +28,76 @@ class OlfactometerChannelType(str, Enum):
 class OlfactometerChannelConfig(BaseModel):
     """Configuration for a single olfactometer channel"""
 
-    channel_index: int = Field(title="Odor channel index")
     channel_type: OlfactometerChannelType = Field(
         default=OlfactometerChannelType.ODOR, title="Olfactometer channel type"
     )
     flow_rate_capacity: Literal[100, 1000] = Field(default=100, title="Flow capacity. mL/min")
-    flow_rate: float = Field(
-        default=100, le=100, title="Target flow rate. mL/min. If channel_type == CARRIER, this value is ignored."
-    )
     odorant: Optional[str] = Field(default=None, title="Odorant name")
     odorant_dilution: Optional[float] = Field(default=None, title="Odorant dilution (%v/v)")
 
 
 class OlfactometerCalibration(DatedCalibration):
     """Olfactometer device configuration model"""
 
-    channel_config: Dict[OlfactometerChannel, OlfactometerChannelConfig] = Field(
-        default={}, description="Configuration of olfactometer channels"
-    )
+    channel0: Optional[OlfactometerChannelConfig] = Field(default=None, title="Configuration for channel 0")
+    channel1: Optional[OlfactometerChannelConfig] = Field(default=None, title="Configuration for channel 1")
+    channel2: Optional[OlfactometerChannelConfig] = Field(default=None, title="Configuration for channel 2")
+    channel3: Optional[OlfactometerChannelConfig] = Field(default=None, title="Configuration for channel 3")
+
+    @model_validator(mode="before")
+    @classmethod
+    def _deserialize_legacy_channel_config(cls, data: Any) -> Any:
+        """Back-compat: accept deprecated `channel_config` on input only."""
+        if not isinstance(data, Mapping):
+            return data
+
+        parsed = dict(data)
+        legacy = parsed.pop("channel_config", None)  # remove so it is never serialized back
+        if not isinstance(legacy, Mapping):
+            return parsed
+
+        valid_indices = {channel.value for channel in OlfactometerChannel}
+
+        for raw_key, config in legacy.items():
+            idx: Optional[int] = None
+
+            # Legacy payload keys are expected to be "0", "1", "2", "3"
+            if isinstance(raw_key, int):
+                idx = raw_key
+            elif isinstance(raw_key, str):
+                key = raw_key.strip().lower()
+                if key.isdigit():
+                    idx = int(key)
+                elif key.startswith("channel") and key[7:].isdigit():
+                    idx = int(key[7:])
+
+            if idx in valid_indices:
+                parsed.setdefault(f"channel{idx}", config)
+
+        return parsed
+
+    def as_dictionary(self) -> Dict[OlfactometerChannel, OlfactometerChannelConfig]:
+        """Return the channel configuration as a dictionary"""
+        return {
+            channel: config
+            for channel in OlfactometerChannel
+            if (config := getattr(self, f"channel{channel.value}")) is not None
+        }
+
+
+class StrictOlfactometerCalibration(OlfactometerCalibration):
+    """Strict olfactometer calibration where all channels must be defined"""
+
+    channel0: OlfactometerChannelConfig = Field(title="Configuration for channel 0")
+    channel1: OlfactometerChannelConfig = Field(title="Configuration for channel 1")
+    channel2: OlfactometerChannelConfig = Field(title="Configuration for channel 2")
+    channel3: OlfactometerChannelConfig = Field(title="Configuration for channel 3")
+
+
+class StrictOlfactometer(HarpOlfactometer):
+    """A calibrated olfactometer device where all channels must be defined"""
+
+    calibration: StrictOlfactometerCalibration = Field(title="Calibration of the olfactometer", validate_default=True)
 
 
 class Olfactometer(HarpOlfactometer):

tests/rig/test_olfactometer.py

@@ -0,0 +1,59 @@
+from aind_behavior_services.rig.olfactometer import OlfactometerCalibration, OlfactometerChannelType
+
+
+def test_olfactometer_calibration_deserializes_legacy_channel_config_and_does_not_serialize_it():
+    legacy_payload = {
+        "channel_config": {
+            "0": {
+                "channel_index": 0,
+                "channel_type": "Odor",
+                "flow_rate_capacity": 100,
+                "flow_rate": 100.0,
+                "odorant": "Amyl Acetate",
+                "odorant_dilution": 1.5,
+            },
+            "1": {
+                "channel_index": 1,
+                "channel_type": "Odor",
+                "flow_rate_capacity": 100,
+                "flow_rate": 100.0,
+                "odorant": "Banana",
+                "odorant_dilution": 1.5,
+            },
+            "2": {
+                "channel_index": 2,
+                "channel_type": "Odor",
+                "flow_rate_capacity": 100,
+                "flow_rate": 100.0,
+                "odorant": "Apple",
+                "odorant_dilution": 1.5,
+            },
+            "3": {
+                "channel_index": 3,
+                "channel_type": "Carrier",
+                "flow_rate_capacity": 1000,
+                "flow_rate": 100.0,
+                "odorant": None,
+                "odorant_dilution": None,
+            },
+        }
+    }
+
+    model = OlfactometerCalibration.model_validate(legacy_payload)
+
+    assert model.channel0 is not None
+    assert model.channel1 is not None
+    assert model.channel2 is not None
+    assert model.channel3 is not None
+
+    assert model.channel0.odorant == "Amyl Acetate"
+    assert model.channel1.odorant == "Banana"
+    assert model.channel2.odorant == "Apple"
+    assert model.channel3.channel_type == OlfactometerChannelType.CARRIER
+
+    dumped = model.model_dump(mode="json")
+    assert "channel_config" not in dumped
+
+    # legacy-only keys should not survive model serialization
+    assert "channel_index" not in dumped["channel0"]
+    assert "flow_rate" not in dumped["channel0"]