Source code for memory_esn._speedups._loader

"""
Selects the compiled Cython kernel when importable, else the NumPy fallback.

Import errors from the compiled extensions (missing build, ABI mismatch, etc.)
are caught and downgraded to a one-time warning so the package still works.
"""

from __future__ import annotations

import warnings

import numpy as _np

from . import _fallback

# --- reservoir update ------------------------------------------------------
try:
    from .esn_reservoir_optimized import reservoir_update_fast as reservoir_update
    USING_CYTHON_RESERVOIR = True
except Exception:  # noqa: BLE001 - any import/ABI failure -> fallback
    reservoir_update = _fallback.reservoir_update
    USING_CYTHON_RESERVOIR = False

# Sparse reservoir update (large, highly-sparse reservoirs). The fused Cython
# CSR kernel does the recurrent sparse matvec inside one nogil loop -- no
# per-timestep Python overhead -- and benchmarks faster than a dense dgemv AND
# than scipy-per-step approaches. Falls back to the NumPy sparse matvec when the
# compiled extension is unavailable. Selected in BaseESN._compute_states when W
# is a scipy CSR matrix.
try:
    from .esn_reservoir_optimized import reservoir_update_sparse as _cython_sparse
except Exception:  # noqa: BLE001
    _cython_sparse = None




def reservoir_update_sparse(input_sequence, W_in, W_reservoir, activation="tanh",
                            leaky=1.0, bias=None, x0=None, noise_sequence=None):
    """Sparse-``W`` reservoir update; fused Cython CSR kernel when available."""
    if _cython_sparse is None:
        return _fallback.reservoir_update(
            input_sequence, W_in, W_reservoir, activation=activation,
            leaky=leaky, bias=bias, x0=x0, noise_sequence=noise_sequence)
    csr = W_reservoir.tocsr()
    return _cython_sparse(
        input_sequence=_np.ascontiguousarray(input_sequence, dtype=_np.float64),
        W_in=_np.ascontiguousarray(W_in, dtype=_np.float64),
        W_data=_np.ascontiguousarray(csr.data, dtype=_np.float64),
        W_indices=_np.ascontiguousarray(csr.indices, dtype=_np.int32),
        W_indptr=_np.ascontiguousarray(csr.indptr, dtype=_np.int32),
        activation=activation, leaky=leaky, bias=bias, x0=x0,
        noise_sequence=noise_sequence)


# --- fractional differencing ----------------------------------------------
try:
    from .fractional_diff import fractional_diff
    USING_CYTHON_FRACDIFF = True
except Exception:  # noqa: BLE001
    fractional_diff = _fallback.fractional_diff
    USING_CYTHON_FRACDIFF = False

# --- MODWT -----------------------------------------------------------------
try:
    from .modwt_fast import modwt
    USING_CYTHON_MODWT = True
except Exception:  # noqa: BLE001
    modwt = _fallback.modwt
    USING_CYTHON_MODWT = False


if not (USING_CYTHON_RESERVOIR and USING_CYTHON_FRACDIFF and USING_CYTHON_MODWT):
    _missing = [
        name
        for name, ok in (
            ("reservoir", USING_CYTHON_RESERVOIR),
            ("fractional_diff", USING_CYTHON_FRACDIFF),
            ("modwt", USING_CYTHON_MODWT),
        )
        if not ok
    ]
    warnings.warn(
        "memory_esn: using NumPy fallback for {} (compiled kernels not found). "
        "For best performance run: python setup.py build_ext --inplace".format(
            ", ".join(_missing)
        ),
        RuntimeWarning,
        stacklevel=2,
    )