diff --git a/pyproject.toml b/pyproject.toml index 6420af1e..1dfc59ea 100644 --- a/pyproject.toml +++ b/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', diff --git a/schema/session.json b/schema/session.json index 40d635a8..8880359b 100644 --- a/schema/session.json +++ b/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" }, diff --git a/src/aind_behavior_services/rig/_harp_gen.py b/src/aind_behavior_services/rig/_harp_gen.py index 80e72b32..e1943bdc 100644 --- a/src/aind_behavior_services/rig/_harp_gen.py +++ b/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", ] diff --git a/src/aind_behavior_services/rig/olfactometer.py b/src/aind_behavior_services/rig/olfactometer.py index 11606d56..47ce3616 100644 --- a/src/aind_behavior_services/rig/olfactometer.py +++ b/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,14 +28,10 @@ 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)") @@ -43,9 +39,65 @@ class OlfactometerChannelConfig(BaseModel): 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): diff --git a/tests/rig/test_olfactometer.py b/tests/rig/test_olfactometer.py new file mode 100644 index 00000000..7d4da04a --- /dev/null +++ b/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"] diff --git a/uv.lock b/uv.lock index af63c7c9..f1c66fe3 100644 --- a/uv.lock +++ b/uv.lock @@ -38,7 +38,7 @@ wheels = [ [[package]] name = "aind-behavior-services" -version = "0.13.7" +version = "0.14.0rc0" source = { editable = "." } dependencies = [ { name = "aind-behavior-curriculum" },