model/mpnn.py

# coding=utf-8
from __future__ import annotations

from typing import TYPE_CHECKING

import torch
from torch import nn
from torch_geometric.data import Data
from torch_geometric.nn import MessagePassing, InstanceNorm

if TYPE_CHECKING:
    from onescience.utils.pdenneval.mpnn_utils import PDE

class Swish(nn.Module):
    """Swish activation function
    """
    def __init__(self, beta=1):
        super(Swish, self).__init__()
        self.beta = beta

    def forward(self, x):
        return x * torch.sigmoid(self.beta*x)
    

class GNN_Layer(MessagePassing):
    """Message passing layer
    """
    def __init__(self,
                 in_features: int,
                 out_features: int,
                 hidden_features: int,
                 time_window: int,
                 spatial_dim: int,
                 n_variables: int):
        """Initialize message passing layers
        Args:
            in_features (int): number of node input features
            out_features (int): number of node output features
            hidden_features (int): number of hidden features
            time_window (int): number of input/output timesteps (temporal bundling)
            spatial_dim (int): number of dimension of spatial domain  
            n_variables (int): number of equation specific parameters used in the solver
        """
        super(GNN_Layer, self).__init__(node_dim=-2, aggr='mean') # node_dim: The axis along which to propagate. (default: -2)
        self.in_features = in_features
        self.out_features = out_features
        self.hidden_features = hidden_features
        assert (spatial_dim == 1 or spatial_dim == 2 or spatial_dim == 3)

        self.message_net_1 = nn.Sequential(nn.Linear(2 * in_features + time_window + spatial_dim + n_variables, hidden_features), 
                                           Swish()
                                           )
        self.message_net_2 = nn.Sequential(nn.Linear(hidden_features, hidden_features), 
                                           Swish()
                                           )
        self.update_net_1 = nn.Sequential(nn.Linear(in_features + hidden_features + n_variables, hidden_features), 
                                          Swish()
                                          )
        self.update_net_2 = nn.Sequential(nn.Linear(hidden_features, out_features), 
                                          Swish()
                                          )
        self.norm = InstanceNorm(hidden_features)

    def forward(self, x, u, pos, variables, edge_index, batch):
        """Propagate messages along edges
        """
        x = self.propagate(edge_index, x=x, u=u, pos=pos, variables=variables)
        x = self.norm(x, batch)
        return x

    def message(self, x_i, x_j, u_i, u_j, pos_i, pos_j, variables_i):
        """Message update following formula 8 of the paper
        """
        message = self.message_net_1(torch.cat((x_i, x_j, u_i - u_j, pos_i - pos_j, variables_i), dim=-1))
        message = self.message_net_2(message)
        return message

    def update(self, message, x, variables):
        """Node update following formula 9 of the paper
        """
        update = self.update_net_1(torch.cat((x, message, variables), dim=-1))
        update = self.update_net_2(update)
        if self.in_features == self.out_features:
            return x + update
        else:
            return update



class MPNN(nn.Module):
    def __init__(self,
                 pde: PDE,
                 time_window: int = 25,
                 hidden_features: int = 128,
                 hidden_layers: int = 6,
                 eq_variables: dict = {}):
        """Initialize MPNN
        It contains 6 MP-PDE layers with skip connections
        The input graph to the forward pass has the shape [batch*n_nodes, time_window].
        The output graph has the shape [batch*n_nodes, time_window].
        Args:
            pde (PDE): PDE to solve
            time_window (int): number of input/output timesteps (temporal bundling)
            hidden features (int): number of hidden features
            hidden_layer (int): number of hidden layers
            eq_variables (dict): dictionary of equation specific parameters
        """
        super(MPNN, self).__init__()
        # arguments
        self.pde = pde
        self.out_features = time_window
        self.hidden_features = hidden_features
        self.hidden_layers = hidden_layers
        self.time_window = time_window
        self.eq_variables = eq_variables

        # encoder
        self.embedding_mlp = nn.Sequential(
            nn.Linear(self.time_window+self.pde.spatial_dim+1+len(self.eq_variables), self.hidden_features),
            Swish(),
            nn.Linear(self.hidden_features, self.hidden_features),
            Swish()
            )

        # processor
        self.gnn_layers = torch.nn.ModuleList(modules=(GNN_Layer(
            in_features=self.hidden_features,
            hidden_features=self.hidden_features,
            out_features=self.hidden_features,
            time_window=self.time_window,
            spatial_dim=self.pde.spatial_dim,
            n_variables=len(self.eq_variables)+1 # (time is treated as equation variable)
            ) for _ in range(self.hidden_layers)))
        
        # decoder
        if(self.time_window==10): # NEW ADD
            self.output_mlp = nn.Sequential(
                nn.Conv1d(1, 8, 18, stride=5),
                Swish(),
                nn.Conv1d(8, 1, 14, stride=1)
                )
        if(self.time_window==20):
            self.output_mlp = nn.Sequential(
                nn.Conv1d(1, 8, 15, stride=4),
                Swish(),
                nn.Conv1d(8, 1, 10, stride=1)
                )
        if (self.time_window==25):
            self.output_mlp = nn.Sequential(
                nn.Conv1d(1, 8, 16, stride=3),
                Swish(),
                nn.Conv1d(8, 1, 14, stride=1)
                )
        if(self.time_window==50):
            self.output_mlp = nn.Sequential(
                nn.Conv1d(1, 8, 12, stride=2),
                Swish(),
                nn.Conv1d(8, 1, 10, stride=1)
                )
            
    def forward(self, data: Data, v: int=0) -> torch.Tensor:
        """Forward pass of MP-PDE solver class.
        The input graph has the shape [batch*n_nodes, time_window, v].
        The output tensor has the shape [batch*n_nodes, time_window].
        Args:
            data (Data): Pytorch Geometric data graph
            v (int): physics variable to solve
        Returns:
            torch.Tensor: predictive solutions with shape [batch*n_nodes, time_window]
        """
        u = data.x[..., v] # [bs*nx, tw]
        x_pos = data.x_pos # [bs*nx, spatial_dim]
        t_pos = data.t_pos # [bs*nx]
        edge_index = data.edge_index # [2, num_edges]
        batch = data.batch # [bs*nx]
        variables = data.variables # [bs*nx, num_variables]

        # normalize temporal and spatial coordinate
        t_pos = t_pos / self.pde.tmax
        for d in range(self.pde.spatial_dim):
            x_pos[:, d] = x_pos[:, d] / (self.pde.spatial_domain[d][1] - self.pde.spatial_domain[d][0])

        # encode
        node_input = torch.cat((u, x_pos, t_pos.unsqueeze(1), variables), dim=-1)
        f = self.embedding_mlp(node_input) # [bs*nx, hidden_dim]

        # process
        for i in range(self.hidden_layers):
            # f = self.gnn_layers[i](f, u, x_pos, variables, edge_index, batch)
            f = self.gnn_layers[i](f, u, x_pos, torch.cat((t_pos.unsqueeze(1), variables), dim=-1), edge_index, batch)

        # decode
        dt = (self.pde.tmax - self.pde.tmin) / self.pde.resolution_t
        dt = (torch.ones(1, self.time_window) * dt).to(f.device) # [1, tw]
        dt = torch.cumsum(dt, dim=1) # [1, tw]
        # [bs*nx, hidden_dim] -> [bs*nx, 1, hidden_dim]) -> conv1d -> [bs*nx, 1, tw] -> [bs*nx, tw] (diff)
        diff = self.output_mlp(f[:, None]).squeeze(1)
        # [bs*nx] -> [tw, bs*nx] -> [bs*nx, tw]
        out = u[:, -1].repeat(self.time_window, 1).transpose(0, 1) + dt * diff
        # print(out.shape)
        return out