{"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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":14441699,"sourceType":"datasetVersion","datasetId":9224635},{"sourceId":14610545,"sourceType":"datasetVersion","datasetId":9331376}],"dockerImageVersionId":31259,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install /kaggle/input/biopython-cp312/biopython-1.86-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl --no-index --find-links=/kaggle/input/biopython-cp312/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:52:20.180371Z","iopub.execute_input":"2026-01-27T08:52:20.180748Z","iopub.status.idle":"2026-01-27T08:52:27.849088Z","shell.execute_reply.started":"2026-01-27T08:52:20.180693Z","shell.execute_reply":"2026-01-27T08:52:27.847549Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nfrom Bio import Align\nfrom Bio.Seq import Seq\nfrom scipy.spatial.transform import Rotation\nfrom scipy.spatial.distance import pdist, squareform\nimport time\nfrom tqdm.auto import tqdm\nimport gc\nimport sqlite3\n\n# 1. Offline installation (using your identified wheel file)\n# !pip install /kaggle/input/biopython-cp312/biopython-1.86-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl --no-index --find-links=/kaggle/input/biopython-cp312/\n\nprint('Modules successfully loaded with PairwiseAligner')\n\nclass Config:\n    DATA_PATH = '/kaggle/input/stanford-rna-3d-folding-2'\n    MAX_RELATIVE_LENGTH_DIFF = 0.3\n    # PairwiseAligner specific settings\n    ALIGNMENT_MATCH = 2.0\n    ALIGNMENT_MISMATCH = -1.0\n    ALIGNMENT_GAP_OPEN = -10.0\n    ALIGNMENT_GAP_EXTEND = -0.5\n    \n    BOND_DISTANCE_TARGET = 6.0\n    BOND_DISTANCE_TOL = 0.5\n    MIN_NONBOND_DISTANCE = 3.8\n    BASE_PAIRING_DISTANCE_IDEAL = 12.5\n    BASE_PAIRING_DISTANCE_RANGE = (8.0, 14.0)\n    HELIX_RADIUS = 10.0\n    HELIX_RISE_PER_BASE = 2.5\n    HELIX_ANGLE_STEP = 0.6\n    NUM_PREDICTIONS = 5\n    RANDOM_SEED_OFFSET = 1000\n    DEFAULT_RELIABILITY = 0.2\n    PAIRING_PROB_THRESHOLD = 0.7\n    STEP_LENGTH_RANGE = (3.5, 4.5)\n    NOISE_SCALE_MIN = 0.01\n    TEMPLATE_WEIGHT = 0.55\n    RANDOM_WEIGHT = 0.45\n\n# Initialize the modern aligner globally for efficiency\naligner = Align.PairwiseAligner()\naligner.mode = 'global'\naligner.match_score = Config.ALIGNMENT_MATCH\naligner.mismatch_score = Config.ALIGNMENT_MISMATCH\naligner.open_gap_score = Config.ALIGNMENT_GAP_OPEN\naligner.extend_gap_score = Config.ALIGNMENT_GAP_EXTEND\n\ntrain = pd.read_csv(f'{Config.DATA_PATH}/train_sequences.csv')\nval = pd.read_csv(f'{Config.DATA_PATH}/validation_sequences.csv')\ntest = pd.read_csv(f'{Config.DATA_PATH}/test_sequences.csv')\ntrain_labels = pd.read_csv(f'{Config.DATA_PATH}/train_labels.csv')\nval_labels = pd.read_csv(f'{Config.DATA_PATH}/validation_labels.csv')\n\ndef extract_structures(labels_data):\n    structures = {}\n    grouped = labels_data.groupby(labels_data['ID'].str.rsplit('_', n=1).str[0])\n    \n    for group_label, group_data in tqdm(grouped, desc=\"Extracting Ground Truth Structures\"):\n        coords_list = group_data.sort_values('resid')[['x_1', 'y_1', 'z_1']].values\n        \n        structures[group_label] = np.array(coords_list, dtype=np.float32)\n    return structures\n\ndef secondary_structure_filter(seq_a, seq_b, min_bp=2):\n    pairing_map = {'A': 'U', 'U': 'A', 'G': 'C', 'C': 'G'}\n    count = 0\n    for i in range(min(len(seq_a), len(seq_b))):\n        if seq_a[i] in pairing_map and seq_b[i] == pairing_map[seq_a[i]]:\n            count += 1\n    return count >= min_bp\n\ndef compute_alignment_score(query_seq, template_seq):\n    \"\"\"\n    Optimized to prevent OverflowError by checking the score \n    directly rather than counting the number of alignments.\n    \"\"\"\n    # 1. Get the alignments object\n    alignments = aligner.align(query_seq, template_seq)\n    \n    # 2. Access the score directly. \n    # This is O(N*M) and does NOT trigger the path-counting overflow.\n    try:\n        raw_score = alignments.score\n    except (AttributeError, OverflowError):\n        # Fallback for older Biopython or weird edge cases\n        if not alignments:\n            return None, 0.0\n        raw_score = alignments[0].score\n\n    # 3. Normalize the score\n    max_len = min(len(query_seq), len(template_seq))\n    score = raw_score / (Config.ALIGNMENT_MATCH * max_len)\n    \n    # 4. Return the first alignment object and the score\n    # alignments[0] works even if len(alignments) overflows\n    return alignments[0], score\n\ndef morph_template_with_interpolation(dest_seq, source_seq, source_coords):\n    alignment, _ = compute_alignment_score(dest_seq, source_seq)\n    if not alignment:\n        return create_default_structure(dest_seq)\n\n    # Convert PairwiseAligner output to string for mapping\n    aligned_dest = alignment[0]\n    aligned_source = alignment[1]\n\n    # Initialize with NaNs so we can identify what needs the Random Walk\n    morphed_coords = np.full((len(dest_seq), 3), np.nan, dtype=np.float32)\n    dest_idx = source_idx = 0\n\n    # 1. Map known coordinates from Template\n    for a, b in zip(aligned_dest, aligned_source):\n        if dest_idx >= len(morphed_coords) or source_idx >= len(source_coords):\n            break\n            \n        if a != '-' and b != '-':\n            morphed_coords[dest_idx] = source_coords[source_idx]\n            dest_idx += 1\n            source_idx += 1\n        elif a != '-' and b == '-':\n            dest_idx += 1\n        elif a == '-' and b != '-':\n            source_idx += 1\n\n    # --- CHANGED SECTION START ---\n    # REPLACED: Global np.interp with Segment-Based Physics Check\n    \n    valid_mask = ~np.isnan(morphed_coords[:, 0])\n    \n    # If we have valid anchors, try to connect them\n    if valid_mask.any():\n        valid_indices = np.where(valid_mask)[0]\n        \n        # Max physical reach per nucleotide (approx 7.0A + 0.5A tolerance)\n        MAX_STRETCH_PER_RESIDUE = 7.5 \n\n        for i in range(len(valid_indices) - 1):\n            idx_start = valid_indices[i]\n            idx_end = valid_indices[i+1]\n            \n            steps = idx_end - idx_start\n            \n            # If residues are consecutive (steps=1), there is no gap to fill.\n            if steps <= 1:\n                continue\n\n            start_coord = morphed_coords[idx_start]\n            end_coord = morphed_coords[idx_end]\n            euclidean_dist = np.linalg.norm(start_coord - end_coord)\n            \n            # KEY LOGIC: Dynamic Threshold\n            # Can we bridge this gap without breaking physics?\n            max_plausible_dist = steps * MAX_STRETCH_PER_RESIDUE\n            \n            if euclidean_dist < max_plausible_dist: \n                # YES: The gap is reachable. Interpolate linearly.\n                gap_range = np.arange(idx_start + 1, idx_end)\n                for dim in range(3):\n                    morphed_coords[gap_range, dim] = np.interp(\n                        gap_range,\n                        [idx_start, idx_end],\n                        [start_coord[dim], end_coord[dim]]\n                    )\n            else:\n                # NO: The template is broken or \"exploded\". \n                # DO NOTHING. Leave as NaNs.\n                # The Random Walk below will naturally grow a new chain here.\n                pass\n    # --- CHANGED SECTION END ---\n\n    # Random Walk for Extrapolation (Handles Ends AND Internal Breaks)\n    # (This section is largely unchanged but now handles the internal NaNs we left behind)\n    for i in range(len(morphed_coords)):\n        if np.isnan(morphed_coords[i, 0]):\n            if i == 0:\n                morphed_coords[i] = [0.0, 0.0, 0.0]\n            else:\n                # Use previous valid atoms to determine direction\n                if i >= 2 and not np.isnan(morphed_coords[i-1, 0]) and not np.isnan(morphed_coords[i-2, 0]):\n                    prev_vec = morphed_coords[i-1] - morphed_coords[i-2]\n                else:\n                    prev_vec = np.array([1.0, 0.0, 0.0])\n                \n                norm = np.linalg.norm(prev_vec)\n                if norm < 1e-6:\n                    prev_vec = np.random.normal(0, 1, 3)\n                else:\n                    prev_vec /= norm\n                \n                # Move a random distance within the allowed step range\n                step = np.random.uniform(*Config.STEP_LENGTH_RANGE)\n                \n                # OPTIONAL: Slight bias towards the next valid anchor (if one exists)\n                # Find the next valid non-NaN index\n                next_valid_idx = -1\n                for k in range(i + 1, len(morphed_coords)):\n                    if not np.isnan(morphed_coords[k, 0]):\n                        next_valid_idx = k\n                        break\n                \n                if next_valid_idx != -1:\n                    # Vector pointing to the target\n                    target_vec = morphed_coords[next_valid_idx] - morphed_coords[i-1]\n                    target_dist = np.linalg.norm(target_vec)\n                    if target_dist > 1e-6:\n                        target_dir = target_vec / target_dist\n                        # Blend the previous direction (inertia) with target direction (bias)\n                        # 80% inertia, 20% bias\n                        prev_vec = 0.8 * prev_vec + 0.2 * target_dir\n                        prev_vec /= np.linalg.norm(prev_vec)\n\n                morphed_coords[i] = morphed_coords[i - 1] + prev_vec * step\n\n    return morphed_coords\n\ndef enhance_geometry(positions, rna_string, reliability=1.0):\n    refined = positions.copy()\n    n = len(rna_string)\n    adj_intensity = 0.8 * (1.0 - min(reliability, 0.8))\n\n    # DEFINITION: What counts as a \"broken\" chain?\n    # If residues are > 12.0 Angstroms apart, we treat them as disconnected.\n    # We will NOT try to pull them together, as that causes explosions.\n    BREAK_THRESHOLD = 12.0 \n\n    # 1. Backbone Distance Constraint\n    for i in range(n - 1):\n        vec = refined[i + 1] - refined[i]\n        dist = np.linalg.norm(vec)\n\n        # --- THE FIX ---\n        # If the gap is massive (e.g. 50A), do not try to close it. \n        # Skip this bond adjustment entirely.\n        if dist > BREAK_THRESHOLD:\n            continue\n        # ----------------\n\n        if abs(dist - Config.BOND_DISTANCE_TARGET) > Config.BOND_DISTANCE_TOL:\n            target = Config.BOND_DISTANCE_TARGET\n            unit_vec = vec / (dist + 1e-10)\n            delta = (target - dist) * adj_intensity\n            refined[i + 1] = refined[i] + unit_vec * (dist + delta)\n\n    # 2. Angle Smoothing\n    if n > 3:\n        for i in range(1, n - 1):\n            # --- THE FIX ---\n            # We must verify that neighbors are physically close before calculating angles.\n            # Calculating an angle across a 50A void produces garbage vectors.\n            \n            vec_prev = refined[i] - refined[i-1]\n            dist_prev = np.linalg.norm(vec_prev)\n            \n            vec_next = refined[i+1] - refined[i]\n            dist_next = np.linalg.norm(vec_next)\n\n            if dist_prev > BREAK_THRESHOLD or dist_next > BREAK_THRESHOLD:\n                continue\n            # ----------------\n\n            cos_theta = np.dot(vec_prev, vec_next) / (dist_prev * dist_next + 1e-10)\n            angle = np.arccos(np.clip(cos_theta, -1.0, 1.0))\n            if angle > 2.5:\n                smoothed = (refined[i-1] + refined[i+1]) / 2\n                refined[i] = refined[i] * 0.3 + smoothed * 0.7\n\n    # 3. Non-bonded Clashes (Standard - No changes needed usually)\n    dist_matrix = squareform(pdist(refined))\n    clashes = (dist_matrix < Config.MIN_NONBOND_DISTANCE) & (dist_matrix > 0)\n    for i in range(n):\n        for j in range(i + 2, n):\n            if clashes[i, j]:\n                vec = refined[j] - refined[i]\n                unit = vec / (np.linalg.norm(vec) + 1e-10)\n                push = (Config.MIN_NONBOND_DISTANCE - dist_matrix[i, j]) * adj_intensity\n                refined[i] -= unit * (push / 2)\n                refined[j] += unit * (push / 2)\n\n    # 4. Base Pairing (Standard)\n    base_pairs = {'A': 'U', 'U': 'A', 'G': 'C', 'C': 'G'}\n    for i in range(n):\n        partner = base_pairs.get(rna_string[i])\n        if not partner:\n            continue\n        for j in range(i + 3, min(i + 20, n)):\n            if rna_string[j] == partner:\n                d = np.linalg.norm(refined[i] - refined[j])\n                if Config.BASE_PAIRING_DISTANCE_RANGE[0] < d < Config.BASE_PAIRING_DISTANCE_RANGE[1]:\n                    ideal = Config.BASE_PAIRING_DISTANCE_IDEAL\n                    delta = (ideal - d) * adj_intensity * 0.3\n                    unit = (refined[j] - refined[i]) / (d + 1e-10)\n                    refined[i] -= unit * (delta / 2)\n                    refined[j] += unit * (delta / 2)\n                    break\n\n    return refined\n\ndef create_default_structure(seq):\n    n = len(seq)\n    coords = np.zeros((n, 3), dtype=np.float32)\n    \n    if n <= 10:\n        for i in range(n):\n            angle = i * 0.3\n            coords[i] = [\n                8.0 * np.cos(angle),\n                8.0 * np.sin(angle),\n                i * 3.0\n            ]\n        return coords\n    \n    base_pairs = {'G': 'C', 'C': 'G', 'A': 'U', 'U': 'A'}\n    paired = {}\n    \n    for i in range(n):\n        if i in paired:\n            continue\n        for j in range(i+4, min(i+15, n)):\n            if j not in paired and base_pairs.get(seq[i]) == seq[j]:\n                paired[i] = j\n                paired[j] = i\n                break\n    \n    current_pos = np.array([0.0, 0.0, 0.0])\n    direction = np.array([1.0, 0.0, 0.0])\n    \n    for i in range(n):\n        if i in paired and paired[i] > i:\n            j = paired[i]\n            if np.random.rand() < 0.6:\n                vec = np.random.normal(0, 1, 3)\n                vec /= np.linalg.norm(vec)\n                coords[i] = current_pos + vec * 10.0\n                coords[j] = current_pos - vec * 10.0\n                current_pos = (coords[i] + coords[j]) / 2\n                direction = np.random.normal(0, 1, 3)\n                direction /= np.linalg.norm(direction)\n            else:\n                coords[i] = current_pos\n                current_pos += direction * np.random.uniform(3.5, 4.5)\n                coords[j] = current_pos\n                current_pos += direction * np.random.uniform(3.5, 4.5)\n        else:\n            coords[i] = current_pos\n            if np.random.rand() < 0.2:\n                rot_angle = np.random.uniform(-0.4, 0.4)\n                rot_axis = np.random.normal(0, 1, 3)\n                rot_axis /= np.linalg.norm(rot_axis)\n                rot = Rotation.from_rotvec(rot_angle * rot_axis)\n                direction = rot.apply(direction)\n            current_pos += direction * np.random.uniform(3.5, 4.5)\n    \n    return coords\n\ndef fabricate_rna_conformation(seq, seed=None):\n    if seed is not None:\n        np.random.seed(seed)\n    \n    n = len(seq)\n    coords = np.zeros((n, 3), dtype=np.float32)\n    base_pairs = {'G': 'C', 'C': 'G', 'A': 'U', 'U': 'A'}\n    \n    if n <= 3:\n        for i in range(n):\n            coords[i] = [i * 6.0, 0.0, 0.0]\n        return coords\n    \n    start_helix = min(4, n // 2)\n    for i in range(start_helix):\n        angle = i * Config.HELIX_ANGLE_STEP * 0.8\n        coords[i] = [\n            Config.HELIX_RADIUS * 0.7 * np.cos(angle),\n            Config.HELIX_RADIUS * 0.7 * np.sin(angle),\n            i * Config.HELIX_RISE_PER_BASE * 1.2\n        ]\n    \n    direction = np.array([0.0, 0.1, 0.995])\n    direction /= np.linalg.norm(direction)\n    \n    for i in range(start_helix, n):\n        current = seq[i]\n        paired = False\n        partner_idx = -1\n        \n        for j in range(max(0, i - 12), i):\n            if seq[j] == base_pairs.get(current, None):\n                paired = True\n                partner_idx = j\n                break\n        \n        if paired and (i - partner_idx <= 8) and (np.random.rand() < Config.PAIRING_PROB_THRESHOLD):\n            partner_pos = coords[partner_idx]\n            if i - partner_idx > 3:\n                vec = coords[i-1] - coords[max(0, i-2)]\n                if np.linalg.norm(vec) > 1e-6:\n                    perp = np.cross(vec, np.array([0, 0, 1]))\n                    perp /= np.linalg.norm(perp) + 1e-10\n                    coords[i] = coords[i-1] + perp * (Config.BASE_PAIRING_DISTANCE_IDEAL * 0.6)\n                else:\n                    random_dir = np.random.normal(0, 1, 3)\n                    random_dir /= np.linalg.norm(random_dir)\n                    coords[i] = coords[i-1] + random_dir * (Config.BASE_PAIRING_DISTANCE_IDEAL * 0.6)\n            else:\n                center_vec = np.mean(coords[max(0, i-3):i], axis=0) - partner_pos\n                if np.linalg.norm(center_vec) > 1e-6:\n                    center_vec /= np.linalg.norm(center_vec)\n                else:\n                    center_vec = np.random.normal(0, 1, 3)\n                    center_vec /= np.linalg.norm(center_vec)\n                distance = Config.BASE_PAIRING_DISTANCE_IDEAL * 0.8\n                coords[i] = partner_pos + center_vec * distance\n            direction = np.random.normal(0, 0.2, 3)\n            direction /= np.linalg.norm(direction) + 1e-10\n        else:\n            if np.random.rand() < 0.25:\n                rot_angle = np.random.uniform(-0.3, 0.3)\n                rot_axis = np.random.normal(0, 1, 3)\n                rot_axis /= np.linalg.norm(rot_axis)\n                rot = Rotation.from_rotvec(rot_angle * rot_axis)\n                direction = rot.apply(direction)\n            else:\n                direction += np.random.normal(0, 0.1, 3)\n                direction /= np.linalg.norm(direction) + 1e-10\n            \n            step = np.random.uniform(*Config.STEP_LENGTH_RANGE)\n            coords[i] = coords[i - 1] + direction * step\n    \n    return coords\n\ndef find_comparable_seqs(query_seq, template_df, structure_dict, date_filter=None, top_k=5):\n    mask = pd.Series([True] * len(template_df), index=template_df.index)\n    \n    if date_filter:\n        candidates = template_df[template_df['temporal_cutoff'] < date_filter]\n    else:\n        t_lens = template_df['sequence'].str.len()\n        len_diff = np.abs(t_lens - len(query_seq)) / np.maximum(t_lens, len(query_seq))\n        mask &= (len_diff <= Config.MAX_RELATIVE_LENGTH_DIFF)\n        \n        candidates = template_df[mask]\n    \n    matches = []\n    print(f'number of templates and target seq length: {len(candidates)}   {len(query_seq)}')\n    for count, (_, row) in enumerate(tqdm(candidates.iterrows(), total=len(candidates), desc=\"Finding template\")):\n        tid = row['target_id']\n        tseq = row['sequence']\n\n        if tid not in structure_dict:\n            continue\n        if abs(len(tseq) - len(query_seq)) / max(len(tseq), len(query_seq)) > Config.MAX_RELATIVE_LENGTH_DIFF:\n            continue\n        if not secondary_structure_filter(query_seq, tseq):\n            continue\n            \n        _, score = compute_alignment_score(query_seq, tseq)\n        matches.append((tid, tseq, score, structure_dict[tid]))\n\n        # count is now properly defined by the enumerate() function\n        if count % 50 == 0:\n            gc.collect()\n\n    matches.sort(key=lambda x: x[2], reverse=True)\n    return matches[:top_k]\n\ndef compute_structures(seq, identifier, template_df, struct_dict, num_outputs=5, date_limit=None):\n    predictions = []\n\n    comparable = find_comparable_seqs(seq, template_df, struct_dict, date_filter=date_limit, top_k=num_outputs)\n    \n    for tid, tseq, sim_score, tcoords in comparable:\n        morphed = morph_template_with_interpolation(seq, tseq, tcoords)\n        \n        # FIX 1: Center the morphed template at the origin immediately\n        # This prevents inheriting large coordinate values from PDB templates\n        morphed = morphed - np.mean(morphed, axis=0)\n        \n        refined = enhance_geometry(morphed, seq, reliability=sim_score)\n        \n        noise_scale = max(Config.NOISE_SCALE_MIN, 0.8 - sim_score)\n        if sim_score > 0.5:\n            noise_scale *= 0.5\n        \n        noise = np.random.normal(0, noise_scale, refined.shape)\n        predictions.append(refined + noise)\n        \n        if len(predictions) >= num_outputs:\n            break\n\n    while len(predictions) < num_outputs:\n        seed = (hash(identifier) % 10000) + len(predictions) * Config.RANDOM_SEED_OFFSET\n        fake = fabricate_rna_conformation(seq, seed=seed)\n        \n        # FIX 2: Center the fabricated structure immediately\n        fake = fake - np.mean(fake, axis=0)\n        \n        refined = enhance_geometry(fake, seq, reliability=Config.DEFAULT_RELIABILITY)\n        \n        # FIX 3: Removed the logic that shifted 'refined' to match 'centroid_pred'.\n        # We want everything centered at 0,0,0 to avoid boundary errors.\n        predictions.append(refined)\n\n    final_predictions = []\n    for pred in predictions:\n        # FINAL SAFETY: Ensure everything is centered and clipped to Kaggle limits\n        centered = pred - np.mean(pred, axis=0)\n        \n        # Clip to be safely within -999.999 to 9999.999\n        clipped = np.clip(centered, -990.0, 9990.0)\n        final_predictions.append(clipped.astype(np.float32))\n        \n    return final_predictions[:num_outputs]\n\nprint('Done defining the key functions.')\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:52:27.852551Z","iopub.execute_input":"2026-01-27T08:52:27.852963Z","iopub.status.idle":"2026-01-27T08:52:41.315388Z","shell.execute_reply.started":"2026-01-27T08:52:27.852902Z","shell.execute_reply":"2026-01-27T08:52:41.314509Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Use Template Database","metadata":{}},{"cell_type":"code","source":"# Path to your generated DB\ndb_path = \"/kaggle/input/train-sample-100length-w-simscore/rna_cif_template.db\"\n\ntry:\n    # 1. Establish connection\n    conn = sqlite3.connect(db_path)\n    \n    # 2. Read the entire table into a DataFrame\n    # We use a standard SQL query to grab everything from the 'residues' table\n    query = \"SELECT * FROM residues\"\n    template_df = pd.read_sql_query(query, conn)\n    \n    # 3. Close the connection\n    conn.close()\n    \n    print(f\"Successfully loaded {len(template_df):,} residues into template_df.\")\n\nexcept Exception as e:\n    print(f\"Error reading database: {e}\")\n\n# Display the first few rows to verify the columns\ntemplate_df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:52:41.31651Z","iopub.execute_input":"2026-01-27T08:52:41.316841Z","iopub.status.idle":"2026-01-27T08:54:04.470137Z","shell.execute_reply.started":"2026-01-27T08:52:41.316813Z","shell.execute_reply":"2026-01-27T08:54:04.469219Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"template_df[template_df['label_alt_id']!='.'][['label_alt_id']].value_counts()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:54:04.471243Z","iopub.execute_input":"2026-01-27T08:54:04.471612Z","iopub.status.idle":"2026-01-27T08:54:05.370942Z","shell.execute_reply.started":"2026-01-27T08:54:04.471571Z","shell.execute_reply":"2026-01-27T08:54:05.370123Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"template_rev_df = template_df[template_df['label_alt_id'].isin(['.', 'A'])]\ntemplate_rev_df['label_alt_id'].value_counts()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:54:05.372117Z","iopub.execute_input":"2026-01-27T08:54:05.372472Z","iopub.status.idle":"2026-01-27T08:54:08.929869Z","shell.execute_reply.started":"2026-01-27T08:54:05.372415Z","shell.execute_reply":"2026-01-27T08:54:08.928902Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1. Rename columns to match your target table\ntemplate_rev_df = template_rev_df.copy().rename(columns={\n    'pdb_id': 'ID',\n    'label_comp_id': 'resname',\n    'label_seq_id': 'resid',\n    'x': 'x_1',\n    'y': 'y_1',\n    'z': 'z_1',\n    'label_asym_id': 'chain'\n})\n\n# 2. Select and reorder only the columns shown in your example\n# We exclude 'group_PDB', 'atom_id', etc., to keep it clean\ntemplate_rev_df = template_rev_df[['ID', 'resname', 'resid', 'x_1', 'y_1', 'z_1', 'chain']]\ntemplate_rev_df['ID'] = template_rev_df['ID'] + 'database'\n\n# 4. Final check of the first few rows\nprint(template_rev_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:54:08.93252Z","iopub.execute_input":"2026-01-27T08:54:08.932832Z","iopub.status.idle":"2026-01-27T08:54:17.901503Z","shell.execute_reply.started":"2026-01-27T08:54:08.932807Z","shell.execute_reply":"2026-01-27T08:54:17.900718Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"template_struct = extract_structures(template_rev_df)\ntype(template_struct)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:54:17.902532Z","iopub.execute_input":"2026-01-27T08:54:17.902852Z","iopub.status.idle":"2026-01-27T08:54:41.277144Z","shell.execute_reply.started":"2026-01-27T08:54:17.902825Z","shell.execute_reply":"2026-01-27T08:54:41.276197Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"template_seq = (\n    template_rev_df.groupby('ID')['resname']\n    .agg(''.join)\n    .reset_index()\n)\ntemplate_seq.columns = ['target_id', 'sequence']\n\ntemplate_seq.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:54:41.278795Z","iopub.execute_input":"2026-01-27T08:54:41.279116Z","iopub.status.idle":"2026-01-27T08:54:43.849167Z","shell.execute_reply.started":"2026-01-27T08:54:41.279085Z","shell.execute_reply":"2026-01-27T08:54:43.848482Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_val_struct = extract_structures(pd.concat([train_labels, val_labels], axis=0))\ntrain_struct = extract_structures(train_labels)\ntemplate_final_struct = template_struct | train_val_struct\n\nprint(f\"Number of structures: {len(template_final_struct)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:54:43.85025Z","iopub.execute_input":"2026-01-27T08:54:43.850665Z","iopub.status.idle":"2026-01-27T08:55:37.02955Z","shell.execute_reply.started":"2026-01-27T08:54:43.850627Z","shell.execute_reply":"2026-01-27T08:55:37.028499Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"template_final_seq = pd.concat([\n    template_seq[['target_id', 'sequence']], \n    train[['target_id', 'sequence']],\n    val[['target_id', 'sequence']]\n    ], ignore_index=True)\n\ntemplate_final_seq.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:55:37.031034Z","iopub.execute_input":"2026-01-27T08:55:37.032112Z","iopub.status.idle":"2026-01-27T08:55:37.052855Z","shell.execute_reply.started":"2026-01-27T08:55:37.032075Z","shell.execute_reply":"2026-01-27T08:55:37.052013Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"template_final_seq.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:55:37.054079Z","iopub.execute_input":"2026-01-27T08:55:37.054464Z","iopub.status.idle":"2026-01-27T08:55:37.078306Z","shell.execute_reply.started":"2026-01-27T08:55:37.054403Z","shell.execute_reply":"2026-01-27T08:55:37.077299Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"g = template_rev_df.groupby(['ID', 'chain'])\n\n# Calculate differences (Current_Coords - Previous_Coords)\ndx = template_rev_df['x_1'] - g['x_1'].shift(1)\ndy = template_rev_df['y_1'] - g['y_1'].shift(1)\ndz = template_rev_df['z_1'] - g['z_1'].shift(1)\n\n# Euclidean distance formula: sqrt(dx^2 + dy^2 + dz^2)\ntemplate_rev_df['dist_to_prev'] = np.sqrt(dx**2 + dy**2 + dz**2)\n\n# Display the result\ntemplate_rev_df[['resname', 'resid', 'dist_to_prev']].head()","metadata":{"execution":{"iopub.status.busy":"2026-01-27T08:55:37.105174Z","iopub.execute_input":"2026-01-27T08:55:37.105535Z","iopub.status.idle":"2026-01-27T08:55:43.223118Z","shell.execute_reply.started":"2026-01-27T08:55:37.105504Z","shell.execute_reply":"2026-01-27T08:55:43.221987Z"}}},{"cell_type":"markdown","source":"train_labels['target_id'] = train_labels['ID'].apply(lambda x: '_'.join(str(x).split('_')[:-1]))\n\nh = train_labels.groupby(['target_id', 'chain'])\n\n# Calculate differences (Current_Coords - Previous_Coords)\ndx = train_labels['x_1'] - h['x_1'].shift(1)\ndy = train_labels['y_1'] - h['y_1'].shift(1)\ndz = train_labels['z_1'] - h['z_1'].shift(1)\n\n# Euclidean distance formula: sqrt(dx^2 + dy^2 + dz^2)\ntrain_labels['dist_to_prev'] = np.sqrt(dx**2 + dy**2 + dz**2)\n\n# Display the result\ntrain_labels[['resname', 'resid', 'dist_to_prev']].head()","metadata":{"execution":{"iopub.status.busy":"2026-01-27T08:55:43.224371Z","iopub.execute_input":"2026-01-27T08:55:43.224715Z","iopub.status.idle":"2026-01-27T08:55:49.461875Z","shell.execute_reply.started":"2026-01-27T08:55:43.224678Z","shell.execute_reply":"2026-01-27T08:55:49.461003Z"}}},{"cell_type":"markdown","source":"abnormal_resid = template_rev_df[template_rev_df['dist_to_prev'].abs() > 10]\nabnormal_resid.shape[0] / len(template_rev_df)","metadata":{"execution":{"iopub.status.busy":"2026-01-27T08:55:49.463863Z","iopub.execute_input":"2026-01-27T08:55:49.464156Z","iopub.status.idle":"2026-01-27T08:55:49.552763Z","shell.execute_reply.started":"2026-01-27T08:55:49.464129Z","shell.execute_reply":"2026-01-27T08:55:49.551542Z"}}},{"cell_type":"markdown","source":"abnormal_resid_train = train_labels[train_labels['dist_to_prev'].abs() > 10]\nabnormal_resid_train.shape[0] / len(train_labels)","metadata":{"execution":{"iopub.status.busy":"2026-01-27T08:55:49.554067Z","iopub.execute_input":"2026-01-27T08:55:49.554473Z","iopub.status.idle":"2026-01-27T08:55:49.604868Z","shell.execute_reply.started":"2026-01-27T08:55:49.554411Z","shell.execute_reply":"2026-01-27T08:55:49.603936Z"}}},{"cell_type":"code","source":"prediction_records = []\n\nfor _, row in test.iterrows():\n    seq_id = row['target_id']\n    sequence = row['sequence']\n    cutoff = row.get('temporal_cutoff', None)\n\n    structs = compute_structures(\n        sequence,\n        seq_id,\n        train,\n        train_struct,\n        num_outputs=Config.NUM_PREDICTIONS\n    )\n\n    for resid in range(len(sequence)):\n        record = {\n            'ID': f\"{seq_id}_{resid + 1}\",\n            'resname': sequence[resid],\n            'resid': resid + 1\n        }\n        for model_idx in range(Config.NUM_PREDICTIONS):\n            record[f'x_{model_idx + 1}'] = structs[model_idx][resid][0]\n            record[f'y_{model_idx + 1}'] = structs[model_idx][resid][1]\n            record[f'z_{model_idx + 1}'] = structs[model_idx][resid][2]\n        prediction_records.append(record)\n\nsubmission = pd.DataFrame(prediction_records)\ncol_order = ['ID', 'resname', 'resid']\nfor m in range(1, Config.NUM_PREDICTIONS + 1):\n    for axis in ['x', 'y', 'z']:\n        col_order.append(f'{axis}_{m}')\n\nsubmission = submission[col_order]\nsubmission.to_csv('submission.csv', index=False)\nsubmission.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T08:57:46.256953Z","iopub.execute_input":"2026-01-27T08:57:46.257369Z","iopub.status.idle":"2026-01-27T08:58:22.953285Z","shell.execute_reply.started":"2026-01-27T08:57:46.25734Z","shell.execute_reply":"2026-01-27T08:58:22.951521Z"}},"outputs":[],"execution_count":null}]}