""" model.py — GNN Architectures for Microplastic Source Attribution ================================================================ Implements three models for node-level concentration regression: 1. GraphSAGE — inductive, scalable neighbourhood aggregation Hamilton, W. et al. (2017). Inductive Representation Learning on Large Graphs. NeurIPS 2017. https://arxiv.org/abs/1706.02216 2. GAT — Graph Attention Network with interpretable attention weights Veličković, P. et al. (2018). Graph Attention Networks. ICLR 2018. https://arxiv.org/abs/1710.10903 3. Classical baseline — graph centrality features + linear regression (Compares traditional graph mining with modern deep GNNs.) Architecture note: attention weights in the GAT head serve as proxy source-contribution scores, analogous to transfer-entropy edge weights used in network connectivity analysis. """ import torch import torch.nn as nn import torch.nn.functional as F from torch_geometric.nn import ( SAGEConv, GATConv, global_mean_pool, ) from torch_geometric.nn import MessagePassing from torch_geometric.utils import add_self_loops, degree import numpy as np # ────────────────────────────────────────────────────────────────────────────── # 1. GraphSAGE Concentration Regressor # ────────────────────────────────────────────────────────────────────────────── class GraphSAGERegressor(nn.Module): """ GraphSAGE for node-level concentration prediction. Architecture: Input → SAGEConv(128) → BN → ReLU → Dropout → SAGEConv(64) → BN → ReLU → Dropout → SAGEConv(32) → BN → ReLU → Linear(1) → scalar log-concentration The model is trained to predict log(concentration) to handle the log-normal distribution of microplastic counts. """ def __init__( self, in_channels: int = 9, hidden_channels: int = 128, out_channels: int = 1, num_layers: int = 3, dropout: float = 0.3, ): super().__init__() self.num_layers = num_layers self.dropout = dropout self.convs = nn.ModuleList() self.bns = nn.ModuleList() dims = [in_channels] + [hidden_channels] * (num_layers - 1) + [32] for i in range(num_layers): self.convs.append(SAGEConv(dims[i], dims[i + 1])) self.bns.append(nn.BatchNorm1d(dims[i + 1])) self.head = nn.Linear(32, out_channels) def forward(self, x, edge_index, edge_attr=None, return_embeddings=False): for i, (conv, bn) in enumerate(zip(self.convs, self.bns)): x = conv(x, edge_index) x = bn(x) x = F.relu(x) if i < self.num_layers - 1: x = F.dropout(x, p=self.dropout, training=self.training) embeddings = x.clone() out = self.head(x) if return_embeddings: return out, embeddings return out # [N, 1] log-concentration predictions # ────────────────────────────────────────────────────────────────────────────── # 2. GAT Concentration Regressor # ────────────────────────────────────────────────────────────────────────────── class GATRegressor(nn.Module): """ Graph Attention Network for node-level concentration prediction. The multi-head attention weights α_ij serve as interpretability signals: higher α_ij between source i and station j → source i contributes more to the concentration at j. This mirrors the transfer entropy / effective connectivity framework used in network science: both methods ask "how much does node A's state influence node B?" Architecture: Input → GAT(heads=8, 64-per-head) → ELU → Dropout → GAT(heads=4, 32-per-head) → ELU → Dropout → GAT(heads=1, 32) → ELU → Linear(1) """ def __init__( self, in_channels: int = 9, hidden_channels: int = 64, out_channels: int = 1, heads: int = 8, dropout: float = 0.3, ): super().__init__() self.dropout = dropout self.conv1 = GATConv( in_channels, hidden_channels, heads=heads, dropout=dropout, concat=True ) self.conv2 = GATConv( hidden_channels * heads, 32, heads=4, dropout=dropout, concat=True ) self.conv3 = GATConv( 32 * 4, 32, heads=1, dropout=dropout, concat=False ) self.bn1 = nn.BatchNorm1d(hidden_channels * heads) self.bn2 = nn.BatchNorm1d(32 * 4) self.bn3 = nn.BatchNorm1d(32) self.head = nn.Linear(32, out_channels) # Store last attention weights for attribution self._last_attention = None def forward(self, x, edge_index, edge_attr=None, return_attention=False): # Layer 1 x, (edge_idx1, alpha1) = self.conv1( x, edge_index, return_attention_weights=True ) x = self.bn1(x) x = F.elu(x) x = F.dropout(x, p=self.dropout, training=self.training) # Layer 2 x, (edge_idx2, alpha2) = self.conv2( x, edge_index, return_attention_weights=True ) x = self.bn2(x) x = F.elu(x) x = F.dropout(x, p=self.dropout, training=self.training) # Layer 3 x, (edge_idx3, alpha3) = self.conv3( x, edge_index, return_attention_weights=True ) x = self.bn3(x) x = F.elu(x) # Store attention weights for attribution (use last layer) self._last_attention = { "edge_index": edge_idx3.detach(), "alpha": alpha3.detach(), } embeddings = x.clone() out = self.head(x) if return_attention: return out, (edge_idx3, alpha3), embeddings return out # [N, 1] def get_attention_weights(self): """Return the last forward pass attention weights.""" return self._last_attention # ────────────────────────────────────────────────────────────────────────────── # 3. Classical Baseline — Graph Centrality + Linear Regression # ────────────────────────────────────────────────────────────────────────────── class ClassicalBaseline: """ Classical graph-mining baseline: 1. Compute graph centrality metrics (in-degree, betweenness, PageRank, closeness) for each node. 2. Concatenate with raw node features. 3. Fit Ridge regression to predict log-concentration. This baseline lets us quantify the value added by GNN message-passing over traditional centrality-based features — a direct comparison that validates the GNN approach (analogous to comparing transfer entropy baselines with GNN-based brain connectivity analysis). """ def __init__(self, alpha: float = 1.0): from sklearn.linear_model import Ridge from sklearn.preprocessing import StandardScaler from sklearn.pipeline import Pipeline self.model = Pipeline([ ("scaler", StandardScaler()), ("ridge", Ridge(alpha=alpha)), ]) self.centrality_features = None self.is_fitted = False def compute_centrality_features(self, G_nx, num_nodes: int) -> np.ndarray: """Compute centrality vectors for all nodes.""" import networkx as nx in_deg = dict(G_nx.in_degree()) out_deg = dict(G_nx.out_degree()) # PageRank try: pr = nx.pagerank(G_nx, alpha=0.85, max_iter=200) except Exception: pr = {n: 1.0 / num_nodes for n in G_nx.nodes()} # Betweenness (sample-based for speed) try: bc = nx.betweenness_centrality(G_nx, k=min(50, num_nodes), normalized=True) except Exception: bc = {n: 0.0 for n in G_nx.nodes()} # Closeness on undirected version try: cl = nx.closeness_centrality(G_nx.to_undirected()) except Exception: cl = {n: 0.0 for n in G_nx.nodes()} feats = np.zeros((num_nodes, 5)) for n in range(num_nodes): feats[n, 0] = in_deg.get(n, 0) feats[n, 1] = out_deg.get(n, 0) feats[n, 2] = pr.get(n, 0) feats[n, 3] = bc.get(n, 0) feats[n, 4] = cl.get(n, 0) self.centrality_features = feats return feats def fit( self, x: np.ndarray, # [N, node_feat_dim] raw node features centrality: np.ndarray, # [N, 5] centrality features y: np.ndarray, # [N] log-concentration targets mask: np.ndarray, # boolean mask — which nodes have labels ): combined = np.concatenate([x[mask], centrality[mask]], axis=1) self.model.fit(combined, y[mask]) self.is_fitted = True def predict( self, x: np.ndarray, centrality: np.ndarray, mask: np.ndarray, ) -> np.ndarray: combined = np.concatenate([x[mask], centrality[mask]], axis=1) return self.model.predict(combined) def score( self, x: np.ndarray, centrality: np.ndarray, y: np.ndarray, mask: np.ndarray, ) -> float: preds = self.predict(x, centrality, mask) return float(np.corrcoef(preds, y[mask])[0, 1] ** 2) # ────────────────────────────────────────────────────────────────────────────── # Utility: build graph-level dataset for node regression # ────────────────────────────────────────────────────────────────────────────── def build_node_regression_targets(df_split, data, station_ids): """ For a given time-period split DataFrame, compute per-station mean log-concentration and return as a tensor aligned with node indices. Returns ------- y : torch.Tensor [N, 1] — log-concentration for station nodes, 0 elsewhere mask : torch.BoolTensor [N] — True for station nodes that have data """ N = data.num_nodes y = torch.zeros(N, 1, dtype=torch.float) mask = torch.zeros(N, dtype=torch.bool) for s_id in station_ids: rows = df_split[df_split["station_id"] == s_id] if len(rows) > 0: mean_log_conc = rows["log_concentration"].mean() y[s_id, 0] = mean_log_conc mask[s_id] = True return y, mask if __name__ == "__main__": # Quick smoke test import torch x = torch.randn(200, 9) edge_index = torch.randint(0, 200, (2, 500)) sage = GraphSAGERegressor(in_channels=9) out_sage = sage(x, edge_index) print(f"GraphSAGE output shape: {out_sage.shape}") gat = GATRegressor(in_channels=9) out_gat = gat(x, edge_index) print(f"GAT output shape: {out_gat.shape}") print("Model smoke test passed.")