{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","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"}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"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\n\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\nimport seaborn as sns\n\nimport warnings\nwarnings.filterwarnings('ignore')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-09T03:42:37.129771Z","iopub.execute_input":"2026-01-09T03:42:37.130126Z","iopub.status.idle":"2026-01-09T03:42:37.616296Z","shell.execute_reply.started":"2026-01-09T03:42:37.130094Z","shell.execute_reply":"2026-01-09T03:42:37.615175Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_seqs = pd.read_csv('/kaggle/input/stanford-rna-3d-folding-2/train_sequences.csv')\nvalid_seqs = pd.read_csv('/kaggle/input/stanford-rna-3d-folding-2/validation_sequences.csv')\ntest_seqs = pd.read_csv('/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv')\ntrain_labels = pd.read_csv('/kaggle/input/stanford-rna-3d-folding-2/train_labels.csv')\nvalid_labels = pd.read_csv('/kaggle/input/stanford-rna-3d-folding-2/validation_labels.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T03:42:39.269039Z","iopub.execute_input":"2026-01-09T03:42:39.270475Z","iopub.status.idle":"2026-01-09T03:42:51.046936Z","shell.execute_reply.started":"2026-01-09T03:42:39.270438Z","shell.execute_reply":"2026-01-09T03:42:51.045907Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"Loaded {len(train_seqs)} training sequences, {len(valid_seqs)} validation sequences, and {len(test_seqs)} test sequences\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T03:42:51.048708Z","iopub.execute_input":"2026-01-09T03:42:51.049055Z","iopub.status.idle":"2026-01-09T03:42:51.054718Z","shell.execute_reply.started":"2026-01-09T03:42:51.049024Z","shell.execute_reply":"2026-01-09T03:42:51.05363Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_base_composition():\n    \"\"\"\n    Plot the nucleotide composition across different datasets.\n    \"\"\"\n    def get_base_counts(seqs_df):\n        bases = {'A': 0, 'C': 0, 'G': 0, 'U': 0}\n        for seq in seqs_df['sequence']:\n            for base in seq:\n                if base in bases:\n                    bases[base] += 1\n        total = sum(bases.values())\n        return {k: v/total * 100 for k, v in bases.items()}\n    \n    train_comp = get_base_counts(train_seqs)\n    valid_comp = get_base_counts(valid_seqs)\n    test_comp = get_base_counts(test_seqs)\n    \n    fig, ax = plt.subplots(figsize=(10, 6))\n    x = np.arange(4)\n    width = 0.25\n    bases = ['A', 'C', 'G', 'U']\n    \n    ax.bar(x - width, [train_comp[b] for b in bases], width, label='Train', alpha=0.8)\n    ax.bar(x, [valid_comp[b] for b in bases], width, label='Validation', alpha=0.8)\n    ax.bar(x + width, [test_comp[b] for b in bases], width, label='Test', alpha=0.8)\n    \n    ax.set_xlabel('Nucleotide Base')\n    ax.set_ylabel('Percentage (%)')\n    ax.set_title('Nucleotide Composition Across Datasets')\n    ax.set_xticks(x)\n    ax.set_xticklabels(bases)\n    ax.legend()\n    ax.grid(True, alpha=0.3, axis='y')\n    plt.tight_layout()\n    plt.show()\n\nplot_base_composition()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Process training labels to create a dictionary mapping target_id to its 3D coordinates\ndef process_labels(labels_df):\n    \"\"\"\n    Process labels dataframe to create a dictionary mapping target_id to coordinates\n    \"\"\"\n    coords_dict = {}\n    \n    # Group by target ID\n    for id_prefix, group in labels_df.groupby(lambda x: labels_df['ID'][x].rsplit('_', 1)[0]):\n        # Extract just the coordinates columns for the first structure (x_1, y_1, z_1)\n        coords = []\n        for _, row in group.sort_values('resid').iterrows():\n            coords.append([row['x_1'], row['y_1'], row['z_1']])\n        \n        coords_dict[id_prefix] = np.array(coords)\n    \n    return coords_dict\n\n# Process training labels\nprint(\"Processing training labels...\")\ntrain_coords_dict = process_labels(train_labels)\nvalid_coords_dict = process_labels(valid_labels)\nprint(f\"Processed coordinates for {len(train_coords_dict)} training structures and {len(valid_coords_dict)} validation structures\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T03:44:48.467367Z","iopub.execute_input":"2026-01-09T03:44:48.468167Z","iopub.status.idle":"2026-01-09T03:51:28.562704Z","shell.execute_reply.started":"2026-01-09T03:44:48.468132Z","shell.execute_reply":"2026-01-09T03:51:28.561792Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_distance_distributions(coords_dict, sample_size=100):\n    \"\"\"\n    Plot distributions of sequential and pairwise distances in training structures.\n    INSERT: After processing training labels (after line ~47)\n    \"\"\"\n    sequential_dists = []\n    all_pairwise_dists = []\n    \n    # Sample structures to avoid computational overload\n    sampled_ids = random.sample(list(coords_dict.keys()), min(sample_size, len(coords_dict)))\n    \n    for target_id in sampled_ids:\n        coords = coords_dict[target_id]\n        \n        # Sequential distances (consecutive nucleotides)\n        for i in range(len(coords) - 1):\n            dist = np.linalg.norm(coords[i+1] - coords[i])\n            sequential_dists.append(dist)\n        \n        # Sample pairwise distances (to avoid memory issues)\n        if len(coords) > 2:\n            sample_indices = random.sample(range(len(coords)), min(20, len(coords)))\n            for i in range(len(sample_indices)):\n                for j in range(i+2, len(sample_indices)):  # Skip consecutive\n                    dist = np.linalg.norm(coords[sample_indices[i]] - coords[sample_indices[j]])\n                    all_pairwise_dists.append(dist)\n    \n    fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n    \n    # Sequential distances\n    axes[0].hist(sequential_dists, bins=50, edgecolor='black', alpha=0.7, color='steelblue')\n    axes[0].axvline(np.mean(sequential_dists), color='red', linestyle='--', \n                   label=f'Mean: {np.mean(sequential_dists):.2f} Å')\n    axes[0].set_xlabel('Distance (Å)')\n    axes[0].set_ylabel('Frequency')\n    axes[0].set_title('Sequential C1\\' Distance Distribution')\n    axes[0].legend()\n    axes[0].grid(True, alpha=0.3)\n    \n    # All pairwise distances\n    axes[1].hist(all_pairwise_dists, bins=50, edgecolor='black', alpha=0.7, color='coral')\n    axes[1].axvline(np.mean(all_pairwise_dists), color='red', linestyle='--', \n                   label=f'Mean: {np.mean(all_pairwise_dists):.2f} Å')\n    axes[1].set_xlabel('Distance (Å)')\n    axes[1].set_ylabel('Frequency')\n    axes[1].set_title('Non-Sequential Pairwise Distance Distribution')\n    axes[1].legend()\n    axes[1].grid(True, alpha=0.3)\n    \n    plt.tight_layout()\n    plt.show()\n    \nplot_distance_distributions(train_coords_dict, sample_size=100)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T03:51:28.564066Z","iopub.execute_input":"2026-01-09T03:51:28.564452Z","iopub.status.idle":"2026-01-09T03:51:29.978783Z","shell.execute_reply.started":"2026-01-09T03:51:28.564429Z","shell.execute_reply":"2026-01-09T03:51:29.977959Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 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]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_template_similarity_distribution(test_seqs_sample, train_seqs_df, train_coords_dict, n_samples=50):\n    \"\"\"\n    Plot distribution of template similarity scores found for test sequences.\n    \"\"\"\n    similarity_scores = []\n    \n    sampled_test = test_seqs_sample.sample(min(n_samples, len(test_seqs_sample)))\n    \n    for _, row in sampled_test.iterrows():\n        sequence = row['sequence']\n        temporal_cutoff = row.get('temporal_cutoff', None)\n        \n        similar_seqs = find_similar_sequences(sequence, train_seqs_df, train_coords_dict, \n                                            temporal_cutoff=temporal_cutoff, top_n=5)\n        \n        for _, _, sim_score, _ in similar_seqs:\n            similarity_scores.append(sim_score)\n    \n    fig, ax = plt.subplots(figsize=(10, 6))\n    ax.hist(similarity_scores, bins=30, edgecolor='black', alpha=0.7, color='mediumseagreen')\n    ax.axvline(np.mean(similarity_scores), color='red', linestyle='--', \n              label=f'Mean: {np.mean(similarity_scores):.3f}')\n    ax.axvline(np.median(similarity_scores), color='orange', linestyle='--', \n              label=f'Median: {np.median(similarity_scores):.3f}')\n    ax.set_xlabel('Template Similarity Score')\n    ax.set_ylabel('Frequency')\n    ax.set_title(f'Distribution of Template Similarity Scores\\n(Sampled from {n_samples} test sequences)')\n    ax.legend()\n    ax.grid(True, alpha=0.3)\n    plt.tight_layout()\n    plt.show()\n\n\ndef plot_prediction_comparison(target_id, sequence, predictions):\n    \"\"\"\n    Plot all 5 predictions for a single target to visualize ensemble diversity.\n    \"\"\"\n    fig = plt.figure(figsize=(15, 10))\n    \n    for i in range(5):\n        ax = fig.add_subplot(2, 3, i+1, projection='3d')\n        coords = predictions[i]\n        \n        # Plot backbone\n        ax.plot(coords[:, 0], coords[:, 1], coords[:, 2], 'b-', alpha=0.6, linewidth=1.5)\n        \n        # Plot residues\n        colors = {'A': 'red', 'C': 'blue', 'G': 'green', 'U': 'orange'}\n        for j, base in enumerate(sequence):\n            color = colors.get(base, 'gray')\n            ax.scatter(coords[j, 0], coords[j, 1], coords[j, 2], \n                      c=color, s=30, alpha=0.7)\n        \n        ax.set_xlabel('X (Å)')\n        ax.set_ylabel('Y (Å)')\n        ax.set_zlabel('Z (Å)')\n        ax.set_title(f'Prediction {i+1}')\n        \n    fig.suptitle(f'5 Structure Predictions for {target_id}\\nSequence length: {len(sequence)} nt', \n                fontsize=14, fontweight='bold')\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# List to store all prediction records\nall_predictions = []\n\n# Set up time tracking\nstart_time = time.time()\ntotal_targets = len(test_seqs)\n\nplot_template_similarity_distribution(test_seqs, train_seqs, train_coords_dict, n_samples=50)\n\n# For each sequence in the test set\nfor idx, row in test_seqs.iterrows():\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    # Progress tracking\n    if idx % 5 == 0:\n        elapsed = time.time() - start_time\n        targets_processed = idx + 1\n        if targets_processed > 0:\n            avg_time_per_target = elapsed / targets_processed\n            est_time_remaining = avg_time_per_target * (total_targets - targets_processed)\n            print(f\"Processing target {targets_processed}/{total_targets}: {target_id} ({len(sequence)} nt), \"\n                  f\"elapsed: {elapsed:.1f}s, est. remaining: {est_time_remaining:.1f}s\")\n    \n    # Generate 5 different structure predictions\n    predictions = predict_rna_structures(sequence, target_id, train_seqs, train_coords_dict, \n                                        n_predictions=5, temporal_cutoff=temporal_cutoff)\n\n    if idx < 2:  # Only plot first 2 to avoid too many plots\n        plot_prediction_comparison(target_id, sequence, predictions)\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('submission.csv', index=False)\nprint(f\"Generated predictions for {len(test_seqs)} RNA sequences\")\nprint(f\"Total runtime: {time.time() - start_time:.1f} seconds\")\nsubmission_df.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}