File size: 4,382 Bytes
28d9c3b | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | from __future__ import annotations
import sys
from collections.abc import Mapping
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(Path(__file__).resolve().parent))
from common import ( # noqa: E402
build_laplace_data,
checkpoint_state,
load_config,
project_path,
relative_l2,
resolve_device,
resolve_dtype,
)
from model.bpinn import build_model, posterior_predict # noqa: E402
DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
def main() -> None:
config_path = DEFAULT_CONFIG.resolve()
config = load_config(config_path)
common = config["common"]
device = resolve_device(str(common["device"]))
dtype = resolve_dtype(str(common["dtype"]))
weight_dir = project_path(common["weight_dir"], PROJECT_ROOT)
result_dir = project_path(common["result_dir"], PROJECT_ROOT)
refined_path = weight_dir / config["training"]["refined_checkpoint_name"]
base_path = weight_dir / config["training"]["checkpoint_name"]
checkpoint_path = refined_path if refined_path.is_file() else base_path
if not checkpoint_path.is_file():
raise FileNotFoundError(
f"checkpoint not found: {checkpoint_path}. Run scripts/train.py first."
)
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
if not isinstance(checkpoint, Mapping):
raise ValueError(f"invalid BPINN checkpoint: {checkpoint_path}")
state, metadata = checkpoint_state(checkpoint)
model_config = metadata.get("model_config", config["model"])
data_config = dict(metadata.get("data_config", config["data"]))
seed = int(metadata.get("seed", common["seed"]))
data = build_laplace_data(data_config, seed, device, dtype)
model = build_model(model_config, dtype=dtype).to(device=device, dtype=dtype)
model.load_state_dict(state, strict=True)
model.eval()
posterior_states = metadata.get("posterior_states") or [state]
mean, standard_deviation, samples = posterior_predict(
model, posterior_states, data["x_test"]
)
mean_numpy = mean.cpu().numpy()
std_numpy = standard_deviation.cpu().numpy()
exact_numpy = data["u_test"].cpu().numpy()
x_numpy = data["x_test"].cpu().numpy()
error = relative_l2(mean_numpy, exact_numpy)
rmse = float(np.sqrt(np.mean((mean_numpy - exact_numpy) ** 2)))
max_error = float(np.max(np.abs(mean_numpy - exact_numpy)))
result_dir.mkdir(parents=True, exist_ok=True)
predictions_path = result_dir / config["inference"]["predictions_name"]
figure_path = result_dir / config["inference"]["figure_name"]
np.savez_compressed(
predictions_path,
x=x_numpy,
exact=exact_numpy,
mean=mean_numpy,
std=std_numpy,
samples=samples.cpu().numpy(),
relative_l2=error,
rmse=rmse,
max_abs_error=max_error,
)
x_axis = x_numpy.reshape(-1)
exact_axis = exact_numpy.reshape(-1)
mean_axis = mean_numpy.reshape(-1)
std_axis = std_numpy.reshape(-1)
absolute_error = np.abs(mean_axis - exact_axis)
figure, axes = plt.subplots(1, 2, figsize=(12, 4))
axes[0].plot(x_axis, exact_axis, "k-", linewidth=1.5, label="Exact")
axes[0].plot(x_axis, mean_axis, "r--", linewidth=1.5, label="BPINN")
if np.any(std_axis > 0):
axes[0].fill_between(
x_axis,
mean_axis - 2.0 * std_axis,
mean_axis + 2.0 * std_axis,
color="red",
alpha=0.2,
label="2 std",
)
axes[0].set_xlabel("x")
axes[0].set_ylabel("u")
axes[0].legend()
axes[1].semilogy(x_axis, np.maximum(absolute_error, 1.0e-16), "b-")
axes[1].set_xlabel("x")
axes[1].set_ylabel("Absolute error")
figure.tight_layout()
figure.savefig(figure_path, dpi=150)
plt.close(figure)
print(f"Config: {config_path}")
print(f"Checkpoint: {checkpoint_path}")
print(f"Device: {device}")
print(f"Posterior states: {len(posterior_states)}")
print(f"Relative L2={error:.6e}, RMSE={rmse:.6e}, MaxAbs={max_error:.6e}")
print(f"Predictions: {predictions_path}")
print(f"Plot: {figure_path}")
if __name__ == "__main__":
main()
|