edmkit 0.0.5__tar.gz → 0.0.7__tar.gz

{edmkit-0.0.5 → edmkit-0.0.7}/PKG-INFO

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

{edmkit-0.0.5 → edmkit-0.0.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.5"
+version = "0.0.7"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
 readme = "README.md"
@@ -14,6 +14,7 @@ dependencies = [
 
 [dependency-groups]
 dev = [
+    "griffe2md>=1.5.0",
     "hypothesis>=6.151.9",
     "pytest>=9.0.2",
     "ruff>=0.15.6",

{edmkit-0.0.5 → edmkit-0.0.7}/src/edmkit/ccm.py

@@ -1,6 +1,5 @@
 from collections.abc import Callable
 from functools import partial
-from typing import TypeAlias
 
 import numpy as np
 
@@ -8,9 +7,9 @@ from edmkit.simplex_projection import simplex_projection
 from edmkit.smap import smap
 from edmkit.types import PredictFunc
 
-SampleFunc: TypeAlias = Callable[[np.ndarray, int], np.ndarray]
+type SampleFunc = Callable[[np.ndarray, int], np.ndarray]
 """SampleFunc is a function that takes (pool, size) and returns a sampled array."""
-AggregateFunc: TypeAlias = Callable[[np.ndarray], float]
+type AggregateFunc = Callable[[np.ndarray], float]
 """AggregateFunc is a function that takes an array of values and returns a single value."""
 
 
@@ -50,7 +49,7 @@ def bootstrap(
         Target time series (potential driver)
     lib_sizes : np.ndarray
         Array of library sizes to test convergence.
-    predict_func : :type: `PredictFunc`
+    predict_func : PredictFunc
         Prediction function with signature (X, Y, Q) -> predictions.
     n_samples : int, default 20
         Number of random samples per library size for bootstrapping.
@@ -58,16 +57,16 @@ def bootstrap(
         1-D array of integer indices from which library members are sampled.
     prediction_pool : np.ndarray
         1-D array of integer indices that are predicted.
-    sample_func : :type: `SampleFunc` | None, default None
+    sample_func : SampleFunc or None, default None
         Function responsible for drawing a library sample of a given size.
         When None, a fresh RNG-backed sampler is created per call.
-    batch_size : int | None, default 10
+    batch_size : int or None, default 10
         If specified, predictions are made in batches to limit memory usage.
 
     Returns
     -------
-    samples : np.ndarray of shape (n_samples, len(lib_sizes))
-        Per-sample correlation coefficients.
+    samples : np.ndarray
+        Per-sample correlation coefficients of shape ``(n_samples, len(lib_sizes))``.
     """
     if sample_func is None:
         sample_func = make_sample_func()
@@ -142,7 +141,7 @@ def ccm(
         Target time series (potential driver)
     lib_sizes : np.ndarray
         Array of library sizes to test convergence.
-    predict_func : :type: `PredictFunc`
+    predict_func : PredictFunc
         Prediction function with signature (X, Y, Q) -> predictions.
         Can be `simplex_projection`, `smap` with partial application, or a custom function.
     n_samples : int, default 100
@@ -151,13 +150,13 @@ def ccm(
         1-D array of integer indices from which library members are sampled.
     prediction_pool : np.ndarray
         1-D array of integer indices that are predicted.
-    sample_func : :type: `SampleFunc` | None, default None
+    sample_func : SampleFunc or None, default None
         Function responsible for drawing a library sample of a given size.
-        It receives `(pool, size)` and returns an array of indices.
+        It receives ``(pool, size)`` and returns an array of indices.
         When None, a fresh RNG-backed sampler is created per call.
-    aggregate_func : :type: `AggregateFunc`, default `np.mean`
+    aggregate_func : AggregateFunc, default np.mean
         Reducer applied to the correlation samples for each library size.
-    batch_size : int | None, default None
+    batch_size : int or None, default None
         If not specified, batch_size == n_samples.
         If specified, predictions are made in batches to limit memory usage.
     Returns

{edmkit-0.0.5 → edmkit-0.0.7}/src/edmkit/embedding.py

@@ -14,13 +14,17 @@ def lagged_embed(x: np.ndarray, tau: int, e: int):
 
     Parameters
     ----------
-        `x` : `np.ndarray` of shape `(N,)`
-        `tau` : `int`
-        `e` : `int`
+    x : np.ndarray
+        1D time series of shape ``(N,)``.
+    tau : int
+        Time delay.
+    e : int
+        Embedding dimension.
 
     Returns
     -------
-        `np.ndarray` of shape `(N - (e - 1) * tau, e)`
+    np.ndarray
+        Embedded array of shape ``(N - (e - 1) * tau, e)``.
 
     Raises
     ------

{edmkit-0.0.5 → edmkit-0.0.7}/src/edmkit/generate/__init__.py

@@ -2,3 +2,5 @@
 from .double_pendulum import double_pendulum, to_xy
 from .lorenz import lorenz
 from .mackey_glass import mackey_glass
+
+__all__ = ["double_pendulum", "to_xy", "lorenz", "mackey_glass"]

edmkit-0.0.7/src/edmkit/generate/double_pendulum.py

@@ -0,0 +1,105 @@
+import numpy as np
+
+
+def double_pendulum(
+    m1: float,
+    m2: float,
+    L1: float,
+    L2: float,
+    g: float,
+    X0: np.ndarray,
+    dt: float,
+    t_max: int,
+):
+    """Generate double pendulum dynamics via forward Euler integration.
+
+    Parameters
+    ----------
+    m1 : float
+        Mass of first pendulum.
+    m2 : float
+        Mass of second pendulum.
+    L1 : float
+        Length of first pendulum.
+    L2 : float
+        Length of second pendulum.
+    g : float
+        Gravitational acceleration.
+    X0 : np.ndarray
+        Initial state ``(theta1, theta2, omega1, omega2)`` of shape ``(4,)``.
+    dt : float
+        Integration time step.
+    t_max : int
+        Maximum time.
+
+    Returns
+    -------
+    t : np.ndarray
+        Time array.
+    X : np.ndarray
+        State trajectory of shape ``(N, 4)``.
+    """
+
+    def f(x: np.ndarray):
+        theta1, theta2, omega1, omega2 = x
+        delta = theta1 - theta2
+
+        denom = 2 * m1 + m2 - m2 * np.cos(2 * delta)
+
+        dtheta1_dt = omega1
+        dtheta2_dt = omega2
+
+        domega1_dt = (
+            -g * (2 * m1 + m2) * np.sin(theta1)
+            - m2 * g * np.sin(theta1 - 2 * theta2)
+            - 2 * np.sin(delta) * m2 * (omega2**2 * L2 + omega1**2 * L1 * np.cos(delta))
+        ) / (L1 * denom)
+
+        domega2_dt = (2 * np.sin(delta) * (omega1**2 * L1 * (m1 + m2) + g * (m1 + m2) * np.cos(theta1) + omega2**2 * L2 * m2 * np.cos(delta))) / (
+            L2 * denom
+        )
+
+        return np.array([dtheta1_dt, dtheta2_dt, domega1_dt, domega2_dt])
+
+    t = np.arange(0, t_max, dt)
+    X = np.zeros((len(t), 4))
+    X[0] = X0
+
+    for i in range(1, len(t)):
+        X[i] = X[i - 1] + dt * f(X[i - 1])
+
+    return t, X
+
+
+def to_xy(L1: float, L2: float, theta1: np.ndarray, theta2: np.ndarray):
+    """Convert double pendulum angles to Cartesian coordinates.
+
+    Parameters
+    ----------
+    L1 : float
+        Length of first pendulum.
+    L2 : float
+        Length of second pendulum.
+    theta1 : np.ndarray
+        Angle of first pendulum.
+    theta2 : np.ndarray
+        Angle of second pendulum.
+
+    Returns
+    -------
+    x1 : np.ndarray
+        x-coordinate of first pendulum.
+    y1 : np.ndarray
+        y-coordinate of first pendulum.
+    x2 : np.ndarray
+        x-coordinate of second pendulum.
+    y2 : np.ndarray
+        y-coordinate of second pendulum.
+    """
+    x1 = L1 * np.sin(theta1)
+    y1 = -L1 * np.cos(theta1)
+
+    x2 = x1 + L2 * np.sin(theta2)
+    y2 = y1 - L2 * np.cos(theta2)
+
+    return x1, y1, x2, y2

edmkit-0.0.7/src/edmkit/generate/lorenz.py

@@ -0,0 +1,47 @@
+import numpy as np
+
+
+def lorenz(
+    sigma: float,
+    rho: float,
+    beta: float,
+    X0: np.ndarray,
+    dt: float,
+    t_max: int,
+):
+    """Generate a Lorenz system trajectory via forward Euler integration.
+
+    Parameters
+    ----------
+    sigma : float
+        Prandtl number (typical: 10).
+    rho : float
+        Rayleigh number (typical: 28).
+    beta : float
+        Geometric factor (typical: 8/3).
+    X0 : np.ndarray
+        Initial condition of shape ``(3,)``.
+    dt : float
+        Integration time step.
+    t_max : int
+        Maximum time.
+
+    Returns
+    -------
+    t : np.ndarray
+        Time array.
+    X : np.ndarray
+        Trajectory of shape ``(N, 3)`` for ``(x, y, z)``.
+    """
+
+    def f(x: np.ndarray):
+        return np.array([[-sigma, sigma, 0], [rho, -1, -x[0]], [0, x[0], -beta]]) @ x
+
+    t = np.arange(0, t_max, dt)
+    X = np.zeros((len(t), 3))
+    X[0] = X0
+
+    for i in range(1, len(t)):
+        X[i] = X[i - 1] + dt * f(X[i - 1])
+
+    return t, X

edmkit-0.0.7/src/edmkit/generate/mackey_glass.py

@@ -0,0 +1,44 @@
+import numpy as np
+
+
+def mackey_glass(tau: float, n: int, beta: float, gamma: float, x0: float, dt: float, t_max: int):
+    """Generate a Mackey-Glass chaotic time series via forward Euler integration.
+
+    Parameters
+    ----------
+    tau : float
+        Delay parameter (typical: 17 for chaos).
+    n : int
+        Nonlinearity exponent (typical: 10).
+    beta : float
+        Feedback strength (typical: 0.2).
+    gamma : float
+        Decay rate (typical: 0.1).
+    x0 : float
+        Initial condition.
+    dt : float
+        Integration time step.
+    t_max : int
+        Maximum time.
+
+    Returns
+    -------
+    t : np.ndarray
+        Time array.
+    x : np.ndarray
+        1D time series.
+    """
+
+    def f(x, x_tau):
+        return beta * x_tau / (1 + x_tau**n) - gamma * x
+
+    t = np.arange(0, t_max, dt)
+    x = np.zeros_like(t)
+
+    tau_idx = int(tau / dt)
+    x[:tau_idx] = x0
+
+    for i in range(tau_idx, len(t)):
+        x[i] = x[i - 1] + dt * f(x[i - 1], x[i - tau_idx])
+
+    return t, x

{edmkit-0.0.5 → edmkit-0.0.7}/src/edmkit/metrics.py

@@ -1,8 +1,8 @@
-from typing import TYPE_CHECKING, Callable, TypeAlias
+from typing import TYPE_CHECKING, Callable
 
 import numpy as np
 
-MetricFunc: TypeAlias = Callable[[np.ndarray, np.ndarray], np.ndarray]
+type MetricFunc = Callable[[np.ndarray, np.ndarray], np.ndarray]
 """MetricFunc is a function that takes (predictions, observations) and returns a metric value."""
 
 

edmkit-0.0.7/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.7/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.7/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``.
+
+    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.5/src/edmkit → edmkit-0.0.7/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
 
 
@@ -21,19 +20,21 @@ def simplex_projection(
 
     Parameters
     ----------
-    `X` : `np.ndarray`
-        The input data
-    `Y` : `np.ndarray`
-        The target data
-    `Q` : `np.ndarray`
-        The query points for which to find the nearest neighbors in `X`.
-    `use_tensor` : `bool`, default `False`
+    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')
+    Q : np.ndarray
+        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 or 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.**
 
     Returns
     -------
-    predictions : `np.ndarray`
+    predictions : np.ndarray
         The predicted values based on the weighted mean of the nearest neighbors in `Y`.
 
     Raises
@@ -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.5 → edmkit-0.0.7}/src/edmkit/smap.py

@@ -21,23 +21,23 @@ def smap(
 
     Parameters
     ----------
-    `X` : `np.ndarray`
+    X : np.ndarray
         The input data
-    `Y` : `np.ndarray`
+    Y : np.ndarray
         The target data
-    `Q` : `np.ndarray`
+    Q : np.ndarray
         The query points for which to make predictions.
-    `theta` : `float`
+    theta : float
         Locality parameter. (0: global linear, >0: local linear)
-    `alpha` : `float`, default `1e-10`
+    alpha : float, default 1e-10
         Regularization parameter to stabilize the inversion.
-    `use_tensor` : `bool`, default `False`
+    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.**
 
     Returns
     -------
-    predictions : `np.ndarray`
+    predictions : np.ndarray
         The predicted values based on the weighted linear regression.
 
     Raises
@@ -137,21 +137,21 @@ def _numpy(
 
     Parameters
     ----------
-    `X` : `np.ndarray`
+    X : np.ndarray
         (N,) or (N, E) or (B, N, E)
-    `Y` : `np.ndarray`
+    Y : np.ndarray
         (N,) or (N, E') or (B, N, E')
-    `Q` : `np.ndarray`
+    Q : np.ndarray
         The query points for which to make predictions.
         (M,) or (M, E) or (B, M, E)
-    `theta` : `float`
+    theta : float
         Locality parameter. (0: global linear, >0: local linear)
-    `alpha` : `float`, default `1e-10`
+    alpha : float, default 1e-10
         Regularization parameter to stabilize the inversion.
 
     Returns
     -------
-    predictions : `np.ndarray`
+    predictions : np.ndarray
         The predicted values based on the weighted linear regression.
         (M, E') or (B, M, E')
 
@@ -245,20 +245,20 @@ def _tensor(
 
     Parameters
     ----------
-    `X` : `np.ndarray`
+    X : np.ndarray
         The input data
-    `Y` : `np.ndarray`
+    Y : np.ndarray
         The target data
-    `Q` : `np.ndarray`
+    Q : np.ndarray
         The query points for which to make predictions.
-    `theta` : `float`
+    theta : float
         Locality parameter. (0: global linear, >0: local linear)
-    `alpha` : `float`, default `1e-10`
+    alpha : float, default 1e-10
         Regularization parameter to stabilize the inversion.
 
     Returns
     -------
-    predictions : `np.ndarray`
+    predictions : np.ndarray
         The predicted values based on the weighted linear regression.
 
     Raises

{edmkit-0.0.5 → edmkit-0.0.7}/src/edmkit/splits.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Callable, NamedTuple, TypeAlias
+from typing import TYPE_CHECKING, Callable, NamedTuple
 
 import numpy as np
 
@@ -10,7 +10,7 @@ class Fold(NamedTuple):
     validation: np.ndarray  # index array
 
 
-SplitFunc: TypeAlias = Callable[[int], list[Fold]]
+type SplitFunc = Callable[[int], list[Fold]]
 
 
 def temporal_fold(

{edmkit-0.0.5 → edmkit-0.0.7}/src/edmkit/util.py

@@ -8,11 +8,13 @@ def pad(As: list[np.ndarray]):
 
     Parameters
     ----------
-        `As` : `list` of `np.ndarray` of shape `(N, D_i)`
+    As : list[np.ndarray]
+        List of arrays of shape ``(N, D_i)``.
 
     Returns
     -------
-        Single `np.ndarray` of shape `(B, N, max(D))` where `B` is `len(As)`
+    np.ndarray
+        Single array of shape ``(B, N, max(D))`` where B is ``len(As)``.
 
     Raises
     ------
@@ -41,21 +43,18 @@ def pairwise_distance(A: Tensor, B: Tensor | None = None) -> Tensor:
 
     Parameters
     ----------
-    `A` : `Tensor` of shape `(N, D)` or `(B, N, D)`
-        - `B`: batch size
-        - `N`: number of points
-        - `D`: dimension of each point
-    `B` : `Tensor` of shape `(M, D)` or `(B, M, D)`
-        - `B`: batch size
-        - `M`: number of points
-        - `D`: dimension of each point
+    A : Tensor
+        Shape ``(N, D)`` or ``(B, N, D)``.
+        ``B`` is batch size, ``N`` is number of points, ``D`` is dimension of each point.
+    B : Tensor
+        Shape ``(M, D)`` or ``(B, M, D)``.
+        ``B`` is batch size, ``M`` is number of points, ``D`` is dimension of each point.
 
     Returns
     -------
-    When `A` is of shape `(N, D)`:
-        `Tensor` of shape `(N, N)` [or `(N, M)`] where the element at position `(i, j)` is the squared Euclidean distance between `A[i]` and `A[j]` [or between `A[i]` and `B[j]`].
-    When `A` is of shape `(B, N, D)`:
-        `Tensor` of shape `(B, N, N)` [or `(B, N, M)`] where the element at position `(b, i, j)` is the squared Euclidean distance between `A[b, i]` and `A[b, j]`.
+    Tensor
+        When `A` is of shape ``(N, D)``: shape ``(N, N)`` [or ``(N, M)``] where the element at position ``(i, j)`` is the squared Euclidean distance between ``A[i]`` and ``A[j]`` [or between ``A[i]`` and ``B[j]``].
+        When `A` is of shape ``(B, N, D)``: shape ``(B, N, N)`` [or ``(B, N, M)``] where the element at position ``(b, i, j)`` is the squared Euclidean distance between ``A[b, i]`` and ``A[b, j]``.
 
     Raises
     ------
@@ -84,21 +83,18 @@ def pairwise_distance_np(A: np.ndarray, B: np.ndarray | None = None) -> np.ndarr
 
     Parameters
     ----------
-    `A` : `np.ndarray` of shape `(N, D)` or `(B, N, D)`
-        - `B`: batch size
-        - `N`: number of points
-        - `D`: dimension of each point
-    `B` : `np.ndarray` of shape `(M, D)` or `(B, M, D)`
-        - `B`: batch size
-        - `M`: number of points
-        - `D`: dimension of each point
+    A : np.ndarray
+        Shape ``(N, D)`` or ``(B, N, D)``.
+        ``B`` is batch size, ``N`` is number of points, ``D`` is dimension of each point.
+    B : np.ndarray
+        Shape ``(M, D)`` or ``(B, M, D)``.
+        ``B`` is batch size, ``M`` is number of points, ``D`` is dimension of each point.
 
     Returns
     -------
-    When `A` is of shape `(N, D)`:
-        `np.ndarray` of shape `(N, N)` [or `(N, M)`] where the element at position `(i, j)` is the squared Euclidean distance between `A[i]` and `A[j]` [or between `A[i]` and `B[j]`].
-    When `A` is of shape `(B, N, D)`:
-        `np.ndarray` of shape `(B, N, N)` [or `(B, N, M)`] where the element at position `(b, i, j)` is the squared Euclidean distance between `A[b, i]` and `A[b, j]`.
+    np.ndarray
+        When `A` is of shape ``(N, D)``: shape ``(N, N)`` [or ``(N, M)``] where the element at position ``(i, j)`` is the squared Euclidean distance between ``A[i]`` and ``A[j]`` [or between ``A[i]`` and ``B[j]``].
+        When `A` is of shape ``(B, N, D)``: shape ``(B, N, N)`` [or ``(B, N, M)``] where the element at position ``(b, i, j)`` is the squared Euclidean distance between ``A[b, i]`` and ``A[b, j]``.
 
     Raises
     ------
@@ -128,12 +124,15 @@ def dtw(A: np.ndarray, B: np.ndarray):
 
     Parameters
     ----------
-        `A` : Tensor of shape `(N,D)`
-        `B` : Tensor of shape `(M,D)`
+    A : np.ndarray
+        Sequence of shape ``(N, D)``.
+    B : np.ndarray
+        Sequence of shape ``(M, D)``.
 
     Returns
     -------
-        distance : float
+    distance : float
+        The DTW distance between the two sequences.
     """
     N: int = A.shape[0]
     M: int = B.shape[0]
@@ -169,13 +168,17 @@ def autocorrelation(x: np.ndarray, max_lag: int, step: int = 1):
 
     Parameters
     ----------
-        `x` : `np.ndarray` The input array for which to compute the autocorrelation.
-        `max_lag` : `int` The maximum lag up to which the autocorrelation is computed.
-        `step` : `int`, `optional` The step size for the lag. Default is 1.
+    x : np.ndarray
+        The input array for which to compute the autocorrelation.
+    max_lag : int
+        The maximum lag up to which the autocorrelation is computed.
+    step : int, optional
+        The step size for the lag. Default is 1.
 
     Returns
     -------
-        `np.ndarray` of shape `(max_lag // step + 1,)` containing the autocorrelation values.
+    np.ndarray
+        Array of shape ``(max_lag // step + 1,)`` containing the autocorrelation values.
     """
     x = x - np.mean(x)
 

edmkit-0.0.5/src/edmkit/generate/double_pendulum.py

@@ -1,52 +0,0 @@
-import numpy as np
-
-
-def double_pendulum(
-    m1: float,
-    m2: float,
-    L1: float,
-    L2: float,
-    g: float,
-    X0: np.ndarray,  # (theta1, theta2, omega1, omega2)
-    dt: float,
-    t_max: int,
-):
-    def f(x: np.ndarray):
-        theta1, theta2, omega1, omega2 = x
-        delta = theta1 - theta2
-
-        denom = 2 * m1 + m2 - m2 * np.cos(2 * delta)
-
-        dtheta1_dt = omega1
-        dtheta2_dt = omega2
-
-        domega1_dt = (
-            -g * (2 * m1 + m2) * np.sin(theta1)
-            - m2 * g * np.sin(theta1 - 2 * theta2)
-            - 2 * np.sin(delta) * m2 * (omega2**2 * L2 + omega1**2 * L1 * np.cos(delta))
-        ) / (L1 * denom)
-
-        domega2_dt = (2 * np.sin(delta) * (omega1**2 * L1 * (m1 + m2) + g * (m1 + m2) * np.cos(theta1) + omega2**2 * L2 * m2 * np.cos(delta))) / (
-            L2 * denom
-        )
-
-        return np.array([dtheta1_dt, dtheta2_dt, domega1_dt, domega2_dt])
-
-    t = np.arange(0, t_max, dt)
-    X = np.zeros((len(t), 4))
-    X[0] = X0
-
-    for i in range(1, len(t)):
-        X[i] = X[i - 1] + dt * f(X[i - 1])
-
-    return t, X
-
-
-def to_xy(L1: float, L2: float, theta1: np.ndarray, theta2: np.ndarray):
-    x1 = L1 * np.sin(theta1)
-    y1 = -L1 * np.cos(theta1)
-
-    x2 = x1 + L2 * np.sin(theta2)
-    y2 = y1 - L2 * np.cos(theta2)
-
-    return x1, y1, x2, y2

edmkit-0.0.5/src/edmkit/generate/lorenz.py

@@ -1,22 +0,0 @@
-import numpy as np
-
-
-def lorenz(
-    sigma: float,
-    rho: float,
-    beta: float,
-    X0: np.ndarray,
-    dt: float,
-    t_max: int,
-):
-    def f(x: np.ndarray):
-        return np.array([[-sigma, sigma, 0], [rho, -1, -x[0]], [0, x[0], -beta]]) @ x
-
-    t = np.arange(0, t_max, dt)
-    X = np.zeros((len(t), 3))
-    X[0] = X0
-
-    for i in range(1, len(t)):
-        X[i] = X[i - 1] + dt * f(X[i - 1])
-
-    return t, X

edmkit-0.0.5/src/edmkit/generate/mackey_glass.py

@@ -1,17 +0,0 @@
-import numpy as np
-
-
-def mackey_glass(tau: float, n: int, beta: float, gamma: float, x0: float, dt: float, t_max: int):
-    def f(x, x_tau):
-        return beta * x_tau / (1 + x_tau**n) - gamma * x
-
-    t = np.arange(0, t_max, dt)
-    x = np.zeros_like(t)
-
-    tau_idx = int(tau / dt)
-    x[:tau_idx] = x0
-
-    for i in range(tau_idx, len(t)):
-        x[i] = x[i - 1] + dt * f(x[i - 1], x[i - tau_idx])
-
-    return t, x