vizflow 0.4.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

vizflow/__init__.py

@@ -5,7 +5,7 @@ Usage:
     import vizflow as vf
 """
 
-__version__ = "0.4.4"
+__version__ = "0.5.0"
 
 from .config import ColumnSchema, Config, get_config, set_config
 from .io import (
@@ -18,5 +18,5 @@ from .io import (
     scan_trades,
 )
 from .market import CN, CRYPTO, Market, Session
-from .ops import aggregate, bin, parse_time
+from .ops import aggregate, bin, forward_return, parse_time
 from .presets import JYAO_V20251114, PRESETS, YLIN_V20251204

vizflow/ops.py

@@ -135,3 +135,95 @@ def aggregate(
     """
     agg_exprs = [expr.alias(name) for name, expr in metrics.items()]
     return df.group_by(group_by).agg(agg_exprs)
+
+
+def _horizon_to_suffix(horizon_seconds: int) -> str:
+    """Convert horizon in seconds to column name suffix.
+
+    Rule: ≤60s → use seconds (60s), >60s → use minutes (3m, 30m)
+    """
+    if horizon_seconds <= 60:
+        return f"{horizon_seconds}s"
+    else:
+        minutes = horizon_seconds // 60
+        return f"{minutes}m"
+
+
+def forward_return(
+    df_trade: pl.LazyFrame,
+    df_alpha: pl.LazyFrame,
+    horizons: list[int],
+    trade_time_col: str = "elapsed_alpha_ts",
+    alpha_time_col: str = "elapsed_ticktime",
+    price_col: str = "mid",
+    symbol_col: str = "ukey",
+) -> pl.LazyFrame:
+    """Merge alpha's future price to trade and calculate forward returns.
+
+    For each trade row:
+    1. Look up alpha price at trade_time + horizon
+    2. Add forward_{price_col}_{horizon} column (the future price)
+    3. Calculate y_{horizon} = (forward_price - current_price) / current_price
+
+    Output column names follow the convention:
+    - ≤60s → forward_mid_60s, y_60s
+    - >60s → forward_mid_3m, y_3m
+
+    Args:
+        df_trade: Trade LazyFrame with trade_time_col and price_col
+        df_alpha: Alpha LazyFrame with alpha_time_col and price_col
+        horizons: List of horizon in seconds, e.g., [60, 180, 1800]
+        trade_time_col: Time column in trade df (default: "elapsed_alpha_ts")
+        alpha_time_col: Time column in alpha df (default: "elapsed_ticktime")
+        price_col: Column name for price in both dfs (default: "mid")
+        symbol_col: Symbol column for grouping (default: "ukey")
+
+    Returns:
+        Trade LazyFrame with forward_* and y_* columns added
+
+    Example:
+        >>> df_trade = vf.parse_time(vf.scan_trade(date), "alpha_ts")
+        >>> df_alpha = vf.parse_time(vf.scan_alpha(date), "ticktime")
+        >>> df = vf.forward_return(df_trade, df_alpha, horizons=[60, 180, 1800])
+        >>> # Creates: forward_mid_60s, forward_mid_3m, forward_mid_30m
+        >>> #          y_60s, y_3m, y_30m
+    """
+    # Collect for asof join
+    trade = df_trade.collect()
+    alpha = df_alpha.collect()
+
+    # Prepare alpha lookup table: (symbol, time) -> price
+    alpha_lookup = alpha.select([
+        pl.col(symbol_col),
+        pl.col(alpha_time_col),
+        pl.col(price_col),
+    ]).sort([symbol_col, alpha_time_col])
+
+    for horizon in horizons:
+        suffix = _horizon_to_suffix(horizon)
+        horizon_ms = horizon * 1000
+        forward_col = f"forward_{price_col}_{suffix}"
+        return_col = f"y_{suffix}"
+
+        # Add target time column for this horizon
+        trade = trade.with_columns(
+            (pl.col(trade_time_col) + horizon_ms).alias("_forward_time")
+        )
+
+        # Asof join: find alpha price at forward_time
+        joined = trade.join_asof(
+            alpha_lookup.rename({alpha_time_col: "_alpha_time", price_col: "_forward_price"}),
+            left_on="_forward_time",
+            right_on="_alpha_time",
+            by=symbol_col,
+            strategy="nearest",
+            tolerance=1000,  # 1 second tolerance
+        )
+
+        # Add forward price and calculate return
+        trade = joined.with_columns([
+            pl.col("_forward_price").alias(forward_col),
+            ((pl.col("_forward_price") - pl.col(price_col)) / pl.col(price_col)).alias(return_col),
+        ]).drop(["_forward_time", "_alpha_time", "_forward_price"])
+
+    return trade.lazy()

{vizflow-0.4.4.dist-info → vizflow-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vizflow
-Version: 0.4.4
+Version: 0.5.0
 Requires-Python: >=3.9
 Requires-Dist: polars>=0.20.0
 Provides-Extra: dev

{vizflow-0.4.4.dist-info → vizflow-0.5.0.dist-info}/RECORD

@@ -1,10 +1,10 @@
-vizflow/__init__.py,sha256=vYhGrFOQoONxBA6TJ08OQwWFd_UzltTsoZNQryNwnko,480
+vizflow/__init__.py,sha256=ZIMYQ-Yzh3eEAd7MSqNA00SlSuj45bEE6NsXM9Qc6O0,496
 vizflow/config.py,sha256=zSZnhdHzgXSqhDenHcHKm4CDGrMpKAdkNNRoUYYF1uc,6530
 vizflow/io.py,sha256=zmN0fFQOTmSBEBKangMExr0Q5mC2gajZM6GgdAyWkw4,8824
 vizflow/market.py,sha256=MtNz_nnZxC66Aq-i2PXEwaFCTknijFWYZUUv6798k2s,2493
-vizflow/ops.py,sha256=NL-Gtv-m_O1hv-0RUb9Wt43916HsQ5tYK_0e_uKR90w,4062
+vizflow/ops.py,sha256=6hKOjJowFrw1b6z4y8Liea9KTp8Fgy6kCGPZ6t15PVk,7426
 vizflow/presets.py,sha256=h91NZoOH4YAx0bbsaNigECf9WOcWh1QZavguunWkaLE,2452
 vizflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vizflow-0.4.4.dist-info/METADATA,sha256=NxMcTJ5fQKbB6GPak9dHARYTp9h0WwcDwxQAhrHcxRU,388
-vizflow-0.4.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-vizflow-0.4.4.dist-info/RECORD,,
+vizflow-0.5.0.dist-info/METADATA,sha256=cIIvBPZo2U6Sp46Wxgyu_tCVPWF4DGsgvapbavEBGl8,388
+vizflow-0.5.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+vizflow-0.5.0.dist-info/RECORD,,