"""Prepare the pre-reserve single-split close data for forecasting-v7.""" from __future__ import annotations import argparse import hashlib import json from dataclasses import dataclass from datetime import date from pathlib import Path from typing import Any import numpy as np ADAPTER_VERSION = "close-distribution-v3.0" REPO_ID = "tmmycruise/autoresearch-market-data" REVISION = "93ec44a9918e42fb5900af1690753d3bf175709e" SELECTION_ASSETS = ("AAPL", "ABBV", "MCD") CONFIRMATION_ASSETS = ("AMD", "MA") ALL_ASSETS = SELECTION_ASSETS + CONFIRMATION_ASSETS DEVELOPMENT_START = date(2016, 8, 8) RESEARCH_END = date(2023, 8, 8) SEALED_START = date(2023, 8, 8) SEALED_END = date(2026, 8, 8) TRAIN_END = date(2021, 8, 8) VALIDATION_END = date(2022, 8, 8) EVALUATION_HORIZONS = np.arange(2, 33, dtype=np.int16) CANONICAL_HORIZONS = np.asarray((2, 4, 8, 16, 32), dtype=np.int16) CLASS_COUNT = 21 MINUTE_NS = 60_000_000_000 BASE_TOKEN_FEATURE_NAMES = ( "close_log_return_since_last_observed_close", "close_return_observed", "log1p_elapsed_wall_clock_minutes_since_last_observed_close", ) CANDIDATE_TOKEN_FEATURE_NAMES = ( "regular_session_progress", "regular_session_progress_sin", "regular_session_progress_cos", "log1p_minutes_since_previous_session_last_observed_close", "log_raw_close", "close_log_return_since_session_first_observed", "close_log_return_since_previous_session_last_observed", ) TOKEN_FEATURE_NAMES = BASE_TOKEN_FEATURE_NAMES + CANDIDATE_TOKEN_FEATURE_NAMES DISCOVERY_FOLDS = { "single_split": { "train": ("2016-08-08", "2021-08-08"), "validation": ("2021-08-08", "2022-08-08"), }, } DEVELOPMENT_TEST = ("2022-08-08", "2023-08-08") @dataclass(frozen=True) class Session: day: date open_ns: int close_ns: int def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _stable_hash(value: Any) -> str: payload = json.dumps(value, sort_keys=True, separators=(",", ":")).encode() return hashlib.sha256(payload).hexdigest() def _save_npz(path: Path, arrays: dict[str, np.ndarray]) -> None: path.parent.mkdir(parents=True, exist_ok=True) np.savez_compressed(path, **arrays) def _timestamp_ns(value: Any) -> int: if value.tzinfo is None: raise ValueError("session timestamp must be timezone-aware") return int(value.timestamp() * 1_000_000_000) def _read_sessions(path: Path) -> list[Session]: import pyarrow.dataset as ds table = ds.dataset(str(path), format="parquet").to_table( columns=[ "session_date", "open_utc", "close_utc", "regular_minutes", ], filter=( (ds.field("session_date") >= DEVELOPMENT_START) & (ds.field("session_date") < RESEARCH_END) & (ds.field("calendar") == "XNYS") ), ) sessions: list[Session] = [] for row in table.to_pylist(): open_ns = _timestamp_ns(row["open_utc"]) close_ns = _timestamp_ns(row["close_utc"]) expected = (close_ns - open_ns) // MINUTE_NS if expected != int(row["regular_minutes"]): raise ValueError( f"calendar duration mismatch on {row['session_date']}" ) if expected < 180 or expected > 390: raise ValueError( f"invalid regular-session length on {row['session_date']}" ) sessions.append( Session( day=row["session_date"], open_ns=open_ns, close_ns=close_ns, ) ) sessions.sort(key=lambda item: item.day) if not sessions or sessions[-1].day >= RESEARCH_END: raise ValueError("session input crossed the research boundary") return sessions def _read_bars(path: Path, ticker: str) -> dict[str, np.ndarray]: import pyarrow.compute as pc import pyarrow.dataset as ds table = ds.dataset( str(path), format="parquet", partitioning=None, ).to_table( columns=[ "ticker", "window_start_ns", "close", "source_date", "adjusted", ], filter=( (ds.field("ticker") == ticker) & (ds.field("source_date") >= DEVELOPMENT_START) & (ds.field("source_date") < RESEARCH_END) ), ) if table.num_rows == 0: raise ValueError(f"no research bars found for {ticker}") if pc.any(table["adjusted"]).as_py(): raise ValueError(f"{ticker} source unexpectedly contains adjusted rows") if pc.max(table["source_date"]).as_py() >= RESEARCH_END: raise ValueError(f"{ticker} bars crossed the research boundary") arrays = { name: table[name].combine_chunks().to_numpy(zero_copy_only=False) for name in ("window_start_ns", "close", "source_date") } order = np.argsort(arrays["window_start_ns"], kind="stable") arrays = {name: values[order] for name, values in arrays.items()} timestamp = arrays["window_start_ns"].astype(np.int64, copy=False) if np.any(timestamp[1:] == timestamp[:-1]): raise ValueError(f"{ticker} contains duplicate minute timestamps") close = arrays["close"].astype(np.float64, copy=False) if np.any(~np.isfinite(close)) or np.any(close <= 0.0): raise ValueError(f"{ticker} contains invalid closes") return arrays def _read_splits(path: Path, ticker: str) -> list[dict[str, Any]]: import pyarrow.dataset as ds table = ds.dataset(str(path), format="parquet").to_table( columns=[ "id", "execution_date", "split_from", "split_to", ], filter=( (ds.field("ticker") == ticker) & (ds.field("execution_date") >= DEVELOPMENT_START) & (ds.field("execution_date") < RESEARCH_END) ), ) result = sorted( table.to_pylist(), key=lambda row: (row["execution_date"], str(row["id"])), ) for row in result: if float(row["split_from"]) <= 0 or float(row["split_to"]) <= 0: raise ValueError(f"invalid {ticker} split event: {row}") return result def _split_adjusted_close( close: np.ndarray, source_date: np.ndarray, splits: list[dict[str, Any]], ) -> np.ndarray: adjusted = np.asarray(close, dtype=np.float64).copy() for event in splits: factor = float(event["split_from"]) / float(event["split_to"]) adjusted[source_date < event["execution_date"]] *= factor if np.any(~np.isfinite(adjusted)) or np.any(adjusted <= 0.0): raise ValueError("split adjustment produced invalid closes") return adjusted def _align_sessions( sessions: list[Session], bars: dict[str, np.ndarray], splits: list[dict[str, Any]], ) -> dict[str, np.ndarray]: timestamp = bars["window_start_ns"].astype(np.int64, copy=False) adjusted_close = _split_adjusted_close( bars["close"], bars["source_date"], splits, ) parts: dict[str, list[np.ndarray]] = { "timestamp_ns": [], "close": [], "raw_close": [], "observed": [], "session_date": [], "minute_of_session": [], "session_length": [], } for session in sessions: expected = np.arange( session.open_ns, session.close_ns, MINUTE_NS, dtype=np.int64, ) positions = np.searchsorted(timestamp, expected) matched = positions < len(timestamp) matched[matched] &= ( timestamp[positions[matched]] == expected[matched] ) session_close = np.full(len(expected), np.nan, dtype=np.float64) raw_session_close = np.full(len(expected), np.nan, dtype=np.float64) session_close[matched] = adjusted_close[positions[matched]] raw_session_close[matched] = bars["close"][positions[matched]] valid = matched & np.isfinite(session_close) & (session_close > 0.0) day = np.datetime64(session.day.isoformat(), "D").astype(np.int32) parts["timestamp_ns"].append(expected) parts["close"].append(session_close) parts["raw_close"].append(raw_session_close) parts["observed"].append(valid) parts["session_date"].append( np.full(len(expected), day, dtype=np.int32) ) parts["minute_of_session"].append( np.arange(len(expected), dtype=np.int16) ) parts["session_length"].append( np.full(len(expected), len(expected), dtype=np.int16) ) return { name: np.concatenate(values) for name, values in parts.items() } def _token_features( grid: dict[str, np.ndarray], ) -> tuple[np.ndarray, np.ndarray]: timestamp = grid["timestamp_ns"] close = grid["close"] raw_close = grid["raw_close"] observed = grid["observed"] features = np.zeros( (len(timestamp), len(TOKEN_FEATURE_NAMES)), dtype=np.float32, ) volatility = np.full(len(timestamp), np.nan, dtype=np.float32) previous_observed = -1 recent_returns: list[float] = [] current_session = None session_first_close = np.nan session_last_close = np.nan previous_session_last_close = np.nan previous_session_last_timestamp = -1 session_gap = 0.0 for index in range(len(timestamp)): session = int(grid["session_date"][index]) if current_session != session: if current_session is not None and np.isfinite(session_last_close): previous_session_last_close = session_last_close previous_session_last_timestamp = int( timestamp[previous_observed] ) current_session = session session_first_close = np.nan session_last_close = np.nan if previous_session_last_timestamp >= 0: gap_minutes = max( 1.0, ( timestamp[index] - previous_session_last_timestamp ) / MINUTE_NS, ) session_gap = np.log1p(gap_minutes) else: session_gap = 0.0 denominator = max(int(grid["session_length"][index]) - 1, 1) progress = float(grid["minute_of_session"][index]) / denominator angle = 2.0 * np.pi * progress features[index, 3] = progress features[index, 4] = np.sin(angle) features[index, 5] = np.cos(angle) features[index, 6] = session_gap if previous_observed >= 0: elapsed = max( 1.0, (timestamp[index] - timestamp[previous_observed]) / MINUTE_NS, ) features[index, 2] = np.log1p(elapsed) if observed[index] and previous_observed >= 0: value = np.log(close[index] / close[previous_observed]) features[index, 0] = value features[index, 1] = 1.0 features[index, 7] = np.log(raw_close[index]) if not np.isfinite(session_first_close): session_first_close = close[index] features[index, 8] = np.log( close[index] / session_first_close ) if np.isfinite(previous_session_last_close): features[index, 9] = np.log( close[index] / previous_session_last_close ) session_last_close = close[index] recent_returns.append(float(value)) if len(recent_returns) > 32: recent_returns.pop(0) if len(recent_returns) >= 8: volatility[index] = np.sqrt( np.mean(np.square(recent_returns)) ) previous_observed = index elif observed[index]: features[index, 1] = 1.0 features[index, 7] = np.log(raw_close[index]) session_first_close = close[index] session_last_close = close[index] if np.isfinite(previous_session_last_close): features[index, 9] = np.log( close[index] / previous_session_last_close ) previous_observed = index elif previous_observed >= 0 and len(recent_returns) >= 8: volatility[index] = np.sqrt( np.mean(np.square(recent_returns)) ) if not np.all(np.isfinite(features)): raise ValueError("token features contain nonfinite values") return features, volatility def _targets(grid: dict[str, np.ndarray]) -> tuple[np.ndarray, np.ndarray]: close = grid["close"] observed = grid["observed"] minute = grid["minute_of_session"].astype(np.int64) length = grid["session_length"].astype(np.int64) y = np.full( (len(close), len(EVALUATION_HORIZONS)), np.nan, dtype=np.float64, ) mask = np.zeros_like(y, dtype=np.bool_) for horizon_index, horizon in enumerate( EVALUATION_HORIZONS.astype(np.int64) ): candidate = np.arange(len(close), dtype=np.int64) + horizon valid = ( observed & (minute + horizon < length) & (candidate < len(close)) ) rows = np.flatnonzero(valid) same_session = ( grid["session_date"][candidate[rows]] == grid["session_date"][rows] ) target_observed = observed[candidate[rows]] rows = rows[same_session & target_observed] y[rows, horizon_index] = np.log( close[candidate[rows]] / close[rows] ) mask[rows, horizon_index] = True return y, mask def _label_bundle( grid: dict[str, np.ndarray], y: np.ndarray, mask: np.ndarray, ) -> dict[str, np.ndarray]: selected = np.any(mask, axis=1) row_index = np.flatnonzero(selected).astype(np.int64) return { "row_index": row_index, "y": y[row_index], "target_mask": mask[row_index], "session_date": grid["session_date"][row_index], "minute_of_session": grid["minute_of_session"][row_index], "session_third": np.minimum( ( 3.0 * grid["minute_of_session"][row_index] / grid["session_length"][row_index] ).astype(np.int8), 2, ), } def _input_paths( study_dir: Path, ticker: str, boundary: dict[str, Any], ) -> tuple[Path, Path, Path]: asset = boundary["assets"][ticker] paths = {} for name in ("bars", "splits"): record = asset["outputs"][name] path = study_dir / record["path"] if _sha256(path) != record["sha256"]: raise ValueError(f"{ticker} {name} source hash mismatch") paths[name] = path session_record = boundary["sessions"] sessions = study_dir / session_record["path"] if _sha256(sessions) != session_record["sha256"]: raise ValueError("session source hash mismatch") return paths["bars"], sessions, paths["splits"] def _asset_output_dir(study_dir: Path, ticker: str) -> Path: if ticker in CONFIRMATION_ASSETS: return study_dir / "data" / "protected" / "confirmation" return study_dir / "data" / "runner" def _scope_configuration( study_dir: Path, scope: str, ) -> tuple[tuple[str, ...], Path, Path, dict[str, Any] | None]: if scope == "selection": return ( SELECTION_ASSETS, study_dir / "data" / "source" / "source-manifest.json", study_dir / "data" / "prepared-manifest.json", None, ) if scope != "confirmation": raise ValueError("scope must be selection or confirmation") marker_path = study_dir / "confirmation" / "CONFIRMATION_OPENED.json" if not marker_path.exists(): raise PermissionError( "confirmation preparation requires the frozen open marker" ) marker = json.loads(marker_path.read_text()) if marker.get("state") != "opened" or not marker.get("freeze_sha256"): raise ValueError("confirmation open marker is invalid") return ( CONFIRMATION_ASSETS, ( study_dir / "data" / "protected_source" / "confirmation-source-manifest.json" ), ( study_dir / "data" / "protected" / "confirmation" / "prepared-manifest.json" ), marker, ) def prepare( study_dir: Path, *, scope: str = "selection", ) -> dict[str, Any]: assets, boundary_path, manifest_path, marker = _scope_configuration( study_dir, scope, ) if not boundary_path.exists(): raise FileNotFoundError( f"forecasting-v4 {scope} isolated source is missing" ) boundary = json.loads(boundary_path.read_text()) if boundary.get("research_end_exclusive") != RESEARCH_END.isoformat(): raise ValueError("source manifest has the wrong research boundary") if boundary.get("revision") != REVISION: raise ValueError("source manifest revision differs from contract") if set(boundary.get("assets", {})) != set(assets): raise ValueError(f"source manifest has the wrong {scope} assets") derived_hashes: dict[str, str] = {} asset_metadata: dict[str, Any] = {} for ticker in assets: bars_path, sessions_path, splits_path = _input_paths( study_dir, ticker, boundary, ) sessions = _read_sessions(sessions_path) bars = _read_bars(bars_path, ticker) splits = _read_splits(splits_path, ticker) grid = _align_sessions(sessions, bars, splits) features, volatility = _token_features(grid) y, target_mask = _targets(grid) labels = _label_bundle(grid, y, target_mask) output_dir = _asset_output_dir(study_dir, ticker) slug = ticker.lower() feature_path = output_dir / f"{slug}_features.npz" label_path = output_dir / f"{slug}_labels.npz" _save_npz( feature_path, { "X": features, "timestamp_ns": grid["timestamp_ns"], "available_at_ns": grid["timestamp_ns"] + MINUTE_NS, "session_date": grid["session_date"], "minute_of_session": grid["minute_of_session"], "session_length": grid["session_length"], "causal_volatility_32": volatility, }, ) _save_npz(label_path, labels) for path in (feature_path, label_path): relative = str(path.relative_to(study_dir)) derived_hashes[relative] = _sha256(path) valid_target_rows = target_mask.sum(axis=0).astype(np.int64) target_times = [] for horizon_index, horizon in enumerate(EVALUATION_HORIZONS): rows = np.flatnonzero(target_mask[:, horizon_index]) target_times.append( int( ( grid["timestamp_ns"][rows] + (int(horizon) + 1) * MINUTE_NS ).max() ) ) asset_metadata[ticker] = { "confirmation": ticker in CONFIRMATION_ASSETS, "feature_rows": int(len(features)), "label_rows": int(len(labels["row_index"])), "maximum_anchor_available_at_ns": int( ( grid["timestamp_ns"][labels["row_index"]] + MINUTE_NS ).max() ), "maximum_target_available_at_ns": int(max(target_times)), "observed_token_fraction": float(np.mean(grid["observed"])), "split_events": [ { "execution_date": row["execution_date"].isoformat(), "id": str(row["id"]), "split_from": float(row["split_from"]), "split_to": float(row["split_to"]), } for row in splits ], "valid_target_rows_by_horizon": { str(int(horizon)): int(count) for horizon, count in zip( EVALUATION_HORIZONS, valid_target_rows, strict=True, ) }, } sealed_ns = int( np.datetime64(SEALED_START.isoformat(), "ns").astype(np.int64) ) research_end_ns = int( np.datetime64(RESEARCH_END.isoformat(), "ns").astype(np.int64) ) for ticker, values in asset_metadata.items(): for key in ( "maximum_anchor_available_at_ns", "maximum_target_available_at_ns", ): if values[key] >= research_end_ns or values[key] >= sealed_ns: raise ValueError( f"{ticker} {key} crossed the research boundary" ) source_snapshot = { "assets": list(assets), "boundary": boundary, "development_start": DEVELOPMENT_START.isoformat(), "repo_id": REPO_ID, "research_end_exclusive": RESEARCH_END.isoformat(), "revision": REVISION, "sealed_end_exclusive": SEALED_END.isoformat(), "sealed_start": SEALED_START.isoformat(), } metadata = { "adapter_version": ADAPTER_VERSION, "assets": asset_metadata, "canonical_horizons": CANONICAL_HORIZONS.tolist(), "class_count": CLASS_COUNT, "confirmation_assets": ( list(CONFIRMATION_ASSETS) if scope == "confirmation" else [] ), "derived_sha256": derived_hashes, "evaluation_horizons": EVALUATION_HORIZONS.tolist(), "folds": DISCOVERY_FOLDS, "scope": scope, "selection_assets": ( list(SELECTION_ASSETS) if scope == "selection" else [] ), "snapshot_sha256": _stable_hash( { "source": source_snapshot, "derived": derived_hashes, "evaluation_horizons": EVALUATION_HORIZONS.tolist(), "canonical_horizons": CANONICAL_HORIZONS.tolist(), "class_count": CLASS_COUNT, "token_features": list(TOKEN_FEATURE_NAMES), } ), "source_snapshot": source_snapshot, "token_feature_names": list(TOKEN_FEATURE_NAMES), } if marker is not None: selection_manifest = json.loads( (study_dir / "data" / "prepared-manifest.json").read_text() ) metadata["freeze_sha256"] = marker["freeze_sha256"] metadata["selection_snapshot_sha256"] = selection_manifest[ "snapshot_sha256" ] manifest_path.parent.mkdir(parents=True, exist_ok=True) manifest_path.write_text( json.dumps(metadata, indent=2, sort_keys=True) + "\n" ) return metadata def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--study-dir", type=Path, required=True) parser.add_argument( "--scope", choices=("selection", "confirmation"), default="selection", ) args = parser.parse_args() print(json.dumps(prepare(**vars(args)), sort_keys=True)) if __name__ == "__main__": main()