model/xihe.py

import math
import os
import torch
import sys
import numpy as np
import torch.nn as nn
from dataclasses import dataclass
from onescience.utils.YParams import YParams
from onescience.models.meta import ModelMetaData
from onescience.modules.embedding.xiheembedding import XiheEmbedding
from onescience.modules.fuser.xihefuse import XiheFuser
from onescience.modules.recovery.xihepatchrecovery import XihePatchRecovery
from onescience.modules.sample.pangudownsample import PanguDownSample
from onescience.modules.sample.xiheupsample import XiheUpSample



@dataclass
class MetaData(ModelMetaData):
    name: str = "Xihe"
    # 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 TensorWithMask:
    def __init__(self, x, mask):
        self.x = x
        self.mask = mask
        self.y=None

   
class Xihe(nn.Module):
    """
    Xihe A PyTorch impl of: `XiHe: A Data-Driven Model for Global Ocean Eddy-Resolving Forecasting`
    https://arxiv.org/abs/2402.02995
    """
    def __init__(
        self,
        config,
        img_size=(2041, 4320),     
        patch_size=(6, 12),        
        window_size=(6, 12),       
        embed_dim=192,
        num_heads=(6, 12, 12, 6),
        in_chans=96,
        depth=1,
        mask_full=None,
        out_chans=94,
        num_groups=32,
        
    ):
        super().__init__()
        self.img_size = config.img_size
        self.patch_size =config.patch_size
        # 正确初始化 mask_full
        self.mask=config.mask
        mask_full = np.load(self.mask)
        self.mask_full = mask_full if mask_full is not None else None
        self.mask_h_w=None
        self.out_chans=config.out_chans
        self.in_chans=config.in_chans
        self.num_groups=config.num_groups
        self.embed_dim=config.embed_dim
        
        self.skip_proj = nn.Linear(2*self.embed_dim, self.embed_dim)


        self.patchembed2d = XiheEmbedding()
        self.patchrecovery2d = XihePatchRecovery()

        # patch 后的 3D 分辨率: (Pl=1, Lat_out, Lon_out)
        H_out = math.ceil(img_size[0] / patch_size[0])
        W_out = math.ceil(img_size[1] / patch_size[1])
        input_resolution = (1, H_out, W_out)
        
        self.mask_h_w=input_resolution
        # 3D 窗口：把 2D 窗口扩成 (1, win_lat, win_lon)
        window_size_3d = (1, window_size[0], window_size[1]) #  window_size (tuple[int]): Window size [pressure levels, latitude, longitude].

        # 防止过拟合，随机丢弃一部分
        if depth > 1:
            drop_path = np.linspace(0, 0.2, depth).tolist()
        else:
            drop_path = 0.0

        self.block1=XiheFuser(dim=self.embed_dim,input_resolution=input_resolution,num_local=1)

        self.downsample = PanguDownSample(
                                    in_dim=self.embed_dim,
                                    input_resolution=(H_out, W_out),
                                    output_resolution=(H_out // 2, W_out // 2))
        
        input_resolution = (1, H_out // 2, W_out // 2)
        self.mask_h_w=input_resolution

        self.block2=XiheFuser(dim=2*self.embed_dim,input_resolution=input_resolution,num_local=2)

        self.block3=XiheFuser(dim=2*self.embed_dim,input_resolution=input_resolution,num_local=2)
        self.block4=XiheFuser(dim=2*self.embed_dim,input_resolution=input_resolution,num_local=2)

        self.upsample = XiheUpSample(in_dim=2*self.embed_dim,out_dim=embed_dim,input_resolution=(H_out // 2, W_out // 2),  output_resolution=(H_out, W_out), )
        input_resolution = (1, H_out, W_out)
        self.block5=XiheFuser(dim=self.embed_dim,input_resolution=input_resolution,num_local=1)
    def change_mask(self,mask_full, x, h_out, w_out):
        #根据当前层特征分辨率，自动生成掩码（海洋=1，陆地=0）
            if not torch.is_tensor(mask_full):
                mask_full = torch.tensor(mask_full, dtype=torch.float32)
            else:
                mask_full = mask_full

            H, W = mask_full.shape
            patch_h = math.ceil(H / h_out)
            patch_w = math.ceil(W / w_out)

            mask_coarse = torch.zeros((h_out, w_out), dtype=torch.float32)
            for i in range(h_out):
                for j in range(w_out):
                    h0, h1 = i * patch_h, min((i + 1) * patch_h, H)
                    w0, w1 = j * patch_w, min((j + 1) * patch_w, W)
                    patch = mask_full[h0:h1, w0:w1]
                    mask_coarse[i, j] = 1.0 if torch.any(patch > 0.5) else 0.0
            
            mask_coarse = mask_coarse.to(x.device, dtype=x.dtype) 
            B = x.shape[0]                
            mask_coarse = mask_coarse.unsqueeze(0).unsqueeze(0).repeat(B, 1, 1, 1) #broadcast
            return mask_coarse  
        
    def forward(self, x: torch.Tensor):     
        x = self.patchembed2d(x)                  # (B, C=embed_dim, H', W')
        x = x.flatten(2).transpose(1, 2)          # (B, N=H'*W', C) 
        B, N, C = x.shape     
        mask_full=self.mask_full        
        
       
        if mask_full is not None:              # mask1
            H_out = math.ceil(self.img_size[0] / self.patch_size[0])
            W_out = math.ceil(self.img_size[1] / self.patch_size[1])
            mask1 = self.change_mask(mask_full, x, h_out=H_out, w_out=W_out)
        else:
            mask1 = None
        
        obj1 = TensorWithMask(x, mask1)
        x=self.block1(obj1)          # (B, N, C) 经过 3D 全局注意力
        x1=x
        x=self.downsample(x)                 # (B, N, C) 经过 2D 下采样
        
        if mask_full is not None:            #  mask2
            _,H_out,W_out = self.mask_h_w
            mask2 = self.change_mask(mask_full, x, h_out=H_out, w_out=W_out)
        else:
            mask2 = None
        obj2 = TensorWithMask(x, mask2)
        x=self.block2(obj2)   
        obj2 = TensorWithMask(x, mask2)
        x=self.block3(obj2)  
        obj2 = TensorWithMask(x, mask2)
        x=self.block4(obj2) 
        x=self.upsample(x) 
        obj1 = TensorWithMask(x, mask1)
        x=self.block5(obj1)
        x_out = torch.cat([x, x1], dim=-1)         # (B, N, 2C)
        x_out = self.skip_proj(x_out)
        # B, N, C = x.shape
        # H_, W_ = 341, 360   # 对应 patch grid 尺寸
        H_ = math.ceil(self.img_size[0] / self.patch_size[0])
        W_ = math.ceil(self.img_size[1] / self.patch_size[1])
        x_out = x_out.transpose(1, 2).reshape(B, C, H_, W_)
        x=self.patchrecovery2d(x_out)
        return x