{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":118765,"databundleVersionId":15231210,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":14604295,"sourceType":"datasetVersion","datasetId":9328538}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install --no-index /kaggle/input/biopython-cp312/biopython-1.86-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\nStanford RNA 3D Folding Part 2 - Starter Code\nThis notebook provides a baseline approach combining:\n1. Template-based modeling (TBM) from PDB structures\n2. Simple structure prediction fallback\n3. Diversity generation for best-of-5 submissions\n\nStrategy: Template search is key to winning based on Part 1 results\n\"\"\"\n\nimport pandas as pd\nimport numpy as np\nimport os\nfrom pathlib import Path\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# For sequence alignment and template search\nfrom Bio import SeqIO, pairwise2\nfrom Bio.Seq import Seq\nfrom Bio.PDB import PDBParser, PDBIO, Select\nfrom Bio.PDB.Polypeptide import is_aa\nimport scipy.spatial.distance as dist\nfrom scipy.spatial.transform import Rotation\nimport json\n\n#==============================================================================\n# CONFIGURATION\n#==============================================================================\n\nclass Config:\n    # Paths (Kaggle competition paths)\n    INPUT_DIR = '/kaggle/input/stanford-rna-3d-folding-2'\n    MSA_DIR = f'{INPUT_DIR}/MSA'\n    PDB_DIR = f'{INPUT_DIR}/PDB_RNA'\n    \n    # Competition parameters\n    NUM_PREDICTIONS = 5  # Must submit 5 structures per target\n    COORD_MIN = -999.999\n    COORD_MAX = 9999.999\n    \n    # Template search parameters\n    MIN_SEQ_IDENTITY = 0.3  # Minimum sequence identity for template\n    MAX_TEMPLATES = 10       # Top N templates to consider\n    \n    # Random seed for reproducibility\n    RANDOM_SEED = 42\n\nnp.random.seed(Config.RANDOM_SEED)\n\n#==============================================================================\n# UTILITY FUNCTIONS\n#==============================================================================\n\ndef parse_fasta(fasta_string):\n    \"\"\"Parse FASTA formatted string into dictionary\"\"\"\n    sequences = {}\n    current_header = None\n    current_seq = []\n    \n    for line in fasta_string.strip().split('\\n'):\n        if line.startswith('>'):\n            if current_header:\n                sequences[current_header] = ''.join(current_seq)\n            current_header = line[1:].split()[0]  # Get first part of header\n            current_seq = []\n        else:\n            current_seq.append(line.strip())\n    \n    if current_header:\n        sequences[current_header] = ''.join(current_seq)\n    \n    return sequences\n\ndef calculate_sequence_identity(seq1, seq2):\n    \"\"\"Calculate sequence identity between two sequences\"\"\"\n    if len(seq1) != len(seq2):\n        # Use pairwise alignment for different lengths\n        alignments = pairwise2.align.globalxx(seq1, seq2)\n        if not alignments:\n            return 0.0\n        alignment = alignments[0]\n        matches = sum(a == b for a, b in zip(alignment[0], alignment[1]) if a != '-' and b != '-')\n        length = max(len(seq1), len(seq2))\n        return matches / length if length > 0 else 0.0\n    else:\n        # Direct comparison for same length\n        matches = sum(a == b for a, b in zip(seq1, seq2))\n        return matches / len(seq1)\n\ndef extract_c1_prime_coords(cif_file, chain_id='A'):\n    \"\"\"Extract C1' atom coordinates from CIF file\"\"\"\n    coords = []\n    residue_info = []\n    \n    try:\n        parser = PDBParser(QUIET=True)\n        structure = parser.get_structure('rna', cif_file)\n        \n        for model in structure:\n            for chain in model:\n                if chain.id != chain_id:\n                    continue\n                for residue in chain:\n                    if residue.id[0] != ' ':  # Skip hetero residues\n                        continue\n                    # Look for C1' atom\n                    if \"C1'\" in residue:\n                        atom = residue[\"C1'\"]\n                        coords.append(atom.coord)\n                        residue_info.append({\n                            'resname': residue.resname,\n                            'resid': residue.id[1],\n                            'chain': chain.id\n                        })\n    except Exception as e:\n        print(f\"Error parsing {cif_file}: {e}\")\n        return None, None\n    \n    if len(coords) == 0:\n        return None, None\n    \n    return np.array(coords), residue_info\n\ndef align_structures(pred_coords, ref_coords):\n    \"\"\"Align predicted coordinates to reference using Kabsch algorithm\"\"\"\n    # Center both structures\n    pred_center = pred_coords.mean(axis=0)\n    ref_center = ref_coords.mean(axis=0)\n    \n    pred_centered = pred_coords - pred_center\n    ref_centered = ref_coords - ref_center\n    \n    # Compute rotation matrix using SVD\n    H = pred_centered.T @ ref_centered\n    U, S, Vt = np.linalg.svd(H)\n    R = Vt.T @ U.T\n    \n    # Ensure proper rotation (det(R) = 1)\n    if np.linalg.det(R) < 0:\n        Vt[-1, :] *= -1\n        R = Vt.T @ U.T\n    \n    # Apply transformation\n    aligned_coords = (R @ pred_centered.T).T + ref_center\n    \n    return aligned_coords, R\n\ndef clip_coordinates(coords):\n    \"\"\"Clip coordinates to competition bounds\"\"\"\n    return np.clip(coords, Config.COORD_MIN, Config.COORD_MAX)\n\n#==============================================================================\n# TEMPLATE SEARCH\n#==============================================================================\n\nclass TemplateSearcher:\n    \"\"\"Search for structural templates in PDB database\"\"\"\n    \n    def __init__(self, pdb_dir, pdb_metadata_path):\n        self.pdb_dir = Path(pdb_dir)\n        self.templates = self._load_pdb_metadata(pdb_metadata_path)\n    \n    def _load_pdb_metadata(self, metadata_path):\n        \"\"\"Load PDB sequences for template search\"\"\"\n        templates = {}\n        \n        # Load from pdb_seqres_NA.fasta if available\n        fasta_path = self.pdb_dir / 'pdb_seqres_NA.fasta'\n        if fasta_path.exists():\n            with open(fasta_path, 'r') as f:\n                current_id = None\n                current_seq = []\n                for line in f:\n                    if line.startswith('>'):\n                        if current_id:\n                            templates[current_id] = ''.join(current_seq)\n                        current_id = line[1:].strip().split()[0]\n                        current_seq = []\n                    else:\n                        current_seq.append(line.strip())\n                if current_id:\n                    templates[current_id] = ''.join(current_seq)\n        \n        return templates\n    \n    def search_templates(self, query_sequence, top_k=10):\n        \"\"\"Find top-k templates by sequence similarity\"\"\"\n        matches = []\n        \n        for template_id, template_seq in self.templates.items():\n            identity = calculate_sequence_identity(query_sequence, template_seq)\n            if identity >= Config.MIN_SEQ_IDENTITY:\n                matches.append({\n                    'template_id': template_id,\n                    'sequence': template_seq,\n                    'identity': identity\n                })\n        \n        # Sort by identity\n        matches.sort(key=lambda x: x['identity'], reverse=True)\n        return matches[:top_k]\n\n#==============================================================================\n# STRUCTURE PREDICTOR\n#==============================================================================\n\nclass RNAStructurePredictor:\n    \"\"\"Main structure prediction class\"\"\"\n    \n    def __init__(self, pdb_dir, pdb_metadata_path=None):\n        self.pdb_dir = Path(pdb_dir)\n        \n        # Initialize template searcher if metadata available\n        if pdb_metadata_path and os.path.exists(pdb_metadata_path):\n            self.template_searcher = TemplateSearcher(pdb_dir, pdb_metadata_path)\n        else:\n            self.template_searcher = None\n            print(\"Warning: Template searcher not initialized\")\n    \n    def predict_from_template(self, sequence, template_coords, template_seq):\n        \"\"\"Build structure from template by mapping sequence to template\"\"\"\n        # Simple approach: direct mapping if sequences align well\n        predictions = []\n        \n        # Generate alignment\n        alignments = pairwise2.align.globalxx(sequence, template_seq)\n        if not alignments:\n            return None\n        \n        alignment = alignments[0]\n        aligned_query = alignment[0]\n        aligned_template = alignment[1]\n        \n        # Map query positions to template coordinates\n        query_idx = 0\n        template_idx = 0\n        mapped_coords = []\n        \n        for i in range(len(aligned_query)):\n            if aligned_query[i] != '-':\n                if aligned_template[i] != '-':\n                    # Both aligned - use template coordinate\n                    mapped_coords.append(template_coords[template_idx])\n                else:\n                    # Gap in template - interpolate or use previous\n                    if len(mapped_coords) > 0:\n                        mapped_coords.append(mapped_coords[-1] + np.random.randn(3) * 0.5)\n                    else:\n                        mapped_coords.append(np.random.randn(3) * 10)\n                query_idx += 1\n            \n            if aligned_template[i] != '-':\n                template_idx += 1\n        \n        return np.array(mapped_coords) if mapped_coords else None\n    \n    def predict_ab_initio(self, sequence):\n        \"\"\"Fallback: Generate random coil structure\"\"\"\n        n_residues = len(sequence)\n        \n        # Simple random walk with RNA-like distances\n        coords = np.zeros((n_residues, 3))\n        \n        # Start at origin\n        coords[0] = [0, 0, 0]\n        \n        # Generate random walk with ~6 Angstrom steps\n        for i in range(1, n_residues):\n            direction = np.random.randn(3)\n            direction = direction / np.linalg.norm(direction)\n            coords[i] = coords[i-1] + direction * 6.0\n        \n        return coords\n    \n    def generate_diverse_predictions(self, base_coords, num_predictions=5):\n        \"\"\"Generate diverse predictions from base structure\"\"\"\n        predictions = [base_coords]\n        \n        for i in range(num_predictions - 1):\n            # Add random perturbations\n            noise_scale = 0.5 + i * 0.3  # Increasing diversity\n            perturbed = base_coords + np.random.randn(*base_coords.shape) * noise_scale\n            \n            # Add rotation\n            rotation = Rotation.random(random_state=Config.RANDOM_SEED + i)\n            center = perturbed.mean(axis=0)\n            perturbed = rotation.apply(perturbed - center) + center\n            \n            predictions.append(perturbed)\n        \n        return predictions\n    \n    def predict(self, target_id, sequence, msa_file=None):\n        \"\"\"Main prediction function\"\"\"\n        print(f\"Predicting structure for {target_id} (length: {len(sequence)})\")\n        \n        # Try template-based modeling\n        best_coords = None\n        \n        if self.template_searcher:\n            templates = self.template_searcher.search_templates(sequence, top_k=Config.MAX_TEMPLATES)\n            \n            if templates:\n                print(f\"  Found {len(templates)} templates, best identity: {templates[0]['identity']:.3f}\")\n                \n                # Try to use best template\n                template = templates[0]\n                template_pdb_id = template['template_id'].split('_')[0]\n                template_cif = self.pdb_dir / f\"{template_pdb_id}.cif\"\n                \n                if template_cif.exists():\n                    template_coords, _ = extract_c1_prime_coords(str(template_cif))\n                    if template_coords is not None:\n                        best_coords = self.predict_from_template(\n                            sequence, template_coords, template['sequence']\n                        )\n        \n        # Fallback to ab initio if template fails\n        if best_coords is None:\n            print(f\"  Using ab initio prediction\")\n            best_coords = self.predict_ab_initio(sequence)\n        \n        # Generate 5 diverse predictions\n        predictions = self.generate_diverse_predictions(best_coords, Config.NUM_PREDICTIONS)\n        \n        # Clip coordinates\n        predictions = [clip_coordinates(pred) for pred in predictions]\n        \n        return predictions\n\n#==============================================================================\n# SUBMISSION GENERATION\n#==============================================================================\n\ndef create_submission(test_sequences_path, output_path='submission.csv'):\n    \"\"\"Create submission file\"\"\"\n    \n    # Load test sequences\n    test_df = pd.read_csv(test_sequences_path)\n    print(f\"Loaded {len(test_df)} test sequences\")\n    \n    # Initialize predictor\n    pdb_metadata = os.path.join(Config.INPUT_DIR, 'PDB_RNA', 'pdb_seqres_NA.fasta')\n    predictor = RNAStructurePredictor(Config.PDB_DIR, pdb_metadata)\n    \n    # Generate predictions\n    all_predictions = []\n    \n    for idx, row in test_df.iterrows():\n        target_id = row['target_id']\n        sequence = row['sequence']\n        \n        # Get MSA file if available\n        msa_file = Path(Config.MSA_DIR) / f\"{target_id}.MSA.fasta\"\n        msa_file = msa_file if msa_file.exists() else None\n        \n        # Predict structures\n        predictions = predictor.predict(target_id, sequence, msa_file)\n        \n        # Format for submission\n        for resid in range(1, len(sequence) + 1):\n            resname = sequence[resid - 1]  # Get nucleotide\n            \n            # Create row with 5 predictions\n            row_data = {\n                'ID': f\"{target_id}_{resid}\",\n                'resname': resname,\n                'resid': resid\n            }\n            \n            # Add coordinates for each prediction\n            for pred_idx in range(Config.NUM_PREDICTIONS):\n                coords = predictions[pred_idx][resid - 1]\n                row_data[f'x_{pred_idx + 1}'] = coords[0]\n                row_data[f'y_{pred_idx + 1}'] = coords[1]\n                row_data[f'z_{pred_idx + 1}'] = coords[2]\n            \n            all_predictions.append(row_data)\n    \n    # Create DataFrame and save\n    submission_df = pd.DataFrame(all_predictions)\n    submission_df.to_csv(output_path, index=False)\n    print(f\"\\nSubmission saved to {output_path}\")\n    print(f\"Shape: {submission_df.shape}\")\n    print(f\"Columns: {list(submission_df.columns)}\")\n    \n    return submission_df\n\n#==============================================================================\n# MAIN\n#==============================================================================\n\nif __name__ == \"__main__\":\n    # Create submission\n    test_path = os.path.join(Config.INPUT_DIR, 'test_sequences.csv')\n    submission = create_submission(test_path, 'submission.csv')\n    \n    print(\"\\n\" + \"=\"*60)\n    print(\"Submission complete!\")\n    print(\"=\"*60)\n    print(\"\\nNext steps to improve:\")\n    print(\"1. Implement better template search (use MMseqs2, HHblits)\")\n    print(\"2. Add deep learning model (RNAPro, RibonanzaNet2)\")\n    print(\"3. Improve diversity generation (clustering, energy minimization)\")\n    print(\"4. Use MSA information more effectively\")\n    print(\"5. Add structure refinement (Rosetta, molecular dynamics)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T12:49:27.261592Z","iopub.execute_input":"2026-02-01T12:49:27.261835Z","execution_failed":"2026-02-01T16:06:35.387Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}