behaviorstim updates issue #578

Author: saskiad

src/aind_data_schema/session.py

@@ -16,7 +16,7 @@
 from aind_data_schema.device import Calibration, Maintenance, RelativePosition, SpoutSide
 from aind_data_schema.imaging.tile import Channel
 from aind_data_schema.stimulus import StimulusEpoch
-from aind_data_schema.utils.units import AngleUnit, FrequencyUnit, MassUnit, PowerUnit, SizeUnit, TimeUnit
+from aind_data_schema.utils.units import AngleUnit, FrequencyUnit, MassUnit, PowerUnit, SizeUnit, TimeUnit, VolumeUnit
 
 
 # Ophys components
@@ -361,4 +361,6 @@ class Session(AindCoreModel):
     )
     stimulus_epochs: Optional[List[StimulusEpoch]] = Field(None, title="Stimulus")
     reward_delivery: Optional[RewardDelivery] = Field(None, title="Reward delivery")
+    reward_consumed_total: Optional[Decimal] = Field(..., title="Total reward consumed (uL)")
+    reward_consumed_unit: VolumeUnit = Field(VolumeUnit.UL, title="Reward consumed unit")
     notes: Optional[str] = Field(None, title="Notes")

src/aind_data_schema/stimulus.py

@@ -112,12 +112,11 @@ class BehaviorStimulation(AindModel):
         title="Performance parameters",
         description="Performance metrics from session",
     )
-    reward_consumed_during_training: Decimal = Field(..., title="Reward consumed during training (uL)")
-    reward_consumed_total: Decimal = Field(..., title="Total reward consumed (uL)")
+    reward_consumed_during_epoch: Decimal = Field(..., title="Reward consumed during training (uL)")
     reward_consumed_unit: VolumeUnit = Field(VolumeUnit.UL, title="Reward consumed unit")
-    trials_total: int = Field(..., title="Total trials")
-    trials_finished: int = Field(..., title="Finished trials")
-    trials_rewarded: int = Field(..., title="Rewarded trials")
+    trials_total: Optional[int] = Field(..., title="Total trials")
+    trials_finished: Optional[int] = Field(..., title="Finished trials")
+    trials_rewarded: Optional[int] = Field(..., title="Rewarded trials")
     notes: Optional[str] = Field(None, title="Notes")