check/update sel for se_atten

deepmd/entrypoints/main.py

@@ -445,6 +445,12 @@ def main_parser() -> argparse.ArgumentParser:
         required=True,
         help="type map",
     )
+    parser_neighbor_stat.add_argument(
+        "--one-type",
+        action="store_true",
+        default=False,
+        help="treat all types as a single type. Used with se_atten descriptor.",
+    )
 
     # --version
     parser.add_argument('--version', action='version', version='DeePMD-kit v%s' % __version__)

deepmd/entrypoints/neighbor_stat.py

@@ -12,6 +12,7 @@ def neighbor_stat(
     system: str,
     rcut: float,
     type_map: List[str],
+    one_type: bool = False,
     **kwargs,
 ):
     """Calculate neighbor statistics.
@@ -24,6 +25,8 @@ def neighbor_stat(
         cutoff radius
     type_map : list[str]
         type map
+    one_type : bool, optional, default=False
+        treat all types as a single type
 
     Examples
     --------
@@ -42,7 +45,7 @@ def neighbor_stat(
         type_map=type_map,
     )
     data.get_batch()
-    nei = NeighborStat(data.get_ntypes(), rcut)
+    nei = NeighborStat(data.get_ntypes(), rcut, one_type=one_type)
     min_nbor_dist, max_nbor_size = nei.get_stat(data)
     log.info("min_nbor_dist: %f" % min_nbor_dist)
     log.info("max_nbor_size: %s" % str(max_nbor_size))

deepmd/entrypoints/train.py

@@ -90,10 +90,7 @@ def train(
 
     jdata = normalize(jdata)
 
-    if jdata['model']['descriptor']['type'] in ['se_atten'] and isinstance(jdata['model']['descriptor']['sel'], list):
-        jdata['model']['descriptor']['sel'] = sum(jdata['model']['descriptor']['sel'])
-
-    if not is_compress and not skip_neighbor_stat and jdata['model']['descriptor']['type'] not in ['se_atten']:
+    if not is_compress and not skip_neighbor_stat:
         jdata = update_sel(jdata)
 
     with open(output, "w") as fp:
@@ -253,7 +250,7 @@ def get_type_map(jdata):
     return jdata['model'].get('type_map', None)
 
 
-def get_nbor_stat(jdata, rcut):
+def get_nbor_stat(jdata, rcut, one_type: bool = False):
     max_rcut = get_rcut(jdata)
     type_map = get_type_map(jdata)
 
@@ -268,7 +265,7 @@ def get_nbor_stat(jdata, rcut):
         map_ntypes = data_ntypes
     ntypes = max([map_ntypes, data_ntypes])
 
-    neistat = NeighborStat(ntypes, rcut)
+    neistat = NeighborStat(ntypes, rcut, one_type=one_type)
 
     min_nbor_dist, max_nbor_size = neistat.get_stat(train_data)
 
@@ -283,8 +280,8 @@ def get_nbor_stat(jdata, rcut):
         dtype = tf.int32)
     return min_nbor_dist, max_nbor_size
 
-def get_sel(jdata, rcut):
-    _, max_nbor_size = get_nbor_stat(jdata, rcut)
+def get_sel(jdata, rcut, one_type: bool = False):
+    _, max_nbor_size = get_nbor_stat(jdata, rcut, one_type=one_type)
     return max_nbor_size
 
 def get_min_nbor_dist(jdata, rcut):
@@ -320,14 +317,20 @@ def wrap_up_4(xx):
 
 
 def update_one_sel(jdata, descriptor):
+    if descriptor['type'] == 'loc_frame':
+        return descriptor
     rcut = descriptor['rcut']
-    tmp_sel = get_sel(jdata, rcut)
+    tmp_sel = get_sel(jdata, rcut, one_type=descriptor['type'] in ('se_atten',))
+    sel = descriptor['sel']
+    if isinstance(sel, int):
+        # convert to list and finnally convert back to int
+        sel = [sel]
     if parse_auto_sel(descriptor['sel']) :
         ratio = parse_auto_sel_ratio(descriptor['sel'])
-        descriptor['sel'] = [int(wrap_up_4(ii * ratio)) for ii in tmp_sel]
+        descriptor['sel'] = sel = [int(wrap_up_4(ii * ratio)) for ii in tmp_sel]
     else:
         # sel is set by user
-        for ii, (tt, dd) in enumerate(zip(tmp_sel, descriptor['sel'])):
+        for ii, (tt, dd) in enumerate(zip(tmp_sel, sel)):
             if dd and tt > dd:
                 # we may skip warning for sel=0, where the user is likely
                 # to exclude such type in the descriptor
@@ -336,6 +339,8 @@ def update_one_sel(jdata, descriptor):
                     "not less than %d, but you set it to %d. The accuracy"
                     " of your model may get worse." %(ii, tt, dd)
                 )
+    if descriptor['type'] in ('se_atten',):
+        descriptor['sel'] = sel = sum(sel)
     return descriptor
 
 
@@ -344,9 +349,8 @@ def update_sel(jdata):
     descrpt_data = jdata['model']['descriptor']
     if descrpt_data['type'] == 'hybrid':
         for ii in range(len(descrpt_data['list'])):
-            if descrpt_data['list'][ii]['type'] != 'loc_frame':
-                descrpt_data['list'][ii] = update_one_sel(jdata, descrpt_data['list'][ii])
-    elif descrpt_data['type'] != 'loc_frame':
+            descrpt_data['list'][ii] = update_one_sel(jdata, descrpt_data['list'][ii])
+    else:
         descrpt_data = update_one_sel(jdata, descrpt_data)
     jdata['model']['descriptor'] = descrpt_data
     return jdata

deepmd/utils/argcheck.py

@@ -243,7 +243,8 @@ def descrpt_hybrid_args():
 def descrpt_se_atten_args():
     doc_sel = 'This parameter set the number of selected neighbors. Note that this parameter is a little different from that in other descriptors. Instead of separating each type of atoms, only the summation matters. And this number is highly related with the efficiency, thus one should not make it too large. Usually 200 or less is enough, far away from the GPU limitation 4096. It can be:\n\n\
     - `int`. The maximum number of neighbor atoms to be considered. We recommend it to be less than 200. \n\n\
-    - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. Only the summation of `sel[i]` matters, and it is recommended to be less than 200.'
+    - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. Only the summation of `sel[i]` matters, and it is recommended to be less than 200.\
+    - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".'
     doc_rcut = 'The cut-off radius.'
     doc_rcut_smth = 'Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`'
     doc_neuron = 'Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.'
@@ -262,7 +263,7 @@ def descrpt_se_atten_args():
     doc_attn_mask = 'Whether to do mask on the diagonal in the attention matrix'
 
     return [
-        Argument("sel", [int, list], optional=True, default=200, doc=doc_sel),
+        Argument("sel", [int, list, str], optional=True, default="auto", doc=doc_sel),
         Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut),
         Argument("rcut_smth", float, optional=True, default=0.5, doc=doc_rcut_smth),
         Argument("neuron", list, optional=True, default=[10, 20, 40], doc=doc_neuron),

deepmd/utils/neighbor_stat.py

@@ -23,15 +23,20 @@ class NeighborStat():
             The num of atom types
     rcut
             The cut-off radius
+    one_type : bool, optional, default=False
+        Treat all types as a single type.
     """
     def __init__(self,
                  ntypes : int,
-                 rcut: float) -> None:
+                 rcut: float,
+                 one_type : bool = False,
+                 ) -> None:
         """
         Constructor
         """
         self.rcut = rcut
         self.ntypes = ntypes
+        self.one_type = one_type
         sub_graph = tf.Graph()
 
         def builder():
@@ -41,10 +46,17 @@ def builder():
             place_holders['type'] = tf.placeholder(tf.int32, [None, None], name='t_type')
             place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name='t_natoms')
             place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name='t_mesh')
+            t_type = place_holders['type']
+            t_natoms = place_holders['natoms_vec']
+            if self.one_type:
+                # all types = 0, natoms_vec = [natoms, natoms, natoms]
+                t_type = tf.zeros_like(t_type, dtype=tf.int32)
+                t_natoms = tf.repeat(t_natoms[0], 3)
+
             _max_nbor_size, _min_nbor_dist \
                 = op_module.neighbor_stat(place_holders['coord'],
-                                         place_holders['type'],
-                                         place_holders['natoms_vec'],
+                                         t_type,
+                                         t_natoms,
                                          place_holders['box'],
                                          place_holders['default_mesh'],
                                          rcut = self.rcut)
@@ -74,7 +86,9 @@ def get_stat(self,
                 A list with ntypes integers, denotes the actual achieved max sel
         """
         self.min_nbor_dist = 100.0
-        self.max_nbor_size = [0] * self.ntypes
+        self.max_nbor_size = [0]
+        if not self.one_type:
+            self.max_nbor_size *= self.ntypes            
 
         def feed():
             for ii in range(len(data.system_dirs)):

source/tests/test_train.py

@@ -29,13 +29,15 @@ def test_update_one_sel(self, sel_mock):
         sel_mock.return_value = [10,20]
         jdata = {}
         descriptor = {
+            'type': 'se_e2_a',
             'rcut': 6,
             'sel': "auto"
         }
         descriptor = update_one_sel(jdata, descriptor)
         # self.assertEqual(descriptor['sel'], [11,22])
         self.assertEqual(descriptor['sel'], [12,24])
         descriptor = {
+            'type': 'se_e2_a',
             'rcut': 6,
             'sel': "auto:1.5"
         }
@@ -113,7 +115,98 @@ def test_update_sel(self, sel_mock):
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-
+    @patch("deepmd.entrypoints.train.get_sel")
+    def test_update_sel_atten_auto(self, sel_mock):
+        sel_mock.return_value = [25]
+        jdata = {
+            'model' : {
+                'descriptor': {
+                    'type' : 'se_atten',
+                    'sel' : "auto",
+                    'rcut': 6,
+                }
+            }
+        }
+        expected_out = {
+            'model' : {
+                'descriptor': {
+                    'type' : 'se_atten',
+                    'sel' : 28,
+                    'rcut': 6,
+                }
+            }
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    @patch("deepmd.entrypoints.train.get_sel")
+    def test_update_sel_atten_int(self, sel_mock):
+        sel_mock.return_value = [25]
+        jdata = {
+            'model' : {
+                'descriptor': {
+                    'type' : 'se_atten',
+                    'sel' : 30,
+                    'rcut': 6,
+                }
+            }
+        }
+        expected_out = {
+            'model' : {
+                'descriptor': {
+                    'type' : 'se_atten',
+                    'sel' : 30,
+                    'rcut': 6,
+                }
+            }
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    @patch("deepmd.entrypoints.train.get_sel")
+    def test_update_sel_atten_list(self, sel_mock):
+        sel_mock.return_value = [25]
+        jdata = {
+            'model' : {
+                'descriptor': {
+                    'type' : 'se_atten',
+                    'sel' : 30,
+                    'rcut': 6,
+                }
+            }
+        }
+        expected_out = {
+            'model' : {
+                'descriptor': {
+                    'type' : 'se_atten',
+                    'sel' : 30,
+                    'rcut': 6,
+                }
+            }
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    def test_skip_loc_frame(self):
+        jdata = {
+            'model' : {
+                'descriptor': {
+                    'type' : 'loc_frame',
+                    'rcut': 6,
+                }
+            }
+        }
+        expected_out = {
+            'model' : {
+                'descriptor': {
+                    'type' : 'loc_frame',
+                    'rcut': 6,
+                }
+            }
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
     def test_wrap_up_4(self):
         self.assertEqual(wrap_up_4(12), 3 * 4)
         self.assertEqual(wrap_up_4(13), 4 * 4)