Add aux support in ml.py for MLP. Fix jit error when saving ml potent…

…ial with save_dmff2tf.py.

dmff/eann/eann.py

@@ -247,7 +247,7 @@ def get_features(self, radial, dr, pairs, buffer_scales, orb_coeff):
         f_cut = cutoff_cosine(dr_norm, self.rc)
         neigh_list = jnp.concatenate((pairs,pairs[:,[1,0]]),axis=0)
         buffer_scales_ = jnp.concatenate((buffer_scales,buffer_scales),axis=0)
-        totneighbour = len(neigh_list)
+        totneighbour = neigh_list.shape[0]
         prefacs = f_cut.reshape(1, -1)
         angular = prefacs
         for ipsin in range(1,self.nipsin+1):
@@ -310,7 +310,8 @@ def get_energy(positions, box, pairs, params):
             self.rs = params['density.rs']
             self.inta = params['density.inta']
 
-            radial_i, radial_j = get_gto(jnp.arange(len(dr_norm)), dr_norm, pairs, self.rc, self.rs, self.inta, self.elem_indices)
+            length_dr_norm = dr_norm.shape[0]
+            radial_i, radial_j = get_gto(jnp.arange(length_dr_norm), dr_norm, pairs, self.rc, self.rs, self.inta, self.elem_indices)
             radial = jnp.concatenate((radial_i,radial_j), axis=0)
             orb_coeff = params['density.params'][self.elem_indices,:] # (48,16)
 

dmff/generators/ml.py

@@ -116,12 +116,19 @@ def createPotential(self, topdata: DMFFTopology, nonbondedMethod, nonbondedCutof
         n_elem, elem_indices = get_elem_indices(self.ommtopology)
         self.model = EANNForce(n_elem, elem_indices, n_gto=self.ngto, nipsin=self.nipsin, rc=self.rc)
         n_layers = self.model.n_layers
-        def potential_fn(positions, box, pairs, params):
+
+        has_aux = False
+        if "has_aux" in kwargs and kwargs["has_aux"]:
+            has_aux = True
+
+        def potential_fn(positions, box, pairs, params, aux=None):
             # convert unit to angstrom
             positions = positions * 10
             box = box * 10
-
-            return self.model.get_energy(positions, box, pairs, params[self.name])
+            if has_aux:
+                return self.model.get_energy(positions, box, pairs, params[self.name]), aux
+            else:
+                return self.model.get_energy(positions, box, pairs, params[self.name])
 
         self._jaxPotential = potential_fn
         return potential_fn