{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":11403143,"isSourceIdPinned":false,"sourceType":"competition"}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!apt-get update -y\n!apt-get install -y build-essential\n!pip install ViennaRNA","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-14T13:30:07.549525Z","iopub.execute_input":"2025-03-14T13:30:07.549789Z","iopub.status.idle":"2025-03-14T13:30:21.138462Z","shell.execute_reply.started":"2025-03-14T13:30:07.549761Z","shell.execute_reply":"2025-03-14T13:30:21.137621Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport os\nimport re\nimport xgboost as xgb\nfrom sklearn.model_selection import train_test_split, GridSearchCV\nfrom sklearn.metrics import mean_squared_error\nfrom sklearn.preprocessing import OneHotEncoder\nfrom scipy.spatial.transform import Rotation\nfrom scipy.spatial.distance import cdist\nimport logging\nfrom tqdm.auto import tqdm\nimport warnings\nfrom collections import Counter\n\nwarnings.filterwarnings('ignore')\n\n\n# 设置日志记录\ndef setup_logging(log_file='rna_folding.log'):\n    \"\"\"配置日志系统\"\"\"\n    logger = logging.getLogger('rna_folding')\n    logger.setLevel(logging.INFO)\n\n    # 清除已有的处理器\n    if logger.handlers:\n        logger.handlers.clear()\n\n    # 文件处理器\n    fh = logging.FileHandler(log_file)\n    fh.setLevel(logging.INFO)\n\n    # 控制台处理器\n    ch = logging.StreamHandler()\n    ch.setLevel(logging.INFO)\n\n    # 格式化器\n    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')\n    fh.setFormatter(formatter)\n    ch.setFormatter(formatter)\n\n    logger.addHandler(fh)\n    logger.addHandler(ch)\n\n    return logger\n\n\nlogger = setup_logging()\n\n\n# 更新的TM-score计算函数\ndef calculate_tm_score(pred_coords, true_coords):\n    \"\"\"\n    计算TM-score (Template Modeling Score)，按照准确公式实现\n\n    TM-score = max⎛⎝⎜⎜1/Lref ∑i=1 to Lalign 1/(1+(di/d0)²)⎞⎠⎟⎟\n\n    参数:\n        pred_coords: 预测结构的坐标, 形状为(n_residues, 3)\n        true_coords: 参考结构的坐标, 形状为(n_residues, 3)\n\n    返回:\n        tm_score: 0到1之间的分数\n    \"\"\"\n    # 检查输入\n    if pred_coords.shape != true_coords.shape:\n        raise ValueError(f\"预测和参考坐标必须有相同的形状。得到 {pred_coords.shape} 和 {true_coords.shape}\")\n\n    # 获取参考结构残基数量\n    Lref = true_coords.shape[0]\n\n    if Lref < 4:  # 太少的残基无法进行有意义的结构对比\n        return 0.0\n\n    # 计算距离归一化因子 d0 (Angstrom)\n    if Lref >= 30:\n        d0 = 0.6 * (Lref - 0.5) ** 0.5 - 2.5\n    elif Lref >= 24:\n        d0 = 0.7\n    elif Lref >= 20:\n        d0 = 0.6\n    elif Lref >= 16:\n        d0 = 0.5\n    elif Lref >= 12:\n        d0 = 0.4\n    else:  # Lref < 12\n        d0 = 0.3\n\n    # 注意: 我们将尝试不同的旋转和平移来最大化TM-score\n    # 首先将结构居中\n    pred_center = np.mean(pred_coords, axis=0)\n    true_center = np.mean(true_coords, axis=0)\n\n    pred_centered = pred_coords - pred_center\n    true_centered = true_coords - true_center\n\n    best_tm_score = 0.0\n\n    # 使用多个初始对齐来寻找全局最大值\n    # 可以考虑使用不同的初始旋转\n    num_rotations = 10  # 可以根据计算资源调整这个值\n\n    # 从不同的随机旋转开始\n    for _ in range(num_rotations):\n        try:\n            # 生成随机旋转\n            if _ == 0:\n                # 第一次迭代使用SVD进行初始对齐\n                covariance = np.dot(pred_centered.T, true_centered)\n                U, S, Vt = np.linalg.svd(covariance)\n                # 确保旋转矩阵具有正确的手性\n                if np.linalg.det(np.dot(U, Vt)) < 0:\n                    U[:, -1] = -U[:, -1]\n                rotation = np.dot(U, Vt)\n            else:\n                # 随后使用随机旋转作为初始点\n                random_angles = np.random.uniform(0, 2 * np.pi, 3)\n                rotation = Rotation.from_euler('xyz', random_angles).as_matrix()\n\n            # 应用旋转\n            pred_rotated = np.dot(pred_centered, rotation)\n\n            # 进行精细优化，迭代改进对齐\n            for _ in range(10):  # 迭代进行优化\n                # 计算当前对齐下的distances\n                distances = np.sqrt(np.sum((pred_rotated - true_centered) ** 2, axis=1))\n\n                # 计算距离的权重 w = 1 / (1 + (di/d0)^2)\n                weights = 1.0 / (1.0 + (distances / d0) ** 2)\n\n                # 计算加权中心\n                weighted_pred_sum = np.sum(weights[:, np.newaxis] * pred_centered, axis=0)\n                weighted_true_sum = np.sum(weights[:, np.newaxis] * true_centered, axis=0)\n                weight_sum = np.sum(weights)\n\n                if weight_sum > 0:\n                    weighted_pred_center = weighted_pred_sum / weight_sum\n                    weighted_true_center = weighted_true_sum / weight_sum\n\n                    pred_recentered = pred_centered - weighted_pred_center\n                    true_recentered = true_centered - weighted_true_center\n\n                    # 计算加权协方差矩阵\n                    weighted_covariance = np.zeros((3, 3))\n                    for i in range(Lref):\n                        weighted_covariance += weights[i] * np.outer(pred_recentered[i], true_recentered[i])\n\n                    U, S, Vt = np.linalg.svd(weighted_covariance)\n                    # 确保旋转矩阵具有正确的手性\n                    if np.linalg.det(np.dot(U, Vt)) < 0:\n                        U[:, -1] = -U[:, -1]\n                    rotation = np.dot(U, Vt)\n\n                    # 应用新的旋转\n                    pred_rotated = np.dot(pred_recentered, rotation) + weighted_true_center\n\n                    # 计算新的距离\n                    distances = np.sqrt(np.sum((pred_rotated - true_centered) ** 2, axis=1))\n\n                    # 计算当前TM-score\n                    tm_score = (1.0 / Lref) * np.sum(1.0 / (1.0 + (distances / d0) ** 2))\n\n                    if tm_score > best_tm_score:\n                        best_tm_score = tm_score\n\n        except np.linalg.LinAlgError as e:\n            logger.error(f\"SVD计算失败: {e}\")\n            continue\n\n    return best_tm_score\n\n\n# 评估多个结构预测\ndef evaluate_tm_scores(predictions, ground_truth):\n    \"\"\"\n    评估多个结构预测并返回每个目标的最佳TM分数\n\n    参数:\n        predictions: 包含预测的DataFrame\n        ground_truth: 包含真值的DataFrame\n\n    返回:\n        字典，键为目标ID，值为最佳TM分数\n    \"\"\"\n    # 提取唯一的目标ID\n    target_ids = np.unique([id.split('_')[0] for id in predictions['ID']])\n\n    tm_scores = {}\n    for target_id in target_ids:\n        # 获取该目标的预测和真值\n        target_preds = predictions[predictions['ID'].str.startswith(target_id)]\n        target_truth = ground_truth[ground_truth['ID'].str.startswith(target_id)]\n\n        if len(target_preds) == 0 or len(target_truth) == 0:\n            continue\n\n        # 确保行顺序一致\n        target_preds = target_preds.sort_values('resid')\n        target_truth = target_truth.sort_values('resid')\n\n        # 计算每个预测结构的TM分数\n        struct_scores = []\n        for i in range(1, 6):  # 5个预测结构\n            try:\n                pred_coords = target_preds[[f'x_{i}', f'y_{i}', f'z_{i}']].values\n                true_coords = target_truth[['x_1', 'y_1', 'z_1']].values\n\n                if pred_coords.shape == true_coords.shape:\n                    tm_score = calculate_tm_score(pred_coords, true_coords)\n                    struct_scores.append(tm_score)\n            except Exception as e:\n                logger.warning(f\"计算{target_id}的结构{i}时出错: {str(e)}\")\n\n        if struct_scores:\n            tm_scores[target_id] = max(struct_scores)\n\n    return tm_scores\n\n\n# 增强的特征工程 - 使用滑动窗口和更多序列特征\ndef enhanced_feat_eng(df):\n    \"\"\"创建增强的RNA序列特征，包括滑动窗口特征\"\"\"\n    # 基本长度特征\n    result = pd.DataFrame()\n\n    # 确保target_id存在\n    if 'target_id' not in df.columns:\n        result['target_id'] = df.index\n    else:\n        result['target_id'] = df['target_id']\n\n    # 计算序列长度\n    result['seq_length'] = df['sequence'].str.len()\n\n    # 计算单个核苷酸计数\n    for base in ['A', 'C', 'U', 'G']:\n        result[f'{base}_cnt'] = df['sequence'].str.count(base)\n        # 添加百分比特征\n        result[f'{base}_pct'] = result[f'{base}_cnt'] / result['seq_length']\n\n    # 计算GC含量\n    result['gc_content'] = (result['G_cnt'] + result['C_cnt']) / result['seq_length']\n    result['au_content'] = (result['A_cnt'] + result['U_cnt']) / result['seq_length']\n    result['gc_au_ratio'] = result['gc_content'] / result['au_content'].replace(0, 0.001)  # 避免除以零\n\n    # 计算二核苷酸组合\n    for base1 in ['A', 'C', 'U', 'G']:\n        for base2 in ['A', 'C', 'U', 'G']:\n            result[f'{base1}{base2}_cnt'] = df['sequence'].str.count(f'{base1}{base2}')\n            # 添加标准化的二核苷酸频率\n            result[f'{base1}{base2}_freq'] = result[f'{base1}{base2}_cnt'] / (result['seq_length'] - 1).clip(lower=1)\n\n    # 计算三核苷酸的频率\n    important_trimers = ['AAA', 'CCC', 'GGG', 'UUU', 'AUG', 'GCA', 'GUA', 'GUC', 'GUG', 'GUU',\n                         'CAG', 'GGC', 'UCA', 'AGU', 'UCG', 'CUG', 'GAC']\n    for trimer in important_trimers:\n        result[f'{trimer}_freq'] = df['sequence'].apply(\n            lambda x: x.count(trimer) / max(1, len(x) - 2) if len(x) > 0 else 0\n        )\n\n    # 计算序列复杂度指标 - 香农熵\n    def shannon_entropy(seq):\n        if not seq:\n            return 0\n        counts = Counter(seq)\n        probs = [count / len(seq) for count in counts.values()]\n        return -sum(p * np.log2(p) for p in probs)\n\n    result['shannon_entropy'] = df['sequence'].apply(shannon_entropy)\n\n    # 计算序列的长度特征\n    result['seq_len_log'] = np.log1p(result['seq_length'])\n    result['seq_len_sqrt'] = np.sqrt(result['seq_length'])\n\n    return result\n\n\n# 创建滑动窗口特征（针对残基级别）\ndef extract_sliding_window_features(df_seq, df_residues, window_size=5):\n    \"\"\"\n    为每个残基创建滑动窗口特征\n\n    参数:\n        df_seq: 包含序列的DataFrame\n        df_residues: 包含单个残基信息的DataFrame\n        window_size: 窗口大小 (奇数)\n\n    返回:\n        包含窗口特征的DataFrame\n    \"\"\"\n    if window_size % 2 == 0:\n        window_size += 1  # 确保窗口大小是奇数\n\n    half_window = window_size // 2\n    result = pd.DataFrame()\n\n    # 确保有目标ID列\n    if 'target_id' not in df_residues.columns:\n        if 'ID' in df_residues.columns:\n            df_residues['target_id'] = df_residues['ID'].str.rsplit('_', n=1).str[0]\n        else:\n            df_residues['target_id'] = df_residues.index\n\n    # 映射序列到目标ID\n    target_to_seq = dict(zip(df_seq['target_id'], df_seq['sequence']))\n\n    # 注意残基位置是从1开始的\n    def get_window_features(row):\n        target_id = row['target_id']\n        resid = int(row['resid'])\n\n        if target_id not in target_to_seq:\n            return pd.Series({f'pos_{i}': 'X' for i in range(-half_window, half_window + 1)})\n\n        seq = target_to_seq[target_id]\n        if not seq:\n            return pd.Series({f'pos_{i}': 'X' for i in range(-half_window, half_window + 1)})\n\n        # 提取窗口\n        features = {}\n        for i in range(-half_window, half_window + 1):\n            pos = resid - 1 + i  # 转为0-based索引\n            if 0 <= pos < len(seq):\n                features[f'pos_{i}'] = seq[pos]\n            else:\n                features[f'pos_{i}'] = 'X'  # 超出序列范围的填充\n\n        # 添加相对位置特征\n        seq_len = len(seq)\n        features['rel_pos'] = resid / seq_len  # 相对位置 (0-1)\n        features['rel_pos_sin'] = np.sin(2 * np.pi * resid / seq_len)  # 周期性位置特征\n        features['rel_pos_cos'] = np.cos(2 * np.pi * resid / seq_len)\n\n        # 添加局部序列特征\n        window_start = max(0, resid - 1 - half_window)\n        window_end = min(len(seq), resid + half_window)\n        window_seq = seq[window_start:window_end]\n\n        # 局部窗口的核苷酸组成\n        for base in ['A', 'C', 'G', 'U']:\n            features[f'window_{base}_cnt'] = window_seq.count(base)\n            features[f'window_{base}_freq'] = features[f'window_{base}_cnt'] / len(window_seq) if window_seq else 0\n\n        # 局部窗口的GC含量\n        features['window_gc_content'] = (features['window_G_cnt'] + features['window_C_cnt']) / len(\n            window_seq) if window_seq else 0\n\n        # 添加当前位置距离序列两端的距离\n        features['dist_to_start'] = resid - 1  # 距离开始\n        features['dist_to_end'] = seq_len - resid  # 距离结束\n\n        return pd.Series(features)\n\n    # 应用到每一行\n    window_features = df_residues.apply(get_window_features, axis=1)\n\n    # 合并结果\n    result = pd.concat([df_residues, window_features], axis=1)\n\n    return result\n\n\n# 解析目标ID\ndef extract_target_id(id_col):\n    \"\"\"从ID列提取目标ID\"\"\"\n    return id_col.str.split('_').str[:2].str.join('_')\n\n\n# 解析测试序列\ndef parse_test_sequences(test_df):\n    \"\"\"解析测试序列并创建提交格式\"\"\"\n    result = []\n\n    for _, row in test_df.iterrows():\n        seq_length = len(row['sequence'])\n        target_id = row['target_id']\n\n        for i in range(seq_length):\n            resname = row['sequence'][i]\n            resid = i + 1\n            result.append({\n                'ID': f\"{target_id}_{resid}\",\n                'resname': resname,\n                'resid': resid\n            })\n\n    return pd.DataFrame(result)\n\n\n# 应用物理约束的后处理\ndef apply_physical_constraints(coords, resnames, max_iterations=50):\n    \"\"\"\n    应用物理约束到RNA结构坐标，专注于键长约束和冲突解决\n\n    参数:\n        coords: 形状为(n_residues, 3)的坐标数组\n        resnames: 每个残基的类型 (A, C, G, U)\n        max_iterations: 能量最小化的最大迭代次数\n\n    返回:\n        优化后的坐标\n    \"\"\"\n    if len(coords) <= 2:\n        return coords  # 残基太少，无法应用约束\n\n    # 标准RNA骨架键长 (单位：埃)\n    backbone_bond_lengths = {\n        'A': 4.0,  # 腺嘌呤\n        'C': 4.0,  # 胞嘧啶\n        'G': 4.05,  # 鸟嘌呤\n        'U': 4.0  # 尿嘧啶\n    }\n\n    # 创建结果坐标的副本\n    optimized = coords.copy()\n\n    # 迭代优化\n    for iter_num in range(max_iterations):\n        # 跟踪这次迭代中的总移动距离\n        total_movement = 0.0\n\n        # 第1步: 应用键长约束\n        for i in range(1, len(coords)):\n            res_type = resnames[i]\n            ideal_length = backbone_bond_lengths.get(res_type, 4.0)\n\n            # 获取与前一个残基的向量\n            vec = optimized[i] - optimized[i - 1]\n            current_length = np.linalg.norm(vec)\n\n            if abs(current_length - ideal_length) > 0.1:  # 只有在显著偏离时才调整\n                # 正则化并设置为理想长度\n                vec = vec / current_length * ideal_length\n                new_pos = optimized[i - 1] + vec\n\n                # 计算移动距离\n                movement = np.linalg.norm(new_pos - optimized[i])\n                total_movement += movement\n\n                # 更新位置\n                optimized[i] = new_pos\n\n        # 第2步: 检测并解决空间冲突\n        # 计算所有残基对之间的距离\n        distances = cdist(optimized, optimized)\n\n        # 设置对角线元素为一个大值，避免检测自身残基\n        np.fill_diagonal(distances, 999.0)\n\n        # 设置最小允许距离\n        min_allowed_distance = 1.5\n\n        # 找到所有距离小于最小允许距离的残基对\n        conflict_indices = np.where(distances < min_allowed_distance)\n\n        for idx in range(len(conflict_indices[0])):\n            i, j = conflict_indices[0][idx], conflict_indices[1][idx]\n\n            # 只处理一次每一对 (i < j)\n            if i >= j:\n                continue\n\n            # 如果是相邻残基，跳过(它们应该已经通过键长约束处理过)\n            if abs(i - j) == 1:\n                continue\n\n            # 计算当前距离和需要移动的向量\n            current_distance = distances[i, j]\n            direction = optimized[j] - optimized[i]\n            direction = direction / current_distance\n\n            # 计算需要移动的距离\n            move_distance = (min_allowed_distance - current_distance) * 0.5\n            move_vec = direction * move_distance\n\n            # 移动两个残基，使它们远离对方\n            new_i = optimized[i] - move_vec\n            new_j = optimized[j] + move_vec\n\n            # 计算移动距离\n            movement_i = np.linalg.norm(new_i - optimized[i])\n            movement_j = np.linalg.norm(new_j - optimized[j])\n            total_movement += movement_i + movement_j\n\n            # 更新位置\n            optimized[i] = new_i\n            optimized[j] = new_j\n\n        # 如果总移动距离小于阈值，认为优化已收敛\n        if total_movement < 0.1:\n            logger.info(f\"结构优化在第{iter_num + 1}次迭代后收敛\")\n            break\n\n        # 如果这是最后一次迭代，记录日志\n        if iter_num == max_iterations - 1:\n            logger.info(f\"结构优化达到最大迭代次数({max_iterations})，最终移动距离: {total_movement:.3f}\")\n\n    return optimized\n\n\n# 改进的后处理函数\ndef post_process_predictions(predictions):\n    \"\"\"对所有预测结构应用后处理，专注于物理约束\"\"\"\n    result = predictions.copy()\n\n    # 获取唯一目标ID\n    target_ids = np.unique([id.rsplit('_', 1)[0] for id in predictions['ID']])\n\n    for target_id in tqdm(target_ids, desc=\"后处理结构\"):\n        # 获取该目标的预测\n        mask = predictions['ID'].str.startswith(target_id)\n        target_preds = predictions[mask].sort_values('resid')\n\n        # 提取残基名称\n        resnames = target_preds['resname'].values\n\n        # 对每个预测结构应用物理约束\n        for i in range(1, 6):  # 5个预测结构\n            coords = target_preds[[f'x_{i}', f'y_{i}', f'z_{i}']].values\n\n            # 应用物理约束\n            try:\n                constrained_coords = apply_physical_constraints(coords, resnames)\n\n                # 更新预测\n                result.loc[mask, f'x_{i}'] = constrained_coords[:, 0]\n                result.loc[mask, f'y_{i}'] = constrained_coords[:, 1]\n                result.loc[mask, f'z_{i}'] = constrained_coords[:, 2]\n            except Exception as e:\n                logger.warning(f\"处理{target_id}结构{i}时出错: {str(e)}\")\n                # 保留原始坐标\n\n    return result\n\n\n# 修复后的RNA结构预测器类，确保特征一致性\nclass RNAStructurePredictor:\n    \"\"\"管理RNA 3D结构预测工作流的类，包含残基特异性模型\"\"\"\n\n    def __init__(self, model_dir=\"./models\"):\n        \"\"\"初始化预测器和模型目录\"\"\"\n        self.model_dir = model_dir\n        self.models = {'A': {}, 'C': {}, 'G': {}, 'U': {}}  # 按残基类型分类的模型\n        self.feature_columns = None\n        self.categorical_columns = ['resname', 'begin_seq', 'end_seq'] + [f'pos_{i}' for i in range(-2, 3)]\n\n        # 新增：存储每个分类特征的所有可能值\n        self.categorical_values = {}\n\n        # 新增：使用OneHotEncoder替代pd.get_dummies\n        self.encoders = {}\n\n        # 如果目录不存在则创建\n        os.makedirs(model_dir, exist_ok=True)\n\n    def load_data(self, train_path, labels_path):\n        \"\"\"加载并准备训练数据，添加滑动窗口特征\"\"\"\n        logger.info(f\"读取训练数据: {train_path}\")\n        train_sequences = pd.read_csv(train_path)\n\n        logger.info(f\"读取标签数据: {labels_path}\")\n        train_labels = pd.read_csv(labels_path)\n\n        # 提取目标ID\n        train_labels['target_id'] = extract_target_id(train_labels['ID'])\n\n        # 创建基本特征\n        logger.info(\"生成基本特征...\")\n        train_sequences_features = enhanced_feat_eng(train_sequences)\n\n        # 合并基本数据\n        train_data = pd.merge(train_labels, train_sequences_features, on='target_id', how='left')\n\n        # 添加滑动窗口特征\n        logger.info(\"生成滑动窗口特征...\")\n        train_data = extract_sliding_window_features(train_sequences, train_data)\n\n        # 处理缺失值\n        logger.info(\"处理缺失值...\")\n        for col in ['x_1', 'y_1', 'z_1']:\n            # 按目标ID和resname分组计算均值\n            group_means = train_data.groupby(['target_id', 'resname'])[col].transform('mean')\n            # 用组均值填充NA值\n            train_data[col] = train_data[col].fillna(group_means)\n\n        # 删除剩余的NA\n        train_data = train_data.dropna()\n        logger.info(f\"清理后的训练数据: {train_data.shape[0]}行, {train_data.shape[1]}列\")\n\n        return train_data\n\n    def prepare_features(self, data, training=False):\n        \"\"\"\n        准备建模特征，确保分类特征正确编码\n\n        参数:\n            data: 输入数据\n            training: 是否处于训练阶段\n        \"\"\"\n        # 确保窗口位置特征被填充\n        for pos in range(-2, 3):\n            pos_col = f'pos_{pos}'\n            if pos_col not in data.columns:\n                data[pos_col] = 'X'\n\n        # 创建结构化特征列\n        structured_data = data.copy()\n\n        # 处理每个分类特征\n        for col in self.categorical_columns:\n            if col not in structured_data.columns:\n                structured_data[col] = 'X'  # 使用默认值\n\n            # 训练阶段：记录所有可能的值\n            if training:\n                # 存储该特征的所有唯一值\n                self.categorical_values[col] = sorted(structured_data[col].unique().tolist())\n\n                # 为每个分类特征创建一个OneHotEncoder\n                self.encoders[col] = OneHotEncoder(sparse_output=False, handle_unknown='ignore')\n                # 使用所有可能的值拟合编码器\n                self.encoders[col].fit(np.array(self.categorical_values[col]).reshape(-1, 1))\n\n            # 获取编码后的特征\n            if col in self.encoders:\n                # 对该列进行one-hot编码\n                encoded = self.encoders[col].transform(structured_data[col].values.reshape(-1, 1))\n\n                # 创建包含编码值的DataFrame\n                feature_names = [f\"{col}_{val}\" for val in self.encoders[col].categories_[0]]\n                encoded_df = pd.DataFrame(encoded, index=structured_data.index, columns=feature_names)\n\n                # 删除原始列并添加编码列\n                structured_data = pd.concat([structured_data.drop(col, axis=1), encoded_df], axis=1)\n\n        # 识别特征列\n        if self.feature_columns is None and training:\n            self.feature_columns = [col for col in structured_data.columns\n                                    if col not in ['ID', 'target_id', 'x_1', 'y_1', 'z_1']]\n\n        # 确保所有特征列都存在\n        for col in self.feature_columns or []:\n            if col not in structured_data.columns:\n                structured_data[col] = 0\n\n        return structured_data[self.feature_columns] if self.feature_columns else structured_data\n\n    def train(self, train_data, valid_fraction=0.2, optimize=False):\n        \"\"\"为每种残基类型训练单独的x, y, z坐标模型\"\"\"\n        # 分离数据按残基类型\n        residue_types = ['A', 'C', 'G', 'U']\n\n        for res_type in residue_types:\n            logger.info(f\"\\n=== 训练残基类型 {res_type} 的模型 ===\")\n\n            # 过滤该残基类型的数据\n            res_data = train_data[train_data['resname'] == res_type]\n\n            if len(res_data) == 0:\n                logger.warning(f\"没有找到残基类型 {res_type} 的数据，跳过训练\")\n                continue\n\n            # 准备特征 - 传入training=True以记录分类特征值\n            X = self.prepare_features(res_data, training=True)\n\n            # 准备目标\n            y_dict = {\n                'x': res_data['x_1'],\n                'y': res_data['y_1'],\n                'z': res_data['z_1']\n            }\n\n            # 数据分割\n            X_train, X_valid, y_train, y_valid = {}, {}, {}, {}\n            for coord in ['x', 'y', 'z']:\n                X_train[coord], X_valid[coord], y_train[coord], y_valid[coord] = train_test_split(\n                    X, y_dict[coord], test_size=valid_fraction, random_state=42\n                )\n\n            # 如需要优化超参数\n            if optimize:\n                logger.info(f\"为残基类型 {res_type} 执行超参数优化...\")\n                param_grid = {\n                    'max_depth': [6, 10, 15],\n                    'learning_rate': [0.05, 0.1, 0.15],\n                    'n_estimators': [300, 500, 700],\n                    'min_child_weight': [1, 3, 5],\n                    'subsample': [0.7, 0.8, 0.9]\n                }\n\n                best_params = {}\n                for coord in ['x', 'y', 'z']:\n                    logger.info(f\"优化残基类型 {res_type} 的 {coord.upper()} 坐标模型...\")\n                    base_model = xgb.XGBRegressor(objective='reg:squarederror', random_state=42)\n                    grid_search = GridSearchCV(\n                        base_model, param_grid, cv=3, scoring='neg_mean_squared_error', verbose=1\n                    )\n                    grid_search.fit(X_train[coord], y_train[coord])\n                    best_params[coord] = grid_search.best_params_\n                    logger.info(f\"残基类型 {res_type} 的 {coord.upper()} 的最佳参数: {grid_search.best_params_}\")\n            else:\n                # 使用默认参数，但为每种残基类型稍微调整\n                best_params = {\n                    'x': {\n                        'max_depth': 12,\n                        'learning_rate': 0.1,\n                        'n_estimators': 500,\n                        'subsample': 0.85,\n                        'colsample_bytree': 0.85\n                    },\n                    'y': {\n                        'max_depth': 12,\n                        'learning_rate': 0.1,\n                        'n_estimators': 500,\n                        'subsample': 0.85,\n                        'colsample_bytree': 0.85\n                    },\n                    'z': {\n                        'max_depth': 12,\n                        'learning_rate': 0.1,\n                        'n_estimators': 500,\n                        'subsample': 0.85,\n                        'colsample_bytree': 0.85\n                    }\n                }\n\n            # 训练每个坐标的独立模型\n            for coord in ['x', 'y', 'z']:\n                logger.info(f\"\\n训练残基类型 {res_type} 的 {coord.upper()} 坐标模型...\")\n\n                params = best_params[coord].copy()\n\n                # 确保这些参数被设置\n                for key in ['objective', 'random_state', 'n_jobs']:\n                    if key not in params:\n                        params[key] = {'objective': 'reg:squarederror',\n                                       'random_state': 42,\n                                       'n_jobs': -1}[key]\n\n                model = xgb.XGBRegressor(**params)\n                model.fit(\n                    X_train[coord], y_train[coord],\n                    eval_set=[(X_train[coord], y_train[coord]),\n                              (X_valid[coord], y_valid[coord])],\n                    eval_metric='rmse',\n                    verbose=100\n                )\n\n                # 保存模型\n                self.models[res_type][coord] = model\n                model.save_model(f\"{self.model_dir}/xgb_{res_type}_{coord}_model.json\")\n\n                # 保存特征名和分类值信息\n                self._save_feature_info()\n\n                # 打印特征重要性\n                feature_importance = pd.DataFrame({\n                    'feature': X.columns,\n                    'importance': model.feature_importances_\n                }).sort_values('importance', ascending=False)\n\n                logger.info(f\"\\n残基类型 {res_type} 的 {coord.upper()} 模型的前10个重要特征:\")\n                logger.info(feature_importance.head(10).to_string())\n\n                # 评估模型\n                valid_rmse = np.sqrt(mean_squared_error(y_valid[coord], model.predict(X_valid[coord])))\n                logger.info(f\"残基类型 {res_type} 的 {coord.upper()} 模型验证RMSE: {valid_rmse:.4f}\")\n\n    def _save_feature_info(self):\n        \"\"\"保存特征列名和分类特征值\"\"\"\n        import json\n\n        # 保存特征列名\n        if self.feature_columns:\n            with open(f\"{self.model_dir}/feature_columns.json\", 'w') as f:\n                json.dump(self.feature_columns, f)\n\n        # 保存分类特征值\n        with open(f\"{self.model_dir}/categorical_values.json\", 'w') as f:\n            json.dump(self.categorical_values, f)\n\n    def load_models(self):\n        \"\"\"从磁盘加载训练好的残基特异性模型和特征信息\"\"\"\n        import json\n\n        # 加载特征列名\n        try:\n            with open(f\"{self.model_dir}/feature_columns.json\", 'r') as f:\n                self.feature_columns = json.load(f)\n        except FileNotFoundError:\n            logger.warning(\"未找到特征列名文件\")\n\n        # 加载分类特征值\n        try:\n            with open(f\"{self.model_dir}/categorical_values.json\", 'r') as f:\n                self.categorical_values = json.load(f)\n\n            # 重新创建编码器\n            for col, values in self.categorical_values.items():\n                self.encoders[col] = OneHotEncoder(sparse_output=False, handle_unknown='ignore')\n                self.encoders[col].fit(np.array(values).reshape(-1, 1))\n        except FileNotFoundError:\n            logger.warning(\"未找到分类特征值文件\")\n\n        # 加载模型\n        for res_type in ['A', 'C', 'G', 'U']:\n            for coord in ['x', 'y', 'z']:\n                model_path = f\"{self.model_dir}/xgb_{res_type}_{coord}_model.json\"\n                if os.path.exists(model_path):\n                    model = xgb.XGBRegressor()\n                    model.load_model(model_path)\n                    self.models[res_type][coord] = model\n                    logger.info(f\"已加载残基类型 {res_type} 的 {coord.upper()} 坐标模型\")\n                else:\n                    logger.warning(f\"在{model_path}未找到残基类型 {res_type} 的 {coord.upper()} 坐标模型\")\n\n        # 检查是否所有模型都加载了\n        loaded_count = sum([len(models) for models in self.models.values()])\n        if loaded_count < 12:  # 4种残基类型 x 3个坐标 = 12个模型\n            logger.warning(f\"只加载了 {loaded_count}/12 个残基特异性模型\")\n\n    def predict(self, test_sequences, num_structures=5):\n        \"\"\"为测试序列生成预测，使用对应残基类型的模型\"\"\"\n        logger.info(\"正在准备测试数据...\")\n\n        # 解析测试序列\n        test_clean = parse_test_sequences(test_sequences)\n\n        # 创建特征\n        test_sequences_features = enhanced_feat_eng(test_sequences)\n\n        # 合并数据\n        test_data = pd.merge(test_clean, test_sequences_features,\n                             left_on=test_clean['ID'].str.rsplit('_', n=1).str[0],\n                             right_on='target_id',\n                             how='left')\n\n        # 添加滑动窗口特征\n        test_data = extract_sliding_window_features(test_sequences, test_data)\n\n        # 初始化预测数据框 - 使用DataFrame而不是字典+列表\n        predictions = pd.DataFrame({\n            'ID': test_data['ID'],\n            'resname': test_data['resname'],\n            'resid': test_data['resid']\n        })\n\n        # 初始化所有坐标列\n        for i in range(1, num_structures + 1):\n            for coord in ['x', 'y', 'z']:\n                predictions[f'{coord}_{i}'] = 0.0\n\n        logger.info(\"生成预测中...\")\n\n        # 按残基类型分组预测\n        for res_type in ['A', 'C', 'G', 'U']:\n            # 过滤该残基类型的数据\n            res_mask = test_data['resname'] == res_type\n            if not res_mask.any():\n                continue\n\n            # 准备特征 - 确保使用一致的特征集\n            X_test_res = self.prepare_features(test_data[res_mask], training=False)\n\n            # 确保所有所需的模型都存在\n            if not all(coord in self.models[res_type] for coord in ['x', 'y', 'z']):\n                logger.warning(f\"残基类型 {res_type} 缺少一些坐标模型，使用其他残基模型作为后备\")\n                # 找到有所有坐标模型的残基类型作为后备\n                backup_res = next(\n                    (r for r in ['A', 'C', 'G', 'U'] if all(coord in self.models[r] for coord in ['x', 'y', 'z'])),\n                    None)\n                if backup_res:\n                    logger.info(f\"使用残基类型 {backup_res} 的模型作为 {res_type} 的后备\")\n                    self.models[res_type] = self.models[backup_res]\n                else:\n                    raise ValueError(\"没有可用的完整残基模型集\")\n\n            # 为每个结构生成预测\n            for i in range(1, num_structures + 1):\n                for coord in ['x', 'y', 'z']:\n                    # 使用残基特异性模型进行预测\n                    if coord in self.models[res_type]:\n                        # 使用DataFrame的loc索引方式更新值\n                        predictions.loc[res_mask, f'{coord}_{i}'] = self.models[res_type][coord].predict(X_test_res)\n                    else:\n                        raise ValueError(f\"残基类型 {res_type} 的 {coord} 坐标模型未加载\")\n\n        # 确保所有预测列为float类型\n        for i in range(1, num_structures + 1):\n            for coord in ['x', 'y', 'z']:\n                col = f'{coord}_{i}'\n                predictions[col] = predictions[col].astype(float)\n\n        return predictions\n\n    def evaluate(self, predictions, ground_truth):\n        \"\"\"使用TM-score评估预测\"\"\"\n        # 按目标ID分组\n        tm_scores = evaluate_tm_scores(predictions, ground_truth)\n\n        # 计算TM-score统计信息\n        if tm_scores:\n            scores_array = np.array(list(tm_scores.values()))\n            avg_tm_score = np.mean(scores_array)\n            median_tm_score = np.median(scores_array)\n            min_tm_score = np.min(scores_array)\n            max_tm_score = np.max(scores_array)\n\n            logger.info(f\"TM-score统计信息:\")\n            logger.info(f\"平均值: {avg_tm_score:.4f}\")\n            logger.info(f\"中位数: {median_tm_score:.4f}\")\n            logger.info(f\"最小值: {min_tm_score:.4f}\")\n            logger.info(f\"最大值: {max_tm_score:.4f}\")\n\n            # 为有价值的分析创建直方图\n            plt.figure(figsize=(10, 6))\n            plt.hist(scores_array, bins=20, alpha=0.7)\n            plt.title('目标TM-score分布')\n            plt.xlabel('TM-score')\n            plt.ylabel('频率')\n            plt.grid(True, alpha=0.3)\n            plt.savefig('tm_score_distribution.png')\n            plt.close()\n\n            return {\n                'avg': avg_tm_score,\n                'median': median_tm_score,\n                'min': min_tm_score,\n                'max': max_tm_score,\n                'scores': tm_scores\n            }\n        else:\n            logger.warning(\"无法计算有效的TM-score\")\n            return None\n\n\n# 主函数\ndef main():\n    \"\"\"主执行函数\"\"\"\n    logger.info(\"开始RNA 3D结构预测\")\n\n    # 数据路径\n    data_dir = '/kaggle/input/stanford-rna-3d-folding'\n    model_dir = './models'\n\n    # 初始化预测器\n    predictor = RNAStructurePredictor(model_dir=model_dir)\n\n    # 尝试从数据目录读取数据文件\n    try:\n        # 检查是否需要训练模型\n        train_model = True\n\n        # 检查是否有任何残基类型的模型已经存在\n        for res_type in ['A', 'C', 'G', 'U']:\n            if all(os.path.exists(f\"{model_dir}/xgb_{res_type}_{coord}_model.json\") for coord in ['x', 'y', 'z']):\n                train_model = False\n                break\n\n        if train_model:\n            logger.info(\"正在加载训练数据...\")\n            train_data = predictor.load_data(\n                f\"{data_dir}/train_sequences.csv\",\n                f\"{data_dir}/train_labels.csv\"\n            )\n\n            logger.info(\"正在训练残基特异性模型...\")\n            predictor.train(train_data, optimize=False)\n        else:\n            logger.info(\"加载预训练的残基特异性模型...\")\n            predictor.load_models()\n\n        # 检查验证数据\n        val_exists = os.path.exists(f\"{data_dir}/validation_sequences.csv\") and \\\n                     os.path.exists(f\"{data_dir}/validation_labels.csv\")\n\n        if val_exists:\n            logger.info(\"加载验证数据...\")\n            validation_sequences = pd.read_csv(f\"{data_dir}/validation_sequences.csv\")\n            validation_labels = pd.read_csv(f\"{data_dir}/validation_labels.csv\")\n\n            logger.info(\"在验证集上评估模型...\")\n            val_predictions = predictor.predict(validation_sequences)\n\n            # 使用改进的后处理\n            logger.info(\"对验证预测应用物理约束...\")\n            val_predictions = post_process_predictions(val_predictions)\n\n            predictor.evaluate(val_predictions, validation_labels)\n\n        # 检查测试数据\n        test_exists = os.path.exists(f\"{data_dir}/test_sequences.csv\")\n\n        if test_exists:\n            logger.info(\"加载测试数据...\")\n            test_sequences = pd.read_csv(f\"{data_dir}/test_sequences.csv\")\n\n            logger.info(\"生成测试集预测...\")\n            predictions = predictor.predict(test_sequences)\n\n            logger.info(\"应用物理约束后处理...\")\n            predictions = post_process_predictions(predictions)\n\n            # 加载样本提交以确保顺序一致\n            if os.path.exists(f\"{data_dir}/sample_submission.csv\"):\n                sample_submission = pd.read_csv(f\"{data_dir}/sample_submission.csv\")\n                sample_submission['sort_order'] = range(len(sample_submission))\n\n                # 合并并根据sort_order排序\n                predictions = pd.merge(\n                    predictions,\n                    sample_submission[['ID', 'sort_order']],\n                    on='ID',\n                    how='left'\n                )\n                predictions = predictions.sort_values('sort_order').drop('sort_order', axis=1)\n\n            # 保存预测\n            predictions.to_csv('submission.csv', index=False)\n            logger.info(\"预测保存至submission.csv\")\n\n    except Exception as e:\n        logger.error(f\"处理过程中出错: {str(e)}\", exc_info=True)\n\n    logger.info(\"完成\")\n\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-14T13:30:21.140143Z","iopub.execute_input":"2025-03-14T13:30:21.140380Z","iopub.status.idle":"2025-03-14T13:36:07.100548Z","shell.execute_reply.started":"2025-03-14T13:30:21.140360Z","shell.execute_reply":"2025-03-14T13:36:07.099837Z"}},"outputs":[],"execution_count":null}]}