From 6b7780e793a36c4303024a89250624bc0f661e45 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 27 Sep 2024 13:59:31 +0800
Subject: [PATCH] provide and use normalized electronic conf embedding

---
 deepmd/dpmodel/utils/type_embed.py     |  4 +++-
 deepmd/utils/econf_embd.py             | 17 +++++++++++++++++
 source/tests/common/test_econf_embd.py |  5 +++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py
index 1b02b61de8..35f2aab617 100644
--- 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
index c29f700245..a9286dd7cd 100644
--- 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
index b27f69c38a..242ea9ca65 100644
--- 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]
+        )