File size: 3,827 Bytes
2f3c9e4 | 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 | """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()
|