TJWaterServerBinary/app/algorithms/burst_localization/leak_signature.py

"""漏损模拟模块。"""

import math
import multiprocessing as mp
import os
import sys

import pandas as pd
import wntr

from app.algorithms._utils import _cleanup_temp_files

_PIPE2LEAKNODE = None
_SIGNATURE_WORKER_DATA = {}


def _make_temp_prefix(tag):
    temp_dir = os.path.abspath(os.path.join("temp", "burst_location"))
    os.makedirs(temp_dir, exist_ok=True)
    safe_tag = str(tag).replace(os.sep, "_").replace(" ", "_")
    return os.path.join(temp_dir, f"{safe_tag}_{os.getpid()}")


def _snapshot_hydraulic_options(wn):
    options = wn.options
    return {
        "demand_model": options.hydraulic.demand_model,
        "duration": float(options.time.duration),
        "hydraulic_timestep": float(options.time.hydraulic_timestep),
        "pattern_timestep": float(options.time.pattern_timestep),
        "report_timestep": float(options.time.report_timestep),
        "required_pressure": float(options.hydraulic.required_pressure),
        "minimum_pressure": float(options.hydraulic.minimum_pressure),
    }


def _apply_hydraulic_options(wn, option_values):
    options = wn.options
    options.hydraulic.demand_model = option_values["demand_model"]
    options.time.duration = option_values["duration"]
    options.time.hydraulic_timestep = option_values["hydraulic_timestep"]
    options.time.pattern_timestep = option_values["pattern_timestep"]
    options.time.report_timestep = option_values["report_timestep"]
    options.hydraulic.required_pressure = option_values["required_pressure"]
    options.hydraulic.minimum_pressure = option_values["minimum_pressure"]


def simple_add_leak(wn, leak_mag, leak_pipe):
    whole_inf = dict()
    leak_pipe_self = wn.get_link(leak_pipe)
    pipe_diameter = leak_pipe_self.diameter
    pipe_length = leak_pipe_self.length
    pipe_roughness = leak_pipe_self.roughness
    pipe_minor_loss = leak_pipe_self.minor_loss
    # pipe_status = leak_pipe_self.status
    # pipe_check_valve = leak_pipe_self.check_valve
    pipe_start_node = leak_pipe_self.start_node_name
    pipe_end_node = leak_pipe_self.end_node_name

    # close the pipe
    # leak_pipe_self.status = 'Closed'
    wn.remove_link(leak_pipe)
    # add the pipe
    add_pipe1 = leak_pipe + "A"
    add_pipe2 = leak_pipe + "B"
    add_node = leak_pipe + "_"

    start_n = wn.get_node(pipe_start_node)
    end_n = wn.get_node(pipe_end_node)
    if start_n.node_type == "Reservoir":
        end_n_elevation = end_n.elevation
        start_n_elevation = end_n_elevation
    elif end_n.node_type == "Reservoir":
        start_n_elevation = start_n.elevation
        end_n_elevation = start_n_elevation
    else:
        end_n_elevation = end_n.elevation
        start_n_elevation = start_n.elevation
    elevation_self = (start_n_elevation + end_n_elevation) / 2
    coordinates_self = (
        (start_n.coordinates[0] + end_n.coordinates[0]) / 2,
        (start_n.coordinates[1] + end_n.coordinates[1]),
    )

    wn.add_junction(
        add_node, base_demand=0, elevation=elevation_self, coordinates=coordinates_self
    )
    leak_node = wn.get_node(add_node)
    wn.add_pipe(
        add_pipe1,
        start_node_name=pipe_start_node,
        end_node_name=add_node,
        length=pipe_length / 2,
        diameter=pipe_diameter,
        roughness=pipe_roughness,
        minor_loss=pipe_minor_loss,
    )
    wn.add_pipe(
        add_pipe2,
        start_node_name=pipe_end_node,
        end_node_name=add_node,
        length=pipe_length / 2,
        diameter=pipe_diameter,
        roughness=pipe_roughness,
        minor_loss=pipe_minor_loss,
    )

    # simulation
    leak_node.add_demand(base=leak_mag, pattern_name="add_leak")

    whole_inf["leak_node_name"] = add_node
    whole_inf["add_pipe1"] = add_pipe1
    whole_inf["add_pipe2"] = add_pipe2
    whole_inf["leak_pipe"] = leak_pipe
    whole_inf["pipe_start_node"] = pipe_start_node
    whole_inf["pipe_end_node"] = pipe_end_node
    whole_inf["pipe_length"] = pipe_length
    whole_inf["pipe_diameter"] = pipe_diameter
    whole_inf["pipe_roughness"] = pipe_roughness
    whole_inf["pipe_minor_loss"] = pipe_diameter
    return wn, whole_inf, add_pipe1


def simple_recover_wn(wn, whole_inf):
    leak_node = wn.get_node(whole_inf["leak_node_name"])

    del leak_node.demand_timeseries_list[-1]
    # update
    wn.remove_link(whole_inf["add_pipe1"])
    wn.remove_link(whole_inf["add_pipe2"])
    wn.remove_node(whole_inf["leak_node_name"])
    # open the pipe
    # leak_pipe_self.status = 'Open'
    wn.add_pipe(
        whole_inf["leak_pipe"],
        start_node_name=whole_inf["pipe_start_node"],
        end_node_name=whole_inf["pipe_end_node"],
        length=whole_inf["pipe_length"],
        diameter=whole_inf["pipe_diameter"],
        roughness=whole_inf["pipe_roughness"],
        minor_loss=whole_inf["pipe_minor_loss"],
    )

    return wn


def disable_all_controls_temporarily(wn):
    """返回(控制名, 控制对象)的列表，之后可用 restore_controls 还原。"""
    removed = []
    # WNTR 的控制都在 wn.control_name_list / wn.get_control / wn.remove_control
    for cname in list(wn.control_name_list):
        ctrl = wn.get_control(cname)
        removed.append((cname, ctrl))
        wn.remove_control(cname)
    return removed


def restore_controls(wn, removed):
    """把先前禁用的控制全部加回去。"""
    for cname, ctrl in removed:
        wn.add_control(cname, ctrl)


def set_pipe2leaknode_mapping(mapping):
    global _PIPE2LEAKNODE
    _PIPE2LEAKNODE = mapping


def _get_or_create_leak_demand_ts(leak_node):
    """
    返回：泄漏专用 demand 的下标 idx。
    若不存在，以 category='leak' 新建一条 base=0.0 的 demand。
    """
    # 先尝试找到已有的 'leak' 分类
    for i, ts in enumerate(leak_node.demand_timeseries_list):
        # WNTR 的 Demand object 存在 category 属性
        if getattr(ts, "category", None) == "leak":
            return i
    # 没有则新建（base=0.0，后续临时改 base_value）
    leak_node.add_demand(base=0.0, pattern_name=None, category="leak")
    return len(leak_node.demand_timeseries_list) - 1


def ensure_mid_node(wn, leak_pipe):
    add_pipe1 = f"{leak_pipe}A"
    add_pipe2 = f"{leak_pipe}B"
    add_node = f"{leak_pipe}__mid"

    if add_node in wn.node_name_list:
        return add_node

    if leak_pipe in wn.link_name_list:
        leak_pipe_self = wn.get_link(leak_pipe)
        pipe_diameter = leak_pipe_self.diameter
        pipe_length = leak_pipe_self.length
        pipe_roughness = leak_pipe_self.roughness
        pipe_minor_loss = leak_pipe_self.minor_loss
        pipe_start_node = leak_pipe_self.start_node_name
        pipe_end_node = leak_pipe_self.end_node_name

        start_n = wn.get_node(pipe_start_node)
        end_n = wn.get_node(pipe_end_node)
        if start_n.node_type == "Reservoir":
            end_elev = end_n.elevation
            start_elev = end_elev
        elif end_n.node_type == "Reservoir":
            start_elev = start_n.elevation
            end_elev = start_elev
        else:
            end_elev = end_n.elevation
            start_elev = start_n.elevation
        elev_mid = (start_elev + end_elev) / 2.0
        x_mid = (start_n.coordinates[0] + end_n.coordinates[0]) / 2.0
        y_mid = (start_n.coordinates[1] + end_n.coordinates[1]) / 2.0

        wn.remove_link(leak_pipe)
        wn.add_junction(
            add_node, base_demand=0.0, elevation=elev_mid, coordinates=(x_mid, y_mid)
        )
        wn.add_pipe(
            add_pipe1,
            start_node_name=pipe_start_node,
            end_node_name=add_node,
            length=pipe_length / 2.0,
            diameter=pipe_diameter,
            roughness=pipe_roughness,
            minor_loss=pipe_minor_loss,
        )
        wn.add_pipe(
            add_pipe2,
            start_node_name=add_node,
            end_node_name=pipe_end_node,
            length=pipe_length / 2.0,
            diameter=pipe_diameter,
            roughness=pipe_roughness,
            minor_loss=pipe_minor_loss,
        )
        return add_node

    # 若 A/B 已存在但中点不在，建议确认网络一致性
    raise KeyError(f"Cannot ensure mid node for pipe '{leak_pipe}'.")


def leak_simulation_pipe_dd_multi_pf(
    wn, leak_mag, leak_pipe, sensor_name, file_prefix=None
):
    """
    优化版：
    - 不再 remove/add link/node
    - 直接在预插入的中点泄漏节点上设置 base_demand = leak_mag；仿真后设回 0
    """
    wn.options.hydraulic.demand_model = "DD"

    # 确保中点节点存在
    leak_node_name = ensure_mid_node(wn, leak_pipe)
    leak_node = wn.get_node(leak_node_name)

    # 拿到泄漏专用的 demand time-series 下标
    leak_idx = _get_or_create_leak_demand_ts(leak_node)
    ts_obj = leak_node.demand_timeseries_list[leak_idx]

    # 记录原值（通常是 0.0）
    orig_base = ts_obj.base_value
    try:
        # 打开泄漏：只改 base_value，不碰 base_demand（只读）
        ts_obj.base_value = float(leak_mag)

        # 仿真
        sim = wntr.sim.EpanetSimulator(wn)
        if file_prefix is None:
            results = sim.run_sim()
        else:
            results = sim.run_sim(file_prefix=file_prefix)

        # 输出（保持列顺序）
        pressure_output = results.node["pressure"].loc[:, sensor_name]
        # flow_output = results.link['flowrate'].loc[:, sensor_f_name]
        return wn, pressure_output
    finally:
        # 关闭泄漏：还原 base_value
        ts_obj.base_value = orig_base
        if file_prefix is not None:
            _cleanup_temp_files(file_prefix)


def prepare_leak_infrastructure(wn, candidate_pipes):
    """
    把 candidate_pipes 每条管段切成两段，并在中点插入一个泄漏节点（base_demand=0）。
    返回一个映射：pipe_id -> leak_node_name
    注意：只做一次；后续仿真通过在该节点设置 base_demand 实现“打开泄漏”，结束后恢复为 0。
    """
    pipe2leaknode = {}
    for leak_pipe in candidate_pipes:
        if leak_pipe in pipe2leaknode:
            continue

        leak_pipe_self = wn.get_link(leak_pipe)
        pipe_diameter = leak_pipe_self.diameter
        pipe_length = leak_pipe_self.length
        pipe_roughness = leak_pipe_self.roughness
        pipe_minor_loss = leak_pipe_self.minor_loss
        pipe_start_node = leak_pipe_self.start_node_name
        pipe_end_node = leak_pipe_self.end_node_name

        # 计算中点高程/坐标（与原逻辑一致）
        start_n = wn.get_node(pipe_start_node)
        end_n = wn.get_node(pipe_end_node)
        if start_n.node_type == "Reservoir":
            end_elev = end_n.elevation
            start_elev = end_elev
        elif end_n.node_type == "Reservoir":
            start_elev = start_n.elevation
            end_elev = start_elev
        else:
            end_elev = end_n.elevation
            start_elev = start_n.elevation
        elev_mid = (start_elev + end_elev) / 2.0
        x_mid = (start_n.coordinates[0] + end_n.coordinates[0]) / 2.0
        y_mid = (start_n.coordinates[1] + end_n.coordinates[1]) / 2.0

        # 先删原管，再加中点与两段半长管（只做一次）
        wn.remove_link(leak_pipe)

        add_pipe1 = f"{leak_pipe}A"
        add_pipe2 = f"{leak_pipe}B"
        add_node = f"{leak_pipe}__mid"  # 唯一命名，后面直接用它当泄漏节点

        wn.add_junction(
            add_node, base_demand=0.0, elevation=elev_mid, coordinates=(x_mid, y_mid)
        )
        wn.add_pipe(
            add_pipe1,
            start_node_name=pipe_start_node,
            end_node_name=add_node,
            length=pipe_length / 2.0,
            diameter=pipe_diameter,
            roughness=pipe_roughness,
            minor_loss=pipe_minor_loss,
        )
        wn.add_pipe(
            add_pipe2,
            start_node_name=add_node,
            end_node_name=pipe_end_node,
            length=pipe_length / 2.0,
            diameter=pipe_diameter,
            roughness=pipe_roughness,
            minor_loss=pipe_minor_loss,
        )

        pipe2leaknode[leak_pipe] = add_node

    return pipe2leaknode


def normal_simulation_pf(
    wn, drive_mode, sensor_name, sensor_f_name, inp_time, require_p, minimum_p
):
    # inp_time = 0
    if drive_mode == "PDD":  # 需水量根据节点压力动态调整
        wn.options.hydraulic.demand_model = "PDD"
        wn.options.hydraulic.required_pressure = require_p
        wn.options.hydraulic.minimum_pressure = minimum_p
    elif drive_mode == "DD":  # 需水量固定，与压力无关
        wn.options.hydraulic.demand_model = "DD"
    sim = wntr.sim.EpanetSimulator(wn)
    results = sim.run_sim()
    pressure_all = results.node["pressure"][sensor_name]
    pressure = pressure_all.iloc[inp_time]
    demand_all = results.node["demand"]
    demand = demand_all.iloc[inp_time]
    sum_demand = cal_sum_demand(demand)
    flow_all = results.link["flowrate"][sensor_f_name]
    flow = flow_all.iloc[inp_time]

    top_sensor = pressure.idxmin()
    basic_p = results.node["pressure"]
    basic_p = basic_p.iloc[inp_time]
    return pressure, flow, basic_p, top_sensor, sum_demand


def normal_simulation_multi_pf(
    wn, drive_mode, sensor_name, sensor_f_name, require_p, minimum_p
):
    # inp_time = 0
    if drive_mode == "PDD":
        wn.options.hydraulic.demand_model = "PDD"
        wn.options.hydraulic.required_pressure = require_p
        wn.options.hydraulic.minimum_pressure = minimum_p
    elif drive_mode == "DD":
        wn.options.hydraulic.demand_model = "DD"
    sim = wntr.sim.EpanetSimulator(wn)
    results = sim.run_sim()
    pressure_all = results.node["pressure"][sensor_name]
    pressure = pressure_all
    demand_all = results.node["demand"]
    demand = demand_all
    flow = results.link["flowrate"][sensor_f_name]
    sum_demand = pd.Series(dtype=object)
    for i in range(len(demand.index)):
        sum_demand[str(demand.index[i])] = cal_sum_demand(demand.iloc[i])
    if type(pressure) == pd.core.series.Series:
        top_sensor = pressure.idxmin()
    else:
        mean_pressure = pressure.mean()
        top_sensor = mean_pressure.idxmin()
    basic_p = results.node["pressure"]
    return pressure, flow, basic_p, top_sensor, sum_demand


def simple_simulation_pf(wn, sensor_name, sensor_f_name, leak_pipe, add_pipe1):
    sim = wntr.sim.EpanetSimulator(wn)
    results = sim.run_sim()
    pressure_all = results.node["pressure"][sensor_name]
    if len(leak_pipe) > 0 and leak_pipe in sensor_f_name:

        f_sensor_name = [add_pipe1 if i == leak_pipe else i for i in sensor_f_name]
        flow_all = results.link["flowrate"][f_sensor_name]
        flow_all.columns = sensor_f_name
    else:
        flow_all = results.link["flowrate"][sensor_f_name]
    return pressure_all, flow_all


def cal_sum_demand(demand):
    sum_demand = 0
    for i in range(len(demand)):
        if demand.iloc[i] > 0:
            sum_demand += demand.iloc[i]
    return sum_demand


def cal_signature_pipe_multi_pf(
    wn,
    leak_mag,
    candidate_center,
    timestep_list,
    sensor_name,
    n_workers=1,
    wn_inp_path=None,
):
    candidate_center_num = len(candidate_center)
    pressure_leak = pd.DataFrame(
        index=pd.MultiIndex.from_product([candidate_center, timestep_list]),
        columns=sensor_name,
    )
    # flow_leak = pd.DataFrame(index=pd.MultiIndex.from_product([candidate_center, timestep_list]),
    #                          columns=sensor_f_name)
    pressure_leak = pressure_leak.sort_index()
    # flow_leak = flow_leak.sort_index()
    can_parallel = (
        n_workers > 1
        and candidate_center_num > 1
        and wn_inp_path is not None
        and len(str(wn_inp_path)) > 0
    )
    if can_parallel:
        option_values = _snapshot_hydraulic_options(wn)
        worker_count = min(n_workers, candidate_center_num)
        start_methods = mp.get_all_start_methods()
        context_name = "spawn" if "spawn" in start_methods else start_methods[0]
        with mp.get_context(context_name).Pool(
            processes=worker_count,
            initializer=_signature_worker_init,
            initargs=(
                str(wn_inp_path),
                float(leak_mag),
                list(sensor_name),
                option_values,
                list(candidate_center),
            ),
        ) as pool:
            for i, (center_name, pressure_array) in enumerate(
                pool.imap(_signature_worker_run_center, candidate_center)
            ):
                pressure_leak.loc[(center_name, slice(None)), :] = pressure_array
                sys.stdout.write("\r" + "已经完成计算" + str(i + 1) + "个特征中心")
    else:
        # Pre-insert all mid-nodes so every simulation sees the same topology
        for center in candidate_center:
            ensure_mid_node(wn, center)
        for i in range(candidate_center_num):
            temp_prefix = _make_temp_prefix(f"sig_{i}")
            wn, pressure_output = leak_simulation_pipe_dd_multi_pf(
                wn,
                leak_mag,
                candidate_center[i],
                sensor_name,
                file_prefix=temp_prefix,
            )
            # leak_or_not_list.append(leak_or_not)
            pressure_leak.loc[(candidate_center[i], slice(None)), :] = (
                pressure_output.to_numpy()
            )
            # flow_leak.loc[candidate_center[i]].loc[:, :] = flow_output
            sys.stdout.write("\r" + "已经完成计算" + str(i + 1) + "个特征中心")
    return pressure_leak, candidate_center


def _signature_worker_init(
    inp_path, leak_mag, sensor_name, option_values, candidate_centers=None
):
    global _SIGNATURE_WORKER_DATA
    wn = wntr.network.WaterNetworkModel(inp_path)
    _apply_hydraulic_options(wn, option_values)
    # Pre-insert ALL mid-nodes so every simulation runs on the same topology,
    # regardless of which worker handles which task.
    if candidate_centers is not None:
        for center in candidate_centers:
            ensure_mid_node(wn, center)
    _SIGNATURE_WORKER_DATA = {
        "wn": wn,
        "leak_mag": leak_mag,
        "sensor_name": sensor_name,
    }


def _signature_worker_run_center(center_name):
    data = _SIGNATURE_WORKER_DATA
    temp_prefix = _make_temp_prefix(f"sig_worker_{center_name}")
    _, pressure_output = leak_simulation_pipe_dd_multi_pf(
        data["wn"],
        data["leak_mag"],
        center_name,
        data["sensor_name"],
        file_prefix=temp_prefix,
    )
    return center_name, pressure_output.to_numpy()


def pick_pipe(all_pipes, pipe_diameter, limited_diameter):
    candidate_pipe = []
    for each_pipe in all_pipes:
        if pipe_diameter[each_pipe] >= limited_diameter:
            candidate_pipe.append(each_pipe)
    return candidate_pipe


def cal_possible_pipe(leak_flow, all_pipe, pipe_diameter):
    basic_pressure = 10  # 基础压力
    discharge_coeff = 0.6  # 经验系数
    break_area_ratio = 1  # 爆管面积比 0.5  1.25
    break_area = leak_flow / (
        discharge_coeff * math.sqrt(2 * basic_pressure * 9.81)
    )  # 爆管面积 m3/h
    """break_area_diameter = math.sqrt(4 * break_area / math.pi)
    min_diameter = (math.ceil(1000 * break_area_diameter / break_area_ratio)) / 1000"""
    break_area_diameter = math.sqrt(
        4 * break_area / math.pi / break_area_ratio
    )  # 爆管直径
    min_diameter = (math.ceil(1000 * break_area_diameter)) / 1000  # 向上取整
    new_all_pipe = pick_pipe(all_pipe, pipe_diameter, min_diameter)
    return new_all_pipe, min_diameter


def extract_links(data, link_types, direction):
    return [
        link
        for res_data in data.values()
        for link_type in link_types
        for link in res_data[link_type][direction]
    ]