<jupyter_start><jupyter_text>Example Workflows - Latent Optimization>Basic latent optimization workflow  Latent Optimization WorkflowsThis notebook shows a basic workflow for optimizing a latent vector relative to a generative model. The focus here is on showing how to set up the code, rather than maximizing performance. For this reason, we will use a simple template and a simple reward function.During latent optimization, we will create several latent vectors. Upon initialisation, these latent vectors will map to compounds with some score. We will then use gradient descent to optimize these vectors to produce high scoring compounds while keeping the generative model constant  Performance NotesThe workflow in this notebook is more CPU-constrained than GPU-constrained due to the need to evaluate samples on CPU. If you have a multi-core machine, it is recommended that you uncomment and run the `set_global_pool` cells in the notebook. This will trigger the use of multiprocessing, which will result in 2-4x speedups.This notebook may run slow on Collab due to CPU limitations.If running on Collab, remember to change the runtime to GPU<jupyter_code>import sys
sys.path.append('..')

from mrl.imports import *
from mrl.core import *
from mrl.chem import *
from mrl.templates.all import *

from mrl.torch_imports import *
from mrl.torch_core import *
from mrl.layers import *
from mrl.dataloaders import *
from mrl.g_models.all import *
from mrl.vocab import *
from mrl.policy_gradient import *
from mrl.train.all import *
from mrl.model_zoo import *<jupyter_output>/home/dmai/miniconda3/envs/mrl/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: to-Python converter for boost::shared_ptr<RDKit::FilterCatalogEntry const> already registered; second conversion method ignored.
  return f(*args, **kwds)<jupyter_text>AgentHere we create the model we want to optimize. We will use the `FP_Cond_LSTM_LM_Small_ZINC` - a LSTM-based conditionl language model pretrained on part of the ZINC database.Note that for latent optimization we specifically need a conditional generative model<jupyter_code>agent = FP_Cond_LSTM_LM_Small_ZINC(drop_scale=0.2,opt_kwargs={'lr':5e-5}, base_model=None)<jupyter_output><empty_output><jupyter_text>Here we freeze the weights of the model so that the model won't be updated during training<jupyter_code>freeze(agent.model)<jupyter_output><empty_output><jupyter_text>TemplateWe will set up a very basic template that will only check compounds for structural validity<jupyter_code>template = Template([ValidityFilter(), 
                     SingleCompoundFilter()],
                    [])

template_cb = TemplateCallback(template, prefilter=True)<jupyter_output><empty_output><jupyter_text>RewardFor the reward, we will load a scikit-learn linear regression model trained to predict affinity against erbB1 using molecular fingerprints.This score function is extremely simple and won't translate well to affinity. It is used as a lightweight example<jupyter_code>class FP_Regression_Score():
    def __init__(self, fname):
        self.model = torch.load(fname)
        self.fp_function = partial(failsafe_fp, fp_function=ECFP6)
        
    def __call__(self, samples):
        mols = to_mols(samples)
        fps = maybe_parallel(self.fp_function, mols)
        fps = [fp_to_array(i) for i in fps]
        x_vals = np.stack(fps)
        preds = self.model.predict(x_vals)
        return preds
    
# if in the repo
reward_function = FP_Regression_Score('../files/erbB1_regression.sklearn')

# if in Collab
# download_files()
# reward_function = FP_Regression_Score('files/erbB1_regression.sklearn')

reward = Reward(reward_function, weight=1.)

aff_reward = RewardCallback(reward, 'aff')<jupyter_output><empty_output><jupyter_text>LatentsHere we show two ways to set up latent vectors to train  Method 1: Initialize Latents as Random VectorsHere we get our initial latents by sampling from a normal distribution. We create 200 latent vectors. During training, these vectors will be randomly sampled to assemble a batch<jupyter_code>latents = torch.randn((500, agent.model.encoder.d_latent))
gen_bs = 1500

sampler = LatentSampler(agent.vocab, agent.model, latents, 'latent', 0, 1., gen_bs, opt_kwargs={'lr':1e-3})

samplers = [sampler]<jupyter_output><empty_output><jupyter_text>Method 2: Initialize Latents from DataHere we grab 5 high scoring samples from the erbB1 training dataset. We then convert these samples into latent vectors. We make several copies of these vectors and add a small amount of random noise.<jupyter_code># if in the repo
df = pd.read_csv('../files/erbB1_affinity_data.csv')

# if in Collab
# download_files()
# df = pd.read_csv('files/erbB1_affinity_data.csv')

df = df[df.neg_log_ic50>10]

smiles = df.sample(n=5).smiles.values

print(reward_function(smiles))
new_ds = agent.dataset.new(smiles)

batch = collate_ds(new_ds)
x,y = batch

latents = agent.model.x_to_latent(to_device(x))

genbatch = 1500
n_latents = 200

latents = latents[np.random.choice(range(latents.shape[0]), n_latents)]
latents = latents + to_device(torch.randn(latents.shape)/100)

sampler = LatentSampler(agent.vocab, agent.model, latents, 'latent', 0, 1., genbatch, opt_kwargs={'lr':1e-3})

samplers = [sampler]<jupyter_output><empty_output><jupyter_text>Loss FunctionWe will use the `PolicyGradient` class, the simplest policy gradient algorithm<jupyter_code>pg = PolicyGradient(discount=True, gamma=0.97)

loss = PolicyLoss(pg, 'PG')

losses = [loss]<jupyter_output><empty_output><jupyter_text>EnvironmentWe create our environment with the objects assembled so far<jupyter_code>env = Environment(agent, template_cb, samplers=samplers, rewards=[aff_reward], losses=losses,
                 cbs=[])<jupyter_output><empty_output><jupyter_text>Train<jupyter_code>set_global_pool(cpus=min(10, os.cpu_count()))
env.fit(128, 90, 500, 25)
env.log.plot_metrics()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08465430016863407, 'fraction_numerical': 0.012141652613827993, 'mean_word_length': 5.481967213114754, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 43, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '619617', 'license': 'MIT', 'n_tokens_mistral': 1924, 'n_tokens_neox': 1825, 'n_words': 569, 'path': 'nbs/tutorials/tutorials.workflows.latent_optimization.ipynb', 'repo_name': 'DarkMatterAI/mrl', 'size': '5930'}