import math import torch import numpy as np from torch import nn from dataclasses import dataclass from onescience.models.meta import ModelMetaData from onescience.modules.embedding.panguembedding import PanguEmbedding from onescience.modules.fuser.pangufuser import PanguFuser from onescience.modules.recovery.pangupatchrecovery import PanguPatchRecovery from onescience.modules.sample.pangudownsample import PanguDownSample from onescience.modules.sample.panguupsample import PanguUpSample @dataclass class MetaData(ModelMetaData): name: str = "Pangu" # Optimization jit: bool = False # ONNX Ops Conflict cuda_graphs: bool = True amp: bool = True # Inference onnx_cpu: bool = False # No FFT op on CPU onnx_gpu: bool = True onnx_runtime: bool = True # Physics informed var_dim: int = 1 func_torch: bool = False auto_grad: bool = False class Pangu(nn.Module): """ Pangu-Weather 的主模型实现。 该模型使用以下组件完成编码、主干特征提取与输出恢复： - `OneEmbedding(style="PanguEmbedding")` - 将 surface 分支与 upper-air 分支分别映射到 patch 特征空间 - `OneFuser(style="PanguFuser")` - 在统一的三维 token 网格上完成主干特征融合 - `OneSample(style="PanguDownSample" / "PanguUpSample")` - 在主干中完成空间尺度变换 - `OneRecovery(style="PanguPatchRecovery")` - 将 patch 级特征恢复为目标物理场 在当前实现中： - 输入包含 7 个 surface 通道，其中 4 个是待预测地表变量，3 个是静态掩码 - 输入包含 5 个 upper-air 变量，每个变量对应 13 个气压层 - surface patch token 会先在 `PressureLevels` 维补成长度为 1 - 然后与 upper-air patch token 沿 `PressureLevels` 维拼接 - 之后统一送入三维 `PanguFuser` 主干 Reference: - `Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast` - https://arxiv.org/abs/2211.02556 Args: img_size (tuple[int, int]): 输入空间尺寸 `(Height, Width)`。 patch_size (tuple[int, int, int]): patch 切分尺寸 `(PatchPressureLevels, PatchHeight, PatchWidth)`。 embed_dim (int): patch embedding 后的特征维度。 num_heads (tuple[int, int, int, int]): 四个主干阶段对应的注意力头数。 window_size (tuple[int, int, int]): 三维窗口大小 `(PressureLevelsWindow, HeightWindow, WidthWindow)`。 """ def __init__( self, img_size=(721, 1440), patch_size=(2, 4, 4), embed_dim=192, num_heads=(6, 12, 12, 6), window_size=(2, 6, 12), ): super().__init__() drop_path = np.linspace(0, 0.2, 8).tolist() # Surface input contains 4 predicted variables and 3 static masks: # topography mask, land-sea mask, and soil type mask. self.patchembed2d = PanguEmbedding( img_size=img_size, patch_size=patch_size[1:], Variables=7, embed_dim=embed_dim, ) self.patchembed3d = PanguEmbedding( img_size=(13, *img_size), patch_size=patch_size, Variables=5, embed_dim=embed_dim, ) patched_input_shape = ( 8, math.ceil(img_size[0] / patch_size[1]), math.ceil(img_size[1] / patch_size[2]), ) self.layer1 = PanguFuser( dim=embed_dim, input_resolution=patched_input_shape, depth=2, num_heads=num_heads[0], window_size=window_size, drop_path=drop_path[:2], ) patched_downsampled_shape = ( 8, math.ceil(patched_input_shape[1] / 2), math.ceil(patched_input_shape[2] / 2), ) self.downsample = PanguDownSample( in_dim=embed_dim, input_resolution=patched_input_shape, output_resolution=patched_downsampled_shape, ) self.layer2 = PanguFuser( dim=embed_dim * 2, input_resolution=patched_downsampled_shape, depth=6, num_heads=num_heads[1], window_size=window_size, drop_path=drop_path[2:], ) self.layer3 = PanguFuser( dim=embed_dim * 2, input_resolution=patched_downsampled_shape, depth=6, num_heads=num_heads[2], window_size=window_size, drop_path=drop_path[2:], ) self.upsample = PanguUpSample( in_dim=embed_dim * 2, out_dim=embed_dim, input_resolution=patched_downsampled_shape, output_resolution=patched_input_shape, ) self.layer4 = PanguFuser( dim=embed_dim, input_resolution=patched_input_shape, depth=2, num_heads=num_heads[3], window_size=window_size, drop_path=drop_path[:2], ) # The recovered surface output contains only the 4 prognostic surface variables. # Static masks are input-only features and are not part of the prediction target. self.patchrecovery2d = PanguPatchRecovery( img_size=(721, 1440), patch_size=(4, 4), in_chans=embed_dim * 2, out_chans=4, ) self.patchrecovery3d = PanguPatchRecovery( img_size=(13, 721, 1440), patch_size=(2, 4, 4), in_chans=embed_dim * 2, out_chans=5, ) def forward(self, x): """ Args: x (torch.Tensor): Input tensor with shape `(Batch, 4 + 3 + 5 * 13, Height, Width)`. Channel layout: - first 4 channels: prognostic surface variables - next 3 channels: static masks - remaining `5 * 13` channels: upper-air variables flattened over pressure levels Returns: tuple[torch.Tensor, torch.Tensor]: - surface output: `(Batch, 4, Height, Width)` - upper-air output: `(Batch, 5, 13, Height, Width)` """ SurfaceInput = x[:, :7, :, :] UpperAirInput = x[:, 7:, :, :].reshape(x.shape[0], 5, 13, x.shape[2], x.shape[3]) SurfaceFeatures = self.patchembed2d(SurfaceInput) UpperAirFeatures = self.patchembed3d(UpperAirInput) CombinedFeatures = torch.concat( [SurfaceFeatures.unsqueeze(2), UpperAirFeatures], dim=2 ) Batch, Channels, PressureLevels, Height, Width = CombinedFeatures.shape Tokens = CombinedFeatures.reshape(Batch, Channels, -1).transpose(1, 2) Tokens = self.layer1(Tokens) SkipTokens = Tokens Tokens = self.downsample(Tokens) Tokens = self.layer2(Tokens) Tokens = self.layer3(Tokens) Tokens = self.upsample(Tokens) Tokens = self.layer4(Tokens) OutputFeatures = torch.concat([Tokens, SkipTokens], dim=-1) OutputFeatures = OutputFeatures.transpose(1, 2).reshape( Batch, -1, PressureLevels, Height, Width ) output_surface = OutputFeatures[:, :, 0, :, :] output_upper_air = OutputFeatures[:, :, 1:, :, :] output_surface = self.patchrecovery2d(output_surface) output_upper_air = self.patchrecovery3d(output_upper_air) return output_surface, output_upper_air