File size: 10,084 Bytes
c92f17c | 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 203 204 205 206 207 208 209 210 211 212 | """Paper-priority MP-PDE model for experiment E3."""
from __future__ import annotations
from typing import Iterable, Sequence, Tuple
import torch
from torch import Tensor, nn
class Swish(nn.Module):
def forward(self, values: Tensor) -> Tensor:
return values * torch.sigmoid(values)
class TwoLayerMLP(nn.Module):
def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
super().__init__()
self.network = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
Swish(),
nn.Linear(hidden_dim, output_dim),
Swish(),
)
def forward(self, values: Tensor) -> Tensor:
return self.network(values)
def periodic_neighbor_indices(num_nodes: int, offsets: Iterable[int], device: torch.device | None = None) -> Tensor:
"""Return source indices [num_nodes, num_neighbors] for each target node."""
offsets_tensor = torch.as_tensor(tuple(offsets), dtype=torch.long, device=device)
if num_nodes < 7:
raise ValueError(f"The six-neighbor periodic graph requires num_nodes>=7, found {num_nodes}")
if offsets_tensor.numel() != 6 or offsets_tensor.unique().numel() != 6 or torch.any(offsets_tensor == 0):
raise ValueError("MP-PDE E3 requires exactly six distinct non-zero neighbor offsets")
target = torch.arange(num_nodes, dtype=torch.long, device=device)[:, None]
return torch.remainder(target + offsets_tensor[None, :], num_nodes)
class MessagePassingLayer(nn.Module):
"""Equation (8)--(9) processor layer with sum aggregation."""
def __init__(self, hidden_dim: int, history_dim: int, parameter_dim: int, affine_norm: bool):
super().__init__()
edge_dim = 2 * hidden_dim + history_dim + 1 + parameter_dim
node_dim = 2 * hidden_dim + parameter_dim
self.edge_mlp = TwoLayerMLP(edge_dim, hidden_dim, hidden_dim)
self.node_mlp = TwoLayerMLP(node_dim, hidden_dim, hidden_dim)
self.norm = nn.InstanceNorm1d(hidden_dim, affine=affine_norm, track_running_stats=False)
def forward(
self,
hidden: Tensor,
history: Tensor,
x: Tensor,
parameters: Tensor,
neighbors: Tensor,
domain_length: float,
) -> Tensor:
batch, num_nodes, hidden_dim = hidden.shape
num_neighbors = neighbors.shape[1]
source_hidden = hidden[:, neighbors, :]
target_hidden = hidden[:, :, None, :].expand(-1, -1, num_neighbors, -1)
source_history = history[:, neighbors, :]
history_difference = history[:, :, None, :] - source_history
source_x = x[:, neighbors]
displacement = x[:, :, None] - source_x
displacement = torch.remainder(displacement + 0.5 * domain_length, domain_length) - 0.5 * domain_length
theta_edges = parameters[:, None, None, :].expand(-1, num_nodes, num_neighbors, -1)
edge_input = torch.cat(
(target_hidden, source_hidden, history_difference, displacement[..., None], theta_edges), dim=-1
)
messages = self.edge_mlp(edge_input)
aggregated = messages.sum(dim=2)
theta_nodes = parameters[:, None, :].expand(-1, num_nodes, -1)
node_update = self.node_mlp(torch.cat((hidden, aggregated, theta_nodes), dim=-1))
return self.norm((hidden + node_update).transpose(1, 2)).transpose(1, 2)
class TemporalDecoder(nn.Module):
def __init__(
self,
hidden_dim: int,
time_window: int,
middle_channels: int = 8,
kernels: Sequence[int] = (16, 26),
strides: Sequence[int] = (3, 1),
):
super().__init__()
if len(kernels) != 2 or len(strides) != 2:
raise ValueError("The paper-priority decoder requires exactly two convolutions")
length_after_first = (hidden_dim - int(kernels[0])) // int(strides[0]) + 1
output_length = (length_after_first - int(kernels[1])) // int(strides[1]) + 1
if output_length != time_window:
raise ValueError(
f"Decoder does not close hidden={hidden_dim} to K={time_window}: output length={output_length}"
)
self.network = nn.Sequential(
nn.Conv1d(1, middle_channels, kernel_size=int(kernels[0]), stride=int(strides[0])),
Swish(),
nn.Conv1d(middle_channels, 1, kernel_size=int(kernels[1]), stride=int(strides[1])),
)
def forward(self, hidden: Tensor) -> Tensor:
batch, num_nodes, hidden_dim = hidden.shape
decoded = self.network(hidden.reshape(batch * num_nodes, 1, hidden_dim))
return decoded.reshape(batch, num_nodes, decoded.shape[-1])
class MPPDESolver(nn.Module):
"""Message-passing neural solver mapping K E3 states to the next K states."""
def __init__(
self,
time_window: int = 25,
hidden_dim: int = 164,
message_passing_layers: int = 6,
neighbor_offsets: Sequence[int] = (-3, -2, -1, 1, 2, 3),
domain_length: float = 16.0,
final_time: float = 4.0,
parameter_maxima: Sequence[float] = (3.0, 0.4, 1.0),
scale_coordinates: bool = True,
scale_parameters: bool = True,
instance_norm_affine: bool = False,
decoder_middle_channels: int = 8,
decoder_kernels: Sequence[int] = (16, 26),
decoder_strides: Sequence[int] = (3, 1),
):
super().__init__()
if time_window <= 0 or hidden_dim <= 0 or message_passing_layers <= 0:
raise ValueError("time_window, hidden_dim, and message_passing_layers must be positive")
self.time_window = int(time_window)
self.hidden_dim = int(hidden_dim)
self.neighbor_offsets: Tuple[int, ...] = tuple(int(value) for value in neighbor_offsets)
self.domain_length = float(domain_length)
self.final_time = float(final_time)
self.scale_coordinates = bool(scale_coordinates)
self.scale_parameters = bool(scale_parameters)
maxima = torch.tensor(tuple(float(value) for value in parameter_maxima), dtype=torch.float32)
if maxima.shape != (3,) or torch.any(maxima <= 0.0):
raise ValueError("parameter_maxima must contain three positive values")
self.register_buffer("parameter_maxima", maxima, persistent=False)
self.encoder = TwoLayerMLP(self.time_window + 1 + 1 + 3, self.hidden_dim, self.hidden_dim)
self.processor = nn.ModuleList(
MessagePassingLayer(self.hidden_dim, self.time_window, 3, instance_norm_affine)
for _ in range(int(message_passing_layers))
)
self.decoder = TemporalDecoder(
self.hidden_dim, self.time_window, decoder_middle_channels, decoder_kernels, decoder_strides
)
def _canonicalize_inputs(
self, history: Tensor, x: Tensor, current_time: Tensor | float, parameters: Tensor
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
if history.ndim != 3:
raise ValueError(f"history must have shape [B,N,K], found {tuple(history.shape)}")
batch, num_nodes, window = history.shape
if window != self.time_window:
raise ValueError(f"Expected history K={self.time_window}, found {window}")
if num_nodes < 7:
raise ValueError(f"MP-PDE requires N>=7, found {num_nodes}")
if parameters.shape != (batch, 3):
raise ValueError(f"parameters must have shape [{batch},3], found {tuple(parameters.shape)}")
if x.ndim == 1:
if x.shape[0] != num_nodes:
raise ValueError(f"x length {x.shape[0]} does not match N={num_nodes}")
x = x[None, :].expand(batch, -1)
elif x.shape != (batch, num_nodes):
raise ValueError(f"x must have shape [N] or [B,N], found {tuple(x.shape)}")
time = torch.as_tensor(current_time, dtype=history.dtype, device=history.device)
if time.ndim == 0:
time = time.expand(batch)
elif time.shape == (batch, 1):
time = time[:, 0]
elif time.shape != (batch,):
raise ValueError(f"current_time must be scalar or shape [B], found {tuple(time.shape)}")
return history, x.to(history), time, parameters.to(history)
def forward(
self,
history: Tensor,
x: Tensor,
current_time: Tensor | float,
parameters: Tensor,
dt: Tensor | float,
*,
return_derivative: bool = False,
) -> Tensor | Tuple[Tensor, Tensor]:
history, x, current_time, parameters = self._canonicalize_inputs(history, x, current_time, parameters)
batch, num_nodes, _ = history.shape
scaled_x = x / self.domain_length if self.scale_coordinates else x
scaled_time = current_time / self.final_time if self.scale_coordinates else current_time
scaled_parameters = parameters / self.parameter_maxima.to(parameters) if self.scale_parameters else parameters
time_feature = scaled_time[:, None, None].expand(-1, num_nodes, 1)
parameter_features = scaled_parameters[:, None, :].expand(-1, num_nodes, -1)
encoded = torch.cat((history, scaled_x[..., None], time_feature, parameter_features), dim=-1)
hidden = self.encoder(encoded)
neighbors = periodic_neighbor_indices(num_nodes, self.neighbor_offsets, history.device)
for layer in self.processor:
hidden = layer(hidden, history, x, scaled_parameters, neighbors, self.domain_length)
derivative = self.decoder(hidden)
step = torch.as_tensor(dt, dtype=history.dtype, device=history.device)
if step.ndim == 0:
step = step.expand(batch)
elif step.shape != (batch,):
raise ValueError(f"dt must be scalar or shape [B], found {tuple(step.shape)}")
offsets = torch.arange(1, self.time_window + 1, dtype=history.dtype, device=history.device)
delta_times = step[:, None, None] * offsets[None, None, :]
prediction = history[:, :, -1:] + delta_times * derivative
return (prediction, derivative) if return_derivative else prediction
|