From 3779959e9150030b70f166e00cf05735f85ca762 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 27 Sep 2024 15:55:27 +0800 Subject: [PATCH] add type hint and doc --- deepmd/utils/econf_embd.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/deepmd/utils/econf_embd.py b/deepmd/utils/econf_embd.py index a9286dd7cd..7f12206ae3 100644 --- a/deepmd/utils/econf_embd.py +++ b/deepmd/utils/econf_embd.py @@ -1,4 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Dict, + List, +) + import numpy as np from mendeleev import ( element, @@ -222,7 +227,10 @@ def make_element_embedding_list_vec( return ret -def make_econf_embedding(types, flatten=True): +def make_econf_embedding( + types: List[str], flatten: bool = True +) -> Dict[str, np.ndarray]: + """Make the electronic configuration embedding.""" all_ret = {} for ii in types: ir = make_element_embedding_list_vec(ii) @@ -232,7 +240,8 @@ def make_econf_embedding(types, flatten=True): return all_ret -def transform_to_spin_rep(res): +def transform_to_spin_rep(res: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: + """Tranform electron occupation of 0/1/2 to -1,-1/-1,1/1,1.""" ret = {} def transform(ii): @@ -252,7 +261,8 @@ def transform(ii): return ret -def print_econf_embedding(res): +def print_econf_embedding(res: Dict[str, np.ndarray]): + """Print electron configuration embedding.""" for kk, vv in res.items(): vvstr = ",".join([str(ii) for ii in vv]) space = " " * (2 - len(kk))