diff --git a/deepmd/common.py b/deepmd/common.py index 0d98f500a5..6a18cda677 100644 --- a/deepmd/common.py +++ b/deepmd/common.py @@ -448,28 +448,6 @@ def expand_sys_str(root_dir: Union[str, Path]) -> List[str]: return matches -def docstring_parameter(*sub: Tuple[str, ...]): - """Add parameters to object docstring. - - Parameters - ---------- - sub: Tuple[str, ...] - list of strings that will be inserted into prepared locations in docstring. - - Notes - ----- - Can be used on both object and classes. - """ - - @wraps - def dec(obj: "_OBJ") -> "_OBJ": - if obj.__doc__ is not None: - obj.__doc__ = obj.__doc__.format(*sub) - return obj - - return dec - - def get_np_precision(precision: "_PRECISION") -> np.dtype: """Get numpy precision constant from string. diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index ac1e1df2d0..ff7549b124 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -3,8 +3,7 @@ from typing import Tuple, List, Dict, Any from deepmd.env import tf -from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, cast_precision -from deepmd.utils.argcheck import list_to_doc +from deepmd.common import get_activation_func, get_precision, cast_precision from deepmd.env import GLOBAL_TF_FLOAT_PRECISION from deepmd.env import GLOBAL_NP_FLOAT_PRECISION from deepmd.env import op_module @@ -88,9 +87,9 @@ class DescrptSeA (DescrptSe): set_davg_zero Set the shift of embedding net input to zero. activation_function - The activation function in the embedding net. Supported options are {0} + The activation function in the embedding net. Supported options are |ACTIVATION_FN| precision - The precision of the embedding net parameters. Supported options are {1} + The precision of the embedding net parameters. Supported options are |PRECISION| uniform_seed Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed @@ -101,7 +100,6 @@ class DescrptSeA (DescrptSe): systems. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441–4451. """ - @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) def __init__ (self, rcut: float, rcut_smth: float, diff --git a/deepmd/descriptor/se_a_ef.py b/deepmd/descriptor/se_a_ef.py index a6c611837b..cfe9c25d46 100644 --- a/deepmd/descriptor/se_a_ef.py +++ b/deepmd/descriptor/se_a_ef.py @@ -2,8 +2,7 @@ from typing import Tuple, List from deepmd.env import tf -from deepmd.common import add_data_requirement,get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter -from deepmd.utils.argcheck import list_to_doc +from deepmd.common import add_data_requirement from deepmd.utils.sess import run_sess from deepmd.env import GLOBAL_TF_FLOAT_PRECISION from deepmd.env import GLOBAL_NP_FLOAT_PRECISION @@ -43,13 +42,12 @@ class DescrptSeAEf (Descriptor): set_davg_zero Set the shift of embedding net input to zero. activation_function - The activation function in the embedding net. Supported options are {0} + The activation function in the embedding net. Supported options are |ACTIVATION_FN| precision - The precision of the embedding net parameters. Supported options are {1} + The precision of the embedding net parameters. Supported options are |PRECISION| uniform_seed Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed """ - @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) def __init__(self, rcut: float, rcut_smth: float, diff --git a/deepmd/descriptor/se_r.py b/deepmd/descriptor/se_r.py index e8116164ca..929007bcf5 100644 --- a/deepmd/descriptor/se_r.py +++ b/deepmd/descriptor/se_r.py @@ -2,8 +2,7 @@ from typing import Tuple, List from deepmd.env import tf -from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, cast_precision -from deepmd.utils.argcheck import list_to_doc +from deepmd.common import get_activation_func, get_precision, cast_precision from deepmd.env import GLOBAL_TF_FLOAT_PRECISION from deepmd.env import GLOBAL_NP_FLOAT_PRECISION from deepmd.env import op_module @@ -46,13 +45,12 @@ class DescrptSeR (DescrptSe): The excluded pairs of types which have no interaction with each other. For example, `[[0, 1]]` means no interaction between type 0 and type 1. activation_function - The activation function in the embedding net. Supported options are {0} + The activation function in the embedding net. Supported options are |ACTIVATION_FN| precision - The precision of the embedding net parameters. Supported options are {1} + The precision of the embedding net parameters. Supported options are |PRECISION| uniform_seed Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed """ - @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) def __init__ (self, rcut: float, rcut_smth: float, diff --git a/deepmd/descriptor/se_t.py b/deepmd/descriptor/se_t.py index 5ba169683a..f9453f17c3 100644 --- a/deepmd/descriptor/se_t.py +++ b/deepmd/descriptor/se_t.py @@ -2,8 +2,7 @@ from typing import Tuple, List from deepmd.env import tf -from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, cast_precision -from deepmd.utils.argcheck import list_to_doc +from deepmd.common import get_activation_func, get_precision, cast_precision from deepmd.env import GLOBAL_TF_FLOAT_PRECISION from deepmd.env import GLOBAL_NP_FLOAT_PRECISION from deepmd.env import op_module @@ -44,13 +43,12 @@ class DescrptSeT (DescrptSe): set_davg_zero Set the shift of embedding net input to zero. activation_function - The activation function in the embedding net. Supported options are {0} + The activation function in the embedding net. Supported options are |ACTIVATION_FN| precision - The precision of the embedding net parameters. Supported options are {1} + The precision of the embedding net parameters. Supported options are |PRECISION| uniform_seed Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed """ - @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) def __init__ (self, rcut: float, rcut_smth: float, diff --git a/deepmd/fit/dipole.py b/deepmd/fit/dipole.py index c0967bb554..383ea17f1f 100644 --- a/deepmd/fit/dipole.py +++ b/deepmd/fit/dipole.py @@ -3,8 +3,7 @@ from typing import Tuple, List from deepmd.env import tf -from deepmd.common import add_data_requirement, get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, cast_precision -from deepmd.utils.argcheck import list_to_doc +from deepmd.common import add_data_requirement, get_activation_func, get_precision, cast_precision from deepmd.utils.network import one_layer, one_layer_rand_seed_shift from deepmd.utils.graph import get_fitting_net_variables_from_graph_def from deepmd.descriptor import DescrptSeA @@ -31,13 +30,12 @@ class DipoleFittingSeA (Fitting) : seed : int Random seed for initializing the network parameters. activation_function : str - The activation function in the embedding net. Supported options are {0} + The activation function in the embedding net. Supported options are |ACTIVATION_FN| precision : str - The precision of the embedding net parameters. Supported options are {1} + The precision of the embedding net parameters. Supported options are |PRECISION| uniform_seed Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed """ - @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) def __init__ (self, descrpt : tf.Tensor, neuron : List[int] = [120,120,120], diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py index 01609f2acc..4084281865 100644 --- a/deepmd/fit/ener.py +++ b/deepmd/fit/ener.py @@ -4,8 +4,7 @@ from packaging.version import Version from deepmd.env import tf -from deepmd.common import add_data_requirement, get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, cast_precision -from deepmd.utils.argcheck import list_to_doc +from deepmd.common import add_data_requirement, get_activation_func, get_precision, cast_precision from deepmd.utils.network import one_layer, one_layer_rand_seed_shift from deepmd.utils.type_embed import embed_atom_type from deepmd.utils.graph import get_fitting_net_variables_from_graph_def, load_graph_def, get_tensor_by_name_from_graph @@ -72,13 +71,12 @@ class EnerFitting (Fitting): atom_ener Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. activation_function - The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are {0} + The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN| precision - The precision of the embedding net parameters. Supported options are {1} + The precision of the embedding net parameters. Supported options are |PRECISION| uniform_seed Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed """ - @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) def __init__ (self, descrpt : tf.Tensor, neuron : List[int] = [120,120,120], diff --git a/deepmd/fit/polar.py b/deepmd/fit/polar.py index 46951aeb06..3f1b7daa6b 100644 --- a/deepmd/fit/polar.py +++ b/deepmd/fit/polar.py @@ -3,8 +3,7 @@ from typing import Tuple, List from deepmd.env import tf -from deepmd.common import add_data_requirement, cast_precision, get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter -from deepmd.utils.argcheck import list_to_doc +from deepmd.common import add_data_requirement, cast_precision, get_activation_func, get_precision from deepmd.utils.network import one_layer, one_layer_rand_seed_shift from deepmd.utils.graph import get_fitting_net_variables_from_graph_def from deepmd.descriptor import DescrptLocFrame @@ -105,8 +104,33 @@ def build (self, class PolarFittingSeA (Fitting) : """ Fit the atomic polarizability with descriptor se_a + + Parameters + ---------- + descrpt : tf.Tensor + The descrptor + neuron : List[int] + Number of neurons in each hidden layer of the fitting net + resnet_dt : bool + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + sel_type : List[int] + The atom types selected to have an atomic polarizability prediction. If is None, all atoms are selected. + fit_diag : bool + Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix. + scale : List[float] + The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i] + diag_shift : List[float] + The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale. + seed : int + Random seed for initializing the network parameters. + activation_function : str + The activation function in the embedding net. Supported options are |ACTIVATION_FN| + precision : str + The precision of the embedding net parameters. Supported options are |PRECISION| + uniform_seed + Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed """ - @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) def __init__ (self, descrpt : tf.Tensor, neuron : List[int] = [120,120,120], @@ -123,32 +147,6 @@ def __init__ (self, ) -> None: """ Constructor - - Parameters - ---------- - descrpt : tf.Tensor - The descrptor - neuron : List[int] - Number of neurons in each hidden layer of the fitting net - resnet_dt : bool - Time-step `dt` in the resnet construction: - y = x + dt * \phi (Wx + b) - sel_type : List[int] - The atom types selected to have an atomic polarizability prediction. If is None, all atoms are selected. - fit_diag : bool - Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix. - scale : List[float] - The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i] - diag_shift : List[float] - The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale. - seed : int - Random seed for initializing the network parameters. - activation_function : str - The activation function in the embedding net. Supported options are {0} - precision : str - The precision of the embedding net parameters. Supported options are {1} - uniform_seed - Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed """ if not isinstance(descrpt, DescrptSeA) : raise RuntimeError('PolarFittingSeA only supports DescrptSeA') @@ -431,11 +429,10 @@ class GlobalPolarFittingSeA () : seed : int Random seed for initializing the network parameters. activation_function : str - The activation function in the embedding net. Supported options are {0} + The activation function in the embedding net. Supported options are |ACTIVATION_FN| precision : str - The precision of the embedding net parameters. Supported options are {1} + The precision of the embedding net parameters. Supported options are |PRECISION| """ - @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) def __init__ (self, descrpt : tf.Tensor, neuron : List[int] = [120,120,120], diff --git a/deepmd/fit/wfc.py b/deepmd/fit/wfc.py index 9b6b217432..564d601cae 100644 --- a/deepmd/fit/wfc.py +++ b/deepmd/fit/wfc.py @@ -3,8 +3,7 @@ from typing import Tuple, List from deepmd.env import tf -from deepmd.common import ClassArg, add_data_requirement, get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter -from deepmd.utils.argcheck import list_to_doc +from deepmd.common import ClassArg, add_data_requirement, get_activation_func, get_precision from deepmd.utils.network import one_layer, one_layer_rand_seed_shift from deepmd.descriptor import DescrptLocFrame diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 917bd406fd..ed1253d171 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -74,8 +74,13 @@ def descrpt_some_descrpt_args(): alias = tuple(alias) return self.__plugin.register((name, alias)) - def get_all_argument(self) -> List[Argument]: + def get_all_argument(self, exclude_hybrid: bool = False) -> List[Argument]: """Get all arguments. + + Parameters + ---------- + exclude_hybrid : bool + exclude hybrid descriptor to prevent circular calls Returns ------- @@ -84,6 +89,8 @@ def get_all_argument(self) -> List[Argument]: """ arguments = [] for (name, alias), metd in self.__plugin.plugins.items(): + if exclude_hybrid and name == "hybrid": + continue arguments.append(Argument(name=name, dtype=dict, sub_fields=metd(), alias=alias)) return arguments @@ -231,7 +238,7 @@ def descrpt_hybrid_args(): ] -def descrpt_variant_type_args(): +def descrpt_variant_type_args(exclude_hybrid: bool = False) -> Variant: link_lf = make_link('loc_frame', 'model/descriptor[loc_frame]') link_se_e2_a = make_link('se_e2_a', 'model/descriptor[se_e2_a]') link_se_e2_r = make_link('se_e2_r', 'model/descriptor[se_e2_r]') diff --git a/deepmd/utils/type_embed.py b/deepmd/utils/type_embed.py index 80808c4e8e..5d5b5c9888 100644 --- a/deepmd/utils/type_embed.py +++ b/deepmd/utils/type_embed.py @@ -10,8 +10,7 @@ from deepmd.utils.network import embedding_net from deepmd.utils.graph import get_type_embedding_net_variables_from_graph_def -from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, get_np_precision -from deepmd.utils.argcheck import list_to_doc +from deepmd.common import get_activation_func, get_precision def embed_atom_type( @@ -64,9 +63,9 @@ class TypeEmbedNet(): Time-step `dt` in the resnet construction: y = x + dt * \phi (Wx + b) activation_function - The activation function in the embedding net. Supported options are {0} + The activation function in the embedding net. Supported options are |ACTIVATION_FN| precision - The precision of the embedding net parameters. Supported options are {1} + The precision of the embedding net parameters. Supported options are |PRECISION| trainable If the weights of embedding net are trainable. seed @@ -74,7 +73,6 @@ class TypeEmbedNet(): uniform_seed Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed """ - @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) def __init__( self, neuron: List[int]=[], diff --git a/doc/conf.py b/doc/conf.py index 5df63bdacb..ea516fb439 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -16,6 +16,8 @@ import recommonmark from recommonmark.transform import AutoStructify from datetime import date +from deepmd.common import ACTIVATION_FN_DICT, PRECISION_DICT +from deepmd.utils.argcheck import list_to_doc def mkindex(dirname): dirname = dirname + "/" @@ -232,6 +234,11 @@ def setup(app): for typing_type in typing.__all__: numpydoc_xref_aliases[typing_type] = "typing.%s" % typing_type +rst_epilog = """ +.. |ACTIVATION_FN| replace:: %s +.. |PRECISION| replace:: %s +""" % (list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + # -- Options for HTML output ------------------------------------------------- # The theme to use for HTML and HTML Help pages. See the documentation for