|
|
import os, json
|
|
|
import torch
|
|
|
import utils
|
|
|
|
|
|
def calc_feats(smi, ms, nls, cfg):
|
|
|
item = {}
|
|
|
item['ms_bins'] = utils.ms_binner(ms, nls,
|
|
|
min_mz=cfg.min_mz,
|
|
|
max_mz=cfg.max_mz,
|
|
|
bin_size=cfg.bin_size,
|
|
|
add_nl=cfg.add_nl,
|
|
|
binary_intn=cfg.binary_intn)
|
|
|
|
|
|
fmcalced = False
|
|
|
if 'fp' in cfg.mol_encoder:
|
|
|
if not 'fm' in cfg.mol_encoder:
|
|
|
item['mol_fps'] = utils.mol_fp_encoder(smi,
|
|
|
tp=cfg.fptype,
|
|
|
nbits=cfg.mol_embedding_dim)
|
|
|
else:
|
|
|
item['mol_fps'], item['mol_fmvec'] = utils.mol_fp_fm_encoder(smi,
|
|
|
tp=cfg.fptype,
|
|
|
nbits=cfg.mol_embedding_dim)
|
|
|
fmcalced = True
|
|
|
if 'gnn' in cfg.mol_encoder:
|
|
|
f = utils.mol_graph_featurizer(smi)
|
|
|
if not f:
|
|
|
return None
|
|
|
item.update(f)
|
|
|
if 'fm' in cfg.mol_encoder and not fmcalced:
|
|
|
item['mol_fmvec'] = utils.smi2fmvec(smi)
|
|
|
|
|
|
return item
|
|
|
|
|
|
class Dataset(torch.utils.data.Dataset):
|
|
|
def __init__(self, inp, cfg):
|
|
|
if type(inp) is str:
|
|
|
self.data = json.load(open(inp))
|
|
|
else:
|
|
|
self.data = inp
|
|
|
|
|
|
self.cfg = cfg
|
|
|
|
|
|
def __getitem__(self, idx):
|
|
|
item = {}
|
|
|
try:
|
|
|
if 'ms_bins' in self.data[idx]:
|
|
|
return self.data[idx]
|
|
|
|
|
|
if 'nls' in self.data[idx]:
|
|
|
nls = self.data[idx]['nls']
|
|
|
else:
|
|
|
nls = []
|
|
|
|
|
|
ms = self.data[idx]['ms']
|
|
|
smi = self.data[idx]['smiles']
|
|
|
|
|
|
item = calc_feats(smi, ms, nls, self.cfg)
|
|
|
|
|
|
except Exception as e:
|
|
|
print('='*50, idx, str(e))
|
|
|
return None
|
|
|
|
|
|
return item
|
|
|
|
|
|
def __len__(self):
|
|
|
return len(self.data)
|
|
|
|
|
|
class DatasetGNNFP(torch.utils.data.Dataset):
|
|
|
def __init__(self, inp, cfg):
|
|
|
if type(inp) is str:
|
|
|
self.data = json.load(open(inp))
|
|
|
else:
|
|
|
self.data = inp
|
|
|
|
|
|
self.cfg = cfg
|
|
|
|
|
|
def __getitem__(self, idx):
|
|
|
try:
|
|
|
smi = self.data[idx]['smiles']
|
|
|
item = {}
|
|
|
item['mol_fps'] = utils.mol_fp_encoder(smi,
|
|
|
tp=self.cfg.fptype,
|
|
|
nbits=self.cfg.mol_embedding_dim)
|
|
|
item.update(utils.mol_graph_featurizer(smi))
|
|
|
except Exception as e:
|
|
|
print('='*50, idx, str(e))
|
|
|
return None
|
|
|
|
|
|
return item
|
|
|
|
|
|
def __len__(self):
|
|
|
return len(self.data)
|
|
|
|
|
|
class PathDataset(torch.utils.data.Dataset):
|
|
|
def __init__(self, pathlist, cfg):
|
|
|
self.fns = pathlist
|
|
|
self.cfg = cfg
|
|
|
self.data = {}
|
|
|
|
|
|
def __getitem__(self, idx):
|
|
|
try:
|
|
|
item = {}
|
|
|
nls = []
|
|
|
if not idx in self.data:
|
|
|
out = self.proc_data(self.fns[idx], self.cfg.energy)
|
|
|
if out is None:
|
|
|
return None
|
|
|
self.data[idx] = out
|
|
|
|
|
|
ms = self.data[idx]['ms']
|
|
|
smi = self.data[idx]['smiles']
|
|
|
|
|
|
item = calc_feats(smi, ms, nls, self.cfg)
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
|
return None
|
|
|
|
|
|
return item
|
|
|
|
|
|
def proc_data(self, fn, energy='Energy1'):
|
|
|
tl = open(fn).readlines()
|
|
|
l = []
|
|
|
try:
|
|
|
flag = False
|
|
|
for i in tl:
|
|
|
if energy in i:
|
|
|
smi = i.split(';')[-2]
|
|
|
flag = True
|
|
|
continue
|
|
|
if 'END IONS' in i:
|
|
|
if flag:
|
|
|
break
|
|
|
if flag:
|
|
|
mz, intn = i.split(' ')
|
|
|
l.append((float(mz), float(intn)))
|
|
|
except:
|
|
|
return None
|
|
|
|
|
|
out = {'ms': l, 'smiles': smi}
|
|
|
return out
|
|
|
|
|
|
def __len__(self):
|
|
|
return len(self.fns) |
