socr / tests /atomic /test_atomic_vaep.py
scfive's picture
Upload 203 files
d6ea71e verified
import pandas as pd
import pytest
import socceraction.atomic.spadl as atomicspadl
import socceraction.atomic.spadl.config as spadlconfig
import socceraction.atomic.vaep.labels as lab
from socceraction.atomic.vaep import AtomicVAEP
from socceraction.atomic.vaep import features as fs
@pytest.fixture
def test_goal_df() -> pd.DataFrame:
return pd.DataFrame(
[spadlconfig.actiontypes.index("shot"), spadlconfig.actiontypes.index("goal")],
columns=["type_id"],
)
def test_atomic_goal_from_shot_label(test_goal_df: pd.DataFrame) -> None:
assert (lab.goal_from_shot(test_goal_df) == pd.DataFrame([[True], [False]], columns=["goal"]))[
"goal"
].all()
@pytest.mark.e2e
def test_predict(sb_worldcup_data: pd.HDFStore) -> None:
# Convert to atomic actions
games = sb_worldcup_data["games"]
atomic_actions = {
game.game_id: atomicspadl.convert_to_atomic(
sb_worldcup_data[f"actions/game_{game.game_id}"]
)
for game in games.itertuples()
}
# Test the vAEP framework on the StatsBomb World Cup data
model = AtomicVAEP(nb_prev_actions=1)
# comppute features and labels
features = pd.concat(
[
model.compute_features(game, atomic_actions[game.game_id])
for game in games.iloc[:-1].itertuples()
]
)
expected_features = set(fs.feature_column_names(model.xfns, model.nb_prev_actions))
assert set(features.columns) == expected_features
labels = pd.concat(
[
model.compute_labels(game, atomic_actions[game.game_id])
for game in games.iloc[:-1].itertuples()
]
)
expected_labels = {"scores", "concedes"}
assert set(labels.columns) == expected_labels
assert len(features) == len(labels)
# fit the model
model.fit(features, labels)
# rate a game
game = games.iloc[-1]
actions = atomic_actions[game.game_id]
ratings = model.rate(game, actions)
expected_rating_columns = {"offensive_value", "defensive_value", "vaep_value"}
assert set(ratings.columns) == expected_rating_columns