File size: 5,778 Bytes
4ebcd24 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """CPU-only export regression using an existing causal batched issue packet.
This checks packaging/field consistency, not forecast skill. It never calls MPS,
changes neural weights, or writes into historical archive/training directories.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import shutil
import sys
from pathlib import Path
from unittest.mock import patch
import numpy as np
import torch
ROOT = Path(__file__).resolve().parents[1]
MODEL = ROOT / "models/trackformer_1_2_field"
sys.path.insert(0, str(MODEL))
import predict
from model import CoreForecaster
from baseline_model import base
from plot_pressure import render_pressure_map
def sha(path):
with Path(path).open("rb") as stream:
return hashlib.file_digest(stream, "sha256").hexdigest()
def verify(batched_packet, weights, work):
work = work.resolve()
if not work.is_relative_to(Path("/Volumes/D")):
raise ValueError("Diagnostic artifacts must remain on D")
work.mkdir(parents=True, exist_ok=True)
directory = work / "inference"
directory.mkdir(exist_ok=True)
for source in MODEL.iterdir():
if source.is_file() and source.suffix in (".py", ".md", ".json"):
shutil.copy2(source, directory / source.name)
shutil.copy2(weights, directory / "weights.pt")
metadata = json.loads((directory / "manifest.json").read_text())
for name, expected in metadata["source_module_sha256"].items():
if sha(directory / name) != expected:
raise ValueError("Frozen neural module changed: " + name)
with np.load(batched_packet, allow_pickle=False) as source:
arrays = {key: np.asarray(source[key])[0] for key in predict.INPUT_SHAPES}
arrays["history_time_ns"] = np.asarray(source["history_time_ns"], dtype=np.int64)
arrays["issue_time_ns"] = np.asarray(arrays["history_time_ns"][-1], dtype=np.int64)
packet = work / "causal_issue_packet.npz"
np.savez_compressed(packet, **arrays)
captured, origins = [], []
class CapturingModel(CoreForecaster):
def step(self, state):
next_state, out = super().step(state)
captured.append(out)
origins.append(next_state["origin"])
return next_state, out
torch.set_num_threads(2)
torch.set_num_interop_threads(1)
with patch.object(predict, "CoreForecaster", CapturingModel):
result = predict.forecast(packet, directory, "cpu")
contract = metadata["data_contract"]
scale, offset = contract["normalization"]["std"][0], contract["normalization"]["mean"][0]
reference = {
"track_lat_lon": torch.stack([o["center"][0] for o in captured]).numpy(),
"central_pressure_hpa": torch.stack([o["pressure"][0] for o in captured]).numpy(),
"basin_mslp_hpa": torch.stack([o["global"][0,0] for o in captured]).numpy()*scale+offset,
"regional_mslp_hpa": torch.stack([o["regional"][0,0] for o in captured]).numpy()*scale+offset,
"maximum_wind_auxiliary_kt": torch.stack([o["vmax"][0] for o in captured]).numpy(),
}
for key, expected in reference.items():
np.testing.assert_array_equal(result[key], expected, err_msg="Original readout changed: " + key)
alignment = []
for i, (out, origin) in enumerate(zip(captured, origins)):
# The original model's own geographic sampler, not a centre/minimum fit.
grid = CoreForecaster.core_grid(None, out["center"][:,0,None,None], out["center"][:,1,None,None], origin)
sampled = base.sample_field(torch.from_numpy(result["core_mslp_hpa"][i])[None,None], grid)[0,0,0,0]
alignment.append(abs(float(sampled) - float(result["central_pressure_hpa"][i])))
np.testing.assert_array_equal(result["core_latitude_deg"][i], out["core_lat"][0].numpy())
np.testing.assert_array_equal(result["core_longitude_deg"][i], out["core_lon"][0].numpy())
np.testing.assert_array_equal(result["core_valid"][i], out["core_valid"][0,0].numpy())
if max(alignment) > 1e-3:
raise ValueError("Exported core does not reproduce its original central-pressure readout")
np.savez_compressed(work / "forecast.npz", **result)
for lead in (6, 120):
render_pressure_map(result, work / f"pressure_{lead:03d}h.png", lead, 4)
receipt = {
"state": "verified_cpu_export_regression",
"scientific_benchmark": False,
"model": "Trackformer 1.2", "members": 1,
"device": "cpu", "forecast_steps": 20,
"issue_time_ns": int(result["issue_time_ns"]),
"core_shape": list(result["core_mslp_hpa"].shape),
"core_min_hpa": float(result["core_mslp_hpa"][result["core_valid"]].min()),
"core_max_hpa": float(result["core_mslp_hpa"][result["core_valid"]].max()),
"central_pressure_alignment_max_error_hpa": max(alignment),
"original_outputs_bit_identical": list(reference),
"coordinates_and_masks_match_forward_outputs": True,
"inference_weights_sha256": sha(weights),
"source_packet_sha256": sha(batched_packet),
"forecast_sha256": sha(work / "forecast.npz"),
"native_detail_history_available": bool(arrays["detail_available"][0]),
"core_is_learned_reconstruction_not_native_observations": True,
}
(work / "verification.json").write_text(json.dumps(receipt, indent=2) + "\n")
print(json.dumps(receipt, indent=2))
if __name__ == "__main__":
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--batched-packet", type=Path, required=True)
parser.add_argument("--weights", type=Path, required=True)
parser.add_argument("--work", type=Path, required=True)
args = parser.parse_args()
verify(args.batched_packet, args.weights, args.work)
|