set specific mesh shapes for mixed type

deepmd/common.py

@@ -205,28 +205,29 @@ def select_idx_map(atom_types: np.ndarray, select_types: np.ndarray) -> np.ndarr
     return np.concatenate(idx_map)
 
 
-# TODO not really sure if the docstring is right the purpose of this is a bit unclear
-def make_default_mesh(test_box: np.ndarray, cell_size: float = 3.0) -> np.ndarray:
-    """Get number of cells of size=`cell_size` fit into average box.
+def make_default_mesh(pbc: bool, mixed_type: bool) -> np.ndarray:
+    """Make mesh.
+
+    Only the size of mesh matters, not the values:
+    * 6 for PBC, no mixed types
+    * 0 for no PBC, no mixed types
+    * 7 for PBC, mixed types
+    * 1 for no PBC, mixed types
 
     Parameters
     ----------
-    test_box : np.ndarray
-        numpy array with cells of shape Nx9
-    cell_size : float, optional
-        length of one cell, by default 3.0
+    pbc : bool
+        if True, the mesh will be made for periodic boundary conditions
+    mixed_type : bool
+        if True, the mesh will be made for mixed types
 
     Returns
     -------
     np.ndarray
-        mesh for supplied boxes, how many cells fit in each direction
+        mesh
     """
-    cell_lengths = np.linalg.norm(test_box.reshape([-1, 3, 3]), axis=2)
-    avg_cell_lengths = np.average(cell_lengths, axis=0)
-    ncell = (avg_cell_lengths / cell_size).astype(np.int32)
-    ncell[ncell < 2] = 2
-    default_mesh = np.zeros(6, dtype=np.int32)
-    default_mesh[3:6] = ncell
+    mesh_size = int(pbc) * 6 + int(mixed_type)
+    default_mesh = np.zeros(mesh_size, dtype=np.int32)
     return default_mesh
 
 

deepmd/infer/data_modifier.py

@@ -343,7 +343,7 @@ def _eval_fv(self, coords, cells, atom_types, ext_f):
         # make natoms_vec and default_mesh
         natoms_vec = self.make_natoms_vec(atom_types)
         assert natoms_vec[0] == natoms
-        default_mesh = make_default_mesh(cells)
+        default_mesh = make_default_mesh(True, False)
 
         # evaluate
         tensor = []

deepmd/infer/deep_dos.py

@@ -373,10 +373,7 @@ def _prepare_feed_dict(
             feed_dict_test[self.t_box] = cells
         else:
             raise RuntimeError
-        if pbc:
-            feed_dict_test[self.t_mesh] = make_default_mesh(cells)
-        else:
-            feed_dict_test[self.t_mesh] = np.array([], dtype=np.int32)
+        feed_dict_test[self.t_mesh] = make_default_mesh(pbc, mixed_type)
         if self.has_fparam:
             feed_dict_test[self.t_fparam] = np.reshape(fparam, [-1])
         if self.has_aparam:

deepmd/infer/deep_pot.py

@@ -440,10 +440,7 @@ def _prepare_feed_dict(
             raise RuntimeError
         if self.has_efield:
             feed_dict_test[self.t_efield] = np.reshape(efield, [-1])
-        if pbc:
-            feed_dict_test[self.t_mesh] = make_default_mesh(cells)
-        else:
-            feed_dict_test[self.t_mesh] = np.array([], dtype=np.int32)
+        feed_dict_test[self.t_mesh] = make_default_mesh(pbc, mixed_type)
         if self.has_fparam:
             feed_dict_test[self.t_fparam] = np.reshape(fparam, [-1])
         if self.has_aparam:

deepmd/infer/deep_tensor.py

@@ -213,10 +213,7 @@ def eval(
             )
         feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
         feed_dict_test[self.t_box] = np.reshape(cells, [-1])
-        if pbc:
-            feed_dict_test[self.t_mesh] = make_default_mesh(cells)
-        else:
-            feed_dict_test[self.t_mesh] = np.array([], dtype=np.int32)
+        feed_dict_test[self.t_mesh] = make_default_mesh(pbc, mixed_type)
 
         if atomic:
             assert (
@@ -340,10 +337,7 @@ def eval_full(
             )
         feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
         feed_dict_test[self.t_box] = np.reshape(cells, [-1])
-        if pbc:
-            feed_dict_test[self.t_mesh] = make_default_mesh(cells)
-        else:
-            feed_dict_test[self.t_mesh] = np.array([], dtype=np.int32)
+        feed_dict_test[self.t_mesh] = make_default_mesh(pbc, mixed_type)
 
         t_out = [self.t_global_tensor, self.t_force, self.t_virial]
         if atomic:

deepmd/utils/data_system.py

@@ -1,13 +1,19 @@
 import collections
 import logging
 import warnings
+from functools import (
+    lru_cache,
+)
 from typing import (
     List,
     Optional,
 )
 
 import numpy as np
 
+from deepmd.common import (
+    make_default_mesh,
+)
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
@@ -218,23 +224,16 @@ def _load_test(self, ntests=-1):
             for nn in test_system_data:
                 self.test_data[nn].append(test_system_data[nn])
 
-    def _make_default_mesh(self):
-        self.default_mesh = []
-        cell_size = np.max(self.rcut)
-        for ii in range(self.nsystems):
-            if self.data_systems[ii].pbc:
-                test_system_data = self.data_systems[ii].get_batch(self.batch_size[ii])
-                self.data_systems[ii].reset_get_batch()
-                # test_system_data = self.data_systems[ii].get_test()
-                avg_box = np.average(test_system_data["box"], axis=0)
-                avg_box = np.reshape(avg_box, [3, 3])
-                ncell = (np.linalg.norm(avg_box, axis=1) / cell_size).astype(np.int32)
-                ncell[ncell < 2] = 2
-                default_mesh = np.zeros(6, dtype=np.int32)
-                default_mesh[3:6] = ncell
-                self.default_mesh.append(default_mesh)
-            else:
-                self.default_mesh.append(np.array([], dtype=np.int32))
+    @property
+    @lru_cache(maxsize=None)
+    def default_mesh(self) -> List[np.ndarray]:
+        """Mesh for each system."""
+        return [
+            make_default_mesh(
+                self.data_systems[ii].pbc, self.data_systems[ii].mixed_type
+            )
+            for ii in range(self.nsystems)
+        ]
 
     def compute_energy_shift(self, rcond=1e-3, key="energy"):
         sys_ener = []
@@ -391,8 +390,6 @@ def get_batch(self, sys_idx: Optional[int] = None) -> dict:
         dict
             The batch data
         """
-        if not hasattr(self, "default_mesh"):
-            self._make_default_mesh()
         if not self.mixed_systems:
             b_data = self.get_batch_standard(sys_idx)
         else:
@@ -505,8 +502,6 @@ def get_test(self, sys_idx: Optional[int] = None, n_test: int = -1):  # deprecia
         n_test
             Number of test data. If set to -1 all test data will be get.
         """
-        if not hasattr(self, "default_mesh"):
-            self._make_default_mesh()
         if not hasattr(self, "test_data"):
             self._load_test(ntests=n_test)
         if sys_idx is not None:

source/api_cc/src/common.cc

@@ -652,9 +652,9 @@
   box_shape.AddDim(9);
   TensorShape mesh_shape;
   if (b_pbc) {
-    mesh_shape.AddDim(6);
+    mesh_shape.AddDim(7);
   } else {
-    mesh_shape.AddDim(0);
+    mesh_shape.AddDim(1);
   }
   TensorShape natoms_shape;
   natoms_shape.AddDim(2 + ntypes);
@@ -723,6 +723,9 @@
     mesh(4 - 1) = 0;
     mesh(5 - 1) = 0;
     mesh(6 - 1) = 0;
+    mesh(7 - 1) = 0;
+  } else {
+    mesh(1 - 1) = 0;
   }
   natoms(0) = nloc;
   natoms(1) = nall;

source/op/descrpt.cc

@@ -124,6 +124,10 @@
     } else if (mesh_tensor.shape().dim_size(0) == 0) {
       // no pbc
       nei_mode = -1;
+    } else if (mesh_tensor.shape().dim_size(0) == 7 ||
+               mesh_tensor.shape().dim_size(0) == 1) {
+      throw deepmd::deepmd_exception(
+          "Mixed types are not supported by this OP.");
     } else {
       throw deepmd::deepmd_exception("invalid mesh tensor");
     }

source/op/descrpt_se_a_ef.cc

@@ -140,6 +140,10 @@
     } else if (mesh_tensor.shape().dim_size(0) == 0) {
       // no pbc
       nei_mode = -1;
+    } else if (mesh_tensor.shape().dim_size(0) == 7 ||
+               mesh_tensor.shape().dim_size(0) == 1) {
+      throw deepmd::deepmd_exception(
+          "Mixed types are not supported by this OP.");
     } else {
       throw deepmd::deepmd_exception("invalid mesh tensor");
     }

source/op/descrpt_se_a_ef_para.cc

@@ -140,6 +140,10 @@
     } else if (mesh_tensor.shape().dim_size(0) == 0) {
       // no pbc
       nei_mode = -1;
+    } else if (mesh_tensor.shape().dim_size(0) == 7 ||
+               mesh_tensor.shape().dim_size(0) == 1) {
+      throw deepmd::deepmd_exception(
+          "Mixed types are not supported by this OP.");
     } else {
       throw deepmd::deepmd_exception("invalid mesh tensor");
     }

source/op/descrpt_se_a_ef_vert.cc

@@ -140,6 +140,10 @@
     } else if (mesh_tensor.shape().dim_size(0) == 0) {
       // no pbc
       nei_mode = -1;
+    } else if (mesh_tensor.shape().dim_size(0) == 7 ||
+               mesh_tensor.shape().dim_size(0) == 1) {
+      throw deepmd::deepmd_exception(
+          "Mixed types are not supported by this OP.");
     } else {
       throw deepmd::deepmd_exception("invalid mesh tensor");
     }

source/op/neighbor_stat.cc

@@ -67,11 +67,13 @@
                 errors::InvalidArgument("number of box should be 9"));
 
     int nei_mode = 0;
-    if (mesh_tensor.shape().dim_size(0) == 6) {
+    if (mesh_tensor.shape().dim_size(0) == 6 ||
+        mesh_tensor.shape().dim_size(0) == 7) {
       // manual copied pbc
       assert(nloc == nall);
       nei_mode = 1;
-    } else if (mesh_tensor.shape().dim_size(0) == 0) {
+    } else if (mesh_tensor.shape().dim_size(0) == 0 ||
+               mesh_tensor.shape().dim_size(0) == 1) {
       // no pbc
       nei_mode = -1;
     } else {

source/op/prod_env_mat_multi_device.cc

@@ -506,6 +506,10 @@
       // no pbc
       assert(nloc == nall);
       nei_mode = -1;
+    } else if (mesh_tensor.shape().dim_size(0) == 7 ||
+               mesh_tensor.shape().dim_size(0) == 1) {
+      throw deepmd::deepmd_exception(
+          "Mixed types are not supported by this OP.");
     } else {
       throw deepmd::deepmd_exception("invalid mesh tensor");
     }
@@ -788,6 +792,10 @@
       // no pbc
       assert(nloc == nall);
       nei_mode = -1;
+    } else if (mesh_tensor.shape().dim_size(0) == 7 ||
+               mesh_tensor.shape().dim_size(0) == 1) {
+      throw deepmd::deepmd_exception(
+          "Mixed types are not supported by this OP.");
     } else {
       throw deepmd::deepmd_exception("invalid mesh tensor");
     }
@@ -1077,12 +1085,14 @@
     if (mesh_tensor.shape().dim_size(0) == 16) {
       // lammps neighbor list
       nei_mode = 3;
-    } else if (mesh_tensor.shape().dim_size(0) == 6) {
+    } else if (mesh_tensor.shape().dim_size(0) == 6 ||
+               mesh_tensor.shape().dim_size(0) == 7) {
       // manual copied pbc
       assert(nloc == nall);
       nei_mode = 1;
       b_nlist_map = true;
-    } else if (mesh_tensor.shape().dim_size(0) == 0) {
+    } else if (mesh_tensor.shape().dim_size(0) == 0 ||
+               mesh_tensor.shape().dim_size(0) == 1) {
       // no pbc
       assert(nloc == nall);
       nei_mode = -1;

source/op/prod_env_mat_multi_device_nvnmd.cc

@@ -364,6 +364,10 @@
       // no pbc
       assert(nloc == nall);
       nei_mode = -1;
+    } else if (mesh_tensor.shape().dim_size(0) == 7 ||
+               mesh_tensor.shape().dim_size(0) == 1) {
+      throw deepmd::deepmd_exception(
+          "Mixed types are not supported by this OP.");
     } else {
       throw deepmd::deepmd_exception("invalid mesh tensor");
     }

source/tests/test_deepmd_data_sys.py

@@ -412,6 +412,11 @@ def test_get_mixed_batch(self):
         ds.add("null", self.test_ndof, atomic=True, must=False)
         random.seed(114514)
         # with this seed, the batch is fixed, with natoms 3, 6, 6
+        # keep the random behavior before #2481
+        for ii in range(ds.nsystems):
+            if ds.data_systems[ii].pbc:
+                ds.data_systems[ii].get_batch(ds.batch_size[ii])
+                ds.data_systems[ii].reset_get_batch()
         data = ds.get_batch()
         np.testing.assert_equal(data["natoms_vec"], np.array([6, 6, 6, 0, 0]))
         np.testing.assert_equal(data["real_natoms_vec"][:, 0], np.array([3, 6, 6]))