kews 0.1.0__py3-none-any.whl

kews/__init__.py

@@ -0,0 +1,119 @@
+"""KEWS: compact Koopman early warning signals."""
+
+from __future__ import annotations
+
+from typing import Any, Literal, Mapping, Optional, Union
+
+import jax
+import jax.numpy as jnp
+from jax.typing import ArrayLike
+
+from ._core import rolling_residuals
+from .observables import _RFF, time_delay
+
+
+def _window_length(n_time: int, window: Union[int, float]) -> int:
+    """Resolve a relative or absolute window length."""
+    length = int(window * n_time) if window <= 1 else int(window)
+    if not 1 < length < n_time:
+        raise ValueError(
+            f"window resolves to {length}; expected a value in (1, {n_time})"
+        )
+    return length
+
+
+def _windows(data: jax.Array, length: int) -> jax.Array:
+    """Return ``(window, time)`` or ``(window, state, time)`` snapshots."""
+    indices = jnp.arange(data.shape[-1] - length + 1)[:, None] + jnp.arange(length)
+    if data.ndim == 1:
+        return data[indices]
+    return jnp.transpose(data[:, indices], (1, 0, 2))
+
+
+def resdmd(
+    data: ArrayLike,
+    window: Union[int, float] = 0.5,
+    observable: Optional[Any] = None,
+    rank: Union[int, float] = 0.9,
+    control: Optional[ArrayLike] = None,
+    *,
+    control_method: Literal["dmdc", "augmented"] = "dmdc",
+    future_input: Literal["future", "zero"] = "future",
+    fit_kwargs: Optional[Mapping[str, Any]] = None,
+) -> jax.Array:
+    """Compute a rolling Koopman early warning signal.
+
+    The same call covers autonomous and controlled dynamics. Leave ``control``
+    unset for autonomous ResDMD or pass an aligned control trajectory to use a
+    control-aware formulation. Each residual is aligned to the right edge of
+    its window; earlier entries are ``NaN``.
+
+    A learned observable is fitted on the first window and then held fixed.
+    The first window also determines the retained rank. Random Fourier
+    features use one seeded basis for the complete trajectory.
+
+    Args:
+        data: Scalar state shaped ``(time,)`` or multivariate state shaped
+            ``(state, time)``.
+        window: Fraction of the trajectory in ``(0, 1)`` or absolute sample
+            count. It must resolve to at least two samples and less than the
+            complete trajectory.
+        observable: Observable created by :mod:`kews.observables`. Defaults to
+            ``time_delay()``.
+        rank: Retained SVD energy fraction in ``(0, 1)`` or an absolute rank.
+        control: Optional input shaped ``(time,)`` or ``(control, time)``.
+        control_method: ``"dmdc"`` fits DMD with raw control in lifted state
+            space. ``"augmented"`` fits ``[phi(x), phi(x)u, u]`` snapshots.
+        future_input: For the augmented method, use observed next inputs with
+            ``"future"`` or replace them with zero using ``"zero"``.
+        fit_kwargs: Training arguments for a learned observable, such as
+            ``epochs``, ``batch_size``, and ``lr`` for ``observables.mlp``.
+
+    Returns:
+        One-dimensional JAX array with the same time length as ``data``.
+        Positions before the first completed window are ``NaN``.
+
+    Raises:
+        ValueError: If array orientation, window, rank, control alignment, or a
+            method option is invalid.
+    """
+    data = jnp.asarray(data)
+    if data.ndim not in (1, 2):
+        raise ValueError("data must have shape (time,) or (state, time)")
+    length = _window_length(data.shape[-1], window)
+
+    control_windows = None
+    if control is not None:
+        control = jnp.asarray(control)
+        if control.ndim not in (1, 2) or control.shape[-1] != data.shape[-1]:
+            raise ValueError(
+                "control must have shape (time,) or (control, time) and match data"
+            )
+        control_windows = _windows(control, length)
+
+    observable = time_delay() if observable is None else observable
+    if not callable(getattr(observable, "transform", None)):
+        raise ValueError("observable must come from kews.observables")
+
+    if isinstance(observable, _RFF):
+        # Lift once to preserve the paper's shared basis and floating-point evaluation.
+        state_windows = _windows(observable.transform(data), length)
+        observable = None
+    else:
+        state_windows = _windows(data, length)
+
+    residuals = rolling_residuals(
+        state_windows,
+        observable,
+        rank,
+        control_windows,
+        control_method,
+        future_input,
+        fit_kwargs,
+    )
+    result = jnp.full((data.shape[-1],), jnp.nan, dtype=jnp.float32)
+    return result.at[length - 1 :].set(residuals)
+
+
+__version__ = "0.1.0"
+__all__ = ["resdmd"]

kews/_core.py

@@ -0,0 +1,441 @@
+"""Koopman and ResDMD numerical kernels.
+
+Snapshot matrices have shape ``(features, time)``; rolling evaluation adds a
+leading window axis. Normalization and numerical tolerances follow the paper
+reference so that refactoring does not change the plotted indicators.
+"""
+
+from __future__ import annotations
+
+from numbers import Integral, Real
+from typing import Any, Literal, Mapping, Optional, Union
+
+import jax
+import jax.numpy as jnp
+
+Rank = Union[int, float]
+
+
+def delay_embed(x: jax.Array, n_delay: int) -> jax.Array:
+    """Stack delayed state copies with time on the final axis."""
+    if not isinstance(n_delay, Integral) or isinstance(n_delay, bool) or n_delay < 1:
+        raise ValueError("n_delay must be a positive integer")
+    x = jnp.atleast_2d(x)
+    n_time = x.shape[1]
+    if n_delay > n_time:
+        raise ValueError(
+            f"n_delay must not exceed the time length ({n_time}), got {n_delay}"
+        )
+    return jnp.vstack(
+        [x[:, offset : n_time - n_delay + offset + 1] for offset in range(n_delay)]
+    )
+
+
+def random_fourier_features(
+    x: jax.Array,
+    n_features: int,
+    scale: float,
+    key: jax.Array,
+) -> jax.Array:
+    """Map state columns to random cosine features."""
+    x = jnp.atleast_2d(x)
+    projection_key, phase_key = jax.random.split(key)
+    weights = jax.random.normal(projection_key, (n_features, x.shape[0])) / scale
+    phase = jax.random.uniform(
+        phase_key, (n_features,), maxval=2 * jnp.pi
+    )
+    return jnp.sqrt(2 / n_features) * jnp.cos(weights @ x + phase[:, None])
+
+
+def _rank(singular_values: jax.Array, low_rank: Rank) -> int:
+    """Resolve an energy fraction or absolute SVD rank to a concrete integer."""
+    n_modes = singular_values.shape[0]
+    if isinstance(low_rank, Integral) and not isinstance(low_rank, bool):
+        if low_rank < 1:
+            raise ValueError("rank must be a fraction in (0, 1) or an integer >= 1")
+        return min(int(low_rank), n_modes)
+    if isinstance(low_rank, Real) and 0 < low_rank < 1:
+        energy = singular_values**2
+        cumulative = jnp.cumsum(energy) / jnp.maximum(jnp.sum(energy), 1e-12)
+        index = int(jax.device_get(jnp.searchsorted(cumulative, low_rank)))
+        return min(index + 1, n_modes)
+    raise ValueError("rank must be a fraction in (0, 1) or an integer >= 1")
+
+
+def _svd_inverse(singular_values: jax.Array) -> jax.Array:
+    """Return a stable diagonal inverse for retained singular values."""
+    return jnp.diag(
+        jnp.where(singular_values > 1e-12, 1.0 / singular_values, 0.0)
+    )
+
+
+def _edmd_matrices(
+    current: jax.Array,
+    future: jax.Array,
+    rank: int,
+) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
+    """Build EDMD bases, Gram matrix, and cross matrix."""
+    # Preserve the paper's snapshot normalization before truncating the SVD.
+    _, singular_values, vh = jnp.linalg.svd(
+        current / current.shape[1], full_matrices=False
+    )
+    values = singular_values[:rank]
+    right = vh[:rank].T
+    inverse = _svd_inverse(values)
+    basis_current = current @ right @ inverse
+    basis_future = future @ right @ inverse
+    gram = basis_current.T @ basis_current
+    gram = (gram + gram.T) / 2
+    cross = basis_current.T @ basis_future
+    return basis_current, basis_future, gram, cross
+
+
+def _mean_residual(
+    gram: jax.Array,
+    cross: jax.Array,
+    future_gram: jax.Array,
+    operator: jax.Array,
+) -> jax.Array:
+    """Average the ResDMD residual over Koopman eigenpairs."""
+    eigenvalues, eigenvectors = jnp.linalg.eig(operator)
+
+    def one(eigenvalue: jax.Array, eigenvector: jax.Array) -> jax.Array:
+        # Residual energy is normalized by this eigenfunction's Gram norm.
+        residual_matrix = (
+            future_gram
+            - eigenvalue * cross.conj().T
+            - eigenvalue.conj() * cross
+            + jnp.abs(eigenvalue) ** 2 * gram
+        )
+        conjugate = eigenvector.conj()
+        numerator = (conjugate @ residual_matrix @ eigenvector).real
+        denominator = (conjugate @ gram @ eigenvector).real
+        return jnp.where(
+            jnp.abs(denominator) > 1e-10,
+            jnp.abs(numerator / denominator),
+            0.0,
+        )
+
+    return jnp.mean(jax.vmap(one)(eigenvalues, eigenvectors.T))
+
+
+def _edmd_residual(current: jax.Array, future: jax.Array, rank: int) -> jax.Array:
+    """Compute one autonomous or augmented EDMD residual."""
+    _, basis_future, gram, cross = _edmd_matrices(current, future, rank)
+    future_gram = basis_future.T @ basis_future
+    future_gram = (future_gram + future_gram.T) / 2
+    operator = jnp.linalg.pinv(gram, rtol=1e-4) @ cross
+    residual = _mean_residual(gram, cross, future_gram, operator)
+    valid = jnp.all(jnp.isfinite(operator)) & jnp.isfinite(residual)
+    return jnp.where(valid, residual, 0.0)
+
+
+def _dmdc(
+    current: jax.Array,
+    future: jax.Array,
+    control: jax.Array,
+    input_rank: int,
+    output_rank: int,
+) -> tuple[jax.Array, jax.Array]:
+    """Fit DMD with control from aligned snapshot matrices."""
+    n_state = current.shape[0]
+    omega = jnp.vstack([current, control])
+    left, singular_values, vh = jnp.linalg.svd(omega, full_matrices=False)
+    left = left[:, :input_rank]
+    projection = future @ vh[:input_rank].T @ _svd_inverse(
+        singular_values[:input_rank]
+    )
+    state_operator = projection @ left[:n_state].T
+    control_operator = projection @ left[n_state:].T
+
+    # Match the reduced-order lift used by the paper reference implementation.
+    output_basis = jnp.linalg.svd(future, full_matrices=False)[0][:, :output_rank]
+    state_operator = (
+        output_basis
+        @ (output_basis.T @ state_operator @ output_basis)
+        @ output_basis.T
+    )
+    control_operator = output_basis @ (output_basis.T @ control_operator)
+    return state_operator, control_operator
+
+
+def _control_residual(
+    current: jax.Array,
+    future: jax.Array,
+    control: jax.Array,
+    input_rank: int,
+    output_rank: int,
+) -> jax.Array:
+    """Compute one DMDc residual in lifted state space."""
+    operator, _ = _dmdc(current, future, control, input_rank, output_rank)
+    gram = current @ current.T
+    cross = current @ future.T
+    future_gram = future @ future.T
+    future_gram = (future_gram + future_gram.T) / 2
+    residual = _mean_residual(gram, cross, future_gram, operator)
+    valid = jnp.all(jnp.isfinite(operator)) & jnp.isfinite(residual)
+    return jnp.where(valid, residual, 0.0)
+
+
+def _augment(
+    lifted: jax.Array,
+    control: jax.Array,
+    future_input: Literal["future", "zero"],
+) -> tuple[jax.Array, jax.Array]:
+    """Build ``[phi(x), phi(x)u, u]`` current and future snapshots."""
+    state_current, state_future = lifted[:, :-1], lifted[:, 1:]
+    control_current = control[:, :-1]
+    if future_input == "future":
+        control_future = control[:, 1:]
+    elif future_input == "zero":
+        control_future = jnp.zeros_like(control_current)
+    else:
+        raise ValueError("future_input must be 'future' or 'zero'")
+
+    interaction_current = jnp.einsum(
+        "it,jt->ijt", state_current, control_current
+    ).reshape((-1, state_current.shape[1]))
+    interaction_future = jnp.einsum(
+        "it,jt->ijt", state_future, control_future
+    ).reshape((-1, state_future.shape[1]))
+    return (
+        jnp.vstack([state_current, interaction_current, control_current]),
+        jnp.vstack([state_future, interaction_future, control_future]),
+    )
+
+
+def _lift(observable: Any, x: jax.Array) -> jax.Array:
+    """Apply an observable, or pass through features lifted before windowing."""
+    return x if observable is None else observable.transform(x)
+
+
+def rolling_residuals(
+    windows: jax.Array,
+    observable: Any,
+    low_rank: Rank,
+    control_windows: Optional[jax.Array] = None,
+    control_method: Literal["dmdc", "augmented"] = "dmdc",
+    future_input: Literal["future", "zero"] = "future",
+    fit_kwargs: Optional[Mapping[str, Any]] = None,
+) -> jax.Array:
+    """Compute one residual for every pre-built rolling window."""
+    if getattr(observable, "fit_on_first_window", False):
+        # Freeze learned features after the initial window to avoid future-data leakage.
+        observable.fit(windows[0], **dict(fit_kwargs or {}))
+
+    first_lifted = _lift(observable, windows[0])
+    if first_lifted.shape[1] < 2:
+        raise ValueError("observable must produce at least two time columns per window")
+
+    if control_windows is None:
+        first_current = first_lifted[:, :-1]
+        singular_values = jnp.linalg.svd(
+            first_current / first_current.shape[1], full_matrices=False
+        )[1]
+        rank = _rank(singular_values, low_rank)
+
+        # Hold rank fixed so later windows use the same approximation dimension.
+        @jax.jit
+        def one(window: jax.Array) -> jax.Array:
+            lifted = _lift(observable, window)
+            return _edmd_residual(lifted[:, :-1], lifted[:, 1:], rank)
+
+        return jax.vmap(one)(windows)
+
+    if control_method not in ("dmdc", "augmented"):
+        raise ValueError("control_method must be 'dmdc' or 'augmented'")
+
+    first_control = control_windows[0]
+    first_control = (
+        first_control.reshape((1, -1))
+        if first_control.ndim == 1
+        else first_control
+    )
+    n_lifted = first_lifted.shape[1]
+
+    # Delay coordinates end at the latest sample; align control to those endpoints.
+    if control_method == "dmdc":
+        first_current, first_future = first_lifted[:, :-1], first_lifted[:, 1:]
+        aligned = first_control[:, first_control.shape[1] - n_lifted : -1]
+        input_rank = _rank(
+            jnp.linalg.svd(jnp.vstack([first_current, aligned]), full_matrices=False)[1],
+            low_rank,
+        )
+        output_rank = _rank(
+            jnp.linalg.svd(first_future, full_matrices=False)[1], low_rank
+        )
+
+        @jax.jit
+        def one(window: jax.Array, control_window: jax.Array) -> jax.Array:
+            lifted = _lift(observable, window)
+            control = (
+                control_window.reshape((1, -1))
+                if control_window.ndim == 1
+                else control_window
+            )
+            aligned = control[:, control.shape[1] - lifted.shape[1] : -1]
+            return _control_residual(
+                lifted[:, :-1],
+                lifted[:, 1:],
+                aligned,
+                input_rank,
+                output_rank,
+            )
+
+        return jax.vmap(one)(windows, control_windows)
+
+    first_aligned = first_control[:, first_control.shape[1] - n_lifted :]
+    first_current, _ = _augment(first_lifted, first_aligned, future_input)
+    rank = _rank(
+        jnp.linalg.svd(
+            first_current / first_current.shape[1], full_matrices=False
+        )[1],
+        low_rank,
+    )
+
+    @jax.jit
+    def one(window: jax.Array, control_window: jax.Array) -> jax.Array:
+        lifted = _lift(observable, window)
+        control = (
+            control_window.reshape((1, -1))
+            if control_window.ndim == 1
+            else control_window
+        )
+        aligned = control[:, control.shape[1] - lifted.shape[1] :]
+        current, future = _augment(lifted, aligned, future_input)
+        return _edmd_residual(current, future, rank)
+
+    return jax.vmap(one)(windows, control_windows)
+
+
+def fit_koopman_edmd(
+    x: jax.Array,
+    observable: Any,
+    low_rank: Rank = 0.9,
+    fit_kwargs: Optional[Mapping[str, Any]] = None,
+) -> dict[str, Any]:
+    """Fit an autonomous Koopman operator in a supplied observable space."""
+    x = jnp.asarray(x)
+    if getattr(observable, "fit_on_first_window", False):
+        observable.fit(x, **dict(fit_kwargs or {}))
+    lifted = _lift(observable, x)
+    if lifted.shape[1] < 2:
+        raise ValueError("observable must produce at least two time columns")
+    current, future = lifted[:, :-1], lifted[:, 1:]
+    singular_values = jnp.linalg.svd(
+        current / current.shape[1], full_matrices=False
+    )[1]
+    rank = _rank(singular_values, low_rank)
+    basis_current, basis_future, gram, cross = _edmd_matrices(
+        current, future, rank
+    )
+    operator = jnp.linalg.lstsq(gram, cross, rcond=1e-10)[0]
+    eigenvalues, eigenvectors = jnp.linalg.eig(operator)
+    return {
+        "K": operator,
+        "eigvals": eigenvalues,
+        "eigvecs": eigenvectors.T,
+        "Psi_X": basis_current.T @ current,
+        "Psi_Y": basis_current.T @ future,
+        "r": rank,
+    }
+
+
+def _aligned_control(x: jax.Array, u: jax.Array) -> jax.Array:
+    """Normalize control orientation and validate raw time alignment."""
+    u = jnp.atleast_2d(u)
+    if u.shape[1] != x.shape[-1]:
+        raise ValueError("control must have the same time length as state")
+    return u
+
+
+def compute_dmdc(
+    x: jax.Array,
+    u: jax.Array,
+    dim_delay: int = 5,
+    low_rank: Rank = 0.9,
+) -> tuple[jax.Array, jax.Array]:
+    """Fit DMDc in delay-coordinate state space."""
+    x = jnp.atleast_1d(x).flatten()
+    u = _aligned_control(x, u)
+    if x.size < dim_delay + 2:
+        raise ValueError(f"data needs at least {dim_delay + 2} time samples")
+    lifted = delay_embed(x, dim_delay)
+    current, future = lifted[:, :-1], lifted[:, 1:]
+    control = u[:, dim_delay - 1 : -1]
+    input_rank = _rank(
+        jnp.linalg.svd(jnp.vstack([current, control]), full_matrices=False)[1],
+        low_rank,
+    )
+    output_rank = _rank(jnp.linalg.svd(future, full_matrices=False)[1], low_rank)
+    return _dmdc(current, future, control, input_rank, output_rank)
+
+
+def compute_edmdwc(
+    x: jax.Array,
+    u: jax.Array,
+    dim_delay: int = 5,
+    low_rank: Rank = 0.9,
+) -> tuple[jax.Array, jax.Array]:
+    """Fit reduced EDMD with control in delay-coordinate state space."""
+    x = jnp.atleast_1d(x).flatten()
+    u = _aligned_control(x, u)
+    if x.size < dim_delay + 2:
+        raise ValueError(f"data needs at least {dim_delay + 2} time samples")
+    lifted = delay_embed(x, dim_delay)
+    current, future = lifted[:, :-1], lifted[:, 1:]
+    control = u[:, dim_delay - 1 : -1]
+    left, singular_values, _ = jnp.linalg.svd(current, full_matrices=False)
+    rank = _rank(singular_values, low_rank)
+    basis = left[:, :rank]
+    current_reduced, future_reduced = basis.T @ current, basis.T @ future
+    fit = future_reduced @ jnp.linalg.pinv(
+        jnp.vstack([current_reduced, control]), rtol=1e-4
+    )
+    return (
+        basis @ fit[:, :rank] @ basis.T,
+        basis @ fit[:, rank:],
+    )
+
+
+def generator_eigenvalues(eigenvalues: jax.Array, dt: float = 1.0) -> jax.Array:
+    """Convert discrete Koopman eigenvalues to infinitesimal-generator values."""
+    if not isinstance(dt, Real) or dt <= 0:
+        raise ValueError("dt must be positive")
+    eigenvalues = jnp.asarray(eigenvalues)
+    safe = jnp.where(jnp.abs(eigenvalues) < 1e-12, jnp.nan + 0j, eigenvalues)
+    return jnp.log(safe) / dt
+
+
+def generator_dominant_mode(
+    eigenvalues: jax.Array, dt: float = 1.0
+) -> dict[str, Any]:
+    """Summarize the slowest non-constant finite generator mode."""
+    if not isinstance(dt, Real) or dt <= 0:
+        raise ValueError("dt must be positive")
+    eigenvalues = jnp.asarray(eigenvalues)
+    finite = jnp.isfinite(eigenvalues.real) & jnp.isfinite(eigenvalues.imag)
+    valid = finite & (jnp.abs(eigenvalues) > 1e-10)
+    if not bool(jnp.any(valid)):
+        return {
+            "decay_rate": float("nan"),
+            "angular_freq": float("nan"),
+            "period": float("nan"),
+            "e_folding_time": float("nan"),
+            "is_oscillatory": False,
+            "dominant_idx": -1,
+        }
+    index = int(jnp.argmin(jnp.where(valid, jnp.abs(eigenvalues.real), jnp.inf)))
+    value = eigenvalues[index]
+    decay = float(value.real)
+    frequency = float(value.imag)
+    period = 2 * jnp.pi / abs(frequency) if abs(frequency) > 1e-10 else jnp.nan
+    period = float(period)
+    return {
+        "decay_rate": decay,
+        "angular_freq": frequency,
+        "period": period,
+        "e_folding_time": 1 / abs(decay) if decay < -1e-10 else float("nan"),
+        "is_oscillatory": bool(2 * dt <= period <= 1e4 * dt),
+        "dominant_idx": index,
+    }

kews/diagnostics.py

@@ -0,0 +1,116 @@
+"""Optional operator and spectrum diagnostics kept outside the main API."""
+
+from __future__ import annotations
+
+from typing import Any, Mapping, Optional, Union
+
+import jax
+
+from ._core import (
+    compute_dmdc as _dmdc,
+    compute_edmdwc as _edmdc,
+    fit_koopman_edmd as _fit_operator,
+    generator_dominant_mode as _dominant_mode,
+    generator_eigenvalues as _generator_eigenvalues,
+)
+from .observables import time_delay
+
+
+def fit_operator(
+    data: jax.Array,
+    observable: Optional[Any] = None,
+    rank: Union[int, float] = 0.9,
+    fit_kwargs: Optional[Mapping[str, Any]] = None,
+) -> dict[str, Any]:
+    """Fit an autonomous EDMD Koopman operator for diagnostics.
+
+    Args:
+        data: State trajectory shaped ``(time,)`` or ``(state, time)``.
+        observable: Observable from :mod:`kews.observables`; defaults to
+            ``time_delay()``.
+        rank: Retained SVD energy fraction or absolute rank.
+        fit_kwargs: Optional training arguments for learned observables.
+
+    Returns:
+        Dictionary containing ``K``, ``eigvals``, row-wise ``eigvecs``,
+        ``Psi_X``, ``Psi_Y``, and the retained rank ``r``.
+    """
+    return _fit_operator(
+        data,
+        time_delay() if observable is None else observable,
+        rank,
+        fit_kwargs,
+    )
+
+
+def dmdc(
+    data: jax.Array,
+    control: jax.Array,
+    n_delay: int = 5,
+    rank: Union[int, float] = 0.9,
+) -> tuple[jax.Array, jax.Array]:
+    """Fit Dynamic Mode Decomposition with control.
+
+    Args:
+        data: Scalar state trajectory shaped ``(time,)``.
+        control: Input trajectory shaped ``(time,)`` or ``(control, time)``.
+        n_delay: Delay-coordinate dimension for state.
+        rank: Retained SVD energy fraction or absolute rank.
+
+    Returns:
+        Tuple ``(A, B)`` for ``x_next = A @ x + B @ u`` in delay space.
+    """
+    return _dmdc(data, control, n_delay, rank)
+
+
+def edmdc(
+    data: jax.Array,
+    control: jax.Array,
+    n_delay: int = 5,
+    rank: Union[int, float] = 0.9,
+) -> tuple[jax.Array, jax.Array]:
+    """Fit reduced EDMD with control in delay-coordinate space.
+
+    Args:
+        data: Scalar state trajectory shaped ``(time,)``.
+        control: Input trajectory shaped ``(time,)`` or ``(control, time)``.
+        n_delay: Delay-coordinate dimension for state.
+        rank: Retained SVD energy fraction or absolute rank.
+
+    Returns:
+        Tuple ``(K, B)`` containing lifted state and control operators.
+    """
+    return _edmdc(data, control, n_delay, rank)
+
+
+def generator_eigenvalues(
+    eigenvalues: jax.Array, dt: float = 1.0
+) -> jax.Array:
+    """Convert discrete Koopman eigenvalues to generator eigenvalues.
+
+    Args:
+        eigenvalues: Complex eigenvalues of a discrete Koopman operator.
+        dt: Positive sampling interval.
+
+    Returns:
+        Complex generator values ``log(eigenvalues) / dt``. Values near zero
+        map to ``NaN``.
+    """
+    return _generator_eigenvalues(eigenvalues, dt)
+
+
+def dominant_mode(eigenvalues: jax.Array, dt: float = 1.0) -> dict[str, Any]:
+    """Summarize the slowest finite non-constant generator mode.
+
+    Args:
+        eigenvalues: Complex generator eigenvalues.
+        dt: Positive sampling interval used for oscillation bounds.
+
+    Returns:
+        Decay rate, angular frequency, period, e-folding time, oscillatory flag,
+        and selected eigenvalue index.
+    """
+    return _dominant_mode(eigenvalues, dt)
+
+
+__all__ = ["dmdc", "dominant_mode", "edmdc", "fit_operator", "generator_eigenvalues"]

kews/observables.py

@@ -0,0 +1,254 @@
+"""Observable factories for :func:`kews.resdmd`."""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Any, Callable, Optional, Sequence
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import optax
+
+from ._core import delay_embed, random_fourier_features
+
+
+def _matrix(x: jax.Array) -> jax.Array:
+    """Orient time-series input as ``(state, time)``."""
+    return jnp.atleast_2d(x)
+
+
+class _TimeDelay:
+    """Time-delay observable returned by :func:`kews.observables.time_delay`."""
+
+    def __init__(self, n_delay: int):
+        if n_delay < 1:
+            raise ValueError("n_delay must be >= 1")
+        self.n_delay = n_delay
+
+    def transform(self, x: jax.Array) -> jax.Array:
+        """Return delayed state columns."""
+        return delay_embed(x, self.n_delay)
+
+
+class _RFF:
+    """Random Fourier observable returned by :func:`kews.observables.rff`."""
+
+    def __init__(self, n_features: int, scale: float, seed: int):
+        if n_features < 1:
+            raise ValueError("n_features must be >= 1")
+        if scale <= 0:
+            raise ValueError("scale must be positive")
+        self.n_features = n_features
+        self.scale = scale
+        self.seed = seed
+
+    def transform(self, x: jax.Array) -> jax.Array:
+        """Return a fixed random cosine feature matrix."""
+        key = jax.random.PRNGKey(self.seed)
+        return random_fourier_features(x, self.n_features, self.scale, key)
+
+
+def _activation(name: str) -> Callable[[jax.Array], jax.Array]:
+    """Resolve a supported activation name."""
+    activations = {
+        "leaky_relu": nn.leaky_relu,
+        "relu": nn.relu,
+        "sigmoid": nn.sigmoid,
+        "tanh": nn.tanh,
+    }
+    if name not in activations:
+        raise ValueError(f"activation must be one of {tuple(activations)}")
+    return activations[name]
+
+
+class _AutoEncoder(nn.Module):
+    """Symmetric dense autoencoder with an exposed encoder.
+
+    Layer creation order preserves the paper's seeded parameter initialization.
+    """
+
+    input_dim: int
+    hidden_layers: tuple[int, ...]
+    bottleneck_dim: int
+    activation: Callable[[jax.Array], jax.Array]
+
+    @nn.compact
+    def encode(self, x: jax.Array) -> jax.Array:
+        """Return bottleneck features shaped ``(features, time)``."""
+        hidden = x.T
+        for index, width in enumerate(self.hidden_layers):
+            hidden = nn.Dense(width, use_bias=False)(hidden)
+            if index < len(self.hidden_layers) - 1:
+                hidden = self.activation(hidden)
+        return nn.Dense(self.bottleneck_dim, use_bias=False)(hidden).T
+
+    @nn.compact
+    def __call__(self, x: jax.Array) -> jax.Array:
+        """Reconstruct input shaped ``(state, time)``."""
+        hidden = self.encode(x).T
+        for index, width in enumerate(reversed(self.hidden_layers)):
+            hidden = nn.Dense(width, use_bias=False)(hidden)
+            if index < len(self.hidden_layers) - 1:
+                hidden = self.activation(hidden)
+        return nn.Dense(self.input_dim, use_bias=False)(hidden).T
+
+
+@lru_cache(maxsize=32)
+def _train_step(module: _AutoEncoder, lr: float):
+    """Reuse compilation across the paper's same-architecture ensemble fits."""
+    optimizer = optax.adamw(lr)
+
+    @jax.jit
+    def step(params: Any, state: Any, batch: jax.Array):
+        def loss(candidate: Any) -> jax.Array:
+            reconstruction = module.apply(candidate, batch.T)
+            return jnp.mean((batch.T - reconstruction) ** 2)
+
+        _, gradients = jax.value_and_grad(loss)(params)
+        updates, state = optimizer.update(gradients, state, params)
+        return optax.apply_updates(params, updates), state
+
+    return step, optimizer
+
+
+class _MLP:
+    """Trainable observable returned by :func:`kews.observables.mlp`."""
+
+    fit_on_first_window = True
+
+    def __init__(
+        self,
+        hidden_layers: Sequence[int],
+        bottleneck_dim: int,
+        activation: str,
+        seed: int,
+    ):
+        if not hidden_layers or any(width < 1 for width in hidden_layers):
+            raise ValueError("hidden_layers must contain positive widths")
+        if bottleneck_dim < 1:
+            raise ValueError("bottleneck_dim must be >= 1")
+        self.hidden_layers = tuple(hidden_layers)
+        self.bottleneck_dim = bottleneck_dim
+        self.activation = _activation(activation)
+        self.seed = seed
+        self._input_dim: Optional[int] = None
+        self._module: Optional[_AutoEncoder] = None
+        self._params: Any = None
+
+    def fit(
+        self,
+        x: jax.Array,
+        *,
+        lr: float = 1e-3,
+        epochs: int = 1000,
+        batch_size: int = 32,
+        key: Optional[jax.Array] = None,
+        shuffle: bool = True,
+    ) -> "_MLP":
+        """Train on time columns using AdamW and return this observable.
+
+        The seed controls both initialization and epoch shuffles. Repeated
+        calls continue from existing parameters with a fresh optimizer state.
+        ``resdmd`` supplies only the first analysis window to this method.
+        """
+        if lr <= 0 or epochs < 1 or batch_size < 1:
+            raise ValueError("lr, epochs, and batch_size must be positive")
+        x = _matrix(x)
+        if self._input_dim not in (None, x.shape[0]):
+            raise ValueError(
+                f"MLP was fitted with {self._input_dim} states, got {x.shape[0]}"
+            )
+        self._input_dim = x.shape[0]
+        self._module = _AutoEncoder(
+            self._input_dim,
+            self.hidden_layers,
+            self.bottleneck_dim,
+            self.activation,
+        )
+        key = jax.random.PRNGKey(self.seed) if key is None else key
+        if self._params is None:
+            self._params = self._module.init(jax.random.PRNGKey(self.seed), x)
+        params = self._params
+        step, optimizer = _train_step(self._module, lr)
+        state = optimizer.init(params)
+
+        samples = x.T
+        batch_size = min(batch_size, samples.shape[0])
+        for _ in range(epochs):
+            epoch_samples = samples
+            if shuffle:
+                key, shuffle_key = jax.random.split(key)
+                # Shuffle the original sample order, preserving the seeded training path.
+                epoch_samples = samples[jax.random.permutation(shuffle_key, samples.shape[0])]
+            for start in range(0, samples.shape[0], batch_size):
+                params, state = step(params, state, epoch_samples[start : start + batch_size])
+        self._params = params
+        return self
+
+    def transform(self, x: jax.Array) -> jax.Array:
+        """Encode a trajectory after fitting."""
+        x = _matrix(x)
+        if self._module is None or self._params is None:
+            raise RuntimeError("MLP must be fitted before transform")
+        if x.shape[0] != self._input_dim:
+            raise ValueError(
+                f"MLP was fitted with {self._input_dim} states, got {x.shape[0]}"
+            )
+        return self._module.apply(self._params, x, method=self._module.encode)
+
+
+def time_delay(n_delay: int = 5) -> _TimeDelay:
+    """Create a delay-coordinate observable.
+
+    Args:
+        n_delay: Number of delayed state copies in each feature vector.
+
+    Returns:
+        An observable accepted by :func:`kews.resdmd`.
+    """
+    return _TimeDelay(n_delay)
+
+
+# Common abbreviation for time-delay features.
+tdf = time_delay
+
+
+def rff(
+    n_features: int = 500, scale: float = 1.0, seed: int = 0
+) -> _RFF:
+    """Create a random Fourier feature observable.
+
+    Args:
+        n_features: Number of random cosine features.
+        scale: Positive RBF-kernel length scale.
+        seed: Random projection seed.
+
+    Returns:
+        An observable accepted by :func:`kews.resdmd`.
+    """
+    return _RFF(n_features, scale, seed)
+
+
+def mlp(
+    hidden_layers: Sequence[int] = (16, 32),
+    bottleneck_dim: int = 3,
+    activation: str = "sigmoid",
+    seed: int = 0,
+) -> _MLP:
+    """Create a learned autoencoder observable.
+
+    Args:
+        hidden_layers: Encoder widths; the decoder mirrors them in reverse.
+        bottleneck_dim: Number of learned observable features.
+        activation: ``"relu"``, ``"tanh"``, ``"sigmoid"``, or
+            ``"leaky_relu"``.
+        seed: Parameter initialization and training seed.
+
+    Returns:
+        A trainable observable accepted by :func:`kews.resdmd`.
+    """
+    return _MLP(hidden_layers, bottleneck_dim, activation, seed)
+
+
+__all__ = ["mlp", "rff", "tdf", "time_delay"]

kews-0.1.0.dist-info/METADATA

@@ -0,0 +1,173 @@
+Metadata-Version: 2.4
+Name: kews
+Version: 0.1.0
+Summary: Koopman operator early warning signals
+Author-email: Juan Nathaniel <jn2808@columbia.edu>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/juannat7/kews
+Project-URL: Issues, https://github.com/juannat7/kews/issues
+Project-URL: Paper, https://arxiv.org/abs/2608.14716
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE.txt
+Requires-Dist: flax==0.12.4
+Requires-Dist: jax==0.9.0.1
+Requires-Dist: jaxlib==0.9.0.1
+Requires-Dist: optax==0.2.6
+Provides-Extra: paper
+Requires-Dist: numpy==2.3.5; extra == "paper"
+Requires-Dist: scipy==1.16.3; extra == "paper"
+Requires-Dist: matplotlib==3.10.8; extra == "paper"
+Requires-Dist: scikit-learn==1.8.0; extra == "paper"
+Requires-Dist: pandas==3.0.0; extra == "paper"
+Requires-Dist: xarray==2026.1.0; extra == "paper"
+Requires-Dist: netCDF4==1.7.3; extra == "paper"
+Requires-Dist: cmocean==4.0.3; extra == "paper"
+Requires-Dist: requests==2.32.5; extra == "paper"
+Requires-Dist: tqdm==4.67.1; extra == "paper"
+Requires-Dist: jupyterlab==4.5.0; extra == "paper"
+Requires-Dist: nbconvert==7.16.6; extra == "paper"
+Requires-Dist: ipykernel==7.1.0; extra == "paper"
+Provides-Extra: dev
+Requires-Dist: build>=1.2; extra == "dev"
+Requires-Dist: twine>=6; extra == "dev"
+Dynamic: license-file
+
+# KEWS: Koopman Early Warning Signals
+
+**KEWS** (pronounced “cues,” as in cues to potential tipping) is a JAX package for Koopman-based early warning signals of bifurcation- and rate-induced tipping.
+
+<div align="center">
+
+[![Python](https://img.shields.io/badge/Python-3.11%2B-3776AB?logo=python&logoColor=white)](https://www.python.org/)
+[![JAX](https://img.shields.io/badge/JAX-0.9-000000)](https://docs.jax.dev/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/juannat7/kews/blob/main/LICENSE.txt)
+[![PyPI version](https://img.shields.io/pypi/v/kews)](https://pypi.org/project/kews/)
+[![PyPI downloads](https://img.shields.io/pypi/dm/kews)](https://pypi.org/project/kews/)
+[![arXiv](https://img.shields.io/badge/arXiv-2608.14716-b31b1b.svg)](https://arxiv.org/abs/2608.14716)
+
+</div>
+
+## Abstract
+
+Abrupt transitions in complex systems are often preceded by early warning signals. However, most indicators rely on the notion of critical slowing down and do not generally extend to rate-induced tipping where transitions can occur without local loss of stability. This is problematic in stochastic, nonautonomous systems where internal variability and time-varying variables interact to shape tipping onset.
+
+We use Koopman operator theory to develop a unified early warning framework for both bifurcation and rate-induced tipping in stochastic systems. Our approach builds on residual Koopman mode decomposition that measures discrepancies between dynamics and their finite-dimensional approximation, and extends it to the control setting by augmenting the observable space with time-varying control variables.
+
+In idealized examples, the resulting indicators recover expected signatures near bifurcation points and improve detection in rate-induced regimes where classical indicators fail. We further show that learned embeddings through deep learning outperform prescribed dictionaries, especially in a high-dimensional setting.
+
+Applied to simulations of the Atlantic Meridional Overturning Circulation, our Koopman-based indicators distinguish tipping from non-tipping trajectories and reveal interpretable spectral signatures prior to critical transition.
+
+## Installation
+
+Python 3.11 or newer is required. Install the library from PyPI:
+
+```bash
+pip install kews
+```
+
+To reproduce the paper or develop the code:
+
+```bash
+git clone https://github.com/juannat7/kews.git
+cd kews
+pip install -e ".[paper,dev]"
+```
+
+Reproduction instructions are in [examples](https://github.com/juannat7/kews/tree/main/examples) for the model systems and [experiments](https://github.com/juannat7/kews/tree/main/experiments) for the coupled AMOC and empirical cases.
+
+## Quickstart
+
+Compute rolling ResKMD residuals from a scalar time series:
+
+```python
+import jax.numpy as jnp
+from jax import random
+
+import kews
+from kews.observables import time_delay
+
+time = jnp.linspace(0, 10 * jnp.pi, 100)
+state = jnp.sin(time) + 0.01 * random.normal(random.key(7), time.shape)
+
+signal = kews.resdmd(
+    state,
+    window=50,                           # Window length in samples (or a fraction, e.g. 0.5)
+    observable=time_delay(n_delay=10),   # Time-delay dictionary with 10 coordinates
+    rank=0.9,                            # Retain 90% of the SVD energy
+)
+
+# One residual per time step; the first 49 entries are NaN.
+print(signal.shape)  # (100,)
+```
+
+## Control-aware residuals
+
+Pass a control time series aligned with the state:
+
+```python
+control = jnp.linspace(-1, 1, state.size)
+controlled_signal = kews.resdmd(
+    state,
+    window=50,
+    observable=time_delay(10),
+    control=control,
+)
+```
+
+## Observable choices
+
+Use fixed random Fourier features instead of delay coordinates:
+
+```python
+from kews.observables import rff
+
+random_signal = kews.resdmd(
+    state,
+    window=50,
+    observable=rff(n_features=100, scale=0.2, seed=7),
+)
+```
+
+The random basis stays fixed across windows. An autoencoder can learn the features:
+
+```python
+from kews.observables import mlp
+
+learned_signal = kews.resdmd(
+    state,
+    window=50,
+    observable=mlp(hidden_layers=(16, 8), bottleneck_dim=3, seed=7),
+    fit_kwargs={"epochs": 200, "batch_size": 32, "lr": 1e-3},
+)
+```
+
+Training uses only the first window; the encoder is then frozen. Create a new observable for each independent fit. Paper notebooks specify their full training settings.
+
+## Spectral diagnostics
+
+```python
+from kews import diagnostics
+
+fit = diagnostics.fit_operator(state, observable=time_delay(10), rank=0.9)
+generator = diagnostics.generator_eigenvalues(fit["eigvals"], dt=1.0)
+dominant = diagnostics.dominant_mode(generator, dt=1.0)
+```
+
+`dt` is the sampling interval in the units of the input data. The mode summary includes decay rate, oscillation period, and e-folding time.
+
+## Citation
+
+```bibtex
+@article{nathaniel2026koopman,
+  title={Koopman early warning signals for bifurcation and rate-induced tipping},
+  author={Nathaniel, Juan and Roesch, Carla and DeSantis, Derek and Kooloth, Parvathi and Fan, Hang and Lucarini, Valerio and Romanou, Anastasia and Gentine, Pierre},
+  journal={arXiv preprint arXiv:2608.14716},
+  year={2026}
+}
+```
+
+Code uses the [MIT license](https://github.com/juannat7/kews/blob/main/LICENSE.txt). External datasets retain their original terms.

kews-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+kews/__init__.py,sha256=WIc9QuyQNB71Iks-KQdfTOx61V5SVcbXWWv9Nvk8Fgg,4576
+kews/_core.py,sha256=gxT-5wfqMIfSXtCG-dp2AO5m-ytdw3qVxXPXFGOlhus,16130
+kews/diagnostics.py,sha256=sD1D6B3kO1rEwDzAZyxQb0VpD_wKvnaGI6g-RpdqR6c,3520
+kews/observables.py,sha256=jq_HlqwVXTJULD0Zbr5VtvaOmAc2WGZ5C1UL91pv3u4,8486
+kews-0.1.0.dist-info/licenses/LICENSE.txt,sha256=D5YkCzfzd33EgQ7LxPRC19TDIfk-K6EMCv6-feQ-6eg,1076
+kews-0.1.0.dist-info/METADATA,sha256=F5Gf9pHOjDTZkvX0mBIIXMT9e_mCtZA-XeRctQRNK6U,6763
+kews-0.1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+kews-0.1.0.dist-info/top_level.txt,sha256=0ENNba47A1oAS2t06IstrNB_TvoPVKq_T6lRHGrrdDo,5
+kews-0.1.0.dist-info/RECORD,,

kews-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

kews-0.1.0.dist-info/licenses/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 KEWS Developer Team
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

kews-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+kews