| """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() |
|
|
