rustat-python-api 0.6.6__py3-none-any.whl → 0.7.1__py3-none-any.whl

rustat_python_api/kernels/__init__.py

@@ -1,3 +1,6 @@
-from .maha import triton_influence
+try:
+    from .maha import triton_influence
+except ImportError:
+    triton_influence = None
 
 __all__ = ["triton_influence"]

rustat_python_api/kernels/maha.py

@@ -1,72 +1,77 @@
 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
+
+try:
+    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
+
+        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
+
+except ImportError:
+    def triton_influence(*args, **kwargs):
+        raise RuntimeError(
+            "Triton not available. This function requires Linux with NVIDIA GPU. "
+            "For CPU fallback, please use the .calculate_influence() method with device='cpu'. "
+            "On macOS, the package installs without GPU acceleration."
+        )

rustat_python_api/pitch_control.py

@@ -7,8 +7,10 @@ import matplotsoccer as mpl
 import torch
 from tqdm import tqdm
 
-from .kernels import triton_influence
-
+try:
+    from .kernels import triton_influence
+except ImportError:
+    triton_influence = None
 
 class PitchControl:
     def __init__(self, tracking: pd.DataFrame, events: pd.DataFrame, ball_data: pd.DataFrame = None):

{rustat_python_api-0.6.6.dist-info → rustat_python_api-0.7.1.dist-info}/METADATA

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

{rustat_python_api-0.6.6.dist-info → rustat_python_api-0.7.1.dist-info}/RECORD

@@ -2,13 +2,13 @@ rustat_python_api/__init__.py,sha256=Ij-PAm2y5ss_XAZhKTZus35cRPLzvXFyIswDa_Iq3rs
 rustat_python_api/config.py,sha256=eMvi1p8Cfvnbp6Cd4bBOwgehVN7thKnaQV5uzWyGZXM,1844
 rustat_python_api/models_api.py,sha256=oHXEqeCupvZwjVEdoxf7W9LP7ELFKA8-9DuRXpQHLno,1701
 rustat_python_api/parser.py,sha256=c5SW9p5N1PLYnza1olZbnvxkWtCAYYty2ifjyaUwXkw,10280
-rustat_python_api/pitch_control.py,sha256=s501qCpvMI7qNQAajDF_NXnlIDMloJQCgUxAfjKGOa0,26282
+rustat_python_api/pitch_control.py,sha256=iJiF0_3j6yh3riV8kTZ_mtJibPqotQZOUgpumSyS-Zw,26338
 rustat_python_api/processing.py,sha256=WES2D77uu3ScpjN0nW8ozatHteCZJQmmF9WpHPgGfJo,2835
 rustat_python_api/urls.py,sha256=iJTD31T6OyXPAhmhViwFXVehrzwsOjBDONA1SIVc_40,1068
-rustat_python_api/kernels/__init__.py,sha256=fQrUfDakSr5IsFV0LqAxvLoEw9Og5PThfux1GQw1Tps,67
-rustat_python_api/kernels/maha.py,sha256=ARGJRJRp9VlX-zQbKipGMUKMMSqIEWg4Pg-skhkYMWc,2184
-rustat_python_api-0.6.6.dist-info/LICENSE,sha256=4Cohqg5p6Mq1xyrzdEX8AvFSA62GSVvapEOr2xK_tgY,57
-rustat_python_api-0.6.6.dist-info/METADATA,sha256=Ykv3P2-fdkRn6wYT25RIotnuX8t2ZPLgW0vYHl8aj7Y,1808
-rustat_python_api-0.6.6.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-rustat_python_api-0.6.6.dist-info/top_level.txt,sha256=VK0hmkKZE9YThxolUcoE6JtGI67NFeKJMBLuet8kI4w,18
-rustat_python_api-0.6.6.dist-info/RECORD,,
+rustat_python_api/kernels/__init__.py,sha256=eFJ-BMY8VcNZSjf3XjOnZf_nfOQ5t-7Lp57DPCHYOo0,124
+rustat_python_api/kernels/maha.py,sha256=k2PqY6VghgER2j9QH8xGYq61JLfPaHjirLXb4aLnjQw,2591
+rustat_python_api-0.7.1.dist-info/LICENSE,sha256=4Cohqg5p6Mq1xyrzdEX8AvFSA62GSVvapEOr2xK_tgY,57
+rustat_python_api-0.7.1.dist-info/METADATA,sha256=Db46gm9pWb-1Zra3BrFFnUI8AU9aLJ4D93XQxnZc6PU,1808
+rustat_python_api-0.7.1.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+rustat_python_api-0.7.1.dist-info/top_level.txt,sha256=VK0hmkKZE9YThxolUcoE6JtGI67NFeKJMBLuet8kI4w,18
+rustat_python_api-0.7.1.dist-info/RECORD,,