{"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"}],"dockerImageVersionId":31234,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom Bio.Seq import Seq\nfrom Bio import pairwise2\nfrom scipy.spatial.transform import Rotation\nfrom scipy.spatial.distance import pdist, squareform\n\n\nclass Config:\n    DATA_PATH = '/kaggle/input/stanford-rna-3d-folding-2'\n    MAX_RELATIVE_LENGTH_DIFF = 0.5\n    ALIGNMENT_MATCH = 2\n    ALIGNMENT_MISMATCH = -1\n    ALIGNMENT_GAP_OPEN = -10\n    ALIGNMENT_GAP_EXTEND = -0.5\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\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    for group_label, group_data in labels_data.groupby(lambda x: labels_data['ID'][x].rsplit('_', 1)[0]):\n        coords_list = []\n        for _, row_data in group_data.sort_values('resid').iterrows():\n            coords_list.append([row_data['x_1'], row_data['y_1'], row_data['z_1']])\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    align_result = pairwise2.align.globalms(\n        Seq(query_seq),\n        template_seq,\n        Config.ALIGNMENT_MATCH,\n        Config.ALIGNMENT_MISMATCH,\n        Config.ALIGNMENT_GAP_OPEN,\n        Config.ALIGNMENT_GAP_EXTEND,\n        one_alignment_only=True\n    )\n    if not align_result:\n        return None, 0.0\n    alignment = align_result[0]\n    score = alignment.score / (2 * min(len(query_seq), len(template_seq)))\n    return alignment, 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    aligned_dest = str(alignment.seqA)\n    aligned_source = str(alignment.seqB)\n\n    morphed_coords = np.full((len(dest_seq), 3), np.nan, dtype=np.float32)\n    dest_idx = source_idx = 0\n\n    for a, b in zip(aligned_dest, aligned_source):\n        if a != '-' and b != '-':\n            if source_idx < len(source_coords):\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    valid_mask = ~np.isnan(morphed_coords[:, 0])\n    if not valid_mask.any():\n        return create_default_structure(dest_seq)\n\n    indices = np.arange(len(morphed_coords))\n    for dim in range(3):\n        morphed_coords[:, dim] = np.interp(\n            indices,\n            indices[valid_mask],\n            morphed_coords[valid_mask, dim],\n            left=morphed_coords[valid_mask, dim][0] if valid_mask.any() else 0.0,\n            right=morphed_coords[valid_mask, dim][-1] if valid_mask.any() else 0.0\n        )\n\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                prev_vec = morphed_coords[i-1] - (morphed_coords[i-2] if i>=2 else np.zeros(3))\n                if np.linalg.norm(prev_vec) < 1e-6:\n                    prev_vec = np.random.normal(0, 1, 3)\n                prev_vec /= (np.linalg.norm(prev_vec) + 1e-10)\n                step = np.random.uniform(*Config.STEP_LENGTH_RANGE)\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    for i in range(n - 1):\n        vec = refined[i + 1] - refined[i]\n        dist = np.linalg.norm(vec)\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    if n > 3:\n        for i in range(1, n - 1):\n            prev = refined[i] - refined[i-1]\n            nxt = refined[i+1] - refined[i]\n            angle = np.arccos(np.dot(prev, nxt) / (np.linalg.norm(prev) * np.linalg.norm(nxt) + 1e-10))\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    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    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    if date_filter:\n        candidates = template_df[template_df['temporal_cutoff'] < date_filter]\n    else:\n        candidates = template_df\n\n    matches = []\n    for _, row in candidates.iterrows():\n        tid = row['target_id']\n        tseq = row['sequence']\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        _, score = compute_alignment_score(query_seq, tseq)\n        matches.append((tid, tseq, score, structure_dict[tid]))\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        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        refined = enhance_geometry(fake, seq, reliability=Config.DEFAULT_RELIABILITY)\n        \n        scale_factor = np.random.uniform(0.8, 1.2)\n        if len(predictions) > 0:\n            centroid_pred = np.mean(predictions[0], axis=0)\n            centroid_new = np.mean(refined, axis=0)\n            refined = centroid_pred + (refined - centroid_new) * scale_factor\n        \n        predictions.append(refined)\n\n    final_predictions = []\n    for pred in predictions:\n        final_predictions.append(pred - np.mean(pred, axis=0))\n    return final_predictions[:num_outputs]\n\ntrain_structs = extract_structures(train_labels)\nval_structs = extract_structures(val_labels)\n\nprediction_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_structs,\n        num_outputs=Config.NUM_PREDICTIONS,\n        date_limit=cutoff\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()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-07T22:41:21.823424Z","iopub.execute_input":"2026-01-07T22:41:21.823822Z","iopub.status.idle":"2026-01-07T22:41:22.4876Z","shell.execute_reply.started":"2026-01-07T22:41:21.82379Z","shell.execute_reply":"2026-01-07T22:41:22.486877Z"}},"outputs":[],"execution_count":null}]}