File size: 1,325 Bytes
d6ea71e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
"""Schema for Atomic-SPADL actions."""
from typing import Any, Optional
import pandera as pa
from pandera.typing import Series
from . import config as spadlconfig
class AtomicSPADLSchema(pa.SchemaModel):
"""Definition of an Atomic-SPADL dataframe."""
game_id: Series[Any] = pa.Field()
original_event_id: Series[Any] = pa.Field(nullable=True)
action_id: Series[int] = pa.Field()
period_id: Series[int] = pa.Field(ge=1, le=5)
time_seconds: Series[float] = pa.Field(ge=0)
team_id: Series[Any] = pa.Field()
player_id: Series[Any] = pa.Field()
x: Series[float] = pa.Field(ge=0, le=spadlconfig.field_length)
y: Series[float] = pa.Field(ge=0, le=spadlconfig.field_width)
dx: Series[float] = pa.Field(ge=-spadlconfig.field_length, le=spadlconfig.field_length)
dy: Series[float] = pa.Field(ge=-spadlconfig.field_width, le=spadlconfig.field_width)
bodypart_id: Series[int] = pa.Field(isin=spadlconfig.bodyparts_df().bodypart_id)
bodypart_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.bodyparts_df().bodypart_name)
type_id: Series[int] = pa.Field(isin=spadlconfig.actiontypes_df().type_id)
type_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.actiontypes_df().type_name)
class Config: # noqa: D106
strict = True
coerce = True
|