{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":118765,"databundleVersionId":15231210,"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/datasets/kami1976/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,"execution":{"iopub.status.busy":"2026-02-11T18:27:22.481283Z","iopub.execute_input":"2026-02-11T18:27:22.481879Z","iopub.status.idle":"2026-02-11T18:27:27.58175Z","shell.execute_reply.started":"2026-02-11T18:27:22.481852Z","shell.execute_reply":"2026-02-11T18:27:27.581096Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# Stanford RNA 3D Folding Part 2 - Enhanced Template-Based Approach\n# ============================================================\n#\n# Team: DTU Compute\n# Department of Applied Mathematics and Computer Science\n# Technical University of Denmark\n#\n# Team Members:\n#   Olaf Yunus Laitinen Imanov (oyli@dtu.dk)\n#   Technical University of Denmark, Kongens Lyngby, Denmark\n#   ORCID: 0009-0006-5184-0810\n#\n#   Derya Umut Kulali (d_u_k@ogr.eskisehir.edu.tr)\n#   Eskisehir Technical University, Eskisehir, Türkiye\n#   ORCID: 0009-0004-8844-6601\n#\n# Approach: Template-based structure prediction with MSA analysis,\n#           enhanced template scoring, and physics-based refinement\n#\n# Dependencies: biopython, numpy, pandas, scipy\n# Runtime: CPU-only, deterministic, < 8 hours\n# ============================================================\n\nimport pandas as pd\nimport numpy as np\nimport random\nimport time\nimport warnings\nimport os\nimport sys\n\nwarnings.filterwarnings(\"ignore\")\n\n# ============================================================\n# SECTION A: DATA LOADING AND INITIALIZATION\n# ============================================================\n\nDATA_PATH = \"/kaggle/input/stanford-rna-3d-folding-2/\"\n\n# Load primary datasets\ntrain_seqs = pd.read_csv(DATA_PATH + \"train_sequences.csv\")\ntest_seqs = pd.read_csv(DATA_PATH + \"test_sequences.csv\")\ntrain_labels = pd.read_csv(DATA_PATH + \"train_labels.csv\")\n\nsys.path.append(os.path.join(DATA_PATH, \"extra\"))\n\n# ============================================================\n# SECTION B: FASTA PARSING UTILITIES\n# ============================================================\n\ndef setup_fasta_parser():\n    \"\"\"\n    Setup FASTA parser with fallback mechanism.\n    Attempts to use competition-provided parser, falls back to custom implementation.\n    \"\"\"\n    try:\n        import typing as _typing\n        import builtins as _builtins\n\n        _builtins.Dict = getattr(_typing, \"Dict\")\n        _builtins.Tuple = getattr(_typing, \"Tuple\")\n        _builtins.List = getattr(_typing, \"List\")\n\n        from parse_fasta_py import parse_fasta as _parse_fasta_raw\n\n        def parse_fasta(fasta_content: str):\n            \"\"\"Normalize parse_fasta output to {chain_id: sequence_string}\"\"\"\n            d = _parse_fasta_raw(fasta_content)\n            out = {}\n            for k, v in d.items():\n                out[k] = v[0] if isinstance(v, tuple) else v\n            return out\n\n        return parse_fasta\n\n    except Exception:\n        def parse_fasta(fasta_content: str):\n            \"\"\"\n            Fallback FASTA parser implementation.\n            Returns: {chain_id: sequence_string}\n            \"\"\"\n            out = {}\n            cur = None\n            seq_parts = []\n            for line in str(fasta_content).splitlines():\n                line = line.strip()\n                if not line:\n                    continue\n                if line.startswith(\">\"):\n                    if cur is not None:\n                        out[cur] = \"\".join(seq_parts)\n                    header = line[1:]\n                    cur = header.split()[0]\n                    seq_parts = []\n                else:\n                    seq_parts.append(line.replace(\" \", \"\"))\n            if cur is not None:\n                out[cur] = \"\".join(seq_parts)\n            return out\n\n        return parse_fasta\n\nparse_fasta = setup_fasta_parser()\n\n# ============================================================\n# SECTION C: MSA LOADING AND COEVOLUTION ANALYSIS\n# ============================================================\n\ndef load_msa(target_id):\n    \"\"\"\n    Load Multiple Sequence Alignment for a target.\n    \n    Args:\n        target_id: Target identifier\n        \n    Returns:\n        MSA alignment object or None if not found\n    \"\"\"\n    msa_path = f\"{DATA_PATH}/MSA/{target_id}.MSA.fasta\"\n    if os.path.exists(msa_path):\n        try:\n            from Bio import AlignIO\n            return AlignIO.read(msa_path, \"fasta\")\n        except Exception:\n            return None\n    return None\n\ndef calculate_mutual_information(msa, position_i, position_j):\n    \"\"\"\n    Calculate Mutual Information between two positions in MSA.\n    Used to detect coevolving residue pairs (potential base pairs).\n    \n    Args:\n        msa: Multiple sequence alignment\n        position_i, position_j: Positions to analyze\n        \n    Returns:\n        Mutual information score (higher indicates stronger coevolution)\n    \"\"\"\n    if msa is None or len(msa) < 5:\n        return 0.0\n    \n    try:\n        chars_i = [str(record.seq[position_i]) for record in msa if len(record.seq) > position_i]\n        chars_j = [str(record.seq[position_j]) for record in msa if len(record.seq) > position_j]\n        \n        if len(chars_i) < 5 or len(chars_j) < 5:\n            return 0.0\n        \n        # Calculate frequencies\n        from collections import Counter\n        freq_i = Counter(chars_i)\n        freq_j = Counter(chars_j)\n        freq_ij = Counter(zip(chars_i, chars_j))\n        \n        n = len(chars_i)\n        mi = 0.0\n        \n        for (ci, cj), count_ij in freq_ij.items():\n            p_ij = count_ij / n\n            p_i = freq_i[ci] / n\n            p_j = freq_j[cj] / n\n            \n            if p_ij > 0 and p_i > 0 and p_j > 0:\n                mi += p_ij * np.log2(p_ij / (p_i * p_j))\n        \n        return max(0.0, mi)\n    except Exception:\n        return 0.0\n\ndef get_coevolution_pairs(target_id, sequence, threshold=0.3):\n    \"\"\"\n    Identify coevolving residue pairs from MSA.\n    These pairs are likely to form base pairs in the 3D structure.\n    \n    Args:\n        target_id: Target identifier\n        sequence: RNA sequence\n        threshold: Minimum MI score to consider\n        \n    Returns:\n        List of (position_i, position_j, score) tuples\n    \"\"\"\n    msa = load_msa(target_id)\n    if msa is None:\n        return []\n    \n    pairs = []\n    seq_len = len(sequence)\n    \n    # Only check pairs with reasonable separation (potential base pairs)\n    for i in range(seq_len):\n        for j in range(i + 4, seq_len):\n            if j - i > 100:\n                continue\n            \n            mi = calculate_mutual_information(msa, i, j)\n            if mi > threshold:\n                pairs.append((i, j, mi))\n    \n    # Sort by MI score descending\n    pairs.sort(key=lambda x: x[2], reverse=True)\n    return pairs\n\n# ============================================================\n# SECTION D: STOICHIOMETRY AND CHAIN SEGMENT PARSING\n# ============================================================\n\ndef parse_stoichiometry(stoich: str):\n    \"\"\"\n    Parse stoichiometry string into list of (chain_id, count) tuples.\n    \n    Args:\n        stoich: Stoichiometry string (e.g., \"A:2;B:1\")\n        \n    Returns:\n        List of (chain_id, count) tuples\n    \"\"\"\n    if pd.isna(stoich) or str(stoich).strip() == \"\":\n        return []\n    out = []\n    for part in str(stoich).split(\";\"):\n        ch, cnt = part.split(\":\")\n        out.append((ch.strip(), int(cnt)))\n    return out\n\ndef get_chain_segments(row):\n    \"\"\"\n    Determine chain segment boundaries within concatenated sequence.\n    \n    Args:\n        row: DataFrame row with sequence, stoichiometry, all_sequences\n        \n    Returns:\n        List of (start, end) tuples for each chain segment\n    \"\"\"\n    seq = row[\"sequence\"]\n    stoich = row.get(\"stoichiometry\", \"\")\n    all_seq = row.get(\"all_sequences\", \"\")\n\n    if pd.isna(stoich) or pd.isna(all_seq) or str(stoich).strip() == \"\" or str(all_seq).strip() == \"\":\n        return [(0, len(seq))]\n\n    try:\n        chain_dict = parse_fasta(all_seq)\n        order = parse_stoichiometry(stoich)\n\n        segs = []\n        pos = 0\n        for ch, cnt in order:\n            base = chain_dict.get(ch)\n            if base is None:\n                return [(0, len(seq))]\n            for _ in range(cnt):\n                L = len(base)\n                segs.append((pos, pos + L))\n                pos += L\n\n        if pos != len(seq):\n            return [(0, len(seq))]\n        return segs\n    except Exception:\n        return [(0, len(seq))]\n\ndef build_segments_map(df):\n    \"\"\"\n    Build segment and stoichiometry maps for all targets in dataframe.\n    \n    Args:\n        df: DataFrame with target sequences\n        \n    Returns:\n        seg_map: {target_id: [(start, end), ...]}\n        stoich_map: {target_id: stoichiometry_string}\n    \"\"\"\n    seg_map = {}\n    stoich_map = {}\n    for _, r in df.iterrows():\n        tid = r[\"target_id\"]\n        seg_map[tid] = get_chain_segments(r)\n        stoich_map[tid] = str(r.get(\"stoichiometry\", \"\") if not pd.isna(r.get(\"stoichiometry\", \"\")) else \"\")\n    return seg_map, stoich_map\n\ntrain_segs_map, train_stoich_map = build_segments_map(train_seqs)\ntest_segs_map, test_stoich_map = build_segments_map(test_seqs)\n\n# ============================================================\n# SECTION E: TEMPLATE COORDINATE EXTRACTION\n# ============================================================\n\ndef process_labels(labels_df: pd.DataFrame):\n    \"\"\"\n    Extract template coordinates from training labels.\n    \n    Args:\n        labels_df: Training labels with coordinates\n        \n    Returns:\n        coords_dict: {target_id: (L, 3) numpy array of coordinates}\n    \"\"\"\n    coords_dict = {}\n    prefixes = labels_df[\"ID\"].str.rsplit(\"_\", n=1).str[0]\n    for id_prefix, group in labels_df.groupby(prefixes, sort=False):\n        coords_dict[id_prefix] = group.sort_values(\"resid\")[[\"x_1\", \"y_1\", \"z_1\"]].values\n    return coords_dict\n\ntrain_coords_dict = process_labels(train_labels)\n\n# ============================================================\n# SECTION F: ENHANCED TEMPLATE QUALITY ASSESSMENT\n# ============================================================\n\ndef get_structure_quality(target_id):\n    \"\"\"\n    Assess template structure quality based on PDB metadata.\n    Higher scores indicate better quality templates.\n    \n    Args:\n        target_id: PDB identifier\n        \n    Returns:\n        Quality score (0.5 to 1.5)\n    \"\"\"\n    # Default quality score\n    base_score = 1.0\n    \n    # Check if we have metadata about this structure\n    # In a full implementation, this would parse:\n    # - Resolution (X-ray/cryo-EM)\n    # - R-factor\n    # - Completeness\n    # - Experimental method\n    \n    # For now, use simple heuristics based on coordinate quality\n    if target_id in train_coords_dict:\n        coords = train_coords_dict[target_id]\n        \n        # Penalize structures with unusual coordinate ranges\n        coord_range = np.ptp(coords, axis=0).max()\n        if coord_range > 500:\n            base_score *= 0.8\n        elif coord_range < 20:\n            base_score *= 0.9\n        \n        # Reward structures with reasonable bond lengths\n        if len(coords) > 1:\n            dists = np.linalg.norm(coords[1:] - coords[:-1], axis=1)\n            mean_dist = np.mean(dists)\n            if 4.5 < mean_dist < 7.5:\n                base_score *= 1.1\n    \n    return np.clip(base_score, 0.5, 1.5)\n\n# ============================================================\n# SECTION G: SEQUENCE ALIGNMENT AND TEMPLATE SEARCH\n# ============================================================\n\nfrom Bio.Align import PairwiseAligner\n\n# Configure pairwise aligner with strict gap penalties\naligner = PairwiseAligner()\naligner.mode = \"global\"\naligner.match_score = 2\naligner.mismatch_score = -1.5\naligner.open_gap_score = -8\naligner.extend_gap_score = -0.4\n\n# Explicit terminal gap penalties to prevent end-gap sliding\naligner.query_left_open_gap_score = -8\naligner.query_left_extend_gap_score = -0.4\naligner.query_right_open_gap_score = -8\naligner.query_right_extend_gap_score = -0.4\naligner.target_left_open_gap_score = -8\naligner.target_left_extend_gap_score = -0.4\naligner.target_right_open_gap_score = -8\naligner.target_right_extend_gap_score = -0.4\n\ndef enhanced_template_scoring(query_seq, template_seq, template_id):\n    \"\"\"\n    Calculate enhanced template score combining sequence similarity,\n    length compatibility, and structure quality.\n    \n    Args:\n        query_seq: Query RNA sequence\n        template_seq: Template RNA sequence\n        template_id: Template identifier\n        \n    Returns:\n        Enhanced similarity score (higher is better)\n    \"\"\"\n    # Base sequence alignment score\n    raw_score = aligner.score(query_seq, template_seq)\n    seq_score = raw_score / (2 * min(len(query_seq), len(template_seq)))\n    \n    # Length similarity bonus (penalize very different lengths)\n    len_ratio = min(len(query_seq), len(template_seq)) / max(len(query_seq), len(template_seq))\n    len_bonus = len_ratio ** 0.5\n    \n    # Template structure quality from PDB metadata\n    quality_bonus = get_structure_quality(template_id)\n    \n    # Combined score with weights\n    combined_score = seq_score * (0.7 + 0.15 * len_bonus + 0.15 * quality_bonus)\n    \n    return combined_score\n\ndef find_similar_sequences(query_seq: str, train_seqs_df: pd.DataFrame, \n                          train_coords_dict: dict, top_n: int = 5):\n    \"\"\"\n    Find most similar training sequences using enhanced scoring.\n    \n    Args:\n        query_seq: Query RNA sequence\n        train_seqs_df: Training sequences dataframe\n        train_coords_dict: Dictionary of template coordinates\n        top_n: Number of top templates to return\n        \n    Returns:\n        List of (target_id, train_seq, score, coordinates) tuples\n    \"\"\"\n    similar_seqs = []\n    \n    for _, row in train_seqs_df.iterrows():\n        target_id = row[\"target_id\"]\n        train_seq = row[\"sequence\"]\n        \n        if target_id not in train_coords_dict:\n            continue\n\n        # Pre-filter by length difference\n        if abs(len(train_seq) - len(query_seq)) / max(len(train_seq), len(query_seq)) > 0.3:\n            continue\n\n        # Calculate enhanced similarity score\n        score = enhanced_template_scoring(query_seq, train_seq, target_id)\n        similar_seqs.append((target_id, train_seq, score, train_coords_dict[target_id]))\n\n    # Sort by enhanced score descending\n    similar_seqs.sort(key=lambda x: x[2], reverse=True)\n    return similar_seqs[:top_n]\n\n# ============================================================\n# SECTION H: ADVANCED COORDINATE INTERPOLATION\n# ============================================================\n\nfrom scipy.interpolate import CubicSpline\n\ndef smooth_interpolate_gaps(coords):\n    \"\"\"\n    Fill gaps in coordinates using cubic spline interpolation.\n    Provides smoother transitions than linear interpolation.\n    \n    Args:\n        coords: (L, 3) array with potential NaN values\n        \n    Returns:\n        coords: (L, 3) array with interpolated values\n    \"\"\"\n    coords = coords.copy()\n    \n    for dim in range(3):\n        # Find valid (non-NaN) indices\n        valid_mask = ~np.isnan(coords[:, dim])\n        valid_idx = np.where(valid_mask)[0]\n        \n        if len(valid_idx) < 4:\n            # Not enough points for cubic spline, use linear\n            for i in range(len(coords)):\n                if np.isnan(coords[i, dim]):\n                    prev_v = next((j for j in range(i - 1, -1, -1) if valid_mask[j]), -1)\n                    next_v = next((j for j in range(i + 1, len(coords)) if valid_mask[j]), -1)\n                    \n                    if prev_v >= 0 and next_v >= 0:\n                        w = (i - prev_v) / (next_v - prev_v)\n                        coords[i, dim] = (1 - w) * coords[prev_v, dim] + w * coords[next_v, dim]\n                    elif prev_v >= 0:\n                        coords[i, dim] = coords[prev_v, dim] + 3.0\n                    elif next_v >= 0:\n                        coords[i, dim] = coords[next_v, dim] - 3.0\n                    else:\n                        coords[i, dim] = i * 3.0\n            continue\n        \n        # Use cubic spline for smooth interpolation\n        try:\n            cs = CubicSpline(valid_idx, coords[valid_idx, dim], bc_type='natural')\n            \n            # Fill NaN positions\n            nan_mask = np.isnan(coords[:, dim])\n            nan_idx = np.where(nan_mask)[0]\n            \n            if len(nan_idx) > 0:\n                coords[nan_idx, dim] = cs(nan_idx)\n                \n        except Exception:\n            # Fallback to linear if spline fails\n            for i in range(len(coords)):\n                if np.isnan(coords[i, dim]):\n                    prev_v = next((j for j in range(i - 1, -1, -1) if valid_mask[j]), -1)\n                    next_v = next((j for j in range(i + 1, len(coords)) if valid_mask[j]), -1)\n                    \n                    if prev_v >= 0 and next_v >= 0:\n                        w = (i - prev_v) / (next_v - prev_v)\n                        coords[i, dim] = (1 - w) * coords[prev_v, dim] + w * coords[next_v, dim]\n    \n    return coords\n\n# ============================================================\n# SECTION I: TEMPLATE COORDINATE TRANSFER\n# ============================================================\n\ndef adapt_template_to_query(query_seq: str, template_seq: str, template_coords: np.ndarray):\n    \"\"\"\n    Transfer template coordinates to query sequence via alignment.\n    Uses cubic spline interpolation for gap filling.\n    \n    Args:\n        query_seq: Query RNA sequence\n        template_seq: Template RNA sequence\n        template_coords: (L_template, 3) coordinate array\n        \n    Returns:\n        new_coords: (L_query, 3) adapted coordinates\n    \"\"\"\n    alignment = next(iter(aligner.align(query_seq, template_seq)))\n    new_coords = np.full((len(query_seq), 3), np.nan, dtype=float)\n\n    # Transfer coordinates for aligned blocks\n    for (q_start, q_end), (t_start, t_end) in zip(*alignment.aligned):\n        t_chunk = template_coords[t_start:t_end]\n        if len(t_chunk) == (q_end - q_start):\n            new_coords[q_start:q_end] = t_chunk\n\n    # Fill gaps using cubic spline interpolation\n    new_coords = smooth_interpolate_gaps(new_coords)\n    \n    return np.nan_to_num(new_coords)\n\n# ============================================================\n# SECTION J: PHYSICS-BASED STRUCTURE REFINEMENT\n# ============================================================\n\ndef apply_base_pair_constraints(coords, coevolution_pairs, target_dist=10.5, strength=0.3):\n    \"\"\"\n    Apply distance constraints for predicted base pairs.\n    \n    Args:\n        coords: (L, 3) coordinate array\n        coevolution_pairs: List of (i, j, score) tuples\n        target_dist: Target distance for base pairs in Angstroms\n        strength: Constraint strength factor\n        \n    Returns:\n        coords: Refined coordinates\n    \"\"\"\n    coords = coords.copy()\n    \n    for (i, j, score) in coevolution_pairs[:20]:\n        if i >= len(coords) or j >= len(coords):\n            continue\n        \n        vec = coords[j] - coords[i]\n        dist = np.linalg.norm(vec) + 1e-6\n        \n        # Apply constraint based on coevolution score\n        weight = strength * min(score, 1.0)\n        scale = (target_dist - dist) / dist\n        adjustment = vec * scale * weight\n        \n        coords[i] -= adjustment * 0.5\n        coords[j] += adjustment * 0.5\n    \n    return coords\n\ndef adaptive_rna_constraints(coordinates: np.ndarray, target_id: str, \n                            confidence: float = 1.0, passes: int = 2):\n    \"\"\"\n    Apply segment-aware RNA structural constraints.\n    Constraints are applied within chain segments only.\n    \n    Constraints applied:\n    - Bond length (i, i+1) target ~5.95 Angstrom\n    - Backbone angle (i, i+2) target ~10.2 Angstrom\n    - Laplacian smoothing\n    - Self-avoidance on subsampled points\n    \n    Args:\n        coordinates: (L, 3) coordinate array\n        target_id: Target identifier for segment lookup\n        confidence: Template confidence (higher = weaker constraints)\n        passes: Number of refinement iterations\n        \n    Returns:\n        coords: Refined coordinates\n    \"\"\"\n    coords = coordinates.copy()\n    segments = test_segs_map.get(target_id, [(0, len(coords))])\n\n    # Calculate constraint strength based on confidence\n    # Lower confidence requires stronger refinement\n    strength = 0.75 * (1.0 - min(confidence, 0.90))\n    strength = max(strength, 0.02)\n\n    for _ in range(passes):\n        for (s, e) in segments:\n            X = coords[s:e]\n            L = e - s\n            if L < 3:\n                coords[s:e] = X\n                continue\n\n            # Constraint 1: Bond length (i, i+1) to ~5.95 Angstrom\n            d = X[1:] - X[:-1]\n            dist = np.linalg.norm(d, axis=1) + 1e-6\n            target = 5.95\n            scale = (target - dist) / dist\n            adj = (d * scale[:, None]) * (0.22 * strength)\n            X[:-1] -= adj\n            X[1:] += adj\n\n            # Constraint 2: Backbone angle (i, i+2) to ~10.2 Angstrom\n            d2 = X[2:] - X[:-2]\n            dist2 = np.linalg.norm(d2, axis=1) + 1e-6\n            target2 = 10.2\n            scale2 = (target2 - dist2) / dist2\n            adj2 = (d2 * scale2[:, None]) * (0.10 * strength)\n            X[:-2] -= adj2\n            X[2:] += adj2\n\n            # Constraint 3: Laplacian smoothing for local geometry\n            lap = 0.5 * (X[:-2] + X[2:]) - X[1:-1]\n            X[1:-1] += (0.06 * strength) * lap\n\n            # Constraint 4: Self-avoidance (subsample for efficiency)\n            if L >= 25:\n                k = min(L, 160) if L > 220 else L\n                idx = np.linspace(0, L - 1, k).astype(int) if k < L else np.arange(L)\n\n                P = X[idx]\n                diff = P[:, None, :] - P[None, :, :]\n                distm = np.linalg.norm(diff, axis=2) + 1e-6\n                sep = np.abs(idx[:, None] - idx[None, :])\n\n                # Only apply for distant residues that are too close\n                mask = (sep > 2) & (distm < 3.2)\n                if np.any(mask):\n                    force = (3.2 - distm) / distm\n                    vec = (diff * force[:, :, None] * mask[:, :, None]).sum(axis=1)\n                    X[idx] += (0.015 * strength) * vec\n\n            coords[s:e] = X\n\n    return coords\n\n# ============================================================\n# SECTION K: STRUCTURAL DIVERSITY GENERATION\n# ============================================================\n\ndef rotation_matrix(axis, angle):\n    \"\"\"\n    Generate 3D rotation matrix using Rodrigues' formula.\n    \n    Args:\n        axis: Rotation axis (3D vector)\n        angle: Rotation angle in radians\n        \n    Returns:\n        3x3 rotation matrix\n    \"\"\"\n    axis = np.asarray(axis, float)\n    axis = axis / (np.linalg.norm(axis) + 1e-12)\n    x, y, z = axis\n    c, s = np.cos(angle), np.sin(angle)\n    C = 1.0 - c\n    return np.array([\n        [c + x * x * C, x * y * C - z * s, x * z * C + y * s],\n        [y * x * C + z * s, c + y * y * C, y * z * C - x * s],\n        [z * x * C - y * s, z * y * C + x * s, c + z * z * C],\n    ], dtype=float)\n\ndef apply_hinge(coords, seg, rng, max_angle_deg=25):\n    \"\"\"\n    Apply hinge motion to a segment for structural diversity.\n    \n    Args:\n        coords: (L, 3) coordinate array\n        seg: (start, end) segment tuple\n        rng: Random number generator\n        max_angle_deg: Maximum hinge angle in degrees\n        \n    Returns:\n        coords: Coordinates with hinge applied\n    \"\"\"\n    s, e = seg\n    L = e - s\n    if L < 30:\n        return coords\n    \n    # Choose random pivot point\n    pivot = s + int(rng.integers(10, L - 10))\n    \n    # Random rotation axis and angle\n    axis = rng.normal(size=3)\n    ang = np.deg2rad(float(rng.uniform(-max_angle_deg, max_angle_deg)))\n    R = rotation_matrix(axis, ang)\n    \n    # Apply rotation to residues after pivot\n    X = coords.copy()\n    p0 = X[pivot].copy()\n    X[pivot + 1:e] = (X[pivot + 1:e] - p0) @ R.T + p0\n    return X\n\ndef jitter_chains(coords, segments, rng, max_angle_deg=12, max_trans=1.5):\n    \"\"\"\n    Apply small random rotations and translations to each chain.\n    Maintains global center of mass.\n    \n    Args:\n        coords: (L, 3) coordinate array\n        segments: List of (start, end) tuples\n        rng: Random number generator\n        max_angle_deg: Maximum rotation angle per chain\n        max_trans: Maximum translation distance\n        \n    Returns:\n        coords: Jittered coordinates\n    \"\"\"\n    X = coords.copy()\n    global_center = X.mean(axis=0, keepdims=True)\n    \n    for (s, e) in segments:\n        # Random rotation\n        axis = rng.normal(size=3)\n        ang = np.deg2rad(float(rng.uniform(-max_angle_deg, max_angle_deg)))\n        R = rotation_matrix(axis, ang)\n        \n        # Random translation\n        shift = rng.normal(size=3)\n        shift = shift / (np.linalg.norm(shift) + 1e-12) * float(rng.uniform(0.0, max_trans))\n        \n        # Apply to chain\n        c = X[s:e].mean(axis=0, keepdims=True)\n        X[s:e] = (X[s:e] - c) @ R.T + c + shift\n    \n    # Restore global center\n    X -= X.mean(axis=0, keepdims=True) - global_center\n    return X\n\ndef smooth_wiggle(coords, segments, rng, amp=0.8):\n    \"\"\"\n    Apply smooth wave-like perturbations to chain segments.\n    \n    Args:\n        coords: (L, 3) coordinate array\n        segments: List of (start, end) tuples\n        rng: Random number generator\n        amp: Amplitude of wiggle\n        \n    Returns:\n        coords: Wiggled coordinates\n    \"\"\"\n    X = coords.copy()\n    \n    for (s, e) in segments:\n        L = e - s\n        if L < 20:\n            continue\n        \n        # Generate smooth displacement using control points\n        n_ctrl = 6\n        ctrl_x = np.linspace(0, L - 1, n_ctrl)\n        ctrl_disp = rng.normal(0, amp, size=(n_ctrl, 3))\n        \n        # Interpolate displacement\n        t = np.arange(L)\n        disp = np.vstack([np.interp(t, ctrl_x, ctrl_disp[:, k]) for k in range(3)]).T\n        X[s:e] += disp\n    \n    return X\n\n# ============================================================\n# SECTION L: MAIN STRUCTURE PREDICTION PIPELINE\n# ============================================================\n\ndef predict_rna_structures(row, train_seqs_df, train_coords_dict, n_predictions=5):\n    \"\"\"\n    Generate multiple diverse structure predictions for a target RNA.\n    \n    Strategy:\n    - Prediction 1: Best template with minimal perturbation\n    - Prediction 2: Best template with Gaussian noise\n    - Prediction 3: Hinge motion applied\n    - Prediction 4: Chain jittering applied\n    - Prediction 5: Smooth wiggle applied\n    \n    Args:\n        row: Test sequence row\n        train_seqs_df: Training sequences dataframe\n        train_coords_dict: Template coordinates dictionary\n        n_predictions: Number of predictions to generate (must be 5)\n        \n    Returns:\n        List of (L, 3) coordinate arrays\n    \"\"\"\n    tid = row[\"target_id\"]\n    seq = row[\"sequence\"]\n\n    # Validate sequence contains only canonical nucleotides\n    assert set(seq).issubset(set(\"ACGU\")), f\"Non-ACGU nucleotides in {tid}\"\n\n    segments = test_segs_map.get(tid, [(0, len(seq))])\n\n    # Search for similar templates (expanded pool for diversity)\n    cands = find_similar_sequences(\n        query_seq=seq,\n        train_seqs_df=train_seqs_df,\n        train_coords_dict=train_coords_dict,\n        top_n=30\n    )\n    \n    # Validate template data integrity\n    assert all(len(c[3]) == len(c[1]) for c in cands), \"Template coords/seq length mismatch\"\n\n    # Get coevolution pairs from MSA\n    coevol_pairs = get_coevolution_pairs(tid, seq)\n\n    predictions = []\n    used_templates = set()\n\n    for i in range(n_predictions):\n        # Deterministic random seed based on row index and prediction number\n        seed = (row.name * 10000000000 + i * 10007) % (2**32)\n        rng = np.random.default_rng(seed)\n\n        # Fallback for targets with no similar templates\n        if not cands:\n            coords = np.zeros((len(seq), 3), dtype=float)\n            for (s, e) in segments:\n                for j in range(s + 1, e):\n                    coords[j] = coords[j - 1] + [5.95, 0, 0]\n            predictions.append(coords)\n            continue\n\n        # Template selection strategy\n        if i == 0:\n            # Prediction 1: Use best template\n            t_id, t_seq, sim, t_coords = cands[0]\n        else:\n            # Predictions 2-5: Sample from top templates with diversity\n            K = min(12, len(cands))\n            sims = np.array([cands[k][2] for k in range(K)], float)\n            \n            # Weight by similarity, penalize already used templates\n            w = np.exp((sims - sims.max()) / 0.08)\n            for k in range(K):\n                if cands[k][0] in used_templates:\n                    w[k] *= 0.10\n            w = w / (w.sum() + 1e-12)\n            \n            k = int(rng.choice(np.arange(K), p=w))\n            t_id, t_seq, sim, t_coords = cands[k]\n\n        used_templates.add(t_id)\n\n        # Transfer template coordinates to query\n        adapted = adapt_template_to_query(query_seq=seq, template_seq=t_seq, template_coords=t_coords)\n\n        # Apply diversity transformations\n        if i == 0:\n            # Prediction 1: Minimal perturbation\n            X = adapted\n        elif i == 1:\n            # Prediction 2: Add Gaussian noise (scaled by template uncertainty)\n            noise_scale = max(0.01, (0.40 - sim) * 0.06)\n            X = adapted + rng.normal(0, noise_scale, adapted.shape)\n        elif i == 2:\n            # Prediction 3: Hinge motion on longest segment\n            longest = max(segments, key=lambda se: se[1] - se[0])\n            X = apply_hinge(adapted, longest, rng, max_angle_deg=22)\n        elif i == 3:\n            # Prediction 4: Chain jittering\n            X = jitter_chains(adapted, segments, rng, max_angle_deg=10, max_trans=1.0)\n        else:\n            # Prediction 5: Smooth wiggle\n            X = smooth_wiggle(adapted, segments, rng, amp=0.7)\n\n        # Apply MSA-based base pair constraints\n        if coevol_pairs:\n            X = apply_base_pair_constraints(X, coevol_pairs, target_dist=10.5, strength=0.3)\n\n        # Physics-based refinement\n        refined = adaptive_rna_constraints(X, tid, confidence=sim, passes=2)\n        predictions.append(refined)\n\n    return predictions\n\n# ============================================================\n# SECTION M: SUBMISSION FILE GENERATION\n# ============================================================\n\ndef generate_submission():\n    \"\"\"\n    Generate submission.csv file with 5 structure predictions per target.\n    \n    Output format:\n    ID, resname, resid, x_1, y_1, z_1, ..., x_5, y_5, z_5\n    \"\"\"\n    all_predictions = []\n    start_time = time.time()\n\n    for idx, row in test_seqs.iterrows():\n        if idx % 10 == 0:\n            elapsed = time.time() - start_time\n            print(f\"Processing target {idx} | Elapsed time: {elapsed:.1f}s\")\n        \n        tid = row[\"target_id\"]\n        seq = row[\"sequence\"]\n\n        # Generate 5 predictions for this target\n        preds = predict_rna_structures(row, train_seqs, train_coords_dict, n_predictions=5)\n\n        # Validate predictions\n        L = len(seq)\n        for p in preds:\n            assert isinstance(p, np.ndarray) and p.shape == (L, 3), \\\n                f\"Invalid prediction shape for {tid}: {getattr(p, 'shape', None)}\"\n            assert np.isfinite(p).all(), f\"Non-finite coordinates in {tid}\"\n\n        # Format predictions for submission\n        for j in range(L):\n            res = {\n                \"ID\": f\"{tid}_{j+1}\",\n                \"resname\": seq[j],\n                \"resid\": j + 1\n            }\n            for i in range(5):\n                res[f\"x_{i+1}\"], res[f\"y_{i+1}\"], res[f\"z_{i+1}\"] = preds[i][j]\n            all_predictions.append(res)\n\n    # Create submission dataframe\n    sub = pd.DataFrame(all_predictions)\n\n    # Define column order\n    cols = [\"ID\", \"resname\", \"resid\"] + [f\"{c}_{i}\" for i in range(1, 6) for c in [\"x\", \"y\", \"z\"]]\n\n    # Clip coordinates to PDB format limits\n    coord_cols = [c for c in cols if c.startswith((\"x_\", \"y_\", \"z_\"))]\n    sub[coord_cols] = sub[coord_cols].clip(-999.999, 9999.999)\n\n    # Save submission file\n    sub[cols].to_csv(\"submission.csv\", index=False)\n    \n    total_time = time.time() - start_time\n    print(f\"\\nSubmission complete!\")\n    print(f\"Total processing time: {total_time:.1f}s\")\n    print(f\"Targets processed: {len(test_seqs)}\")\n    print(f\"Output file: submission.csv\")\n\n# ============================================================\n# SECTION N: EXECUTION\n# ============================================================\n\nif __name__ == \"__main__\":\n    print(\"=\" * 60)\n    print(\"Stanford RNA 3D Folding Part 2 - Enhanced Pipeline\")\n    print(\"Team: DTU Compute\")\n    print(\"=\" * 60)\n    print()\n    \n    generate_submission()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T18:28:33.540066Z","iopub.execute_input":"2026-02-11T18:28:33.540426Z"}},"outputs":[],"execution_count":null}]}