Mirrors_Framework/deepmind-research
1
0
Fork 0
mirror of https://github.com/google-deepmind/deepmind-research.git synced 2026-10-06 06:59:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Release of "learned_free_energy_estimation".

Browse Source 
PiperOrigin-RevId: 336058713
This commit is contained in:
Peter Wirnsberger
2020-10-12 16:05:37 +01:00
committed by Saran Tunyasuvunakool
parent c82e368e0e
commit 1d86410c90
7 changed files with 900936 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -24,6 +24,7 @@ https://deepmind.com/research/publications/
## Projects
* [Targeted free energy estimation via learned mappings](learned_free_energy_estimation), Journal of Chemical Physics 2020
* [Learning to Simulate Complex Physics with Graph Networks](learning_to_simulate), ICML 2020
* [Physically Embedded Planning Problems](physics_planning_games)
* [PolyGen: PolyGen: An Autoregressive Generative Model of 3D Meshes](polygen), ICML 2020
learned_free_energy_estimation/README.md
+109
View File
@@ -0,0 +1,109 @@
# Targeted free energy estimation via learned mappings
This repository contains supporting data for our publication
([arXiv](https://arxiv.org/abs/2002.04913)). Here, we provide
- molecular dynamics (MD) datasets underlying the results reported in our paper,
- a LAMMPS input script to generate these datasets, and
- the data plotted in Fig. 5 of our paper to facilitate comparison.
## Abstract
Free energy perturbation (FEP) was proposed by Zwanzig more than six decades ago
as a method to estimate free energy differences, and has since inspired a huge
body of related methods that use it as an integral building block. Being an
importance sampling based estimator, however, FEP suffers from a severe
limitation: the requirement of sufficient overlap between distributions.
One strategy to mitigate this problem, called Targeted Free Energy Perturbation,
uses a high-dimensional mapping in configuration space to increase overlap of
the underlying distributions. Despite its potential, this method has attracted
only limited attention due to the formidable challenge of formulating a
tractable mapping. Here, we cast Targeted FEP as a machine learning problem in
which the mapping is parameterized as a neural network that is optimized so as
to increase overlap. We develop a new model architecture that respects
permutational and periodic symmetries often encountered in atomistic simulations
and test our method on a fully-periodic solvation system. We demonstrate that
our method leads to a substantial variance reduction in free energy estimates
when compared against baselines, without requiring any additional data.
## Dataset
We generated the datasets using the open-source MD package
[LAMMPS](https://lammps.sandia.gov). The prototypical solvation problem of study
consists of a solute particle immersed in a liquid comprising 125 solvent
particles. The solvent-solvent interactions are modelled using a Lennard-Jones
potential and the solute-solvent interactions via a Weeks-Chandler-Andersen
(WCA) potential. Further simulation details can be found in the LAMMPS script
provided (see below) and in our [paper](https://arxiv.org/abs/2002.04913)
(see Sec. 4 and Appendix B).
### Download
You can download the compressed datasets (~3.8GB) using the command:
> wget https://storage.googleapis.com/learned_free_energy_estimation/learned_free_energy_estimation_datasets.tar.bz2
or by copying the above link directly into your browser.
Once the archive `learned_free_energy_estimation_datasets.tar.bz2` is
downloaded, you can extract it with the command:
> tar -xvf learned_free_energy_estimation_datasets.tar.bz2
### Data format
The archive contains a total of 40 files:
- 10 train datasets for ensemble *A* (`ensemble_a_train_<<index>>.dat`),
- 10 train datasets for ensemble *B* (`ensemble_b_train_<<index>>.dat`),
- 10 test datasets for ensemble *A* (`ensemble_a_test_<<index>>.dat`) and
- 10 test datasets for ensemble *B* (`ensemble_b_test_<<index>>.dat`).
Each file is text-based and stored in a LAMMPS compatible format (see [dump command](https://lammps.sandia.gov/doc/dump.html)). Train datasets contain 90k records
each and test datasets contain 10k records, totalling 1M records for each
ensemble.
Each record contains 135 lines and is structured as follows:
- lines 1-9: Header information.
- lines 10-135: A matrix with shape `[126, 5]` containing the
- `id` (column 1),
- `type` (column 2) and
- `x, y, z` coordinates (columns 3-5)
of all particles.
For information on how the data was generated and partitioned into the final
datasets we refer to Sec. 4 and Appendix B of our [paper](https://arxiv.org/abs/2002.04913).
## LAMMPS script
The file `lammps.dat` contains a sample input script to generate data from
ensemble *A*. You can generate data from ensemble *B* by updating the value of
the solute radius, as suggested in the inline comment. For more information on
how the datasets were post-processed and partitioned, we refer to Sec. 4 and
Appendix B of our [paper](https://arxiv.org/abs/2002.04913).
## Figures
The subdirectory `figures` contains 4 files:
- `figure_5a_work_values.dat`: contains data underlying the histogram of work values in Fig. 5a.
- `figure_5b_df_bar.dat`: contains the BAR estimate of dF in Fig. 5b.
- `figure_5b_df_lbar.dat`: contains the LBAR estimate of dF in Fig. 5b.
- `figure_5b_df_mbar.dat`: contains the MBAR estimate of dF in Fig. 5b.
## Reference
If you find this repository helpful for your research, please cite our publication:
```
@article{Wirnsberger2020,
title={Targeted free energy estimation via learned mappings},
author={Wirnsberger, Peter and Ballard, Andrew J and Papamakarios, George and
Abercrombie, Stuart and Racanière, Sébastien and Pritzel, Alexander and
Jimenez Rezende, Danilo and Blundell, Charles}
journal={Journal of Chemical Physics},
vol={153},
year={2020},
doi={10.1063/5.0018903}
}
```
## Disclaimer
This is not an official Google product.
learned_free_energy_estimation/figures/figure_5a_work_values.dat
+900005
View File
File diff suppressed because it is too large Load Diff
learned_free_energy_estimation/figures/figure_5b_df_bar.dat
+254
View File
File diff suppressed because it is too large Load Diff
learned_free_energy_estimation/figures/figure_5b_df_lbar.dat
+254
View File
File diff suppressed because it is too large Load Diff
learned_free_energy_estimation/figures/figure_5b_df_mbar.dat
+184
View File
@@ -0,0 +1,184 @@
# This file contains the free energy difference estimate (dF) obtained with MBAR.
# column 1: samples per stage
# column 2: running estimate of dF
# column 3: error estimate of dF
5.000000e+02 3.348947e+01 6.862397e-02
1.000000e+03 3.347387e+01 4.854148e-02
1.500000e+03 3.347623e+01 3.967133e-02
2.000000e+03 3.348673e+01 3.439447e-02
2.500000e+03 3.347544e+01 3.075069e-02
3.000000e+03 3.348640e+01 2.809994e-02
3.500000e+03 3.347364e+01 2.599614e-02
4.000000e+03 3.346986e+01 2.431905e-02
4.500000e+03 3.346961e+01 2.291930e-02
5.000000e+03 3.346484e+01 2.173656e-02
5.500000e+03 3.346934e+01 2.072029e-02
6.000000e+03 3.346451e+01 1.984469e-02
6.500000e+03 3.346112e+01 1.906542e-02
7.000000e+03 3.346448e+01 1.837210e-02
7.500000e+03 3.345465e+01 1.775237e-02
8.000000e+03 3.345397e+01 1.718846e-02
8.500000e+03 3.345652e+01 1.667432e-02
9.000000e+03 3.345805e+01 1.620102e-02
9.500000e+03 3.346146e+01 1.577490e-02
1.000000e+04 3.346144e+01 1.537617e-02
1.050000e+04 3.345814e+01 1.500688e-02
1.100000e+04 3.345771e+01 1.466039e-02
1.150000e+04 3.346071e+01 1.433827e-02
1.200000e+04 3.346027e+01 1.403640e-02
1.250000e+04 3.345883e+01 1.375421e-02
1.300000e+04 3.345853e+01 1.348973e-02
1.350000e+04 3.346009e+01 1.323468e-02
1.400000e+04 3.346103e+01 1.299762e-02
1.450000e+04 3.345552e+01 1.276921e-02
1.500000e+04 3.345334e+01 1.255442e-02
1.550000e+04 3.345482e+01 1.235281e-02
1.600000e+04 3.345732e+01 1.215894e-02
1.650000e+04 3.345797e+01 1.197288e-02
1.700000e+04 3.345815e+01 1.179542e-02
1.750000e+04 3.345590e+01 1.162467e-02
1.800000e+04 3.345629e+01 1.146072e-02
1.850000e+04 3.345542e+01 1.130306e-02
1.900000e+04 3.345807e+01 1.115385e-02
1.950000e+04 3.345522e+01 1.101022e-02
2.000000e+04 3.345209e+01 1.087127e-02
2.050000e+04 3.345170e+01 1.073895e-02
2.100000e+04 3.345297e+01 1.061029e-02
2.150000e+04 3.345232e+01 1.048628e-02
2.200000e+04 3.345287e+01 1.036741e-02
2.250000e+04 3.345100e+01 1.025204e-02
2.300000e+04 3.345211e+01 1.014036e-02
2.350000e+04 3.345150e+01 1.003258e-02
2.400000e+04 3.345267e+01 9.927045e-03
2.450000e+04 3.345188e+01 9.824914e-03
2.500000e+04 3.344938e+01 9.725246e-03
2.550000e+04 3.345166e+01 9.629378e-03
2.600000e+04 3.345270e+01 9.535662e-03
2.650000e+04 3.345214e+01 9.445890e-03
2.700000e+04 3.345417e+01 9.358481e-03
2.750000e+04 3.345247e+01 9.272714e-03
2.800000e+04 3.345346e+01 9.188538e-03
2.850000e+04 3.345430e+01 9.108273e-03
2.900000e+04 3.345416e+01 9.029007e-03
2.950000e+04 3.345309e+01 8.952522e-03
3.000000e+04 3.345190e+01 8.878327e-03
3.050000e+04 3.344897e+01 8.805188e-03
3.100000e+04 3.344964e+01 8.732210e-03
3.150000e+04 3.344939e+01 8.663034e-03
3.200000e+04 3.344794e+01 8.594927e-03
3.250000e+04 3.344955e+01 8.528876e-03
3.300000e+04 3.345027e+01 8.463777e-03
3.350000e+04 3.344999e+01 8.400015e-03
3.400000e+04 3.344984e+01 8.338387e-03
3.450000e+04 3.345104e+01 8.277393e-03
3.500000e+04 3.344932e+01 8.218733e-03
3.550000e+04 3.344877e+01 8.160845e-03
3.600000e+04 3.344790e+01 8.104316e-03
3.650000e+04 3.344738e+01 8.048099e-03
3.700000e+04 3.344686e+01 7.993584e-03
3.750000e+04 3.344745e+01 7.940271e-03
3.800000e+04 3.344820e+01 7.888244e-03
3.850000e+04 3.344765e+01 7.836854e-03
3.900000e+04 3.344642e+01 7.786979e-03
3.950000e+04 3.344496e+01 7.737392e-03
4.000000e+04 3.344476e+01 7.688966e-03
4.050000e+04 3.344575e+01 7.641278e-03
4.100000e+04 3.344529e+01 7.594414e-03
4.150000e+04 3.344507e+01 7.548341e-03
4.200000e+04 3.344515e+01 7.503361e-03
4.250000e+04 3.344542e+01 7.458615e-03
4.300000e+04 3.344475e+01 7.415060e-03
4.350000e+04 3.344495e+01 7.372151e-03
4.400000e+04 3.344539e+01 7.330419e-03
4.450000e+04 3.344489e+01 7.289283e-03
4.500000e+04 3.344518e+01 7.248814e-03
4.550000e+04 3.344483e+01 7.208634e-03
4.600000e+04 3.344505e+01 7.169400e-03
4.650000e+04 3.344549e+01 7.130802e-03
4.700000e+04 3.344656e+01 7.092623e-03
4.750000e+04 3.344620e+01 7.055154e-03
4.800000e+04 3.344553e+01 7.017677e-03
4.850000e+04 3.344510e+01 6.981500e-03
4.900000e+04 3.344430e+01 6.945682e-03
4.950000e+04 3.344375e+01 6.910502e-03
5.000000e+04 3.344317e+01 6.875701e-03
5.050000e+04 3.344380e+01 6.841235e-03
5.100000e+04 3.344336e+01 6.807528e-03
5.150000e+04 3.344314e+01 6.774260e-03
5.200000e+04 3.344290e+01 6.741623e-03
5.250000e+04 3.344230e+01 6.709539e-03
5.300000e+04 3.344314e+01 6.678291e-03
5.350000e+04 3.344342e+01 6.647072e-03
5.400000e+04 3.344318e+01 6.616171e-03
5.450000e+04 3.344379e+01 6.585470e-03
5.500000e+04 3.344381e+01 6.555845e-03
5.550000e+04 3.344434e+01 6.526642e-03
5.600000e+04 3.344501e+01 6.497285e-03
5.650000e+04 3.344619e+01 6.468459e-03
5.700000e+04 3.344623e+01 6.439974e-03
5.750000e+04 3.344664e+01 6.412096e-03
5.800000e+04 3.344615e+01 6.384088e-03
5.850000e+04 3.344599e+01 6.356764e-03
5.900000e+04 3.344632e+01 6.330156e-03
5.950000e+04 3.344659e+01 6.303682e-03
6.000000e+04 3.344593e+01 6.277147e-03
6.050000e+04 3.344578e+01 6.251312e-03
6.100000e+04 3.344642e+01 6.225851e-03
6.150000e+04 3.344603e+01 6.200300e-03
6.200000e+04 3.344555e+01 6.175259e-03
6.250000e+04 3.344555e+01 6.150557e-03
6.300000e+04 3.344542e+01 6.126220e-03
6.350000e+04 3.344515e+01 6.101949e-03
6.400000e+04 3.344525e+01 6.077983e-03
6.450000e+04 3.344496e+01 6.054339e-03
6.500000e+04 3.344501e+01 6.030827e-03
6.550000e+04 3.344471e+01 6.007627e-03
6.600000e+04 3.344506e+01 5.984742e-03
6.650000e+04 3.344514e+01 5.962256e-03
6.700000e+04 3.344494e+01 5.940045e-03
6.750000e+04 3.344566e+01 5.917884e-03
6.800000e+04 3.344522e+01 5.895819e-03
6.850000e+04 3.344488e+01 5.874460e-03
6.900000e+04 3.344457e+01 5.853076e-03
6.950000e+04 3.344488e+01 5.831963e-03
7.000000e+04 3.344486e+01 5.811116e-03
7.050000e+04 3.344525e+01 5.790804e-03
7.100000e+04 3.344566e+01 5.770541e-03
7.150000e+04 3.344601e+01 5.750545e-03
7.200000e+04 3.344603e+01 5.730280e-03
7.250000e+04 3.344619e+01 5.710471e-03
7.300000e+04 3.344545e+01 5.690757e-03
7.350000e+04 3.344417e+01 5.671172e-03
7.400000e+04 3.344337e+01 5.651868e-03
7.450000e+04 3.344334e+01 5.632921e-03
7.500000e+04 3.344244e+01 5.613913e-03
7.550000e+04 3.344261e+01 5.595240e-03
7.600000e+04 3.344314e+01 5.576832e-03
7.650000e+04 3.344388e+01 5.558491e-03
7.700000e+04 3.344386e+01 5.540588e-03
7.750000e+04 3.344375e+01 5.522747e-03
7.800000e+04 3.344326e+01 5.505254e-03
7.850000e+04 3.344329e+01 5.487653e-03
7.900000e+04 3.344302e+01 5.470438e-03
7.950000e+04 3.344398e+01 5.453349e-03
8.000000e+04 3.344413e+01 5.436500e-03
8.050000e+04 3.344421e+01 5.419629e-03
8.100000e+04 3.344489e+01 5.402948e-03
8.150000e+04 3.344496e+01 5.386430e-03
8.200000e+04 3.344543e+01 5.369879e-03
8.250000e+04 3.344561e+01 5.353712e-03
8.300000e+04 3.344542e+01 5.337605e-03
8.350000e+04 3.344541e+01 5.321604e-03
8.400000e+04 3.344607e+01 5.305934e-03
8.450000e+04 3.344640e+01 5.290176e-03
8.500000e+04 3.344655e+01 5.274630e-03
8.550000e+04 3.344652e+01 5.259196e-03
8.600000e+04 3.344677e+01 5.243956e-03
8.650000e+04 3.344692e+01 5.228882e-03
8.700000e+04 3.344680e+01 5.213822e-03
8.750000e+04 3.344745e+01 5.198685e-03
8.800000e+04 3.344757e+01 5.183810e-03
8.850000e+04 3.344768e+01 5.169171e-03
8.900000e+04 3.344788e+01 5.154574e-03
8.950000e+04 3.344715e+01 5.140136e-03
9.000000e+04 3.344682e+01 5.125559e-03
learned_free_energy_estimation/lammps.dat
+129
View File
@@ -0,0 +1,129 @@
# Simulate a WCA solute particle immersed into a Lennard-Jones solvent.
#
# We first generate an initial configuration of particles located at random
# positions. We then minimise the system and equilibrate it to the target
# temperature. During the production run, we sample configurations at regular
# intervals.
################################################################################
# Specify input parameters.
################################################################################
# We use the same box dimensions and temperature as in Jarzynski (2002),
# Phys. Rev. E 65, 046122, but convert these quantities to reduced units.
# T=300 K, eps=0.1854 kcal/mol, k_boltzmann=0.001987204 kcal/mol/K.
variable solute_radius equal 9.2/3.542
variable box_length equal 22.28/3.542
variable temperature_reference equal 300*0.001987204/0.1854
variable temperature_damp equal 0.5
variable seed equal 1234
variable num_solvent_particles equal 125
variable dump_frequency equal 5000
variable timestep equal 0.002
variable runtime_equi equal 50000
variable runtime_prod equal 500000
variable steps_equilibration equal floor(${runtime_equi}/${timestep})
variable steps_production equal floor(${runtime_prod}/${timestep})
################################################################################
# Create solvent particles at random positions and place solute at the origin.
################################################################################
units lj
atom_style atomic
dimension 3
boundary p p p
variable box_length_half equal ${box_length}/2.0
region domain block -${box_length_half} ${box_length_half} &
-${box_length_half} ${box_length_half} &
-${box_length_half} ${box_length_half} &
units box
create_box 2 domain
create_atoms 2 random 1 ${seed} domain
create_atoms 1 random ${num_solvent_particles} ${seed} domain
set type 2 x 0.0 y 0.0 z 0.0
group solute type 2
group solvent type 1
# Initialise particle masses.
mass 1 1
mass 2 10000000
################################################################################
# Define pair-style.
################################################################################
variable sigma_solute equal ${solute_radius}
variable radius_cutoff_solute equal 1.12246*${sigma_solute}
variable radius_cutoff equal ${box_length}/2.0
pair_style lj/cut ${radius_cutoff}
# Lennard-Jones interaction between solvent particles.
pair_coeff 1 1 1.0 1.0
# WCA interaction between solute and solvent particles.
pair_coeff 1 2 1.0 ${sigma_solute} ${radius_cutoff_solute}
# No interaction between solute particles.
pair_coeff 2 2 0.0 0.0
# Shift energies to be zero at the cutoff and update neighbor list settings.
pair_modify shift yes
neighbor 0.3 bin
neigh_modify delay 5
################################################################################
# Perform minimisation.
################################################################################
fix freeze solute setforce 0.0 0.0 0.0
minimize 1.0e-4 1.0e-6 100 1000
################################################################################
# Perform equilibration run.
################################################################################
reset_timestep 0
compute temperature_solvent solvent temp
compute kinetic_energy all ke
compute potential_energy all pe
variable total_energy equal c_kinetic_energy+c_potential_energy
# Compute centre-of-mass velocity and solute position for monitoring.
variable vcmx equal "vcm(all,x)"
variable vcmy equal "vcm(all,y)"
variable vcmz equal "vcm(all,z)"
variable vcm2 equal v_vcmx*v_vcmx+v_vcmy*v_vcmy+v_vcmz*v_vcmz
compute position_solute solute com
# Specify terminal output.
thermo_style custom step temp c_temperature_solvent c_potential_energy &
c_kinetic_energy v_total_energy press v_vcm2 &
c_position_solute[1] c_position_solute[2] c_position_solute[3]
thermo_modify norm no
thermo 1000
# Specify integration timestep.
timestep ${timestep}
# Apply Langevin thermostat to solvent particles.
fix flangevin solvent langevin ${temperature_reference} &
${temperature_reference} ${temperature_damp} ${seed} zero yes
fix fnve all nve
run ${steps_equilibration}
################################################################################
# Production run.
################################################################################
reset_timestep 0
dump fDumpTrajectory all custom ${dump_frequency} &
trajectory.dat id type x y z
dump_modify fDumpTrajectory sort id format float %.8g
fix fDumpEnergy all ave/time ${dump_frequency} 1 ${dump_frequency} &
c_potential_energy file energy.dat format %.8g
run ${steps_production}