{"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":11553390,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":10855324,"sourceType":"datasetVersion","datasetId":6742586},{"sourceId":11278691,"sourceType":"datasetVersion","datasetId":7051341},{"sourceId":11279607,"sourceType":"datasetVersion","datasetId":7051942},{"sourceId":11469248,"sourceType":"datasetVersion","datasetId":7187409},{"sourceId":311741,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":264400,"modelId":285488}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nimport torch\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nimport random\nimport pickle\nimport os\nimport sys","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-04-19T05:33:09.644424Z","iopub.execute_input":"2025-04-19T05:33:09.644705Z","iopub.status.idle":"2025-04-19T05:33:14.101224Z","shell.execute_reply.started":"2025-04-19T05:33:09.644670Z","shell.execute_reply":"2025-04-19T05:33:14.100260Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"config = {\n    \"seed\": 0,\n    \"cutoff_date\": \"2020-01-01\",\n    \"test_cutoff_date\": \"2022-05-01\",\n    \"max_len\": 384,\n    \"batch_size\": 1,\n    \"learning_rate\": 1e-4,\n    \"weight_decay\": 0.0,\n    \"mixed_precision\": \"bf16\",\n    \"model_config_path\": \"../working/configs/pairwise.yaml\",  # Adjust path as needed\n    \"epochs\": 10,\n    \"cos_epoch\": 5,\n    \"loss_power_scale\": 1.0,\n    \"max_cycles\": 1,\n    \"grad_clip\": 0.1,\n    \"gradient_accumulation_steps\": 1,\n    \"d_clamp\": 30,\n    \"max_len_filter\": 9999999,\n    \"structural_violation_epoch\": 50,\n    \"balance_weight\": False,\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T05:33:15.198485Z","iopub.execute_input":"2025-04-19T05:33:15.198833Z","iopub.status.idle":"2025-04-19T05:33:15.203331Z","shell.execute_reply.started":"2025-04-19T05:33:15.198801Z","shell.execute_reply":"2025-04-19T05:33:15.202322Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_data=pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/test_sequences.csv\")\ntest_data.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T05:33:17.579283Z","iopub.execute_input":"2025-04-19T05:33:17.579583Z","iopub.status.idle":"2025-04-19T05:33:17.611171Z","shell.execute_reply.started":"2025-04-19T05:33:17.579561Z","shell.execute_reply":"2025-04-19T05:33:17.610271Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\n\nclass RNADataset(Dataset):\n    def __init__(self,data):\n        self.data=data\n        self.tokens={nt:i for i,nt in enumerate('ACGU')}\n\n    def __len__(self):\n        return len(self.data)\n    \n    def __getitem__(self, idx):\n        sequence=[self.tokens[nt] for nt in (self.data.loc[idx,'sequence'])]\n        sequence=np.array(sequence)\n        sequence=torch.tensor(sequence)\n\n\n\n\n        return {'sequence':sequence}\n\ntest_dataset=RNADataset(test_data)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T05:33:20.326011Z","iopub.execute_input":"2025-04-19T05:33:20.326331Z","iopub.status.idle":"2025-04-19T05:33:20.331658Z","shell.execute_reply.started":"2025-04-19T05:33:20.326300Z","shell.execute_reply":"2025-04-19T05:33:20.330802Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sys.path.append(\"/kaggle/input/ribonanzanet2/pytorch/alpha/1\")\n\nimport torch.nn as nn\nfrom Network import *\n\nclass SinusoidalPosEmb(nn.Module):\n    def __init__(self, dim):\n        super().__init__()\n        self.dim = dim\n\n    def forward(self, x):\n        device = x.device\n        half_dim = self.dim // 2\n        emb = math.log(10000) / (half_dim - 1)\n        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)\n        emb = x[:, None] * emb[None, :]\n        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)\n        return emb\n\nclass finetuned_RibonanzaNet(RibonanzaNet):\n    def __init__(self, rnet_config, config, pretrained=False):\n        rnet_config.dropout=0.1\n        rnet_config.use_grad_checkpoint=True\n        super(finetuned_RibonanzaNet, self).__init__(rnet_config)\n        if pretrained:\n            self.load_state_dict(torch.load(config.pretrained_weight_path,map_location='cpu'))\n        # self.ct_predictor=nn.Sequential(nn.Linear(64,256),\n        #                                 nn.ReLU(),\n        #                                 nn.Linear(256,64),\n        #                                 nn.ReLU(),\n        #                                 nn.Linear(64,1)) \n        self.dropout=nn.Dropout(0.0)\n\n        decoder_dim=config.decoder_dim\n        self.structure_module=[SimpleStructureModule(d_model=decoder_dim, nhead=config.decoder_nhead, \n                 dim_feedforward=decoder_dim*4, pairwise_dimension=rnet_config.pairwise_dimension, dropout=0.0) for i in range(config.decoder_num_layers)]\n        self.structure_module=nn.ModuleList(self.structure_module)\n\n        self.xyz_embedder=nn.Linear(3,decoder_dim)\n        self.xyz_norm=nn.LayerNorm(decoder_dim)\n        self.xyz_predictor=nn.Linear(decoder_dim,3)\n        \n        self.adaptor=nn.Sequential(nn.Linear(rnet_config.ninp,decoder_dim),nn.LayerNorm(decoder_dim))\n\n        self.distogram_predictor=nn.Sequential(nn.LayerNorm(rnet_config.pairwise_dimension),\n                                                nn.Linear(rnet_config.pairwise_dimension,40))\n\n        self.time_embedder=SinusoidalPosEmb(decoder_dim)\n\n        self.time_mlp=nn.Sequential(nn.Linear(decoder_dim,decoder_dim),\n                                    nn.ReLU(),  \n                                    nn.Linear(decoder_dim,decoder_dim))\n        self.time_norm=nn.LayerNorm(decoder_dim)\n\n        self.distance2pairwise=nn.Linear(1,rnet_config.pairwise_dimension,bias=False)\n\n        self.pair_mlp=nn.Sequential(nn.Linear(rnet_config.pairwise_dimension,rnet_config.pairwise_dimension),\n                                    nn.ReLU(),\n                                    nn.Linear(rnet_config.pairwise_dimension,rnet_config.pairwise_dimension))\n\n\n        #hyperparameters for diffusion\n        self.n_times = config.n_times\n\n        #self.model = model\n        \n        # define linear variance schedule(betas)\n        beta_1, beta_T = config.beta_min, config.beta_max\n        betas = torch.linspace(start=beta_1, end=beta_T, steps=config.n_times)#.to(device) # follows DDPM paper\n        self.sqrt_betas = torch.sqrt(betas)\n                                     \n        # define alpha for forward diffusion kernel\n        self.alphas = 1 - betas\n        self.sqrt_alphas = torch.sqrt(self.alphas)\n        alpha_bars = torch.cumprod(self.alphas, dim=0)\n        self.sqrt_one_minus_alpha_bars = torch.sqrt(1-alpha_bars)\n        self.sqrt_alpha_bars = torch.sqrt(alpha_bars)\n\n        self.data_std=config.data_std\n\n\n    def custom(self, module):\n        def custom_forward(*inputs):\n            inputs = module(*inputs)\n            return inputs\n        return custom_forward\n    \n    def embed_pair_distance(self,inputs):\n        pairwise_features,xyz=inputs\n        distance_matrix=xyz[:,None,:,:]-xyz[:,:,None,:]\n        distance_matrix=(distance_matrix**2).sum(-1).clip(2,37**2).sqrt()\n        distance_matrix=distance_matrix[:,:,:,None]\n        pairwise_features=pairwise_features+self.distance2pairwise(distance_matrix)\n\n        return pairwise_features\n\n    def forward(self,src,xyz,t):\n        \n        #with torch.no_grad():\n        sequence_features, pairwise_features=self.get_embeddings(src, torch.ones_like(src).long().to(src.device))\n        \n        distogram=self.distogram_predictor(pairwise_features)\n\n        sequence_features=self.adaptor(sequence_features)\n\n        decoder_batch_size=xyz.shape[0]\n        sequence_features=sequence_features.repeat(decoder_batch_size,1,1)\n        \n\n        pairwise_features=pairwise_features.expand(decoder_batch_size,-1,-1,-1)\n\n        pairwise_features= checkpoint.checkpoint(self.custom(self.embed_pair_distance), [pairwise_features,xyz],use_reentrant=False)\n\n        time_embed=self.time_embedder(t).unsqueeze(1)\n        tgt=self.xyz_norm(sequence_features+self.xyz_embedder(xyz)+time_embed)\n\n        tgt=self.time_norm(tgt+self.time_mlp(tgt))\n\n        for layer in self.structure_module:\n            #tgt=layer([tgt, sequence_features,pairwise_features,xyz,None])\n            tgt=checkpoint.checkpoint(self.custom(layer),\n            [tgt, sequence_features,pairwise_features,xyz,None],\n            use_reentrant=False)\n            # xyz=xyz+self.xyz_predictor(sequence_features).squeeze(0)\n            # xyzs.append(xyz)\n            #print(sequence_features.shape)\n        \n        xyz=self.xyz_predictor(tgt).squeeze(0)\n        #.squeeze(0)\n\n        return xyz, distogram\n    \n\n    def denoise(self,sequence_features,pairwise_features,xyz,t):\n        decoder_batch_size=xyz.shape[0]\n        sequence_features=sequence_features.expand(decoder_batch_size,-1,-1)\n        pairwise_features=pairwise_features.expand(decoder_batch_size,-1,-1,-1)\n\n        pairwise_features=self.embed_pair_distance([pairwise_features,xyz])\n\n        sequence_features=self.adaptor(sequence_features)\n        time_embed=self.time_embedder(t).unsqueeze(1)\n        tgt=self.xyz_norm(sequence_features+self.xyz_embedder(xyz)+time_embed)\n        tgt=self.time_norm(tgt+self.time_mlp(tgt))\n        #xyz_batch_size=xyz.shape[0]\n        \n\n\n        for layer in self.structure_module:\n            tgt=layer([tgt, sequence_features,pairwise_features,xyz,None])\n            # xyz=xyz+self.xyz_predictor(sequence_features).squeeze(0)\n            # xyzs.append(xyz)\n            #print(sequence_features.shape)\n        xyz=self.xyz_predictor(tgt).squeeze(0)\n        # print(xyz.shape)\n        # exit()\n        return xyz\n\n\n    def extract(self, a, t, x_shape):\n        \"\"\"\n            from lucidrains' implementation\n                https://github.com/lucidrains/denoising-diffusion-pytorch/blob/beb2f2d8dd9b4f2bd5be4719f37082fe061ee450/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py#L376\n        \"\"\"\n        b, *_ = t.shape\n        out = a.gather(-1, t)\n        return out.reshape(b, *((1,) * (len(x_shape) - 1)))\n    \n    def scale_to_minus_one_to_one(self, x):\n        # according to the DDPMs paper, normalization seems to be crucial to train reverse process network\n        return x * 2 - 1\n    \n    def reverse_scale_to_zero_to_one(self, x):\n        return (x + 1) * 0.5\n    \n    def make_noisy(self, x_zeros, t): \n        # assume we get raw data, so center and scale by 35\n        x_zeros = x_zeros - torch.nanmean(x_zeros,1,keepdim=True)\n        x_zeros = x_zeros/self.data_std\n        #rotate randomly\n        x_zeros = random_rotation_point_cloud_torch_batch(x_zeros)\n\n\n        # perturb x_0 into x_t (i.e., take x_0 samples into forward diffusion kernels)\n        epsilon = torch.randn_like(x_zeros).to(x_zeros.device)\n        \n        sqrt_alpha_bar = self.extract(self.sqrt_alpha_bars.to(x_zeros.device), t, x_zeros.shape)\n        sqrt_one_minus_alpha_bar = self.extract(self.sqrt_one_minus_alpha_bars.to(x_zeros.device), t, x_zeros.shape)\n        \n        # Let's make noisy sample!: i.e., Forward process with fixed variance schedule\n        #      i.e., sqrt(alpha_bar_t) * x_zero + sqrt(1-alpha_bar_t) * epsilon\n        noisy_sample = x_zeros * sqrt_alpha_bar + epsilon * sqrt_one_minus_alpha_bar\n    \n        return noisy_sample.detach(), epsilon\n    \n    \n    # def forward(self, x_zeros):\n    #     x_zeros = self.scale_to_minus_one_to_one(x_zeros)\n        \n    #     B, _, _, _ = x_zeros.shape\n        \n    #     # (1) randomly choose diffusion time-step\n    #     t = torch.randint(low=0, high=self.n_times, size=(B,)).long().to(x_zeros.device)\n        \n    #     # (2) forward diffusion process: perturb x_zeros with fixed variance schedule\n    #     perturbed_images, epsilon = self.make_noisy(x_zeros, t)\n        \n    #     # (3) predict epsilon(noise) given perturbed data at diffusion-timestep t.\n    #     pred_epsilon = self.model(perturbed_images, t)\n        \n    #     return perturbed_images, epsilon, pred_epsilon\n    \n    \n    def denoise_at_t(self, x_t, sequence_features, pairwise_features, timestep, t):\n        B, _, _ = x_t.shape\n        if t > 1:\n            z = torch.randn_like(x_t).to(sequence_features.device)\n        else:\n            z = torch.zeros_like(x_t).to(sequence_features.device)\n        \n        # at inference, we use predicted noise(epsilon) to restore perturbed data sample.\n        epsilon_pred = self.denoise(sequence_features, pairwise_features, x_t, timestep)\n        \n        alpha = self.extract(self.alphas.to(x_t.device), timestep, x_t.shape)\n        sqrt_alpha = self.extract(self.sqrt_alphas.to(x_t.device), timestep, x_t.shape)\n        sqrt_one_minus_alpha_bar = self.extract(self.sqrt_one_minus_alpha_bars.to(x_t.device), timestep, x_t.shape)\n        sqrt_beta = self.extract(self.sqrt_betas.to(x_t.device), timestep, x_t.shape)\n        \n        # denoise at time t, utilizing predicted noise\n        x_t_minus_1 = 1 / sqrt_alpha * (x_t - (1-alpha)/sqrt_one_minus_alpha_bar*epsilon_pred) + sqrt_beta*z\n        \n        return x_t_minus_1#.clamp(-1., 1)\n                \n    def sample(self, src, N):\n        # start from random noise vector, NxLx3\n        x_t = torch.randn((N, src.shape[1], 3)).to(src.device)\n        \n        # autoregressively denoise from x_T to x_0\n        #     i.e., generate image from noise, x_T\n\n        #first get conditioning\n        sequence_features, pairwise_features=self.get_embeddings(src, torch.ones_like(src).long().to(src.device))\n        # sequence_features=sequence_features.expand(N,-1,-1)\n        # pairwise_features=pairwise_features.expand(N,-1,-1,-1)\n        distogram=self.distogram_predictor(pairwise_features).squeeze()\n        distogram=distogram.squeeze()[:,:,2:40]*torch.arange(2,40).float().cuda() \n        distogram=distogram.sum(-1)  \n\n        for t in range(self.n_times-1, -1, -1):\n            timestep = torch.tensor([t]).repeat_interleave(N, dim=0).long().to(src.device)\n            x_t = self.denoise_at_t(x_t, sequence_features, pairwise_features, timestep, t)\n        \n        # denormalize x_0 into 0 ~ 1 ranged values.\n        #x_0 = self.reverse_scale_to_zero_to_one(x_t)\n        x_0 = x_t * self.data_std\n        return x_0, distogram\n\n\n\n\nclass SimpleStructureModule(nn.Module):\n\n    def __init__(self, d_model, nhead, \n                 dim_feedforward, pairwise_dimension, dropout=0.1,\n                 ):\n        super(SimpleStructureModule, self).__init__()\n        #self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)\n        self.self_attn = MultiHeadAttention(d_model, nhead, d_model//nhead, d_model//nhead, dropout=dropout)\n        #self.cross_attn = MultiHeadAttention(d_model, nhead, d_model//nhead, d_model//nhead, dropout=dropout)\n\n        self.linear1 = nn.Linear(d_model, dim_feedforward)\n        self.dropout = nn.Dropout(dropout)\n        self.linear2 = nn.Linear(dim_feedforward, d_model)\n\n        self.norm1 = nn.LayerNorm(d_model)\n        self.norm2 = nn.LayerNorm(d_model)\n        self.dropout1 = nn.Dropout(dropout)\n        self.dropout2 = nn.Dropout(dropout)\n\n        self.pairwise2heads=nn.Linear(pairwise_dimension,nhead,bias=False)\n        self.pairwise_norm=nn.LayerNorm(pairwise_dimension)\n\n        #self.distance2heads=nn.Linear(1,nhead,bias=False)\n        #self.pairwise_norm=nn.LayerNorm(pairwise_dimension)\n\n        self.activation = nn.GELU()\n\n        \n    def custom(self, module):\n        def custom_forward(*inputs):\n            inputs = module(*inputs)\n            return inputs\n        return custom_forward\n\n    def forward(self, input):\n        tgt , src,  pairwise_features, pred_t, src_mask = input\n        \n        #src = src*src_mask.float().unsqueeze(-1)\n\n        pairwise_bias=self.pairwise2heads(self.pairwise_norm(pairwise_features)).permute(0,3,1,2)\n\n        \n\n\n        #print(pairwise_bias.shape,distance_bias.shape)\n\n        #pairwise_bias=pairwise_bias+distance_bias\n\n\n        res=tgt\n        tgt,attention_weights = self.self_attn(tgt, tgt, tgt, mask=pairwise_bias, src_mask=src_mask)\n        tgt = res + self.dropout1(tgt)\n        tgt = self.norm1(tgt)\n\n        # print(tgt.shape,src.shape)\n        # exit()\n\n        res=tgt\n        tgt = self.linear2(self.dropout(self.activation(self.linear1(tgt))))\n        tgt = res + self.dropout2(tgt)\n        tgt = self.norm2(tgt)\n\n\n        return tgt\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T05:38:03.687942Z","iopub.execute_input":"2025-04-19T05:38:03.688250Z","iopub.status.idle":"2025-04-19T05:38:03.715352Z","shell.execute_reply.started":"2025-04-19T05:38:03.688225Z","shell.execute_reply":"2025-04-19T05:38:03.714473Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import yaml\n\nclass Config:\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        config = yaml.safe_load(file)\n    return Config(**config)\n\n\ndiffusion_config=load_config_from_yaml(\"/kaggle/input/ribonanzanet2-ddpm-v2/diffusion_config.yaml\")\nrnet_config=load_config_from_yaml(\"/kaggle/input/ribonanzanet2/pytorch/alpha/1/pairwise.yaml\")\n\nmodel=finetuned_RibonanzaNet(rnet_config,diffusion_config).cuda()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T05:38:08.297942Z","iopub.execute_input":"2025-04-19T05:38:08.298273Z","iopub.status.idle":"2025-04-19T05:38:10.763265Z","shell.execute_reply.started":"2025-04-19T05:38:08.298246Z","shell.execute_reply":"2025-04-19T05:38:10.762368Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"state_dict=torch.load(\"/kaggle/input/ribonanzanet2-ddpm-v2/RibonanzaNet-DDPM-v2.pt\",map_location='cpu')\n\n#get rid of module. from ddp state dict\nnew_state_dict={}\n\nfor key in state_dict:\n    new_state_dict[key[7:]]=state_dict[key]\n\nmodel.load_state_dict(new_state_dict)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T05:38:19.543494Z","iopub.execute_input":"2025-04-19T05:38:19.543920Z","iopub.status.idle":"2025-04-19T05:38:24.782122Z","shell.execute_reply.started":"2025-04-19T05:38:19.543879Z","shell.execute_reply":"2025-04-19T05:38:24.781270Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tqdm import tqdm\nmodel.eval()\npreds=[]\nfor i in tqdm(range(len(test_dataset))):\n    src=test_dataset[i]['sequence'].long()\n    src=src.unsqueeze(0).cuda()\n    target_id=test_data.loc[i,'target_id']\n\n    #tmp=[]\n    predicted_dm=[]\n    #for _ in range(5):\n    with torch.no_grad():\n        xyz,distogram=model.sample(src,5)\n\n    preds.append(xyz.cpu().numpy())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T05:38:31.245057Z","iopub.execute_input":"2025-04-19T05:38:31.245377Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ID=[]\nresname=[]\nresid=[]\nx=[]\ny=[]\nz=[]\n\ndata=[]\n\nfor i in range(len(test_data)):\n    #print(test_data.loc[i])\n\n    \n    for j in range(len(test_data.loc[i,'sequence'])):\n        # ID.append(test_data.loc[i,'sequence_id']+f\"_{j+1}\")\n        # resname.append(test_data.loc[i,'sequence'][j])\n        # resid.append(j+1) # 1 indexed\n        row=[test_data.loc[i,'target_id']+f\"_{j+1}\",\n             test_data.loc[i,'sequence'][j],\n             j+1]\n\n        for k in range(5):\n            for kk in range(3):\n                row.append(preds[i][k][j][kk])\n        data.append(row)\n\ncolumns=['ID','resname','resid']\nfor i in range(1,6):\n    columns+=[f\"x_{i}\"]\n    columns+=[f\"y_{i}\"]\n    columns+=[f\"z_{i}\"]\n\n\nsubmission=pd.DataFrame(data,columns=columns)\n\n\nsubmission\nsubmission.to_csv('submission.csv',index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-04T20:16:40.438902Z","iopub.execute_input":"2025-04-04T20:16:40.439268Z","iopub.status.idle":"2025-04-04T20:16:40.582276Z","shell.execute_reply.started":"2025-04-04T20:16:40.439239Z","shell.execute_reply":"2025-04-04T20:16:40.581409Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#score val\nimport pandas as pd\nimport pandas.api.types\nimport os\nimport re\n\n# Function to parse TMscore output\ndef parse_tmscore_output(output):\n    result = {}\n\n    # Extract TM-score based on length of reference structure (second)\n    tm_score_match = re.findall(r\"TM-score=\\s+([\\d.]+)\", output)[1]\n    result['TM-score'] = float(tm_score_match) if tm_score_match else None\n\n    return result\n\ndef write_pdb_line(atom_name, atom_serial, residue_name, chain_id, residue_num, x_coord, y_coord, z_coord, occupancy=1.0, b_factor=0.0, atom_type='P'):\n    \"\"\"\n    Writes a single line of PDB format based on provided atom information. \n    \n    Args:\n        atom_name (str): Name of the atom (e.g., \"N\", \"CA\").\n        atom_serial (int): Atom serial number.\n        residue_name (str): Residue name (e.g., \"ALA\"). \n        chain_id (str): Chain identifier. \n        residue_num (int): Residue number. \n        x_coord (float): X coordinate.\n        y_coord (float): Y coordinate.\n        z_coord (float): Z coordinate.\n        occupancy (float, optional): Occupancy value (default: 1.0). \n        b_factor (float, optional): B-factor value (default: 0.0). \n    \n    Returns:\n        str: A single line of PDB string.\n    \"\"\"\n    line = f\"ATOM  {atom_serial:>5d}  {atom_name:<5s} {residue_name:<3s} {residue_num:>3d}    {x_coord:>8.3f}{y_coord:>8.3f}{z_coord:>8.3f}{occupancy:>6.2f}{b_factor:>6.2f}           {atom_type}\\n\"\n    return line\n\ndef write2pdb(df, xyz_id, pdb_path):\n    resolved_cnt=0\n    with open(pdb_path, \"w\") as pdb_file:\n        for _, row in df.iterrows():\n            x_coord=row[f\"x_{xyz_id}\"]\n            y_coord=row[f\"y_{xyz_id}\"]\n            z_coord=row[f\"z_{xyz_id}\"]\n\n            if x_coord>-1e17 and y_coord>-1e17 and z_coord>-1e17:\n            #if True:\n                resolved_cnt+=1\n                pdb_line = write_pdb_line(\n                    atom_name=\"C1'\", \n                    atom_serial=int(row[\"resid\"]), \n                    residue_name=row['resname'], \n                    chain_id='0', \n                    residue_num=int(row[\"resid\"]), \n                    x_coord=x_coord, \n                    y_coord=y_coord, \n                    z_coord=z_coord,\n                    atom_type=\"C\"\n                )\n                pdb_file.write(pdb_line)\n    return resolved_cnt\n\n\ndef score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str) -> float:\n    '''\n    Computes the TM-score between predicted and native RNA structures using USalign.\n\n    This function evaluates the structural similarity of RNA predictions to native structures\n    by computing the TM-score. It uses USalign, a structural alignment tool, to compare\n    the predicted structures with the native structures.\n\n    Workflow:\n    1. Copies the USalign binary to the working directory and grants execution permissions.\n    2. Extracts the `pdb_id` from the `ID` column of both the solution and submission DataFrames.\n    3. Iterates over each unique `pdb_id`, grouping the native and predicted structures.\n    4. Writes PDB files for native and predicted structures.\n    5. Runs USalign on each predicted-native pair and extracts the TM-score.\n    6. Computes the highest TM-score per target and returns aggregated results.\n\n    Args:\n        solution (pd.DataFrame): A DataFrame containing the native RNA structures.\n        submission (pd.DataFrame): A DataFrame containing the predicted RNA structures.\n        row_id_column_name (str): The name of the column containing unique row identifiers.\n\n    Returns:\n        tuple:\n            - results (list): The highest TM-score for each `pdb_id`.\n            - results_per_sub (list): TM-scores for each predicted-native pair.\n            - outputs (list): Raw output logs from USalign for debugging.\n    '''\n\n    os.system(\"cp /kaggle/input/usalign/USalign /kaggle/working/\")\n    os.system(\"sudo chmod u+x /kaggle/working//USalign\")\n\n\n    # Extract pdb_id from ID (pdb_resid)\n    solution[\"pdb_id\"] = solution[\"ID\"].apply(lambda x: x.split(\"_\")[0])\n    submission[\"pdb_id\"] = submission[\"ID\"].apply(lambda x: x.split(\"_\")[0])\n\n    #fix pdb_ids comment out later\n    # solution.loc[solution['pdb_id']==\"R1138v1\",'pdb_id']='R1138'\n    # solution.loc[solution['pdb_id']==\"R1117\",'pdb_id']='R1117v2'\n    \n    results=[]\n    outputs=[]\n    results_per_sub=[]\n    # Iterate through each pdb_id and generate PDB files for both clean and corrupted data\n    for pdb_id, group_native in solution.groupby(\"pdb_id\"):\n        group_predicted = submission[submission[\"pdb_id\"] == pdb_id]\n        #print(group_native,group_predicted)\n        # Define output file paths\n        # clean_pdb_path = os.path.join(output_folder, f\"{pdb_id}_C3_clean.pdb\")\n        # corrupted_pdb_path = os.path.join(output_folder, f\"{pdb_id}_C3_corrupted.pdb\")\n        native_pdb=f'native.pdb'\n        predicted_pdb=f'predicted.pdb'\n\n        all_scores=[]\n        for pred_cnt in range(1,6):\n            tmp=[]\n            for native_cnt in range(1,41):\n                # Write solution PDB\n                resolved_cnt=write2pdb(group_native, native_cnt, native_pdb)\n                \n                # Write predicted PDB\n                _=write2pdb(group_predicted, pred_cnt, predicted_pdb)\n\n                if resolved_cnt>0:\n                    command = f\"/kaggle/working/USalign {predicted_pdb} {native_pdb} -atom \\\" C1'\\\"\"\n                    output = os.popen(command).read()\n                    outputs.append(output)\n                    parsed_data = parse_tmscore_output(output)\n                    tmp.append(parsed_data['TM-score'])\n                    \n            all_scores.append(max(tmp))\n        # print(output)\n        # stop\n        print(pdb_id)\n        print(all_scores)\n        results_per_sub.append(all_scores)\n        results.append(max(all_scores))\n    \n    print(results)\n    #return sum(results)/len(results), outputs\n    return results, results_per_sub, outputs\n    #return outputs\n\nif 'R1107' in set(test_data['target_id']):\n    solution=pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/validation_labels.csv\")\n    \n    scores,results_per_sub,outputs=score(solution,submission,'ID')\n    print(np.mean(scores))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-04T20:19:02.897819Z","iopub.execute_input":"2025-04-04T20:19:02.898106Z","iopub.status.idle":"2025-04-04T20:20:32.568722Z","shell.execute_reply.started":"2025-04-04T20:19:02.898085Z","shell.execute_reply":"2025-04-04T20:20:32.567861Z"}},"outputs":[],"execution_count":null}]}