Files
TJWaterServer/api/s36_service_area.py
2023-04-30 18:43:44 +08:00

120 lines
4.2 KiB
Python

import sys
import json
from queue import Queue
from .database import *
from .s0_base import get_node_links, get_link_nodes
sys.path.append('..')
from epanet.epanet import run_project
def calculate_service_area(name: str, time_index: int = 0) -> dict[str, Any]:
inp = json.loads(run_project(name))
time_count = len(inp['node_results'][0]['result'])
if time_index >= time_count:
return {}
sources : dict[str, list[str]] = {}
for node_result in inp['node_results']:
result = node_result['result'][time_index]
if result['demand'] < 0:
sources[node_result['node']] = []
link_flows: dict[str, float] = {}
for link_result in inp['link_results']:
result = link_result['result'][time_index]
link_flows[link_result['link']] = float(result['flow'])
# build source to nodes map
for source in sources:
queue = Queue()
queue.put(source)
while not queue.empty():
cursor = queue.get()
if cursor not in sources[source]:
sources[source].append(cursor)
links = get_node_links(name, cursor)
for link in links:
node1, node2 = get_link_nodes(name, link)
if node1 == cursor and link_flows[link] > 0:
queue.put(node2)
elif node2 == cursor and link_flows[link] < 0:
queue.put(node1)
# calculation concentration
concentration_map: dict[str, dict[str, float]] = {}
node_wip: list[str] = []
for source, nodes in sources.items():
for node in nodes:
if node not in concentration_map:
concentration_map[node] = {}
concentration_map[node][source] = 0.0
if node not in node_wip:
node_wip.append(node)
# if only one source, done
for node, concentrations in concentration_map.items():
if len(concentrations) == 1:
node_wip.remove(node)
for key in concentrations.keys():
concentration_map[node][key] = 1.0
node_upstream : dict[str, list[tuple[str, str]]] = {}
for node in node_wip:
if node not in node_upstream:
node_upstream[node] = []
links = get_node_links(name, node)
for link in links:
node1, node2 = get_link_nodes(name, link)
if node2 == node and link_flows[link] > 0:
node_upstream[node].append((link, node1))
elif node1 == node and link_flows[link] < 0:
node_upstream[node].append((link, node2))
while len(node_wip) != 0:
done = []
for node in node_wip:
up_link_nodes = node_upstream[node]
ready = True
for link_node in up_link_nodes:
if link_node in node_wip:
ready = False
break
if ready:
for link_node in up_link_nodes:
for source, concentration in concentration_map[link_node[1]].items():
concentration_map[node][source] += concentration * abs(link_flows[link_node[0]])
# normalize
sum = 0.0
for source, concentration in concentration_map[node].items():
sum += concentration
for source in concentration_map[node].keys():
concentration_map[node][source] /= sum
done.append(node)
for node in done:
node_wip.remove(node)
source_to_main_node: dict[str, list[str]] = {}
for node, value in concentration_map.items():
max_source = ''
max_concentration = 0.0
for s, c in value.items():
if c > max_concentration:
max_concentration = c
max_source = s
if max_source not in source_to_main_node:
source_to_main_node[max_source] = []
source_to_main_node[max_source].append(node)
sas: list[dict[str, Any]] = []
for source, nodes in source_to_main_node.items():
sas.append({ 'source': source, 'nodes': nodes })
return { 'service_areas' : sas, 'concentrations': concentration_map }