provide and use normalized electronic conf embedding

iProzd · Sep 27, 2024 · 6b7780e · 6b7780e
1 parent b61d65f
commit 6b7780e
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py
@@ -207,7 +207,9 @@ def change_type_map(
 def get_econf_tebd(type_map, precision: str = "default"):
     from deepmd.utils.econf_embd import (
         ECONF_DIM,
-        electronic_configuration_embedding,
+    )
+    from deepmd.utils.econf_embd import (
+        normalized_electronic_configuration_embedding as electronic_configuration_embedding,
     )
     from deepmd.utils.econf_embd import type_map as periodic_table
 

diff --git a/deepmd/utils/econf_embd.py b/deepmd/utils/econf_embd.py
@@ -4,6 +4,13 @@
     element,
 )
 
+__all__ = [
+    "electronic_configuration_embedding",
+    "normalized_electronic_configuration_embedding",
+    "make_econf_embedding",
+    "transform_to_spin_rep",
+]
+
 ###
 # made by command
 # ret = make_econf_embedding(type_map, flatten=True)
@@ -173,6 +180,16 @@
 ECONF_DIM = electronic_configuration_embedding[type_map[0]].shape[0]
 
 
+def normalize_vec_length(res):
+    scale = 1.0 / np.sqrt(ECONF_DIM)
+    return {kk: scale * vv for kk, vv in res.items()}
+
+
+normalized_electronic_configuration_embedding = normalize_vec_length(
+    electronic_configuration_embedding
+)
+
+
 def make_empty_list_vec():
     ret = {}
     for kk in conf_keys:

diff --git a/source/tests/common/test_econf_embd.py b/source/tests/common/test_econf_embd.py
@@ -4,6 +4,7 @@
 from deepmd.utils.econf_embd import (
     electronic_configuration_embedding,
     make_econf_embedding,
+    normalized_electronic_configuration_embedding,
     transform_to_spin_rep,
 )
 
@@ -55,3 +56,7 @@ def test_dict(self):
         expected_res = [1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,-1,1,-1,1,-1,1,-1,1,1,1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1]
         # fmt: on
         self.assertEqual(list(res), expected_res)
+        res = normalized_electronic_configuration_embedding["Fe"]
+        self.assertEqual(
+            list(res), [ii / len(expected_res) ** 0.5 for ii in expected_res]
+        )