edmkit 0.0.4__tar.gz → 0.0.6__tar.gz

{edmkit-0.0.4 → edmkit-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: edmkit
-Version: 0.0.4
+Version: 0.0.6
 Summary: Simple EDM (Empirical Dynamic Modeling) library
 Author: FUJISHIGE TEMMA
 Author-email: FUJISHIGE TEMMA <tenma.x0@gmail.com>

{edmkit-0.0.4 → edmkit-0.0.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.4"
+version = "0.0.6"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
 readme = "README.md"

{edmkit-0.0.4 → edmkit-0.0.6}/src/edmkit/embedding.py

@@ -195,8 +195,11 @@ def select(
 ) -> tuple[int, int, float]:
     """Select best (E, tau) from scan results.
 
-    Aggregates over the fold axis (axis=2) with nanmean, then
-    finds the (E, tau) combination with the highest mean score.
+    Ranks each (E, tau) by ``mean - SE`` where SE is the standard error
+    of the mean across folds.  This penalises combinations whose scores
+    vary widely across folds (unstable predictions) and those with fewer
+    valid folds (less certainty), favouring parameters we are *confident*
+    perform well.
 
     Parameters
     ----------
@@ -210,15 +213,31 @@ def select(
     Returns
     -------
     (best_E, best_tau, best_score)
+        ``best_score`` is the mean over folds (not the adjusted value)
+        so that it remains directly interpretable.
     """
-    valid_counts = np.sum(~np.isnan(scores), axis=2)
-    summed_scores = np.nansum(scores, axis=2)
+    K = np.sum(~np.isnan(scores), axis=2)
+    nan_out = np.full(scores.shape[:2], np.nan)
+
     mean_scores = np.divide(
-        summed_scores,
-        valid_counts,
-        out=np.full(summed_scores.shape, np.nan, dtype=float),
-        where=valid_counts > 0,
+        np.nansum(scores, axis=2),
+        K,
+        out=nan_out.copy(),
+        where=K > 0,
     )
-    flat_idx = int(np.nanargmax(mean_scores))
-    e_idx, t_idx = np.unravel_index(flat_idx, mean_scores.shape)
+
+    # SE = sqrt(var / K) = sqrt(sum_sq / (K * (K - 1)))
+    sum_sq = np.nansum((scores - mean_scores[:, :, None]) ** 2, axis=2)
+    se = np.sqrt(
+        np.divide(
+            sum_sq,
+            K * np.maximum(K - 1, 1),
+            out=np.zeros_like(nan_out),
+            where=K > 1,
+        )
+    )
+
+    adjusted = mean_scores - se
+    flat_idx = int(np.nanargmax(adjusted))
+    e_idx, t_idx = np.unravel_index(flat_idx, adjusted.shape)
     return E[e_idx], tau[t_idx], float(mean_scores[e_idx, t_idx])

edmkit-0.0.6/src/edmkit/simplex_projection/__init__.py

@@ -0,0 +1,5 @@
+from .knn import knn
+from .loo import loo
+from .simplex_projection import simplex_projection
+
+__all__ = ["knn", "loo", "simplex_projection"]

edmkit-0.0.6/src/edmkit/simplex_projection/knn.py

@@ -0,0 +1,40 @@
+import numpy as np
+from scipy.spatial import KDTree
+from usearch.index import Index
+
+
+def knn(X: np.ndarray, Q: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
+    """Find the k-nearest neighbors of `Q` in `X` using either `usearch` or `scipy.spatial.KDTree` depending on the size and dimensionality of the data.
+
+    Parameters
+    ----------
+    `X` : `np.ndarray`
+        The input data (N, E)
+    `Q` : `np.ndarray`
+        The query points (M, E)
+    `k` : `int`
+        The number of nearest neighbors to find (typically E+1 for simplex projection).
+
+    Returns
+    -------
+    distances : `np.ndarray`
+        The distances from each query point in `Q` to its k nearest neighbors in `X` (M, k)
+    indices : `np.ndarray`
+        The indices of the k nearest neighbors in `X` for each query point in `Q` (M, k)
+    """
+
+    N, E = X.shape
+
+    if N < k:
+        raise ValueError(f"Not enough points in X to find {k} neighbors, got N={N}")
+
+    if E >= 15 and N >= 10_000:
+        index = Index(ndim=E, metric="l2sq")
+        index.add(np.arange(len(X)), np.ascontiguousarray(X, dtype=np.float32))
+        matches = index.search(np.ascontiguousarray(Q, dtype=np.float32), k)
+        distances = np.atleast_2d(np.sqrt(np.asarray(matches.distances)))
+        indices = np.atleast_2d(np.asarray(matches.keys).astype(np.intp))
+        return distances, indices
+    else:
+        tree = KDTree(X)
+        return tree.query(Q, k=k)

edmkit-0.0.6/src/edmkit/simplex_projection/loo.py

@@ -0,0 +1,115 @@
+import numpy as np
+
+from edmkit.simplex_projection.knn import knn
+
+
+def loo(
+    X: np.ndarray,
+    Y: np.ndarray,
+    *,
+    theiler_window: int,
+) -> np.ndarray:
+    """
+    Leave-one-out simplex projection: predict each point in `X` from its neighbors, excluding temporally close points.
+
+    Equivalent to ``simplex_projection(X, Y, X)`` with Theiler window exclusion,
+    but with the correct temporal index handling.
+
+    Parameters
+    ----------
+    `X` : `np.ndarray`
+        The input data of shape (N,) or (N, E) or (B, N, E).
+    `Y` : `np.ndarray`
+        The target data of shape (N,) or (N, E') or (B, N, E').
+    `theiler_window` : `int`
+        Theiler window half-width. Library points ``j`` where
+        ``|i - j| <= theiler_window`` are excluded when predicting point ``i``.
+        For lagged embedding, use ``(E - 1) * tau + n_ahead``.
+
+    Returns
+    -------
+    predictions : `np.ndarray`
+        The predicted values of shape (N,) or (N, E') or (B, N, E').
+
+    Raises
+    ------
+    ValueError
+        - If the input arrays `X` and `Y` do not have the same number of points.
+        - If there are not enough library points outside the Theiler window.
+    """
+    # ensure 2D or 3D arrays
+    if X.ndim == 1:
+        X = X[:, None]
+    if Y.ndim == 1:
+        Y = Y[:, None]
+
+    # X (N, E), Y (N, E')
+    if X.ndim == 2 and Y.ndim == 2:
+        N, E = X.shape
+        if Y.shape[0] != N:
+            raise ValueError(f"X and Y must have the same length, got X.shape={X.shape} and Y.shape={Y.shape}")
+
+        k: int = E + 1
+        n_exclude = 2 * theiler_window + 1
+
+        if N - n_exclude < k:
+            raise ValueError(
+                f"Not enough library points outside Theiler window: need at least k={k} points, but only {N - n_exclude} available, N={N}, theiler_window={theiler_window}"
+            )
+
+        distances, indices = knn(X, X, k + n_exclude)
+
+        distances = np.where(np.abs(indices - np.arange(N)[:, None]) <= theiler_window, np.inf, distances)
+
+        top_k = np.argsort(distances, axis=1)[:, :k]
+        distances = np.take_along_axis(distances, top_k, axis=1)
+        indices = np.take_along_axis(indices, top_k, axis=1)
+
+        Y_neighbors = Y[indices]  # (N, k, E')
+
+        # clamp to avoid division by zero
+        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)
+
+        return predictions.squeeze()  # (N,) or (N, E')
+    # X (B, N, E), Y (B, N, E')
+    elif X.ndim == 3 and Y.ndim == 3:
+        B, N, E = X.shape
+        if Y.shape[0] != B or Y.shape[1] != N:
+            raise ValueError(f"batch size and length of X and Y must match, got X.shape={X.shape} and Y.shape={Y.shape}")
+
+        k: int = E + 1
+        n_exclude = 2 * theiler_window + 1
+
+        if N - n_exclude < k:
+            raise ValueError(
+                f"Not enough library points outside Theiler window: need at least k={k} points, but only {N - n_exclude} available, N={N}, theiler_window={theiler_window}"
+            )
+
+        distances = np.empty((B, N, k + n_exclude))
+        indices = np.empty((B, N, k + n_exclude), dtype=np.intp)
+        for b in range(B):
+            distances[b], indices[b] = knn(X[b], X[b], k + n_exclude)
+
+        distances = np.where(np.abs(indices - np.arange(N)[None, :, None]) <= theiler_window, np.inf, distances)
+
+        top_k = np.argsort(distances, axis=2)[:, :, :k]
+        distances = np.take_along_axis(distances, top_k, axis=2)
+        indices = np.take_along_axis(indices, top_k, axis=2)
+
+        batch_idx = np.arange(B)[:, None, None]  # (B, 1, 1)
+        Y_neighbors = Y[batch_idx, indices]  # (B, N, k, E')
+
+        # clamp to avoid division by zero
+        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')
+
+        return predictions
+    else:
+        raise ValueError(f"X and Y must both be 2D or both be 3D arrays, got X.ndim={X.ndim}, Y.ndim={Y.ndim}")

{edmkit-0.0.4/src/edmkit → edmkit-0.0.6/src/edmkit/simplex_projection}/simplex_projection.py

@@ -1,10 +1,9 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-from scipy.spatial import KDTree
 from tinygrad import Tensor, dtypes
-from usearch.index import Index
 
+from edmkit.simplex_projection.knn import knn
 from edmkit.util import pairwise_distance
 
 
@@ -22,11 +21,13 @@ def simplex_projection(
     Parameters
     ----------
     `X` : `np.ndarray`
-        The input data
+        The input data of shape (N,) or (N, E) or (B, N, E)
     `Y` : `np.ndarray`
-        The target data
+        The target data of shape (N,) or (N, E') or (B, N, E')
     `Q` : `np.ndarray`
-        The query points for which to find the nearest neighbors in `X`.
+        The query points of shape (M,) or (M, E) or (B, M, E) for which to find the nearest neighbors in `X`.
+    `mask` : `np.ndarray | None`
+        Boolean mask of shape (N,) or (B, N) indicating which library points to include when finding nearest neighbors for the queries in `Q`.
     `use_tensor` : `bool`, default `False`
         Whether to use `tinygrad.Tensor` for computation.
         **This may be slower than the NumPy implementation in most cases for now.**
@@ -78,43 +79,6 @@ def simplex_projection(
     return _numpy(X, Y, Q, mask=mask) if not use_tensor else _tensor(X, Y, Q, mask=mask)
 
 
-def knn(X: np.ndarray, Q: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
-    """Find the k-nearest neighbors of `Q` in `X` using either `usearch` or `scipy.spatial.KDTree` depending on the size and dimensionality of the data.
-
-    Parameters
-    ----------
-    `X` : `np.ndarray`
-        The input data (N, E)
-    `Q` : `np.ndarray`
-        The query points (M, E)
-    `k` : `int`
-        The number of nearest neighbors to find (typically E+1 for simplex projection).
-
-    Returns
-    -------
-    distances : `np.ndarray`
-        The distances from each query point in `Q` to its k nearest neighbors in `X` (M, k)
-    indices : `np.ndarray`
-        The indices of the k nearest neighbors in `X` for each query point in `Q` (M, k)
-    """
-
-    N, E = X.shape
-
-    if N < k:
-        raise ValueError(f"Not enough points in X to find {k} neighbors, got N={N}")
-
-    if E >= 15 and N >= 10_000:
-        index = Index(ndim=E, metric="l2sq")
-        index.add(np.arange(len(X)), np.ascontiguousarray(X, dtype=np.float32))
-        matches = index.search(np.ascontiguousarray(Q, dtype=np.float32), k)
-        distances = np.atleast_2d(np.sqrt(np.asarray(matches.distances)))
-        indices = np.atleast_2d(np.asarray(matches.keys).astype(np.intp))
-        return distances, indices
-    else:
-        tree = KDTree(X)
-        return tree.query(Q, k=k)
-
-
 def _numpy(
     X: np.ndarray,
     Y: np.ndarray,
@@ -122,30 +86,6 @@ def _numpy(
     *,
     mask: np.ndarray | None = None,
 ):
-    """
-    Perform simplex projection from `X` to `Y` using the nearest neighbors of the points specified by `Q`.
-
-    Parameters
-    ----------
-    `X` : `np.ndarray`
-        (N,) or (N, E) or (B, N, E)
-    `Y` : `np.ndarray`
-        (N,) or (N, E') or (B, N, E')
-    `Q` : `np.ndarray`
-        The query points for which to find the nearest neighbors in `X`.
-        (M,) or (M, E) or (B, M, E)
-
-    Returns
-    -------
-    predictions : `np.ndarray`
-        The predicted values based on the weighted mean of the nearest neighbors in `Y`.
-        (M,) or (M, E') or (B, M, E')
-
-    Raises
-    ------
-    ValueError
-        - If the input arrays `X` and `Y` do not have the same number of points.
-    """
     # ensure 2D or 3D arrays
     if X.ndim == 1:
         X = X[:, None]
@@ -154,19 +94,19 @@ def _numpy(
     if Q.ndim == 1:
         Q = Q[:, None]
 
-    # X (N, E), Y (N, E'), Q (M, E)
+    # X (N, E), Y (N, E'), Q (M, E), mask (N,)
     if X.ndim == 2 and Y.ndim == 2 and Q.ndim == 2:
-        if X.shape[0] != Y.shape[0]:
+        N, E = X.shape
+        if Y.shape[0] != N:
             raise ValueError(f"X and Y must have the same length, got X.shape={X.shape} and Y.shape={Y.shape}")
 
-        k: int = X.shape[1] + 1
+        k: int = E + 1
 
         if mask is not None:
             X = X[mask]
             Y = Y[mask]
 
         distances, indices = knn(X, Q, k)
-
         Y_neighbors = Y[indices]  # (M, k, E')
 
         # clamp to avoid division by zero
@@ -221,30 +161,6 @@ def _tensor(
     *,
     mask: np.ndarray | None = None,
 ):
-    """
-    Perform simplex projection from `X` to `Y` using the nearest neighbors of the points specified by `Q`.
-
-    Parameters
-    ----------
-    `X` : `np.ndarray`
-        (N,) or (N, E) or (B, N, E)
-    `Y` : `np.ndarray`
-        (N,) or (N, E') or (B, N, E')
-    `Q` : `np.ndarray`
-        The query points for which to find the nearest neighbors in `X`.
-        (M,) or (M, E) or (B, M, E)
-
-    Returns
-    -------
-    predictions : `np.ndarray`
-        The predicted values based on the weighted mean of the nearest neighbors in `Y`.
-        (M,) or (M, E') or (B, M, E')
-
-    Raises
-    ------
-    ValueError
-        - If the input arrays `X` and `Y` do not have the same number of points.
-    """
     if X.ndim == 1:
         X = X[:, None]
     if Y.ndim == 1:
@@ -297,29 +213,10 @@ def _tensor(
 
         distances, indices = D.topk(k, dim=2, largest=False, sorted_=True)  # (B, M, k)
 
-        # --- Neighbor lookup -------------------------------------------------------
-        # Purpose:
-        #   `indices` contains the k-nearest-neighbor indices in X for each batch and each query point. (B, M, k)
-        #   However, we need to gather the corresponding Y values from (B, N, E'),
-        #   and tinygrad currently doesn’t support a batched gather operation like PyTorch does.
-        #   Therefore, we flatten the batch dimension so we can perform a single gather
-        #   from a flattened (B*N, E') tensor.
-        #
-        # Notation:
-        #   B = batch size, M = number of query points, N = number of library points,
-        #   k = number of neighbors, E' = output dimension
-        #
-        # Steps:
-        #   1) Create per-batch offsets [0*N, 1*N, ..., (B-1)*N]
-        #   2) Add these offsets to the neighbor indices (B, M, k)
-        #      -> converts them to flattened indices relative to (B*N)
-        #   3) Reshape Y into (B*N, E') and gather using the flattened indices
-        #   4) Reshape the gathered results back to (B, M, k, E') to continue computation
-        offsets = Tensor.arange(B, dtype=dtypes.int32).reshape(B, 1, 1) * N  # (B,1,1) create per-batch offsets spaced by N
-        flat_indices = (indices + offsets).reshape(B * Q.shape[1], k)  # (B*M, k) flatten batch and query dimensions
-        Y_flat = Y_tensor.reshape(B * N, Y_tensor.shape[-1])  # (B*N, E') flatten batch and library points
-        Y_neighbors = Y_flat[flat_indices].reshape(B, Q.shape[1], k, Y_tensor.shape[-1])  # (B, M, k, E') restore shape
-        # ---------------------------------------------------------------------------
+        offsets = Tensor.arange(B, dtype=dtypes.int32).reshape(B, 1, 1) * N
+        flat_indices = (indices + offsets).reshape(B * Q.shape[1], k)
+        Y_flat = Y_tensor.reshape(B * N, Y_tensor.shape[-1])
+        Y_neighbors = Y_flat[flat_indices].reshape(B, Q.shape[1], k, Y_tensor.shape[-1])
 
         d_min = distances[:, :, :1].clip(min_=1e-6)  # (B, M, 1)
         weights: Tensor = (-distances / d_min).exp()  # (B, M, k)

{edmkit-0.0.4 → edmkit-0.0.6}/src/edmkit/splits.py

@@ -103,6 +103,8 @@ def expanding_folds(
 
     if stride is None:
         stride = validation_size
+    if stride <= 0:
+        raise ValueError(f"stride must be positive, got {stride}")
 
     folds: list[Fold] = []
     validation_start = initial_train_size + gap
@@ -156,6 +158,8 @@ def sliding_folds(
 
     if stride is None:
         stride = validation_size
+    if stride <= 0:
+        raise ValueError(f"stride must be positive, got {stride}")
 
     folds: list[Fold] = []
     validation_start = train_size + gap