update eann.py

.github/workflows/ut.yml

@@ -19,10 +19,17 @@ jobs:
       - name: Install Dependencies
         run: |
           source $CONDA/bin/activate
+<<<<<<< HEAD
+          conda create -n dmff -y python=${{ matrix.python-version }} numpy openmm==7.7.0 pytest rdkit biopandas openbabel mdtraj ambertools -c conda-forge
+          conda activate dmff
+          pip install --upgrade pip
+          pip install jax jaxlib jaxopt networkx parmed pymbar==4.0.1 chex==0.1.4 tqdm
+=======
           conda create -n dmff -y python=${{ matrix.python-version }} numpy openmm==7.7.0 pytest rdkit openbabel mdtraj ambertools -c conda-forge
           conda activate dmff
           pip install --upgrade pip
           pip install jax jaxlib jaxopt networkx parmed pymbar==4.0.1 optax tqdm
+>>>>>>> upstream/devel
       - name: Install DMFF
         run: |
           source $CONDA/bin/activate dmff && pip install .

.gitignore

@@ -789,6 +789,9 @@ FodyWeavers.xsd
 # debugging ipynb
 debug.ipynb
 test.xml
+<<<<<<< HEAD
+=======
 
 # PyCharm Cache
-.idea/
+.idea/
+>>>>>>> upstream/devel

dmff/admp/pme.py

@@ -13,18 +13,33 @@
 from .settings import POL_CONV, MAX_N_POL
 from .recip import generate_pme_recip, Ck_1
 from .multipole import (
+<<<<<<< HEAD
+    C1_c2h, 
+=======
     C1_c2h,
     C1_h2c,
     C2_h2c,
+>>>>>>> upstream/devel
     convert_cart2harm,
     rot_ind_global2local,
     rot_global2local,
     rot_local2global,
 )
+<<<<<<< HEAD
+from .spatial import (
+    v_pbc_shift, 
+    generate_construct_local_frames, 
+    build_quasi_internal
+)
+from .pairwise import (
+    distribute_scalar, 
+    distribute_v3, 
+=======
 from .spatial import v_pbc_shift, generate_construct_local_frames, build_quasi_internal
 from .pairwise import (
     distribute_scalar,
     distribute_v3,
+>>>>>>> upstream/devel
     distribute_multipoles,
     distribute_matrix,
 )
@@ -38,6 +53,12 @@ class ADMPPmeForce:
     This is a convenient wrapper for multipolar PME calculations
     It wrapps all the environment parameters of multipolar PME calculation
     The so called "environment paramters" means parameters that do not need to be differentiable
+<<<<<<< HEAD
+    '''
+
+    def __init__(self, box, axis_type, axis_indices, rc, ethresh, lmax, lpol=False, lpme=True, steps_pol=None, has_aux=False):
+        '''
+=======
     """
 
     def __init__(
@@ -55,6 +76,7 @@ def __init__(
         has_aux=False,
     ):
         """
+>>>>>>> upstream/devel
         Initialize the ADMPPmeForce calculator.
 
         Input:
@@ -106,8 +128,11 @@ def __init__(
         # self.n_atoms = int(covalent_map.shape[0]) # len(axis_type)
         self.n_atoms = len(axis_type)
         self.has_aux = has_aux
+<<<<<<< HEAD
+=======
 
         self.Q_local_to_global = self.generate_Q_global_function()
+>>>>>>> upstream/devel
 
         # setup calculators
         self.refresh_calculators()
@@ -118,6 +143,13 @@ def generate_get_energy(self):
         if not self.lpol:
 
             def get_energy(positions, box, pairs, Q_local, mScales):
+<<<<<<< HEAD
+                return energy_pme(positions, box, pairs,
+                                Q_local, None, None, None,
+                                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 energy_pme(
                     positions,
                     box,
@@ -140,6 +172,7 @@ def get_energy(positions, box, pairs, Q_local, mScales):
                     lpme=self.lpme,
                 )
 
+>>>>>>> upstream/devel
             return get_energy
         else:
             # this is the bare energy calculator, with Uind as explicit input
@@ -184,6 +217,23 @@ def energy_fn(
 
             # this is the wrapper that include a Uind optimizer
             def get_energy(
+<<<<<<< HEAD
+                    positions, box, pairs, 
+                    Q_local, pol, tholes, mScales, pScales, dScales, 
+                    U_init = self.U_ind, aux = None):
+                self.U_ind, self.lconverg, self.n_cycle = self.optimize_Uind(
+                        positions, box, pairs, Q_local, pol, tholes, 
+                        mScales, pScales, dScales, 
+                        U_init=U_init, steps_pol=self.steps_pol)
+                # here we rely on Feynman-Hellman theorem, drop the term dV/dU*dU/dr !
+                # self.U_ind = jax.lax.stop_gradient(U_ind)
+                energy = self.energy_fn(positions, box, pairs, Q_local, self.U_ind, pol, tholes, mScales, pScales, dScales)
+                if aux is not None:
+                    aux["U_ind"] = self.U_ind
+                    return energy, aux
+                else:
+                    return energy
+=======
                 positions,
                 box,
                 pairs,
@@ -232,6 +282,7 @@ def get_energy(
                 else:
                     return energy
 
+>>>>>>> upstream/devel
             return get_energy
 
     def generate_esp(self):
@@ -935,10 +986,15 @@ def pme_real_kernel(
     Output:
         energy:
             float, realspace interaction energy between the sites
+<<<<<<< HEAD
+    '''
+    cc, cd, dd_m0, dd_m1, cq, dq_m0, dq_m1, qq_m0, qq_m1, qq_m2 = calc_e_perm(dr, mscales, kappa, lmax)
+=======
     """
     cc, cd, dd_m0, dd_m1, cq, dq_m0, dq_m1, qq_m0, qq_m1, qq_m2 = calc_e_perm(
         dr, mscales, kappa, lmax
     )
+>>>>>>> upstream/devel
     if lpol:
         cud, dud_m0, dud_m1, udq_m0, udq_m1, udud_m0, udud_m1 = calc_e_ind(
             dr, thole1, thole2, dmp, pscales, dscales, kappa, lmax

dmff/admp/qeq.py

@@ -5,8 +5,11 @@
 from ..classical.inter import CoulNoCutoffForce, CoulombPMEForce
 from typing import Tuple, List
 from ..settings import PRECISION
+<<<<<<< HEAD
+=======
 from .pme import energy_pme
 from .recip import generate_pme_recip, Ck_1
+>>>>>>> upstream/devel
 
 if PRECISION == "double":
     CONST_0 = jnp.array(0, dtype=jnp.float64)
@@ -17,6 +20,16 @@
 
 try:
     import jaxopt
+<<<<<<< HEAD
+    try:
+        from jaxopt import Broyden
+        JAXOPT_OLD = False
+    except ImportError:
+        JAXOPT_OLD = True
+        print("jaxopt is too old. The QEQ potential function cannot be jitted. Please update jaxopt to the latest version for speed concern.")
+except ImportError:
+    print("jaxopt not found, QEQ cannot be used.")
+=======
 
     try:
         from jaxopt import Broyden
@@ -31,6 +44,7 @@
 except ImportError:
     import warnings
     warnings.warn("jaxopt not found, QEQ cannot be used.")
+>>>>>>> upstream/devel
 import jax
 
 from jax.scipy.special import erf, erfc
@@ -106,9 +120,13 @@ def E_sr2(pos, box, pairs, q, eta, ds, buffer_scales):
 
 @jit_condition()
 def E_sr3(pos, box, pairs, q, eta, ds, buffer_scales):
+<<<<<<< HEAD
+    etasqrt = jnp.sqrt(eta[pairs[:, 0]] ** 2 + eta[pairs[:, 1]] ** 2)
+=======
     etasqrt = jnp.sqrt(
         eta[pairs[:, 0]] ** 2 + eta[pairs[:, 1]] ** 2 + 1e-64
     )  # add eta to avoid division by zero
+>>>>>>> upstream/devel
     epiece = eta_piecewise(etasqrt, ds)
     pre_pair = -epiece * DIELECTRIC
     pre_self = etainv_piecewise(eta) / (jnp.sqrt(2 * jnp.pi)) * DIELECTRIC
@@ -164,12 +182,20 @@ def ds_pairs(positions, box, pairs, pbc_flag):
     if pbc_flag is False:
         dr = pos1 - pos2
     else:
+<<<<<<< HEAD
+        box_inv = jnp.linalg.inv(box)
+=======
         box_inv = jnp.linalg.inv(box + jnp.eye(3) * 1e-36)
+>>>>>>> upstream/devel
         dpos = pos1 - pos2
         dpos = dpos.dot(box_inv)
         dpos -= jnp.floor(dpos + 0.5)
         dr = dpos.dot(box)
+<<<<<<< HEAD
+    ds = jnp.linalg.norm(dr, axis=1)
+=======
     ds = jnp.linalg.norm(dr + 1e-64, axis=1)  # add eta to avoid division by zero
+>>>>>>> upstream/devel
     return ds
 
 
@@ -251,6 +277,14 @@ def __init__(
             raise ValueError("damp_mod must be 1, 2 or 3")
 
         if pbc_flag:
+<<<<<<< HEAD
+            force = CoulombPMEForce(r_cut, kappa, K)
+            self.kappa = kappa
+        else:
+            force = CoulNoCutoffForce()
+            self.kappa = 1.0
+        self.coul_energy = force.generate_get_energy()
+=======
             pme_recip_fn = generate_pme_recip(
                 Ck_fn=Ck_1,
                 kappa=kappa / 10,
@@ -317,6 +351,7 @@ def coul_energy(positions, box, pairs, q, mscales):
             self.kappa = 0.0
 
         self.coul_energy = coul_energy
+>>>>>>> upstream/devel
 
     def generate_get_energy(self):
         @jit_condition()
@@ -401,5 +436,9 @@ def get_energy(positions, box, pairs, mscales, eta, chi, J, aux=None):
                 return energy, aux
             else:
                 return energy
+<<<<<<< HEAD
+
+=======
 
+>>>>>>> upstream/devel
         return get_energy

dmff/api/graph.py

@@ -7,6 +7,11 @@
     import simtk.openmm.app as app
     import simtk.unit as unit
 from typing import Dict, Tuple, List
+<<<<<<< HEAD
+from rdkit import Chem
+from rdkit.Chem import AllChem
+
+=======
 try:
     from rdkit import Chem
     from rdkit.Chem import AllChem
@@ -21,12 +26,17 @@ def is_same_list(l1, l2):
         if l1[nn] != l2[nn]:
             return False
     return True
+>>>>>>> upstream/devel
 
 def matchTemplate(graph, template):
     if graph.number_of_nodes() != template.number_of_nodes():
         # print("Node with different number of nodes.")
         return False, {}, {}
 
+<<<<<<< HEAD
+    def match_func(n1, n2):
+        return n1["element"] == n2["element"] and n1["external_bond"] == n2["external_bond"]
+=======
     name_graph = sorted([i[1]['name'] for i in graph.nodes.data()])
     name_template = sorted([i[1]['name'] for i in template.nodes.data()])
 
@@ -36,6 +46,7 @@ def match_func(n1, n2):
     else:
         def match_func(n1, n2):
             return n1["element"] == n2["element"] and n1["external_bond"] == n2["external_bond"]
+>>>>>>> upstream/devel
 
     def edge_match(e1, e2):
         if len(e1) == 0 and len(e2) == 0:
@@ -118,7 +129,11 @@ def graph2top(graph: nx.Graph) -> app.Topology:
     return top
 
 
+<<<<<<< HEAD
+def top2rdmol(top: app.Topology, indices: List[int]) -> Chem.rdchem.Mol:
+=======
 def top2rdmol(top: app.Topology, indices: List[int]):
+>>>>>>> upstream/devel
     rdmol = Chem.Mol()
     emol = Chem.EditableMol(rdmol)
     idx2ridx = {}

dmff/api/paramset.py

@@ -24,11 +24,15 @@ class ParamSet:
         Converts all parameters to jax arrays.
     """
 
+<<<<<<< HEAD
+    def __init__(self, data: Dict[str, Union[Dict[str, jnp.ndarray], jnp.ndarray]] = None, mask: Dict[str, Union[Dict[str, jnp.ndarray], jnp.ndarray]] = None):
+=======
     def __init__(
         self,
         data: Dict[str, Union[Dict[str, jnp.ndarray], jnp.ndarray]] = None,
         mask: Dict[str, Union[Dict[str, jnp.ndarray], jnp.ndarray]] = None,
     ):
+>>>>>>> upstream/devel
         """
         Initializes a new ParamSet object.
 
@@ -56,13 +60,17 @@ def addField(self, field: str) -> None:
         self.parameters[field] = {}
         self.mask[field] = {}
 
+<<<<<<< HEAD
+    def addParameter(self, values: jnp.ndarray, name: str, field: str = None, mask: jnp.ndarray = None) -> None:
+=======
     def addParameter(
         self,
         values: jnp.ndarray,
         name: str,
         field: str = None,
         mask: jnp.ndarray = None,
     ) -> None:
+>>>>>>> upstream/devel
         """
         Adds a new parameter to the parameters and mask dictionaries.
 
@@ -97,7 +105,12 @@ def to_jax(self) -> None:
         for key1 in self.parameters:
             if isinstance(self.parameters[key1], dict):
                 for key2 in self.parameters[key1]:
+<<<<<<< HEAD
+                    self.parameters[key1][key2] = jnp.array(
+                        self.parameters[key1][key2])
+=======
                     self.parameters[key1][key2] = jnp.array(self.parameters[key1][key2])
+>>>>>>> upstream/devel
             else:
                 self.parameters[key1] = jnp.array(self.parameters[key1])
 
@@ -117,12 +130,15 @@ def __getitem__(self, key: str) -> Union[Dict[str, jnp.ndarray], jnp.ndarray]:
         """
         return self.parameters[key]
 
+<<<<<<< HEAD
+=======
     def update_mask(self, gradients):
         gradients = jax.tree_map(
             lambda g, m: jnp.where(jnp.abs(m - 1.0) > 1e-5, g, 0.0), gradients, self.mask
         )
         return gradients
 
+>>>>>>> upstream/devel
 
 def flatten_paramset(prmset: ParamSet) -> tuple:
     """
@@ -160,4 +176,9 @@ def unflatten_paramset(aux_data: Dict, contents: tuple) -> ParamSet:
     return ParamSet(data=contents[0], mask=aux_data)
 
 
+<<<<<<< HEAD
+jax.tree_util.register_pytree_node(
+    ParamSet, flatten_paramset, unflatten_paramset)
+=======
 jax.tree_util.register_pytree_node(ParamSet, flatten_paramset, unflatten_paramset)
+>>>>>>> upstream/devel