File size: 3,073 Bytes
87f2bd3 | 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 | """Restore every LOSO fold and emit both methods' 21 complete response maps."""
import sys
from pathlib import Path
import numpy as np
import torch
import yaml
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.climemu_s2l import FEATURE_COUNT, FORMAT_VERSION, GRID_SHAPE, DualRidge, SharedKernelGPR
def load_checkpoint(path):
try:
return torch.load(path, map_location="cpu", weights_only=False)
except TypeError:
return torch.load(path, map_location="cpu")
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
checkpoint = load_checkpoint(ROOT / config["paths"]["checkpoint"])
if checkpoint["format_version"] != FORMAT_VERSION or tuple(checkpoint["grid_shape"]) != GRID_SHAPE:
raise ValueError("checkpoint protocol mismatch")
x = checkpoint["short_response"].double()
y = checkpoint["long_response"].double()
required = {"model", "model_config", "format_version"}
if not required.issubset(checkpoint):
raise ValueError(f"checkpoint is missing standard fields: {sorted(required - checkpoint.keys())}")
predictions = {"ridge": [], "gpr": []}
for expected_fold, state in enumerate(checkpoint["model"]["folds"]):
if state["fold"] != expected_fold or state["held_out_scenario_id"] != checkpoint["scenario_ids"][expected_fold]:
raise ValueError("fold/scenario identity mismatch")
training = state["train_indices"]
ridge = DualRidge(float(state["ridge_alpha"])).fit(x[training], y[training])
gpr = SharedKernelGPR(state["gpr_kernel_mode"], config["model"]["gpr"]["jitter"])
gpr.load_hyperparameters(state["gpr_hyperparameters"]).restore_posterior(x[training], y[training])
predictions["ridge"].append(ridge.predict(x[expected_fold:expected_fold + 1]).detach().numpy()[0])
predictions["gpr"].append(gpr.predict(x[expected_fold:expected_fold + 1]).detach().numpy()[0])
ridge = np.asarray(predictions["ridge"], dtype=np.float32).reshape(21, *GRID_SHAPE)
gpr = np.asarray(predictions["gpr"], dtype=np.float32).reshape(21, *GRID_SHAPE)
if ridge.shape != (21, 145, 192) or not np.isfinite(gpr).all():
raise ValueError("inference did not produce 21 finite full-grid fields")
source = np.load(ROOT / config["data"]["path"])
output = ROOT / config["paths"]["inference"]
output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output, format_version=np.array(FORMAT_VERSION),
scenario_ids=np.asarray(checkpoint["scenario_ids"]),
latitude_deg=source["latitude_deg"], longitude_deg=source["longitude_deg"],
short_response=x.numpy().reshape(21, *GRID_SHAPE).astype(np.float32),
long_response=y.numpy().reshape(21, *GRID_SHAPE).astype(np.float32),
ridge_prediction=ridge, gpr_prediction=gpr)
print(f"predictions={output.relative_to(ROOT)} methods=2 shape={ridge.shape}")
if __name__ == "__main__":
main()
|