Release DMFF v0.2.0 (#71)

* feat(classical): offer tools to organize xml forcefield and match parameters. * fix(tests): change API to be consistent with dmff * feat(dmff): Add class Potential to keep OpenMM and DMFF potentials * fix(api): auto-update parameter before writing new xml file. * refactor: admp api.py using old xmlholder api * fix(fftree): make the name of API easier to be understand * update * update: admp related api * fix: fix bug in admp api * doc: programming style convention about typing and numpy style docstring * fix(tests): Add attribute to choose improper order * fix(classical): remove double check of cutoff distance in potential function. * fix(classical): avoid creating empty dict in impr parameters * fix(inter): add 1e-12 as eps value in jnp.sqrt function * fix(inter): use more tiny eps for numerical consistent * update: frontend docs * fix typo: fix typo in dev_guide * Temprary commit before test * Debugging new api for admp * peg system benchmark checked out * Update examples * Fix bug: was not reading the XZ component of multipole * remove jit for nbl.update * Modify setup.py * Require jax version smaller than 0.3.16, which is incompatible with current version of jax_md (0.1.29) * Update demo notebook in water example * Update documents and the notebook in classical FF example * Fix bug in nblist module * change allocate * fix: fix nblist jit-related bug * fix: test sequence in test_nblist.py * Fix bug in Slater ISA short range interaction * update: all tests pass * clean up code * feat: brutal Freud-based nblist * Bug fix in parsing internal xmls (#62) 1. Add NeighborFreud class to use freud package to obtain neighborlist. 2. Fix bug of parsing openmm internal force field xmls 3. Handle unmatched torsion params * settle down dependencies version * rename `nmax` to `capacity_multiplier`(consistent with jax_md); remove property:distance and dr in neighborList due to catastrophic bug * add freud-analysis require * remove overflow judgement so update in nblist can be jitted * Adapt examples to new nblist api * add md_ipi in examples to run classical MD for bulk water * modified md_ipi * modified md_ipi * new * new * new_2 * Refactor: decompose api.py and generators to separate files * feat(MBAR): Add differentiable MBAR impl * bug fix: record the name of matched template * Add unit test for MBAR estimator * feat: Estimate free energy of an extra state * set specific version no. for jax in installation guide * Update MBAR Estimator API * Update unit test for the latest MBAR API * Update github workflow to support new unit test. * Update requirement of pymbar * fix package including problem * Change cell vector to prevent numerical problem (ceil(12.0 / 1.0) would jump between 12 and 13) * Update settings.py * remove unused imports * fix "==" in requirements * update (gitignore): hmtff cache * update (api): docstring in createPotential * Update optax transforms for force field parameters * Update __init__.py * Update LJ jax force API in LennardJonesForce generator * Let energy function in TargetState use the whole trajectory instead of using single frame * Update MBAR UT to fit API change * Remove the using of numpy to make free energy & weight differentiable * Update __init__.py * Set precision again in dmff.mbar after importing pymbar * Increase MBAR numerical stability * Update an example of MBAR-based optimization * Correct the typo in document * Add openmm sample state. * Update genOptimizer to support multiple optimizers * Set default pressure to be 0 (NVT) * Add Gitee_mirror (#67) * Fix Mirror CI/CD (#68) * Add Gitee_mirror * Fix mirror CI/CD Co-authored-by: WangXinyan940 <[email protected]> * Add tutorial_utils for demo usage (#69) * Hook jax force to generator for intra potentials * Remove the requirement of mdtraj * Let NeighborList.update return (N, 3) pairs with colv_map information * Update nblist.py * bugfix: correctly recognize if the LJForce card uses type or class * SMIRKS-based typing scheme (#72) * add support for SMIRKS, bcc and virtual site * add dependencies in ut workflow * add interface to obtain topmat * code clean * fix multiple matches_dict var * add tests for building top * add tests for smirks/bcc/vsite * add dockerfile for dev env * add examples of smirks * add usage of smirks in doc * rm api doc in _date/_version * fix typo * update (doc): architecture img * update (README): citation and code structure * update (doc): bad rendered latex * update (doc): typo in usage * update (doc): smirks in XML * Hotfix in BCC parametrization (#73) * update (smirks): change bcc err to warning * hotfix: incorrect bcc matching * Save covalent map to Potential object & make energy function generator (#74) * Update python version requirement to 3.8 * Save meta data in Potential object * Add support to save meta data in Hamiltonian * Create an attribute to save meta data * Use potential.meta["cov_map"] in unittests * Build a function generator to calculate energies of a trajectory * Change API function name * fix meta data for admp generators * Make buildEnergyFunction cleaner Co-authored-by: KuangYu <[email protected]> Co-authored-by: Yingze Wang <[email protected]> Co-authored-by: Roy-Kid <[email protected]> Co-authored-by: KuangYu <[email protected]> Co-authored-by: Jichen Li <[email protected]> Co-authored-by: crone <[email protected]> Co-authored-by: Yuzhi Zhang <[email protected]>
deepmodeling · Dec 1, 2022 · 2fa71cf · 2fa71cf
1 parent af4b1f2
commit 2fa71cf
Show file tree

Hide file tree

Showing 172 changed files with 74,244 additions and 31,648 deletions.
diff --git a/.github/workflows/mirror_gitee.yml b/.github/workflows/mirror_gitee.yml
@@ -0,0 +1,19 @@
+name: Mirror to Gitee Repo
+
+on: [ push, delete, create ]
+
+# Ensures that only one mirror task will run at a time.
+concurrency:
+  group: git-mirror
+
+jobs:
+  git-mirror:
+    if: github.repository_owner == 'deepmodeling'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: wearerequired/git-mirror-action@v1
+        env:
+          SSH_PRIVATE_KEY: ${{ secrets.SYNC_GITEE_PRIVATE_KEY }}
+        with:
+          source-repo: "https://github.com/deepmodeling/dmff.git"
+          destination-repo: "[email protected]:deepmodeling/DMFF.git"
diff --git a/.github/workflows/ut.yml b/.github/workflows/ut.yml
@@ -22,8 +22,9 @@ jobs:
           $CONDA/bin/conda update -n base -c defaults conda
           conda install pip
           conda update pip
-          conda install numpy openmm pytest -c conda-forge
+          conda install numpy openmm pytest rdkit biopandas openbabel -c conda-forge
           pip install jax jax_md
+          pip install mdtraj==1.9.7 pymbar==4.0.1
       - name: Install DMFF
         run: |
           source $CONDA/bin/activate

diff --git a/.gitignore b/.gitignore
@@ -781,4 +781,7 @@ FodyWeavers.xsd
 *.acpype/
 
 */_date.py
-*/_version.py
+*/_version.py
+
+# hmtff cache
+*.hmtff/
diff --git a/README.md b/README.md
@@ -1,9 +1,20 @@
 # 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)
+
+## 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.
 
 The behavior of organic molecular systems (e.g., protein folding, polymer structure, etc.) is often determined by a complex effect of many different types of interactions. The existing organic molecular force fields are mainly empirically fitted and their performance relies heavily on error cancellation. Therefore, the transferability and the prediction power of these force fields are insufficient. For new molecules, the parameter fitting process requires essential manual intervention and can be quite cumbersome. In order to automate the parametrization process and increase the robustness of the model, it is necessary to apply modern AI techniques in conventional force field development. This project serves for this purpose by utilizing the automatic differentiable programming technique to develop a codebase, which allows a more convenient incorporation of modern AI optimization techniques. It also helps the realization of many exciting functions including (but not limited to): hybrid machine learning/force field models and parameter optimization based on trajectory.
 
+### 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.`
+
 ## User Guide
 
 + [1. Introduction](docs/user_guide/introduction.md)
@@ -18,9 +29,20 @@ The behavior of organic molecular systems (e.g., protein folding, polymer struct
 + [3. Coding conventions](docs/dev_guide/convention.md)
 + [4. Document writing](docs/dev_guide/write_docs.md)
 
-## Modules
-+ [1. ADMP](docs/modules/admp.md)
+## 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/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.
 
 ## Support and Contribution
 

diff --git a/dmff/__init__.py b/dmff/__init__.py
@@ -1,3 +1,4 @@
 from .settings import *
-from .common.nblist import NeighborList
-from .api import Hamiltonian
+from .common.nblist import NeighborList, NeighborListFreud
+from .api import Hamiltonian
+from .generators import *
diff --git a/dmff/admp/disp_pme.py b/dmff/admp/disp_pme.py
@@ -17,8 +17,8 @@ class ADMPDispPmeForce:
     The so called "environment paramters" means parameters that do not need to be differentiable
     '''
 
-    def __init__(self, box, covalent_map, rc, ethresh, pmax, lpme=True):
-        self.covalent_map = covalent_map
+    def __init__(self, box, rc, ethresh, pmax, lpme=True):
+
         self.rc = rc
         self.ethresh = ethresh
         self.pmax = pmax
@@ -44,7 +44,7 @@ def __init__(self, box, covalent_map, rc, ethresh, pmax, lpme=True):
     def generate_get_energy(self):
         def get_energy(positions, box, pairs, c_list, mScales):
             return energy_disp_pme(positions, box, pairs, 
-                                  c_list, mScales, self.covalent_map,
+                                  c_list, mScales,
                                   self.kappa, self.K1, self.K2, self.K3, self.pmax,
                                   self.d6_recip, self.d8_recip, self.d10_recip, lpme=self.lpme)
         return get_energy
@@ -78,7 +78,7 @@ def refresh_calculators(self):
 
 
 def energy_disp_pme(positions, box, pairs,
-        c_list, mScales, covalent_map,
+        c_list, mScales,
         kappa, K1, K2, K3, pmax, 
         recip_fn6, recip_fn8, recip_fn10, lpme=True):
     '''
@@ -90,7 +90,7 @@ def energy_disp_pme(positions, box, pairs,
         box:
             3 * 3: box, axes arranged in row
         pairs:
-            Np * 2: interacting pair indices
+            Np * 3: interacting pair indices and topology distance
         c_list:
             Na * (pmax-4)/2: atomic dispersion coefficients
         mScales:
@@ -115,7 +115,7 @@ def energy_disp_pme(positions, box, pairs,
     if lpme is False:
         kappa = 0
 
-    ene_real = disp_pme_real(positions, box, pairs, c_list, mScales, covalent_map, kappa, pmax)
+    ene_real = disp_pme_real(positions, box, pairs, c_list, mScales, kappa, pmax)
 
     if lpme:
         ene_recip = recip_fn6(positions, box, c_list[:, 0, jnp.newaxis])
@@ -132,7 +132,7 @@ def energy_disp_pme(positions, box, pairs,
 
 def disp_pme_real(positions, box, pairs, 
         c_list, 
-        mScales, covalent_map, 
+        mScales, 
         kappa, pmax):
     '''
     This function calculates the dispersion real space energy
@@ -144,7 +144,7 @@ def disp_pme_real(positions, box, pairs,
         box:
             3 * 3: box, axes arranged in row
         pairs:
-            Np * 2: interacting pair indices
+            Np * 3: interacting pair indices and topology distance
         c_list:
             Na * (pmax-4)/2: atomic dispersion coefficients
         mScales:
@@ -162,16 +162,16 @@ def disp_pme_real(positions, box, pairs,
 
     # expand pairwise parameters
     # pairs = pairs[pairs[:, 0] < pairs[:, 1]]
-    pairs = regularize_pairs(pairs)
+    pairs = pairs.at[:, :2].set(regularize_pairs(pairs[:, :2]))
 
     box_inv = jnp.linalg.inv(box)
 
     ri = distribute_v3(positions, pairs[:, 0])
     rj = distribute_v3(positions, pairs[:, 1])
-    nbonds = covalent_map[pairs[:, 0], pairs[:, 1]]
+    nbonds = pairs[:, 2]
     mscales = distribute_scalar(mScales, nbonds-1)
 
-    buffer_scales = pair_buffer_scales(pairs)
+    buffer_scales = pair_buffer_scales(pairs[:, :2])
     mscales = mscales * buffer_scales
 
     ci = distribute_dispcoeff(c_list, pairs[:, 0])

diff --git a/dmff/admp/pairwise.py b/dmff/admp/pairwise.py
@@ -44,7 +44,7 @@ def distribute_dispcoeff(c_list, index):
 def distribute_matrix(multipoles,index1,index2):
     return multipoles[index1,index2]
 
-def generate_pairwise_interaction(pair_int_kernel, covalent_map, static_args):
+def generate_pairwise_interaction(pair_int_kernel, static_args):
     '''
     This is a calculator generator for pairwise interaction 
 
@@ -53,9 +53,6 @@ def generate_pairwise_interaction(pair_int_kernel, covalent_map, static_args):
             function type (dr, m, p1i, p1j, p2i, p2j) -> energy : the vectorized kernel function, 
             dr is the distance, m is the topological scaling factor, p1i, p1j, p2i, p2j are pairwise parameters
 
-        covalent_map:
-            Na * Na, int: the covalent_map matrix that marks the topological distances between atoms
-
         static_args:
             dict: a dictionary that stores all static global parameters (such as lmax, kappa, etc)
 
@@ -67,13 +64,14 @@ def generate_pairwise_interaction(pair_int_kernel, covalent_map, static_args):
     '''
 
     def pair_int(positions, box, pairs, mScales, *atomic_params):
-        pairs = regularize_pairs(pairs)
+        # pairs = regularize_pairs(pairs)
+        pairs = pairs.at[:, :2].set(regularize_pairs(pairs[:, :2]))
 
         ri = distribute_v3(positions, pairs[:, 0])
         rj = distribute_v3(positions, pairs[:, 1])
         # ri = positions[pairs[:, 0]]
         # rj = positions[pairs[:, 1]]
-        nbonds = covalent_map[pairs[:, 0], pairs[:, 1]]
+        nbonds = pairs[:, 3]
         mscales = distribute_scalar(mScales, nbonds-1)
 
         buffer_scales = pair_buffer_scales(pairs)
@@ -114,8 +112,7 @@ def TT_damping_qq_c6_kernel(dr, m, ai, aj, bi, bj, qi, qj, ci, cj):
     exp_br = jnp.exp(-br)
     f = 2625.5 * a * exp_br \
         + (-2625.5) * exp_br * (1+br) * q / r \
-        + exp_br*(1+br+br2/2+br3/6+br4/24+br5/120+br6/720) * c / dr**6
-
+        + exp_br*(1+br+br2/2+br3/6+br4/24+br5/120+br6/720) * c / dr**6 
     return f * m
 
 

diff --git a/dmff/admp/pme.py b/dmff/admp/pme.py
@@ -42,7 +42,7 @@ class ADMPPmeForce:
     The so called "environment paramters" means parameters that do not need to be differentiable
     '''
 
-    def __init__(self, box, axis_type, axis_indices, covalent_map, rc, ethresh, lmax, lpol=False, lpme=True, steps_pol=None):
+    def __init__(self, box, axis_type, axis_indices, rc, ethresh, lmax, lpol=False, lpme=True, steps_pol=None):
         '''
         Initialize the ADMPPmeForce calculator.
 
@@ -51,8 +51,6 @@ def __init__(self, box, axis_type, axis_indices, covalent_map, rc, ethresh, lmax
                 (3, 3) float, box size in row
             axis_type:
                 (na,) int, types of local axis (bisector, z-then-x etc.)
-            covalent_map:
-                (na, na) int, covalent map matrix, labels the topological distances between atoms
             rc: 
                 float: cutoff distance
             ethresh: 
@@ -91,10 +89,10 @@ def __init__(self, box, axis_type, axis_indices, covalent_map, rc, ethresh, lmax
             self.K2 = K2
             self.K3 = K3
         self.pme_order = 6
-        self.covalent_map = covalent_map
         self.lpol = lpol
         self.steps_pol = steps_pol
-        self.n_atoms = int(covalent_map.shape[0]) # len(axis_type)
+        # self.n_atoms = int(covalent_map.shape[0]) # len(axis_type)
+        self.n_atoms = len(axis_type)
 
         # setup calculators
         self.refresh_calculators()
@@ -107,7 +105,7 @@ def generate_get_energy(self):
             def get_energy(positions, box, pairs, Q_local, mScales):
                 return energy_pme(positions, box, pairs,
                                  Q_local, None, None, None,
-                                 mScales, None, None, self.covalent_map,
+                                 mScales, None, None, 
                                  self.construct_local_frames, self.pme_recip,
                                  self.kappa, self.K1, self.K2, self.K3, self.lmax, False, lpme=self.lpme)
             return get_energy
@@ -116,7 +114,7 @@ def get_energy(positions, box, pairs, Q_local, mScales):
             def energy_fn(positions, box, pairs, Q_local, Uind_global, pol, tholes, mScales, pScales, dScales):
                 return energy_pme(positions, box, pairs,
                                  Q_local, Uind_global, pol, tholes,
-                                 mScales, pScales, dScales, self.covalent_map,
+                                 mScales, pScales, dScales, 
                                  self.construct_local_frames, self.pme_recip,
                                  self.kappa, self.K1, self.K2, self.K3, self.lmax, True, lpme=self.lpme)
             self.energy_fn = energy_fn
@@ -284,7 +282,7 @@ def setup_ewald_parameters(
 # @jit_condition(static_argnums=())
 def energy_pme(positions, box, pairs,
         Q_local, Uind_global, pol, tholes,
-        mScales, pScales, dScales, covalent_map, 
+        mScales, pScales, dScales, 
         construct_local_frame_fn, pme_recip_fn, kappa, K1, K2, K3, lmax, lpol, lpme=True):
     '''
     This is the top-level wrapper for multipole PME
@@ -306,7 +304,7 @@ def energy_pme(positions, box, pairs,
             (Nexcl,): multipole-multipole interaction exclusion scalings: 1-2, 1-3 ...
             for permanent-permanent, permanent-induced, induced-induced interactions
         pairs:
-            Np * 2: interacting pair indices
+            Np * 3: interacting pair indices and topology distance
         covalent_map:
             Na * Na: topological distances between atoms, if i, j are topologically distant, then covalent_map[i, j] == 0
         construct_local_frame_fn:
@@ -353,26 +351,24 @@ def energy_pme(positions, box, pairs,
 
     if lpol:
         ene_real = pme_real(positions, box, pairs, Q_global, U_ind, pol, tholes, 
-                           mScales, pScales, dScales, covalent_map, kappa, lmax, True)
+                           mScales, pScales, dScales, kappa, lmax, True)
     else:
         ene_real = pme_real(positions, box, pairs, Q_global, None, None, None,
-                           mScales, None, None, covalent_map, kappa, lmax, False)
+                           mScales, None, None, kappa, lmax, False)
 
     if lpme:
         ene_recip = pme_recip_fn(positions, box, Q_global_tot)
         ene_self = pme_self(Q_global_tot, kappa, lmax)
 
         if lpol:
             ene_self += pol_penalty(U_ind, pol)
-
         return ene_real + ene_recip + ene_self
 
     else:
         if lpol:
             ene_self = pol_penalty(U_ind, pol)
         else:
             ene_self = 0.0
-
         return ene_real + ene_self
 
 
@@ -747,7 +743,7 @@ def pme_real_kernel(dr, qiQI, qiQJ, qiUindI, qiUindJ, thole1, thole2, dmp, mscal
 # @jit_condition(static_argnums=(7))
 def pme_real(positions, box, pairs, 
         Q_global, Uind_global, pol, tholes,
-        mScales, pScales, dScales, covalent_map, 
+        mScales, pScales, dScales,
         kappa, lmax, lpol):
     '''
     This is the real space PME calculate function
@@ -763,7 +759,7 @@ def pme_real(positions, box, pairs,
         box:
             3 * 3: box, axes arranged in row
         pairs:
-            Np * 2: interacting pair indices
+            Np * 3: interacting pair indices and topology distance
         Q_global:
             Na * (l+1)**2: harmonics multipoles of each atom, in global frame
         Uind_global:
@@ -786,14 +782,14 @@ def pme_real(positions, box, pairs,
     Output:
         ene: pme realspace energy
     '''
-    pairs = regularize_pairs(pairs)
-    buffer_scales = pair_buffer_scales(pairs)
+    pairs = pairs.at[:, :2].set(regularize_pairs(pairs[:, :2]))
+    buffer_scales = pair_buffer_scales(pairs[:, :2])
     box_inv = jnp.linalg.inv(box)
     r1 = distribute_v3(positions, pairs[:, 0])
     r2 = distribute_v3(positions, pairs[:, 1])
     Q_extendi = distribute_multipoles(Q_global, pairs[:, 0])
     Q_extendj = distribute_multipoles(Q_global, pairs[:, 1])
-    nbonds = distribute_matrix(covalent_map,pairs[:, 0],pairs[:, 1])
+    nbonds = pairs[:, 2]
     #nbonds = covalent_map[pairs[:, 0], pairs[:, 1]]
     indices = nbonds-1
     mscales = distribute_scalar(mScales, indices)