rustat-python-api 0.6.1__tar.gz → 0.6.2__tar.gz

{rustat-python-api-0.6.1/rustat_python_api.egg-info → rustat-python-api-0.6.2}/PKG-INFO

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

rustat-python-api-0.6.2/rustat_python_api/kernels/__init__.py

@@ -0,0 +1,3 @@
+from .maha import triton_influence
+
+__all__ = ["triton_influence"]

rustat-python-api-0.6.2/rustat_python_api/kernels/maha.py

@@ -0,0 +1,72 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def influence_kernel(
+    MU,           # (F,P,2)
+    SIGMA_INV,    # (F,P,2,2)
+    LOCS,         # (N,2)  – read-only
+    OUT,          # (F,N)
+    F, P, N,      # sizes
+    BLOCK_N: tl.constexpr,      # кол-во точек на блок
+):
+    pid_f = tl.program_id(0)                    # какой frame
+    offs_n = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)   # chunk точек
+    mask_n = offs_n < N
+
+    # ----- загрузим координаты сетки (N,2) -----
+    loc_ptr = LOCS + offs_n[:, None] * 2        # strides=(2,)
+    x = tl.load(loc_ptr, mask=mask_n[:, None])  # (BLOCK_N,2)
+
+    # аккум для суммирования влияний
+    acc = tl.zeros([BLOCK_N], dtype=tl.float32)
+
+    for p in tl.static_range(0, 32):
+        is_valid = p < P
+
+        # --- pointers ---
+        mu_ptr   = MU + pid_f * P * 2 + p * 2
+        s_ptr    = SIGMA_INV + pid_f * P * 4 + p * 4
+
+        # --- scalar loads ---
+        mu0 = tl.load(mu_ptr + 0, mask=is_valid)
+        mu1 = tl.load(mu_ptr + 1, mask=is_valid)
+
+        s00 = tl.load(s_ptr + 0, mask=is_valid)
+        s01 = tl.load(s_ptr + 1, mask=is_valid)
+        s10 = tl.load(s_ptr + 2, mask=is_valid)
+        s11 = tl.load(s_ptr + 3, mask=is_valid)
+
+        # --- grid coords (also scalar) ---
+        x0 = tl.load(LOCS + offs_n * 2 + 0, mask=mask_n)   # (BLOCK_N,)
+        x1 = tl.load(LOCS + offs_n * 2 + 1, mask=mask_n)
+
+        diff0 = x0 - mu0
+        diff1 = x1 - mu1
+
+        tmp0  = diff0 * s00 + diff1 * s01
+        tmp1  = diff0 * s10 + diff1 * s11
+        maha  = diff0 * tmp0 + diff1 * tmp1
+
+        acc += tl.where(is_valid, tl.exp(-0.5 * maha), 0.0)
+
+    # запишем
+    out_ptr = OUT + pid_f*N + offs_n
+    tl.store(out_ptr, acc, mask=mask_n)
+
+
+def triton_influence(mu, sigma_inv, locs, BLOCK_N=64):
+    F, P, _ = mu.shape
+    N = locs.shape[0]
+    out = torch.empty((F, N), device=mu.device, dtype=mu.dtype)
+
+    grid = (F, triton.cdiv(N, BLOCK_N))
+    influence_kernel[grid](
+        mu, sigma_inv, locs, out,
+        F, P, N,
+        BLOCK_N=BLOCK_N,
+        num_warps=4
+    )
+    return out

{rustat-python-api-0.6.1 → rustat-python-api-0.6.2}/rustat_python_api/pitch_control.py

@@ -7,6 +7,8 @@ import matplotsoccer as mpl
 import torch
 from tqdm import tqdm
 
+from .kernels import triton_influence
+
 
 class PitchControl:
     def __init__(self, tracking: pd.DataFrame, events: pd.DataFrame, ball_data: pd.DataFrame = None):
@@ -320,15 +322,19 @@ class PitchControl:
         diff = locs.view(1, -1, 2)  # (1,N,2)
         diff = diff - mu.unsqueeze(1)   # (P,N,2)
 
-        maha = torch.einsum('pni,pij,pnj->pn', diff, Sigma_inv, diff)  # (P,N)
-
-        # Replace NaNs that arise from invalid player positions with large value
-        # so their influence tends to zero after exponent, then eliminate residual NaNs.
-        maha = torch.nan_to_num(maha, nan=1e9, posinf=1e9, neginf=1e9)
+        if device == "cuda":
+            out = triton_influence(
+                mu.unsqueeze(0), Sigma_inv.unsqueeze(0),
+                locs, BLOCK_N=64
+            )[0]  # (N,)
 
-        out = torch.exp(-0.5 * maha)  # (P,N)
+            return out
+        else:
+            maha = torch.einsum('pni,pij,pnj->pn', diff, Sigma_inv, diff)  # (P,N)
+            maha = torch.nan_to_num(maha, nan=1e9, posinf=1e9, neginf=1e9)
+            out = torch.exp(-0.5 * maha)  # (P,N)
 
-        return out.sum(dim=0)  # sum over players
+            return out.sum(dim=0)  # sum over players
 
     def _batch_team_influence_frames_pt(
         self,
@@ -385,15 +391,16 @@ class PitchControl:
         diff = locs.view(1, 1, -1, 2)  # (1,1,N,2)
         diff = diff - mu.unsqueeze(2)   # (F,P,N,2)
 
-        maha = torch.einsum('fpni,fpij,fpnj->fpn', diff, Sigma_inv, diff)  # (F,P,N)
+        if device == "cuda":
+            out = triton_influence(mu, Sigma_inv, locs, BLOCK_N=64)  # (F,N)
 
-        # Replace NaNs that arise from invalid player positions with large value
-        # so their influence tends to zero after exponent, then eliminate residual NaNs.
-        maha = torch.nan_to_num(maha, nan=1e9, posinf=1e9, neginf=1e9)
-
-        out = torch.exp(-0.5 * maha)  # (F,P,N)
+            return out
+        else:
+            maha = torch.einsum('fpni,fpij,fpnj->fpn', diff, Sigma_inv, diff)  # (F,P,N)
+            maha = torch.nan_to_num(maha, nan=1e9, posinf=1e9, neginf=1e9)
+            out = torch.exp(-0.5 * maha)  # (F,P,N)
 
-        return out.sum(dim=1)  # sum over players
+            return out.sum(dim=1)  # sum over players
 
     @staticmethod
     def _stack_team_frames(players: list[np.ndarray], frames: np.ndarray, device: str, dtype: torch.dtype):

{rustat-python-api-0.6.1 → rustat-python-api-0.6.2/rustat_python_api.egg-info}/PKG-INFO

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

{rustat-python-api-0.6.1 → rustat-python-api-0.6.2}/rustat_python_api.egg-info/SOURCES.txt

@@ -13,4 +13,6 @@ rustat_python_api.egg-info/PKG-INFO
 rustat_python_api.egg-info/SOURCES.txt
 rustat_python_api.egg-info/dependency_links.txt
 rustat_python_api.egg-info/requires.txt
-rustat_python_api.egg-info/top_level.txt
+rustat_python_api.egg-info/top_level.txt
+rustat_python_api/kernels/__init__.py
+rustat_python_api/kernels/maha.py

{rustat-python-api-0.6.1 → rustat-python-api-0.6.2}/rustat_python_api.egg-info/requires.txt

@@ -5,3 +5,4 @@ scipy==1.14.1
 matplotlib
 matplotsoccer
 torch
+triton==3.0.0

{rustat-python-api-0.6.1 → rustat-python-api-0.6.2}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='rustat-python-api',
-    version='0.6.1',
+    version='0.6.2',
     description='A Python wrapper for RuStat API',
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',
@@ -18,7 +18,8 @@ setup(
         'scipy==1.14.1',
         'matplotlib',
         'matplotsoccer',
-        'torch'
+        'torch',
+        'triton==3.0.0',
     ],
     classifiers=[
         'Programming Language :: Python :: 3',