WQ_GUI/src/utils/band_math.py

import pandas as pd
import numpy as np
import re


class BandMathCalculator:
    def __init__(self, csv_file):
        """
        初始化计算器
        csv_file: 包含光谱反射率的CSV文件路径
        """
        self.df = pd.read_csv(csv_file)
        self.wavelengths = self._extract_wavelengths()

    def _extract_wavelengths(self):
        """从列名中提取波长信息"""
        wavelengths = []
        for col in self.df.columns:
            # 尝试从列名中提取数字（波长）
            numbers = re.findall(r'\d+\.?\d*', str(col))
            if numbers:
                wavelengths.append(float(numbers[0]))
            else:
                wavelengths.append(None)
        return wavelengths

    def _find_closest_wavelength(self, target_wavelength):
        """找到最接近目标波长的列索引"""
        valid_indices = [i for i, wl in enumerate(self.wavelengths) if wl is not None]
        if not valid_indices:
            raise ValueError("未找到有效的波长列")

        # 计算与目标波长的差值
        differences = [abs(self.wavelengths[i] - target_wavelength) for i in valid_indices]
        min_diff_index = np.argmin(differences)
        closest_index = valid_indices[min_diff_index]
        closest_wavelength = self.wavelengths[closest_index]

        print(
            f"目标波长 {target_wavelength}nm -> 最接近波长 {closest_wavelength}nm (列: {self.df.columns[closest_index]})")
        return closest_index

    def _parse_expression(self, expression):
        """解析表达式，提取所有波段变量 - 支持大小写"""
        # 匹配 w或W后面跟着数字的格式的变量
        pattern = r'[wW](\d+\.?\d*)'
        matches = re.findall(pattern, expression)
        return matches  # 返回字符串列表，如 ['686', '672', '715', '672']

    def _create_substitution_dict(self, variables, row_index=0):
        """创建变量替换字典 - 支持大小写"""
        substitution_dict = {}
        for var in variables:
            wavelength = float(var)  # 将字符串转换为浮点数
            col_index = self._find_closest_wavelength(wavelength)
            value = self.df.iloc[row_index, col_index]
            # 同时添加小写和大写版本的变量
            substitution_dict[f'w{var}'] = value
            substitution_dict[f'W{var}'] = value
        return substitution_dict

    def calculate(self, expression, row_index=0):
        """
        计算自定义波段表达式

        参数:
        expression: 波段计算表达式，如 'chl=w560/w760'
        row_index: 要计算的数据行索引，默认为第0行

        返回:
        计算结果
        """
        try:
            # 提取表达式中的计算部分
            if '=' in expression:
                # 如果包含赋值，只取等号右边的计算部分
                calc_part = expression.split('=')[1].strip()
                var_name = expression.split('=')[0].strip()
            else:
                calc_part = expression.strip()
                var_name = None

            # 解析变量
            variables = self._parse_expression(calc_part)
            print(f"解析到的波长变量: {variables}")

            # 创建替换字典
            sub_dict = self._create_substitution_dict(variables, row_index)
            print(f"变量值: {sub_dict}")

            # 替换表达式中的变量 - 使用安全的字符串替换
            calc_expression = calc_part
            for var_pattern, value in sub_dict.items():
                # 确保替换完整的变量名，避免部分匹配
                calc_expression = re.sub(r'\b' + re.escape(var_pattern) + r'\b', f"({value})", calc_expression)

            print(f"计算表达式: {calc_expression}")

            # 【新增安全防护】引入 numpy 命名空间，让 eval 引擎安全识别 nan 与 inf
            import numpy as np
            try:
                # 即使 calc_expression 含有纯字符 nan，也能被 np.nan 安全接管
                result = eval(calc_expression, {"__builtins__": None}, {"nan": np.nan, "inf": np.inf, "np": np})
            except Exception as e:
                print(f"⚠️ 警告：公式计算异常 ({e})，该点赋值为 nan")
                result = np.nan

            # 返回结果
            if var_name:
                return {var_name: result}
            else:
                return result

        except Exception as e:
            print(f"计算错误: {e}")
            import traceback
            traceback.print_exc()
            return None

    def calculate_all_rows(self, expression):
        """为所有行计算表达式"""
        results = []
        for i in range(len(self.df)):
            print(f"\n--- 计算第 {i} 行 ---")
            result = self.calculate(expression, i)
            if result is not None:
                if isinstance(result, dict):
                    results.append(list(result.values())[0])
                else:
                    results.append(result)
            else:
                # 如果计算失败，添加NaN值以保持结果数量一致
                results.append(np.nan)
                print(f"第 {i} 行计算失败，使用NaN填充")
        return results

    def process_formulas_from_csv(self, formula_csv_file, formula_names=None, output_file=None):
        """
        从公式CSV文件中批量计算并添加到数据文件中
        
        参数:
        formula_csv_file: 公式CSV文件路径，第一列为公式名称，第三列为具体公式
        formula_names: 要计算的公式名称列表，如果为None则计算所有公式
        output_file: 输出文件路径，如果为None则自动生成
        
        返回:
        包含计算结果的新DataFrame
        """
        # 读取公式CSV文件
        try:
            formulas_df = pd.read_csv(formula_csv_file)
            print(f"读取到 {len(formulas_df)} 个公式")
            
            # 检查CSV格式，假设第一列为公式名称，第三列为具体公式
            if len(formulas_df.columns) < 3:
                raise ValueError("公式CSV文件需要至少3列")
            
            formula_name_col = formulas_df.columns[0]  # 第一列：公式名称
            formula_expr_col = formulas_df.columns[2]  # 第三列：具体公式
            
            # 创建结果DataFrame的副本
            result_df = self.df.copy()
            
            # 如果指定了公式名称，则只计算这些公式
            if formula_names is not None:
                if isinstance(formula_names, str):
                    formula_names = [formula_names]  # 转换为列表
                
                # 筛选出指定的公式
                selected_formulas = formulas_df[formulas_df[formula_name_col].isin(formula_names)]
                print(f"找到 {len(selected_formulas)} 个指定的公式")
                
                if len(selected_formulas) == 0:
                    print(f"警告: 未找到指定的公式: {formula_names}")
                    return result_df
                
                formulas_to_process = selected_formulas
            else:
                # 计算所有公式
                formulas_to_process = formulas_df
            
            # 为每个公式计算所有行
            for _, row in formulas_to_process.iterrows():
                formula_name = row[formula_name_col]
                formula_expr = row[formula_expr_col]
                
                if pd.isna(formula_name) or pd.isna(formula_expr):
                    print(f"跳过空公式: {row}")
                    continue
                    
                print(f"\n计算公式: {formula_name} = {formula_expr}")
                
                # 计算所有行的结果
                results = self.calculate_all_rows(formula_expr)
                
                # 将结果添加到DataFrame
                result_df[formula_name] = results
                print(f"公式 '{formula_name}' 计算完成，添加到数据中")
            
            # 保存结果
            if output_file is None:
                # 自动生成输出文件名
                import os
                base_name = os.path.splitext(os.path.basename(formula_csv_file))[0]
                output_file = f"band_math_results_{base_name}.csv"
            
            result_df.to_csv(output_file, index=False)
            print(f"结果已保存到: {output_file}")
            
            return result_df
            
        except Exception as e:
            print(f"处理公式CSV文件时出错: {e}")
            import traceback
            traceback.print_exc()
            return None


# 更新使用示例
if __name__ == "__main__":
    # 创建计算器实例
    calculator = BandMathCalculator(r"E:\code\WQ\pipeline_result\work_dir\5_training_spectra\training_spectra.csv")

    # 示例1: 计算所有公式
    # result_df = calculator.process_formulas_from_csv(r"E:\code\WQ\封装\sub\水质参数.csv", "enhanced_data.csv")
    
    # 示例2: 计算指定公式
    result_df = calculator.process_formulas_from_csv(
        r"E:\code\WQ\封装\sub\水质参数.csv", 
        formula_names=["BGA_Am09KBBI", "BGA_Be162B643sub629"],
        output_file=r"E:\code\WQ\pipeline_result\work_dir\5_training_spectra\enhanced_data.csv"
    )