set explict ntyps

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

source/tests/tf/test_data_large_batch.py

@@ -118,6 +118,7 @@ def test_data_mixed_type(self):
         fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True)
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             resnet_dt=typeebd_param["resnet_dt"],
             activation_function=None,
@@ -316,6 +317,7 @@ def test_stripped_data_mixed_type(self):
         fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True)
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             resnet_dt=typeebd_param["resnet_dt"],
             activation_function=None,
@@ -514,6 +516,7 @@ def test_compressible_data_mixed_type(self):
         fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True)
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             resnet_dt=typeebd_param["resnet_dt"],
             activation_function=None,

source/tests/tf/test_descrpt_hybrid.py

@@ -63,6 +63,7 @@ def test_descriptor_hybrid(self):
 
         # init models
         typeebd = TypeEmbedNet(
+            ntypes=ntypes,
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],

source/tests/tf/test_descrpt_se_a_type.py

@@ -60,6 +60,7 @@ def test_descriptor_two_sides(self):
 
         # init models
         typeebd = TypeEmbedNet(
+            ntypes=ntypes,
             neuron=typeebd_param["neuron"],
             resnet_dt=typeebd_param["resnet_dt"],
             seed=typeebd_param["seed"],
@@ -224,6 +225,7 @@ def test_descriptor_one_side(self):
 
         # init models
         typeebd = TypeEmbedNet(
+            ntypes=ntypes,
             neuron=typeebd_param["neuron"],
             resnet_dt=typeebd_param["resnet_dt"],
             seed=typeebd_param["seed"],

source/tests/tf/test_descrpt_se_atten.py

@@ -69,6 +69,7 @@ def test_descriptor_two_sides(self):
 
         # init models
         typeebd = TypeEmbedNet(
+            ntypes=ntypes,
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],
@@ -246,6 +247,7 @@ def test_descriptor_one_side(self):
 
         # init models
         typeebd = TypeEmbedNet(
+            ntypes=ntypes,
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],
@@ -423,6 +425,7 @@ def test_stripped_type_embedding_descriptor_two_sides(self):
 
         # init models
         typeebd = TypeEmbedNet(
+            ntypes=ntypes,
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],
@@ -594,6 +597,7 @@ def test_compressible_descriptor_two_sides(self):
 
         # init models
         typeebd = TypeEmbedNet(
+            ntypes=ntypes,
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],

source/tests/tf/test_dipole_se_a_tebd.py

@@ -72,6 +72,7 @@ def test_model(self):
         fitting = DipoleFittingSeA(**jdata["model"]["fitting_net"], uniform_seed=True)
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             resnet_dt=typeebd_param["resnet_dt"],
             seed=typeebd_param["seed"],

source/tests/tf/test_model_se_a.py

@@ -299,7 +299,10 @@ def test_model_atom_ener_type_embedding(self):
         test_data = data.get_test()
         numb_test = 1
 
-        typeebd = TypeEmbedNet(**jdata["model"]["type_embeding"])
+        typeebd = TypeEmbedNet(
+            ntypes=len(jdata["model"]["descriptor"]["type_map"]),
+            **jdata["model"]["type_embeding"],
+        )
         jdata["model"]["descriptor"].pop("type", None)
         descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True)
         jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes()

source/tests/tf/test_model_se_a_ebd_v2.py

@@ -60,6 +60,7 @@ def test_model(self):
         jdata["model"]["type_embedding"]["seed"] = 1
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=len(jdata["model"]["descriptor"]["type_map"]),
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],

source/tests/tf/test_model_se_a_type.py

@@ -61,6 +61,7 @@ def test_model(self):
         fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True)
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             resnet_dt=typeebd_param["resnet_dt"],
             seed=typeebd_param["seed"],

source/tests/tf/test_model_se_atten.py

@@ -73,6 +73,7 @@ def test_model(self):
         fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True)
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],
@@ -232,6 +233,7 @@ def test_exclude_types(self):
         descrpt = DescrptSeAtten(ntypes=ntypes, **jdata["model"]["descriptor"])
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],
@@ -300,6 +302,7 @@ def test_compressible_model(self):
         fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True)
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],
@@ -461,6 +464,7 @@ def test_compressible_exclude_types(self):
         jdata["model"]["descriptor"]["stripped_type_embedding"] = True
         jdata["model"]["descriptor"]["attn_layer"] = 0
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],
@@ -529,6 +533,7 @@ def test_stripped_type_embedding_model(self):
         fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True)
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],
@@ -693,6 +698,7 @@ def test_stripped_type_embedding_exclude_types(self):
         jdata["model"]["descriptor"]["stripped_type_embedding"] = True
         jdata["model"]["descriptor"]["attn_layer"] = 2
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],
@@ -769,6 +775,7 @@ def test_smoothness_of_stripped_type_embedding_smooth_model(self):
         fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True)
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             activation_function=None,
             resnet_dt=typeebd_param["resnet_dt"],

source/tests/tf/test_polar_se_a_tebd.py

@@ -71,6 +71,7 @@ def test_model(self):
         fitting = PolarFittingSeA(**jdata["model"]["fitting_net"], uniform_seed=True)
         typeebd_param = jdata["model"]["type_embedding"]
         typeebd = TypeEmbedNet(
+            ntypes=descrpt.get_ntypes(),
             neuron=typeebd_param["neuron"],
             resnet_dt=typeebd_param["resnet_dt"],
             seed=typeebd_param["seed"],

source/tests/tf/test_type_embed.py

@@ -28,7 +28,7 @@ def test_embed_atom_type(self):
         np.testing.assert_almost_equal(atom_embed, expected_out, 10)
 
     def test_type_embed_net(self):
-        ten = TypeEmbedNet([2, 4, 8], seed=1, uniform_seed=True)
+        ten = TypeEmbedNet(ntypes=2, neuron=[2, 4, 8], seed=1, uniform_seed=True)
         type_embedding = ten.build(2)
         sess = self.cached_session().__enter__()
         sess.run(tf.global_variables_initializer())