|
| 1 | +from pathlib import Path |
| 2 | +from unravel.soccer import SoccerGraphConverterPolars, KloppyPolarsDataset |
| 3 | +from unravel.utils import ( |
| 4 | + dummy_labels, |
| 5 | + dummy_graph_ids, |
| 6 | + CustomSpektralDataset, |
| 7 | +) |
| 8 | + |
| 9 | +from kloppy import skillcorner |
| 10 | +from kloppy.domain import Ground, TrackingDataset, Orientation |
| 11 | +from typing import List, Dict |
| 12 | + |
| 13 | +from spektral.data import Graph |
| 14 | + |
| 15 | +import pytest |
| 16 | + |
| 17 | +import numpy as np |
| 18 | + |
| 19 | + |
| 20 | +class TestKloppyPolarsData: |
| 21 | + @pytest.fixture |
| 22 | + def match_data(self, base_dir: Path) -> str: |
| 23 | + return base_dir / "files" / "skillcorner_match_data.json" |
| 24 | + |
| 25 | + @pytest.fixture |
| 26 | + def structured_data(self, base_dir: Path) -> str: |
| 27 | + return base_dir / "files" / "skillcorner_structured_data.json.gz" |
| 28 | + |
| 29 | + @pytest.fixture() |
| 30 | + def kloppy_dataset(self, match_data: str, structured_data: str) -> TrackingDataset: |
| 31 | + return skillcorner.load( |
| 32 | + raw_data=structured_data, |
| 33 | + meta_data=match_data, |
| 34 | + coordinates="tracab", |
| 35 | + include_empty_frames=False, |
| 36 | + limit=500, |
| 37 | + ) |
| 38 | + |
| 39 | + @pytest.fixture() |
| 40 | + def kloppy_polars_dataset( |
| 41 | + self, kloppy_dataset: TrackingDataset |
| 42 | + ) -> KloppyPolarsDataset: |
| 43 | + dataset = KloppyPolarsDataset( |
| 44 | + kloppy_dataset=kloppy_dataset, |
| 45 | + ball_carrier_threshold=25.0, |
| 46 | + ) |
| 47 | + dataset.load() |
| 48 | + dataset.add_dummy_labels(by=["game_id", "frame_id"]) |
| 49 | + dataset.add_graph_ids(by=["game_id", "frame_id"]) |
| 50 | + return dataset |
| 51 | + |
| 52 | + @pytest.fixture() |
| 53 | + def spc_padding( |
| 54 | + self, kloppy_polars_dataset: KloppyPolarsDataset |
| 55 | + ) -> SoccerGraphConverterPolars: |
| 56 | + return SoccerGraphConverterPolars( |
| 57 | + dataset=kloppy_polars_dataset, |
| 58 | + chunk_size=2_0000, |
| 59 | + non_potential_receiver_node_value=0.1, |
| 60 | + max_player_speed=12.0, |
| 61 | + max_player_acceleration=12.0, |
| 62 | + max_ball_speed=13.5, |
| 63 | + max_ball_acceleration=100, |
| 64 | + self_loop_ball=True, |
| 65 | + adjacency_matrix_connect_type="ball", |
| 66 | + adjacency_matrix_type="split_by_team", |
| 67 | + label_type="binary", |
| 68 | + defending_team_node_value=0.0, |
| 69 | + random_seed=False, |
| 70 | + pad=True, |
| 71 | + verbose=False, |
| 72 | + ) |
| 73 | + |
| 74 | + @pytest.fixture() |
| 75 | + def soccer_polars_converter( |
| 76 | + self, kloppy_polars_dataset: KloppyPolarsDataset |
| 77 | + ) -> SoccerGraphConverterPolars: |
| 78 | + |
| 79 | + return SoccerGraphConverterPolars( |
| 80 | + dataset=kloppy_polars_dataset, |
| 81 | + chunk_size=2_0000, |
| 82 | + non_potential_receiver_node_value=0.1, |
| 83 | + max_player_speed=12.0, |
| 84 | + max_player_acceleration=12.0, |
| 85 | + max_ball_speed=13.5, |
| 86 | + max_ball_acceleration=100, |
| 87 | + self_loop_ball=True, |
| 88 | + adjacency_matrix_connect_type="ball", |
| 89 | + adjacency_matrix_type="split_by_team", |
| 90 | + label_type="binary", |
| 91 | + defending_team_node_value=0.0, |
| 92 | + random_seed=False, |
| 93 | + pad=False, |
| 94 | + verbose=False, |
| 95 | + ) |
| 96 | + |
| 97 | + def test_padding(self, spc_padding: SoccerGraphConverterPolars): |
| 98 | + spektral_graphs = spc_padding.to_spektral_graphs() |
| 99 | + |
| 100 | + assert 1 == 1 |
| 101 | + |
| 102 | + data = spektral_graphs |
| 103 | + assert len(data) == 384 |
| 104 | + assert isinstance(data[0], Graph) |
| 105 | + |
| 106 | + def test_to_spektral_graph( |
| 107 | + self, soccer_polars_converter: SoccerGraphConverterPolars |
| 108 | + ): |
| 109 | + """ |
| 110 | + Test navigating (next/prev) through events |
| 111 | + """ |
| 112 | + spektral_graphs = soccer_polars_converter.to_spektral_graphs() |
| 113 | + |
| 114 | + assert 1 == 1 |
| 115 | + |
| 116 | + data = spektral_graphs |
| 117 | + assert data[0].id == "2417-1529" |
| 118 | + assert len(data) == 489 |
| 119 | + assert isinstance(data[0], Graph) |
| 120 | + |
| 121 | + x = data[0].x |
| 122 | + n_players = x.shape[0] |
| 123 | + assert x.shape == (n_players, 15) |
| 124 | + assert 0.4524340998288571 == pytest.approx(x[0, 0], abs=1e-5) |
| 125 | + assert 0.9948105277764999 == pytest.approx(x[0, 4], abs=1e-5) |
| 126 | + assert 0.2941671698429814 == pytest.approx(x[8, 2], abs=1e-5) |
| 127 | + |
| 128 | + e = data[0].e |
| 129 | + assert e.shape == (129, 6) |
| 130 | + assert 0.0 == pytest.approx(e[0, 0], abs=1e-5) |
| 131 | + assert 0.5 == pytest.approx(e[0, 4], abs=1e-5) |
| 132 | + assert 0.7140882876637022 == pytest.approx(e[8, 2], abs=1e-5) |
| 133 | + |
| 134 | + a = data[0].a |
| 135 | + assert a.shape == (n_players, n_players) |
| 136 | + assert 1.0 == pytest.approx(a[0, 0], abs=1e-5) |
| 137 | + assert 1.0 == pytest.approx(a[0, 4], abs=1e-5) |
| 138 | + assert 0.0 == pytest.approx(a[8, 2], abs=1e-5) |
| 139 | + |
| 140 | + dataset = CustomSpektralDataset(graphs=spektral_graphs) |
| 141 | + N, F, S, n_out, n = dataset.dimensions() |
| 142 | + assert N == 20 |
| 143 | + assert F == 15 |
| 144 | + assert S == 6 |
| 145 | + assert n_out == 1 |
| 146 | + assert n == 489 |
| 147 | + |
| 148 | + train, test, val = dataset.split_test_train_validation( |
| 149 | + split_train=4, |
| 150 | + split_test=1, |
| 151 | + split_validation=1, |
| 152 | + by_graph_id=True, |
| 153 | + random_seed=42, |
| 154 | + ) |
| 155 | + assert train.n_graphs == 326 |
| 156 | + assert test.n_graphs == 81 |
| 157 | + assert val.n_graphs == 82 |
| 158 | + |
| 159 | + train, test, val = dataset.split_test_train_validation( |
| 160 | + split_train=4, |
| 161 | + split_test=1, |
| 162 | + split_validation=1, |
| 163 | + by_graph_id=False, |
| 164 | + random_seed=42, |
| 165 | + ) |
| 166 | + assert train.n_graphs == 326 |
| 167 | + assert test.n_graphs == 81 |
| 168 | + assert val.n_graphs == 82 |
| 169 | + |
| 170 | + train, test = dataset.split_test_train( |
| 171 | + split_train=4, split_test=1, by_graph_id=False, random_seed=42 |
| 172 | + ) |
| 173 | + assert train.n_graphs == 391 |
| 174 | + assert test.n_graphs == 98 |
| 175 | + |
| 176 | + train, test = dataset.split_test_train( |
| 177 | + split_train=4, split_test=5, by_graph_id=False, random_seed=42 |
| 178 | + ) |
| 179 | + assert train.n_graphs == 217 |
| 180 | + assert test.n_graphs == 272 |
| 181 | + |
| 182 | + with pytest.raises( |
| 183 | + NotImplementedError, |
| 184 | + match="Make sure split_train > split_test >= split_validation, other behaviour is not supported when by_graph_id is True...", |
| 185 | + ): |
| 186 | + dataset.split_test_train( |
| 187 | + split_train=4, split_test=5, by_graph_id=True, random_seed=42 |
| 188 | + ) |
0 commit comments