File size: 2,693 Bytes
20cdc88 | 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 | """Fit statistical CME networks and the precipitation-constraint GP."""
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 (FORMAT_VERSION, LaggedPartialCorrelationCME,
PrecipitationConstraintGP, asymmetric_f1)
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
np.random.seed(int(config["seed"]))
data = np.load(ROOT / config["paths"]["dataset"])
if str(data["format_version"]) != config["data"]["format_version"]:
raise ValueError("incompatible synthetic data format")
estimator = LaggedPartialCorrelationCME(**config["model"])
references = [estimator.fit(series) for series in data["reference_series"]]
model_networks, scores = [], []
for model_series in data["model_series"]:
networks = [estimator.fit(series) for series in model_series]
model_networks.append(networks)
scores.append(float(np.mean([asymmetric_f1(ref, net, int(config["evaluation"]["lag_tolerance"]))["f1"]
for ref, net in zip(references, networks)])))
gp = PrecipitationConstraintGP(int(config["seed"])).fit(np.asarray(scores), data["delta_precipitation"])
checkpoint = {
"version": FORMAT_VERSION, "config": config, "network_config": estimator.config(),
"reference_networks": [network.state_dict() for network in references],
"model_networks": [[network.state_dict() for network in networks] for networks in model_networks],
"model_f1": torch.tensor(scores), "gp": gp.state_dict(),
"metadata": {"method": "lagged target-history conditional regression ParCorr approximation",
"gradient_training": False, "tensor_layout": "source,target,lag", "time_step_days": 3}
}
checkpoint_path = ROOT / config["paths"]["checkpoint"]
metrics_path = ROOT / config["paths"]["training_metrics"]
checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
metrics_path.parent.mkdir(parents=True, exist_ok=True)
torch.save(checkpoint, checkpoint_path)
metrics_path.write_text(json.dumps({"model_f1": scores, "gp_kernel": str(gp.model.kernel_),
"paper_alpha": config["paper_model"]["alpha"],
"engineering_alpha": config["model"]["alpha"]}, indent=2) + "\n")
print(f"checkpoint={checkpoint_path.relative_to(ROOT)} networks={len(references) * (1 + len(model_networks))} gp_samples={len(scores)}")
if __name__ == "__main__":
main()
|