edmkit 0.0.9__tar.gz → 0.0.11__tar.gz

{edmkit-0.0.9 → edmkit-0.0.11}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.3
 Name: edmkit
-Version: 0.0.9
+Version: 0.0.11
 Summary: Simple EDM (Empirical Dynamic Modeling) library
 Author: FUJISHIGE TEMMA
 Author-email: FUJISHIGE TEMMA <tenma.x0@gmail.com>
-Requires-Dist: kdtree-rs>=0.1.0
-Requires-Dist: numpy>=2.4.3
+Requires-Dist: kdtree-rs>=0.3.0
+Requires-Dist: numpy>=2.4.6
 Requires-Dist: scipy>=1.17.1
-Requires-Dist: tinygrad>=0.11.0
+Requires-Dist: tinygrad>=0.13.0
 Requires-Python: >=3.13
 Description-Content-Type: text/markdown
 

{edmkit-0.0.9 → edmkit-0.0.11}/pyproject.toml

@@ -1,15 +1,15 @@
 [project]
 name = "edmkit"
-version = "0.0.9"
+version = "0.0.11"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
 readme = "README.md"
 requires-python = ">= 3.13"
 dependencies = [
-    "kdtree-rs>=0.1.0",
-    "numpy>=2.4.3",
+    "kdtree-rs>=0.3.0",
+    "numpy>=2.4.6",
     "scipy>=1.17.1",
-    "tinygrad>=0.11.0",
+    "tinygrad>=0.13.0",
 ]
 
 [dependency-groups]

{edmkit-0.0.9 → edmkit-0.0.11}/src/edmkit/simplex_projection/loo.py

@@ -71,8 +71,8 @@ def loo(
         d_min = np.fmax(distances.min(axis=1, keepdims=True), 1e-6)  # (N, 1)
         weights = np.exp(-distances / d_min)  # (N, k)
 
-        weighted_sum = np.sum(weights[..., None] * Y_neighbors, axis=1)
-        predictions = weighted_sum / np.sum(weights, axis=1, keepdims=True)
+        weighted_sum = np.matmul(weights[:, None, :], Y_neighbors).squeeze(-2)  # (N, E')
+        predictions = weighted_sum / weights.sum(axis=1, keepdims=True)
 
         return predictions.squeeze()  # (N,) or (N, E')
     # X (B, N, E), Y (B, N, E')
@@ -107,8 +107,9 @@ def loo(
         d_min = np.fmax(distances.min(axis=2, keepdims=True), 1e-6)  # (B, N, 1)
         weights = np.exp(-distances / d_min)  # (B, N, k)
 
-        weighted_sum = np.sum(weights[..., None] * Y_neighbors, axis=2)  # (B, N, E')
-        predictions = weighted_sum / np.sum(weights, axis=2, keepdims=True)  # (B, N, E')
+        # matmul avoids the (B, N, k, E') temporary that `weights[..., None] * Y_neighbors` would build
+        weighted_sum = np.matmul(weights[..., None, :], Y_neighbors).squeeze(-2)  # (B, N, E')
+        predictions = weighted_sum / weights.sum(axis=2, keepdims=True)  # (B, N, E')
 
         return predictions
     else:

{edmkit-0.0.9 → edmkit-0.0.11}/src/edmkit/simplex_projection/simplex_projection.py

@@ -1,7 +1,6 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-from tinygrad import Tensor, dtypes
 
 from edmkit.simplex_projection.knn import knn
 from edmkit.util import pairwise_distance
@@ -113,8 +112,8 @@ def _numpy(
         d_min = np.fmax(distances.min(axis=1, keepdims=True), 1e-6)  # (M, 1)
         weights = np.exp(-distances / d_min)  # (M, k)
 
-        weighted_sum = np.sum(weights[..., None] * Y_neighbors, axis=1)
-        predictions = weighted_sum / np.sum(weights, axis=1, keepdims=True)
+        weighted_sum = np.matmul(weights[:, None, :], Y_neighbors).squeeze(-2)  # (M, E')
+        predictions = weighted_sum / weights.sum(axis=1, keepdims=True)
 
         return predictions.squeeze()  # (M,) or (M, E')
     # X (B, N, E), Y (B, N, E'), Q (B, M, E)
@@ -146,8 +145,8 @@ def _numpy(
         d_min = np.fmax(distances.min(axis=2, keepdims=True), 1e-6)  # (B, M, 1)
         weights = np.exp(-distances / d_min)  # (B, M, k)
 
-        weighted_sum = np.sum(weights[..., None] * Y_neighbors, axis=2)  # (B, M, E')
-        predictions = weighted_sum / np.sum(weights, axis=2, keepdims=True)  # (B, M, E')
+        weighted_sum = np.matmul(weights[..., None, :], Y_neighbors).squeeze(-2)  # (B, M, E')
+        predictions = weighted_sum / weights.sum(axis=2, keepdims=True)  # (B, M, E')
 
         return predictions
     else:
@@ -161,6 +160,10 @@ def _tensor(
     *,
     mask: np.ndarray | None = None,
 ):
+    # Lazy import: tinygrad starts a per-CPU async-executor pool at import time
+    # Loading it only when use_tensor=True keeps the numpy path free of that scheduler pressure.
+    from tinygrad import Tensor, dtypes
+
     if X.ndim == 1:
         X = X[:, None]
     if Y.ndim == 1:

{edmkit-0.0.9 → edmkit-0.0.11}/src/edmkit/util.py

@@ -1,6 +1,12 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import numpy as np
 from scipy.spatial.distance import cdist
-from tinygrad import Tensor
+
+if TYPE_CHECKING:
+    from tinygrad import Tensor
 
 
 def pad(As: list[np.ndarray]):