Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

docs: move arg docs from Variant to sub Arguments #4369

Merged
merged 1 commit into from
Nov 19, 2024
Merged

docs: move arg docs from Variant to sub Arguments #4369

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 39 additions & 42 deletions deepmd/utils/argcheck.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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."
Expand All @@ -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\
Expand Down Expand Up @@ -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\
Expand Down Expand Up @@ -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"
Expand All @@ -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\
Expand Down Expand Up @@ -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"

Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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 = (
Expand Down Expand Up @@ -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\
Expand Down Expand Up @@ -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",
Expand All @@ -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."
Expand Down Expand Up @@ -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."
Expand Down Expand Up @@ -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.'
Expand Down Expand Up @@ -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.'
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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,
Expand Down