166 lines
8.1 KiB
Python
166 lines
8.1 KiB
Python
"""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
|
|
from app.infra.db.postgresql.hydraulic_playback import HydraulicGeometryRepository
|
|
from app.infra.db.timescaledb.repositories.hydraulic_playback import HydraulicPlaybackRepository
|
|
|
|
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("<I", len(encoded)) + encoded + values.tobytes()
|
|
|
|
|
|
def mercator(longitude: float, latitude: float) -> 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:
|
|
nodes, pipes = await HydraulicGeometryRepository.get_rows(conn)
|
|
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 for row in HydraulicPlaybackRepository.iter_node_pressures(conn, times):
|
|
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])
|
|
async for row in HydraulicPlaybackRepository.iter_pipe_metrics(conn, times):
|
|
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)
|