Docs refine

README.md

@@ -8,15 +8,15 @@ The behavior of organic molecular systems (e.g., protein folding, polymer struct
 
 + [1. Introduction](docs/user_guide/introduction.md)
 + [2. Installation](docs/user_guide/installation.md)
-+ [3. Compute energy and forces](docs/user_guide/compute.md)
-+ [4. Compute gradients with auto differentiable framework](docs/user_guide/auto_diff.md)
-+ [5. Theories](docs/user_guide/theory.md)
-+ [6. Introduction to force field xml files](docs/user_guide/xml_spec.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)
 
 ## Developer Guide
 + [1. Introduction](docs/dev_guide/introduction.md)
-+ [2. Architecture](docs/dev_guide/arch.md)
-+ [3. Convention](docs/dev_guide/convention.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)
 
 ## Modules
 + [1. ADMP](docs/modules/admp.md)

docs/dev_guide/arch.md

@@ -1,4 +1,4 @@
-# Architecture of DMFF
+# 2. Software architecture
 
 ![arch](../assets/arch.png)
 
@@ -19,7 +19,7 @@ The backend module is usually an automatically differentiable calculator built w
 
 The structures of the frontend and the backend modules will be introduced in detail in below.
 
-## How Frontend Works
+## 2.1 Frontend
 
 Frontend modules are stored in `api.py`. `Hamiltonian` class is the top-level class exposed to users by DMFF. 
 `Hamiltonian` class reads the path of the XML file, parses the XML file, and calls different frontend modules 
@@ -38,7 +38,7 @@ and the rounded box represents the internal operation logic of OpenMM when execu
 
 ![openmm_workflow](../assets/opemm_workflow.svg)
 
-### Hamiltonian Class
+### 2.1.1 Hamiltonian Class
 
 The `Hamiltonian` class is the top-level frontend module, which inherits the 
 [forcefield class](https://github.com/openmm/openmm/blob/master/wrappers/python/openmm/app/forcefield.py) in OpenMM. 
@@ -74,7 +74,7 @@ corresponding `parseElement` method, then calls it to initialize the `generator`
 parameters from the XML file. You can access all the generators by the `getGenerators()` method in Hamiltonian. 
 
 
-### Generator Class
+### 2.1.2 Generator Class
 
 
 The generator class is in charge of input file analysis, molecular topology construction, atom classification, 
@@ -324,9 +324,9 @@ finfo = XMLNodeInfo("HarmonicBondForce")
 ```
 
 
-## How Backend Works
+## 2.2 Backend
 
-### Force Class
+### 2.2.1 Force Class
 
 Force class is the backend module that wraps the calculator function. 
 It does not rely on OpenMM and can be very flexible. For instance, 

docs/dev_guide/convention.md

@@ -1,22 +1,26 @@
-# Code Convention
+# 3. Coding conventions
 
 In this section, you will learn:
-    - How is DMFF organized
-    - 
 
-## code organization
+- How DMFF is organized
+- Programming style recommended to follow in DMFF development
+
+## 3.1 Code Organization
 
 The root directory of DMFF has following sub-directories:
 
-    - `dmff`: source code of project
-    - `docs`: documents in markdown
-    - `examples`: examples can be run independently
-    - `tests`: unit and integration tests
+- `dmff`: source code of project
+- `docs`: documents in markdown
+- `examples`: examples can be run independently
+- `tests`: unit and integration tests
 
 Under the `dmff`, there are several files and sub-directory:
 
-    - `api.py`: store all the frontend modules
-    - `settings.py`: global settings 
-    - `utils.py`: helper functions
-    — each sub-directory represents a set of potential form, e.g. `admp` is Automatic Differentiable Multipolar Polarizable, `classical` is differentiable GAFF forcefield.
+- `api.py`: API (frontend modules)
+- `settings.py`: global settings 
+- `utils.py`: basic functions
+- each sub-directory represents a set of potential form, e.g. `admp` stands for Automatic Differentiable Multipolar Polarizable force field, and `classical` is the differentiable implementation of classical fixed-charge force field.
+
+## 3.2 Programming Style
 
+TBA

docs/dev_guide/write_docs.md

@@ -1,10 +1,10 @@
-# Write related docs
+# 4. Document writing
 
 The most important thing after you implement your code in DMFF is to write the docs. Good documentation can help users use your module correctly and allow other maintainers to improve functions better.
 
 The documentation of DMFF use [MKDocs](https://www.mkdocs.org/) framework. 
 
-## install MKDocs
+## 4.1 Install MKDocs
 
 Before we start to write our docs, run the following command in termianl to install MKDocs:
 
@@ -18,11 +18,11 @@ To support latex rendering, we also need markdown extension:
 pip install pymdown-extensions
 ```
 
-## Write your docs
+## 4.2 Write your docs
 
 According to the existing document architecture, create a new markdown file in the appropriate directory. Write it! If you need to insert picture, upload the picture in `assets` directory, and use `![_placeholder](relative/path/to/this/file)` syntax to insert picture.
 
-## Preview you docs
+## 4.3 Preview you docs
 
 MkDocs comes with a built-in dev-server that lets you preview your documentation as you work on it. Make sure you're in the same directory as the `mkdocs.yml` configuration file, and then start the server by running the mkdocs serve command:
 

docs/user_guide/introduction.md

@@ -1,13 +1,13 @@
 # 1. Introduction
 
-In this user guide, you will learn how to:
+In this user guide, you will learn:
 
-- Install DMFF
-- Compute energy and force
-- Auto differentiate potential functions
-- Couple DMFF with MD engine
+- [DMFF Installation](./installation.md) 
+- [Basic usage](./usage.md) of DMFF, including how to compute energy, forces and parametric gradients
+- [How to write XML format force field](./xml_spec.md)
+- [Theoretical background](./theory.md) of various force field models 
 
-The first thing you should know is that DMFF is not an actual force field (such as OPLS or AMBER), but a differentiable implementation of various force field (or "potential") functional forms. It contains many modules:
+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:
 
 - ADMP module: Automatic Differentiable Multipolar Polarizable potential (MPID like potentials)
 - Classical module: implements classical force fields (OPLS or GAFF like potentials)

docs/user_guide/theory.md

@@ -1,4 +1,4 @@
-# Theoretical background
+# 5. Theory
 
 DMFF project aims to implement organic molecular force fields using a differentiable programming framework, such that derivatives with respect to atomic positions, box shape, and force field parameters can be easily computed. It contains different modules, dealing with different types of force field terms. Currently, there are two primary modules:
 
@@ -23,7 +23,7 @@ where $q_i$ is the charge of atom $i$.
 More complex (and supposedly more accurate) force field can be obtained by including more multipoles with higher orders. Some force fields, such as MPID, goes as high as octupoles. Currently in DMFF, we support up to quadrupoles:
 
 $$
-V=\sum_{tu} \hat{Q}_t^A T^{AB}_{tu} \hat{Q}_u^B
+V=\sum_{tu} Q_t^A T^{AB}_{tu} Q_u^B
 $$
 
 where $Q_t^A$ represents the t-component of the multipole moment of atom A. Note there are two (equivalent) ways to define multipole moments: cartesian and spherical harmonics. Cartesian representation is over-complete but with a simpler definition, while spherical harmonics are easier to use in real calculations. In the user API, we use cartesian representation, in consistent with the AMOEBA and the MPID plugins in OpenMM. However, spherical harmonics are always used in the computation kernel, and we assume all components are arranged in the following order:
@@ -46,7 +46,7 @@ Different to charges, the definition of multipole moments depends on the coordin
 
 ## Polarization Interaction
 
-DMFF supports polarizable force fields, in which the dipole moment of the atom can respond to the change of the external electric field. In practice, each atom has not only permanent multipoles $\hat{Q}_t$, but also induced dipoles $U_{ind}$. The induced dipole-induced dipole and induced dipole-permanent multipole interactions needs to be damped at short-range to avoid polarization catastrophe. In DMFF, we use the Thole damping scheme identical to MPID (ref 6), which introduces a damping width ($a_i$) for each atom $i$. The damping function is then computed and applied to the corresponding interaction tensor. Taking $U_{ind}$-permanent charge interaction as an example, the definition of damping function is:
+DMFF supports polarizable force fields, in which the dipole moment of the atom can respond to the change of the external electric field. In practice, each atom has not only permanent multipoles $Q_t$, but also induced dipoles $U_{ind}$. The induced dipole-induced dipole and induced dipole-permanent multipole interactions needs to be damped at short-range to avoid polarization catastrophe. In DMFF, we use the Thole damping scheme identical to MPID (ref 6), which introduces a damping width ($a_i$) for each atom $i$. The damping function is then computed and applied to the corresponding interaction tensor. Taking $U_{ind}$-permanent charge interaction as an example, the definition of damping function is:
 
 $$
 \displaylines{

docs/user_guide/usage.md

@@ -0,0 +1,103 @@
+# 3. Basic usage
+## 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 pot in potentials:
+    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 `param` attribute of force generators.
+```python
+nbparam = ff.getGenerators()[3].params
+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([
+    [10.0,  0.0,  0.0],
+    [ 0.0, 10.0,  0.0],
+    [ 0.0,  0.0, 10.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[3]
+nbene = nbfunc(positions, box, pairs, nbparam)
+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 = ff.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.
+```
+param_grad_func = jax.grad(nbfunc, argnums=-1)
+pgrad = param_grad_func(positions, box, pairs, nbparam)
+print(pgrad["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  ]
+```

docs/user_guide/xml_spec.md

@@ -1,4 +1,4 @@
-# How to write XML file
+# 4. XML format force field
 
 The design of openmm force field file is quite modular and convenient to use. Unfortunately, only limited documentations are available right now to explain the details of the file format. Here, the format and the meaning of the OpenMM XML file are sorted in details in below.
 

mkdocs.yml

@@ -4,16 +4,16 @@ nav:
   - User Guide:
     - 1. Introduction: user_guide/introduction.md
     - 2. Installation: user_guide/installation.md
-    - 3. Compute energy and force: user_guide/compute.md
-    - 4. Auto-diff: user_guide/auto_diff.md
-    - 5. Theory: user_guide/theory.md
-    - 6. couple with MD: user_guide/couple.md
-    - 7. Introduction to force field xml file: user_guide/xml_spec.md
-
+    - 3. Basic usage: user_guide/usage.md
+    - 4. XML format force field: user_guide/xml_spec.md
+    - 5. Theory: user_guide/theory.md  
+
   - Developer Guide:
     - 1. Introduction: dev_guide/introduction.md
-    - 2. Architecture: dev_guide/arch.md
-    - 3. Convention: dev_guide/convention.md
+    - 2. Software architecture: dev_guide/arch.md
+    - 3. Coding conventions: dev_guide/convention.md
+    - 4. Document Writing: dev_guide/write_docs.md
+
   - Modules:
     - ADMP: 
       - Introduction: admp/readme.md