| | 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) |
