"""Train one forecasting-v7 seed and export its immutable development sheet.""" from __future__ import annotations import argparse import hashlib import json import os import platform import random import sys import time from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any, Iterable import numpy as np import torch from torch import nn from project.data.close_distribution_v3 import ( ADAPTER_VERSION, CANONICAL_HORIZONS, CLASS_COUNT, DEVELOPMENT_START, DEVELOPMENT_TEST, EVALUATION_HORIZONS, MINUTE_NS, RESEARCH_END, SELECTION_ASSETS, TOKEN_FEATURE_NAMES, TRAIN_END, VALIDATION_END, ) from project.data.reference_level_features import ( materialize_reference_features, ) from project.evaluators.close_distribution_v2 import ( ClassMap, class_map_to_arrays, encode_returns, fit_class_map, ) from project.experiments.forecasting_v4.runner import ( FEATURE_SETS, _context, _load_npz, _masked_nll, _normalization, _optimizer, _ordinal, ) from project.experiments.forecasting_v5.runner import ( _asset_balanced_loss_part, _balanced_epoch_chunks, ) from project.models.close_distribution_v4 import ( TypedAuxiliaryTransformer, build_typed_auxiliary_model, ) RUNNER_VERSION = "forecasting-v7.0" FEATURE_SET = "session_time" FEATURE_NAMES = FEATURE_SETS[FEATURE_SET] AUXILIARY_FEATURE_NAMES = ( "prior_5_session_net_return_bps", "prior_5_session_rms_observed_return_bps", "prior_5_session_range_width_bps", "prior_5_session_terminal_drawdown_bps", "prior_5_session_available", "prior_20_session_net_return_bps", "prior_20_session_rms_observed_return_bps", "prior_20_session_range_width_bps", "prior_20_session_terminal_drawdown_bps", "prior_20_session_available", "distance_to_session_open_bps", "distance_to_session_max_close_so_far_bps", "distance_to_session_min_close_so_far_bps", "session_range_position_so_far", "session_history_available", ) AUXILIARY_MASK_INDICES = (4, 9, 14) DEFAULT_EPOCHS = 16 DEFAULT_CONTEXT_LENGTH = 128 DEFAULT_TARGET_PARAMETERS = 25_000 DEFAULT_ROTARY_BASE = 16.0 DEFAULT_BATCH_SIZE = 512 DEFAULT_GRADIENT_ACCUMULATION = 8 DEFAULT_INFERENCE_BATCH_SIZE = 1_024 PROVENANCE_PATHS = ( "project/data/close_distribution_v3.py", "project/data/reference_level_features.py", "project/evaluators/close_distribution_v2.py", "project/experiments/forecasting_v4/runner.py", "project/experiments/forecasting_v5/runner.py", "project/experiments/forecasting_v7/runner.py", "project/models/close_distribution_v3.py", "project/models/close_distribution_v4.py", "project/studies/forecasting_v7.md", "requirements.txt", ) @dataclass class AssetPhaseData: asset: str asset_id: int features: dict[str, np.ndarray] bundle: dict[str, np.ndarray] center: np.ndarray scale: np.ndarray auxiliary_center: np.ndarray auxiliary_scale: np.ndarray auxiliary: np.ndarray class_map: ClassMap train_rows: np.ndarray train_labels: np.ndarray train_target_mask: np.ndarray validation_rows: np.ndarray validation_labels: np.ndarray validation_target_mask: np.ndarray 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 _configure(seed: int, threads: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.set_num_threads(threads) torch.use_deterministic_algorithms(True) def _selected_feature_indices() -> tuple[int, ...]: if tuple(FEATURE_NAMES) != ( "close_log_return_since_last_observed_close", "close_return_observed", "log1p_elapsed_wall_clock_minutes_since_last_observed_close", "regular_session_progress", "regular_session_progress_sin", "regular_session_progress_cos", ): raise ValueError("the frozen v7 sequence feature set changed") return tuple(TOKEN_FEATURE_NAMES.index(name) for name in FEATURE_NAMES) def _canonical_indices() -> tuple[int, ...]: lookup = { int(value): index for index, value in enumerate(EVALUATION_HORIZONS) } return tuple(lookup[int(value)] for value in CANONICAL_HORIZONS) def _load_selection_data( study_dir: Path, ) -> tuple[ dict[str, Any], dict[str, tuple[dict[str, np.ndarray], dict[str, np.ndarray]]], ]: manifest_path = study_dir / "data" / "prepared-manifest.json" manifest = json.loads(manifest_path.read_text()) if manifest.get("adapter_version") != ADAPTER_VERSION: raise ValueError("v7 prepared data has the wrong adapter version") if manifest.get("scope") != "selection": raise ValueError("v7 accepts only selection-scope prepared data") if tuple(manifest.get("selection_assets", ())) != SELECTION_ASSETS: raise ValueError("v7 selection assets changed") source = manifest.get("source_snapshot", {}) if source.get("research_end_exclusive") != RESEARCH_END.isoformat(): raise ValueError("v7 prepared data has the wrong research boundary") if source.get("sealed_start") != RESEARCH_END.isoformat(): raise ValueError("v7 sealed boundary is inconsistent") loaded = {} for asset in SELECTION_ASSETS: slug = asset.lower() feature_path = study_dir / "data" / "runner" / f"{slug}_features.npz" label_path = study_dir / "data" / "runner" / f"{slug}_labels.npz" for path in (feature_path, label_path): relative = str(path.relative_to(study_dir)) expected = manifest["derived_sha256"].get(relative) if expected is None or _sha256(path) != expected: raise ValueError(f"v7 prepared hash mismatch: {relative}") loaded[asset] = (_load_npz(feature_path), _load_npz(label_path)) return manifest, loaded def _raw_auxiliary(features: dict[str, np.ndarray]) -> np.ndarray: bundle = materialize_reference_features(features) result = np.zeros( (len(bundle.generic_values), len(AUXILIARY_FEATURE_NAMES)), dtype=np.float32, ) result[:, :10] = bundle.generic_values[:, 5:15] result[:, 10:15] = bundle.session_values if not np.isfinite(result).all(): raise ValueError("v7 auxiliary features contain nonfinite values") return result def _fit_auxiliary_scaler( values: np.ndarray, train_rows: np.ndarray, ) -> tuple[np.ndarray, np.ndarray]: center = np.zeros(values.shape[1], dtype=np.float64) scale = np.ones(values.shape[1], dtype=np.float64) selected = values[train_rows].astype(np.float64, copy=False) for window_index, mask_index in enumerate(AUXILIARY_MASK_INDICES): available = selected[:, mask_index] > 0.5 if not np.any(available): continue start = window_index * 5 for feature_index in range(start, start + 4): observed = selected[available, feature_index] median = float(np.median(observed)) q25, q75 = np.quantile(observed, (0.25, 0.75)) robust_scale = float((q75 - q25) / 1.349) if not np.isfinite(robust_scale) or robust_scale < 1e-8: robust_scale = float(np.std(observed)) if not np.isfinite(robust_scale) or robust_scale < 1e-8: robust_scale = 1.0 center[feature_index] = median scale[feature_index] = robust_scale return center.astype(np.float32), scale.astype(np.float32) def _scale_auxiliary( values: np.ndarray, center: np.ndarray, scale: np.ndarray, ) -> np.ndarray: result = np.zeros_like(values, dtype=np.float32) for window_index, mask_index in enumerate(AUXILIARY_MASK_INDICES): available = values[:, mask_index] > 0.5 start = window_index * 5 result[available, start : start + 4] = ( values[available, start : start + 4] - center[None, start : start + 4] ) / scale[None, start : start + 4] result[:, mask_index] = values[:, mask_index] if not np.isfinite(result).all(): raise ValueError("scaled v7 auxiliary features contain nonfinite values") return result def _date_mask( dates: np.ndarray, interval: tuple[str, str], ) -> np.ndarray: return ( (dates >= _ordinal(interval[0])) & (dates < _ordinal(interval[1])) ) def _prepare_phase_data( *, asset: str, asset_id: int, features: dict[str, np.ndarray], bundle: dict[str, np.ndarray], train_interval: tuple[str, str], validation_interval: tuple[str, str] | None, selected_indices: tuple[int, ...], canonical_indices: tuple[int, ...], max_train_rows_per_asset: int, ) -> AssetPhaseData: dates = bundle["session_date"] train = _date_mask(dates, train_interval) if not np.any(train): raise ValueError(f"{asset} has no v7 training rows") train_rows = bundle["row_index"][train] train_y = bundle["y"][train] train_target_mask = bundle["target_mask"][train] if ( max_train_rows_per_asset > 0 and len(train_rows) > max_train_rows_per_asset ): selected = np.linspace( 0, len(train_rows) - 1, max_train_rows_per_asset, dtype=np.int64, ) train_rows = train_rows[selected] train_y = train_y[selected] train_target_mask = train_target_mask[selected] class_map = fit_class_map( train_y, train_target_mask, class_count=CLASS_COUNT, ) all_train_labels = encode_returns( train_y, train_target_mask, class_map.edges, ) center, scale = _normalization( features, train_end=train_interval[1], selected_indices=selected_indices, ) raw_auxiliary = _raw_auxiliary(features) auxiliary_center, auxiliary_scale = _fit_auxiliary_scaler( raw_auxiliary, train_rows, ) auxiliary = _scale_auxiliary( raw_auxiliary, auxiliary_center, auxiliary_scale, ) if validation_interval is None: validation_rows = np.empty(0, dtype=np.int64) validation_labels = np.empty( (0, len(canonical_indices)), dtype=np.int16, ) validation_target_mask = np.empty( (0, len(canonical_indices)), dtype=np.bool_, ) else: validation = _date_mask(dates, validation_interval) validation_rows = bundle["row_index"][validation] validation_target_mask = bundle["target_mask"][validation][ :, canonical_indices ] validation_labels = encode_returns( bundle["y"][validation], bundle["target_mask"][validation], class_map.edges, )[:, canonical_indices] if not len(validation_rows): raise ValueError(f"{asset} has no v7 validation rows") return AssetPhaseData( asset=asset, asset_id=asset_id, features=features, bundle=bundle, center=center, scale=scale, auxiliary_center=auxiliary_center, auxiliary_scale=auxiliary_scale, auxiliary=auxiliary, class_map=class_map, train_rows=train_rows, train_labels=all_train_labels[:, canonical_indices], train_target_mask=train_target_mask[:, canonical_indices], validation_rows=validation_rows, validation_labels=validation_labels, validation_target_mask=validation_target_mask, ) def _build_model( checkpoint: dict[str, Any], *, device: torch.device, ) -> TypedAuxiliaryTransformer: if tuple(checkpoint.get("assets", ())) != SELECTION_ASSETS: raise ValueError("locked checkpoint has the wrong assets") if tuple(checkpoint.get("feature_names", ())) != tuple(FEATURE_NAMES): raise ValueError("locked checkpoint has the wrong sequence features") if tuple(checkpoint.get("supervised_horizons", ())) != tuple( int(value) for value in CANONICAL_HORIZONS ): raise ValueError("locked checkpoint has the wrong horizons") if checkpoint.get("asset_conditioning") != "early_add": raise ValueError("locked checkpoint is not early-add") if int(checkpoint.get("context_length", -1)) != DEFAULT_CONTEXT_LENGTH: raise ValueError("locked checkpoint has the wrong context length") if float(checkpoint.get("rotary_base", -1.0)) != DEFAULT_ROTARY_BASE: raise ValueError("locked checkpoint has the wrong RoPE base") model = build_typed_auxiliary_model( context_length=DEFAULT_CONTEXT_LENGTH, channels=len(FEATURE_NAMES) + 1, output_horizons=tuple( float(value) for value in CANONICAL_HORIZONS ), classes=CLASS_COUNT, target_parameters=DEFAULT_TARGET_PARAMETERS, rotary_base=DEFAULT_ROTARY_BASE, asset_count=len(SELECTION_ASSETS), asset_conditioning="early_add", auxiliary_dim=len(AUXILIARY_FEATURE_NAMES), ) incompatible = model.load_state_dict( checkpoint["state_dict"], strict=False, ) if set(incompatible.missing_keys) != {"auxiliary_affine.weight"}: raise ValueError( f"unexpected missing checkpoint keys: {incompatible.missing_keys}" ) if incompatible.unexpected_keys: raise ValueError( f"unexpected checkpoint keys: {incompatible.unexpected_keys}" ) return model.to(device) def _training_microbatch( data: tuple[AssetPhaseData, ...], asset_ids: np.ndarray, positions: np.ndarray, *, selected_indices: tuple[int, ...], ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, ]: tokens = [] padding = [] auxiliary = [] labels = [] masks = [] ids = [] for item in data: selected = asset_ids == item.asset_id if not np.any(selected): continue local = positions[selected] rows = item.train_rows[local] context, context_padding = _context( item.features["X"], rows, context_length=DEFAULT_CONTEXT_LENGTH, selected_indices=selected_indices, center=item.center, scale=item.scale, ) tokens.append(context) padding.append(context_padding) auxiliary.append(torch.from_numpy(item.auxiliary[rows])) labels.append( torch.from_numpy( item.train_labels[local].astype(np.int64, copy=False) ) ) masks.append( torch.from_numpy( item.train_target_mask[local].astype( np.float32, copy=False, ) ) ) ids.append( torch.full( (len(local),), item.asset_id, dtype=torch.long, ) ) return ( torch.cat(tokens), torch.cat(padding), torch.cat(auxiliary), torch.cat(labels), torch.cat(masks), torch.cat(ids), ) def _predict( model: TypedAuxiliaryTransformer, item: AssetPhaseData, *, rows: np.ndarray, selected_indices: tuple[int, ...], batch_size: int, device: torch.device, query_horizons: Iterable[int] | None, ) -> np.ndarray: horizons = ( tuple(int(value) for value in model.output_horizons.cpu().numpy()) if query_horizons is None else tuple(int(value) for value in query_horizons) ) result = np.empty( (len(rows), len(horizons), CLASS_COUNT), dtype=np.float32, ) query = ( None if query_horizons is None else torch.tensor(horizons, dtype=torch.float32, device=device) ) model.eval() with torch.inference_mode(): for start in range(0, len(rows), batch_size): stop = min(start + batch_size, len(rows)) local_rows = rows[start:stop] tokens, padding = _context( item.features["X"], local_rows, context_length=DEFAULT_CONTEXT_LENGTH, selected_indices=selected_indices, center=item.center, scale=item.scale, ) asset_id = torch.full( (stop - start,), item.asset_id, dtype=torch.long, device=device, ) logits = model( tokens.to(device), padding.to(device), asset_id, query, torch.from_numpy(item.auxiliary[local_rows]).to(device), ) result[start:stop] = torch.softmax(logits, dim=2).cpu().numpy() if not np.isfinite(result).all(): raise ValueError("v7 prediction contains nonfinite values") if not np.allclose(result.sum(axis=2), 1.0, atol=2e-6, rtol=0.0): raise ValueError("v7 probabilities do not sum to one") return result def _validation_nll( model: TypedAuxiliaryTransformer, data: tuple[AssetPhaseData, ...], *, selected_indices: tuple[int, ...], inference_batch_size: int, device: torch.device, ) -> tuple[float, dict[str, float]]: by_asset = {} for item in data: probabilities = _predict( model, item, rows=item.validation_rows, selected_indices=selected_indices, batch_size=inference_batch_size, device=device, query_horizons=None, ) by_asset[item.asset] = _masked_nll( probabilities, item.validation_labels, item.validation_target_mask, ) return float(np.mean(list(by_asset.values()))), by_asset def _train( model: TypedAuxiliaryTransformer, data: tuple[AssetPhaseData, ...], *, seed: int, epochs: int, batch_size: int, gradient_accumulation_steps: int, inference_batch_size: int, selected_indices: tuple[int, ...], device: torch.device, select_on_validation: bool, ) -> dict[str, Any]: optimizer = _optimizer(model, "adamw_constant") effective_batch_size = batch_size * gradient_accumulation_steps generator = np.random.default_rng(seed + 10_007) train_history = [] validation_history = [] best_epoch = 0 best_nll = float("inf") best_state: dict[str, torch.Tensor] | None = None updates_per_epoch = None total_updates = 0 for epoch in range(epochs): epoch_started = time.monotonic() model.train() updates = _balanced_epoch_chunks( tuple(len(item.train_rows) for item in data), effective_batch_size=effective_batch_size, generator=generator, ) if updates_per_epoch is None: updates_per_epoch = len(updates) elif updates_per_epoch != len(updates): raise AssertionError("v7 updates per epoch changed") epoch_loss_sum = np.zeros( (len(data), len(CANONICAL_HORIZONS)), dtype=np.float64, ) epoch_target_count = np.zeros_like( epoch_loss_sum, dtype=np.int64, ) for chunks in updates: update_asset_ids = np.concatenate( [ np.full(len(chunk), index, dtype=np.int64) for index, chunk in enumerate(chunks) ] ) update_positions = np.concatenate(chunks) denominators = np.stack( [ item.train_target_mask[chunks[index]].sum(axis=0) for index, item in enumerate(data) ] ).astype(np.float32) if np.any(denominators == 0): raise ValueError( "a v7 asset-balanced update has no target for a horizon" ) denominator_tensor = torch.from_numpy(denominators).to(device) optimizer.zero_grad(set_to_none=True) for start in range(0, len(update_positions), batch_size): stop = min(start + batch_size, len(update_positions)) ( tokens, padding, auxiliary, labels, target_mask, asset_id, ) = _training_microbatch( data, update_asset_ids[start:stop], update_positions[start:stop], selected_indices=selected_indices, ) tokens = tokens.to(device) padding = padding.to(device) auxiliary = auxiliary.to(device) labels = labels.to(device) target_mask = target_mask.to(device) asset_id = asset_id.to(device) logits = model( tokens, padding, asset_id, auxiliary=auxiliary, ) loss = _asset_balanced_loss_part( logits, labels, target_mask, asset_id, denominator_tensor, ) loss.backward() with torch.no_grad(): losses = nn.functional.cross_entropy( logits.reshape(-1, CLASS_COUNT), labels.clamp_min(0).reshape(-1), reduction="none", ).reshape(labels.shape) for asset_index in range(len(data)): selected = asset_id == asset_index epoch_loss_sum[asset_index] += ( losses[selected] * target_mask[selected] ).sum(dim=0).cpu().numpy() epoch_target_count[asset_index] += ( target_mask[selected] .sum(dim=0) .cpu() .numpy() .astype(np.int64) ) nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_updates += 1 train_nll = float( np.mean(epoch_loss_sum / epoch_target_count) ) train_history.append(train_nll) record: dict[str, Any] = { "duration_seconds": time.monotonic() - epoch_started, "epoch": epoch + 1, "train_nll": train_nll, } if select_on_validation: equal_asset_nll, by_asset = _validation_nll( model, data, selected_indices=selected_indices, inference_batch_size=inference_batch_size, device=device, ) record["equal_asset_validation_nll"] = equal_asset_nll record["validation_nll_by_asset"] = by_asset if equal_asset_nll < best_nll: best_nll = equal_asset_nll best_epoch = epoch + 1 best_state = { name: value.detach().cpu().clone() for name, value in model.state_dict().items() } validation_history.append(record) if select_on_validation: if best_state is None: raise AssertionError("v7 validation selected no checkpoint") model.load_state_dict(best_state) else: best_epoch = epochs if updates_per_epoch is None: raise AssertionError("v7 training produced no updates") return { "best_epoch": best_epoch, "best_equal_asset_validation_nll": ( best_nll if select_on_validation else None ), "completed_fixed_epoch_budget": len(train_history) == epochs, "epoch_history": validation_history, "train_nll": train_history, "total_updates": total_updates, "updates_per_epoch": updates_per_epoch, } def _raw_close_vector( study_dir: Path, manifest: dict[str, Any], item: AssetPhaseData, ) -> np.ndarray: import pyarrow.parquet as pq record = manifest["source_snapshot"]["boundary"]["assets"][ item.asset ]["outputs"]["bars"] path = study_dir / record["path"] if _sha256(path) != record["sha256"]: raise ValueError(f"{item.asset} raw source hash changed") table = pq.read_table( path, columns=["window_start_ns", "close"], ) timestamp = table["window_start_ns"].to_numpy(zero_copy_only=False) close = table["close"].to_numpy(zero_copy_only=False).astype(np.float64) order = np.argsort(timestamp, kind="stable") timestamp = timestamp[order] close = close[order] grid_timestamp = item.features["timestamp_ns"] positions = np.searchsorted(timestamp, grid_timestamp) matched = positions < len(timestamp) matched[matched] &= timestamp[positions[matched]] == grid_timestamp[matched] result = np.full(len(grid_timestamp), np.nan, dtype=np.float64) result[matched] = close[positions[matched]] observed = item.features["X"][:, 1] > 0.5 if np.any(observed & ~np.isfinite(result)): raise ValueError(f"{item.asset} observed row lacks a raw close") return result def _write_prediction_sheet( model: TypedAuxiliaryTransformer, data: tuple[AssetPhaseData, ...], *, study_dir: Path, manifest: dict[str, Any], output_path: Path, seed: int, selected_indices: tuple[int, ...], inference_batch_size: int, device: torch.device, max_development_rows_per_asset: int, ) -> dict[str, Any]: import pyarrow as pa import pyarrow.parquet as pq probability_names = tuple( f"p{index:02d}" for index in range(CLASS_COUNT) ) schema_fields = [ pa.field("stable_row_id", pa.uint64()), pa.field("seed", pa.int16()), pa.field("asset", pa.string()), pa.field("asset_id", pa.int8()), pa.field("anchor_row_index", pa.int64()), pa.field("anchor_timestamp_ns", pa.int64()), pa.field("available_at_ns", pa.int64()), pa.field("session_date", pa.date32()), pa.field("minute_of_session", pa.int16()), pa.field("horizon_minutes", pa.int16()), pa.field("target_valid", pa.bool_()), pa.field("target_timestamp_ns", pa.int64()), pa.field("target_available_at_ns", pa.int64()), pa.field("anchor_raw_close", pa.float64()), pa.field("target_raw_close", pa.float64()), pa.field("realized_log_return", pa.float64()), pa.field("target_class", pa.int16()), pa.field("expected_return", pa.float32()), pa.field("expected_positive_probability", pa.float32()), pa.field("entropy", pa.float32()), pa.field("max_probability", pa.float32()), pa.field("top_class", pa.int8()), ] schema_fields.extend( pa.field(name, pa.float32()) for name in probability_names ) schema = pa.schema(schema_fields) temporary = output_path.with_suffix(output_path.suffix + ".partial") if output_path.exists() or temporary.exists(): raise FileExistsError(f"prediction sheet already exists: {output_path}") output_path.parent.mkdir(parents=True, exist_ok=True) writer = pq.ParquetWriter( temporary, schema, compression="zstd", use_dictionary=("asset",), ) horizons = EVALUATION_HORIZONS.astype(np.int64) horizon_count = len(horizons) by_asset: dict[str, Any] = {} total_rows = 0 try: for item in data: test = _date_mask(item.bundle["session_date"], DEVELOPMENT_TEST) test_positions = np.flatnonzero(test) if ( max_development_rows_per_asset > 0 and len(test_positions) > max_development_rows_per_asset ): selected = np.linspace( 0, len(test_positions) - 1, max_development_rows_per_asset, dtype=np.int64, ) test_positions = test_positions[selected] rows = item.bundle["row_index"][test_positions] returns = item.bundle["y"][test_positions] target_mask = item.bundle["target_mask"][test_positions] target_labels = encode_returns( returns, target_mask, item.class_map.edges, ) raw_close = _raw_close_vector(study_dir, manifest, item) nll_sum = np.zeros(horizon_count, dtype=np.float64) nll_count = np.zeros(horizon_count, dtype=np.int64) asset_output_rows = 0 for start in range(0, len(rows), inference_batch_size): stop = min(start + inference_batch_size, len(rows)) local_rows = rows[start:stop] probabilities = _predict( model, item, rows=local_rows, selected_indices=selected_indices, batch_size=inference_batch_size, device=device, query_horizons=EVALUATION_HORIZONS, ) count = len(local_rows) flat_count = count * horizon_count flat_probabilities = probabilities.reshape( flat_count, CLASS_COUNT, ) local_mask = target_mask[start:stop] local_returns = returns[start:stop] local_labels = target_labels[start:stop] for horizon_index in range(horizon_count): valid = local_mask[:, horizon_index] if np.any(valid): selected_probability = probabilities[ valid, horizon_index, local_labels[valid, horizon_index], ] nll_sum[horizon_index] += float( -np.log( np.clip( selected_probability, 1e-12, 1.0, ) ).sum() ) nll_count[horizon_index] += int(np.sum(valid)) repeated_rows = np.repeat(local_rows, horizon_count) tiled_horizons = np.tile(horizons, count) repeated_mask = local_mask.reshape(-1) candidate = ( local_rows[:, None] + horizons[None, :] ).reshape(-1) safe_candidate = np.minimum( candidate, len(item.features["timestamp_ns"]) - 1, ) target_timestamp = item.features["timestamp_ns"][ safe_candidate ].astype(np.int64, copy=True) target_timestamp[~repeated_mask] = -1 target_close = raw_close[safe_candidate].copy() target_close[~repeated_mask] = np.nan anchor_close = np.repeat( raw_close[local_rows], horizon_count, ) flat_returns = local_returns.reshape(-1) flat_labels = local_labels.reshape(-1) expected_return = np.einsum( "nhc,hc->nh", probabilities.astype(np.float64), item.class_map.return_means, ).reshape(-1) expected_positive = np.einsum( "nhc,hc->nh", probabilities.astype(np.float64), item.class_map.positive_rates, ).reshape(-1) entropy = -np.sum( probabilities * np.log(np.clip(probabilities, 1e-12, 1.0)), axis=2, ).reshape(-1) stable = ( np.uint64(item.asset_id + 1) * np.uint64(1 << 56) + repeated_rows.astype(np.uint64) ) columns: dict[str, pa.Array] = { "stable_row_id": pa.array(stable, type=pa.uint64()), "seed": pa.array( np.full(flat_count, seed, dtype=np.int16) ), "asset": pa.array([item.asset] * flat_count), "asset_id": pa.array( np.full( flat_count, item.asset_id, dtype=np.int8, ) ), "anchor_row_index": pa.array(repeated_rows), "anchor_timestamp_ns": pa.array( np.repeat( item.features["timestamp_ns"][local_rows], horizon_count, ) ), "available_at_ns": pa.array( np.repeat( item.features["available_at_ns"][local_rows], horizon_count, ) ), "session_date": pa.array( np.repeat( item.features["session_date"][local_rows], horizon_count, ), type=pa.date32(), ), "minute_of_session": pa.array( np.repeat( item.features["minute_of_session"][local_rows], horizon_count, ) ), "horizon_minutes": pa.array( tiled_horizons.astype(np.int16) ), "target_valid": pa.array(repeated_mask), "target_timestamp_ns": pa.array( target_timestamp, mask=~repeated_mask, ), "target_available_at_ns": pa.array( target_timestamp + MINUTE_NS, mask=~repeated_mask, ), "anchor_raw_close": pa.array(anchor_close), "target_raw_close": pa.array( target_close, mask=~repeated_mask, ), "realized_log_return": pa.array( flat_returns, mask=~repeated_mask, ), "target_class": pa.array( flat_labels, mask=~repeated_mask, ), "expected_return": pa.array( expected_return.astype(np.float32) ), "expected_positive_probability": pa.array( expected_positive.astype(np.float32) ), "entropy": pa.array(entropy.astype(np.float32)), "max_probability": pa.array( flat_probabilities.max(axis=1).astype(np.float32) ), "top_class": pa.array( flat_probabilities.argmax(axis=1).astype(np.int8) ), } for class_index, name in enumerate(probability_names): columns[name] = pa.array( flat_probabilities[:, class_index] ) writer.write_table( pa.Table.from_pydict(columns, schema=schema), row_group_size=flat_count, ) asset_output_rows += flat_count total_rows += flat_count if np.any(nll_count == 0): raise ValueError( f"{item.asset} development sheet lacks a horizon" ) by_horizon = nll_sum / nll_count by_asset[item.asset] = { "anchor_rows": int(len(rows)), "macro_nll": float(np.mean(by_horizon)), "nll_by_horizon": { str(int(horizon)): float(by_horizon[index]) for index, horizon in enumerate(EVALUATION_HORIZONS) }, "output_rows": asset_output_rows, "valid_targets_by_horizon": { str(int(horizon)): int(nll_count[index]) for index, horizon in enumerate(EVALUATION_HORIZONS) }, } finally: writer.close() os.replace(temporary, output_path) return { "by_asset": by_asset, "equal_asset_macro_nll": float( np.mean([value["macro_nll"] for value in by_asset.values()]) ), "output_rows": total_rows, "parquet_bytes": output_path.stat().st_size, "parquet_sha256": _sha256(output_path), } def _write_calibration( path: Path, data: tuple[AssetPhaseData, ...], ) -> str: if path.exists(): raise FileExistsError(f"calibration artifact exists: {path}") arrays: dict[str, np.ndarray] = { "evaluation_horizons": EVALUATION_HORIZONS, "canonical_horizons": CANONICAL_HORIZONS, } for item in data: prefix = item.asset.lower() arrays[f"{prefix}_token_center"] = item.center arrays[f"{prefix}_token_scale"] = item.scale arrays[f"{prefix}_auxiliary_center"] = item.auxiliary_center arrays[f"{prefix}_auxiliary_scale"] = item.auxiliary_scale for name, value in class_map_to_arrays(item.class_map).items(): arrays[f"{prefix}_{name}"] = value np.savez_compressed(path, **arrays) return _sha256(path) def _run_manifest( *, config: dict[str, Any], project_root: Path, ) -> dict[str, Any]: return { "command": sys.argv, "config": config, "created_at_utc": datetime.now(timezone.utc).isoformat(), "environment": { "machine": platform.machine(), "numpy": np.__version__, "platform": platform.platform(), "python": sys.version, "torch": torch.__version__, "torch_mps_available": torch.backends.mps.is_available(), }, "provenance_sha256": { relative: _sha256(project_root / relative) for relative in PROVENANCE_PATHS }, "runner_version": RUNNER_VERSION, } def run( *, study_dir: Path, checkpoint_dir: Path, output_dir: Path, seed: int, epochs: int = DEFAULT_EPOCHS, batch_size: int = DEFAULT_BATCH_SIZE, gradient_accumulation_steps: int = DEFAULT_GRADIENT_ACCUMULATION, inference_batch_size: int = DEFAULT_INFERENCE_BATCH_SIZE, torch_threads: int = 6, device_name: str = "mps", max_train_rows_per_asset: int = 0, max_development_rows_per_asset: int = 0, ) -> dict[str, Any]: started = time.monotonic() if seed not in (0, 1, 2): raise ValueError("v7 seed must be 0, 1, or 2") if epochs <= 0 or batch_size <= 0 or inference_batch_size <= 0: raise ValueError("v7 training dimensions must be positive") if gradient_accumulation_steps <= 0: raise ValueError("v7 gradient accumulation must be positive") if max_train_rows_per_asset < 0 or max_development_rows_per_asset < 0: raise ValueError("v7 row caps cannot be negative") if ( batch_size * gradient_accumulation_steps != DEFAULT_BATCH_SIZE * DEFAULT_GRADIENT_ACCUMULATION ): raise ValueError("v7 effective batch size is frozen at 4096") if device_name == "mps" and not torch.backends.mps.is_available(): raise ValueError("v7 requested unavailable MPS") checkpoint_results_path = checkpoint_dir / "results.json" checkpoint_manifest_path = checkpoint_dir / "run_manifest.json" checkpoint_path = checkpoint_dir / "discovery_2_model.pt" checkpoint_results = json.loads(checkpoint_results_path.read_text()) checkpoint_manifest = json.loads(checkpoint_manifest_path.read_text()) if int(checkpoint_results.get("seed", -1)) != seed: raise ValueError("v7 seed and locked checkpoint seed differ") expected_checkpoint = { "asset_conditioning": "early_add", "context_length": DEFAULT_CONTEXT_LENGTH, "feature_set": FEATURE_SET, "max_train_rows_per_asset": 100_000, "rotary_base": DEFAULT_ROTARY_BASE, "supervised_horizons": [ int(value) for value in CANONICAL_HORIZONS ], "target_parameters": DEFAULT_TARGET_PARAMETERS, } for key, expected in expected_checkpoint.items(): if checkpoint_results.get(key) != expected: raise ValueError( f"locked checkpoint setting {key} differs from v7" ) manifest_config = checkpoint_manifest.get("config", {}) if manifest_config.get("seed") != seed: raise ValueError("locked checkpoint manifest has the wrong seed") manifest, loaded = _load_selection_data(study_dir) selected_indices = _selected_feature_indices() canonical_indices = _canonical_indices() checkpoint = torch.load( checkpoint_path, map_location="cpu", weights_only=False, ) config = { "assets": list(SELECTION_ASSETS), "auxiliary_feature_names": list(AUXILIARY_FEATURE_NAMES), "batch_size": batch_size, "checkpoint_dir": str(checkpoint_dir), "checkpoint_model_sha256": _sha256(checkpoint_path), "checkpoint_results_sha256": _sha256(checkpoint_results_path), "context_length": DEFAULT_CONTEXT_LENGTH, "development_test": list(DEVELOPMENT_TEST), "device_name": device_name, "epochs": epochs, "evaluation_horizons": EVALUATION_HORIZONS.tolist(), "feature_names": list(FEATURE_NAMES), "gradient_accumulation_steps": gradient_accumulation_steps, "inference_batch_size": inference_batch_size, "max_development_rows_per_asset": ( max_development_rows_per_asset ), "max_train_rows_per_asset": max_train_rows_per_asset, "optimizer_recipe": "adamw_constant", "refit_interval": [ DEVELOPMENT_START.isoformat(), VALIDATION_END.isoformat(), ], "rotary_base": DEFAULT_ROTARY_BASE, "seed": seed, "snapshot_sha256": manifest["snapshot_sha256"], "supervised_horizons": CANONICAL_HORIZONS.tolist(), "target_parameters": DEFAULT_TARGET_PARAMETERS, "torch_threads": torch_threads, "train_interval": [ DEVELOPMENT_START.isoformat(), TRAIN_END.isoformat(), ], "validation_interval": [ TRAIN_END.isoformat(), VALIDATION_END.isoformat(), ], } output_dir.mkdir(parents=True, exist_ok=True) manifest_path = output_dir / "run_manifest.json" results_path = output_dir / "results.json" prediction_path = output_dir / f"development_predictions_seed{seed}.parquet" checkpoint_output_path = output_dir / f"model_seed{seed}.pt" calibration_path = output_dir / f"calibration_seed{seed}.npz" guarded = ( manifest_path, results_path, prediction_path, checkpoint_output_path, calibration_path, ) if any(path.exists() for path in guarded): raise FileExistsError("v7 output directory is not immutable-empty") project_root = Path(__file__).resolve().parents[3] manifest_path.write_text( json.dumps( _run_manifest(config=config, project_root=project_root), indent=2, sort_keys=True, ) + "\n" ) _configure(seed, torch_threads) device = torch.device(device_name) selection_data = tuple( _prepare_phase_data( asset=asset, asset_id=asset_id, features=loaded[asset][0], bundle=loaded[asset][1], train_interval=( DEVELOPMENT_START.isoformat(), TRAIN_END.isoformat(), ), validation_interval=( TRAIN_END.isoformat(), VALIDATION_END.isoformat(), ), selected_indices=selected_indices, canonical_indices=canonical_indices, max_train_rows_per_asset=max_train_rows_per_asset, ) for asset_id, asset in enumerate(SELECTION_ASSETS) ) selection_model = _build_model(checkpoint, device=device) selection_training = _train( selection_model, selection_data, seed=seed, epochs=epochs, batch_size=batch_size, gradient_accumulation_steps=gradient_accumulation_steps, inference_batch_size=inference_batch_size, selected_indices=selected_indices, device=device, select_on_validation=True, ) selected_epoch = int(selection_training["best_epoch"]) del selection_model, selection_data if device.type == "mps": torch.mps.empty_cache() refit_data = tuple( _prepare_phase_data( asset=asset, asset_id=asset_id, features=loaded[asset][0], bundle=loaded[asset][1], train_interval=( DEVELOPMENT_START.isoformat(), VALIDATION_END.isoformat(), ), validation_interval=None, selected_indices=selected_indices, canonical_indices=canonical_indices, max_train_rows_per_asset=max_train_rows_per_asset, ) for asset_id, asset in enumerate(SELECTION_ASSETS) ) refit_model = _build_model(checkpoint, device=device) refit_training = _train( refit_model, refit_data, seed=seed + 70_001, epochs=selected_epoch, batch_size=batch_size, gradient_accumulation_steps=gradient_accumulation_steps, inference_batch_size=inference_batch_size, selected_indices=selected_indices, device=device, select_on_validation=False, ) torch.save( { "assets": list(SELECTION_ASSETS), "auxiliary_feature_names": AUXILIARY_FEATURE_NAMES, "context_length": DEFAULT_CONTEXT_LENGTH, "feature_names": FEATURE_NAMES, "refit_interval": config["refit_interval"], "rotary_base": DEFAULT_ROTARY_BASE, "seed": seed, "selected_epoch": selected_epoch, "state_dict": { name: value.detach().cpu() for name, value in refit_model.state_dict().items() }, "supervised_horizons": tuple( int(value) for value in CANONICAL_HORIZONS ), }, checkpoint_output_path, ) calibration_sha256 = _write_calibration( calibration_path, refit_data, ) development = _write_prediction_sheet( refit_model, refit_data, study_dir=study_dir, manifest=manifest, output_path=prediction_path, seed=seed, selected_indices=selected_indices, inference_batch_size=inference_batch_size, device=device, max_development_rows_per_asset=max_development_rows_per_asset, ) result = { **config, "calibration_sha256": calibration_sha256, "checkpoint_sha256": _sha256(checkpoint_output_path), "development_test_metrics": development, "duration_seconds": time.monotonic() - started, "parameter_count": sum( parameter.numel() for parameter in refit_model.parameters() if parameter.requires_grad ), "prediction_path": str(prediction_path), "refit_training": refit_training, "run_manifest_sha256": _sha256(manifest_path), "selected_epoch": selected_epoch, "selection_training": selection_training, } results_path.write_text( json.dumps(result, indent=2, sort_keys=True) + "\n" ) return result def main() -> None: parser = argparse.ArgumentParser( description="Train one immutable forecasting-v7 ensemble seed." ) parser.add_argument("--study-dir", type=Path, required=True) parser.add_argument("--checkpoint-dir", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--seed", type=int, choices=(0, 1, 2), required=True) parser.add_argument("--epochs", type=int, default=DEFAULT_EPOCHS) parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE) parser.add_argument( "--gradient-accumulation-steps", type=int, default=DEFAULT_GRADIENT_ACCUMULATION, ) parser.add_argument( "--inference-batch-size", type=int, default=DEFAULT_INFERENCE_BATCH_SIZE, ) parser.add_argument("--torch-threads", type=int, default=6) parser.add_argument("--device-name", default="mps") parser.add_argument("--max-train-rows-per-asset", type=int, default=0) parser.add_argument( "--max-development-rows-per-asset", type=int, default=0, ) args = parser.parse_args() print(json.dumps(run(**vars(args)), sort_keys=True)) if __name__ == "__main__": main()