import logging from functools import partial from pathlib import Path import hydra import matplotlib import numpy as np import torch from matplotlib import animation from matplotlib import pyplot as plt from common import load_config from onescience.datapipes.cfd import DeepMindLagrangianDatapipe, graph_update from onescience.distributed.manager import DistributedManager from onescience.launch.utils import load_checkpoint matplotlib.use("Agg") TYPE_TO_COLOR = { 0: "green", 3: "black", 5: "blue", 6: "gold", 7: "magenta", } def setup_logging(rank: int): level = logging.INFO if rank == 0 else logging.WARNING logging.basicConfig(level=level, format="%(asctime)s - %(levelname)s - %(message)s") return logging.getLogger("lagrangian_mgn.inference") def resolve_device(requested: str): if requested == "cpu": return torch.device("cpu") if requested in ("cuda", "gpu"): if not torch.cuda.is_available(): raise RuntimeError("Config requested CUDA, but torch.cuda.is_available() is false.") return torch.device("cuda:0") return torch.device("cuda:0" if torch.cuda.is_available() else "cpu") class MGNRollout: def __init__(self, cfg, device, logger): if int(cfg.test.batch_size) != 1: raise ValueError(f"Only batch size 1 is supported, got {cfg.test.batch_size}") self.cfg = cfg self.device = device self.logger = logger self.frame_skip = int(cfg.inference.frame_skip) self.datapipe = DeepMindLagrangianDatapipe(cfg, distributed=False) self.dataloader = self.datapipe.test_dataloader() self.dataset = self.datapipe.test_dataset self.dim = self.dataset.dim self.radius = self.dataset.radius self.dt = self.dataset.dt self.bounds = self.dataset.bounds self.num_history = self.dataset.num_history self.num_node_types = self.dataset.num_node_types self.model = hydra.utils.instantiate(cfg.model) if cfg.compile.enabled: self.model = torch.compile(self.model, **cfg.compile.args) self.model = self.model.to(device) epoch = load_checkpoint(cfg.resume_dir, models=self.model, device=device) if epoch == 0: logger.warning("No checkpoint found in %s; running with random weights.", cfg.resume_dir) self.model.eval() def compute_boundary_feature(self, position): dist = torch.cat( [position - self.bounds[0], self.bounds[1] - position], dim=-1 ) feat = torch.exp(-(dist**2) / self.radius**2) feat[dist > self.radius] = 0 return feat def boundary_clamp(self, position): min_bound = self.bounds[0] + 1e-3 max_bound = self.bounds[1] - 1e-3 return torch.clamp(position, min=min_bound, max=max_bound) def pack_inputs(self, position, vel_history, node_type): bound_feat = self.compute_boundary_feature(position) vel_hist = vel_history.permute(1, 0, 2).flatten(1) return torch.cat((position, vel_hist, bound_feat, node_type), dim=-1) @torch.inference_mode() def predict(self): pred_pos, gt_pos, node_type = [], [], None for graph in self.dataloader: graph = graph.to(self.device) if graph.ndata["t"][0].item() == 0: if pred_pos: yield torch.stack(pred_pos), torch.stack(gt_pos), node_type pred_pos, gt_pos = [], [] position, vel_history, node_type = self.dataset.unpack_inputs(graph) position = position.clone() vel_history = vel_history.clone() pred_pos.append(position) gt_pos.append(position) graph.ndata["x"] = self.pack_inputs(position, vel_history, node_type) graph.ndata["pos"] = position graph_update(graph, self.radius) acceleration = self.model(graph.ndata["x"], graph.edata["x"], graph) next_pos, next_vel = self.dataset.time_integrator( position=position, velocity=vel_history[-1], acceleration=acceleration, dt=self.dt, denormalize=True, ) next_pos = self.boundary_clamp(next_pos) next_vel_norm = ( next_vel - self.dataset.vel_mean.to(self.device) ) / self.dataset.vel_std.to(self.device) vel_history = torch.cat((vel_history[1:], next_vel_norm.unsqueeze(0)), dim=0) position = next_pos pred_pos.append(position) gt_next_pos, _, _ = self.dataset.unpack_targets(graph) gt_pos.append(gt_next_pos) if pred_pos: yield torch.stack(pred_pos), torch.stack(gt_pos), node_type def init_animation(subplot_kw=None): return plt.subplots(1, 2, figsize=(12, 6), subplot_kw=subplot_kw) def plot_particles_2d(ax, title, position, node_color, bounds): ax.cla() ax.set_aspect("equal") ax.scatter(position[:, 0], position[:, 1], c=node_color, s=10) ax.set_xlim(bounds[0], bounds[1]) ax.set_ylim(bounds[0], bounds[1]) ax.set_title(title) def plot_particles_3d(ax, title, position, node_color, bounds): ax.cla() ax.scatter(position[:, 2], position[:, 0], position[:, 1], c=node_color, s=10) ax.set_xlim(bounds[0], bounds[1]) ax.set_ylim(bounds[0], bounds[1]) ax.set_zlim(bounds[0], bounds[1]) ax.set_title(title) def animate_frame(num, plotter, fig, ax1, ax2, pred, gt, node_color, bounds, frame_skip): num *= frame_skip if num >= len(pred): return plotter(ax1, "Prediction", pred[num], node_color, bounds) plotter(ax2, "Ground Truth", gt[num], node_color, bounds) fig.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95, wspace=0.1) def save_animation(pred, gt, node_color, bounds, frame_skip, output_path, dim): if dim == 2: fig, (ax1, ax2) = init_animation() plotter = plot_particles_2d elif dim == 3: fig, (ax1, ax2) = init_animation(subplot_kw={"projection": "3d"}) plotter = plot_particles_3d else: raise ValueError(f"Unsupported dim: {dim}") ani = animation.FuncAnimation( fig, partial( animate_frame, plotter=plotter, fig=fig, ax1=ax1, ax2=ax2, pred=pred, gt=gt, node_color=node_color, bounds=bounds, frame_skip=frame_skip, ), frames=max((pred.shape[0] - 1) // frame_skip, 1), interval=100, ) ani.save(output_path, writer="pillow", fps=10) plt.close(fig) def main(): cfg = load_config() output_dir = Path(cfg.inference.output_dir) output_dir.mkdir(parents=True, exist_ok=True) DistributedManager.initialize() manager = DistributedManager() logger = setup_logging(manager.rank) device = resolve_device(str(cfg.test.device)) logger.info("Starting rollout on %s", device) rollout = MGNRollout(cfg, device, logger) mse_list = [] for index, (pred_pos, gt_pos, node_type) in enumerate(rollout.predict()): pred = pred_pos.cpu().numpy() gt = gt_pos.cpu().numpy() node_type_np = node_type.cpu().numpy() mse = float(np.mean((pred - gt) ** 2)) mse_list.append(mse) np.savez(output_dir / f"sequence_{index}.npz", prediction=pred, target=gt) logger.info("Sequence %s MSE %.4e", index, mse) if cfg.inference.save_animations: node_color = [ TYPE_TO_COLOR.get(node_id, "gray") for node_id in np.argmax(node_type_np, axis=1) ] save_animation( pred, gt, node_color, rollout.bounds, rollout.frame_skip, output_dir / f"animation_{index}.gif", int(cfg.dim), ) if index + 1 >= int(cfg.inference.max_sequences): break np.savez(output_dir / "rollout_metrics.npz", mse=np.asarray(mse_list, dtype=np.float32)) logger.info("Average MSE %.4e", float(np.mean(mse_list)) if mse_list else float("nan")) manager.cleanup() if __name__ == "__main__": main()