| """Draw forecasts from a trained model: for every initial condition of a split, ``n_samples`` samples of x_T. |
| |
| python -m bifurcation.sample --config CONFIG --data DIR --checkpoint FILE --out DIR [--split test] [--gids 5,16] |
| [--n-samples 10000] [--n-steps N] [--seed-base 777] [--chunk 2000] [--truth-bank DIR] [--device cuda] |
| |
| CONFIG is a run configuration (records/<system>/runs/<model>/config.yaml, or the config_resolved.yaml of a run trained |
| here). The EMA weights of the checkpoint are used when the method keeps an EMA. ``--n-steps`` overrides the sampler |
| resolution (steps per segment for Cascaded SI, SDE steps for Direct SI, Heun steps for EDM). The noise of condition |
| ``gid`` comes from the stream seeded with ``seed_base * 100003 + gid`` and is drawn in chunks of ``--chunk`` samples. |
| ``--truth-bank`` supplies x_0 of test conditions that are not in the package (waves, NLKG). |
| |
| Output: ``OUT/<variant>/group_<gid>.npz`` with ``gen_xT_norm`` (n_samples, D), ``x0_norm`` (D,), ``true_xT_norm`` |
| (the package's truth members at the horizon) and ``gid``; states are in the package's normalized units. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| from datetime import datetime, timezone |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from .data.loaders import get_split_set |
| from .registry import build_method, build_model |
| from .utils import load_yaml |
|
|
|
|
| def parse_args(argv=None): |
| p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| p.add_argument("--config", required=True, help="run configuration (YAML)") |
| p.add_argument("--data", required=True, help="dataset package directory") |
| p.add_argument("--checkpoint", required=True) |
| p.add_argument("--out", required=True, help="output directory") |
| p.add_argument("--split", default="test", choices=("train", "val", "test")) |
| p.add_argument("--gids", default=None, help="comma-separated subset of the split's conditions") |
| p.add_argument("--n-samples", type=int, default=10000) |
| p.add_argument("--n-steps", type=int, default=None, help="sampler resolution override") |
| p.add_argument("--seed-base", type=int, default=777) |
| p.add_argument("--chunk", type=int, default=None, help="draw the samples of a condition in chunks of this size") |
| p.add_argument("--truth-bank", default=None, help="truth bank with the x_0 of test conditions absent from the package") |
| p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| return p.parse_args(argv) |
|
|
|
|
| def _sha256(path: Path) -> str: |
| h = hashlib.sha256() |
| with open(path, "rb") as f: |
| for chunk in iter(lambda: f.read(1 << 24), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def main(argv=None) -> int: |
| args = parse_args(argv) |
| cfg = load_yaml(args.config) |
| cfg["data"]["packaged_dataset_path"] = str(Path(args.data).resolve()) + "/" |
| device = str(args.device) |
|
|
| allow_tf32 = cfg.get("training", {}).get("allow_tf32") |
| if allow_tf32 is not None: |
| torch.backends.cuda.matmul.allow_tf32 = bool(allow_tf32) |
| torch.backends.cudnn.allow_tf32 = bool(allow_tf32) |
| seed = int(cfg.get("seed", cfg["data"].get("seed", 0))) |
|
|
| ckpt_path = Path(args.checkpoint).resolve() |
| ckpt = torch.load(ckpt_path, map_location=device) |
| model = build_model(cfg["model"], cfg["method"], cfg["data"]) |
| ema_halflife = cfg.get("training", {}).get("ema_halflife_kimg", cfg["method"].get("ema_halflife_kimg")) |
| method_uses_ema = ema_halflife is not None and float(ema_halflife) > 0 |
| if "ema_state_dict" in ckpt and method_uses_ema: |
| model.load_state_dict(ckpt["ema_state_dict"]) |
| weights = "ema" |
| else: |
| model.load_state_dict(ckpt["model_state_dict"]) |
| weights = "raw" |
| model.to(device).eval() |
| print(f"checkpoint: {ckpt_path.name} (step {ckpt.get('step')}, {weights} weights)") |
|
|
| if args.n_steps is not None: |
| sampling = cfg["method"]["sampling"] |
| for key in ("n_sde_steps", "n_sde_steps_per_seg", "n_steps"): |
| if key in sampling: |
| break |
| else: |
| raise KeyError("--n-steps: method.sampling has none of n_sde_steps / n_sde_steps_per_seg / n_steps") |
| print(f"sampler resolution: method.sampling.{key} {sampling[key]} -> {args.n_steps}") |
| sampling[key] = args.n_steps |
| method = build_method(cfg["method"]) |
|
|
| split_x0, truth_clouds, gids = get_split_set(cfg, args.split, seed=seed, x0_bank=args.truth_bank) |
| if args.gids is not None: |
| wanted = sorted(int(g) for g in args.gids.split(",")) |
| missing = [g for g in wanted if g not in gids] |
| if missing: |
| raise ValueError(f"--gids not in split '{args.split}': {missing}") |
| sel = [gids.index(g) for g in wanted] |
| split_x0, truth_clouds, gids = split_x0[sel], truth_clouds[sel], wanted |
| split_x0, truth_clouds = split_x0.to(device), truth_clouds.to(device) |
| n_gids = split_x0.shape[0] |
| print(f"split {args.split}: {n_gids} conditions, {args.n_samples} samples each") |
|
|
| out = Path(args.out).resolve() |
| out.mkdir(parents=True, exist_ok=True) |
| split_offset = {"train": 1000, "val": 2000, "test": 3000}[args.split] |
| variants = method.variants_for_eval() |
| for v_idx, (label, sample_kwargs) in enumerate(variants.items()): |
| bundle_dir = out / label |
| bundle_dir.mkdir(parents=True, exist_ok=True) |
| torch.manual_seed(seed + split_offset + 100 * v_idx) |
| with torch.no_grad(): |
| for i in range(n_gids): |
| x0_i, gid = split_x0[i], int(gids[i]) |
| torch.manual_seed(int(args.seed_base) * 100003 + gid) |
| if args.chunk is not None and args.n_samples > int(args.chunk): |
| parts, left = [], int(args.n_samples) |
| while left > 0: |
| m = min(int(args.chunk), left) |
| parts.append(method.sample(model, x0_i, n_samples=m, **sample_kwargs).squeeze(0)) |
| left -= m |
| gen_cloud = torch.cat(parts, dim=0) |
| else: |
| gen_cloud = method.sample(model, x0_i, n_samples=int(args.n_samples), **sample_kwargs).squeeze(0) |
| np.savez(bundle_dir / f"group_{gid:04d}.npz", true_xT_norm=truth_clouds[i].cpu().numpy(), |
| gen_xT_norm=gen_cloud.cpu().numpy(), x0_norm=x0_i.cpu().numpy(), gid=np.int64(gid)) |
| print(f" {label}: condition {gid} ({i + 1}/{n_gids})", flush=True) |
|
|
| record = {"config": str(Path(args.config).resolve()), "checkpoint": str(ckpt_path), "checkpoint_sha256": _sha256(ckpt_path), |
| "weights": weights, "split": args.split, "gids": list(gids), "n_samples": int(args.n_samples), |
| "sampling": {k: cfg["method"]["sampling"].get(k) for k in ("n_sde_steps", "n_sde_steps_per_seg", "n_steps", |
| "n_segments", "noise_scale") if k in cfg["method"]["sampling"]}, |
| "seed_base": int(args.seed_base), "chunk": args.chunk, "variants": list(variants), "device": device, |
| "allow_tf32": allow_tf32, "at": datetime.now(timezone.utc).isoformat(timespec="seconds")} |
| (out / "SAMPLES.json").write_text(json.dumps(record, indent=1) + "\n") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
