adds autowater logic

schema/aind_behavior_dynamic_foraging.json

@@ -433,6 +433,34 @@
       "title": "AuditorySecondaryReinforcer",
       "type": "object"
     },
+    "AutoWaterParameters": {
+      "properties": {
+        "min_ignored_trials": {
+          "default": 3,
+          "description": "Minimum consecutive ignored trials before auto water is triggered.",
+          "minimum": 0,
+          "title": "Min Ignored Trials",
+          "type": "integer"
+        },
+        "min_unrewarded_trials": {
+          "default": 3,
+          "description": "Minimum consecutive unrewarded trials before auto water is triggered.",
+          "minimum": 0,
+          "title": "Min Unrewarded Trials",
+          "type": "integer"
+        },
+        "reward_fraction": {
+          "default": 0.8,
+          "description": "Fraction of full reward volume delivered during auto water (0=none, 1=full).",
+          "maximum": 1,
+          "minimum": 0,
+          "title": "Reward Fraction",
+          "type": "number"
+        }
+      },
+      "title": "AutoWaterParameters",
+      "type": "object"
+    },
     "Axis": {
       "description": "Motor axis available",
       "enum": [
@@ -958,6 +986,22 @@
           },
           "description": "Parameters defining the reward probability structure."
         },
+        "autowater_parameters": {
+          "default": {
+            "min_ignored_trials": 3,
+            "min_unrewarded_trials": 3,
+            "reward_fraction": 0.8
+          },
+          "description": "Auto water settings. If set, free water is delivered when the animal exceeds the ignored or unrewarded trial thresholds.",
+          "oneOf": [
+            {
+              "$ref": "#/$defs/AutoWaterParameters"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        },
         "is_baiting": {
           "default": false,
           "description": "Whether uncollected rewards carry over to the next trial.",

src/Extensions/AindBehaviorDynamicForaging.Generated.cs

@@ -1060,6 +1060,116 @@ public override string ToString()
     }
 
 
+    [System.CodeDom.Compiler.GeneratedCodeAttribute("Bonsai.Sgen", "0.7.2.0 (Newtonsoft.Json v13.0.0.0)")]
+    [Bonsai.WorkflowElementCategoryAttribute(Bonsai.ElementCategory.Source)]
+    [Bonsai.CombinatorAttribute(MethodName="Generate")]
+    public partial class AutoWaterParameters
+    {
+
+        private int _minIgnoredTrials;
+
+        private int _minUnrewardedTrials;
+
+        private double _rewardFraction;
+
+        public AutoWaterParameters()
+        {
+            _minIgnoredTrials = 3;
+            _minUnrewardedTrials = 3;
+            _rewardFraction = 0.8D;
+        }
+
+        protected AutoWaterParameters(AutoWaterParameters other)
+        {
+            _minIgnoredTrials = other._minIgnoredTrials;
+            _minUnrewardedTrials = other._minUnrewardedTrials;
+            _rewardFraction = other._rewardFraction;
+        }
+
+        /// <summary>
+        /// Minimum consecutive ignored trials before auto water is triggered.
+        /// </summary>
+        [Newtonsoft.Json.JsonPropertyAttribute("min_ignored_trials")]
+        [System.ComponentModel.DescriptionAttribute("Minimum consecutive ignored trials before auto water is triggered.")]
+        public int MinIgnoredTrials
+        {
+            get
+            {
+                return _minIgnoredTrials;
+            }
+            set
+            {
+                _minIgnoredTrials = value;
+            }
+        }
+
+        /// <summary>
+        /// Minimum consecutive unrewarded trials before auto water is triggered.
+        /// </summary>
+        [Newtonsoft.Json.JsonPropertyAttribute("min_unrewarded_trials")]
+        [System.ComponentModel.DescriptionAttribute("Minimum consecutive unrewarded trials before auto water is triggered.")]
+        public int MinUnrewardedTrials
+        {
+            get
+            {
+                return _minUnrewardedTrials;
+            }
+            set
+            {
+                _minUnrewardedTrials = value;
+            }
+        }
+
+        /// <summary>
+        /// Fraction of full reward volume delivered during auto water (0=none, 1=full).
+        /// </summary>
+        [Newtonsoft.Json.JsonPropertyAttribute("reward_fraction")]
+        [System.ComponentModel.DescriptionAttribute("Fraction of full reward volume delivered during auto water (0=none, 1=full).")]
+        public double RewardFraction
+        {
+            get
+            {
+                return _rewardFraction;
+            }
+            set
+            {
+                _rewardFraction = value;
+            }
+        }
+
+        public System.IObservable<AutoWaterParameters> Generate()
+        {
+            return System.Reactive.Linq.Observable.Defer(() => System.Reactive.Linq.Observable.Return(new AutoWaterParameters(this)));
+        }
+
+        public System.IObservable<AutoWaterParameters> Generate<TSource>(System.IObservable<TSource> source)
+        {
+            return System.Reactive.Linq.Observable.Select(source, _ => new AutoWaterParameters(this));
+        }
+
+        protected virtual bool PrintMembers(System.Text.StringBuilder stringBuilder)
+        {
+            stringBuilder.Append("MinIgnoredTrials = " + _minIgnoredTrials + ", ");
+            stringBuilder.Append("MinUnrewardedTrials = " + _minUnrewardedTrials + ", ");
+            stringBuilder.Append("RewardFraction = " + _rewardFraction);
+            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>
     /// Motor axis available
     /// </summary>
@@ -2416,6 +2526,8 @@ public partial class CoupledTrialGeneratorSpec : TrialGeneratorSpec
 
         private RewardProbabilityParameters _rewardProbabilityParameters;
 
+        private AutoWaterParameters _autowaterParameters;
+
         private bool _isBaiting;
 
         private CoupledTrialGenerationEndConditions _trialGenerationEndParameters;
@@ -2434,6 +2546,7 @@ public CoupledTrialGeneratorSpec()
             _minBlockReward = 1;
             _kernelSize = 2;
             _rewardProbabilityParameters = new RewardProbabilityParameters();
+            _autowaterParameters = new AutoWaterParameters();
             _isBaiting = false;
             _trialGenerationEndParameters = new CoupledTrialGenerationEndConditions();
             _behaviorStabilityParameters = new BehaviorStabilityParameters();
@@ -2451,6 +2564,7 @@ protected CoupledTrialGeneratorSpec(CoupledTrialGeneratorSpec other) :
             _minBlockReward = other._minBlockReward;
             _kernelSize = other._kernelSize;
             _rewardProbabilityParameters = other._rewardProbabilityParameters;
+            _autowaterParameters = other._autowaterParameters;
             _isBaiting = other._isBaiting;
             _trialGenerationEndParameters = other._trialGenerationEndParameters;
             _behaviorStabilityParameters = other._behaviorStabilityParameters;
@@ -2594,6 +2708,25 @@ public RewardProbabilityParameters RewardProbabilityParameters
             }
         }
 
+        /// <summary>
+        /// Auto water settings. If set, free water is delivered when the animal exceeds the ignored or unrewarded trial thresholds.
+        /// </summary>
+        [System.Xml.Serialization.XmlIgnoreAttribute()]
+        [Newtonsoft.Json.JsonPropertyAttribute("autowater_parameters")]
+        [System.ComponentModel.DescriptionAttribute("Auto water settings. If set, free water is delivered when the animal exceeds the " +
+            "ignored or unrewarded trial thresholds.")]
+        public AutoWaterParameters AutowaterParameters
+        {
+            get
+            {
+                return _autowaterParameters;
+            }
+            set
+            {
+                _autowaterParameters = value;
+            }
+        }
+
         /// <summary>
         /// Whether uncollected rewards carry over to the next trial.
         /// </summary>
@@ -2690,6 +2823,7 @@ protected override bool PrintMembers(System.Text.StringBuilder stringBuilder)
             stringBuilder.Append("MinBlockReward = " + _minBlockReward + ", ");
             stringBuilder.Append("KernelSize = " + _kernelSize + ", ");
             stringBuilder.Append("RewardProbabilityParameters = " + _rewardProbabilityParameters + ", ");
+            stringBuilder.Append("AutowaterParameters = " + _autowaterParameters + ", ");
             stringBuilder.Append("IsBaiting = " + _isBaiting + ", ");
             stringBuilder.Append("TrialGenerationEndParameters = " + _trialGenerationEndParameters + ", ");
             stringBuilder.Append("BehaviorStabilityParameters = " + _behaviorStabilityParameters + ", ");
@@ -9508,6 +9642,11 @@ public System.IObservable<string> Process(System.IObservable<AuditorySecondaryRe
             return Process<AuditorySecondaryReinforcer>(source);
         }
 
+        public System.IObservable<string> Process(System.IObservable<AutoWaterParameters> source)
+        {
+            return Process<AutoWaterParameters>(source);
+        }
+
         public System.IObservable<string> Process(System.IObservable<AxisConfiguration> source)
         {
             return Process<AxisConfiguration>(source);
@@ -9833,6 +9972,7 @@ public System.IObservable<string> Process(System.IObservable<ScalingParameters2>
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<AindManipulator>))]
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<AindManipulatorCalibration>))]
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<AuditorySecondaryReinforcer>))]
+    [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<AutoWaterParameters>))]
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<AxisConfiguration>))]
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<BaseModel>))]
     [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping<BehaviorStabilityParameters>))]

src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py

@@ -20,6 +20,21 @@
 logger = logging.getLogger(__name__)
 
 
+class AutoWaterParameters(BaseModel):
+    min_ignored_trials: int = Field(
+        default=3, ge=0, description="Minimum consecutive ignored trials before auto water is triggered."
+    )
+    min_unrewarded_trials: int = Field(
+        default=3, ge=0, description="Minimum consecutive unrewarded trials before auto water is triggered."
+    )
+    reward_fraction: float = Field(
+        default=0.8,
+        ge=0,
+        le=1,
+        description="Fraction of full reward volume delivered during auto water (0=none, 1=full).",
+    )
+
+
 class RewardProbabilityParameters(BaseModel):
     """Defines the reward probability structure for a dynamic foraging task.
 
@@ -91,6 +106,12 @@ class BlockBasedTrialGeneratorSpec(BaseTrialGeneratorSpecModel):
         default=RewardProbabilityParameters(), description="Parameters defining the reward probability structure."
     )
 
+    autowater_parameters: Optional[AutoWaterParameters] = Field(
+        default=AutoWaterParameters(),
+        validate_default=True,
+        description="Auto water settings. If set, free water is delivered when the animal exceeds the ignored or unrewarded trial thresholds.",
+    )
+
     is_baiting: bool = Field(default=False, description="Whether uncollected rewards carry over to the next trial.")
 
     def create_generator(self) -> "BlockBasedTrialGenerator":
@@ -155,29 +176,57 @@ def next(self) -> Trial | None:
         iti = draw_sample(self.spec.inter_trial_interval_duration)
         quiescent = draw_sample(self.spec.quiescent_duration)
 
-        p_reward_left = self.block.p_left_reward
-        p_reward_right = self.block.p_right_reward
+        # determine baiting
+        random_numbers = np.random.random(2)
+        is_left_baited = self.block.p_left_reward > random_numbers[0]
+        is_right_baited = self.block.p_right_reward > random_numbers[1]
 
         if self.spec.is_baiting:
-            random_numbers = np.random.random(2)
-
-            is_left_baited = self.block.p_left_reward > random_numbers[0] or self.is_left_baited
-            logger.debug(f"Left baited: {is_left_baited}")
-            p_reward_left = 1 if is_left_baited else p_reward_left
-
-            is_right_baited = self.block.p_right_reward > random_numbers[1] or self.is_right_baited
-            logger.debug(f"Right baited: {is_left_baited}")
-            p_reward_right = 1 if is_right_baited else p_reward_right
-
+            is_left_baited = is_left_baited or self.is_left_baited
+            is_right_baited = is_right_baited or self.is_right_baited
+            logger.debug(f"Left baited: {is_left_baited}, Right baited: {is_right_baited}")
+
+        # determine autowater
+        is_autowater_trial = self._are_autowater_conditions_met()
+        is_left_autowater = is_left_baited and is_autowater_trial
+        is_right_autowater = is_right_baited and is_autowater_trial
+
+        p_reward_left = 1 if (is_left_baited or is_left_autowater) else self.block.p_left_reward
+        p_reward_right = 1 if (is_right_baited or is_right_autowater) else self.block.p_right_reward
+
         return Trial(
             p_reward_left=p_reward_left,
             p_reward_right=p_reward_right,
             reward_consumption_duration=self.spec.reward_consumption_duration,
             response_deadline_duration=self.spec.response_duration,
             quiescence_period_duration=quiescent,
             inter_trial_interval_duration=iti,
+            is_auto_response_right=is_right_autowater, 
         )
 
+    def _are_autowater_conditions_met(self) -> bool:
+        """Checks whether autowater should be given.
+
+        Returns:
+            True if autowater conditions are met, False otherwise.
+        """
+
+        if self.spec.autowater_parameters is None:  # autowater disabled
+            return False
+
+        min_ignore = self.spec.autowater_parameters.min_ignored_trials
+        min_unreward = self.spec.autowater_parameters.min_unrewarded_trials
+
+        is_ignored = [choice is None for choice in self.is_right_choice_history]
+        if all(is_ignored[-min_ignore:]):
+            return True
+
+        is_unrewarded = [not reward for reward in self.reward_history]
+        if all(is_unrewarded[-min_unreward:]):
+            return True
+
+        return False
+
     @abstractmethod
     def _are_end_conditions_met(self) -> bool:
         """Checks whether the session should end.
@@ -193,7 +242,7 @@ def _generate_next_block(
         reward_pairs: list[list[float, float]],
         base_reward_sum: float,
         block_len: Union[UniformDistribution, ExponentialDistribution],
-        current_block: Optional[None] = None,
+        current_block: Optional[Block] = None,
     ) -> Block:
         """Generates the next block, avoiding repeating the current block's side bias.