| | import torch |
| | import os, glob |
| | import pandas as pd |
| | import numpy as np |
| | from src.data.esm.sdk.api import ESMProtein |
| | from src.utils.utils import pmap_multi |
| |
|
| |
|
| | def read_data(aa_seq, pdb_path, target, unique_id, position=None, msa_aas=None): |
| | try: |
| | if unique_id is None: |
| | unique_id = str(hash(aa_seq)) |
| | |
| | if position: |
| | position = torch.tensor([int(ele) for ele in position.split(",")]) |
| | |
| | if pdb_path: |
| | structure = ESMProtein.from_pdb(pdb_path) |
| | coordinates = structure.coordinates |
| | else: |
| | coordinates = None |
| |
|
| | return { |
| | 'name': target, |
| | 'seq': aa_seq, |
| | 'X': coordinates, |
| | 'label': torch.tensor([0]), |
| | 'unique_id': unique_id, |
| | 'pdb_path': pdb_path, |
| | 'position': position, |
| | 'msa_aas': msa_aas |
| | } |
| | except: |
| | return None |
| | |
| |
|
| | class MSADataset(torch.utils.data.Dataset): |
| | def __init__( |
| | self, |
| | msa_csv_path: str, |
| | type: str = "center" |
| | ): |
| | self.msa_csv_path = msa_csv_path |
| | msa_df = pd.read_csv(self.msa_csv_path) |
| | |
| | if type == "center": |
| | msa_df = msa_df[msa_df["type"] == type] |
| | else: |
| | center_df = msa_df[msa_df["type"] == "center"] |
| | id_seq_dict = dict(zip(center_df['unique_id'], center_df['aa_seq_ori'])) |
| | msa_df = msa_df[msa_df["type"] == type] |
| |
|
| | def compute_position_diff(x): |
| | seq_with_gap = x["aa_seq_ori"] |
| | center_seq = id_seq_dict[x["target"]] |
| | |
| | pos_diff, msa_pos_aa = [], [] |
| | center_idx = 0 |
| |
|
| | for i, aa in enumerate(seq_with_gap): |
| | if aa == "-": |
| | continue |
| |
|
| | if center_idx >= len(center_seq): |
| | break |
| |
|
| | if aa.upper() != center_seq[center_idx].upper(): |
| | pos_diff.append(str(i)) |
| | msa_pos_aa.append(aa.upper()) |
| | center_idx += 1 |
| |
|
| | return ",".join(pos_diff) |
| | |
| | def compute_aa_diff(x): |
| | seq_with_gap = x["aa_seq_ori"] |
| | center_seq = id_seq_dict[x["target"]] |
| | |
| | pos_diff, msa_pos_aa = [], [] |
| | center_idx = 0 |
| |
|
| | for i, aa in enumerate(seq_with_gap): |
| | if aa == "-": |
| | continue |
| |
|
| | if center_idx >= len(center_seq): |
| | break |
| |
|
| | if aa.upper() != center_seq[center_idx].upper(): |
| | pos_diff.append(str(i)) |
| | msa_pos_aa.append(aa.upper()) |
| | center_idx += 1 |
| |
|
| | return "".join(msa_pos_aa) |
| | |
| | def compute_aa_diff(x): |
| | seq_with_gap = x["aa_seq_ori"] |
| | center_seq = id_seq_dict[x["target"]] |
| | |
| | pos_diff, msa_pos_aa = [], [] |
| | center_idx = 0 |
| |
|
| | for i, aa in enumerate(seq_with_gap): |
| | if aa == "-": |
| | continue |
| |
|
| | if center_idx >= len(center_seq): |
| | break |
| |
|
| | if aa.upper() != center_seq[center_idx].upper(): |
| | pos_diff.append(str(i)) |
| | msa_pos_aa.append(aa.upper()) |
| | center_idx += 1 |
| |
|
| | return "".join(msa_pos_aa) |
| |
|
| | msa_df["position"] = msa_df.apply(compute_position_diff, axis=1) |
| | msa_df["msa_aas"] = msa_df.apply(compute_aa_diff, axis=1) |
| | path_list = [] |
| | for i in range(len(msa_df)): |
| | path_list.append( |
| | ( |
| | msa_df.iloc[i].get('aa_seq'), |
| | msa_df.iloc[i].get('pdb_path'), |
| | msa_df.iloc[i].get('target'), |
| | msa_df.iloc[i].get('unique_id'), |
| | msa_df.iloc[i].get('position'), |
| | msa_df.iloc[i].get('msa_aas'), |
| | ) |
| | ) |
| | |
| | self.data = pmap_multi(read_data, path_list, n_jobs=8) |
| | self.data = [d for d in self.data if d is not None] |
| |
|
| | def __len__(self): |
| | return len(self.data) |
| | |
| | def __getitem__( |
| | self, |
| | idx |
| | ): |
| | return self.data[idx] |
| | |
| |
|
| | if __name__ == "__main__": |
| | msa_data_center = MSADataset( |
| | msa_csv_path = "/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/zeroshot/msa/msa_samples_zeroshot_w_pdb.csv", |
| | type="center" |
| | ) |
| | print(f"length of msa dataset: {len(msa_data_center)}...") |
| |
|
| | msa_data_msa = MSADataset( |
| | msa_csv_path = "/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/zeroshot/msa/msa_samples_zeroshot_w_pdb.csv", |
| | type="msa" |
| | ) |
| | print(f"length of msa dataset: {len(msa_data_msa)}...") |
| |
|
