Merge pull request #206 from AllenNeuralDynamics/feat-update-session

First prerelease v0.13

Author: bruno-f-cruz

examples/aind_manipulator.py

@@ -1,50 +1,32 @@
 import os
 
 from aind_behavior_services.base import get_commit_hash
-from aind_behavior_services.calibration import aind_manipulator as m
+from aind_behavior_services.rig import aind_manipulator as m
 from aind_behavior_services.session import AindBehaviorSessionModel
 from aind_behavior_services.utils import utcnow
 
 calibration_data = m.AindManipulatorCalibration(
-    device_name="AindManipulatorCalibration",
-    input=m.AindManipulatorCalibrationInput(
-        full_step_to_mm=m.ManipulatorPosition(x=0.01, y1=0.01, y2=0.01, z=0.01),
-        axis_configuration=[
-            m.AxisConfiguration(axis=m.Axis.X),
-            m.AxisConfiguration(axis=m.Axis.Y1),
-            m.AxisConfiguration(axis=m.Axis.Y2),
-            m.AxisConfiguration(axis=m.Axis.Z),
-        ],
-        homing_order=[m.Axis.Y2, m.Axis.Y1, m.Axis.X, m.Axis.Z],
-        initial_position=m.ManipulatorPosition(y1=0, y2=0, x=0, z=10000),
-    ),
-    output=m.AindManipulatorCalibrationOutput(),
-    date=utcnow(),
+    full_step_to_mm=m.ManipulatorPosition(x=0.01, y1=0.01, y2=0.01, z=0.01),
+    axis_configuration=[
+        m.AxisConfiguration(axis=m.Axis.X),
+        m.AxisConfiguration(axis=m.Axis.Y1),
+        m.AxisConfiguration(axis=m.Axis.Y2),
+        m.AxisConfiguration(axis=m.Axis.Z),
+    ],
+    homing_order=[m.Axis.Y2, m.Axis.Y1, m.Axis.X, m.Axis.Z],
+    initial_position=m.ManipulatorPosition(y1=0, y2=0, x=0, z=10000),
 )
 
-calibration_logic = m.CalibrationLogic(task_parameters=m.CalibrationParameters())
-
 
 calibration_session = AindBehaviorSessionModel(
-    root_path="C:\\Data",
     allow_dirty_repo=False,
     experiment="AindManipulatorCalibration",
     date=utcnow(),
     subject="AindManipulator",
-    experiment_version="manipulator_control",
     commit_hash=get_commit_hash(),
 )
 
-rig = m.CalibrationRig(
-    manipulator=m.AindManipulatorDevice(port_name="COM4", calibration=calibration_data),
-    rig_name="AindManipulatorRig",
-)
-
 seed_path = "local/aind_manipulator_{suffix}.json"
 os.makedirs(os.path.dirname(seed_path), exist_ok=True)
-with open(seed_path.format(suffix="calibration_logic"), "w") as f:
-    f.write(calibration_logic.model_dump_json(indent=3))
 with open(seed_path.format(suffix="session"), "w") as f:
     f.write(calibration_session.model_dump_json(indent=3))
-with open(seed_path.format(suffix="rig"), "w") as f:
-    f.write(rig.model_dump_json(indent=3))

examples/load_cells.py

@@ -1,28 +1,17 @@
-from aind_behavior_services.calibration import load_cells as lc
-from aind_behavior_services.utils import utcnow
+from aind_behavior_services.rig import load_cells as lc
 
-lc0 = lc.LoadCellCalibrationInput(
+lc0 = lc.LoadCellChannelCalibration(
     channel=0,
-    offset_measurement=[lc.MeasuredOffset(offset=100, baseline=0.1)],
-    weight_measurement=[lc.MeasuredWeight(weight=0.1, baseline=0.1)],
-)
-lc1 = lc.LoadCellCalibrationInput(
-    channel=1,
-    offset_measurement=[lc.MeasuredOffset(offset=100, baseline=0.1)],
-    weight_measurement=[lc.MeasuredWeight(weight=0.1, baseline=0.1)],
+    baseline=0.1,
+    offset=100,
+    slope=2.0,
 )
 
-lc_calibration_input = lc.LoadCellsCalibrationInput(channels=[lc1, lc0])
-lc_calibration_output = lc.LoadCellsCalibrationOutput(
-    channels=[
-        lc.LoadCellCalibrationOutput(channel=0, offset=6, baseline=1000, weight_lookup=[]),
-        lc.LoadCellCalibrationOutput(channel=1, offset=6, baseline=1000, weight_lookup=[]),
-    ]
+lc1 = lc.LoadCellChannelCalibration(
+    channel=1,
+    baseline=0.2,
+    offset=150,
+    slope=1.5,
 )
 
-calibration = lc.LoadCellsCalibration(
-    input=lc_calibration_input,
-    output=lc_calibration_output,
-    device_name="LoadCells",
-    date=utcnow(),
-)
+lc_calibration_input = lc.LoadCellsCalibration(channels=[lc1, lc0])

examples/treadmill.py

@@ -1,16 +1,8 @@
-from aind_behavior_services.calibration import treadmill
-from aind_behavior_services.utils import utcnow
+from aind_behavior_services.rig import treadmill
 
-treadmill_calibration = treadmill.TreadmillCalibrationOutput(
+treadmill_calibration = treadmill.TreadmillCalibration(
     wheel_diameter=10,
     pulses_per_revolution=10000,
     invert_direction=False,
     brake_lookup_calibration=[[0, 0], [0.5, 32768], [1, 65535]],
 )
-
-calibration = treadmill.TreadmillCalibration(
-    device_name="Treadmill",
-    input=treadmill.TreadmillCalibrationInput(),
-    output=treadmill_calibration,
-    date=utcnow(),
-)

examples/water_valve.py

@@ -1,5 +1,4 @@
-from aind_behavior_services.calibration import water_valve as wv
-from aind_behavior_services.utils import utcnow
+from aind_behavior_services.rig import water_valve as wv
 
 
 def linear_model(time, slope, offset):
@@ -16,20 +15,4 @@ def linear_model(time, slope, offset):
     for t in zip(_delta_times, _water_weights)
 ]
 
-
-_outputs = wv.WaterValveCalibrationOutput(
-    interval_average={interval: volume for interval, volume in zip(_delta_times, _water_weights)},
-    slope=_slope,
-    offset=_offset,
-    r2=1.0,
-    valid_domain=[value for value in _delta_times],
-)
-
-input = wv.WaterValveCalibrationInput(measurements=_inputs)
-
-calibration = wv.WaterValveCalibration(
-    input=input,
-    output=input.calibrate_output(),
-    device_name="WaterValve",
-    date=utcnow(),
-)
+calibration = wv.calibrate_water_valves(_inputs)

pyproject.toml

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

src/_generators/json_schema.py

@@ -5,9 +5,9 @@
 
 from pydantic import BaseModel
 
-from aind_behavior_services.calibration import aind_manipulator
 from aind_behavior_services.data_types import DataTypes
 from aind_behavior_services.message_protocol import MessageProtocol
+from aind_behavior_services.rig.aind_manipulator import AindManipulator
 from aind_behavior_services.session import AindBehaviorSessionModel
 from aind_behavior_services.utils import export_schema, pascal_to_snake_case
 
@@ -34,7 +34,7 @@ def main():
         ToGenerateJsonSchema(model_name="AindBehaviorSessionModel", model=AindBehaviorSessionModel, remove_root=False),
         ToGenerateJsonSchema(model_name="DataTypes", model=DataTypes, remove_root=True),
         ToGenerateJsonSchema(model_name="MessageProtocol", model=MessageProtocol, remove_root=True),
-        ToGenerateJsonSchema(model_name="AindManipulator", model=aind_manipulator.CalibrationRig, remove_root=False),
+        ToGenerateJsonSchema(model_name="AindManipulator", model=AindManipulator, remove_root=False),
     )
 
     for m in models:

src/_generators/templates/harp.j2

@@ -10,7 +10,6 @@ from ._base import Device
 
 class _HarpDeviceBase(Device):
     who_am_i: Optional[int] = Field(default=None, le=9999, ge=0, description="Device WhoAmI")
-    device_type: str
     serial_number: Optional[str] = Field(default=None, description="Device serial number")
     port_name: str = Field(..., description="Device port name")