Chytorch RxnMap

Semisupervised Model trained on USPTO and Pistachio datasets.

Publication (Preprint) with details.

Installation

Use pip install chytorch-rxnmap to install release version.

Or pip install . in source code directory to install DEV version.

Perform Atom-to-atom mapping

AAM integrated into chython package and available as reaction object method. See chython documentation here.

from chython import smiles

r = smiles('OC(=O)C(=C)C=C.C=CC#N>>OC(=O)C1=CCCC(C1)C#N')
r.reset_mapping()
print(format(r, 'm'))
>> [C:2]([C:4](=[CH2:5])[CH:6]=[CH2:7])(=[O:3])[OH:1].[CH2:8]=[CH:9][C:10]#[N:11]>>[O:3]=[C:2]([OH:1])[C:4]=1[CH2:5][CH:9]([C:10]#[N:11])[CH2:8][CH2:7][CH:6]=1

Pretrained model

To load pretrained model use:

from chytorch.zoo.rxnmap import Model  
model = Model.pretrained()

To prepare data-loader use:

from chython import SMILESRead

data = []
for r in SMILESRead('data.smi'):
    r.canonicalize()  # fix aromaticity and functional groups
    data.append(r.pack())  # store in compressed format

dl = model.prepare_dataloader(data, batch_size=20)

To get embeddings use:

for b in dl:
    e = model(b)

Note: embeddings contain: cls embedding, [unusable molecular embedding, list of atoms embeddings] * n. Where n is the number of molecules in reaction equation.

To extract aggregated embedding, use cls embedding x = e[:, 0].

To extract atoms-only embeddings, use masking:

• x = e[b[3] > 1] - for all atoms
• x = e[b[3] == 2] - for reactants only
• x = e[b[3] == 3] - for products only

To get all-to-all tokens attention matrix:

for b in dl:
    a = model(b, mapping_task=True)

Training new model

from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.plugins import DDPPlugin

callback = ModelCheckpoint(monitor='trn_loss_tot', save_weights_only=True, save_top_k=3, save_last=True,
                           every_n_train_steps=10000)
trainer = Trainer(gpus=-1, precision=16, max_steps=1000000, callbacks=[callback],
                  strategy=DDPPlugin(find_unused_parameters=False))

model = Model(lr_warmup=1e4, lr_period=5e5, lr_max=1e-4, lr_decrease_coef=.01, masking_rate=.15, **kwargs)
# lr_warmup=1e4, lr_period=5e5, lr_max=1e-4, lr_decrease_coef=.01 - see chytorch.optim.lr_scheduler.WarmUpCosine. 
# kwargs - see chytorch.nn.ReactionEncoder.
# masking_rate - probability of token masking.
trainer.fit(model, dl)