import torch import torch._dynamo import torch.nn as nn import torch.nn.functional as F from torchdiffeq import odeint from typing import Tuple, Optional, List import matplotlib.pyplot as plt # ---------------------------------------------------------------------- # 1. Spatial mixing: multi‑directional information propagation # (can be replaced later by TensorNetworkPropagator or GraphConv) # ---------------------------------------------------------------------- class SpatialPropagator(nn.Module): """ A ConvGRU that iteratively processes a hidden field. Uses 3×3 convolutions → information flows in all 8 spatial directions. """ def __init__(self, hidden_dim: int, kernel_size: int = 3): super().__init__() self.hidden_dim = hidden_dim self.conv_z = nn.Conv2d(hidden_dim*2, hidden_dim, kernel_size, padding=kernel_size//2) self.conv_r = nn.Conv2d(hidden_dim*2, hidden_dim, kernel_size, padding=kernel_size//2) self.conv_h = nn.Conv2d(hidden_dim*2, hidden_dim, kernel_size, padding=kernel_size//2) def forward(self, x: torch.Tensor, h: torch.Tensor) -> torch.Tensor: """ Args: x: input signal at current step (B, hidden_dim, H, W) h: previous hidden state (B, hidden_dim, H, W) Returns: new hidden state (B, hidden_dim, H, W) """ cat = torch.cat([x, h], dim=1) z = torch.sigmoid(self.conv_z(cat)) r = torch.sigmoid(self.conv_r(cat)) cat_r = torch.cat([x, r * h], dim=1) h_tilde = torch.tanh(self.conv_h(cat_r)) h_new = (1 - z) * h + z * h_tilde return h_new # ---------------------------------------------------------------------- # 2. Neural ODE derivative function: time + space dynamics # ---------------------------------------------------------------------- class ODEFunc(nn.Module): """ Defines dh/dt = f(h, t) with spatial interactions. The function `f` uses a SpatialPropagator and a small MLP. """ def __init__(self, hidden_dim: int, spatial_mixer: nn.Module): super().__init__() self.hidden_dim = hidden_dim self.spatial = spatial_mixer # any module that takes (x, h) -> new h self.time_mlp = nn.Sequential( nn.Linear(1, hidden_dim), nn.Tanh(), nn.Linear(hidden_dim, hidden_dim) ) def forward(self, t: torch.Tensor, h: torch.Tensor) -> torch.Tensor: """ h shape: (B, hidden_dim, H, W) t: scalar time (broadcastable) """ # Encode time as a per‑pixel modulation B, C, H, W = h.shape t_vec = t.view(1, 1).expand(B, 1) # (B,1) t_feat = self.time_mlp(t_vec) # (B, hidden_dim) t_feat = t_feat.view(B, C, 1, 1).expand(-1, -1, H, W) # Combine time feature with current state as "input signal" x = h + t_feat dh = self.spatial(x, h) - h # residual update return dh # ---------------------------------------------------------------------- # 3. Collapse module: from field to forecast path # ---------------------------------------------------------------------- class Collapser(nn.Module): """ Soft‑attention over the spatial field, producing a single forecast vector. Also learns a "confidence" that can later be fed into information‑physics criteria. """ def __init__(self, hidden_dim: int, forecast_dim: int): super().__init__() self.forecast_dim = forecast_dim # Key/query for spatial attention self.query = nn.Linear(hidden_dim, hidden_dim, bias=False) self.key = nn.Conv2d(hidden_dim, hidden_dim, 1) # Decoder: from attended hidden to output self.decoder = nn.Sequential( nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, forecast_dim) ) # Optional: learnable initial query (or can be derived from history) self.init_query = nn.Parameter(torch.zeros(1, 1, hidden_dim)) def forward(self, h: torch.Tensor, query: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: """ Args: h: field (B, hidden_dim, H, W) query: (B, 1, hidden_dim) or None (uses learned init) Returns: forecast: (B, forecast_dim) attn_weights: (B, H, W) – for analysis (e.g. Fisher info) """ B, C, H, W = h.shape if query is None: query = self.init_query.expand(B, -1, -1) # (B,1,C) # Flatten spatial dims h_flat = h.view(B, C, H*W).transpose(1,2) # (B, N, C) keys = self.key(h).view(B, C, H*W).transpose(1,2) # (B, N, C) # Attention scores scores = torch.bmm(query, keys.transpose(1,2)) # (B, 1, N) attn = F.softmax(scores / (C ** 0.5), dim=-1) # (B, 1, N) # Weighted sum of hidden vectors attended = torch.bmm(attn, h_flat).squeeze(1) # (B, C) forecast = self.decoder(attended) # (B, forecast_dim) return forecast, attn.view(B, H, W) @torch._dynamo.disable def call_odeint(func, h0, t_span): return odeint(func, h0, t_span, method='dopri5') # ---------------------------------------------------------------------- # 4. Full model: NeuralODEFieldCollapser # ---------------------------------------------------------------------- class NeuralODEFieldCollapser(nn.Module): """ Spatio‑temporal forecaster using: - Encoder: maps input frames to initial hidden field - Neural ODE: propagates field forward in continuous time - Collapser: condenses field into a single forecast at each output step """ def __init__(self, input_channels: int, hidden_dim: int, forecast_dim: int, ode_hidden: int = 64): super().__init__() self.hidden_dim = hidden_dim self.forecast_dim = forecast_dim # Map raw input → hidden field self.encoder = nn.Sequential( nn.Conv2d(input_channels, hidden_dim, 3, padding=1), nn.GroupNorm(8, hidden_dim), nn.ReLU(), nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), ) # ODE function + integrator self.spatial_prop = SpatialPropagator(hidden_dim) self.ode_func = ODEFunc(hidden_dim, self.spatial_prop) self.ode_hidden = ode_hidden # for time‑integration # Collapse module (can be swapped with Fisher‑guided version later) self.collapser = Collapser(hidden_dim, forecast_dim) # Optional: a small MLP to produce the initial attention query from history self.query_net = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(hidden_dim, hidden_dim) ) def encode_initial_state(self, x: torch.Tensor) -> torch.Tensor: """ x: (B, C_in, H, W) – one frame or a stack of past frames. Returns initial hidden field (B, hidden_dim, H, W). """ return self.encoder(x) def forward(self, past_frames: torch.Tensor, t_span: torch.Tensor, forecast_steps: int) -> Tuple[torch.Tensor, List[torch.Tensor]]: """ Args: past_frames: (B, T_in, C_in, H, W) – observed window t_span: 1D tensor of integration times, e.g. [0, 0.5, 1.0, ...] forecast_steps: number of future times to predict Returns: forecasts: (B, forecast_steps, forecast_dim) attn_maps: list of attention maps (B, H, W) for each forecast step """ B, T_in, C, H, W = past_frames.shape # Encode the latest frame (or fuse multiple) into hidden field # For simplicity: use last frame only; later can be improved last_frame = past_frames[:, -1] # (B, C, H, W) h0 = self.encode_initial_state(last_frame) # (B, hidden_dim, H, W) # Initial query from whole initial field init_query = self.query_net(h0).unsqueeze(1) # (B, 1, hidden_dim) # Step through time using ODE integrator (continuous) # t_span should be a 1D tensor of sorted times including the final forecast horizon h_traj = call_odeint(self.ode_func, h0, t_span) # (T, B, Hdim, H, W) # We only care about times that correspond to output steps. # For simplicity assume t_span has length forecast_steps+1 (including t=0) # h_traj[0] = h0, h_traj[1:] are future states forecasts = [] attn_maps = [] query = init_query for step in range(1, forecast_steps+1): h_step = h_traj[step] # (B, Hdim, H, W) fc, attn = self.collapser(h_step, query) forecasts.append(fc) attn_maps.append(attn) # Update query: simple momentum from previous attended state # (In full version, use KLDivergenceMinimizer or Fisher update) attended_h = torch.bmm(attn.view(B, 1, H*W), h_step.view(B, self.hidden_dim, H*W).transpose(1,2)).squeeze(1) query = 0.9 * query + 0.1 * attended_h.unsqueeze(1) forecasts = torch.stack(forecasts, dim=1) # (B, steps, forecast_dim) return forecasts, attn_maps # ---------------------------------------------------------------------- # 5. Incremental training loop (skeleton) # ---------------------------------------------------------------------- def train_step(model, batch, optimizer, device): past, future, t_span = batch past = past.to(device) future = future.to(device) t_span = t_span.to(device) forecast_steps = future.shape[1] pred, _ = model(past, t_span, forecast_steps) loss = F.mse_loss(pred, future) optimizer.zero_grad() loss.backward() optimizer.step() return loss.item() # ---------------------------------------------------------------------- # Example usage (not run here) # ---------------------------------------------------------------------- if __name__ == "__main__": # Hyperparameters input_channels = 1 # e.g. temperature hidden_dim = 32 forecast_dim = 1 # scaler per grid point? Could be H*W if full frame prediction ode_hidden = 64 # For full frame prediction, forecast_dim = H*W; decoder becomes Conv. # Here we show a simple 1D forecast per grid cell: forecast_dim = 1, # but then we'd need to apply this to every pixel. Real implementation would # output a whole frame with ConvDecoder instead of Collapser. # For demonstration, we use a scalar summary (e.g. average temperature). model = NeuralODEFieldCollapser(input_channels, hidden_dim, forecast_dim) model = torch.compile(model) # fast compilation, PyTorch 2.0 compiler # Dummy data with a trend/pattern (e.g. sine wave) to make training meaningful B, T, C, H, W = 4, 5, 1, 16, 16 t_span = torch.linspace(0, 1, steps=4) # 0, 0.33, 0.66, 1.0 -> 3 future states + start # Simple synthetic heat/diffusion-like trend: past is sine, future is sine continuing past = torch.sin(torch.linspace(0, 3.14, T).view(1, T, 1, 1, 1).expand(B, T, C, H, W)) + 0.1 * torch.randn(B, T, C, H, W) # future is [B, 3, 1] - let's make it a decaying trend from the past future = torch.stack([ torch.sin(torch.tensor(3.14 + 0.5 * (step + 1))) * torch.ones(B, 1) for step in range(3) ], dim=1) + 0.05 * torch.randn(B, 3, 1) optimizer = torch.optim.Adam(model.parameters(), lr=1e-2) # Run a short training loop print("Starting training...") losses = [] for step in range(100): loss = train_step(model, (past, future, t_span), optimizer, torch.device('cpu')) losses.append(loss) if (step + 1) % 20 == 0: print(f"Step {step+1:3d}/100 | Loss: {loss:.4f}") # Get final predictions model.eval() with torch.no_grad(): pred, _ = model(past, t_span, 3) # Plot loss and forecasting result fig, axs = plt.subplots(1, 2, figsize=(12, 5)) # Loss Curve axs[0].plot(losses, label='Train MSE', color='purple', lw=2) axs[0].set_title('Training Loss Curve') axs[0].set_xlabel('Steps') axs[0].set_ylabel('MSE Loss') axs[0].grid(True) axs[0].legend() # Predictions vs. Ground Truth (for the first batch item) time_steps = [1, 2, 3] axs[1].plot(time_steps, future[0, :, 0].cpu().numpy(), 'o--', label='Ground Truth', color='black', alpha=0.8) axs[1].plot(time_steps, pred[0, :, 0].cpu().numpy(), 's-', label='Forecast Prediction', color='crimson') axs[1].set_xticks(time_steps) axs[1].set_title('Single-Instance Forecast (t=1..3)') axs[1].set_xlabel('Forecast Steps') axs[1].set_ylabel('Value') axs[1].grid(True) axs[1].legend() plt.tight_layout() plt.savefig('ex03_plot.png') print("Saved training and forecasting plot to 'ex03_plot.png'")