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:
244
models/spatiotemporal_gcn/sampler.py
Normal file
244
models/spatiotemporal_gcn/sampler.py
Normal file
@@ -0,0 +1,244 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
GraphSAINT-style Sampler for PyTorch Geometric.
|
||||
|
||||
Mini-batch sampler for large graphs (140k+ nodes) using neighbor sampling.
|
||||
Compatible with base PyG installation (no torch-sparse or pyg-lib required).
|
||||
|
||||
Usage:
|
||||
from models.spatiotemporal_gcn.sampler import GraphSAINTSampler
|
||||
|
||||
sampler = GraphSAINTSampler(
|
||||
data=data,
|
||||
batch_size=256,
|
||||
num_neighbors=[256, 128, 64]
|
||||
)
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torch_geometric.data import Data
|
||||
from torch_geometric.utils import subgraph
|
||||
|
||||
|
||||
class GraphSAINTDataset(Dataset):
|
||||
"""Dataset that samples node indices for mini-batching."""
|
||||
|
||||
def __init__(self, num_nodes: int, num_steps: int = 10):
|
||||
self.num_nodes = num_nodes
|
||||
self.num_steps = num_steps
|
||||
|
||||
def __len__(self):
|
||||
return self.num_steps
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return torch.randint(0, self.num_nodes, (1,))
|
||||
|
||||
|
||||
class GraphSAINTSampler:
|
||||
"""
|
||||
GraphSAINT-style mini-batch sampler for large graphs.
|
||||
|
||||
Implements neighbor sampling to create subgraphs that fit in GPU memory.
|
||||
For each batch, samples seed nodes and their multi-hop neighbors.
|
||||
|
||||
Args:
|
||||
data: Full graph with edge_index and node features.
|
||||
batch_size: Seed nodes per batch (default: 256).
|
||||
num_neighbors: Neighbors per layer [layer0, layer1, ...].
|
||||
Default: [256, 128, 64] for 3-layer GCN.
|
||||
num_steps: Batches per epoch (default: 10).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data: Data,
|
||||
batch_size: int = 256,
|
||||
num_neighbors: list = None,
|
||||
num_steps: int = 10,
|
||||
):
|
||||
if num_neighbors is None:
|
||||
num_neighbors = [256, 128, 64]
|
||||
|
||||
self.data = data
|
||||
self.batch_size = batch_size
|
||||
self.num_neighbors = num_neighbors
|
||||
self.num_steps = num_steps
|
||||
self.num_nodes = data.num_nodes
|
||||
self.edge_index = data.edge_index
|
||||
|
||||
if data.num_nodes > 100000:
|
||||
print(f"Sampler for large graph: {data.num_nodes:,} nodes")
|
||||
print(f" Batch size: {batch_size}")
|
||||
print(f" Layer depths: {num_neighbors}")
|
||||
|
||||
def _sample_neighbors(self, seed_nodes: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Sample multi-hop neighbors for seed nodes.
|
||||
|
||||
Args:
|
||||
seed_nodes: Initial node indices.
|
||||
|
||||
Returns:
|
||||
All sampled node indices (seed + neighbors).
|
||||
"""
|
||||
sampled = seed_nodes.unique()
|
||||
|
||||
for num_neighbors in self.num_neighbors:
|
||||
if len(sampled) == 0:
|
||||
break
|
||||
|
||||
mask = torch.isin(self.edge_index[0], sampled)
|
||||
neighbor_edges = self.edge_index[:, mask]
|
||||
|
||||
if neighbor_edges.shape[1] == 0:
|
||||
break
|
||||
|
||||
neighbors = neighbor_edges[1]
|
||||
|
||||
if len(neighbors) > num_neighbors:
|
||||
neighbors = neighbors[torch.randperm(len(neighbors))[:num_neighbors]]
|
||||
|
||||
sampled = torch.cat([sampled, neighbors]).unique()
|
||||
|
||||
return sampled
|
||||
|
||||
def _create_subgraph(self, node_indices: torch.Tensor) -> Data:
|
||||
edge_index, _, edge_mask = subgraph(
|
||||
node_indices,
|
||||
self.edge_index,
|
||||
relabel_nodes=True,
|
||||
return_edge_mask=True,
|
||||
)
|
||||
|
||||
subgraph_data = Data(
|
||||
x=self.data.x[node_indices],
|
||||
edge_index=edge_index,
|
||||
n_id=node_indices,
|
||||
)
|
||||
|
||||
if hasattr(self.data, 'y') and self.data.y is not None:
|
||||
subgraph_data.y = self.data.y[node_indices]
|
||||
|
||||
return subgraph_data
|
||||
|
||||
def __iter__(self):
|
||||
for _ in range(self.num_steps):
|
||||
seed_nodes = torch.randint(0, self.num_nodes, (self.batch_size,))
|
||||
sampled_nodes = self._sample_neighbors(seed_nodes)
|
||||
batch = self._create_subgraph(sampled_nodes)
|
||||
yield batch
|
||||
|
||||
def __len__(self):
|
||||
return self.num_steps
|
||||
|
||||
|
||||
class GraphSAINTConfig:
|
||||
"""Configuration for GraphSAINT-style sampling."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
batch_size: int = 256,
|
||||
num_neighbors: list = None,
|
||||
num_steps: int = 10,
|
||||
):
|
||||
self.batch_size = batch_size
|
||||
self.num_neighbors = num_neighbors if num_neighbors is not None else [256, 128, 64]
|
||||
self.num_steps = num_steps
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
f"GraphSAINTConfig(\n"
|
||||
f" batch_size={self.batch_size},\n"
|
||||
f" num_neighbors={self.num_neighbors},\n"
|
||||
f" num_steps={self.num_steps}\n"
|
||||
f")"
|
||||
)
|
||||
|
||||
|
||||
def create_graph_saint_loader(
|
||||
data: Data,
|
||||
batch_size: int = 256,
|
||||
num_neighbors: list = None,
|
||||
num_steps: int = 10,
|
||||
):
|
||||
"""
|
||||
Create a GraphSAINT-style sampler for large graph training.
|
||||
|
||||
Args:
|
||||
data: Full graph data with edge_index and features.
|
||||
batch_size: Seed nodes per batch (default: 256).
|
||||
num_neighbors: Layer-wise neighbor counts (default: [256, 128, 64]).
|
||||
num_steps: Batches per epoch (default: 10).
|
||||
|
||||
Returns:
|
||||
GraphSAINTSampler: Mini-batch iterator.
|
||||
"""
|
||||
return GraphSAINTSampler(
|
||||
data=data,
|
||||
batch_size=batch_size,
|
||||
num_neighbors=num_neighbors,
|
||||
num_steps=num_steps,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
"""Example usage with dummy data."""
|
||||
print("=" * 60)
|
||||
print("GraphSAINT-style Sampler Demo")
|
||||
print("=" * 60)
|
||||
|
||||
print("\nCreating dummy graph (10k nodes)...")
|
||||
N = 10000
|
||||
num_features = 48
|
||||
|
||||
edge_index = torch.randint(0, N, (2, N * 3))
|
||||
x = torch.randn(N, num_features)
|
||||
y = torch.randint(0, 3, (N,))
|
||||
|
||||
data = Data(x=x, y=y, edge_index=edge_index)
|
||||
print(f" Nodes: {data.num_nodes:,}")
|
||||
print(f" Edges: {data.num_edges:,}")
|
||||
print(f" Features: {data.num_node_features}")
|
||||
|
||||
print("\nCreating sampler...")
|
||||
config = GraphSAINTConfig(
|
||||
batch_size=256,
|
||||
num_neighbors=[256, 128, 64],
|
||||
num_steps=5,
|
||||
)
|
||||
print(config)
|
||||
|
||||
loader = create_graph_saint_loader(
|
||||
data=data,
|
||||
batch_size=config.batch_size,
|
||||
num_neighbors=config.num_neighbors,
|
||||
num_steps=config.num_steps,
|
||||
)
|
||||
|
||||
print(f"\nIterating through {len(loader)} batches...")
|
||||
for i, batch in enumerate(loader):
|
||||
print(f" Batch {i+1}/{len(loader)}:")
|
||||
print(f" Nodes: {batch.num_nodes:,}")
|
||||
print(f" Edges: {batch.num_edges:,}")
|
||||
print(f" Features: {batch.x.shape}")
|
||||
print(f" Node IDs: {batch.n_id.shape}")
|
||||
|
||||
if i >= 2:
|
||||
break
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Sampler ready for training!")
|
||||
print("=" * 60)
|
||||
|
||||
print("\nFor your 140k node graph:")
|
||||
print(" 1. Load graph: data = load_your_graph()")
|
||||
print(" 2. Create loader: loader = create_graph_saint_loader(data, batch_size=256)")
|
||||
print(" 3. Train: for batch in loader: out = model(batch.x, batch.edge_index)")
|
||||
print("\nRecommended for 4GB GPU:")
|
||||
print(" - batch_size: 256")
|
||||
print(" - num_neighbors: [256, 128, 64]")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user