diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml index 2082e7e4cc..d98f8ca58d 100644 --- a/.github/workflows/test_cc.yml +++ b/.github/workflows/test_cc.yml @@ -56,7 +56,7 @@ jobs: TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp - LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib + LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib:${{ github.workspace }}/libtorch/lib if: ${{ !matrix.check_memleak }} # test ipi - run: pytest --cov=deepmd source/ipi/tests @@ -65,7 +65,7 @@ jobs: TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 PATH: ${{ github.workspace }}/dp_test/bin:$PATH - LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib + LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib:${{ github.workspace }}/libtorch/lib if: ${{ !matrix.check_memleak }} - uses: codecov/codecov-action@v4 env: diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index 0d934e6d77..915d983663 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -38,6 +38,8 @@ jobs: with: useLocalCache: true useCloudCache: false + - name: Install wget and unzip + run: apt-get update && apt-get install -y wget unzip - run: | wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \ && sudo dpkg -i cuda-keyring_1.0-1_all.deb \ @@ -53,7 +55,13 @@ jobs: DP_ENABLE_NATIVE_OPTIMIZATION: 1 - run: dp --version - run: python -m pytest source/tests --durations=0 - - run: source/install/test_cc_local.sh + - name: Download libtorch + run: | + wget https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.1%2Bcu121.zip -O libtorch.zip + unzip libtorch.zip + - run: | + export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch + source/install/test_cc_local.sh env: OMP_NUM_THREADS: 1 TF_INTRA_OP_PARALLELISM_THREADS: 1 @@ -63,7 +71,7 @@ jobs: DP_VARIANT: cuda DP_USE_MPICH2: 1 - run: | - export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/dp_test/lib:$CUDA_PATH/lib64:$LD_LIBRARY_PATH + export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/dp_test/lib:$GITHUB_WORKSPACE/libtorch/lib:$CUDA_PATH/lib64:$LD_LIBRARY_PATH export PATH=$GITHUB_WORKSPACE/dp_test/bin:$PATH python -m pytest source/lmp/tests python -m pytest source/ipi/tests diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 178b286e79..cd349749fa 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -17,6 +17,9 @@ from deepmd.dpmodel.output_def import ( FittingOutputDef, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from .base_atomic_model import ( BaseAtomicModel, @@ -132,6 +135,7 @@ def serialize(self) -> dict: return { "@class": "Model", "type": "standard", + "@version": 1, "type_map": self.type_map, "descriptor": self.descriptor.serialize(), "fitting": self.fitting.serialize(), @@ -140,6 +144,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "DPAtomicModel": data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class") data.pop("type") descriptor_obj = BaseDescriptor.deserialize(data["descriptor"]) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index e1130eaf45..6d8aea499e 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -19,6 +19,9 @@ get_multiple_nlist_key, nlist_distinguish_types, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from ..output_def import ( FittingOutputDef, @@ -185,6 +188,7 @@ def serialize(models) -> dict: return { "@class": "Model", "type": "linear", + "@version": 1, "models": [model.serialize() for model in models], "model_name": [model.__class__.__name__ for model in models], } @@ -192,6 +196,7 @@ def serialize(models) -> dict: @staticmethod def deserialize(data) -> List[BaseAtomicModel]: data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class") data.pop("type") model_names = data["model_name"] @@ -271,6 +276,7 @@ def serialize(self) -> dict: return { "@class": "Model", "type": "zbl", + "@version": 1, "models": LinearAtomicModel.serialize([self.dp_model, self.zbl_model]), "sw_rmin": self.sw_rmin, "sw_rmax": self.sw_rmax, @@ -280,6 +286,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "DPZBLLinearAtomicModel": data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class") data.pop("type") sw_rmin = data["sw_rmin"] diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index dc3dfaf2ed..ddece80f2d 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -16,6 +16,9 @@ from deepmd.utils.pair_tab import ( PairTab, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from .base_atomic_model import ( BaseAtomicModel, @@ -109,6 +112,7 @@ def serialize(self) -> dict: return { "@class": "Model", "type": "pairtab", + "@version": 1, "tab": self.tab.serialize(), "rcut": self.rcut, "sel": self.sel, @@ -117,6 +121,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "PairTabAtomicModel": data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class") data.pop("type") rcut = data["rcut"] diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index a28215c35a..b102933ac9 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -9,6 +9,9 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.version import ( + check_version_compatibility, +) try: from deepmd._version import version as __version__ @@ -345,6 +348,7 @@ def serialize(self) -> dict: return { "@class": "Descriptor", "type": "se_e2_a", + "@version": 1, "rcut": self.rcut, "rcut_smth": self.rcut_smth, "sel": self.sel, @@ -371,6 +375,7 @@ def serialize(self) -> dict: def deserialize(cls, data: dict) -> "DescrptSeA": """Deserialize from dict.""" data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class", None) data.pop("type", None) variables = data.pop("@variables") diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index 77e43f7d85..5973c55353 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -4,6 +4,9 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.version import ( + check_version_compatibility, +) try: from deepmd._version import version as __version__ @@ -282,6 +285,7 @@ def serialize(self) -> dict: return { "@class": "Descriptor", "type": "se_r", + "@version": 1, "rcut": self.rcut, "rcut_smth": self.rcut_smth, "sel": self.sel, @@ -307,6 +311,7 @@ def serialize(self) -> dict: def deserialize(cls, data: dict) -> "DescrptSeR": """Deserialize from dict.""" data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class", None) data.pop("type", None) variables = data.pop("@variables") diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 152836e928..752a550849 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -21,6 +21,9 @@ FittingNet, NetworkCollection, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from .base_fitting import ( BaseFitting, @@ -210,6 +213,7 @@ def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", + "@version": 1, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -241,6 +245,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class") data.pop("type") variables = data.pop("@variables") diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index 2133bc4889..feb3355e77 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -20,6 +20,10 @@ import h5py import numpy as np +from deepmd.utils.version import ( + check_version_compatibility, +) + try: from deepmd._version import version as __version__ except ImportError: @@ -189,6 +193,8 @@ def serialize(self) -> dict: "idt": self.idt, } return { + "@class": "Layer", + "@version": 1, "bias": self.b is not None, "use_timestep": self.idt is not None, "activation_function": self.activation_function, @@ -208,6 +214,8 @@ def deserialize(cls, data: dict) -> "NativeLayer": The dict to deserialize from. """ data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) variables = data.pop("@variables") assert variables["w"] is not None and len(variables["w"].shape) == 2 num_in, num_out = variables["w"].shape @@ -349,7 +357,11 @@ def serialize(self) -> dict: dict The serialized network. """ - return {"layers": [layer.serialize() for layer in self.layers]} + return { + "@class": "NN", + "@version": 1, + "layers": [layer.serialize() for layer in self.layers], + } @classmethod def deserialize(cls, data: dict) -> "NN": @@ -360,6 +372,9 @@ def deserialize(cls, data: dict) -> "NN": data : dict The dict to deserialize from. """ + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) return cls(data["layers"]) def __getitem__(self, key): @@ -471,6 +486,8 @@ def serialize(self) -> dict: The serialized network. """ return { + "@class": "EmbeddingNetwork", + "@version": 1, "in_dim": self.in_dim, "neuron": self.neuron.copy(), "activation_function": self.activation_function, @@ -490,6 +507,8 @@ def deserialize(cls, data: dict) -> "EmbeddingNet": The dict to deserialize from. """ data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) layers = data.pop("layers") obj = cls(**data) super(EN, obj).__init__(layers) @@ -566,6 +585,8 @@ def serialize(self) -> dict: The serialized network. """ return { + "@class": "FittingNetwork", + "@version": 1, "in_dim": self.in_dim, "out_dim": self.out_dim, "neuron": self.neuron.copy(), @@ -586,6 +607,8 @@ def deserialize(cls, data: dict) -> "FittingNet": The dict to deserialize from. """ data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) layers = data.pop("layers") obj = cls(**data) T_Network.__init__(obj, layers) @@ -688,6 +711,8 @@ def serialize(self) -> dict: network_type_map_inv = {v: k for k, v in self.NETWORK_TYPE_MAP.items()} network_type_name = network_type_map_inv[self.network_type] return { + "@class": "NetworkCollection", + "@version": 1, "ndim": self.ndim, "ntypes": self.ntypes, "network_type": network_type_name, @@ -703,4 +728,7 @@ def deserialize(cls, data: dict) -> "NetworkCollection": data : dict The dict to deserialize from. """ + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) return cls(**data) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 35d170cdab..de964b88b9 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -472,6 +472,33 @@ def _standard_input(self, coords, cells, atom_types, fparam, aparam, mixed_type) aparam = np.array(aparam) natoms, nframes = self._get_natoms_and_nframes(coords, atom_types, mixed_type) atom_types = self._expande_atype(atom_types, nframes, mixed_type) + coords = coords.reshape(nframes, natoms, 3) + if cells is not None: + cells = cells.reshape(nframes, 3, 3) + if fparam is not None: + fdim = self.get_dim_fparam() + if fparam.size == nframes * fdim: + fparam = np.reshape(fparam, [nframes, fdim]) + elif fparam.size == fdim: + fparam = np.tile(fparam.reshape([-1]), [nframes, 1]) + else: + raise RuntimeError( + "got wrong size of frame param, should be either %d x %d or %d" + % (nframes, fdim, fdim) + ) + if aparam is not None: + fdim = self.get_dim_aparam() + if aparam.size == nframes * natoms * fdim: + aparam = np.reshape(aparam, [nframes, natoms * fdim]) + elif aparam.size == natoms * fdim: + aparam = np.tile(aparam.reshape([-1]), [nframes, 1]) + elif aparam.size == fdim: + aparam = np.tile(aparam.reshape([-1]), [nframes, natoms]) + else: + raise RuntimeError( + "got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d" + % (nframes, natoms, fdim, natoms, fdim, fdim) + ) return coords, cells, atom_types, fparam, aparam, nframes, natoms def get_sel_type(self) -> List[int]: diff --git a/deepmd/main.py b/deepmd/main.py index 4d2d62ed14..df5c99bb2d 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -226,7 +226,7 @@ def main_parser() -> argparse.ArgumentParser: "--init-frz-model", type=str, default=None, - help="(Supported backend: TensorFlow) Initialize the training from the frozen model.", + help="Initialize the training from the frozen model.", ) parser_train_subgroup.add_argument( "-t", diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 212a6824e7..a317cea6a9 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -65,6 +65,7 @@ def get_trainer( finetune_model=None, model_branch="", force_load=False, + init_frz_model=None, ): # Initialize DDP local_rank = os.environ.get("LOCAL_RANK") @@ -200,6 +201,7 @@ def prepare_trainer_input_single( finetune_model=finetune_model, force_load=force_load, shared_links=shared_links, + init_frz_model=init_frz_model, ) return trainer @@ -243,6 +245,7 @@ def train(FLAGS): FLAGS.finetune, FLAGS.model_branch, FLAGS.force_load, + FLAGS.init_frz_model, ) trainer.run() diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index b13a968a61..f75052166b 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -54,6 +54,9 @@ DEVICE, GLOBAL_PT_FLOAT_PRECISION, ) +from deepmd.pt.utils.utils import ( + to_torch_tensor, +) if TYPE_CHECKING: import ase.neighborlist @@ -228,8 +231,6 @@ def eval( The output of the evaluation. The keys are the names of the output variables, and the values are the corresponding output arrays. """ - if fparam is not None or aparam is not None: - raise NotImplementedError # convert all of the input to numpy array atom_types = np.array(atom_types, dtype=np.int32) coords = np.array(coords) @@ -240,7 +241,12 @@ def eval( ) request_defs = self._get_request_defs(atomic) out = self._eval_func(self._eval_model, numb_test, natoms)( - coords, cells, atom_types, request_defs + coords, + cells, + atom_types, + fparam, + aparam, + request_defs, ) return dict( zip( @@ -330,6 +336,8 @@ def _eval_model( coords: np.ndarray, cells: Optional[np.ndarray], atom_types: np.ndarray, + fparam: Optional[np.ndarray], + aparam: Optional[np.ndarray], request_defs: List[OutputVariableDef], ): model = self.dp.to(DEVICE) @@ -355,12 +363,26 @@ def _eval_model( ) else: box_input = None - + if fparam is not None: + fparam_input = to_torch_tensor(fparam.reshape(-1, self.get_dim_fparam())) + else: + fparam_input = None + if aparam is not None: + aparam_input = to_torch_tensor( + aparam.reshape(-1, natoms, self.get_dim_aparam()) + ) + else: + aparam_input = None do_atomic_virial = any( x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs ) batch_output = model( - coord_input, type_input, box=box_input, do_atomic_virial=do_atomic_virial + coord_input, + type_input, + box=box_input, + do_atomic_virial=do_atomic_virial, + fparam=fparam_input, + aparam=aparam_input, ) if isinstance(batch_output, tuple): batch_output = batch_output[0] diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 881ea4c97d..d2c1743d30 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -24,6 +24,9 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from .base_atomic_model import ( BaseAtomicModel, @@ -95,6 +98,7 @@ def serialize(self) -> dict: return { "@class": "Model", "type": "standard", + "@version": 1, "type_map": self.type_map, "descriptor": self.descriptor.serialize(), "fitting": self.fitting_net.serialize(), @@ -103,6 +107,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "DPAtomicModel": data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) descriptor_obj = BaseDescriptor.deserialize(data["descriptor"]) fitting_obj = BaseFitting.deserialize(data["fitting"]) obj = cls(descriptor_obj, fitting_obj, type_map=data["type_map"]) diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 68ff303d64..52f5f1d13c 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import copy import sys from abc import ( abstractmethod, @@ -24,6 +25,9 @@ get_multiple_nlist_key, nlist_distinguish_types, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from .base_atomic_model import ( BaseAtomicModel, @@ -206,6 +210,7 @@ def fitting_output_def(self) -> FittingOutputDef: def serialize(models) -> dict: return { "@class": "Model", + "@version": 1, "type": "linear", "models": [model.serialize() for model in models], "model_name": [model.__class__.__name__ for model in models], @@ -213,6 +218,8 @@ def serialize(models) -> dict: @staticmethod def deserialize(data) -> List[BaseAtomicModel]: + data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) model_names = data["model_name"] models = [ getattr(sys.modules[__name__], name).deserialize(model) @@ -303,6 +310,7 @@ def serialize(self) -> dict: return { "@class": "Model", "type": "zbl", + "@version": 1, "models": LinearAtomicModel.serialize([self.dp_model, self.zbl_model]), "sw_rmin": self.sw_rmin, "sw_rmax": self.sw_rmax, @@ -311,6 +319,8 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "DPZBLLinearAtomicModel": + data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) sw_rmin = data["sw_rmin"] sw_rmax = data["sw_rmax"] smin_alpha = data["smin_alpha"] diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index 86bfe98c36..c0b7c65d7a 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import copy from typing import ( Dict, List, @@ -15,6 +16,9 @@ from deepmd.utils.pair_tab import ( PairTab, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from .base_atomic_model import ( BaseAtomicModel, @@ -124,6 +128,7 @@ def serialize(self) -> dict: return { "@class": "Model", "type": "pairtab", + "@version": 1, "tab": self.tab.serialize(), "rcut": self.rcut, "sel": self.sel, @@ -131,6 +136,8 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "PairTabAtomicModel": + data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) rcut = data["rcut"] sel = data["sel"] tab = PairTab.deserialize(data["tab"]) diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index e693116cf4..b1df56a004 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -304,7 +304,7 @@ def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None) } for item in merged ] - descrpt.compute_input_stats(merged_tmp) + descrpt.compute_input_stats(merged_tmp, path) def serialize(self) -> dict: """Serialize the obj to dict.""" diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 033d640ad8..6c29636d6d 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -31,6 +31,9 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.version import ( + check_version_compatibility, +) try: from typing import ( @@ -182,6 +185,7 @@ def serialize(self) -> dict: return { "@class": "Descriptor", "type": "se_e2_a", + "@version": 1, "rcut": obj.rcut, "rcut_smth": obj.rcut_smth, "sel": obj.sel, @@ -208,6 +212,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "DescrptSeA": data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class", None) data.pop("type", None) variables = data.pop("@variables") diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index c685640426..bdb7dafe73 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -36,6 +36,9 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from .base_descriptor import ( BaseDescriptor, @@ -277,6 +280,7 @@ def serialize(self) -> dict: return { "@class": "Descriptor", "type": "se_r", + "@version": 1, "rcut": self.rcut, "rcut_smth": self.rcut_smth, "sel": self.sel, @@ -302,6 +306,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "DescrptSeR": data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) variables = data.pop("@variables") embeddings = data.pop("embeddings") env_mat = data.pop("env_mat") diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 0dc9ae20af..b823a051f5 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -20,7 +20,7 @@ BaseDescriptor, ) from deepmd.pt.model.task import ( - Fitting, + BaseFitting, ) from .dp_model import ( @@ -61,7 +61,7 @@ def get_zbl_model(model_params): fitting_net["out_dim"] = descriptor.get_dim_emb() if "ener" in fitting_net["type"]: fitting_net["return_energy"] = True - fitting = Fitting(**fitting_net) + fitting = BaseFitting(**fitting_net) dp_model = DPAtomicModel(descriptor, fitting, type_map=model_params["type_map"]) # pairtab filepath = model_params["use_srtab"] @@ -97,9 +97,8 @@ def get_model(model_params): fitting_net["out_dim"] = descriptor.get_dim_emb() if "ener" in fitting_net["type"]: fitting_net["return_energy"] = True - fitting = Fitting(**fitting_net) - - model = EnergyModel(descriptor, fitting, type_map=model_params["type_map"]) + fitting = BaseFitting(**fitting_net) + model = DPModel(descriptor, fitting, type_map=model_params["type_map"]) model.model_def_script = json.dumps(model_params) return model diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index 5410f518d1..79c129334a 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -10,6 +10,7 @@ ) from deepmd.pt.model.task.ener import ( EnergyFittingNet, + EnergyFittingNetDirect, ) from deepmd.pt.model.task.polarizability import ( PolarFittingNet, @@ -36,7 +37,9 @@ def __new__(cls, descriptor, fitting, *args, **kwargs): # according to the fitting network to decide the type of the model if cls is DPModel: # map fitting to model - if isinstance(fitting, EnergyFittingNet): + if isinstance(fitting, EnergyFittingNet) or isinstance( + fitting, EnergyFittingNetDirect + ): cls = EnergyModel elif isinstance(fitting, DipoleFittingNet): cls = DipoleModel diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py index 0f5e27aea9..e32d2f307d 100644 --- a/deepmd/pt/model/model/model.py +++ b/deepmd/pt/model/model/model.py @@ -59,9 +59,9 @@ # in DPAtomicModel (and other classes), but this requires the developer aware # of it when developing it... class BaseModel(make_base_model()): - def __init__(self): + def __init__(self, *args, **kwargs): """Construct a basic model for different tasks.""" - super().__init__() + super().__init__(*args, **kwargs) def compute_or_load_stat( self, diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index bff3dd93bc..9df3a5fb32 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -132,14 +132,6 @@ def output_def(self) -> FittingOutputDef: ] ) - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the fitting. - Return a list of statistic names needed, such as "bias_atom_e". - """ - return [] - def forward( self, descriptor: torch.Tensor, diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 8479111819..ff7ae6f8ec 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -28,8 +28,11 @@ from deepmd.pt.utils.env import ( DEFAULT_PRECISION, ) -from deepmd.pt.utils.stat import ( - compute_output_bias, +from deepmd.pt.utils.utils import ( + to_numpy_array, +) +from deepmd.utils.out_stat import ( + compute_stats_from_redu, ) from deepmd.utils.path import ( DPPath, @@ -135,16 +138,8 @@ def serialize(self) -> dict: data["atom_ener"] = self.atom_ener return data - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the fitting. - Return a list of statistic names needed, such as "bias_atom_e". - """ - return ["bias_atom_e"] - def compute_output_stats(self, merged, stat_file_path: Optional[DPPath] = None): - energy = [item["energy"] for item in merged] + energy = [item[self.var_name] for item in merged] data_mixed_type = "real_natoms_vec" in merged[0] if data_mixed_type: input_natoms = [item["real_natoms_vec"] for item in merged] @@ -155,7 +150,22 @@ def compute_output_stats(self, merged, stat_file_path: Optional[DPPath] = None): if stat_file_path is not None and stat_file_path.is_file(): bias_atom_e = stat_file_path.load_numpy() else: - bias_atom_e = compute_output_bias(energy, input_natoms, rcond=self.rcond) + # shape: (nframes, ndim) + merged_energy = to_numpy_array(torch.cat(energy)) + # shape: (nframes, ntypes) + merged_natoms = to_numpy_array(torch.cat(input_natoms)[:, 2:]) + if self.atom_ener is not None and len(self.atom_ener) > 0: + assigned_atom_ener = np.array( + [ee if ee is not None else np.nan for ee in self.atom_ener] + ) + else: + assigned_atom_ener = None + bias_atom_e, _ = compute_stats_from_redu( + merged_energy, + merged_natoms, + assigned_bias=assigned_atom_ener, + rcond=self.rcond, + ) if stat_file_path is not None: stat_file_path.save_numpy(bias_atom_e) assert all(x is not None for x in [bias_atom_e]) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 6c395d3800..20876d9be7 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -43,6 +43,9 @@ to_numpy_array, to_torch_tensor, ) +from deepmd.utils.version import ( + check_version_compatibility, +) dtype = env.GLOBAL_PT_FLOAT_PRECISION device = env.DEVICE @@ -89,14 +92,6 @@ def share_params(self, base_class, shared_level, resume=False): else: raise NotImplementedError - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the fitting. - Return a list of statistic names needed, such as "bias_atom_e". - """ - raise NotImplementedError("data_stat_key is not implemented!") - def change_energy_bias( self, config, model, old_type_map, new_type_map, bias_shift="delta", ntest=10 ): @@ -367,6 +362,7 @@ def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", + "@version": 1, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -404,6 +400,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) variables = data.pop("@variables") nets = data.pop("nets") obj = cls(**data) diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index 13b0d56e31..1bc4798c48 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -160,14 +160,6 @@ def output_def(self) -> FittingOutputDef: ] ) - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the fitting. - Return a list of statistic names needed, such as "bias_atom_e". - """ - return [] - def forward( self, descriptor: torch.Tensor, diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 5a783e412b..152c69a444 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -75,6 +75,7 @@ def __init__( finetune_model=None, force_load=False, shared_links=None, + init_frz_model=None, ): """Construct a DeePMD trainer. @@ -271,7 +272,7 @@ def get_loss(loss_params, start_lr, _ntypes): self.warmup_steps = training_params.get("warmup_steps", 0) self.gradient_max_norm = training_params.get("gradient_max_norm", 0.0) assert ( - self.num_steps - self.warmup_steps > 0 + self.num_steps - self.warmup_steps > 0 or self.warmup_steps == 0 ), "Warm up steps must be less than total training steps!" if self.multi_task and config.get("learning_rate_dict", None) is not None: self.lr_exp = {} @@ -394,6 +395,9 @@ def get_loss(loss_params, start_lr, _ntypes): ntest=ntest, bias_shift=model_params.get("bias_shift", "delta"), ) + if init_frz_model is not None: + frz_model = torch.jit.load(init_frz_model, map_location=DEVICE) + self.model.load_state_dict(frz_model.state_dict()) # Set trainable params self.wrapper.set_trainable_params() @@ -724,6 +728,15 @@ def log_loss_valid(_task_key="Default"): if ( self.rank == 0 or dist.get_rank() == 0 ): # Handle the case if rank 0 aborted and re-assigned + if self.num_steps == 0: + # when num_steps is 0, the checkpoint is never not saved + self.latest_model = Path(self.save_ckpt + "-0.pt") + self.save_model(self.latest_model, lr=0, step=0) + log.info(f"Saved model to {self.latest_model}") + symlink_prefix_files(self.latest_model.stem, self.save_ckpt) + with open("checkpoint", "w") as f: + f.write(str(self.latest_model)) + if JIT: pth_model_path = ( "frozen_model.pth" # We use .pth to denote the frozen model @@ -759,9 +772,10 @@ def get_data(self, is_train=True, task_key="Default"): batch_data = next(iter(self.training_data)) except StopIteration: # Refresh the status of the dataloader to start from a new epoch - self.training_data = BufferedIterator( - iter(self.training_dataloader) - ) + with torch.device("cpu"): + self.training_data = BufferedIterator( + iter(self.training_dataloader) + ) batch_data = next(iter(self.training_data)) else: try: diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index 2207f111a0..74b4a83ce7 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -164,6 +164,8 @@ def forward( task_key: Optional[torch.Tensor] = None, inference_only=False, do_atomic_virial=False, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, ): if not self.multi_task: task_key = "Default" @@ -172,7 +174,12 @@ def forward( task_key is not None ), f"Multitask model must specify the inference task! Supported tasks are {list(self.model.keys())}." model_pred = self.model[task_key]( - coord, atype, box=box, do_atomic_virial=do_atomic_virial + coord, + atype, + box=box, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, ) natoms = atype.shape[-1] if not self.inference_only and not inference_only: diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py index 3af03bda97..cd2943e6a8 100644 --- a/deepmd/pt/utils/env_mat_stat.py +++ b/deepmd/pt/utils/env_mat_stat.py @@ -80,7 +80,7 @@ def iter( Parameters ---------- data : List[Dict[str, torch.Tensor]] - The environment matrix. + The data. Yields ------ @@ -101,6 +101,14 @@ def iter( dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE, ) + if self.last_dim == 4: + radial_only = False + elif self.last_dim == 1: + radial_only = True + else: + raise ValueError( + "last_dim should be 1 for raial-only or 4 for full descriptor." + ) for system in data: coord, atype, box, natoms = ( system["coord"], @@ -130,6 +138,7 @@ def iter( self.descriptor.get_rcut(), # TODO: export rcut_smth from DescriptorBlock self.descriptor.rcut_smth, + radial_only, ) # reshape to nframes * nloc at the atom level, # so nframes/mixed_type do not matter diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 38f71d6994..4c769f019e 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -import numpy as np import torch log = logging.getLogger(__name__) @@ -57,23 +56,3 @@ def make_stat_input(datasets, dataloaders, nbatches): sys_stat[key] = sys_stat_list lst.append(sys_stat) return lst - - -def compute_output_bias(energy, natoms, rcond=None): - """Update output bias for fitting net. - - Args: - - energy: Batched energy with shape [nframes, 1]. - - natoms: Batched atom statisics with shape [self.ntypes+2]. - - Returns - ------- - - energy_coef: Average enery per atom for each element. - """ - for i in range(len(energy)): - energy[i] = energy[i].mean(dim=0, keepdim=True) - natoms[i] = natoms[i].double().mean(dim=0, keepdim=True) - sys_ener = torch.cat(energy).cpu() - sys_tynatom = torch.cat(natoms)[:, 2:].cpu() - energy_coef, _, _, _ = np.linalg.lstsq(sys_tynatom, sys_ener, rcond) - return energy_coef diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py index e1b7258c63..0e15ba13a8 100644 --- a/deepmd/tf/descriptor/se_a.py +++ b/deepmd/tf/descriptor/se_a.py @@ -65,6 +65,9 @@ from deepmd.tf.utils.type_embed import ( embed_atom_type, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from .descriptor import ( Descriptor, @@ -1368,6 +1371,7 @@ def deserialize(cls, data: dict, suffix: str = ""): if cls is not DescrptSeA: raise NotImplementedError("Not implemented in class %s" % cls.__name__) data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class", None) data.pop("type", None) embedding_net_variables = cls.deserialize_network( @@ -1422,6 +1426,7 @@ def serialize(self, suffix: str = "") -> dict: return { "@class": "Descriptor", "type": "se_e2_a", + "@version": 1, "rcut": self.rcut_r, "rcut_smth": self.rcut_r_smth, "sel": self.sel_a, diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index 1a12befdf0..ba1a261390 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -38,6 +38,9 @@ from deepmd.tf.utils.tabulate import ( DPTabulate, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from .descriptor import ( Descriptor, @@ -720,6 +723,7 @@ def deserialize(cls, data: dict, suffix: str = ""): if cls is not DescrptSeR: raise NotImplementedError("Not implemented in class %s" % cls.__name__) data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) embedding_net_variables = cls.deserialize_network( data.pop("embeddings"), suffix=suffix ) @@ -763,6 +767,7 @@ def serialize(self, suffix: str = "") -> dict: return { "@class": "Descriptor", "type": "se_r", + "@version": 1, "rcut": self.rcut, "rcut_smth": self.rcut_smth, "sel": self.sel_r, diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 3557d00aa0..f503789308 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -30,6 +30,9 @@ one_layer, one_layer_rand_seed_shift, ) +from deepmd.utils.version import ( + check_version_compatibility, +) @Fitting.register("dipole") @@ -346,6 +349,7 @@ def serialize(self, suffix: str) -> dict: data = { "@class": "Fitting", "type": "dipole", + "@version": 1, "var_name": "dipole", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -388,6 +392,8 @@ def deserialize(cls, data: dict, suffix: str): Model The deserialized model """ + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( data["nets"], diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index e8681f47ea..0cc5a7df62 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -43,6 +43,9 @@ from deepmd.tf.utils.network import ( one_layer_rand_seed_shift, ) +from deepmd.utils.out_stat import ( + compute_stats_from_redu, +) log = logging.getLogger(__name__) @@ -225,8 +228,10 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False): sys_tynatom = np.reshape(sys_tynatom, [nsys, -1]) sys_tynatom = sys_tynatom[:, 2:] - dos_shift, resd, rank, s_value = np.linalg.lstsq( - sys_tynatom, sys_dos, rcond=rcond + dos_shift, _ = compute_stats_from_redu( + sys_dos, + sys_tynatom, + rcond=rcond, ) return dos_shift diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 0cdd1a1676..a842df50bd 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -53,6 +53,12 @@ from deepmd.tf.utils.spin import ( Spin, ) +from deepmd.utils.out_stat import ( + compute_stats_from_redu, +) +from deepmd.utils.version import ( + check_version_compatibility, +) if TYPE_CHECKING: pass @@ -292,21 +298,17 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False): # In this situation, we directly use these assigned energies instead of computing stats. # This will make the loss decrease quickly assigned_atom_ener = np.array( - [ee for ee in self.atom_ener_v if ee is not None] + [ee if ee is not None else np.nan for ee in self.atom_ener_v] ) - assigned_ener_idx = [ - ii for ii, ee in enumerate(self.atom_ener_v) if ee is not None - ] - # np.dot out size: nframe - sys_ener -= np.dot(sys_tynatom[:, assigned_ener_idx], assigned_atom_ener) - sys_tynatom[:, assigned_ener_idx] = 0.0 - energy_shift, resd, rank, s_value = np.linalg.lstsq( - sys_tynatom, sys_ener, rcond=rcond + else: + assigned_atom_ener = None + energy_shift, _ = compute_stats_from_redu( + sys_ener.reshape(-1, 1), + sys_tynatom, + assigned_bias=assigned_atom_ener, + rcond=rcond, ) - if len(self.atom_ener) > 0: - for ii in assigned_ener_idx: - energy_shift[ii] = self.atom_ener_v[ii] - return energy_shift + return energy_shift.ravel() def compute_input_stats(self, all_stat: dict, protection: float = 1e-2) -> None: """Compute the input statistics. @@ -959,6 +961,8 @@ def deserialize(cls, data: dict, suffix: str = ""): Model The deserialized model """ + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( data["nets"], @@ -984,6 +988,7 @@ def serialize(self, suffix: str = "") -> dict: data = { "@class": "Fitting", "type": "ener", + "@version": 1, "var_name": "energy", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index f5cebf9a39..7ac31809f3 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -34,6 +34,9 @@ one_layer, one_layer_rand_seed_shift, ) +from deepmd.utils.version import ( + check_version_compatibility, +) @Fitting.register("polar") @@ -148,16 +151,14 @@ def get_out_size(self) -> int: """Get the output size. Should be 9.""" return 9 - def compute_input_stats(self, all_stat, protection=1e-2): - """Compute the input statistics. + def compute_output_stats(self, all_stat): + """Compute the output statistics. Parameters ---------- all_stat Dictionary of inputs. can be prepared by model.make_stat_input - protection - Divided-by-zero protection """ if "polarizability" not in all_stat.keys(): self.avgeig = np.zeros([9]) @@ -536,6 +537,7 @@ def serialize(self, suffix: str) -> dict: data = { "@class": "Fitting", "type": "polar", + "@version": 1, "var_name": "polar", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -581,6 +583,8 @@ def deserialize(cls, data: dict, suffix: str): Model The deserialized model """ + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( data["nets"], diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 2ae2879226..889f7ccc4d 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -63,6 +63,9 @@ from deepmd.utils.plugin import ( make_plugin_registry, ) +from deepmd.utils.version import ( + check_version_compatibility, +) class Model(ABC, make_plugin_registry("model")): @@ -778,7 +781,7 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": The deserialized descriptor """ data = copy.deepcopy(data) - + check_version_compatibility(data.pop("@version", 1), 1, 1) descriptor = Descriptor.deserialize(data.pop("descriptor"), suffix=suffix) fitting = Fitting.deserialize(data.pop("fitting"), suffix=suffix) return cls( @@ -807,6 +810,7 @@ def serialize(self, suffix: str = "") -> dict: return { "@class": "Model", "type": "standard", + "@version": 1, "type_map": self.type_map, "descriptor": self.descrpt.serialize(suffix=suffix), "fitting": self.fitting.serialize(suffix=suffix), diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 20111558cf..592b1f9748 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -22,6 +22,9 @@ from deepmd.utils.data import ( DeepmdData, ) +from deepmd.utils.out_stat import ( + compute_stats_from_redu, +) log = logging.getLogger(__name__) @@ -248,10 +251,12 @@ def compute_energy_shift(self, rcond=None, key="energy"): sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION) sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1]) sys_tynatom = sys_tynatom[:, 2:] - energy_shift, resd, rank, s_value = np.linalg.lstsq( - sys_tynatom, sys_ener, rcond=rcond + energy_shift, _ = compute_stats_from_redu( + sys_ener.reshape(-1, 1), + sys_tynatom, + rcond=rcond, ) - return energy_shift + return energy_shift.ravel() def add_dict(self, adict: dict) -> None: """Add items to the data system by a `dict`. diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py new file mode 100644 index 0000000000..8f68e32417 --- /dev/null +++ b/deepmd/utils/out_stat.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Output statistics.""" +from typing import ( + Optional, + Tuple, +) + +import numpy as np + + +def compute_stats_from_redu( + output_redu: np.ndarray, + natoms: np.ndarray, + assigned_bias: Optional[np.ndarray] = None, + rcond: Optional[float] = None, +) -> Tuple[np.ndarray, np.ndarray]: + """Compute the output statistics. + + Given the reduced output value and the number of atoms for each atom, + compute the least-squares solution as the atomic output bais and std. + + Parameters + ---------- + output_redu + The reduced output value, shape is [nframes, ndim]. + natoms + The number of atoms for each atom, shape is [nframes, ntypes]. + assigned_bias + The assigned output bias, shape is [ntypes, ndim]. Set to nan + if not assigned. + rcond + Cut-off ratio for small singular values of a. + + Returns + ------- + np.ndarray + The computed output bias, shape is [ntypes, ndim]. + np.ndarray + The computed output std, shape is [ntypes, ndim]. + """ + output_redu = np.array(output_redu) + natoms = np.array(natoms) + # check shape + assert output_redu.ndim == 2 + assert natoms.ndim == 2 + assert output_redu.shape[0] == natoms.shape[0] # nframes + if assigned_bias is not None: + assigned_bias = np.array(assigned_bias).reshape( + natoms.shape[1], output_redu.shape[1] + ) + # compute output bias + if assigned_bias is not None: + # Atomic energies stats are incorrect if atomic energies are assigned. + # In this situation, we directly use these assigned energies instead of computing stats. + # This will make the loss decrease quickly + assigned_bias_atom_mask = ~np.isnan(assigned_bias).any(axis=1) + # assigned_bias_masked: nmask, ndim + assigned_bias_masked = assigned_bias[assigned_bias_atom_mask] + # assigned_bias_natoms: nframes, nmask + assigned_bias_natoms = natoms[:, assigned_bias_atom_mask] + # output_redu: nframes, ndim + output_redu -= np.einsum( + "ij,jk->ik", assigned_bias_natoms, assigned_bias_masked + ) + # remove assigned atom + natoms[:, assigned_bias_atom_mask] = 0 + + # computed_output_bias: ntypes, ndim + computed_output_bias, _, _, _ = np.linalg.lstsq(natoms, output_redu, rcond=rcond) + if assigned_bias is not None: + # add back assigned atom; this might not be required + computed_output_bias[assigned_bias_atom_mask] = assigned_bias_masked + # rest_redu: nframes, ndim + rest_redu = output_redu - np.einsum("ij,jk->ik", natoms, computed_output_bias) + output_std = rest_redu.std(axis=0) + return computed_output_bias, output_std + + +def compute_stats_from_atomic( + output: np.ndarray, + atype: np.ndarray, +) -> Tuple[np.ndarray, np.ndarray]: + """Compute the output statistics. + + Given the output value and the type of atoms, + compute the atomic output bais and std. + + Parameters + ---------- + output + The output value, shape is [nframes, nloc, ndim]. + atype + The type of atoms, shape is [nframes, nloc]. + + Returns + ------- + np.ndarray + The computed output bias, shape is [ntypes, ndim]. + np.ndarray + The computed output std, shape is [ntypes, ndim]. + """ + output = np.array(output) + atype = np.array(atype) + # check shape + assert output.ndim == 3 + assert atype.ndim == 2 + assert output.shape[:2] == atype.shape + # compute output bias + nframes, nloc, ndim = output.shape + ntypes = atype.max() + 1 + output_bias = np.zeros((ntypes, ndim)) + output_std = np.zeros((ntypes, ndim)) + for type_i in range(ntypes): + mask = atype == type_i + output_bias[type_i] = output[mask].mean(axis=0) + output_std[type_i] = output[mask].std(axis=0) + return output_bias, output_std diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index b807354171..1b397a3cfa 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -12,6 +12,10 @@ CubicSpline, ) +from deepmd.utils.version import ( + check_version_compatibility, +) + log = logging.getLogger(__name__) @@ -72,6 +76,8 @@ def reinit(self, filename: str, rcut: Optional[float] = None) -> None: def serialize(self) -> dict: return { + "@class": "PairTab", + "@version": 1, "rmin": self.rmin, "rmax": self.rmax, "hh": self.hh, @@ -87,6 +93,9 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "PairTab": + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class") variables = data.pop("@variables") tab = PairTab(None, None) tab.vdata = variables["vdata"] diff --git a/deepmd/utils/path.py b/deepmd/utils/path.py index c9a7cd8554..79361b6c23 100644 --- a/deepmd/utils/path.py +++ b/deepmd/utils/path.py @@ -355,6 +355,7 @@ def save_numpy(self, arr: np.ndarray) -> None: if self._name in self._keys: del self.root[self._name] self.root.create_dataset(self._name, data=arr) + self.root.flush() def glob(self, pattern: str) -> List["DPPath"]: """Search path using the glob pattern. diff --git a/deepmd/utils/version.py b/deepmd/utils/version.py new file mode 100644 index 0000000000..a0b479778d --- /dev/null +++ b/deepmd/utils/version.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +def check_version_compatibility( + current_version: int, + maximum_supported_version: int, + minimal_supported_version: int = 1, +): + """Check if the current version is compatible with the supported versions. + + Parameters + ---------- + current_version : int + The current version. + maximum_supported_version : int + The maximum supported version. + minimal_supported_version : int, optional + The minimal supported version. Default is 1. + + Raises + ------ + ValueError + If the current version is not compatible with the supported versions. + """ + if not minimal_supported_version <= current_version <= maximum_supported_version: + raise ValueError( + f"Current version {current_version} is not compatible with supported versions " + f"[{minimal_supported_version}, {maximum_supported_version}]." + ) diff --git a/examples/water/dpa2/input_torch.json b/examples/water/dpa2/input_torch.json index 9d783b35d5..108e75df62 100644 --- a/examples/water/dpa2/input_torch.json +++ b/examples/water/dpa2/input_torch.json @@ -1,18 +1,13 @@ { "_comment": "that's all", "model": { - "type_embedding": { - "neuron": [ - 8 - ], - "tebd_input_mode": "concat" - }, "type_map": [ "O", "H" ], "descriptor": { "type": "dpa2", + "tebd_dim": 8, "repinit_rcut": 9.0, "repinit_rcut_smth": 8.0, "repinit_nsel": 120, @@ -74,6 +69,7 @@ "_comment": " that's all" }, "training": { + "stat_file": "./dpa2", "training_data": { "systems": [ "../data/data_0", diff --git a/examples/water/se_atten/input_torch.json b/examples/water/se_atten/input_torch.json index 7da3d64164..bc948cc2a0 100644 --- a/examples/water/se_atten/input_torch.json +++ b/examples/water/se_atten/input_torch.json @@ -15,6 +15,7 @@ 50, 100 ], + "tebd_dim": 8, "axis_neuron": 16, "attn": 128, "attn_layer": 2, @@ -59,6 +60,7 @@ "_comment": " that's all" }, "training": { + "stat_file": "./dpa1", "training_data": { "systems": [ "../data/data_0", diff --git a/examples/water/se_e2_a/input_torch.json b/examples/water/se_e2_a/input_torch.json index 053a721a44..c686b49d45 100644 --- a/examples/water/se_e2_a/input_torch.json +++ b/examples/water/se_e2_a/input_torch.json @@ -51,6 +51,7 @@ "_comment": " that's all" }, "training": { + "stat_file": "./se_e2_a", "training_data": { "systems": [ "../data/data_0", diff --git a/source/api_cc/include/DeepPotPT.h b/source/api_cc/include/DeepPotPT.h new file mode 100644 index 0000000000..1b757069c3 --- /dev/null +++ b/source/api_cc/include/DeepPotPT.h @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#include + +#include "DeepPot.h" +#include "commonPT.h" + +namespace deepmd { +/** + * @brief PyTorch implementation for Deep Potential. + **/ +class DeepPotPT : public DeepPotBase { + public: + /** + * @brief DP constructor without initialization. + **/ + DeepPotPT(); + ~DeepPotPT(); + /** + * @brief DP constructor with initialization. + * @param[in] model The name of the frozen model file. + * @param[in] gpu_rank The GPU rank. Default is 0. + * @param[in] file_content The content of the model file. If it is not empty, + *DP will read from the string instead of the file. + **/ + DeepPotPT(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = ""); + /** + * @brief Initialize the DP. + * @param[in] model The name of the frozen model file. + * @param[in] gpu_rank The GPU rank. Default is 0. + * @param[in] file_content The content of the model file. If it is not empty, + *DP will read from the string instead of the file. + **/ + void init(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = ""); + + private: + /** + * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[out] atom_energy The atomic energy. + * @param[out] atom_virial The atomic virial. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam The frame parameter. The array can be of size : + * nframes x dim_fparam. + * dim_fparam. Then all frames are assumed to be provided with the same + *fparam. + * @param[in] aparam The atomic parameter The array can be of size : + * nframes x natoms x dim_aparam. + * natoms x dim_aparam. Then all frames are assumed to be provided with the + *same aparam. + **/ + template + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box); + // const std::vector& fparam = std::vector(), + // const std::vector& aparam = std::vector()); + /** + * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[out] atom_energy The atomic energy. + * @param[out] atom_virial The atomic virial. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] nghost The number of ghost atoms. + * @param[in] lmp_list The input neighbour list. + * @param[in] ago Update the internal neighbour list if ago is 0. + * @param[in] fparam The frame parameter. The array can be of size : + * nframes x dim_fparam. + * dim_fparam. Then all frames are assumed to be provided with the same + *fparam. + * @param[in] aparam The atomic parameter The array can be of size : + * nframes x natoms x dim_aparam. + * natoms x dim_aparam. Then all frames are assumed to be provided with the + *same aparam. + **/ + template + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + // const int nghost, + const InputNlist& lmp_list, + const int& ago); + // const std::vector& fparam = std::vector(), + // const std::vector& aparam = std::vector()); + /** + * @brief Evaluate the energy, force, and virial with the mixed type + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[in] nframes The number of frames. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The array should be of size nframes x + *natoms. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam The frame parameter. The array can be of size : + * nframes x dim_fparam. + * dim_fparam. Then all frames are assumed to be provided with the same + *fparam. + * @param[in] aparam The atomic parameter The array can be of size : + * nframes x natoms x dim_aparam. + * natoms x dim_aparam. Then all frames are assumed to be provided with the + *same aparam. + **/ + template + void compute_mixed_type( + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + /** + * @brief Evaluate the energy, force, and virial with the mixed type + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[out] atom_energy The atomic energy. + * @param[out] atom_virial The atomic virial. + * @param[in] nframes The number of frames. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The array should be of size nframes x + *natoms. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam The frame parameter. The array can be of size : + * nframes x dim_fparam. + * dim_fparam. Then all frames are assumed to be provided with the same + *fparam. + * @param[in] aparam The atomic parameter The array can be of size : + * nframes x natoms x dim_aparam. + * natoms x dim_aparam. Then all frames are assumed to be provided with the + *same aparam. + **/ + template + void compute_mixed_type( + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + + public: + /** + * @brief Get the cutoff radius. + * @return The cutoff radius. + **/ + double cutoff() const { + assert(inited); + return rcut; + }; + /** + * @brief Get the number of types. + * @return The number of types. + **/ + int numb_types() const { + assert(inited); + return ntypes; + }; + /** + * @brief Get the number of types with spin. + * @return The number of types with spin. + **/ + int numb_types_spin() const { + assert(inited); + return ntypes_spin; + }; + /** + * @brief Get the dimension of the frame parameter. + * @return The dimension of the frame parameter. + **/ + int dim_fparam() const { + assert(inited); + return dfparam; + }; + /** + * @brief Get the dimension of the atomic parameter. + * @return The dimension of the atomic parameter. + **/ + int dim_aparam() const { + assert(inited); + return daparam; + }; + /** + * @brief Get the type map (element name of the atom types) of this model. + * @param[out] type_map The type map of this model. + **/ + void get_type_map(std::string& type_map); + + /** + * @brief Get whether the atom dimension of aparam is nall instead of fparam. + * @param[out] aparam_nall whether the atom dimension of aparam is nall + *instead of fparam. + **/ + bool is_aparam_nall() const { + assert(inited); + return aparam_nall; + }; + + // forward to template class + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew_mixed_type( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew_mixed_type( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + + private: + int num_intra_nthreads, num_inter_nthreads; + bool inited; + int ntypes; + int ntypes_spin; + int dfparam; + int daparam; + bool aparam_nall; + // copy neighbor list info from host + torch::jit::script::Module module; + double rcut; + NeighborListDataPT nlist_data; + int max_num_neighbors; + int gpu_id; + bool gpu_enabled; + at::Tensor firstneigh_tensor; +}; + +} // namespace deepmd diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h new file mode 100644 index 0000000000..57ffd5b295 --- /dev/null +++ b/source/api_cc/include/commonPT.h @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#include + +#include +#include +#include +#include + +#include "neighbor_list.h" +namespace deepmd { +struct NeighborListDataPT { + /// Array stores the core region atom's index + std::vector ilist; + /// Array stores the core region atom's neighbor index + std::vector jlist; + /// Array stores the number of neighbors of core region atoms + std::vector numneigh; + /// Array stores the the location of the first neighbor of core region atoms + std::vector firstneigh; + + public: + void copy_from_nlist(const InputNlist& inlist, int& max_num_neighbors); +}; +} // namespace deepmd diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index c598549844..442e2d90cc 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -10,6 +10,9 @@ #ifdef BUILD_TENSORFLOW #include "DeepPotTF.h" #endif +#ifdef BUILD_PYTORCH +#include "DeepPotPT.h" +#endif #include "device.h" using namespace deepmd; @@ -34,8 +37,14 @@ void DeepPot::init(const std::string& model, << std::endl; return; } - // TODO: To implement detect_backend - DPBackend backend = deepmd::DPBackend::TensorFlow; + DPBackend backend; + if (model.length() >= 4 && model.substr(model.length() - 4) == ".pth") { + backend = deepmd::DPBackend::PyTorch; + } else if (model.length() >= 3 && model.substr(model.length() - 3) == ".pb") { + backend = deepmd::DPBackend::TensorFlow; + } else { + throw deepmd::deepmd_exception("Unsupported model file format"); + } if (deepmd::DPBackend::TensorFlow == backend) { #ifdef BUILD_TENSORFLOW dp = std::make_shared(model, gpu_rank, file_content); @@ -43,7 +52,11 @@ void DeepPot::init(const std::string& model, throw deepmd::deepmd_exception("TensorFlow backend is not built"); #endif } else if (deepmd::DPBackend::PyTorch == backend) { - throw deepmd::deepmd_exception("PyTorch backend is not supported yet"); +#ifdef BUILD_PYTORCH + dp = std::make_shared(model, gpu_rank, file_content); +#else + throw deepmd::deepmd_exception("PyTorch backend is not built"); +#endif } else if (deepmd::DPBackend::Paddle == backend) { throw deepmd::deepmd_exception("PaddlePaddle backend is not supported yet"); } else { diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index c94fb4247b..f05e27b9b2 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/source/api_cc/src/DeepPotPT.cc @@ -1,8 +1,358 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #ifdef BUILD_PYTORCH -#include +#include "DeepPotPT.h" -void test_function_please_remove_after_torch_is_actually_used() { - torch::Tensor tensor = torch::rand({2, 3}); +#include "common.h" +using namespace deepmd; +DeepPotPT::DeepPotPT() : inited(false) {} +DeepPotPT::DeepPotPT(const std::string& model, + const int& gpu_rank, + const std::string& file_content) + : inited(false) { + try { + init(model, gpu_rank, file_content); + } catch (...) { + // Clean up and rethrow, as the destructor will not be called + throw; + } +} +void DeepPotPT::init(const std::string& model, + const int& gpu_rank, + const std::string& file_content) { + if (inited) { + std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " + "nothing at the second call of initializer" + << std::endl; + return; + } + gpu_id = gpu_rank; + torch::Device device(torch::kCUDA, gpu_rank); + gpu_enabled = torch::cuda::is_available(); + if (!gpu_enabled) { + device = torch::Device(torch::kCPU); + std::cout << "load model from: " << model << " to cpu " << gpu_rank + << std::endl; + } else { + std::cout << "load model from: " << model << " to gpu " << gpu_rank + << std::endl; + } + module = torch::jit::load(model, device); + + torch::jit::FusionStrategy strategy; + strategy = {{torch::jit::FusionBehavior::DYNAMIC, 10}}; + torch::jit::setFusionStrategy(strategy); + + get_env_nthreads(num_intra_nthreads, + num_inter_nthreads); // need to be fixed as + // DP_INTRA_OP_PARALLELISM_THREADS + if (num_inter_nthreads) { + try { + at::set_num_interop_threads(num_inter_nthreads); + } catch (...) { + } + } + if (num_intra_nthreads) { + try { + at::set_num_threads(num_intra_nthreads); + } catch (...) { + } + } + + auto rcut_ = module.run_method("get_rcut").toDouble(); + rcut = static_cast(rcut_); + ntypes = 0; + ntypes_spin = 0; + dfparam = 0; + daparam = 0; + aparam_nall = false; + inited = true; +} +DeepPotPT::~DeepPotPT() {} + +template +void DeepPotPT::compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const InputNlist& lmp_list, + const int& ago) { + torch::Device device(torch::kCUDA, gpu_id); + if (!gpu_enabled) { + device = torch::Device(torch::kCPU); + } + std::vector coord_wrapped = coord; + int natoms = atype.size(); + auto options = torch::TensorOptions().dtype(torch::kFloat64); + torch::ScalarType floatType = torch::kFloat64; + if (std::is_same_v) { + options = torch::TensorOptions().dtype(torch::kFloat32); + floatType = torch::kFloat32; + } + auto int_options = torch::TensorOptions().dtype(torch::kInt64); + auto int32_options = torch::TensorOptions().dtype(torch::kInt32); + at::Tensor coord_wrapped_Tensor = + torch::from_blob(coord_wrapped.data(), {1, natoms, 3}, options) + .to(device); + std::vector atype_64(atype.begin(), atype.end()); + at::Tensor atype_Tensor = + torch::from_blob(atype_64.data(), {1, natoms}, int_options).to(device); + if (ago == 0) { + nlist_data.copy_from_nlist(lmp_list, max_num_neighbors); + } + at::Tensor firstneigh = + torch::from_blob(nlist_data.jlist.data(), + {1, lmp_list.inum, max_num_neighbors}, int32_options); + firstneigh_tensor = firstneigh.to(torch::kInt64).to(device); + bool do_atom_virial_tensor = true; + c10::optional optional_tensor; + c10::Dict outputs = + module + .run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor, + firstneigh_tensor, optional_tensor, optional_tensor, + optional_tensor, do_atom_virial_tensor) + .toGenericDict(); + c10::IValue energy_ = outputs.at("energy"); + c10::IValue force_ = outputs.at("extended_force"); + c10::IValue virial_ = outputs.at("virial"); + c10::IValue atom_virial_ = outputs.at("extended_virial"); + c10::IValue atom_energy_ = outputs.at("atom_energy"); + torch::Tensor flat_energy_ = energy_.toTensor().view({-1}); + torch::Tensor cpu_energy_ = flat_energy_.to(torch::kCPU); + ener.assign(cpu_energy_.data_ptr(), + cpu_energy_.data_ptr() + cpu_energy_.numel()); + torch::Tensor flat_atom_energy_ = + atom_energy_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_atom_energy_ = flat_atom_energy_.to(torch::kCPU); + atom_energy.resize(natoms, 0.0); // resize to nall to be consistenet with TF. + atom_energy.assign( + cpu_atom_energy_.data_ptr(), + cpu_atom_energy_.data_ptr() + cpu_atom_energy_.numel()); + torch::Tensor flat_force_ = force_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_force_ = flat_force_.to(torch::kCPU); + force.assign(cpu_force_.data_ptr(), + cpu_force_.data_ptr() + cpu_force_.numel()); + torch::Tensor flat_virial_ = virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); + virial.assign(cpu_virial_.data_ptr(), + cpu_virial_.data_ptr() + cpu_virial_.numel()); + torch::Tensor flat_atom_virial_ = + atom_virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU); + atom_virial.assign( + cpu_atom_virial_.data_ptr(), + cpu_atom_virial_.data_ptr() + cpu_atom_virial_.numel()); +} +template void DeepPotPT::compute>( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const InputNlist& lmp_list, + const int& ago); +template void DeepPotPT::compute>( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const InputNlist& lmp_list, + const int& ago); +template +void DeepPotPT::compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box) { + torch::Device device(torch::kCUDA, gpu_id); + if (!gpu_enabled) { + device = torch::Device(torch::kCPU); + } + std::vector coord_wrapped = coord; + int natoms = atype.size(); + auto options = torch::TensorOptions().dtype(torch::kFloat64); + torch::ScalarType floatType = torch::kFloat64; + if (std::is_same_v) { + options = torch::TensorOptions().dtype(torch::kFloat32); + floatType = torch::kFloat32; + } + auto int_options = torch::TensorOptions().dtype(torch::kInt64); + std::vector inputs; + at::Tensor coord_wrapped_Tensor = + torch::from_blob(coord_wrapped.data(), {1, natoms, 3}, options) + .to(device); + inputs.push_back(coord_wrapped_Tensor); + std::vector atype_64(atype.begin(), atype.end()); + at::Tensor atype_Tensor = + torch::from_blob(atype_64.data(), {1, natoms}, int_options).to(device); + inputs.push_back(atype_Tensor); + c10::optional box_Tensor; + if (!box.empty()) { + box_Tensor = + torch::from_blob(const_cast(box.data()), {1, 9}, options) + .to(device); + } + inputs.push_back(box_Tensor); + c10::optional fparam_tensor; + inputs.push_back(fparam_tensor); + c10::optional aparam_tensor; + inputs.push_back(aparam_tensor); + bool do_atom_virial_tensor = true; + inputs.push_back(do_atom_virial_tensor); + c10::Dict outputs = + module.forward(inputs).toGenericDict(); + c10::IValue energy_ = outputs.at("energy"); + c10::IValue force_ = outputs.at("force"); + c10::IValue virial_ = outputs.at("virial"); + c10::IValue atom_virial_ = outputs.at("atom_virial"); + c10::IValue atom_energy_ = outputs.at("atom_energy"); + torch::Tensor flat_energy_ = energy_.toTensor().view({-1}); + torch::Tensor cpu_energy_ = flat_energy_.to(torch::kCPU); + ener.assign(cpu_energy_.data_ptr(), + cpu_energy_.data_ptr() + cpu_energy_.numel()); + torch::Tensor flat_atom_energy_ = + atom_energy_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_atom_energy_ = flat_atom_energy_.to(torch::kCPU); + atom_energy.assign( + cpu_atom_energy_.data_ptr(), + cpu_atom_energy_.data_ptr() + cpu_atom_energy_.numel()); + torch::Tensor flat_force_ = force_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_force_ = flat_force_.to(torch::kCPU); + force.assign(cpu_force_.data_ptr(), + cpu_force_.data_ptr() + cpu_force_.numel()); + torch::Tensor flat_virial_ = virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); + virial.assign(cpu_virial_.data_ptr(), + cpu_virial_.data_ptr() + cpu_virial_.numel()); + torch::Tensor flat_atom_virial_ = + atom_virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU); + atom_virial.assign( + cpu_atom_virial_.data_ptr(), + cpu_atom_virial_.data_ptr() + cpu_atom_virial_.numel()); +} + +template void DeepPotPT::compute>( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box); +template void DeepPotPT::compute>( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box); +void DeepPotPT::get_type_map(std::string& type_map) { + auto ret = module.run_method("get_type_map").toList(); + for (const torch::IValue& element : ret) { + type_map += torch::str(element); // Convert each element to a string + type_map += " "; // Add a space between elements + } +} + +// forward to template method +void DeepPotPT::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box); +} +void DeepPotPT::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box); +} +void DeepPotPT::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + // TODO: atomic compute unsupported + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + inlist, ago); +} +void DeepPotPT::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + inlist, ago); +} +void DeepPotPT::computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + throw deepmd::deepmd_exception("computew_mixed_type is not implemented"); +} +void DeepPotPT::computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + throw deepmd::deepmd_exception("computew_mixed_type is not implemented"); } #endif diff --git a/source/api_cc/src/commonPT.cc b/source/api_cc/src/commonPT.cc new file mode 100644 index 0000000000..4ed3b21fe8 --- /dev/null +++ b/source/api_cc/src/commonPT.cc @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#ifdef BUILD_PYTORCH +#include "commonPT.h" +using namespace deepmd; +void NeighborListDataPT::copy_from_nlist(const InputNlist& inlist, + int& max_num_neighbors) { + int inum = inlist.inum; + ilist.resize(inum); + numneigh.resize(inum); + memcpy(&ilist[0], inlist.ilist, inum * sizeof(int)); + int* max_element = std::max_element(inlist.numneigh, inlist.numneigh + inum); + max_num_neighbors = *max_element; + unsigned long nlist_size = (unsigned long)inum * max_num_neighbors; + jlist.resize(nlist_size); + jlist.assign(nlist_size, -1); + for (int ii = 0; ii < inum; ++ii) { + int jnum = inlist.numneigh[ii]; + numneigh[ii] = inlist.numneigh[ii]; + memcpy(&jlist[(unsigned long)ii * max_num_neighbors], inlist.firstneigh[ii], + jnum * sizeof(int)); + } +} +#endif diff --git a/source/api_cc/tests/test_deeppot_pt.cc b/source/api_cc/tests/test_deeppot_pt.cc new file mode 100644 index 0000000000..e0e90ac75c --- /dev/null +++ b/source/api_cc/tests/test_deeppot_pt.cc @@ -0,0 +1,625 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "DeepPot.h" +#include "neighbor_list.h" +#include "test_utils.h" + +template +class TestInferDeepPotAPt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + std::vector expected_e = { + + -93.016873944029, -185.923296645958, -185.927096544970, + -93.019371018039, -185.926179995548, -185.924351901852}; + std::vector expected_f = { + + 0.006277522211, -0.001117962774, 0.000618580445, 0.009928999655, + 0.003026035654, -0.006941982227, 0.000667853212, -0.002449963843, + 0.006506463508, -0.007284129115, 0.000530662205, -0.000028806821, + 0.000068097781, 0.006121331983, -0.009019754602, -0.009658343745, + -0.006110103225, 0.008865499697}; + std::vector expected_v = { + -0.000155238009, 0.000116605516, -0.007869862476, 0.000465578340, + 0.008182547185, -0.002398713212, -0.008112887338, -0.002423738425, + 0.007210716605, -0.019203504012, 0.001724938709, 0.009909211091, + 0.001153857542, -0.001600015103, -0.000560024090, 0.010727836276, + -0.001034836404, -0.007973454377, -0.021517399106, -0.004064359664, + 0.004866398692, -0.003360038617, -0.007241406162, 0.005920941051, + 0.004899151657, 0.006290788591, -0.006478820311, 0.001921504710, + 0.001313470921, -0.000304091236, 0.001684345981, 0.004124109256, + -0.006396084465, -0.000701095618, -0.006356507032, 0.009818550859, + -0.015230664587, -0.000110244376, 0.000690319396, 0.000045953023, + -0.005726548770, 0.008769818495, -0.000572380210, 0.008860603423, + -0.013819348050, -0.021227082558, -0.004977781343, 0.006646239696, + -0.005987066507, -0.002767831232, 0.003746502525, 0.007697590397, + 0.003746130152, -0.005172634748}; + int natoms; + double expected_tot_e; + std::vector expected_tot_v; + + deepmd::DeepPot dp; + + void SetUp() override { + std::string file_name = "../../tests/infer/deeppot_sea.pth"; + + dp.init(file_name); + + natoms = expected_e.size(); + EXPECT_EQ(natoms * 3, expected_f.size()); + EXPECT_EQ(natoms * 9, expected_v.size()); + expected_tot_e = 0.; + expected_tot_v.resize(9); + std::fill(expected_tot_v.begin(), expected_tot_v.end(), 0.); + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e += expected_e[ii]; + } + for (int ii = 0; ii < natoms; ++ii) { + for (int dd = 0; dd < 9; ++dd) { + expected_tot_v[dd] += expected_v[ii * 9 + dd]; + } + } + }; + + void TearDown() override { remove("deeppot.pb"); }; +}; + +TYPED_TEST_SUITE(TestInferDeepPotAPt, ValueTypes); + +TYPED_TEST(TestInferDeepPotAPt, cpu_build_nlist) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotAPt, cpu_build_nlist_numfv) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + class MyModel : public EnergyModelTest { + deepmd::DeepPot& mydp; + const std::vector atype; + + public: + MyModel(deepmd::DeepPot& dp_, const std::vector& atype_) + : mydp(dp_), atype(atype_){}; + virtual void compute(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& box) { + mydp.compute(ener, force, virial, coord, atype, box); + } + }; + MyModel model(dp, atype); + model.test_f(coord, box); + model.test_v(coord, box); + std::vector box_(box); + box_[1] -= 0.4; + model.test_f(coord, box_); + model.test_v(coord, box_); + box_[2] += 0.5; + model.test_f(coord, box_); + model.test_v(coord, box_); + box_[4] += 0.2; + model.test_f(coord, box_); + model.test_v(coord, box_); + box_[3] -= 0.3; + model.test_f(coord, box_); + model.test_v(coord, box_); + box_[6] -= 0.7; + model.test_f(coord, box_); + model.test_v(coord, box_); + box_[7] += 0.6; + model.test_f(coord, box_); + model.test_v(coord, box_); +} + +TYPED_TEST(TestInferDeepPotAPt, cpu_build_nlist_atomic) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + double ener; + std::vector force, virial, atom_ener, atom_vir; + dp.compute(ener, force, virial, atom_ener, atom_vir, coord, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + EXPECT_EQ(atom_ener.size(), natoms); + EXPECT_EQ(atom_vir.size(), natoms * 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); + } + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotAPt, cpu_lmp_nlist) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + float rc = dp.cutoff(); + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector > nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc); + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + double ener; + std::vector force_, virial; + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0); + std::vector force; + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + + ener = 0.; + std::fill(force_.begin(), force_.end(), 0.0); + std::fill(virial.begin(), virial.end(), 0.0); + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 1); + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotAPt, cpu_lmp_nlist_atomic) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + float rc = dp.cutoff(); + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector > nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc); + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + double ener; + std::vector force_, atom_ener_, atom_vir_, virial; + std::vector force, atom_ener, atom_vir; + dp.compute(ener, force_, virial, atom_ener_, atom_vir_, coord_cpy, atype_cpy, + box, nall - nloc, inlist, 0); + _fold_back(force, force_, mapping, nloc, nall, 3); + _fold_back(atom_ener, atom_ener_, mapping, nloc, nall, 1); + _fold_back(atom_vir, atom_vir_, mapping, nloc, nall, 9); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + EXPECT_EQ(atom_ener.size(), natoms); + EXPECT_EQ(atom_vir.size(), natoms * 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); + } + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); + } + + ener = 0.; + std::fill(force_.begin(), force_.end(), 0.0); + std::fill(virial.begin(), virial.end(), 0.0); + std::fill(atom_ener_.begin(), atom_ener_.end(), 0.0); + std::fill(atom_vir_.begin(), atom_vir_.end(), 0.0); + dp.compute(ener, force_, virial, atom_ener_, atom_vir_, coord_cpy, atype_cpy, + box, nall - nloc, inlist, 1); + _fold_back(force, force_, mapping, nloc, nall, 3); + _fold_back(atom_ener, atom_ener_, mapping, nloc, nall, 1); + _fold_back(atom_vir, atom_vir_, mapping, nloc, nall, 9); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + EXPECT_EQ(atom_ener.size(), natoms); + EXPECT_EQ(atom_vir.size(), natoms * 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); + } + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotAPt, cpu_lmp_nlist_2rc) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + float rc = dp.cutoff(); + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector > nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc * 2); + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + double ener; + std::vector force_(nall * 3, 0.0), virial(9, 0.0); + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0); + std::vector force; + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + + ener = 0.; + std::fill(force_.begin(), force_.end(), 0.0); + std::fill(virial.begin(), virial.end(), 0.0); + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 1); + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotAPt, cpu_lmp_nlist_type_sel) { + GTEST_SKIP() << "Skipping this test for unsupported"; + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + float rc = dp.cutoff(); + + // add vir atoms + int nvir = 2; + std::vector coord_vir(nvir * 3); + std::vector atype_vir(nvir, 2); + for (int ii = 0; ii < nvir; ++ii) { + coord_vir[ii] = coord[ii]; + } + coord.insert(coord.begin(), coord_vir.begin(), coord_vir.end()); + atype.insert(atype.begin(), atype_vir.begin(), atype_vir.end()); + natoms += nvir; + std::vector expected_f_vir(nvir * 3, 0.0); + expected_f.insert(expected_f.begin(), expected_f_vir.begin(), + expected_f_vir.end()); + + // build nlist + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector > nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc); + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + // dp compute + double ener; + std::vector force_(nall * 3, 0.0), virial(9, 0.0); + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0); + // fold back + std::vector force; + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotAPt, cpu_lmp_nlist_type_sel_atomic) { + GTEST_SKIP() << "Skipping this test for unsupported"; + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + float rc = dp.cutoff(); + + // add vir atoms + int nvir = 2; + std::vector coord_vir(nvir * 3); + std::vector atype_vir(nvir, 2); + for (int ii = 0; ii < nvir; ++ii) { + coord_vir[ii] = coord[ii]; + } + coord.insert(coord.begin(), coord_vir.begin(), coord_vir.end()); + atype.insert(atype.begin(), atype_vir.begin(), atype_vir.end()); + natoms += nvir; + std::vector expected_f_vir(nvir * 3, 0.0); + expected_f.insert(expected_f.begin(), expected_f_vir.begin(), + expected_f_vir.end()); + + // build nlist + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector > nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc); + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + // dp compute + double ener; + std::vector force_(nall * 3, 0.0), virial(9, 0.0), atomic_energy, + atomic_virial; + dp.compute(ener, force_, virial, atomic_energy, atomic_virial, coord_cpy, + atype_cpy, box, nall - nloc, inlist, 0); + // fold back + std::vector force; + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotAPt, print_summary) { + deepmd::DeepPot& dp = this->dp; + dp.print_summary(""); +} + +template +class TestInferDeepPotAPtNoPbc : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {}; + std::vector expected_e = {-93.003304908874, -185.915806542480, + -185.928116717624, -93.017934934346, + -185.924393412278, -185.923906740801}; + std::vector expected_f = { + 0.000868182637, -0.000363698132, -0.000657003077, -0.000868182637, + 0.000363698132, 0.000657003077, 0.007932614680, -0.001003609844, + 0.000737731722, -0.003883788858, 0.000686896282, -0.000578400682, + 0.004064895086, 0.006115547962, -0.008747097814, -0.008113720908, + -0.005798834400, 0.008587766774}; + std::vector expected_v = { + 0.007762485364, -0.003251851977, -0.005874313248, -0.003251851977, + 0.001362262315, 0.002460860955, -0.005874313248, 0.002460860955, + 0.004445426242, -0.007120030212, 0.002982715359, 0.005388130971, + 0.002982715359, -0.001249515894, -0.002257190002, 0.005388130971, + -0.002257190002, -0.004077504519, -0.015805863589, 0.001952684835, + -0.001522876482, 0.001796574704, -0.000358803950, 0.000369710813, + -0.001108943040, 0.000332585300, -0.000395481309, 0.008873525623, + 0.001919112114, -0.001486235522, 0.002002929532, 0.004222469272, + -0.006517211126, -0.001656192522, -0.006501210045, 0.010118622295, + -0.006548889778, -0.000465126991, 0.001002876603, 0.000240398734, + -0.005794489784, 0.008940685179, -0.000121727685, 0.008931999051, + -0.013852797563, -0.017962955675, -0.004645050453, 0.006214692837, + -0.005278283465, -0.002662692758, 0.003618275905, 0.007095320684, + 0.003648086464, -0.005023397513}; + int natoms; + double expected_tot_e; + std::vector expected_tot_v; + + deepmd::DeepPot dp; + + void SetUp() override { + std::string file_name = "../../tests/infer/deeppot_sea.pth"; + dp.init(file_name); + + natoms = expected_e.size(); + EXPECT_EQ(natoms * 3, expected_f.size()); + EXPECT_EQ(natoms * 9, expected_v.size()); + expected_tot_e = 0.; + expected_tot_v.resize(9); + std::fill(expected_tot_v.begin(), expected_tot_v.end(), 0.); + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e += expected_e[ii]; + } + for (int ii = 0; ii < natoms; ++ii) { + for (int dd = 0; dd < 9; ++dd) { + expected_tot_v[dd] += expected_v[ii * 9 + dd]; + } + } + }; + + void TearDown() override { remove("deeppot.pb"); }; +}; + +TYPED_TEST_SUITE(TestInferDeepPotAPtNoPbc, ValueTypes); + +TYPED_TEST(TestInferDeepPotAPtNoPbc, cpu_build_nlist) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 22d22a27f6..73aa74ed90 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -18,7 +18,15 @@ INSTALL_PREFIX=${SCRIPT_PATH}/../../dp_test BUILD_TMP_DIR=${SCRIPT_PATH}/../build_tests mkdir -p ${BUILD_TMP_DIR} cd ${BUILD_TMP_DIR} -cmake -DINSTALL_TENSORFLOW=FALSE -DUSE_TF_PYTHON_LIBS=TRUE -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} -DBUILD_TESTING:BOOL=TRUE -DLAMMPS_VERSION=stable_2Aug2023_update2 ${CUDA_ARGS} .. +cmake \ + -D ENABLE_TENSORFLOW=TRUE \ + -D ENABLE_PYTORCH=TRUE \ + -D INSTALL_TENSORFLOW=FALSE \ + -D USE_TF_PYTHON_LIBS=TRUE \ + -D CMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} \ + -D BUILD_TESTING:BOOL=TRUE \ + -D LAMMPS_VERSION=stable_2Aug2023_update2 \ + ${CUDA_ARGS} .. cmake --build . -j${NPROC} cmake --install . ctest --output-on-failure diff --git a/source/ipi/tests/test_driver.py b/source/ipi/tests/test_driver.py index 1b2e1dd951..b0fbf53b01 100644 --- a/source/ipi/tests/test_driver.py +++ b/source/ipi/tests/test_driver.py @@ -251,3 +251,108 @@ def test_normalize_coords(self): ) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1) np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) + + +class TestDPIPIPt(TestDPIPI): + @classmethod + def setUpClass(cls): + cls.model_file = str(tests_path / "infer" / "deeppot_sea.pth") + + def setUp(self): + super().setUp() + + self.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + self.expected_e = np.array( + [ + -93.016873944029, + -185.923296645958, + -185.927096544970, + -93.019371018039, + -185.926179995548, + -185.924351901852, + ] + ) + self.expected_f = np.array( + [ + 0.006277522211, + -0.001117962774, + 0.000618580445, + 0.009928999655, + 0.003026035654, + -0.006941982227, + 0.000667853212, + -0.002449963843, + 0.006506463508, + -0.007284129115, + 0.000530662205, + -0.000028806821, + 0.000068097781, + 0.006121331983, + -0.009019754602, + -0.009658343745, + -0.006110103225, + 0.008865499697, + ] + ) + self.expected_v = np.array( + [ + -0.000155238009, + 0.000116605516, + -0.007869862476, + 0.000465578340, + 0.008182547185, + -0.002398713212, + -0.008112887338, + -0.002423738425, + 0.007210716605, + -0.019203504012, + 0.001724938709, + 0.009909211091, + 0.001153857542, + -0.001600015103, + -0.000560024090, + 0.010727836276, + -0.001034836404, + -0.007973454377, + -0.021517399106, + -0.004064359664, + 0.004866398692, + -0.003360038617, + -0.007241406162, + 0.005920941051, + 0.004899151657, + 0.006290788591, + -0.006478820311, + 0.001921504710, + 0.001313470921, + -0.000304091236, + 0.001684345981, + 0.004124109256, + -0.006396084465, + -0.000701095618, + -0.006356507032, + 0.009818550859, + -0.015230664587, + -0.000110244376, + 0.000690319396, + 0.000045953023, + -0.005726548770, + 0.008769818495, + -0.000572380210, + 0.008860603423, + -0.013819348050, + -0.021227082558, + -0.004977781343, + 0.006646239696, + -0.005987066507, + -0.002767831232, + 0.003746502525, + 0.007697590397, + 0.003746130152, + -0.005172634748, + ] + ) + + @classmethod + def tearDownClass(cls): + cls.dp = None diff --git a/source/lmp/tests/test_lammps_pt.py b/source/lmp/tests/test_lammps_pt.py new file mode 100644 index 0000000000..bf1ef97e2b --- /dev/null +++ b/source/lmp/tests/test_lammps_pt.py @@ -0,0 +1,695 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import os +import subprocess as sp +import sys +from pathlib import ( + Path, +) + +import constants +import numpy as np +import pytest +from lammps import ( + PyLammps, +) +from write_lmp_data import ( + write_lmp_data, +) + +pbtxt_file2 = ( + Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot-1.pbtxt" +) +pb_file = Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_sea.pth" +pb_file2 = Path(__file__).parent / "graph2.pb" +system_file = Path(__file__).parent.parent.parent / "tests" +data_file = Path(__file__).parent / "data.lmp" +data_file_si = Path(__file__).parent / "data.si" +data_type_map_file = Path(__file__).parent / "data_type_map.lmp" +md_file = Path(__file__).parent / "md.out" + +# this is as the same as python and c++ tests, test_deeppot_a.py +expected_ae = np.array( + [ + -93.016873944029, + -185.923296645958, + -185.927096544970, + -93.019371018039, + -185.926179995548, + -185.924351901852, + ] +) +expected_e = np.sum(expected_ae) +expected_f = np.array( + [ + 0.006277522211, + -0.001117962774, + 0.000618580445, + 0.009928999655, + 0.003026035654, + -0.006941982227, + 0.000667853212, + -0.002449963843, + 0.006506463508, + -0.007284129115, + 0.000530662205, + -0.000028806821, + 0.000068097781, + 0.006121331983, + -0.009019754602, + -0.009658343745, + -0.006110103225, + 0.008865499697, + ] +).reshape(6, 3) + +expected_f2 = np.array( + [ + [-0.6454949, 1.72457783, 0.18897958], + [1.68936514, -0.36995299, -1.36044464], + [-1.09902692, -1.35487928, 1.17416702], + [1.68426111, -0.50835585, 0.98340415], + [0.05771758, 1.12515818, -1.77561531], + [-1.686822, -0.61654789, 0.78950921], + ] +) + +expected_v = -np.array( + [ + -0.000155238009, + 0.000116605516, + -0.007869862476, + 0.000465578340, + 0.008182547185, + -0.002398713212, + -0.008112887338, + -0.002423738425, + 0.007210716605, + -0.019203504012, + 0.001724938709, + 0.009909211091, + 0.001153857542, + -0.001600015103, + -0.000560024090, + 0.010727836276, + -0.001034836404, + -0.007973454377, + -0.021517399106, + -0.004064359664, + 0.004866398692, + -0.003360038617, + -0.007241406162, + 0.005920941051, + 0.004899151657, + 0.006290788591, + -0.006478820311, + 0.001921504710, + 0.001313470921, + -0.000304091236, + 0.001684345981, + 0.004124109256, + -0.006396084465, + -0.000701095618, + -0.006356507032, + 0.009818550859, + -0.015230664587, + -0.000110244376, + 0.000690319396, + 0.000045953023, + -0.005726548770, + 0.008769818495, + -0.000572380210, + 0.008860603423, + -0.013819348050, + -0.021227082558, + -0.004977781343, + 0.006646239696, + -0.005987066507, + -0.002767831232, + 0.003746502525, + 0.007697590397, + 0.003746130152, + -0.005172634748, + ] +).reshape(6, 9) +expected_v2 = -np.array( + [ + [ + -0.70008436, + -0.06399891, + 0.63678391, + -0.07642171, + -0.70580035, + 0.20506145, + 0.64098364, + 0.20305781, + -0.57906794, + ], + [ + -0.6372635, + 0.14315552, + 0.51952246, + 0.04604049, + -0.06003681, + -0.02688702, + 0.54489318, + -0.10951559, + -0.43730539, + ], + [ + -0.25090748, + -0.37466262, + 0.34085833, + -0.26690852, + -0.37676917, + 0.29080825, + 0.31600481, + 0.37558276, + -0.33251064, + ], + [ + -0.80195614, + -0.10273138, + 0.06935364, + -0.10429256, + -0.29693811, + 0.45643496, + 0.07247872, + 0.45604679, + -0.71048816, + ], + [ + -0.03840668, + -0.07680205, + 0.10940472, + -0.02374189, + -0.27610266, + 0.4336071, + 0.02465248, + 0.4290638, + -0.67496763, + ], + [ + -0.61475065, + -0.21163135, + 0.26652929, + -0.26134659, + -0.11560267, + 0.15415902, + 0.34343952, + 0.1589482, + -0.21370642, + ], + ] +).reshape(6, 9) + +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +coord = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +type_OH = np.array([1, 2, 2, 1, 2, 2]) +type_HO = np.array([2, 1, 1, 2, 1, 1]) + + +sp.check_output( + "{} -m deepmd convert-from pbtxt -i {} -o {}".format( + sys.executable, + pbtxt_file2.resolve(), + pb_file2.resolve(), + ).split() +) + + +def setup_module(): + write_lmp_data(box, coord, type_OH, data_file) + write_lmp_data(box, coord, type_HO, data_type_map_file) + write_lmp_data( + box * constants.dist_metal2si, + coord * constants.dist_metal2si, + type_OH, + data_file_si, + ) + + +def teardown_module(): + os.remove(data_file) + os.remove(data_type_map_file) + + +def _lammps(data_file, units="metal") -> PyLammps: + lammps = PyLammps() + lammps.units(units) + lammps.boundary("p p p") + lammps.atom_style("atomic") + if units == "metal" or units == "real": + lammps.neighbor("2.0 bin") + elif units == "si": + lammps.neighbor("2.0e-10 bin") + else: + raise ValueError("units should be metal, real, or si") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + if units == "metal" or units == "real": + lammps.mass("1 16") + lammps.mass("2 2") + elif units == "si": + lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) + lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) + else: + raise ValueError("units should be metal, real, or si") + if units == "metal": + lammps.timestep(0.0005) + elif units == "real": + lammps.timestep(0.5) + elif units == "si": + lammps.timestep(5e-16) + else: + raise ValueError("units should be metal, real, or si") + lammps.fix("1 all nve") + return lammps + + +@pytest.fixture +def lammps(): + lmp = _lammps(data_file=data_file) + yield lmp + lmp.close() + + +@pytest.fixture +def lammps_type_map(): + lmp = _lammps(data_file=data_type_map_file) + yield lmp + lmp.close() + + +@pytest.fixture +def lammps_real(): + lmp = _lammps(data_file=data_file, units="real") + yield lmp + lmp.close() + + +@pytest.fixture +def lammps_si(): + lmp = _lammps(data_file=data_file_si, units="si") + yield lmp + lmp.close() + + +def test_pair_deepmd(lammps): + lammps.pair_style(f"deepmd {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f[lammps.atoms[ii].id - 1] + ) + lammps.run(1) + + +def test_pair_deepmd_virial(lammps): + lammps.pair_style(f"deepmd {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.compute("virial all centroid/stress/atom NULL pair") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.dump( + "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) + ) + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f[lammps.atoms[ii].id - 1] + ) + idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + for ii in range(9): + assert np.array( + lammps.variables[f"virial{ii}"].value + ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) + + +def test_pair_deepmd_model_devi(lammps): + lammps.pair_style( + "deepmd {} {} out_file {} out_freq 1 atomic".format( + pb_file.resolve(), pb_file2.resolve(), md_file.resolve() + ) + ) + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f[lammps.atoms[ii].id - 1] + ) + # load model devi + md = np.loadtxt(md_file.resolve()) + expected_md_f = np.linalg.norm(np.std([expected_f, expected_f2], axis=0), axis=1) + assert md[7:] == pytest.approx(expected_md_f) + assert md[4] == pytest.approx(np.max(expected_md_f)) + assert md[5] == pytest.approx(np.min(expected_md_f)) + assert md[6] == pytest.approx(np.mean(expected_md_f)) + expected_md_v = ( + np.std([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0) / 6 + ) + assert md[1] == pytest.approx(np.max(expected_md_v)) + assert md[2] == pytest.approx(np.min(expected_md_v)) + assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) + + +def test_pair_deepmd_model_devi_virial(lammps): + lammps.pair_style( + "deepmd {} {} out_file {} out_freq 1 atomic".format( + pb_file.resolve(), pb_file2.resolve(), md_file.resolve() + ) + ) + lammps.pair_coeff("* *") + lammps.compute("virial all centroid/stress/atom NULL pair") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.dump( + "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) + ) + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f[lammps.atoms[ii].id - 1] + ) + idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + for ii in range(9): + assert np.array( + lammps.variables[f"virial{ii}"].value + ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) + # load model devi + md = np.loadtxt(md_file.resolve()) + expected_md_f = np.linalg.norm(np.std([expected_f, expected_f2], axis=0), axis=1) + assert md[7:] == pytest.approx(expected_md_f) + assert md[4] == pytest.approx(np.max(expected_md_f)) + assert md[5] == pytest.approx(np.min(expected_md_f)) + assert md[6] == pytest.approx(np.mean(expected_md_f)) + expected_md_v = ( + np.std([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0) / 6 + ) + assert md[1] == pytest.approx(np.max(expected_md_v)) + assert md[2] == pytest.approx(np.min(expected_md_v)) + assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) + + +def test_pair_deepmd_model_devi_atomic_relative(lammps): + relative = 1.0 + lammps.pair_style( + "deepmd {} {} out_file {} out_freq 1 atomic relative {}".format( + pb_file.resolve(), pb_file2.resolve(), md_file.resolve(), relative + ) + ) + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f[lammps.atoms[ii].id - 1] + ) + # load model devi + md = np.loadtxt(md_file.resolve()) + norm = np.linalg.norm(np.mean([expected_f, expected_f2], axis=0), axis=1) + expected_md_f = np.linalg.norm(np.std([expected_f, expected_f2], axis=0), axis=1) + expected_md_f /= norm + relative + assert md[7:] == pytest.approx(expected_md_f) + assert md[4] == pytest.approx(np.max(expected_md_f)) + assert md[5] == pytest.approx(np.min(expected_md_f)) + assert md[6] == pytest.approx(np.mean(expected_md_f)) + expected_md_v = ( + np.std([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0) / 6 + ) + assert md[1] == pytest.approx(np.max(expected_md_v)) + assert md[2] == pytest.approx(np.min(expected_md_v)) + assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) + + +def test_pair_deepmd_model_devi_atomic_relative_v(lammps): + relative = 1.0 + lammps.pair_style( + "deepmd {} {} out_file {} out_freq 1 atomic relative_v {}".format( + pb_file.resolve(), pb_file2.resolve(), md_file.resolve(), relative + ) + ) + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f[lammps.atoms[ii].id - 1] + ) + md = np.loadtxt(md_file.resolve()) + expected_md_f = np.linalg.norm(np.std([expected_f, expected_f2], axis=0), axis=1) + assert md[7:] == pytest.approx(expected_md_f) + assert md[4] == pytest.approx(np.max(expected_md_f)) + assert md[5] == pytest.approx(np.min(expected_md_f)) + assert md[6] == pytest.approx(np.mean(expected_md_f)) + expected_md_v = ( + np.std([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0) / 6 + ) + norm = ( + np.abs( + np.mean([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0) + ) + / 6 + ) + expected_md_v /= norm + relative + assert md[1] == pytest.approx(np.max(expected_md_v)) + assert md[2] == pytest.approx(np.min(expected_md_v)) + assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) + + +def test_pair_deepmd_type_map(lammps_type_map): + lammps_type_map.pair_style(f"deepmd {pb_file.resolve()}") + lammps_type_map.pair_coeff("* * H O") + lammps_type_map.run(0) + assert lammps_type_map.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps_type_map.atoms[ii].force == pytest.approx( + expected_f[lammps_type_map.atoms[ii].id - 1] + ) + lammps_type_map.run(1) + + +def test_pair_deepmd_real(lammps_real): + lammps_real.pair_style(f"deepmd {pb_file.resolve()}") + lammps_real.pair_coeff("* *") + lammps_real.run(0) + assert lammps_real.eval("pe") == pytest.approx( + expected_e * constants.ener_metal2real + ) + for ii in range(6): + assert lammps_real.atoms[ii].force == pytest.approx( + expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real + ) + lammps_real.run(1) + + +def test_pair_deepmd_virial_real(lammps_real): + lammps_real.pair_style(f"deepmd {pb_file.resolve()}") + lammps_real.pair_coeff("* *") + lammps_real.compute("virial all centroid/stress/atom NULL pair") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.dump( + "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) + ) + lammps_real.run(0) + assert lammps_real.eval("pe") == pytest.approx( + expected_e * constants.ener_metal2real + ) + for ii in range(6): + assert lammps_real.atoms[ii].force == pytest.approx( + expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real + ) + idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + for ii in range(9): + assert np.array( + lammps_real.variables[f"virial{ii}"].value + ) / constants.nktv2p_real == pytest.approx( + expected_v[idx_map, ii] * constants.ener_metal2real + ) + + +def test_pair_deepmd_model_devi_real(lammps_real): + lammps_real.pair_style( + "deepmd {} {} out_file {} out_freq 1 atomic".format( + pb_file.resolve(), pb_file2.resolve(), md_file.resolve() + ) + ) + lammps_real.pair_coeff("* *") + lammps_real.run(0) + assert lammps_real.eval("pe") == pytest.approx( + expected_e * constants.ener_metal2real + ) + for ii in range(6): + assert lammps_real.atoms[ii].force == pytest.approx( + expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real + ) + # load model devi + md = np.loadtxt(md_file.resolve()) + expected_md_f = np.linalg.norm(np.std([expected_f, expected_f2], axis=0), axis=1) + assert md[7:] == pytest.approx(expected_md_f * constants.force_metal2real) + assert md[4] == pytest.approx(np.max(expected_md_f) * constants.force_metal2real) + assert md[5] == pytest.approx(np.min(expected_md_f) * constants.force_metal2real) + assert md[6] == pytest.approx(np.mean(expected_md_f) * constants.force_metal2real) + expected_md_v = ( + np.std([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0) / 6 + ) + assert md[1] == pytest.approx(np.max(expected_md_v) * constants.ener_metal2real) + assert md[2] == pytest.approx(np.min(expected_md_v) * constants.ener_metal2real) + assert md[3] == pytest.approx( + np.sqrt(np.mean(np.square(expected_md_v))) * constants.ener_metal2real + ) + + +def test_pair_deepmd_model_devi_virial_real(lammps_real): + lammps_real.pair_style( + "deepmd {} {} out_file {} out_freq 1 atomic".format( + pb_file.resolve(), pb_file2.resolve(), md_file.resolve() + ) + ) + lammps_real.pair_coeff("* *") + lammps_real.compute("virial all centroid/stress/atom NULL pair") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.dump( + "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) + ) + lammps_real.run(0) + assert lammps_real.eval("pe") == pytest.approx( + expected_e * constants.ener_metal2real + ) + for ii in range(6): + assert lammps_real.atoms[ii].force == pytest.approx( + expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real + ) + idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + for ii in range(9): + assert np.array( + lammps_real.variables[f"virial{ii}"].value + ) / constants.nktv2p_real == pytest.approx( + expected_v[idx_map, ii] * constants.ener_metal2real + ) + # load model devi + md = np.loadtxt(md_file.resolve()) + expected_md_f = np.linalg.norm(np.std([expected_f, expected_f2], axis=0), axis=1) + assert md[7:] == pytest.approx(expected_md_f * constants.force_metal2real) + assert md[4] == pytest.approx(np.max(expected_md_f) * constants.force_metal2real) + assert md[5] == pytest.approx(np.min(expected_md_f) * constants.force_metal2real) + assert md[6] == pytest.approx(np.mean(expected_md_f) * constants.force_metal2real) + expected_md_v = ( + np.std([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0) / 6 + ) + assert md[1] == pytest.approx(np.max(expected_md_v) * constants.ener_metal2real) + assert md[2] == pytest.approx(np.min(expected_md_v) * constants.ener_metal2real) + assert md[3] == pytest.approx( + np.sqrt(np.mean(np.square(expected_md_v))) * constants.ener_metal2real + ) + + +def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): + relative = 1.0 + lammps_real.pair_style( + "deepmd {} {} out_file {} out_freq 1 atomic relative {}".format( + pb_file.resolve(), + pb_file2.resolve(), + md_file.resolve(), + relative * constants.force_metal2real, + ) + ) + lammps_real.pair_coeff("* *") + lammps_real.run(0) + assert lammps_real.eval("pe") == pytest.approx( + expected_e * constants.ener_metal2real + ) + for ii in range(6): + assert lammps_real.atoms[ii].force == pytest.approx( + expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real + ) + # load model devi + md = np.loadtxt(md_file.resolve()) + norm = np.linalg.norm(np.mean([expected_f, expected_f2], axis=0), axis=1) + expected_md_f = np.linalg.norm(np.std([expected_f, expected_f2], axis=0), axis=1) + expected_md_f /= norm + relative + assert md[7:] == pytest.approx(expected_md_f * constants.force_metal2real) + assert md[4] == pytest.approx(np.max(expected_md_f) * constants.force_metal2real) + assert md[5] == pytest.approx(np.min(expected_md_f) * constants.force_metal2real) + assert md[6] == pytest.approx(np.mean(expected_md_f) * constants.force_metal2real) + expected_md_v = ( + np.std([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0) / 6 + ) + assert md[1] == pytest.approx(np.max(expected_md_v) * constants.ener_metal2real) + assert md[2] == pytest.approx(np.min(expected_md_v) * constants.ener_metal2real) + assert md[3] == pytest.approx( + np.sqrt(np.mean(np.square(expected_md_v))) * constants.ener_metal2real + ) + + +def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): + relative = 1.0 + lammps_real.pair_style( + "deepmd {} {} out_file {} out_freq 1 atomic relative_v {}".format( + pb_file.resolve(), + pb_file2.resolve(), + md_file.resolve(), + relative * constants.ener_metal2real, + ) + ) + lammps_real.pair_coeff("* *") + lammps_real.run(0) + assert lammps_real.eval("pe") == pytest.approx( + expected_e * constants.ener_metal2real + ) + for ii in range(6): + assert lammps_real.atoms[ii].force == pytest.approx( + expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real + ) + md = np.loadtxt(md_file.resolve()) + expected_md_f = np.linalg.norm(np.std([expected_f, expected_f2], axis=0), axis=1) + assert md[7:] == pytest.approx(expected_md_f * constants.force_metal2real) + assert md[4] == pytest.approx(np.max(expected_md_f) * constants.force_metal2real) + assert md[5] == pytest.approx(np.min(expected_md_f) * constants.force_metal2real) + assert md[6] == pytest.approx(np.mean(expected_md_f) * constants.force_metal2real) + expected_md_v = ( + np.std([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0) / 6 + ) + norm = ( + np.abs( + np.mean([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0) + ) + / 6 + ) + expected_md_v /= norm + relative + assert md[1] == pytest.approx(np.max(expected_md_v) * constants.ener_metal2real) + assert md[2] == pytest.approx(np.min(expected_md_v) * constants.ener_metal2real) + assert md[3] == pytest.approx( + np.sqrt(np.mean(np.square(expected_md_v))) * constants.ener_metal2real + ) + + +def test_pair_deepmd_si(lammps_si): + lammps_si.pair_style(f"deepmd {pb_file.resolve()}") + lammps_si.pair_coeff("* *") + lammps_si.run(0) + assert lammps_si.eval("pe") == pytest.approx(expected_e * constants.ener_metal2si) + for ii in range(6): + assert lammps_si.atoms[ii].force == pytest.approx( + expected_f[lammps_si.atoms[ii].id - 1] * constants.force_metal2si + ) + lammps_si.run(1) diff --git a/source/tests/common/test_out_stat.py b/source/tests/common/test_out_stat.py new file mode 100644 index 0000000000..c0cfc25071 --- /dev/null +++ b/source/tests/common/test_out_stat.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.utils.out_stat import ( + compute_stats_from_atomic, + compute_stats_from_redu, +) + + +class TestOutStat(unittest.TestCase): + def setUp(self) -> None: + rng = np.random.default_rng(20240227) + ndim = 5 + nframes = 1000 + ntypes = 3 + nloc = 1000 + self.atype = rng.integers(0, ntypes, size=(nframes, nloc)) + # compute the number of atoms for each type in each frame + self.natoms = np.zeros((nframes, ntypes), dtype=np.int64) + for i in range(ntypes): + self.natoms[:, i] = (self.atype == i).sum(axis=1) + self.mean = rng.random((ntypes, ndim)) * 1e4 + self.std = rng.random((ntypes, ndim)) * 1e-3 + + # generate random output + self.output = rng.normal( + loc=self.mean[self.atype, :], + scale=self.std[self.atype, :], + size=(nframes, nloc, ndim), + ) + self.output_redu = self.output.sum(axis=1) + + return super().setUp() + + def test_compute_stats_from_redu(self): + bias, std = compute_stats_from_redu(self.output_redu, self.natoms) + np.testing.assert_allclose(bias, self.mean, rtol=1e-7) + reference_std = np.array( + [ + 0.01700638138272794, + 0.01954897296228177, + 0.020281857747683162, + 0.010741237959989648, + 0.020258211828681347, + ] + ) + np.testing.assert_allclose( + std, + reference_std, + rtol=1e-7, + ) + # ensure the sum is close + np.testing.assert_allclose( + self.output_redu, + self.natoms @ bias, + rtol=1e-7, + ) + + def test_compute_stats_from_redu_with_assigned_bias(self): + assigned_bias = np.full_like(self.mean, np.nan) + assigned_bias[0] = self.mean[0] + bias, std = compute_stats_from_redu( + self.output_redu, + self.natoms, + assigned_bias=assigned_bias, + ) + np.testing.assert_allclose(bias, self.mean, rtol=1e-7) + np.testing.assert_allclose(bias[0], self.mean[0], rtol=1e-14) + reference_std = np.array( + [ + 0.017015794087883902, + 0.019549011723239484, + 0.020285565914828625, + 0.01074124012073672, + 0.020283557003416414, + ] + ) + np.testing.assert_allclose( + std, + reference_std, + rtol=1e-7, + ) + # ensure the sum is close + np.testing.assert_allclose( + self.output_redu, + self.natoms @ bias, + rtol=1e-7, + ) + + def test_compute_stats_from_atomic(self): + bias, std = compute_stats_from_atomic(self.output, self.atype) + np.testing.assert_allclose(bias, self.mean) + reference_std = np.array( + [ + [ + 0.0005452949516910239, + 0.000686732800598535, + 0.00089423457667224, + 7.818017989121455e-05, + 0.0004758637035637342, + ], + [ + 2.0610161678825724e-05, + 0.0007728218734771541, + 0.0004754659308165858, + 0.0001809007655290948, + 0.0008187364708029638, + ], + [ + 0.0007935836092665254, + 0.00031176505013516624, + 0.0005469653430009186, + 0.0005652240916389281, + 0.0006087722080071852, + ], + ] + ) + np.testing.assert_allclose( + std, + reference_std, + rtol=1e-7, + ) diff --git a/source/tests/infer/deeppot_sea.pth b/source/tests/infer/deeppot_sea.pth new file mode 100644 index 0000000000..98aaa8a2ad Binary files /dev/null and b/source/tests/infer/deeppot_sea.pth differ diff --git a/source/tests/infer/fparam_aparam.pbtxt b/source/tests/infer/fparam_aparam.pbtxt index a89596961e..8c2e884090 100644 --- a/source/tests/infer/fparam_aparam.pbtxt +++ b/source/tests/infer/fparam_aparam.pbtxt @@ -35,7 +35,7 @@ node { dtype: DT_STRING tensor_shape { } - string_val: "{\"model\":{\"data_stat_nbatch\":1,\"descriptor\":{\"type\":\"se_e2_a\",\"sel\":[60],\"rcut_smth\":1.8,\"rcut\":6.0,\"neuron\":[5,10,20],\"resnet_dt\":false,\"axis_neuron\":8,\"seed\":1,\"activation_function\":\"tanh\",\"type_one_side\":false,\"precision\":\"default\",\"trainable\":true,\"exclude_types\":[],\"set_davg_zero\":false},\"fitting_net\":{\"neuron\":[5,5,5],\"resnet_dt\":true,\"numb_fparam\":1,\"numb_aparam\":1,\"seed\":1,\"type\":\"ener\",\"activation_function\":\"tanh\",\"precision\":\"default\",\"trainable\":true,\"rcond\":0.001,\"atom_ener\":[],\"use_aparam_as_mask\":false},\"data_stat_protect\":0.01,\"data_bias_nsample\":10},\"loss\":{\"start_pref_e\":0.02,\"limit_pref_e\":1,\"start_pref_f\":1000,\"limit_pref_f\":1,\"start_pref_v\":0,\"limit_pref_v\":0,\"type\":\"ener\",\"start_pref_ae\":0.0,\"limit_pref_ae\":0.0,\"start_pref_pf\":0.0,\"limit_pref_pf\":0.0,\"enable_atom_ener_coeff\":false},\"learning_rate\":{\"start_lr\":0.001,\"stop_lr\":3e-08,\"decay_steps\":5000,\"scale_by_worker\":\"linear\",\"type\":\"exp\"},\"training\":{\"training_data\":{\"systems\":[\"../data/e3000_i2000/\",\"../data/e8000_i2000/\"],\"set_prefix\":\"set\",\"batch_size\":1,\"auto_prob\":\"prob_sys_size\",\"sys_probs\":null},\"seed\":1,\"disp_file\":\"lcurve.out\",\"disp_freq\":100,\"save_freq\":1000,\"save_ckpt\":\"model.ckpt\",\"disp_training\":true,\"time_training\":true,\"profiling\":false,\"profiling_file\":\"timeline.json\",\"numb_steps\":1000,\"validation_data\":null,\"enable_profiler\":false,\"tensorboard\":false,\"tensorboard_log_dir\":\"log\",\"tensorboard_freq\":1}}" + string_val: "{\"model\":{\"data_stat_nbatch\":1,\"type_map\":[\"O\"],\"descriptor\":{\"type\":\"se_e2_a\",\"sel\":[60],\"rcut_smth\":1.8,\"rcut\":6.0,\"neuron\":[5,10,20],\"resnet_dt\":false,\"axis_neuron\":8,\"seed\":1,\"activation_function\":\"tanh\",\"type_one_side\":false,\"precision\":\"default\",\"trainable\":true,\"exclude_types\":[],\"set_davg_zero\":false},\"fitting_net\":{\"neuron\":[5,5,5],\"resnet_dt\":true,\"numb_fparam\":1,\"numb_aparam\":1,\"seed\":1,\"type\":\"ener\",\"activation_function\":\"tanh\",\"precision\":\"default\",\"trainable\":true,\"rcond\":0.001,\"atom_ener\":[],\"use_aparam_as_mask\":false},\"data_stat_protect\":0.01,\"data_bias_nsample\":10},\"loss\":{\"start_pref_e\":0.02,\"limit_pref_e\":1,\"start_pref_f\":1000,\"limit_pref_f\":1,\"start_pref_v\":0,\"limit_pref_v\":0,\"type\":\"ener\",\"start_pref_ae\":0.0,\"limit_pref_ae\":0.0,\"start_pref_pf\":0.0,\"limit_pref_pf\":0.0,\"enable_atom_ener_coeff\":false},\"learning_rate\":{\"start_lr\":0.001,\"stop_lr\":3e-08,\"decay_steps\":5000,\"scale_by_worker\":\"linear\",\"type\":\"exp\"},\"training\":{\"training_data\":{\"systems\":[\"../data/e3000_i2000/\",\"../data/e8000_i2000/\"],\"set_prefix\":\"set\",\"batch_size\":1,\"auto_prob\":\"prob_sys_size\",\"sys_probs\":null},\"seed\":1,\"disp_file\":\"lcurve.out\",\"disp_freq\":100,\"save_freq\":1000,\"save_ckpt\":\"model.ckpt\",\"disp_training\":true,\"time_training\":true,\"profiling\":false,\"profiling_file\":\"timeline.json\",\"numb_steps\":1000,\"validation_data\":null,\"enable_profiler\":false,\"tensorboard\":false,\"tensorboard_log_dir\":\"log\",\"tensorboard_freq\":1}}" } } } diff --git a/source/tests/infer/fparam_aparam.pth b/source/tests/infer/fparam_aparam.pth new file mode 100644 index 0000000000..7b0204cdd3 Binary files /dev/null and b/source/tests/infer/fparam_aparam.pth differ diff --git a/source/tests/pt/model/test_deeppot.py b/source/tests/pt/model/test_deeppot.py index 334206a2b0..102e1f6b0c 100644 --- a/source/tests/pt/model/test_deeppot.py +++ b/source/tests/pt/model/test_deeppot.py @@ -23,6 +23,10 @@ DeepPot, ) +from ...tf.test_deeppot_a import ( + FparamAparamCommonTest, +) + class TestDeepPot(unittest.TestCase): def setUp(self): @@ -123,3 +127,21 @@ def setUp(self): @unittest.mock.patch("deepmd.pt.infer.deep_eval.DEVICE", torch.device("cpu")) def test_dp_test_cpu(self): self.test_dp_test() + + +class TestFparamAparamPT(FparamAparamCommonTest, unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.dp = DeepPot( + str(Path(__file__).parent.parent.parent / "infer/fparam_aparam.pth") + ) + + def setUp(self): + super().setUp() + # For unclear reason, the precision is only 1e-7 + # not sure if it is expected... + self.places = 1e-7 + + @classmethod + def tearDownClass(cls): + pass diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py index c999f06863..5b8b6c9251 100644 --- a/source/tests/pt/model/test_descriptor_se_r.py +++ b/source/tests/pt/model/test_descriptor_se_r.py @@ -15,6 +15,9 @@ from deepmd.pt.utils.env import ( PRECISION_DICT, ) +from deepmd.pt.utils.env_mat_stat import ( + EnvMatStatSe, +) from .test_env_mat import ( TestCaseSingleFrameWithNlist, @@ -103,13 +106,61 @@ def test_consistency( err_msg=err_msg, ) + def test_load_stat(self): + rng = np.random.default_rng() + _, _, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 1)) + dstd = rng.normal(size=(self.nt, nnei, 1)) + dstd = 0.1 + np.abs(dstd) + + for idt, prec in itertools.product( + [False, True], + ["float64", "float32"], + ): + dtype = PRECISION_DICT[prec] + + # sea new impl + dd0 = DescrptSeR( + self.rcut, + self.rcut_smth, + self.sel, + precision=prec, + resnet_dt=idt, + old_impl=False, + ) + dd0.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) + dd0.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) + dd1 = DescrptSeR.deserialize(dd0.serialize()) + dd1.compute_input_stats( + [ + { + "r0": None, + "coord": torch.from_numpy(self.coord_ext) + .reshape(-1, self.nall, 3) + .to(env.DEVICE), + "atype": torch.from_numpy(self.atype_ext).to(env.DEVICE), + "box": None, + "natoms": self.nall, + } + ] + ) + + with self.assertRaises(ValueError) as cm: + ev = EnvMatStatSe(dd1) + ev.last_dim = 3 + ev.load_or_compute_stats([]) + self.assertEqual( + "last_dim should be 1 for raial-only or 4 for full descriptor.", + str(cm.exception), + ) + def test_jit( self, ): rng = np.random.default_rng() _, _, nnei = self.nlist.shape - davg = rng.normal(size=(self.nt, nnei, 4)) - dstd = rng.normal(size=(self.nt, nnei, 4)) + davg = rng.normal(size=(self.nt, nnei, 1)) + dstd = rng.normal(size=(self.nt, nnei, 1)) dstd = 0.1 + np.abs(dstd) for idt, prec in itertools.product( diff --git a/source/tests/pt/test_init_frz_model.py b/source/tests/pt/test_init_frz_model.py new file mode 100644 index 0000000000..d156eddc41 --- /dev/null +++ b/source/tests/pt/test_init_frz_model.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import unittest +from argparse import ( + Namespace, +) +from copy import ( + deepcopy, +) +from pathlib import ( + Path, +) + +import numpy as np + +from deepmd.pt.entrypoints.main import ( + freeze, + get_trainer, +) +from deepmd.pt.infer.deep_eval import ( + DeepPot, +) + + +class TestInitFrzModel(unittest.TestCase): + def setUp(self): + input_json = str(Path(__file__).parent / "water/se_atten.json") + with open(input_json) as f: + config = json.load(f) + config["training"]["numb_steps"] = 1 + config["training"]["save_freq"] = 1 + config["learning_rate"]["start_lr"] = 1.0 + config["training"]["training_data"]["systems"] = [ + str(Path(__file__).parent / "water/data/single") + ] + config["training"]["validation_data"]["systems"] = [ + str(Path(__file__).parent / "water/data/single") + ] + + self.models = [] + for imodel in range(2): + if imodel == 1: + config["training"]["numb_steps"] = 0 + trainer = get_trainer(deepcopy(config), init_frz_model=self.models[-1]) + else: + trainer = get_trainer(deepcopy(config)) + trainer.run() + + frozen_model = f"frozen_model{imodel}.pth" + ns = Namespace( + model="model.pt", + output=frozen_model, + head=None, + ) + freeze(ns) + self.models.append(frozen_model) + + def test_dp_test(self): + dp1 = DeepPot(str(self.models[0])) + dp2 = DeepPot(str(self.models[1])) + cell = np.array( + [ + 5.122106549439247480e00, + 4.016537340154059388e-01, + 6.951654033828678081e-01, + 4.016537340154059388e-01, + 6.112136112297989143e00, + 8.178091365465004481e-01, + 6.951654033828678081e-01, + 8.178091365465004481e-01, + 6.159552512682983760e00, + ] + ).reshape(1, 3, 3) + coord = np.array( + [ + 2.978060152121375648e00, + 3.588469695887098077e00, + 2.792459820604495491e00, + 3.895592322591093115e00, + 2.712091020667753760e00, + 1.366836847133650501e00, + 9.955616170888935690e-01, + 4.121324820711413039e00, + 1.817239061889086571e00, + 3.553661462345699906e00, + 5.313046969500791583e00, + 6.635182659098815883e00, + 6.088601018589653080e00, + 6.575011420004332585e00, + 6.825240650611076099e00, + ] + ).reshape(1, -1, 3) + atype = np.array([0, 0, 0, 1, 1]).reshape(1, -1) + + e1, f1, v1, ae1, av1 = dp1.eval(coord, cell, atype, atomic=True) + e2, f2, v2, ae2, av2 = dp2.eval(coord, cell, atype, atomic=True) + np.testing.assert_allclose(e1, e2, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(f1, f2, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(v1, v2, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(ae1, ae2, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(av1, av2, rtol=1e-10, atol=1e-10) diff --git a/source/tests/pt/test_stat.py b/source/tests/pt/test_stat.py index 1e3c707d6f..98d4e59d95 100644 --- a/source/tests/pt/test_stat.py +++ b/source/tests/pt/test_stat.py @@ -20,15 +20,15 @@ from deepmd.pt.model.descriptor.dpa1 import ( DescrptDPA1, ) +from deepmd.pt.model.task.ener import ( + EnergyFittingNet, +) from deepmd.pt.utils import ( env, ) from deepmd.pt.utils.dataloader import ( DpLoaderSet, ) -from deepmd.pt.utils.stat import ( - compute_output_bias, -) from deepmd.pt.utils.stat import make_stat_input as my_make from deepmd.tf.common import ( expand_sys_str, @@ -145,9 +145,14 @@ def my_merge(energy, natoms): dp_fn = EnerFitting( self.dp_d.get_ntypes(), self.dp_d.get_dim_out(), self.n_neuron ) - dp_fn.compute_output_stats(self.dp_sampled) - bias_atom_e = compute_output_bias(energy, natoms) - self.assertTrue(np.allclose(dp_fn.bias_atom_e, bias_atom_e[:, 0])) + dp_fn.compute_output_stats(self.dp_sampled, mixed_type=self.mixed_type) + pt_fn = EnergyFittingNet( + self.dp_d.get_ntypes(), self.dp_d.get_dim_out(), self.n_neuron + ) + pt_fn.compute_output_stats(self.my_sampled) + np.testing.assert_allclose( + dp_fn.bias_atom_e, pt_fn.bias_atom_e.detach().cpu().numpy().ravel() + ) # temporarily delete this function for performance of seeds in tf and pytorch may be different """ diff --git a/source/tests/tf/common.py b/source/tests/tf/common.py index a83397c11c..0bcb29b4b5 100644 --- a/source/tests/tf/common.py +++ b/source/tests/tf/common.py @@ -17,6 +17,9 @@ tf, ) from deepmd.tf.utils import random as dp_random +from deepmd.utils.out_stat import ( + compute_stats_from_redu, +) if GLOBAL_NP_FLOAT_PRECISION == np.float32: global_default_fv_hh = 1e-2 @@ -1041,10 +1044,12 @@ def compute_energy_shift(self): sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION) sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1]) sys_tynatom = sys_tynatom[:, 2:] - energy_shift, resd, rank, s_value = np.linalg.lstsq( - sys_tynatom, sys_ener, rcond=None + energy_shift, _ = compute_stats_from_redu( + sys_ener.reshape(-1, 1), + sys_tynatom, + rcond=None, ) - return energy_shift + return energy_shift.ravel() def process_sys_weights(self, sys_weights): sys_weights = np.array(sys_weights) diff --git a/source/tests/tf/test_deeppot_a.py b/source/tests/tf/test_deeppot_a.py index af060aca1c..9b4d64282f 100644 --- a/source/tests/tf/test_deeppot_a.py +++ b/source/tests/tf/test_deeppot_a.py @@ -894,17 +894,9 @@ def test_eval_typeebd(self): np.testing.assert_almost_equal(eval_typeebd, expected_typeebd, default_places) -class TestFparamAparam(unittest.TestCase): +class FparamAparamCommonTest: """Test fparam and aparam.""" - @classmethod - def setUpClass(cls): - convert_pbtxt_to_pb( - str(infer_path / os.path.join("fparam_aparam.pbtxt")), - "fparam_aparam.pb", - ) - cls.dp = DeepPot("fparam_aparam.pb") - def setUp(self): self.coords = np.array( [ @@ -1022,15 +1014,11 @@ def setUp(self): 2.875323131744185121e-02, ] ) - - @classmethod - def tearDownClass(cls): - os.remove("fparam_aparam.pb") - cls.dp = None + self.places = default_places def test_attrs(self): self.assertEqual(self.dp.get_ntypes(), 1) - self.assertAlmostEqual(self.dp.get_rcut(), 6.0, places=default_places) + self.assertAlmostEqual(self.dp.get_rcut(), 6.0, places=self.places) self.assertEqual(self.dp.get_dim_fparam(), 1) self.assertEqual(self.dp.get_dim_aparam(), 1) @@ -1050,13 +1038,11 @@ def test_1frame(self): self.assertEqual(ff.shape, (nframes, natoms, 3)) self.assertEqual(vv.shape, (nframes, 9)) # check values - np.testing.assert_almost_equal( - ff.ravel(), self.expected_f.ravel(), default_places - ) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), self.places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1) - np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), self.places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1) - np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), self.places) def test_1frame_atm(self): ee, ff, vv, ae, av = self.dp.eval( @@ -1076,19 +1062,13 @@ def test_1frame_atm(self): self.assertEqual(ae.shape, (nframes, natoms, 1)) self.assertEqual(av.shape, (nframes, natoms, 9)) # check values - np.testing.assert_almost_equal( - ff.ravel(), self.expected_f.ravel(), default_places - ) - np.testing.assert_almost_equal( - ae.ravel(), self.expected_e.ravel(), default_places - ) - np.testing.assert_almost_equal( - av.ravel(), self.expected_v.ravel(), default_places - ) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), self.places) + np.testing.assert_almost_equal(ae.ravel(), self.expected_e.ravel(), self.places) + np.testing.assert_almost_equal(av.ravel(), self.expected_v.ravel(), self.places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1) - np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), self.places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1) - np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), self.places) def test_2frame_atm_single_param(self): coords2 = np.concatenate((self.coords, self.coords)) @@ -1113,13 +1093,13 @@ def test_2frame_atm_single_param(self): expected_f = np.concatenate((self.expected_f, self.expected_f), axis=0) expected_e = np.concatenate((self.expected_e, self.expected_e), axis=0) expected_v = np.concatenate((self.expected_v, self.expected_v), axis=0) - np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(), default_places) - np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(), default_places) - np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(), default_places) + np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(), self.places) + np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(), self.places) + np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(), self.places) expected_se = np.sum(expected_e.reshape([nframes, -1]), axis=1) - np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), self.places) expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis=1) - np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), self.places) def test_2frame_atm_all_param(self): coords2 = np.concatenate((self.coords, self.coords)) @@ -1144,13 +1124,28 @@ def test_2frame_atm_all_param(self): expected_f = np.concatenate((self.expected_f, self.expected_f), axis=0) expected_e = np.concatenate((self.expected_e, self.expected_e), axis=0) expected_v = np.concatenate((self.expected_v, self.expected_v), axis=0) - np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(), default_places) - np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(), default_places) - np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(), default_places) + np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(), self.places) + np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(), self.places) + np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(), self.places) expected_se = np.sum(expected_e.reshape([nframes, -1]), axis=1) - np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), self.places) expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis=1) - np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), self.places) + + +class TestFparamAparam(FparamAparamCommonTest, unittest.TestCase): + @classmethod + def setUpClass(cls): + convert_pbtxt_to_pb( + str(infer_path / os.path.join("fparam_aparam.pbtxt")), + "fparam_aparam.pb", + ) + cls.dp = DeepPot("fparam_aparam.pb") + + @classmethod + def tearDownClass(cls): + os.remove("fparam_aparam.pb") + cls.dp = None class TestDeepPotAPBCNeighborList(TestDeepPotAPBC):