class method for creating empty dataframe based on pydantic model

iblrig/base_choice_world.py

@@ -99,26 +99,7 @@ def __init__(self, *args, delay_secs=0, **kwargs):
         self.block_num = -1
         self.block_trial_num = -1
         # init the tables, there are 2 of them: a trials table and a ambient sensor data table
-        self.trials_table = pd.DataFrame(
-            {
-                'contrast': np.zeros(NTRIALS_INIT) * np.NaN,
-                'position': np.zeros(NTRIALS_INIT) * np.NaN,
-                'quiescent_period': np.zeros(NTRIALS_INIT) * np.NaN,
-                'response_side': np.zeros(NTRIALS_INIT, dtype=np.int8),
-                'response_time': np.zeros(NTRIALS_INIT) * np.NaN,
-                'reward_amount': np.zeros(NTRIALS_INIT) * np.NaN,
-                'reward_valve_time': np.zeros(NTRIALS_INIT) * np.NaN,
-                'stim_angle': np.zeros(NTRIALS_INIT) * np.NaN,
-                'stim_freq': np.zeros(NTRIALS_INIT) * np.NaN,
-                'stim_gain': np.zeros(NTRIALS_INIT) * np.NaN,
-                'stim_phase': np.zeros(NTRIALS_INIT) * np.NaN,
-                'stim_reverse': np.zeros(NTRIALS_INIT, dtype=bool),
-                'stim_sigma': np.zeros(NTRIALS_INIT) * np.NaN,
-                'trial_correct': np.zeros(NTRIALS_INIT, dtype=bool),
-                'trial_num': np.zeros(NTRIALS_INIT, dtype=np.int16),
-                'pause_duration': np.zeros(NTRIALS_INIT, dtype=float),
-            }
-        )
+        self.trials_table = TrialData.prepare_dataframe(NTRIALS_INIT)
 
         self.ambient_sensor_table = pd.DataFrame(
             {

iblrig/pydantic_definitions.py

@@ -3,6 +3,8 @@
 from pathlib import Path
 from typing import Annotated, Literal
 
+import numpy as np
+import pandas as pd
 from annotated_types import Ge, Le
 from pydantic import (
     AnyUrl,
@@ -20,7 +22,7 @@
 
 from iblrig.constants import BASE_PATH
 
-FilePath = Annotated[FilePath, PlainSerializer(lambda s: str(s), return_type=str)]
+ExistingFilePath = Annotated[FilePath, PlainSerializer(lambda s: str(s), return_type=str)]
 """Validate that path exists and is file. Cast to str upon save."""
 
 BehaviourInputPort = Annotated[int, Ge(1), Le(4)]
@@ -108,7 +110,7 @@ class HardwareSettingsRotaryEncoder(BunchModel):
 
 
 class HardwareSettingsScreen(BunchModel):
-    DISPLAY_IDX: int = Field(gte=0, lte=1)  # -1 = Default, 0 = First, 1 = Second, 2 = Third, etc
+    DISPLAY_IDX: int = Field(ge=0, le=1)  # -1 = Default, 0 = First, 1 = Second, 2 = Third, etc
     SCREEN_FREQ_TARGET: int = Field(gt=0)
     SCREEN_FREQ_TEST_DATE: date | None = None
     SCREEN_FREQ_TEST_STATUS: str | None = None
@@ -161,12 +163,12 @@ class HardwareSettingsCamera(BunchModel):
 
 
 class HardwareSettingsCameraWorkflow(BunchModel):
-    setup: FilePath | None = Field(
+    setup: ExistingFilePath | None = Field(
         title='Optional camera setup workflow',
         default=None,
         description='An optional path to the camera setup Bonsai workflow.',
     )
-    recording: FilePath = Field(
+    recording: ExistingFilePath = Field(
         title='Camera recording workflow', description='The path to the Bonsai workflow for camera recording.'
     )
 
@@ -201,14 +203,21 @@ class HardwareSettings(BunchModel):
     VERSION: str
 
 
-class TrialData(BaseModel):
+class TrialDataModel(BaseModel):
     # allow adding extra fields
     model_config = ConfigDict(extra='allow')
 
+    @classmethod
+    def prepare_dataframe(cls, n_rows: int) -> pd.DataFrame:
+        dtypes = {field: field_info.annotation for field, field_info in cls.model_fields.items()}
+        return pd.DataFrame(np.zeros((n_rows, len(dtypes))), columns=dtypes.keys()).astype(dtypes)
+
+
+class TrialData(TrialDataModel):
     contrast: Annotated[float, Ge(0.0), Le(1.0)]
-    position: int
+    position: float
     quiescent_period: Annotated[float, Ge(0.0)]
-    response_side: Literal[-1, 0, 1]
+    response_side: Annotated[int, Ge(-1), Le(1)]
     response_time: Annotated[float, Ge(0.0)]
     reward_amount: Annotated[float, Ge(0.0)]
     reward_valve_time: Annotated[float, Ge(0.0)]