| import torch |
| from torch import nn |
|
|
|
|
| class StableNNPhys(nn.Module): |
| def __init__(self, hidden_size=32, input_size=71, output_size=68): |
| super().__init__() |
| self.hidden_size = hidden_size |
| self.input_size = input_size |
| self.output_size = output_size |
| self.hidden = nn.Linear(input_size, hidden_size) |
| self.output = nn.Linear(hidden_size, output_size) |
| self.bypass = nn.Linear(input_size, output_size) |
| nn.init.zeros_(self.output.weight) |
| nn.init.zeros_(self.output.bias) |
| nn.init.zeros_(self.bypass.weight) |
| nn.init.zeros_(self.bypass.bias) |
|
|
| def forward(self, x): |
| return self.output(torch.relu(self.hidden(x))) + self.bypass(x) |
|
|
|
|
| def rollout(model, initial_state, surface, advection, state_mean, state_std, |
| tendency_mean, tendency_std, dt_seconds=10800.0): |
| """Integrate advection trapezoidally, then neural physics with Euler.""" |
| states = [initial_state] |
| physics = [] |
| state = initial_state |
| for step in range(surface.shape[1]): |
| adv_now = advection[:, step] |
| adv_next = advection[:, min(step + 1, advection.shape[1] - 1)] |
| forced = state + 0.5 * dt_seconds * (adv_now + adv_next) |
| surface_scaled = surface[:, step] / surface.new_tensor([100.0, 100.0, 1000.0]) |
| features = torch.cat(((forced - state_mean) / state_std, surface_scaled), dim=-1) |
| tendency = model(features) * tendency_std + tendency_mean |
| state = forced + dt_seconds * tendency |
| physics.append(tendency) |
| states.append(state) |
| return torch.stack(states, dim=1), torch.stack(physics, dim=1) |
|
|
|
|
| def rollout_loss(prediction, target, layer_mass, mode="paper"): |
| error = torch.abs(prediction[:, 1:] - target[:, 1:]) |
| if mode == "paper": |
| weights = torch.cat((layer_mass, layer_mass), dim=-1) |
| weights = weights / weights.mean(dim=-1, keepdim=True) |
| return (error * weights[:, None]).mean() |
| if mode == "official_v0_3": |
| scale = target[:, 1:].std(dim=(0, 1), unbiased=False).clamp_min(1e-6) |
| return (error / scale).mean() |
| raise ValueError(f"Unknown loss mode: {mode}") |
|
|
