rustat-python-api 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

rustat_python_api/parser.py

@@ -150,8 +150,6 @@ class RuStatParser:
 
             df = processing(df)
 
-            df["event_chron_id"] = range(1, len(df) + 1)
-
         return (df, subs) if return_subs else df
 
     def get_tracking(self, match_id: int) -> pd.DataFrame | None:

rustat_python_api/pitch_control.py

@@ -597,6 +597,8 @@ class PitchControl:
         verbose: bool = False,
         use_fp16: bool = True,
         return_teams: str = "both",
+        include_players: list | None = None,
+        exclude_players: list | None = None,
     ):
         """Torch-accelerated pitch-control calculation.
 
@@ -621,12 +623,16 @@ class PitchControl:
         else:
             xx_t, yy_t, locs_t = self._grid_cache[key]
 
+        # --- filter players ---
+        home_dict = self._filter_players(self.locs_home[half], include_players, exclude_players)
+        away_dict = self._filter_players(self.locs_away[half], include_players, exclude_players)
+
         # --- vectorised influence computation ---
         Zh = self._batch_influence_pt(
-            self.locs_home[half], locs_t, tp, half, device, dtype
+            home_dict, locs_t, tp, half, device, dtype
         )
         Za = self._batch_influence_pt(
-            self.locs_away[half], locs_t, tp, half, device, dtype
+            away_dict, locs_t, tp, half, device, dtype
         )
 
         # Return raw team influence or combined logistic transformation
@@ -640,6 +646,14 @@ class PitchControl:
         # Convert to numpy for downstream plotting
         return res_t.cpu().numpy(), xx_t.cpu().numpy(), yy_t.cpu().numpy()
 
+    @staticmethod
+    def _filter_players(player_dict, include_players=None, exclude_players=None):
+        if include_players is not None:
+            player_dict = {k: v for k, v in player_dict.items() if k in include_players}
+        if exclude_players is not None:
+            player_dict = {k: v for k, v in player_dict.items() if k not in exclude_players}
+        return player_dict
+
     def _fit_np(
         self,
         half: int,
@@ -647,6 +661,8 @@ class PitchControl:
         dt: int,
         verbose: bool = False,
         return_teams: str = "both",
+        include_players: list | None = None,
+        exclude_players: list | None = None,
     ) -> tuple:
         x = np.linspace(0, 105, dt)
         y = np.linspace(0, 68, dt)
@@ -657,11 +673,12 @@ class PitchControl:
 
         locations = np.c_[xx.flatten(), yy.flatten()]
 
-        for k in self.locs_home[half].keys():
-            # if len(self.locs_home[half][k]) >= tp:
+        home_dict = self._filter_players(self.locs_home[half], include_players, exclude_players)
+        away_dict = self._filter_players(self.locs_away[half], include_players, exclude_players)
+
+        for k in home_dict.keys():
             Zh += self.influence_np(k, locations, tp, "h", half, verbose)
-        for k in self.locs_away[half].keys():
-            # if len(self.locs_away[half][k]) >= tp:
+        for k in away_dict.keys():
             Za += self.influence_np(k, locations, tp, "a", half, verbose)
 
         Zh = Zh.reshape((dt, dt))
@@ -687,11 +704,24 @@ class PitchControl:
         verbose: bool = False,
         use_fp16: bool = True,
         return_teams: str = "both",
+        include_players: list | None = None,
+        exclude_players: list | None = None,
     ):
-        """Selects NumPy or PyTorch backend depending on `backend`."""
+        """Selects NumPy or PyTorch backend depending on `backend`.
+
+        Parameters
+        ----------
+        include_players : list of player_ids, optional
+            If provided, only these players are included in the computation.
+        exclude_players : list of player_ids, optional
+            If provided, these players are excluded from the computation.
+        """
         match backend:
             case "np" | "numpy":
-                return self._fit_np(half, tp, dt, verbose, return_teams)
+                return self._fit_np(
+                    half, tp, dt, verbose, return_teams,
+                    include_players, exclude_players,
+                )
             case "torch" | "pt":
                 return self._fit_pt(
                     half,
@@ -701,6 +731,8 @@ class PitchControl:
                     verbose=verbose,
                     use_fp16=use_fp16,
                     return_teams=return_teams,
+                    include_players=include_players,
+                    exclude_players=exclude_players,
                 )
             case _:
                 raise ValueError(f"Unknown backend '{backend}'. Use 'np' or 'torch'.")

rustat_python_api/processing.py

@@ -1,5 +1,5 @@
-import pandas as pd
 import numpy as np
+import pandas as pd
 
 from .config import columns, id2type
 
@@ -13,6 +13,7 @@ def process_list(x: pd.Series):
     else:
         return lst
 
+
 def take_last(x: pd.Series):
     lst = x.dropna().tolist()
     if lst:
@@ -21,45 +22,49 @@ def take_last(x: pd.Series):
 
 
 def gluing(df: pd.DataFrame) -> pd.DataFrame:
-    cols = ['player_id', 'half', 'second', 'pos_x', 'pos_y']
+    cols = ["player_id", "half", "second", "pos_x", "pos_y"]
 
     agg_rules = {}
 
     for col_name in df.columns:
         if col_name not in cols:
-            if col_name in ['action_name', 'action_id']:
+            if col_name in ["action_name", "action_id"]:
                 agg_rules[col_name] = process_list
             else:
                 agg_rules[col_name] = take_last
 
     df_gb = df.groupby(cols).agg(agg_rules).reset_index()
 
-    df_gb['pos_dest_nan'] = (df_gb['pos_dest_x'].isna() & df_gb['pos_dest_y'].isna()).astype(int)
-    df_gb = df_gb.sort_values(by=['half', 'second', 'possession_number', 'pos_dest_nan']).reset_index(drop=True)
+    df_gb["pos_dest_nan"] = (
+        df_gb["pos_dest_x"].isna() & df_gb["pos_dest_y"].isna()
+    ).astype(int)
+    df_gb = df_gb.sort_values(
+        by=["half", "second", "possession_number", "pos_dest_nan"]
+    ).reset_index(drop=True)
     return df_gb
 
 
 def add_reciever(glued_df: pd.DataFrame) -> pd.DataFrame:
     df = glued_df.copy()
-    df['receiver_id'] = df['player_id'].shift(-1)
-    df['receiver_name'] = df['player_name'].shift(-1)
+    df["receiver_id"] = df["player_id"].shift(-1)
+    df["receiver_name"] = df["player_name"].shift(-1)
 
     mask = (
-            (df['action_name'] == 'Ball receiving')
-            & (df['pos_x'] == df['pos_dest_x'].shift(1))
-            & (df['pos_y'] == df['pos_dest_y'].shift(1))
-            & (df['team_id'] == df['team_id'].shift(1))
-            & (df['player_id'] != df['player_id'].shift(1))
-            & (df['possession_number'] == df['possession_number'].shift(1))
+        (df["action_name"] == "Ball receiving")
+        & (df["pos_x"] == df["pos_dest_x"].shift(1))
+        & (df["pos_y"] == df["pos_dest_y"].shift(1))
+        & (df["team_id"] == df["team_id"].shift(1))
+        & (df["player_id"] != df["player_id"].shift(1))
+        & (df["possession_number"] == df["possession_number"].shift(1))
     )
 
     idx = df[mask].index
-    remaining_idx = df.drop(idx-1).index
+    remaining_idx = df.drop(idx - 1).index
 
-    df.loc[remaining_idx, 'receiver_id'] = np.nan
-    df.loc[remaining_idx, 'receiver_name'] = np.nan
+    df.loc[remaining_idx, "receiver_id"] = np.nan
+    df.loc[remaining_idx, "receiver_name"] = np.nan
 
-    df = df[df['action_name'] != 'Ball receiving'].reset_index(drop=True)
+    df = df[df["action_name"] != "Ball receiving"].reset_index(drop=True)
 
     return df
 
@@ -69,20 +74,23 @@ def filter_data(df: pd.DataFrame) -> pd.DataFrame:
         if column not in df.columns:
             df[column] = np.nan
 
-    return df[(~df['possession_number'].isna()) | (df['second'] != 0)][columns].reset_index(drop=True)
+    return df[(~df["possession_number"].isna()) | (df["second"] != 0)][
+        columns
+    ].reset_index(drop=True)
 
 
 def tagging(df: pd.DataFrame) -> pd.DataFrame:
-    df = df.rename(columns={'action_name': 'sub_tags', 'action_id': 'sub_tags_ids'})
-    df['sub_tags'] = df['sub_tags'].apply(lambda x: x if isinstance(x, list) else [x])
-    df['sub_tags_ids'] = df['sub_tags_ids'].apply(
-        lambda x:
-        list(set([int(t) // 1000 for t in x]))
-        if isinstance(x, list)
-        else [int(x) // 1000]
+    df = df.rename(columns={"action_name": "sub_tags", "action_id": "sub_tags_ids"})
+    df["sub_tags"] = df["sub_tags"].apply(lambda x: x if isinstance(x, list) else [x])
+    df["sub_tags_ids"] = df["sub_tags_ids"].apply(
+        lambda x: (
+            list(set([int(t) // 1000 for t in x]))
+            if isinstance(x, list)
+            else [int(x) // 1000]
+        )
     )
-    df['sub_tags_ids'] = df['sub_tags_ids'].apply(lambda x: [id2type[t] for t in x])
-    df = df.rename(columns={'sub_tags_ids': 'tags'})
+    df["sub_tags_ids"] = df["sub_tags_ids"].apply(lambda x: [id2type[t] for t in x])
+    df = df.rename(columns={"sub_tags_ids": "tags"})
 
     return df
 
@@ -93,4 +101,6 @@ def processing(df: pd.DataFrame) -> pd.DataFrame:
     df = filter_data(df)
     df = tagging(df)
 
+    df["event_chron_id"] = range(1, len(df) + 1)
+
     return df

{rustat_python_api-0.8.0.dist-info → rustat_python_api-0.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: rustat-python-api
-Version: 0.8.0
+Version: 0.9.0
 Summary: A Python wrapper for RuStat API
 Home-page: https://github.com/dailydaniel/rustat-python-api
 Author: Daniel Zholkovsky

{rustat_python_api-0.8.0.dist-info → rustat_python_api-0.9.0.dist-info}/RECORD

@@ -1,9 +1,9 @@
 rustat_python_api/__init__.py,sha256=Ij-PAm2y5ss_XAZhKTZus35cRPLzvXFyIswDa_Iq3rs,164
 rustat_python_api/config.py,sha256=eMvi1p8Cfvnbp6Cd4bBOwgehVN7thKnaQV5uzWyGZXM,1844
 rustat_python_api/models_api.py,sha256=oHXEqeCupvZwjVEdoxf7W9LP7ELFKA8-9DuRXpQHLno,1701
-rustat_python_api/parser.py,sha256=7bvD9uQICNXbg--KNkUT7I4WSoMpgGcJVxPGunFfDWA,10921
-rustat_python_api/pitch_control.py,sha256=KN7c1tUWsygFe8WMBX2ADZdGPP600bcqmE84QZOFH6Q,30247
-rustat_python_api/processing.py,sha256=WES2D77uu3ScpjN0nW8ozatHteCZJQmmF9WpHPgGfJo,2835
+rustat_python_api/parser.py,sha256=hMPZER6CwPk_dm8GkLulq6G0_jwOyRiZjbbDw3eO1Rc,10863
+rustat_python_api/pitch_control.py,sha256=XSakx2KxDAvv17qFoBtsBBIEYGn3OqI_3q3Xfq3_OoA,31688
+rustat_python_api/processing.py,sha256=sjZdjSs2BKyM_f7T_Utu1CJSF1l-V8UsWrAj7dhqyNg,2930
 rustat_python_api/urls.py,sha256=iJTD31T6OyXPAhmhViwFXVehrzwsOjBDONA1SIVc_40,1068
 rustat_python_api/kernels/__init__.py,sha256=eFJ-BMY8VcNZSjf3XjOnZf_nfOQ5t-7Lp57DPCHYOo0,124
 rustat_python_api/kernels/maha.py,sha256=k2PqY6VghgER2j9QH8xGYq61JLfPaHjirLXb4aLnjQw,2591
@@ -11,8 +11,8 @@ rustat_python_api/matching/__init__.py,sha256=QIqkft_NVban8p0rxJsim3x4defZN3eQBq
 rustat_python_api/matching/dataloader.py,sha256=FZ95ZC5Z-6FrZEVy-d5S1rXt5nIBwYMqTT0hvFlc4mI,2753
 rustat_python_api/matching/pc_adder.py,sha256=oP84iplIqmqF1jFDN9ernkjKw7wgsmWbuOVPEsoTNWM,7582
 rustat_python_api/matching/tr_adder.py,sha256=0sAjYfoLitScDhg2yjh2PIKWTseoaSeUa5chPnVau-Q,16025
-rustat_python_api-0.8.0.dist-info/LICENSE,sha256=4Cohqg5p6Mq1xyrzdEX8AvFSA62GSVvapEOr2xK_tgY,57
-rustat_python_api-0.8.0.dist-info/METADATA,sha256=Kx3REXPxgkqyq-UfsW9jI3gnNBT7mfAXXZNJYdC_w0U,1920
-rustat_python_api-0.8.0.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-rustat_python_api-0.8.0.dist-info/top_level.txt,sha256=VK0hmkKZE9YThxolUcoE6JtGI67NFeKJMBLuet8kI4w,18
-rustat_python_api-0.8.0.dist-info/RECORD,,
+rustat_python_api-0.9.0.dist-info/LICENSE,sha256=4Cohqg5p6Mq1xyrzdEX8AvFSA62GSVvapEOr2xK_tgY,57
+rustat_python_api-0.9.0.dist-info/METADATA,sha256=9_I5OFJ6m9tWDX4g5elAfjNgXpAIbLyJ6kdrVHEm9aE,1920
+rustat_python_api-0.9.0.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+rustat_python_api-0.9.0.dist-info/top_level.txt,sha256=VK0hmkKZE9YThxolUcoE6JtGI67NFeKJMBLuet8kI4w,18
+rustat_python_api-0.9.0.dist-info/RECORD,,