feat: Initial CBPOA commit — 武汉儿童呼吸疾病风险评估系统
Context: Build a spatial risk assessment system correlating air quality data with children's respiratory disease incidence across Wuhan. Approach: FastAPI backend serving PostGIS spatial queries, React frontend with Deck.gl maps, and a PyTorch SpatialTemporalGCN pipeline for multi-day (1d/3d/7d) risk prediction. Changes: - backend/ — FastAPI API with auth (JWT), alerts, risk analysis, geocoded case data, grid statistics, and report endpoints - frontend/ — React dashboard with interactive risk maps, alert monitoring, district comparison charts, and timeline player - models/ — SpatialTemporalGCN model with trained weights and ONNX export for inference - scripts/ — ETL pipeline for weather + medical data, grid generation, feature engineering, training, and daily inference - deploy/ — Docker Compose configs for backend, frontend, and MLflow - docs/ — API docs, deployment guide, user guide, and code review Impact: Enables spatial risk visualization, alert monitoring, and ML-driven health risk forecasting for environmental health teams.
This commit is contained in:
149
models/spatiotemporal_gcn/model.py
Normal file
149
models/spatiotemporal_gcn/model.py
Normal file
@@ -0,0 +1,149 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Spatial-Temporal Transformer + GCN Model for Wuhan Respiratory Disease Risk Prediction.
|
||||
Architecture per PRD acceptance criteria:
|
||||
- Temporal Transformer: 3 layers, 4 heads
|
||||
- GCN: 2 layers [GCNConv(48, 128) → ReLU → Dropout(0.2) → GCNConv(128, 64)]
|
||||
- Input: [N, T, 48] node features, [N, N] adjacency
|
||||
- Output: [N, 3] risk values (1-day, 3-day, 7-day)
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch_geometric.nn import GCNConv
|
||||
from torch_geometric.utils import add_self_loops
|
||||
|
||||
|
||||
class SpatialTemporalGCN(nn.Module):
|
||||
"""
|
||||
Spatial-Temporal Graph Convolutional Network with Transformer encoder.
|
||||
|
||||
Args:
|
||||
node_features (int): Number of input node features (default: 48)
|
||||
temporal_heads (int): Number of attention heads in Transformer (default: 4)
|
||||
temporal_layers (int): Number of Transformer layers (default: 3)
|
||||
gcn_hidden (int): Hidden dimension for GCN layers (default: 128)
|
||||
gcn_output (int): Output dimension of GCN (default: 64)
|
||||
dropout (float): Dropout rate (default: 0.2)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
node_features: int = 48,
|
||||
temporal_heads: int = 4,
|
||||
temporal_layers: int = 3,
|
||||
gcn_hidden: int = 128,
|
||||
gcn_output: int = 64,
|
||||
dropout: float = 0.2,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Temporal Transformer encoder
|
||||
encoder_layer = nn.TransformerEncoderLayer(
|
||||
d_model=node_features,
|
||||
nhead=temporal_heads,
|
||||
dim_feedforward=node_features * 4,
|
||||
dropout=dropout,
|
||||
activation='gelu',
|
||||
batch_first=True,
|
||||
norm_first=True,
|
||||
)
|
||||
self.temporal_transformer = nn.TransformerEncoder(
|
||||
encoder_layer,
|
||||
num_layers=temporal_layers,
|
||||
)
|
||||
|
||||
# GCN layers
|
||||
self.conv1 = GCNConv(node_features, gcn_hidden)
|
||||
self.conv2 = GCNConv(gcn_hidden, gcn_output)
|
||||
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.relu = nn.ReLU()
|
||||
|
||||
# Output head: 3 risk horizons (1-day, 3-day, 7-day)
|
||||
self.risk_head = nn.Linear(gcn_output, 3)
|
||||
|
||||
def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass.
|
||||
|
||||
Args:
|
||||
x: Node features [N, T, 48] — N nodes, T time steps, 48 features
|
||||
edge_index: Graph connectivity [2, E]
|
||||
|
||||
Returns:
|
||||
Risk predictions [N, 3] — 1-day, 3-day, 7-day risk
|
||||
"""
|
||||
N, T, F = x.shape
|
||||
|
||||
# Temporal Transformer: process each node's time series
|
||||
# Input [N, T, 48] → Transformer → [N, T, 48]
|
||||
x_temporal = self.temporal_transformer(x)
|
||||
|
||||
# Take the last time step as the spatial representation
|
||||
x_spatial = x_temporal[:, -1, :] # [N, 48]
|
||||
|
||||
# Add self-loops for GCN
|
||||
edge_index, _ = add_self_loops(edge_index, num_nodes=N)
|
||||
|
||||
# GCN layer 1: [N, 48] → [N, 128]
|
||||
x_gcn = self.conv1(x_spatial, edge_index)
|
||||
x_gcn = self.relu(x_gcn)
|
||||
x_gcn = self.dropout(x_gcn)
|
||||
|
||||
# GCN layer 2: [N, 128] → [N, 64]
|
||||
x_gcn = self.conv2(x_gcn, edge_index)
|
||||
x_gcn = self.relu(x_gcn)
|
||||
x_gcn = self.dropout(x_gcn)
|
||||
|
||||
# Risk prediction head: [N, 64] → [N, 3]
|
||||
risk = self.risk_head(x_gcn)
|
||||
|
||||
# Clamp output to [0, 1] range (risk probability)
|
||||
risk = torch.sigmoid(risk)
|
||||
|
||||
return risk
|
||||
|
||||
|
||||
def export_onnx(model, output_path: str, node_features: int = 48):
|
||||
"""Export model to ONNX format for inference."""
|
||||
model.eval()
|
||||
N = 512 # Dummy batch size for export
|
||||
|
||||
# Dummy inputs matching expected shapes
|
||||
dummy_x = torch.randn(N, 14, node_features) # [N, T=14, 48]
|
||||
dummy_edge_index = torch.randint(0, N, (2, N * 4)) # Sparse edges
|
||||
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(dummy_x, dummy_edge_index),
|
||||
output_path,
|
||||
input_names=['node_features', 'edge_index'],
|
||||
output_names=['risk'],
|
||||
dynamic_axes={
|
||||
'node_features': {0: 'num_nodes'},
|
||||
'edge_index': {1: 'num_edges'},
|
||||
'risk': {0: 'num_nodes'},
|
||||
},
|
||||
opset_version=17,
|
||||
)
|
||||
print(f"ONNX model exported to {output_path}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# Quick forward pass test on dummy data
|
||||
model = SpatialTemporalGCN()
|
||||
|
||||
# Dummy input: [512 nodes, 14 time steps, 48 features]
|
||||
N, T, F = 512, 14, 48
|
||||
x = torch.randn(N, T, F)
|
||||
edge_index = torch.randint(0, N, (2, N * 4))
|
||||
|
||||
risk = model(x, edge_index)
|
||||
print(f"Input: {x.shape}")
|
||||
print(f"Edge index: {edge_index.shape}")
|
||||
print(f"Output risk: {risk.shape} — 1d:{risk[:,0].mean():.3f}, 3d:{risk[:,1].mean():.3f}, 7d:{risk[:,2].mean():.3f}")
|
||||
|
||||
# ONNX export
|
||||
export_onnx(model, 'models/spatiotemporal_gcn/model_1_3_7.onnx')
|
||||
Reference in New Issue
Block a user