import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchdiffeq import odeint import numpy as np import math # ---------------------------------------------------------------------- # 1. Spatial Message Passing (BeliefPropagationSolver, TensorNetworkPropagator) # Aggregates information from all four cardinal directions on a 2D grid. # ---------------------------------------------------------------------- class SpatialPropagator(nn.Module): """ Lightweight graph convolution on a regular 2D grid. Implements multi‑directional message passing: +x, -x, +y, -y. """ def __init__(self, latent_dim, hidden_dim=64): super().__init__() self.message_mlp = nn.Sequential( nn.Linear(latent_dim * 2, hidden_dim), # pairwise interaction nn.ReLU(), nn.Linear(hidden_dim, latent_dim) ) self.update_mlp = nn.Sequential( nn.Linear(latent_dim + latent_dim, hidden_dim), # self + aggregated nn.ReLU(), nn.Linear(hidden_dim, latent_dim) ) def forward(self, h, grid_shape): """ h: [B, C, H, W] feature map of latent state Returns updated h """ B, C, H, W = h.shape # Pad to handle boundaries (circular or reflect – here circular for simplicity) h_pad = F.pad(h, (1,1,1,1), mode='circular') # Extract neighbours: left, right, up, down h_left = h_pad[:, :, 1:H+1, 0:W] # -x h_right = h_pad[:, :, 1:H+1, 2:W+2] # +x h_up = h_pad[:, :, 0:H, 1:W+1] # +y h_down = h_pad[:, :, 2:H+2, 1:W+1] # -y # Compute messages from each direction def message(neighbour): cat = torch.cat([h, neighbour], dim=1) # [B, 2C, H, W] # Permute to apply MLP per location cat = cat.permute(0,2,3,1).reshape(-1, 2*C) msg = self.message_mlp(cat).reshape(B, H, W, C).permute(0,3,1,2) return msg m_left = message(h_left) m_right = message(h_right) m_up = message(h_up) m_down = message(h_down) # Aggregate: sum messages agg = m_left + m_right + m_up + m_down # Update using self-connection and aggregated message cat = torch.cat([h, agg], dim=1).permute(0,2,3,1).reshape(-1, 2*C) h_new = self.update_mlp(cat).reshape(B, H, W, C).permute(0,3,1,2) return h_new # ---------------------------------------------------------------------- # 2. Neural ODE function (MLP regressor) # This is the "thinking function" φ. # ---------------------------------------------------------------------- class ODEFunc(nn.Module): """ MLP that takes (t, state) and returns dstate/dt. The state is a latent vector per spatial location (flattened). Spatial coordinates are embedded to break translation invariance if desired. """ def __init__(self, latent_dim, hidden_dim=128): super().__init__() self.net = nn.Sequential( nn.Linear(latent_dim + 1, hidden_dim), # +1 for time t nn.Tanh(), nn.Linear(hidden_dim, latent_dim) ) # Learnable embedding for spatial positions (optional) self.pos_emb = None # can be added later def forward(self, t, state): # state shape: [B, C, H, W] B, C, H, W = state.shape state_flat = state.permute(0,2,3,1).reshape(-1, C) # [B*H*W, C] t_tensor = t * torch.ones(B*H*W, 1, device=state.device) inp = torch.cat([t_tensor, state_flat], dim=1) dstate = self.net(inp).reshape(B, H, W, C).permute(0,3,1,2) return dstate # ---------------------------------------------------------------------- # 3. Prior model (MaxEntReconstructor) – a simple historical average dynamics # Provides a baseline path for the KL regularizer. # ---------------------------------------------------------------------- class PriorDynamics: """ Learns a smooth temporal prior from historical data. For simplicity, we approximate the prior as independent Gaussian at each time. In a full version this would be trained via MaxEnt on past trajectories. """ def __init__(self, latent_dim, grid_shape): self.latent_dim = latent_dim self.grid_shape = grid_shape self.prior_mean = None # will be set from data self.prior_logvar = None def set_prior(self, mean, logvar): self.prior_mean = mean self.prior_logvar = logvar def kl_divergence(self, h): """ Compute KL( N(h, I) || N(prior_mean, diag(exp(prior_logvar))) ) Assumes h is a sample from model's path (for simplicity we treat h as mean). We will use a stochastic interpretation: model outputs mean, var=1. """ if self.prior_mean is None: return 0.0 var_prior = torch.exp(self.prior_logvar) kl = 0.5 * (var_prior + (h - self.prior_mean)**2 - 1 - self.prior_logvar).sum(dim=[1,2,3]).mean() return kl # ---------------------------------------------------------------------- # 4. Main Collapser Model # ---------------------------------------------------------------------- class NeuralODEFieldCollapser(nn.Module): """ Encodes a sequence of past frames into an initial hidden field, propagates spatially, evolves with ODE, and collapses onto a forecast path. """ def __init__(self, input_channels, latent_dim, hidden_dim=128, num_propagate=2): super().__init__() self.input_channels = input_channels self.latent_dim = latent_dim self.num_propagate = num_propagate # Encoder: maps a sequence of frames (T_in, C_in, H, W) to initial state (C, H, W) self.encoder = nn.Sequential( nn.Conv3d(input_channels, 32, kernel_size=(3,3,3), padding=(1,1,1)), nn.ReLU(), nn.Conv3d(32, latent_dim, kernel_size=(3,3,3), padding=(0,1,1)), # reduce time to 1 nn.ReLU() ) # Spatial propagation self.propagator = SpatialPropagator(latent_dim, hidden_dim) # ODE function and solver settings self.ode_func = ODEFunc(latent_dim, hidden_dim) self.ode_solver = 'dopri5' self.ode_options = {'rtol': 1e-3, 'atol': 1e-4} # Decoder: from latent state to output value (e.g., 1 channel) self.decoder = nn.Sequential( nn.Conv2d(latent_dim, 64, kernel_size=3, padding=1), nn.ReLU(), nn.Conv2d(64, 1, kernel_size=3, padding=1) ) # Prior model (simplified – we'll just use a fixed prior for demo) self.prior = PriorDynamics(latent_dim, (1,1)) # grid_shape set later self.beta = 0.01 # KL regularization weight def set_prior_from_data(self, dataset): # In practice, compute prior mean/logvar over historical latent states # Here dummy: zeros prior_mean = torch.zeros(1, self.latent_dim, 1, 1) prior_logvar = torch.zeros(1, self.latent_dim, 1, 1) # log(1)=0 self.prior.set_prior(prior_mean, prior_logvar) def forward(self, x_past, t_span): """ x_past: [B, T_in, C_in, H, W] (C_in=1 for scalar fields) t_span: [T_out] times to evaluate the forecast (including start? we start from t=0 after encoding) Returns: predictions at t_span times, shape [B, T_out, 1, H, W] """ B, T_in, C_in, H, W = x_past.shape # Encode past sequence into initial state at t0 # Permute to [B, C_in, T_in, H, W] x = x_past.permute(0,2,1,3,4) h0 = self.encoder(x).squeeze(2) # [B, latent_dim, H, W] # Store grid shape for prior self.prior.grid_shape = (H, W) # Spatial propagation (multi‑directional information fusion) for _ in range(self.num_propagate): h0 = self.propagator(h0, (H, W)) # Evolve with Neural ODE # odeint expects state as tensor with first dimension = batch*... but works with any shape. # We'll keep [B, C, H, W] and pass a wrapper that flattens time dimension. # torchdiffeq.odeint can handle arbitrary shape as long as func returns same shape. t_span_tensor = t_span.to(x_past.device) h_t = odeint(self.ode_func, h0, t_span_tensor, method=self.ode_solver, options=self.ode_options) # h_t: [T_out, B, C, H, W] # Decode each time step preds = [] for i in range(h_t.size(0)): pred = self.decoder(h_t[i]) # [B, 1, H, W] preds.append(pred.unsqueeze(1)) # add time dim predictions = torch.cat(preds, dim=1) # [B, T_out, 1, H, W] return predictions, h_t # also return latent states for loss def collapse_loss(self, h_t, predictions, targets, t_span): """ Compute a physics‑inspired loss: 1. MSE prediction error 2. KL divergence from prior dynamics (information bottleneck) 3. Fisher sharpness term (encourages collapse, approximated by gradient norm) """ # 1. Prediction loss mse = F.mse_loss(predictions, targets) # 2. KL to prior (average over time steps) kl = 0.0 for i in range(h_t.size(0)): kl += self.prior.kl_divergence(h_t[i]) kl /= h_t.size(0) # 3. Fisher‑information sharpness: penalize flatness? Actually maximize Fisher # Approximate Fisher as norm of gradient of log‑likelihood w.r.t path parameters. # We use a simple surrogate: variance of predictions (sharp peaks have low variance). # Here we encourage small variance (collapse) by penalising output entropy. # We'll compute per‑pixel variance over time as a proxy for sharpness. var = predictions.var(dim=1).mean() # average over batch, spatial # Lower variance → sharper forecast (collapse). We want to minimise var, # but we add it as a small penalty (or we could maximise Fisher by minimising entropy). # We'll just use var as a regularizer. # Jarzynski‑type regularisation: penalise large jumps (energy dissipation) diff_h = (h_t[1:] - h_t[:-1]).pow(2).mean() total_loss = mse + self.beta * kl + 0.001 * var + 0.0001 * diff_h return total_loss, {'mse': mse.item(), 'kl': kl.item(), 'var': var.item(), 'diff': diff_h.item()} # ---------------------------------------------------------------------- # 5. Synthetic dataset: 2D heat equation with a nonlinear source # Provides spatio‑temporal data that needs direction‑aware forecasting. # ---------------------------------------------------------------------- def generate_heat_equation_data(num_samples=200, T=20, dt=0.1, grid_size=16): """ Solve ∂u/∂t = ∇²u + 0.1*u*(1-u) on a grid. Returns sequences of shape [N, T, 1, H, W]. We'll sample random initial conditions and noise. """ dx = 1.0 / (grid_size - 1) dt_sim = dt steps = T data = [] for _ in range(num_samples): u = torch.rand(1, grid_size, grid_size) * 0.5 # [0,0.5] u_seq = [u.clone()] for t in range(1, steps): # Laplacian using finite difference u_pad = F.pad(u, (1,1,1,1), mode='circular') laplacian = (u_pad[:, 2:, 1:-1] + u_pad[:, :-2, 1:-1] + u_pad[:, 1:-1, 2:] + u_pad[:, 1:-1, :-2] - 4*u) / (dx**2) nonlinear = 0.1 * u * (1 - u) u = u + dt_sim * (laplacian + nonlinear) + 0.01 * torch.randn_like(u) u_seq.append(u.clone()) seq = torch.stack(u_seq, dim=1) # [T, C, H, W] data.append(seq.unsqueeze(0)) # [1, T, C, H, W] data = torch.cat(data, dim=0) return data # ---------------------------------------------------------------------- # 6. Training loop # ---------------------------------------------------------------------- def train(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Hyperparameters batch_size = 16 latent_dim = 32 hidden_dim = 128 input_channels = 1 T_past = 10 T_future = 10 grid_size = 16 lr = 1e-3 epochs = 50 # Generate data full_data = generate_heat_equation_data(num_samples=500, T=T_past+T_future, grid_size=grid_size) full_data = full_data.to(device) # Split train_data = full_data[:400] val_data = full_data[400:] # Model model = NeuralODEFieldCollapser(input_channels, latent_dim, hidden_dim, num_propagate=2).to(device) optimizer = optim.Adam(model.parameters(), lr=lr) # Set a trivial prior (could be pretrained on historical data) model.set_prior_from_data(None) # Time span for ODE: we integrate from t=0 to t=T_future*dt (relative time) dt = 0.1 t_span = torch.arange(0, T_future*dt, dt, device=device) # length T_future for epoch in range(epochs): model.train() total_loss = 0.0 # Shuffle idx = torch.randperm(len(train_data)) for start in range(0, len(train_data), batch_size): end = start + batch_size batch_idx = idx[start:end] x_past = train_data[batch_idx, :T_past] # [B, T_past, C, H, W] y_future = train_data[batch_idx, T_past:] # [B, T_future, C, H, W] optimizer.zero_grad() preds, h_t = model(x_past, t_span) loss, loss_dict = model.collapse_loss(h_t, preds, y_future, t_span) loss.backward() optimizer.step() total_loss += loss.item() * (end-start) avg_loss = total_loss / len(train_data) # Validation with torch.no_grad(): x_val = val_data[:, :T_past] y_val = val_data[:, T_past:] pred_val, _ = model(x_val, t_span) val_mse = F.mse_loss(pred_val, y_val).item() print(f"Epoch {epoch:3d} | Train Loss: {avg_loss:.4f} | Val MSE: {val_mse:.4f}") return model if __name__ == "__main__": model = train()