File size: 2,162 Bytes
7f71cfd | 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 | 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}")
|