Homec4science
Diffusion master_thesis_Lee (master) Browse

TWtest/enviPath_python/prediction_utils.py
0ea7cdc539ebmaster

Last Change
Raw File

prediction_utils.py
View Options

from enviPath_python.objects import *
from datetime import datetime
import pytz
tz = pytz.timezone('Europe/Berlin')
def upload_envipath_pathway(result, user, pkg):
source = result[0]
time = datetime.now(tz).strftime("%Y-%m-%d-%H:%M:%S")
pkg.add_compound(smiles=source['smiles'],name=source['name'])
pathway = Pathway.create(pkg, smiles=source['smiles'], name=user + '-' + source['name'] + ' ' + time, root_node_only=True)
# Add the observed degradation product as a second node
for TP in result[1:]:
# print('adding to pw:', TP['name'], TP['smiles'], TP['generation'], TP['parent'])
pathway.add_node(smiles=TP['smiles'], node_name=TP['name'], node_depth=TP['generation'])
# check for the case of multiple parents:
for parent in TP['parent'].split(','):
pathway.add_node(smiles=parent)
pathway.add_edge(smirks='{}>>{}'.format(parent, TP['smiles']))
print('New pathway created for {}: {}'.format(source['name'], pathway.id))
return {'name': source['name'], 'id': pathway.id}
def clean_result(result_dict): # also add names
result_list = list(result_dict.values())
result_list.sort(key=lambda x: x['generation']) # to make sure that source compound is first
result_list.sort(reverse=True, key=lambda x: x['combined_probability'])
# get name of source compound
source_name = result_list[0]['name']
TP_count = 0
for res in result_list[1:]:
TP_count += 1
res['name'] = 'TP_{}_{}'.format(source_name, TP_count)
return result_list
def load_input(input_path):
f = open(input_path)
input_list = []
for line in f:
if line not in ['', '\n']:
line_split = line.rstrip().split('\t')
input_list.append({'smiles': line_split[0], 'name': line_split[1]})
#todo: canonicalize input structure
return input_list
def print_result(open_file, result, pathway = None):
open_file.write('///')
if pathway:
open_file.write(' Pathway name: {}, Pathway id: {}'.format(pathway['name'], pathway['id'])) # signifies new compound entry
open_file.write('\n')
header= ('SMILES\tname\tcombined_probability\trules\tgeneration\tprobability\tparent\n')
open_file.write(header)
for TP in result:
open_file.write('{}\t{}\t{}\t{}\t{}\t{}\t{}\n'.format(TP['smiles'],
TP['name'],
TP['combined_probability'],
TP['rules'],
TP['generation'],
TP['probability'],
TP['parent'])
)
def predict_TPs(model, input_smiles, input_name, limnode, limdepth, threshold, greedy):
print('\n### PREDICT TPs FOR COMPOUND {} ###\n'.format(input_name))
validated_TPs = {} # container for resulting predictions
queued_items = [{'probability': 1, 'combined_probability': 1, 'smiles': input_smiles, 'generation': 0, 'parent': '', 'rules': '', 'name': input_name}]
queue = [input_smiles] # queue is updated after each cycle to have top TP first, list of smiles
depth = 0
while len(validated_TPs) <= limnode and depth < limdepth:
if len(queue) == 0:
print('\nEmpty queue - The exploration of has converged at {} predicted TPs'.format(len(validated_TPs)))
return validated_TPs # stop TP prediction
smiles = queue.pop(0) # get top item in queue
data = queued_items.pop(0) # remove data from queued items
result_list = expand_smiles(model,smiles) # create children
TP_dict = result_to_compound_dict(result_list)
TP_dict = filter_TP_dict(TP_dict, threshold)
queue, queued_items, validated_TPs = update_queue(queue, queued_items, validated_TPs, TP_dict, data, greedy)
validated_TPs[smiles] = data
depth = get_largestDepth(validated_TPs, depth)
return validated_TPs
def filter_TP_dict(TP_dict, threshold):
new_dict = dict()
for key in TP_dict:
if TP_dict[key]['probability'] >= threshold:
new_dict[key] = TP_dict[key]
else:
print("Remove TP with probability:",TP_dict[key]['probability'])
return new_dict
def get_largestDepth(validated_TPs, depth):
max_depth = depth
for key in validated_TPs.keys():
max_depth = max(max_depth, validated_TPs[key]['generation'])
return max_depth
# Functions
def expand_smiles(model, smiles):
res = model.classify_smiles(smiles)
# sort by probability
res.sort(reverse=True, key=lambda x: x['probability'])
return res
def result_to_compound_dict(result):
compound_dict = {}
for r in result:
probability = float(r['probability'])
for product_smiles in r['products']:
if product_smiles not in compound_dict.keys():
compound_dict[product_smiles] = {'rules' : r['name'], 'probability': probability, 'smiles': product_smiles}
else:
# check if there's a rule with better probability
if probability > compound_dict[product_smiles]['probability']:
# update probability and rules associated to this probability
compound_dict[product_smiles]['probability'] = probability
compound_dict[product_smiles]['rules'] = r['name']
return compound_dict
def update_queue(_queue,_queued_items, _validated_TPs, _TPs, _parent_data, greedy):
parent_probability = _parent_data['combined_probability']
parent_generation = _parent_data['generation']
parent_smiles = _parent_data['smiles']
queue_before = len(_queue)
for smiles in _TPs:
data = _TPs[smiles]
# If the probability is 0 , we don't consider the TP further
this_probability = data['probability']
if this_probability == 0:
continue
# add combined probability
this_combined_probability = parent_probability * this_probability
this_generation = parent_generation + 1
rules = data['rules']
# first, check if compound already in validated. if yes, update
if smiles in _validated_TPs.keys():
_validated_TPs[smiles] = update_compound_entry(_validated_TPs[smiles],
this_combined_probability, rules,
this_generation, parent_smiles)
# next, check if compound is already in queue. if yes, update
elif smiles in _queue:
index = _queue.index(smiles)
assert smiles == _queued_items[index]['smiles'], 'smiles {} does not match smiles in {}'.format(smiles, _queued_items[index])
_queued_items[index] = update_compound_entry(_queued_items[index],
this_combined_probability, rules,
this_generation, parent_smiles)
# else, add new item to queue
else:
data['combined_probability'] = this_combined_probability
data['generation'] = this_generation
data['parent'] = parent_smiles
_queued_items.append(data)
_queue.append(smiles)
assert len(_queued_items) == len(_queue)
# order dict by combined probability
if greedy:
_queued_items.sort(reverse=True, key=lambda x: x['combined_probability'])
queue_after = len(_queue)
print('Added {} smiles to queue'.format(queue_after - queue_before))
new_queue = [] # resetting queue
[new_queue.append(x['smiles']) for x in _queued_items]
print ('New queue for compound', parent_smiles)
for q in new_queue:
print(q, _queued_items[new_queue.index(q)]['combined_probability'])
return new_queue, _queued_items, _validated_TPs
def update_compound_entry(compound_entry, this_combined_probability, rules, this_generation, parent_smiles):
if compound_entry['combined_probability'] < this_combined_probability:
compound_entry['combined_probability'] = this_combined_probability
compound_entry['rules'] = rules
compound_entry['generation'] = this_generation
compound_entry['parent'] = parent_smiles
elif compound_entry['combined_probability'] == this_combined_probability:
compound_entry['rules'] += ',{}'.format(rules)
compound_entry['parent'] += ',{}'.format(parent_smiles)
return compound_entry
c4science · Help