source/project/experiments/forecasting_v5/runner.py

"""Run shared-weight multi-asset close-distribution experiments."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
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_v2 import (
    CANONICAL_HORIZONS,
    CLASS_COUNT,
    DISCOVERY_FOLDS,
    EVALUATION_HORIZONS,
    SELECTION_ASSETS,
)
from project.evaluators.close_distribution_v2 import (
    ClassMap,
    class_map_to_arrays,
    encode_returns,
    fit_class_map,
    metrics,
    prior_probabilities,
    score_arrays,
)
from project.experiments.forecasting_v4.runner import (
    FEATURE_SETS,
    _context,
    _feature_indices,
    _horizon_indices,
    _load_npz,
    _masked_nll,
    _normalization,
    _optimizer,
    _optimizer_settings,
    _ordinal,
    _parse_horizons,
    _session_scores,
)
from project.models.close_distribution_v3 import (
    ASSET_CONDITIONING_MODES,
    build_shared_model,
)


PROVENANCE_PATHS = (
    "project/data/close_distribution_v2.py",
    "project/evaluators/close_distribution_v2.py",
    "project/experiments/forecasting_v4/runner.py",
    "project/experiments/forecasting_v5/runner.py",
    "project/models/close_distribution_v2.py",
    "project/models/close_distribution_v3.py",
    "project/studies/forecasting_v5.md",
    "requirements.txt",
)


@dataclass
class AssetFoldData:
    asset: str
    asset_id: int
    features: dict[str, np.ndarray]
    bundle: dict[str, np.ndarray]
    center: np.ndarray
    scale: np.ndarray
    class_map: ClassMap
    train_rows: np.ndarray
    train_labels: np.ndarray
    train_target_mask: np.ndarray
    diagnostic_rows: np.ndarray
    diagnostic_labels: np.ndarray
    diagnostic_target_mask: np.ndarray
    validation_rows: np.ndarray
    validation_y: np.ndarray
    validation_target_mask: np.ndarray
    validation_session_date: np.ndarray
    diagnostic_purge_rows: 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 _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 _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("scope") != "selection":
        raise ValueError("v5 requires the selection-only prepared manifest")
    if tuple(manifest.get("selection_assets", ())) != SELECTION_ASSETS:
        raise ValueError("selection asset cohort changed")
    data_dir = study_dir / "data" / "runner"
    loaded = {}
    for asset in SELECTION_ASSETS:
        slug = asset.lower()
        feature_path = data_dir / f"{slug}_features.npz"
        label_path = data_dir / f"{slug}_labels.npz"
        for path in (feature_path, label_path):
            relative = str(path.relative_to(study_dir))
            if _sha256(path) != manifest["derived_sha256"][relative]:
                raise ValueError(f"prepared hash mismatch: {relative}")
        loaded[asset] = (_load_npz(feature_path), _load_npz(label_path))
    return manifest, loaded


def _prepare_asset_fold(
    *,
    asset: str,
    asset_id: int,
    features: dict[str, np.ndarray],
    bundle: dict[str, np.ndarray],
    split: dict[str, tuple[str, str]],
    supervised_indices: tuple[int, ...],
    selected_indices: tuple[int, ...],
    max_train_rows_per_asset: int,
    diagnostic_tail_fraction: float,
) -> AssetFoldData:
    dates = bundle["session_date"]
    train = (
        (dates >= _ordinal(split["train"][0]))
        & (dates < _ordinal(split["train"][1]))
    )
    validation = (
        (dates >= _ordinal(split["validation"][0]))
        & (dates < _ordinal(split["validation"][1]))
    )
    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]

    diagnostic_rows = np.empty(0, dtype=train_rows.dtype)
    diagnostic_y = np.empty((0, train_y.shape[1]), dtype=train_y.dtype)
    diagnostic_target_mask = np.empty(
        (0, train_target_mask.shape[1]),
        dtype=np.bool_,
    )
    diagnostic_purge_rows = 0
    normalization_max_row = None
    if diagnostic_tail_fraction:
        diagnostic_count = max(
            1,
            round(len(train_rows) * diagnostic_tail_fraction),
        )
        diagnostic_rows = train_rows[-diagnostic_count:]
        diagnostic_y = train_y[-diagnostic_count:]
        diagnostic_target_mask = train_target_mask[-diagnostic_count:]
        train_rows = train_rows[:-diagnostic_count]
        train_y = train_y[:-diagnostic_count]
        train_target_mask = train_target_mask[:-diagnostic_count]
        normalization_max_row = int(diagnostic_rows[0])
        unpurged_count = len(train_rows)
        keep = train_rows + int(EVALUATION_HORIZONS.max()) < int(
            diagnostic_rows[0]
        )
        train_rows = train_rows[keep]
        train_y = train_y[keep]
        train_target_mask = train_target_mask[keep]
        diagnostic_purge_rows = unpurged_count - len(train_rows)

    class_map = fit_class_map(
        train_y,
        train_target_mask,
        class_count=CLASS_COUNT,
    )
    center, scale = _normalization(
        features,
        train_end=split["train"][1],
        selected_indices=selected_indices,
        max_row_exclusive=normalization_max_row,
    )
    all_train_labels = encode_returns(
        train_y,
        train_target_mask,
        class_map.edges,
    )
    all_diagnostic_labels = encode_returns(
        diagnostic_y,
        diagnostic_target_mask,
        class_map.edges,
    )
    return AssetFoldData(
        asset=asset,
        asset_id=asset_id,
        features=features,
        bundle=bundle,
        center=center,
        scale=scale,
        class_map=class_map,
        train_rows=train_rows,
        train_labels=all_train_labels[:, supervised_indices],
        train_target_mask=train_target_mask[:, supervised_indices],
        diagnostic_rows=diagnostic_rows,
        diagnostic_labels=all_diagnostic_labels[:, supervised_indices],
        diagnostic_target_mask=diagnostic_target_mask[:, supervised_indices],
        validation_rows=bundle["row_index"][validation],
        validation_y=bundle["y"][validation],
        validation_target_mask=bundle["target_mask"][validation],
        validation_session_date=bundle["session_date"][validation],
        diagnostic_purge_rows=diagnostic_purge_rows,
    )


def _balanced_epoch_chunks(
    row_counts: tuple[int, ...],
    *,
    effective_batch_size: int,
    generator: np.random.Generator,
) -> list[tuple[np.ndarray, ...]]:
    if not row_counts or any(value <= 0 for value in row_counts):
        raise ValueError("every asset must have positive training rows")
    rows_per_asset_update = effective_batch_size // len(row_counts)
    if rows_per_asset_update <= 0:
        raise ValueError("effective batch is smaller than the asset cohort")
    updates = math.ceil(max(row_counts) / rows_per_asset_update)
    per_asset = [
        np.array_split(generator.permutation(count), updates)
        for count in row_counts
    ]
    if any(any(not len(chunk) for chunk in chunks) for chunks in per_asset):
        raise ValueError("asset-balanced update contains an empty asset slice")
    return [
        tuple(chunks[update] for chunks in per_asset)
        for update in range(updates)
    ]


def _asset_balanced_loss_part(
    logits: torch.Tensor,
    labels: torch.Tensor,
    mask: torch.Tensor,
    asset_id: torch.Tensor,
    denominators: torch.Tensor,
) -> torch.Tensor:
    if denominators.ndim != 2:
        raise ValueError("asset/horizon denominators must be two-dimensional")
    losses = nn.functional.cross_entropy(
        logits.reshape(-1, logits.shape[-1]),
        labels.clamp_min(0).reshape(-1),
        reduction="none",
    ).reshape(labels.shape)
    result = losses.new_zeros(())
    for index in range(denominators.shape[0]):
        selected = asset_id == index
        horizon_sum = (losses[selected] * mask[selected]).sum(dim=0)
        result = result + torch.sum(horizon_sum / denominators[index])
    return result / denominators.numel()


def _training_microbatch(
    data: tuple[AssetFoldData, ...],
    asset_ids: np.ndarray,
    positions: np.ndarray,
    *,
    context_length: int,
    selected_indices: tuple[int, ...],
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
    tokens = []
    padding = []
    labels = []
    masks = []
    ids = []
    for asset_data in data:
        selected = asset_ids == asset_data.asset_id
        if not np.any(selected):
            continue
        local = positions[selected]
        context, context_padding = _context(
            asset_data.features["X"],
            asset_data.train_rows[local],
            context_length=context_length,
            selected_indices=selected_indices,
            center=asset_data.center,
            scale=asset_data.scale,
        )
        tokens.append(context)
        padding.append(context_padding)
        labels.append(
            torch.from_numpy(
                asset_data.train_labels[local].astype(
                    np.int64,
                    copy=False,
                )
            )
        )
        masks.append(
            torch.from_numpy(
                asset_data.train_target_mask[local].astype(
                    np.float32,
                    copy=False,
                )
            )
        )
        ids.append(
            torch.full(
                (len(local),),
                asset_data.asset_id,
                dtype=torch.long,
            )
        )
    return (
        torch.cat(tokens),
        torch.cat(padding),
        torch.cat(labels),
        torch.cat(masks),
        torch.cat(ids),
    )


def _predict_asset(
    model: nn.Module,
    asset_data: AssetFoldData,
    *,
    rows: np.ndarray,
    context_length: int,
    selected_indices: tuple[int, ...],
    batch_size: int,
    device: torch.device,
    query_horizons: tuple[int, ...] | None,
) -> np.ndarray:
    horizon_count = (
        len(query_horizons)
        if query_horizons is not None
        else len(model.output_horizons)
    )
    result = np.empty(
        (len(rows), horizon_count, CLASS_COUNT),
        dtype=np.float64,
    )
    query = (
        torch.tensor(query_horizons, dtype=torch.float32, device=device)
        if query_horizons is not None
        else None
    )
    model.eval()
    with torch.inference_mode():
        for start in range(0, len(rows), batch_size):
            stop = min(start + batch_size, len(rows))
            tokens, padding = _context(
                asset_data.features["X"],
                rows[start:stop],
                context_length=context_length,
                selected_indices=selected_indices,
                center=asset_data.center,
                scale=asset_data.scale,
            )
            asset_id = torch.full(
                (stop - start,),
                asset_data.asset_id,
                dtype=torch.long,
                device=device,
            )
            logits = model(
                tokens.to(device),
                padding.to(device),
                asset_id,
                query,
            )
            result[start:stop] = torch.softmax(
                logits.cpu().to(torch.float64),
                dim=2,
            ).numpy()
    return result


def _train_model(
    model: nn.Module,
    data: tuple[AssetFoldData, ...],
    *,
    context_length: int,
    selected_indices: tuple[int, ...],
    seed: int,
    epochs: int,
    batch_size: int,
    gradient_accumulation_steps: int,
    optimizer_recipe: str,
    inference_batch_size: int,
    device: torch.device,
) -> dict[str, Any]:
    optimizer = _optimizer(model, optimizer_recipe)
    effective_batch_size = batch_size * gradient_accumulation_steps
    generator = np.random.default_rng(seed + 10_007)
    histories: list[float] = []
    diagnostic_history: list[float] = []
    best_epoch = 0
    best_diagnostic_nll = float("inf")
    best_state: dict[str, torch.Tensor] | None = None
    update_count = 0
    expected_updates_per_epoch = None

    for epoch in range(epochs):
        model.train()
        updates = _balanced_epoch_chunks(
            tuple(len(item.train_rows) for item in data),
            effective_batch_size=effective_batch_size,
            generator=generator,
        )
        if expected_updates_per_epoch is None:
            expected_updates_per_epoch = len(updates)
        elif len(updates) != expected_updates_per_epoch:
            raise AssertionError("updates per epoch changed")
        epoch_loss_sum = np.zeros(
            (len(data), data[0].train_labels.shape[1]),
            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(
                [
                    asset_data.train_target_mask[chunk].sum(axis=0)
                    for asset_data, chunk in zip(data, chunks, strict=True)
                ]
            ).astype(np.float32)
            if np.any(denominators == 0):
                raise ValueError(
                    "an 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,
                    labels,
                    mask,
                    asset_id,
                ) = _training_microbatch(
                    data,
                    update_asset_ids[start:stop],
                    update_positions[start:stop],
                    context_length=context_length,
                    selected_indices=selected_indices,
                )
                tokens = tokens.to(device)
                padding = padding.to(device)
                labels = labels.to(device)
                mask = mask.to(device)
                asset_id = asset_id.to(device)
                logits = model(tokens, padding, asset_id)
                loss = _asset_balanced_loss_part(
                    logits,
                    labels,
                    mask,
                    asset_id,
                    denominator_tensor,
                )
                loss.backward()
                with torch.no_grad():
                    item_losses = nn.functional.cross_entropy(
                        logits.reshape(-1, logits.shape[-1]),
                        labels.clamp_min(0).reshape(-1),
                        reduction="none",
                    ).reshape(labels.shape)
                    for index in range(len(data)):
                        selected = asset_id == index
                        epoch_loss_sum[index] += (
                            item_losses[selected] * mask[selected]
                        ).sum(dim=0).cpu().numpy()
                        epoch_target_count[index] += (
                            mask[selected]
                            .sum(dim=0)
                            .cpu()
                            .numpy()
                            .astype(np.int64)
                        )
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            update_count += 1

        histories.append(
            float(np.mean(epoch_loss_sum / epoch_target_count))
        )
        if all(len(item.diagnostic_rows) for item in data):
            diagnostic_nll = []
            for asset_data in data:
                probabilities = _predict_asset(
                    model,
                    asset_data,
                    rows=asset_data.diagnostic_rows,
                    context_length=context_length,
                    selected_indices=selected_indices,
                    batch_size=inference_batch_size,
                    device=device,
                    query_horizons=None,
                )
                diagnostic_nll.append(
                    _masked_nll(
                        probabilities,
                        asset_data.diagnostic_labels,
                        asset_data.diagnostic_target_mask,
                    )
                )
            equal_asset_nll = float(np.mean(diagnostic_nll))
            diagnostic_history.append(equal_asset_nll)
            if equal_asset_nll < best_diagnostic_nll:
                best_diagnostic_nll = equal_asset_nll
                best_epoch = epoch + 1
                best_state = {
                    name: value.detach().cpu().clone()
                    for name, value in model.state_dict().items()
                }

    if best_state is not None:
        model.load_state_dict(best_state)
    else:
        best_epoch = epochs
    if expected_updates_per_epoch is None:
        raise AssertionError("training produced no updates")
    return {
        "best_epoch": best_epoch,
        "completed_fixed_update_budget": (
            update_count == expected_updates_per_epoch * epochs
        ),
        "diagnostic_tail_nll": diagnostic_history,
        "train_loss": histories,
        "updates": update_count,
        "updates_per_epoch": expected_updates_per_epoch,
    }


def _run_manifest(config: dict[str, Any]) -> dict[str, Any]:
    project_root = Path(__file__).resolve().parents[3]
    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
        },
    }


def run(
    *,
    study_dir: Path,
    output_dir: Path,
    asset_conditioning: str,
    supervised_horizons: str | Iterable[int] = "2,4,8,16,32",
    context_length: int = 128,
    seed: int = 0,
    epochs: int = 16,
    batch_size: int = 256,
    inference_batch_size: int = 512,
    target_parameters: int = 25_000,
    torch_threads: int = 6,
    max_train_rows_per_asset: int = 100_000,
    device_name: str = "mps",
    gradient_accumulation_steps: int = 16,
    optimizer_recipe: str = "adamw_constant",
    rotary_base: float = 16.0,
    feature_set: str = "session_time",
    diagnostic_tail_fraction: float = 0.1,
) -> dict[str, Any]:
    started = time.monotonic()
    horizons = _parse_horizons(supervised_horizons)
    if asset_conditioning not in ASSET_CONDITIONING_MODES:
        raise ValueError("unknown asset conditioning mode")
    if feature_set not in FEATURE_SETS:
        raise ValueError("unknown feature set")
    if optimizer_recipe != "adamw_constant":
        raise ValueError(
            "v5 currently implements only the frozen AdamW constant recipe"
        )
    if horizons[0] != 2 or horizons[-1] != 32:
        raise ValueError("supervised horizons must include 2 and 32")
    if epochs <= 0 or batch_size <= 0 or context_length <= 0:
        raise ValueError("training dimensions must be positive")
    if gradient_accumulation_steps <= 0:
        raise ValueError("gradient accumulation must be positive")
    if not 0.0 <= diagnostic_tail_fraction < 0.5:
        raise ValueError("diagnostic fraction must lie in [0, 0.5)")
    if device_name == "mps" and not torch.backends.mps.is_available():
        raise ValueError("MPS was requested but is unavailable")

    manifest, loaded = _load_selection_data(study_dir)
    selected_indices = _feature_indices(feature_set)
    supervised_indices = _horizon_indices(horizons)
    canonical_indices = _horizon_indices(CANONICAL_HORIZONS)
    dense_unseen_horizons = tuple(
        int(value)
        for value in EVALUATION_HORIZONS
        if value not in set(CANONICAL_HORIZONS.tolist())
    )
    dense_unseen_indices = _horizon_indices(dense_unseen_horizons)
    learning_rate, weight_decay = _optimizer_settings(optimizer_recipe)
    config = {
        "asset_conditioning": asset_conditioning,
        "assets": list(SELECTION_ASSETS),
        "batch_size": batch_size,
        "context_length": context_length,
        "device_name": device_name,
        "diagnostic_tail_fraction": diagnostic_tail_fraction,
        "epochs": epochs,
        "feature_names": list(FEATURE_SETS[feature_set]),
        "feature_set": feature_set,
        "gradient_accumulation_steps": gradient_accumulation_steps,
        "inference_batch_size": inference_batch_size,
        "learning_rate": learning_rate,
        "max_train_rows_per_asset": max_train_rows_per_asset,
        "model_name": "shared_transformer_rotary",
        "optimizer_recipe": optimizer_recipe,
        "rotary_base": rotary_base,
        "seed": seed,
        "snapshot_sha256": manifest["snapshot_sha256"],
        "supervised_horizons": list(horizons),
        "target_parameters": target_parameters,
        "torch_threads": torch_threads,
        "weight_decay": weight_decay,
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    results_path = output_dir / "results.json"
    run_manifest_path = output_dir / "run_manifest.json"
    if results_path.exists() or run_manifest_path.exists():
        raise FileExistsError("experiment output already exists")
    run_manifest_path.write_text(
        json.dumps(_run_manifest(config), indent=2, sort_keys=True) + "\n"
    )
    run_manifest_sha256 = _sha256(run_manifest_path)

    _configure(seed, torch_threads)
    device = torch.device(device_name)
    fold_results: dict[str, Any] = {}
    for fold_name, split in DISCOVERY_FOLDS.items():
        fold_started = time.monotonic()
        data = tuple(
            _prepare_asset_fold(
                asset=asset,
                asset_id=asset_id,
                features=loaded[asset][0],
                bundle=loaded[asset][1],
                split=split,
                supervised_indices=supervised_indices,
                selected_indices=selected_indices,
                max_train_rows_per_asset=max_train_rows_per_asset,
                diagnostic_tail_fraction=diagnostic_tail_fraction,
            )
            for asset_id, asset in enumerate(SELECTION_ASSETS)
        )
        model = build_shared_model(
            context_length=context_length,
            channels=len(selected_indices) + 1,
            output_horizons=horizons,
            classes=CLASS_COUNT,
            target_parameters=target_parameters,
            rotary_base=rotary_base,
            asset_count=len(SELECTION_ASSETS),
            asset_conditioning=asset_conditioning,
        ).to(device)
        parameter_count = sum(
            parameter.numel()
            for parameter in model.parameters()
            if parameter.requires_grad
        )
        core_parameter_count = parameter_count - sum(
            parameter.numel()
            for name, parameter in model.named_parameters()
            if name.startswith(("asset_embedding", "asset_affine"))
        )
        training = _train_model(
            model,
            data,
            context_length=context_length,
            selected_indices=selected_indices,
            seed=seed,
            epochs=epochs,
            batch_size=batch_size,
            gradient_accumulation_steps=gradient_accumulation_steps,
            optimizer_recipe=optimizer_recipe,
            inference_batch_size=inference_batch_size,
            device=device,
        )

        asset_results = {}
        for asset_data in data:
            probabilities = _predict_asset(
                model,
                asset_data,
                rows=asset_data.validation_rows,
                context_length=context_length,
                selected_indices=selected_indices,
                batch_size=inference_batch_size,
                device=device,
                query_horizons=tuple(
                    int(value) for value in EVALUATION_HORIZONS
                ),
            )
            scores = score_arrays(
                asset_data.validation_y,
                asset_data.validation_target_mask,
                probabilities,
                class_map=asset_data.class_map,
            )
            prior = prior_probabilities(
                len(asset_data.validation_y),
                asset_data.class_map,
            )
            prior_scores = score_arrays(
                asset_data.validation_y,
                asset_data.validation_target_mask,
                prior,
                class_map=asset_data.class_map,
            )
            metric_slices = {
                "canonical": metrics(
                    asset_data.validation_y,
                    asset_data.validation_target_mask,
                    probabilities,
                    class_map=asset_data.class_map,
                    horizon_indices=canonical_indices,
                    score_values=scores,
                ),
                "dense_unseen": metrics(
                    asset_data.validation_y,
                    asset_data.validation_target_mask,
                    probabilities,
                    class_map=asset_data.class_map,
                    horizon_indices=dense_unseen_indices,
                    score_values=scores,
                ),
                "prior_canonical": metrics(
                    asset_data.validation_y,
                    asset_data.validation_target_mask,
                    prior,
                    class_map=asset_data.class_map,
                    horizon_indices=canonical_indices,
                    score_values=prior_scores,
                ),
            }
            np.savez_compressed(
                output_dir
                / f"{fold_name}_{asset_data.asset.lower()}_session_scores.npz",
                **_session_scores(
                    asset_data.validation_y,
                    asset_data.validation_target_mask,
                    probabilities,
                    asset_data.validation_session_date,
                    class_map=asset_data.class_map,
                    score_values=scores,
                ),
            )
            asset_results[asset_data.asset] = {
                "class_map_rows": int(len(asset_data.train_rows)),
                "diagnostic_purge_rows": (
                    asset_data.diagnostic_purge_rows
                ),
                "diagnostic_tail_rows": int(
                    len(asset_data.diagnostic_rows)
                ),
                "metric_slices": metric_slices,
                "train_rows": int(len(asset_data.train_rows)),
                "train_valid_targets_by_supervised_horizon": {
                    str(horizon): int(count)
                    for horizon, count in zip(
                        horizons,
                        asset_data.train_target_mask.sum(axis=0),
                        strict=True,
                    )
                },
                "validation_rows": int(len(asset_data.validation_rows)),
            }

        torch.save(
            {
                "asset_calibration": {
                    item.asset: {
                        "center": item.center,
                        "scale": item.scale,
                        **class_map_to_arrays(item.class_map),
                    }
                    for item in data
                },
                "asset_conditioning": asset_conditioning,
                "assets": list(SELECTION_ASSETS),
                "context_length": context_length,
                "feature_names": FEATURE_SETS[feature_set],
                "rotary_base": rotary_base,
                "state_dict": model.state_dict(),
                "supervised_horizons": horizons,
            },
            output_dir / f"{fold_name}_model.pt",
        )
        fold_results[fold_name] = {
            **training,
            "assets": asset_results,
            "checkpoint_selection": (
                "lowest_equal_asset_inner_tail_nll"
                if diagnostic_tail_fraction
                else "final_epoch"
            ),
            "conditioning_parameter_count": (
                parameter_count - core_parameter_count
            ),
            "core_parameter_count": core_parameter_count,
            "duration_seconds": time.monotonic() - fold_started,
            "estimated_madds_per_example": int(model.estimated_madds),
            "parameter_count": parameter_count,
        }
        del model
        if device.type == "mps":
            torch.mps.empty_cache()

    by_asset = {}
    for asset in SELECTION_ASSETS:
        canonical = float(
            np.mean(
                [
                    fold_results[fold]["assets"][asset]["metric_slices"][
                        "canonical"
                    ]["macro_nll"]
                    for fold in DISCOVERY_FOLDS
                ]
            )
        )
        dense = float(
            np.mean(
                [
                    fold_results[fold]["assets"][asset]["metric_slices"][
                        "dense_unseen"
                    ]["macro_nll"]
                    for fold in DISCOVERY_FOLDS
                ]
            )
        )
        by_asset[asset] = {
            "mean_canonical_macro_nll": canonical,
            "mean_dense_unseen_macro_nll": dense,
        }
    summary = {
        **config,
        "by_asset": by_asset,
        "duration_seconds": time.monotonic() - started,
        "equal_asset_mean_canonical_macro_nll": float(
            np.mean(
                [
                    value["mean_canonical_macro_nll"]
                    for value in by_asset.values()
                ]
            )
        ),
        "equal_asset_mean_dense_unseen_macro_nll": float(
            np.mean(
                [
                    value["mean_dense_unseen_macro_nll"]
                    for value in by_asset.values()
                ]
            )
        ),
        "folds": fold_results,
        "run_manifest_sha256": run_manifest_sha256,
    }
    results_path.write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n"
    )
    return summary


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--study-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument(
        "--asset-conditioning",
        choices=ASSET_CONDITIONING_MODES,
        required=True,
    )
    parser.add_argument(
        "--supervised-horizons",
        default="2,4,8,16,32",
    )
    parser.add_argument("--context-length", type=int, default=128)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--epochs", type=int, default=16)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--inference-batch-size", type=int, default=512)
    parser.add_argument("--target-parameters", type=int, default=25_000)
    parser.add_argument("--torch-threads", type=int, default=6)
    parser.add_argument(
        "--max-train-rows-per-asset",
        type=int,
        default=100_000,
    )
    parser.add_argument(
        "--device-name",
        choices=("cpu", "mps"),
        default="mps",
    )
    parser.add_argument(
        "--gradient-accumulation-steps",
        type=int,
        default=16,
    )
    parser.add_argument(
        "--optimizer-recipe",
        choices=("adamw_constant",),
        default="adamw_constant",
    )
    parser.add_argument("--rotary-base", type=float, default=16.0)
    parser.add_argument(
        "--feature-set",
        choices=tuple(FEATURE_SETS),
        default="session_time",
    )
    parser.add_argument(
        "--diagnostic-tail-fraction",
        type=float,
        default=0.1,
    )
    args = parser.parse_args()
    print(json.dumps(run(**vars(args)), sort_keys=True))


if __name__ == "__main__":
    main()