factorflow 0.1.0__py3-none-any.whl

factorflow/__init__.py

@@ -0,0 +1,59 @@
+"""factorflow: point-in-time factor pipelines that run the same in batch and live.
+
+Bars are rows; the newest row may be *open* (provisional) and is revised by ticks until it
+is sealed. Batch and incremental execution give bit-identical results for sealed rows.
+"""
+
+from factorflow.bars import BarState, Trades, build_time_bars
+from factorflow.batch import BatchEngine, BatchResult
+from factorflow.events import (
+    AdaptiveDaily,
+    AsofSampler,
+    EventBar,
+    EventBarBuilder,
+    EventBars,
+    EventPipeline,
+    Fixed,
+    build_event_bars,
+    sample_asof,
+    sample_panel,
+)
+from factorflow.graph import Pipeline, compile_pipeline
+from factorflow.grid import BarGrid, to_ns
+from factorflow.live import LiveEngine, Snapshot
+from factorflow.panel import Panel
+from factorflow.terms import BarData, Classifier, Column, Factor, Filter, Term, where
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "AdaptiveDaily",
+    "AsofSampler",
+    "BarData",
+    "BarGrid",
+    "BarState",
+    "BatchEngine",
+    "BatchResult",
+    "Classifier",
+    "Column",
+    "EventBar",
+    "EventBarBuilder",
+    "EventBars",
+    "EventPipeline",
+    "Factor",
+    "Filter",
+    "Fixed",
+    "LiveEngine",
+    "Panel",
+    "Pipeline",
+    "Snapshot",
+    "Term",
+    "Trades",
+    "build_event_bars",
+    "build_time_bars",
+    "compile_pipeline",
+    "sample_asof",
+    "sample_panel",
+    "to_ns",
+    "where",
+]

factorflow/bars.py

@@ -0,0 +1,246 @@
+"""Clock-time bar construction (trades -> OHLCV bars on a :class:`BarGrid`).
+
+One accumulator, :class:`BarState`, is shared by the batch builder and the live engine,
+so a bar built live from a trade stream is identical to the batch-built bar.
+
+Bar rules (docs/BARS.md):
+
+* a trade at ``ts`` belongs to the first bar with ``end > ts`` (bars are ``[start, end)``,
+  labeled by ``end``);
+* a bar with no trades carries the previous close forward: ``open = high = low = close =
+  previous close`` and ``volume = dollar = n_ticks = 0`` (NaN before an asset's first trade);
+* ``has_trade`` is ``n_ticks > 0``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+
+from factorflow.panel import Panel
+
+FIELDS = ("open", "high", "low", "close", "volume", "dollar", "n_ticks")
+
+
+@dataclass
+class Trades:
+    """A trade stream: parallel arrays sorted by ``ts`` (int64 ns UTC).
+
+    Parameters
+    ----------
+    asset : array of int
+        Index into the asset list passed to :func:`build_time_bars`.
+    ts : array of int64
+        Trade time, nanoseconds UTC, non-decreasing.
+    price : array of float
+        Trade price, in the same units the caller uses everywhere else.
+    size : array of float
+        Trade size. The sign is not special: volume bars add ``size`` as given.
+
+    Raises
+    ------
+    ValueError
+        If ``ts`` goes backwards.
+    """
+
+    asset: np.ndarray  # int asset index
+    ts: np.ndarray
+    price: np.ndarray
+    size: np.ndarray
+
+    def __post_init__(self) -> None:
+        """Cast the arrays and require non-decreasing timestamps.
+
+        Raises
+        ------
+        ValueError
+            If a timestamp is earlier than the one before it.
+        """
+        self.asset = np.asarray(self.asset, dtype=np.int64)
+        self.ts = np.asarray(self.ts, dtype=np.int64)
+        self.price = np.asarray(self.price, dtype=np.float64)
+        self.size = np.asarray(self.size, dtype=np.float64)
+        if self.ts.size > 1 and not np.all(np.diff(self.ts) >= 0):
+            raise ValueError("trades must be sorted by ts")
+
+    def __len__(self) -> int:
+        """Return the number of trades."""
+        return int(self.ts.size)
+
+
+class BarState:
+    """The open (provisional) bar for ``N`` assets.
+
+    Batch (:func:`build_time_bars`) and live (:class:`~factorflow.live.LiveEngine`)
+    both fold trades through this object, so a bar built either way is the same
+    object updated by the same rules.
+
+    Parameters
+    ----------
+    n : int
+        Number of assets.
+    """
+
+    def __init__(self, n: int):
+        """Allocate one open bar and start it from a missing previous close.
+
+        Parameters
+        ----------
+        n : int
+            Number of assets.
+        """
+        self.n = n
+        self.f = {k: np.full(n, np.nan) for k in FIELDS}
+        self.start(np.full(n, np.nan))
+
+    def start(self, prev_close: np.ndarray) -> None:
+        """Open a new bar from the previous close.
+
+        Parameters
+        ----------
+        prev_close : ndarray of shape (N,)
+            Close to carry forward. OHLC start at this value. Volume, dollar
+            value, and trade count start at 0. A NaN previous close stays NaN
+            until the first trade, which is an asset that has not traded yet.
+        """
+        pc = np.asarray(prev_close, dtype=np.float64)
+        for k in ("open", "high", "low", "close"):
+            self.f[k][:] = pc
+        for k in ("volume", "dollar", "n_ticks"):
+            self.f[k][:] = 0.0
+
+    def tick(self, j: int, price: float, size: float) -> None:
+        """Apply one trade to asset ``j``.
+
+        Parameters
+        ----------
+        j : int
+            Asset column.
+        price : float
+            Trade price. The first trade in the bar sets open, high, and low.
+            Later trades expand high and low and replace the close.
+        size : float
+            Added to volume. Dollar volume increases by ``price * size``.
+            The trade count increases by 1.
+        """
+        f = self.f
+        if f["n_ticks"][j] == 0:
+            f["open"][j] = f["high"][j] = f["low"][j] = price
+        else:
+            if price > f["high"][j]:
+                f["high"][j] = price
+            if price < f["low"][j]:
+                f["low"][j] = price
+        f["close"][j] = price
+        f["volume"][j] += size
+        f["dollar"][j] += price * size
+        f["n_ticks"][j] += 1
+
+    def set_bar(
+        self,
+        j: int,
+        open: float,
+        high: float,
+        low: float,
+        close: float,
+        volume: float,
+        dollar: float | None = None,
+        n_ticks: float | None = None,
+    ) -> None:
+        """Replace asset ``j``'s open bar with values the caller already aggregated.
+
+        Parameters
+        ----------
+        j : int
+            Asset column.
+        open, high, low, close, volume : float
+            The in-progress bar. Prices are in the caller's price units.
+            ``volume`` is in the caller's size units.
+        dollar : float or None, optional
+            Price times size accumulated so far. ``None`` stores ``close * volume``.
+        n_ticks : float or None, optional
+            Trade count. ``None`` stores 1 when ``volume > 0``, else 0.
+
+        Notes
+        -----
+        This is the path for a feed that publishes the current bar instead of
+        individual trades. It does not check that high and low bracket the close.
+        """
+        f = self.f
+        f["open"][j], f["high"][j], f["low"][j], f["close"][j] = open, high, low, close
+        f["volume"][j] = volume
+        f["dollar"][j] = close * volume if dollar is None else dollar
+        f["n_ticks"][j] = (1.0 if volume > 0 else 0.0) if n_ticks is None else n_ticks
+
+    def row(self, name: str) -> np.ndarray:
+        """Return one field of the open bar.
+
+        Parameters
+        ----------
+        name : str
+            ``"open"``, ``"high"``, ``"low"``, ``"close"``, ``"volume"``,
+            ``"dollar"``, ``"n_ticks"``, or ``"has_trade"``. ``has_trade`` is
+            ``n_ticks > 0`` and is not stored.
+
+        Returns
+        -------
+        ndarray of shape (N,)
+            A view of the field, except ``has_trade``, which is a new bool array.
+        """
+        if name == "has_trade":
+            return self.f["n_ticks"] > 0
+        return self.f[name]
+
+
+def build_time_bars(
+    trades: Trades,
+    assets: Sequence[str],
+    starts: np.ndarray,
+    ends: np.ndarray,
+    prev_close: np.ndarray | None = None,
+) -> Panel:
+    """Aggregate trades into clock bars.
+
+    Parameters
+    ----------
+    trades : Trades
+        Sorted by timestamp. ``trades.asset`` indexes ``assets``. A trade
+        belongs to the first bar with ``end > ts``.
+    assets : sequence of str
+        Asset names, length ``N``.
+    starts, ends : ndarray of int64
+        Bar bounds, nanoseconds UTC, from :meth:`BarGrid.bars`. Same length.
+    prev_close : ndarray of shape (N,) or None, optional
+        Close before the first bar. ``None`` starts from NaN, so an asset
+        with no trade in the first bars stays missing.
+
+    Returns
+    -------
+    Panel
+        Columns ``open``, ``high``, ``low``, ``close``, ``volume``, ``dollar``,
+        ``n_ticks``, and bool ``has_trade``. An empty bar repeats the previous
+        close and stores zero volume. ``times`` is ``ends``.
+
+    Notes
+    -----
+    The loop is the same :class:`BarState` the live engine ticks. Trades must
+    be in time order. Out-of-order trades are not repaired here.
+    """
+    ends = np.asarray(ends, dtype=np.int64)
+    t, n = len(ends), len(assets)
+    out = {k: np.empty((t, n)) for k in FIELDS}
+    st = BarState(n)
+    st.start(np.full(n, np.nan) if prev_close is None else prev_close)
+    bar_of = np.searchsorted(ends, trades.ts, side="right")
+    cuts = np.searchsorted(bar_of, np.arange(t + 1), side="left")
+    for r in range(t):
+        for i in range(cuts[r], cuts[r + 1]):
+            st.tick(int(trades.asset[i]), float(trades.price[i]), float(trades.size[i]))
+        for k in FIELDS:
+            out[k][r] = st.f[k]
+        st.start(st.f["close"].copy())
+    cols: dict[str, np.ndarray] = dict(out)
+    cols["has_trade"] = out["n_ticks"] > 0
+    return Panel(ends, tuple(assets), cols, np.asarray(starts, dtype=np.int64))

factorflow/batch.py

@@ -0,0 +1,221 @@
+"""Batch executor: run a compiled pipeline over a whole :class:`~factorflow.panel.Panel`.
+
+Batch and live share kernels. For a windowed term, this module prepends
+``window - 1`` missing rows (the warmup) and calls ``compute`` once on the whole
+history. The live engine instead calls the same ``compute`` on exactly ``window``
+rows and keeps row 0 of the result. The arithmetic is the same; only the slice
+of history differs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import polars as pl
+
+from factorflow.graph import Graph, Pipeline, compile_pipeline
+from factorflow.panel import Panel
+from factorflow.terms import Term, _Input
+
+
+def _pad(x: np.ndarray, rows: int, missing) -> np.ndarray:
+    """Prepend ``rows`` copies of ``missing`` so a window of length ``w`` can start at row 0.
+
+    The first ``w - 1`` output rows then see a window that contains the missing
+    value and come out missing. That is the warmup. ``rows`` is 0 for a term
+    whose window is 1.
+    """
+    if rows <= 0:
+        return x
+    pad = np.full((rows, *x.shape[1:]), missing, dtype=x.dtype)
+    return np.concatenate([pad, x], axis=0)
+
+
+def input_array(panel: Panel, t: Term) -> np.ndarray:
+    """Read one input term out of a panel.
+
+    Parameters
+    ----------
+    panel : Panel
+        History the pipeline is running on.
+    t : Term
+        An input term (a :class:`~factorflow.terms.Column` node). ``t.name``
+        is the column name.
+
+    Returns
+    -------
+    ndarray of shape (T, N)
+        Contiguous array in the term's dtype.
+
+    Raises
+    ------
+    KeyError
+        If the panel has no column of that name.
+    """
+    name = t.name  # type: ignore[attr-defined]
+    if name not in panel.columns:
+        raise KeyError(f"panel has no column {name!r} (needed by the pipeline)")
+    return np.ascontiguousarray(panel.columns[name], dtype=t.dtype)
+
+
+def compute_node(
+    t: Term, blocks: list[np.ndarray], prev: np.ndarray | None
+) -> tuple[np.ndarray, np.ndarray | None]:
+    """Evaluate one term on a whole history.
+
+    Parameters
+    ----------
+    t : Term
+        The node to run. Not an input term.
+    blocks : list of ndarray
+        One ``(T, N)`` array per input, already aligned to the panel.
+    prev : ndarray of shape (N,) or None
+        Previous output for a recursive term. Batch starts with ``None``,
+        so an EWMA starts at the first finite observation.
+
+    Returns
+    -------
+    values : ndarray of shape (T, N)
+        One output row per panel row.
+    last : ndarray of shape (N,) or None
+        The last row, kept as recursive state. ``None`` if the term is not
+        recursive or the history is empty.
+
+    Notes
+    -----
+    Row-wise terms (cross-section, EWMA) loop one row at a time, which is the
+    same loop the live engine runs for a single open row. Windowed terms are
+    padded by ``window - 1`` missing rows and then handed to ``compute``.
+    """
+    n_rows = blocks[0].shape[0]
+    if t.rowwise:
+        out = np.empty(blocks[0].shape, dtype=t.dtype)
+        p = prev
+        for r in range(n_rows):
+            row = t.compute_row([b[r : r + 1] for b in blocks], p)
+            out[r] = row
+            p = row if t.recursive else None
+        return out, (out[-1].copy() if (t.recursive and n_rows) else prev)
+    w = t.window
+    padded = [_pad(b, w - 1, inp.missing) for b, inp in zip(blocks, t.inputs, strict=True)]
+    return np.asarray(t.compute(padded), dtype=t.dtype), None
+
+
+@dataclass
+class BatchResult:
+    """Everything :meth:`BatchEngine.run` computed.
+
+    Parameters
+    ----------
+    times : ndarray of int64
+        Bar ends, nanoseconds UTC. Same as the panel.
+    assets : tuple of str
+        Asset names, column order.
+    values : dict of str to ndarray
+        Named pipeline outputs, each ``(T, N)``. These are the full arrays,
+        including rows a screen would drop and including warmup NaNs.
+    screen : ndarray of bool, shape (T, N), or None
+        The pipeline screen, or ``None`` if the pipeline has no screen.
+    nodes : dict of Term to ndarray
+        Every node's ``(T, N)`` values, including inputs. The live engine
+        warms its ring buffers from this.
+    recursive_state : dict of Term to ndarray or None
+        Last sealed row of each recursive term.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import factorflow as ff
+    >>> close = np.arange(6, dtype=float).reshape(3, 2)
+    >>> panel = ff.Panel(np.arange(3), ("A", "B"), {"close": close})
+    >>> out = ff.BatchEngine().run(ff.Pipeline({"lag": ff.Column("close").shift(1)}), panel)
+    >>> out.values["lag"][2].tolist()
+    [2.0, 3.0]
+    """
+
+    times: np.ndarray
+    assets: tuple[str, ...]
+    values: dict[str, np.ndarray]
+    screen: np.ndarray | None
+    nodes: dict[Term, np.ndarray]
+    recursive_state: dict[Term, np.ndarray | None]
+
+    def frame(self, open_last: bool = False) -> pl.DataFrame:
+        """Stack the screened outputs into a long frame.
+
+        Parameters
+        ----------
+        open_last : bool, optional
+            If ``True``, the last bar is flagged ``is_open``. Use that when the
+            last row was an open bar treated as final. Default is ``False``
+            (every row sealed).
+
+        Returns
+        -------
+        polars.DataFrame
+            Columns ``ts`` (UTC datetime, ns), ``asset``, ``is_open``, then
+            each output. Rows that fail the screen are omitted. With no screen,
+            every asset and bar is present.
+        """
+        t, n = len(self.times), len(self.assets)
+        keep = self.screen if self.screen is not None else np.ones((t, n), dtype=bool)
+        ti, ai = np.nonzero(keep)
+        is_open = np.zeros(ti.shape, dtype=bool)
+        if open_last and t:
+            is_open = ti == t - 1
+        data = {
+            "ts": pl.Series(self.times[ti]).cast(pl.Datetime("ns", "UTC")),
+            "asset": [self.assets[i] for i in ai],
+            "is_open": is_open,
+        }
+        for k, v in self.values.items():
+            data[k] = v[ti, ai]
+        return pl.DataFrame(data)
+
+
+class BatchEngine:
+    """Run a pipeline on a whole panel in one pass.
+
+    There is no constructor state. :meth:`run` compiles the pipeline if needed
+    and evaluates nodes in height order.
+    """
+
+    def run(self, pipeline: Pipeline | Graph, panel: Panel) -> BatchResult:
+        """Evaluate ``pipeline`` on every row of ``panel``.
+
+        Parameters
+        ----------
+        pipeline : Pipeline or Graph
+            A :class:`~factorflow.graph.Pipeline`, or the graph from
+            :func:`~factorflow.graph.compile_pipeline`.
+        panel : Panel
+            Must contain every input column the pipeline reads.
+
+        Returns
+        -------
+        BatchResult
+            Full ``(T, N)`` outputs. The first ``window - 1`` rows of a
+            windowed factor are missing.
+
+        Raises
+        ------
+        KeyError
+            If a required column is absent.
+        TypeError
+            If a pipeline column is a bare constant. Raised while compiling.
+        """
+        g = pipeline if isinstance(pipeline, Graph) else compile_pipeline(pipeline)
+        vals: dict[Term, np.ndarray] = {}
+        rec: dict[Term, np.ndarray | None] = {}
+        for t in g.order:
+            if isinstance(t, _Input):
+                vals[t] = input_array(panel, t)
+                continue
+            out, last = compute_node(t, [vals[i] for i in t.inputs], None)
+            vals[t] = out
+            if t.recursive:
+                rec[t] = last
+        outs = {k: vals[v] for k, v in g.outputs.items()}
+        screen = vals[g.screen] if g.screen is not None else None
+        return BatchResult(panel.times, panel.assets, outs, screen, vals, rec)