diff --git a/README.md b/README.md index a3dd85dc4..4ad11df2c 100644 --- a/README.md +++ b/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) diff --git a/docs/dev_guide/arch.md b/docs/dev_guide/arch.md index 30f5f5b4a..c2f7b7661 100644 --- a/docs/dev_guide/arch.md +++ b/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, diff --git a/docs/dev_guide/convention.md b/docs/dev_guide/convention.md index 6aa83f78a..d4d18f455 100644 --- a/docs/dev_guide/convention.md +++ b/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 \ No newline at end of file diff --git a/docs/dev_guide/write_docs.md b/docs/dev_guide/write_docs.md index 3c6019c1d..bb6e59b0b 100644 --- a/docs/dev_guide/write_docs.md +++ b/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: diff --git a/docs/user_guide/introduction.md b/docs/user_guide/introduction.md index fc8033ec1..b78a08c96 100644 --- a/docs/user_guide/introduction.md +++ b/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) diff --git a/docs/user_guide/theory.md b/docs/user_guide/theory.md index 8a9e06e9f..720476852 100644 --- a/docs/user_guide/theory.md +++ b/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{ diff --git a/docs/user_guide/usage.md b/docs/user_guide/usage.md new file mode 100644 index 000000000..972ff106d --- /dev/null +++ b/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. +``` +.potential_fn at 0x112504af0> +.potential_fn at 0x1124cd820> +.potential_fn at 0x18509b790> +.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 ] +``` \ No newline at end of file diff --git a/docs/user_guide/xml_spec.md b/docs/user_guide/xml_spec.md index c9ac7c747..5304e3141 100644 --- a/docs/user_guide/xml_spec.md +++ b/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. diff --git a/mkdocs.yml b/mkdocs.yml index 6833919a7..ea1b37ba5 100644 --- a/mkdocs.yml +++ b/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