edmkit 0.0.3__tar.gz → 0.0.4__tar.gz

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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.3"
+version = "0.0.4"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
 readme = "README.md"
@@ -23,6 +23,7 @@ dev = [
 
 [tool.pytest.ini_options]
 addopts = "-v --tb=short"
+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\"')",

{edmkit-0.0.3 → edmkit-0.0.4}/src/edmkit/ccm.py

@@ -1,27 +1,27 @@
-"""Convergent Cross Mapping (CCM) for causality detection in time series."""
-
 from collections.abc import Callable
 from functools import partial
+from typing import TypeAlias
 
 import numpy as np
 
 from edmkit.simplex_projection import simplex_projection
 from edmkit.smap import smap
+from edmkit.types import PredictFunc
 
-# NOTE: Module-level RNG is shared mutable state. When test order changes,
-# the sequence of draws from this RNG changes too, causing non-determinism.
-# Tests should always pass an explicit sampler (e.g. via make_seeded_sampler)
-# rather than relying on default_sampler.
-_rng = np.random.default_rng(42)
+SampleFunc: TypeAlias = 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]
+"""AggregateFunc is a function that takes an array of values and returns a single value."""
 
 
-def default_sampler(pool: np.ndarray, size: int) -> np.ndarray:
-    return _rng.choice(pool, size=size, replace=True)
+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)
 
-PredictFunc = Callable[[np.ndarray, np.ndarray, np.ndarray], np.ndarray]
-Sampler = Callable[[np.ndarray, int], np.ndarray]
-Aggregator = Callable[[np.ndarray], float]
+    return sample_func
 
 
 def bootstrap(
@@ -33,7 +33,7 @@ def bootstrap(
     *,
     library_pool: np.ndarray,
     prediction_pool: np.ndarray,
-    sampler: Sampler = default_sampler,
+    sample_func: SampleFunc | None = None,
     batch_size: int | None = 10,
 ) -> np.ndarray:
     """
@@ -51,15 +51,16 @@ def bootstrap(
     lib_sizes : np.ndarray
         Array of library sizes to test convergence.
     predict_func : :type: `PredictFunc`
-        Prediction function with signature (lib_X, lib_Y, pred_X) -> predictions.
+        Prediction function with signature (X, Y, Q) -> predictions.
     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.
-    sampler : :type: `Sampler`, default `default_sampler`
+    sample_func : :type: `SampleFunc` | 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
         If specified, predictions are made in batches to limit memory usage.
 
@@ -68,6 +69,9 @@ def bootstrap(
     samples : np.ndarray of shape (n_samples, len(lib_sizes))
         Per-sample correlation coefficients.
     """
+    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 not callable(predict_func):
@@ -86,7 +90,7 @@ def bootstrap(
         Y = Y[:, None]
 
     prediction_indices = np.tile(prediction_pool, (batch_size, 1))
-    query_points = X[prediction_indices]
+    Q = X[prediction_indices]
     actual = Y[prediction_indices]
 
     samples = np.zeros((n_samples, len(lib_sizes)))
@@ -96,12 +100,12 @@ def bootstrap(
         while remaining > 0:
             batch = min(batch_size, remaining)
 
-            library_indices = np.vstack([sampler(library_pool, lib_size) for _ in range(batch)])
+            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, query_points[:batch])
+            predictions = predict_func(lib_X, lib_Y, Q[:batch])
 
             offset = n_samples - remaining
             samples[offset : offset + batch, i] = pearson_correlation(predictions, actual[:batch])
@@ -119,8 +123,8 @@ def ccm(
     *,
     library_pool: np.ndarray,
     prediction_pool: np.ndarray,
-    sampler: Sampler = default_sampler,
-    aggregator: Aggregator = np.mean,
+    sample_func: SampleFunc | None = None,
+    aggregate_func: AggregateFunc = np.mean,
     batch_size: int | None = 10,
 ) -> np.ndarray:
     """
@@ -139,7 +143,7 @@ def ccm(
     lib_sizes : np.ndarray
         Array of library sizes to test convergence.
     predict_func : :type: `PredictFunc`
-        Prediction function with signature (lib_X, lib_Y, pred_X) -> predictions.
+        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
         Number of random samples per library size for bootstrapping.
@@ -147,10 +151,11 @@ 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.
-    sampler : :type: `Sampler`, default `default_sampler`
+    sample_func : :type: `SampleFunc` | None, default None
         Function responsible for drawing a library sample of a given size.
         It receives `(pool, size)` and returns an array of indices.
-    aggregator : :type: `Aggregator`, default `np.mean`
+        When None, a fresh RNG-backed sampler is created per call.
+    aggregate_func : :type: `AggregateFunc`, default `np.mean`
         Reducer applied to the correlation samples for each library size.
     batch_size : int | None, default None
         If not specified, batch_size == n_samples.
@@ -167,7 +172,7 @@ def ccm(
         - If `lib_sizes` contains non-positive values.
         - If `predict_func` is not callable.
         - If `n_samples` is not positive.
-        - If `aggregator` is not callable.
+        - If `aggregate_func` is not callable.
         - If `library_pool` or `prediction_pool` is invalid.
 
     Notes
@@ -231,8 +236,8 @@ def ccm(
     )
     ```
     """
-    if aggregator is None or not callable(aggregator):
-        raise ValueError("aggregator must be a callable")
+    if aggregate_func is None or not callable(aggregate_func):
+        raise ValueError("aggregate_func must be a callable")
 
     samples = bootstrap(
         X,
@@ -242,11 +247,11 @@ def ccm(
         n_samples,
         library_pool=library_pool,
         prediction_pool=prediction_pool,
-        sampler=sampler,
+        sample_func=sample_func,
         batch_size=batch_size,
     )
 
-    return np.array([aggregator(samples[:, i]) for i in range(samples.shape[1])])
+    return np.array([aggregate_func(samples[:, i]) for i in range(samples.shape[1])])
 
 
 def pearson_correlation(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
@@ -275,7 +280,9 @@ def pearson_correlation(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
     cov = ((X - mean_X) * (Y - mean_Y)).mean(axis=1)
     std_X = X.std(axis=1)
     std_Y = Y.std(axis=1)
-    correlation = cov / (std_X * std_Y)
+    denom = std_X * std_Y
+    safe_denom = np.where(denom > 0, denom, 1.0)
+    correlation = np.where(denom > 0, cov / safe_denom, 0.0)
 
     return correlation.squeeze()
 
@@ -289,8 +296,8 @@ def with_simplex_projection(
     *,
     library_pool: np.ndarray,
     prediction_pool: np.ndarray,
-    sampler: Sampler = default_sampler,
-    aggregator: Aggregator = np.mean,
+    sample_func: SampleFunc | None = None,
+    aggregate_func: AggregateFunc = np.mean,
 ) -> np.ndarray:
     """
     Perform Convergent Cross Mapping using simplex projection.
@@ -314,10 +321,10 @@ def with_simplex_projection(
         Indices that can be used to draw library samples. Defaults to the full range.
     prediction_pool : np.ndarray, optional
         Indices that should be predicted (leave-one-out over this set). Defaults to the full range.
-    sampler : callable, optional
+    sample_func : callable, optional
         Function responsible for drawing a library sample of a given size.
-        Falls back to `default_sampler` (bootstrap with replacement) when omitted.
-    aggregator : callable, optional
+        When omitted, a fresh RNG-backed sampler is created per call.
+    aggregate_func : callable, optional
         Reducer applied to the correlation samples for each library size.
         Falls back to `np.mean` when omitted.
     Returns
@@ -381,8 +388,8 @@ def with_simplex_projection(
         n_samples=n_samples,
         library_pool=library_pool,
         prediction_pool=prediction_pool,
-        sampler=sampler,
-        aggregator=aggregator,
+        sample_func=sample_func,
+        aggregate_func=aggregate_func,
     )
 
 
@@ -397,8 +404,8 @@ def with_smap(
     *,
     library_pool: np.ndarray,
     prediction_pool: np.ndarray,
-    sampler: Sampler = default_sampler,
-    aggregator: Aggregator = np.mean,
+    sample_func: SampleFunc | None = None,
+    aggregate_func: AggregateFunc = np.mean,
 ) -> np.ndarray:
     """
     Perform Convergent Cross Mapping using S-Map (local linear regression).
@@ -426,10 +433,10 @@ def with_smap(
         Indices that can be used to draw library samples. Defaults to the full range.
     prediction_pool : np.ndarray, optional
         Indices that should be predicted (leave-one-out over this set). Defaults to the full range.
-    sampler : callable, optional
+    sample_func : callable, optional
         Function responsible for drawing a library sample of a given size.
-        Falls back to `default_sampler` (bootstrap with replacement) when omitted.
-    aggregator : callable, optional
+        When omitted, a fresh RNG-backed sampler is created per call.
+    aggregate_func : callable, optional
         Reducer applied to the correlation samples for each library size.
         Falls back to `np.mean` when omitted.
     Returns
@@ -494,6 +501,6 @@ def with_smap(
         n_samples=n_samples,
         library_pool=library_pool,
         prediction_pool=prediction_pool,
-        sampler=sampler,
-        aggregator=aggregator,
+        sample_func=sample_func,
+        aggregate_func=aggregate_func,
     )

edmkit-0.0.4/src/edmkit/embedding.py

@@ -0,0 +1,224 @@
+from functools import partial
+from itertools import product
+
+import numpy as np
+
+from edmkit.metrics import MetricFunc, mean_rho
+from edmkit.simplex_projection import simplex_projection
+from edmkit.splits import SplitFunc, sliding_folds
+from edmkit.types import PredictFunc
+
+
+def lagged_embed(x: np.ndarray, tau: int, e: int):
+    """Lagged embedding of a time series `x`.
+
+    Parameters
+    ----------
+        `x` : `np.ndarray` of shape `(N,)`
+        `tau` : `int`
+        `e` : `int`
+
+    Returns
+    -------
+        `np.ndarray` of shape `(N - (e - 1) * tau, e)`
+
+    Raises
+    ------
+    ValueError
+        - If `x` is not a 1D array.
+        - If `tau` or `e` is not positive.
+        - If `e * tau >= len(x)`.
+
+    Notes
+    -----
+    - While open to interpretation, it's generally more intuitive to consider the embedding as starting from the `(e - 1) * tau`th element of the original time series and ending at the `len(x) - 1`th element (the last value), rather than starting from the 0th element and ending at `len(x) - 1 - (e - 1) * tau`.
+    - This distinction reflects whether we think of "attaching past values to the present" or "attaching future values to the present". The information content of the result is the same either way.
+    - The use of `reversed` in the implementation emphasizes this perspective.
+
+    Examples
+    --------
+    ```
+    import numpy as np
+    from edm.embedding import lagged_embed
+
+    x = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+    tau = 2
+    e = 3
+
+    E = lagged_embed(x, tau, e)
+    print(E)
+    print(E.shape)
+    # [[4 2 0]
+    #  [5 3 1]
+    #  [6 4 2]
+    #  [7 5 3]
+    #  [8 6 4]
+    #  [9 7 5]]
+    # (6, 3)
+    ```
+    """
+    if not len(x.shape) == 1:
+        raise ValueError(f"X must be a 1D array, got x.shape={x.shape}")
+    if tau <= 0 or e <= 0:
+        raise ValueError(f"tau and e must be positive, got tau={tau}, e={e}")
+    if (e - 1) * tau >= x.shape[0]:
+        raise ValueError(f"e and tau must satisfy `(e - 1) * tau < len(X)`, got e={e}, tau={tau}")
+
+    return np.array([x[tau * (e - 1) :]] + [x[tau * i : -tau * ((e - 1) - i)] for i in reversed(range(e - 1))]).transpose()
+
+
+def scan(
+    x: np.ndarray,
+    Y: np.ndarray | None = None,
+    *,
+    E: list[int],
+    tau: list[int],
+    n_ahead: int = 1,
+    split: SplitFunc | None = None,
+    predict: PredictFunc | None = None,
+    metric: MetricFunc | None = None,
+) -> np.ndarray:
+    """Grid search over (E, tau) with cross-validation.
+
+    Parameters
+    ----------
+    x : np.ndarray, shape (N,)
+        Time series to embed.
+    Y : np.ndarray or None, shape (N,) or (N, M)
+        Prediction target. If None, self-prediction (Y = x).
+    E : list[int]
+        Embedding dimension candidates.
+    tau : list[int]
+        Time delay candidates.
+    n_ahead : int
+        Prediction horizon (steps ahead).
+    split : SplitFunc or None
+        Callable ``(n: int) -> list[Fold]``. Defaults to sliding_folds.
+    predict : PredictFunc or None
+        Prediction function. Defaults to ``simplex_projection``.
+    metric : MetricFunc or None
+        Evaluation metric. Defaults to ``mean_rho``.
+
+    Returns
+    -------
+    scores : np.ndarray, shape (len(E), len(tau), K_max)
+        Per-fold CV metric for each (E, tau) combination.
+        K_max is the maximum number of folds across all E values.
+        Entries where the fold does not exist are NaN.
+    """
+    N = len(x)
+
+    if Y is None:
+        Y = x
+    if predict is None:
+        predict = simplex_projection
+    if metric is None:
+        metric = mean_rho
+    if split is None:
+        split = partial(
+            sliding_folds,
+            train_size=max(N // 5, 2),
+            validation_size=max(N // 10, 1),
+        )
+
+    if Y.ndim == 1:
+        Y = Y[:, None]
+
+    n_tau = len(tau)
+    tau_max = max(tau)
+    n_targets = Y.shape[1]
+
+    # collect ndarrays of shape (n_tau, n_folds) for each E, then pack into a single ndarray at the end
+    results: list[np.ndarray | None] = []
+
+    for e in E:
+        k = e + 1
+        max_lag = (e - 1) * tau_max
+        n_usable = N - max_lag - n_ahead
+
+        if n_usable < 2:
+            results.append(None)
+            continue
+
+        embeddings = [lagged_embed(x, t, e)[-(n_usable + n_ahead) : -n_ahead] for t in tau]
+
+        Y_aligned = Y[max_lag + n_ahead : N]
+
+        folds = split(n_usable)
+        folds = [fold for fold in folds if len(fold.train) >= k]  # ensure at least k points
+        n_folds = len(folds)
+
+        if n_folds == 0:
+            results.append(None)
+            continue
+
+        validation_size = len(folds[0].validation)  # now only support fixed validation size across folds, which simplifies batching
+        max_train_size = max(len(fold.train) for fold in folds)
+        batch_size = n_tau * n_folds
+
+        X_batch = np.zeros((batch_size, max_train_size, e))
+        Y_batch = np.zeros((batch_size, max_train_size, n_targets))
+        mask = np.zeros((batch_size, max_train_size), dtype=bool)
+        Q = np.empty((batch_size, validation_size, e))
+        Y_validation = np.empty((batch_size, validation_size, n_targets))
+
+        for batch_idx, (tau_idx, fold_idx) in enumerate(product(range(n_tau), range(n_folds))):
+            X = embeddings[tau_idx]
+            fold = folds[fold_idx]
+
+            n_train = len(fold.train)
+
+            X_batch[batch_idx, :n_train] = X[fold.train]
+            Y_batch[batch_idx, :n_train] = Y_aligned[fold.train]
+            Q[batch_idx] = X[fold.validation]
+            mask[batch_idx, :n_train] = True
+            Y_validation[batch_idx] = Y_aligned[fold.validation]
+
+        predictions = predict(X_batch, Y_batch, Q, mask=None if mask.all() else mask)
+        batch_result = metric(predictions, Y_validation)
+        results.append(batch_result.reshape(n_tau, n_folds))
+
+    K_max = max((r.shape[1] for r in results if r is not None), default=0)  # max(len(folds)) for all E values, or 0 if no valid folds
+    scores = np.full((len(E), n_tau, K_max), np.nan)
+    for batch_idx, batch_result in enumerate(results):
+        if batch_result is not None:
+            scores[batch_idx, :, : batch_result.shape[1]] = batch_result
+
+    return scores
+
+
+def select(
+    scores: np.ndarray,
+    *,
+    E: list[int],
+    tau: list[int],
+) -> tuple[int, int, float]:
+    """Select best (E, tau) from scan results.
+
+    Aggregates over the fold axis (axis=2) with nanmean, then
+    finds the (E, tau) combination with the highest mean score.
+
+    Parameters
+    ----------
+    scores : np.ndarray, shape (len(E), len(tau), K_max)
+        Output of ``scan``.
+    E : list[int]
+        Embedding dimension candidates (same as passed to ``scan``).
+    tau : list[int]
+        Time delay candidates (same as passed to ``scan``).
+
+    Returns
+    -------
+    (best_E, best_tau, best_score)
+    """
+    valid_counts = np.sum(~np.isnan(scores), axis=2)
+    summed_scores = np.nansum(scores, axis=2)
+    mean_scores = np.divide(
+        summed_scores,
+        valid_counts,
+        out=np.full(summed_scores.shape, np.nan, dtype=float),
+        where=valid_counts > 0,
+    )
+    flat_idx = int(np.nanargmax(mean_scores))
+    e_idx, t_idx = np.unravel_index(flat_idx, mean_scores.shape)
+    return E[e_idx], tau[t_idx], float(mean_scores[e_idx, t_idx])

edmkit-0.0.4/src/edmkit/metrics.py

@@ -0,0 +1,130 @@
+from typing import TYPE_CHECKING, Callable, TypeAlias
+
+import numpy as np
+
+MetricFunc: TypeAlias = Callable[[np.ndarray, np.ndarray], np.ndarray]
+"""MetricFunc is a function that takes (predictions, observations) and returns a metric value."""
+
+
+def validate_and_promote(
+    predictions: np.ndarray,
+    observations: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Validate shape match and promote 1D to 2D."""
+    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")
+    if predictions.ndim == 1:
+        predictions = predictions[:, None]
+        observations = observations[:, None]
+    return predictions, observations
+
+
+def rhos(
+    predictions: np.ndarray,
+    observations: np.ndarray,
+) -> np.ndarray:
+    """Pearson correlation per dimension.
+
+    Parameters
+    ----------
+    predictions : np.ndarray
+        ``(N,)``, ``(N, D)``, or ``(B, N, D)``.
+    observations : np.ndarray
+        Same shape as predictions.
+
+    Returns
+    -------
+    np.ndarray
+        ``(1,)`` for 1D input, ``(D,)`` for 2D, ``(B, D)`` for 3D.
+    """
+    predictions, observations = validate_and_promote(predictions, observations)
+
+    p_centered = predictions - predictions.mean(axis=-2, keepdims=True)
+    o_centered = observations - observations.mean(axis=-2, keepdims=True)
+    num = (p_centered * o_centered).sum(axis=-2)
+    denom = np.sqrt((p_centered**2).sum(axis=-2) * (o_centered**2).sum(axis=-2))
+    safe_denom = np.where(denom > 0, denom, 1.0)
+
+    return np.where(denom > 0, num / safe_denom, 0.0)
+
+
+def mean_rho(
+    predictions: np.ndarray,
+    observations: np.ndarray,
+) -> np.ndarray:
+    """Mean Pearson correlation.
+
+    Parameters
+    ----------
+    predictions : np.ndarray
+        ``(N,)``, ``(N, D)``, or ``(B, N, D)``.
+    observations : np.ndarray
+        Same shape as predictions.
+
+    Returns
+    -------
+    np.ndarray
+        ``()`` for 1D/2D input, ``(B,)`` for 3D input.
+    """
+    return rhos(predictions, observations).mean(axis=-1)
+
+
+def rmse(
+    predictions: np.ndarray,
+    observations: np.ndarray,
+) -> np.ndarray:
+    """Root Mean Squared Error.
+
+    Parameters
+    ----------
+    predictions : np.ndarray
+        ``(N,)``, ``(N, D)``, or ``(B, N, D)``.
+    observations : np.ndarray
+        Same shape as *predictions*.
+
+    Returns
+    -------
+    np.ndarray
+        ``()`` for 1D/2D input, ``(B,)`` for 3D input.
+    """
+    predictions, observations = validate_and_promote(predictions, observations)
+
+    # 2D: (N, D) -> (N,) -> ()
+    # 3D: (B, N, D) -> (B, N) -> (B,)
+    return np.sqrt(((predictions - observations) ** 2).mean(axis=-1).mean(axis=-1))
+
+
+def mae(
+    predictions: np.ndarray,
+    observations: np.ndarray,
+) -> np.ndarray:
+    """Mean Absolute Error.
+
+    Parameters
+    ----------
+    predictions : np.ndarray
+        ``(N,)``, ``(N, D)``, or ``(B, N, D)``.
+    observations : np.ndarray
+        Same shape as predictions.
+
+    Returns
+    -------
+    np.ndarray
+        ``()`` for 1D/2D input, ``(B,)`` for 3D input.
+    """
+    predictions, observations = validate_and_promote(predictions, observations)
+
+    # 2D: (N, D) -> (N,) -> ()
+    # 3D: (B, N, D) -> (B, N) -> (B,)
+    return np.abs(predictions - observations).mean(axis=-1).mean(axis=-1)
+
+
+if TYPE_CHECKING:
+    func: MetricFunc
+
+    func = rhos
+    func = mean_rho
+    func = rmse
+    func = mae