model/Graph_UNet.py

import torch
import torch.nn as nn
import torch_geometric.nn as nng
import random
from onescience.modules.mlp.MLP import StandardMLP
from onescience.modules.sample.SpatialGraphDownsample import SpatialGraphDownsample
from onescience.modules.sample.SpatialGraphUpsample import SpatialGraphUpsample

class Model(nn.Module):
    """
    Graph U-Net 模型。
    
    基于图神经网络的 U-Net 结构，包含 Encoder-Decoder 和 Skip Connections。
    使用 SpatialGraphDownsample 进行图池化，使用 SpatialGraphUpsample 进行反池化。
    """
    def __init__(
        self,
        args,
        device,
        pool="random",
        scale=5,
        list_r=[0.05, 0.2, 0.5, 1, 10],
        pool_ratio=[0.5, 0.5, 0.5, 0.5, 0.5],
        max_neighbors=64,
        layer="SAGE",
        head=2,
    ):
        super(Model, self).__init__()
        self.__name__ = "Graph_UNet"
        
        # 参数绑定
        self.L = scale
        self.layer = layer
        self.pool_type = pool
        self.pool_ratio = pool_ratio
        self.list_r = list_r
        self.size_hidden = args.n_hidden
        self.dim_enc = args.n_hidden
        self.bn_bool = True
        self.res = False
        self.head = head
        self.activation = nn.ReLU()
        
        self.encoder = StandardMLP(
            input_dim=args.fun_dim,
            output_dim=args.n_hidden,
            hidden_dims=[args.n_hidden * 2],
            activation=args.act,
            use_bias=True
        )
        
        self.decoder = StandardMLP(
            input_dim=args.n_hidden,
            output_dim=args.out_dim,
            hidden_dims=[args.n_hidden * 2],
            activation=args.act,
            use_bias=True
        )

        # Down Path Layers
        self.down_convs = nn.ModuleList()
        self.down_samples = nn.ModuleList()
        self.down_bns = nn.ModuleList()

        # Level 0 (Initial)
        self._add_conv_layer(self.down_convs, self.dim_enc, self.size_hidden)
        if self.bn_bool:
            self._add_bn_layer(self.down_bns, self.size_hidden)

        # Level 1 to L-1
        current_dim = self.size_hidden
        for n in range(self.L - 1):
            self.down_samples.append(
                SpatialGraphDownsample(
                    in_channels=current_dim,
                    ratio=self.pool_ratio[n],
                    r=self.list_r[n],
                    max_num_neighbors=max_neighbors,
                    pool_method=self.pool_type
                )
            )
            
            # Conv Layer
            in_c = current_dim
            out_c = 2 * current_dim if layer == "SAGE" else current_dim
            self._add_conv_layer(self.down_convs, in_c, out_c)
            current_dim = out_c
            
            if self.bn_bool:
                self._add_bn_layer(self.down_bns, current_dim)

        # Up Path Layers
        self.up_convs = nn.ModuleList()
        
        # --- 3. Upsample Module ---
        self.up_sampler = SpatialGraphUpsample()
        
        self.up_bns = nn.ModuleList()
        
        curr_h_init = args.n_hidden
        
        # Up Layer 0 (Top Layer)
        if self.layer == "SAGE":
            self.up_convs.append(nng.SAGEConv(3 * curr_h_init, self.dim_enc))
            curr_h_init = 2 * curr_h_init
        elif self.layer == "GAT":
            self.up_convs.append(nng.GATConv(2 * self.head * curr_h_init, self.dim_enc, heads=2, concat=False))
        
        if self.bn_bool:
             self.up_bns.append(nng.BatchNorm(self.dim_enc, track_running_stats=False))

        # Up Layer 1 to L-1 (Middle Layers)
        for n in range(1, self.L - 1):
            if self.layer == "SAGE":
                self.up_convs.append(nng.SAGEConv(3 * curr_h_init, curr_h_init))
                bn_dim = curr_h_init
                curr_h_init = 2 * curr_h_init
            elif self.layer == "GAT":
                self.up_convs.append(nng.GATConv(2 * self.head * curr_h_init, curr_h_init, heads=2, concat=True))
                bn_dim = curr_h_init * 2 # GAT concat=True
            
            if self.bn_bool:
                self.up_bns.append(nng.BatchNorm(bn_dim, track_running_stats=False))

    def _add_conv_layer(self, module_list, in_c, out_c):
        if self.layer == "SAGE":
            module_list.append(nng.SAGEConv(in_c, out_c))
        elif self.layer == "GAT":
            module_list.append(nng.GATConv(in_c, out_c, heads=self.head, concat=True, add_self_loops=False))

    def _add_bn_layer(self, module_list, in_c):
        dim = in_c * self.head if self.layer == "GAT" else in_c
        module_list.append(nng.BatchNorm(dim, track_running_stats=False))

    def forward(self, x, fx, T=None, geo=None):
        if geo is None: raise ValueError("Edge index required")
        if fx.dim() == 3: fx = fx.squeeze(0)
        if geo.dim() == 3: edge_index = geo.squeeze(0)
        else: edge_index = geo
        
        # Encoder
        z = self.encoder(fx)
        if self.res: z_res = z.clone()

        # Downsampling Path
        skip_connections = [] 
        pos_history = []     
        edge_index_history = [edge_index.clone()]
        
        # Level 0 Conv
        z = self.down_convs[0](z, edge_index)
        if self.bn_bool: z = self.down_bns[0](z)
        z = self.activation(z)
        
        skip_connections.append(z.clone())
        
        # Assuming x contains coords in first 2 columns as per original code logic
        current_pos = x[:, :2] 
        pos_history.append(current_pos.clone())

        # Levels 1 to L-1
        for n in range(self.L - 1):
            z, current_pos, edge_index, _ = self.down_samples[n](z, current_pos, edge_index)
            
            pos_history.append(current_pos.clone())
            edge_index_history.append(edge_index.clone())

            z = self.down_convs[n+1](z, edge_index)
            if self.bn_bool: z = self.down_bns[n+1](z)
            z = self.activation(z)
            
            skip_connections.append(z.clone())
        
        # Up Path
        for n in range(self.L - 1, 0, -1):
            layer_idx = n - 1
            
            pos_low = pos_history[n]
            pos_high = pos_history[n-1]
            z_skip = skip_connections[n-1]
            
            target_edge_index = edge_index_history[n-1]

            z = self.up_sampler(z, pos_low, pos_high)
            
            z = torch.cat([z, z_skip], dim=1)
            
            z = self.up_convs[layer_idx](z, target_edge_index)
            
            if self.bn_bool: 
                z = self.up_bns[layer_idx](z)
            
            if n != 1:
                z = self.activation(z)

        # Decoder
        if self.res: z = z + z_res
        z = self.decoder(z)
        return z.unsqueeze(0)