prediction.py

import torch
from torch_geometric.loader import DataLoader
from torch_geometric.data import Dataset, Data
from Bio.PDB import PDBParser
import pandas as pd
from sklearn.preprocessing import OneHotEncoder
from challenge_1 import load_features, TransformerBinary
import os
import pickle
from torch_geometric.nn import radius_graph
import numpy as np



class TestResidueGraphDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
        self.pdb_dir = root
        self.parser = PDBParser(QUIET=True)
        self.dssp_features = load_features(r"merged_test")  # DSSP-Features
        self.encoder = OneHotEncoder(sparse_output=False)
        self._fit_encoder()

    def _fit_encoder(self):
        # Fit the encoder on all possible residues (20 standard amino acids)
        resnames = [['ALA'], ['ARG'], ['ASN'], ['ASP'], ['CYS'], ['GLN'], ['GLU'], ['GLY'],
                    ['HIS'], ['ILE'], ['LEU'], ['LYS'], ['MET'], ['PHE'], ['PRO'], ['SER'],
                    ['THR'], ['TRP'], ['TYR'], ['VAL']]
        self.encoder.fit(resnames)

    def len(self):
        return len(os.listdir(self.pdb_dir))

    def __getitem__(self, idx):
        pdb_path = os.path.join(self.pdb_dir, f"{idx}_protein.pdb")
        structure = self.parser.get_structure(str(idx), pdb_path)
        

        nodes = []
        positions = {}
        resnames = []
        seq_indices = []
        coords_list = []
        dssp_feature_list = []
        
        for model in structure:
            for chain in model:
                for residue in chain:
                        
                    if residue.id[0] != " ":  # skip heteroatoms/water
                        continue
                    chain_id = chain.id
                    res_id = residue.id[1]
                    icode = residue.id[2]
                    #node_id = f"{chain.id}_{res_id}_{icode}"
                    node_id = f"{chain_id}_{res_id}"
                            
                     # DSSP-Features
                    dssp_residue_key = f"{chain_id}_{res_id}"
                    dssp_features_for_residue = self.dssp_features.get(dssp_residue_key, [0.0, 0.0, 1.0, 0.0, 0.0, 0.0])
                    dssp_feature_list.append(dssp_features_for_residue)  
                  
                    coords = [atom.coord for atom in residue.get_atoms()]
                    centroid = np.mean(coords, axis=0)
                    positions[node_id] = centroid
                    resnames.append([residue.get_resname()])
                    nodes.append(node_id)
                    coords_list.append(centroid.tolist())  # Save 3D position

                    # Use residue index as normalized position in chain
                    seq_indices.append([res_id / 1000.0])
                
      
        edge_index = radius_graph(torch.tensor(coords_list), r=8)
        
        with open(r"esm_test\{0}_protein.pkl".format(idx), "rb") as f:
            esm_emb = pickle.load(f)
            
        esm_features_np = np.array(esm_emb)
       
        res_coords = torch.tensor(coords_list, dtype=torch.float)
        res_coords = res_coords - res_coords.mean(dim=0, keepdim=True)
        
                
        features = np.concatenate([res_coords, seq_indices, dssp_feature_list, esm_features_np], axis=1)
        x = torch.tensor(features, dtype=torch.float)
        data = Data(x=x, edge_index=edge_index)
        data.nodes = nodes
        data.idx = idx
        return data


@torch.no_grad()
def evaluate(model, dataloader, output_csv="predictions.csv"):
    model.eval()
    predictions = []

    for batch in dataloader:
        batch = batch.to(device)

        # Forward pass
        out = model(batch.x, batch.edge_index)
        probs = torch.sigmoid(out)  # [num_nodes]
        preds = (probs > 0.5).cpu().numpy()
        batch_indices = batch.batch.cpu().numpy()

        data_list = batch.to_data_list()

        for i, data in enumerate(data_list):
            node_indices = np.where(batch_indices == i)[0]
            pred_for_graph = preds[node_indices]
            node_ids = data.nodes

            # Get predicted residue identifiers
            pred_residues = [node_ids[j] for j, p in enumerate(pred_for_graph) if p == 1]

            dataset_idx = int(data.idx.cpu())
            prediction_line = f"{dataset_idx}, " + " ".join(pred_residues)
            predictions.append(prediction_line)

    # Save to CSV
    with open(output_csv, "w") as f:
        f.write("id,prediction\n")
        for line in sorted(predictions, key=lambda x: int(x.split(",")[0])):
            f.write(line + "\n")

    print(f"\nPredictions saved to {output_csv}")



if __name__ == "__main__":
    dataset = TestResidueGraphDataset(root=r"./test")
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print(device)

    in_channels = dataset[0].x.size(1)
    model = TransformerBinary(in_channels, hidden_channels=256).to(device)
    model.load_state_dict(torch.load("modelv1.pth", weights_only=True))
    model.eval()
    test_loader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=2)
    evaluate(model, test_loader)