from __future__ import annotations import sys from pathlib import Path from typing import Any import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import torch from matplotlib.ticker import FormatStrFormatter PROJECT_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(PROJECT_ROOT)) sys.path.insert(0, str(Path(__file__).resolve().parent)) LOCAL_WORKSPACE = PROJECT_ROOT.parents[1] if (LOCAL_WORKSPACE / "onescience" / "src" / "onescience" / "datapipes").is_dir(): sys.path.insert(0, str(LOCAL_WORKSPACE)) from common import ( # noqa: E402 load_config, project_path, relative_l2, resolve_device, rollout, to_attr_dict, ) from model.factformer import FactFormer2D # noqa: E402 from onescience.datapipes.cfd import KolmogorovFlow2DDatapipe # noqa: E402 from onescience.distributed.manager import DistributedManager # noqa: E402 DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml" def prepare_config(config: dict[str, Any]) -> None: datapipe = config["datapipe"] datapipe["source"]["data_dir"] = str( project_path(datapipe["source"]["data_dir"], PROJECT_ROOT).resolve() ) stats_file = datapipe["data"].get("stats_file") if stats_file: datapipe["data"]["stats_file"] = str( project_path(stats_file, PROJECT_ROOT).resolve() ) def checkpoint_datapipe_config(checkpoint: dict[str, Any]) -> dict[str, Any]: if "datapipe_config" in checkpoint: return checkpoint["datapipe_config"] if "data_config" in checkpoint: return checkpoint["data_config"] raise KeyError("Checkpoint does not contain a datapipe configuration") def validate_checkpoint_config( datapipe_config: dict[str, Any], checkpoint: dict[str, Any] ) -> None: checkpoint_data = checkpoint_datapipe_config(checkpoint)["data"] current_data = datapipe_config["data"] for key in ( "resolution", "interval", "t_in", "t_out", "out_dim", "normalize", ): if current_data[key] != checkpoint_data[key]: raise ValueError( f"Config datapipe.data.{key}={current_data[key]} does not match " f"checkpoint value {checkpoint_data[key]}" ) def build_model( model_config: dict[str, Any], spatial_shape: tuple[int, int] ) -> FactFormer2D: return FactFormer2D( in_dim=int(model_config["in_dim"]), out_dim=int(model_config["out_dim"]), spatial_shape=spatial_shape, hidden_dim=int(model_config["hidden_dim"]), depth=int(model_config["depth"]), heads=int(model_config["heads"]), mlp_ratio=int(model_config["mlp_ratio"]), dropout=float(model_config["dropout"]), activation=str(model_config["activation"]), include_pos=bool(model_config["include_pos"]), space_dim=int(model_config["space_dim"]), latent_multiplier=float(model_config["latent_multiplier"]), max_latent_steps=int(model_config["max_latent_steps"]), ) def save_visualization( prediction: torch.Tensor, target: torch.Tensor, spatial_shape: tuple[int, int], output_path: Path, ) -> None: prediction = prediction.detach().cpu() target = target.detach().cpu() height, width = spatial_shape frame_count = prediction.shape[-1] frames = sorted({0, frame_count // 2, frame_count - 1}) sample_prediction = prediction[0].reshape(height, width, frame_count) sample_target = target[0].reshape(height, width, frame_count) rows = ( ("Target", sample_target), ("Prediction", sample_prediction), ("Abs Error", (sample_prediction - sample_target).abs()), ) figure, axes = plt.subplots( 3, len(frames), figsize=(4 * len(frames), 9), squeeze=False ) for row_index, (title, values) in enumerate(rows): for column_index, frame in enumerate(frames): axis = axes[row_index][column_index] image = axis.imshow(values[..., frame], origin="lower", cmap="viridis") axis.set_title(f"{title} t+{frame + 1}") axis.set_xticks([]) axis.set_yticks([]) colorbar = figure.colorbar(image, ax=axis, fraction=0.046, pad=0.04) colorbar.formatter = FormatStrFormatter("%.3g") colorbar.update_ticks() figure.tight_layout() figure.savefig(output_path, dpi=150) plt.close(figure) def main() -> None: config_path = DEFAULT_CONFIG.resolve() config = load_config(config_path) prepare_config(config) training = config["training"] inference = config["inference"] datapipe_config = config["datapipe"] num_samples = int(inference["num_samples"]) if num_samples < 1: raise ValueError("inference.num_samples must be positive") checkpoint_path = ( project_path(training["weight_dir"], PROJECT_ROOT).resolve() / str(training["checkpoint_name"]) ) if not checkpoint_path.exists(): raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) validate_checkpoint_config(datapipe_config, checkpoint) model_config = dict(checkpoint["model_config"]) spatial_shape = tuple(int(value) for value in checkpoint["spatial_shape"]) device = resolve_device(str(config["common"]["device"])) result_dir = project_path(inference["result_dir"], PROJECT_ROOT).resolve() print(f"Config: {config_path}") print(f"Checkpoint: {checkpoint_path}") print( "Data: " f"{Path(datapipe_config['source']['data_dir']) / datapipe_config['source']['file_name']}" ) print(f"Device: {device}") DistributedManager.initialize() try: datapipe = KolmogorovFlow2DDatapipe( to_attr_dict(datapipe_config), distributed=False, normalizer_state=checkpoint.get("normalizer"), ) test_loader, _ = datapipe.test_dataloader() if tuple(datapipe.spatial_shape) != spatial_shape: raise ValueError( f"Checkpoint grid {spatial_shape} does not match data grid " f"{datapipe.spatial_shape}" ) model = build_model(model_config, spatial_shape).to(device) model.load_state_dict(checkpoint["model_state"]) model.eval() batch = next(iter(test_loader)) pos = batch["pos"].to(device) state = batch["x"].to(device) target = batch["y"].to(device) with torch.no_grad(): prediction = rollout( model, pos, state, int(datapipe_config["data"]["t_out"]), int(datapipe_config["data"]["out_dim"]), int(model_config["max_latent_steps"]), ) prediction = datapipe.decode_solution(prediction) target = datapipe.decode_solution(target) error = relative_l2(prediction, target).mean().item() sample_count = min(num_samples, prediction.shape[0]) result_dir.mkdir(parents=True, exist_ok=True) tensor_path = result_dir / str(inference["tensor_name"]) image_path = result_dir / str(inference["figure_name"]) torch.save( { "prediction": prediction[:sample_count].cpu(), "target": target[:sample_count].cpu(), "spatial_shape": spatial_shape, "relative_l2": error, }, tensor_path, ) save_visualization(prediction[:1], target[:1], spatial_shape, image_path) print(f"Relative L2: {error:.6e}") print(f"Saved tensor: {tensor_path}") print(f"Saved figure: {image_path}") finally: DistributedManager.cleanup() if __name__ == "__main__": main()