diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml index 8b78f4cd0b..0f9fc61acd 100644 --- a/.github/workflows/test_python.yml +++ b/.github/workflows/test_python.yml @@ -42,17 +42,18 @@ jobs: - name: Get durations from cache uses: actions/cache@v4 with: - path: test_durations + path: .test_durations # the key must never match, even when restarting workflows, as that # will cause durations to get out of sync between groups, the # combined durations will be loaded if available - key: test-durations-split-${{ github.run_id }}-${{ github.run_number}}-${{ matrix.python }}-${{ matrix.group }} + key: test2-durations-split-${{ github.run_id }}-${{ github.run_number}}-${{ matrix.python }}-${{ matrix.group }} restore-keys: | - test-durations-combined-${{ matrix.python }}-${{ github.sha }} - test-durations-combined-${{ matrix.python }} - - run: pytest --cov=deepmd source/tests --durations=0 --splits 6 --group ${{ matrix.group }} --store-durations --durations-path=.test_durations_${{ matrix.group }} --splitting-algorithm least_duration + test2-durations-combined-${{ matrix.python }}-${{ github.sha }} + test2-durations-combined-${{ matrix.python }} + - run: pytest --cov=deepmd source/tests --durations=0 --splits 6 --group ${{ matrix.group }} --store-durations --durations-path=.test_durations --splitting-algorithm least_duration env: NUM_WORKERS: 0 + - run: mv .test_durations .test_durations_${{ matrix.group }} - name: Upload partial durations uses: actions/upload-artifact@v4 with: @@ -77,15 +78,15 @@ jobs: # key won't match during the first run for the given commit, but # restore-key will if there's a previous stored durations file, # so cache will both be loaded and stored - key: test-durations-combined-${{ matrix.python }}-${{ github.sha }} - restore-keys: test-durations-combined-${{ matrix.python }} + key: test2-durations-combined-${{ matrix.python }}-${{ github.sha }} + restore-keys: test2-durations-combined-${{ matrix.python }} - name: Download artifacts uses: actions/download-artifact@v4 with: pattern: split-${{ matrix.python }}-* merge-multiple: true - name: Combine test durations - run: jq '. + input' .test_durations_* > .test_durations + run: jq -s add .test_durations_* > .test_durations pass: name: Pass testing Python needs: [testpython, update_durations] diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0e5671d6c2..fb022030fd 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -29,7 +29,7 @@ repos: exclude: ^source/3rdparty - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.4.4 + rev: v0.4.5 hooks: - id: ruff args: ["--fix"] diff --git a/CITATIONS.bib b/CITATIONS.bib index a38c22deab..d5524a14f6 100644 --- a/CITATIONS.bib +++ b/CITATIONS.bib @@ -110,19 +110,22 @@ @article{Wang_NuclFusion_2022_v62_p126013 doi = {10.1088/1741-4326/ac888b}, } -@misc{Zhang_2022_DPA1, - annote = {attention-based descriptor}, +@article{Zhang_NpjComputMater_2024_v10_p94, + annote = {DPA-1, attention-based descriptor}, author = { - Zhang, Duo and Bi, Hangrui and Dai, Fu-Zhi and Jiang, Wanrun and Zhang, - Linfeng and Wang, Han + Duo Zhang and Hangrui Bi and Fu-Zhi Dai and Wanrun Jiang and Xinzijian Liu + and Linfeng Zhang and Han Wang }, title = { - {DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular - Simulation} + {Pretraining of attention-based deep learning potential model for molecular + simulation} }, - publisher = {arXiv}, - year = 2022, - doi = {10.48550/arXiv.2208.08236}, + journal = {Npj Comput. Mater}, + year = 2024, + volume = 10, + issue = 1, + pages = 94, + doi = {10.1038/s41524-024-01278-7}, } @misc{Zhang_2023_DPA2, diff --git a/deepmd/common.py b/deepmd/common.py index 098bb0ed11..f58634f224 100644 --- a/deepmd/common.py +++ b/deepmd/common.py @@ -16,7 +16,6 @@ Any, Dict, List, - Optional, Set, TypeVar, Union, @@ -39,11 +38,8 @@ ) __all__ = [ - "data_requirement", - "add_data_requirement", "select_idx_map", "make_default_mesh", - "j_must_have", "j_loader", "expand_sys_str", "get_np_precision", @@ -78,64 +74,6 @@ ) -# TODO: refactor data_requirement to make it not a global variable -# this is not a good way to do things. This is some global variable to which -# anyone can write and there is no good way to keep track of the changes -data_requirement = {} - - -def add_data_requirement( - key: str, - ndof: int, - atomic: bool = False, - must: bool = False, - high_prec: bool = False, - type_sel: Optional[bool] = None, - repeat: int = 1, - default: float = 0.0, - dtype: Optional[np.dtype] = None, - output_natoms_for_type_sel: bool = False, -): - """Specify data requirements for training. - - Parameters - ---------- - key : str - type of data stored in corresponding `*.npy` file e.g. `forces` or `energy` - ndof : int - number of the degrees of freedom, this is tied to `atomic` parameter e.g. forces - have `atomic=True` and `ndof=3` - atomic : bool, optional - specifies whwther the `ndof` keyworrd applies to per atom quantity or not, - by default False - must : bool, optional - specifi if the `*.npy` data file must exist, by default False - high_prec : bool, optional - if true load data to `np.float64` else `np.float32`, by default False - type_sel : bool, optional - select only certain type of atoms, by default None - repeat : int, optional - if specify repaeat data `repeat` times, by default 1 - default : float, optional, default=0. - default value of data - dtype : np.dtype, optional - the dtype of data, overwrites `high_prec` if provided - output_natoms_for_type_sel : bool, optional - if True and type_sel is True, the atomic dimension will be natoms instead of nsel - """ - data_requirement[key] = { - "ndof": ndof, - "atomic": atomic, - "must": must, - "high_prec": high_prec, - "type_sel": type_sel, - "repeat": repeat, - "default": default, - "dtype": dtype, - "output_natoms_for_type_sel": output_natoms_for_type_sel, - } - - def select_idx_map(atom_types: np.ndarray, select_types: np.ndarray) -> np.ndarray: """Build map of indices for element supplied element types from all atoms list. @@ -188,15 +126,20 @@ def make_default_mesh(pbc: bool, mixed_type: bool) -> np.ndarray: return default_mesh -# TODO: rename j_must_have to j_deprecated and only warn about deprecated keys -# maybe rename this to j_deprecated and only warn about deprecated keys, -# if the deprecated_key argument is left empty function puppose is only custom -# error since dict[key] already raises KeyError when the key is missing -def j_must_have( +def j_deprecated( jdata: Dict[str, "_DICT_VAL"], key: str, deprecated_key: List[str] = [] ) -> "_DICT_VAL": """Assert that supplied dictionary conaines specified key. + Parameters + ---------- + jdata : Dict[str, _DICT_VAL] + dictionary to check + key : str + key to check + deprecated_key : List[str], optional + list of deprecated keys, by default [] + Returns ------- _DICT_VAL @@ -247,7 +190,6 @@ def j_loader(filename: Union[str, Path]) -> Dict[str, Any]: raise TypeError("config file must be json, or yaml/yml") -# TODO port expand_sys_str completely to pathlib when all callers are ported def expand_sys_str(root_dir: Union[str, Path]) -> List[str]: """Recursively iterate over directories taking those that contain `type.raw` file. diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 94aad59c70..001e7b0fbe 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -9,6 +9,9 @@ import numpy as np +from deepmd.dpmodel.common import ( + NativeOP, +) from deepmd.dpmodel.output_def import ( FittingOutputDef, OutputVariableDef, @@ -25,7 +28,7 @@ BaseAtomicModel_ = make_base_atomic_model(np.ndarray) -class BaseAtomicModel(BaseAtomicModel_): +class BaseAtomicModel(BaseAtomicModel_, NativeOP): def __init__( self, type_map: List[str], @@ -192,6 +195,24 @@ def forward_common_atomic( return ret_dict + def call( + self, + extended_coord: np.ndarray, + extended_atype: np.ndarray, + nlist: np.ndarray, + mapping: Optional[np.ndarray] = None, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + ) -> Dict[str, np.ndarray]: + return self.forward_common_atomic( + extended_coord, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + ) + def serialize(self) -> dict: return { "type_map": self.type_map, diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 6ce5f7b6ce..66f35f151f 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -855,10 +855,8 @@ def call( else: raise NotImplementedError - input_r = dmatrix.reshape(-1, nnei, 4)[:, :, 1:4] / np.maximum( - np.linalg.norm( - dmatrix.reshape(-1, nnei, 4)[:, :, 1:4], axis=-1, keepdims=True - ), + input_r = rr.reshape(-1, nnei, 4)[:, :, 1:4] / np.maximum( + np.linalg.norm(rr.reshape(-1, nnei, 4)[:, :, 1:4], axis=-1, keepdims=True), 1e-12, ) gg = self.dpa1_attention( diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 68889ad331..7993f10abd 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -236,6 +236,8 @@ def call_lower( model_predict = self.output_type_cast(model_predict, input_prec) return model_predict + forward_lower = call_lower + def input_type_cast( self, coord: np.ndarray, @@ -473,4 +475,8 @@ def atomic_output_def(self) -> FittingOutputDef: """Get the output def of the atomic model.""" return self.atomic_model.atomic_output_def() + def get_ntypes(self) -> int: + """Get the number of types.""" + return len(self.get_type_map()) + return CM diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index ca8b18023b..018f50f1a5 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -9,10 +9,42 @@ import numpy as np from .region import ( + normalize_coord, to_face_distance, ) +def extend_input_and_build_neighbor_list( + coord, + atype, + rcut: float, + sel: List[int], + mixed_types: bool = False, + box: Optional[np.ndarray] = None, +): + nframes, nloc = atype.shape[:2] + if box is not None: + coord_normalized = normalize_coord( + coord.reshape(nframes, nloc, 3), + box.reshape(nframes, 3, 3), + ) + else: + coord_normalized = coord + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, atype, box, rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + nloc, + rcut, + sel, + distinguish_types=(not mixed_types), + ) + extended_coord = extended_coord.reshape(nframes, -1, 3) + return extended_coord, extended_atype, mapping, nlist + + ## translated from torch implemantation by chatgpt def build_neighbor_list( coord: np.ndarray, diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index f3272b0a3e..863a33d680 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -256,6 +256,26 @@ def forward_common_atomic( return ret_dict + def forward( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + comm_dict: Optional[Dict[str, torch.Tensor]] = None, + ) -> Dict[str, torch.Tensor]: + return self.forward_common_atomic( + extended_coord, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + comm_dict=comm_dict, + ) + def serialize(self) -> dict: return { "type_map": self.type_map, diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 2ffcb62ff9..a59eaca409 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -559,7 +559,7 @@ def forward( raise NotImplementedError input_r = torch.nn.functional.normalize( - dmatrix.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1 + rr.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1 ) gg = self.dpa1_attention( gg, nlist_mask, input_r=input_r, sw=sw diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py index a3279f1727..9eaea16c3e 100644 --- a/deepmd/pt/utils/env_mat_stat.py +++ b/deepmd/pt/utils/env_mat_stat.py @@ -141,7 +141,6 @@ def iter( zero_mean, one_stddev, self.descriptor.get_rcut(), - # TODO: export rcut_smth from DescriptorBlock self.descriptor.get_rcut_smth(), radial_only, protection=self.descriptor.get_env_protection(), diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py index cdee6e3722..a24a5aef72 100644 --- a/deepmd/pt/utils/nlist.py +++ b/deepmd/pt/utils/nlist.py @@ -326,7 +326,7 @@ def extend_coord_with_ghosts( # +1: central cell nbuff = torch.ceil(rcut / to_face).to(torch.long) # 3 - nbuff = torch.max(nbuff, dim=0, keepdim=False).values + nbuff = torch.amax(nbuff, dim=0) # faster than torch.max nbuff_cpu = nbuff.cpu() xi = torch.arange(-nbuff_cpu[0], nbuff_cpu[0] + 1, 1, device="cpu") yi = torch.arange(-nbuff_cpu[1], nbuff_cpu[1] + 1, 1, device="cpu") diff --git a/deepmd/tf/common.py b/deepmd/tf/common.py index 06be22a2ee..9244b4f7ee 100644 --- a/deepmd/tf/common.py +++ b/deepmd/tf/common.py @@ -23,12 +23,9 @@ from deepmd.common import ( VALID_ACTIVATION, VALID_PRECISION, - add_data_requirement, - data_requirement, expand_sys_str, get_np_precision, j_loader, - j_must_have, make_default_mesh, select_idx_map, ) @@ -47,11 +44,8 @@ __all__ = [ # from deepmd.common - "data_requirement", - "add_data_requirement", "select_idx_map", "make_default_mesh", - "j_must_have", "j_loader", "expand_sys_str", "get_np_precision", @@ -291,8 +285,6 @@ def wrapper(self, *args, **kwargs): def clear_session(): """Reset all state generated by DeePMD-kit.""" tf.reset_default_graph() - # TODO: remove this line when data_requirement is not a global variable - data_requirement.clear() _TF_VERSION = Version(TF_VERSION) if _TF_VERSION < Version("2.4.0"): tf.train.experimental.disable_mixed_precision_graph_rewrite() diff --git a/deepmd/tf/descriptor/descriptor.py b/deepmd/tf/descriptor/descriptor.py index fd6f6729e8..fabaf78c85 100644 --- a/deepmd/tf/descriptor/descriptor.py +++ b/deepmd/tf/descriptor/descriptor.py @@ -23,6 +23,9 @@ from deepmd.tf.utils import ( PluginVariant, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from deepmd.utils.plugin import ( make_plugin_registry, ) @@ -512,3 +515,8 @@ def serialize(self, suffix: str = "") -> dict: Name suffix to identify this descriptor """ raise NotImplementedError(f"Not implemented in class {self.__name__}") + + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + return [] diff --git a/deepmd/tf/descriptor/se_a_ebd.py b/deepmd/tf/descriptor/se_a_ebd.py index f252bf114c..c558cd285e 100644 --- a/deepmd/tf/descriptor/se_a_ebd.py +++ b/deepmd/tf/descriptor/se_a_ebd.py @@ -6,9 +6,6 @@ import numpy as np -from deepmd.tf.common import ( - add_data_requirement, -) from deepmd.tf.env import ( GLOBAL_TF_FLOAT_PRECISION, op_module, @@ -18,6 +15,9 @@ embedding_net, one_layer, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from .descriptor import ( Descriptor, @@ -110,8 +110,6 @@ def __init__( self.type_nlayer = type_nlayer self.type_one_side = type_one_side self.numb_aparam = numb_aparam - if self.numb_aparam > 0: - add_data_requirement("aparam", 3, atomic=True, must=True, high_prec=False) def build( self, @@ -600,3 +598,15 @@ def _ebd_filter( result = tf.reshape(result, [-1, outputs_size_2 * outputs_size]) return result, qmat + + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + data_requirement = super().input_requirement + if self.numb_aparam > 0: + data_requirement.append( + DataRequirementItem( + "aparam", 3, atomic=True, must=True, high_prec=False + ) + ) + return data_requirement diff --git a/deepmd/tf/descriptor/se_a_ef.py b/deepmd/tf/descriptor/se_a_ef.py index f1201d30fb..5a9020a6e6 100644 --- a/deepmd/tf/descriptor/se_a_ef.py +++ b/deepmd/tf/descriptor/se_a_ef.py @@ -7,9 +7,6 @@ import numpy as np -from deepmd.tf.common import ( - add_data_requirement, -) from deepmd.tf.env import ( GLOBAL_NP_FLOAT_PRECISION, GLOBAL_TF_FLOAT_PRECISION, @@ -20,6 +17,9 @@ from deepmd.tf.utils.sess import ( run_sess, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from .descriptor import ( Descriptor, @@ -361,8 +361,6 @@ def __init__( self.dstd = None self.davg = None - add_data_requirement("efield", 3, atomic=True, must=True, high_prec=False) - self.place_holders = {} avg_zero = np.zeros([self.ntypes, self.ndescrpt]).astype( GLOBAL_NP_FLOAT_PRECISION @@ -586,3 +584,12 @@ def _compute_dstats_sys_smth( sysr2.append(sumr2) sysa2.append(suma2) return sysr, sysr2, sysa, sysa2, sysn + + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + data_requirement = super().input_requirement + data_requirement.append( + DataRequirementItem("efield", 3, atomic=True, must=True, high_prec=False) + ) + return data_requirement diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index 2fef038f7d..3c4decbe8c 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -15,7 +15,6 @@ from deepmd.tf.common import ( j_loader, - j_must_have, ) from deepmd.tf.env import ( reset_default_tf_session_config, @@ -195,6 +194,7 @@ def _do_work(jdata: Dict[str, Any], run_opt: RunOptions, is_compress: bool = Fal train_data = get_data( jdata["training"]["training_data"], rcut, ipt_type_map, modifier ) + train_data.add_data_requirements(model.data_requirements) train_data.print_summary("training") if jdata["training"].get("validation_data", None) is not None: valid_data = get_data( @@ -203,13 +203,14 @@ def _do_work(jdata: Dict[str, Any], run_opt: RunOptions, is_compress: bool = Fal train_data.type_map, modifier, ) + valid_data.add_data_requirements(model.data_requirements) valid_data.print_summary("validation") else: if modifier is not None: modifier.build_fv_graph() # get training info - stop_batch = j_must_have(jdata["training"], "numb_steps") + stop_batch = jdata["training"]["numb_steps"] origin_type_map = jdata["model"].get("origin_type_map", None) if ( origin_type_map is not None and not origin_type_map diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index d967ce03d0..bc5180b60a 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -8,7 +8,6 @@ import numpy as np from deepmd.tf.common import ( - add_data_requirement, cast_precision, get_activation_func, get_precision, @@ -43,6 +42,9 @@ from deepmd.tf.utils.network import ( one_layer_rand_seed_shift, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from deepmd.utils.out_stat import ( compute_stats_from_redu, ) @@ -151,18 +153,9 @@ def __init__( self.useBN = False self.bias_dos = np.zeros((self.ntypes, self.numb_dos), dtype=np.float64) - # data requirement - if self.numb_fparam > 0: - add_data_requirement( - "fparam", self.numb_fparam, atomic=False, must=True, high_prec=False - ) self.fparam_avg = None self.fparam_std = None self.fparam_inv_std = None - if self.numb_aparam > 0: - add_data_requirement( - "aparam", self.numb_aparam, atomic=True, must=True, high_prec=False - ) self.aparam_avg = None self.aparam_std = None self.aparam_inv_std = None @@ -738,3 +731,21 @@ def serialize(self, suffix: str = "") -> dict: }, } return data + + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + data_requirement = [] + if self.numb_fparam > 0: + data_requirement.append( + DataRequirementItem( + "fparam", self.numb_fparam, atomic=False, must=True, high_prec=False + ) + ) + if self.numb_aparam > 0: + data_requirement.append( + DataRequirementItem( + "aparam", self.numb_aparam, atomic=True, must=True, high_prec=False + ) + ) + return data_requirement diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 873f7258db..a1eb916a1c 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -9,7 +9,6 @@ import numpy as np from deepmd.tf.common import ( - add_data_requirement, cast_precision, get_activation_func, get_precision, @@ -53,6 +52,9 @@ from deepmd.tf.utils.spin import ( Spin, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from deepmd.utils.finetune import ( change_energy_bias_lower, ) @@ -218,18 +220,9 @@ def __init__( self.atom_ener.append(None) self.useBN = False self.bias_atom_e = np.zeros(self.ntypes, dtype=np.float64) - # data requirement - if self.numb_fparam > 0: - add_data_requirement( - "fparam", self.numb_fparam, atomic=False, must=True, high_prec=False - ) self.fparam_avg = None self.fparam_std = None self.fparam_inv_std = None - if self.numb_aparam > 0: - add_data_requirement( - "aparam", self.numb_aparam, atomic=True, must=True, high_prec=False - ) self.aparam_avg = None self.aparam_std = None self.aparam_inv_std = None @@ -939,3 +932,21 @@ def serialize(self, suffix: str = "") -> dict: }, } return data + + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + data_requirement = [] + if self.numb_fparam > 0: + data_requirement.append( + DataRequirementItem( + "fparam", self.numb_fparam, atomic=False, must=True, high_prec=False + ) + ) + if self.numb_aparam > 0: + data_requirement.append( + DataRequirementItem( + "aparam", self.numb_aparam, atomic=True, must=True, high_prec=False + ) + ) + return data_requirement diff --git a/deepmd/tf/fit/fitting.py b/deepmd/tf/fit/fitting.py index d2aebd8f97..9190261187 100644 --- a/deepmd/tf/fit/fitting.py +++ b/deepmd/tf/fit/fitting.py @@ -25,6 +25,9 @@ from deepmd.tf.utils import ( PluginVariant, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from deepmd.utils.plugin import ( make_plugin_registry, ) @@ -252,3 +255,8 @@ def deserialize_network(cls, data: dict, suffix: str = "") -> dict: # prevent keyError fitting_net_variables[f"{layer_name}{key}{suffix}/idt"] = 0.0 return fitting_net_variables + + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + return [] diff --git a/deepmd/tf/loss/dos.py b/deepmd/tf/loss/dos.py index 763e75638f..385d2484a8 100644 --- a/deepmd/tf/loss/dos.py +++ b/deepmd/tf/loss/dos.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + List, +) + import numpy as np -from deepmd.tf.common import ( - add_data_requirement, -) from deepmd.tf.env import ( global_cvt_2_ener_float, global_cvt_2_tf_float, @@ -12,6 +13,9 @@ from deepmd.tf.utils.sess import ( run_sess, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from .loss import ( Loss, @@ -56,13 +60,6 @@ def __init__( self.has_cdf = self.start_pref_cdf != 0.0 or self.limit_pref_cdf != 0.0 self.has_ados = self.start_pref_ados != 0.0 or self.limit_pref_ados != 0.0 self.has_acdf = self.start_pref_acdf != 0.0 or self.limit_pref_acdf != 0.0 - # data required - add_data_requirement( - "dos", self.numb_dos, atomic=False, must=True, high_prec=True - ) - add_data_requirement( - "atom_dos", self.numb_dos, atomic=True, must=False, high_prec=True - ) def build(self, learning_rate, natoms, model_dict, label_dict, suffix): dos = model_dict["dos"] @@ -212,3 +209,20 @@ def eval(self, sess, feed_dict, natoms): results["rmse_acdf"] = np.sqrt(error_acdf) return results + + @property + def label_requirement(self) -> List[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + data_requirements = [] + # data required + data_requirements.append( + DataRequirementItem( + "dos", self.numb_dos, atomic=False, must=True, high_prec=True + ) + ) + data_requirements.append( + DataRequirementItem( + "atom_dos", self.numb_dos, atomic=True, must=False, high_prec=True + ) + ) + return data_requirements diff --git a/deepmd/tf/loss/ener.py b/deepmd/tf/loss/ener.py index baa4aa3e02..7ecb185818 100644 --- a/deepmd/tf/loss/ener.py +++ b/deepmd/tf/loss/ener.py @@ -1,13 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( + List, Optional, ) import numpy as np -from deepmd.tf.common import ( - add_data_requirement, -) from deepmd.tf.env import ( global_cvt_2_ener_float, global_cvt_2_tf_float, @@ -16,6 +14,9 @@ from deepmd.tf.utils.sess import ( run_sess, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from .loss import ( Loss, @@ -111,32 +112,6 @@ def __init__( raise RuntimeError( "When generalized force loss is used, the dimension of generalized coordinates should be larger than 0" ) - # data required - add_data_requirement("energy", 1, atomic=False, must=False, high_prec=True) - add_data_requirement("force", 3, atomic=True, must=False, high_prec=False) - add_data_requirement("virial", 9, atomic=False, must=False, high_prec=False) - add_data_requirement("atom_ener", 1, atomic=True, must=False, high_prec=False) - add_data_requirement( - "atom_pref", 1, atomic=True, must=False, high_prec=False, repeat=3 - ) - # drdq: the partial derivative of atomic coordinates w.r.t. generalized coordinates - if self.has_gf > 0: - add_data_requirement( - "drdq", - self.numb_generalized_coord * 3, - atomic=True, - must=False, - high_prec=False, - ) - if self.enable_atom_ener_coeff: - add_data_requirement( - "atom_ener_coeff", - 1, - atomic=True, - must=False, - high_prec=False, - default=1.0, - ) def build(self, learning_rate, natoms, model_dict, label_dict, suffix): energy = model_dict["energy"] @@ -380,6 +355,54 @@ def eval(self, sess, feed_dict, natoms): results["rmse_gf"] = np.sqrt(error_gf) return results + @property + def label_requirement(self) -> List[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + data_requirements = [] + # data required + data_requirements.append( + DataRequirementItem("energy", 1, atomic=False, must=False, high_prec=True) + ) + data_requirements.append( + DataRequirementItem("force", 3, atomic=True, must=False, high_prec=False) + ) + data_requirements.append( + DataRequirementItem("virial", 9, atomic=False, must=False, high_prec=False) + ) + data_requirements.append( + DataRequirementItem( + "atom_ener", 1, atomic=True, must=False, high_prec=False + ) + ) + data_requirements.append( + DataRequirementItem( + "atom_pref", 1, atomic=True, must=False, high_prec=False, repeat=3 + ) + ) + # drdq: the partial derivative of atomic coordinates w.r.t. generalized coordinates + if self.has_gf > 0: + data_requirements.append( + DataRequirementItem( + "drdq", + self.numb_generalized_coord * 3, + atomic=True, + must=False, + high_prec=False, + ) + ) + if self.enable_atom_ener_coeff: + data_requirements.append( + DataRequirementItem( + "atom_ener_coeff", + 1, + atomic=True, + must=False, + high_prec=False, + default=1.0, + ) + ) + return data_requirements + class EnerSpinLoss(Loss): def __init__( @@ -422,23 +445,6 @@ def __init__( self.has_fm = self.start_pref_fm != 0.0 or self.limit_pref_fm != 0.0 self.has_v = self.start_pref_v != 0.0 or self.limit_pref_v != 0.0 self.has_ae = self.start_pref_ae != 0.0 or self.limit_pref_ae != 0.0 - # data required - add_data_requirement("energy", 1, atomic=False, must=False, high_prec=True) - add_data_requirement("force", 3, atomic=True, must=False, high_prec=False) - add_data_requirement("virial", 9, atomic=False, must=False, high_prec=False) - add_data_requirement("atom_ener", 1, atomic=True, must=False, high_prec=False) - add_data_requirement( - "atom_pref", 1, atomic=True, must=False, high_prec=False, repeat=3 - ) - if self.enable_atom_ener_coeff: - add_data_requirement( - "atom_ener_coeff", - 1, - atomic=True, - must=False, - high_prec=False, - default=1.0, - ) def build(self, learning_rate, natoms, model_dict, label_dict, suffix): energy_pred = model_dict["energy"] @@ -719,6 +725,43 @@ def print_on_training( return print_str + @property + def label_requirement(self) -> List[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + data_requirements = [] + # data required + data_requirements.append( + DataRequirementItem("energy", 1, atomic=False, must=False, high_prec=True) + ) + data_requirements.append( + DataRequirementItem("force", 3, atomic=True, must=False, high_prec=False) + ) + data_requirements.append( + DataRequirementItem("virial", 9, atomic=False, must=False, high_prec=False) + ) + data_requirements.append( + DataRequirementItem( + "atom_ener", 1, atomic=True, must=False, high_prec=False + ) + ) + data_requirements.append( + DataRequirementItem( + "atom_pref", 1, atomic=True, must=False, high_prec=False, repeat=3 + ) + ) + if self.enable_atom_ener_coeff: + data_requirements.append( + DataRequirementItem( + "atom_ener_coeff", + 1, + atomic=True, + must=False, + high_prec=False, + default=1.0, + ) + ) + return data_requirements + class EnerDipoleLoss(Loss): def __init__( @@ -734,11 +777,6 @@ def __init__( self.limit_pref_e = limit_pref_e self.start_pref_ed = start_pref_ed self.limit_pref_ed = limit_pref_ed - # data required - add_data_requirement("energy", 1, atomic=False, must=True, high_prec=True) - add_data_requirement( - "energy_dipole", 3, atomic=False, must=True, high_prec=False - ) def build(self, learning_rate, natoms, model_dict, label_dict, suffix): coord = model_dict["coord"] @@ -832,3 +870,18 @@ def eval(self, sess, feed_dict, natoms): "rmse_ed": np.sqrt(error_ed), } return results + + @property + def label_requirement(self) -> List[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + data_requirements = [] + # data required + data_requirements.append( + DataRequirementItem("energy", 1, atomic=False, must=False, high_prec=True) + ) + data_requirements.append( + DataRequirementItem( + "energy_dipole", 3, atomic=False, must=True, high_prec=False + ) + ) + return data_requirements diff --git a/deepmd/tf/loss/loss.py b/deepmd/tf/loss/loss.py index 327aea5230..ca90c2eb64 100644 --- a/deepmd/tf/loss/loss.py +++ b/deepmd/tf/loss/loss.py @@ -5,6 +5,7 @@ ) from typing import ( Dict, + List, Tuple, ) @@ -13,6 +14,9 @@ from deepmd.tf.env import ( tf, ) +from deepmd.utils.data import ( + DataRequirementItem, +) class Loss(metaclass=ABCMeta): @@ -91,3 +95,8 @@ def display_if_exist(loss: tf.Tensor, find_property: float) -> tf.Tensor: lambda: loss, lambda: tf.cast(np.nan, dtype=loss.dtype), ) + + @property + @abstractmethod + def label_requirement(self) -> List[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" diff --git a/deepmd/tf/loss/tensor.py b/deepmd/tf/loss/tensor.py index 6a0eb30a44..d3a1b95369 100644 --- a/deepmd/tf/loss/tensor.py +++ b/deepmd/tf/loss/tensor.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + List, +) + import numpy as np -from deepmd.tf.common import ( - add_data_requirement, -) from deepmd.tf.env import ( global_cvt_2_tf_float, tf, @@ -11,6 +12,9 @@ from deepmd.tf.utils.sess import ( run_sess, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from .loss import ( Loss, @@ -50,24 +54,6 @@ def __init__(self, jdata, **kwarg): "Can not assian zero weight both to `pref` and `pref_atomic`" ) - # data required - add_data_requirement( - "atomic_" + self.label_name, - self.tensor_size, - atomic=True, - must=False, - high_prec=False, - type_sel=self.type_sel, - ) - add_data_requirement( - self.label_name, - self.tensor_size, - atomic=False, - must=False, - high_prec=False, - type_sel=self.type_sel, - ) - def build(self, learning_rate, natoms, model_dict, label_dict, suffix): polar_hat = label_dict[self.label_name] atomic_polar_hat = label_dict["atom_" + self.label_name] @@ -154,3 +140,30 @@ def eval(self, sess, feed_dict, natoms): if self.global_weight > 0.0: results["rmse_gl"] = np.sqrt(error_gl) / atoms return results + + @property + def label_requirement(self) -> List[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + data_requirements = [] + # data required + data_requirements.append( + DataRequirementItem( + "atomic_" + self.label_name, + self.tensor_size, + atomic=True, + must=False, + high_prec=False, + type_sel=self.type_sel, + ) + ) + data_requirements.append( + DataRequirementItem( + self.label_name, + self.tensor_size, + atomic=False, + must=False, + high_prec=False, + type_sel=self.type_sel, + ) + ) + return data_requirements diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index c8b9b2d6f5..66aaff8189 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -16,6 +16,9 @@ from deepmd.tf.utils.data_system import ( DeepmdDataSystem, ) +from deepmd.tf.utils.pair_tab import ( + PairTab, +) from deepmd.tf.utils.spin import ( Spin, ) @@ -106,6 +109,16 @@ def __init__( self.numb_fparam = self.fitting.get_numb_fparam() self.numb_aparam = self.fitting.get_numb_aparam() + self.srtab_name = use_srtab + if self.srtab_name is not None: + self.srtab = PairTab(self.srtab_name, rcut=self.get_rcut()) + self.smin_alpha = smin_alpha + self.sw_rmin = sw_rmin + self.sw_rmax = sw_rmax + self.srtab_add_bias = srtab_add_bias + else: + self.srtab = None + def get_rcut(self): return self.rcut diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index 86676bfe0b..fa28b2cc58 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -6,6 +6,7 @@ Enum, ) from typing import ( + List, Optional, Union, ) @@ -34,6 +35,9 @@ get_tensor_by_name_from_graph, load_graph_def, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from .model import ( Model, @@ -261,3 +265,23 @@ def serialize(self, suffix: str = "") -> dict: @classmethod def deserialize(cls, data: dict, suffix: str = ""): raise RuntimeError("Should not touch here.") + + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + data_requirement = [] + numb_fparam = self.model.get_dim_fparam() + numb_aparam = self.model.get_dim_aparam() + if numb_fparam > 0: + data_requirement.append( + DataRequirementItem( + "fparam", numb_fparam, atomic=False, must=True, high_prec=False + ) + ) + if numb_aparam > 0: + data_requirement.append( + DataRequirementItem( + "aparam", numb_aparam, atomic=True, must=True, high_prec=False + ) + ) + return data_requirement diff --git a/deepmd/tf/model/linear.py b/deepmd/tf/model/linear.py index ae1b0b5c78..26bc382569 100644 --- a/deepmd/tf/model/linear.py +++ b/deepmd/tf/model/linear.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import operator from enum import ( Enum, ) from functools import ( lru_cache, + reduce, ) from typing import ( List, @@ -22,6 +24,9 @@ from deepmd.tf.loss.loss import ( Loss, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from .model import ( Model, @@ -145,6 +150,13 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): ] return local_jdata_cpy + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + return reduce( + operator.iadd, [model.input_requirement for model in self.models], [] + ) + @Model.register("linear_ener") class LinearEnergyModel(LinearModel): diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 5914dcf48d..194750a9d7 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -53,15 +53,15 @@ from deepmd.tf.utils.graph import ( load_graph_def, ) -from deepmd.tf.utils.pair_tab import ( - PairTab, -) from deepmd.tf.utils.spin import ( Spin, ) from deepmd.tf.utils.type_embed import ( TypeEmbedNet, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from deepmd.utils.plugin import ( make_plugin_registry, ) @@ -116,11 +116,6 @@ def __init__( data_stat_nbatch: int = 10, data_bias_nsample: int = 10, data_stat_protect: float = 1e-2, - use_srtab: Optional[str] = None, - smin_alpha: Optional[float] = None, - sw_rmin: Optional[float] = None, - sw_rmax: Optional[float] = None, - srtab_add_bias: bool = True, spin: Optional[Spin] = None, compress: Optional[dict] = None, **kwargs, @@ -142,15 +137,6 @@ def __init__( self.data_stat_nbatch = data_stat_nbatch self.data_bias_nsample = data_bias_nsample self.data_stat_protect = data_stat_protect - self.srtab_name = use_srtab - if self.srtab_name is not None: - self.srtab = PairTab(self.srtab_name) - self.smin_alpha = smin_alpha - self.sw_rmin = sw_rmin - self.sw_rmax = sw_rmax - self.srtab_add_bias = srtab_add_bias - else: - self.srtab = None def get_type_map(self) -> list: """Get the type map.""" @@ -588,6 +574,11 @@ def serialize(self, suffix: str = "") -> dict: """ raise NotImplementedError(f"Not implemented in class {self.__name__}") + @property + @abstractmethod + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + @Model.register("standard") class StandardModel(Model): @@ -842,3 +833,8 @@ def serialize(self, suffix: str = "") -> dict: "out_std": np.ones([1, ntypes, dict_fit["dim_out"]]), }, } + + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + return self.descrpt.input_requirement + self.fitting.input_requirement diff --git a/deepmd/tf/model/pairtab.py b/deepmd/tf/model/pairtab.py index 3cc1114f81..4f71dcd76e 100644 --- a/deepmd/tf/model/pairtab.py +++ b/deepmd/tf/model/pairtab.py @@ -32,6 +32,9 @@ from deepmd.tf.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data import ( + DataRequirementItem, +) @Model.register("pairtab") @@ -66,7 +69,7 @@ def __init__( ): super().__init__() self.tab_file = tab_file - self.tab = PairTab(self.tab_file) + self.tab = PairTab(self.tab_file, rcut=rcut) self.ntypes = self.tab.ntypes self.rcut = rcut if isinstance(sel, int): @@ -286,3 +289,8 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: """ local_jdata_cpy = local_jdata.copy() return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True) + + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + return [] diff --git a/deepmd/tf/model/pairwise_dprc.py b/deepmd/tf/model/pairwise_dprc.py index 6b0e95e88b..3d61dfd339 100644 --- a/deepmd/tf/model/pairwise_dprc.py +++ b/deepmd/tf/model/pairwise_dprc.py @@ -7,7 +7,6 @@ ) from deepmd.tf.common import ( - add_data_requirement, make_default_mesh, ) from deepmd.tf.env import ( @@ -34,6 +33,9 @@ from deepmd.tf.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data import ( + DataRequirementItem, +) @Model.register("pairwise_dprc") @@ -103,7 +105,6 @@ def __init__( type_embedding=self.typeebd, compress=compress, ) - add_data_requirement("aparam", 1, atomic=True, must=True, high_prec=False) self.rcut = max(self.qm_model.get_rcut(), self.qmmm_model.get_rcut()) def build( @@ -424,6 +425,15 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): UpdateSel().get_min_nbor_dist(global_jdata, 6.0) return local_jdata + @property + def input_requirement(self) -> List[DataRequirementItem]: + """Return data requirements needed for the model input.""" + data_requirement = [] + data_requirement.append( + DataRequirementItem("aparam", 1, atomic=True, must=True, high_prec=False) + ) + return data_requirement + def gather_placeholder( params: tf.Tensor, indices: tf.Tensor, placeholder: float = 0.0, **kwargs diff --git a/deepmd/tf/train/trainer.py b/deepmd/tf/train/trainer.py index 855b2ee722..60a468be3e 100644 --- a/deepmd/tf/train/trainer.py +++ b/deepmd/tf/train/trainer.py @@ -28,9 +28,7 @@ format_training_message_per_task, ) from deepmd.tf.common import ( - data_requirement, get_precision, - j_must_have, ) from deepmd.tf.env import ( GLOBAL_ENER_FLOAT_PRECISION, @@ -63,6 +61,9 @@ from deepmd.tf.utils.sess import ( run_sess, ) +from deepmd.utils.data import ( + DataRequirementItem, +) log = logging.getLogger(__name__) @@ -89,7 +90,7 @@ def __init__(self, jdata, run_opt, is_compress=False): def _init_param(self, jdata): # model config - model_param = j_must_have(jdata, "model") + model_param = jdata["model"] # nvnmd self.nvnmd_param = jdata.get("nvnmd", {}) @@ -121,7 +122,7 @@ def get_lr_and_coef(lr_param): return lr, scale_lr_coef # learning rate - lr_param = j_must_have(jdata, "learning_rate") + lr_param = jdata["learning_rate"] self.lr, self.scale_lr_coef = get_lr_and_coef(lr_param) # loss # infer loss type by fitting_type @@ -270,7 +271,7 @@ def _build_network(self, data, suffix=""): self.place_holders[kk] = tf.placeholder( GLOBAL_TF_FLOAT_PRECISION, [None], "t_" + kk ) - self._get_place_holders(data_requirement) + self._get_place_holders({rr.key: rr.dict for rr in self.data_requirements}) else: self._get_place_holders(data.get_data_dict()) @@ -860,6 +861,10 @@ def _change_energy_bias( bias_adjust_mode=bias_adjust_mode, ) + @property + def data_requirements(self) -> List[DataRequirementItem]: + return self.model.input_requirement + self.loss.label_requirement + class DatasetLoader: """Generate an OP that loads the training data from the given DeepmdDataSystem. diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index b35d9833d5..30971c7256 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -61,6 +61,11 @@ def __init__(self, initial_batch_size: int = 1024, factor: float = 2.0) -> None: self.maximum_working_batch_size = initial_batch_size if self.is_gpu_available(): self.minimal_not_working_batch_size = 2**31 + log.info( + "If you encounter the error 'an illegal memory access was encountered', this may be due to a TensorFlow issue. " + "To avoid this, set the environment variable DP_INFER_BATCH_SIZE to a smaller value than the last adjusted batch size. " + "The environment variable DP_INFER_BATCH_SIZE controls the inference batch size (nframes * natoms). " + ) else: self.minimal_not_working_batch_size = ( self.maximum_working_batch_size + 1 diff --git a/deepmd/utils/compat.py b/deepmd/utils/compat.py index 11f5e639dc..edd01b8291 100644 --- a/deepmd/utils/compat.py +++ b/deepmd/utils/compat.py @@ -17,7 +17,7 @@ import numpy as np from deepmd.common import ( - j_must_have, + j_deprecated, ) @@ -127,8 +127,8 @@ def _smth_descriptor(jdata: Dict[str, Any]) -> Dict[str, Any]: descriptor["sel"] = jdata["sel_a"] _jcopy(jdata, descriptor, ("rcut",)) descriptor["rcut_smth"] = jdata.get("rcut_smth", descriptor["rcut"]) - descriptor["neuron"] = j_must_have(jdata, "filter_neuron") - descriptor["axis_neuron"] = j_must_have(jdata, "axis_neuron", ["n_axis_neuron"]) + descriptor["neuron"] = jdata["filter_neuron"] + descriptor["axis_neuron"] = j_deprecated(jdata, "axis_neuron", ["n_axis_neuron"]) descriptor["resnet_dt"] = False if "resnet_dt" in jdata: descriptor["resnet_dt"] = jdata["filter_resnet_dt"] @@ -154,7 +154,7 @@ def _fitting_net(jdata: Dict[str, Any]) -> Dict[str, Any]: seed = jdata.get("seed", None) if seed is not None: fitting_net["seed"] = seed - fitting_net["neuron"] = j_must_have(jdata, "fitting_neuron", ["n_neuron"]) + fitting_net["neuron"] = j_deprecated(jdata, "fitting_neuron", ["n_neuron"]) fitting_net["resnet_dt"] = True if "resnet_dt" in jdata: fitting_net["resnet_dt"] = jdata["resnet_dt"] @@ -237,16 +237,16 @@ def _training(jdata: Dict[str, Any]) -> Dict[str, Any]: training["disp_file"] = "lcurve.out" if "disp_file" in jdata: training["disp_file"] = jdata["disp_file"] - training["disp_freq"] = j_must_have(jdata, "disp_freq") - training["numb_test"] = j_must_have(jdata, "numb_test") - training["save_freq"] = j_must_have(jdata, "save_freq") - training["save_ckpt"] = j_must_have(jdata, "save_ckpt") - training["disp_training"] = j_must_have(jdata, "disp_training") - training["time_training"] = j_must_have(jdata, "time_training") + training["disp_freq"] = jdata["disp_freq"] + training["numb_test"] = jdata["numb_test"] + training["save_freq"] = jdata["save_freq"] + training["save_ckpt"] = jdata["save_ckpt"] + training["disp_training"] = jdata["disp_training"] + training["time_training"] = jdata["time_training"] if "profiling" in jdata: training["profiling"] = jdata["profiling"] if training["profiling"]: - training["profiling_file"] = j_must_have(jdata, "profiling_file") + training["profiling_file"] = jdata["profiling_file"] return training @@ -378,7 +378,7 @@ def is_deepmd_v0_input(jdata): return "model" not in jdata.keys() def is_deepmd_v1_input(jdata): - return "systems" in j_must_have(jdata, "training").keys() + return "systems" in jdata["training"].keys() if is_deepmd_v0_input(jdata): jdata = convert_input_v0_v1(jdata, warning, None) diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index cd0e414b5f..91782d898f 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -787,3 +787,11 @@ def __getitem__(self, key: str): if key not in self.dict: raise KeyError(key) return self.dict[key] + + def __eq__(self, __value: object) -> bool: + if not isinstance(__value, DataRequirementItem): + return False + return self.dict == __value.dict + + def __repr__(self) -> str: + return f"DataRequirementItem({self.dict})" diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 693845f8d0..f02398278f 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -17,15 +17,14 @@ import deepmd.utils.random as dp_random from deepmd.common import ( - data_requirement, expand_sys_str, - j_must_have, make_default_mesh, ) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) from deepmd.utils.data import ( + DataRequirementItem, DeepmdData, ) from deepmd.utils.out_stat import ( @@ -267,7 +266,7 @@ def compute_energy_shift(self, rcond=None, key="energy"): ) return energy_shift.ravel() - def add_dict(self, adict: dict) -> None: + def add_dict(self, adict: Dict[str, Dict[str, Any]]) -> None: """Add items to the data system by a `dict`. `adict` should have items like .. code-block:: python. @@ -299,6 +298,12 @@ def add_dict(self, adict: dict) -> None: ), ) + def add_data_requirements( + self, data_requirements: List[DataRequirementItem] + ) -> None: + """Add items to the data system by a list of `DataRequirementItem`.""" + self.add_dict({rr.key: rr.dict for rr in data_requirements}) + def add( self, key: str, @@ -786,10 +791,10 @@ def get_data( DeepmdDataSystem The data system """ - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] systems = process_systems(systems) - batch_size = j_must_have(jdata, "batch_size") + batch_size = jdata["batch_size"] sys_probs = jdata.get("sys_probs", None) auto_prob = jdata.get("auto_prob", "prob_sys_size") optional_type_map = not multi_task_mode @@ -807,6 +812,5 @@ def get_data( sys_probs=sys_probs, auto_prob_style=auto_prob, ) - data.add_dict(data_requirement) return data diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 1b397a3cfa..4335a32372 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -31,6 +31,8 @@ class PairTab: The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly. + rcut : float, optional + cutoff raduis for the tabulated potential """ def __init__(self, filename: str, rcut: Optional[float] = None) -> None: @@ -49,6 +51,8 @@ def reinit(self, filename: str, rcut: Optional[float] = None) -> None: The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly. + rcut : float, optional + cutoff raduis for the tabulated potential """ if filename is None: self.tab_info, self.tab_data = None, None diff --git a/doc/credits.rst b/doc/credits.rst index 64880d9035..1b39dc1e0e 100644 --- a/doc/credits.rst +++ b/doc/credits.rst @@ -47,7 +47,7 @@ Cite DeePMD-kit and methods .. bibliography:: :filter: False - Zhang_2022_DPA1 + Zhang_NpjComputMater_2024_v10_p94 - If DPA-2 descriptor (`dpa2`) is used, diff --git a/pyproject.toml b/pyproject.toml index 80d5ad9ee9..7703ce71f9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -351,6 +351,7 @@ banned-module-level-imports = [ "deepmd/pt/**" = ["TID253"] "source/tests/tf/**" = ["TID253"] "source/tests/pt/**" = ["TID253"] +"source/tests/universal/pt/**" = ["TID253"] "source/ipi/tests/**" = ["TID253"] "source/lmp/tests/**" = ["TID253"] "**/*.ipynb" = ["T20"] # printing in a nb file is expected diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 027b42825b..d4fbdd3363 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -219,7 +219,9 @@ void PairDeepMD::make_fparam_from_compute(vector &fparam) { int icompute = modify->find_compute(compute_fparam_id); Compute *compute = modify->compute[icompute]; - assert(compute); + if (!compute) { + error->all(FLERR, "compute id is not found: " + compute_fparam_id); + } fparam.resize(dim_fparam); if (dim_fparam == 1) { @@ -246,7 +248,9 @@ void PairDeepMD::make_aparam_from_compute(vector &aparam) { int icompute = modify->find_compute(compute_aparam_id); Compute *compute = modify->compute[icompute]; - assert(compute); + if (!compute) { + error->all(FLERR, "compute id is not found: " + compute_aparam_id); + } int nlocal = atom->nlocal; aparam.resize(static_cast(dim_aparam) * nlocal); @@ -277,7 +281,9 @@ void PairDeepMD::make_ttm_fparam(vector &fparam) { ttm_fix = dynamic_cast(modify->fix[ii]); } } - assert(ttm_fix); + if (!ttm_fix) { + error->all(FLERR, "fix ttm id is not found: " + ttm_fix_id); + } fparam.resize(dim_fparam); @@ -316,7 +322,9 @@ void PairDeepMD::make_ttm_aparam(vector &daparam) { ttm_fix = dynamic_cast(modify->fix[ii]); } } - assert(ttm_fix); + if (!ttm_fix) { + error->all(FLERR, "fix ttm id is not found: " + ttm_fix_id); + } // modify double **x = atom->x; int *mask = atom->mask; diff --git a/source/tests/consistent/descriptor/test_dpa1.py b/source/tests/consistent/descriptor/test_dpa1.py index 5e9ea01602..8bd9af6d5f 100644 --- a/source/tests/consistent/descriptor/test_dpa1.py +++ b/source/tests/consistent/descriptor/test_dpa1.py @@ -201,7 +201,6 @@ def skip_tf(self) -> bool: precision, use_econf_tebd, ) = self.param - # TODO (excluded_types != [] and attn_layer > 0) need fix return ( CommonTest.skip_tf or ( @@ -209,7 +208,6 @@ def skip_tf(self) -> bool: or smooth_type_embedding or not normalize or temperature != 1.0 - or (excluded_types != [] and attn_layer > 0) or (type_one_side and tebd_input_mode == "strip") # not consistent yet ) or self.is_meaningless_zero_attention_layer_tests( diff --git a/source/tests/consistent/descriptor/test_dpa2.py b/source/tests/consistent/descriptor/test_dpa2.py index 25887aa4e5..b24274c9d7 100644 --- a/source/tests/consistent/descriptor/test_dpa2.py +++ b/source/tests/consistent/descriptor/test_dpa2.py @@ -395,7 +395,7 @@ def atol(self) -> float: use_econf_tebd, ) = self.param if precision == "float64": - return 1e-8 + return 1e-6 # need to fix in the future, see issue https://github.com/deepmodeling/deepmd-kit/issues/3786 elif precision == "float32": return 1e-4 else: diff --git a/source/tests/pt/model/test_model.py b/source/tests/pt/model/test_model.py index 6eb460e808..8fdbdaf413 100644 --- a/source/tests/pt/model/test_model.py +++ b/source/tests/pt/model/test_model.py @@ -31,7 +31,6 @@ ) from deepmd.pt.utils.learning_rate import LearningRateExp as MyLRExp from deepmd.tf.common import ( - data_requirement, expand_sys_str, ) from deepmd.tf.descriptor import DescrptSeA as DescrptSeA_tf @@ -114,6 +113,8 @@ def get_intermediate_state(self, num_steps=1): dp_loss = self._get_dp_loss() dp_lr = self._get_dp_lr() dp_ds = self._get_dp_dataset() + dp_ds.add_data_requirements(dp_model.input_requirement) + dp_ds.add_data_requirements(dp_loss.label_requirement) dp_model.data_stat(dp_ds) # Build graph @@ -188,7 +189,6 @@ def _get_dp_dataset(self): type_map=self.type_map, trn_all_set=True, ) - data.add_dict(data_requirement) return data def _get_dp_model(self): diff --git a/source/tests/pt/model/test_polarizability_fitting.py b/source/tests/pt/model/test_polarizability_fitting.py index a061780f45..84d6bd91ab 100644 --- a/source/tests/pt/model/test_polarizability_fitting.py +++ b/source/tests/pt/model/test_polarizability_fitting.py @@ -269,12 +269,7 @@ def test_permu(self): def test_trans(self): atype = self.atype.reshape(1, 5) - coord_s = torch.matmul( - torch.remainder( - torch.matmul(self.coord + self.shift, torch.linalg.inv(self.cell)), 1.0 - ), - self.cell, - ) + coord_s = self.coord + self.shift for fit_diag, scale in itertools.product([True, False], [None, self.scale]): ft0 = PolarFittingNet( self.nt, diff --git a/source/tests/tf/test_data_large_batch.py b/source/tests/tf/test_data_large_batch.py index 1232f8b1db..1b19d664dd 100644 --- a/source/tests/tf/test_data_large_batch.py +++ b/source/tests/tf/test_data_large_batch.py @@ -5,9 +5,6 @@ import numpy as np from packaging.version import parse as parse_version -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeAtten, ) @@ -50,11 +47,11 @@ def test_data_mixed_type(self): jfile = "water_se_atten_mixed_type.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - type_map = j_must_have(jdata["model"], "type_map") + rcut = jdata["model"]["descriptor"]["rcut"] + type_map = jdata["model"]["type_map"] data = DeepmdDataSystem(systems, batch_size, test_size, rcut, type_map=type_map) data_requirement = { @@ -248,11 +245,11 @@ def test_stripped_data_mixed_type(self): jfile = "water_se_atten_mixed_type.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - type_map = j_must_have(jdata["model"], "type_map") + rcut = jdata["model"]["descriptor"]["rcut"] + type_map = jdata["model"]["type_map"] data = DeepmdDataSystem(systems, batch_size, test_size, rcut, type_map=type_map) data_requirement = { @@ -446,11 +443,11 @@ def test_compressible_data_mixed_type(self): jfile = "water_se_atten_mixed_type.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - type_map = j_must_have(jdata["model"], "type_map") + rcut = jdata["model"]["descriptor"]["rcut"] + type_map = jdata["model"]["type_map"] data = DeepmdDataSystem(systems, batch_size, test_size, rcut, type_map=type_map) data_requirement = { diff --git a/source/tests/tf/test_data_modifier.py b/source/tests/tf/test_data_modifier.py index eae3155cb8..890de86034 100644 --- a/source/tests/tf/test_data_modifier.py +++ b/source/tests/tf/test_data_modifier.py @@ -3,10 +3,6 @@ import numpy as np -from deepmd.tf.common import ( - data_requirement, - j_must_have, -) from deepmd.tf.env import ( GLOBAL_NP_FLOAT_PRECISION, tf, @@ -62,16 +58,16 @@ def _setUp(self): rcut = model.model.get_rcut() # init data system - systems = j_must_have(jdata["training"], "systems") + systems = jdata["training"]["systems"] # systems[0] = tests_path / systems[0] systems = [tests_path / ii for ii in systems] set_pfx = "set" - batch_size = j_must_have(jdata["training"], "batch_size") - test_size = j_must_have(jdata["training"], "numb_test") + batch_size = jdata["training"]["batch_size"] + test_size = jdata["training"]["numb_test"] data = DeepmdDataSystem( systems, batch_size, test_size, rcut, set_prefix=set_pfx ) - data.add_dict(data_requirement) + data.add_data_requirements(model.data_requirements) # clear the default graph tf.reset_default_graph() diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index 1fdd8b4967..b79f1b5706 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -4,10 +4,6 @@ import numpy as np -from deepmd.tf.common import ( - data_requirement, - j_must_have, -) from deepmd.tf.env import ( GLOBAL_NP_FLOAT_PRECISION, tf, @@ -65,14 +61,14 @@ def _setUp(self): rcut = model.model.get_rcut() # init data system - systems = j_must_have(jdata["training"], "systems") + systems = jdata["training"]["systems"] set_pfx = "set" - batch_size = j_must_have(jdata["training"], "batch_size") - test_size = j_must_have(jdata["training"], "numb_test") + batch_size = jdata["training"]["batch_size"] + test_size = jdata["training"]["numb_test"] data = DeepmdDataSystem( systems, batch_size, test_size, rcut, set_prefix=set_pfx ) - data.add_dict(data_requirement) + data.add_data_requirements(model.data_requirements) # clear the default graph tf.reset_default_graph() diff --git a/source/tests/tf/test_data_requirement.py b/source/tests/tf/test_data_requirement.py deleted file mode 100644 index e825bc3f92..0000000000 --- a/source/tests/tf/test_data_requirement.py +++ /dev/null @@ -1,19 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import unittest - -from deepmd.tf.common import ( - add_data_requirement, - data_requirement, -) - - -class TestDataRequirement(unittest.TestCase): - def test_add(self): - add_data_requirement("test", 3) - self.assertEqual(data_requirement["test"]["ndof"], 3) - self.assertEqual(data_requirement["test"]["atomic"], False) - self.assertEqual(data_requirement["test"]["must"], False) - self.assertEqual(data_requirement["test"]["high_prec"], False) - self.assertEqual(data_requirement["test"]["repeat"], 1) - self.assertEqual(data_requirement["test"]["default"], 0.0) - self.assertEqual(data_requirement["test"]["output_natoms_for_type_sel"], False) diff --git a/source/tests/tf/test_descrpt_hybrid.py b/source/tests/tf/test_descrpt_hybrid.py index 4023521fca..76c0932316 100644 --- a/source/tests/tf/test_descrpt_hybrid.py +++ b/source/tests/tf/test_descrpt_hybrid.py @@ -4,9 +4,6 @@ import numpy as np from packaging.version import parse as parse_version -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptHybrid, ) @@ -40,10 +37,8 @@ def test_descriptor_hybrid(self): jfile = "water_hybrid.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") batch_size = 2 test_size = 1 rcut = 6 diff --git a/source/tests/tf/test_descrpt_se_a_mask.py b/source/tests/tf/test_descrpt_se_a_mask.py index 7f2461a6f4..b495072c7d 100644 --- a/source/tests/tf/test_descrpt_se_a_mask.py +++ b/source/tests/tf/test_descrpt_se_a_mask.py @@ -3,9 +3,6 @@ import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeAMask, ) @@ -231,12 +228,12 @@ def test_descriptor_se_a_mask(self): jdata["training"]["validation_data"]["systems"] = [ str(tests_path / "data_dp_mask") ] - systems = j_must_have(jdata["training"]["validation_data"], "systems") + systems = jdata["training"]["validation_data"]["systems"] set_pfx = "set" batch_size = 2 test_size = 1 rcut = 20.0 # For DataSystem interface compatibility, not used in this test. - sel = j_must_have(jdata["model"]["descriptor"], "sel") + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(sel) total_atom_num = np.cumsum(sel)[-1] diff --git a/source/tests/tf/test_descrpt_se_a_type.py b/source/tests/tf/test_descrpt_se_a_type.py index d321183c54..5bc4680d6d 100644 --- a/source/tests/tf/test_descrpt_se_a_type.py +++ b/source/tests/tf/test_descrpt_se_a_type.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -33,15 +30,12 @@ def test_descriptor_two_sides(self): jfile = "water_se_a_type.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") batch_size = 2 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - sel = j_must_have(jdata["model"]["descriptor"], "sel") + rcut = jdata["model"]["descriptor"]["rcut"] + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(sel) data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -197,15 +191,14 @@ def test_descriptor_one_side(self): jfile = "water_se_a_type.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - sel = j_must_have(jdata["model"]["descriptor"], "sel") + rcut = jdata["model"]["descriptor"]["rcut"] + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(sel) data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_descrpt_se_atten.py b/source/tests/tf/test_descrpt_se_atten.py index 610788d32e..84325cadc9 100644 --- a/source/tests/tf/test_descrpt_se_atten.py +++ b/source/tests/tf/test_descrpt_se_atten.py @@ -5,9 +5,6 @@ import numpy as np from packaging.version import parse as parse_version -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeAtten, ) @@ -42,15 +39,14 @@ def test_descriptor_two_sides(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 2 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - sel = j_must_have(jdata["model"]["descriptor"], "sel") + rcut = jdata["model"]["descriptor"]["rcut"] + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(jdata["model"]["type_map"]) data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -219,15 +215,14 @@ def test_descriptor_one_side(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - sel = j_must_have(jdata["model"]["descriptor"], "sel") + rcut = jdata["model"]["descriptor"]["rcut"] + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(jdata["model"]["type_map"]) data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -397,15 +392,14 @@ def test_stripped_type_embedding_descriptor_two_sides(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 2 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - sel = j_must_have(jdata["model"]["descriptor"], "sel") + rcut = jdata["model"]["descriptor"]["rcut"] + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(jdata["model"]["type_map"]) data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -568,15 +562,14 @@ def test_compressible_descriptor_two_sides(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 2 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - sel = j_must_have(jdata["model"]["descriptor"], "sel") + rcut = jdata["model"]["descriptor"]["rcut"] + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(jdata["model"]["type_map"]) data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_dipole_se_a.py b/source/tests/tf/test_dipole_se_a.py index 5764861ca3..77e8fc4f0a 100644 --- a/source/tests/tf/test_dipole_se_a.py +++ b/source/tests/tf/test_dipole_se_a.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -38,14 +35,13 @@ def test_model(self): jfile = "polar_se_a.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_dipole_se_a_tebd.py b/source/tests/tf/test_dipole_se_a_tebd.py index 769a715346..66be544e78 100644 --- a/source/tests/tf/test_dipole_se_a_tebd.py +++ b/source/tests/tf/test_dipole_se_a_tebd.py @@ -4,9 +4,6 @@ import numpy as np from packaging.version import parse as parse_version -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -48,14 +45,13 @@ def test_model(self): jfile = "polar_se_a_tebd.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_fitting_dos.py b/source/tests/tf/test_fitting_dos.py index c414c543fe..e66f0078c6 100644 --- a/source/tests/tf/test_fitting_dos.py +++ b/source/tests/tf/test_fitting_dos.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -33,15 +30,14 @@ def test_fitting(self): jfile = "train_dos.json" jdata = j_loader(jfile) - systems = j_must_have(jdata["training"], "systems") + systems = jdata["training"]["systems"] set_pfx = "set" - batch_size = j_must_have(jdata["training"], "batch_size") - test_size = j_must_have(jdata["training"], "numb_test") + batch_size = jdata["training"]["batch_size"] + test_size = jdata["training"]["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata["training"], "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - sel = j_must_have(jdata["model"]["descriptor"], "sel") + rcut = jdata["model"]["descriptor"]["rcut"] + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(sel) data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_fitting_ener_type.py b/source/tests/tf/test_fitting_ener_type.py index 5c3027924b..94ee99c386 100644 --- a/source/tests/tf/test_fitting_ener_type.py +++ b/source/tests/tf/test_fitting_ener_type.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -33,15 +30,14 @@ def test_fitting(self): jfile = "water_se_a_type.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - sel = j_must_have(jdata["model"]["descriptor"], "sel") + rcut = jdata["model"]["descriptor"]["rcut"] + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(sel) data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_loss_gf.py b/source/tests/tf/test_loss_gf.py index 78e5404e03..116b98b649 100644 --- a/source/tests/tf/test_loss_gf.py +++ b/source/tests/tf/test_loss_gf.py @@ -5,6 +5,9 @@ from deepmd.tf.loss import ( EnerStdLoss, ) +from deepmd.utils.data import ( + DataRequirementItem, +) class TestLossGf(tf.test.TestCase): @@ -26,6 +29,62 @@ def setUp(self): numb_generalized_coord=2, ) + def test_label_requirements(self): + """Test label_requirements are expected.""" + self.assertCountEqual( + self.loss.label_requirement, + [ + DataRequirementItem( + "energy", + 1, + atomic=False, + must=False, + high_prec=True, + repeat=1, + ), + DataRequirementItem( + "force", + 3, + atomic=True, + must=False, + high_prec=False, + repeat=1, + ), + DataRequirementItem( + "virial", + 9, + atomic=False, + must=False, + high_prec=False, + repeat=1, + ), + DataRequirementItem( + "atom_pref", + 1, + atomic=True, + must=False, + high_prec=False, + repeat=3, + ), + DataRequirementItem( + "atom_ener", + 1, + atomic=True, + must=False, + high_prec=False, + repeat=1, + ), + DataRequirementItem( + "drdq", + 2 * 3, + atomic=True, + must=False, + high_prec=False, + repeat=1, + ), + ], + ) + def test_build_loss(self): natoms = tf.constant([6, 6]) model_dict = { diff --git a/source/tests/tf/test_model_dos.py b/source/tests/tf/test_model_dos.py index 6ba1d73a0c..fc6e4fdb2d 100644 --- a/source/tests/tf/test_model_dos.py +++ b/source/tests/tf/test_model_dos.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -40,14 +37,13 @@ def test_model(self): jfile = "train_dos.json" jdata = j_loader(jfile) - systems = j_must_have(jdata["training"], "systems") + systems = jdata["training"]["systems"] set_pfx = "set" - batch_size = j_must_have(jdata["training"], "batch_size") - test_size = j_must_have(jdata["training"], "numb_test") + batch_size = jdata["training"]["batch_size"] + test_size = jdata["training"]["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata["training"], "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_model_loc_frame.py b/source/tests/tf/test_model_loc_frame.py index 91d38bab52..a043d16677 100644 --- a/source/tests/tf/test_model_loc_frame.py +++ b/source/tests/tf/test_model_loc_frame.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptLocFrame, ) @@ -35,14 +32,14 @@ def setUp(self): def test_model(self): jfile = "water.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + stop_batch = jdata["stop_batch"] + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_model_pairtab.py b/source/tests/tf/test_model_pairtab.py index 5835d67343..5caeb0a053 100644 --- a/source/tests/tf/test_model_pairtab.py +++ b/source/tests/tf/test_model_pairtab.py @@ -2,9 +2,6 @@ import numpy as np import scipy.spatial.distance -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.env import ( tf, ) @@ -30,7 +27,7 @@ def setUp(self): def test_model(self): jfile = "water.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" batch_size = 1 test_size = 1 @@ -42,7 +39,7 @@ def test_model(self): "rcut": 6, "sel": [6], } - rcut = j_must_have(jdata["model"], "rcut") + rcut = jdata["model"]["rcut"] def pair_pot(r: float): # LJ, as exmaple diff --git a/source/tests/tf/test_model_se_a.py b/source/tests/tf/test_model_se_a.py index ad2c1b7ced..1d67ef5fab 100644 --- a/source/tests/tf/test_model_se_a.py +++ b/source/tests/tf/test_model_se_a.py @@ -2,9 +2,6 @@ import dpdata import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -59,14 +56,13 @@ def test_model_atom_ener(self): sys.data["forces"] = np.zeros([nframes, natoms, 3]) sys.to_deepmd_npy("system", prec=np.float64) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) test_data = data.get_test() @@ -140,14 +136,13 @@ def test_model(self): jfile = "water_se_a.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -265,6 +260,9 @@ def test_model(self): np.testing.assert_almost_equal(f, reff, places) np.testing.assert_almost_equal(v, refv, places) + # test input requirement for the model + self.assertCountEqual(model.input_requirement, []) + def test_model_atom_ener_type_embedding(self): """Test atom ener with type embedding.""" jfile = "water_se_a.json" @@ -286,14 +284,13 @@ def test_model_atom_ener_type_embedding(self): sys.data["forces"] = np.zeros([nframes, natoms, 3]) sys.to_deepmd_npy("system", prec=np.float64) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) test_data = data.get_test() diff --git a/source/tests/tf/test_model_se_a_aparam.py b/source/tests/tf/test_model_se_a_aparam.py index e44e1c8c9f..fcc7dd52d1 100644 --- a/source/tests/tf/test_model_se_a_aparam.py +++ b/source/tests/tf/test_model_se_a_aparam.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -16,6 +13,9 @@ from deepmd.tf.model import ( EnerModel, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from .common import ( DataSystem, @@ -35,14 +35,13 @@ def setUp(self): def test_model(self): jfile = "water_se_a_aparam.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -165,3 +164,9 @@ def test_model(self): np.testing.assert_almost_equal(e, refe, places) np.testing.assert_almost_equal(f, reff, places) np.testing.assert_almost_equal(v, refv, places) + + # test input requirement for the model + self.assertCountEqual( + model.input_requirement, + [DataRequirementItem("aparam", 2, atomic=True, must=True, high_prec=False)], + ) diff --git a/source/tests/tf/test_model_se_a_ebd.py b/source/tests/tf/test_model_se_a_ebd.py index 8f91295f26..9f48076de8 100644 --- a/source/tests/tf/test_model_se_a_ebd.py +++ b/source/tests/tf/test_model_se_a_ebd.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor.se_a_ebd import ( DescrptSeAEbd, ) @@ -36,14 +33,14 @@ def test_model(self): jfile = "water_se_a_ebd.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + stop_batch = jdata["stop_batch"] + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_model_se_a_ebd_v2.py b/source/tests/tf/test_model_se_a_ebd_v2.py index 6156455402..1aa4fdb92b 100644 --- a/source/tests/tf/test_model_se_a_ebd_v2.py +++ b/source/tests/tf/test_model_se_a_ebd_v2.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor.se_a_ebd_v2 import ( DescrptSeAEbdV2, ) @@ -39,14 +36,13 @@ def test_model(self): jfile = "water_se_a_ebd.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_model_se_a_fparam.py b/source/tests/tf/test_model_se_a_fparam.py index ce31f94488..7529873e83 100644 --- a/source/tests/tf/test_model_se_a_fparam.py +++ b/source/tests/tf/test_model_se_a_fparam.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -16,6 +13,9 @@ from deepmd.tf.model import ( EnerModel, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from .common import ( DataSystem, @@ -36,14 +36,13 @@ def test_model(self): jfile = "water_se_a_fparam.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -166,3 +165,13 @@ def test_model(self): np.testing.assert_almost_equal(e, refe, places) np.testing.assert_almost_equal(f, reff, places) np.testing.assert_almost_equal(v, refv, places) + + # test input requirement for the model + self.assertCountEqual( + model.input_requirement, + [ + DataRequirementItem( + "fparam", 2, atomic=False, must=True, high_prec=False + ) + ], + ) diff --git a/source/tests/tf/test_model_se_a_srtab.py b/source/tests/tf/test_model_se_a_srtab.py index ae192b09ec..dddd007006 100644 --- a/source/tests/tf/test_model_se_a_srtab.py +++ b/source/tests/tf/test_model_se_a_srtab.py @@ -3,9 +3,6 @@ import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -53,14 +50,13 @@ def test_model(self): jfile = "water_se_a_srtab.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_model_se_a_type.py b/source/tests/tf/test_model_se_a_type.py index f92fdaa2a7..b0f5da6b7e 100644 --- a/source/tests/tf/test_model_se_a_type.py +++ b/source/tests/tf/test_model_se_a_type.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -39,14 +36,13 @@ def test_model(self): jfile = "water_se_a_type.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_model_se_atten.py b/source/tests/tf/test_model_se_atten.py index e1143d8c2e..06945fab9c 100644 --- a/source/tests/tf/test_model_se_atten.py +++ b/source/tests/tf/test_model_se_atten.py @@ -5,9 +5,6 @@ import numpy as np from packaging.version import parse as parse_version -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeAtten, ) @@ -50,14 +47,12 @@ def test_model(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -204,13 +199,11 @@ def test_exclude_types(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] ntypes = 2 data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -277,14 +270,11 @@ def test_compressible_model(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -433,13 +423,12 @@ def test_compressible_exclude_types(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] ntypes = 2 data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -508,14 +497,12 @@ def test_stripped_type_embedding_model(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -667,13 +654,13 @@ def test_stripped_type_embedding_exclude_types(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] ntypes = 2 data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -747,14 +734,11 @@ def test_smoothness_of_stripped_type_embedding_smooth_model(self): jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -896,11 +880,11 @@ def test_smoothness_of_stripped_type_embedding_smooth_model_excluded_types(self) jfile = "water_se_atten.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_model_se_r.py b/source/tests/tf/test_model_se_r.py index 90976b9d47..628a59d320 100644 --- a/source/tests/tf/test_model_se_r.py +++ b/source/tests/tf/test_model_se_r.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeR, ) @@ -36,14 +33,14 @@ def test_model(self): jfile = "water_se_r.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + stop_batch = jdata["stop_batch"] + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_model_se_t.py b/source/tests/tf/test_model_se_t.py index 1948c8b54b..6167e2e224 100644 --- a/source/tests/tf/test_model_se_t.py +++ b/source/tests/tf/test_model_se_t.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeT, ) @@ -36,14 +33,11 @@ def test_model(self): jfile = "water_se_t.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_model_spin.py b/source/tests/tf/test_model_spin.py index 65a866b021..0e9c0678c2 100644 --- a/source/tests/tf/test_model_spin.py +++ b/source/tests/tf/test_model_spin.py @@ -3,9 +3,6 @@ import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -48,18 +45,16 @@ def test_model_spin(self): # set system information set_pfx = "set" - batch_size = j_must_have(jdata["training"]["training_data"], "batch_size") batch_size = 2 - test_size = j_must_have(jdata["training"]["validation_data"], "numb_btch") - stop_batch = j_must_have(jdata["training"], "numb_steps") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + test_size = jdata["training"]["validation_data"]["numb_btch"] + rcut = jdata["model"]["descriptor"]["rcut"] jdata["training"]["training_data"]["systems"] = [ str(tests_path / "model_spin/") ] jdata["training"]["validation_data"]["systems"] = [ str(tests_path / "model_spin/") ] - systems = j_must_have(jdata["training"]["training_data"], "systems") + systems = jdata["training"]["training_data"]["systems"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) test_data = data.get_test() diff --git a/source/tests/tf/test_pairwise_dprc.py b/source/tests/tf/test_pairwise_dprc.py index 38b8d8b775..3a5836ce45 100644 --- a/source/tests/tf/test_pairwise_dprc.py +++ b/source/tests/tf/test_pairwise_dprc.py @@ -13,7 +13,6 @@ ) from deepmd.tf.common import ( j_loader, - j_must_have, ) from deepmd.tf.env import ( GLOBAL_ENER_FLOAT_PRECISION, @@ -34,6 +33,9 @@ from deepmd.tf.utils.sess import ( run_sess, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from .common import ( run_dp, @@ -434,7 +436,7 @@ def test_model_ener(self): idxs = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]) np.save("system/set.000/aparam.npy", idxs) - systems = j_must_have(jdata["training"]["training_data"], "systems") + systems = jdata["training"]["training_data"]["systems"] batch_size = 1 test_size = 1 rcut = model.get_rcut() @@ -523,6 +525,12 @@ def test_model_ener(self): self.assertAllClose(e[0], 0.189075, 1e-6) self.assertAllClose(f[0, 0], 0.060047, 1e-6) + # test input requirement for the model + self.assertCountEqual( + model.input_requirement, + [DataRequirementItem("aparam", 1, atomic=True, must=True, high_prec=False)], + ) + def test_nloc(self): jfile = tests_path / "pairwise_dprc.json" jdata = j_loader(jfile) @@ -616,7 +624,7 @@ def test_nloc(self): idxs = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]) np.save("system/set.000/aparam.npy", idxs) - systems = j_must_have(jdata["training"]["training_data"], "systems") + systems = jdata["training"]["training_data"]["systems"] batch_size = 1 test_size = 1 rcut = model.get_rcut() diff --git a/source/tests/tf/test_polar_se_a.py b/source/tests/tf/test_polar_se_a.py index 63f4bfa589..f71ead0d0f 100644 --- a/source/tests/tf/test_polar_se_a.py +++ b/source/tests/tf/test_polar_se_a.py @@ -5,9 +5,6 @@ import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -45,14 +42,11 @@ def test_model(self): jfile = "polar_se_a.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -229,7 +223,7 @@ def test_data_stat(self): batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DeepmdDataSystem(systems, batch_size, test_size, rcut) data.add( diff --git a/source/tests/tf/test_polar_se_a_tebd.py b/source/tests/tf/test_polar_se_a_tebd.py index f9f4954d74..2abaedd40a 100644 --- a/source/tests/tf/test_polar_se_a_tebd.py +++ b/source/tests/tf/test_polar_se_a_tebd.py @@ -4,9 +4,6 @@ import numpy as np from packaging.version import parse as parse_version -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( DescrptSeA, ) @@ -48,14 +45,12 @@ def test_model(self): jfile = "polar_se_a_tebd.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - stop_batch = j_must_have(jdata, "stop_batch") - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + rcut = jdata["model"]["descriptor"]["rcut"] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_type_one_side.py b/source/tests/tf/test_type_one_side.py index e2f992f212..2ab7cc03a7 100644 --- a/source/tests/tf/test_type_one_side.py +++ b/source/tests/tf/test_type_one_side.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.descriptor import ( Descriptor, ) @@ -41,14 +38,14 @@ def test_descriptor_one_side_exclude_types(self): jfile = "water_se_a.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - sel = j_must_have(jdata["model"]["descriptor"], "sel") + rcut = jdata["model"]["descriptor"]["rcut"] + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(sel) data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) @@ -148,14 +145,14 @@ def test_se_r_one_side_exclude_types(self): jfile = "water_se_r.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] set_pfx = "set" - batch_size = j_must_have(jdata, "batch_size") - test_size = j_must_have(jdata, "numb_test") + batch_size = jdata["batch_size"] + test_size = jdata["numb_test"] batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - sel = j_must_have(jdata["model"]["descriptor"], "sel") + rcut = jdata["model"]["descriptor"]["rcut"] + sel = jdata["model"]["descriptor"]["sel"] ntypes = len(sel) data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) diff --git a/source/tests/tf/test_virtual_type.py b/source/tests/tf/test_virtual_type.py index a3e87a35ed..40ec2674dd 100644 --- a/source/tests/tf/test_virtual_type.py +++ b/source/tests/tf/test_virtual_type.py @@ -6,9 +6,6 @@ import numpy as np -from deepmd.tf.common import ( - j_must_have, -) from deepmd.tf.infer import ( DeepPot, ) @@ -130,11 +127,11 @@ def test_data_mixed_type(self): jfile = "water_se_atten_mixed_type.json" jdata = j_loader(jfile) - systems = j_must_have(jdata, "systems") + systems = jdata["systems"] batch_size = 1 test_size = 1 - rcut = j_must_have(jdata["model"]["descriptor"], "rcut") - type_map = j_must_have(jdata["model"], "type_map") + rcut = jdata["model"]["descriptor"]["rcut"] + type_map = jdata["model"]["type_map"] data = DeepmdDataSystem(systems, batch_size, test_size, rcut, type_map=type_map) data.get_batch() diff --git a/source/tests/universal/__init__.py b/source/tests/universal/__init__.py new file mode 100644 index 0000000000..3c8d925dcc --- /dev/null +++ b/source/tests/universal/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Universal tests for the project.""" diff --git a/source/tests/universal/common/__init__.py b/source/tests/universal/common/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/common/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/common/backend.py b/source/tests/universal/common/backend.py new file mode 100644 index 0000000000..d5747b77b7 --- /dev/null +++ b/source/tests/universal/common/backend.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Common test case.""" + +from abc import ( + ABC, + abstractmethod, +) + + +class BackendTestCase(ABC): + """Backend test case.""" + + module: object + """Module to test.""" + + @property + @abstractmethod + def modules_to_test(self) -> list: + pass + + @abstractmethod + def forward_wrapper(self, x): + pass diff --git a/source/tests/universal/common/cases/__init__.py b/source/tests/universal/common/cases/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/common/cases/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/common/cases/atomic_model/__init__.py b/source/tests/universal/common/cases/atomic_model/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/common/cases/atomic_model/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/common/cases/atomic_model/ener_model.py b/source/tests/universal/common/cases/atomic_model/ener_model.py new file mode 100644 index 0000000000..0f1daaf87b --- /dev/null +++ b/source/tests/universal/common/cases/atomic_model/ener_model.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + + +from .utils import ( + AtomicModelTestCase, +) + + +class EnerAtomicModelTest(AtomicModelTestCase): + def setUp(self) -> None: + self.expected_rcut = 5.0 + self.expected_type_map = ["foo", "bar"] + self.expected_dim_fparam = 0 + self.expected_dim_aparam = 0 + self.expected_sel_type = [0, 1] + self.expected_aparam_nall = False + self.expected_model_output_type = ["energy", "mask"] + self.expected_sel = [8, 12] diff --git a/source/tests/universal/common/cases/atomic_model/utils.py b/source/tests/universal/common/cases/atomic_model/utils.py new file mode 100644 index 0000000000..3b5fc64fda --- /dev/null +++ b/source/tests/universal/common/cases/atomic_model/utils.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, + Callable, + List, +) + +import numpy as np + +from deepmd.dpmodel.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + + +class AtomicModelTestCase: + """Common test case for atomic model.""" + + expected_type_map: List[str] + """Expected type map.""" + expected_rcut: float + """Expected cut-off radius.""" + expected_dim_fparam: int + """Expected number (dimension) of frame parameters.""" + expected_dim_aparam: int + """Expected number (dimension) of atomic parameters.""" + expected_sel_type: List[int] + """Expected selected atom types.""" + expected_aparam_nall: bool + """Expected shape of atomic parameters.""" + expected_model_output_type: List[str] + """Expected output type for the model.""" + expected_sel: List[int] + """Expected number of neighbors.""" + forward_wrapper: Callable[[Any], Any] + """Calss wrapper for forward method.""" + + def test_get_type_map(self): + """Test get_type_map.""" + for module in self.modules_to_test: + self.assertEqual(module.get_type_map(), self.expected_type_map) + + def test_get_rcut(self): + """Test get_rcut.""" + for module in self.modules_to_test: + self.assertAlmostEqual(module.get_rcut(), self.expected_rcut) + + def test_get_dim_fparam(self): + """Test get_dim_fparam.""" + for module in self.modules_to_test: + self.assertEqual(module.get_dim_fparam(), self.expected_dim_fparam) + + def test_get_dim_aparam(self): + """Test get_dim_aparam.""" + for module in self.modules_to_test: + self.assertEqual(module.get_dim_aparam(), self.expected_dim_aparam) + + def test_get_sel_type(self): + """Test get_sel_type.""" + for module in self.modules_to_test: + self.assertEqual(module.get_sel_type(), self.expected_sel_type) + + def test_is_aparam_nall(self): + """Test is_aparam_nall.""" + for module in self.modules_to_test: + self.assertEqual(module.is_aparam_nall(), self.expected_aparam_nall) + + def test_get_nnei(self): + """Test get_nnei.""" + expected_nnei = sum(self.expected_sel) + for module in self.modules_to_test: + self.assertEqual(module.get_nnei(), expected_nnei) + + def test_get_ntypes(self): + """Test get_ntypes.""" + for module in self.modules_to_test: + self.assertEqual(module.get_ntypes(), len(self.expected_type_map)) + + def test_forward(self): + """Test forward.""" + nf = 1 + coord = np.array( + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + ], + dtype=np.float64, + ).reshape([nf, -1]) + atype = np.array([0, 0, 1], dtype=int).reshape([nf, -1]) + cell = 6.0 * np.eye(3).reshape([nf, 9]) + coord_ext, atype_ext, mapping, nlist = extend_input_and_build_neighbor_list( + coord, + atype, + self.expected_rcut, + self.expected_sel, + mixed_types=True, + box=cell, + ) + ret_lower = [] + for module in self.modules_to_test: + module = self.forward_wrapper(module) + + ret_lower.append(module(coord_ext, atype_ext, nlist)) + for kk in ret_lower[0].keys(): + subret = [] + for rr in ret_lower: + if rr is not None: + subret.append(rr[kk]) + if len(subret): + for ii, rr in enumerate(subret[1:]): + if subret[0] is None: + assert rr is None + else: + np.testing.assert_allclose( + subret[0], rr, err_msg=f"compare {kk} between 0 and {ii}" + ) diff --git a/source/tests/universal/common/cases/model/__init__.py b/source/tests/universal/common/cases/model/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/common/cases/model/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/common/cases/model/ener_model.py b/source/tests/universal/common/cases/model/ener_model.py new file mode 100644 index 0000000000..35d44f9784 --- /dev/null +++ b/source/tests/universal/common/cases/model/ener_model.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + + +from .utils import ( + ModelTestCase, +) + + +class EnerModelTest(ModelTestCase): + def setUp(self) -> None: + self.expected_rcut = 5.0 + self.expected_type_map = ["foo", "bar"] + self.expected_dim_fparam = 0 + self.expected_dim_aparam = 0 + self.expected_sel_type = [0, 1] + self.expected_aparam_nall = False + self.expected_model_output_type = ["energy", "mask"] + self.expected_sel = [8, 12] diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py new file mode 100644 index 0000000000..d67ac8e80d --- /dev/null +++ b/source/tests/universal/common/cases/model/utils.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, + Callable, + List, +) + +import numpy as np + +from deepmd.dpmodel.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + + +class ModelTestCase: + """Common test case for model.""" + + expected_type_map: List[str] + """Expected type map.""" + expected_rcut: float + """Expected cut-off radius.""" + expected_dim_fparam: int + """Expected number (dimension) of frame parameters.""" + expected_dim_aparam: int + """Expected number (dimension) of atomic parameters.""" + expected_sel_type: List[int] + """Expected selected atom types.""" + expected_aparam_nall: bool + """Expected shape of atomic parameters.""" + expected_model_output_type: List[str] + """Expected output type for the model.""" + expected_sel: List[int] + """Expected number of neighbors.""" + forward_wrapper: Callable[[Any], Any] + """Calss wrapper for forward method.""" + + def test_get_type_map(self): + """Test get_type_map.""" + for module in self.modules_to_test: + self.assertEqual(module.get_type_map(), self.expected_type_map) + + def test_get_rcut(self): + """Test get_rcut.""" + for module in self.modules_to_test: + self.assertAlmostEqual(module.get_rcut(), self.expected_rcut) + + def test_get_dim_fparam(self): + """Test get_dim_fparam.""" + for module in self.modules_to_test: + self.assertEqual(module.get_dim_fparam(), self.expected_dim_fparam) + + def test_get_dim_aparam(self): + """Test get_dim_aparam.""" + for module in self.modules_to_test: + self.assertEqual(module.get_dim_aparam(), self.expected_dim_aparam) + + def test_get_sel_type(self): + """Test get_sel_type.""" + for module in self.modules_to_test: + self.assertEqual(module.get_sel_type(), self.expected_sel_type) + + def test_is_aparam_nall(self): + """Test is_aparam_nall.""" + for module in self.modules_to_test: + self.assertEqual(module.is_aparam_nall(), self.expected_aparam_nall) + + def test_model_output_type(self): + """Test model_output_type.""" + for module in self.modules_to_test: + self.assertEqual( + module.model_output_type(), self.expected_model_output_type + ) + + def test_get_nnei(self): + """Test get_nnei.""" + expected_nnei = sum(self.expected_sel) + for module in self.modules_to_test: + self.assertEqual(module.get_nnei(), expected_nnei) + + def test_get_ntypes(self): + """Test get_ntypes.""" + for module in self.modules_to_test: + self.assertEqual(module.get_ntypes(), len(self.expected_type_map)) + + def test_forward(self): + """Test forward and forward_lower.""" + nf = 1 + coord = np.array( + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + ], + dtype=np.float64, + ).reshape([nf, -1]) + atype = np.array([0, 0, 1], dtype=int).reshape([nf, -1]) + cell = 6.0 * np.eye(3).reshape([nf, 9]) + coord_ext, atype_ext, mapping, nlist = extend_input_and_build_neighbor_list( + coord, + atype, + self.expected_rcut, + self.expected_sel, + mixed_types=True, + box=cell, + ) + ret = [] + ret_lower = [] + for module in self.modules_to_test: + module = self.forward_wrapper(module) + ret.append(module(coord, atype, cell)) + + ret_lower.append(module.forward_lower(coord_ext, atype_ext, nlist)) + for kk in ret[0].keys(): + subret = [] + for rr in ret: + if rr is not None: + subret.append(rr[kk]) + if len(subret): + for ii, rr in enumerate(subret[1:]): + if subret[0] is None: + assert rr is None + else: + np.testing.assert_allclose( + subret[0], rr, err_msg=f"compare {kk} between 0 and {ii}" + ) + for kk in ret_lower[0].keys(): + subret = [] + for rr in ret_lower: + if rr is not None: + subret.append(rr[kk]) + if len(subret): + for ii, rr in enumerate(subret[1:]): + if subret[0] is None: + assert rr is None + else: + np.testing.assert_allclose( + subret[0], rr, err_msg=f"compare {kk} between 0 and {ii}" + ) + same_keys = set(ret[0].keys()) & set(ret_lower[0].keys()) + self.assertTrue(same_keys) + for key in same_keys: + for rr in ret: + if rr[key] is not None: + rr1 = rr[key] + break + else: + continue + for rr in ret_lower: + if rr[key] is not None: + rr2 = rr[key] + break + else: + continue + np.testing.assert_allclose(rr1, rr2) diff --git a/source/tests/universal/dpmodel/__init__.py b/source/tests/universal/dpmodel/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/dpmodel/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/dpmodel/atomc_model/__init__.py b/source/tests/universal/dpmodel/atomc_model/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/dpmodel/atomc_model/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/dpmodel/atomc_model/test_ener_atomic_model.py b/source/tests/universal/dpmodel/atomc_model/test_ener_atomic_model.py new file mode 100644 index 0000000000..6cf4598646 --- /dev/null +++ b/source/tests/universal/dpmodel/atomc_model/test_ener_atomic_model.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +from deepmd.dpmodel.atomic_model.dp_atomic_model import ( + DPAtomicModel, +) +from deepmd.dpmodel.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting.ener_fitting import ( + EnergyFittingNet, +) + +from ...common.cases.atomic_model.ener_model import ( + EnerAtomicModelTest, +) +from ..backend import ( + DPTestCase, +) + + +class TestEnergyAtomicModelDP(unittest.TestCase, EnerAtomicModelTest, DPTestCase): + def setUp(self): + EnerAtomicModelTest.setUp(self) + ds = DescrptSeA( + rcut=self.expected_rcut, + rcut_smth=self.expected_rcut / 2, + sel=self.expected_sel, + ) + ft = EnergyFittingNet( + ntypes=len(self.expected_type_map), + dim_descrpt=ds.get_dim_out(), + mixed_types=ds.mixed_types(), + ) + self.module = DPAtomicModel( + ds, + ft, + type_map=self.expected_type_map, + ) diff --git a/source/tests/universal/dpmodel/backend.py b/source/tests/universal/dpmodel/backend.py new file mode 100644 index 0000000000..61982fea98 --- /dev/null +++ b/source/tests/universal/dpmodel/backend.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.common import ( + NativeOP, +) + +from ..common.backend import ( + BackendTestCase, +) + + +class DPTestCase(BackendTestCase): + """Common test case.""" + + module: NativeOP + """DP module to test.""" + + def forward_wrapper(self, x): + return x + + @property + def deserialized_module(self): + return self.module.deserialize(self.module.serialize()) + + @property + def modules_to_test(self): + modules = [ + self.module, + self.deserialized_module, + ] + return modules diff --git a/source/tests/universal/dpmodel/model/__init__.py b/source/tests/universal/dpmodel/model/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/dpmodel/model/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/dpmodel/model/test_ener_model.py b/source/tests/universal/dpmodel/model/test_ener_model.py new file mode 100644 index 0000000000..506564260f --- /dev/null +++ b/source/tests/universal/dpmodel/model/test_ener_model.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +from deepmd.dpmodel.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting.ener_fitting import ( + EnergyFittingNet, +) +from deepmd.dpmodel.model.ener_model import ( + EnergyModel, +) + +from ...common.cases.model.ener_model import ( + EnerModelTest, +) +from ..backend import ( + DPTestCase, +) + + +class TestEnergyModelDP(unittest.TestCase, EnerModelTest, DPTestCase): + def setUp(self): + EnerModelTest.setUp(self) + ds = DescrptSeA( + rcut=self.expected_rcut, + rcut_smth=self.expected_rcut / 2, + sel=self.expected_sel, + ) + ft = EnergyFittingNet( + ntypes=len(self.expected_type_map), + dim_descrpt=ds.get_dim_out(), + mixed_types=ds.mixed_types(), + ) + self.module = EnergyModel( + ds, + ft, + type_map=self.expected_type_map, + ) diff --git a/source/tests/universal/pt/__init__.py b/source/tests/universal/pt/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/pt/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/pt/atomc_model/__init__.py b/source/tests/universal/pt/atomc_model/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/pt/atomc_model/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/pt/atomc_model/test_ener_atomic_model.py b/source/tests/universal/pt/atomc_model/test_ener_atomic_model.py new file mode 100644 index 0000000000..5ba3be0fad --- /dev/null +++ b/source/tests/universal/pt/atomc_model/test_ener_atomic_model.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +from deepmd.pt.model.atomic_model.dp_atomic_model import ( + DPAtomicModel, +) +from deepmd.pt.model.descriptor.se_a import ( + DescrptSeA, +) +from deepmd.pt.model.task.ener import ( + EnergyFittingNet, +) + +from ...common.cases.atomic_model.ener_model import ( + EnerAtomicModelTest, +) +from ..backend import ( + PTTestCase, +) + + +class TestEnergyAtomicModelDP(unittest.TestCase, EnerAtomicModelTest, PTTestCase): + def setUp(self): + EnerAtomicModelTest.setUp(self) + ds = DescrptSeA( + rcut=self.expected_rcut, + rcut_smth=self.expected_rcut / 2, + sel=self.expected_sel, + ) + ft = EnergyFittingNet( + ntypes=len(self.expected_type_map), + dim_descrpt=ds.get_dim_out(), + mixed_types=ds.mixed_types(), + ) + self.module = DPAtomicModel( + ds, + ft, + type_map=self.expected_type_map, + ) diff --git a/source/tests/universal/pt/backend.py b/source/tests/universal/pt/backend.py new file mode 100644 index 0000000000..61110a0cc6 --- /dev/null +++ b/source/tests/universal/pt/backend.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import torch + +from deepmd.pt.utils.utils import ( + to_numpy_array, + to_torch_tensor, +) + + +class PTTestCase: + """Common test case.""" + + module: "torch.nn.Module" + """PT module to test.""" + + @property + def script_module(self): + return torch.jit.script(self.module) + + @property + def deserialized_module(self): + return self.module.deserialize(self.module.serialize()) + + @property + def modules_to_test(self): + modules = [ + self.module, + self.deserialized_module, + ] + return modules + + def test_jit(self): + self.script_module + + def forward_wrapper(self, module): + def create_wrapper_method(method): + def wrapper_method(self, *args, **kwargs): + # convert to torch tensor + args = [to_torch_tensor(arg) for arg in args] + kwargs = {k: to_torch_tensor(v) for k, v in kwargs.items()} + # forward + output = method(*args, **kwargs) + # convert to numpy array + if isinstance(output, tuple): + output = tuple(to_numpy_array(o) for o in output) + elif isinstance(output, dict): + output = {k: to_numpy_array(v) for k, v in output.items()} + else: + output = to_numpy_array(output) + return output + + return wrapper_method + + class wrapper_module: + __call__ = create_wrapper_method(module.__call__) + if hasattr(module, "forward_lower"): + forward_lower = create_wrapper_method(module.forward_lower) + + return wrapper_module() diff --git a/source/tests/universal/pt/model/__init__.py b/source/tests/universal/pt/model/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/pt/model/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/pt/model/test_ener_model.py b/source/tests/universal/pt/model/test_ener_model.py new file mode 100644 index 0000000000..af5d77d5b4 --- /dev/null +++ b/source/tests/universal/pt/model/test_ener_model.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +from deepmd.pt.model.descriptor.se_a import ( + DescrptSeA, +) +from deepmd.pt.model.model.ener_model import ( + EnergyModel, +) +from deepmd.pt.model.task.ener import ( + EnergyFittingNet, +) + +from ...common.cases.model.ener_model import ( + EnerModelTest, +) +from ..backend import ( + PTTestCase, +) + + +class TestEnergyModelDP(unittest.TestCase, EnerModelTest, PTTestCase): + @property + def modules_to_test(self): + # for Model, we can test script module API + modules = [ + *PTTestCase.modules_to_test.fget(self), + self.script_module, + ] + return modules + + def setUp(self): + EnerModelTest.setUp(self) + ds = DescrptSeA( + rcut=self.expected_rcut, + rcut_smth=self.expected_rcut / 2, + sel=self.expected_sel, + ) + ft = EnergyFittingNet( + ntypes=len(self.expected_type_map), + dim_descrpt=ds.get_dim_out(), + mixed_types=ds.mixed_types(), + ) + self.module = EnergyModel( + ds, + ft, + type_map=self.expected_type_map, + )