You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
93 lines
2.7 KiB
Python
93 lines
2.7 KiB
Python
2 months ago
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
|
||
|
|
|
||
|
|
class ResidualBlock(nn.Module):
|
||
|
|
def __init__(self, dim: int, dropout: float = 0.0):
|
||
|
|
super().__init__()
|
||
|
|
self.net = nn.Sequential(
|
||
|
|
nn.LayerNorm(dim),
|
||
|
|
nn.Linear(dim, dim),
|
||
|
|
nn.GELU(),
|
||
|
|
nn.Dropout(dropout) if dropout > 0 else nn.Identity(),
|
||
|
|
nn.Linear(dim, dim),
|
||
|
|
)
|
||
|
|
self.act = nn.GELU()
|
||
|
|
|
||
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
|
|
return self.act(x + self.net(x))
|
||
|
|
|
||
|
|
|
||
|
|
class TimeConditionedSurrogate(nn.Module):
|
||
|
|
"""Point-wise forward surrogate.
|
||
|
|
|
||
|
|
f(params, full_schedule, loglog_time_point) -> [log_pressure, log_derivative]
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
param_dim: int,
|
||
|
|
schedule_dim: int,
|
||
|
|
time_dim: int,
|
||
|
|
hidden_dim: int = 256,
|
||
|
|
n_blocks: int = 4,
|
||
|
|
dropout: float = 0.05,
|
||
|
|
use_schedule: bool = True,
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
self.use_schedule = bool(use_schedule)
|
||
|
|
|
||
|
|
self.param_encoder = nn.Sequential(
|
||
|
|
nn.Linear(param_dim, hidden_dim),
|
||
|
|
nn.LayerNorm(hidden_dim),
|
||
|
|
nn.GELU(),
|
||
|
|
)
|
||
|
|
self.time_encoder = nn.Sequential(
|
||
|
|
nn.Linear(time_dim, hidden_dim // 2),
|
||
|
|
nn.LayerNorm(hidden_dim // 2),
|
||
|
|
nn.GELU(),
|
||
|
|
)
|
||
|
|
if self.use_schedule:
|
||
|
|
self.schedule_encoder = nn.Sequential(
|
||
|
|
nn.Linear(schedule_dim, hidden_dim),
|
||
|
|
nn.LayerNorm(hidden_dim),
|
||
|
|
nn.GELU(),
|
||
|
|
)
|
||
|
|
fusion_dim = hidden_dim * 2 + hidden_dim // 2
|
||
|
|
else:
|
||
|
|
self.schedule_encoder = None
|
||
|
|
fusion_dim = hidden_dim + hidden_dim // 2
|
||
|
|
|
||
|
|
self.input_proj = nn.Sequential(
|
||
|
|
nn.Linear(fusion_dim, hidden_dim),
|
||
|
|
nn.GELU(),
|
||
|
|
)
|
||
|
|
self.blocks = nn.Sequential(*[ResidualBlock(hidden_dim, dropout=dropout) for _ in range(int(n_blocks))])
|
||
|
|
self.head = nn.Sequential(
|
||
|
|
nn.LayerNorm(hidden_dim),
|
||
|
|
nn.Linear(hidden_dim, hidden_dim // 2),
|
||
|
|
nn.GELU(),
|
||
|
|
nn.Linear(hidden_dim // 2, 2),
|
||
|
|
)
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self,
|
||
|
|
params_x: torch.Tensor,
|
||
|
|
time_x: torch.Tensor,
|
||
|
|
schedule_x: torch.Tensor | None = None,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
p = self.param_encoder(params_x)
|
||
|
|
t = self.time_encoder(time_x)
|
||
|
|
if self.use_schedule:
|
||
|
|
if schedule_x is None:
|
||
|
|
raise ValueError("use_schedule=True but schedule_x is None")
|
||
|
|
s = self.schedule_encoder(schedule_x)
|
||
|
|
x = torch.cat([p, s, t], dim=-1)
|
||
|
|
else:
|
||
|
|
x = torch.cat([p, t], dim=-1)
|
||
|
|
x = self.input_proj(x)
|
||
|
|
x = self.blocks(x)
|
||
|
|
return self.head(x)
|