File size: 2,693 Bytes
4742cab |
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 |
from itertools import accumulate
import numpy as np
import torch
from torch.utils.data import Dataset
class ProcessedLigandPocketDataset(Dataset):
def __init__(self, npz_path, center=True, transform=None):
self.transform = transform
with np.load(npz_path, allow_pickle=True) as f:
data = {key: val for key, val in f.items()}
# split data based on mask
self.data = {}
for (k, v) in data.items():
if k == 'names' or k == 'receptors':
self.data[k] = v
continue
sections = np.where(np.diff(data['lig_mask']))[0] + 1 \
if 'lig' in k \
else np.where(np.diff(data['pocket_mask']))[0] + 1
self.data[k] = [torch.from_numpy(x) for x in np.split(v, sections)]
# add number of nodes for convenience
if k == 'lig_mask':
self.data['num_lig_atoms'] = \
torch.tensor([len(x) for x in self.data['lig_mask']])
elif k == 'pocket_mask':
self.data['num_pocket_nodes'] = \
torch.tensor([len(x) for x in self.data['pocket_mask']])
if center:
for i in range(len(self.data['lig_coords'])):
mean = (self.data['lig_coords'][i].sum(0) +
self.data['pocket_coords'][i].sum(0)) / \
(len(self.data['lig_coords'][i]) + len(self.data['pocket_coords'][i]))
self.data['lig_coords'][i] = self.data['lig_coords'][i] - mean
self.data['pocket_coords'][i] = self.data['pocket_coords'][i] - mean
def __len__(self):
return len(self.data['names'])
def __getitem__(self, idx):
data = {key: val[idx] for key, val in self.data.items()}
if self.transform is not None:
data = self.transform(data)
return data
@staticmethod
def collate_fn(batch):
out = {}
for prop in batch[0].keys():
if prop == 'names' or prop == 'receptors':
out[prop] = [x[prop] for x in batch]
elif prop == 'num_lig_atoms' or prop == 'num_pocket_nodes' \
or prop == 'num_virtual_atoms':
out[prop] = torch.tensor([x[prop] for x in batch])
elif 'mask' in prop:
# make sure indices in batch start at zero (needed for
# torch_scatter)
out[prop] = torch.cat([i * torch.ones(len(x[prop]))
for i, x in enumerate(batch)], dim=0)
else:
out[prop] = torch.cat([x[prop] for x in batch], dim=0)
return out
|