Support calculate distribution

api/__init__.py

@@ -143,6 +143,7 @@ from .s32_region_util import get_nodes_in_boundary, get_nodes_in_region, calcula
 from .s33_region import get_region_schema, get_region, set_region, add_region, delete_region
 
 from .s34_water_distribution import DISTRIBUTION_TYPE_ADD, DISTRIBUTION_TYPE_OVERRIDE
+from .s34_water_distribution import calculate_demand_to_nodes, calculate_demand_to_region
 from .s34_water_distribution import distribute_demand_to_nodes, distribute_demand_to_region
 
 from .s35_district_metering_area import PARTITION_TYPE_RB, PARTITION_TYPE_KWAY

api/s34_water_distribution.py

@@ -4,9 +4,45 @@ from .s9_demands import get_demand
 from .s32_region_util import Topology, get_nodes_in_region
 from .batch_exe import execute_batch_command
 
+
 DISTRIBUTION_TYPE_ADD = 'ADD'
 DISTRIBUTION_TYPE_OVERRIDE = 'OVERRIDE'
 
+
+def calculate_demand_to_nodes(name: str, demand: float, nodes: list[str]) -> dict[str, float]:
+    if len(nodes) == 0 or demand == 0.0:
+        return {}
+
+    topology = Topology(name, nodes)
+    t_nodes = topology.nodes()
+    t_links = topology.links()
+
+    length_sum = 0.0
+    for value in t_links.values():
+        length_sum += abs(value['length'])
+
+    if length_sum <= 0.0:
+        return {}
+
+    demand_per_length = demand / length_sum
+
+    result: dict[str, float] = {}
+    for node, value in t_nodes.items():
+        if not is_junction(name, node):
+            continue
+        demand_per_node = 0.0
+        for link in value['links']:
+            demand_per_node += abs(t_links[link]['length']) * demand_per_length * 0.5
+        result[node] = demand_per_node
+
+    return result
+
+
+def calculate_demand_to_region(name: str, demand: float, region: str) -> dict[str, float]:
+    nodes = get_nodes_in_region(name, region)
+    return calculate_demand_to_nodes(name, demand, nodes)
+
+
 def distribute_demand_to_nodes(name: str, demand: float, nodes: list[str], type: str = DISTRIBUTION_TYPE_ADD) -> ChangeSet:
     if len(nodes) == 0 or demand == 0.0:
         return ChangeSet()