File size: 1,034 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 34 |
import socceraction.atomic.spadl as spadlcfg
from pandas import testing as tm
from pandera.typing import DataFrame
from socceraction.atomic.spadl import AtomicSPADLSchema
from socceraction.atomic.vaep import features as fs
xfns = [
fs.actiontype,
fs.actiontype_onehot,
fs.bodypart,
fs.bodypart_detailed,
fs.bodypart_onehot,
fs.bodypart_detailed_onehot,
fs.team,
fs.time,
fs.time_delta,
fs.location,
fs.polar,
fs.movement_polar,
fs.direction,
fs.goalscore,
]
def test_same_index(atomic_spadl_actions: DataFrame[AtomicSPADLSchema]) -> None:
"""The feature generators should not change the index of the input dataframe."""
atomic_spadl_actions.index += 10
game_actions_with_names = spadlcfg.add_names(atomic_spadl_actions)
gamestates = fs.gamestates(game_actions_with_names, 3)
gamestates = fs.play_left_to_right(gamestates, 782)
for fn in xfns:
features = fn(gamestates)
tm.assert_index_equal(features.index, atomic_spadl_actions.index)
|