{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.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":118765,"databundleVersionId":15231210,"sourceType":"competition"},{"sourceId":10855324,"sourceType":"datasetVersion","datasetId":6742586},{"sourceId":11118830,"sourceType":"datasetVersion","datasetId":6933267},{"sourceId":14519720,"sourceType":"datasetVersion","datasetId":9271415},{"sourceId":14534919,"sourceType":"datasetVersion","datasetId":9283271},{"sourceId":311741,"sourceType":"modelInstanceVersion","modelInstanceId":264400,"modelId":285488}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# V2: imports recomputed ccd-cache","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Reference\n\n- https://www.kaggle.com/code/theoviel/stanford-rna-3d-folding-pt2-rnapro-inference","metadata":{}},{"cell_type":"code","source":"import os\nIS_SCORING_RUN = os.environ.get('KAGGLE_IS_COMPETITION_RERUN')\nprint(IS_SCORING_RUN)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T12:38:17.706437Z","iopub.execute_input":"2026-01-17T12:38:17.706758Z","iopub.status.idle":"2026-01-17T12:38:17.711385Z","shell.execute_reply.started":"2026-01-17T12:38:17.706726Z","shell.execute_reply":"2026-01-17T12:38:17.710612Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/input/rnapro-src/RNAPro .\n!cp /kaggle/input/rnapro-src/rnapro-private-best-500m.ckpt .","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T12:38:17.886495Z","iopub.execute_input":"2026-01-17T12:38:17.887164Z","iopub.status.idle":"2026-01-17T12:39:00.247825Z","shell.execute_reply.started":"2026-01-17T12:38:17.887138Z","shell.execute_reply":"2026-01-17T12:39:00.247017Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cd RNAPro","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T12:39:00.249435Z","iopub.execute_input":"2026-01-17T12:39:00.249687Z","iopub.status.idle":"2026-01-17T12:39:00.257112Z","shell.execute_reply.started":"2026-01-17T12:39:00.249658Z","shell.execute_reply":"2026-01-17T12:39:00.256455Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pwd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T12:39:00.25801Z","iopub.execute_input":"2026-01-17T12:39:00.25831Z","iopub.status.idle":"2026-01-17T12:39:00.271905Z","shell.execute_reply.started":"2026-01-17T12:39:00.258283Z","shell.execute_reply":"2026-01-17T12:39:00.271294Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pip install -e . --no-deps","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T12:39:00.27339Z","iopub.execute_input":"2026-01-17T12:39:00.273565Z","iopub.status.idle":"2026-01-17T12:39:05.177589Z","shell.execute_reply.started":"2026-01-17T12:39:00.273549Z","shell.execute_reply":"2026-01-17T12:39:05.176836Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cd ..","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T12:39:05.178647Z","iopub.execute_input":"2026-01-17T12:39:05.17887Z","iopub.status.idle":"2026-01-17T12:39:05.183931Z","shell.execute_reply.started":"2026-01-17T12:39:05.178844Z","shell.execute_reply":"2026-01-17T12:39:05.183291Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pwd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T12:39:05.184779Z","iopub.execute_input":"2026-01-17T12:39:05.185023Z","iopub.status.idle":"2026-01-17T12:39:07.828088Z","shell.execute_reply.started":"2026-01-17T12:39:05.185004Z","shell.execute_reply":"2026-01-17T12:39:07.82735Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"TBM\n---\n","metadata":{}},{"cell_type":"code","source":"# Cell 1: Imports and Setup\nimport os\nimport pandas as pd\nimport numpy as np\nfrom scipy.spatial.transform import Rotation as R\nfrom scipy.spatial import distance_matrix\nimport random\nimport time\nimport warnings\n\nfrom tqdm.auto import tqdm\nwarnings.filterwarnings('ignore')\nfrom Bio.Align import PairwiseAligner\nfrom Bio.Seq import Seq\n\n# Enable tqdm for pandas operations\ntqdm.pandas()\n    \n# Initialize global aligner with RNA-appropriate parameters\nALIGNER = PairwiseAligner()\nALIGNER.mode = 'global'\nALIGNER.match_score = 2.1\nALIGNER.mismatch_score = -1\nALIGNER.open_gap_score = -10\nALIGNER.extend_gap_score = -0.5\n\nseed = 21\nnp.random.seed(seed)\nrandom.seed(seed)\n\n# Cell 2: Load Data\nBASE_PATH = '/kaggle/input/stanford-rna-3d-folding-2'\n\n# Load with progress indication\ntest_seqs = pd.read_csv(f'{BASE_PATH}/test_sequences.csv')\ntrain_seqs = pd.read_csv(f'{BASE_PATH}/train_sequences.csv')\nvalidation_seqs = pd.read_csv(f'{BASE_PATH}/validation_sequences.csv')\ntrain_labels = pd.read_csv(f'{BASE_PATH}/train_labels.csv', low_memory=False)\nvalidation_labels = pd.read_csv(f'{BASE_PATH}/validation_labels.csv')\nsample_submission = pd.read_csv(f'{BASE_PATH}/sample_submission.csv')\n\nprint(f\"✓ Loaded {len(train_seqs)} training sequences\")\nprint(f\"✓ Loaded {len(validation_seqs)} validation sequences\") \nprint(f\"✓ Loaded {len(test_seqs)} test sequences\")\n\n# Cell 3: Process Training Labels to Coordinate Dictionary\ndef process_labels(labels_df, use_first_model_only=True):\n    \"\"\"\n    Process labels dataframe to create a dictionary mapping target_id to coordinates.\n    Vectorized implementation for improved performance.\n    \n    Args:\n        labels_df: DataFrame with ID, resid, x_1, y_1, z_1, etc.\n        use_first_model_only: If True, only extract first model coordinates\n        \n    Returns:\n        Dictionary mapping target_id to numpy array of coordinates\n    \"\"\"\n    print(\"Extracting target IDs...\")\n    # Vectorized target_id extraction\n    labels_df = labels_df.copy()\n    labels_df['target_id'] = labels_df['ID'].str.rsplit('_', n=1).str[0]\n    \n    # Sort once for all groups\n    labels_df = labels_df.sort_values(['target_id', 'resid'])\n    \n    # Vectorized coordinate extraction\n    print(\"Extracting coordinates...\")\n    coord_cols = ['x_1', 'y_1', 'z_1']\n    \n    # Replace placeholder values with NaN in one operation\n    coords_array = labels_df[coord_cols].values.copy()\n    coords_array[coords_array < -1e6] = np.nan\n    labels_df[coord_cols] = coords_array\n    \n    # Group and convert to dictionary with progress bar\n    print(\"Grouping by target_id...\")\n    coords_dict = {}\n    \n    grouped = labels_df.groupby('target_id', sort=False)\n    for target_id, group in tqdm(grouped, desc=\"Processing structures\", total=len(grouped)):\n        # Directly extract coordinates as numpy array (already sorted)\n        coords_dict[target_id] = group[coord_cols].values\n    \n    return coords_dict\n\n\ntrain_coords_dict = process_labels(train_labels)\nvalid_coords_dict = process_labels(validation_labels)\n\n# Cell 4: Sequence Alignment and Template Finding\ndef get_alignment_score(query_seq, template_seq):\n    alignments = ALIGNER.align(query_seq, template_seq)\n    best_alignment = next(iter(alignments), None)\n    \n    if best_alignment is None:\n        return 0.0\n    \n    # Normalize score by theoretical max (perfect match of shorter sequence)\n    max_possible = 2.1 * min(len(query_seq), len(template_seq))\n    normalized_score = best_alignment.score / max_possible\n    \n    return min(normalized_score, 1.0)  # Cap at 1.0\n\n\ndef get_aligned_sequences(query_seq, template_seq):\n    alignments = ALIGNER.align(query_seq, template_seq)\n    best_alignment = next(iter(alignments), None)\n    \n    if best_alignment is None:\n        return None, None\n    \n    # Use the robust method that builds from alignment coordinates\n    return build_aligned_sequences(query_seq, template_seq, best_alignment)\n\n\ndef build_aligned_sequences(query_seq, template_seq, alignment):\n    \"\"\"\n    Build aligned sequences with gaps from alignment coordinates.\n    Uses alignment.aligned property which returns tuples of (start, end) ranges.\n    \"\"\"\n    # Get aligned blocks: alignment.aligned returns ((query_ranges), (template_ranges))\n    query_ranges, template_ranges = alignment.aligned\n    \n    # If no aligned blocks, return None\n    if len(query_ranges) == 0:\n        return None, None\n    \n    aligned_query = []\n    aligned_template = []\n    \n    query_pos = 0\n    template_pos = 0\n    \n    for (q_start, q_end), (t_start, t_end) in zip(query_ranges, template_ranges):\n        # Add gaps for unaligned query residues (query has residues, template doesn't)\n        while query_pos < q_start:\n            aligned_query.append(query_seq[query_pos])\n            aligned_template.append('-')\n            query_pos += 1\n        \n        # Add gaps for unaligned template residues (template has residues, query doesn't)\n        while template_pos < t_start:\n            aligned_query.append('-')\n            aligned_template.append(template_seq[template_pos])\n            template_pos += 1\n        \n        # Add aligned region (both have residues)\n        block_len = q_end - q_start  # Should equal t_end - t_start\n        for i in range(block_len):\n            aligned_query.append(query_seq[q_start + i])\n            aligned_template.append(template_seq[t_start + i])\n        \n        query_pos = q_end\n        template_pos = t_end\n    \n    # Add any remaining unaligned query residues at the end\n    while query_pos < len(query_seq):\n        aligned_query.append(query_seq[query_pos])\n        aligned_template.append('-')\n        query_pos += 1\n    \n    # Add any remaining unaligned template residues at the end\n    while template_pos < len(template_seq):\n        aligned_query.append('-')\n        aligned_template.append(template_seq[template_pos])\n        template_pos += 1\n    \n    return ''.join(aligned_query), ''.join(aligned_template)\n\n\ndef find_similar_sequences(query_seq, train_seqs_df, train_coords_dict, \n                          temporal_cutoff=None, top_n=5):\n    \"\"\"\n    Find sequences in the training data similar to the query sequence.\n    Uses modern PairwiseAligner for sequence comparison.\n    \n    Args:\n        query_seq: The RNA sequence to find templates for\n        train_seqs_df: DataFrame containing training sequences\n        train_coords_dict: Dictionary mapping target_ids to their 3D coordinates\n        temporal_cutoff: Only consider training sequences published before this date\n        top_n: Number of top templates to return\n        \n    Returns:\n        List of (target_id, sequence, similarity_score, coordinates) tuples\n    \"\"\"\n    similar_seqs = []\n    \n    # Filter training sequences by temporal cutoff if provided\n    if temporal_cutoff:\n        filtered_train_seqs = train_seqs_df[train_seqs_df['temporal_cutoff'] < temporal_cutoff]\n    else:\n        filtered_train_seqs = train_seqs_df\n    \n    for _, row in filtered_train_seqs.iterrows():\n        target_id = row['target_id']\n        train_seq = row['sequence']\n        \n        # Skip if coordinates not available\n        if target_id not in train_coords_dict:\n            continue\n        \n        # Skip if sequence length difference is too large (>50%)\n        len_diff = abs(len(train_seq) - len(query_seq)) / max(len(train_seq), len(query_seq))\n        if len_diff > 0.5:\n            continue\n        \n        # Calculate similarity score using new aligner\n        similarity_score = get_alignment_score(query_seq, train_seq)\n        \n        similar_seqs.append((target_id, train_seq, similarity_score, train_coords_dict[target_id]))\n    \n    # Sort by similarity score (higher is better)\n    similar_seqs.sort(key=lambda x: x[2], reverse=True)\n    return similar_seqs[:top_n]\n\n# Cell 5: Template Adaptation Functions\ndef adapt_template_to_query(query_seq, template_seq, template_coords):\n    # Get aligned sequences\n    aligned_query, aligned_template = get_aligned_sequences(query_seq, template_seq)\n    \n    if aligned_query is None or aligned_template is None:\n        return adapt_template_simple(query_seq, template_seq, template_coords)\n    \n    # Initialize coordinates for query sequence\n    query_coords = np.zeros((len(query_seq), 3))\n    query_coords.fill(np.nan)\n    \n    # Map template coordinates to query based on alignment\n    query_idx = 0\n    template_idx = 0\n    \n    for i in range(len(aligned_query)):\n        query_char = aligned_query[i]\n        template_char = aligned_template[i]\n        \n        if query_char != '-' and template_char != '-':\n            # Both aligned - copy template coordinate to query\n            if query_idx < len(query_seq) and template_idx < len(template_coords):\n                # Handle NaN coordinates in template\n                if not np.any(np.isnan(template_coords[template_idx])):\n                    query_coords[query_idx] = template_coords[template_idx]\n            template_idx += 1\n            query_idx += 1\n        elif query_char != '-' and template_char == '-':\n            # Gap in template - query residue has no template coord\n            query_idx += 1\n        elif query_char == '-' and template_char != '-':\n            # Gap in query - skip template residue\n            template_idx += 1\n    \n    # Fill in gaps by interpolation\n    query_coords = fill_coordinate_gaps(query_coords)\n    \n    return query_coords\n\n\ndef adapt_template_simple(query_seq, template_seq, template_coords):\n    \"\"\"Simple template adaptation without Biopython alignment.\"\"\"\n    query_coords = np.zeros((len(query_seq), 3))\n    \n    # Simple mapping based on position\n    scale = len(template_coords) / len(query_seq)\n    for i in range(len(query_seq)):\n        template_idx = int(i * scale)\n        template_idx = min(template_idx, len(template_coords) - 1)\n        if not np.any(np.isnan(template_coords[template_idx])):\n            query_coords[i] = template_coords[template_idx]\n        else:\n            query_coords[i] = [np.nan, np.nan, np.nan]\n    \n    # Fill gaps\n    query_coords = fill_coordinate_gaps(query_coords)\n    return query_coords\n\n\ndef fill_coordinate_gaps(coords):\n    \"\"\"Fill NaN gaps in coordinates by interpolation.\"\"\"\n    n = len(coords)\n    typical_step = 4.0\n    \n    # First pass: interpolate between valid points\n    for i in range(n):\n        if np.isnan(coords[i, 0]):\n            prev_valid = next((j for j in range(i-1, -1, -1) if not np.isnan(coords[j, 0])), -1)\n            next_valid = next((j for j in range(i+1, n) if not np.isnan(coords[j, 0])), -1)\n            \n            if prev_valid >= 0 and next_valid >= 0:\n                weight = (i - prev_valid) / (next_valid - prev_valid)\n                coords[i] = (1 - weight) * coords[prev_valid] + weight * coords[next_valid]\n    \n    # Second pass: handle remaining NaNs at edges\n    for i in range(n):\n        if np.isnan(coords[i, 0]):\n            if i == 0:\n                first_valid = next((j for j in range(1, n) if not np.isnan(coords[j, 0])), -1)\n                if first_valid >= 0:\n                    for j in range(first_valid - 1, -1, -1):\n                        direction = np.random.normal(0, 1, 3)\n                        direction = direction / (np.linalg.norm(direction) + 1e-10) * typical_step\n                        coords[j] = coords[j + 1] - direction\n                else:\n                    coords = generate_basic_structure_coords(n)\n                    break\n            else:\n                prev_valid = next((j for j in range(i-1, -1, -1) if not np.isnan(coords[j, 0])), -1)\n                if prev_valid >= 0:\n                    direction = np.random.normal(0, 1, 3)\n                    direction = direction / (np.linalg.norm(direction) + 1e-10) * typical_step\n                    coords[i] = coords[prev_valid] + direction\n    \n    # Final cleanup\n    coords = np.nan_to_num(coords)\n    return coords\n\n\ndef generate_basic_structure(sequence):\n    \"\"\"Generate a simple helical structure.\"\"\"\n    return generate_basic_structure_coords(len(sequence))\n\n\ndef generate_basic_structure_coords(n_residues):\n    \"\"\"Generate basic helical coordinates.\"\"\"\n    coords = np.zeros((n_residues, 3))\n    for i in range(n_residues):\n        angle = i * 0.6\n        coords[i] = [10.0 * np.cos(angle), 10.0 * np.sin(angle), i * 2.5]\n    return coords\n\n# Cell 6: RNA Geometric Constraints\ndef adaptive_rna_constraints(coordinates, sequence, confidence=1.0):\n    \"\"\"\n    Apply RNA geometric constraints with adaptive strength based on confidence.\n    \"\"\"\n    refined_coords = coordinates.copy()\n    n_residues = len(sequence)\n    \n    constraint_strength = 0.8 * (1.0 - min(confidence, 0.8))\n    \n    # Sequential distance constraints\n    seq_min_dist, seq_max_dist = 5.5, 6.5\n    \n    for i in range(n_residues - 1):\n        current_pos = refined_coords[i]\n        next_pos = refined_coords[i + 1]\n        current_dist = np.linalg.norm(next_pos - current_pos)\n        \n        if current_dist < seq_min_dist or current_dist > seq_max_dist:\n            target_dist = (seq_min_dist + seq_max_dist) / 2\n            direction = next_pos - current_pos\n            direction = direction / (np.linalg.norm(direction) + 1e-10)\n            adjustment = (target_dist - current_dist) * constraint_strength\n            refined_coords[i + 1] = current_pos + direction * (current_dist + adjustment)\n    \n    # Steric clash prevention\n    min_allowed_distance = 3.8\n    dist_matrix = distance_matrix(refined_coords, refined_coords)\n    severe_clashes = np.where((dist_matrix < min_allowed_distance) & (dist_matrix > 0))\n    \n    for idx in range(len(severe_clashes[0])):\n        i, j = severe_clashes[0][idx], severe_clashes[1][idx]\n        if abs(i - j) <= 1 or i >= j:\n            continue\n        \n        pos_i, pos_j = refined_coords[i], refined_coords[j]\n        current_dist = dist_matrix[i, j]\n        direction = pos_j - pos_i\n        direction = direction / (np.linalg.norm(direction) + 1e-10)\n        adjustment = (min_allowed_distance - current_dist) * constraint_strength\n        refined_coords[i] = pos_i - direction * (adjustment / 2)\n        refined_coords[j] = pos_j + direction * (adjustment / 2)\n    \n    return refined_coords\n\n# Cell 7: De Novo Structure Generation\ndef generate_rna_structure(sequence, seed=None):\n    \"\"\"Generate a more realistic RNA structure prediction.\"\"\"\n    if seed is not None:\n        np.random.seed(seed)\n        random.seed(seed)\n    \n    n_residues = len(sequence)\n    coordinates = np.zeros((n_residues, 3))\n    \n    # Initialize first residues\n    for i in range(min(3, n_residues)):\n        angle = i * 0.6\n        coordinates[i] = [10.0 * np.cos(angle), 10.0 * np.sin(angle), i * 2.5]\n    \n    current_direction = np.array([0.0, 0.0, 1.0])\n    complementary = {'G': 'C', 'C': 'G', 'A': 'U', 'U': 'A'}\n    \n    for i in range(3, n_residues):\n        current_base = sequence[i]\n        has_pair = False\n        pair_idx = -1\n        \n        window_size = min(i, 15)\n        for j in range(i - window_size, i):\n            if j >= 0 and sequence[j] == complementary.get(current_base, 'X'):\n                has_pair = True\n                pair_idx = j\n                break\n        \n        if has_pair and i - pair_idx <= 10 and random.random() < 0.7:\n            pair_pos = coordinates[pair_idx]\n            random_offset = np.random.normal(0, 1, 3) * 2.0\n            base_pair_distance = 10.0 + random.uniform(-1.0, 1.0)\n            \n            center = np.mean(coordinates[:i], axis=0)\n            direction = center - pair_pos\n            direction = direction / (np.linalg.norm(direction) + 1e-10)\n            \n            coordinates[i] = pair_pos + direction * base_pair_distance + random_offset\n            current_direction = np.random.normal(0, 0.3, 3)\n            current_direction = current_direction / (np.linalg.norm(current_direction) + 1e-10)\n        else:\n            if random.random() < 0.3:\n                angle = random.uniform(0.2, 0.6)\n                axis = np.random.normal(0, 1, 3)\n                axis = axis / (np.linalg.norm(axis) + 1e-10)\n                rotation = R.from_rotvec(angle * axis)\n                current_direction = rotation.apply(current_direction)\n            else:\n                current_direction += np.random.normal(0, 0.15, 3)\n                current_direction = current_direction / (np.linalg.norm(current_direction) + 1e-10)\n            \n            step_size = random.uniform(3.5, 4.5)\n            coordinates[i] = coordinates[i - 1] + step_size * current_direction\n    \n    return coordinates\n\n# Cell 8: Main Prediction Function\ndef predict_rna_structures(sequence, target_id, train_seqs_df, train_coords_dict, \n                          n_predictions=5, temporal_cutoff=None):\n    \"\"\"Generate multiple structure predictions for an RNA sequence.\"\"\"\n    predictions = []\n    \n    # Find similar sequences\n    similar_seqs = find_similar_sequences(\n        sequence, train_seqs_df, train_coords_dict,\n        temporal_cutoff=temporal_cutoff, top_n=n_predictions\n    )\n    \n    # Use templates if found\n    if similar_seqs:\n        for template_id, template_seq, similarity, template_coords in similar_seqs:\n            adapted_coords = adapt_template_to_query(sequence, template_seq, template_coords)\n            \n            if adapted_coords is not None:\n                refined_coords = adaptive_rna_constraints(adapted_coords, sequence, confidence=similarity)\n                random_scale = max(0.05, 0.8 - similarity)\n                randomized_coords = refined_coords + np.random.normal(0, random_scale, refined_coords.shape)\n                predictions.append(randomized_coords)\n                \n                if len(predictions) >= n_predictions:\n                    break\n    \n    # Fill remaining with de novo structures\n    while len(predictions) < n_predictions:\n        seed_value = hash(target_id) % 10000 + len(predictions) * 1000\n        de_novo_coords = generate_rna_structure(sequence, seed=seed_value)\n        refined_de_novo = adaptive_rna_constraints(de_novo_coords, sequence, confidence=0.2)\n        predictions.append(refined_de_novo)\n    \n    return predictions[:n_predictions]\n\n# Cell 9: Generate Predictions\ndef generate_predictions_for_dataset(seqs_df, train_seqs_df, train_coords_dict, dataset_name=\"dataset\"):\n    \"\"\"Generate predictions for a given sequence dataset.\"\"\"\n    all_predictions = []\n    start_time = time.time()\n    total_targets = len(seqs_df)\n    \n    print(f\"\\n=== Generating Predictions for {dataset_name} ({total_targets} sequences) ===\")\n    \n    for idx, row in seqs_df.iterrows():\n        target_id = row['target_id']\n        sequence = row['sequence']\n        temporal_cutoff = row.get('temporal_cutoff', None)\n        \n        if idx % 5 == 0:\n            elapsed = time.time() - start_time\n            if idx > 0:\n                est_remaining = elapsed / (idx + 1) * (total_targets - idx - 1)\n                print(f\"Processing {idx+1}/{total_targets}: {target_id} ({len(sequence)} nt), \"\n                      f\"elapsed: {elapsed:.1f}s, remaining: {est_remaining:.1f}s\")\n            else:\n                print(f\"Processing {idx+1}/{total_targets}: {target_id} ({len(sequence)} nt)\")\n        \n        predictions = predict_rna_structures(\n            sequence, target_id, train_seqs_df, train_coords_dict,\n            n_predictions=5, temporal_cutoff=temporal_cutoff\n        )\n        \n        for j in range(len(sequence)):\n            pred_row = {\n                'ID': f\"{target_id}_{j+1}\",\n                'resname': sequence[j],\n                'resid': j + 1\n            }\n            for i in range(5):\n                pred_row[f'x_{i+1}'] = predictions[i][j][0]\n                pred_row[f'y_{i+1}'] = predictions[i][j][1]\n                pred_row[f'z_{i+1}'] = predictions[i][j][2]\n            all_predictions.append(pred_row)\n    \n    # Create DataFrame\n    df = pd.DataFrame(all_predictions)\n    column_order = ['ID', 'resname', 'resid']\n    for i in range(1, 6):\n        for coord in ['x', 'y', 'z']:\n            column_order.append(f'{coord}_{i}')\n    df = df[column_order]\n    \n    print(f\"\\nGenerated predictions for {total_targets} sequences\")\n    print(f\"Total runtime: {time.time() - start_time:.1f} seconds\")\n    print(f\"Output shape: {df.shape}\")\n    \n    return df\n\n# Cell 10: Generate Predictions for TEST Set (Submission)\n# Generate test predictions for submission\ntest_pred_df = generate_predictions_for_dataset(\n    test_seqs, train_seqs, train_coords_dict,\n    dataset_name=\"Test\"\n)\ntest_pred_df.head()\n\n# Cell 11: Save Submission\ntest_pred_df.to_csv('submission_tbm.csv', index=False)\nprint(\"\\n=== Submission file saved ===\")\nprint(f\"submission.csv shape: {test_pred_df.shape}\")\nprint(f\"\\nFirst few rows:\")\ntest_pred_df.head(10)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sub_tbm = pd.read_csv(\"/kaggle/working/submission_tbm.csv\")\nsub_tbm","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pwd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T09:55:31.612382Z","iopub.execute_input":"2026-01-17T09:55:31.6127Z","iopub.status.idle":"2026-01-17T09:55:31.617472Z","shell.execute_reply.started":"2026-01-17T09:55:31.612666Z","shell.execute_reply":"2026-01-17T09:55:31.61673Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!ls","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T09:55:31.618293Z","iopub.execute_input":"2026-01-17T09:55:31.618587Z","iopub.status.idle":"2026-01-17T09:55:31.755283Z","shell.execute_reply.started":"2026-01-17T09:55:31.618542Z","shell.execute_reply":"2026-01-17T09:55:31.754614Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"continue RNAPro\n---  \n\n","metadata":{}},{"cell_type":"code","source":"cd RNAPro","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T09:55:31.757752Z","iopub.execute_input":"2026-01-17T09:55:31.757973Z","iopub.status.idle":"2026-01-17T09:55:31.762893Z","shell.execute_reply.started":"2026-01-17T09:55:31.757947Z","shell.execute_reply":"2026-01-17T09:55:31.762296Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python preprocess/convert_templates_to_pt_files.py --input_csv /kaggle/working/submission_tbm.csv --output_name templates.pt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T09:55:31.763612Z","iopub.execute_input":"2026-01-17T09:55:31.763826Z","iopub.status.idle":"2026-01-17T09:55:38.052643Z","shell.execute_reply.started":"2026-01-17T09:55:31.763804Z","shell.execute_reply":"2026-01-17T09:55:38.051828Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DIST = \"/kaggle/working/RNAPro/release_data/ccd_cache/\"\n!mkdir -p $DIST","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T09:55:38.053901Z","iopub.execute_input":"2026-01-17T09:55:38.054221Z","iopub.status.idle":"2026-01-17T09:55:38.179877Z","shell.execute_reply.started":"2026-01-17T09:55:38.05418Z","shell.execute_reply":"2026-01-17T09:55:38.17901Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Recomputed the cdd cache using python preprocess/gen_ccd_cache.py\n# Exported as dataset here: https://www.kaggle.com/datasets/jaejohn/rnapro-ccd-cache\n\n# You will need the following packages to recompute it on our own\n\n# pip install gemmi\n# pip install pdbeccdutils","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T09:55:38.181137Z","iopub.execute_input":"2026-01-17T09:55:38.181798Z","iopub.status.idle":"2026-01-17T09:55:38.185236Z","shell.execute_reply.started":"2026-01-17T09:55:38.181765Z","shell.execute_reply":"2026-01-17T09:55:38.184637Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# updated file paths\n!cp /kaggle/input/rnapro-ccd-cache/ccd_cache/components.cif $DIST\n!cp /kaggle/input/rnapro-ccd-cache/ccd_cache/components.cif.rdkit_mol.pkl $DIST","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Inference  \n---","metadata":{}},{"cell_type":"code","source":"# %%python\nimport pandas as pd\ndf = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\")\nif not IS_SCORING_RUN:\n    df = df.head(5)\ndf.to_csv('/kaggle/working/sample_sequences.csv', index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T09:55:42.896241Z","iopub.execute_input":"2026-01-17T09:55:42.896847Z","iopub.status.idle":"2026-01-17T09:55:42.907083Z","shell.execute_reply.started":"2026-01-17T09:55:42.896814Z","shell.execute_reply":"2026-01-17T09:55:42.906502Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile runner/inference.py\n# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n# SPDX-License-Identifier: Apache-2.0\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport shutil\nimport logging\nimport traceback\nimport warnings\nimport argparse\nfrom contextlib import nullcontext\nfrom os.path import join as opjoin\nfrom typing import Any, Mapping\n\nimport json\nimport torch\nimport pandas as pd\nimport numpy as np\nfrom biotite.structure.io import pdbx\n\nfrom configs.configs_base import configs as configs_base\nfrom configs.configs_data import data_configs\nfrom configs.configs_inference import inference_configs\nfrom runner.dumper import DataDumper\n\nfrom rnapro.config import parse_sys_args\nfrom rnapro.config.config import ConfigManager, ArgumentNotSet\nfrom rnapro.data.infer_data_pipeline import get_inference_dataloader\nfrom rnapro.model.RNAPro import RNAPro\nfrom rnapro.utils.distributed import DIST_WRAPPER\nfrom rnapro.utils.seed import seed_everything\nfrom rnapro.utils.torch_utils import to_device\n\nwarnings.filterwarnings(\"ignore\", category=UserWarning)\nwarnings.filterwarnings(\"ignore\", category=FutureWarning)\nwarnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n\nlogger = logging.getLogger(__name__)\n\n# Silence all info logging\nlogging.basicConfig(level=logging.WARNING)\n# Silence dataloader logging specifically\nlogging.getLogger(\"rnapro.data\").setLevel(logging.WARNING)\nlogging.getLogger(\"rnapro\").setLevel(logging.WARNING)\n\n\ndef parse_configs(\n    configs: dict, arg_str: str = None, fill_required_with_null: bool = False\n):\n    \"\"\"\n    Parses and merges configuration settings from a dictionary and command-line arguments.\n\n    Args:\n        configs (dict): A dictionary containing initial configuration settings.\n        arg_str (str, optional): A string representing command-line arguments. Defaults to None.\n        fill_required_with_null (bool, optional):\n            A boolean flag indicating whether required values should be filled with `None` if not provided. Defaults to False.\n\n    Returns:\n        ConfigDict: The merged configuration dictionary.\n    \"\"\"\n    manager = ConfigManager(configs, fill_required_with_null=fill_required_with_null)\n    parser = argparse.ArgumentParser()\n\n    # This is new\n    parser.add_argument(\n        \"--max_len\",\n        type=int,\n        default=10000,\n        required=False,\n        help=\"Maximum length of the sequence. Longer sequences will be skipped during inference\"\n    )\n\n    # Register arguments\n    for key, (\n        dtype,\n        default_value,\n        allow_none,\n        required,\n    ) in manager.config_infos.items():\n        # All config use str type, strings will be converted to real dtype later\n        parser.add_argument(\n            \"--\" + key, type=str, default=ArgumentNotSet(), required=required\n        )\n    # Merge user commandline pargs with default ones\n    merged_configs = manager.merge_configs(\n        vars(parser.parse_args(arg_str.split())) if arg_str else {}\n    )\n\n    max_len = parser.parse_args(arg_str.split()).max_len\n    merged_configs.max_len = max_len\n\n    return merged_configs\n\n\nclass dotdict(dict):\n    __setattr__ = dict.__setitem__\n    __delattr__ = dict.__delitem__\n\n    def __getattr__(self, name):\n        try:\n            return self[name]\n        except KeyError:\n            raise AttributeError(name)\n\n\nclass InferenceRunner(object):\n    def __init__(self, configs: Any) -> None:\n        self.configs = configs\n        self.init_env()\n        self.init_basics()\n        self.init_model()\n        self.load_checkpoint()\n        self.init_dumper(\n            need_atom_confidence=configs.need_atom_confidence,\n            sorted_by_ranking_score=configs.sorted_by_ranking_score,\n        )\n\n    def init_env(self) -> None:\n        self.print(\n            f\"Distributed environment: world size: {DIST_WRAPPER.world_size}, \"\n            + f\"global rank: {DIST_WRAPPER.rank}, local rank: {DIST_WRAPPER.local_rank}\"\n        )\n        self.use_cuda = torch.cuda.device_count() > 0\n        if self.use_cuda:\n            self.device = torch.device(\"cuda:{}\".format(DIST_WRAPPER.local_rank))\n            os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n            all_gpu_ids = \",\".join(str(x) for x in range(torch.cuda.device_count()))\n            devices = os.getenv(\"CUDA_VISIBLE_DEVICES\", all_gpu_ids)\n            logging.info(\n                f\"LOCAL_RANK: {DIST_WRAPPER.local_rank} - CUDA_VISIBLE_DEVICES: [{devices}]\"\n            )\n            torch.cuda.set_device(self.device)\n        else:\n            self.device = torch.device(\"cpu\")\n        if self.configs.use_deepspeed_evo_attention:\n            env = os.getenv(\"CUTLASS_PATH\", None)\n            self.print(f\"env: {env}\")\n            assert (\n                env is not None\n            ), \"if use ds4sci, set `CUTLASS_PATH` env as https://www.deepspeed.ai/tutorials/ds4sci_evoformerattention/\"\n            if env is not None:\n                logging.info(\n                    \"The kernels will be compiled when DS4Sci_EvoformerAttention is called for the first time.\"\n                )\n        use_fastlayernorm = os.getenv(\"LAYERNORM_TYPE\", None)\n        if use_fastlayernorm == \"fast_layernorm\":\n            logging.info(\n                \"The kernels will be compiled when fast_layernorm is called for the first time.\"\n            )\n\n        logging.info(\"Finished init ENV.\")\n\n    def init_basics(self) -> None:\n        self.dump_dir = self.configs.dump_dir\n        self.error_dir = opjoin(self.dump_dir, \"ERR\")\n        os.makedirs(self.dump_dir, exist_ok=True)\n        os.makedirs(self.error_dir, exist_ok=True)\n\n    def init_model(self) -> None:\n        self.model = RNAPro(self.configs).to(self.device)\n        num_params = sum(p.numel() for p in self.model.parameters())\n        self.print(f\"Total number of parameters: {num_params:,}\")\n\n    def load_checkpoint(self) -> None:\n        checkpoint_path = self.configs.load_checkpoint_path\n\n        if not os.path.exists(checkpoint_path):\n            raise Exception(f\"Given checkpoint path not exist [{checkpoint_path}]\")\n        self.print(\n            f\"Loading from {checkpoint_path}, strict: {self.configs.load_strict}\"\n        )\n        checkpoint = torch.load(checkpoint_path, self.device)\n\n        sample_key = [k for k in checkpoint[\"model\"].keys()][0]\n        # self.print(f\"Sampled key: {sample_key}\")\n        if sample_key.startswith(\"module.\"):  # DDP checkpoint has module. prefix\n            checkpoint[\"model\"] = {\n                k[len(\"module.\"):]: v for k, v in checkpoint[\"model\"].items()\n            }\n        self.model.load_state_dict(\n            state_dict=checkpoint[\"model\"],\n            strict=True,\n        )\n        self.model.eval()\n        # self.print(\"Finish loading checkpoint.\")\n\n    def init_dumper(\n        self, need_atom_confidence: bool = False, sorted_by_ranking_score: bool = True\n    ):\n        self.dumper = DataDumper(\n            base_dir=self.dump_dir,\n            need_atom_confidence=need_atom_confidence,\n            sorted_by_ranking_score=sorted_by_ranking_score,\n        )\n\n    def print_dict(self, d):\n        for k, v in d.items():\n            if isinstance(v, torch.Tensor):\n                print(f\"{k}: \", v.shape)\n            else:\n                pass\n                # print(f\"{k}: {v}\")\n\n    # Adapted from runner.train.Trainer.evaluate\n    @torch.no_grad()\n    def predict(self, data: Mapping[str, Mapping[str, Any]]) -> dict[str, torch.Tensor]:\n        eval_precision = {\n            \"fp32\": torch.float32,\n            \"bf16\": torch.bfloat16,\n            \"fp16\": torch.float16,\n        }[self.configs.dtype]\n        # print(\"eval_precision: \", eval_precision)\n        enable_amp = (\n            torch.autocast(device_type=\"cuda\", dtype=eval_precision)\n            if torch.cuda.is_available()\n            else nullcontext()\n        )\n        #         print('input_feature_dict: ', self.print_dict(data[\"input_feature_dict\"]))\n        #         exit(0)\n\n        data = to_device(data, self.device)\n        with enable_amp:\n            prediction, _, _ = self.model(\n                input_feature_dict=data[\"input_feature_dict\"],\n                label_full_dict=None,\n                label_dict=None,\n                mode=\"inference\",\n            )\n\n        return prediction\n\n    def print(self, msg: str):\n        if DIST_WRAPPER.rank == 0:\n            # logger.info(msg)\n            print(msg)\n\n    def update_model_configs(self, new_configs: Any) -> None:\n        self.model.configs = new_configs\n\n\ndef update_inference_configs(configs: Any, N_token: int):\n    # Setting the default inference configs for different N_token and N_atom\n    # when N_token is larger than 3000, the default config might OOM even on a\n    # A100 80G GPUS,\n    if N_token > 3840:\n        configs.skip_amp.confidence_head = False\n        configs.skip_amp.sample_diffusion = False\n    elif N_token > 2560:\n        configs.skip_amp.confidence_head = False\n        configs.skip_amp.sample_diffusion = True\n    else:\n        configs.skip_amp.confidence_head = True\n        configs.skip_amp.sample_diffusion = True\n    return configs\n\n\ndef infer_predict(runner: InferenceRunner, configs: Any) -> None:\n    # Data\n    # logger.info(f\"Loading data from {configs.input_json_path}\")\n    try:\n        dataloader = get_inference_dataloader(configs=configs)\n    except Exception as e:\n        error_message = f\"{e}:\\n{traceback.format_exc()}\"\n        logger.info(error_message)\n        with open(opjoin(runner.error_dir, \"error.txt\"), \"a\") as f:\n            f.write(error_message)\n        return\n\n    num_data = len(dataloader.dataset)\n    for seed in configs.seeds:\n        seed_everything(seed=seed, deterministic=configs.deterministic)\n        for batch in dataloader:\n            try:\n                data, atom_array, data_error_message = batch[0]\n                sample_name = data[\"sample_name\"]\n\n                if len(data_error_message) > 0:\n                    logger.info(data_error_message)\n                    with open(opjoin(runner.error_dir, f\"{sample_name}.txt\"), \"a\") as f:\n                        f.write(data_error_message)\n                    continue\n\n                logger.info(\n                    (\n                        f\"[Rank {DIST_WRAPPER.rank} ({data['sample_index'] + 1}/{num_data})] {sample_name}: \"\n                        f\"N_asym {data['N_asym'].item()}, N_token {data['N_token'].item()}, \"\n                        f\"N_atom {data['N_atom'].item()}, N_msa {data['N_msa'].item()}\"\n                    )\n                )\n                new_configs = update_inference_configs(configs, data[\"N_token\"].item())\n                runner.update_model_configs(new_configs)\n                prediction = runner.predict(data)\n                runner.dumper.dump(\n                    dataset_name=\"\",\n                    pdb_id=sample_name,\n                    seed=seed,\n                    pred_dict=prediction,\n                    atom_array=atom_array,\n                    entity_poly_type=data[\"entity_poly_type\"],\n                )\n\n                logger.info(\n                    f\"[Rank {DIST_WRAPPER.rank}] {data['sample_name']} succeeded - \"\n                    f\"Results saved to {configs.dump_dir}\"\n                )\n                torch.cuda.empty_cache()\n            except Exception as e:\n                error_message = f\"[Rank {DIST_WRAPPER.rank}]{data['sample_name']} {e}:\\n{traceback.format_exc()}\"\n                logger.info(error_message)\n                # Save error info\n                with open(opjoin(runner.error_dir, f\"{sample_name}.txt\"), \"a\") as f:\n                    f.write(error_message)\n                if hasattr(torch.cuda, \"empty_cache\"):\n                    torch.cuda.empty_cache()\n\n\n# data helper\ndef make_dummy_solution(valid_df):\n    solution = dotdict()\n    for i, row in valid_df.iterrows():\n        target_id = row.target_id\n        sequence = row.sequence\n        solution[target_id] = dotdict(\n            target_id=target_id,\n            sequence=sequence,\n            coord=[],\n        )\n    return solution\n\n\ndef solution_to_submit_df(solution):\n    submit_df = []\n    for k, s in solution.items():\n        df = coord_to_df(s.sequence, s.coord, s.target_id)\n        submit_df.append(df)\n\n    submit_df = pd.concat(submit_df)\n    return submit_df\n\n\ndef coord_to_df(sequence, coord, target_id):\n    L = len(sequence)\n    df = pd.DataFrame()\n    df[\"ID\"] = [f\"{target_id}_{i + 1}\" for i in range(L)]\n    df[\"resname\"] = [s for s in sequence]\n    df[\"resid\"] = [i + 1 for i in range(L)]\n\n    num_coord = len(coord)\n    for j in range(num_coord):\n        df[f\"x_{j+1}\"] = coord[j][:, 0]\n        df[f\"y_{j+1}\"] = coord[j][:, 1]\n        df[f\"z_{j+1}\"] = coord[j][:, 2]\n    return df\n\n\ndef main(configs: Any) -> None:\n    # Runner\n    runner = InferenceRunner(configs)\n    infer_predict(runner, configs)\n\n\ndef create_input_json(sequence, target_id):\n    # print(\"input_no_msa\")\n    input_json = [\n        {\n            \"sequences\": [\n                {\n                    \"rnaSequence\": {\n                        \"sequence\": sequence,\n                        \"count\": 1,\n                    }\n                }\n            ],\n            \"name\": target_id,\n        }\n    ]\n    return input_json\n\n\ndef extract_c1_coordinates(cif_file_path):\n    try:\n        # Read the CIF file using the correct biotite method\n        with open(cif_file_path, \"r\") as f:\n            cif_data = pdbx.CIFFile.read(f)\n\n        # Get structure from CIF data\n        atom_array = pdbx.get_structure(cif_data, model=1)\n\n        # Clean atom names and find C1' atoms\n        atom_names_clean = np.char.strip(atom_array.atom_name.astype(str))\n        mask_c1 = atom_names_clean == \"C1'\"\n        c1_atoms = atom_array[mask_c1]\n\n        if len(c1_atoms) == 0:\n            print(f\"Warning: No C1' atoms found in {cif_file_path}\")\n            return None\n\n        # Sort by residue ID and return coordinates\n        sort_indices = np.argsort(c1_atoms.res_id)\n        c1_atoms_sorted = c1_atoms[sort_indices]\n        c1_coords = c1_atoms_sorted.coord\n\n        return c1_coords\n    except Exception as e:\n        print(f\"Error extracting C1' coordinates from {cif_file_path}: {e}\")\n        return None\n\n\ndef process_sequence(sequence, target_id, temp_dir):\n    # Create input JSON\n    input_json = create_input_json(sequence, target_id)\n\n    # Save JSON to temporary file\n    os.makedirs(temp_dir, exist_ok=True)\n    input_json_path = os.path.join(temp_dir, f\"{target_id}_input.json\")\n    with open(input_json_path, \"w\") as f:\n        json.dump(input_json, f, indent=4)\n\n\ndef run_ptx(target_id, sequence, configs, solution, template_idx, runner):\n    # Create directories\n    temp_dir = f\"./{configs.dump_dir}/input\"  # Same as in kaggle_inference.py\n    output_dir = f\"./{configs.dump_dir}/output\"  # Same as in kaggle_inference.py\n    os.makedirs(temp_dir, exist_ok=True)\n    os.makedirs(output_dir, exist_ok=True)\n\n    process_sequence(sequence=sequence, target_id=target_id, temp_dir=temp_dir)\n    configs.input_json_path = os.path.join(temp_dir, f\"{target_id}_input.json\")\n    configs.template_idx = int(template_idx)\n\n    infer_predict(runner, configs)\n\n    cif_file_path = (\n        f\"{configs.dump_dir}/{target_id}/seed_42/predictions/{target_id}_sample_0.cif\"\n    )\n    cif_new_path = f\"{configs.dump_dir}/{target_id}/seed_42/predictions/{target_id}_sample_{template_idx}_new.cif\"\n    shutil.copy(cif_file_path, cif_new_path)\n    coord = extract_c1_coordinates(cif_file_path)\n    if coord is None:\n        coord = np.zeros((len(sequence), 3), dtype=np.float32)\n    elif coord.shape[0] < (len(sequence)):\n        pad_len = len(sequence) - coord.shape[0]\n        pad = np.zeros((pad_len, 3), dtype=np.float32)\n        coord = np.concatenate([coord, pad], axis=0)\n    solution[target_id].coord.append(coord)\n\n\ndef run() -> None:\n    LOG_FORMAT = \"%(asctime)s,%(msecs)-3d %(levelname)-8s [%(filename)s:%(lineno)s %(funcName)s] %(message)s\"\n    logging.basicConfig(\n        format=LOG_FORMAT,\n        level=logging.WARNING,\n        datefmt=\"%Y-%m-%d %H:%M:%S\",\n        filemode=\"w\",\n    )\n    # Silence dataloader and rnapro module logging\n    logging.getLogger(\"rnapro.data\").setLevel(logging.WARNING)\n    logging.getLogger(\"rnapro\").setLevel(logging.WARNING)\n    configs_base[\"use_deepspeed_evo_attention\"] = (\n        os.environ.get(\"USE_DEEPSPEED_EVO_ATTENTION\", False) == \"true\"\n    )\n    configs = {**configs_base, **{\"data\": data_configs}, **inference_configs}\n    configs = parse_configs(\n        configs=configs,\n        arg_str=parse_sys_args(),\n        fill_required_with_null=True,\n    )\n\n    valid_df = pd.read_csv(configs.sequences_csv)\n    print(f\"\\n -> Loaded {len(valid_df)} sequence(s)\")\n\n    # Build model and load checkpoint once before looping over sequences\n\n    print('\\n -> Building model and loading checkpoint')\n    runner = InferenceRunner(configs)\n    print('\\n -> Done, starting inference...')\n\n    solution = make_dummy_solution(valid_df)\n    for idx, row in valid_df.iterrows():\n        print(f\"\\n -> Sequence {row.target_id}: {row.sequence}\")\n\n        if len(row.sequence) > configs.max_len:\n            print(f'Sequence is too long ({len(row.sequence)} > {configs.max_len}), skipping')\n            for template_idx in range(5):\n                coord = np.zeros((len(row.sequence), 3), dtype=np.float32)\n                solution[row.target_id].coord.append(coord)\n            continue\n\n        try:\n            target_id = row.target_id\n            sequence = row.sequence\n            for template_idx in range(5):\n                print()\n                run_ptx(\n                    target_id=target_id,\n                    sequence=sequence,\n                    configs=configs,\n                    solution=solution,\n                    template_idx=template_idx,\n                    runner=runner,\n                )\n        except Exception as e:\n            print(f\"Error processing {row.target_id}: {e}\")\n            continue\n\n    print('\\n\\n -> Inference done ! Saving to submission.csv')\n    submit_df = solution_to_submit_df(solution)\n    submit_df = submit_df.fillna(0.0)\n    submit_df.to_csv(\"./submission.csv\", index=False)\n\n\nif __name__ == \"__main__\":\n    run()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T09:56:08.738357Z","iopub.execute_input":"2026-01-17T09:56:08.738681Z","iopub.status.idle":"2026-01-17T09:56:08.75229Z","shell.execute_reply.started":"2026-01-17T09:56:08.738652Z","shell.execute_reply":"2026-01-17T09:56:08.751601Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile rnapro_inference_kaggle.sh\n\nexport LAYERNORM_TYPE=torch # fast_layernorm, torch\n\n\n# Inference parameters (RNAPro)\nSEED=42\nN_SAMPLE=1\nN_STEP=200\nN_CYCLE=10\n\n# Paths\nDUMP_DIR=\"../output\"\n# Set a valid checkpoint file path below\nCHECKPOINT_PATH=\"../rnapro-private-best-500m.ckpt\"\n\n# Template/MSA settings\nTEMPLATE_DATA=\"./release_data/kaggle/templates.pt\"\n# Note: template_idx supports 5 choices and maps to top-k:\n# 0->top1, 1->top2, 2->top3, 3->top4, 4->top5\nTEMPLATE_IDX=0\nRNA_MSA_DIR=\"/kaggle/input/stanford-rna-3d-folding-2/MSA\"\n\n# SEQUENCES_CSV=\"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\"\nSEQUENCES_CSV=\"/kaggle/working/sample_sequences.csv\"\n\n# RibonanzaNet2 path (keep as-is per request)\nRIBONANZA_PATH=\"/kaggle/input/ribonanzanet2/pytorch/alpha/1/\"\n\n# Model selection: keep to an existing key to align defaults (N_step=200, N_cycle=10)\nMODEL_NAME=\"rnapro_base\"\nmkdir -p \"${DUMP_DIR}\"\n\npython3 runner/inference.py \\\n    --model_name \"${MODEL_NAME}\" \\\n    --seeds ${SEED} \\\n    --dump_dir \"${DUMP_DIR}\" \\\n    --load_checkpoint_path \"${CHECKPOINT_PATH}\" \\\n    --use_msa true \\\n    --use_template \"ca_precomputed\" \\\n    --model.use_template \"ca_precomputed\" \\\n    --model.use_RibonanzaNet2 true \\\n    --model.template_embedder.n_blocks 2 \\\n    --model.ribonanza_net_path \"${RIBONANZA_PATH}\" \\\n    --template_data \"${TEMPLATE_DATA}\" \\\n    --template_idx ${TEMPLATE_IDX} \\\n    --rna_msa_dir \"${RNA_MSA_DIR}\" \\\n    --model.N_cycle ${N_CYCLE} \\\n    --sample_diffusion.N_sample ${N_SAMPLE} \\\n    --sample_diffusion.N_step ${N_STEP} \\\n    --load_strict true \\\n    --num_workers 0 \\\n    --triangle_attention \"torch\" \\\n    --triangle_multiplicative \"torch\" \\\n    --sequences_csv \"${SEQUENCES_CSV}\" \\\n    --max_len 1000\n\n\n# --triangle_attention supports 'triattention', 'cuequivariance', 'deepspeed', 'torch'\n# --triangle_multiplicative supports 'cuequivariance', 'torch'\n# --max_len 1000: Sequences longer than max_len will be skipped to avoid oom","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T09:56:09.76588Z","iopub.execute_input":"2026-01-17T09:56:09.766383Z","iopub.status.idle":"2026-01-17T09:56:09.771668Z","shell.execute_reply.started":"2026-01-17T09:56:09.766357Z","shell.execute_reply":"2026-01-17T09:56:09.771072Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!bash ./rnapro_inference_kaggle.sh","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T09:56:10.466806Z","iopub.execute_input":"2026-01-17T09:56:10.467452Z","iopub.status.idle":"2026-01-17T10:07:51.562684Z","shell.execute_reply.started":"2026-01-17T09:56:10.467427Z","shell.execute_reply":"2026-01-17T10:07:51.56185Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!mv submission.csv ..","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T10:07:51.564722Z","iopub.execute_input":"2026-01-17T10:07:51.565047Z","iopub.status.idle":"2026-01-17T10:07:51.689405Z","shell.execute_reply.started":"2026-01-17T10:07:51.56502Z","shell.execute_reply":"2026-01-17T10:07:51.688709Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cd ..","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T10:07:51.69073Z","iopub.execute_input":"2026-01-17T10:07:51.690997Z","iopub.status.idle":"2026-01-17T10:07:51.697019Z","shell.execute_reply.started":"2026-01-17T10:07:51.690969Z","shell.execute_reply":"2026-01-17T10:07:51.696313Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!head submission.csv","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T10:07:51.698544Z","iopub.execute_input":"2026-01-17T10:07:51.698863Z","iopub.status.idle":"2026-01-17T10:07:51.828618Z","shell.execute_reply.started":"2026-01-17T10:07:51.698832Z","shell.execute_reply":"2026-01-17T10:07:51.827969Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nsub = pd.read_csv(\"/kaggle/working/submission.csv\")\nsub","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T10:07:51.829715Z","iopub.execute_input":"2026-01-17T10:07:51.82992Z","iopub.status.idle":"2026-01-17T10:07:51.873465Z","shell.execute_reply.started":"2026-01-17T10:07:51.829897Z","shell.execute_reply":"2026-01-17T10:07:51.872886Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\n\ndf_tbm = pd.read_csv(\"submission_tbm.csv\")\ndf_rnapro = pd.read_csv(\"submission.csv\")\ndf_seqs = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\")\nlong_targets = df_seqs[df_seqs['sequence'].str.len() > 1000]['target_id'].values\n\nprint(f\"Targets to replace with TBM (len > 1000): {long_targets}\")\n\nmask_long = df_rnapro['ID'].apply(lambda x: any(str(x).startswith(t + \"_\") for t in long_targets))\n\nif mask_long.sum() > 0:\n    print(f\"Replacing {mask_long.sum()} residues with TBM predictions...\")\n    \n    # Set index to ID for easy alignment\n    df_rnapro_idx = df_rnapro.set_index('ID')\n    df_tbm_idx = df_tbm.set_index('ID')\n    \n    # Update rows in RNAPro df with TBM df for the specific IDs\n    # This works if indices match\n    ids_to_update = df_rnapro_idx[mask_long.values].index\n    \n    # Check if these IDs exist in TBM file\n    valid_ids = [i for i in ids_to_update if i in df_tbm_idx.index]\n    \n    df_rnapro_idx.loc[valid_ids] = df_tbm_idx.loc[valid_ids]\n    \n    # Reset index to get ID column back\n    df_final = df_rnapro_idx.reset_index()\n    \n    # Save merged\n    df_final.to_csv(\"submission.csv\", index=False)\n    print(\"Merged submission saved to submission.csv\")\nelse:\n    print(\"No long targets found to replace. Keeping RNAPro submission as is.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sub = pd.read_csv(\"/kaggle/working/submission.csv\")\nsub","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Evaluation  \n---","metadata":{}},{"cell_type":"code","source":"%%writefile /kaggle/working/metric.py\n#!/usr/bin/env python\n# coding: utf-8\n\n# In[ ]:\n\n\nimport os\nimport re\nimport math\nimport pandas as pd\nfrom pathlib import Path\nimport shutil\nimport sys\nimport csv\n\n# ---------------------\n# Helper: parse USalign output\n# ---------------------\ndef parse_tmscore_output(output: str) -> float:\n    matches = re.findall(r'TM-score=\\s+([\\d.]+)', output)\n    if len(matches) < 2:\n        raise ValueError('No TM score found in USalign output')\n    return float(matches[1])\n\n# ---------------------\n# PDB writers\n# ---------------------\n\ndef sanitize(xyz):\n    MIN_COORD=-999.999\n    MAX_COORD=9999.999\n    return min(max(xyz,MIN_COORD),MAX_COORD)\n\ndef write_target_line(atom_name, atom_serial, residue_name, chain_id, residue_num,\n                      x_coord, y_coord, z_coord, occupancy=1.0, b_factor=0.0, atom_type='P') -> str:\n    return f'ATOM  {atom_serial:>5d}  {atom_name:4s}{residue_name:>3s} {chain_id:1s}{residue_num:>4d}    {sanitize(x_coord):>8.3f}{sanitize(y_coord):>8.3f}{sanitize(z_coord):>8.3f}{occupancy:>6.2f}{b_factor:>6.2f}           {atom_type}\\n'\n    \ndef write2pdb(df: pd.DataFrame, xyz_id: int, target_path: str) -> int:\n    \"\"\"\n    Write single-chain PDB (chain 'A') using row['resid'] as residue_num.\n    Raises exceptions on invalid data.\n    \"\"\"\n    resolved_cnt = 0\n    with open(target_path, 'w') as fh:\n        for _, row in df.iterrows():\n            x = row[f'x_{xyz_id}']\n            y = row[f'y_{xyz_id}']\n            z = row[f'z_{xyz_id}']\n            if x > -1e6 and y > -1e6 and z > -1e6:\n                resolved_cnt += 1\n                resid_num = int(row['resid'])\n                fh.write(write_target_line(\"C1'\", resid_num, row['resname'], 'A', resid_num, x, y, z, atom_type='C'))\n    return resolved_cnt\n\ndef write2pdb_singlechain_native(df_native: pd.DataFrame, xyz_id: int, target_path: str) -> int:\n    \"\"\"\n    Write native single-chain using row['resid'] as residue numbers.\n    Assumes all required columns exist and are valid.\n    \"\"\"\n    df_sorted = df_native.copy()\n    df_sorted['__resid_int'] = df_sorted['resid'].astype(int)\n    df_sorted = df_sorted.sort_values('__resid_int').reset_index(drop=True)\n\n    resolved_cnt = 0\n    with open(target_path, 'w') as fh:\n        for _, row in df_sorted.iterrows():\n            x = row[f'x_{xyz_id}']\n            y = row[f'y_{xyz_id}']\n            z = row[f'z_{xyz_id}']\n            if x > -1e6 and y > -1e6 and z > -1e6:\n                resolved_cnt += 1\n                resid_num = int(row['resid'])\n                fh.write(write_target_line(\"C1'\", resid_num, row['resname'], 'A', resid_num, x, y, z, atom_type='C'))\n    return resolved_cnt\n\ndef write2pdb_multichain_from_solution(df_solution: pd.DataFrame, xyz_id: int, target_path: str) -> int:\n    \"\"\"\n    Write multi-chain PDB for native solution using columns 'chain' and 'copy' to assign chain letters.\n    Expects 'resid' convertible to int and chain/copy present. No fallbacks.\n    \"\"\"\n    df_sorted = df_solution.copy()\n    df_sorted['__resid_int'] = df_sorted['resid'].astype(int)\n    df_sorted = df_sorted.sort_values('__resid_int')\n\n    chain_map = {}\n    next_ord = ord('A')\n    written = 0\n    with open(target_path, 'w') as fh:\n        for _, row in df_sorted.iterrows():\n            x = row[f'x_{xyz_id}']\n            y = row[f'y_{xyz_id}']\n            z = row[f'z_{xyz_id}']\n            if not (x > -1e6 and y > -1e6 and z > -1e6):\n                continue\n            chain_val = row['chain']\n            copy_key = int(row['copy'])\n            g = (str(chain_val), copy_key)\n            if g not in chain_map:\n                if next_ord <= ord('Z'):\n                    ch = chr(next_ord)\n                else:\n                    ov = next_ord - ord('Z') - 1\n                    if ov < 26:\n                        ch = chr(ord('a') + ov)\n                    else:\n                        ch = chr(ord('0') + (ov - 26) % 10)\n                chain_map[g] = ch\n                next_ord += 1\n            chain_id = chain_map[g]\n            written += 1\n            resid_num = int(row['resid'])\n            fh.write(write_target_line(\"C1'\", resid_num, row['resname'], chain_id, resid_num, x, y, z, atom_type='C'))\n    return written\n\ndef write2pdb_multichain_from_groups(df_pred: pd.DataFrame, xyz_id: int, target_path: str, groups_list) -> (int, list):\n    \"\"\"\n    Write predicted multichain PDB based on a positional groups_list (tuple per residue: (chain, copy)).\n    Requires groups_list length == number of residues in df_pred (after sorting).\n    Returns (written_count, chain_letters_per_res).\n    \"\"\"\n    df_sorted = df_pred.copy()\n    df_sorted['__resid_int'] = df_sorted['resid'].astype(int)\n    df_sorted = df_sorted.sort_values('__resid_int').reset_index(drop=True)\n\n    if groups_list is None or len(groups_list) != len(df_sorted):\n        raise ValueError(\"groups_list must be provided and match number of residues in predicted df\")\n\n    chain_map = {}\n    next_ord = ord('A')\n    chain_letters = []\n    written = 0\n    with open(target_path, 'w') as fh:\n        for idx, row in df_sorted.iterrows():\n            g = groups_list[idx]\n            if isinstance(g, tuple):\n                gkey = (str(g[0]), int(g[1]))\n            else:\n                gkey = (str(g), None)\n            if gkey not in chain_map:\n                if next_ord <= ord('Z'):\n                    ch = chr(next_ord)\n                else:\n                    ov = next_ord - ord('Z') - 1\n                    if ov < 26:\n                        ch = chr(ord('a') + ov)\n                    else:\n                        ch = chr(ord('0') + (ov - 26) % 10)\n                chain_map[gkey] = ch\n                next_ord += 1\n            chain_id = chain_map[gkey]\n            chain_letters.append(chain_id)\n            x = row[f'x_{xyz_id}']\n            y = row[f'y_{xyz_id}']\n            z = row[f'z_{xyz_id}']\n            if x > -1e6 and y > -1e6 and z > -1e6:\n                written += 1\n                resid_num = int(row['resid'])\n                fh.write(write_target_line(\"C1'\", resid_num, row['resname'], chain_id, resid_num, x, y, z, atom_type='C'))\n    return written, chain_letters\n\ndef write2pdb_singlechain_permuted_pred(df_pred: pd.DataFrame, xyz_id: int, permuted_indices: list, target_path: str) -> int:\n    \"\"\"\n    Create single-chain PDB by concatenating predicted residues in permuted_indices order.\n    Output residue numbers are sequential starting at 1 and increase for every permuted position.\n    Raises exception if indices out of range.\n    \"\"\"\n    df_sorted = df_pred.copy()\n    df_sorted['__resid_int'] = df_sorted['resid'].astype(int)\n    df_sorted = df_sorted.sort_values('__resid_int').reset_index(drop=True)\n\n    written = 0\n    next_res = 1\n    with open(target_path, 'w') as fh:\n        for idx in permuted_indices:\n            if idx < 0 or idx >= len(df_sorted):\n                # strict behavior: raise error for invalid index\n                raise IndexError(f\"permuted index {idx} out of range for predicted residues\")\n            row = df_sorted.iloc[idx]\n            x = row[f'x_{xyz_id}']\n            y = row[f'y_{xyz_id}']\n            z = row[f'z_{xyz_id}']\n            out_resnum = next_res\n            if x > -1e6 and y > -1e6 and z > -1e6:\n                written += 1\n                fh.write(write_target_line(\"C1'\", out_resnum, row['resname'], 'A', out_resnum, x, y, z, atom_type='C'))\n            next_res += 1\n    return written\n\n# ---------------------\n# USalign wrappers\n# ---------------------\ndef run_usalign_raw(predicted_pdb: str, native_pdb: str, usalign_bin='USalign', align_sequence=False, tmscore=None) -> str:\n    cmd = f'{usalign_bin} {predicted_pdb} {native_pdb} -atom \" C1\\'\"'\n    if tmscore is not None:\n        cmd += f' -TMscore {tmscore}'\n        if int(tmscore) == 0:\n            cmd += ' -mm 1 -ter 0'\n    elif not align_sequence:\n        cmd += ' -TMscore 1'\n    return os.popen(cmd).read()\n\ndef parse_usalign_chain_orders(output: str):\n    \"\"\"\n    Parse USalign output for both Structure_1 and Structure_2 chain lists.\n    Returns (chain_list_structure1, chain_list_structure2).\n    Raises if parsing fails to find either line.\n    \"\"\"\n    chain1 = None\n    chain2 = None\n    for line in output.splitlines():\n        line = line.strip()\n        if line.startswith('Name of Structure_1:'):\n            parts = line.split(':')\n            clist = []\n            for part in parts[2:]:\n                token = part.strip()\n                if token == '':\n                    continue\n                token0 = token.split()[0]\n                last = token0.split(',')[-1]\n                ch = re.sub(r'[^A-Za-z0-9]', '', last)\n                if ch:\n                    clist.append(ch)\n            chain1 = clist\n        elif line.startswith('Name of Structure_2:'):\n            parts = line.split(':')\n            clist = []\n            for part in parts[2:]:\n                token = part.strip()\n                if token == '':\n                    continue\n                token0 = token.split()[0]\n                last = token0.split(',')[-1]\n                ch = re.sub(r'[^A-Za-z0-9]', '', last)\n                if ch:\n                    clist.append(ch)\n            chain2 = clist\n    if chain1 is None or chain2 is None:\n        raise ValueError(\"Failed to parse chain orders from USalign output\")\n    return chain1, chain2\n\n# ---------------------\n# Main scoring function (no try/except, no fallbacks)\n# ---------------------\ndef score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str, usalign_bin_hint: str = None) -> float:\n    \"\"\"\n    Enhanced scoring with chain-permutation handling for multicopy targets.\n    This version contains no try/except blocks and will raise on any error.\n    \"\"\"\n    # determine usalign binary\n    if usalign_bin_hint:\n        usalign_bin = usalign_bin_hint\n    else:\n        if os.path.exists('/kaggle/input/usalign/USalign') and not os.path.exists('/kaggle/working/USalign'):\n            shutil.copy2('/kaggle/input/usalign/USalign', '/kaggle/working/USalign')\n            os.chmod('/kaggle/working/USalign', 0o755)\n        usalign_bin = '/kaggle/working/USalign' if os.path.exists('/kaggle/working/USalign') else 'USalign'\n\n    sol = solution.copy()\n    sub = submission.copy()\n    sol['target_id'] = sol['ID'].apply(lambda x: '_'.join(str(x).split('_')[:-1]))\n    sub['target_id'] = sub['ID'].apply(lambda x: '_'.join(str(x).split('_')[:-1]))\n\n    results = []\n\n    for target_id, group_native in sol.groupby('target_id'):\n        group_predicted = sub[sub['target_id'] == target_id]\n        has_chain_copy = ('chain' in group_native.columns) and ('copy' in group_native.columns)\n        is_multicopy = has_chain_copy and (group_native['copy'].astype(float).max() > 1)\n\n        # precompute native models that have coords\n        native_with_coords = []\n        for native_cnt in range(1, 41):\n            native_pdb = f'native_{target_id}_{native_cnt}.pdb'\n            resolved_native = write2pdb(group_native, native_cnt, native_pdb)\n            if resolved_native > 0:\n                native_with_coords.append(native_cnt)\n            else:\n                if os.path.exists(native_pdb):\n                    os.remove(native_pdb)\n\n        if not native_with_coords:\n            raise ValueError(f\"No native models with coordinates for target {target_id}\")\n\n        best_per_pred = []\n        for pred_cnt in range(1, 6):\n            if not is_multicopy:\n                predicted_pdb = f'predicted_{target_id}_{pred_cnt}.pdb'\n                resolved_pred = write2pdb(group_predicted, pred_cnt, predicted_pdb)\n                if resolved_pred <= 2:\n                    #print(f\"Predicted model {pred_cnt} for target {target_id} has insufficient coordinates\")\n                    best_per_pred.append( 0.0 )\n                    continue\n                \n                scores = []\n                for native_cnt in native_with_coords:\n                    native_pdb = f'native_{target_id}_{native_cnt}.pdb'\n                    out = run_usalign_raw(predicted_pdb, native_pdb, usalign_bin=usalign_bin, align_sequence=False, tmscore=1)\n                    s = parse_tmscore_output(out)\n                    scores.append(s)\n                best_per_pred.append(max(scores))\n\n            else:\n                # multicopy\n                # strict: require chain and copy columns convertible\n                gn_sorted = group_native.copy()\n                gn_sorted['__resid_int'] = gn_sorted['resid'].astype(int)\n                gn_sorted = gn_sorted.sort_values('__resid_int').reset_index(drop=True)\n                groups_list = []\n                for _, r in gn_sorted.iterrows():\n                    chain_val = r['chain']\n                    copy_i = int(r['copy'])\n                    groups_list.append((chain_val, copy_i))\n\n                # predicted multichain - groups_list must match predicted residue count or error\n                dfp_sorted = group_predicted.copy()\n                dfp_sorted['__resid_int'] = dfp_sorted['resid'].astype(int)\n                dfp_sorted = dfp_sorted.sort_values('__resid_int').reset_index(drop=True)\n                if len(groups_list) != len(dfp_sorted):\n                    raise ValueError(f\"groups_list length ({len(groups_list)}) does not match predicted residue count ({len(dfp_sorted)}) for target {target_id}\")\n\n                predicted_multi_pdb = f'pred_multi_{target_id}_{pred_cnt}.pdb'\n                resolved_pred_multi, pred_chain_letters = write2pdb_multichain_from_groups(group_predicted, pred_cnt, predicted_multi_pdb, groups_list)\n                if resolved_pred_multi == 0:\n                    #print(f\"Predicted multi model {pred_cnt} for target {target_id} has no coordinates\")\n                    best_per_pred.append( 0.0 )\n                    continue\n\n                scores = []\n                for native_cnt in native_with_coords:\n                    native_multi_pdb = f'native_multi_{target_id}_{native_cnt}.pdb'\n                    resolved_native_multi = write2pdb_multichain_from_solution(group_native, native_cnt, native_multi_pdb)\n                    if resolved_native_multi == 0:\n                        continue\n\n                    raw_out = run_usalign_raw(predicted_multi_pdb, native_multi_pdb, usalign_bin=usalign_bin, align_sequence=True, tmscore=0)\n                    chain1, chain2 = parse_usalign_chain_orders(raw_out)  # will raise if parsing fails\n\n                    # build native->pred mapping chain2[i] -> chain1[i]\n                    native_to_pred = {n_ch: p_ch for n_ch, p_ch in zip(chain2, chain1)}\n\n                    # canonical native order = chain2 unique in order seen\n                    #native_chain_order = []\n                    #for ch in chain2:\n                    #    if ch not in native_chain_order:\n                    #        native_chain_order.append(ch)\n                    native_chain_order = list(native_to_pred.keys())\n                    native_chain_order.sort() # this is critical...\n\n                    # predicted chain order by following native chain A,B,...\n                    pred_chain_order = [native_to_pred[n_ch] for n_ch in native_chain_order if native_to_pred.get(n_ch) is not None]\n\n                    # construct pred_positions_by_chain\n                    pred_positions_by_chain = {}\n                    for idx, ch in enumerate(pred_chain_letters):\n                        if ch is None:\n                            continue\n                        pred_positions_by_chain.setdefault(ch, []).append(idx)\n\n                    # require that each chain in pred_chain_order exists in pred_positions_by_chain\n                    pred_chain_order = [p for p in pred_chain_order if p in pred_positions_by_chain]\n\n                    # form permuted indices by concatenation\n                    permuted_indices = []\n                    for ch in pred_chain_order:\n                        permuted_indices.extend(pred_positions_by_chain[ch])\n                    # append any remaining\n                    for idx in range(len(pred_chain_letters)):\n                        if idx not in permuted_indices:\n                            permuted_indices.append(idx)\n\n                    # write permuted single-chain predicted and native single-chain\n                    pred_single_perm = f'pred_permuted_{target_id}_{pred_cnt}_{native_cnt}.pdb'\n                    written_pred_single = write2pdb_singlechain_permuted_pred(group_predicted, pred_cnt, permuted_indices, pred_single_perm)\n                    native_single = f'native_single_{target_id}_{native_cnt}.pdb'\n                    written_native = write2pdb_singlechain_native(group_native, native_cnt, native_single)\n\n                    if written_pred_single <= 2 or written_native <= 2:\n                        raise ValueError(f\"Insufficient residues after permutation for target {target_id}, pred {pred_cnt}, native {native_cnt}\")\n\n                    out = run_usalign_raw(pred_single_perm, native_single, usalign_bin=usalign_bin, align_sequence=False, tmscore=1)\n                    score_final = parse_tmscore_output(out)\n                    scores.append(score_final)\n\n                best_per_pred.append(max(scores))\n\n        results.append(max(best_per_pred))\n\n    if not results:\n        pass\n        #raise ValueError(\"No targets scored\")\n    return float(sum(results) / len(results)) if len(results)>0 else 0.0","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2026-01-17T10:07:57.961773Z","iopub.execute_input":"2026-01-17T10:07:57.962482Z","iopub.status.idle":"2026-01-17T10:07:57.974834Z","shell.execute_reply.started":"2026-01-17T10:07:57.962454Z","shell.execute_reply":"2026-01-17T10:07:57.97408Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import runpy\nmodule_globals = runpy.run_path(\"/kaggle/working/metric.py\")\nscore = module_globals['score']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T10:07:58.040366Z","iopub.execute_input":"2026-01-17T10:07:58.040775Z","iopub.status.idle":"2026-01-17T10:07:58.048914Z","shell.execute_reply.started":"2026-01-17T10:07:58.040754Z","shell.execute_reply":"2026-01-17T10:07:58.04822Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not IS_SCORING_RUN:\n    import pandas as pd\n    sub = pd.read_csv('/kaggle/working/submission.csv')\n    sol = pd.read_csv('/kaggle/input/stanford-rna-3d-folding-2/validation_labels.csv')\n\n    sub['target_id'] = sub['ID'].apply(lambda x: '_'.join(str(x).split('_')[:-1]))\n    sol['target_id'] = sol['ID'].apply(lambda x: '_'.join(str(x).split('_')[:-1]))\n    \n    # Get unique targets from submission\n    sub_targets = sub['target_id'].unique()\n    \n    results = []\n    for target_id in sub_targets:\n        group_native = sol[sol['target_id'] == target_id]\n        group_predicted = sub[sub['target_id'] == target_id]\n        result = score(group_native, group_predicted, 'ID')\n        print(f\"{target_id}: {result:.4f}\")\n        results.append(result)\n    \n    print(f\"\\nMean score: {sum(results)/len(results):.4f} (n={len(results)})\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-17T10:07:59.696325Z","iopub.execute_input":"2026-01-17T10:07:59.696975Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}