""" End-to-end runner, rebuilt against explicit contracts (see provenance.py): 1. Data: get_dataset() [REAL, hard-fails] or get_synthetic_dataset() [explicit opt-in] — never a silent fallback between them. 2. Validation: trajectory lengths checked BEFORE training starts. 3. Dataset reuse: DatasetRegistry.claim() blocks retraining on a dataset already CONSUMED by a prior run — required because multiple contributors will supply datasets over time. 4. Checkpointing: CheckpointStore — content-addressed (sha256 of config+code+dataset), atomic write, and a human-readable meta.json sidecar. Run (real data required by default): python -m src.run_full --data-root ./data/real Run against synthetic data (explicit opt-in, for smoke-testing only): python -m src.run_full --synthetic """ from __future__ import annotations import os, sys, argparse sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import torch from torch.utils.data import DataLoader import numpy as np from src.data_real import get_dataset, get_synthetic_dataset from src.data_pbdb import get_pbdb_dataset, DEFAULT_TAXON_GROUPS from src.normalization import FieldNormalizer from src.model import MultiScaleEncoder, HierarchicalHyperbolicPredictor, HyperbolicCritic from src.physics_losses import combined_physics_loss from src.env import MultiStepPoincareEnv from src.ppo import PoincareActor, PPOTrainer from src.provenance import ( DatasetRegistry, CheckpointStore, hash_dataset, hash_code, validate_trajectory_lengths, DatasetAlreadyUsedError, DatasetInProgressError, ) from src.config import BEST_HPARAMS as BEST, WINDOW SRC_DIR = os.path.dirname(os.path.abspath(__file__)) def collate(batch): return torch.stack([b["fields"] for b in batch]) def supervised_pretrain(model, norm, ds, device, epochs=5): loader = DataLoader(ds, batch_size=BEST["batch_size"], shuffle=True, collate_fn=collate) opt = torch.optim.Adam(model.parameters(), lr=BEST["lr"]) ps = BEST["pred_steps"] w = BEST["w_phys"] for ep in range(epochs): total, n = 0.0, 0 for batch in loader: B, T, C, H, W = batch.shape batch = batch.to(device) flat = norm.transform(batch.view(B * T, C, H, W)).view(B, T, C, H, W) x = flat[:, :WINDOW] with torch.no_grad(): tgt = torch.stack([model.encode(flat[:, WINDOW + s]) for s in range(ps)], 1) pred = model(x) loss = model.hyperbolic_loss(pred, tgt) + combined_physics_loss( flat[:, : WINDOW + ps], w_smooth=w, w_temp=w, w_cons=w * 0.5 ) if not torch.isfinite(loss): raise RuntimeError( f"[NON_FINITE_LOSS] loss became {loss.item()} at epoch {ep+1}; " f"stopping rather than silently continuing with a corrupted model." ) opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() total += loss.item() n += 1 print(f" Pretrain epoch {ep+1}/{epochs} loss={total/max(n,1):.4f}") return model def main(): parser = argparse.ArgumentParser() parser.add_argument("--data-root", action="append", default=None, help="Directory containing real .hdf5/.h5 Well files. " "May be repeated. Default: ./data/real, ./data/well") parser.add_argument("--synthetic", action="store_true", help="Explicit opt-in to synthetic data (smoke test only).") parser.add_argument("--pbdb", action="store_true", help="Explicit opt-in to real PBDB fossil-occurrence data " "(spatiotemporal occurrence/diversity density fields). " "Requires network access to paleobiodb.org.") parser.add_argument("--pbdb-taxa", nargs="+", default=None, help=f"Taxon groups to fetch from PBDB. Default: {list(DEFAULT_TAXON_GROUPS)}") parser.add_argument("--experiment-id", default=None, help="Human label for this run. Default: auto-generated.") parser.add_argument("--allow-dataset-reuse", action="store_true", help="Explicit override to retrain on an already-CONSUMED " "dataset. Off by default — reuse is blocked.") args = parser.parse_args() device = "cpu" print("=" * 64) print("Full pipeline: data -> hierarchical Poincare -> physics -> PPO") print("Optuna best HPs:", BEST) print("=" * 64) # ---- 1. Data: explicit, no silent fallback ----------------------- if args.synthetic: ds, provenance = get_synthetic_dataset(max_samples=128, n_steps=14) elif args.pbdb: ds, provenance = get_pbdb_dataset( taxon_groups=args.pbdb_taxa or DEFAULT_TAXON_GROUPS, ) else: ds, provenance = get_dataset( max_samples=128, n_steps=14, search_roots=args.data_root, ) print(f"[data] provenance={provenance} size={len(ds)}") # ---- 2. Validate BEFORE training, not mid-loop -------------------- required_length = WINDOW + BEST["pred_steps"] validate_trajectory_lengths(ds, required_length=required_length) print(f"[validate] all sampled trajectories >= {required_length} steps: OK") # ---- 3. Dataset-reuse registry ------------------------------------- dataset_hash = hash_dataset(ds, sample_cap=64) code_hash = hash_code(SRC_DIR) experiment_id = args.experiment_id or f"run_full:{dataset_hash[:8]}:{code_hash[:8]}" registry = DatasetRegistry(registry_dir="registry/datasets") if args.allow_dataset_reuse: status = registry.status(dataset_hash) if status and status["status"] == "CONSUMED": print(f"[registry] WARNING: explicit override — retraining on " f"already-CONSUMED dataset {dataset_hash[:12]}") registry.allow_retry(dataset_hash) try: registry.claim(dataset_hash, experiment_id) except (DatasetAlreadyUsedError, DatasetInProgressError) as e: print(f"[registry] BLOCKED: {e}") raise try: # ---- 4. Normalizer --------------------------------------------- samples = [] for i in range(min(48, len(ds))): item = ds[i] samples.append(item["fields"] if isinstance(item, dict) else item) data = torch.stack(samples) norm = FieldNormalizer(mode="zscore").fit(data) print("[norm] fitted") # ---- 5. Hierarchical model -------------------------------------- enc = MultiScaleEncoder(hidden=BEST["hidden"], out_dim=8) model = HierarchicalHyperbolicPredictor( enc, c=BEST["curvature"], pred_steps=BEST["pred_steps"], levels=BEST["levels"] ).to(device) print("\n--- Supervised pre-training with physics priors ---") model = supervised_pretrain(model, norm, ds, device, epochs=4) # ---- 6. PPO with hyperbolic critic ------------------------------- print("\n--- PPO fine-tuning with hyperbolic critic ---") env = MultiStepPoincareEnv( dataset=ds, normalizer=norm, encoder=model.encoder, poincare_module=model.poincare, window=WINDOW, horizon=BEST["pred_steps"], device=device, ) actor = PoincareActor(obs_dim=8, action_dim=8, hidden=64) critic = HyperbolicCritic(c=BEST["curvature"]) ppo = PPOTrainer(actor, critic, model.poincare, lr=BEST["lr"], device=device) returns = [] for update in range(12): rollout = ppo.collect_rollout(env, n_steps=48) loss = ppo.update(rollout, n_epochs=3, batch_size=16) ep_ret = float(np.sum(rollout["rewards"])) returns.append(ep_ret) if (update + 1) % 3 == 0: print(f" PPO update {update+1}/12 loss={loss:.4f} rollout_return={ep_ret:.3f}") print(f" Mean return (last 4): {np.mean(returns[-4:]):.3f}") # ---- 7. Content-addressed, atomic checkpoint --------------------- store = CheckpointStore(checkpoints_dir="checkpoints") result = store.save( model_state={ "model": model.state_dict(), "actor": actor.state_dict(), "critic": critic.state_dict(), }, config=BEST, dataset_hash=dataset_hash, code_hash=code_hash, data_provenance=provenance, extra={ "normalizer": norm.state_dict(), "ppo_returns": returns, "experiment_id": experiment_id, }, ) print(f"\n[checkpoint] {result['outcome_code']} -> {result['path']}") registry.mark_consumed(dataset_hash) print(f"[registry] dataset {dataset_hash[:12]} marked CONSUMED " f"(future runs on this exact data will be blocked by default)") except Exception as e: registry.mark_failed(dataset_hash, error_detail=str(e)) print(f"[registry] dataset {dataset_hash[:12]} marked FAILED: {e}") raise print("\nDone.") if __name__ == "__main__": main()