{"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":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210,"isSourceIdPinned":false}],"dockerImageVersionId":31328,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-03-24T14:39:40.361168Z","iopub.execute_input":"2026-03-24T14:39:40.361465Z","iopub.status.idle":"2026-03-24T14:40:01.65214Z","shell.execute_reply.started":"2026-03-24T14:39:40.361444Z","shell.execute_reply":"2026-03-24T14:40:01.650577Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport math\nimport numpy as np\nimport pandas as pd\n\nDATA_ROOT = \"/kaggle/input/competitions/stanford-rna-3d-folding-2\"\n\nTRAIN_SEQ_PATH = os.path.join(DATA_ROOT, \"train_sequences.csv\")\nTRAIN_LBL_PATH = os.path.join(DATA_ROOT, \"train_labels.csv\")\nTEST_SEQ_PATH  = os.path.join(DATA_ROOT, \"test_sequences.csv\")\nSAMPLE_SUB_PATH = os.path.join(DATA_ROOT, \"sample_submission.csv\")\n\nprint(\"Check input root exists:\", os.path.exists(DATA_ROOT))\nprint(\"Files in data root:\")\nprint(os.listdir(DATA_ROOT))\n\nprint(\"Loading files...\")\ntrain_seq = pd.read_csv(TRAIN_SEQ_PATH)\ntest_seq = pd.read_csv(TEST_SEQ_PATH)\nsample_sub = pd.read_csv(SAMPLE_SUB_PATH)\n\ntrain_lbl = pd.read_csv(\n    TRAIN_LBL_PATH,\n    usecols=[\"ID\", \"resname\", \"resid\", \"x_1\", \"y_1\", \"z_1\"]\n)\n\nprint(\"train_seq shape:\", train_seq.shape)\nprint(\"train_lbl shape:\", train_lbl.shape)\nprint(\"test_seq shape :\", test_seq.shape)\nprint(\"sample_sub shape:\", sample_sub.shape)\n# =========================================================\n# Helper functions\n# =========================================================\ndef safe_norm(x, axis=-1, keepdims=False, eps=1e-12):\n    return np.sqrt(np.sum(x * x, axis=axis, keepdims=keepdims) + eps)\n\ndef unit_vector(x, eps=1e-12):\n    n = safe_norm(x, axis=-1, keepdims=True, eps=eps)\n    return x / n\n\ndef angle_3pts(a, b, c):\n    \"\"\"\n    angle ABC in radians\n    \"\"\"\n    ba = a - b\n    bc = c - b\n    ba_u = unit_vector(ba)\n    bc_u = unit_vector(bc)\n    cosang = np.sum(ba_u * bc_u, axis=-1)\n    cosang = np.clip(cosang, -1.0, 1.0)\n    return np.arccos(cosang)\n\ndef dihedral_4pts(p0, p1, p2, p3):\n    \"\"\"\n    torsion angle for four points in radians\n    \"\"\"\n    b0 = p1 - p0\n    b1 = p2 - p1\n    b2 = p3 - p2\n\n    b1_u = unit_vector(b1)\n\n    v = b0 - np.sum(b0 * b1_u, axis=-1, keepdims=True) * b1_u\n    w = b2 - np.sum(b2 * b1_u, axis=-1, keepdims=True) * b1_u\n\n    x = np.sum(v * w, axis=-1)\n    y = np.sum(np.cross(b1_u, v) * w, axis=-1)\n    return np.arctan2(y, x)\n\ndef circular_mean_std(angles):\n    \"\"\"\n    angles: np.array in radians\n    returns circular mean and approx circular std\n    \"\"\"\n    if len(angles) == 0:\n        return 0.0, 1.0\n\n    s = np.mean(np.sin(angles))\n    c = np.mean(np.cos(angles))\n    mean_angle = np.arctan2(s, c)\n\n    # unwrap around mean for a rough std\n    diff = np.angle(np.exp(1j * (angles - mean_angle)))\n    std = np.std(diff)\n    return mean_angle, std\n\ndef robust_mean_std(x, default_mean, default_std):\n    x = np.asarray(x, dtype=np.float64)\n    x = x[np.isfinite(x)]\n    if len(x) == 0:\n        return default_mean, default_std\n    return float(np.mean(x)), float(np.std(x) + 1e-6)\n\ndef place_next_point(a, b, c, bond_length, bond_angle, dihedral_angle):\n    \"\"\"\n    Given previous three points a,b,c, place next point d using NeRF-like construction.\n    bond_length: |d-c|\n    bond_angle: angle between (b-c) and (d-c)\n    dihedral_angle: torsion for a-b-c-d\n    \"\"\"\n    bc = c - b\n    bc_u = bc / (np.linalg.norm(bc) + 1e-12)\n\n    n = np.cross(b - a, bc_u)\n    n_norm = np.linalg.norm(n)\n\n    # 如果前三点接近共线，构造一个稳定法向量\n    if n_norm < 1e-8:\n        ref = np.array([1.0, 0.0, 0.0], dtype=np.float64)\n        if abs(np.dot(ref, bc_u)) > 0.9:\n            ref = np.array([0.0, 1.0, 0.0], dtype=np.float64)\n        n = np.cross(ref, bc_u)\n        n_norm = np.linalg.norm(n)\n\n    n_u = n / (n_norm + 1e-12)\n    m = np.cross(n_u, bc_u)\n\n    d = c + bond_length * (\n        -np.cos(bond_angle) * bc_u\n        + np.sin(bond_angle) * (\n            np.cos(dihedral_angle) * m\n            + np.sin(dihedral_angle) * n_u\n        )\n    )\n    return d\n\n# =========================================================\n# Parse target_id from label ID\n# Example: 157D_1 -> 157D\n# =========================================================\nprint(\"Parsing target ids from train_labels...\")\ntrain_lbl[\"target_id\"] = train_lbl[\"ID\"].str.rsplit(\"_\", n=1).str[0]\n\n# 只保留有效坐标\nvalid_xyz = train_lbl[[\"x_1\", \"y_1\", \"z_1\"]].notna().all(axis=1)\ntrain_lbl = train_lbl.loc[valid_xyz].copy()\n\ncoords = train_lbl[[\"x_1\", \"y_1\", \"z_1\"]].to_numpy(dtype=np.float64)\nbases = train_lbl[\"resname\"].astype(str).to_numpy()\ntarget_ids = train_lbl[\"target_id\"].astype(str).to_numpy()\n\nprint(\"usable train label rows:\", len(train_lbl))\n\n# =========================================================\n# Collect internal-coordinate statistics from train data\n# =========================================================\n# We condition:\n# 1) bond length on dinucleotide prev+curr\n# 2) bond angle on center nucleotide\n# 3) dihedral on center pair seq[i-2] + seq[i-1]\n#\n# Fallbacks:\n# global length / angle / dihedral stats\n\nlength_by_dinuc = {}\nangle_by_center = {}\ndihedral_by_pair = {}\n\nglobal_lengths = []\nglobal_angles = []\nglobal_dihedrals = []\n\nn = len(train_lbl)\n\n# ---- bond lengths: i-1 -> i\nsame_prev = target_ids[1:] == target_ids[:-1]\nprev_b = bases[:-1][same_prev]\ncurr_b = bases[1:][same_prev]\n\nvec_len = coords[1:] - coords[:-1]\ndist = np.linalg.norm(vec_len, axis=1)[same_prev]\n\nfor key, val in zip(prev_b + curr_b, dist):\n    length_by_dinuc.setdefault(key, []).append(float(val))\nglobal_lengths.extend(dist.tolist())\n\n# ---- bond angles: i-2, i-1, i\nsame_triple = (\n    (target_ids[2:] == target_ids[1:-1]) &\n    (target_ids[1:-1] == target_ids[:-2])\n)\n\na = coords[:-2][same_triple]\nb = coords[1:-1][same_triple]\nc = coords[2:][same_triple]\ncenter_base = bases[1:-1][same_triple]\n\nang = angle_3pts(a, b, c)\n\nfor key, val in zip(center_base, ang):\n    angle_by_center.setdefault(key, []).append(float(val))\nglobal_angles.extend(ang.tolist())\n\n# ---- dihedral: i-3, i-2, i-1, i\nsame_quad = (\n    (target_ids[3:] == target_ids[2:-1]) &\n    (target_ids[2:-1] == target_ids[1:-2]) &\n    (target_ids[1:-2] == target_ids[:-3])\n)\n\np0 = coords[:-3][same_quad]\np1 = coords[1:-2][same_quad]\np2 = coords[2:-1][same_quad]\np3 = coords[3:][same_quad]\n\npair_key = (bases[1:-2][same_quad] + bases[2:-1][same_quad])\ndih = dihedral_4pts(p0, p1, p2, p3)\n\nfor key, val in zip(pair_key, dih):\n    dihedral_by_pair.setdefault(key, []).append(float(val))\nglobal_dihedrals.extend(dih.tolist())\n\nprint(\"Collected:\")\nprint(\"  dinuc length keys   :\", len(length_by_dinuc))\nprint(\"  center angle keys   :\", len(angle_by_center))\nprint(\"  pair dihedral keys  :\", len(dihedral_by_pair))\n\n# =========================================================\n# Build summary statistics\n# =========================================================\nglobal_length_mean, global_length_std = robust_mean_std(global_lengths, 6.0, 0.5)\nglobal_angle_mean, global_angle_std   = robust_mean_std(global_angles, 2.0, 0.3)\nglobal_dih_mean, global_dih_std       = circular_mean_std(np.array(global_dihedrals, dtype=np.float64))\n\nlength_stats = {}\nfor k, vals in length_by_dinuc.items():\n    length_stats[k] = robust_mean_std(vals, global_length_mean, global_length_std)\n\nangle_stats = {}\nfor k, vals in angle_by_center.items():\n    angle_stats[k] = robust_mean_std(vals, global_angle_mean, global_angle_std)\n\ndihedral_stats = {}\nfor k, vals in dihedral_by_pair.items():\n    mean_k, std_k = circular_mean_std(np.array(vals, dtype=np.float64))\n    dihedral_stats[k] = (mean_k, std_k)\n\nprint(\"Global stats:\")\nprint(\"  bond length mean/std :\", global_length_mean, global_length_std)\nprint(\"  bond angle  mean/std :\", global_angle_mean, global_angle_std)\nprint(\"  dihedral    mean/std :\", global_dih_mean, global_dih_std)\n\n# =========================================================\n# Geometry prediction\n# =========================================================\ndef get_length(prev_base, curr_base):\n    return length_stats.get(prev_base + curr_base, (global_length_mean, global_length_std))\n\ndef get_angle(center_base):\n    return angle_stats.get(center_base, (global_angle_mean, global_angle_std))\n\ndef get_dihedral(pair_key):\n    return dihedral_stats.get(pair_key, (global_dih_mean, global_dih_std))\n\ndef build_single_structure(sequence, dihedral_scale=0.0):\n    \"\"\"\n    Return coordinates of shape (L, 3)\n    dihedral_scale controls how far we move away from mean torsion.\n    \"\"\"\n    seq = list(sequence)\n    L = len(seq)\n\n    if L == 0:\n        return np.zeros((0, 3), dtype=np.float64)\n\n    if L == 1:\n        return np.array([[0.0, 0.0, 0.0]], dtype=np.float64)\n\n    # first bond\n    d01_mean, _ = get_length(seq[0], seq[1])\n\n    if L == 2:\n        return np.array([\n            [0.0, 0.0, 0.0],\n            [d01_mean, 0.0, 0.0]\n        ], dtype=np.float64)\n\n    # second bond and first angle\n    d12_mean, _ = get_length(seq[1], seq[2])\n    a1_mean, _ = get_angle(seq[1])\n\n    p0 = np.array([0.0, 0.0, 0.0], dtype=np.float64)\n    p1 = np.array([d01_mean, 0.0, 0.0], dtype=np.float64)\n\n    # place p2 in xy plane\n    p2 = p1 + d12_mean * np.array([\n        -np.cos(a1_mean),\n         np.sin(a1_mean),\n         0.0\n    ], dtype=np.float64)\n\n    pts = [p0, p1, p2]\n\n    for i in range(3, L):\n        # length between i-1 and i\n        bond_mean, _ = get_length(seq[i - 1], seq[i])\n\n        # angle at i-1\n        angle_mean, _ = get_angle(seq[i - 1])\n\n        # dihedral for (i-3, i-2, i-1, i), keyed by center pair (i-2, i-1)\n        dih_key = seq[i - 2] + seq[i - 1]\n        dih_mean, dih_std = get_dihedral(dih_key)\n\n        dih_val = dih_mean + dihedral_scale * dih_std\n\n        new_pt = place_next_point(\n            pts[i - 3],\n            pts[i - 2],\n            pts[i - 1],\n            bond_mean,\n            angle_mean,\n            dih_val\n        )\n        pts.append(new_pt)\n\n    xyz = np.stack(pts, axis=0)\n\n    # Center coordinates for numerical stability\n    xyz = xyz - xyz.mean(axis=0, keepdims=True)\n\n    return xyz\n\ndef build_five_structures(sequence):\n    \"\"\"\n    Produce 5 candidate conformations.\n    \"\"\"\n    scales = [-1.5, -0.75, 0.0, 0.75, 1.5]\n    preds = [build_single_structure(sequence, dihedral_scale=s) for s in scales]\n    return preds\n\n# =========================================================\n# Build predictions for test set\n# =========================================================\nprint(\"Predicting test structures...\")\n\npred_frames = []\n\nfor idx, row in test_seq.iterrows():\n    target_id = row[\"target_id\"]\n    sequence = row[\"sequence\"]\n\n    preds = build_five_structures(sequence)\n\n    L = len(sequence)\n    out = pd.DataFrame({\n        \"ID\": [f\"{target_id}_{i+1}\" for i in range(L)],\n        \"resname\": list(sequence),\n        \"resid\": np.arange(1, L + 1, dtype=np.int32)\n    })\n\n    for k in range(5):\n        out[f\"x_{k+1}\"] = preds[k][:, 0]\n        out[f\"y_{k+1}\"] = preds[k][:, 1]\n        out[f\"z_{k+1}\"] = preds[k][:, 2]\n\n    pred_frames.append(out)\n\n    if (idx + 1) % 5 == 0 or (idx + 1) == len(test_seq):\n        print(f\"  finished {idx+1}/{len(test_seq)} targets\")\n\npred_df = pd.concat(pred_frames, axis=0, ignore_index=True)\n\n# =========================================================\n# Align to sample_submission row order\n# =========================================================\nsubmission = sample_sub[[\"ID\", \"resname\", \"resid\"]].merge(\n    pred_df,\n    on=[\"ID\", \"resname\", \"resid\"],\n    how=\"left\"\n)\n\n# 理论上不应该有缺失；保险起见补 0\ncoord_cols = [c for c in submission.columns if c not in [\"ID\", \"resname\", \"resid\"]]\nsubmission[coord_cols] = submission[coord_cols].fillna(0.0)\n\n# 列顺序严格按 sample submission\nsubmission = submission[\n    [\"ID\", \"resname\", \"resid\"] +\n    [f\"{axis}_{k}\" for k in range(1, 6) for axis in [\"x\", \"y\", \"z\"]]\n]\n\nsubmission.to_csv(\"/kaggle/working/submission.csv\", index=False)\n\nimport os\nprint(\"Done. submission.csv saved.\")\nprint(\"exists:\", os.path.exists(\"/kaggle/working/submission.csv\"))\nprint(\"working files:\", os.listdir(\"/kaggle/working\"))\nprint(submission.head())\nprint(\"submission shape:\", submission.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-24T14:40:01.653938Z","iopub.execute_input":"2026-03-24T14:40:01.654175Z","iopub.status.idle":"2026-03-24T14:40:37.503889Z","shell.execute_reply.started":"2026-03-24T14:40:01.654155Z","shell.execute_reply":"2026-03-24T14:40:37.502964Z"}},"outputs":[],"execution_count":null}]}