#!/usr/bin/env python3 """Build a private multi-region, multi-area and multi-time design from screening evidence.""" from __future__ import annotations import argparse import asyncio import json import math import os import sys from datetime import datetime, timezone from pathlib import Path from uuid import uuid4 def require_isolatable_pipe(topology: list[str], pipe_id: str, sources: list[str], closed: list[str]) -> dict: from app.algorithms.valve_isolation.topology_search import valve_isolation_analysis result = valve_isolation_analysis(topology, pipe_id, accident_element_type="link", water_source_nodes=sources, closed_links=closed) if not result["isolatable"] or not result["source_check_performed"] or not result["must_close_valves"]: raise ValueError("selected region cannot be isolated from water sources") return result def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--screening", type=Path, required=True) parser.add_argument("--additional-screening", type=Path, action="append", default=[]) parser.add_argument("--pipe", action="append", required=True, help="New pipe ID; the existing primary is retained") parser.add_argument("--allow-weak", action="store_true", help="Include weak-signal physical bursts without treating them as no-burst controls") parser.add_argument("--case-start", type=int, default=9) parser.add_argument("--output-root", type=Path, default=Path("data/benchmark-regional")) args = parser.parse_args() sys.path.insert(0, str(Path(__file__).resolve().parents[1])) os.umask(0o077) from scripts.benchmark.calibration import save_json, digest from scripts.benchmark.simulation_design import BurstSimulationDesign, load_simulation_design from scripts.check_burst_simulation_design import readonly, resolve_project, validate_database, validate_tool_inputs original = load_simulation_design(Path("resources/benchmarks/tjwater2-v2-burst-simulation.json")) screening = json.loads(args.screening.read_text()) candidates = {row["pipe_id"]: row for row in screening["results"]} for path in args.additional_screening: extra = json.loads(path.read_text()) if extra["model_sha256"] != screening["model_sha256"]: raise ValueError("screening model fingerprints do not match") candidates.update({row["pipe_id"]: row for row in extra["results"]}) requested = list(dict.fromkeys(args.pipe)) if not 1 <= len(requested) <= 3 or any(pipe not in candidates for pipe in requested): raise ValueError("select one to three screened pipes") for pipe in requested: if not candidates[pipe].get("possible_candidate") and not args.allow_weak: raise ValueError("selected pipe has not passed single-frame screening") if candidates[pipe].get("status") != "completed": raise ValueError("selected pipe screening simulation did not complete") folder = (args.output_root / (datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + "-" + uuid4().hex[:8])).resolve() folder.mkdir(parents=True, mode=0o700) biz, _ = asyncio.run(resolve_project(original)) with readonly(biz) as connection: links = connection.execute( "SELECT l.id,l.start_node_id,l.end_node_id,l.link_type," "COALESCE(s.status::text,p.status::text,'OPEN') status FROM network.links l " "LEFT JOIN network.pipes p ON p.link_id=l.id " "LEFT JOIN network.link_initial_settings s ON s.link_id=l.id" ).fetchall() sources = [row["id"] for row in connection.execute( "SELECT id FROM network.nodes WHERE node_type IN ('reservoir','tank')" )] devices = connection.execute( "SELECT device_id,device_type,node_id,link_id,measurement_unit FROM asset.scada_devices " "WHERE device_type IN ('pressure','pipe_flow') AND device_id<>'547063'" ).fetchall() positions = {row["id"]: row for row in connection.execute( "SELECT l.id,l.start_node_id,l.end_node_id,p.diameter,p.length," "ST_X(ST_Transform(g.geom,4326)) longitude,ST_Y(ST_Transform(g.geom,4326)) latitude " "FROM network.links l JOIN network.pipes p ON p.link_id=l.id " "JOIN gis.node_geometries g ON g.node_id=l.end_node_id WHERE l.id=ANY(%s)", ([original.pipes[0].pipe_id, *requested],), )} pipe_geometries = {row["id"]: json.loads(row["geometry"]) for row in connection.execute( "SELECT id,ST_AsGeoJSON(ST_Transform(geom,4326)) geometry FROM gis.pipes WHERE id=ANY(%s)", ([original.pipes[0].pipe_id, *requested],), )} topology = [f"{row['id']}:{row['link_type']}:{row['start_node_id']}:{row['end_node_id']}" for row in links] closed = [row["id"] for row in links if str(row["status"]).upper() == "CLOSED"] by_device = {row["device_id"]: row for row in devices} pipes = [original.pipes[0].model_dump(mode="json")] for pipe_id in requested: candidate = candidates[pipe_id] position = positions[pipe_id] pressure_ids = [key for key in candidate["pressure_drop_percent"] if key in by_device and by_device[key]["device_type"] == "pressure"][:5] # Keep the plant outlets plus the two most responsive other flow meters. flow_ids = ["542017", "486033", "542121"] for key, _ in sorted(candidate["flow_change_percent"].items(), key=lambda pair: pair[1], reverse=True): if key in by_device and key not in flow_ids: flow_ids.append(key) if len(flow_ids) >= 5: break pressure = [{"device_id": key, "node_id": by_device[key]["node_id"]} for key in pressure_ids] flow = [{"device_id": key, "link_id": by_device[key]["link_id"]} for key in flow_ids] area = math.pi * (position["diameter"] / 10) ** 2 / 4 pipes.append({ "role": "secondary", "pipe_id": pipe_id, "diameter_mm": position["diameter"], "length_m": position["length"], "start_node_id": position["start_node_id"], "emitter_node_id": position["end_node_id"], "longitude": position["longitude"], "latitude": position["latitude"], "candidate_orifice_areas_cm2": {"light": round(area / 32, 6), "medium": round(area / 16, 6), "severe": round(area / 8, 6)}, "pressure_bindings": pressure, "flow_bindings": flow, "boundary_valve_ids": [], "output_node_ids": sorted({position["start_node_id"], position["end_node_id"], *[item["node_id"] for item in pressure]}), "output_link_ids": [], }) isolations = {} for pipe in pipes: pipe_id = pipe["pipe_id"] result = require_isolatable_pipe(topology, pipe_id, sources, closed) pipe["boundary_valve_ids"] = result["must_close_valves"] pipe["output_link_ids"] = sorted({pipe_id, *result["must_close_valves"], *[item["link_id"] for item in pipe["flow_bindings"]]}) pipe.update(longitude=positions[pipe_id]["longitude"], latitude=positions[pipe_id]["latitude"]) isolations[pipe_id] = result payload = original.model_dump(mode="json") times = [f"2026-10-08T{hour:02d}:00:00+08:00" for hour in (2, 8, 12, 20)] payload.update(pipes=pipes, duration_seconds=18000, event_times=times) payload["selection_rule"]["expected_sample_count"] = 21 if not 1 <= args.case_start <= 84: raise ValueError("case start must leave room for sixteen distinct Bxx cases") payload["cases"] = [{"case_id": f"B{args.case_start + p * len(times) + t:02d}", "pipe_id": pipe["pipe_id"], "occurred_at": at} for p, pipe in enumerate(pipes) for t, at in enumerate(times)] design = BurstSimulationDesign.model_validate(payload) validate_tool_inputs(design) with readonly(biz) as connection: validate_database(connection, design, template=False) path = folder / "design.json" save_json(path, design.model_dump(mode="json")) save_json(folder / "selection.json", {"source_effect": "read-only", "screening_sha256": digest(args.screening), "additional_screening_sha256": [digest(path) for path in args.additional_screening], "weak_signal_cases_included": args.allow_weak, "assessment_note": "a missed pressure rule does not change physical burst ground truth", "model_sha256": screening["model_sha256"], "design_sha256": digest(path), "isolations": isolations, "coordinates": "WGS84 burst endpoint, not pipe midpoint"}) valve_ids = sorted({valve for pipe in pipes for valve in pipe["boundary_valve_ids"]}) with readonly(biz) as connection: valve_geometries = {row["id"]: json.loads(row["geometry"]) for row in connection.execute( "SELECT l.id,ST_AsGeoJSON(ST_Transform(ST_LineInterpolatePoint(ST_MakeLine(a.geom,b.geom),0.5),4326)) geometry " "FROM network.links l JOIN gis.node_geometries a ON a.node_id=l.start_node_id " "JOIN gis.node_geometries b ON b.node_id=l.end_node_id WHERE l.id=ANY(%s)", (valve_ids,), )} features = [] for index, pipe in enumerate(pipes, 1): properties = {"region": f"R{index}", "pipe_id": pipe["pipe_id"], "diameter_mm": pipe["diameter_mm"]} features.append({"type": "Feature", "properties": {**properties, "kind": "burst"}, "geometry": {"type": "Point", "coordinates": [pipe["longitude"], pipe["latitude"]]}}) features.append({"type": "Feature", "properties": {**properties, "kind": "pipe"}, "geometry": pipe_geometries[pipe["pipe_id"]]}) for valve in pipe["boundary_valve_ids"]: features.append({"type": "Feature", "properties": {**properties, "kind": "valve", "valve_id": valve}, "geometry": valve_geometries[valve]}) save_json(folder / "locations.geojson", {"type": "FeatureCollection", "features": features}) print(f"regional design ready: {len(pipes)} regions x 3 areas x {len(times)} times = {len(pipes) * 3 * len(times)} physical cases; {path}") return 0 if __name__ == "__main__": raise SystemExit(main())