""" graph_builder.py — Hydrological Flow Graph Builder for Coastal Georgia Rivers ============================================================================== Builds a synthetic but geographically realistic directed flow graph for the Ogeechee, Savannah, and Altamaha river systems in coastal Georgia (USA). Graph structure (200 nodes total): - 50 sampling stations (measurement sites) - 100 candidate sources (factories, urban runoff, agricultural runoff) - 50 river junctions (confluence / routing nodes) Geography grounded in real HydroSHEDS-style drainage networks. Outputs a PyTorch Geometric Data object and a GeoJSON file. References: - Lehner, B. et al. (2008). New global hydrography derived from spaceborne elevation data. Geophysical Research Letters, 35(10). https://www.hydrosheds.org/ - Wang, M. et al. (2019). Deep Graph Library. arXiv:1909.01315. """ import json import math import random from pathlib import Path import networkx as nx import numpy as np import torch from torch_geometric.data import Data from torch_geometric.utils import from_networkx # --------------------------------------------------------------------------- # Reproducibility # --------------------------------------------------------------------------- SEED = 42 random.seed(SEED) np.random.seed(SEED) torch.manual_seed(SEED) # --------------------------------------------------------------------------- # Real geographic anchor points for coastal Georgia river systems # (lat, lon) — sourced from USGS StreamStats / NHD Plus # --------------------------------------------------------------------------- # Savannah River corridor (Hardeeville, SC / Savannah, GA to Augusta, GA) SAVANNAH_ANCHORS = [ (32.0835, -81.0998), # Savannah River mouth (Port of Savannah area) (32.1200, -81.2000), # Lower Savannah (32.3000, -81.4500), # Mid Savannah (32.5500, -81.7000), # Upper mid Savannah (32.8000, -81.9000), # Sylvania area (33.0000, -81.9800), # Statesboro / Millen area (33.3500, -82.0500), # Waynesboro area (33.4800, -82.0800), # Augusta confluence ] # Ogeechee River corridor OGEECHEE_ANCHORS = [ (31.8600, -81.1800), # Ogeechee mouth / Ossabaw Sound (31.9600, -81.3500), # Lower Ogeechee (32.0800, -81.6000), # Richmond Hill area (32.2000, -81.8500), # Mid Ogeechee (32.5000, -82.0000), # Statesboro area (32.6500, -82.2000), # Upper mid Ogeechee (32.8500, -82.5000), # Upper Ogeechee / Sandersville area ] # Altamaha River corridor ALTAMAHA_ANCHORS = [ (31.3300, -81.3000), # Altamaha mouth / Wolf Island (31.4500, -81.5000), # Lower Altamaha (31.6500, -81.7000), # Everett City area (31.8500, -81.9000), # Jesup area (31.9500, -82.1500), # Baxley area (32.1000, -82.4000), # Hazlehurst area (32.2000, -82.6500), # Douglas area (31.7500, -82.3000), # Ocmulgee tributary junction (31.6500, -82.5000), # Ocmulgee lower reach ] # Major coastal Georgia cities / industrial zones for source locations SOURCE_ANCHORS = { "savannah_port_industrial": (32.0835, -81.0998), "pooler_manufacturing": (32.1156, -81.2484), "statesboro_urban": (32.4487, -81.7831), "augusta_industrial": (33.4735, -82.0105), "brunswick_industrial": (31.1499, -81.4915), "jesup_paper_mill": (31.6077, -81.8849), "douglas_agri": (31.5087, -82.8510), "baxley_forestry": (31.7805, -82.3485), "waynesboro_agri": (33.0898, -82.0136), "sylvania_agri": (32.7479, -81.6370), "reidsville_urban": (32.0843, -82.1151), "vidalia_agri": (32.2174, -82.4135), "claxton_agri": (32.1601, -81.9071), "glennville_agri": (31.9324, -81.9239), "darien_coastal": (31.3716, -81.4315), "macon_industrial": (32.8407, -83.6324), "hinesville_urban": (31.8468, -81.5962), "richmond_hill_suburban": (31.9249, -81.3001), } def _interpolate_points(anchors: list, n: int) -> list: """Linearly interpolate n points along a sequence of anchor (lat, lon) pairs.""" total_dist = 0.0 segments = [] for i in range(len(anchors) - 1): lat1, lon1 = anchors[i] lat2, lon2 = anchors[i + 1] d = math.hypot(lat2 - lat1, lon2 - lon1) segments.append((d, anchors[i], anchors[i + 1])) total_dist += d points = [] cumulative = 0.0 seg_idx = 0 for k in range(n): t_global = (k / (n - 1)) * total_dist if n > 1 else 0 while seg_idx < len(segments) - 1 and cumulative + segments[seg_idx][0] < t_global: cumulative += segments[seg_idx][0] seg_idx += 1 seg_len, (lat1, lon1), (lat2, lon2) = segments[seg_idx] if seg_len > 0: t_local = (t_global - cumulative) / seg_len else: t_local = 0.0 t_local = max(0.0, min(1.0, t_local)) lat = lat1 + t_local * (lat2 - lat1) lon = lon1 + t_local * (lon2 - lon1) points.append((lat, lon)) return points def _elevation_from_coords(lat: float, lon: float) -> float: """ Approximate elevation (meters) from coordinates using a simple gradient model: coastal GA rises ~1 m per ~0.1° northward/westward from sea level at the coast (lat ~31.3°, lon ~-81.0°). """ base_lat, base_lon = 31.3, -81.0 elev = (lat - base_lat) * 15.0 + (-(lon) - 81.0) * 8.0 return max(0.0, elev + np.random.normal(0, 2)) def _population_density(node_type: str, lat: float, lon: float) -> float: """ Approximate population density (persons/km²) based on proximity to known urban centers in coastal Georgia. """ urban_centers = [ (32.0835, -81.0998, 1200), # Savannah (33.4735, -82.0105, 800), # Augusta (31.1499, -81.4915, 600), # Brunswick (32.4487, -81.7831, 400), # Statesboro (31.8468, -81.5962, 300), # Hinesville ] if node_type == "junction": return 0.0 density = 20.0 # rural baseline for clat, clon, pop in urban_centers: dist = math.hypot(lat - clat, lon - clon) if dist < 0.5: density = max(density, pop * math.exp(-dist / 0.2)) return density + abs(np.random.normal(0, density * 0.1)) # --------------------------------------------------------------------------- # Node-type encoding # --------------------------------------------------------------------------- NODE_TYPES = ["station", "factory", "urban_runoff", "agricultural_runoff", "junction"] # one-hot positions: 0=station, 1=factory, 2=urban_runoff, 3=agri_runoff, 4=junction TYPE_TO_IDX = {t: i for i, t in enumerate(NODE_TYPES)} def _one_hot(idx: int, n: int = 5) -> list: v = [0.0] * n v[idx] = 1.0 return v # --------------------------------------------------------------------------- # Main graph builder # --------------------------------------------------------------------------- def build_graph(save_dir: str = None) -> Data: """ Build and return the coastal Georgia hydrological flow graph. Returns ------- data : torch_geometric.data.Data Graph with node features x, edge_index, edge_attr, and metadata. """ G = nx.DiGraph() node_meta = {} # node_id -> dict of metadata node_id = 0 # ── 1. Sampling Stations (50) ────────────────────────────────────────── # Distribute across three rivers sav_stations = _interpolate_points(SAVANNAH_ANCHORS, 18) oge_stations = _interpolate_points(OGEECHEE_ANCHORS, 16) alt_stations = _interpolate_points(ALTAMAHA_ANCHORS, 16) station_ids = [] for pts, river in [(sav_stations, "savannah"), (oge_stations, "ogeechee"), (alt_stations, "altamaha")]: for lat, lon in pts: jitter_lat = lat + np.random.normal(0, 0.005) jitter_lon = lon + np.random.normal(0, 0.005) ntype = "station" meta = { "node_id": node_id, "node_type": ntype, "type_idx": TYPE_TO_IDX[ntype], "lat": jitter_lat, "lon": jitter_lon, "elevation": _elevation_from_coords(jitter_lat, jitter_lon), "population_density": _population_density(ntype, jitter_lat, jitter_lon), "river": river, "label": f"station_{node_id:03d}", } G.add_node(node_id, **meta) node_meta[node_id] = meta station_ids.append(node_id) node_id += 1 # ── 2. Candidate Sources (100) ───────────────────────────────────────── # Mix: 30 factories, 35 urban runoff, 35 agricultural runoff source_ids = [] # Factories near industrial anchors factory_anchors = [ SOURCE_ANCHORS["savannah_port_industrial"], SOURCE_ANCHORS["pooler_manufacturing"], SOURCE_ANCHORS["jesup_paper_mill"], SOURCE_ANCHORS["augusta_industrial"], SOURCE_ANCHORS["brunswick_industrial"], ] for i in range(30): base = factory_anchors[i % len(factory_anchors)] lat = base[0] + np.random.normal(0, 0.05) lon = base[1] + np.random.normal(0, 0.05) ntype = "factory" meta = { "node_id": node_id, "node_type": ntype, "type_idx": TYPE_TO_IDX[ntype], "lat": lat, "lon": lon, "elevation": _elevation_from_coords(lat, lon) + np.random.uniform(0, 10), "population_density": _population_density(ntype, lat, lon), "river": "mixed", "label": f"factory_{i:03d}", } G.add_node(node_id, **meta) node_meta[node_id] = meta source_ids.append(node_id) node_id += 1 # Urban runoff near cities urban_anchors = [ SOURCE_ANCHORS["savannah_port_industrial"], SOURCE_ANCHORS["statesboro_urban"], SOURCE_ANCHORS["hinesville_urban"], SOURCE_ANCHORS["richmond_hill_suburban"], SOURCE_ANCHORS["reidsville_urban"], ] for i in range(35): base = urban_anchors[i % len(urban_anchors)] lat = base[0] + np.random.normal(0, 0.08) lon = base[1] + np.random.normal(0, 0.08) ntype = "urban_runoff" meta = { "node_id": node_id, "node_type": ntype, "type_idx": TYPE_TO_IDX[ntype], "lat": lat, "lon": lon, "elevation": _elevation_from_coords(lat, lon) + np.random.uniform(0, 5), "population_density": _population_density(ntype, lat, lon), "river": "mixed", "label": f"urban_runoff_{i:03d}", } G.add_node(node_id, **meta) node_meta[node_id] = meta source_ids.append(node_id) node_id += 1 # Agricultural runoff agri_anchors = [ SOURCE_ANCHORS["sylvania_agri"], SOURCE_ANCHORS["vidalia_agri"], SOURCE_ANCHORS["claxton_agri"], SOURCE_ANCHORS["glennville_agri"], SOURCE_ANCHORS["baxley_forestry"], SOURCE_ANCHORS["douglas_agri"], SOURCE_ANCHORS["waynesboro_agri"], ] for i in range(35): base = agri_anchors[i % len(agri_anchors)] lat = base[0] + np.random.normal(0, 0.10) lon = base[1] + np.random.normal(0, 0.10) ntype = "agricultural_runoff" meta = { "node_id": node_id, "node_type": ntype, "type_idx": TYPE_TO_IDX[ntype], "lat": lat, "lon": lon, "elevation": _elevation_from_coords(lat, lon) + np.random.uniform(0, 3), "population_density": 0.0, "river": "mixed", "label": f"agri_runoff_{i:03d}", } G.add_node(node_id, **meta) node_meta[node_id] = meta source_ids.append(node_id) node_id += 1 # ── 3. River Junctions (50) ──────────────────────────────────────────── junction_locs = ( _interpolate_points(SAVANNAH_ANCHORS, 18)[:15] + _interpolate_points(OGEECHEE_ANCHORS, 14)[:12] + _interpolate_points(ALTAMAHA_ANCHORS, 16)[:12] + [ (32.25, -81.90), # Ogeechee / Canoochee confluence (31.80, -82.10), # Altamaha / Ocmulgee confluence (31.55, -81.40), # Altamaha lower (31.40, -81.38), # Altamaha delta (32.05, -81.12), # Savannah lower (32.30, -81.55), # Savannah mid (32.70, -81.80), # Savannah upper mid (33.10, -81.95), # Savannah upper (31.95, -81.65), # Ogeechee mid (32.60, -82.10), # Ogeechee upper (31.70, -82.05), # Altamaha inner ] ) # Ensure exactly 50 junctions junction_locs = junction_locs[:50] while len(junction_locs) < 50: lat = random.uniform(31.3, 33.5) lon = random.uniform(-83.0, -81.0) junction_locs.append((lat, lon)) junction_ids = [] for i, (lat, lon) in enumerate(junction_locs): ntype = "junction" meta = { "node_id": node_id, "node_type": ntype, "type_idx": TYPE_TO_IDX[ntype], "lat": lat, "lon": lon, "elevation": _elevation_from_coords(lat, lon), "population_density": 0.0, "river": "junction", "label": f"junction_{i:03d}", } G.add_node(node_id, **meta) node_meta[node_id] = meta junction_ids.append(node_id) node_id += 1 assert node_id == 200, f"Expected 200 nodes, got {node_id}" # ── 4. Build Directed Edges (downstream flow) ────────────────────────── def haversine(lat1, lon1, lat2, lon2): R = 6371.0 dlat = math.radians(lat2 - lat1) dlon = math.radians(lon2 - lon1) a = math.sin(dlat / 2) ** 2 + math.cos(math.radians(lat1)) * math.cos(math.radians(lat2)) * math.sin(dlon / 2) ** 2 return R * 2 * math.asin(math.sqrt(a)) def flow_rate(elev_up: float, elev_down: float, dist_km: float) -> float: """Synthetic flow rate (m³/s) — higher gradient = higher flow.""" grad = max(0, elev_up - elev_down) / max(dist_km, 0.1) return 0.5 + grad * 50 + abs(np.random.normal(0, 0.5)) def wind_correlation(lat1, lon1, lat2, lon2) -> float: """Prevailing wind in coastal GA is SW→NE; correlation if edge aligns.""" dlat = lat2 - lat1 dlon = lon2 - lon1 # Prevailing direction: NE (dlat>0, dlon>0) from SW origin angle = math.atan2(dlat, dlon) # radians preferred = math.radians(45) # NE corr = math.cos(angle - preferred) return max(0.0, corr) + abs(np.random.normal(0, 0.05)) added_edges = set() def add_edge(u: int, v: int): if u == v or (u, v) in added_edges: return meta_u = node_meta[u] meta_v = node_meta[v] dist = haversine(meta_u["lat"], meta_u["lon"], meta_v["lat"], meta_v["lon"]) fr = flow_rate(meta_u["elevation"], meta_v["elevation"], dist) wc = wind_correlation(meta_u["lat"], meta_u["lon"], meta_v["lat"], meta_v["lon"]) G.add_edge(u, v, flow_rate=fr, distance=dist, wind_correlation=wc) added_edges.add((u, v)) # --- Station-to-station edges along each river (longitudinal flow) --- sav_stn = station_ids[:18] # Savannah: upstream (high index) → downstream (low index lat-wise) oge_stn = station_ids[18:34] alt_stn = station_ids[34:50] # Sort stations by latitude descending (upstream=high lat → downstream=low lat) for river_stations in [sav_stn, oge_stn, alt_stn]: ordered = sorted(river_stations, key=lambda n: -node_meta[n]["lat"]) for i in range(len(ordered) - 1): add_edge(ordered[i], ordered[i + 1]) # --- Junctions connect rivers and route flow --- # Connect some junctions into the station chains for jid in junction_ids: jlat = node_meta[jid]["lat"] jlon = node_meta[jid]["lon"] # Find nearby stations and connect junction → downstream station nearby = sorted( station_ids, key=lambda s: haversine(jlat, jlon, node_meta[s]["lat"], node_meta[s]["lon"]) ) # Connect to 1-2 nearest downstream stations for stn in nearby[:2]: slat = node_meta[stn]["lat"] # Junction → station if junction is upstream (higher lat/elev) if node_meta[jid]["elevation"] >= node_meta[stn]["elevation"]: add_edge(jid, stn) else: add_edge(stn, jid) # --- Sources connect to nearest junction or station (sources → downstream nodes) --- all_routing = junction_ids + station_ids for src_id in source_ids: slat = node_meta[src_id]["lat"] slon = node_meta[src_id]["lon"] dists = [(haversine(slat, slon, node_meta[n]["lat"], node_meta[n]["lon"]), n) for n in all_routing] dists.sort() # Connect to 1-3 nearest routing nodes for d, n in dists[:3]: if d < 50.0: # within 50 km # Sources always flow downstream add_edge(src_id, n) # --- Inter-junction edges (tributary merges) --- for i, jid in enumerate(junction_ids): jlat = node_meta[jid]["lat"] jlon = node_meta[jid]["lon"] dists = [(haversine(jlat, jlon, node_meta[jid2]["lat"], node_meta[jid2]["lon"]), jid2) for jid2 in junction_ids if jid2 != jid] dists.sort() for d, jid2 in dists[:2]: if d < 80.0: # Higher elevation flows to lower if node_meta[jid]["elevation"] > node_meta[jid2]["elevation"]: add_edge(jid, jid2) else: add_edge(jid2, jid) # --- Ensure strong connectivity: Add a few cross-river edges --- add_edge(sav_stn[-1], oge_stn[-1]) # Savannah outlet → Ogeechee outlet area add_edge(oge_stn[-1], alt_stn[-1]) # Ogeechee outlet → Altamaha outlet area print(f"Graph built: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges") # ── 5. Assemble node features ────────────────────────────────────────── # Feature vector per node: # [lat, lon, elevation, population_density, one_hot×5] = 9 features node_features = [] node_labels = [] node_types_list = [] node_lats = [] node_lons = [] for nid in range(200): m = node_meta[nid] oh = _one_hot(m["type_idx"], 5) # Normalize lat/lon to [0,1] range for the study area norm_lat = (m["lat"] - 31.3) / (33.5 - 31.3) norm_lon = (m["lon"] - (-83.0)) / ((-81.0) - (-83.0)) norm_elev = m["elevation"] / 100.0 # max ~100m in study area norm_pop = m["population_density"] / 1500.0 feat = [norm_lat, norm_lon, norm_elev, norm_pop] + oh node_features.append(feat) node_labels.append(m["label"]) node_types_list.append(m["node_type"]) node_lats.append(m["lat"]) node_lons.append(m["lon"]) x = torch.tensor(node_features, dtype=torch.float) # [200, 9] # ── 6. Assemble edge tensors ─────────────────────────────────────────── edges = list(G.edges(data=True)) src_list = [e[0] for e in edges] dst_list = [e[1] for e in edges] edge_index = torch.tensor([src_list, dst_list], dtype=torch.long) # Edge features: [flow_rate, distance, wind_correlation] # Normalise: flow 0-100 m³/s, distance 0-200 km, wind corr 0-1 edge_feats = [] for _, _, attr in edges: ef = [ min(attr["flow_rate"], 100.0) / 100.0, min(attr["distance"], 200.0) / 200.0, min(attr["wind_correlation"], 1.0), ] edge_feats.append(ef) edge_attr = torch.tensor(edge_feats, dtype=torch.float) # [E, 3] # ── 7. Create PyG Data object ────────────────────────────────────────── data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr) data.num_nodes = 200 data.node_labels = node_labels data.node_types = node_types_list data.node_lats = node_lats data.node_lons = node_lons data.station_ids = station_ids data.source_ids = source_ids data.junction_ids = junction_ids # ── 8. Save outputs ─────────────────────────────────────────────────── if save_dir is not None: save_path = Path(save_dir) save_path.mkdir(parents=True, exist_ok=True) # PyG graph torch.save(data, save_path / "flow_graph.pt") print(f"Saved PyG graph → {save_path / 'flow_graph.pt'}") # GeoJSON for dashboard geojson = { "type": "FeatureCollection", "features": [] } for nid in range(200): m = node_meta[nid] feat = { "type": "Feature", "geometry": { "type": "Point", "coordinates": [m["lon"], m["lat"]] }, "properties": { "id": nid, "label": m["label"], "node_type": m["node_type"], "elevation": round(m["elevation"], 2), "population_density": round(m["population_density"], 2), "river": m["river"], } } geojson["features"].append(feat) # Add edges as LineStrings for u, v, attr in edges: mu, mv = node_meta[u], node_meta[v] feat = { "type": "Feature", "geometry": { "type": "LineString", "coordinates": [[mu["lon"], mu["lat"]], [mv["lon"], mv["lat"]]] }, "properties": { "from": u, "to": v, "flow_rate": round(attr["flow_rate"], 3), "distance_km": round(attr["distance"], 3), "wind_correlation": round(attr["wind_correlation"], 3), } } geojson["features"].append(feat) geo_path = save_path / "flow_graph.geojson" with open(geo_path, "w") as f: json.dump(geojson, f) print(f"Saved GeoJSON → {geo_path}") # Node metadata CSV import csv csv_path = save_path / "node_metadata.csv" with open(csv_path, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=[ "node_id", "label", "node_type", "lat", "lon", "elevation", "population_density", "river" ]) writer.writeheader() for nid in range(200): m = node_meta[nid] writer.writerow({ "node_id": m["node_id"], "label": m["label"], "node_type": m["node_type"], "lat": round(m["lat"], 6), "lon": round(m["lon"], 6), "elevation": round(m["elevation"], 2), "population_density": round(m["population_density"], 2), "river": m["river"], }) print(f"Saved node metadata → {csv_path}") return data, node_meta, G if __name__ == "__main__": data, node_meta, G = build_graph( save_dir="/home/user/workspace/MicroPlastiNet/data/processed/m3" ) print(f"\nNode feature shape: {data.x.shape}") print(f"Edge index shape: {data.edge_index.shape}") print(f"Edge attr shape: {data.edge_attr.shape}") print(f"Station IDs (first 5): {data.station_ids[:5]}") print(f"Source IDs (first 5): {data.source_ids[:5]}")