Spaces:
Running on Zero
Running on Zero
File size: 7,784 Bytes
093748c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | from setup.configProcessor import get_id_list_path, get_embeddings_path, get_sequence_path, get_label_path, use_cache
from data_processing.structure_processing import get_connectivity, get_dssp_features
from misc.enums import LabelType
import numpy as np
import h5py
import fastapy
import torch
import copy
from torch_geometric.data import Dataset, Data
class BindingResidueDataset(Dataset):
def __init__(self, structure_cutoff):
super().__init__('.', None, None, None)
self.structure_cutoff = structure_cutoff
self.protein_ids = self.get_id_list()
self.embeddings = self.get_embeddings()
self.sequences = self.get_sequences()
self.connectivity_matrices = get_connectivity(self.structure_cutoff, self.protein_ids, self.sequences)
self.dssp_features = get_dssp_features(self.protein_ids)
def len(self):
return len(self.protein_ids)
def get(self, item):
protein_id = self.protein_ids[item]
protein_graph_edges = np.array(self.get_connectivity_matrix(item)["backbone"].nonzero())
protein_graph_cutoff_edges = np.array(self.get_connectivity_matrix(item)["cutoff"].nonzero())
# retrieve distances of edges below structure cutoff
edge_attributes = \
self.get_connectivity_matrix(item)["distance"][protein_graph_cutoff_edges[0], protein_graph_cutoff_edges[1]]
# divide by distance cutoff to avoid div by zero error for linalg.norm
edge_attributes = (edge_attributes / self.structure_cutoff)
# invert distance to achieve weights: distance zero has highest weight
edge_features = (1 - edge_attributes)
protein_graph = Data(
x=torch.Tensor(self.embeddings[protein_id]),
edge_index=torch.LongTensor(protein_graph_edges),
edge_index_cutoff=torch.LongTensor(protein_graph_cutoff_edges),
edge_features=torch.Tensor(edge_features),
dssp_features=torch.tensor(self.dssp_features[protein_id]).float()
)
return protein_graph, protein_id
def get_embedding(self, item):
protein_id = self.protein_ids[item]
return self.embeddings[protein_id]
def get_sequence(self, item):
protein_id = self.protein_ids[item]
return self.sequences[protein_id]
def get_connectivity_matrix(self, item):
protein_id = self.protein_ids[item]
return self.connectivity_matrices[protein_id]
def train_val_split(self, train_ids, val_ids):
train_subset = self.create_subset(train_ids)
val_subset = self.create_subset(val_ids)
return train_subset, val_subset
def create_subset(self, subset_ids):
subset = copy.deepcopy(self)
self.collect_subset_data(subset_ids, subset)
return subset
def collect_subset_data(self, subset_ids, subset):
subset.protein_ids = [protein_id for protein_id in subset_ids]
subset.embeddings = dict(
zip(
subset.protein_ids,
[self.embeddings[protein_id] for protein_id in subset.protein_ids]
)
)
subset.sequences = dict(
zip(
subset.protein_ids,
[self.sequences[protein_id] for protein_id in subset.protein_ids]
)
)
subset.connectivity_matrices = dict(
zip(
subset.protein_ids,
[self.connectivity_matrices[protein_id] for protein_id in subset.protein_ids]
)
)
@staticmethod
def get_id_list():
id_list = []
with open(get_id_list_path(), 'r') as fh:
lines = fh.readlines()
for line in lines:
protein_id = line.strip()
if protein_id != "":
id_list.append(protein_id)
return id_list
def get_embeddings(self):
embeddings = {}
embedding_h5 = h5py.File(get_embeddings_path(), 'r')
for key, value in embedding_h5.items():
if key in self.protein_ids:
embeddings[key] = np.array(value[:, :], dtype=np.float32)
return embeddings
def get_sequences(self):
protein_sequences = {}
for record in fastapy.parse(get_sequence_path()):
if record.id in self.protein_ids:
protein_sequences[record.id] = record.seq
return protein_sequences
class BindingResidueDatasetWithLabels(BindingResidueDataset):
def __init__(self, structure_cutoff):
super(BindingResidueDatasetWithLabels, self).__init__(structure_cutoff)
self.labels = self.get_labels(self.sequences)
def get(self, item):
protein_id = self.protein_ids[item]
protein_graph_edges = np.array(self.get_connectivity_matrix(item)["backbone"].nonzero())
protein_graph_cutoff_edges = np.array(self.get_connectivity_matrix(item)["cutoff"].nonzero())
edge_attributes = \
self.get_connectivity_matrix(item)["distance"][protein_graph_cutoff_edges[0], protein_graph_cutoff_edges[1]]
# divide by distance cutoff to avoid div by zero error for linalg.norm
edge_attributes = (edge_attributes / self.structure_cutoff)
# invert distance to achieve weights: distance zero has highest weight
edge_features = (1 - edge_attributes)
protein_graph = Data(
x=torch.Tensor(self.get_embedding(item)),
edge_index=torch.LongTensor(protein_graph_edges),
edge_index_cutoff=torch.LongTensor(protein_graph_cutoff_edges),
edge_features=torch.Tensor(edge_features),
y=torch.Tensor(self.labels[protein_id]),
dssp_features=torch.tensor(self.dssp_features[protein_id]).float()
)
return protein_graph, protein_id
def collect_subset_data(self, subset_ids, subset):
super().collect_subset_data(subset_ids, subset)
subset.labels = dict(
zip(
subset.protein_ids,
[self.labels[protein_id] for protein_id in subset.protein_ids]
)
)
def get_labels(self, protein_sequences):
labels = {}
metal_labels = self.generate_label_dict(get_label_path(LabelType.METAL))
small_labels = self.generate_label_dict(get_label_path(LabelType.SMALL))
nuclear_labels = self.generate_label_dict(get_label_path(LabelType.NUCLEAR))
for protein_id in protein_sequences.keys():
length = len(protein_sequences[protein_id])
label_array = np.zeros((length, 3))
if protein_id in metal_labels:
for metal_label in metal_labels[protein_id]:
label_array[metal_label, LabelType.METAL.value] = 1
if protein_id in small_labels:
for small_label in small_labels[protein_id]:
label_array[small_label, LabelType.SMALL.value] = 1
if protein_id in nuclear_labels:
for nuclear_label in nuclear_labels[protein_id]:
label_array[nuclear_label, LabelType.NUCLEAR.value] = 1
labels[protein_id] = label_array
return labels
@staticmethod
def generate_label_dict(file_path):
labels = {}
with open(file_path, 'r') as fh:
lines = fh.readlines()
for line in lines:
line = line.rstrip()
parts = line.split('\t')
if len(parts) == 2:
protein_id = parts[0]
# label file is 1 indexed so we need to shift all entries by 1
binding_residues = [int(entry) - 1 for entry in parts[1].split(',')]
labels[protein_id] = binding_residues
return labels
|