File size: 4,165 Bytes
c059069 | 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 | """Run chunked inference over every point of all complete coarse grids."""
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.rf_climparam import FORMAT_VERSION, load_models
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 precipitation(targets, seconds):
# qT + qp tendencies are column-integrated with fixed positive layer mass weights.
moisture = targets[:, 48:96] + targets[:, 96:144]
weights = np.linspace(1.2, 0.2, 48, dtype=np.float32)
return np.maximum(-(moisture * weights).sum(1) * seconds, 0.0).astype(np.float32)
def predict_chunked(model, inputs, chunk_size):
prediction = np.empty((len(inputs), len(model.statistics["output_mean"])), dtype=np.float32)
for start in range(0, len(inputs), chunk_size):
stop = min(start + chunk_size, len(inputs))
prediction[start:stop] = model.predict(inputs[start:stop])
return prediction
def infer_field(source, pair, seconds, chunk_size, prefix):
ny, nx = map(int, source["grid_shape"])
count = ny * nx
linear = source["grid_row"].astype(np.int64) * nx + source["grid_column"].astype(np.int64)
if len(source["tend_inputs"]) != count or not np.array_equal(linear, np.arange(count)):
raise ValueError(f"{prefix} is not one complete reversibly flattened grid")
tend_prediction = predict_chunked(pair["rf_tend"], source["tend_inputs"], chunk_size)
diff_prediction = predict_chunked(pair["rf_diff"], source["diff_inputs"], chunk_size)
values = {
f"{prefix}_grid_shape": np.asarray((ny, nx), dtype=np.int32),
f"{prefix}_grid_row": source["grid_row"], f"{prefix}_grid_column": source["grid_column"],
f"{prefix}_time_hours": source["time_hours"],
f"{prefix}_latitude_deg": source["latitude_deg"],
f"{prefix}_longitude_deg": source["longitude_deg"],
f"{prefix}_tend_targets": source["tend_targets"],
f"{prefix}_tend_predictions": tend_prediction,
f"{prefix}_diff_targets": source["diff_targets"],
f"{prefix}_diff_predictions": diff_prediction,
f"{prefix}_precipitation_target_3h": precipitation(source["tend_targets"], seconds),
f"{prefix}_precipitation_prediction_3h": precipitation(tend_prediction, seconds),
}
if not all(np.isfinite(value).all() for value in values.values()):
raise ValueError(f"non-finite prediction data for {prefix}")
return values
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
checkpoint = load_checkpoint(ROOT / config["paths"]["checkpoint"])
models = load_models(checkpoint)
output = ROOT / config["paths"]["inference"]
output.parent.mkdir(parents=True, exist_ok=True)
seconds = int(config["evaluation"]["precipitation_seconds"])
chunk_size = int(config["evaluation"]["prediction_chunk_size"])
if chunk_size < 1:
raise ValueError("prediction_chunk_size must be positive")
scales = list(config["data"]["scales"])
packed = {}
variable_names = None
for scale in scales:
source = np.load(ROOT / config["data"]["root"] / f"{scale}.npz")
packed.update(infer_field(source, models[scale], seconds, chunk_size, scale))
if variable_names is None:
variable_names = {name: source[name] for name in ("tend_input_names", "tend_output_names",
"diff_input_names", "diff_output_names")}
source = np.load(ROOT / config["data"]["root"] / "x32_online_proxy.npz")
packed.update(infer_field(source, models["x32"], seconds, chunk_size, "online_x32_native"))
np.savez_compressed(output, format_version=np.array(FORMAT_VERSION), scales=np.asarray(scales),
online_scale=np.array("x32"), **variable_names, **packed)
print(f"predictions={output.relative_to(ROOT)} checkpoint_model={checkpoint['model_name']}")
if __name__ == "__main__":
main()
|