{"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":"none","dataSources":[{"sourceId":118765,"databundleVersionId":15231210,"sourceType":"competition"},{"sourceId":14445506,"sourceType":"datasetVersion","datasetId":9227295}],"dockerImageVersionId":31234,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#!/usr/bin/env python\n# coding: utf-8\n\n\"\"\"\nStanford RNA 3D Folding Part 2 - Advanced Submission (Target: 0.7+ TM-score)\n============================================================================\n\nThis notebook implements a comprehensive approach to RNA structure prediction:\n1. MSA-based contact prediction using coevolution analysis\n2. Secondary structure prediction\n3. Distance geometry-based 3D modeling\n4. Ensemble generation with geometric diversity\n\"\"\"\n\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom scipy.spatial.distance import pdist, squareform\nfrom scipy.optimize import minimize\nfrom collections import defaultdict, Counter\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# ============================================================================\n# CONFIGURATION\n# ============================================================================\n\nclass Config:\n    # Paths\n    INPUT_DIR = Path('/kaggle/input/stanford-rna-3d-folding-2')\n    MSA_DIR = INPUT_DIR / 'MSA'\n    TEST_FILE = 'test_sequences.csv'\n    OUTPUT_FILE = 'submission.csv'\n    \n    # Prediction settings\n    NUM_PREDICTIONS = 5\n    MIN_MSA_DEPTH = 10  # Minimum MSA sequences for contact prediction\n    TOP_CONTACTS_FRAC = 1.0  # Use top contacts\n    \n    # Structure parameters\n    C1_C1_DISTANCE = 5.9  # Average C1'-C1' distance in RNA\n    BOND_LENGTH_STD = 0.3\n    \n    # Watson-Crick base pair distances\n    WC_DISTANCES = {\n        ('A', 'U'): 10.7,\n        ('U', 'A'): 10.7,\n        ('G', 'C'): 10.5,\n        ('C', 'G'): 10.5,\n        ('G', 'U'): 11.0,\n        ('U', 'G'): 11.0,\n    }\n\n# ============================================================================\n# MSA PROCESSING AND CONTACT PREDICTION\n# ============================================================================\n\ndef parse_fasta(fasta_file):\n    \"\"\"Parse FASTA file into sequences\"\"\"\n    sequences = []\n    current_seq = []\n    \n    try:\n        with open(fasta_file, 'r') as f:\n            for line in f:\n                line = line.strip()\n                if line.startswith('>'):\n                    if current_seq:\n                        sequences.append(''.join(current_seq))\n                        current_seq = []\n                else:\n                    current_seq.append(line)\n            \n            if current_seq:\n                sequences.append(''.join(current_seq))\n    except Exception as e:\n        msg = \"Warning: Could not read MSA file {}: {}\".format(fasta_file, e)\n        print(msg)\n        return []\n    \n    return sequences\n\ndef compute_mutual_information(msa):\n    \"\"\"Compute mutual information matrix from MSA\"\"\"\n    n_seqs = len(msa)\n    if n_seqs < 10:\n        return None\n    \n    seq_len = len(msa[0])\n    mi_matrix = np.zeros((seq_len, seq_len))\n    \n    # Convert sequences to numerical representation\n    bases = 'ACGU-'\n    base_to_idx = {b: i for i, b in enumerate(bases)}\n    \n    # Convert MSA to numerical matrix\n    msa_numeric = np.zeros((n_seqs, seq_len), dtype=int)\n    for i, seq in enumerate(msa):\n        for j, base in enumerate(seq):\n            msa_numeric[i, j] = base_to_idx.get(base, 4)\n    \n    # Compute MI for each pair of positions\n    for i in range(seq_len):\n        for j in range(i + 1, seq_len):\n            # Get columns\n            col_i = msa_numeric[:, i]\n            col_j = msa_numeric[:, j]\n            \n            # Compute joint and marginal frequencies\n            joint_freq = np.zeros((5, 5))\n            for ii, jj in zip(col_i, col_j):\n                joint_freq[ii, jj] += 1\n            joint_freq /= n_seqs\n            \n            freq_i = np.bincount(col_i, minlength=5) / n_seqs\n            freq_j = np.bincount(col_j, minlength=5) / n_seqs\n            \n            # Compute MI\n            mi = 0.0\n            for ii in range(5):\n                for jj in range(5):\n                    if joint_freq[ii, jj] > 0:\n                        mi += joint_freq[ii, jj] * np.log(\n                            joint_freq[ii, jj] / (freq_i[ii] * freq_j[jj] + 1e-10) + 1e-10\n                        )\n            \n            mi_matrix[i, j] = mi\n            mi_matrix[j, i] = mi\n    \n    return mi_matrix\n\ndef apply_apc(mi_matrix):\n    \"\"\"Apply Average Product Correction to MI matrix\"\"\"\n    if mi_matrix is None:\n        return None\n    \n    n = mi_matrix.shape[0]\n    apc_matrix = mi_matrix.copy()\n    \n    # Compute mean MI for each position\n    mean_i = np.mean(mi_matrix, axis=1)\n    mean_all = np.mean(mi_matrix)\n    \n    # Apply APC correction\n    for i in range(n):\n        for j in range(n):\n            if i != j:\n                apc_matrix[i, j] = mi_matrix[i, j] - (mean_i[i] * mean_i[j]) / mean_all\n    \n    return apc_matrix\n\ndef predict_contacts_from_msa(msa, sequence, top_frac=1.0):\n    \"\"\"Predict contacts from MSA using MI+APC\"\"\"\n    if len(msa) < 10:\n        return []\n    \n    # Compute MI matrix\n    mi_matrix = compute_mutual_information(msa)\n    if mi_matrix is None:\n        return []\n    \n    # Apply APC\n    contact_matrix = apply_apc(mi_matrix)\n    if contact_matrix is None:\n        return []\n    \n    # Extract contacts (excluding local interactions)\n    contacts = []\n    seq_len = len(sequence)\n    \n    for i in range(seq_len):\n        for j in range(i + 6, seq_len):  # Exclude local contacts\n            score = contact_matrix[i, j]\n            contacts.append((i, j, score))\n    \n    # Sort by score and take top contacts\n    contacts.sort(key=lambda x: -x[2])\n    n_contacts = int(seq_len * top_frac)\n    top_contacts = contacts[:max(n_contacts, seq_len // 2)]\n    \n    return top_contacts\n\n# ============================================================================\n# SECONDARY STRUCTURE PREDICTION\n# ============================================================================\n\ndef predict_secondary_structure(sequence):\n    \"\"\"Predict secondary structure using simple base pairing rules\"\"\"\n    n = len(sequence)\n    pairs = []\n    \n    # Simple Nussinov-like algorithm\n    dp = np.zeros((n, n), dtype=int)\n    traceback = {}\n    \n    def can_pair(i, j):\n        \"\"\"Check if bases can form Watson-Crick or wobble pair\"\"\"\n        pair = (sequence[i], sequence[j])\n        return pair in [('A', 'U'), ('U', 'A'), ('G', 'C'), ('C', 'G'), \n                       ('G', 'U'), ('U', 'G')]\n    \n    # Fill DP table\n    for length in range(5, n + 1):\n        for i in range(n - length + 1):\n            j = i + length - 1\n            \n            # Case 1: j is unpaired\n            dp[i][j] = dp[i][j-1] if j > 0 else 0\n            \n            # Case 2: j pairs with some k\n            for k in range(i, j):\n                if can_pair(k, j) and j - k >= 4:\n                    score = (1 if k > 0 else 0)\n                    left = dp[i][k-1] if k > i else 0\n                    inside = dp[k+1][j-1] if k + 1 < j else 0\n                    \n                    new_score = left + inside + score\n                    if new_score > dp[i][j]:\n                        dp[i][j] = new_score\n                        traceback[(i, j)] = k\n    \n    # Traceback to get pairs\n    def traceback_pairs(i, j):\n        if i >= j:\n            return\n        \n        if (i, j) in traceback:\n            k = traceback[(i, j)]\n            pairs.append((k, j))\n            if k > i:\n                traceback_pairs(i, k - 1)\n            if k + 1 < j:\n                traceback_pairs(k + 1, j - 1)\n        elif j > 0:\n            traceback_pairs(i, j - 1)\n    \n    traceback_pairs(0, n - 1)\n    \n    return pairs\n\n# ============================================================================\n# 3D STRUCTURE BUILDING\n# ============================================================================\n\ndef build_distance_matrix(sequence, contacts, base_pairs):\n    \"\"\"Build target distance matrix from contacts and base pairs\"\"\"\n    n = len(sequence)\n    dist_matrix = np.full((n, n), 100.0)  # Initialize with large distances\n    \n    # Set diagonal\n    np.fill_diagonal(dist_matrix, 0.0)\n    \n    # Backbone connectivity\n    for i in range(n - 1):\n        dist_matrix[i, i + 1] = Config.C1_C1_DISTANCE\n        dist_matrix[i + 1, i] = Config.C1_C1_DISTANCE\n    \n    # Base pair distances\n    for i, j in base_pairs:\n        pair = (sequence[i], sequence[j])\n        if pair in Config.WC_DISTANCES:\n            dist = Config.WC_DISTANCES[pair]\n            dist_matrix[i, j] = dist\n            dist_matrix[j, i] = dist\n    \n    # Contact-derived distances\n    for i, j, score in contacts[:len(sequence)]:\n        # Use score to weight the distance\n        base_dist = 12.0  # Average contact distance\n        weighted_dist = base_dist * (1.0 / (1.0 + score))\n        if dist_matrix[i, j] > weighted_dist:\n            dist_matrix[i, j] = weighted_dist\n            dist_matrix[j, i] = weighted_dist\n    \n    return dist_matrix\n\ndef distance_geometry_embedding(dist_matrix, dim=3):\n    \"\"\"Convert distance matrix to 3D coordinates using MDS-like approach\"\"\"\n    n = dist_matrix.shape[0]\n    \n    # Use simple distance geometry with optimization\n    # Initialize with random coordinates\n    coords = np.random.randn(n, dim) * 10.0\n    \n    # Place first few residues in a line\n    for i in range(min(3, n)):\n        coords[i] = np.array([i * 5.9, 0, 0])\n    \n    def objective(flat_coords):\n        coords = flat_coords.reshape(n, dim)\n        error = 0.0\n        \n        # Pairwise distance errors\n        for i in range(n):\n            for j in range(i + 1, n):\n                target_dist = dist_matrix[i, j]\n                if target_dist < 50:  # Only consider reasonable distances\n                    actual_dist = np.linalg.norm(coords[i] - coords[j])\n                    weight = 1.0 / (target_dist + 1.0)\n                    error += weight * (actual_dist - target_dist) ** 2\n        \n        return error\n    \n    # Optimize\n    try:\n        flat_coords = coords.flatten()\n        result = minimize(objective, flat_coords, method='L-BFGS-B', \n                         options={'maxiter': 1000, 'ftol': 1e-6})\n        coords = result.x.reshape(n, dim)\n    except:\n        pass\n    \n    return coords\n\ndef add_noise_to_structure(coords, noise_level=1.0):\n    \"\"\"Add Gaussian noise to create diverse ensemble members\"\"\"\n    noisy_coords = coords + np.random.randn(*coords.shape) * noise_level\n    return noisy_coords\n\ndef refine_structure(coords, dist_matrix, iterations=100):\n    \"\"\"Refine structure using gradient descent\"\"\"\n    n = len(coords)\n    learning_rate = 0.01\n    \n    for iteration in range(iterations):\n        forces = np.zeros_like(coords)\n        \n        # Compute forces from distance constraints\n        for i in range(n):\n            for j in range(i + 1, n):\n                target_dist = dist_matrix[i, j]\n                if target_dist < 50:\n                    vec = coords[j] - coords[i]\n                    actual_dist = np.linalg.norm(vec)\n                    if actual_dist > 0:\n                        direction = vec / actual_dist\n                        force_mag = (actual_dist - target_dist) * 0.1\n                        forces[i] += direction * force_mag\n                        forces[j] -= direction * force_mag\n        \n        # Update coordinates\n        coords += forces * learning_rate\n        \n        # Decay learning rate\n        learning_rate *= 0.99\n    \n    return coords\n\n# ============================================================================\n# RNA STRUCTURE PREDICTOR\n# ============================================================================\n\nclass RNAPredictor:\n    def __init__(self, config):\n        self.config = config\n    \n    def load_msa(self, target_id):\n        \"\"\"Load MSA for target\"\"\"\n        msa_file = self.config.MSA_DIR / \"{}.MSA.fasta\".format(target_id)\n        if msa_file.exists():\n            return parse_fasta(msa_file)\n        return []\n    \n    def predict_structure(self, target_id, sequence):\n        \"\"\"Predict 3D structure for a sequence\"\"\"\n        print(\"  Predicting {} (length: {})\".format(target_id, len(sequence)))\n        \n        # Load MSA\n        msa = self.load_msa(target_id)\n        print(\"    MSA depth: {}\".format(len(msa)))\n        \n        # Predict contacts from MSA\n        contacts = []\n        if len(msa) >= self.config.MIN_MSA_DEPTH:\n            contacts = predict_contacts_from_msa(msa, sequence, \n                                                self.config.TOP_CONTACTS_FRAC)\n            print(\"    Predicted {} contacts\".format(len(contacts)))\n        \n        # Predict secondary structure\n        base_pairs = predict_secondary_structure(sequence)\n        print(\"    Predicted {} base pairs\".format(len(base_pairs)))\n        \n        # Build distance matrix\n        dist_matrix = build_distance_matrix(sequence, contacts, base_pairs)\n        \n        # Generate ensemble of structures\n        structures = []\n        for i in range(self.config.NUM_PREDICTIONS):\n            # Initial embedding\n            coords = distance_geometry_embedding(dist_matrix, dim=3)\n            \n            # Add diversity\n            if i > 0:\n                coords = add_noise_to_structure(coords, noise_level=0.5 + i * 0.5)\n            \n            # Refine structure\n            coords = refine_structure(coords, dist_matrix, iterations=50 + i * 20)\n            \n            # Clip coordinates\n            coords = np.clip(coords, -999.999, 9999.999)\n            \n            structures.append(coords)\n        \n        return structures\n    \n    def predict_all(self, test_df):\n        \"\"\"Predict structures for all test sequences\"\"\"\n        predictions = {}\n        \n        for idx, row in test_df.iterrows():\n            target_id = row['target_id']\n            sequence = row['sequence']\n            \n            print(\"\\n[{}/{}] {}\".format(idx+1, len(test_df), target_id))\n            structures = self.predict_structure(target_id, sequence)\n            predictions[target_id] = structures\n        \n        return predictions\n\n# ============================================================================\n# SUBMISSION CREATION\n# ============================================================================\n\ndef create_submission(predictions, test_df, config):\n    \"\"\"Create submission file\"\"\"\n    rows = []\n    \n    for _, row in test_df.iterrows():\n        target_id = row['target_id']\n        sequence = row['sequence']\n        \n        if target_id not in predictions:\n            continue\n        \n        structures = predictions[target_id]\n        \n        for res_idx, nucleotide in enumerate(sequence):\n            res_id = res_idx + 1\n            \n            row_data = {\n                'ID': '{}_{}'.format(target_id, res_id),\n                'resname': nucleotide,\n                'resid': res_id\n            }\n            \n            # Add coordinates from all 5 predictions\n            for pred_idx in range(config.NUM_PREDICTIONS):\n                coords = structures[pred_idx][res_idx]\n                row_data['x_{}'.format(pred_idx+1)] = float(coords[0])\n                row_data['y_{}'.format(pred_idx+1)] = float(coords[1])\n                row_data['z_{}'.format(pred_idx+1)] = float(coords[2])\n            \n            rows.append(row_data)\n    \n    df = pd.DataFrame(rows)\n    \n    # Ensure correct column order\n    coord_cols = []\n    for i in range(1, config.NUM_PREDICTIONS + 1):\n        coord_cols.extend(['x_{}'.format(i), 'y_{}'.format(i), 'z_{}'.format(i)])\n    \n    columns = ['ID', 'resname', 'resid'] + coord_cols\n    df = df[columns]\n    \n    return df\n\n# ============================================================================\n# MAIN EXECUTION\n# ============================================================================\n\ndef main():\n    print(\"=\" * 80)\n    print(\"Stanford RNA 3D Folding Part 2 - Advanced Submission\")\n    print(\"=\" * 80)\n    \n    config = Config()\n    \n    # Load test data\n    print(\"\\nLoading test data...\")\n    test_df = pd.read_csv(config.INPUT_DIR / config.TEST_FILE)\n    print(\"Loaded {} test sequences\".format(len(test_df)))\n    \n    # Initialize predictor\n    print(\"\\nInitializing predictor...\")\n    predictor = RNAPredictor(config)\n    \n    # Generate predictions\n    print(\"\\nGenerating predictions...\")\n    predictions = predictor.predict_all(test_df)\n    \n    # Create submission\n    print(\"\\nCreating submission file...\")\n    submission_df = create_submission(predictions, test_df, config)\n    \n    # Save submission\n    submission_df.to_csv(config.OUTPUT_FILE, index=False)\n    print(\"\\nSubmission saved to {}\".format(config.OUTPUT_FILE))\n    print(\"Total rows: {}\".format(len(submission_df)))\n    print(\"\\n\" + \"=\" * 80)\n    print(\"✓ Submission ready!\")\n    print(\"=\" * 80)\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:56:44.112381Z","iopub.execute_input":"2026-01-27T08:56:44.112666Z","iopub.status.idle":"2026-01-27T09:00:30.256391Z","shell.execute_reply.started":"2026-01-27T08:56:44.112642Z","shell.execute_reply":"2026-01-27T09:00:30.255525Z"}},"outputs":[],"execution_count":null}]}