Release DMFF v0.2.0

.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: "git@gitee.com:deepmodeling/DMFF.git"

.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

.gitignore

@@ -781,4 +781,7 @@ FodyWeavers.xsd
 *.acpype/
 
 */_date.py
-*/_version.py
+*/_version.py
+
+# hmtff cache
+*.hmtff/

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
 

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 *

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])

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
 
 

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)