Merge pull request #30 from UnravelSports/fix/sample_rate

fix sample rate

Author: UnravelSports

tests/test_bigdb.py

@@ -200,7 +200,7 @@ def test_settings(self, gnnc_non_default, non_default_arguments):
         assert 1 == 1
 
         data = spektral_graphs
-        assert len(data) == 132
+        assert len(data) == 130
         assert isinstance(data[0], Graph)
 
         assert settings.pitch_dimensions.pitch_length == 120.0

unravel/american_football/graphs/graph_converter.py

@@ -63,7 +63,7 @@ def __init__(
             else dataset._graph_id_column
         )
 
-        self.sample = 1.0 / kwargs.get("sample_rate", 1.0)
+        self.sample_rate = kwargs.get("sample_rate", None)
         self.chunk_size = chunk_size
         self.attacking_non_qb_node_value = attacking_non_qb_node_value
         self.graph_feature_cols = graph_feature_cols
@@ -73,6 +73,27 @@ def __init__(
 
         self._sport_specific_checks()
 
+        self._sample()
+        self._shuffle()
+
+    def _sample(self):
+        if self.sample_rate is None:
+            return
+        else:
+            self.dataset = self.dataset.filter(
+                pl.col(Column.FRAME_ID) % (1.0 / self.sample_rate) == 0
+            )
+
+    def _shuffle(self):
+        if isinstance(self.settings.random_seed, int):
+            self.dataset = self.dataset.sample(
+                fraction=1.0, seed=self.settings.random_seed
+            )
+        elif self.settings.random_seed == True:
+            self.dataset = self.dataset.sample(fraction=1.0)
+        else:
+            pass
+
     def _sport_specific_checks(self):
         def __remove_with_missing_values(min_object_count: int = 10):
             cs = (
@@ -338,7 +359,6 @@ def process_chunk(chunk: pl.DataFrame) -> List[dict]:
                     "id": chunk[self.graph_id_column][i],
                 }
                 for i in range(len(chunk))
-                if i % self.sample == 0
             ]
 
         graph_df = self._convert()

unravel/soccer/graphs/graph_converter_pl.py

@@ -80,8 +80,6 @@ def __post_init__(self):
         if not isinstance(self.dataset, KloppyPolarsDataset):
             raise ValueError("dataset should be of type KloppyPolarsDataset...")
 
-        self.sample = 1.0 if self.sample_rate is None else 1.0 / self.sample_rate
-
         self.pitch_dimensions: MetricPitchDimensions = (
             self.dataset.settings.pitch_dimensions
         )
@@ -114,8 +112,17 @@ def __post_init__(self):
         else:
             self.dataset = self._remove_incomplete_frames()
 
+        self._sample()
         self._shuffle()
 
+    def _sample(self):
+        if self.sample_rate is None:
+            return
+        else:
+            self.dataset = self.dataset.filter(
+                pl.col(Column.FRAME_ID) % (1.0 / self.sample_rate) == 0
+            )
+
     def _verify_feature_funcs(self, funcs, feature_type: Literal["edge", "node"]):
         for i, func in enumerate(funcs):
             # Check if it has the attributes added by the decorator
@@ -523,7 +530,6 @@ def process_chunk(chunk: pl.DataFrame) -> List[dict]:
                     "id": chunk[self.graph_id_column][i],
                 }
                 for i in range(len(chunk))
-                if i % self.sample == 0
             ]
 
         graph_df = self._convert()