File size: 2,931 Bytes
e2703dc ede74c0 e2703dc ede74c0 e2703dc ede74c0 e2703dc ede74c0 e2703dc ede74c0 e2703dc ede74c0 e2703dc ede74c0 e2703dc ede74c0 e2703dc ede74c0 e2703dc ede74c0 e2703dc ede74c0 | 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 | from __future__ import annotations
import argparse
import numpy as np
import torch
from onescience.utils.pinnsformer_util import get_data, make_time_sequence
from common import (
build_model,
exact_reaction_solution,
load_config,
project_path,
relative_errors,
select_device,
)
def main() -> None:
parser = argparse.ArgumentParser(description="Run PINNsformer inference for the 1D reaction equation.")
parser.add_argument("--config", default=None, help="Path to config.yaml.")
parser.add_argument("--checkpoint", default=None, help="Path to a trained checkpoint.")
parser.add_argument("--device", default=None, help="Override runtime.device.")
args = parser.parse_args()
cfg = load_config(args.config)
device = select_device(args.device or cfg["runtime"]["device"])
checkpoint = project_path(args.checkpoint or cfg["training"]["checkpoint"])
if not checkpoint.exists():
raise FileNotFoundError(f"Checkpoint not found: {checkpoint}. Run scripts/train.py first.")
model = build_model(cfg).to(device)
try:
state = torch.load(checkpoint, map_location=device, weights_only=True)
except TypeError:
state = torch.load(checkpoint, map_location=device)
model.load_state_dict(state.get("model_state_dict", state))
model.eval()
data_cfg = cfg["data"]
res, _, _, _, _ = get_data(
data_cfg["x_range"],
data_cfg["t_range"],
int(data_cfg["x_num"]),
int(data_cfg["t_num"]),
)
seq = make_time_sequence(
res,
num_step=int(data_cfg["sequence"]["num_step"]),
step=float(data_cfg["sequence"]["step"]),
)
res_tensor = torch.tensor(seq, dtype=torch.float32, requires_grad=False, device=device)
x_test, t_test = res_tensor[:, :, 0:1], res_tensor[:, :, 1:2]
with torch.no_grad():
pred = model(x_test, t_test)[:, 0:1].detach().cpu().numpy().reshape(
int(data_cfg["t_num"]),
int(data_cfg["x_num"]),
)
x = res[:, 0].reshape(int(data_cfg["t_num"]), int(data_cfg["x_num"]))
t = res[:, 1].reshape(int(data_cfg["t_num"]), int(data_cfg["x_num"]))
target = exact_reaction_solution(x, t, cfg).astype(np.float32)
metrics = relative_errors(pred, target)
output_path = project_path(cfg["paths"]["prediction"])
output_path.parent.mkdir(parents=True, exist_ok=True)
np.savez(
output_path,
x=x.astype(np.float32),
t=t.astype(np.float32),
prediction=pred.astype(np.float32),
target=target,
absolute_error=np.abs(pred - target).astype(np.float32),
relative_l1=metrics["relative_l1"],
relative_l2=metrics["relative_l2"],
)
print(f"prediction saved to {output_path}")
print(f"relative L1 error: {metrics['relative_l1']:.6f}")
print(f"relative L2 error: {metrics['relative_l2']:.6f}")
if __name__ == "__main__":
main()
|