Upload folder using huggingface_hub

a3441f1 verified 4 months ago

8.2 kB

# 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

	# 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]
	# [bsnx, hidden_dim] -> [bsnx, 1, hidden_dim]) -> conv1d -> [bsnx, 1, tw] -> [bsnx, tw] (diff)
	diff = self.output_mlp(f[:, None]).squeeze(1)
	# [bsnx] -> [tw, bsnx] -> [bs*nx, tw]
	out = u[:, -1].repeat(self.time_window, 1).transpose(0, 1) + dt * diff
	# print(out.shape)
	return out