Yuxuan Yan
109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
"""CLI for daily feature computation."""
|
|
|
|
import os
|
|
|
|
import click
|
|
import pandas as pd
|
|
|
|
from pipeline.features.compute import compute_feature
|
|
from pipeline.features.registry import available_features, load_feature_module
|
|
|
|
|
|
@click.group(name="feature")
|
|
def feature():
|
|
"""Compute daily feature parquet files from minute bars."""
|
|
|
|
|
|
def _coerce(value: str):
|
|
"""Best-effort coercion of a CLI string to int, then float, else str."""
|
|
for cast in (int, float):
|
|
try:
|
|
return cast(value)
|
|
except ValueError:
|
|
continue
|
|
return value
|
|
|
|
|
|
def _parse_params(pairs: tuple[str, ...]) -> dict:
|
|
"""Parse repeated ``name=value`` options into a params dict."""
|
|
params: dict = {}
|
|
for pair in pairs:
|
|
if "=" not in pair:
|
|
raise click.BadParameter(f"--param must be name=value, got '{pair}'")
|
|
key, value = pair.split("=", 1)
|
|
params[key.strip()] = _coerce(value.strip())
|
|
return params
|
|
|
|
|
|
@feature.command("list")
|
|
@click.option(
|
|
"--feature-module", "feature_modules", multiple=True,
|
|
help="External module(s) to import first (dotted path or .py file)",
|
|
)
|
|
def list_(feature_modules):
|
|
"""List the registered feature types."""
|
|
for spec in feature_modules:
|
|
load_feature_module(spec)
|
|
for name in available_features():
|
|
click.echo(name)
|
|
|
|
|
|
@feature.command("compute")
|
|
@click.option("--minute-path", required=True, help="Path to minute parquet dataset/file")
|
|
@click.option("--daily-path", default=None, help="Optional daily data parquet for alignment")
|
|
@click.option("--feature-type", required=True, help="Registry key of the feature class")
|
|
@click.option("--feature-name", required=True, help="Name for this feature run/output file")
|
|
@click.option("--output-dir", default="features", help="Directory to save feature parquet")
|
|
@click.option(
|
|
"--feature-module", "feature_modules", multiple=True,
|
|
help="External module(s) to import so their features register (dotted path or .py file)",
|
|
)
|
|
@click.option(
|
|
"--param", "extra_params", multiple=True,
|
|
help="Extra feature constructor param as name=value (repeatable)",
|
|
)
|
|
def compute(
|
|
minute_path,
|
|
daily_path,
|
|
feature_type,
|
|
feature_name,
|
|
output_dir,
|
|
feature_modules,
|
|
extra_params,
|
|
):
|
|
"""Compute one daily feature file from raw minute bars."""
|
|
for spec in feature_modules:
|
|
load_feature_module(spec)
|
|
|
|
options = available_features()
|
|
if feature_type not in options:
|
|
raise click.BadParameter(
|
|
f"Unknown feature-type '{feature_type}'. Available: {options}. "
|
|
f"Use --feature-module to register an external feature.",
|
|
param_hint="--feature-type",
|
|
)
|
|
|
|
minute = pd.read_parquet(minute_path)
|
|
click.echo(f"Loaded minute bars: {len(minute):,} rows from {minute_path}")
|
|
|
|
daily = None
|
|
if daily_path:
|
|
daily = pd.read_parquet(daily_path)
|
|
click.echo(f"Loaded daily data: {len(daily):,} rows from {daily_path}")
|
|
|
|
result = compute_feature(
|
|
minute=minute,
|
|
daily=daily,
|
|
feature_type=feature_type,
|
|
**_parse_params(extra_params),
|
|
)
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
out_path = f"{output_dir}/{feature_name}.pq"
|
|
result.to_parquet(out_path, index=False)
|
|
feature_cols = [col for col in result.columns if col not in ("symbol_id", "date")]
|
|
click.echo(
|
|
f"Saved feature: {out_path} ({len(result):,} rows, "
|
|
f"{len(feature_cols)} columns)"
|
|
)
|