DualFlow/models/uf_rl/rl_model/DQN/dqn_model/dqn_statebuilder.py

from typing import Dict, Optional
from dataclasses import replace

import numpy as np
import pandas as pd
from sklearn.metrics import r2_score

from env.env_params import UFState

from uf_data_process.load import UFConfigLoader
from uf_data_process.label import UFEventClassifier, PostBackwashInletMarker
from uf_data_process.filter import ConstantFlowFilter,EventQualityFilter, InletSegmentFilter,FlowOutlierFilter
from uf_data_process.calculate import UFResistanceCalculator, UFResistanceAnalyzer
from uf_data_process.fit import ShortTermCycleFoulingFitter, LongTermFoulingFitter

class DQNStateBuilder:
    """
    在 DQN 决策前构建状态的工具类

    相关数据：
        * CSV1 = 上一完整化学周期
        * CSV2 = 新周期初始进水段
        * CSV3 = 新周期预测进水段
    """

    def __init__(self, config_path: str):
        """
        Parameters
        ----------
        config_path : str
            uf_analyze_config.yaml 路径
        """
        self.cfg = UFConfigLoader(config_path)
        uf_cfg = self.cfg.uf
        params = self.cfg.params


        self.units = uf_cfg.get("units", ["UF1", "UF2", "UF3", "UF4"])
        self.stable_inlet_code = uf_cfg.get("stable_inlet_code", [24.0, 26.0])

        column_formats = uf_cfg.get("column_formats", {})
        self.ctrl_format = column_formats.get("ctrl_col", "C.M.{unit}_DB@word_control")
        self.flow_format = column_formats.get("flow_col", "C.M.{unit}_FT_JS@out")
        self.tmp_format = column_formats.get("tmp_col", "C.M.{unit}_DB@press_PV")
        self.per_format = column_formats.get("per_col", "{unit}Per")
        self.temp_col = column_formats.get("temp_col", "C.M.RO_TT_ZJS@out")


        # 过滤器
        self.min_points = params.get("min_points", 20)
        self.initial_points = params.get("initial_points", 10)

        self.quality_filter = EventQualityFilter(min_points=self.min_points)
        self.flow_filter = FlowOutlierFilter(n_sigma=3)
        self.initial_label = PostBackwashInletMarker(n_points=self.initial_points)

        # 阻力计算器
        self.res_calc = UFResistanceCalculator(self.units, area_m2=uf_cfg["area_m2"], scale_factor=params.get("scale_factor", 1e10))
        self.segment_head_n = params.get("segment_head_n", 10)
        self.segment_tail_n = params.get("segment_tail_n", 10)

    # ======================================================================
    # 对外主接口
    # ======================================================================
    def build_from_csv_pair(
            self,
            unit_name,
            uf_state_default,
            state_bounds,
            prev_cycle_csv: str,
            init_cycle_csv: str,
            predict_cycle_csv: Optional[str] = None,
    ) -> UFState:
        """
        使用【上一完整化学周期 CSV】+【当前周期初始 CSV】+ 【当前周期预测 CSV】构建 UFState
        """

        df_prev = pd.read_csv(prev_cycle_csv)
        df_init = pd.read_csv(init_cycle_csv)

        # predict_csv 允许为空
        df_predict = None
        if predict_cycle_csv is not None:
            df_predict = pd.read_csv(predict_cycle_csv)

        # 分别处理两个 CSV
        prev_features = self._analyze_previous_cycle_csv(
            df_prev,
            unit_name,
            uf_state_default,
            state_bounds
        )

        init_features = self._analyze_init_cycle_csv(
            df_init,
            unit_name,
            uf_state_default,
            state_bounds
        )

        # 化学清洗去除阻力（上一周期末 - 当前初始）
        ceb_removal = max(
            prev_features["R_end"] - init_features["R_start"],
            0.0
        )

        # 默认直接使用上一周期 nuK
        corrected_nuk = prev_features["nuK"]

        # 如果提供了 predict_csv，则进行在线修正
        if df_predict is not None:
            corrected_nuk = self._correct_nuk_with_predict(
                df_predict=df_predict,
                unit_name=unit_name,
                R_start=init_features["R_start"],
                q_mean=init_features["q_mean"],
                temp_celsius=init_features["temp_mean"],
                base_nuk=prev_features["nuK"],
                uf_state_default=uf_state_default,
            )

        # 构建 UFState
        current_state = replace(
            uf_state_default,
            TMP=init_features["tmp_mean"],
            q_UF=init_features["q_mean"],
            temp=init_features["temp_mean"],
            R=init_features["R_start"],
            nuK=corrected_nuk,
            slope=prev_features["slope"],
            power=prev_features["power"],
            ceb_removal=ceb_removal,
        )

        return current_state

    # ======================================================================
    # 上一完整化学周期分析
    # ======================================================================

    def _analyze_previous_cycle_csv(
        self,
        df,
        unit_name,
        uf_state_default,
        state_bounds
    ) -> Dict[str, float]:
        """
        上一完整化学周期分析逻辑

        步骤：
        1. 事件标注
        2. 进水段过滤（质量过滤）
        3. 膜阻力计算
        4. 提取周期末稳定阻力
        5. 拟合 nuK
        6. 拟合长期不可逆污染（slope / power）
        """

        ctrl_col = self.ctrl_format.format(unit=unit_name)
        flow_col = self.flow_format.format(unit=unit_name)
        tmp_col = self.tmp_format.format(unit=unit_name)

        # 事件标注
        event_clf = UFEventClassifier(unit_name, self.cfg.uf["inlet_codes"],
                                      self.cfg.uf["physical_bw_code"], self.cfg.uf["chemical_bw_code"],
                                      ctrl_col)
        df_unit = event_clf.classify(df)  # 产生 event_type 列
        df_unit_mark = self.initial_label.mark(df_unit)  # 标记反冲洗事件后的前 N 个进水点
        seg_df = event_clf.segment(df_unit_mark)  # 根据 event_type 列编号事件段落

        # 对 seg_df 进行按 segment 分组后逐段过滤：
        const_flow_filter = ConstantFlowFilter(flow_col=flow_col, repeat_len=20)
        segments = const_flow_filter.filter(seg_df)  # 去除出现网络错误的进水段
        segments = self.quality_filter.filter(segments)  # 去除时间过短的进水段
        # 提取稳定进水段
        stable_extractor = InletSegmentFilter(ctrl_col, stable_codes=self.stable_inlet_code, min_points=self.min_points)
        stable_segments = stable_extractor.extract(segments)  # 提取稳定进水数据

        if len(stable_segments) == 0:
            raise ValueError("上一周期无有效稳定进水段，无法构建状态，请使用run_dqn_decide.py")

        # 膜阻力计算
        stable_segments = self.res_calc.calculate_for_segments(
            stable_segments,
            temp_col=self.temp_col,
            flow_col=flow_col,
            tmp_col=tmp_col,
        )

        # -------- 膜阻力统计 --------
        res_col = f"{unit_name}_R_scaled"
        ura = UFResistanceAnalyzer(
            resistance_col=res_col,
            head_n=self.segment_head_n,
            tail_n=self.segment_tail_n
        )
        stable_segments = ura.analyze_segments(stable_segments)
        df_all = stable_segments[-1]
        R_end = df_all["R_scaled_end"].iloc[0]

        # ===== 确保 time 为 datetime =====
        for i, seg in enumerate(stable_segments):
            if not pd.api.types.is_datetime64_any_dtype(seg["time"]):
                seg = seg.copy()
                seg["time"] = pd.to_datetime(seg["time"], errors="coerce")
                seg = seg.dropna(subset=["time"])
                stable_segments[i] = seg

        # -------- 5️⃣ 短期污染拟合（nuK）--------
        st_fitter = ShortTermCycleFoulingFitter(unit_name)
        nuK, st_r2 = st_fitter.fit_cycle(stable_segments)

        if (
                pd.isna(nuK)
                or pd.isna(st_r2)
                or not np.isfinite(nuK)
                or st_r2 < 0.4
        ):
            nuK = uf_state_default.nuK

        nuK = float(
            np.clip(
                nuK,
                state_bounds.nuK_min,
                state_bounds.nuK_max,
            )
        )

        # -------- 6️⃣ 长期不可逆污染拟合 --------
        lt_fitter = LongTermFoulingFitter(unit_name)
        slope, power, lt_r2 = lt_fitter.fit_cycle(stable_segments)
        if (
                pd.isna(slope)
                or pd.isna(power)
                or pd.isna(lt_r2)
                or not np.isfinite(slope)
                or not np.isfinite(power)
                or lt_r2 < 0.4
        ):
            slope = uf_state_default.slope
            power = uf_state_default.power

        return {
            "R_end": R_end,
            "nuK": float(nuK),
            "slope": float(slope),
            "power": float(power),
        }

    # ======================================================================
    # 当前周期初始进水段分析
    # ======================================================================

    def _analyze_init_cycle_csv(
        self,
        df,
        unit_name,
        uf_state_default,
        state_bounds
    ) -> Dict[str, float]:
        """
        当前周期初始进水段分析
        """

        flow_col = self.flow_format.format(unit=unit_name)
        tmp_col = self.tmp_format.format(unit=unit_name)
        temp_col = self.temp_col
        res_col = f"{unit_name}_R_scaled"

        segments = [df]
        segments = self.res_calc.calculate_for_segments(
            segments,
            temp_col=self.temp_col,
            flow_col=flow_col,
            tmp_col=tmp_col,
        )
        df = segments[-1]

        return {
            "q_mean": float(df[flow_col].mean()),
            "tmp_mean": float(df[tmp_col].mean()),
            "temp_mean": float(df[temp_col].mean()),
            "R_start": float(df[res_col].mean()),
        }

    def _correct_nuk_with_predict(
            self,
            df_predict,
            unit_name,
            R_start,
            q_mean,
            temp_celsius,
            base_nuk,
            uf_state_default,
    ):
        """
        使用 predict_csv 中预测的渗透率，
        对短期污染参数 nuK 进行在线修正。
        """

        per_col = self.per_format.format(unit=unit_name)
        if per_col not in df_predict.columns:
            return base_nuk
        per = pd.to_numeric(
            df_predict[per_col],
            errors="coerce"
        )
        mask = (
                np.isfinite(per)
                & (per > 0)
        )
        per = per[mask]

        if len(per) < 5:
            return base_nuk


        mu = self.res_calc.xishan_viscosity(temp_celsius)
        R = 3.6e11 / (mu * per)
        R_scaled = R / self.res_calc.scale_factor
        R_scaled = (R_scaled- R_scaled.iloc[0]+ R_start)
        t = np.arange(len(R_scaled)) * 60.0
        J_const = q_mean / self.res_calc.A / 3600.0

        if not np.isfinite(J_const) or J_const <= 0:
            return base_nuk

        x = J_const * t

        try:
            coef = np.polyfit(x, R_scaled, 1)
            predict_nuk = float(coef[0])
            pred = np.polyval(coef, x)
            r2 = r2_score(R_scaled, pred)

        except Exception:
            return base_nuk

        if (
                not np.isfinite(r2)
                or not np.isfinite(predict_nuk)
                or predict_nuk <= 0
                or r2 < 0.4
        ):
            return base_nuk

        corrected_nuk = (
                0.7 * base_nuk
                + 0.3 * predict_nuk
        )

        if (
                not np.isfinite(corrected_nuk)
                or corrected_nuk <= 0
        ):
            return uf_state_default.nuK

        return float(corrected_nuk)