Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """Inference script for MHD Edge Surrogate Model.""" | |
| import argparse | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) | |
| from mhd_model.dataset import ( | |
| MhdEdgeDataset, | |
| DataStats, | |
| apply_transform, | |
| CHANNEL_NAMES, | |
| ) | |
| from mhd_model.model import create_model | |
| def denormalize_to_raw(pred_norm, stats): | |
| """Convert normalized predictions back to raw values.""" | |
| pred_transformed = stats.denormalize_outputs(pred_norm) | |
| pred_raw = apply_transform( | |
| pred_transformed, stats.channel_transforms, | |
| asinh_scale=stats.asinh_scale, inverse=True, | |
| ) | |
| return pred_raw | |
| def load_model_and_stats(checkpoint_path, device="cpu"): | |
| """Load model and stats from checkpoint. | |
| Returns: | |
| model, stats, args_dict | |
| """ | |
| ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False) | |
| args_dict = ckpt.get("args", {}) | |
| stats = DataStats.from_dict(ckpt["stats"]) | |
| model = create_model( | |
| hidden_dim=args_dict.get("hidden_dim", 256), | |
| num_heads=args_dict.get("num_heads", 8), | |
| num_latent_tokens=args_dict.get("num_latent_tokens", 32), | |
| num_decoder_layers=args_dict.get("num_decoder_layers", 1), | |
| ).to(device) | |
| model.load_state_dict(ckpt["model_state_dict"]) | |
| model.eval() | |
| return model, stats, args_dict | |
| def predict(model, stats, inputs, query_coords, device="cpu"): | |
| """Run inference for a single sample. | |
| Args: | |
| model: Trained CfdDecoderOnly model | |
| stats: DataStats for normalization | |
| inputs: dict with keys {Ip, c_frac, d_scale, n_core, p_input} or tensor (5,) | |
| query_coords: (66, 26, 2) or (N, 2) R,Z coordinates | |
| device: torch device | |
| Returns: | |
| dict with predicted fields: | |
| te: (66, 26), ti: (66, 26), rad: (66, 26), | |
| ni: (66, 26, 2), up: (66, 26, 2), v2ce_0: (66, 26) | |
| """ | |
| from mhd_model.model import MhdCfdDecoderOnly | |
| # Prepare input params | |
| if isinstance(inputs, dict): | |
| input_value = torch.tensor([ | |
| inputs["Ip"], inputs["c_frac"], inputs["d_scale"], | |
| inputs["n_core"], inputs["p_input"], | |
| ], dtype=torch.float32) | |
| else: | |
| input_value = torch.as_tensor(inputs, dtype=torch.float32) | |
| input_value = stats.normalize_inputs(input_value).unsqueeze(0).to(device) # (1, 5) | |
| # Prepare query coords | |
| coords = torch.as_tensor(query_coords, dtype=torch.float32) | |
| if coords.dim() == 3: # (66, 26, 2) | |
| spatial_shape = coords.shape[:2] | |
| coords = coords.reshape(-1, 2) | |
| else: | |
| spatial_shape = None | |
| coords_norm = stats.normalize_coords(coords).unsqueeze(0).to(device) # (1, N, 2) | |
| N = coords_norm.shape[1] | |
| unbatch_idx, unbatch_select = MhdCfdDecoderOnly.create_unbatch_indices_from_lengths( | |
| [N], device=device, | |
| ) | |
| with torch.no_grad(): | |
| outputs = model( | |
| params=input_value, | |
| query_pos=coords_norm, | |
| unbatch_idx=unbatch_idx, | |
| unbatch_select=unbatch_select, | |
| ) | |
| pred_norm = outputs["pred"] # (N, 8) or (1, N, 8) | |
| if pred_norm.dim() == 2: | |
| pred_norm = pred_norm.unsqueeze(0) | |
| pred_raw = denormalize_to_raw(pred_norm, stats).squeeze(0).cpu() # (N, 8) | |
| # Reshape to spatial grid | |
| if spatial_shape is not None: | |
| pred_raw = pred_raw.reshape(*spatial_shape, 8) | |
| return { | |
| "te": pred_raw[..., 0].numpy(), # (66, 26) in eV | |
| "ti": pred_raw[..., 1].numpy(), # (66, 26) in eV | |
| "rad": pred_raw[..., 2].numpy(), # (66, 26) | |
| "ni": pred_raw[..., 3:5].numpy(), # (66, 26, 2) | |
| "up": pred_raw[..., 5:7].numpy(), # (66, 26, 2) | |
| "v2ce_0": pred_raw[..., 7].numpy(), # (66, 26) | |
| } | |
| return {"all_channels": pred_raw.numpy()} # (N, 8) | |
| def predict_batch(model, stats, inputs_batch, coords_batch, device="cpu"): | |
| """Run inference for a batch of samples. | |
| Args: | |
| inputs_batch: (B, 5) tensor of input parameters | |
| coords_batch: (B, 66, 26, 2) or (B, N, 2) coordinates | |
| Returns: | |
| pred_raw: (B, N, 8) raw predictions | |
| """ | |
| from mhd_model.model import MhdCfdDecoderOnly | |
| inputs_batch = torch.as_tensor(inputs_batch, dtype=torch.float32) | |
| coords_batch = torch.as_tensor(coords_batch, dtype=torch.float32) | |
| if coords_batch.dim() == 4: # (B, 66, 26, 2) | |
| B = coords_batch.shape[0] | |
| coords_batch = coords_batch.reshape(B, -1, 2) | |
| input_norm = stats.normalize_inputs(inputs_batch).to(device) | |
| coords_norm = stats.normalize_coords(coords_batch).to(device) | |
| B, N, _ = coords_norm.shape | |
| unbatch_idx, unbatch_select = MhdCfdDecoderOnly.create_unbatch_indices_from_lengths( | |
| [N] * B, device=device, | |
| ) | |
| with torch.no_grad(): | |
| outputs = model( | |
| params=input_norm, | |
| query_pos=coords_norm, | |
| unbatch_idx=unbatch_idx, | |
| unbatch_select=unbatch_select, | |
| ) | |
| pred_norm = outputs["pred"] | |
| if pred_norm.dim() == 2: | |
| pred_norm = pred_norm.view(B, N, -1) | |
| pred_raw = denormalize_to_raw(pred_norm, stats).cpu() | |
| return pred_raw.numpy() | |
| def visualize_comparison(pred, true_values, coords, sample_idx=0, save_path=None): | |
| """Create contour plots comparing predicted vs true fields. | |
| Args: | |
| pred: dict from predict() with field arrays | |
| true_values: (66, 26, 9) raw ground truth | |
| coords: (66, 26, 2) R,Z coordinates | |
| sample_idx: sample index for title | |
| save_path: path to save figure (optional) | |
| """ | |
| import matplotlib.pyplot as plt | |
| R = coords[..., 0] | |
| Z = coords[..., 1] | |
| scalar_fields = ["te", "ti", "rad"] | |
| fig, axes = plt.subplots(len(scalar_fields), 3, figsize=(18, 5 * len(scalar_fields))) | |
| true_map = { | |
| "te": true_values[..., 0], | |
| "ti": true_values[..., 1], | |
| "rad": true_values[..., 2], | |
| } | |
| for row, fname in enumerate(scalar_fields): | |
| pred_field = pred[fname] | |
| true_field = true_map[fname] | |
| error = np.abs(pred_field - true_field) | |
| # True | |
| ax = axes[row, 0] | |
| c = ax.contourf(R, Z, true_field, levels=50, cmap="viridis") | |
| plt.colorbar(c, ax=ax) | |
| ax.set_title(f"{fname} - True") | |
| ax.set_xlabel("R [m]") | |
| ax.set_ylabel("Z [m]") | |
| # Predicted | |
| ax = axes[row, 1] | |
| c = ax.contourf(R, Z, pred_field, levels=50, cmap="viridis") | |
| plt.colorbar(c, ax=ax) | |
| ax.set_title(f"{fname} - Predicted") | |
| ax.set_xlabel("R [m]") | |
| # Error | |
| ax = axes[row, 2] | |
| c = ax.contourf(R, Z, error, levels=50, cmap="magma") | |
| plt.colorbar(c, ax=ax) | |
| ax.set_title(f"{fname} - |Error|") | |
| ax.set_xlabel("R [m]") | |
| plt.suptitle(f"Sample {sample_idx}", fontsize=16) | |
| plt.tight_layout() | |
| if save_path: | |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") | |
| print(f"Saved figure to {save_path}") | |
| else: | |
| plt.show() | |
| plt.close(fig) | |
| def main(): | |
| parser = argparse.ArgumentParser(description="MHD Edge Surrogate Inference") | |
| parser.add_argument("--checkpoint", type=Path, required=True) | |
| parser.add_argument("--hdf5-path", type=Path, | |
| default=Path.home() / "UEDGE_MHD_KSTAR_dataset.hdf5") | |
| parser.add_argument("--sample-idx", type=int, default=0, | |
| help="Sample index from test set to evaluate") | |
| parser.add_argument("--output-dir", type=Path, default=Path("inference_outputs")) | |
| parser.add_argument("--device", type=str, | |
| default="cuda" if torch.cuda.is_available() else "cpu") | |
| args = parser.parse_args() | |
| device = torch.device(args.device) | |
| args.output_dir.mkdir(parents=True, exist_ok=True) | |
| # Load model | |
| print(f"Loading checkpoint: {args.checkpoint}") | |
| model, stats, ckpt_args = load_model_and_stats(args.checkpoint, device) | |
| # Load test dataset | |
| test_ds = MhdEdgeDataset(args.hdf5_path, split="test", stats=stats) | |
| print(f"Test set: {len(test_ds)} samples") | |
| # Get sample | |
| idx = args.sample_idx | |
| sample = test_ds[idx] | |
| real_idx = test_ds.indices[idx] | |
| # Get raw coords (before normalization) | |
| coords_raw = torch.from_numpy(test_ds.coords[real_idx]) # (66, 26, 2) | |
| inputs_raw = torch.from_numpy(test_ds.inputs[real_idx]) # (5,) | |
| print(f"\nSample {idx} (dataset index {real_idx}):") | |
| print(f" Inputs: {dict(zip(['Ip','c_frac','d_scale','n_core','p_input'], inputs_raw.tolist()))}") | |
| # Predict | |
| pred = predict(model, stats, inputs_raw, coords_raw, device=str(device)) | |
| # Ground truth | |
| true_values = test_ds.outputs[real_idx] # (66, 26, 8) | |
| # Per-channel errors | |
| print("\nPer-channel errors:") | |
| for c, name in enumerate(CHANNEL_NAMES): | |
| true_c = true_values[..., c] | |
| if name in pred: | |
| pred_c = pred[name] | |
| elif name in ["ni_0", "ni_1"]: | |
| pred_c = pred["ni"][..., int(name[-1])] | |
| elif name in ["up_0", "up_1"]: | |
| pred_c = pred["up"][..., int(name[-1])] | |
| else: | |
| pred_c = pred[name] | |
| mae = np.abs(pred_c - true_c).mean() | |
| rel = np.abs(pred_c - true_c).mean() / (np.abs(true_c).mean() + 1e-30) | |
| print(f" {name:8s}: MAE={mae:.4e} Rel={rel:.4f}") | |
| # Visualize | |
| save_path = args.output_dir / f"comparison_sample_{idx}.png" | |
| visualize_comparison(pred, true_values, coords_raw.numpy(), idx, save_path) | |
| if __name__ == "__main__": | |
| main() | |