File size: 8,195 Bytes
e2703dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | # 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
|