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 _parse_coeff(self, coeff_str: str) -> np.ndarray: """ 将 Coefficient 字符串解析为 numpy 多项式系数数组。 - "1.0" → [1.0] - "a,b,c" → [a, b, c] （多项式，从高次到低次，供 np.polyval 直接使用） - "1.0,2.0" → [1.0, 2.0] （线性 y = a*x + b） """ s = str(coeff_str).strip() if not pd.isna(coeff_str) else "" if s in ("", "nan", "None"): return np.array([1.0]) parts = [float(x.strip()) for x in s.split(",")] return np.array(parts) def process_formulas_from_csv(self, formula_csv_file, formula_names=None, output_file=None): """ 从公式CSV文件中批量计算并添加到数据文件中。 支持两种 CSV 格式： - 新版（6列）：Formula_Name, Category, Formula_Type, Formula, Coefficient, Reference - 旧版（≥3列）：第一列=公式名，第三列=表达式（兼容处理） Formula_Type 为 'concentration' 时，计算比值后额外通过 np.polyval 应用 Coefficient。 参数: formula_csv_file: 公式CSV文件路径 formula_names: 要计算的公式名称列表，如果为None则计算所有公式 output_file: 输出文件路径，如果为None则自动生成 返回: 包含计算结果的新DataFrame """ try: formulas_df = pd.read_csv(formula_csv_file) print(f"读取到 {len(formulas_df)} 个公式") has_new_format = set(["Formula_Name", "Formula_Type", "Formula", "Coefficient"]).issubset( set(formulas_df.columns) ) if has_new_format: name_col = "Formula_Name" type_col = "Formula_Type" expr_col = "Formula" coeff_col = "Coefficient" else: name_col = formulas_df.columns[0] type_col = None expr_col = formulas_df.columns[2] coeff_col = None result_df = self.df.copy() if formula_names is not None: if isinstance(formula_names, str): formula_names = [formula_names] selected = formulas_df[formulas_df[name_col].isin(formula_names)] print(f"找到 {len(selected)} 个指定公式") if len(selected) == 0: print(f"警告: 未找到指定公式: {formula_names}") return result_df formulas_to_process = selected else: formulas_to_process = formulas_df for _, row in formulas_to_process.iterrows(): formula_name = row[name_col] formula_expr = row[expr_col] if pd.isna(formula_name) or pd.isna(formula_expr): print(f"跳过空公式: {row.to_dict()}") continue ftype = str(row[type_col]).strip().lower() if type_col and not pd.isna(row.get(type_col)) else "ratio" coeff_str = str(row[coeff_col]).strip() if coeff_col and not pd.isna(row.get(coeff_col)) else "1.0" print(f"\n计算公式: {formula_name} = {formula_expr} [type={ftype}, coeff={coeff_str}]") results = self.calculate_all_rows(formula_expr) if ftype == "concentration": coeff = self._parse_coeff(coeff_str) results = np.polyval(coeff, np.array(results)) 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" )