A Keras implementation of Aspuru-Guzik's molecular autoencoder paper

Abstract from the paper We report a method to convert discrete representations of molecules to and from a multidimensional continuous representation. This generative model allows efficient search and optimization through open-ended spaces of chemical compounds. We train deep neural networks on hundreds of thousands of existing chemical structures to construct two coupled functions: an encoder and a decoder. The encoder converts the discrete representation of a molecule into a real-valued continuous vector, and the decoder converts these continuous vectors back to the discrete representation from this latent space. Continuous representations allow us to automatically generate novel chemical structures by performing simple operations in the latent space, such as decoding random vectors, perturbing known chemical structures, or interpolating between molecules. Continuous representations also allow the use of powerful gradient-based optimization to efficiently guide the search for optimized functional compounds. We demonstrate our method in the design of drug-like molecules as well as organic light-emitting diodes. Link to the paper arXiv

Requirements

Install using

pip install -r requirements.txt

docker build .

The docker container can also be built different TensorFlow binary, for example in order to use GPU:

docker build --build-arg TF_BINARY_URL=https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-0.11.0rc1-cp27-none-linux_x86_64.whl .

You'll need to ensure the proper CUDA libraries are installed for this version to work.

Getting the datasets

A small 50k molecule dataset is included in

data/smiles_50k.h5

data/smiles_500k.h5

smiles_500k.h5

data/model_500k.h5

All h5 files in this repo by git-lfs rather than included directly in the repo.

To download original datasets to work with, you can use the

download_dataset.py

• python download_dataset.py --dataset zinc12

• python download_dataset.py --dataset chembl22

• python download_dataset.py --uri http://my-domain.com/my-file.csv --outfile data/my-file.csv

Preparing the data

To train the network you need a lot of SMILES strings. The

preprocess.py

structure

• Normalizes the length of each string to 120 by appending whitespace as needed.
• Builds a list of the unique characters used in the dataset. (The "charset")
• Substitutes each character in each SMILES string with the integer ID of its location in the charset.
• Converts each character position to a one-hot vector of len(charset).
• Saves this matrix to the specified output file.

Example:

python preprocess.py data/smiles_50k.h5 data/processed.h5

Training the network

The preprocessed data can be fed into the

train.py

python train.py data/processed.h5 model.h5 --epochs 20

If a model file already exists it will be opened and resumed. If it doesn't exist, it will be created.

By default, the latent space is 292-D per the paper, and is configurable with the

--latent_dim

--latent_dim

sample.py

Sampling from a trained model

The

sample.py

--target

Examples:

python sample.py data/processed.h5 model.h5 --target autoencoder

python sample.py data/processed.h5 model.h5 --target encoder --save_h5 encoded.h5
python sample.py target/encoded.h5 model.h5 --target decoder

Performance

After 30 epochs on a 500,000 molecule extract from ChEMBL 21 (~7 hours on a NVIDIA GTX 1080), I'm seeing a loss of 0.26 and a reconstruction accuracy of 0.98.

Projecting the dataset onto 2D latent space gives a figure that looks pretty reasonably like Figure 3 from the paper, though there are some strange striations and it's not quite as well spread out as the examples in the paper.