Merge pull request #460 from AllenNeuralDynamics/feat-add-water

Add water calculation to launcher and aind-data-schema

Author: bruno-f-cruz

src/DataSchemas/aind_behavior_vr_foraging.json

@@ -5242,7 +5242,7 @@
           "description": "The optional reward specification of the virtual site",
           "oneOf": [
             {
-              "$ref": "#/$defs/RewardSpecification"
+              "$ref": "#/$defs/VirtualSiteRewardSpecification"
             },
             {
               "type": "null"
@@ -5328,6 +5328,33 @@
       "title": "VirtualSiteLabels",
       "type": "string"
     },
+    "VirtualSiteRewardSpecification": {
+      "description": "THIS CLASS IS NOT MEANT TO BE DIRECTLY INSTANTIATED.\nSpecifies reward parameters and behavior for a virtual site.\n\nNote: This class is primarily used internally for runtime site generation\nand is not meant to be directly instantiated in task configuration DSL.",
+      "properties": {
+        "operant_logic": {
+          "$ref": "#/$defs/OperantLogic",
+          "description": "The optional operant logic of the reward"
+        },
+        "delay": {
+          "$ref": "#/$defs/Distribution",
+          "default": {
+            "family": "Scalar",
+            "distribution_parameters": {
+              "family": "Scalar",
+              "value": 0.0
+            },
+            "truncation_parameters": null,
+            "scaling_parameters": null
+          },
+          "description": "The optional distribution where the delay to reward will be drawn from"
+        }
+      },
+      "required": [
+        "operant_logic"
+      ],
+      "title": "VirtualSiteRewardSpecification",
+      "type": "object"
+    },
     "VisualCorridor": {
       "description": "Defines a visual corridor segment in the VR environment.\n\nVisual corridors are the basic building blocks of the VR environment,\ndefining spatial regions with specific textures, dimensions, and positions.",
       "properties": {

src/Extensions/AddRewardSite.bonsai

@@ -5,7 +5,6 @@
                  xmlns:p1="clr-namespace:AindVrForagingDataSchema;assembly=Extensions"
                  xmlns:sys="clr-namespace:System;assembly=mscorlib"
                  xmlns:num="clr-namespace:Bonsai.Numerics;assembly=Bonsai.Numerics"
-                 xmlns:p2="clr-namespace:;assembly=Extensions"
                  xmlns="https://bonsai-rx.org/2018/workflow">
   <Workflow>
     <Nodes>
@@ -114,67 +113,17 @@
                 <Property Name="RenderSpecification" />
               </PropertyMappings>
             </Expression>
-            <Expression xsi:type="GroupWorkflow">
-              <Name>ReconstructReward</Name>
-              <Workflow>
-                <Nodes>
-                  <Expression xsi:type="SubscribeSubject">
-                    <Name>ActivePatch</Name>
-                  </Expression>
-                  <Expression xsi:type="PropertyMapping">
-                    <PropertyMappings>
-                      <Property Name="OperantLogic" Selector="RewardSpecification.OperantLogic" />
-                      <Property Name="Delay" Selector="RewardSpecification.Delay" />
-                      <Property Name="RewardFunction" Selector="RewardSpecification.RewardFunction" />
-                    </PropertyMappings>
-                  </Expression>
-                  <Expression xsi:type="Combinator">
-                    <Combinator xsi:type="p2:NullDistribution" />
-                  </Expression>
-                  <Expression xsi:type="PropertyMapping">
-                    <PropertyMappings>
-                      <Property Name="Amount" />
-                      <Property Name="Probability" />
-                      <Property Name="Available" />
-                    </PropertyMappings>
-                  </Expression>
-                  <Expression xsi:type="Combinator">
-                    <Combinator xsi:type="p1:RewardSpecification" />
-                  </Expression>
-                  <Expression xsi:type="WorkflowOutput" />
-                </Nodes>
-                <Edges>
-                  <Edge From="0" To="1" Label="Source1" />
-                  <Edge From="1" To="4" Label="Source1" />
-                  <Edge From="2" To="3" Label="Source1" />
-                  <Edge From="3" To="4" Label="Source2" />
-                  <Edge From="4" To="5" Label="Source1" />
-                </Edges>
-              </Workflow>
-            </Expression>
             <Expression xsi:type="SubscribeSubject">
               <Name>ActivePatch</Name>
             </Expression>
-            <Expression xsi:type="MemberSelector">
-              <Selector>RewardSpecification</Selector>
-            </Expression>
-            <Expression xsi:type="IncludeWorkflow" Path="AllenNeuralDynamics.Core:FilterNotNull.bonsai" />
-            <Expression xsi:type="Combinator">
-              <Combinator xsi:type="rx:SubscribeWhen" />
-            </Expression>
-            <Expression xsi:type="SubscribeSubject">
-              <Name>ActivePatch</Name>
-            </Expression>
-            <Expression xsi:type="MemberSelector">
-              <Selector>RewardSpecification</Selector>
-            </Expression>
-            <Expression xsi:type="Combinator">
-              <Combinator xsi:type="rx:Merge" />
+            <Expression xsi:type="PropertyMapping">
+              <PropertyMappings>
+                <Property Name="OperantLogic" Selector="RewardSpecification.OperantLogic" />
+                <Property Name="Delay" Selector="RewardSpecification.Delay" />
+              </PropertyMappings>
             </Expression>
             <Expression xsi:type="Combinator">
-              <Combinator xsi:type="rx:Take">
-                <rx:Count>1</rx:Count>
-              </Combinator>
+              <Combinator xsi:type="p1:VirtualSiteRewardSpecification" />
             </Expression>
             <Expression xsi:type="PropertyMapping">
               <PropertyMappings>
@@ -215,26 +164,20 @@
             <Edge From="15" To="16" Label="Source4" />
             <Edge From="16" To="17" Label="Source1" />
             <Edge From="17" To="18" Label="Source1" />
-            <Edge From="18" To="36" Label="Source1" />
+            <Edge From="18" To="30" Label="Source1" />
             <Edge From="19" To="20" Label="Source1" />
             <Edge From="20" To="21" Label="Source1" />
-            <Edge From="21" To="36" Label="Source2" />
+            <Edge From="21" To="30" Label="Source2" />
             <Edge From="22" To="23" Label="Source1" />
             <Edge From="23" To="24" Label="Source1" />
-            <Edge From="24" To="36" Label="Source3" />
-            <Edge From="25" To="29" Label="Source1" />
+            <Edge From="24" To="30" Label="Source3" />
+            <Edge From="25" To="26" Label="Source1" />
             <Edge From="26" To="27" Label="Source1" />
             <Edge From="27" To="28" Label="Source1" />
-            <Edge From="28" To="29" Label="Source2" />
-            <Edge From="29" To="32" Label="Source1" />
+            <Edge From="28" To="30" Label="Source4" />
+            <Edge From="29" To="30" Label="Source5" />
             <Edge From="30" To="31" Label="Source1" />
             <Edge From="31" To="32" Label="Source2" />
-            <Edge From="32" To="33" Label="Source1" />
-            <Edge From="33" To="34" Label="Source1" />
-            <Edge From="34" To="36" Label="Source4" />
-            <Edge From="35" To="36" Label="Source5" />
-            <Edge From="36" To="37" Label="Source1" />
-            <Edge From="37" To="38" Label="Source2" />
           </Edges>
         </Workflow>
       </Expression>

src/Extensions/AindBehaviorVrForaging.Generated.cs

@@ -13795,7 +13795,7 @@ public partial class VirtualSite
 
         private OdorSpecification _odorSpecification;
 
-        private RewardSpecification _rewardSpecification;
+        private VirtualSiteRewardSpecification _rewardSpecification;
 
         private RenderSpecification _renderSpecification;
 
@@ -13915,7 +13915,7 @@ public OdorSpecification OdorSpecification
         [System.Xml.Serialization.XmlIgnoreAttribute()]
         [Newtonsoft.Json.JsonPropertyAttribute("reward_specification")]
         [System.ComponentModel.DescriptionAttribute("The optional reward specification of the virtual site")]
-        public RewardSpecification RewardSpecification
+        public VirtualSiteRewardSpecification RewardSpecification
         {
             get
             {
@@ -14176,6 +14176,108 @@ public enum VirtualSiteLabels
     }
 
 
+    /// <summary>
+    /// THIS CLASS IS NOT MEANT TO BE DIRECTLY INSTANTIATED.
+    ///Specifies reward parameters and behavior for a virtual site.
+    ///
+    ///Note: This class is primarily used internally for runtime site generation
+    ///and is not meant to be directly instantiated in task configuration DSL.
+    /// </summary>
+    [System.CodeDom.Compiler.GeneratedCodeAttribute("Bonsai.Sgen", "0.6.1.0 (Newtonsoft.Json v13.0.0.0)")]
+    [System.ComponentModel.DescriptionAttribute(@"THIS CLASS IS NOT MEANT TO BE DIRECTLY INSTANTIATED.
+    Specifies reward parameters and behavior for a virtual site.
+
+    Note: This class is primarily used internally for runtime site generation
+    and is not meant to be directly instantiated in task configuration DSL.")]
+    [Bonsai.WorkflowElementCategoryAttribute(Bonsai.ElementCategory.Source)]
+    [Bonsai.CombinatorAttribute(MethodName="Generate")]
+    public partial class VirtualSiteRewardSpecification
+    {
+    
+        private OperantLogic _operantLogic;
+    
+        private Distribution _delay;
+    
+        public VirtualSiteRewardSpecification()
+        {
+            _operantLogic = new OperantLogic();
+            _delay = new Distribution();
+        }
+    
+        protected VirtualSiteRewardSpecification(VirtualSiteRewardSpecification other)
+        {
+            _operantLogic = other._operantLogic;
+            _delay = other._delay;
+        }
+    
+        /// <summary>
+        /// The optional operant logic of the reward
+        /// </summary>
+        [System.Xml.Serialization.XmlIgnoreAttribute()]
+        [Newtonsoft.Json.JsonPropertyAttribute("operant_logic", Required=Newtonsoft.Json.Required.Always)]
+        [System.ComponentModel.DescriptionAttribute("The optional operant logic of the reward")]
+        public OperantLogic OperantLogic
+        {
+            get
+            {
+                return _operantLogic;
+            }
+            set
+            {
+                _operantLogic = value;
+            }
+        }
+    
+        /// <summary>
+        /// The optional distribution where the delay to reward will be drawn from
+        /// </summary>
+        [System.Xml.Serialization.XmlIgnoreAttribute()]
+        [Newtonsoft.Json.JsonPropertyAttribute("delay")]
+        [System.ComponentModel.DescriptionAttribute("The optional distribution where the delay to reward will be drawn from")]
+        public Distribution Delay
+        {
+            get
+            {
+                return _delay;
+            }
+            set
+            {
+                _delay = value;
+            }
+        }
+    
+        public System.IObservable<VirtualSiteRewardSpecification> Generate()
+        {
+            return System.Reactive.Linq.Observable.Defer(() => System.Reactive.Linq.Observable.Return(new VirtualSiteRewardSpecification(this)));
+        }
+    
+        public System.IObservable<VirtualSiteRewardSpecification> Generate<TSource>(System.IObservable<TSource> source)
+        {
+            return System.Reactive.Linq.Observable.Select(source, _ => new VirtualSiteRewardSpecification(this));
+        }
+    
+        protected virtual bool PrintMembers(System.Text.StringBuilder stringBuilder)
+        {
+            stringBuilder.Append("OperantLogic = " + _operantLogic + ", ");
+            stringBuilder.Append("Delay = " + _delay);
+            return true;
+        }
+    
+        public override string ToString()
+        {
+            System.Text.StringBuilder stringBuilder = new System.Text.StringBuilder();
+            stringBuilder.Append(GetType().Name);
+            stringBuilder.Append(" { ");
+            if (PrintMembers(stringBuilder))
+            {
+                stringBuilder.Append(" ");
+            }
+            stringBuilder.Append("}");
+            return stringBuilder.ToString();
+        }
+    }
+
+
     /// <summary>
     /// Defines a visual corridor segment in the VR environment.
     ///
@@ -17138,6 +17240,11 @@ public System.IObservable<string> Process(System.IObservable<VirtualSiteGenerato
             return Process<VirtualSiteGenerator>(source);
         }
 
+        public System.IObservable<string> Process(System.IObservable<VirtualSiteRewardSpecification> source)
+        {
+            return Process<VirtualSiteRewardSpecification>(source);
+        }
+
         public System.IObservable<string> Process(System.IObservable<VisualCorridor> source)
         {
             return Process<VisualCorridor>(source);
@@ -17346,6 +17453,7 @@ public System.IObservable<string> Process(System.IObservable<ScalingParameters2>
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<VideoWriterOpenCv>))]
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<VirtualSite>))]
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<VirtualSiteGenerator>))]
+    [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<VirtualSiteRewardSpecification>))]
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<VisualCorridor>))]
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<WallTextures>))]
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<WaterValveCalibration>))]

src/Extensions/TriggeredCameras.bonsai

@@ -1,5 +1,5 @@
 <?xml version="1.0" encoding="utf-8"?>
-<WorkflowBuilder Version="2.8.5"
+<WorkflowBuilder Version="2.9.0"
                  xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
                  xmlns:rx="clr-namespace:Bonsai.Reactive;assembly=Bonsai.Core"
                  xmlns:p1="clr-namespace:AindVrForagingDataSchema;assembly=Extensions"
@@ -142,7 +142,7 @@
               <VideoExtension>mp4</VideoExtension>
               <OutputArguments>-vf "scale=out_color_matrix=bt709:out_range=full" -c:v h264_nvenc -pix_fmt nv12 -color_range full -colorspace bt709 -color_trc linear -tune hq -preset p4 -rc vbr -cq 12 -b:v 0M -maxrate 700M -bufsize 350M</OutputArguments>
               <InputArguments>-colorspace rgb -color_primaries bt709 -color_trc linear</InputArguments>
-              <Verbosity>Verbose</Verbosity>
+              <Verbosity>Warning</Verbosity>
               <FrameRate>0</FrameRate>
             </Expression>
           </Nodes>

src/aind_behavior_vr_foraging/cli.py

@@ -1,14 +1,46 @@
+import os
 import typing as t
+from pathlib import Path
 
-from pydantic import Field, RootModel
-from pydantic_settings import BaseSettings, CliApp, CliSubCommand
+from pydantic import AwareDatetime, Field, RootModel
+from pydantic_settings import BaseSettings, CliApp, CliPositionalArg, CliSubCommand
 
 from aind_behavior_vr_foraging import __semver__, regenerate
-from aind_behavior_vr_foraging.data_mappers import DataMapperCli
-from aind_behavior_vr_foraging.data_qc import DataQcCli
 from aind_behavior_vr_foraging.launcher import ClabeCli
 
 
+class DataMapperCli(BaseSettings, cli_kebab_case=True):
+    data_path: os.PathLike = Field(description="Path to the session data directory.")
+    repo_path: os.PathLike = Field(
+        default=Path("."), description="Path to the repository. By default it will use the current directory."
+    )
+    curriculum_suggestion: t.Optional[os.PathLike] = Field(
+        default=None, description="Path to curriculum suggestion file."
+    )
+    session_end_time: AwareDatetime | None = Field(
+        default=None,
+        description="End time of the session in ISO format. If not provided, will use the time the data mapping is run.",
+    )
+
+    def cli_cmd(self):
+        from aind_behavior_vr_foraging.data_mappers import cli_cmd
+
+        return cli_cmd(self)
+
+
+class DataQcCli(BaseSettings, cli_kebab_case=True):
+    data_path: CliPositionalArg[os.PathLike] = Field(description="Path to the session data directory.")
+    version: str = Field(default=__semver__, description="Version of the dataset.")
+    report_path: Path | None = Field(
+        default=None, description="Path to save the Html QC report. If not provided, report is not saved."
+    )
+
+    def cli_cmd(self):
+        from aind_behavior_vr_foraging.data_qc import cli_cmd
+
+        return cli_cmd(self)
+
+
 class VersionCli(RootModel):
     root: t.Any
 

src/aind_behavior_vr_foraging/data_contract/utils.py

@@ -0,0 +1,25 @@
+import os
+from typing import Optional
+
+
+def calculate_consumed_water(session_path: os.PathLike) -> Optional[float]:
+    """Calculate the total volume of water consumed during a session.
+
+    Args:
+        session_path (os.PathLike): Path to the session directory.
+
+    Returns:
+        Optional[float]: Total volume of water consumed in milliliters, or None if unavailable.
+    """
+    from aind_behavior_vr_foraging.data_contract import dataset
+
+    reward = dataset(session_path)["Behavior"]["SoftwareEvents"]["GiveReward"].load()
+    extra = dataset(session_path)["Behavior"]["SoftwareEvents"]["ForceGiveReward"].load()
+    total = 0
+    if reward.has_data is False and extra.has_data is False:
+        return None
+    if reward.has_data:
+        total += reward.data["data"].sum() * 1e-3
+    if extra.has_data:
+        total += extra.data["data"].sum() * 1e-3
+    return total