"""Versioned static geometry and compact, indexed hydraulic playback frames.""" from array import array import asyncio from dataclasses import dataclass from datetime import datetime, timedelta import hashlib import json import math import struct import sys from typing import Any from psycopg import AsyncConnection from app.infra.db.postgresql.network_settings import NetworkSettingsRepository WORLD_METERS = 40075016.68557849 MAX_FRAME_COUNT = 8 def packet(header: dict[str, Any], values: array) -> bytes: """Four-byte aligned JSON header followed by little-endian Float32 values.""" encoded = json.dumps(header, separators=(",", ":"), ensure_ascii=False).encode() encoded += b" " * (-len(encoded) % 4) if sys.byteorder != "little": values = array("f", values) values.byteswap() return struct.pack(" tuple[float, float]: latitude = max(-85.05112878, min(85.05112878, latitude)) return ((longitude + 180) / 360, (1 - math.asinh(math.tan(math.radians(latitude))) / math.pi) / 2) @dataclass(frozen=True) class PlaybackGeometry: version: str node_ids: tuple[str, ...] pipe_ids: tuple[str, ...] node_index: dict[str, int] pipe_index: dict[str, int] content: bytes flow_factor: float pressure_factor: float velocity_factor: float def result_factors(units: dict[str, str]) -> tuple[float, float, float]: """Match EPANET outputs to m³/h, metres of water, and m/s.""" flow = units["flow"].upper() flow_factors = {"LPS": 3.6, "LPM": .06, "MLD": 1000 / 24, "CMH": 1, "CMD": 1 / 24, "CFS": .028316846592 * 3600, "GPM": .003785411784 * 60, "MGD": 3785.411784 / 24, "IMGD": 4546.09 / 24, "AFD": 1233.48183754752 / 24} pressure = units["pressure"].upper() pressure_factors = {"M": 1, "METERS": 1, "PSI": .7032496149, "KPA": .1019716213} if flow not in flow_factors or pressure not in pressure_factors: raise ValueError("不支持的模拟结果单位") return (flow_factors[flow], pressure_factors[pressure], .3048 if flow in {"CFS", "GPM", "MGD", "IMGD", "AFD"} else 1) def line_paths(geometry: dict) -> list: if geometry["type"] == "LineString": return [geometry["coordinates"]] if geometry["type"] == "MultiLineString": return geometry["coordinates"] if geometry["type"] == "GeometryCollection": # PostGIS LineMerge returns an empty collection for zero-length lines. return [path for child in geometry["geometries"] for path in line_paths(child)] raise ValueError("管线几何必须为线类型") def build_geometry(nodes: list[dict], pipes: list[dict], units: dict[str, str]) -> PlaybackGeometry: if not nodes or not pipes: raise ValueError("管网没有可发布的节点或管线几何") node_ids = tuple(str(row["id"]) for row in nodes) pipe_ids = tuple(str(row["id"]) for row in pipes) node_positions = [mercator(float(row["longitude"]), float(row["latitude"])) for row in nodes] positions_by_id = dict(zip(node_ids, node_positions)) origin = [(min(p[i] for p in node_positions) + max(p[i] for p in node_positions)) / 2 for i in (0, 1)] values = array("f", (coordinate - origin[axis] for position in node_positions for axis, coordinate in enumerate(position))) segment_count = 0 for index, row in enumerate(pipes): geometry = row["geometry"] if isinstance(geometry, str): geometry = json.loads(geometry) paths = line_paths(geometry) start_node = positions_by_id.get(str(row.get("start_node_id", ""))) if start_node and paths and paths[0] and paths[-1]: first = mercator(*paths[0][0][:2]) last = mercator(*paths[-1][-1][:2]) if math.dist(last, start_node) < math.dist(first, start_node): paths = [list(reversed(path)) for path in reversed(paths)] distance = 0.0 for path in paths: positions = [mercator(float(p[0]), float(p[1])) for p in path] for start, end in zip(positions, positions[1:]): length = math.hypot(end[0] - start[0], end[1] - start[1]) if length <= 0: continue # Instanced segment: endpoints relative to origin, pipe index, # cumulative distance, segment length, physical diameter (mm). values.extend((start[0] - origin[0], start[1] - origin[1], end[0] - origin[0], end[1] - origin[1], index, distance, length, float(row["diameter"] or 0))) distance += length segment_count += 1 factors = result_factors(units) descriptor = json.dumps([node_ids, pipe_ids, origin, units], separators=(",", ":")).encode() version = hashlib.sha256(descriptor + values.tobytes()).hexdigest()[:24] header = {"kind": "geometry", "protocol": 1, "modelVersion": version, "nodeIds": node_ids, "pipeIds": pipe_ids, "origin": origin, "segmentCount": segment_count, "units": {"pressure": "m", "flow": "m³/h", "velocity": "m/s"}} return PlaybackGeometry(version, node_ids, pipe_ids, {key: i for i, key in enumerate(node_ids)}, {key: i for i, key in enumerate(pipe_ids)}, packet(header, values), *factors) async def read_geometry(conn: AsyncConnection) -> PlaybackGeometry: async with conn.cursor() as cursor: await cursor.execute(""" SELECT n.id, ST_X(ST_Transform(g.geom,4326)) AS longitude, ST_Y(ST_Transform(g.geom,4326)) AS latitude FROM gis.network_nodes n JOIN gis.node_geometries g ON g.node_id=n.id WHERE g.geom IS NOT NULL AND NOT ST_IsEmpty(g.geom) ORDER BY n.id """) nodes = await cursor.fetchall() await cursor.execute(""" SELECT id, start_node_id, diameter, ST_AsGeoJSON(ST_Transform(ST_LineMerge(geom),4326))::json AS geometry FROM gis.pipes WHERE geom IS NOT NULL AND NOT ST_IsEmpty(geom) ORDER BY id """) pipes = await cursor.fetchall() units = await NetworkSettingsRepository.get_result_units(conn) return await asyncio.to_thread(build_geometry, nodes, pipes, units) def finite(value: Any, factor: float) -> float: if value is None: return math.nan result = float(value) * factor return result if math.isfinite(result) and abs(result) <= 3.402823466e38 else math.nan async def read_frames(conn: AsyncConnection, geometry: PlaybackGeometry, start: datetime, frame_count: int, step_seconds: int) -> bytes: times = [start + timedelta(seconds=step_seconds * i) for i in range(frame_count)] node_count, pipe_count = len(geometry.node_ids), len(geometry.pipe_ids) stride = node_count + 2 * pipe_count values = array("f", [math.nan]) * (stride * frame_count) time_index = {value: i for i, value in enumerate(times)} valid_counts = {metric: [0] * frame_count for metric in ("pressure", "flow", "velocity")} async with conn.cursor() as cursor: await cursor.execute(""" SELECT time, node_id, pressure FROM realtime.node_results WHERE time = ANY(%s) ORDER BY time, node_id """, (times,)) async for row in cursor: index = geometry.node_index.get(str(row["node_id"])) if index is not None: frame = time_index[row["time"]] values[frame * stride + index] = finite(row["pressure"], geometry.pressure_factor) valid_counts["pressure"][frame] += math.isfinite(values[frame * stride + index]) await cursor.execute(""" SELECT time, link_id, flow, velocity FROM realtime.link_results WHERE time = ANY(%s) ORDER BY time, link_id """, (times,)) async for row in cursor: index = geometry.pipe_index.get(str(row["link_id"])) if index is not None: frame = time_index[row["time"]] offset = frame * stride + node_count values[offset + index] = finite(row["flow"], geometry.flow_factor) values[offset + pipe_count + index] = finite(row["velocity"], geometry.velocity_factor) valid_counts["flow"][frame] += math.isfinite(values[offset + index]) valid_counts["velocity"][frame] += math.isfinite(values[offset + pipe_count + index]) return packet({"kind": "frames", "protocol": 1, "modelVersion": geometry.version, "nodeCount": node_count, "pipeCount": pipe_count, "times": [int(t.timestamp() * 1000) for t in times], "validCounts": valid_counts, "available": [valid_counts["pressure"][i] > 0 and valid_counts["flow"][i] > 0 for i in range(frame_count)]}, values)