{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":87793,"databundleVersionId":12276181,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport re\nimport subprocess\nimport pandas as pd\nimport numpy as np # For argmax\n\ndef parse_tmscore_output(output: str) -> float:\n    \"\"\"\n    Extracts TM-score from USalign output, expecting the second score\n    (normalized by length of reference structure).\n    \"\"\"\n    tm_score_matches = re.findall(r'TM-score=\\s+([\\d.]+)', output)\n    if not tm_score_matches:\n        raise ValueError('No TM-score pattern found in USalign output.')\n    if len(tm_score_matches) < 2:\n        raise ValueError(f'Expected at least 2 TM-scores in USalign output, found {len(tm_score_matches)}. Output: {output[:500]}')\n    return float(tm_score_matches[1])\n\n\ndef write_target_line(\n    atom_name: str, atom_serial: int, residue_name: str, chain_id: str,\n    residue_num: int, x_coord: float, y_coord: float, z_coord: float,\n    occupancy: float = 1.0, b_factor: float = 0.0, atom_type: str = 'P'\n) -> str:\n    return f'ATOM  {atom_serial:>5d}  {atom_name:<5s} {residue_name:<3s} {residue_num:>3d}    {x_coord:>8.3f}{y_coord:>8.3f}{z_coord:>8.3f}{occupancy:>6.2f}{b_factor:>6.2f}           {atom_type}\\n'\n\n\ndef write2pdb(df: pd.DataFrame, xyz_id: int, target_path: str) -> int:\n    resolved_cnt = 0\n    try:\n        # Ensure DataFrame is not empty and required columns for the given xyz_id exist\n        coord_cols = [f'x_{xyz_id}', f'y_{xyz_id}', f'z_{xyz_id}']\n        if df.empty or not all(col in df.columns for col in coord_cols):\n            # print(f\"Warning: DataFrame is empty or missing coordinate columns for xyz_id {xyz_id} for PDB {target_path}.\")\n            return 0\n            \n        with open(target_path, 'w') as target_file:\n            for _, row in df.iterrows():\n                # Check if coordinate values are present and valid for the current row\n                if any(pd.isna(row[col]) for col in coord_cols):\n                    # print(f\"Skipping row in PDB {target_path} due to NaN coordinates for xyz_id {xyz_id}: resid {row.get('resid', 'N/A')}\")\n                    continue\n\n                x_coord = row[f'x_{xyz_id}']\n                y_coord = row[f'y_{xyz_id}']\n                z_coord = row[f'z_{xyz_id}']\n\n                if x_coord > -1e17 and y_coord > -1e17 and z_coord > -1e17: # Check for sentinel values\n                    resolved_cnt += 1\n                    target_line_str = write_target_line(\n                        atom_name=\"C1'\",\n                        atom_serial=int(row['resid']),\n                        residue_name=row['resname'],\n                        chain_id='0',\n                        residue_num=int(row['resid']),\n                        x_coord=x_coord,\n                        y_coord=y_coord,\n                        z_coord=z_coord,\n                        atom_type='C',\n                    )\n                    target_file.write(target_line_str)\n    except KeyError as e:\n        # This might happen if columns like 'resid' or 'resname' are missing, or if xyz_id leads to missing x_,y_,z_ columns\n        print(f\"Error writing PDB {target_path}: Missing column {e}. Ensure DataFrame has required columns (resid, resname, x_{xyz_id}, y_{xyz_id}, z_{xyz_id}).\")\n        return 0\n    except Exception as e:\n        print(f\"An unexpected error occurred while writing PDB {target_path}: {e}\")\n        return 0\n    return resolved_cnt\n\n\ndef score(solution_df: pd.DataFrame, submission_df: pd.DataFrame, usalign_executable_path: str, verbose: bool = False) -> float:\n    \"\"\"\n    Computes the average of best-of-5 TM-scores for predicted RNA structures against native structures.\n\n    Args:\n        solution_df (pd.DataFrame): DataFrame with native structures.\n        submission_df (pd.DataFrame): DataFrame with predicted structures (5 models).\n        usalign_executable_path (str): Path to the USalign executable.\n        verbose (bool): If True, prints detailed scores for each target and model.\n\n    Returns:\n        float: The average of the highest TM-scores per target.\n    \"\"\"\n    if not os.path.exists(usalign_executable_path):\n        print(f\"Error: USalign executable not found at '{usalign_executable_path}'. Cannot calculate scores.\")\n        return 0.0\n    if not os.access(usalign_executable_path, os.X_OK):\n        print(f\"Error: USalign executable at '{usalign_executable_path}' is not executable. Please check permissions.\")\n        return 0.0\n\n    # Standardize target_id column name\n    solution_target_id_col = 'target_id'\n    submission_target_id_col = 'target_id'\n\n    if solution_target_id_col not in solution_df.columns:\n        solution_df[solution_target_id_col] = solution_df['ID'].apply(lambda x: x.split('_')[0])\n    if submission_target_id_col not in submission_df.columns:\n        submission_df[submission_target_id_col] = submission_df['ID'].apply(lambda x: x.split('_')[0])\n\n    all_target_best_tm_scores = []\n    temp_dir = \"_temp_pdb_scoring\"\n    os.makedirs(temp_dir, exist_ok=True)\n\n    if verbose:\n        print(\"\\n--- Detailed Scoring ---\")\n\n    for target_id, group_native_full in solution_df.groupby(solution_target_id_col):\n        # Filter submission data for the current target_id\n        group_predicted_full = submission_df[submission_df[submission_target_id_col] == target_id]\n\n        if group_predicted_full.empty:\n            if verbose:\n                print(f\"\\nTarget: {target_id}\")\n                print(\"  Status: No predictions found in submission data. Skipping.\")\n            all_target_best_tm_scores.append(0.0) # Append 0 for missing targets as per original behavior\n            continue\n        \n        if verbose:\n            print(f\"\\nTarget: {target_id}\")\n\n        native_pdb_path = os.path.join(temp_dir, f\"{target_id}_native.pdb\")\n        num_native_atoms = write2pdb(group_native_full, 1, native_pdb_path) # Native uses xyz_id=1\n\n        if num_native_atoms == 0:\n            if verbose:\n                print(f\"  Status: Native PDB for {target_id} is empty or could not be written. TM-score: 0.0\")\n            all_target_best_tm_scores.append(0.0)\n            if os.path.exists(native_pdb_path): os.remove(native_pdb_path)\n            continue\n\n        current_target_tm_scores = []\n        for pred_model_idx in range(1, 6): # Models 1 to 5\n            predicted_pdb_path = os.path.join(temp_dir, f\"{target_id}_predicted_model_{pred_model_idx}.pdb\")\n            num_pred_atoms = write2pdb(group_predicted_full, pred_model_idx, predicted_pdb_path)\n            \n            tm_score_for_model = 0.0\n            if num_pred_atoms == 0:\n                if verbose:\n                    print(f\"  Model {pred_model_idx}: Predicted PDB empty. TM-score: 0.0\")\n            else:\n                command = [usalign_executable_path, predicted_pdb_path, native_pdb_path, \"-atom\", \" C1'\"]\n                try:\n                    result = subprocess.run(command, capture_output=True, text=True, check=False, timeout=120)\n                    if result.returncode == 0 and result.stdout:\n                        tm_score_for_model = parse_tmscore_output(result.stdout)\n                    elif verbose: # Print errors only in verbose mode if USalign fails\n                        print(f\"  Model {pred_model_idx}: USalign execution failed or no output.\")\n                        # print(f\"    Return Code: {result.returncode}, Stdout: '{result.stdout[:100]}...', Stderr: '{result.stderr[:100]}...'\")\n                except subprocess.TimeoutExpired:\n                    if verbose: print(f\"  Model {pred_model_idx}: USalign timed out.\")\n                except ValueError as e:\n                    if verbose: print(f\"  Model {pred_model_idx}: Error parsing USalign output: {e}\")\n                except Exception as e:\n                    if verbose: print(f\"  Model {pred_model_idx}: Unexpected error: {e}\")\n            \n            current_target_tm_scores.append(tm_score_for_model)\n            if verbose:\n                 print(f\"  Model {pred_model_idx}: TM-score = {tm_score_for_model:.4f}\")\n            \n            if os.path.exists(predicted_pdb_path): os.remove(predicted_pdb_path)\n        \n        if os.path.exists(native_pdb_path): os.remove(native_pdb_path)\n\n        if current_target_tm_scores:\n            best_tm_for_target = max(current_target_tm_scores)\n            best_model_index = np.argmax(current_target_tm_scores) + 1 # 1-indexed\n            all_target_best_tm_scores.append(best_tm_for_target)\n            if verbose:\n                print(f\"  Best TM-score for {target_id}: {best_tm_for_target:.4f} (from Model {best_model_index})\")\n        else:\n            all_target_best_tm_scores.append(0.0)\n            if verbose:\n                print(f\"  Status: No valid TM-scores obtained for {target_id}.\")\n    \n    # Cleanup temp_dir\n    try:\n        for f_name in os.listdir(temp_dir):\n            os.remove(os.path.join(temp_dir, f_name))\n        os.rmdir(temp_dir)\n    except OSError as e:\n        print(f\"Warning: Could not completely remove temporary directory {temp_dir}: {e}\")\n\n    if not all_target_best_tm_scores:\n        if verbose: print(\"\\nNo TM-scores calculated for any target.\")\n        return 0.0\n    \n    final_average_score = sum(all_target_best_tm_scores) / len(all_target_best_tm_scores)\n    if verbose:\n        print(f\"\\n--- Scoring Summary ---\")\n        print(f\"Number of targets processed: {len(all_target_best_tm_scores)}\")\n        print(f\"Average Best-of-5 TM-score: {final_average_score:.4f}\")\n    return final_average_score\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"    import pandas as pd\n    submission_df = pd.read_csv(\"/kaggle/working/validation.csv\")\n\n    solution_df = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/validation_labels.csv\") # Or your local path\n    \n    os.system('cp //kaggle/input/usalign/USalign /kaggle/working/')\n    os.system('sudo chmod u+x /kaggle/working//USalign')\n    usalign_exe = \"/kaggle/working//USalign\" \n\n    final_score = score(solution_df, submission_df, usalign_exe, verbose=True)\n    print(f\"Competition Score (Average Best-of-5 TM-score): {final_score:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}