edmkit 0.0.12__tar.gz → 0.0.13__tar.gz

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

@@ -1,14 +1,14 @@
 Metadata-Version: 2.3
 Name: edmkit
-Version: 0.0.12
+Version: 0.0.13
 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.1
 Requires-Dist: scipy>=1.17.1
-Requires-Dist: tinygrad>=0.13.0
-Requires-Dist: tinysolve>=0.2.0
+Requires-Dist: tinygrad>=0.14.0
+Requires-Dist: tinysolve>=0.2.1
 Requires-Python: >=3.13
 Description-Content-Type: text/markdown
 

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

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.12"
+version = "0.0.13"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 readme = "README.md"
 requires-python = ">= 3.13"
@@ -8,8 +8,8 @@ dependencies = [
     "kdtree-rs>=0.3.1",
     "numpy>=2.5.1",
     "scipy>=1.17.1",
-    "tinygrad>=0.13.0",
-    "tinysolve>=0.2.0",
+    "tinygrad>=0.14.0",
+    "tinysolve>=0.2.1",
 ]
 
 [[project.authors]]

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

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.12"
+version = "0.0.13"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
 readme = "README.md"
@@ -9,8 +9,8 @@ dependencies = [
     "kdtree-rs>=0.3.1",
     "numpy>=2.5.1",
     "scipy>=1.17.1",
-    "tinygrad>=0.13.0",
-    "tinysolve>=0.2.0",
+    "tinygrad>=0.14.0",
+    "tinysolve>=0.2.1",
 ]
 
 [dependency-groups]

edmkit-0.0.13/src/edmkit/ccm.py

@@ -0,0 +1,274 @@
+from collections.abc import Callable
+
+import numpy as np
+
+from edmkit.metrics import MetricFunc
+from edmkit.types import PredictFunc
+
+__all__ = [
+    "AggregateFunc",
+    "SampleFunc",
+    "bootstrap",
+    "ccm",
+    "make_sample_func",
+]
+
+type SampleFunc = Callable[[np.ndarray, int], np.ndarray]
+"""SampleFunc is a function that takes (pool, size) and returns a sampled array."""
+type AggregateFunc = Callable[[np.ndarray], float]
+"""AggregateFunc is a function that takes an array of values and returns a single value."""
+
+
+def make_sample_func(seed: int | None = 42) -> SampleFunc:
+    """Create a sample function with its own independent RNG."""
+    rng = np.random.default_rng(seed)
+
+    def sample_func(pool: np.ndarray, size: int) -> np.ndarray:
+        return rng.choice(pool, size=size, replace=True)
+
+    return sample_func
+
+
+def bootstrap(
+    X: np.ndarray,
+    Y: np.ndarray,
+    lib_sizes: np.ndarray,
+    predict_func: PredictFunc,
+    metric_func: MetricFunc,
+    n_samples: int = 20,
+    *,
+    library_pool: np.ndarray,
+    prediction_pool: np.ndarray,
+    sample_func: SampleFunc | None = None,
+    batch_size: int | None = 20,
+) -> np.ndarray:
+    """
+    Perform Convergent Cross Mapping and return per-sample scores.
+
+    Same as :func:`ccm` but returns the raw per-sample scores instead of
+    aggregating them.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Library delay vectors of shape ``(T,)`` or ``(T, E)`` (potential response).
+    Y : np.ndarray
+        Target values of shape ``(T,)`` or ``(T, E')`` (potential driver).
+        A 2D `Y` cross-maps to a delay vector; the metric scores all columns jointly.
+    lib_sizes : np.ndarray
+        Array of library sizes to test convergence.
+    predict_func : PredictFunc
+        Prediction function with signature (X, Y, Q) -> predictions.
+    metric_func : MetricFunc
+        Metric function with signature (predictions, observations) -> metric value.
+    n_samples : int, default 20
+        Number of random samples per library size for bootstrapping.
+    library_pool : np.ndarray
+        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 : 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 or None, default 20
+        If specified, predictions are made in batches to limit memory usage.
+
+    Returns
+    -------
+    samples : np.ndarray
+        Per-sample skill scores of shape ``(n_samples, len(lib_sizes))``.
+    """
+    if sample_func is None:
+        sample_func = make_sample_func()
+
+    if X.shape[0] != Y.shape[0]:
+        raise ValueError(f"X and Y must have same length, got {X.shape[0]} and {Y.shape[0]}")
+    if predict_func is None or not callable(predict_func):
+        raise ValueError(f"predict_func must be callable, got {type(predict_func)}")
+    if n_samples <= 0:
+        raise ValueError(f"n_samples must be positive, got {n_samples}")
+
+    if batch_size is None:
+        batch_size = n_samples
+    else:
+        batch_size = min(batch_size, n_samples)
+
+    if X.ndim == 1:
+        X = X[:, None]
+    if Y.ndim == 1:
+        Y = Y[:, None]
+
+    prediction_indices = np.tile(prediction_pool, (batch_size, 1))
+    Q = X[prediction_indices]
+    actual = Y[prediction_indices]
+
+    samples = np.zeros((n_samples, len(lib_sizes)))
+
+    for i, lib_size in enumerate(lib_sizes):
+        remaining = n_samples
+        while remaining > 0:
+            batch = min(batch_size, remaining)
+
+            library_indices = np.vstack([sample_func(library_pool, lib_size) for _ in range(batch)])
+
+            lib_X = X[library_indices]
+            lib_Y = Y[library_indices]
+
+            predictions = predict_func(lib_X, lib_Y, Q[:batch])
+
+            offset = n_samples - remaining
+            samples[offset : offset + batch, i] = metric_func(predictions, actual[:batch])
+            remaining -= batch
+
+    return samples
+
+
+def ccm(
+    X: np.ndarray,
+    Y: np.ndarray,
+    lib_sizes: np.ndarray,
+    predict_func: PredictFunc,
+    metric_func: MetricFunc,
+    n_samples: int = 20,
+    *,
+    library_pool: np.ndarray,
+    prediction_pool: np.ndarray,
+    sample_func: SampleFunc | None = None,
+    aggregate_func: AggregateFunc = np.mean,
+    batch_size: int | None = 20,
+) -> np.ndarray:
+    """
+    Perform Convergent Cross Mapping using a custom prediction function.
+
+    CCM tests for causality from X to Y by using the attractor reconstructed from Y
+    to predict values of X. If X causes Y, then Y's attractor contains information
+    about X, allowing cross-mapping from Y to X.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Library delay vectors of shape ``(T,)`` or ``(T, E)`` (potential response).
+    Y : np.ndarray
+        Target values of shape ``(T,)`` or ``(T, E')`` (potential driver).
+        A 2D `Y` cross-maps to a delay vector; the metric scores all columns jointly.
+    lib_sizes : np.ndarray
+        Array of library sizes to test convergence.
+    predict_func : PredictFunc
+        Prediction function with signature (X, Y, Q) -> predictions.
+        Can be `simplex_projection`, `smap` with partial application, or a custom function.
+    metric_func : MetricFunc
+        Metric function with signature (predictions, observations) -> metric value.
+    n_samples : int, default 100
+        Number of random samples per library size for bootstrapping.
+    library_pool : np.ndarray
+        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 : 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.
+        When None, a fresh RNG-backed sampler is created per call.
+    aggregate_func : AggregateFunc, default np.mean
+        Reducer applied to the skill samples for each library size.
+    batch_size : int or None, default 20
+        If not specified, batch_size == n_samples.
+        If specified, predictions are made in batches to limit memory usage.
+    Returns
+    -------
+    correlations : np.ndarray
+        Mean correlation coefficient for each library size.
+
+    Raises
+    ------
+    ValueError
+        - If `X` and `Y` have different lengths.
+        - If `lib_sizes` contains non-positive values.
+        - If `predict_func` is not callable.
+        - If `n_samples` is not positive.
+        - If `aggregate_func` is not callable.
+        - If `library_pool` or `prediction_pool` is invalid.
+
+    Notes
+    -----
+    - Higher correlation at larger library sizes indicates convergence and suggests X influences Y (X -> Y causality)
+    - Convergence is the key signature of causality in CCM
+    - The method uses Y's attractor to predict X (cross-mapping)
+
+    Examples
+    --------
+    ```python
+    from functools import partial
+
+    import numpy as np
+
+    from edmkit.ccm import ccm
+    from edmkit.embedding import embed
+    from edmkit.metrics import pearson_correlation
+    from edmkit.simplex_projection import simplex_projection
+    from edmkit.smap import smap
+
+    # Generate coupled logistic maps (X drives Y)
+    N = 1000
+    rx, ry, Bxy = 3.8, 3.5, 0.02
+    X = np.zeros(N)
+    Y = np.zeros(N)
+    X[0], Y[0] = 0.4, 0.2
+    for i in range(1, N):
+        X[i] = X[i - 1] * (rx - rx * X[i - 1])
+        Y[i] = Y[i - 1] * (ry - ry * Y[i - 1]) + Bxy * X[i - 1]
+
+    tau = 1
+    E = 2
+
+    # To test X -> Y causality, cross-map from Y's attractor to X
+    coordinates = np.array([[0, tau * j] for j in range(E)])
+    Y_embedding, _ = embed(coordinates, Y)
+    shift = tau * (E - 1)
+    X_aligned = X[shift:]
+
+    library_pool = np.arange(Y_embedding.shape[0] // 2)
+    prediction_pool = np.arange(Y_embedding.shape[0] // 2, Y_embedding.shape[0])
+
+    # logarithmic within range 10 to max library size
+    lib_sizes = np.logspace(np.log10(10), np.log10(library_pool[-1]), num=5, dtype=int)
+
+    # Using simplex projection
+    correlations = ccm(
+        Y_embedding,
+        X_aligned,
+        lib_sizes=lib_sizes,
+        predict_func=simplex_projection,
+        metric_func=pearson_correlation,
+        library_pool=library_pool,
+        prediction_pool=prediction_pool,
+    )
+
+    # Using S-Map with partial application
+    correlations = ccm(
+        Y_embedding,
+        X_aligned,
+        lib_sizes=lib_sizes,
+        predict_func=partial(smap, theta=2.0, alpha=1e-10),
+        metric_func=pearson_correlation,
+        library_pool=library_pool,
+        prediction_pool=prediction_pool,
+    )
+    ```
+    """
+    if aggregate_func is None or not callable(aggregate_func):
+        raise ValueError("aggregate_func must be a callable")
+
+    samples = bootstrap(
+        X,
+        Y,
+        lib_sizes,
+        predict_func,
+        metric_func,
+        n_samples,
+        library_pool=library_pool,
+        prediction_pool=prediction_pool,
+        sample_func=sample_func,
+        batch_size=batch_size,
+    )
+
+    return np.array([aggregate_func(samples[:, i]) for i in range(samples.shape[1])])

edmkit-0.0.13/src/edmkit/embedding.py

@@ -0,0 +1,85 @@
+import numpy as np
+
+__all__ = ["embed"]
+
+
+def embed(coordinates: np.ndarray, x: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Delay vectors for an arbitrary choice of variables and lags.
+
+    Parameters
+    ----------
+    coordinates : np.ndarray
+        Integer array of shape ``(E, 2)``. Row ``j`` is ``(variable, lag)``, selecting
+        observation ``x[t - lag, variable]`` for the delay vector at time ``t``.
+        A negative lag selects an observation taken after time ``t``.
+    x : np.ndarray
+        Observations of shape ``(T, d)``, or ``(T,)``.
+
+    Returns
+    -------
+    X : np.ndarray
+        Delay vectors of shape ``(N, E)``, one row per time in `t`.
+    t : np.ndarray
+        Times of shape ``(N,)`` at which every coordinate is available, that is
+        ``lags.max() <= t <= T - 1 + lags.min()``. A time outside ``[0, T - 1]`` is
+        included when every coordinate still reads an observation inside `x`.
+
+    Raises
+    ------
+    ValueError
+        - If `coordinates` is not an integer array of shape ``(E, 2)`` with ``E >= 1``.
+        - If `x` is not a 1D or 2D array.
+        - If a variable index lies outside `x`.
+        - If no time has every coordinate available.
+
+    Examples
+    --------
+    ```python
+    import numpy as np
+
+    from edmkit.embedding import embed
+
+    x = np.arange(20).reshape(10, 2)  # ten times, two variables
+
+    # Classical univariate delay coordinates.
+    tau, E = 2, 3
+    coordinates = np.array([[0, tau * j] for j in range(E)])
+    X, t = embed(coordinates, x)
+    print(X)
+    print(t)
+    # [[ 8  4  0]
+    #  [10  6  2]
+    #  [12  8  4]
+    #  [14 10  6]
+    #  [16 12  8]
+    #  [18 14 10]]
+    # [4 5 6 7 8 9]
+
+    # One step ahead of the same coordinates, as iterated forecasting needs.
+    Y, t_ahead = embed(coordinates - [0, 1], x)
+    ```
+    """
+    coordinates = np.asarray(coordinates)
+    if coordinates.ndim != 2 or coordinates.shape[-1] != 2:
+        raise ValueError(f"coordinates must have shape (E, 2), got coordinates.shape={coordinates.shape}")
+    if coordinates.shape[0] < 1:
+        raise ValueError("coordinates must select at least one delay coordinate")
+    if not np.issubdtype(coordinates.dtype, np.integer):
+        raise ValueError(f"coordinates must be an integer array, got coordinates.dtype={coordinates.dtype}")
+    if x.ndim == 1:
+        x = x[:, None]
+    if x.ndim != 2:
+        raise ValueError(f"x must have shape (T,) or (T, d), got x.shape={x.shape}")
+
+    variables, lags = coordinates[:, 0], coordinates[:, 1]
+    if variables.min() < 0 or variables.max() >= x.shape[1]:
+        raise ValueError(f"variable indices must lie in [0, {x.shape[1] - 1}], got {variables.min()} to {variables.max()}")
+
+    # A delay vector at time t needs 0 <= t - lag <= T - 1 for every coordinate.
+    first = int(lags.max())
+    last = x.shape[0] - 1 + int(lags.min())
+    if last < first:
+        raise ValueError(f"no time has every coordinate available: lags span {int(lags.min())} to {int(lags.max())} for T={x.shape[0]}")
+
+    t = np.arange(first, last + 1)
+    return x[t[:, None] - lags[None, :], variables[None, :]], t

{edmkit-0.0.12 → edmkit-0.0.13}/src/edmkit/generate/__init__.py

@@ -1,6 +1,5 @@
-# ruff: noqa: F401
 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"]
+__all__ = ["double_pendulum", "lorenz", "mackey_glass", "to_xy"]

edmkit-0.0.13/src/edmkit/metrics/__init__.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol, overload
+
+import numpy as np
+
+from edmkit.metrics.mae import mae
+from edmkit.metrics.pearson_correlation import pearson_correlation
+from edmkit.metrics.rmse import rmse
+
+__all__ = ["MetricFunc", "mae", "pearson_correlation", "rmse"]
+
+if TYPE_CHECKING:
+    from tinygrad import Tensor
+
+
+class MetricFunc(Protocol):
+    """Metric function protocol.
+
+    Accepts predictions and observations of the same shape and returns a metric value.
+    """
+
+    @overload
+    def __call__(self, predictions: np.ndarray, observations: np.ndarray) -> np.ndarray: ...
+
+    @overload
+    def __call__(self, predictions: Tensor, observations: Tensor) -> Tensor: ...
+
+
+if TYPE_CHECKING:
+    func: MetricFunc
+
+    func = pearson_correlation
+    func = rmse
+    func = mae

edmkit-0.0.13/src/edmkit/metrics/mae.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, overload
+
+import numpy as np
+
+__all__ = ["mae"]
+
+if TYPE_CHECKING:
+    from tinygrad import Tensor
+
+
+@overload
+def mae(predictions: np.ndarray, observations: np.ndarray) -> np.ndarray: ...
+
+
+@overload
+def mae(predictions: Tensor, observations: Tensor) -> Tensor: ...
+
+
+def mae(predictions, observations):
+    """Mean Absolute Error.
+
+    Parameters
+    ----------
+    predictions : np.ndarray or Tensor
+        ``(N,)``, ``(N, D)``, or ``(B, N, D)``.
+    observations : np.ndarray or Tensor
+        Same shape as `predictions`.
+
+    Returns
+    -------
+    np.ndarray or Tensor
+        ``()`` for 1D/2D input, ``(B,)`` for 3D input.
+
+    Raises
+    ------
+    ValueError
+        - If `predictions` and `observations` have different shapes.
+        - If the inputs are not 1D, 2D, or 3D.
+    """
+    if isinstance(predictions, np.ndarray):
+        return _numpy(predictions, observations)
+
+    return _tensor(predictions, observations)
+
+
+def _numpy(predictions: np.ndarray, observations: np.ndarray) -> np.ndarray:
+    if predictions.shape != observations.shape:
+        raise ValueError(f"Shape mismatch: predictions {predictions.shape} vs observations {observations.shape}")
+    if predictions.ndim not in (1, 2, 3):
+        raise ValueError(f"Expected 1D, 2D, or 3D arrays, got {predictions.ndim}D")
+
+    # treat a 1D series as a single target dimension
+    if predictions.ndim == 1:
+        predictions = predictions[:, None]
+        observations = observations[:, None]
+
+    # 2D: (N, D) -> (N,) -> ()
+    # 3D: (B, N, D) -> (B, N) -> (B,)
+    return np.abs(predictions - observations).mean(axis=-1).mean(axis=-1)
+
+
+def _tensor(predictions: Tensor, observations: Tensor) -> Tensor:
+    if predictions.shape != observations.shape:
+        raise ValueError(f"Shape mismatch: predictions {predictions.shape} vs observations {observations.shape}")
+    if predictions.ndim not in (1, 2, 3):
+        raise ValueError(f"Expected 1D, 2D, or 3D tensors, got {predictions.ndim}D")
+
+    # treat a 1D series as a single target dimension
+    if predictions.ndim == 1:
+        predictions = predictions.unsqueeze(-1)
+        observations = observations.unsqueeze(-1)
+
+    # 2D: (N, D) -> (N,) -> ()
+    # 3D: (B, N, D) -> (B, N) -> (B,)
+    return (predictions - observations).abs().mean(axis=-1).mean(axis=-1)

edmkit-0.0.13/src/edmkit/metrics/pearson_correlation.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, overload
+
+import numpy as np
+
+__all__ = ["pearson_correlation"]
+
+if TYPE_CHECKING:
+    from tinygrad import Tensor
+
+
+@overload
+def pearson_correlation(predictions: np.ndarray, observations: np.ndarray) -> np.ndarray: ...
+
+
+@overload
+def pearson_correlation(predictions: Tensor, observations: Tensor) -> Tensor: ...
+
+
+def pearson_correlation(predictions, observations):
+    """Mean Pearson correlation over the target dimensions.
+
+    The correlation is computed per target dimension along the sample axis,
+    then averaged over the dimensions.
+
+    Parameters
+    ----------
+    predictions : np.ndarray or Tensor
+        ``(N,)``, ``(N, D)``, or ``(B, N, D)``.
+    observations : np.ndarray or Tensor
+        Same shape as `predictions`.
+
+    Returns
+    -------
+    np.ndarray or Tensor
+        ``()`` for 1D/2D input, ``(B,)`` for 3D input.
+
+    Raises
+    ------
+    ValueError
+        - If `predictions` and `observations` have different shapes.
+        - If the inputs are not 1D, 2D, or 3D.
+    """
+    if isinstance(predictions, np.ndarray):
+        return _numpy(predictions, observations)
+
+    return _tensor(predictions, observations)
+
+
+def _numpy(predictions: np.ndarray, observations: np.ndarray) -> np.ndarray:
+    if predictions.shape != observations.shape:
+        raise ValueError(f"Shape mismatch: predictions {predictions.shape} vs observations {observations.shape}")
+    if predictions.ndim not in (1, 2, 3):
+        raise ValueError(f"Expected 1D, 2D, or 3D arrays, got {predictions.ndim}D")
+
+    # treat a 1D series as a single target dimension
+    if predictions.ndim == 1:
+        predictions = predictions[:, None]
+        observations = observations[:, None]
+
+    predictions_centered = predictions - predictions.mean(axis=-2, keepdims=True)
+    observations_centered = observations - observations.mean(axis=-2, keepdims=True)
+
+    numerator = (predictions_centered * observations_centered).sum(axis=-2)
+    denominator = np.sqrt((predictions_centered**2).sum(axis=-2) * (observations_centered**2).sum(axis=-2))
+
+    denominator = np.where(denominator > 0, denominator, np.nan)
+
+    # 2D: (D,) -> ()
+    # 3D: (B, D) -> (B,)
+    return (numerator / denominator).mean(axis=-1)
+
+
+def _tensor(predictions: Tensor, observations: Tensor) -> Tensor:
+    if predictions.shape != observations.shape:
+        raise ValueError(f"Shape mismatch: predictions {predictions.shape} vs observations {observations.shape}")
+    if predictions.ndim not in (1, 2, 3):
+        raise ValueError(f"Expected 1D, 2D, or 3D tensors, got {predictions.ndim}D")
+
+    # treat a 1D series as a single target dimension
+    if predictions.ndim == 1:
+        predictions = predictions.unsqueeze(-1)
+        observations = observations.unsqueeze(-1)
+
+    predictions_centered = predictions - predictions.mean(axis=-2, keepdim=True)
+    observations_centered = observations - observations.mean(axis=-2, keepdim=True)
+
+    numerator = (predictions_centered * observations_centered).sum(axis=-2)
+    denominator = predictions_centered.pow(2).sum(axis=-2) * observations_centered.pow(2).sum(axis=-2)
+
+    denominator = (denominator > 0).where(denominator.clamp(min_=1e-12).sqrt(), float("nan"))
+
+    # 2D: (D,) -> ()
+    # 3D: (B, D) -> (B,)
+    return (numerator / denominator).mean(axis=-1)