From 947df8595695e07b218b7465213c7b4d02ffa828 Mon Sep 17 00:00:00 2001 From: Florian Rau Date: Fri, 30 Aug 2024 18:46:04 +0100 Subject: [PATCH] class method for creating empty dataframe based on pydantic model --- iblrig/base_choice_world.py | 21 +-------------------- iblrig/pydantic_definitions.py | 23 ++++++++++++++++------- 2 files changed, 17 insertions(+), 27 deletions(-) diff --git a/iblrig/base_choice_world.py b/iblrig/base_choice_world.py index 99c5ad957..fc976de18 100644 --- a/iblrig/base_choice_world.py +++ b/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( { diff --git a/iblrig/pydantic_definitions.py b/iblrig/pydantic_definitions.py index 17ead10d2..5f14daa44 100644 --- a/iblrig/pydantic_definitions.py +++ b/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)]