Merge pull request #254 from AllenNeuralDynamics/feat-implement-tread…

…mill

Implement treadmill brake control

docs/conf.py

@@ -5,12 +5,12 @@
 
 import os
 import sys
+
 import aind_behavior_vr_foraging.rig
 import aind_behavior_vr_foraging.task_logic
 import erdantic as erd
 from pydantic import BaseModel
 
-
 sys.path.insert(0, os.path.abspath("../src/DataSchemas"))
 import aind_behavior_vr_foraging
 
@@ -69,6 +69,7 @@
 
 # -- Options for linkcode extension ---------------------------------------
 
+
 def linkcode_resolve(domain, info):
     if domain != "py":
         return None
@@ -80,6 +81,7 @@ def linkcode_resolve(domain, info):
 
 # -- Class diagram generation
 
+
 def export_model_diagram(model: BaseModel, root: str = "_static") -> None:
     diagram = erd.create(model)
     diagram.draw(f"{root}/{model.__name__}.svg")

examples/examples.py

@@ -19,7 +19,10 @@
 )
 from aind_behavior_services.session import AindBehaviorSessionModel
 from aind_behavior_vr_foraging.rig import AindVrForagingRig, RigCalibration, Treadmill
-from aind_behavior_vr_foraging.task_logic import AindVrForagingTaskLogic, AindVrForagingTaskParameters
+from aind_behavior_vr_foraging.task_logic import (
+    AindVrForagingTaskLogic,
+    AindVrForagingTaskParameters,
+)
 
 
 def main():

pyproject.toml

@@ -14,7 +14,7 @@ readme = "README.md"
 dynamic = ["version"]
 
 dependencies = [
-  "aind_behavior_services@git+https://github.com/AllenNeuralDynamics/Aind.Behavior.Services@0.7.5",
+  "aind_behavior_services@git+https://github.com/AllenNeuralDynamics/Aind.Behavior.Services@0.7.8",
 ]
 
 [project.optional-dependencies]

scripts/regenerate.py

@@ -4,7 +4,11 @@
 import aind_behavior_vr_foraging.rig
 import aind_behavior_vr_foraging.task_logic
 from aind_behavior_services.session import AindBehaviorSessionModel
-from aind_behavior_services.utils import convert_pydantic_to_bonsai, pascal_to_snake_case, snake_to_pascal_case
+from aind_behavior_services.utils import (
+    convert_pydantic_to_bonsai,
+    pascal_to_snake_case,
+    snake_to_pascal_case,
+)
 
 SCHEMA_ROOT = Path("./src/DataSchemas/")
 EXTENSIONS_ROOT = Path("./src/Extensions/")

src/DataSchemas/aind_behavior_vr_foraging/rig.py

@@ -2,26 +2,33 @@
 from __future__ import annotations
 
 # Import core types
-from typing import Annotated, Literal, Optional, Union
+from typing import Annotated, List, Literal, Optional
 
 import aind_behavior_services.calibration.olfactometer as oc
 import aind_behavior_services.calibration.water_valve as wvc
 import aind_behavior_services.rig as rig
 from aind_behavior_services.rig import AindBehaviorRigModel
-from pydantic import BaseModel, Field, RootModel
+from pydantic import BaseModel, Field
 
 __version__ = "0.3.0"
 
-TreadmillSettings = rig.Treadmill
 
+ValuePair = Annotated[List[float], Field(min_length=2, max_length=2, description="A tuple of two values")]
 
-class TreadmillBoard(RootModel):
-    root: Annotated[Union[rig.HarpTreadmill, rig.HarpBehavior], Field(discriminator="who_am_i")]
+
+class Treadmill(rig.Treadmill):
+    break_lookup_calibration: List[ValuePair] = Field(
+        default=[[0, 0], [1, 65535]],
+        validate_default=True,
+        min_length=2,
+        min=0,
+        description="Break lookup calibration. Each Tuple is (0-1 (percent), 0-full-scale). \
+            Values are linearly interpolated",
+    )
 
 
-class Treadmill(BaseModel):
-    harp_board: TreadmillBoard = Field(..., description="The board to be used as a treadmill input")
-    settings: rig.Treadmill = Field(default=rig.Treadmill(), description="Treadmill settings")
+class HarpTreadmill(rig.HarpTreadmill):
+    calibration: Treadmill = Field(Treadmill(), description="Treadmill calibration settings", validate_default=True)
 
 
 class RigCalibration(BaseModel):
@@ -42,7 +49,7 @@ class AindVrForagingRig(AindBehaviorRigModel):
     harp_lickometer: rig.HarpLickometer = Field(..., description="Harp lickometer")
     harp_clock_generator: rig.HarpClockGenerator = Field(..., description="Harp clock generator")
     harp_analog_input: Optional[rig.HarpAnalogInput] = Field(default=None, description="Harp analog input")
-    treadmill: Treadmill = Field(..., description="Treadmill settings")
+    harp_treadmill: HarpTreadmill = Field(..., description="Harp treadmill")
     harp_sniff_detector: Optional[rig.HarpSniffDetector] = Field(None, description="Sniff detector settings")
     screen: rig.Screen = Field(default=rig.Screen(), description="Screen settings")
     calibration: RigCalibration = Field(..., description="Calibration models")

src/DataSchemas/aind_behavior_vr_foraging/task_logic.py

@@ -173,7 +173,7 @@ class RenderSpecification(BaseModel):
 
 class VirtualSiteGenerator(BaseModel):
     render_specification: RenderSpecification = Field(
-        default=RenderSpecification(), description="Contrast of the environment"
+        default=RenderSpecification(), description="Contrast of the environment", validate_default=True
     )
     label: VirtualSiteLabels = Field(default=VirtualSiteLabels.UNSPECIFIED, description="Label of the virtual site")
     length_distribution: distributions.Distribution = Field(
@@ -183,16 +183,31 @@ class VirtualSiteGenerator(BaseModel):
 
 class VirtualSiteGeneration(BaseModel):
     inter_site: VirtualSiteGenerator = Field(
-        default=VirtualSiteGenerator(), description="Generator of the inter-site virtual sites"
+        default=VirtualSiteGenerator(label=VirtualSiteLabels.INTERSITE),
+        validate_default=True,
+        description="Generator of the inter-site virtual sites",
     )
     inter_patch: VirtualSiteGenerator = Field(
-        default=VirtualSiteGenerator(), description="Generator of the inter-patch virtual sites"
+        default=VirtualSiteGenerator(label=VirtualSiteLabels.INTERPATCH),
+        validate_default=True,
+        description="Generator of the inter-patch virtual sites",
     )
-    post_patch: VirtualSiteGenerator = Field(
-        default=VirtualSiteGenerator(), description="Generator of the post-inter-patch virtual sites"
+    post_patch: Optional[VirtualSiteGenerator] = Field(
+        default=None,
+        validate_default=True,
+        description="Generator of the post-patch virtual sites",
     )
     reward_site: VirtualSiteGenerator = Field(
-        default=VirtualSiteGenerator(), description="Generator of the reward-site virtual sites"
+        default=VirtualSiteGenerator(label=VirtualSiteLabels.REWARDSITE),
+        validate_default=True,
+        description="Generator of the reward-site virtual sites",
+    )
+
+
+class TreadmillSpecification(BaseModel):
+    friction: Optional[distributions.Distribution] = Field(
+        default=None,
+        description="Friction of the treadmill (0-1). The drawn value must be between 0 and 1",
     )
 
 
@@ -210,6 +225,9 @@ class VirtualSite(BaseModel):
     render_specification: RenderSpecification = Field(
         default=RenderSpecification(), description="The optional render specification of the virtual site"
     )
+    treadmill_specification: Optional[TreadmillSpecification] = Field(
+        default=None, description="Treadmill specification"
+    )
 
 
 class PatchStatistics(BaseModel):
@@ -222,7 +240,7 @@ class PatchStatistics(BaseModel):
         default=None, description="The optional reward specification of the patch"
     )
     virtual_site_generation: VirtualSiteGeneration = Field(
-        default=VirtualSiteGeneration(), description="Virtual site generation specification"
+        default=VirtualSiteGeneration(), validate_default=True, description="Virtual site generation specification"
     )
 
 

src/DataSchemas/aind_vr_foraging_rig.json

@@ -526,16 +526,27 @@
           ]
         },
         "calibration": {
-          "default": null,
-          "description": "Calibration",
-          "oneOf": [
-            {
-              "$ref": "#/definitions/BaseModel"
-            },
+          "allOf": [
             {
-              "type": "null"
+              "$ref": "#/definitions/Treadmill"
             }
-          ]
+          ],
+          "default": {
+            "wheel_diameter": 15.0,
+            "pulses_per_revolution": 28800,
+            "invert_direction": false,
+            "break_lookup_calibration": [
+              [
+                0.0,
+                0.0
+              ],
+              [
+                1.0,
+                65535.0
+              ]
+            ]
+          },
+          "description": "Treadmill calibration settings"
         },
         "who_am_i": {
           "const": 1402,
@@ -963,23 +974,57 @@
       "title": "SpinnakerCamera",
       "type": "object"
     },
-    "TreadmillBoard": {
-      "discriminator": {
-        "mapping": {
-          "1216": "#/definitions/HarpBehavior",
-          "1402": "#/definitions/HarpTreadmill"
+    "Treadmill": {
+      "properties": {
+        "wheel_diameter": {
+          "default": 15,
+          "description": "Wheel diameter",
+          "minimum": 0.0,
+          "title": "Wheel Diameter",
+          "type": "number"
         },
-        "propertyName": "who_am_i"
-      },
-      "oneOf": [
-        {
-          "$ref": "#/definitions/HarpTreadmill"
+        "pulses_per_revolution": {
+          "default": 28800,
+          "description": "Pulses per revolution",
+          "minimum": 1,
+          "title": "Pulses Per Revolution",
+          "type": "integer"
         },
-        {
-          "$ref": "#/definitions/HarpBehavior"
+        "invert_direction": {
+          "default": false,
+          "description": "Invert direction",
+          "title": "Invert Direction",
+          "type": "boolean"
+        },
+        "break_lookup_calibration": {
+          "default": [
+            [
+              0,
+              0
+            ],
+            [
+              1,
+              65535
+            ]
+          ],
+          "description": "Break lookup calibration. Each Tuple is (0-1 (percent), 0-full-scale).             Values are linearly interpolated",
+          "items": {
+            "description": "A tuple of two values",
+            "items": {
+              "type": "number"
+            },
+            "maxItems": 2,
+            "minItems": 2,
+            "type": "array"
+          },
+          "min": 0,
+          "minItems": 2,
+          "title": "Break Lookup Calibration",
+          "type": "array"
         }
-      ],
-      "title": "TreadmillBoard"
+      },
+      "title": "Treadmill",
+      "type": "object"
     },
     "VideoWriter": {
       "discriminator": {
@@ -1266,62 +1311,6 @@
       },
       "title": "WebCamera",
       "type": "object"
-    },
-    "aind_behavior_services__rig__Treadmill": {
-      "properties": {
-        "wheel_diameter": {
-          "default": 15,
-          "description": "Wheel diameter",
-          "minimum": 0.0,
-          "title": "Wheel Diameter",
-          "type": "number"
-        },
-        "pulses_per_revolution": {
-          "default": 28800,
-          "description": "Pulses per revolution",
-          "minimum": 1,
-          "title": "Pulses Per Revolution",
-          "type": "integer"
-        },
-        "invert_direction": {
-          "default": false,
-          "description": "Invert direction",
-          "title": "Invert Direction",
-          "type": "boolean"
-        }
-      },
-      "title": "Treadmill",
-      "type": "object"
-    },
-    "aind_behavior_vr_foraging__rig__Treadmill": {
-      "properties": {
-        "harp_board": {
-          "allOf": [
-            {
-              "$ref": "#/definitions/TreadmillBoard"
-            }
-          ],
-          "description": "The board to be used as a treadmill input"
-        },
-        "settings": {
-          "allOf": [
-            {
-              "$ref": "#/definitions/aind_behavior_services__rig__Treadmill"
-            }
-          ],
-          "default": {
-            "wheel_diameter": 15.0,
-            "pulses_per_revolution": 28800,
-            "invert_direction": false
-          },
-          "description": "Treadmill settings"
-        }
-      },
-      "required": [
-        "harp_board"
-      ],
-      "title": "Treadmill",
-      "type": "object"
     }
   },
   "properties": {
@@ -1405,13 +1394,13 @@
         }
       ]
     },
-    "treadmill": {
+    "harp_treadmill": {
       "allOf": [
         {
-          "$ref": "#/definitions/aind_behavior_vr_foraging__rig__Treadmill"
+          "$ref": "#/definitions/HarpTreadmill"
         }
       ],
-      "description": "Treadmill settings"
+      "description": "Harp treadmill"
     },
     "harp_sniff_detector": {
       "default": null,
@@ -1461,7 +1450,7 @@
     "harp_olfactometer",
     "harp_lickometer",
     "harp_clock_generator",
-    "treadmill",
+    "harp_treadmill",
     "calibration"
   ],
   "title": "AindVrForagingRig",