build neighbor list with external Python program

Fix deepmodeling#2877

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/infer/deep_eval.py

@@ -56,6 +56,7 @@ def __init__(
         default_tf_graph: bool = False,
         auto_batch_size: Union[bool, int, AutoBatchSize] = False,
         input_map: Optional[dict] = None,
+        neighbor_list=None,
     ):
         self.graph = self._load_graph(
             model_file,
@@ -86,6 +87,11 @@ def __init__(
         else:
             raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize")
 
+        self.neighbor_list = neighbor_list
+        if neighbor_list is not None and self.auto_batch_size is not None:
+            # force nbatch = 1
+            self.auto_batch_size.maximum_working_batch_size = 1
+
     @property
     @lru_cache(maxsize=None)
     def model_type(self) -> str:
@@ -360,3 +366,89 @@ def eval_typeebd(self) -> np.ndarray:
         t_typeebd = self._get_tensor("t_typeebd:0")
         [typeebd] = run_sess(self.sess, [t_typeebd], feed_dict={})
         return typeebd
+
+    def build_neighbor_list(
+        self,
+        coords: np.ndarray,
+        cell: Optional[np.ndarray],
+        atype: np.ndarray,
+        imap: np.ndarray,
+        neighbor_list,
+    ):
+        """Make the mesh with neighbor list for a single frame.
+
+        Parameters
+        ----------
+        coords : np.ndarray
+            The coordinates of atoms. Should be of shape [natoms, 3]
+        cell : Optional[np.ndarray]
+            The cell of the system. Should be of shape [3, 3]
+        atype : np.ndarray
+            The type of atoms. Should be of shape [natoms]
+        imap : np.ndarray
+            The index map of atoms. Should be of shape [natoms]
+        neighbor_list : ase.neighborlist.NewPrimitiveNeighborList
+            ASE neighbor list. The following method or attribute will be
+            used/set: bothways, self_interaction, update, build, first_neigh,
+            pair_second, offset_vec.
+
+        Returns
+        -------
+        natoms_vec : np.ndarray
+            The number of atoms. This tensor has the length of Ntypes + 2
+            natoms[0]: nloc
+            natoms[1]: nall
+            natoms[i]: 2 <= i < Ntypes+2, number of type i atoms for nloc
+        coords : np.ndarray
+            The coordinates of atoms, including ghost atoms. Should be of
+            shape [nframes, nall, 3]
+        atype : np.ndarray
+            The type of atoms, including ghost atoms. Should be of shape [nall]
+        mesh : np.ndarray
+            The mesh in nei_mode=4.
+        imap : np.ndarray
+            The index map of atoms. Should be of shape [nall]
+        """
+        pbc = np.repeat(cell is not None, 3)
+        cell = cell.reshape(3, 3)
+        positions = coords.reshape(-1, 3)
+        neighbor_list.bothways = True
+        neighbor_list.self_interaction = False
+        if neighbor_list.update(pbc, cell, positions):
+            neighbor_list.build(pbc, cell, positions)
+        first_neigh = neighbor_list.first_neigh.copy()
+        pair_second = neighbor_list.pair_second.copy()
+        offset_vec = neighbor_list.offset_vec.copy()
+        # get out-of-box neighbors
+        out_mask = np.any(offset_vec != 0, axis=1)
+        out_idx = pair_second[out_mask]
+        out_offset = offset_vec[out_mask]
+        out_coords = positions[out_idx] + out_offset.dot(cell)
+        atype = np.array(atype, dtype=int)
+        out_atype = atype[out_idx]
+
+        nloc = positions.shape[0]
+        nghost = out_idx.size
+        all_coords = np.concatenate((positions, out_coords), axis=0)
+        all_atype = np.concatenate((atype, out_atype), axis=0)
+        # convert neighbor indexes
+        pair_second[out_mask] = np.arange(nloc, nloc + nghost)
+        # get the mesh
+        mesh = np.zeros(16 + nloc * 2 + pair_second.size, dtype=int)
+        mesh[0] = nloc
+        # ilist
+        mesh[16 : 16 + nloc] = np.arange(nloc)
+        # numnei
+        mesh[16 + nloc : 16 + nloc * 2] = first_neigh[1:] - first_neigh[:-1]
+        # jlist
+        mesh[16 + nloc * 2 :] = pair_second
+
+        # natoms_vec
+        natoms_vec = np.zeros(self.ntypes + 2).astype(int)
+        natoms_vec[0] = nloc
+        natoms_vec[1] = nloc + nghost
+        for ii in range(self.ntypes):
+            natoms_vec[ii + 2] = np.count_nonzero(atype == ii)
+        # imap append ghost atoms
+        imap = np.concatenate((imap, np.arange(nloc, nloc + nghost)))
+        return natoms_vec, all_coords, all_atype, mesh, imap

deepmd/infer/deep_pot.py

@@ -78,6 +78,7 @@ def __init__(
         default_tf_graph: bool = False,
         auto_batch_size: Union[bool, int, AutoBatchSize] = True,
         input_map: Optional[dict] = None,
+        neighbor_list=None,
     ) -> None:
         # add these tensors on top of what is defined by DeepTensor Class
         # use this in favor of dict update to move attribute from class to
@@ -112,6 +113,7 @@ def __init__(
             default_tf_graph=default_tf_graph,
             auto_batch_size=auto_batch_size,
             input_map=input_map,
+            neighbor_list=neighbor_list,
         )
 
         # load optional tensors
@@ -479,8 +481,18 @@ def _prepare_feed_dict(
             aparam = np.reshape(aparam, [nframes, natoms * fdim])
 
         # make natoms_vec and default_mesh
-        natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type)
-        assert natoms_vec[0] == natoms
+        if self.neighbor_list is not None:
+            natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type)
+            assert natoms_vec[0] == natoms
+            mesh = make_default_mesh(pbc, mixed_type)
+        else:
+            natoms_vec, coords, atom_types, mesh, imap = self.build_neighbor_list(
+                coords,
+                cells if cells is not None else None,
+                atom_types,
+                imap,
+                self.neighbor_list,
+            )
 
         # evaluate
         feed_dict_test = {}
@@ -501,7 +513,7 @@ def _prepare_feed_dict(
             raise RuntimeError
         if self.has_efield:
             feed_dict_test[self.t_efield] = np.reshape(efield, [-1])
-        feed_dict_test[self.t_mesh] = make_default_mesh(pbc, mixed_type)
+        feed_dict_test[self.t_mesh] = mesh
         if self.has_fparam:
             feed_dict_test[self.t_fparam] = np.reshape(fparam, [-1])
         if self.has_aparam:
@@ -526,6 +538,9 @@ def _eval_inner(
             coords, cells, atom_types, fparam, aparam, efield, mixed_type=mixed_type
         )
 
+        nloc = natoms_vec[0]
+        nall = natoms_vec[1]
+
         t_out = [self.t_energy, self.t_force, self.t_virial]
         if atomic:
             t_out += [self.t_ae, self.t_av]
@@ -556,11 +571,15 @@ def _eval_inner(
             av = self.reverse_map(np.reshape(av, [nframes, -1, 9]), imap)
 
         energy = np.reshape(energy, [nframes, 1])
-        force = np.reshape(force, [nframes, natoms, 3])
+        force = np.reshape(force, [nframes, nall, 3])
+        if nloc < nall:
+            force = force[:, :nloc, :]
         virial = np.reshape(virial, [nframes, 9])
         if atomic:
             ae = np.reshape(ae, [nframes, natoms_real, 1])
-            av = np.reshape(av, [nframes, natoms, 9])
+            av = np.reshape(av, [nframes, nall, 9])
+            if nloc < nall:
+                av = av[:, :nloc, :]
             return energy, force, virial, ae, av
         else:
             return energy, force, virial

source/tests/test_deeppot_a.py

@@ -3,6 +3,7 @@
 import shutil
 import unittest
 
+import ase.neighborlist
 import numpy as np
 from common import (
     run_dp,
@@ -1096,3 +1097,21 @@ def test_2frame_atm_all_param(self):
         np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places)
         expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis=1)
         np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places)
+
+
+class TestDeepPotAPBCNeighborList(TestDeepPotAPBC):
+    @classmethod
+    def setUpClass(cls):
+        convert_pbtxt_to_pb(
+            str(tests_path / os.path.join("infer", "deeppot.pbtxt")), "deeppot.pb"
+        )
+        cls.dp = DeepPot(
+            "deeppot.pb",
+            neighbor_list=ase.neighborlist.NewPrimitiveNeighborList(
+                cutoffs=6, bothways=True
+            ),
+        )
+
+    @unittest.skip("Zero atoms not supported")
+    def test_zero_input(self):
+        pass