From b9303c76b7433ddac54921ea0eb31d1bddc32183 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 19 Nov 2024 00:20:33 -0500 Subject: [PATCH] docs: move arg docs from Variant to sub Arguments (#4369) See https://github.com/deepmodeling/dargs/pull/82 ## Summary by CodeRabbit - **Documentation** - Enhanced documentation for various descriptor arguments and fitting configurations, improving clarity and usability. - Added detailed descriptions for descriptors related to the smooth edition of Deep Potential. - Updated registration of descriptor arguments to include new documentation strings for better guidance on functionality and expected parameters. Signed-off-by: Jinzhe Zeng --- deepmd/utils/argcheck.py | 81 +++++++++++++++++++--------------------- 1 file changed, 39 insertions(+), 42 deletions(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 60c85123d7..469a9ada46 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -40,6 +40,21 @@ doc_only_tf_supported = "(Supported Backend: TensorFlow) " doc_only_pt_supported = "(Supported Backend: PyTorch) " +# descriptors +doc_loc_frame = "Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame." +doc_se_e2_a = "Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor." +doc_se_e2_r = "Used by the smooth edition of Deep Potential. Only the distance between atoms is used to construct the descriptor." +doc_se_e3 = "Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Three-body embedding will be used by this descriptor." +doc_se_a_tpe = "Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Type embedding will be used by this descriptor." +doc_se_atten = "Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Attention mechanism will be used by this descriptor." +doc_se_atten_v2 = "Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Attention mechanism with new modifications will be used by this descriptor." +doc_se_a_mask = "Used by the smooth edition of Deep Potential. It can accept a variable number of atoms in a frame (Non-PBC system). *aparam* are required as an indicator matrix for the real/virtual sign of input atoms." +doc_hybrid = "Concatenate of a list of descriptors as a new descriptor." +# fitting +doc_ener = "Fit an energy model (potential energy surface)." +doc_dos = "Fit a density of states model. The total density of states / site-projected density of states labels should be provided by `dos.npy` or `atom_dos.npy` in each data system. The file has number of frames lines and number of energy grid columns (times number of atoms in `atom_dos.npy`). See `loss` parameter." +doc_dipole = "Fit an atomic dipole model. Global dipole labels or atomic dipole labels for all the selected atoms (see `sel_type`) should be provided by `dipole.npy` in each data system. The file either has number of frames lines and 3 times of number of selected atoms columns, or has number of frames lines and 3 columns. See `loss` parameter." +doc_polar = "Fit an atomic polarizability model. Global polarizazbility labels or atomic polarizability labels for all the selected atoms (see `sel_type`) should be provided by `polarizability.npy` in each data system. The file with has number of frames lines and 9 times of number of selected atoms columns, or has number of frames lines and 9 columns. See `loss` parameter." def list_to_doc(xx): @@ -229,7 +244,7 @@ def get_all_argument(self, exclude_hybrid: bool = False) -> list[Argument]: descrpt_args_plugin = ArgsPlugin() -@descrpt_args_plugin.register("loc_frame", doc=doc_only_tf_supported) +@descrpt_args_plugin.register("loc_frame", doc=doc_only_tf_supported + doc_loc_frame) def descrpt_local_frame_args(): doc_sel_a = "A list of integers. The length of the list should be the same as the number of atom types in the system. `sel_a[i]` gives the selected number of type-i neighbors. The full relative coordinates of the neighbors are used by the descriptor." doc_sel_r = "A list of integers. The length of the list should be the same as the number of atom types in the system. `sel_r[i]` gives the selected number of type-i neighbors. Only relative distance of the neighbors are used by the descriptor. sel_a[i] + sel_r[i] is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius." @@ -250,7 +265,7 @@ def descrpt_local_frame_args(): ] -@descrpt_args_plugin.register("se_e2_a", alias=["se_a"]) +@descrpt_args_plugin.register("se_e2_a", alias=["se_a"], doc=doc_se_e2_a) def descrpt_se_a_args(): doc_sel = 'This parameter set the number of selected neighbors for each type of atom. It can be:\n\n\ - `list[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ @@ -318,7 +333,9 @@ def descrpt_se_a_args(): ] -@descrpt_args_plugin.register("se_e3", alias=["se_at", "se_a_3be", "se_t"]) +@descrpt_args_plugin.register( + "se_e3", alias=["se_at", "se_a_3be", "se_t"], doc=doc_se_e3 +) def descrpt_se_t_args(): doc_sel = 'This parameter set the number of selected neighbors for each type of atom. It can be:\n\n\ - `list[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ @@ -373,7 +390,9 @@ def descrpt_se_t_args(): ] -@descrpt_args_plugin.register("se_a_tpe", alias=["se_a_ebd"], doc=doc_only_tf_supported) +@descrpt_args_plugin.register( + "se_a_tpe", alias=["se_a_ebd"], doc=doc_only_tf_supported + doc_se_a_tpe +) def descrpt_se_a_tpe_args(): doc_type_nchanl = "number of channels for type embedding" doc_type_nlayer = "number of hidden layers of type embedding net" @@ -387,7 +406,7 @@ def descrpt_se_a_tpe_args(): ] -@descrpt_args_plugin.register("se_e2_r", alias=["se_r"]) +@descrpt_args_plugin.register("se_e2_r", alias=["se_r"], doc=doc_se_e2_r) def descrpt_se_r_args(): doc_sel = 'This parameter set the number of selected neighbors for each type of atom. It can be:\n\n\ - `list[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ @@ -446,7 +465,7 @@ def descrpt_se_r_args(): ] -@descrpt_args_plugin.register("hybrid") +@descrpt_args_plugin.register("hybrid", doc=doc_hybrid) def descrpt_hybrid_args(): doc_list = "A list of descriptor definitions" @@ -538,7 +557,7 @@ def descrpt_se_atten_common_args(): ] -@descrpt_args_plugin.register("se_atten", alias=["dpa1"]) +@descrpt_args_plugin.register("se_atten", alias=["dpa1"], doc=doc_se_atten) def descrpt_se_atten_args(): doc_smooth_type_embedding = f"Whether to use smooth process in attention weights calculation. {doc_only_tf_supported} When using stripped type embedding, whether to dot smooth factor on the network output of type embedding to keep the network smooth, instead of setting `set_davg_zero` to be True." doc_set_davg_zero = "Set the normalization average to zero. This option should be set when `se_atten` descriptor or `atom_ener` in the energy fitting is used" @@ -774,7 +793,7 @@ def descrpt_se_e3_tebd_args(): ] -@descrpt_args_plugin.register("se_atten_v2") +@descrpt_args_plugin.register("se_atten_v2", doc=doc_se_atten_v2) def descrpt_se_atten_v2_args(): doc_set_davg_zero = "Set the normalization average to zero. This option should be set when `se_atten` descriptor or `atom_ener` in the energy fitting is used" doc_trainable_ln = ( @@ -1343,7 +1362,7 @@ def descrpt_se_a_ebd_v2_args(): return descrpt_se_a_args() -@descrpt_args_plugin.register("se_a_mask", doc=doc_only_tf_supported) +@descrpt_args_plugin.register("se_a_mask", doc=doc_only_tf_supported + doc_se_a_mask) def descrpt_se_a_mask_args(): doc_sel = 'This parameter sets the number of selected neighbors for each type of atom. It can be:\n\n\ - `list[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ @@ -1397,27 +1416,7 @@ def descrpt_se_a_mask_args(): def descrpt_variant_type_args(exclude_hybrid: bool = False) -> Variant: - link_lf = make_link("loc_frame", "model[standard]/descriptor[loc_frame]") - link_se_e2_a = make_link("se_e2_a", "model[standard]/descriptor[se_e2_a]") - link_se_e2_r = make_link("se_e2_r", "model[standard]/descriptor[se_e2_r]") - link_se_e3 = make_link("se_e3", "model[standard]/descriptor[se_e3]") - link_se_a_tpe = make_link("se_a_tpe", "model[standard]/descriptor[se_a_tpe]") - link_hybrid = make_link("hybrid", "model[standard]/descriptor[hybrid]") - link_se_atten = make_link("se_atten", "model[standard]/descriptor[se_atten]") - link_se_atten_v2 = make_link( - "se_atten_v2", "model[standard]/descriptor[se_atten_v2]" - ) - link_se_a_mask = make_link("se_a_mask", "model[standard]/descriptor[se_a_mask]") - doc_descrpt_type = f"The type of the descriptor. See explanation below. \n\n\ -- {link_lf}: Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.\n\n\ -- {link_se_e2_a}: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor.\n\n\ -- {link_se_e2_r}: Used by the smooth edition of Deep Potential. Only the distance between atoms is used to construct the descriptor.\n\n\ -- {link_se_e3}: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Three-body embedding will be used by this descriptor.\n\n\ -- {link_se_a_tpe}: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Type embedding will be used by this descriptor.\n\n\ -- {link_se_atten}: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Attention mechanism will be used by this descriptor.\n\n\ -- {link_se_atten_v2}: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Attention mechanism with new modifications will be used by this descriptor.\n\n\ -- {link_se_a_mask}: Used by the smooth edition of Deep Potential. It can accept a variable number of atoms in a frame (Non-PBC system). *aparam* are required as an indicator matrix for the real/virtual sign of input atoms. \n\n\ -- {link_hybrid}: Concatenate of a list of descriptors as a new descriptor." + doc_descrpt_type = "The type of the descriptor." return Variant( "type", @@ -1430,7 +1429,7 @@ def descrpt_variant_type_args(exclude_hybrid: bool = False) -> Variant: fitting_args_plugin = ArgsPlugin() -@fitting_args_plugin.register("ener") +@fitting_args_plugin.register("ener", doc=doc_ener) def fitting_ener(): doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." @@ -1506,7 +1505,7 @@ def fitting_ener(): ] -@fitting_args_plugin.register("dos") +@fitting_args_plugin.register("dos", doc=doc_dos) def fitting_dos(): doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." @@ -1594,7 +1593,7 @@ def fitting_property(): ] -@fitting_args_plugin.register("polar") +@fitting_args_plugin.register("polar", doc=doc_polar) def fitting_polar(): doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -1648,7 +1647,7 @@ def fitting_polar(): # return fitting_polar() -@fitting_args_plugin.register("dipole") +@fitting_args_plugin.register("dipole", doc=doc_dipole) def fitting_dipole(): doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -1687,11 +1686,7 @@ def fitting_dipole(): # YWolfeee: Delete global polar mode, merge it into polar mode and use loss setting to support. def fitting_variant_type_args(): - doc_descrpt_type = "The type of the fitting. See explanation below. \n\n\ -- `ener`: Fit an energy model (potential energy surface).\n\n\ -- `dos` : Fit a density of states model. The total density of states / site-projected density of states labels should be provided by `dos.npy` or `atom_dos.npy` in each data system. The file has number of frames lines and number of energy grid columns (times number of atoms in `atom_dos.npy`). See `loss` parameter. \n\n\ -- `dipole`: Fit an atomic dipole model. Global dipole labels or atomic dipole labels for all the selected atoms (see `sel_type`) should be provided by `dipole.npy` in each data system. The file either has number of frames lines and 3 times of number of selected atoms columns, or has number of frames lines and 3 columns. See `loss` parameter.\n\n\ -- `polar`: Fit an atomic polarizability model. Global polarizazbility labels or atomic polarizability labels for all the selected atoms (see `sel_type`) should be provided by `polarizability.npy` in each data system. The file with has number of frames lines and 9 times of number of selected atoms columns, or has number of frames lines and 9 columns. See `loss` parameter.\n\n" + doc_descrpt_type = "The type of the fitting." return Variant( "type", @@ -1722,12 +1717,14 @@ def modifier_dipole_charge(): def modifier_variant_type_args(): - doc_modifier_type = "The type of modifier. See explanation below.\n\n\ --`dipole_charge`: Use WFCC to model the electronic structure of the system. Correct the long-range interaction" + doc_modifier_type = "The type of modifier." + doc_dipole_charge = "Use WFCC to model the electronic structure of the system. Correct the long-range interaction." return Variant( "type", [ - Argument("dipole_charge", dict, modifier_dipole_charge()), + Argument( + "dipole_charge", dict, modifier_dipole_charge(), doc=doc_dipole_charge + ), ], optional=False, doc=doc_modifier_type,