| """Reload the CME checkpoint and export complete network and projection outputs.""" |
|
|
| import json |
| 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.causalmodelevaluation import CausalNetwork, FORMAT_VERSION, PrecipitationConstraintGP |
|
|
|
|
| def stack_networks(states, field): |
| return np.stack([[getattr(CausalNetwork.from_state_dict(state), field).numpy() for state in group] for group in states]) |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location="cpu", weights_only=False) |
| if checkpoint["version"] != FORMAT_VERSION: |
| raise ValueError("unsupported checkpoint version") |
| reference_states = [checkpoint["reference_networks"]] |
| model_states = checkpoint["model_networks"] |
| f1 = checkpoint["model_f1"].numpy() |
| gp = PrecipitationConstraintGP.from_state_dict(checkpoint["gp"], int(config["seed"])) |
| query = np.sort(np.unique(np.append(f1, float(config["evaluation"]["reference_f1_for_projection"])))) |
| mean, lower, upper = gp.predict(query) |
| source = np.load(ROOT / config["paths"]["dataset"]) |
| output = ROOT / config["paths"]["inference"] |
| output.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed( |
| output, edges=stack_networks(model_states, "edges"), pvalues=stack_networks(model_states, "pvalues"), |
| mci=stack_networks(model_states, "mci"), reference_edges=stack_networks(reference_states, "edges")[0], |
| reference_pvalues=stack_networks(reference_states, "pvalues")[0], |
| reference_mci=stack_networks(reference_states, "mci")[0], model_f1=f1, |
| reference_precipitation=source["reference_precipitation"], model_precipitation=source["model_precipitation"], |
| delta_precipitation=source["delta_precipitation"], latitude_degrees=source["latitude_degrees"], |
| longitude_degrees=source["longitude_degrees"], gp_query_f1=query, gp_mean_delta_precipitation=mean, |
| gp_lower_95=lower, gp_upper_95=upper, seasons=source["seasons"], |
| network_metadata=np.asarray(json.dumps(checkpoint["metadata"])), |
| projection_metadata=np.asarray(json.dumps({"kernel": str(gp.model.kernel_), "confidence": 0.95, |
| "input": "CME asymmetric F1", "output": "delta precipitation"})), |
| format_version=np.asarray(FORMAT_VERSION)) |
| print(f"inference={output.relative_to(ROOT)} edges={stack_networks(model_states, 'edges').shape} reference={stack_networks(reference_states, 'edges')[0].shape}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
