{"cells":[{"cell_type":"markdown","metadata":{},"source":"# Stanford RNA 3D Folding 2 — Template Matching v1 | WandB Offline Sync via kaggle-wandb-sync\n\n## Synced with kaggle-wandb-sync\n\nThis notebook uses **[kaggle-wandb-sync](https://pypi.org/project/kaggle-wandb-sync/)** to track experiments with Weights & Biases — even with internet disabled.\n\nSince Kaggle competition notebooks run with internet **disabled**, you can't push W&B metrics in real time.\n`kaggle-wandb-sync` solves this with a simple offline sync pipeline:\n\n```\nNotebook (WANDB_MODE=offline)\n    → W&B logs saved to /kaggle/working/wandb/\n    → kaggle kernels output  (download via GitHub Actions)\n    → wandb sync             (push to W&B cloud)\n```\n\n### How to use\n\n```bash\npip install kaggle-wandb-sync\n\n# All-in-one: push notebook → poll → download output → wandb sync\nexport WANDB_API_KEY=your_api_key\nkaggle-wandb-sync run stanford-rna-3d-folding-2/\n\n# Or step by step:\nkaggle-wandb-sync push  stanford-rna-3d-folding-2/\nkaggle-wandb-sync poll  yasunorim/stanford-rna-3d-folding-2-baseline\nkaggle-wandb-sync output yasunorim/stanford-rna-3d-folding-2-baseline\nkaggle-wandb-sync sync  ./kaggle_output\n```\n\n→ **kaggle-wandb-sync GitHub**: https://github.com/yasumorishima/kaggle-wandb-sync  \n→ **kaggle-wandb-sync PyPI**: https://pypi.org/project/kaggle-wandb-sync/  \n→ **Notebook source**: https://github.com/yasumorishima/kaggle-competitions/tree/main/stanford-rna-3d-folding-2\n\n---\n\n## Competition Overview\n\n**Task**: Predict the 3D structure of RNA molecules from sequence alone.  \n**Metric**: TM-score (best of 5 predictions per sequence, averaged across targets)  \n**Output**: x, y, z coordinates of the C1' atom for every nucleotide — 5 structures per sequence.\n\n**This notebook (template-v1)**: Use training data 3D coordinates as templates.  \nFor each test sequence, find the 5 training structures with the closest sequence length, then interpolate coordinates to match the test length. Real RNA structures provide much better starting points than idealized helices."},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# W&B must be set to offline BEFORE importing wandb\nimport os\nos.environ['WANDB_MODE'] = 'offline'\nos.environ['WANDB_PROJECT'] = 'stanford-rna-3d-folding-2'\nos.environ['WANDB_RUN_GROUP'] = 'template'\n\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom scipy.interpolate import interp1d\nimport warnings\nwarnings.filterwarnings('ignore')\n\nimport wandb\n\nprint('Libraries loaded.')\nprint(f'WANDB_MODE: {os.environ[\"WANDB_MODE\"]}')"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# --- Data path detection ---\nINPUT_ROOT = Path('/kaggle/input')\nSLUG = 'stanford-rna-3d-folding-2'\n\nprint('=== /kaggle/input/ structure ===')\nfor p in sorted(INPUT_ROOT.iterdir()):\n    print(f'  {p.name}/')\n    for sub in sorted(p.iterdir())[:5]:\n        print(f'    {sub.name}')\n\n# Auto-detect DATA_DIR\nDATA_DIR = None\nfor p in INPUT_ROOT.rglob('test_sequences.csv'):\n    DATA_DIR = p.parent\n    break\n\nif DATA_DIR is None:\n    raise FileNotFoundError(f'test_sequences.csv not found under {INPUT_ROOT}')\n\nprint(f'\\nDATA_DIR: {DATA_DIR}')\nprint('\\nCSV files:')\nfor f in sorted(DATA_DIR.glob('*.csv')):\n    print(f'  {f.name}')"},{"cell_type":"markdown","metadata":{},"source":"## 1. Load Data (Test + Training)"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"test_df = pd.read_csv(DATA_DIR / 'test_sequences.csv')\nsample_sub = pd.read_csv(DATA_DIR / 'sample_submission.csv')\ntrain_seq_df = pd.read_csv(DATA_DIR / 'train_sequences.csv')\ntrain_labels_df = pd.read_csv(DATA_DIR / 'train_labels.csv')\n\nprint(f'Test sequences:     {len(test_df)} rows')\nprint(f'Sample submission:  {len(sample_sub)} rows')\nprint(f'Train sequences:    {len(train_seq_df)} rows')\nprint(f'Train labels:       {len(train_labels_df)} rows')\nprint(f'\\nTrain sequences columns: {list(train_seq_df.columns)}')\nprint(f'Train labels columns:    {list(train_labels_df.columns)}')\ntrain_seq_df.head(3)"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Build training structure library: {target_id: (length, xyz_array)}\n# train_labels has columns like: ID, resname, resid, x_1, y_1, z_1\n# We use the first structure (x_1, y_1, z_1) as the representative coords\n\n# Detect coordinate columns\nlabel_cols = train_labels_df.columns.tolist()\nprint(f'Train labels columns: {label_cols}')\n\n# Extract target_id from ID column (format: targetid_resid)\ntrain_labels_df['_target'] = train_labels_df['ID'].str.rsplit('_', n=1).str[0]\n\n# Build library of training structures\ntrain_structures = {}\nfor target_id, group in train_labels_df.groupby('_target', sort=False):\n    coords = group[['x_1', 'y_1', 'z_1']].values.astype(np.float64)\n    # Skip structures with NaN coordinates\n    if np.isnan(coords).any():\n        nan_frac = np.isnan(coords).mean()\n        if nan_frac > 0.5:\n            continue\n        # For partial NaN, interpolate missing values\n        for col in range(3):\n            mask = np.isnan(coords[:, col])\n            if mask.any() and not mask.all():\n                valid = np.where(~mask)[0]\n                coords[mask, col] = np.interp(\n                    np.where(mask)[0], valid, coords[valid, col]\n                )\n    train_structures[target_id] = coords\n\n# Compute lengths\ntrain_lengths = {tid: len(c) for tid, c in train_structures.items()}\ntrain_ids = np.array(list(train_lengths.keys()))\ntrain_lens = np.array(list(train_lengths.values()))\n\nprint(f'\\nTraining structures loaded: {len(train_structures)}')\nprint(f'Length range: {train_lens.min()} - {train_lens.max()} (mean: {train_lens.mean():.1f})')\n\n# Test sequence lengths (from sample_submission)\nsample_sub['_target'] = sample_sub['ID'].str.rsplit('_', n=1).str[0]\ntest_lengths = sample_sub.groupby('_target', sort=False).size()\nprint(f'\\nTest sequences: {len(test_lengths)}')\nprint(f'Test length range: {test_lengths.min()} - {test_lengths.max()} (mean: {test_lengths.mean():.1f})')"},{"cell_type":"markdown","metadata":{},"source":"## 2. W&B Initialization"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"SEQ_COL = 'sequence' if 'sequence' in test_df.columns else test_df.columns[1]\nseq_lengths = test_df[SEQ_COL].str.len()\n\nrun = wandb.init(\n    project='stanford-rna-3d-folding-2',\n    name='template-v1',\n    config={\n        'approach': 'template_matching',\n        'n_structures': 5,\n        'n_templates': 5,\n        'interpolation': 'linear',\n        'n_train_structures': len(train_structures),\n    }\n)\n\nwandb.log({\n    'n_test_sequences': len(test_lengths),\n    'seq_len_min': int(test_lengths.min()),\n    'seq_len_max': int(test_lengths.max()),\n    'seq_len_mean': float(test_lengths.mean()),\n    'n_train_structures': len(train_structures),\n    'train_len_min': int(train_lens.min()),\n    'train_len_max': int(train_lens.max()),\n    'train_len_mean': float(train_lens.mean()),\n})\n\nprint(f'W&B run: {run.name} (mode: {os.environ[\"WANDB_MODE\"]})')"},{"cell_type":"markdown","metadata":{},"source":"## 3. Prediction — Template Matching with Interpolation\n\nFor each test sequence:\n1. Find the 5 training structures with the closest sequence length\n2. If the training structure has a different length, use linear interpolation (scipy) to resize the 3D coordinates to match the test length\n3. Each of the 5 templates becomes one of the 5 required structure predictions\n\nThis leverages real RNA folding patterns from the training data, which should produce much better predictions than idealized helices. The diversity across 5 different templates also provides meaningful structural variety for best-of-5 TM-score evaluation."},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"N_TEMPLATES = 5\n\n\ndef interpolate_coords(coords: np.ndarray, target_len: int) -> np.ndarray:\n    \"\"\"Resize 3D coordinates from source length to target length using linear interpolation.\n\n    Args:\n        coords: (source_len, 3) array of x, y, z coordinates\n        target_len: desired output length\n\n    Returns:\n        (target_len, 3) array of interpolated coordinates\n    \"\"\"\n    source_len = len(coords)\n    if source_len == target_len:\n        return coords.copy()\n\n    # Normalize positions to [0, 1] for both source and target\n    src_positions = np.linspace(0, 1, source_len)\n    tgt_positions = np.linspace(0, 1, target_len)\n\n    result = np.zeros((target_len, 3))\n    for dim in range(3):\n        f = interp1d(src_positions, coords[:, dim], kind='linear')\n        result[:, dim] = f(tgt_positions)\n\n    return result\n\n\ndef find_closest_templates(target_len: int, n: int = 5) -> list:\n    \"\"\"Find n training structures with lengths closest to target_len.\n\n    Returns list of (train_target_id, length_diff) tuples.\n    \"\"\"\n    diffs = np.abs(train_lens - target_len)\n    # Get indices of n smallest differences\n    closest_idx = np.argsort(diffs)[:n]\n    return [(train_ids[i], diffs[i]) for i in closest_idx]\n\n\n# Use sample_submission as template\nsubmission = sample_sub.copy()\n\ntemplate_stats = []\n\nfor target_id, group in submission.groupby('_target', sort=False):\n    test_len = len(group)\n    idx = group.index\n\n    # Find 5 closest training structures\n    templates = find_closest_templates(test_len, N_TEMPLATES)\n\n    for s, (train_tid, len_diff) in enumerate(templates, 1):\n        train_coords = train_structures[train_tid]\n        # Interpolate to match test length\n        pred_coords = interpolate_coords(train_coords, test_len)\n\n        submission.loc[idx, f'x_{s}'] = pred_coords[:, 0].round(3)\n        submission.loc[idx, f'y_{s}'] = pred_coords[:, 1].round(3)\n        submission.loc[idx, f'z_{s}'] = pred_coords[:, 2].round(3)\n\n    template_stats.append({\n        'target_id': target_id,\n        'test_len': test_len,\n        'templates': [(tid, int(d)) for tid, d in templates],\n        'avg_len_diff': float(np.mean([d for _, d in templates])),\n    })\n\nsubmission = submission.drop(columns=['_target'])\n\n# Log template matching stats\navg_diff = np.mean([s['avg_len_diff'] for s in template_stats])\nmax_diff = max(s['avg_len_diff'] for s in template_stats)\nwandb.log({\n    'avg_template_len_diff': avg_diff,\n    'max_template_len_diff': max_diff,\n})\n\nprint(f'Submission shape: {submission.shape}')\nprint(f'Average template length difference: {avg_diff:.1f}')\nprint(f'Max average length difference: {max_diff:.1f}')\nprint(f'\\nTemplate assignments (first 5):')\nfor s in template_stats[:5]:\n    t_info = ', '.join([f'{tid}(±{d})' for tid, d in s['templates']])\n    print(f'  {s[\"target_id\"]} (len={s[\"test_len\"]}): {t_info}')\n\nsubmission.head(3)"},{"cell_type":"markdown","metadata":{},"source":"## 4. Save Submission"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"OUTPUT_PATH = Path('/kaggle/working/submission.csv')\nsubmission.to_csv(OUTPUT_PATH, index=False)\nprint(f'Saved: {OUTPUT_PATH}  ({OUTPUT_PATH.stat().st_size / 1024:.1f} KB)')\n\nwandb.log({\n    'n_submission_rows': len(submission),\n    'submission_columns': len(submission.columns),\n})\n\n# Log template assignment table\ntemplate_table = wandb.Table(\n    columns=['target_id', 'test_len', 'template_1', 'template_2', 'template_3', 'template_4', 'template_5', 'avg_len_diff'],\n    data=[\n        [\n            s['target_id'], s['test_len'],\n            *[f\"{tid}(±{d})\" for tid, d in s['templates']],\n            round(s['avg_len_diff'], 1),\n        ]\n        for s in template_stats\n    ],\n)\nwandb.log({'template_assignments': template_table})\n\nwandb.finish()\nprint('\\nW&B run finished (offline). Sync with:')\nprint('  kaggle-wandb-sync run stanford-rna-3d-folding-2/ --skip-push')"},{"cell_type":"markdown","metadata":{},"source":"## Summary\n\n| Step | Detail |\n|---|---|\n| Approach | Template matching — closest-length training structures |\n| Templates | 5 per test sequence (closest by sequence length) |\n| Interpolation | Linear (scipy) to resize coordinates |\n| W&B run | template-v1 |\n| W&B mode | offline → synced via kaggle-wandb-sync |\n\n### Why template matching?\n\nThe previous approach (A-form helix + noise) places all residues on a single helix, which doesn't reflect actual RNA folding. Real RNA structures from the training data contain stems, loops, hairpins, and other motifs. By using training structures as templates:\n- We get realistic 3D geometries instead of idealized helices\n- Different templates provide meaningful structural diversity\n- Linear interpolation preserves the overall shape while adjusting to the target length\n\n### Sync W&B runs after execution\n\n```bash\n# After kaggle kernels push:\nkaggle-wandb-sync run stanford-rna-3d-folding-2/\n\n# Or to re-sync without re-running:\nkaggle-wandb-sync run stanford-rna-3d-folding-2/ --skip-push\n```"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.0"}},"nbformat":4,"nbformat_minor":4}