| import json |
| import logging |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| from model import HeteroGNS |
| from onescience.datapipes.cfd import BENODatapipe |
| from onescience.distributed.manager import DistributedManager |
| from onescience.utils.YParams import YParams |
| from onescience.utils.beno.util import record_data, to_np_array |
| from onescience.utils.beno.utilities import LpLoss, plot_data |
|
|
|
|
| logging.basicConfig(level=logging.INFO) |
| LOGGER = logging.getLogger("BENO_Inference") |
|
|
|
|
| def set_default_data_env(): |
| os.environ.setdefault("ONESCIENCE_BENO_DATA_DIR", str(PROJECT_ROOT / "data")) |
|
|
|
|
| def resolve_path(path_value): |
| path = Path(path_value) |
| return path if path.is_absolute() else PROJECT_ROOT / path |
|
|
|
|
| def load_config(): |
| set_default_data_env() |
| cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root") |
| cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir)) |
| cfg.datapipe.source.cache_dir = str(resolve_path(cfg.datapipe.source.cache_dir)) |
| cfg.inference.checkpoint_dir = str(resolve_path(cfg.inference.checkpoint_dir)) |
| cfg.inference.result_dir = str(resolve_path(cfg.inference.result_dir)) |
| cfg.inference.picture_dir = str(resolve_path(cfg.inference.picture_dir)) |
| cfg.inference.metrics_path = str(resolve_path(cfg.inference.metrics_path)) |
| if cfg.inference.checkpoint_path: |
| cfg.inference.checkpoint_path = str(resolve_path(cfg.inference.checkpoint_path)) |
| return cfg |
|
|
|
|
| def activation_from_name(name): |
| name = str(name).lower() |
| if name == "relu": |
| return nn.ReLU |
| if name == "elu": |
| return nn.ELU |
| if name == "leakyrelu": |
| return nn.LeakyReLU |
| return nn.SiLU |
|
|
|
|
| def build_model(model_cfg): |
| return HeteroGNS( |
| nnode_in_features=model_cfg.nnode_in_features, |
| nnode_out_features=model_cfg.nnode_out_features, |
| nedge_in_features=model_cfg.nedge_in_features, |
| latent_dim=model_cfg.get("latent_dim", model_cfg.get("width", 128)), |
| nmessage_passing_steps=model_cfg.get("nmessage_passing_steps", 10), |
| nmlp_layers=model_cfg.nmlp_layers, |
| mlp_hidden_dim=model_cfg.get("mlp_hidden_dim", model_cfg.get("width", 128)), |
| activation=activation_from_name(model_cfg.act), |
| boundary_dim=model_cfg.boundary_dim, |
| trans_layer=model_cfg.trans_layer, |
| ) |
|
|
|
|
| def select_device(device_name, dist): |
| device_name = str(device_name).lower() |
| if device_name == "cpu": |
| return torch.device("cpu") |
| if device_name == "cuda": |
| if not torch.cuda.is_available(): |
| raise RuntimeError("inference.device is cuda, but CUDA is not available.") |
| return torch.device(f"cuda:{dist.local_rank}") |
| return dist.device |
|
|
|
|
| def find_checkpoint(cfg): |
| if cfg.inference.checkpoint_path: |
| path = Path(cfg.inference.checkpoint_path) |
| return path if path.exists() else None |
| checkpoints = sorted(Path(cfg.inference.checkpoint_dir).glob("model_epoch_*.pt")) |
| return checkpoints[-1] if checkpoints else None |
|
|
|
|
| def main(): |
| cfg = load_config() |
| DistributedManager.initialize() |
| dist = DistributedManager() |
| device = select_device(cfg.inference.get("device", "auto"), dist) |
|
|
| datapipe = BENODatapipe(cfg, distributed=(dist.world_size > 1)) |
| test_loader, _ = datapipe.test_dataloader() |
| u_normalizer = datapipe.u_normalizer.to(device) |
| a_normalizer = datapipe.a_normalizer.to(device) |
| resolution = int(cfg.datapipe.data.resolution) |
|
|
| model = build_model(cfg.model).to(device) |
| checkpoint = find_checkpoint(cfg) |
| if checkpoint is not None: |
| LOGGER.info("Loading checkpoint: %s", checkpoint) |
| state_dict = torch.load(checkpoint, map_location=device) |
| model.load_state_dict(state_dict) |
| else: |
| LOGGER.warning("No checkpoint found. Running with random weights.") |
| model.eval() |
|
|
| result_dir = Path(cfg.inference.result_dir) |
| picture_dir = Path(cfg.inference.picture_dir) |
| result_dir.mkdir(parents=True, exist_ok=True) |
| picture_dir.mkdir(parents=True, exist_ok=True) |
|
|
| myloss = LpLoss(size_average=False) |
| analysis_record = {} |
| out_all = np.array([]) |
| label_all = np.array([]) |
| a_ori_all = np.array([]) |
| mask_all = np.array([]) |
| grid_all = np.array([]) |
|
|
| with torch.no_grad(): |
| for sample_id, data in enumerate(test_loader): |
| data = data.to(device) |
| out_indomain = model(data) |
|
|
| full_out = torch.zeros((resolution * resolution, 1), device=device) |
| full_label = torch.zeros((resolution * resolution), device=device) |
| full_input_a = torch.zeros((resolution * resolution, 10), device=device) |
| full_grid = torch.zeros((resolution * resolution, 2), device=device) |
| indices = data["G1"].sample_idx |
|
|
| full_out[indices] = out_indomain |
| full_label[indices] = data["G1+2"].y |
| full_input_a[indices, :] = data["G1"].x |
| full_grid[indices, :] = data["G1"].x[:, :2] |
|
|
| pred_decoded = u_normalizer.decode(full_out.view(1, -1)) |
| a_decoded = a_normalizer.decode(full_input_a[:, 2].view(1, -1)) |
| label_reshaped = full_label.view(1, -1) |
|
|
| cell_state_full = torch.zeros((1, resolution * resolution), device=device) |
| cell_state_full[0, :] = data["G1"].cell_state |
|
|
| l2_item = myloss(pred_decoded, label_reshaped).item() |
| mae_item = nn.L1Loss()(pred_decoded, label_reshaped).item() |
| record_data(analysis_record, [l2_item, mae_item], ["L2", "MAE"]) |
|
|
| np.savez( |
| result_dir / f"sample_{sample_id:04d}.npz", |
| predict=to_np_array(pred_decoded).reshape(resolution, resolution), |
| label=to_np_array(label_reshaped).reshape(resolution, resolution), |
| forcing=to_np_array(a_decoded).reshape(resolution, resolution), |
| mask=to_np_array(cell_state_full).reshape(resolution, resolution), |
| ) |
|
|
| out_all = np.append(out_all, to_np_array(pred_decoded)) |
| label_all = np.append(label_all, to_np_array(label_reshaped)) |
| a_ori_all = np.append(a_ori_all, to_np_array(a_decoded)) |
| mask_all = np.append(mask_all, to_np_array(cell_state_full)) |
| grid_all = np.append(grid_all, to_np_array(full_grid.unsqueeze(0))) |
|
|
| metrics = { |
| "mean_l2": float(np.mean(analysis_record["L2"])), |
| "std_l2": float(np.std(analysis_record["L2"])), |
| "mean_mae": float(np.mean(analysis_record["MAE"])), |
| "num_samples": int(len(analysis_record["L2"])), |
| } |
| metrics_path = Path(cfg.inference.metrics_path) |
| metrics_path.parent.mkdir(parents=True, exist_ok=True) |
| metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8") |
| LOGGER.info("Metrics: %s", metrics) |
|
|
| plot_samples = min(int(cfg.inference.num_visualize), metrics["num_samples"]) |
| if plot_samples >= 2: |
| plot_path = picture_dir / "forcing_solution_comparison.png" |
| plot_data( |
| predict_term=out_all, |
| true_term=label_all, |
| forcing_term=a_ori_all, |
| forcing_mask=mask_all, |
| grid_info=grid_all, |
| resolution=resolution, |
| num_samples=plot_samples, |
| interpolation="bilinear", |
| save_path=str(plot_path), |
| ) |
| LOGGER.info("Plot saved to %s", plot_path) |
| elif plot_samples == 1: |
| LOGGER.info("Skipping plot_data for one sample because the runtime helper expects a 2D axes grid.") |
|
|
| dist.cleanup() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
