| """Run chunked inference on the full 251 x 251 x 44 x 3 paper-resolution domain.""" |
|
|
| import importlib.util |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def load_model_module(): |
| spec = importlib.util.spec_from_file_location("pinn_tc_model", ROOT / "model/pinn-tc.py") |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| def coordinates_for_indices(indices, y_axis, x_axis, times, pressures): |
| x_count, p_count, t_count = len(x_axis), len(pressures), len(times) |
| y_index = indices // (x_count * p_count * t_count) |
| remainder = indices % (x_count * p_count * t_count) |
| x_index = remainder // (p_count * t_count) |
| remainder %= p_count * t_count |
| p_index = remainder // t_count |
| t_index = remainder % t_count |
| return np.column_stack((y_axis[y_index], x_axis[x_index], times[t_index], pressures[p_index])).astype(np.float32) |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| module = load_model_module() |
| requested = config["runtime"]["device"] |
| device = torch.device("cuda" if requested == "auto" and torch.cuda.is_available() else ("cpu" if requested == "auto" else requested)) |
| checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=False) |
| if checkpoint.get("format_version") != config["data"]["format_version"]: |
| raise ValueError("checkpoint format_version mismatch") |
| if checkpoint["input_order"] != list(module.INPUT_ORDER) or checkpoint["output_order"] != list(module.OUTPUT_ORDER): |
| raise ValueError("checkpoint coordinate or variable protocol mismatch") |
| model = module.PINNTC(**checkpoint["model_config"]).to(device) |
| model.load_state_dict(checkpoint["model_state_dict"]); model.eval() |
| data = config["data"] |
| y_axis = np.linspace(-data["horizontal_extent_m"], data["horizontal_extent_m"], data["grid_points"], dtype=np.float32) |
| x_axis = y_axis.copy() |
| times = np.asarray(data["observation_times_s"], dtype=np.float32) |
| pressures = np.linspace(data["pressure_min_pa"], data["pressure_max_pa"], data["pressure_levels"], dtype=np.float32) |
| shape = (len(y_axis), len(x_axis), len(pressures), len(times), 4) |
| expected = (251, 251, 44, 3, 4) |
| if shape != expected: |
| raise ValueError(f"dense grid must remain {expected}, got {shape}") |
| output = ROOT / config["paths"]["prediction"] |
| output.parent.mkdir(parents=True, exist_ok=True) |
| temporary = output.with_suffix(".work.npy") |
| prediction = np.lib.format.open_memmap(temporary, mode="w+", dtype=np.float32, shape=shape) |
| flat = prediction.reshape(-1, 4) |
| chunk_size = int(config["runtime"]["inference_chunk_size"]) |
| with torch.inference_mode(): |
| for start in range(0, len(flat), chunk_size): |
| stop = min(start + chunk_size, len(flat)) |
| coordinates = coordinates_for_indices(np.arange(start, stop, dtype=np.int64), y_axis, x_axis, times, pressures) |
| values = model(torch.from_numpy(coordinates).to(device)).cpu().numpy().astype(np.float32) |
| if not np.isfinite(values).all(): |
| raise FloatingPointError(f"non-finite prediction in rows {start}:{stop}") |
| flat[start:stop] = values |
| prediction.flush() |
| np.savez(output, predictions=prediction, y=y_axis, x=x_axis, time=times, pressure=pressures, |
| input_order=np.asarray(module.INPUT_ORDER), output_order=np.asarray(module.OUTPUT_ORDER), |
| format_version=np.asarray(config["data"]["format_version"]), layout=np.asarray("YXPTV")) |
| del prediction |
| temporary.unlink() |
| print(f"predictions={output.relative_to(ROOT)} shape={shape} bytes={output.stat().st_size}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
