graphabl/dataset.py

import os
import networkx as nx
import torch
from Bio.PDB import PDBParser, is_aa
from Bio.PDB.Polypeptide import three_to_one
from torch.utils.data import Dataset
import pandas as pd
import numpy as np
from tqdm import tqdm
from Bio.Align import PairwiseAligner
from torch_geometric.data import Data
from torch_geometric.utils import from_networkx
from Bio.SeqUtils.ProtParam import ProteinAnalysis


aligner = PairwiseAligner()
# Define the scoring system
aligner.mode = 'global'  # Global alignment
aligner.match_score = 1
aligner.mismatch_score = -1
aligner.open_gap_score = -0.5
aligner.extend_gap_score = -0.1


def calculate_similarity(seq_gen, seq_tem):

    # Perform the alignment
    score = aligner.score(seq_gen, seq_tem)
    return score / len(seq_gen)


def parse_ll37_testset(path='metadata/LL37_v0.csv'):
    files = open(path, 'r').readlines()
    seqs = []
    for file in files:
        seqs.append(str(file.strip().split(',')[-1].strip().upper()))
    return seqs


def gmean(labels):
    # Assuming 'labels' is a NumPy array of shape (n, m)
    # Calculate the geometric mean across axis 1

    # Take the product of elements along axis=1
    product = np.prod(labels, axis=1)

    # nth root of product, where n is the number of elements along axis=1
    geometric_means = product ** (1.0 / labels.shape[1])
    return geometric_means


def clamp(n, smallest, largest):
    return sorted([smallest, n, largest])[1]


ATOMS = {'H': 1, 'C': 12, 'N': 14, 'O': 16, 'S': 30}
ATOMS_R = {'H': 1, 'C': 1.5, 'N': 1.5, 'O': 1.5, 'S': 2}
AMINO_ACID_WATER = {'A': 255, 'V': 255, 'P': 255, 'F': 255, 'W': 255, 'I': 255, 'L': 255, 'G': 155, 'M': 155,
                    'Y': 55, 'S': 55, 'T': 55, 'C': 55, 'N': 55, 'Q': 55, 'D': 55, 'E': 55, 'K': 55, 'R': 55, 'H': 55}
AMINO_ACID_CHARGE = {'D': 55, 'E': 55, 'A': 155, 'V': 155, 'P': 155, 'F': 155, 'W': 155, 'I': 155, 'L': 155, 'G': 155,
                     'M': 155, 'Y': 155, 'S': 155, 'T': 155, 'C': 155, 'N': 155, 'Q': 155, 'K': 255, 'R': 255, 'H': 255}
# AMAs = {'G': 20, 'A': 1, 'V': 2, 'L': 3, 'I': 4, 'P': 5, 'F': 6, 'Y': 7, 'W': 8, 'S': 9, 'T': 10, 'C': 11,
#         'M': 12, 'N': 13, 'Q': 14, 'D': 15, 'E': 16, 'K': 17, 'R': 18, 'H': 19}
AMAs = {'G': 20, 'A': 1, 'V': 2, 'L': 3, 'I': 4, 'P': 5, 'F': 6, 'Y': 7, 'W': 8, 'S': 9, 'T': 10, 'C': 11,
        'M': 12, 'N': 13, 'Q': 14, 'D': 15, 'E': 16, 'K': 17, 'R': 18, 'H': 19, 'X': 21}


def pdb_parser(structure):
    """

    """
    voxel = np.zeros((4, 64, 64, 64), dtype=np.int8)
    id = ''
    seq_str = ''
    for i in structure[0]:
        id = i.id
    chain = structure[0][id]
    for res in chain:
        if is_aa(res.get_resname(), standard=True):
            resname = res.get_resname()
            amino = three_to_one(resname)
            seq_str += str(amino)
            ATOM_WATER = AMINO_ACID_WATER[amino]
            ATOM_CHARGE = AMINO_ACID_CHARGE[amino]
            ATOM_CATEGORY = AMAs[amino] * 20

            for i in res:
                if i.id not in ATOMS.keys():
                    continue
                x, y, z = i.get_coord()
                if abs(x) > 32:
                    x = clamp(x, -31, 31)
                if abs(y) > 32:
                    y = clamp(x, -31, 31)
                if abs(z) > 32:
                    z = clamp(x, -31, 31)
                x_i, y_i, z_i = int(x) + 32, int(y) + 32, int(z) + 32
                ATOM_WEIGHT = ATOMS[i.id]
                ATOM_R = ATOMS_R[i.id]

                if ATOM_R <= 1.5:
                    voxel[0, x_i - 1:x_i + 1, y_i - 1:y_i + 1, z_i - 1:z_i + 1] = ATOM_WEIGHT
                    voxel[1, x_i - 1:x_i + 1, y_i - 1:y_i + 1, z_i - 1:z_i + 1] = ATOM_WATER
                    voxel[2, x_i - 1:x_i + 1, y_i - 1:y_i + 1, z_i - 1:z_i + 1] = ATOM_CHARGE
                    voxel[3, x_i - 1:x_i + 1, y_i - 1:y_i + 1, z_i - 1:z_i + 1] = ATOM_CATEGORY
                else:
                    voxel[0, x_i - ATOM_R: x_i + ATOM_R, x_i - ATOM_R: x_i + ATOM_R,
                    x_i - ATOM_R: x_i + ATOM_R] = ATOM_WEIGHT
                    voxel[1, x_i - ATOM_R: x_i + ATOM_R, x_i - ATOM_R: x_i + ATOM_R,
                    x_i - ATOM_R: x_i + ATOM_R] = ATOM_WATER
                    voxel[2, x_i - ATOM_R: x_i + ATOM_R, x_i - ATOM_R: x_i + ATOM_R,
                    x_i - ATOM_R: x_i + ATOM_R] = ATOM_CHARGE
                    voxel[3, x_i - ATOM_R: x_i + ATOM_R, x_i - ATOM_R: x_i + ATOM_R,
                    x_i - ATOM_R: x_i + ATOM_R] = ATOM_CATEGORY
    return voxel


def read_scoring_functions(pdb):
    scoring = False
    profile = []
    for line in open(pdb):
        if line.startswith("#END_POSE_ENERGIES_TABLE"):
            scoring = False
        if scoring:
            data = [float(v) for v in line.split()[1:]]
            profile.append(data)
        if line.startswith("pose"):
            scoring = True
    return profile


def load_aa_features(feature_path):
    aa_features = {}
    for line in open(feature_path):
        line = line.strip().split()
        aa, features = line[0], [float(feature) for feature in line[1:]]
        aa_features[aa] = features
    return aa_features


class PairData(Data):
    def __init__(self, edge_index_s, x_s):
        super(PairData, self).__init__()
        self.edge_index_s = edge_index_s
        self.x_s = x_s

    def __inc__(self, key, value, *args):
        if key == 'edge_index_s':
            return self.x_s.size(0)
        if key == 'wide_nodes':
            return self.x_s.num_nodes
        else:
            return super().__inc__(key, value, *args)
        

def calculate_property(seq):
    # print(seq)
    analysed_seq = ProteinAnalysis(seq)
    aa_counts = analysed_seq.count_amino_acids()
    aliphatic_index = ((aa_counts['A'] + 2.9 * aa_counts['V'] + 3.9 * (aa_counts['I'] + aa_counts['L'])) / len(seq))

    charged_amino_acids = ['D', 'E', 'R', 'K', 'H']
    total_charged = sum(aa_counts[aa] for aa in charged_amino_acids)
    charge_density = total_charged / len(seq)
    alpha_helix, beta_helix, turn_helix = analysed_seq.secondary_structure_fraction()

    return list(
        [round(analysed_seq.gravy(), 3) * 10, round(aliphatic_index, 3) * 10, round(analysed_seq.aromaticity(), 3) * 10,
         round(analysed_seq.instability_index(), 3), round(alpha_helix * 10, 3), round(beta_helix * 10, 3),
         round(turn_helix * 10, 3), round(analysed_seq.charge_at_pH(7), 3), round(analysed_seq.isoelectric_point(), 3),
         round(charge_density, 3) * 10])


def construct_graph(structure, seq, gt, pdb_path, threshold, max_length):
    # Initialize an empty graph
    G = nx.Graph()

    # Check if chain 'A' is present in the structure, return None if not
    if 'A' not in structure[0]:
        return None
    chain = structure[0]['A']

    # Add nodes to the graph for each amino acid in the chain with their 3D positions
    for res in chain:
        if is_aa(res.get_resname(), standard=True):
            resname = res.get_resname()
            alpha_c_atom = res["CA"]
            x, y, z = alpha_c_atom.get_coord()  # Get 3D coordinates
            G.add_node(res.id[1], name=resname, position=(x, y, z))

    # Add edges between nodes if the distance between alpha carbons is <= 5Å
    for i, m in enumerate(G.nodes):
        for n in list(G.nodes)[i + 1:]:
            distance = chain[m]["CA"] - chain[n]["CA"]
            if distance <= 5:
                G.add_edge(m, n, weight=5 / distance)

    # Load scoring functions and features
    scoring = read_scoring_functions(pdb_path)
    # print(pdb_path)
    aa_features = load_aa_features('metadata/features.txt')

    # Assign features to nodes in the graph, including the 3D position
    for node_id, node_data in G.nodes(data=True):
        res = chain[node_id]
        aa_feature = aa_features[three_to_one(res.get_resname())]
        energy_feature = scoring[node_id - 1]
        position_feature = list(node_data['position'])  # Convert position tuple to list
        node_data['x'] = aa_feature + energy_feature + position_feature
        # node_data['x'][20:40]=[0 for i in range(20)]
        # node_data['x'][40:43]=[0 for i in range(3)]

    # Convert graph to data format used in ML models
    G = nx.convert_node_labels_to_integers(G)
    data_wt = from_networkx(G)
    data_graph = PairData(data_wt.edge_index, data_wt.x)
    data_graph.gt = (gt < threshold).astype(int)
    global_properties = calculate_property(seq)  

    # Encode the sequence and pad it to a fixed length
    seq_emb = [AMAs[char] for char in seq] + [0] * (max_length - len(seq))

    # print(seq_emb)
    data_graph.seq = seq_emb
    data_graph.global_f = global_properties
    return data_graph


def remove_features(x, features_to_remove):
    # 创建一个掩码来保留特征
    num_features = x.size(1)
    mask = torch.ones(num_features, dtype=torch.bool)
    mask[features_to_remove] = False
    
    # 更新节点特征矩阵
    x = x[:, mask]
    
    return x


class MDataset(Dataset):
    '''
    x: Features.
    y: Targets, if none, do prediction.
    '''

    def __init__(self, threshold=32, mode='train', max_length=50, pdb_src='af', data_ver='0922', exclude_feature=None):
        self.num_classes = 6
        self.num_features = 43
        exclude_flag = True
        if exclude_feature == 0:
            self.num_features = 23
            exclude_index = torch.arange(0, 20)
        elif exclude_feature == 1:
            self.num_features = 23
            exclude_index = torch.arange(20, 40)
        elif exclude_feature == 2:
            self.num_features = 40
            exclude_index = torch.arange(40, 43)
        else:
            exclude_flag = False
        exclude_list = pd.read_csv('../metadata/data_simi.csv', encoding="unicode_escape")['Seq'].str.upper().str.strip().tolist()
        exclude_filter = False

        p = PDBParser(QUIET=True)

        if mode == 'train':
            all_data = pd.read_csv(f'../metadata/data_{data_ver}_i.csv', encoding="unicode_escape").values
            exclude_filter = True
        elif mode == 'qlx':
            all_data = pd.read_csv('../metadata/data_qlx.csv', encoding="unicode_escape").values
        elif mode == 'saap':
            all_data = pd.read_csv('../metadata/data_saap.csv', encoding="unicode_escape").values
        else:
            raise NotImplementedError
        idx_list, seq_list, labels = all_data[:, 0], all_data[:, 1], np.concatenate((all_data[:, 4:9], all_data[:, 10:11]), axis=1)

        filter_idx_list = []
        seq_new_list = []
        label_list = []
        for idx in range(len(idx_list)):
            seq = seq_list[idx].upper().strip()
            if 'X' in seq or 'B' in seq or 'J' in seq or 'Z' in seq or 'U' in seq or 'O' in seq or len(seq) > max_length or len(seq) < 6:
                continue
            if exclude_filter:
                if seq in exclude_list:
                    continue

            filter_idx_list.append(idx)
            seq_new_list.append(seq)
            label_list.append(labels[idx])

        if pdb_src == 'af':
            pdb_root = '../pdb/pdb_af/' 
        elif pdb_src == 'hf':
            pdb_root = '../pdb/pdb_dbassp/'
        else:
            raise NotImplementedError
        self.data_list = []
        for i in tqdm(range(len(filter_idx_list))):
                idx = filter_idx_list[i]
                seq = seq_new_list[i]
                label = label_list[i]

                # if os.path.exists("./pdb/pdb_dbassp/" + seq + ".pdb"):
                #     pdb_path = "./pdb/pdb_dbassp/" + seq + ".pdb"
                # elif os.path.exists("./pdb/pdb_gen/" + seq + ".pdb"):
                #     pdb_path = "./pdb/pdb_gen/" + seq + ".pdb"
                if os.path.exists(pdb_root + seq + ".pdb"):
                    pdb_path = pdb_root + seq + ".pdb"
                else:
                    continue
                    # raise FileNotFoundError

                structure = p.get_structure(idx, pdb_path)
                voxel1 = construct_graph(structure, seq, label, pdb_path, threshold, max_length)
                if exclude_flag:
                    voxel1.x_s = remove_features(voxel1.x_s, exclude_index)
                self.data_list.append(voxel1)

    def __getitem__(self, idx):
        voxel1 = self.data_list[idx]

        return voxel1

    def __len__(self):
        return len(self.data_list)
    

# class SDataset(Dataset):
#     '''
#     x: Features.
#     y: Targets, if none, do prediction.
#     '''

#     def __init__(self, threshold=32, mode='train', task='0', max_length=50, pdb_src='af', data_ver='0922'):
#         self.num_classes = 1
#         self.task = int(task)
#         exclude_list = pd.read_csv('metadata/data_simi.csv', encoding="unicode_escape")['Seq'].str.upper().str.strip().tolist()
#         exclude_filter = False

#         p = PDBParser(QUIET=True)

#         if mode == 'train':
#             all_data = pd.read_csv(f'metadata/data_{data_ver}_i.csv', encoding="unicode_escape").values
#             exclude_filter = True
#         elif mode == 'qlx':
#             all_data = pd.read_csv('metadata/data_qlx.csv', encoding="unicode_escape").values
#         elif mode == 'saap':
#             all_data = pd.read_csv('metadata/data_saap.csv', encoding="unicode_escape").values
#         else:
#             raise NotImplementedError
#         idx_list, seq_list, labels = all_data[:, 0], all_data[:, 1], np.concatenate((all_data[:, 4:9], all_data[:, 10:11]), axis=1)
#         labels = labels[:, self.task:self.task+1]

#         filter_idx_list = []
#         seq_new_list = []
#         label_list = []
#         for idx in range(len(idx_list)):
#             seq = seq_list[idx].upper().strip()
#             if 'X' in seq or 'B' in seq or 'J' in seq or 'Z' in seq or 'U' in seq or 'O' in seq or len(seq) > max_length or len(seq) < 6:
#                 continue
#             if exclude_filter:
#                 if seq in exclude_list:
#                     continue

#             filter_idx_list.append(idx)
#             seq_new_list.append(seq)
#             label_list.append(labels[idx])

#         if pdb_src == 'af':
#             pdb_root = './pdb/pdb_af/' 
#         elif pdb_src == 'hf':
#             pdb_root = './pdb/pdb_dbassp/'
#         else:
#             raise NotImplementedError
#         self.data_list = []
#         for i in tqdm(range(len(filter_idx_list))):
#                 idx = filter_idx_list[i]
#                 seq = seq_new_list[i]
#                 label = label_list[i]

#                 # if os.path.exists("./pdb/pdb_dbassp/" + seq + ".pdb"):
#                 #     pdb_path = "./pdb/pdb_dbassp/" + seq + ".pdb"
#                 # elif os.path.exists("./pdb/pdb_gen/" + seq + ".pdb"):
#                 #     pdb_path = "./pdb/pdb_gen/" + seq + ".pdb"
#                 if os.path.exists(pdb_root + seq + ".pdb"):
#                     pdb_path = pdb_root + seq + ".pdb"
#                 else:
#                     continue
#                     # raise FileNotFoundError

#                 structure = p.get_structure(idx, pdb_path)
#                 voxel1 = construct_graph(structure, seq, label, pdb_path, threshold, max_length)
#                 self.data_list.append(voxel1)

#     def __getitem__(self, idx):
#         voxel1 = self.data_list[idx]

#         return voxel1

#     def __len__(self):
#         return len(self.data_list)