Merge branch 'deepmodeling:wangxy/v1.0.0-devel' into wangxy/v1.0.0-devel

docs/user_guide/2.installation.md

@@ -8,9 +8,9 @@ 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

docs/user_guide/4.2ADMPPmeForce.md

@@ -7,8 +7,9 @@ ADMPmeForce provides a support to multipolar polarizable electrostatic energy ca
 ### 1.1 Multipole Expansion
 
 The electrostatic interaction between two atoms can be described using multipole expansion, in which the electron cloud of an atom can be expanded as a series of multipole moments including charges, dipoles, quadrupoles, and octupoles etc. If only the charges (zero-moment) are considered, it is reduced to the point charge model in classical force fields:
+
 $$
-V=\sum_{ij} \frac{q_i q_j}{r_{ij}}
+V=\sum_{ij} \\frac{q_i q_j}{r_{ij}}
 $$
 
 where $q_i$ is the charge of atom $i$.
@@ -23,7 +24,7 @@ where $Q_t^A$ represents the $t$-component of the multipole moment of atom $A$.
 
 $$0, 10, 1c, 1s, 20, 21c, 21s, 22c, 22s, ...$$
 
-The $T_{tu}^{AB}$ represents the interaction tensor between multipoles. The mathematical expression of these tensors can be found in the appendix F of Ref 1. The user can also find the conversion rule between different representations in Ref 1 & 5.
+The $T_{tu}^{AB}$ represents the interaction tensor between multipoles. The mathematical expression of these tensors can be found in the appendix F of Ref 1. The user can also find the conversion rule between different representations in Ref 1.
 
 ### 1.2 Coordinate System for Multipoles
 
@@ -37,7 +38,7 @@ Different to charges, the definition of multipole moments depends on local coord
 
 ### 1.3 Polarization
 
-ADMPPmeForce 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:
+ADMPPmeForce 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 5), 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{
@@ -47,7 +48,7 @@ u=r_{ij}/\left(\alpha_i \alpha_j\right)^{1/6}
 }
 $$
 
-Other damping functions between multipole moments can be found in Ref 6, table I. 
+Other damping functions between multipole moments can be found in Ref 5, table I. 
 
 It is noted that the atomic damping parameter $a=a_i+a_j$ is only effective on topological neighboring pairs (with $pscale = 0$), while a default value of $a_{default}$ is set for all other pairs. In DMFF, the atomic $a_i$ is specified via the xml API, while $a_{default}$  is controlled by the `DEFAULT_THOLE_WIDTH` variable, which is set to 5.0 by default. This variable can be changed in the following way:
 
@@ -113,7 +114,7 @@ E = E_{real}+E_{recip}+E_{self}
 }
 $$
 
-As for multipolar PME and dispersion PME, the users and developers are referred to Ref 2, 3, and 5 for mathematical details.
+As for multipolar PME and dispersion PME, the users and developers are referred to Ref 2, 3 for mathematical details.
 
 The key parameters in PME include:
 
@@ -143,8 +144,7 @@ In the current version, the dispersion PME calculator uses the same parameters a
 2. [The Multipolar Ewald paper in JCTC:  J. Chem. Theory Comput. 2015, 11, 2, 436–450](https://pubs.acs.org/doi/abs/10.1021/ct5007983)
 3. [The dispersion Ewald/PME](https://aip.scitation.org/doi/pdf/10.1063/1.470117)
 4. [Frenkel & Smit book](https://www.elsevier.com/books/understanding-molecular-simulation/frenkel/978-0-12-267351-1)
-5. Note on multipole ewald. [link](multipole_pme.md)
-6. [MPID Reference](https://doi.org/10.1063/1.4984113)
+5. [MPID Reference](https://doi.org/10.1063/1.4984113)
 
 
 ## 2. Frontend 
@@ -237,21 +237,21 @@ The important tags in the frontend includes:
 For each atom type, there is an `<Atom/>` node, with the following tags:
 
 * `type`: the label for the atomtype.
-* `C6, C8, C10`: dispersion terms, defined in `e*nm^n`.
+* `C6, C8, C10`: dispersion terms, defined in `kj/mol/nm^n`.
 
 To create a DispADMP PME potential function, one needs to do the following in python:
 
 ```python
 H = Hamiltonian('forcefield.xml')
-pots = H.createPotential(pdb.topology, nonbondedCutoff=rc*unit.nanometer, nonbondedMethod=app.PME, ethresh=5e-4, step_pol=5)
+pots = H.createPotential(pdb.topology, nonbondedCutoff=rc*unit.nanometer, nonbondedMethod=app.PME, ethresh=5e-4)
 ```
 
 Then the potential function, the parameters, and the generator can be accessed as:
 
 ```python
 efunc_pme = pots.dmff_potentials["ADMPDispPmeForce"]
 params_pme = H.getParameters()["ADMPDispPmeForce"]
-generator_pme = H.getGenerators()[1] # if ADMPDispPmeForce is the Second force node in xml
+generator_pme = H.getGenerators()[0] # if ADMPDispPmeForce is the First force node in xml
 ```
 The keyword `ethresh` in `createPotential` would impact the behavior of the function:
 
@@ -408,7 +408,7 @@ The backend of the ADMPDisp PME energy is an `ADMPDispPmeForce` object. It conta
 ***ATTRIBUTES:***
 * `ethresh`: float, accuracy for PME setups
 * `kappa`: float, the range-separation parameter $\kappa$ used in PME calculation
-* `lpme`: bool, use PME or not (sohuld be TRUE 99.9% of time)
+* `lpme`: bool, use PME or not (should be TRUE 99.9% of time)
 * `n_atoms`: number of atoms
 * `pme_order`: order of PME interpolation, now only support 6
 * `rc`: real space cutoffs
@@ -442,3 +442,4 @@ The backend of the ADMPDisp PME energy is an `ADMPDispPmeForce` object. It conta
 * `get_forces`
 
   Same as `get_energy`, only return forces per atom, instead of total energy, equivalent to `jax.grad(get_energy)`
+

docs/user_guide/4.4MLForce.md

@@ -32,7 +32,7 @@ $$
 \phi_{l_{x} y_{l} y_{z}}^{\alpha, r_{s}}=x^{l_{x}} y^{l_{y}} z^{l_{z}} \exp \left(-\alpha\left|r-r_{s}\right|^{2}\right)
 $$
 
-where each atom is taken as the origin, $r=(x,y,z)$ constitutes the coordinate vector of an electron, $r$ is the norm of the vector,$α$ and $r_s$ are parameters that determine radial distributions of atomic orbitals, ${l_x+l_y+l_z=L}$ specifies the orbital angular momentum ($L$), e.g., $L$ = 0, 1, and 2, correspond to the s, p, and d orbitals, respectively. In this representation, the embedded density of atom $i$ can be taken as the square of the linear combination of atomic orbitals from neighboring atoms, in a similar spirit as that in Hartree−Fock (HF) and densityfunctional theory (DFT). This would generate a scalar $ρ^i$ value for the embedding atom $i$, as used in the EAM, which has been proven to offer insufficient representability for the total energyand can be improved by including the gradients of density.
+where each atom is taken as the origin, $r=(x,y,z)$ constitutes the coordinate vector of an electron, $r$ is the norm of the vector, $α$ and $r_s$ are parameters that determine radial distributions of atomic orbitals, ${l_x+l_y+l_z=L}$ specifies the orbital angular momentum ($L$), e.g., $L$ = 0, 1, and 2, correspond to the s, p, and d orbitals, respectively. In this representation, the embedded density of atom $i$ can be taken as the square of the linear combination of atomic orbitals from neighboring atoms, in a similar spirit as that in Hartree−Fock (HF) and densityfunctional theory (DFT). This would generate a scalar $ρ^i$ value for the embedding atom $i$, as used in the EAM, which has been proven to offer insufficient representability for the total energyand can be improved by including the gradients of density.
 
 The other detail can be found in References.
 
@@ -144,7 +144,7 @@ The backend of the SGNN energy is an `MolGNNForce` object. It contains the follo
 * `sizes`: [tuple, tuple], optional, the sizes (numbers of hidden neurons) of the network before and after message passing, default = [(40, 20, 20), (20, 10)]
 * `nn`: int, optional size of the subgraphs, i.e., how many neighbors to include around the central bond, default = 1
 * `sigma`: float, optional, final scaling factor of the energy. default = 162.13039087945623
-* `mu`: float, optional, a constant shift, the final total energy would be ${(E_{NN} + \mu) * \sigma}
+* `mu`: float, optional, a constant shift, the final total energy would be ${(E_{NN} + \mu)}*{\sigma}$
 * `seed`: int, optional, the seed for random number generator, default 12345
 
 ***METHODS***

docs/user_guide/4.5Optimization.md

@@ -41,8 +41,8 @@ multiTrans = MultiTransform(your_parameter_tree)
 - For each part of the parameter tree you want to optimize differently, set an optimizer. For example:
 
 ```python
-multiTrans["Label/Path1"] = genOptimizer(learning_rate=lr1, clip=clip1)
-multiTrans["Label/Path2"] = genOptimizer(learning_rate=lr2, clip=clip2)
+multiTrans["Force1"]["Parameter1"] = genOptimizer(learning_rate=lr1, clip=clip1)
+multiTrans["Force1"]["Parameter2"] = genOptimizer(learning_rate=lr2, clip=clip2)
 ```
 
 - Finalize MultiTransform
@@ -65,19 +65,18 @@ traj = md.load("init.dcd", top="box.pdb")[50:]
 grad_transform = optax.multi_transform(multiTrans.transforms, multiTrans.labels)
 ```
 
-- Mask Integer Parameters (If Needed)
+- Mask Parameters (If Needed)
 
-- If you have parameters in your tree that shouldn't be updated, create a mask and then mask your transformation:
+If you have parameters in your tree that shouldn't be updated, create a mask and then mask your transformation:
 
 ```python
-mask = jax.tree_util.tree_map(lambda x: x.dtype != jnp.int32 and x.dtype != int, your_parameter_tree)
-grad_transform = optax.masked(grad_transform, mask)
+grad_transform = optax.masked(grad_transform, hamiltonian.getParameters().mask)
 ```
 
 - Initialize Optimization State:
 
 ```python
-opt_state = grad_transform.init(your_parameter_tree)
+opt_state = grad_transform.init(hamiltonian.getParameters().parameters)
 ```
 
 By following the above steps, you'll have a `grad_transform` that can handle complex parameter trees and an optimization state `opt_state` ready for your optimization routine.