Fix examples

AllenNeuralDynamics · Aug 30, 2024 · b7164a0 · b7164a0
1 parent a34061b
commit b7164a0
Showing 1 changed file with 12 additions and 7 deletions.
diff --git a/examples/examples.py b/examples/examples.py
@@ -21,6 +21,12 @@
     OlfactometerChannelConfig,
     OlfactometerChannelType,
 )
+from aind_behavior_services.calibration.treadmill import (
+    Treadmill,
+    TreadmillCalibration,
+    TreadmillCalibrationInput,
+    TreadmillCalibrationOutput,
+)
 from aind_behavior_services.calibration.water_valve import (
     Measurement,
     WaterValveCalibration,
@@ -36,9 +42,7 @@
     HarpLickometer,
     HarpOlfactometer,
     HarpSniffDetector,
-    HarpTreadmill,
     RigCalibration,
-    Treadmill,
 )
 from aind_behavior_vr_foraging.task_logic import (
     AindVrForagingTaskLogic,
@@ -151,12 +155,13 @@ def mock_rig() -> AindVrForagingRig:
         harp_clock_generator=HarpClockGenerator(port_name="COM6"),
         harp_analog_input=None,
         harp_sniff_detector=HarpSniffDetector(port_name="COM7"),
-        harp_treadmill=HarpTreadmill(
+        harp_treadmill=Treadmill(
             port_name="COM8",
-            calibration=Treadmill(
-                wheel_diameter=15,
-                pulses_per_revolution=28800,
-                break_lookup_calibration=[[0, 0], [1, 65535]],
+            calibration=TreadmillCalibration(
+                input=TreadmillCalibrationInput(),
+                output=TreadmillCalibrationOutput(
+                    wheel_diameter=15, pulses_per_revolution=28800, brake_lookup_calibration=[[0, 0], [1, 65535]]
+                ),
             ),
         ),
         manipulator=AindManipulatorDevice(port_name="COM9", calibration=manipulator_calibration),