{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":118765,"databundleVersionId":15231210,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":11118830,"sourceType":"datasetVersion","datasetId":6933267},{"sourceId":11969392,"sourceType":"datasetVersion","datasetId":7526656},{"sourceId":14519720,"sourceType":"datasetVersion","datasetId":9271415},{"sourceId":291833135,"sourceType":"kernelVersion"},{"sourceId":292115982,"sourceType":"kernelVersion"},{"sourceId":311741,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":264400,"modelId":285488}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":3061.935344,"end_time":"2026-01-15T23:04:47.93547","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2026-01-15T22:13:46.000126","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 🧬RNAPro: An accurate RNA structure prediction model by Kaggle synthesis\n\n**About:**\nThis notebooks shows how to use RNAPro for offline inference\n\n⭐ The code is available on GitHub ⭐\n> https://github.com/NVIDIA-Digital-Bio/RNAPro","metadata":{"papermill":{"duration":0.006278,"end_time":"2026-01-15T22:13:48.680723","exception":false,"start_time":"2026-01-15T22:13:48.674445","status":"completed"},"tags":[]}},{"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-21T10:15:07.72448Z","iopub.execute_input":"2026-01-21T10:15:07.724753Z","iopub.status.idle":"2026-01-21T10:15:07.731989Z","shell.execute_reply.started":"2026-01-21T10:15:07.724733Z","shell.execute_reply":"2026-01-21T10:15:07.731336Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%python\nimport os\nimport pandas as pd\ndf = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\")\n\nif not os.environ.get('KAGGLE_IS_COMPETITION_RERUN', False):\n    df = df.head(3)\n\ndf.to_csv('/kaggle/working/sequences.csv', index=False)\nprint(f'-> Using {len(df)} sequences for inference')","metadata":{"execution":{"iopub.status.busy":"2026-01-21T10:15:07.733533Z","iopub.execute_input":"2026-01-21T10:15:07.733789Z","iopub.status.idle":"2026-01-21T10:15:08.351462Z","shell.execute_reply.started":"2026-01-21T10:15:07.733769Z","shell.execute_reply":"2026-01-21T10:15:08.350692Z"},"papermill":{"duration":0.824995,"end_time":"2026-01-15T22:52:35.171136","exception":false,"start_time":"2026-01-15T22:52:34.346141","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Templates\n\n> - https://www.kaggle.com/code/theoviel/stanford-rna-3d-folding-pt2-templates\n> - https://www.kaggle.com/code/jaejohn/rna-3d-folds-tbm-only-approach/notebook\n> - https://www.kaggle.com/code/nihilisticneuralnet/stanford-rna-folding-2-template-based-approach","metadata":{}},{"cell_type":"code","source":"!pip install /kaggle/input/parasail/biopython-1.85-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:15:08.352225Z","iopub.execute_input":"2026-01-21T10:15:08.352455Z","iopub.status.idle":"2026-01-21T10:15:17.445511Z","shell.execute_reply.started":"2026-01-21T10:15:08.352427Z","shell.execute_reply":"2026-01-21T10:15:17.444662Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!mkdir -p templates","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:15:17.446721Z","iopub.execute_input":"2026-01-21T10:15:17.447035Z","iopub.status.idle":"2026-01-21T10:15:17.562741Z","shell.execute_reply.started":"2026-01-21T10:15:17.447001Z","shell.execute_reply":"2026-01-21T10:15:17.561807Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cd templates","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:15:17.564579Z","iopub.execute_input":"2026-01-21T10:15:17.564797Z","iopub.status.idle":"2026-01-21T10:15:17.570022Z","shell.execute_reply.started":"2026-01-21T10:15:17.564772Z","shell.execute_reply":"2026-01-21T10:15:17.569403Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Utils","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom scipy.spatial.transform import Rotation as R\nimport random\nfrom Bio import pairwise2\nfrom Bio.Seq import Seq\nimport time\nfrom sklearn.preprocessing import normalize\nfrom scipy.spatial import distance_matrix\nfrom tqdm.notebook import tqdm\n\n\n# 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\n# Function to find similar sequences in the training data\ndef find_similar_sequences(query_seq, train_seqs_df, train_coords_dict, temporal_cutoff=None, top_n=5):\n    \"\"\"\n    Find sequences in the training data that are similar to the query sequence,\n    respecting temporal cutoffs if provided.\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    # Convert query to Seq object for alignment\n    query_seq_obj = Seq(query_seq)\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 is too different in length (more than 50% difference)\n        if abs(len(train_seq) - len(query_seq)) / max(len(train_seq), len(query_seq)) > 0.5:\n            continue\n        \n        # Perform sequence alignment\n        # We're using global alignment with gap open penalty of -10 and gap extension penalty of -0.5\n        alignments = pairwise2.align.globalms(query_seq_obj, train_seq, 2, -1, -10, -0.5, one_alignment_only=True)\n        \n        if alignments:\n            alignment = alignments[0]\n            similarity_score = alignment.score / (2 * min(len(query_seq), len(train_seq)))\n            similar_seqs.append((target_id, train_seq, similarity_score, train_coords_dict[target_id]))\n    \n    # Sort by similarity score (higher is better) and return top N\n    similar_seqs.sort(key=lambda x: x[2], reverse=True)\n    return similar_seqs[:top_n]\n\n\ndef adaptive_rna_constraints(coordinates, sequence, confidence=1.0):\n    \"\"\"\n    Apply RNA geometric constraints with adaptive strength based on confidence.\n    \n    Args:\n        coordinates: Array of 3D coordinates, shape (len(sequence), 3)\n        sequence: RNA sequence string\n        confidence: Float between 0 and 1 indicating confidence in template \n                   (higher = less constraint enforcement)\n        \n    Returns:\n        Refined coordinates\n    \"\"\"\n    # Make a copy of coordinates to refine\n    refined_coords = coordinates.copy()\n    n_residues = len(sequence)\n    \n    # Calculate constraint strength (inverse of confidence)\n    # High confidence templates receive gentler constraints\n    constraint_strength = 0.8 * (1.0 - min(confidence, 0.8))\n    \n    # 1. Sequential distance constraints (consecutive nucleotides)\n    # More flexible distance range (statistical distribution from PDB)\n    seq_min_dist = 5.5  # Minimum sequential distance\n    seq_max_dist = 6.5  # Maximum sequential distance\n    \n    for i in range(n_residues - 1):\n        current_pos = refined_coords[i]\n        next_pos = refined_coords[i+1]\n        \n        # Calculate current distance\n        current_dist = np.linalg.norm(next_pos - current_pos)\n        \n        # Only adjust if significantly outside expected range\n        if current_dist < seq_min_dist or current_dist > seq_max_dist:\n            # Calculate target distance (midpoint of range)\n            target_dist = (seq_min_dist + seq_max_dist) / 2\n            \n            # Get direction vector\n            direction = next_pos - current_pos\n            direction = direction / (np.linalg.norm(direction) + 1e-10)\n            \n            # Apply partial adjustment based on constraint strength\n            adjustment = (target_dist - current_dist) * constraint_strength\n            \n            # Only adjust the next position to preserve the overall fold\n            refined_coords[i+1] = current_pos + direction * (current_dist + adjustment)\n    \n    # 2. Steric clash prevention (more conservative)\n    min_allowed_distance = 3.8  # Minimum distance between non-consecutive C1' atoms\n    \n    # Calculate all pairwise distances\n    dist_matrix = distance_matrix(refined_coords, refined_coords)\n    \n    # Find severe clashes (atoms too close)\n    severe_clashes = np.where((dist_matrix < min_allowed_distance) & (dist_matrix > 0))\n    \n    # Fix severe clashes\n    for idx in range(len(severe_clashes[0])):\n        i, j = severe_clashes[0][idx], severe_clashes[1][idx]\n        \n        # Skip consecutive nucleotides and previously processed pairs\n        if abs(i - j) <= 1 or i >= j:\n            continue\n            \n        # Get current positions and distance\n        pos_i = refined_coords[i]\n        pos_j = refined_coords[j]\n        current_dist = dist_matrix[i, j]\n        \n        # Calculate necessary adjustment but scale by constraint strength\n        direction = pos_j - pos_i\n        direction = direction / (np.linalg.norm(direction) + 1e-10)\n        \n        # Calculate partial adjustment\n        adjustment = (min_allowed_distance - current_dist) * constraint_strength\n        \n        # Move points apart\n        refined_coords[i] = pos_i - direction * (adjustment / 2)\n        refined_coords[j] = pos_j + direction * (adjustment / 2)\n    \n    # 3. Very light base-pair constraining (if confidence is low)\n    if constraint_strength > 0.3:  # Only apply if template confidence is low\n        # Simple Watson-Crick base pairs\n        pairs = {'A': 'U', 'U': 'A', 'G': 'C', 'C': 'G'}\n        \n        # Scan for potential base pairs\n        for i in range(n_residues):\n            base_i = sequence[i]\n            complement = pairs.get(base_i)\n            \n            if not complement:\n                continue\n                \n            # Look for complementary bases within a reasonable range\n            for j in range(i + 3, min(i + 20, n_residues)):\n                if sequence[j] == complement:\n                    # Calculate current distance\n                    current_dist = np.linalg.norm(refined_coords[i] - refined_coords[j])\n                    \n                    # Only consider if distance suggests potential pairing\n                    if 8.0 < current_dist < 14.0:\n                        # Target 10.5Å as generic base-pair C1'-C1' distance\n                        target_dist = 10.5\n                        \n                        # Calculate very gentle adjustment (scaled by constraint_strength)\n                        adjustment = (target_dist - current_dist) * (constraint_strength * 0.3)\n                        \n                        # Get direction vector\n                        direction = refined_coords[j] - refined_coords[i]\n                        direction = direction / (np.linalg.norm(direction) + 1e-10)\n                        \n                        # Apply very gentle adjustment to both positions\n                        refined_coords[i] = refined_coords[i] - direction * (adjustment / 2)\n                        refined_coords[j] = refined_coords[j] + direction * (adjustment / 2)\n                        \n                        # Only consider one potential pair per base (closest match)\n                        break\n    \n    return refined_coords\n\n\ndef adapt_template_to_query(query_seq, template_seq, template_coords, alignment=None):\n    \"\"\"\n    Adapt template coordinates to fit the query sequence based on sequence alignment,\n    with robust handling of gaps and NaN values.\n    \n    Args:\n        query_seq: The target RNA sequence\n        template_seq: The template RNA sequence\n        template_coords: 3D coordinates of the template\n        alignment: Optional pre-computed alignment\n        \n    Returns:\n        Array of adapted coordinates for the query sequence\n    \"\"\"\n    if alignment is None:\n        # Perform sequence alignment\n        from Bio.Seq import Seq\n        from Bio import pairwise2\n        \n        query_seq_obj = Seq(query_seq)\n        template_seq_obj = Seq(template_seq)\n        alignments = pairwise2.align.globalms(query_seq_obj, template_seq_obj, 2, -1, -10, -0.5, one_alignment_only=True)\n        \n        if not alignments:\n            # If alignment fails, return a basic structure\n            return generate_basic_structure(query_seq)\n            \n        alignment = alignments[0]\n    \n    # Extract aligned sequences\n    aligned_query = alignment.seqA\n    aligned_template = alignment.seqB\n    \n    # Initialize coordinates for query sequence\n    query_coords = np.zeros((len(query_seq), 3))\n    query_coords.fill(np.nan)  # Fill with NaN initially to track unmapped positions\n    \n    # Map template coordinates to query\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 sequences have a residue at this position\n            if template_idx < len(template_coords):\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 - will handle these gaps in post-processing\n            query_idx += 1\n        elif query_char == '-' and template_char != '-':\n            # Gap in query, skip this template position\n            template_idx += 1\n    \n    # Post-processing step: Fill in any NaN positions\n    # First pass: handle isolated NaNs surrounded by valid coordinates\n    for i in range(len(query_coords)):\n        if np.isnan(query_coords[i, 0]):\n            # Find nearest valid coordinates before and after\n            prev_valid = -1\n            for j in range(i-1, -1, -1):\n                if not np.isnan(query_coords[j, 0]):\n                    prev_valid = j\n                    break\n                    \n            next_valid = -1\n            for j in range(i+1, len(query_coords)):\n                if not np.isnan(query_coords[j, 0]):\n                    next_valid = j\n                    break\n            \n            # Interpolate if we have valid points on both sides\n            if prev_valid >= 0 and next_valid >= 0:\n                # Linear interpolation\n                weight = (i - prev_valid) / (next_valid - prev_valid)\n                query_coords[i] = (1 - weight) * query_coords[prev_valid] + weight * query_coords[next_valid]\n    \n    # Second pass: handle remaining NaNs (at start, end, or long stretches)\n    typical_step = 4.0  # Typical nucleotide distance in Angstroms\n    \n    for i in range(len(query_coords)):\n        if np.isnan(query_coords[i, 0]):\n            # If this is the first position with NaN\n            if i == 0:\n                # Look for the first valid position\n                first_valid = -1\n                for j in range(1, len(query_coords)):\n                    if not np.isnan(query_coords[j, 0]):\n                        first_valid = j\n                        break\n                \n                if first_valid >= 0:\n                    # Work backwards from the first valid position\n                    for j in range(first_valid-1, -1, -1):\n                        # Create a random direction for the step\n                        direction = np.random.normal(0, 1, 3)\n                        direction = direction / (np.linalg.norm(direction) + 1e-10) * typical_step\n                        query_coords[j] = query_coords[j+1] - direction\n                else:\n                    # No valid positions at all, create a basic structure\n                    for j in range(len(query_coords)):\n                        angle = j * 0.6\n                        query_coords[j] = [10.0 * np.cos(angle), 10.0 * np.sin(angle), j * 2.5]\n                    break  # Exit the loop as we've filled all positions\n            else:\n                # Not the first position, look for previous valid position\n                prev_valid = -1\n                for j in range(i-1, -1, -1):\n                    if not np.isnan(query_coords[j, 0]):\n                        prev_valid = j\n                        break\n                \n                if prev_valid >= 0:\n                    # We have a valid previous position\n                    if prev_valid > 0 and not np.isnan(query_coords[prev_valid-1, 0]):\n                        # We have two previous valid positions, continue the direction\n                        direction = query_coords[prev_valid] - query_coords[prev_valid-1]\n                        direction = direction / (np.linalg.norm(direction) + 1e-10) * typical_step\n                        query_coords[i] = query_coords[prev_valid] + direction\n                    else:\n                        # Just one previous valid position, create a random step\n                        direction = np.random.normal(0, 1, 3)\n                        direction = direction / (np.linalg.norm(direction) + 1e-10) * typical_step\n                        query_coords[i] = query_coords[prev_valid] + direction\n                else:\n                    # No valid previous positions (unusual case), start from origin\n                    query_coords[i] = np.array([0.0, 0.0, 0.0]) + np.random.normal(0, 1, 3) * i\n    \n    # Verify no NaNs remain\n    if np.isnan(query_coords).any():\n        # Final fallback: replace any remaining NaNs with zeros\n        query_coords = np.nan_to_num(query_coords)\n    \n    return query_coords\n\ndef generate_basic_structure(sequence):\n    \"\"\"\n    Generate a simple helical structure for a sequence when alignment fails.\n    \n    Args:\n        sequence: RNA sequence string\n        \n    Returns:\n        Array of 3D coordinates\n    \"\"\"\n    n_residues = len(sequence)\n    coordinates = np.zeros((n_residues, 3))\n    \n    # Create a simple helix\n    radius = 10.0\n    rise_per_residue = 2.5\n    angle_per_residue = 0.6\n    \n    for i in range(n_residues):\n        angle = i * angle_per_residue\n        coordinates[i] = [\n            radius * np.cos(angle), \n            radius * np.sin(angle), \n            i * rise_per_residue\n        ]\n    \n    return coordinates\n\n\n# Function to create a more realistic RNA structure when no good templates are found\ndef generate_rna_structure(sequence, seed=None):\n    \"\"\"\n    Generate a more realistic RNA structure prediction for a given sequence.\n    \n    Args:\n        sequence: RNA sequence (string of A, C, G, U)\n        seed: Random seed for reproducibility\n        \n    Returns:\n        coordinates: Array of shape (len(sequence), 3) with x, y, z coordinates\n    \"\"\"\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 the first few residues in a helix\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    # Add more complex folding patterns\n    current_direction = np.array([0.0, 0.0, 1.0])  # Start moving along z-axis\n    \n    # Define base-pairing tendencies (G-C and A-U pairs)\n    for i in range(3, n_residues):\n        # Check for potential base-pairing in the sequence\n        has_pair = False\n        pair_idx = -1\n        \n        # Simple detection of complementary bases (G-C, A-U)\n        complementary = {'G': 'C', 'C': 'G', 'A': 'U', 'U': 'A'}\n        current_base = sequence[i]\n        \n        # Look for potential base-pairing within a window before the current position\n        window_size = min(i, 15)  # Look back up to 15 bases\n        for j in range(i-window_size, i):\n            if j >= 0 and sequence[j] == complementary.get(current_base, 'X'):\n                # Found a potential pair\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            # Try to create a base-pair by positioning this nucleotide near its pair\n            pair_pos = coordinates[pair_idx]\n            \n            # Create a position that's roughly opposite to the pair\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            # Calculate a vector from base-pair toward center of structure\n            center = np.mean(coordinates[:i], axis=0)\n            direction = center - pair_pos\n            direction = direction / (np.linalg.norm(direction) + 1e-10)\n            \n            # Position new nucleotide in the general direction of the \"center\"\n            coordinates[i] = pair_pos + direction * base_pair_distance + random_offset\n            \n            # Update direction for next nucleotide\n            current_direction = np.random.normal(0, 0.3, 3)\n            current_direction = current_direction / (np.linalg.norm(current_direction) + 1e-10)\n            \n        else:\n            # No base-pairing detected, continue with the current fold direction\n            # Randomly rotate current direction to simulate RNA flexibility\n            if random.random() < 0.3:\n                # More significant direction change\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                # Small random changes in direction\n                current_direction += np.random.normal(0, 0.15, 3)\n                current_direction = current_direction / (np.linalg.norm(current_direction) + 1e-10)\n            \n            # Distance between consecutive nucleotides (3.5-4.5Å is typical)\n            step_size = random.uniform(3.5, 4.5)\n            \n            # Update position\n            coordinates[i] = coordinates[i-1] + step_size * current_direction\n    \n    return coordinates\n\n\ndef predict_rna_structures(sequence, target_id, train_seqs_df, train_coords_dict, n_predictions=5, temporal_cutoff=None):\n    predictions = []\n    \n    # Find similar sequences in the training data\n    similar_seqs = find_similar_sequences(sequence, train_seqs_df, train_coords_dict, \n                                         temporal_cutoff=temporal_cutoff, top_n=n_predictions)\n    \n    # If we found any similar sequences, use them as templates\n    if similar_seqs:\n        for i, (template_id, template_seq, similarity, template_coords) in enumerate(similar_seqs):\n            # Adapt template coordinates to the query sequence\n            adapted_coords = adapt_template_to_query(sequence, template_seq, template_coords)\n            \n            if adapted_coords is not None:\n                # Apply adaptive constraints based on template similarity\n                # For high similarity templates, apply very gentle constraints\n                refined_coords = adaptive_rna_constraints(adapted_coords, sequence, confidence=similarity)\n                \n                # Add some randomness (less for better templates)\n                random_scale = max(0.05, 0.8 - similarity)  # Reduced randomness\n                randomized_coords = refined_coords.copy()\n                randomized_coords += np.random.normal(0, random_scale, randomized_coords.shape)\n                \n                predictions.append(randomized_coords)\n                \n                if len(predictions) >= n_predictions:\n                    break\n    \n    # If we don't have enough predictions from templates, generate 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        \n        # Apply stronger constraints to de novo structures (lower confidence)\n        refined_de_novo = adaptive_rna_constraints(de_novo_coords, sequence, confidence=0.2)\n        \n        predictions.append(refined_de_novo)\n    \n    return predictions[:n_predictions]\n","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-input":true,"execution":{"iopub.status.busy":"2026-01-21T10:21:52.728891Z","iopub.execute_input":"2026-01-21T10:21:52.729255Z","iopub.status.idle":"2026-01-21T10:21:52.76867Z","shell.execute_reply.started":"2026-01-21T10:21:52.729223Z","shell.execute_reply":"2026-01-21T10:21:52.767734Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\n\n# Function to extend the original dataset with new records from v2\ndef extend_dataset(original_df, v2_df, key_columns, dataset_name):\n    print(f\"Extending {dataset_name}...\")\n    print(f\"  Original size: {len(original_df)} rows\")\n    print(f\"  v2 size: {len(v2_df)} rows\")\n    \n    # Create a composite key for identification if multiple key columns\n    if isinstance(key_columns, list) and len(key_columns) > 1:\n        original_df['temp_key'] = original_df[key_columns].astype(str).agg('_'.join, axis=1)\n        v2_df['temp_key'] = v2_df[key_columns].astype(str).agg('_'.join, axis=1)\n        key_for_identification = 'temp_key'\n    else:\n        key_for_identification = key_columns[0] if isinstance(key_columns, list) else key_columns\n    \n    # Identify unique records in each dataset\n    original_keys = set(original_df[key_for_identification])\n    v2_keys = set(v2_df[key_for_identification])\n    \n    # Calculate stats\n    keys_only_in_original = original_keys - v2_keys\n    keys_only_in_v2 = v2_keys - original_keys \n    common_keys = original_keys.intersection(v2_keys)\n    \n    print(f\"  Keys only in original: {len(keys_only_in_original)}\")\n    print(f\"  Keys only in v2: {len(keys_only_in_v2)}\")\n    print(f\"  Common keys: {len(common_keys)}\")\n    \n    # Create a mask to filter v2 records that don't exist in original\n    new_records_mask = ~v2_df[key_for_identification].isin(original_keys)\n    new_records = v2_df[new_records_mask].copy()\n    \n    # Drop temporary key if it was created\n    if key_for_identification == 'temp_key':\n        new_records.drop('temp_key', axis=1, inplace=True)\n        original_df.drop('temp_key', axis=1, inplace=True)\n    \n    # Combine original with new records from v2\n    extended_df = pd.concat([original_df, new_records], ignore_index=True)\n    \n    # Report final sizes\n    print(f\"  New records added: {len(new_records)}\")\n    print(f\"  Extended dataset size: {len(extended_df)} rows\")\n    print(f\"  Verification - All original keys in extended dataset: {set(original_df[key_columns[0] if isinstance(key_columns, list) else key_columns]).issubset(set(extended_df[key_columns[0] if isinstance(key_columns, list) else key_columns]))}\")\n    \n    # Check for missing values in key columns\n    for col in extended_df.columns:\n        original_missing = original_df[col].isnull().sum()\n        extended_missing = extended_df[col].isnull().sum()\n        if original_missing > 0 or extended_missing > 0:\n            print(f\"  Column '{col}': Missing values - Original: {original_missing}, Extended: {extended_missing}\")\n    \n    # Clean up\n    if key_for_identification == 'temp_key' and 'temp_key' in v2_df.columns:\n        v2_df.drop('temp_key', axis=1, inplace=True)\n        \n    return extended_df","metadata":{"trusted":true,"_kg_hide-input":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Load","metadata":{}},{"cell_type":"code","source":"seed = 21\nnp.random.seed(seed)\nrandom.seed(seed)\n\nBASE_PATH = '/kaggle/input/stanford-rna-3d-folding-2'\n\ntest_seqs = pd.read_csv(f'/kaggle/working/sequences.csv')\ntrain_seqs = pd.read_csv(f'{BASE_PATH}/train_sequences.csv')\ntrain_labels = pd.read_csv(f'{BASE_PATH}/train_labels.csv', low_memory=False)\n    \nprint(f\"✓ Loaded {len(train_seqs)} training sequences\")\nprint(f\"✓ Loaded {len(test_seqs)} test sequences\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:47:44.098754Z","iopub.execute_input":"2026-01-21T10:47:44.099387Z","iopub.status.idle":"2026-01-21T10:47:55.674804Z","shell.execute_reply.started":"2026-01-21T10:47:44.099358Z","shell.execute_reply":"2026-01-21T10:47:55.674129Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Extend train sequences","metadata":{}},{"cell_type":"code","source":"EXTEND = True","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:47:55.676153Z","iopub.execute_input":"2026-01-21T10:47:55.676453Z","iopub.status.idle":"2026-01-21T10:47:55.679879Z","shell.execute_reply.started":"2026-01-21T10:47:55.676428Z","shell.execute_reply":"2026-01-21T10:47:55.67913Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if EXTEND:\n    train_seqs_v2 = pd.read_csv('/kaggle/input/rna-cif-to-csv/rna_sequences.csv')\n    train_labels_v2 = pd.read_csv('/kaggle/input/rna-cif-to-csv/rna_coordinates.csv')\n    \n    # 1. Extend train_seqs with train_seqs_v2\n    print(\"\\n\" + \"=\"*50)\n    print(\"EXTENDING SEQUENCE DATASETS\")\n    print(\"=\"*50)\n    train_seqs_extended = extend_dataset(\n        train_seqs, \n        train_seqs_v2,\n        ['target_id'],  # Using target_id as the unique identifier\n        \"train_seqs\"\n    )\n    \n    # 2. Extend train_labels with train_labels_v2\n    print(\"\\n\" + \"=\"*50)\n    print(\"EXTENDING LABELS DATASETS\")\n    print(\"=\"*50)\n    # For labels, we need a composite key of ID and resid\n    train_labels_extended = extend_dataset(\n        train_labels,\n        train_labels_v2,\n        ['ID', 'resid'],  # Using composite key\n        \"train_labels\"\n    )\n    \n    # Verify relationships between extended datasets\n    print(\"\\n\" + \"=\"*50)\n    print(\"VERIFYING RELATIONSHIPS\")\n    print(\"=\"*50)\n    \n    # Check if all sequence IDs have corresponding labels\n    seq_ids = set(train_seqs_extended['target_id'].unique())\n    label_ids = set(train_labels_extended['ID'].unique())\n    \n    seq_ids_with_labels = seq_ids.intersection(label_ids)\n    seq_ids_without_labels = seq_ids - label_ids\n    \n    print(f\"Total unique sequence IDs: {len(seq_ids)}\")\n    print(f\"Sequence IDs with corresponding labels: {len(seq_ids_with_labels)} ({len(seq_ids_with_labels)/len(seq_ids)*100:.2f}%)\")\n    print(f\"Sequence IDs without corresponding labels: {len(seq_ids_without_labels)} ({len(seq_ids_without_labels)/len(seq_ids)*100:.2f}%)\")\n    \n    if len(seq_ids_without_labels) > 0:\n        print(\"Sample of sequence IDs without labels (up to 5):\")\n        print(list(seq_ids_without_labels)[:5])\n    \n    # Print summary of extended datasets\n    print(\"\\n\" + \"=\"*50)\n    print(\"SUMMARY OF EXTENDED DATASETS\")\n    print(\"=\"*50)\n    \n    print(f\"Original train_seqs: {len(train_seqs)} rows\")\n    print(f\"Original train_labels: {len(train_labels)} rows\")\n    print(f\"Extended train_seqs: {len(train_seqs_extended)} rows (+{len(train_seqs_extended)-len(train_seqs)})\")\n    print(f\"Extended train_labels: {len(train_labels_extended)} rows (+{len(train_labels_extended)-len(train_labels)})\")\n    \n    # Save the extended datasets (uncomment to save)\n    # train_seqs_extended.to_csv('train_seqs_combined.csv', index=False)\n    # train_labels_extended.to_csv('train_labels_combined.csv', index=False)\n    \n    train_seqs = train_seqs_extended\n    train_labels = train_labels_extended","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:47:55.680819Z","iopub.execute_input":"2026-01-21T10:47:55.681107Z","iopub.status.idle":"2026-01-21T10:50:10.015627Z","shell.execute_reply.started":"2026-01-21T10:47:55.681076Z","shell.execute_reply":"2026-01-21T10:50:10.014954Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Main","metadata":{}},{"cell_type":"code","source":"train_coords_dict = process_labels(train_labels)\n\nall_predictions = []\ntotal_targets = len(test_seqs)\n\n# For each sequence in the test set\nfor idx, row in tqdm(test_seqs.iterrows(), total=len(test_seqs)):\n    target_id = row['target_id']\n    sequence = row['sequence']\n    temporal_cutoff = row['temporal_cutoff'] if 'temporal_cutoff' in row else None\n    \n    # Generate 5 different structure predictions\n    predictions = predict_rna_structures(\n        sequence,\n        target_id,\n        train_seqs,\n        train_coords_dict, \n        n_predictions=5,\n        temporal_cutoff=temporal_cutoff\n    )\n    \n    # For each residue in the sequence\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        \n        # Add coordinates from all 5 predictions\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        \n        all_predictions.append(pred_row)\n\n# Create DataFrame with predictions\nsubmission_df = pd.DataFrame(all_predictions)\n\n# Ensure the submission file has the correct format\ncolumn_order = ['ID', 'resname', 'resid']\nfor i in range(1, 6):\n    for coord in ['x', 'y', 'z']:\n        column_order.append(f'{coord}_{i}')\nsubmission_df = submission_df[column_order]\n\n# Save the submission file\nsubmission_df.to_csv('/kaggle/working/templates_tbm.csv', index=False)\nprint(f\"Generated predictions for {len(test_seqs)} RNA sequences\")\nsubmission_df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:43:02.47741Z","iopub.execute_input":"2026-01-21T10:43:02.478026Z","iopub.status.idle":"2026-01-21T10:44:12.70783Z","shell.execute_reply.started":"2026-01-21T10:43:02.477999Z","shell.execute_reply":"2026-01-21T10:44:12.707244Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## RNAPro setup","metadata":{}},{"cell_type":"code","source":"cd /kaggle/working/","metadata":{"papermill":{"duration":0.239616,"end_time":"2026-01-15T22:48:48.537648","exception":false,"start_time":"2026-01-15T22:48:48.298032","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:23:46.092666Z","iopub.execute_input":"2026-01-21T10:23:46.093468Z","iopub.status.idle":"2026-01-21T10:23:46.098283Z","shell.execute_reply.started":"2026-01-21T10:23:46.093439Z","shell.execute_reply":"2026-01-21T10:23:46.09756Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Weights","metadata":{"papermill":{"duration":0.004687,"end_time":"2026-01-15T22:13:48.700197","exception":false,"start_time":"2026-01-15T22:13:48.69551","status":"completed"},"tags":[]}},{"cell_type":"code","source":"!cp /kaggle/input/rnapro-src/rnapro-private-best-500m.ckpt ./","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2026-01-21T10:23:46.896493Z","iopub.execute_input":"2026-01-21T10:23:46.897128Z","iopub.status.idle":"2026-01-21T10:24:33.392093Z","shell.execute_reply.started":"2026-01-21T10:23:46.897102Z","shell.execute_reply":"2026-01-21T10:24:33.391273Z"},"papermill":{"duration":1975.401231,"end_time":"2026-01-15T22:46:44.106287","exception":false,"start_time":"2026-01-15T22:13:48.705056","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Code\n- Uses a fork of the original repo, adapted for Kaggle inference","metadata":{"papermill":{"duration":0.167493,"end_time":"2026-01-15T22:46:44.436316","exception":false,"start_time":"2026-01-15T22:46:44.268823","status":"completed"},"tags":[]}},{"cell_type":"code","source":"!cp -r /kaggle/input/rnapro-src/RNAPro ./","metadata":{"execution":{"iopub.status.busy":"2026-01-21T10:24:33.393835Z","iopub.execute_input":"2026-01-21T10:24:33.394092Z","iopub.status.idle":"2026-01-21T10:24:34.369501Z","shell.execute_reply.started":"2026-01-21T10:24:33.394064Z","shell.execute_reply":"2026-01-21T10:24:34.368756Z"},"papermill":{"duration":1.070913,"end_time":"2026-01-15T22:46:45.672274","exception":false,"start_time":"2026-01-15T22:46:44.601361","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Installation","metadata":{"papermill":{"duration":0.216564,"end_time":"2026-01-15T22:48:47.98686","exception":false,"start_time":"2026-01-15T22:48:47.770296","status":"completed"},"tags":[]}},{"cell_type":"code","source":"pip install -r RNAPro/requirements.txt --find-links=/kaggle/input/stanford-rna-3d-folding-pt2-rnapro-inference/wheels/ --no-index --no-build-isolation","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2026-01-21T10:24:34.370676Z","iopub.execute_input":"2026-01-21T10:24:34.370992Z","iopub.status.idle":"2026-01-21T10:27:47.464706Z","shell.execute_reply.started":"2026-01-21T10:24:34.370962Z","shell.execute_reply":"2026-01-21T10:27:47.463803Z"},"papermill":{"duration":206.03569,"end_time":"2026-01-15T22:52:14.805078","exception":false,"start_time":"2026-01-15T22:48:48.769388","status":"completed"},"scrolled":true,"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cd RNAPro","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:27:47.466937Z","iopub.execute_input":"2026-01-21T10:27:47.467236Z","iopub.status.idle":"2026-01-21T10:27:47.473044Z","shell.execute_reply.started":"2026-01-21T10:27:47.467177Z","shell.execute_reply":"2026-01-21T10:27:47.471936Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pip install -e . --no-deps","metadata":{"execution":{"iopub.status.busy":"2026-01-21T10:27:47.474212Z","iopub.execute_input":"2026-01-21T10:27:47.474535Z","iopub.status.idle":"2026-01-21T10:27:52.216941Z","shell.execute_reply.started":"2026-01-21T10:27:47.474501Z","shell.execute_reply":"2026-01-21T10:27:52.215938Z"},"papermill":{"duration":5.048747,"end_time":"2026-01-15T22:52:20.595672","exception":false,"start_time":"2026-01-15T22:52:15.546925","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Templates","metadata":{"papermill":{"duration":0.217386,"end_time":"2026-01-15T22:52:21.051653","exception":false,"start_time":"2026-01-15T22:52:20.834267","status":"completed"},"tags":[]}},{"cell_type":"code","source":"!python preprocess/convert_templates_to_pt_files.py --input_csv /kaggle/working/templates_tbm.csv --output_name templates.pt","metadata":{"papermill":{"duration":3.887097,"end_time":"2026-01-15T22:52:25.157531","exception":false,"start_time":"2026-01-15T22:52:21.270434","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:27:52.218185Z","iopub.execute_input":"2026-01-21T10:27:52.218471Z","iopub.status.idle":"2026-01-21T10:27:55.023896Z","shell.execute_reply.started":"2026-01-21T10:27:52.218443Z","shell.execute_reply":"2026-01-21T10:27:55.022965Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### cdd cache\n\nYou can either \n- Recompute the cdd cache using `python preprocess/gen_ccd_cache.py`\n- Use the precomputed files from the proteinix external dataset.\n\nRunning the script might be better since the resulting file will include the latest information, but requires some time.","metadata":{"papermill":{"duration":0.214123,"end_time":"2026-01-15T22:52:25.673436","exception":false,"start_time":"2026-01-15T22:52:25.459313","status":"completed"},"tags":[]}},{"cell_type":"code","source":"DIST = \"/kaggle/working/RNAPro/release_data/ccd_cache/\"\n!mkdir -p $DIST","metadata":{"execution":{"iopub.status.busy":"2026-01-21T10:27:55.025265Z","iopub.execute_input":"2026-01-21T10:27:55.025535Z","iopub.status.idle":"2026-01-21T10:27:55.210255Z","shell.execute_reply.started":"2026-01-21T10:27:55.025508Z","shell.execute_reply":"2026-01-21T10:27:55.209499Z"},"papermill":{"duration":0.336228,"end_time":"2026-01-15T22:52:26.226391","exception":false,"start_time":"2026-01-15T22:52:25.890163","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !python preprocess/gen_ccd_cache.py -c $DIST","metadata":{"execution":{"iopub.status.busy":"2026-01-21T10:27:55.211721Z","iopub.execute_input":"2026-01-21T10:27:55.212336Z","iopub.status.idle":"2026-01-21T10:27:55.215553Z","shell.execute_reply.started":"2026-01-21T10:27:55.212304Z","shell.execute_reply":"2026-01-21T10:27:55.214954Z"},"papermill":{"duration":0.223126,"end_time":"2026-01-15T22:52:26.665028","exception":false,"start_time":"2026-01-15T22:52:26.441902","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp /kaggle/input/protenix-checkpoints/components.v20240608.cif $DIST\n!cp /kaggle/input/protenix-checkpoints/components.v20240608.cif.rdkit_mol.pkl $DIST","metadata":{"execution":{"iopub.status.busy":"2026-01-21T10:27:55.216586Z","iopub.execute_input":"2026-01-21T10:27:55.216856Z","iopub.status.idle":"2026-01-21T10:27:58.736468Z","shell.execute_reply.started":"2026-01-21T10:27:55.21683Z","shell.execute_reply":"2026-01-21T10:27:58.73569Z"},"papermill":{"duration":5.839296,"end_time":"2026-01-15T22:52:32.72341","exception":false,"start_time":"2026-01-15T22:52:26.884114","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Inference\n","metadata":{"papermill":{"duration":0.308629,"end_time":"2026-01-15T22:52:33.251757","exception":false,"start_time":"2026-01-15T22:52:32.943128","status":"completed"},"tags":[]}},{"cell_type":"code","source":"cd /kaggle/working/RNAPro","metadata":{"execution":{"iopub.status.busy":"2026-01-21T10:27:58.739227Z","iopub.execute_input":"2026-01-21T10:27:58.739459Z","iopub.status.idle":"2026-01-21T10:27:58.744264Z","shell.execute_reply.started":"2026-01-21T10:27:58.739434Z","shell.execute_reply":"2026-01-21T10:27:58.743471Z"},"papermill":{"duration":0.224458,"end_time":"2026-01-15T22:52:33.695085","exception":false,"start_time":"2026-01-15T22:52:33.470627","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Inference script\n- Where the parameters are specified","metadata":{"papermill":{"duration":0.223898,"end_time":"2026-01-15T22:52:36.592426","exception":false,"start_time":"2026-01-15T22:52:36.368528","status":"completed"},"tags":[]}},{"cell_type":"code","source":"%%writefile rnapro_inference_kaggle.sh\n\nexport LAYERNORM_TYPE=torch # fast_layernorm, torch\n\n# Inference parameters (RNAPro)\nSEED=42\nN_SAMPLE=1\nN_STEP=200\nN_CYCLE=10\n\n# Save folder\nDUMP_DIR=\"../output\"\n\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\n\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/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    --logger \"print\" \\\n    --n_templates_inf 1\n\n# Notes:\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\n# --logger handles logging of the inference runner, supports \"logging\", \"print\"\n# --n_templates_inf sets the number of inferences to do with different template combinations","metadata":{"execution":{"iopub.status.busy":"2026-01-21T10:27:58.745641Z","iopub.execute_input":"2026-01-21T10:27:58.745924Z","iopub.status.idle":"2026-01-21T10:27:58.766819Z","shell.execute_reply.started":"2026-01-21T10:27:58.745892Z","shell.execute_reply":"2026-01-21T10:27:58.766049Z"},"papermill":{"duration":0.225419,"end_time":"2026-01-15T22:52:37.032055","exception":false,"start_time":"2026-01-15T22:52:36.806636","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Run !","metadata":{"papermill":{"duration":0.304625,"end_time":"2026-01-15T22:52:37.573143","exception":false,"start_time":"2026-01-15T22:52:37.268518","status":"completed"},"tags":[]}},{"cell_type":"code","source":"!bash ./rnapro_inference_kaggle.sh","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.status.busy":"2026-01-21T10:27:58.767823Z","iopub.execute_input":"2026-01-21T10:27:58.768052Z","iopub.status.idle":"2026-01-21T10:31:14.321601Z","shell.execute_reply.started":"2026-01-21T10:27:58.768033Z","shell.execute_reply":"2026-01-21T10:31:14.320593Z"},"papermill":{"duration":727.648981,"end_time":"2026-01-15T23:04:45.440608","exception":false,"start_time":"2026-01-15T22:52:37.791627","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!head submission.csv","metadata":{"papermill":{"duration":0.415157,"end_time":"2026-01-15T23:04:47.071737","exception":false,"start_time":"2026-01-15T23:04:46.65658","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:31:14.32305Z","iopub.execute_input":"2026-01-21T10:31:14.323836Z","iopub.status.idle":"2026-01-21T10:31:14.512318Z","shell.execute_reply.started":"2026-01-21T10:31:14.3238Z","shell.execute_reply":"2026-01-21T10:31:14.511134Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Replace long sequences\n- Credits: https://www.kaggle.com/code/jaejohn/rnapro-inference-with-tbm","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\n\ndf_tbm = pd.read_csv(\"/kaggle/working/templates_tbm.csv\")\ndf_rnapro = pd.read_csv(\"submission.csv\")\ndf_seqs = pd.read_csv(\"/kaggle/working/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(\n    lambda x: any(str(x).startswith(t + \"_\") for t in long_targets)\n)\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,"execution":{"iopub.status.busy":"2026-01-21T10:31:14.51374Z","iopub.execute_input":"2026-01-21T10:31:14.514006Z","iopub.status.idle":"2026-01-21T10:31:14.532987Z","shell.execute_reply.started":"2026-01-21T10:31:14.513979Z","shell.execute_reply":"2026-01-21T10:31:14.532247Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Clean everything","metadata":{}},{"cell_type":"code","source":"!mv submission.csv ..","metadata":{"execution":{"iopub.status.busy":"2026-01-21T10:31:14.533904Z","iopub.execute_input":"2026-01-21T10:31:14.534151Z","iopub.status.idle":"2026-01-21T10:31:14.727347Z","shell.execute_reply.started":"2026-01-21T10:31:14.534131Z","shell.execute_reply":"2026-01-21T10:31:14.726523Z"},"papermill":{"duration":0.339126,"end_time":"2026-01-15T23:04:45.996398","exception":false,"start_time":"2026-01-15T23:04:45.657272","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cd ..","metadata":{"execution":{"iopub.status.busy":"2026-01-21T10:31:14.728665Z","iopub.execute_input":"2026-01-21T10:31:14.7294Z","iopub.status.idle":"2026-01-21T10:31:14.734082Z","shell.execute_reply.started":"2026-01-21T10:31:14.729369Z","shell.execute_reply":"2026-01-21T10:31:14.733406Z"},"papermill":{"duration":0.223984,"end_time":"2026-01-15T23:04:46.43761","exception":false,"start_time":"2026-01-15T23:04:46.213626","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!rm -r output RNAPro *.ckpt templates","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:31:14.734942Z","iopub.execute_input":"2026-01-21T10:31:14.735176Z","iopub.status.idle":"2026-01-21T10:31:15.424086Z","shell.execute_reply.started":"2026-01-21T10:31:14.735149Z","shell.execute_reply":"2026-01-21T10:31:15.423294Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Done !","metadata":{"papermill":{"duration":0.215086,"end_time":"2026-01-15T23:04:47.501983","exception":false,"start_time":"2026-01-15T23:04:47.286897","status":"completed"},"tags":[]}}]}