vizflow 0.5.3__py3-none-any.whl → 0.5.5__py3-none-any.whl

vizflow/__init__.py

@@ -5,7 +5,7 @@ Usage:
     import vizflow as vf
 """
 
-__version__ = "0.5.3"
+__version__ = "0.5.5"
 
 from .config import Config, get_config, set_config
 from .io import (
@@ -14,9 +14,11 @@ from .io import (
     scan_alphas,
     scan_trade,
     scan_trades,
+    scan_univ,
+    scan_univs,
 )
 from .market import CN, CRYPTO, Market, Session
-from .ops import aggregate, bin, forward_return, parse_time
+from .ops import aggregate, bin, forward_return, mark_to_close, parse_time, sign_by_side
 from .schema_evolution import (
     JYAO_V20251114,
     SCHEMAS,

vizflow/config.py

@@ -62,6 +62,8 @@ class Config:
     alpha_pattern: str = "alpha_{date}.feather"
     trade_dir: Path | None = None
     trade_pattern: str = "trade_{date}.feather"
+    univ_dir: Path | None = None
+    univ_pattern: str = "{date}.csv"
     calendar_path: Path | None = None
 
     # === Output Paths ===
@@ -75,6 +77,7 @@ class Config:
     # Can be a string (schema name) or SchemaEvolution instance
     trade_schema: str | SchemaEvolution | None = None
     alpha_schema: str | SchemaEvolution | None = None
+    univ_schema: str | SchemaEvolution | None = None
 
     # === Aggregation ===
     binwidths: dict[str, float] = field(default_factory=dict)
@@ -87,13 +90,16 @@ class Config:
     def __post_init__(self):
         """Convert string paths to Path objects.
 
-        Note: String values for path fields (alpha_dir, trade_dir, calendar_path,
-        replay_dir, aggregate_dir) are automatically converted to Path objects.
+        Note: String values for path fields (alpha_dir, trade_dir, univ_dir,
+        calendar_path, replay_dir, aggregate_dir) are automatically converted
+        to Path objects.
         """
         if isinstance(self.alpha_dir, str):
             self.alpha_dir = Path(self.alpha_dir)
         if isinstance(self.trade_dir, str):
             self.trade_dir = Path(self.trade_dir)
+        if isinstance(self.univ_dir, str):
+            self.univ_dir = Path(self.univ_dir)
         if isinstance(self.calendar_path, str):
             self.calendar_path = Path(self.calendar_path)
         if isinstance(self.replay_dir, str):
@@ -135,6 +141,23 @@ class Config:
             raise ValueError("trade_dir is not set in Config")
         return self.trade_dir / self.trade_pattern.format(date=date)
 
+    def get_univ_path(self, date: str) -> Path:
+        """Get universe file path for a date.
+
+        Args:
+            date: Date string, e.g. "20241001"
+
+        Returns:
+            Full path to univ file
+
+        Raises:
+            ValueError: If univ_dir is not set or date format is invalid
+        """
+        _validate_date(date)
+        if self.univ_dir is None:
+            raise ValueError("univ_dir is not set in Config")
+        return self.univ_dir / self.univ_pattern.format(date=date)
+
     def get_replay_path(self, date: str, suffix: str = ".parquet") -> Path:
         """Get replay output file path for a date (FIFO results).
 

vizflow/io.py

@@ -2,12 +2,42 @@
 
 from __future__ import annotations
 
+import re
+from pathlib import Path
+
 import polars as pl
 
 from .config import Config, get_config
 from .schema_evolution import SchemaEvolution, get_schema
 
 
+def _extract_date_from_path(path: Path, pattern: str) -> str | None:
+    """Extract date from filename using pattern.
+
+    Args:
+        path: File path.
+        pattern: Pattern with {date} placeholder, e.g., "{date}.meords"
+
+    Returns:
+        Extracted date string, or None if no match or no {date} in pattern.
+
+    Example:
+        >>> _extract_date_from_path(Path("data/11110101.meords"), "{date}.meords")
+        "11110101"
+    """
+    # If pattern has no {date} placeholder, return None
+    if "{date}" not in pattern:
+        return None
+
+    # Convert pattern to regex: "{date}" -> "(?P<date>.+)"
+    # Escape other special chars first
+    regex_pattern = re.escape(pattern).replace(r"\{date\}", r"(?P<date>[^/]+)")
+    match = re.search(regex_pattern, path.name)
+    if match:
+        return match.group("date")
+    return None
+
+
 def _resolve_schema(
     schema_ref: str | SchemaEvolution | None,
 ) -> SchemaEvolution | None:
@@ -141,13 +171,15 @@ def scan_trade(date: str, config: Config | None = None) -> pl.LazyFrame:
 
 
 def scan_trades(config: Config | None = None) -> pl.LazyFrame:
-    """Scan all trade files with schema evolution.
+    """Scan all trade files with schema evolution and data_date column.
+
+    Extracts date from each filename using the pattern and adds a "data_date" column.
 
     Args:
         config: Config to use, or get_config() if None
 
     Returns:
-        LazyFrame with schema evolution applied
+        LazyFrame with schema evolution applied and data_date column added
 
     Raises:
         ValueError: If trade_dir is not set or no files found
@@ -159,7 +191,7 @@ def scan_trades(config: Config | None = None) -> pl.LazyFrame:
         ...     trade_schema="ylin_v20251204",
         ... )
         >>> vf.set_config(config)
-        >>> df = vf.scan_trades()
+        >>> df = vf.scan_trades()  # Has "data_date" column
     """
     config = config or get_config()
     if config.trade_dir is None:
@@ -172,14 +204,19 @@ def scan_trades(config: Config | None = None) -> pl.LazyFrame:
 
     schema = _resolve_schema(config.trade_schema)
 
-    # Concatenate all files using lazy scanning
-    dfs = [_scan_file(f, schema=schema) for f in files]
-    df = pl.concat(dfs)
+    # Scan each file, apply schema evolution, and add data_date column
+    # Schema evolution must be applied per-file BEFORE concat to ensure matching schemas
+    dfs = []
+    for f in files:
+        df = _scan_file(f, schema=schema)
+        if schema:
+            df = _apply_schema_evolution(df, schema)
+        date = _extract_date_from_path(f, config.trade_pattern)
+        if date:
+            df = df.with_columns(pl.lit(date).alias("data_date"))
+        dfs.append(df)
 
-    if schema:
-        df = _apply_schema_evolution(df, schema)
-
-    return df
+    return pl.concat(dfs)
 
 
 def scan_alpha(date: str, config: Config | None = None) -> pl.LazyFrame:
@@ -235,15 +272,87 @@ def scan_alphas(config: Config | None = None) -> pl.LazyFrame:
 
     schema = _resolve_schema(config.alpha_schema)
 
-    dfs = [_scan_file(f, schema=schema) for f in files]
-    df = pl.concat(dfs)
+    # Apply schema evolution per-file BEFORE concat to ensure matching schemas
+    # Note: Alpha files already have data_date column, no need to add it
+    dfs = []
+    for f in files:
+        df = _scan_file(f, schema=schema)
+        if schema:
+            df = _apply_schema_evolution(df, schema)
+        dfs.append(df)
+
+    return pl.concat(dfs)
+
+
+def scan_univ(date: str, config: Config | None = None) -> pl.LazyFrame:
+    """Scan single date universe file with schema evolution.
+
+    Args:
+        date: Date string, e.g. "20241001"
+        config: Config to use, or get_config() if None
+
+    Returns:
+        LazyFrame with schema evolution applied
+
+    Example:
+        >>> config = vf.Config(
+        ...     univ_dir=Path("/data/jyao/univ"),
+        ...     univ_pattern="{date}.csv",
+        ... )
+        >>> vf.set_config(config)
+        >>> df = vf.scan_univ("20241001")
+    """
+    config = config or get_config()
+    path = config.get_univ_path(date)
+    schema = _resolve_schema(config.univ_schema)
 
+    df = _scan_file(path, schema=schema)
     if schema:
         df = _apply_schema_evolution(df, schema)
 
     return df
 
 
+def scan_univs(config: Config | None = None) -> pl.LazyFrame:
+    """Scan all universe files with schema evolution and data_date column.
+
+    Extracts date from each filename using the pattern and adds a "data_date" column.
+
+    Args:
+        config: Config to use, or get_config() if None
+
+    Returns:
+        LazyFrame with schema evolution applied and data_date column added
+
+    Raises:
+        ValueError: If univ_dir is not set or no files found
+    """
+    config = config or get_config()
+    if config.univ_dir is None:
+        raise ValueError("univ_dir is not set in Config")
+
+    pattern = config.univ_pattern.replace("{date}", "*")
+    files = sorted(config.univ_dir.glob(pattern))
+    if not files:
+        raise ValueError(f"No files found matching {pattern} in {config.univ_dir}")
+
+    schema = _resolve_schema(config.univ_schema)
+
+    # Scan each file, apply schema evolution, and add data_date column
+    # Schema evolution must be applied per-file BEFORE concat to ensure matching schemas
+    dfs = []
+    for f in files:
+        df = _scan_file(f, schema=schema)
+        if schema:
+            df = _apply_schema_evolution(df, schema)
+        date = _extract_date_from_path(f, config.univ_pattern)
+        if date:
+            df = df.with_columns(pl.lit(date).alias("data_date"))
+        dfs.append(df)
+
+    return pl.concat(dfs)
+
+
 def load_calendar(config: Config | None = None) -> pl.DataFrame:
     """Load trading calendar.
 

vizflow/ops.py

@@ -235,3 +235,102 @@ def forward_return(
         ]).drop(["_forward_time", "_alpha_time", "_forward_price"])
 
     return trade.lazy()
+
+
+def mark_to_close(
+    df_trade: pl.LazyFrame,
+    df_univ: pl.LazyFrame,
+    mid_col: str = "mid",
+    close_col: str = "close",
+    symbol_col: str = "ukey",
+) -> pl.LazyFrame:
+    """Add mark-to-close return column.
+
+    Joins trade with universe data to get close price, then calculates:
+    y_close = (close - mid) / mid
+
+    Args:
+        df_trade: Trade LazyFrame with mid_col and symbol_col
+        df_univ: Universe LazyFrame with close_col and symbol_col
+        mid_col: Column name for mid price in trade df (default: "mid")
+        close_col: Column name for close price in univ df (default: "close")
+        symbol_col: Symbol column for joining (default: "ukey")
+
+    Returns:
+        Trade LazyFrame with y_close column added
+
+    Example:
+        >>> df_trade = vf.scan_trade(date)
+        >>> df_univ = vf.scan_univ(date)
+        >>> df = vf.mark_to_close(df_trade, df_univ)
+        >>> # Creates: y_close
+    """
+    # Select only needed columns from univ to avoid column conflicts
+    univ_cols = [symbol_col, close_col]
+    # Check if data_date exists in both for multi-day joining
+    trade_schema = df_trade.collect_schema()
+    univ_schema = df_univ.collect_schema()
+
+    if "data_date" in trade_schema.names() and "data_date" in univ_schema.names():
+        # Multi-day case: join on symbol and date
+        join_cols = [symbol_col, "data_date"]
+        univ_cols.append("data_date")
+    else:
+        # Single-day case: join on symbol only
+        join_cols = [symbol_col]
+
+    # Join with univ to get close price
+    df = df_trade.join(
+        df_univ.select(univ_cols),
+        on=join_cols,
+        how="left",
+    )
+
+    # Calculate return (guard against zero mid)
+    df = df.with_columns(
+        pl.when(pl.col(mid_col) != 0)
+        .then((pl.col(close_col) - pl.col(mid_col)) / pl.col(mid_col))
+        .otherwise(pl.lit(None))
+        .alias("y_close")
+    )
+
+    return df
+
+
+def sign_by_side(
+    df: pl.LazyFrame,
+    cols: list[str],
+    side_col: str = "order_side",
+) -> pl.LazyFrame:
+    """Sign return columns by order side.
+
+    For Buy trades: keep sign as-is (price going up = positive = good)
+    For Sell trades: negate sign (price going down = positive = good)
+
+    This makes all returns have consistent interpretation:
+    positive = favorable price move for that order side
+
+    Args:
+        df: LazyFrame with return columns
+        cols: List of column names to sign (e.g., ["y_10m", "y_30m", "y_close"])
+        side_col: Column containing order side (default: "order_side")
+            Expected values: "Buy" or "Sell"
+
+    Returns:
+        LazyFrame with signed return columns
+
+    Example:
+        >>> df = vf.sign_by_side(df, cols=["y_10m", "y_30m", "y_close"])
+        >>> # All y_* columns now have positive = favorable for that side
+    """
+    signed_exprs = []
+    for col in cols:
+        signed = (
+            pl.when(pl.col(side_col) == "Sell")
+            .then(-pl.col(col))
+            .otherwise(pl.col(col))
+            .alias(col)
+        )
+        signed_exprs.append(signed)
+
+    return df.with_columns(signed_exprs)

{vizflow-0.5.3.dist-info → vizflow-0.5.5.dist-info}/METADATA

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

vizflow-0.5.5.dist-info/RECORD

@@ -0,0 +1,10 @@
+vizflow/__init__.py,sha256=_jP6bp5U2iBFkPGTSLNwvOQay9XU_opNdrylF22iR9s,589
+vizflow/config.py,sha256=nPZPXlqQbaY8u_FAdtPShvb0mdx3e2TRaQ2CILzliAU,7192
+vizflow/io.py,sha256=qaILo7JYHRILZv7RgPdwltOkCew5ac3nmwOTn1p2ha4,11351
+vizflow/market.py,sha256=MtNz_nnZxC66Aq-i2PXEwaFCTknijFWYZUUv6798k2s,2493
+vizflow/ops.py,sha256=oR44HYKrfaXLh0SmbfXefl714UESSIC5lTNJBrR1kto,10858
+vizflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vizflow/schema_evolution.py,sha256=puwuuJ8LAzelHq1JTwQD968J7GYLvgAsCaHJnDHzu4U,14025
+vizflow-0.5.5.dist-info/METADATA,sha256=m6Zj_2AjsyBr1Tvq7uf20dzH9X9fNrFwpgM6wrDpcVg,388
+vizflow-0.5.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+vizflow-0.5.5.dist-info/RECORD,,

vizflow-0.5.3.dist-info/RECORD

@@ -1,10 +0,0 @@
-vizflow/__init__.py,sha256=pO3pTaKj7LNyuQlTrCHau9nBymhNGdmiImgJXH5lxt4,529
-vizflow/config.py,sha256=y4vRvjVTa1H5AdQf0q_XhYr-3EBDJst8BJq52ODN3uk,6456
-vizflow/io.py,sha256=eheqyLHGiSh69erxMk98FK-GYycbSheqkrIYRYGFy3A,7687
-vizflow/market.py,sha256=MtNz_nnZxC66Aq-i2PXEwaFCTknijFWYZUUv6798k2s,2493
-vizflow/ops.py,sha256=4UwxOTPhvZ1_4PI3pcxbXfLAYsn1Ecj6nyBtBBr7KS8,7761
-vizflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vizflow/schema_evolution.py,sha256=puwuuJ8LAzelHq1JTwQD968J7GYLvgAsCaHJnDHzu4U,14025
-vizflow-0.5.3.dist-info/METADATA,sha256=e8LTSGVMiHLNcTjMaAO11Lj5xYa6etsoG4yBuxI8E64,388
-vizflow-0.5.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-vizflow-0.5.3.dist-info/RECORD,,