Merge branch 'wangxy/v1.0.0-devel' of https://github.com/KuangYu/DMFF …

…into pr/130

README.md

@@ -1,6 +1,6 @@
 # DMFF
 
-[![doi:10.26434/chemrxiv-2022-2c7gv](https://img.shields.io/badge/DOI-10.26434%2Fchemrxiv--2022--2c7gv-blue)](https://doi.org/10.26434/chemrxiv-2022-2c7gv)
+[![DOI: 10.1021/acs.jctc.2c01297](https://img.shields.io/badge/DOI-10.1021%2Facs.jctc.2c01297-blue)](https://pubs.acs.org/doi/10.1021/acs.jctc.2c01297)
 
 ## About DMFF
 
@@ -10,24 +10,31 @@ The behavior of organic molecular systems (e.g., protein folding, polymer struct
 
 ### License and credits
 
-The project DMFF is licensed under [GNU LGPL v3.0](LICENSE). If you use this code in any future publications, please cite this using `Wang X, Li J, Yang L, Chen F, Wang Y, Chang J, et al. DMFF: An Open-Source Automatic
-Differentiable Platform for Molecular Force Field
-Development and Molecular Dynamics
-Simulation. ChemRxiv. Cambridge: Cambridge Open Engage; 2022;  This content is a preprint and has not been peer-reviewed.`
+The project DMFF is licensed under [GNU LGPL v3.0](LICENSE). If you use this code in any future publications, please cite this using `Xinyan Wang, Jichen Li, Lan Yang, Feiyang Chen, Yingze Wang, Junhan Chang, Junmin Chen, Wei Feng, Linfeng Zhang, and Kuang Yu 
+Journal of Chemical Theory and Computation 2023 19 (17), 5897-5909
+DOI: 10.1021/acs.jctc.2c01297`
 
 ## User Guide
 
-+ [1. Introduction](docs/user_guide/introduction.md)
-+ [2. Installation](docs/user_guide/installation.md)
-+ [3. Basic usage](docs/user_guide/usage.md)
-+ [4. XML format force field](docs/user_guide/xml_spec.md)
-+ [5. Theory](docs/user_guide/theory.md)
++ [1. Introduction](docs/user_guide/1.introduction.md)
++ [2. Installation](docs/user_guide/2.installation.md)
++ [3. Basic Usage](docs/user_guide/3.usage.md)
++ [4. Modules](docs/user_guide/4.modules.md)
+  + [Classical](docs/user_guide/4.1classical.md)
+  + [ADMP](docs/user_guide/4.2ADMPPmeForce.md)
+  + [Qeq](docs/user_guide/4.3ADMPQeqForce.md)
+  + [Machine Learning](docs/user_guide/4.4MLForce.md)
+  + [Optimization](docs/user_guide/4.5Optimization.md)
+  + [Mbar Estimator](docs/user_guide/4.6MBAR.md)
+  + [OpenMM Plugin](docs/user_guide/4.7OpenMMplugin.md)
++ [5. Advanced examples](docs/user_guide/DMFF_example.ipynb)
 
 ## Developer Guide
 + [1. Introduction](docs/dev_guide/introduction.md)
 + [2. Software architecture](docs/dev_guide/arch.md)
 + [3. Coding conventions](docs/dev_guide/convention.md)
 + [4. Document writing](docs/dev_guide/write_docs.md)
++ [5. An example for developing: how to write a generator?](docs/dev_guide/generator.ipynb)
 
 ## Code Structure
 
@@ -38,11 +45,15 @@ The code is organized as follows:
 + `package`: files for constructing packages or images, such as conda recipe and docker files.
 + `tests`: unit tests.
 + `dmff`: DMFF python codes
++ `dmff/api`: source code of application programming interface of DMFF.
 + `dmff/admp`: source code of automatic differentiable multipolar polarizable (ADMP) force field module.
 + `dmff/classical`: source code of classical force field module.
 + `dmff/common`: source code of common functions, such as neighbor list.
-+ `dmff/generators`: source code of force generators.
 + `dmff/sgnn`: source of subgragh neural network force field model.
++ `dmff/eann`: source of embedded atom neural network force field model.
++ `dmff/generators`: source code of force generators.
++ `dmff/operators`: source code of operators.
+
 
 ## Support and Contribution
 

docs/dev_guide/introduction.md

@@ -19,3 +19,6 @@ In the *Implementation of New Potentials* section, we will show how to implement
 In the *Document Writing* section, we will talk about how to write docs for your new module. DMFF is a collection of force field calculators, and each force field calculator has its own parameters and may be invoked in different ways. Therefore, to make it easy to use and easy to maintain, the developers are required to document the theroy behind the code and the user interface of the code before PR. 
 
 In the *Checklist before PR* section is what you should do before you submit your PR to Github. In this section, you will learn how to write unit tests, check format, and add proper commentsin your code.
+
+Finally, we will provide a case study, guiding you step-by-step on *how to write a generator*.
++ [An example for developing: how to write a generator?](./generator.ipynb)

docs/index.md

@@ -1,7 +1,7 @@
 # DMFF
 
-[![doi:10.26434/chemrxiv-2022-2c7gv](https://img.shields.io/badge/DOI-10.26434%2Fchemrxiv--2022--2c7gv-blue)](https://doi.org/10.26434/chemrxiv-2022-2c7gv)
-
+[![DOI: 10.1021/acs.jctc.2c01297](https://img.shields.io/badge/DOI-10.1021%2Facs.jctc.2c01297-blue)](https://pubs.acs.org/doi/10.1021/acs.jctc.2c01297)
+ 
 ## About DMFF
 
 **DMFF** (**D**ifferentiable **M**olecular **F**orce **F**ield) is a Jax-based python package that provides a full differentiable implementation of molecular force field models. This project aims to establish an extensible codebase to minimize the efforts in force field parameterization, and to ease the force and virial tensor evaluations for advanced complicated potentials (e.g., polarizable models with geometry-dependent atomic parameters). Currently, this project mainly focuses on the molecular systems such as: water, biological macromolecules (peptides, proteins, nucleic acids), organic polymers, and small organic molecules (organic electrolyte, drug-like molecules) etc. We support both the conventional point charge models (OPLS and AMBER like) and multipolar polarizable models (AMOEBA and MPID like). The entire project is backed by the XLA technique in JAX, thus can be "jitted" and run in GPU devices much more efficiently compared to normal python codes.
@@ -10,38 +10,49 @@ The behavior of organic molecular systems (e.g., protein folding, polymer struct
 
 ### License and credits
 
-The project DMFF is licensed under [GNU LGPL v3.0](LICENSE). If you use this code in any future publications, please cite this using `Wang X, Li J, Yang L, Chen F, Wang Y, Chang J, et al. DMFF: An Open-Source Automatic
-Differentiable Platform for Molecular Force Field
-Development and Molecular Dynamics
-Simulation. ChemRxiv. Cambridge: Cambridge Open Engage; 2022;  This content is a preprint and has not been peer-reviewed.
+The project DMFF is licensed under [GNU LGPL v3.0](LICENSE). If you use this code in any future publications, please cite this using `Xinyan Wang, Jichen Li, Lan Yang, Feiyang Chen, Yingze Wang, Junhan Chang, Junmin Chen, Wei Feng, Linfeng Zhang, and Kuang Yu 
+Journal of Chemical Theory and Computation 2023 19 (17), 5897-5909
+DOI: 10.1021/acs.jctc.2c01297`
 
 ## User Guide
 
-+ [1. Introduction](user_guide/introduction.md)
-+ [2. Installation](user_guide/installation.md)
-+ [3. Basic usage](user_guide/usage.md)
-+ [4. XML format force field](user_guide/xml_spec.md)
-+ [5. Theory](user_guide/theory.md)
++ [1. Introduction](./user_guide/1.introduction.md)
++ [2. Installation](./user_guide/2.installation.md)
++ [3. Basic Usage](./user_guide/3.usage.md)
++ [4. Modules](./user_guide/4.modules.md)
+  + [Classical](./user_guide/4.1classical.md)
+  + [ADMP](./user_guide/4.2ADMPPmeForce.md)
+  + [Qeq](./user_guide/4.3ADMPQeqForce.md)
+  + [Machine Learning](./user_guide/4.4MLForce.md)
+  + [Optimization](./user_guide/4.5Optimization.md)
+  + [Mbar Estimator](./user_guide/4.6MBAR.md)
+  + [OpenMM Plugin](./user_guide/4.7OpenMMplugin.md)
++ [5. Advanced examples](./user_guide/DMFF_example.ipynb)
 
 ## Developer Guide
-+ [1. Introduction](dev_guide/introduction.md)
-+ [2. Software architecture](dev_guide/arch.md)
-+ [3. Coding conventions](dev_guide/convention.md)
-+ [4. Document writing](dev_guide/write_docs.md)
++ [1. Introduction](./dev_guide/introduction.md)
++ [2. Software architecture](./dev_guide/arch.md)
++ [3. Coding conventions](./dev_guide/convention.md)
++ [4. Document writing](./dev_guide/write_docs.md)
++ [5. An example for developing: how to write a generator?](./dev_guide/generator.ipynb)
 
 ## Code Structure
+
 The code is organized as follows:
 
 + `examples`: demos presented in Jupyter Notebook.
 + `docs`: documentation.
 + `package`: files for constructing packages or images, such as conda recipe and docker files.
 + `tests`: unit tests.
 + `dmff`: DMFF python codes
++ `dmff/api`: source code of application programming interface of DMFF.
 + `dmff/admp`: source code of automatic differentiable multipolar polarizable (ADMP) force field module.
 + `dmff/classical`: source code of classical force field module.
 + `dmff/common`: source code of common functions, such as neighbor list.
-+ `dmff/generators`: source code of force generators.
 + `dmff/sgnn`: source of subgragh neural network force field model.
++ `dmff/eann`: source of embedded atom neural network force field model.
++ `dmff/generators`: source code of force generators.
++ `dmff/operators`: source code of operators.
 
 
 ## Support and Contribution

docs/user_guide/1.introduction.md

@@ -0,0 +1,22 @@
+# 1. Introduction
+
+In this user guide, you will learn:
+
+- [DMFF Installation](./2.installation.md) 
+- [Basic Usage](./3.usage.md) of DMFF, including how to compute energy, forces and parametric gradients
+- [New modules release](./4.modules.md)
+- [Advanced examples](./DMFF_example.ipynb) of DMFF, which can help you quickly get started with DMFF
+
+
+The first thing you should know is that DMFF is not an actual force field model (such as OPLS or AMBER), but a differentiable implementation of various force field (or "potential") functional forms. It contains following modules:
+
+- Classical module: implements classical force fields (OPLS or GAFF like potentials)
+- ADMP module: Automatic Differentiable Multipolar Polarizable potential (MPID like potentials)
+- Qeq module: supports to coulombic energy calculation for constant potential model and constant charge model.
+- ML module: Machine Learning force field include sgnn and eann, implementing subgragh neural network model for intramolecular interactions
+- Optimization module： Implements automatic optimization of force field parameters.
+- MBAR Estimator module: Achieves differentiable estimation of thermodynamic quantities and trajectory reweighting through MBAR.
+- OpenMM DMFF Plugin module
+
+Each module implements a particular form of force field, which takes a unique set of input parameters, usually provided in a XML file, or other feature you need. With DMFF, one can easily compute the energy as well as energy gradients including: forces, virial tensors, and gradients with respect to force field parameters etc.
+

docs/user_guide/installation.md → docs/user_guide/2.installation.md

@@ -3,13 +3,14 @@
 + Create conda environment:
 ```
 conda create -n dmff python=3.9 --yes
+conda activate dmff
 ```
 + Install [jax](https://github.com/google/jax) (select the correct cuda version, see more details in the Jax installation guide):
 ```bash
 # CPU version
-pip install "jax[cpu]==0.4.19"
+pip install "jax[cpu]==0.4.14"
 # GPU version
-pip install "jax[cuda11_local]==0.4.19" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install "jax[cuda11_local]==0.4.14" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 + Install [mdtraj](https://github.com/mdtraj/mdtraj), [optax](https://github.com/deepmind/optax), [jaxopt](https://github.com/google/jaxopt) and [pymbar](https://github.com/choderalab/pymbar):
 ```bash
@@ -25,12 +26,9 @@ conda install -c conda-forge openmm==7.7.0
 conda install -c conda-forge rdkit
 ```
 ## 2.2 Install DMFF from Source Code
-One can download the DMFF source code from github:
+One can download the DMFF source code from github and install it using `pip`. :
 ```bash
 git clone https://github.com/deepmodeling/DMFF.git
-```
-Then you may install DMFF using `pip`:
-```bash
 cd DMFF
 pip install . --user
 ```

docs/user_guide/3.usage.md

@@ -0,0 +1,106 @@
+# 3. Basic usage
+This chapter will introduce some basic usage of DMFF. All scripts can be found in `examples/` directory in which Jupyter notebook-based demos are provided. 
+## 3.1 Compute energy
+DMFF uses OpenMM to parse input files, including coordinates file, topology specification file and force field parameter file. Then, the core class `Hamiltonian` inherited from `openmm.ForceField` will be initialized and the method `createPotential` will be called to create differentiable potential energy functions for different energy terms. Take parametrzing an organic moleclue with GAFF2 force field as an example:
+```python
+import jax
+import jax.numpy as jnp
+import openmm.app as app
+import openmm.unit as unit
+from dmff import Hamiltonian, NeighborList
+
+app.Topology.loadBondDefinitions("lig-top.xml")
+pdb = app.PDBFile("lig.pdb")
+ff = Hamiltonian("gaff-2.11.xml", "lig-prm.xml")
+potentials = ff.createPotential(pdb.topology)
+for k in potentials.dmff_potentials.keys():
+    pot = potentials.dmff_potentials[k]
+    print(pot)
+```
+In this example, `lig.pdb` is the PDB file containing atomic coordinates, and `lig-top.xml` specifying bond connections within a molecule and this information is required by `openmm.app` to generate molecular topology. Note that this file is not always required, if bond conncections are defined in .pdb file by `CONNECT` keyword. `gaff-2.11.xml` contains GAFF2 force field parameters (bonds, angles, torsion and vdW), and `lig-prm.xml` contains atomic partial charges (GAFF2 requests a user-defined charge assignment process). This xml format is compatitable with OpenMM definitions, and a detailed description can be found in [OpenMM user guide](`http://docs.openmm.org/latest/userguide/application/05_creating_ffs.html`) or [XML-format force fields](./xml_spec.md) section.
+
+If you run this script in `examples/classical`, you will get the following output.
+```
+<function HarmonicBondJaxGenerator.createForce.<locals>.potential_fn at 0x112504af0>
+<function HarmonicAngleJaxGenerator.createForce.<locals>.potential_fn at 0x1124cd820>
+<function PeriodicTorsionJaxGenerator.createForce.<locals>.potential_fn at 0x18509b790>
+<function NonbondJaxGenerator.createForce.<locals>.potential_fn at 0x18509baf0>
+```
+The force field parameters are stored as a Python dict in the Hamiltonian.
+```python
+params = ff.getParameters()
+nbparam = params['NonbondedForce']
+nbparam
+```
+
+```
+{
+    'sigma': DeviceArray([0.33152124, ...], dtype=float32),
+    'epsilon': DeviceArray([0.4133792, ...], dtype=float32),
+    'epsfix': DeviceArray([], dtype=float32),
+    'sigfix': DeviceArray([], dtype=float32),
+    'charge': DeviceArray([-0.75401515, ...], dtype=float32),
+    'coulomb14scale': DeviceArray([0.8333333], dtype=float32),
+    'lj14scale': DeviceArray([0.5], dtype=float32)
+}
+```
+
+
+Each generated function will read coordinates, box, pairs and force field parameters as inputs.
+```python
+positions = jnp.array(pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer))
+box = jnp.array([
+    [ 1.0,  0.0,  0.0],
+    [ 0.0,  1.0,  0.0],
+    [ 0.0,  0.0,  1.0]
+])
+nbList = NeighborList(box, rc=4)
+nbList.allocate(positions)
+pairs = nbList.pairs
+```
+Note that in order to take advantages of the auto-differentiable implementation in JAX, the input arrays have to be `jax.numpy.ndarray`, otherwise DMFF will raise an error. `pairs` is a $N\times 2$ integer array in which each row specifying atoms condsidered as neighbors within `rcut`. As shown above, this can be calculated with `dmff.NeighborList` class which is supported by `jax_md`. 
+
+The potential energy function will give energy (a scalar, in kJ/mol) as output:
+```python
+nbfunc = potentials.dmff_potentials['NonbondedForce']
+nbene = nbfunc(positions, box, pairs, params)
+print(nbene)
+```
+If everything works fine, you will get `-425.41412` as a result. In addition, you can also use `getPotentialFunc()` and `getParameters()` to obtain the whole potential energy function and force field parameter set, instead of seperated functions for different energy terms.
+```python
+efunc = potentials.getPotentialFunc()
+params = ff.getParameters()
+totene = efunc(positions, box, pairs, params)
+```
+
+## 3.2 Compute forces
+Different from conventional programming frameworks, explicit definition of atomic force calculation functions are no longer needed. Instead, the forces can be evaluated in an automatic manner with `jax.grad`.
+```
+pos_grad_func = jax.grad(efunc, argnums=0)
+force = -pos_grad_func(positions, box, pairs, params)
+```
+
+## 3.3 Compute parametric gradients
+Similarly, the derivative of energy with regard to force field parameters can also be computed easily.
+```python
+param_grad_func = jax.grad(nbfunc, argnums=-1)
+pgrad = param_grad_func(positions, box, pairs, params)
+print(pgrad["NonbondedForce"]["charge"])
+```
+
+```python
+[ 652.7753      55.108738   729.36115   -171.4929     502.70837
+  -44.917206   129.63994   -142.31796   -149.62088    453.21503
+   46.372574   140.15303    575.488      461.46902    294.4358
+  335.25153     27.828705   671.3637     390.8903     519.6835
+  220.51129    238.7695     229.97302    210.58838    231.8734
+  196.40994    237.08563     35.663574   457.76416     77.4798
+  256.54382    402.2121     611.9573     440.8465     -52.09662
+  421.86688    592.46265    237.98883    110.286194   150.65375
+  218.61087    240.20477   -211.85376    150.7331     310.89404
+  208.65228   -139.23026   -168.8883     114.3645       3.7261353
+  399.6282     298.28455    422.06445    526.18463    521.27563
+  575.85767    606.74744    394.40845    549.84033    556.4724
+  485.1427     512.1267     558.55896    560.4667     562.812
+  333.74194  ]
+```