{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.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":87793,"databundleVersionId":11228175,"sourceType":"competition"},{"sourceId":7639698,"sourceType":"datasetVersion","datasetId":4299272},{"sourceId":8318191,"sourceType":"datasetVersion","datasetId":4459124}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install einops","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\")\n\nimport os\nimport sys\nimport random\nimport pickle\nimport yaml\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport matplotlib.pyplot as plt\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 1. CONFIG & SEED","metadata":{}},{"cell_type":"code","source":"def set_seed(seed: int):\n    \"\"\"Set a random seed for Python, NumPy, PyTorch (CPU & GPU) to ensure reproducibility.\"\"\"\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\n# Example configuration (you can load this from a YAML, JSON, etc.)\nconfig = {\n    \"seed\": 42,\n    \"cutoff_date\": \"2020-01-01\",\n    \"test_cutoff_date\": \"2022-05-01\",\n    \"max_len\": 384,\n    \"batch_size\": 1,\n    \"model_config_path\": \"/kaggle/input/ribonanzanet2d-final/configs/pairwise.yaml\",\n    \"max_len_filter\": 9999999,\n    \"min_len_filter\": 10,\n    \n    \"train_sequences_path\": \"/kaggle/input/stanford-rna-3d-folding/train_sequences.csv\",\n    \"train_labels_path\": \"/kaggle/input/stanford-rna-3d-folding/train_labels.csv\",\n    \"pretrained_weights_path\": \"/kaggle/input/ribonanzanet-weights/RibonanzaNet.pt\",\n    \"save_weights_name\": \"RibonanzaNet-3D.pt\",\n    \"save_weights_final\": \"RibonanzaNet-3D-final.pt\",\n}\n\n# Set the seed for reproducibility\nset_seed(config[\"seed\"])","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2. DATA LOADING & PREPARATION","metadata":{}},{"cell_type":"code","source":"# Load CSVs\ntrain_sequences = pd.read_csv(config[\"train_sequences_path\"])\ntrain_labels = pd.read_csv(config[\"train_labels_path\"])\n\n# Create a pdb_id field\ntrain_labels[\"pdb_id\"] = train_labels[\"ID\"].apply(\n    lambda x: x.split(\"_\")[0] + \"_\" + x.split(\"_\")[1]\n)\n\n# Collect xyz data for each sequence\nall_xyz = []\nfor pdb_id in tqdm(train_sequences[\"target_id\"], desc=\"Collecting XYZ data\"):\n    df = train_labels[train_labels[\"pdb_id\"] == pdb_id]\n    xyz = df[[\"x_1\", \"y_1\", \"z_1\"]].to_numpy().astype(\"float32\")\n    xyz[xyz < -1e17] = float(\"nan\")\n    all_xyz.append(xyz)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3. DATA FILTERING","metadata":{}},{"cell_type":"code","source":"valid_indices = []\nmax_len_seen = 0\n\nfor i, xyz in enumerate(all_xyz):\n    # Track the maximum length\n    if len(xyz) > max_len_seen:\n        max_len_seen = len(xyz)\n\n    nan_ratio = np.isnan(xyz).mean()\n    seq_len = len(xyz)\n    # Keep sequence if it meets criteria\n    if (nan_ratio <= 0.5) and (config[\"min_len_filter\"] < seq_len < config[\"max_len_filter\"]):\n        valid_indices.append(i)\n\nprint(f\"Longest sequence in train: {max_len_seen}\")\n\n# Filter sequences & xyz based on valid_indices\ntrain_sequences = train_sequences.loc[valid_indices].reset_index(drop=True)\nall_xyz = [all_xyz[i] for i in valid_indices]\n\n# Prepare final data dictionary\ndata = {\n    \"sequence\": train_sequences[\"sequence\"].tolist(),\n    \"temporal_cutoff\": train_sequences[\"temporal_cutoff\"].tolist(),\n    \"description\": train_sequences[\"description\"].tolist(),\n    \"all_sequences\": train_sequences[\"all_sequences\"].tolist(),\n    \"xyz\": all_xyz,\n}","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 4. TRAIN / VAL SPLIT","metadata":{}},{"cell_type":"code","source":"cutoff_date = pd.Timestamp(config[\"cutoff_date\"])\ntest_cutoff_date = pd.Timestamp(config[\"test_cutoff_date\"])\n\ntrain_indices = [i for i, date_str in enumerate(data[\"temporal_cutoff\"]) if pd.Timestamp(date_str) <= cutoff_date]\ntest_indices = [i for i, date_str in enumerate(data[\"temporal_cutoff\"]) if cutoff_date < pd.Timestamp(date_str) <= test_cutoff_date]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 5. DATASET & DATALOADER","metadata":{}},{"cell_type":"code","source":"class RNA3D_Dataset(Dataset):\n    \"\"\"\n    A PyTorch Dataset for 3D RNA structures.\n    \"\"\"\n    def __init__(self, indices, data_dict, max_len=384):\n        self.indices = indices\n        self.data = data_dict\n        self.max_len = max_len\n        self.nt_to_idx = {nt: i for i, nt in enumerate(\"ACGU\")}\n\n    def __len__(self):\n        return len(self.indices)\n    \n    def __getitem__(self, idx):\n        data_idx = self.indices[idx]\n        # Convert nucleotides to integer tokens\n        sequence = [self.nt_to_idx[nt] for nt in self.data[\"sequence\"][data_idx]]\n        sequence = torch.tensor(sequence, dtype=torch.long)\n        # Convert xyz to torch tensor\n        xyz = torch.tensor(self.data[\"xyz\"][data_idx], dtype=torch.float32)\n\n        # If sequence is longer than max_len, randomly crop\n        if len(sequence) > self.max_len:\n            crop_start = np.random.randint(len(sequence) - self.max_len)\n            crop_end = crop_start + self.max_len\n            sequence = sequence[crop_start:crop_end]\n            xyz = xyz[crop_start:crop_end]\n\n        return {\"sequence\": sequence, \"xyz\": xyz}\n\ntrain_dataset = RNA3D_Dataset(train_indices, data, max_len=config[\"max_len\"])\nval_dataset = RNA3D_Dataset(test_indices, data, max_len=config[\"max_len\"])\n\ntrain_loader = DataLoader(train_dataset, batch_size=config[\"batch_size\"], shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=config[\"batch_size\"], shuffle=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 6. MODEL, CONFIG CLASSES & HELPER FUNCTIONS","metadata":{}},{"cell_type":"code","source":"sys.path.append(\"/kaggle/input/ribonanzanet2d-final\")\n\nfrom Network import RibonanzaNet\n\nclass Config:\n    \"\"\"Simple Config class that can load from a dict or YAML.\"\"\"\n    def __init__(self, **entries):\n        self.__dict__.update(entries)\n        self.entries = entries\n\n    def print(self):\n        print(self.entries)\n\ndef load_config_from_yaml(file_path):\n    with open(file_path, 'r') as file:\n        cfg = yaml.safe_load(file)\n    return Config(**cfg)\n\nclass FinetunedRibonanzaNet(RibonanzaNet):\n    \"\"\"\n    A finetuned version of RibonanzaNet adapted for predicting 3D coordinates.\n    \"\"\"\n    def __init__(self, config_obj, pretrained=False, dropout=0.1):\n        # Modify config dropout before super init, if needed\n        config_obj.dropout = dropout\n        super(FinetunedRibonanzaNet, self).__init__(config_obj)\n\n        # Load pretrained weights if requested\n        if pretrained:\n            self.load_state_dict(\n                torch.load(config[\"pretrained_weights_path\"], map_location=\"cpu\")\n            )\n\n        self.dropout = nn.Dropout(p=0.0)\n        self.xyz_predictor = nn.Linear(256, 3)\n\n    def forward(self, src):\n        \"\"\"Forward pass to predict 3D XYZ coordinates.\"\"\"\n        # get_embeddings returns (sequence_features, *some_other_outputs)\n        sequence_features, _ = self.get_embeddings(\n            src, torch.ones_like(src).long().to(src.device)\n        )\n        xyz_pred = self.xyz_predictor(sequence_features)\n        return xyz_pred\n\n# Instantiate the model\nmodel_cfg = load_config_from_yaml(config[\"model_config_path\"])\nmodel = FinetunedRibonanzaNet(model_cfg, pretrained=True).cuda()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 7. LOSS FUNCTIONS","metadata":{}},{"cell_type":"code","source":"def calculate_distance_matrix(X, Y, epsilon=1e-4):\n    \"\"\"\n    Calculate pairwise distances between every point in X and every point in Y.\n    Shape: (len(X), len(Y))\n    \"\"\"\n    return ((X[:, None] - Y[None, :])**2 + epsilon).sum(dim=-1).sqrt()\n\ndef dRMSD(pred_x, pred_y, gt_x, gt_y, epsilon=1e-4, Z=10, d_clamp=None):\n    \"\"\"\n    Distance-based RMSD.\n    pred_x, pred_y: predicted coordinates (usually the same tensor for X and Y).\n    gt_x, gt_y: ground truth coordinates.\n    \"\"\"\n    pred_dm = calculate_distance_matrix(pred_x, pred_y)\n    gt_dm = calculate_distance_matrix(gt_x, gt_y)\n\n    mask = ~torch.isnan(gt_dm)\n    mask[torch.eye(mask.shape[0], device=mask.device).bool()] = False\n\n    diff_sq = (pred_dm[mask] - gt_dm[mask])**2 + epsilon\n    if d_clamp is not None:\n        diff_sq = diff_sq.clamp(max=d_clamp**2)\n\n    return diff_sq.sqrt().mean() / Z\n\ndef local_dRMSD(pred_x, pred_y, gt_x, gt_y, epsilon=1e-4, Z=10, d_clamp=30):\n    \"\"\"\n    Local distance-based RMSD, ignoring distances above a clamp threshold.\n    \"\"\"\n    pred_dm = calculate_distance_matrix(pred_x, pred_y)\n    gt_dm = calculate_distance_matrix(gt_x, gt_y)\n\n    mask = (~torch.isnan(gt_dm)) & (gt_dm < d_clamp)\n    mask[torch.eye(mask.shape[0], device=mask.device).bool()] = False\n\n    diff_sq = (pred_dm[mask] - gt_dm[mask])**2 + epsilon\n    return diff_sq.sqrt().mean() / Z\n\ndef dRMAE(pred_x, pred_y, gt_x, gt_y, epsilon=1e-4, Z=10):\n    \"\"\"\n    Distance-based Mean Absolute Error.\n    \"\"\"\n    pred_dm = calculate_distance_matrix(pred_x, pred_y)\n    gt_dm = calculate_distance_matrix(gt_x, gt_y)\n\n    mask = ~torch.isnan(gt_dm)\n    mask[torch.eye(mask.shape[0], device=mask.device).bool()] = False\n\n    diff = torch.abs(pred_dm[mask] - gt_dm[mask])\n    return diff.mean() / Z\n\ndef align_svd_mae(input_coords, target_coords, Z=10):\n    \"\"\"\n    Align input_coords to target_coords via SVD (Kabsch algorithm) and compute MAE.\n    \"\"\"\n    assert input_coords.shape == target_coords.shape, \"Input and target must have the same shape\"\n\n    # Create mask for valid points\n    mask = ~torch.isnan(target_coords.sum(dim=-1))\n    input_coords = input_coords[mask]\n    target_coords = target_coords[mask]\n    \n    # Compute centroids\n    centroid_input = input_coords.mean(dim=0, keepdim=True)\n    centroid_target = target_coords.mean(dim=0, keepdim=True)\n\n    # Center the points\n    input_centered = input_coords - centroid_input\n    target_centered = target_coords - centroid_target\n\n    # Compute covariance matrix\n    cov_matrix = input_centered.T @ target_centered\n\n    # SVD to find optimal rotation\n    U, S, Vt = torch.svd(cov_matrix)\n    R = Vt @ U.T\n\n    # Ensure a proper rotation (determinant R == 1)\n    if torch.det(R) < 0:\n        Vt_adj = Vt.clone()   # Clone to avoid in-place modification issues\n        Vt_adj[-1, :] = -Vt_adj[-1, :]\n        R = Vt_adj @ U.T\n\n    # Rotate input and compute mean absolute error\n    aligned_input = (input_centered @ R.T) + centroid_target\n    return torch.abs(aligned_input - target_coords).mean() / Z","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 8. TRAINING LOOP","metadata":{}},{"cell_type":"code","source":"def train_model(model, train_dl, val_dl, epochs=50, cos_epoch=35, lr=3e-4, clip=1):\n    \"\"\"Train the model with a CosineAnnealingLR after `cos_epoch` epochs.\"\"\"\n    optimizer = torch.optim.AdamW(model.parameters(), weight_decay=0.0, lr=lr)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer,\n        T_max=(epochs - cos_epoch) * len(train_dl),\n    )\n\n    best_val_loss = float(\"inf\")\n    best_preds = None\n\n    for epoch in range(epochs):\n        model.train()\n        train_pbar = tqdm(train_dl, desc=f\"Training Epoch {epoch+1}/{epochs}\")\n        running_loss = 0.0\n\n        for idx, batch in enumerate(train_pbar):\n            sequence = batch[\"sequence\"].cuda()\n            gt_xyz = batch[\"xyz\"].squeeze().cuda()\n\n            pred_xyz = model(sequence).squeeze()\n\n            # Combine two distance-based losses\n            loss = dRMAE(pred_xyz, pred_xyz, gt_xyz, gt_xyz) + align_svd_mae(pred_xyz, gt_xyz)\n            loss.backward()\n\n            # Gradient clipping\n            torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n            optimizer.step()\n            optimizer.zero_grad()\n\n            if (epoch + 1) > cos_epoch:\n                scheduler.step()\n\n            running_loss += loss.item()\n            avg_loss = running_loss / (idx + 1)\n            train_pbar.set_description(f\"Epoch {epoch+1} | Loss: {avg_loss:.4f}\")\n\n        # Validation\n        model.eval()\n        val_loss = 0.0\n        val_preds = []\n        with torch.no_grad():\n            for idx, batch in enumerate(val_dl):\n                sequence = batch[\"sequence\"].cuda()\n                gt_xyz = batch[\"xyz\"].squeeze().cuda()\n\n                pred_xyz = model(sequence).squeeze()\n                loss = dRMAE(pred_xyz, pred_xyz, gt_xyz, gt_xyz)\n                val_loss += loss.item()\n\n                val_preds.append((gt_xyz.cpu().numpy(), pred_xyz.cpu().numpy()))\n\n            val_loss /= len(val_dl)\n            print(f\"Validation Loss (Epoch {epoch+1}): {val_loss:.4f}\")\n\n            # Check for improvement\n            if val_loss < best_val_loss:\n                best_val_loss = val_loss\n                best_preds = val_preds\n                torch.save(model.state_dict(), config[\"save_weights_name\"])\n                print(f\"  -> New best model saved at epoch {epoch+1}\")\n\n    # Save final model\n    torch.save(model.state_dict(), config[\"save_weights_final\"])\n    return best_val_loss, best_preds","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 9. RUN TRAINING","metadata":{}},{"cell_type":"code","source":"if __name__ == \"__main__\":\n    best_loss, best_predictions = train_model(\n        model=model,\n        train_dl=train_loader,\n        val_dl=val_loader,\n        epochs=50,         # or config[\"epochs\"]\n        cos_epoch=35,      # or config[\"cos_epoch\"]\n        lr=3e-4,\n        clip=1\n    )\n    print(f\"Best Validation Loss: {best_loss:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}