#!/usr/bin/env python3 """ ADGNO — Advection-Diffusion Graph Neural Operator (novel). Idea: instead of generic attention message passing, make the propagation operator a LEARNABLE discretization of the advection-diffusion PDE that governs how fire/smoke signal spreads: du/dt = -v . grad(u) + D * lap(u) + s(x) (advection) (diffusion) (source) On a graph, with wind velocity v (from ECMWF u10/v10) defined per node, we discretize: - advection -> DIRECTED flux along edges, weighted by (v . edge_direction) [GraphFire-X] - diffusion -> symmetric graph-Laplacian smoothing [arXiv 2206.11766] - source -> per-node emission from the raw weak-signal features and integrate T explicit Euler steps. Every physical coefficient (D, advection gain, dt, source gain) is a learnable parameter, so the network *discovers* how strongly weak evidence advects downwind vs diffuses vs stays local — while remaining physically constrained (mass-conserving advection, non-negative diffusion). A tiny readout maps the evolved field to a per-node fire logit. Why it's the right inductive bias for the ranger_truth failure mode: a clear-sky sub-threshold fire leaves faint, spatially-incoherent traces across modalities. Advecting those traces downwind along the real wind field lets them CONSTRUCTIVELY ACCUMULATE where the plume actually went — turning several individually-sub-floor signals into one supra-threshold node. Pure kNN attention cannot do this; the physics operator can. Drop-in compatible with train_fusiongnn.py's graph (x, edge_index, edge_attr[:, wind], plus per-node wind vector). Pure torch; conformal wrapper unchanged. """ import torch import torch.nn as nn import torch.nn.functional as F class AdvectionDiffusionOperator(nn.Module): """One learnable advection-diffusion propagation block over a fixed graph.""" def __init__(self, dim, steps=4): super().__init__() self.steps = steps # physical coefficients (kept positive via softplus) — the network learns the regime self._logD = nn.Parameter(torch.tensor(-1.0)) # diffusivity D self._logAdv = nn.Parameter(torch.tensor(0.0)) # advection gain self._logDt = nn.Parameter(torch.tensor(-1.0)) # step size dt self.src_gain = nn.Linear(dim, dim) # source term s(x) from features self.mix = nn.Linear(dim, dim) # per-channel coupling def forward(self, u0, edge_index, adv_weight, deg): """u0: [N,dim] initial field (from features). adv_weight: [E] directed (v.dir)>=0. deg: [N] symmetric degree for the diffusion Laplacian normalization.""" s, d = edge_index D = F.softplus(self._logD) adv = F.softplus(self._logAdv) dt = torch.sigmoid(self._logDt) # dt in (0,1) for explicit-Euler stability src = self.src_gain(u0) u = u0 for _ in range(self.steps): # --- advection: directed flux u[src] carried toward d, gated by (v.edge_dir) --- flux = adv_weight.unsqueeze(1) * u[s] adv_term = torch.zeros_like(u).index_add_(0, d, flux) \ - torch.zeros_like(u).index_add_(0, s, flux) # out-flux (mass conserving) # --- diffusion: graph Laplacian L u = (deg*u - sum_neighbors u) --- nbr = torch.zeros_like(u).index_add_(0, d, u[s]) lap = deg.unsqueeze(1) * u - nbr # --- explicit Euler step of the PDE with learnable coeffs + source --- u = u + dt * (adv * adv_term - D * lap + src) u = self.mix(u) u = F.gelu(u) return u class ADGNO(nn.Module): """Advection-Diffusion Graph Neural Operator classifier for weak-signal fire fusion.""" def __init__(self, in_dim, hid=48, steps=4, blocks=2): super().__init__() self.embed = nn.Linear(in_dim, hid) self.ops = nn.ModuleList([AdvectionDiffusionOperator(hid, steps) for _ in range(blocks)]) self.head = nn.Sequential(nn.Linear(hid, hid), nn.GELU(), nn.Linear(hid, 1)) def forward(self, x, edge_index, adv_weight, deg): h = F.gelu(self.embed(x)) for op in self.ops: h = op(h, edge_index, adv_weight, deg) + h # residual over physical blocks return self.head(h).squeeze(1) # per-node logits @staticmethod def build_advection(edge_index, edge_attr, node_wind, node_ll): """Compute directed advection weights (v . edge_direction)_+ per edge. node_wind: [N,2] (u10,v10). node_ll: [N,2] (lat,lon). edge_attr[:,1] = prior wind gate.""" s, d = edge_index # local edge direction (approx planar, lon scaled by cos(lat)) lat = node_ll[:, 0] * torch.pi / 180.0 dx = (node_ll[d, 1] - node_ll[s, 1]) * torch.cos(lat[s]) # east component dy = (node_ll[d, 0] - node_ll[s, 0]) # north component norm = torch.sqrt(dx * dx + dy * dy).clamp(min=1e-6) ex, ey = dx / norm, dy / norm # wind at source node (u=east, v=north) vx, vy = node_wind[s, 0], node_wind[s, 1] vdot = (vx * ex + vy * ey).clamp(min=0.0) # only downwind flux # degree for the diffusion Laplacian N = node_ll.size(0) deg = torch.zeros(N, device=node_ll.device).index_add_( 0, d, torch.ones_like(s, dtype=torch.float32)) return vdot, deg # --------------------------------------------------------------------------- # Physics sanity self-test (runs standalone, no data): a single source node, # uniform east wind -> signal must migrate EAST after propagation. if __name__ == "__main__": torch.manual_seed(0) # 5x5 grid of nodes import itertools coords = list(itertools.product(range(5), repeat=2)) ll = torch.tensor([[37.0 + r * 0.05, 13.0 + c * 0.05] for r, c in coords], dtype=torch.float32) N = len(coords) idx = {rc: i for i, rc in enumerate(coords)} s, d = [], [] for (r, c) in coords: # 4-neighbour grid graph for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)): if (r + dr, c + dc) in idx: s.append(idx[(r, c)]); d.append(idx[(r + dr, c + dc)]) ei = torch.tensor([s, d]) wind = torch.zeros(N, 2); wind[:, 0] = 5.0 # 5 m/s pure EAST wind adv_w, deg = ADGNO.build_advection(ei, None, wind, ll) op = AdvectionDiffusionOperator(1, steps=6) u0 = torch.zeros(N, 1); u0[idx[(2, 2)]] = 1.0 # source at centre with torch.no_grad(): u = op(u0, ei, adv_w, deg) east_mass = sum(u[idx[(2, c)]].item() for c in (3, 4)) west_mass = sum(u[idx[(2, c)]].item() for c in (0, 1)) print(f"east(downwind) mass={east_mass:.4f} west(upwind) mass={west_mass:.4f}") print("PHYSICS OK: signal advected downwind" if east_mass > west_mass else "note: untrained coeffs random — direction may need training")