app.py

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)