edmkit 0.0.13__tar.gz → 0.0.15__tar.gz

{edmkit-0.0.13 → edmkit-0.0.15}/PKG-INFO

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

{edmkit-0.0.13 → edmkit-0.0.15}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.13"
+version = "0.0.15"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 readme = "README.md"
 requires-python = ">= 3.13"
@@ -20,6 +20,7 @@ email = "tenma.x0@gmail.com"
 dev = [
     "griffe2md>=1.5.0",
     "hypothesis>=6.151.9",
+    "matplotlib>=3.11.2",
     "pytest>=9.1.1",
     "ruff>=0.15.20",
     "ty>=0.0.56",

{edmkit-0.0.13 → edmkit-0.0.15}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.13"
+version = "0.0.15"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
 readme = "README.md"
@@ -17,6 +17,7 @@ dependencies = [
 dev = [
     "griffe2md>=1.5.0",
     "hypothesis>=6.151.9",
+    "matplotlib>=3.11.2",
     "pytest>=9.1.1",
     "ruff>=0.15.20",
     "ty>=0.0.56",

edmkit-0.0.15/src/edmkit/generate/__init__.py

@@ -0,0 +1,10 @@
+from .ctrnn import ctrnn
+from .double_pendulum import double_pendulum, to_xy
+from .duffing import duffing
+from .lorenz import lorenz
+from .lorenz96 import lorenz96
+from .mackey_glass import mackey_glass
+from .rk45 import rk45
+from .rossler import rossler
+
+__all__ = ["ctrnn", "double_pendulum", "duffing", "lorenz", "lorenz96", "mackey_glass", "rk45", "rossler", "to_xy"]

edmkit-0.0.15/src/edmkit/generate/ctrnn.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+__all__ = ["ctrnn"]
+
+
+def ctrnn(g: float, J: np.ndarray):
+    """Vector field of a continuous-time recurrent neural network ``dx/dt = -x + g J tanh(x)``.
+
+    Parameters
+    ----------
+    g : float
+        Gain (chaos for ``g > 1`` when ``J`` is large and random with variance ``1 / N``).
+    J : np.ndarray
+        Connectivity of shape ``(N, N)``; ``J[i, j]`` is the weight with which unit ``j`` drives unit ``i``.
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state of shape ``(N,)``.
+    """
+
+    def f(t: float, x: np.ndarray):
+        return -x + g * J @ np.tanh(x)
+
+    return f

{edmkit-0.0.13 → edmkit-0.0.15}/src/edmkit/generate/double_pendulum.py

@@ -3,17 +3,8 @@ import numpy as np
 __all__ = ["double_pendulum", "to_xy"]
 
 
-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.
+def double_pendulum(m1: float, m2: float, L1: float, L2: float, g: float):
+    """Vector field of the double pendulum.
 
     Parameters
     ----------
@@ -27,22 +18,14 @@ def double_pendulum(
         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)``.
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(theta1, theta2, omega1, omega2)`` of shape ``(4,)``.
     """
 
-    def f(x: np.ndarray):
+    def f(t: float, x: np.ndarray):
         theta1, theta2, omega1, omega2 = x
         delta = theta1 - theta2
 
@@ -63,14 +46,7 @@ def double_pendulum(
 
         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
+    return f
 
 
 def to_xy(L1: float, L2: float, theta1: np.ndarray, theta2: np.ndarray):

edmkit-0.0.15/src/edmkit/generate/duffing.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+__all__ = ["duffing"]
+
+
+def duffing(alpha: float, beta: float, delta: float, gamma: float, omega: float):
+    """Vector field of the forced Duffing oscillator.
+
+    Parameters
+    ----------
+    alpha : float
+        Linear stiffness (typical: -1).
+    beta : float
+        Cubic stiffness (typical: 1).
+    delta : float
+        Damping coefficient (typical: 0.3).
+    gamma : float
+        Forcing amplitude (typical: 0.5 for chaos).
+    omega : float
+        Forcing angular frequency (typical: 1.2).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(x, v)`` of shape ``(2,)``.
+    """
+
+    def f(t: float, x: np.ndarray):
+        return np.array([[0, 1], [-alpha - beta * x[0] ** 2, -delta]]) @ x + np.array([0, gamma * np.cos(omega * t)])
+
+    return f

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

@@ -0,0 +1,27 @@
+import numpy as np
+
+__all__ = ["lorenz"]
+
+
+def lorenz(sigma: float, rho: float, beta: float):
+    """Vector field of the Lorenz system.
+
+    Parameters
+    ----------
+    sigma : float
+        Prandtl number (typical: 10).
+    rho : float
+        Rayleigh number (typical: 28).
+    beta : float
+        Geometric factor (typical: 8/3).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(x, y, z)`` of shape ``(3,)``.
+    """
+
+    def f(t: float, x: np.ndarray):
+        return np.array([[-sigma, sigma, 0], [rho, -1, -x[0]], [0, x[0], -beta]]) @ x
+
+    return f

edmkit-0.0.15/src/edmkit/generate/lorenz96.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+__all__ = ["lorenz96"]
+
+
+def lorenz96(F: float):
+    """Vector field of the Lorenz 96 system.
+
+    Parameters
+    ----------
+    F : float
+        Forcing constant (typical: 8 for chaos).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state of shape ``(D,)`` with ``D >= 4``.
+    """
+
+    def f(t: float, x: np.ndarray):
+        # p[i + 2] is x[i] on the ring, so this is (x[i + 1] - x[i - 2]) * x[i - 1] - x[i] + F
+        p = np.concatenate([x[-2:], x, x[:1]])
+        return (p[3:] - p[:-3]) * p[1:-2] - x + F
+
+    return f

edmkit-0.0.15/src/edmkit/generate/rk45.py

@@ -0,0 +1,85 @@
+from collections.abc import Callable
+
+import numpy as np
+
+__all__ = ["rk45"]
+
+# Dormand-Prince 5(4)
+A = np.array(
+    [
+        [0, 0, 0, 0, 0, 0],
+        [1 / 5, 0, 0, 0, 0, 0],
+        [3 / 40, 9 / 40, 0, 0, 0, 0],
+        [44 / 45, -56 / 15, 32 / 9, 0, 0, 0],
+        [19372 / 6561, -25360 / 2187, 64448 / 6561, -212 / 729, 0, 0],
+        [9017 / 3168, -355 / 33, 46732 / 5247, 49 / 176, -5103 / 18656, 0],
+        [35 / 384, 0, 500 / 1113, 125 / 192, -2187 / 6784, 11 / 84],
+    ]
+)
+C = A.sum(axis=1)
+E = np.array([71 / 57600, 0, -71 / 16695, 71 / 1920, -17253 / 339200, 22 / 525, -1 / 40])
+
+
+def rk45(
+    f: Callable[[float, np.ndarray], np.ndarray],
+    X0: np.ndarray,
+    dt: float,
+    t_max: int,
+    rtol: float = 1e-6,
+    atol: float = 1e-9,
+):
+    """Integrate ``dx/dt = f(t, x)`` via the adaptive Dormand-Prince RK45 method.
+
+    The step size adapts between samples, while the output stays on the uniform grid of spacing ``dt``.
+
+    Parameters
+    ----------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side of the system.
+    X0 : np.ndarray
+        Initial condition of shape ``(D,)``.
+    dt : float
+        Sampling time step.
+    t_max : int
+        Maximum time.
+    rtol : float
+        Relative tolerance of the local error.
+    atol : float
+        Absolute tolerance of the local error.
+
+    Returns
+    -------
+    t : np.ndarray
+        Time array.
+    X : np.ndarray
+        Trajectory of shape ``(N, D)``.
+    """
+    t = np.arange(0, t_max, dt)
+    X = np.zeros((len(t), len(X0)))
+    X[0] = X0
+
+    K = np.zeros((7, len(X0)))
+    h = dt
+
+    for i in range(1, len(t)):
+        s, x = t[i - 1], X[i - 1]
+
+        while s < t[i]:
+            last = h >= t[i] - s
+            step = t[i] - s if last else h
+
+            for j in range(7):
+                K[j] = f(s + C[j] * step, x + step * (A[j, :j] @ K[:j]))
+
+            x_new = x + step * (A[6] @ K[:6])
+            error = np.max(np.abs(step * (E @ K)) / (atol + rtol * np.maximum(np.abs(x), np.abs(x_new))))
+            accepted = error <= 1
+
+            if accepted:
+                s, x = t[i] if last else s + step, x_new
+            if not (accepted and last):
+                h = step * min(10, max(0.2, 0.9 * error**-0.2)) if error > 0 else step * 10
+
+        X[i] = x
+
+    return t, X

edmkit-0.0.15/src/edmkit/generate/rossler.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+__all__ = ["rossler"]
+
+
+def rossler(a: float, b: float, c: float):
+    """Vector field of the Rössler system.
+
+    Parameters
+    ----------
+    a : float
+        Spiral growth rate (typical: 0.2).
+    b : float
+        Constant injection into z (typical: 0.2).
+    c : float
+        Folding threshold (typical: 5.7).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(x, y, z)`` of shape ``(3,)``.
+    """
+
+    def f(t: float, x: np.ndarray):
+        return np.array([[0, -1, -1], [1, a, 0], [0, 0, x[0] - c]]) @ x + np.array([0, 0, b])
+
+    return f

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

@@ -0,0 +1,7 @@
+from .knn import knn as knn
+from .linear_operator import linear_operator
+from .simplex_projection import simplex_projection
+from .soft import soft_simplex_projection
+from .soft_linear_operator import soft_linear_operator
+
+__all__ = ["linear_operator", "simplex_projection", "soft_linear_operator", "soft_simplex_projection"]

edmkit-0.0.15/src/edmkit/simplex_projection/linear_operator.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, overload
+
+import numpy as np
+
+from edmkit.simplex_projection.knn import knn
+from edmkit.util import pairwise_distance
+
+__all__ = ["linear_operator"]
+
+if TYPE_CHECKING:
+    from tinygrad import Tensor
+
+
+@overload
+def linear_operator(
+    X: np.ndarray,
+    Q: np.ndarray,
+    *,
+    k: int | None = None,
+    mask: np.ndarray | None = None,
+) -> np.ndarray: ...
+
+
+@overload
+def linear_operator(
+    X: Tensor,
+    Q: Tensor,
+    *,
+    k: int | None = None,
+    mask: Tensor | None = None,
+) -> Tensor: ...
+
+
+def linear_operator(
+    X,
+    Q,
+    *,
+    k=None,
+    mask=None,
+):
+    """
+    Build the linear operator `S` of simplex projection, such that ``S @ Y`` equals ``simplex_projection(X, Y, Q)`` for any target `Y`.
+
+    Parameters
+    ----------
+    X : np.ndarray or Tensor
+        The input data of shape (N,) or (N, E) or (B, N, E)
+    Q : np.ndarray or Tensor
+        The query points of shape (M,) or (M, E) or (B, M, E) for which to find the nearest neighbors in `X`.
+    k : int or None, default None
+        The number of nearest neighbors to use. If None, uses E + 1, where E is the dimension of `X`.
+    mask : np.ndarray or Tensor or None
+        Boolean mask of shape (M, N) or (B, M, N) indicating, for each query in `Q`, which library points to include when finding nearest neighbors.
+
+    Returns
+    -------
+    S : np.ndarray or Tensor
+        The operator of shape (M, N) or (B, M, N).
+
+    Raises
+    ------
+    ValueError
+        - If `k` is not positive.
+        - If the batch size or dimension of `X` and `Q` do not match.
+        - If `mask` does not have shape (M, N) or (B, M, N).
+        - If fewer than `k` library points are unmasked for some query.
+
+    Examples
+    --------
+    ```python
+    import numpy as np
+
+    from edmkit.simplex_projection import linear_operator, simplex_projection
+
+    rng = np.random.default_rng(0)
+    X = rng.normal(size=(200, 3))
+    Q = rng.normal(size=(50, 3))
+    Y = rng.normal(size=(200, 4))
+
+    assert np.allclose(linear_operator(X, Q) @ Y, simplex_projection(X, Y, Q))
+    ```
+    """
+    if isinstance(X, np.ndarray):
+        return _numpy(X, Q, k=k, mask=mask)
+
+    return _tensor(X, Q, k=k, mask=mask)
+
+
+def _numpy(
+    X: np.ndarray,
+    Q: np.ndarray,
+    *,
+    k: int | None = None,
+    mask: np.ndarray | None = None,
+) -> np.ndarray:
+    # ensure at least 2D
+    if X.ndim == 1:
+        X = X[:, None]
+    if Q.ndim == 1:
+        Q = Q[:, None]
+
+    if not (X.ndim == Q.ndim and X.ndim in (2, 3)):
+        raise ValueError(f"X and Q must both be 2D or both be 3D arrays, got X.ndim={X.ndim}, Q.ndim={Q.ndim}")
+    if Q.shape[:-2] != X.shape[:-2] or Q.shape[-1] != X.shape[-1]:
+        raise ValueError(f"batch size and dimension of X and Q must match, got X.shape={X.shape} and Q.shape={Q.shape}")
+    if mask is not None and mask.shape != Q.shape[:-1] + X.shape[-2:-1]:
+        raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
+
+    return_ndim = X.ndim
+
+    # treat the unbatched case as a single batch
+    if X.ndim == 2:
+        X, Q = X[None], Q[None]
+        if mask is not None:
+            mask = mask[None]
+
+    B, N, E = X.shape
+    M = Q.shape[1]
+    k = E + 1 if k is None else k
+    if k <= 0:
+        raise ValueError(f"k must be positive, got k={k}")
+
+    distances = np.empty((B, M, k))
+    indices = np.empty((B, M, k), dtype=np.intp)
+    for b in range(B):
+        if mask is None:
+            distances[b], indices[b] = knn(X[b], Q[b], k)
+        else:  # over-fetch, then drop the masked-out candidates
+            n_exclude = int((~mask[b]).sum(axis=-1).max())
+            if N - n_exclude < k:
+                raise ValueError(
+                    f"Not enough unmasked points in X to find {k} neighbors, got N={N} with up to {n_exclude} points masked out per query"
+                )
+            _distances, _indices = knn(X[b], Q[b], k + n_exclude)
+            _distances = np.where(np.take_along_axis(mask[b], _indices, axis=-1), _distances, np.inf)
+            top_k = np.argsort(_distances, axis=-1)[..., :k]
+            distances[b] = np.take_along_axis(_distances, top_k, axis=-1)
+            indices[b] = np.take_along_axis(_indices, top_k, axis=-1)
+
+    # clamp to avoid division by zero
+    d_min = np.fmax(distances.min(axis=-1, keepdims=True), 1e-6)  # (B, M, 1)
+    weights = np.exp(-distances / d_min)  # (B, M, k)
+    weights /= weights.sum(axis=-1, keepdims=True)
+
+    S = np.zeros((B, M, N))
+    np.put_along_axis(S, indices, weights, axis=-1)
+
+    if return_ndim == 2:
+        return S[0]  # (M, N)
+    return S  # (B, M, N)
+
+
+def _tensor(
+    X: Tensor,
+    Q: Tensor,
+    *,
+    k: int | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    # ensure at least 2D
+    if X.ndim == 1:
+        X = X[:, None]
+    if Q.ndim == 1:
+        Q = Q[:, None]
+
+    if not (X.ndim == Q.ndim and X.ndim in (2, 3)):
+        raise ValueError(f"X and Q must both be 2D or both be 3D tensors, got X.ndim={X.ndim}, Q.ndim={Q.ndim}")
+    if Q.shape[:-2] != X.shape[:-2] or Q.shape[-1] != X.shape[-1]:
+        raise ValueError(f"batch size and dimension of X and Q must match, got X.shape={X.shape} and Q.shape={Q.shape}")
+    if mask is not None and mask.shape != Q.shape[:-1] + X.shape[-2:-1]:
+        raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
+
+    return_ndim = X.ndim
+
+    # treat the unbatched case as a single batch
+    if X.ndim == 2:
+        X, Q = X[None], Q[None]
+        if mask is not None:
+            mask = mask[None]
+
+    B, N, E = (int(X.shape[0]), int(X.shape[1]), int(X.shape[2]))
+    M = int(Q.shape[1])
+    k = E + 1 if k is None else k
+    if k <= 0:
+        raise ValueError(f"k must be positive, got k={k}")
+
+    # clamp to avoid NaN gradient
+    D = pairwise_distance(Q, X).clamp(min_=1e-12).sqrt()  # (B, M, N)
+
+    if mask is not None:
+        D = mask.where(D, float("inf"))
+
+    distances, indices = D.topk(k, dim=-1, largest=False, sorted_=True)  # (B, M, k)
+
+    # neighbors are sorted, so the first is the nearest; clamp to avoid division by zero
+    d_min = distances[..., :1].clamp(min_=1e-6)  # (B, M, 1)
+    weights = (-distances / d_min).exp()  # (B, M, k)
+    weights = weights / weights.sum(axis=-1, keepdim=True)
+
+    S = weights.zeros(B, M, N, dtype=weights.dtype, device=weights.device).scatter(-1, indices, weights)  # (B, M, N)
+
+    if return_ndim == 2:
+        return S[0]  # (M, N)
+    return S  # (B, M, N)

{edmkit-0.0.13 → edmkit-0.0.15}/src/edmkit/simplex_projection/simplex_projection.py

@@ -63,12 +63,14 @@ def simplex_projection(
     -------
     predictions : np.ndarray or Tensor
         The predicted values based on the weighted mean of the nearest neighbors in `Y`.
+        Shape (M,) if `Y` is 1D, (M, E') if `Y` is 2D, or (B, M, E') if `Y` is 3D.
 
     Raises
     ------
     ValueError
         - If `k` is not positive.
         - If the input arrays `X` and `Y` do not have the same number of points.
+        - If the batch size or dimension of `X` and `Q` do not match.
         - If `mask` does not have shape (M, N) or (B, M, N).
         - If fewer than `k` library points are unmasked for some query.
 
@@ -121,6 +123,8 @@ def _numpy(
     k: int | None = None,
     mask: np.ndarray | None = None,
 ):
+    return_ndim = Y.ndim
+
     # ensure at least 2D
     if X.ndim == 1:
         X = X[:, None]
@@ -139,8 +143,7 @@ def _numpy(
         raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
 
     # treat the unbatched case as a single batch
-    batched = X.ndim == 3
-    if not batched:
+    if X.ndim == 2:
         X, Y, Q = X[None], Y[None], Q[None]
         if mask is not None:
             mask = mask[None]
@@ -177,9 +180,11 @@ def _numpy(
 
     predictions = np.matmul(weights[..., None, :], Y_neighbors).squeeze(-2) / weights.sum(axis=-1, keepdims=True)  # (B, M, E')
 
-    if not batched:
-        return predictions.squeeze()  # (M,) or (M, E')
-    return predictions
+    if return_ndim == 1:
+        return predictions[0, :, 0]  # (M,)
+    if return_ndim == 2:
+        return predictions[0]  # (M, E')
+    return predictions  # (B, M, E')
 
 
 def _tensor(
@@ -194,6 +199,8 @@ def _tensor(
     # Loading it only when needed keeps the numpy path free of that scheduler pressure.
     from tinygrad import Tensor, dtypes
 
+    return_ndim = Y.ndim
+
     # ensure at least 2D
     if X.ndim == 1:
         X = X[:, None]
@@ -212,8 +219,7 @@ def _tensor(
         raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
 
     # treat the unbatched case as a single batch
-    batched = X.ndim == 3
-    if not batched:
+    if X.ndim == 2:
         X, Y, Q = X[None], Y[None], Q[None]
         if mask is not None:
             mask = mask[None]
@@ -242,9 +248,11 @@ def _tensor(
 
     predictions: Tensor = (weights.unsqueeze(-1) * Y_neighbors).sum(axis=-2) / weights.sum(axis=-1, keepdim=True)  # (B, M, E')
 
-    if not batched:
-        return predictions.squeeze()  # (M,) or (M, E')
-    return predictions
+    if return_ndim == 1:
+        return predictions[0, :, 0]  # (M,)
+    if return_ndim == 2:
+        return predictions[0]  # (M, E')
+    return predictions  # (B, M, E')
 
 
 if TYPE_CHECKING: