import streamlit as st import numpy as np import pandas as pd import pickle import pygad from tqdm.auto import tqdm from VQGAE.models import VQGAE, OrderingNetwork from CGRtools.containers import QueryContainer from VQGAE.utils import frag_counts_to_inds, restore_order, decode_molecules # define groups to filter allene = QueryContainer() allene.add_atom("C") allene.add_atom("A") allene.add_atom("A") allene.add_bond(1, 2, 2) allene.add_bond(1, 3, 2) peroxide_charge = QueryContainer() peroxide_charge.add_atom("O", charge=-1) peroxide_charge.add_atom("O") peroxide_charge.add_bond(1, 2, 1) peroxide = QueryContainer() peroxide.add_atom("O") peroxide.add_atom("O") peroxide.add_bond(1, 2, 1) def tanimoto_kernel(x, y): """ "The Tanimoto coefficient is a measure of the similarity between two sets. It is defined as the size of the intersection divided by the size of the union of the sample sets." The Tanimoto coefficient is also known as the Jaccard index Adoppted from https://github.com/cimm-kzn/CIMtools/blob/master/CIMtools/metrics/pairwise.py :param x: 2D array of features. :param y: 2D array of features. :return: The Tanimoto coefficient between the two arrays. """ x_dot = np.dot(x, y.T) x2 = (x ** 2).sum(axis=1) y2 = (y ** 2).sum(axis=1) len_x2 = len(x2) len_y2 = len(y2) result = x_dot / (np.array([x2] * len_y2).T + np.array([y2] * len_x2) - x_dot) result[np.isnan(result)] = 0 return result def rescoring(vqgae_latents): frag_counts = np.array(vqgae_latents) rf_scores = rf_model.predict_proba(frag_counts)[:, 1] similarity_scores = tanimoto_kernel(frag_counts, X).max(-1) frag_inds = frag_counts_to_inds(frag_counts, max_atoms=51) _, ordering_scores = restore_order(frag_inds, ordering_model) return rf_scores.tolist(), similarity_scores.tolist(), ordering_scores def fitness_func_batch(ga_instance, solutions, solutions_indices): frag_counts = np.array(solutions) # prediction of activity by random forest rf_score = rf_model.predict_proba(frag_counts)[:, 1] # size penalty if molecule too small mol_size = frag_counts.sum(-1).astype(np.int64) size_penalty = np.where(mol_size < 18, -1.0, 0.) # adding dissimilarity so it generates different solutions dissimilarity_score = 1 - tanimoto_kernel(frag_counts, X).max(-1) dissimilarity_score += np.where(dissimilarity_score == 0, -5, 0) # prediction of ordering score frag_inds = frag_counts_to_inds(frag_counts, max_atoms=51) _, ordering_scores = restore_order(frag_inds, ordering_model) ordering_scores = np.array(ordering_scores) # full fitness function fitness = 0.5 * rf_score + 0.3 * dissimilarity_score + size_penalty + 0.2 * ordering_scores return fitness.tolist() def on_generation_progress(ga): pbar.update(1) @st.cache_data def load_data(batch_size): X = np.load("saved_model/tubulin_qsar_class_train_data_vqgae.npz")["x"] Y = np.load("saved_model/tubulin_qsar_class_train_data_vqgae.npz")["y"] with open("saved_model/rf_class_train_tubulin.pickle", "rb") as inp: rf_model = pickle.load(inp) vqgae_model = VQGAE.load_from_checkpoint( "saved_model/vqgae.ckpt", task="decode", batch_size=batch_size, map_location="cpu" ) vqgae_model = vqgae_model.eval() ordering_model = OrderingNetwork.load_from_checkpoint( "saved_model/ordering_network.ckpt", batch_size=batch_size, map_location="cpu" ) ordering_model = ordering_model.eval() return X, Y, rf_model, vqgae_model, ordering_model st.title('Inverse QSAR of Tubulin inhibitors in colchicine site with VQGAE') data_load_state = st.text('Loading data...') batch_size = 500 X, Y, rf_model, vqgae_model, ordering_model = load_data(batch_size) data_load_state.text("Done! (using st.cache_data)") # initial_pop = X # # num_parents_mating = int(initial_pop.shape[0] * 0.33 // 10 * 10) # keep_parents = int(num_parents_mating * 0.66 // 10 * 10) # print(num_parents_mating, keep_parents) # # num_generations = 30 # with tqdm(total=num_generations) as pbar: # ga_instance = pygad.GA( # fitness_func=fitness_func_batch, # on_generation=on_generation_progress, # initial_population=initial_pop, # num_genes=initial_pop.shape[-1], # fitness_batch_size=batch_size, # num_generations=num_generations, # num_parents_mating=num_parents_mating, # parent_selection_type="rws", # crossover_type="single_point", # mutation_type="adaptive", # mutation_percent_genes=[10, 5], # # https://pygad.readthedocs.io/en/latest/pygad.html#use-adaptive-mutation-in-pygad # save_best_solutions=False, # save_solutions=True, # keep_elitism=0, # turn it off to make keep_parents work # keep_parents=keep_parents, # 2/3 of num_parents_mating # # parallel_processing=['process', 5], # suppress_warnings=True, # random_seed=42, # gene_type=int # ) # ga_instance.run() # # solutions = ga_instance.solutions # solutions = list(set(tuple(s) for s in solutions)) # print(len(solutions)) # # scores = {"rf_score": [], "similarity_score": [], "ordering_score": []} # for i in tqdm(range(len(solutions) // 100 + 1)): # solution = solutions[i * 100: (i + 1) * 100] # rf_score, similarity_score, ordering_score = rescoring(solution) # scores["rf_score"].extend(rf_score) # scores["similarity_score"].extend(similarity_score) # scores["ordering_score"].extend(ordering_score) # # sc_df = pd.DataFrame(scores) # # chosen_gen = sc_df[(sc_df["similarity_score"] < 0.95) & (sc_df["rf_score"] > 0.5) & (sc_df["ordering_score"] > 0.7)] # # chosen_ids = chosen_gen.index.to_list() # chosen_solutions = np.array([solutions[ind] for ind in chosen_ids]) # gen_frag_inds = frag_counts_to_inds(chosen_solutions, max_atoms=51) # # gen_molecules = [] # results = {"score": [], "valid": []} # for i in tqdm(range(gen_frag_inds.shape[0] // batch_size + 1)): # inputs = gen_frag_inds[i * batch_size: (i + 1) * batch_size] # canon_order_inds, scores = restore_order( # frag_inds=inputs, # ordering_model=ordering_model, # ) # molecules, validity = decode_molecules( # ordered_frag_inds=canon_order_inds, # vqgae_model=vqgae_model # ) # gen_molecules.extend(molecules) # results["score"].extend(scores) # results["valid"].extend([1 if i else 0 for i in validity]) # # gen_stats = pd.DataFrame(results) # full_stats = pd.concat([chosen_gen.reset_index(), gen_stats], axis=1, ignore_index=False) # valid_gen_stats = full_stats[full_stats.valid == 1] # valid_gen_mols = [] # for i, record in zip(list(valid_gen_stats.index), valid_gen_stats.to_dict("records")): # mol = gen_molecules[i] # mol.meta.update({ # "rf_score": record["rf_score"], # "similarity_score": record["similarity_score"], # "ordering_score": record["ordering_score"], # }) # valid_gen_mols.append(mol) # # filtered_gen_mols = [] # for mol in valid_gen_mols: # is_frag = allene < mol or peroxide_charge < mol or peroxide < mol # is_macro = False # for ring in mol.sssr: # if len(ring) > 8 or len(ring) < 4: # is_macro = True # break # if not is_frag and not is_macro: # filtered_gen_mols.append(mol)