edmkit 0.0.15__tar.gz → 0.0.17__tar.gz

edmkit-0.0.17/PKG-INFO

@@ -0,0 +1,77 @@
+Metadata-Version: 2.3
+Name: edmkit
+Version: 0.0.17
+Summary: Simple EDM (Empirical Dynamic Modeling) library
+Author: FUJISHIGE TEMMA
+Author-email: FUJISHIGE TEMMA <tenma.x0@gmail.com>
+Requires-Dist: kdtree-rs>=0.3.1
+Requires-Dist: numpy>=2.5.4
+Requires-Dist: scipy>=1.18.1
+Requires-Dist: tinygrad>=0.14.0
+Requires-Dist: tinysolve>=0.3.0
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+
+# edmkit
+
+This library is a collection of tools and utilities that are useful for Empirical Data Modeling (EDM) and related tasks. The library is designed to be fast and lightweight, and easy to use.
+
+[benchmarks](./benchmarks)
+
+![](./benchmarks/benchmarks.png)
+
+[compare.py](./benchmarks/compare.py)
+
+![](./benchmarks/compare.png)
+
+| simplex N | edmkit [ms] | pyEDM [ms] | speedup |
+|---:|---:|---:|---:|
+| 500 | 0.2 | 1.4 | 8x |
+| 1,000 | 0.4 | 2.3 | 5x |
+| 2,000 | 0.9 | 4.2 | 5x |
+| 5,000 | 2.3 | 10.7 | 5x |
+| 10,000 | 4.9 | 22.8 | 5x |
+| 20,000 | 10.3 | 46.8 | 5x |
+
+| smap N | edmkit [ms] | pyEDM [ms] | speedup |
+|---:|---:|---:|---:|
+| 500 | 1.9 | 38.0 | 20x |
+| 1,000 | 7.5 | 132.1 | 18x |
+| 2,000 | 21.6 | 522.6 | 24x |
+| 5,000 | 183.7 | 3,364 | 18x |
+
+| ccm N | edmkit [ms] | pyEDM [ms] | speedup | edmkit parallel [ms] | pyEDM parallel [ms] | speedup |
+|---:|---:|---:|---:|---:|---:|---:|
+| 500 | 116.7 | 148.4 | 1.3x | 22.0 | 64.5 | 2.9x |
+| 1,000 | 197.1 | 316.8 | 1.6x | 31.8 | 89.7 | 2.8x |
+| 2,000 | 382.8 | 696.8 | 1.8x | 58.6 | 133.8 | 2.3x |
+| 5,000 | 788.0 | 1,863 | 2.4x | 108.9 | 283.5 | 2.6x |
+| 10,000 | 1,470 | 4,166 | 2.8x | 197.0 | 558.9 | 2.8x |
+
+
+## Installation
+
+To install the library, you can use pip:
+
+```bash
+pip install edmkit
+```
+
+Or you can also use [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv add edmkit
+```
+
+## Testing
+
+```bash
+# Run all tests
+uv run pytest
+
+# Skip slow tests
+uv run pytest -m "not slow"
+
+# Run with CI profile (more hypothesis examples)
+HYPOTHESIS_PROFILE=ci uv run pytest
+```

edmkit-0.0.17/README.md

@@ -0,0 +1,63 @@
+# edmkit
+
+This library is a collection of tools and utilities that are useful for Empirical Data Modeling (EDM) and related tasks. The library is designed to be fast and lightweight, and easy to use.
+
+[benchmarks](./benchmarks)
+
+![](./benchmarks/benchmarks.png)
+
+[compare.py](./benchmarks/compare.py)
+
+![](./benchmarks/compare.png)
+
+| simplex N | edmkit [ms] | pyEDM [ms] | speedup |
+|---:|---:|---:|---:|
+| 500 | 0.2 | 1.4 | 8x |
+| 1,000 | 0.4 | 2.3 | 5x |
+| 2,000 | 0.9 | 4.2 | 5x |
+| 5,000 | 2.3 | 10.7 | 5x |
+| 10,000 | 4.9 | 22.8 | 5x |
+| 20,000 | 10.3 | 46.8 | 5x |
+
+| smap N | edmkit [ms] | pyEDM [ms] | speedup |
+|---:|---:|---:|---:|
+| 500 | 1.9 | 38.0 | 20x |
+| 1,000 | 7.5 | 132.1 | 18x |
+| 2,000 | 21.6 | 522.6 | 24x |
+| 5,000 | 183.7 | 3,364 | 18x |
+
+| ccm N | edmkit [ms] | pyEDM [ms] | speedup | edmkit parallel [ms] | pyEDM parallel [ms] | speedup |
+|---:|---:|---:|---:|---:|---:|---:|
+| 500 | 116.7 | 148.4 | 1.3x | 22.0 | 64.5 | 2.9x |
+| 1,000 | 197.1 | 316.8 | 1.6x | 31.8 | 89.7 | 2.8x |
+| 2,000 | 382.8 | 696.8 | 1.8x | 58.6 | 133.8 | 2.3x |
+| 5,000 | 788.0 | 1,863 | 2.4x | 108.9 | 283.5 | 2.6x |
+| 10,000 | 1,470 | 4,166 | 2.8x | 197.0 | 558.9 | 2.8x |
+
+
+## Installation
+
+To install the library, you can use pip:
+
+```bash
+pip install edmkit
+```
+
+Or you can also use [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv add edmkit
+```
+
+## Testing
+
+```bash
+# Run all tests
+uv run pytest
+
+# Skip slow tests
+uv run pytest -m "not slow"
+
+# Run with CI profile (more hypothesis examples)
+HYPOTHESIS_PROFILE=ci uv run pytest
+```

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

@@ -1,15 +1,15 @@
 [project]
 name = "edmkit"
-version = "0.0.15"
+version = "0.0.17"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 readme = "README.md"
 requires-python = ">= 3.13"
 dependencies = [
     "kdtree-rs>=0.3.1",
-    "numpy>=2.5.1",
-    "scipy>=1.17.1",
+    "numpy>=2.5.4",
+    "scipy>=1.18.1",
     "tinygrad>=0.14.0",
-    "tinysolve>=0.2.1",
+    "tinysolve>=0.3.0",
 ]
 
 [[project.authors]]
@@ -17,13 +17,19 @@ name = "FUJISHIGE TEMMA"
 email = "tenma.x0@gmail.com"
 
 [dependency-groups]
+benchmark = [
+    "pandas>=3.1.0",
+    "pyedm>=2.5.7",
+]
 dev = [
-    "griffe2md>=1.5.0",
-    "hypothesis>=6.151.9",
+    "griffe2md>=1.5.1",
+    "hypothesis>=6.170.14",
     "matplotlib>=3.11.2",
-    "pytest>=9.1.1",
-    "ruff>=0.15.20",
-    "ty>=0.0.56",
+    "pytest>=9.2.0",
+    "pytest-benchmark>=5.3.1",
+    "pywavelets>=1.10.0",
+    "ruff>=0.17.3",
+    "ty>=0.0.95",
 ]
 
 [tool.uv]
@@ -35,12 +41,19 @@ namespace = true
 
 [tool.pytest.ini_options]
 addopts = "-v --tb=short"
+testpaths = ["tests"]
+python_files = [
+    "test_*.py",
+    "bench_*.py",
+]
 filterwarnings = ["error"]
 markers = [
     """slow: marks tests as slow (deselect with '-m "not slow"')""",
     """gpu: marks tests requiring tinygrad tensor backend (deselect with '-m "not gpu"')""",
+    "dense: marks benchmarks that build an (M, N) matrix",
+    """huge: marks benchmarks whose N is too large for an (M, N) matrix in memory (deselect with '-m "not (dense and huge)"')""",
 ]
 
 [build-system]
-requires = ["uv_build>=0.10.9,<0.11.0"]
+requires = ["uv_build>=0.13.6,<0.14.0"]
 build-backend = "uv_build"

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

@@ -1,26 +1,32 @@
 [project]
 name = "edmkit"
-version = "0.0.15"
+version = "0.0.17"
 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.3.1",
-    "numpy>=2.5.1",
-    "scipy>=1.17.1",
+    "numpy>=2.5.4",
+    "scipy>=1.18.1",
     "tinygrad>=0.14.0",
-    "tinysolve>=0.2.1",
+    "tinysolve>=0.3.0",
 ]
 
 [dependency-groups]
+benchmark = [
+    "pandas>=3.1.0",
+    "pyedm>=2.5.7",
+]
 dev = [
-    "griffe2md>=1.5.0",
-    "hypothesis>=6.151.9",
+    "griffe2md>=1.5.1",
+    "hypothesis>=6.170.14",
     "matplotlib>=3.11.2",
-    "pytest>=9.1.1",
-    "ruff>=0.15.20",
-    "ty>=0.0.56",
+    "pytest>=9.2.0",
+    "pytest-benchmark>=5.3.1",
+    "pywavelets>=1.10.0",
+    "ruff>=0.17.3",
+    "ty>=0.0.95",
 ]
 
 [tool.uv]
@@ -30,15 +36,19 @@ required-environments = [
 
 [tool.pytest.ini_options]
 addopts = "-v --tb=short"
+testpaths = ["tests"]
+python_files = ["test_*.py", "bench_*.py"]
 filterwarnings = ["error"]
 markers = [
     "slow: marks tests as slow (deselect with '-m \"not slow\"')",
     "gpu: marks tests requiring tinygrad tensor backend (deselect with '-m \"not gpu\"')",
+    "dense: marks benchmarks that build an (M, N) matrix",
+    "huge: marks benchmarks whose N is too large for an (M, N) matrix in memory (deselect with '-m \"not (dense and huge)\"')",
 ]
 
 
 [build-system]
-requires = ["uv_build>=0.10.9,<0.11.0"]
+requires = ["uv_build>=0.13.6,<0.14.0"]
 build-backend = "uv_build"
 
 [tool.uv.build-backend]

edmkit-0.0.17/src/edmkit/ccm.py

@@ -0,0 +1,168 @@
+from collections.abc import Callable, Iterator
+from functools import partial
+
+import numpy as np
+
+from edmkit.types import Mask, PredictFunc
+
+__all__ = ["bootstrap"]
+
+
+def bootstrap(
+    X: np.ndarray,
+    Y: np.ndarray,
+    Q: np.ndarray,
+    lib_sizes: np.ndarray,
+    *,
+    predict_func: PredictFunc,
+    n_samples: int = 20,
+    replace: bool = True,
+    seed: int | None = 42,
+    mask: Mask | None = None,
+    batch_size: int | None = 20,
+) -> Iterator[tuple[tuple[int, slice], Callable[[], np.ndarray]]]:
+    """
+    Iterate `predict_func(X[library], Y[library], Q)` for `n_samples` times for each library size in `lib_sizes`.
+
+    By gathering the results of the returned tasks, you can get `(len(lib_sizes), n_samples, M, E')` array of predictions.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        The input data of shape (N,) or (N, E)
+    Y : np.ndarray
+        The target data of shape (N,) or (N, E'), sampled together with `X`.
+    Q : np.ndarray
+        The query points of shape (M,) or (M, E), shared by all sampled libraries as a read-only view.
+    lib_sizes : np.ndarray
+        The 1D array of library sizes, yielded in the given order. Sizes may exceed N only if `replace` is True.
+    predict_func : PredictFunc
+        The batched prediction function, such as `simplex_projection` or ``functools.partial(smap, theta=...)``.
+    n_samples : int, default 20
+        The number of libraries to sample for each size in `lib_sizes`.
+    replace : bool, default True
+        Whether to sample the library points with replacement.
+    seed : int or None, default 42
+        The seed of the random generator. If None, uses fresh entropy.
+    mask : Mask or None
+        ``mask(i, j)`` indicating, for each query in `Q`, which points of `X` to include. Use `theiler_window` for temporal exclusion.
+    batch_size : int or None, default 20
+        The maximum number of libraries per task. If None, uses `n_samples`.
+
+    Yields
+    ------
+    index : tuple[int, slice]
+        The position of the task's libraries in an array of shape (len(`lib_sizes`), `n_samples`), such as the scores.
+    samples : Callable[[], np.ndarray]
+        The task that returns the predictions on B sampled libraries, of shape (B, M, E'), where B <= `batch_size` and E' = 1 if `Y` is 1D.
+
+    Raises
+    ------
+    ValueError
+        - If `X`, `Y`, or `Q` is empty or not 1D or 2D.
+        - If the input arrays `X` and `Y` do not have the same number of points.
+        - If the dimension of `X` and `Q` do not match.
+        - If `lib_sizes` is not a nonempty 1D array of positive integers.
+        - If some size in `lib_sizes` exceeds N when `replace` is False.
+        - If `n_samples` or `batch_size` is not positive.
+        - If `predict_func` raises when a task is executed.
+
+    Examples
+    --------
+    ```python
+    from concurrent.futures import ThreadPoolExecutor
+
+    import numpy as np
+
+    from edmkit.ccm import bootstrap
+    from edmkit.embedding import embed
+    from edmkit.metrics import pearson_correlation
+    from edmkit.simplex_projection import simplex_projection
+
+    # Generate a simple time series (logistic map)
+    N = 300
+    x = np.zeros(N)
+    x[0] = 0.4
+    for i in range(1, N):
+        x[i] = 3.9 * x[i - 1] * (1 - x[i - 1])
+
+    tau = 2
+    E = 3
+
+    coordinates = np.array([[0, tau * j] for j in range(E)])
+    embedding, _ = embed(coordinates, x)
+    shift = tau * (E - 1)
+
+    pool_size = 200
+    Tp = 1
+    X = embedding[:pool_size - shift]
+    Y = x[shift + Tp : pool_size + Tp]
+    Q = embedding[pool_size - shift : -Tp]
+    actual = x[pool_size + Tp :, None]
+
+    lib_sizes = np.array([10, 50, 190])
+    n_samples = 100
+
+
+    def score(samples):
+        predictions = samples()  # (B, M, 1)
+        return pearson_correlation(predictions, np.broadcast_to(actual, predictions.shape))
+
+
+    scores = np.empty((len(lib_sizes), n_samples))
+    tasks = bootstrap(X, Y, Q, lib_sizes, predict_func=simplex_projection, n_samples=n_samples)
+    indices, samples = zip(*tasks)
+    with ThreadPoolExecutor() as executor:
+        for index, result in zip(indices, executor.map(score, samples)):
+            scores[index] = result
+
+    for lib_size, correlation in zip(lib_sizes, scores.mean(axis=1)):
+        print(f"Correlation (L={lib_size}): {correlation:.3f}")
+    ```
+    """
+    if X.ndim == 1:
+        X = X[:, None]
+    if Y.ndim == 1:
+        Y = Y[:, None]
+    if Q.ndim == 1:
+        Q = Q[:, None]
+
+    if X.ndim != 2 or Y.ndim != 2 or Q.ndim != 2:
+        raise ValueError(f"X, Y, and Q must be 1D or 2D arrays, got X.ndim={X.ndim}, Y.ndim={Y.ndim}, Q.ndim={Q.ndim}")
+    if X.size == 0 or Y.size == 0 or Q.size == 0:
+        raise ValueError("X, Y, and Q must be nonempty")
+    if X.shape[0] != Y.shape[0]:
+        raise ValueError(f"X and Y must have the same length, got X.shape={X.shape}, Y.shape={Y.shape}")
+    if X.shape[1] != Q.shape[1]:
+        raise ValueError(f"X and Q must have the same dimension, got X.shape={X.shape}, Q.shape={Q.shape}")
+    if lib_sizes.ndim != 1 or lib_sizes.size == 0 or not np.issubdtype(lib_sizes.dtype, np.integer) or np.any(lib_sizes <= 0):
+        raise ValueError("lib_sizes must be a nonempty 1D integer array of positive sizes")
+    if not replace and np.any(lib_sizes > len(X)):
+        raise ValueError(f"lib_sizes cannot exceed N={len(X)} when replace is False")
+    batch_size = n_samples if batch_size is None else batch_size
+    if n_samples <= 0 or batch_size <= 0:
+        raise ValueError(f"n_samples and batch_size must be positive, got n_samples={n_samples}, batch_size={batch_size}")
+
+    rng = np.random.default_rng(seed)
+    for i, lib_size in enumerate(lib_sizes):
+        for start in range(0, n_samples, batch_size):
+            batch = min(batch_size, n_samples - start)
+            if replace:
+                indices = rng.choice(len(X), size=(batch, lib_size))
+            else:
+                indices = np.stack([rng.choice(len(X), size=lib_size, replace=False) for _ in range(batch)])
+            yield (i, slice(start, start + batch)), partial(samples, X, Y, Q, indices, predict_func=predict_func, mask=mask)
+
+
+def samples(
+    X: np.ndarray,
+    Y: np.ndarray,
+    Q: np.ndarray,
+    indices: np.ndarray,
+    *,
+    predict_func: PredictFunc,
+    mask: Mask | None,
+) -> np.ndarray:
+    queries = np.broadcast_to(Q, (len(indices), *Q.shape))
+    # position j of library b is the point indices[b, j] of X
+    return predict_func(X[indices], Y[indices], queries, mask=None if mask is None else (lambda b, i, j: mask(i, indices[b, j])))

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

@@ -0,0 +1,40 @@
+from .chua import chua
+from .ctrnn import ctrnn
+from .double_pendulum import double_pendulum, to_xy
+from .duffing import duffing
+from .euler_spiral import euler_spiral
+from .fitzhugh_nagumo import fitzhugh_nagumo
+from .hindmarsh_rose import hindmarsh_rose
+from .kuramoto import kuramoto
+from .lorenz import lorenz
+from .lorenz96 import lorenz96
+from .lotka_volterra import lotka_volterra
+from .mackey_glass import mackey_glass
+from .morris_lecar import morris_lecar
+from .pyloric import pyloric
+from .rk45 import rk45
+from .rossler import rossler
+from .rulkov_map import rulkov_map
+from .three_scroll import three_scroll
+
+__all__ = [
+    "chua",
+    "ctrnn",
+    "double_pendulum",
+    "duffing",
+    "euler_spiral",
+    "fitzhugh_nagumo",
+    "hindmarsh_rose",
+    "kuramoto",
+    "lorenz",
+    "lorenz96",
+    "lotka_volterra",
+    "mackey_glass",
+    "morris_lecar",
+    "pyloric",
+    "rk45",
+    "rossler",
+    "rulkov_map",
+    "three_scroll",
+    "to_xy",
+]

edmkit-0.0.17/src/edmkit/generate/chua.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+__all__ = ["chua"]
+
+
+def chua(alpha: float, beta: float, m0: float, m1: float):
+    """Vector field of Chua's circuit, with the piecewise-linear diode ``h(x) = m1 x + (m0 - m1) (|x + 1| - |x - 1|) / 2``.
+
+    Parameters
+    ----------
+    alpha : float
+        Ratio of capacitances (typical: 15.6).
+    beta : float
+        Ratio of inductance to resistance (typical: 28).
+    m0 : float
+        Diode slope for ``|x| < 1`` (typical: -8/7).
+    m1 : float
+        Diode slope for ``|x| > 1`` (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,)``, or ``(B, 3)`` for a batch.
+    """
+
+    def f(t: float, s: np.ndarray):
+        x, y, z = np.moveaxis(s, -1, 0)
+        h = m1 * x + (m0 - m1) * np.clip(x, -1, 1)  # (|x + 1| - |x - 1|) / 2 is clip(x, -1, 1)
+        return np.stack([alpha * (y - x - h), x - y + z, -beta * y], axis=-1)
+
+    return f

{edmkit-0.0.15 → edmkit-0.0.17}/src/edmkit/generate/ctrnn.py

@@ -16,10 +16,10 @@ def ctrnn(g: float, J: np.ndarray):
     Returns
     -------
     f : Callable[[float, np.ndarray], np.ndarray]
-        Right-hand side ``f(t, x)`` for a state of shape ``(N,)``.
+        Right-hand side ``f(t, x)`` for a state of shape ``(N,)``, or ``(B, N)`` for a batch.
     """
 
     def f(t: float, x: np.ndarray):
-        return -x + g * J @ np.tanh(x)
+        return -x + g * np.tanh(x) @ J.T
 
     return f

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

@@ -22,11 +22,11 @@ def double_pendulum(m1: float, m2: float, L1: float, L2: float, g: float):
     Returns
     -------
     f : Callable[[float, np.ndarray], np.ndarray]
-        Right-hand side ``f(t, x)`` for a state ``(theta1, theta2, omega1, omega2)`` of shape ``(4,)``.
+        Right-hand side ``f(t, x)`` for a state ``(theta1, theta2, omega1, omega2)`` of shape ``(4,)``, or ``(B, 4)`` for a batch.
     """
 
     def f(t: float, x: np.ndarray):
-        theta1, theta2, omega1, omega2 = x
+        theta1, theta2, omega1, omega2 = np.moveaxis(x, -1, 0)
         delta = theta1 - theta2
 
         denom = 2 * m1 + m2 - m2 * np.cos(2 * delta)
@@ -44,7 +44,7 @@ def double_pendulum(m1: float, m2: float, L1: float, L2: float, g: float):
             L2 * denom
         )
 
-        return np.array([dtheta1_dt, dtheta2_dt, domega1_dt, domega2_dt])
+        return np.stack([dtheta1_dt, dtheta2_dt, domega1_dt, domega2_dt], axis=-1)
 
     return f
 

{edmkit-0.0.15 → edmkit-0.0.17}/src/edmkit/generate/duffing.py

@@ -22,10 +22,11 @@ def duffing(alpha: float, beta: float, delta: float, gamma: float, omega: float)
     Returns
     -------
     f : Callable[[float, np.ndarray], np.ndarray]
-        Right-hand side ``f(t, x)`` for a state ``(x, v)`` of shape ``(2,)``.
+        Right-hand side ``f(t, x)`` for a state ``(x, v)`` of shape ``(2,)``, or ``(B, 2)`` for a batch.
     """
 
-    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)])
+    def f(t: float, s: np.ndarray):
+        x, v = np.moveaxis(s, -1, 0)
+        return np.stack([v, (-alpha - beta * x**2) * x - delta * v + gamma * np.cos(omega * t)], axis=-1)
 
     return f

edmkit-0.0.17/src/edmkit/generate/euler_spiral.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+__all__ = ["euler_spiral"]
+
+
+def euler_spiral(a: float):
+    """Vector field of the Euler spiral (clothoid).
+
+    Parameters
+    ----------
+    a : float
+        Rate at which the curvature grows (typical: 1).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(x, y, theta)`` of shape ``(3,)``, or ``(B, 3)`` for a batch.
+    """
+
+    def f(t: float, x: np.ndarray):
+        theta = x[..., 2]
+        return np.stack([np.cos(theta), np.sin(theta), np.full_like(theta, a * t)], axis=-1)
+
+    return f

edmkit-0.0.17/src/edmkit/generate/fitzhugh_nagumo.py

@@ -0,0 +1,30 @@
+import numpy as np
+
+__all__ = ["fitzhugh_nagumo"]
+
+
+def fitzhugh_nagumo(a: float, b: float, epsilon: float, I: float):
+    """Vector field of the FitzHugh-Nagumo neuron model.
+
+    Parameters
+    ----------
+    a : float
+        Offset of the recovery variable (typical: 0.7).
+    b : float
+        Damping of the recovery variable (typical: 0.8).
+    epsilon : float
+        Time scale of the recovery variable (typical: 0.08).
+    I : float
+        Applied current (typical: 0.5 for spiking).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(v, w)`` of shape ``(2,)``, or ``(B, 2)`` for a batch.
+    """
+
+    def f(t: float, x: np.ndarray):
+        v, w = np.moveaxis(x, -1, 0)
+        return np.stack([v - v**3 / 3 - w + I, epsilon * (v + a - b * w)], axis=-1)
+
+    return f

edmkit-0.0.17/src/edmkit/generate/hindmarsh_rose.py

@@ -0,0 +1,38 @@
+import numpy as np
+
+__all__ = ["hindmarsh_rose"]
+
+
+def hindmarsh_rose(a: float, b: float, c: float, d: float, r: float, s: float, x_R: float, I: float):
+    """Vector field of the Hindmarsh-Rose neuron model.
+
+    Parameters
+    ----------
+    a : float
+        Cubic coefficient of the membrane potential (typical: 1).
+    b : float
+        Quadratic coefficient of the membrane potential (typical: 3).
+    c : float
+        Constant drive of the recovery variable (typical: 1).
+    d : float
+        Quadratic coupling of the potential into the recovery variable (typical: 5).
+    r : float
+        Time scale of the slow adaptation current (typical: 0.006).
+    s : float
+        Adaptation strength (typical: 4).
+    x_R : float
+        Resting potential (typical: -1.6).
+    I : float
+        Applied current (typical: 3.25 for chaotic bursting).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(x, y, z)`` of shape ``(3,)``, or ``(B, 3)`` for a batch: membrane potential, recovery variable and adaptation current.
+    """
+
+    def f(t: float, x: np.ndarray):
+        v, w, u = np.moveaxis(x, -1, 0)
+        return np.stack([w - a * v**3 + b * v**2 - u + I, c - d * v**2 - w, r * (s * (v - x_R) - u)], axis=-1)
+
+    return f

edmkit-0.0.17/src/edmkit/generate/kuramoto.py

@@ -0,0 +1,26 @@
+import numpy as np
+
+__all__ = ["kuramoto"]
+
+
+def kuramoto(omega: np.ndarray, K: np.ndarray):
+    """Vector field of the Kuramoto model ``dθ_i/dt = ω_i + Σ_j K_ij sin(θ_j - θ_i)``.
+
+    Parameters
+    ----------
+    omega : np.ndarray
+        Natural frequencies of shape ``(N,)``.
+    K : np.ndarray
+        Coupling of shape ``(N, N)``; ``K[i, j]`` is the strength with which oscillator ``j`` drives oscillator ``i`` (typical: ``k / N`` everywhere).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for unwrapped phases of shape ``(N,)``, or ``(B, N)`` for a batch.
+    """
+
+    def f(t: float, x: np.ndarray):
+        s, c = np.sin(x), np.cos(x)
+        return omega + c * (s @ K.T) - s * (c @ K.T)
+
+    return f