- Notifications
You must be signed in to change notification settings - Fork0
Conditional Molecular Structure Generation
License
Ferg-Lab/molgen
Folders and files
Name | Name | Last commit message | Last commit date | |
---|---|---|---|---|
Repository files navigation
This package implements Generative Adversarial Networks (GANs) and Denoising Diffusion Probabilistic Models (DDPMs) for generative tasks such as conditional molecular structure generation.
To usemolgen
, you will need an environment with the following packages:
- Python 3.7+
- PyTorch
- PyTorch Lightning
- Einops
For running and visualizing examples:
Once you have these packages installed, you can installmolgen
in the same environment using
$ pip install -e .
Once installed, you can use the package. This example trains a WGANGP to reproduce the alanine dipeptide backbone atoms conditioned on the backbone diedral angles (examples
directory.
frommolgen.modelsimportWGANGPfrompathlibimportPathimportmdtajasmdimporttorchimportnumpyasnp# load datapdb_fname='examples/data/alanine-dipeptide-nowater.pdb'trj_fnames= [str(i)foriinPath('examples/data/').glob('alanine-dipeptide-*-250ns-nowater.xtc')]trjs= [md.load(t,top=pdb_fname).center_coordinates()fortintrj_fnames]# process xyz coordinates and conditioning variablesxyz=list()phi_psi=list()fortrjintrjs:t_backbone=trj.atom_slice(trj.top.select('backbone')).center_coordinates()n=trj.xyz.shape[0]_,phi=md.compute_phi(trj)_,psi=md.compute_psi(trj)xyz.append(torch.tensor(t_backbone.xyz.reshape(n,-1)).float())phi_psi.append(torch.tensor(np.concatenate((phi,psi),-1)).float())# ininstantiate the modelmodel=WGANGP(xyz[0].shape[1],phi_psi[0].shape[1])# fit the modelmodel.fit(xyz,phi_psi,max_epochs=25)# Generate synthetic configurationsxyz_gen=model.generate(torch.cat(phi_psi))xyz_gen=xyz_gen.reshape(xyz_gen.size(0),-1,3)# Save model checkpointmodel.save('ADP.ckpt')# Load from checkpointmodel=WGANGP.load_from_checkpoint('ADP.ckpt')
Supports both generators based on both Generative Adversarial Networks (GANs) and Denoising Diffusion Probabilistic Models (DDPMs). The example above uses GANs, DDPMs support an equivalent API -- for example,
frommolgen.modelsimportDDPMmodel=DDPM(....)
Code for the DDPM models are taken from:https://github.com/lucidrains/denoising-diffusion-pytorch (version 1.0.5)
Copyright (c) 2023, Kirill Shmilovich
Project based on theComputational Molecular Science Python Cookiecutter version 1.1.
About
Conditional Molecular Structure Generation
Resources
License
Code of conduct
Uh oh!
There was an error while loading.Please reload this page.
Stars
Watchers
Forks
Releases
Packages0
Uh oh!
There was an error while loading.Please reload this page.