Clean code

unravel/soccer/graphs/graph_converter.py

@@ -2,7 +2,6 @@
 import sys
 from copy import deepcopy
 import json
-#from .. import __version__
 
 import warnings
 
@@ -115,7 +114,7 @@ class GraphConverter:
     node_features: NodeFeatureSet = field(
         init=False, repr=False, default_factory=NodeFeatureSet
     )
-    
+
     edge_features: EdgeFeatureSet = field(
         init=False, repr=False, default_factory=EdgeFeatureSet
     )
@@ -280,7 +279,7 @@ def to_graph_frames(self) -> dict:
                         graph_id=graph_id,
                         settings=self.settings,
                         node_features=self.node_features,
-                        edge_features=self.edge_features
+                        edge_features=self.edge_features,
                     )
                     if gnn_frame.graph_data:
                         self.graph_frames.append(gnn_frame)
@@ -325,16 +324,22 @@ def to_pickle(self, file_path: str) -> None:
         with gzip.open(file_path, "wb") as file:
             data = [x.graph_data for x in self.graph_frames]
             pickle.dump(data, file)
-            
+
     def export_settings(self) -> None:
-        file_path = 'settings.json'
+        file_path = "settings.json"
         data = {
             "__version__": "0.1.2",
-            "node_features": [func_name for func_name,_,_ in self.node_features.get_features()],
-            "edge_features": [func_name for func_name,_,_ in self.edge_features.get_features()],
-            "graph_settings": self.settings.to_dict()
-            }
+            "node_features": [
+                func_name for func_name, _, _ in self.node_features.get_features()
+            ],
+            "edge_features": [
+                func_name for func_name, _, _ in self.edge_features.get_features()
+            ],
+            "graph_settings": self.settings.to_dict(),
+        }
         print(data)
 
-        with open(file_path, 'w') as json_file:
+        with open(file_path, "w") as json_file:
             json.dump(data, json_file, indent=4)
+
+        return

unravel/utils/features/edge_feature_set.py

@@ -1,68 +1,142 @@
 import numpy as np
 
-from .utils import (
-    normalize_coords,
-    unit_vector,
-    normalize_speed,
-    normalize_angles,
-    normalize_distance,
-    normalize_sincos
-)
+from .utils import normalize_speed, normalize_distance, normalize_sincos
+
 
 class EdgeFeatureSet:
     """
     To manage and store edge feature functions configured by user
     """
+
     def __init__(self):
         self.edge_feature_functions = []
-    
+
     def add_dist_matrix(self, normed: bool = True):
         if normed:
-            self.edge_feature_functions.append(("normalize_dist", normalize_distance, ["distances_between_players", "max_dist_to_player"]))
+            self.edge_feature_functions.append(
+                (
+                    "normalize_dist",
+                    normalize_distance,
+                    ["distances_between_players", "max_dist_to_player"],
+                )
+            )
         else:
-            self.edge_feature_functions.append(("raw_dist", lambda dist: dist, ["distances_between_players"]))
-
+            self.edge_feature_functions.append(
+                ("raw_dist", lambda dist: dist, ["distances_between_players"])
+            )
+
         return self
-    
+
     def add_speed_diff_matrix(self, normed: bool = True):
         if normed:
-            self.edge_feature_functions.append(("normalize_speed_diff", lambda a_spe, h_spe, max_speed: np.nan_to_num(normalize_speed(a_spe[None, :], max_speed) - normalize_speed(h_spe[:, None], max_speed)), ['a_spe', 'h_spe', 'max_speed']))
+            self.edge_feature_functions.append(
+                (
+                    "normalize_speed_diff",
+                    lambda a_spe, h_spe, max_speed: np.nan_to_num(
+                        normalize_speed(a_spe[None, :], max_speed)
+                        - normalize_speed(h_spe[:, None], max_speed)
+                    ),
+                    ["a_spe", "h_spe", "max_speed"],
+                )
+            )
         else:
-            self.edge_feature_functions.append(("raw_speed_diff", lambda a_spe, h_spe: np.nan_to_num(a_spe[None, :] - h_spe[:, None]), ['a_spe', 'h_spe']))
-
+            self.edge_feature_functions.append(
+                (
+                    "raw_speed_diff",
+                    lambda a_spe, h_spe: np.nan_to_num(a_spe[None, :] - h_spe[:, None]),
+                    ["a_spe", "h_spe"],
+                )
+            )
+
         return self
-    
+
     def add_pos_cos_matrix(self, normed: bool = True):
         if normed:
-            self.edge_feature_functions.append(("normalise_cos_pos", lambda angle_pos_matrix: normalize_sincos(np.nan_to_num(np.cos(angle_pos_matrix))), ['angle_pos_matrix']))
+            self.edge_feature_functions.append(
+                (
+                    "normalise_cos_pos",
+                    lambda angle_pos_matrix: normalize_sincos(
+                        np.nan_to_num(np.cos(angle_pos_matrix))
+                    ),
+                    ["angle_pos_matrix"],
+                )
+            )
         else:
-            self.edge_feature_functions.append(("raw_cos_pos", lambda angle_pos_matrix: np.nan_to_num(np.cos(angle_pos_matrix)), ['angle_pos_matrix']))
-
+            self.edge_feature_functions.append(
+                (
+                    "raw_cos_pos",
+                    lambda angle_pos_matrix: np.nan_to_num(np.cos(angle_pos_matrix)),
+                    ["angle_pos_matrix"],
+                )
+            )
+
         return self
-    
+
     def add_pos_sin_matrix(self, normed: bool = True):
         if normed:
-            self.edge_feature_functions.append(("normalise_sin_pos", lambda angle_pos_matrix: normalize_sincos(np.nan_to_num(np.sin(angle_pos_matrix))), ['angle_pos_matrix']))
+            self.edge_feature_functions.append(
+                (
+                    "normalise_sin_pos",
+                    lambda angle_pos_matrix: normalize_sincos(
+                        np.nan_to_num(np.sin(angle_pos_matrix))
+                    ),
+                    ["angle_pos_matrix"],
+                )
+            )
         else:
-            self.edge_feature_functions.append(("raw_sin_pos", lambda angle_pos_matrix: np.nan_to_num(np.sin(angle_pos_matrix)), ['angle_pos_matrix']))
-
+            self.edge_feature_functions.append(
+                (
+                    "raw_sin_pos",
+                    lambda angle_pos_matrix: np.nan_to_num(np.sin(angle_pos_matrix)),
+                    ["angle_pos_matrix"],
+                )
+            )
+
         return self
-    
+
     def add_vel_cos_matrix(self, normed: bool = True):
         if normed:
-            self.edge_feature_functions.append(("normalise_cos_vel", lambda angle_vel_matrix: normalize_sincos(np.nan_to_num(np.cos(angle_vel_matrix))), ['angle_vel_matrix']))
+            self.edge_feature_functions.append(
+                (
+                    "normalise_cos_vel",
+                    lambda angle_vel_matrix: normalize_sincos(
+                        np.nan_to_num(np.cos(angle_vel_matrix))
+                    ),
+                    ["angle_vel_matrix"],
+                )
+            )
         else:
-            self.edge_feature_functions.append(("raw_cos_vel", lambda angle_vel_matrix: np.nan_to_num(np.cos(angle_vel_matrix)), ['angle_vel_matrix']))
-
+            self.edge_feature_functions.append(
+                (
+                    "raw_cos_vel",
+                    lambda angle_vel_matrix: np.nan_to_num(np.cos(angle_vel_matrix)),
+                    ["angle_vel_matrix"],
+                )
+            )
+
         return self
-    
+
     def add_vel_sin_matrix(self, normed: bool = True):
         if normed:
-            self.edge_feature_functions.append(("normalise_sin_vel", lambda angle_vel_matrix: normalize_sincos(np.nan_to_num(np.sin(angle_vel_matrix))), ['angle_vel_matrix']))
+            self.edge_feature_functions.append(
+                (
+                    "normalise_sin_vel",
+                    lambda angle_vel_matrix: normalize_sincos(
+                        np.nan_to_num(np.sin(angle_vel_matrix))
+                    ),
+                    ["angle_vel_matrix"],
+                )
+            )
         else:
-            self.edge_feature_functions.append(("raw_sin_vel", lambda angle_vel_matrix: np.nan_to_num(np.sin(angle_vel_matrix)), ['angle_vel_matrix']))
-
+            self.edge_feature_functions.append(
+                (
+                    "raw_sin_vel",
+                    lambda angle_vel_matrix: np.nan_to_num(np.sin(angle_vel_matrix)),
+                    ["angle_vel_matrix"],
+                )
+            )
+
         return self
-    
+
     def get_features(self):
-        return self.edge_feature_functions
+        return self.edge_feature_functions

unravel/utils/features/edge_features.py

@@ -28,7 +28,7 @@ def edge_features(
     # reshape it to a (<=23**2, ) matrix and then mask all values that are 0 in `a` (nz)
     # then we concat all the features into a single (nz, n_edge_features) matrix
     """
-    #print(function_list)
+    # print(function_list)
     max_dist_to_player = np.sqrt(
         pitch_dimensions.pitch_length**2 + pitch_dimensions.pitch_width**2
     )
@@ -99,15 +99,15 @@ def edge_features(
     #         reindex(vel_sin_matrix, non_zero_idxs, len_a),
     #     ]
     # )
-    #Flexible feature set
+    # Flexible feature set
     all_params = {
         "distances_between_players": distances_between_players,
         "max_dist_to_player": max_dist_to_player,
         "a_spe": a_spe,
         "h_spe": h_spe,
         "max_speed": max(max_player_speed, max_ball_speed),
         "angle_pos_matrix": angle_pos_matrix,
-        "angle_vel_matrix": angle_vel_matrix
+        "angle_vel_matrix": angle_vel_matrix,
     }
     e_features = []
     computed_values = {}
@@ -131,10 +131,10 @@ def edge_features(
                 computed_values[func_name] = 0
                 e_features.append(0)
         except Exception as e:
-                print(f"Error while executing function '{func_name}': {e}")
-                computed_values[func_name] = None
-                e_features.append(None)
-    
+            print(f"Error while executing function '{func_name}': {e}")
+            computed_values[func_name] = None
+            e_features.append(None)
+
     e_tuple = list(e_features)
 
     if delaunay_adjacency_matrix is not None:
@@ -144,6 +144,8 @@ def edge_features(
         e_tuple.extend(extra_tuple)
 
     e = np.concatenate(e_tuple, axis=1)
-    print(function_list)
-    print(np.nan_to_num(e))
+    # print(function_list)
+    # print(np.nan_to_num(e))
+    # with open('output_file2.txt', 'w') as file:
+    #     file.write(np.array2string(np.nan_to_num(e)))
     return np.nan_to_num(e)

unravel/utils/features/node_features.py

@@ -153,13 +153,15 @@ def ball_features(ball):
 
     # compute ball features
     b_features = ball_features(ball)
-    print(b_features)
+    # print(b_features)
     X = np.append(ap_features, dp_features, axis=0)
 
     if include_ball_node:
         X = np.append(X, b_features, axis=0)
 
     # convert np.NaN to 0 (zero)
     X = np.nan_to_num(X)
-    #print(X)
+    # print(X)
+    # with open('output_file2.txt', 'w') as file:
+    #     file.write(np.array2string(X))
     return X

unravel/utils/objects/graph_frame.py

@@ -137,7 +137,7 @@ def _edge_features(self, A, A_delaunay):
             self.settings.pitch_dimensions,
             A,
             A_delaunay,
-            function_list=self.edge_features.get_features()
+            function_list=self.edge_features.get_features(),
         )
 
     def _quality_check(self, X, E):

unravel/utils/objects/graph_settings.py

@@ -150,9 +150,9 @@ def to_dict(self):
             "pad": self.pad,
             "verbose": self.verbose,
             "pitch_dimensions": self._serialize_pitch_dimensions(),
-            "pad_settings": self.pad_settings
+            "pad_settings": self.pad_settings,
         }
-        
+
     def _serialize_pitch_dimensions(self):
         return {
             "pitch_length": self.pitch_dimensions.pitch_length,
@@ -161,4 +161,4 @@ def _serialize_pitch_dimensions(self):
             "min_x": self.pitch_dimensions.x_dim.min,
             "max_y": self.pitch_dimensions.y_dim.max,
             "min_y": self.pitch_dimensions.y_dim.min,
-        }
+        }