扩展指标

2026-07-09 05:08:16 +08:00
commit 308c46ab9a
537 changed files with 152299 additions and 0 deletions
@@ -0,0 +1,251 @@
+"""
+ferro_ta.chunked — Chunked / out-of-core processing.
+====================================================
+
+Run ferro-ta indicators on data that is too large to fit in memory by
+processing it in overlapping chunks.  Each chunk contains a warm-up prefix
+(``overlap`` bars) from the previous chunk so that indicator state is
+correct.  After computing the indicator, the warm-up prefix is discarded and
+the resulting arrays are concatenated.
+
+Functions
+---------
+chunk_apply(fn, series, chunk_size, overlap, **fn_kwargs)
+    Run a single-input indicator function on a large series in chunks.
+
+make_chunk_ranges(n, chunk_size, overlap)
+    Return (start, end) index pairs for chunked processing.
+
+trim_overlap(chunk_out, overlap)
+    Discard the first *overlap* elements from an array.
+
+stitch_chunks(chunks)
+    Concatenate trimmed chunk outputs into one array.
+
+Rust backend
+------------
+    ferro_ta._ferro_ta.make_chunk_ranges
+    ferro_ta._ferro_ta.trim_overlap
+    ferro_ta._ferro_ta.stitch_chunks
+    ferro_ta._ferro_ta.chunk_apply_close_indicator
+
+Notes
+-----
+Indicators that rely on the full history (e.g. HT_TRENDLINE) cannot
+produce exact results in chunked mode; the approximation improves with
+larger ``overlap`` values.  Indicators with a finite look-back period
+(SMA, EMA, RSI, etc.) are exact when ``overlap >= timeperiod - 1``.
+
+For very large datasets or distributed execution, the optional Dask
+integration (``dask.dataframe.map_partitions``) can be used directly
+by passing any ferro-ta indicator function.  See the example in the
+docstring of ``chunk_apply``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+
+from ferro_ta._ferro_ta import (
+    chunk_apply_close_indicator as _rust_chunk_apply_close_indicator,
+)
+from ferro_ta._ferro_ta import (
+    make_chunk_ranges as _rust_make_chunk_ranges,
+)
+from ferro_ta._ferro_ta import (
+    stitch_chunks as _rust_stitch_chunks,
+)
+from ferro_ta._ferro_ta import (
+    trim_overlap as _rust_trim_overlap,
+)
+from ferro_ta._utils import _to_f64
+
+__all__ = [
+    "chunk_apply",
+    "make_chunk_ranges",
+    "trim_overlap",
+    "stitch_chunks",
+]
+
+_FASTPATH_DEFAULT_PERIODS: dict[str, int] = {
+    "SMA": 30,
+    "EMA": 30,
+    "RSI": 14,
+}
+
+
+def _resolve_chunk_fastpath(
+    fn: Callable[..., Any], fn_kwargs: dict[str, Any]
+) -> tuple[str, int] | None:
+    name = getattr(fn, "__name__", "").upper()
+    if name not in _FASTPATH_DEFAULT_PERIODS:
+        return None
+    if set(fn_kwargs) - {"timeperiod"}:
+        return None
+    raw = fn_kwargs.get("timeperiod", _FASTPATH_DEFAULT_PERIODS[name])
+    if not isinstance(raw, int):
+        return None
+    return name, int(raw)
+
+
+def make_chunk_ranges(
+    n: int,
+    chunk_size: int,
+    overlap: int,
+) -> NDArray[np.int64]:
+    """Compute start/end index pairs for chunked processing.
+
+    Parameters
+    ----------
+    n          : int — total length of the series
+    chunk_size : int — desired output bars per chunk (>= 1)
+    overlap    : int — warm-up bars prepended to each chunk (>= 0)
+
+    Returns
+    -------
+    numpy.ndarray of int64 with shape (n_chunks, 2) — each row is
+    ``[start_index, end_index)`` of the slice to pass to the indicator.
+
+    Examples
+    --------
+    >>> from ferro_ta.data.chunked import make_chunk_ranges
+    >>> make_chunk_ranges(10, 4, 2)
+    array([[ 0,  6],
+           [ 4, 10]])
+    """
+    raw = np.asarray(
+        _rust_make_chunk_ranges(int(n), int(chunk_size), int(overlap)),
+        dtype=np.int64,
+    )
+    if len(raw) == 0:
+        return raw.reshape(0, 2)
+    return raw.reshape(-1, 2)
+
+
+def trim_overlap(
+    chunk_out: ArrayLike,
+    overlap: int,
+) -> NDArray[np.float64]:
+    """Discard the first *overlap* elements from a chunk's indicator output.
+
+    Parameters
+    ----------
+    chunk_out : array-like — indicator output for a chunk
+    overlap   : int — number of leading warm-up elements to discard
+
+    Returns
+    -------
+    numpy.ndarray of float64 — the remaining elements
+    """
+    arr = np.ascontiguousarray(_to_f64(chunk_out))
+    return np.asarray(_rust_trim_overlap(arr, int(overlap)), dtype=np.float64)
+
+
+def stitch_chunks(
+    chunks: list[ArrayLike],
+) -> NDArray[np.float64]:
+    """Concatenate trimmed chunk outputs into a single array.
+
+    Parameters
+    ----------
+    chunks : list of array-like — trimmed indicator outputs
+
+    Returns
+    -------
+    numpy.ndarray of float64 — full concatenated result
+    """
+    converted = [np.ascontiguousarray(_to_f64(c)) for c in chunks]
+    return np.asarray(_rust_stitch_chunks(converted), dtype=np.float64)
+
+
+def chunk_apply(
+    fn: Callable[..., Any],
+    series: ArrayLike,
+    chunk_size: int = 10_000,
+    overlap: int = 100,
+    **fn_kwargs: Any,
+) -> NDArray[np.float64]:
+    """Run a 1-D indicator function on a large series in overlapping chunks.
+
+    Parameters
+    ----------
+    fn         : callable — indicator function with signature ``fn(series, **kwargs)``
+        that accepts a 1-D numpy array and returns a 1-D numpy array of the
+        same length.  Examples: ``ferro_ta.SMA``, ``ferro_ta.RSI``.
+    series     : array-like — the full (possibly large) input series
+    chunk_size : int — output bars per chunk (default 10 000).  Tune this
+        for memory/performance.
+    overlap    : int — warm-up bars prepended to each chunk (default 100).
+        Set to at least ``timeperiod - 1`` for the indicator to be accurate.
+    **fn_kwargs : extra keyword arguments forwarded to *fn* on every chunk.
+
+    Returns
+    -------
+    numpy.ndarray of float64 — full indicator output over the entire series.
+
+    Notes
+    -----
+    For Dask DataFrames, call ``dask.dataframe.map_partitions`` directly::
+
+        import dask.dataframe as dd
+        from ferro_ta import RSI
+
+        ddf = dd.from_pandas(pd.Series(close), npartitions=4)
+        result = ddf.map_partitions(lambda s: pd.Series(RSI(s.values)))
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from ferro_ta import SMA
+    >>> from ferro_ta.data.chunked import chunk_apply
+    >>> rng = np.random.default_rng(0)
+    >>> big_series = rng.standard_normal(50_000).cumsum() + 100
+    >>> out = chunk_apply(SMA, big_series, chunk_size=5000, overlap=30,
+    ...                   timeperiod=20)
+    >>> out.shape
+    (50000,)
+    """
+    s = _to_f64(series)
+    n = len(s)
+    if n == 0:
+        return np.empty(0, dtype=np.float64)
+
+    fastpath = _resolve_chunk_fastpath(fn, fn_kwargs)
+    if fastpath is not None:
+        indicator, timeperiod = fastpath
+        return np.asarray(
+            _rust_chunk_apply_close_indicator(
+                np.ascontiguousarray(s),
+                indicator,
+                int(timeperiod),
+                int(chunk_size),
+                int(overlap),
+            ),
+            dtype=np.float64,
+        )
+
+    ranges = make_chunk_ranges(n, chunk_size, overlap)
+    if len(ranges) == 0:
+        result = fn(s, **fn_kwargs)
+        return np.asarray(result, dtype=np.float64)
+
+    trimmed_chunks: list[NDArray[np.float64]] = []
+
+    for i, (start, end) in enumerate(ranges):
+        chunk = s[int(start) : int(end)]
+        result = fn(chunk, **fn_kwargs)
+        result_arr = np.asarray(result, dtype=np.float64)
+
+        # Determine how many leading bars to discard:
+        # - first chunk: keep everything (no prior overlap)
+        # - subsequent chunks: discard the leading `overlap` bars
+        discard = 0 if i == 0 else int(overlap)
+
+        trimmed = trim_overlap(result_arr, discard)
+        trimmed_chunks.append(trimmed)
+
+    return stitch_chunks(trimmed_chunks)  # type: ignore[arg-type]