From 430dfa9ded456a5c096f4979a0981f53fe78d260 Mon Sep 17 00:00:00 2001 From: zhenyu <76582286+wangzyphysics@users.noreply.github.com> Date: Wed, 6 Nov 2024 17:23:18 +0800 Subject: [PATCH] fix: make sure `head` can be used in `DeepPot` (#4312) Following the example mentioned [here](https://www.aissquare.com/models/detail?pageType=models&name=DPA-2.3.0-v3.0.0b4&id=279) , I first select a `head` and put it into `DPCalculator` but encounter the error like this: ```shell (/public/home/mzq001/soft/deepmd-kit) [mzq001@login01 ~]$ python test.py To get the best performance, it is recommended to adjust the number of threads by setting the environment variables OMP_NUM_THREADS, DP_INTRA_OP_PARALLELISM_THREADS, and DP_INTER_OP_PARALLELISM_THREADS. See https://deepmd.rtfd.io/parallelism/ for more information. Traceback (most recent call last): File "/public/home/mzq001/test.py", line 4, in dp = DPCalculator("DPA2_medium_28_10M_beta4.pt", head="H2O_H2O-PD") ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/public/home/mzq001/soft/deepmd-kit/lib/python3.11/site-packages/deepmd/calculator.py", line 92, in __init__ self.dp = DeepPot(str(Path(model).resolve()), neighbor_list=neighbor_list) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/public/home/mzq001/soft/deepmd-kit/lib/python3.11/site-packages/deepmd/infer/deep_eval.py", line 334, in __init__ self.deep_eval = DeepEvalBackend( ^^^^^^^^^^^^^^^^ File "/public/home/mzq001/soft/deepmd-kit/lib/python3.11/site-packages/deepmd/pt/infer/deep_eval.py", line 121, in __init__ head is not None AssertionError: Head must be set for multitask model! Available heads are: ['Domains_Alloy', 'Domains_Anode', 'Domains_Cluster', 'Domains_Drug', 'Domains_FerroEle', 'Domains_OC2M', 'Domains_SSE-PBE', 'Domains_SemiCond', 'H2O_H2O-PD', 'Metals_AgAu-PBE', 'Metals_AlMgCu', 'Metals_Cu', 'Metals_Sn', 'Metals_Ti', 'Metals_V', 'Metals_W', 'Others_C12H26', 'Others_HfO2', 'Domains_ANI', 'Domains_SSE-PBESol', 'Domains_Transition1x', 'H2O_H2O-DPLR', 'H2O_H2O-PBE0TS-MD', 'H2O_H2O-PBE0TS', 'H2O_H2O-SCAN0', 'Metals_AgAu-PBED3', 'Others_In2Se3', 'MP_traj_v024_alldata_mixu'] ``` ```python ## Compute potential energy from ase import Atoms from deepmd.calculator import DP as DPCalculator dp = DPCalculator("DPA2_medium_28_10M_beta4.pt", head="H2O_H2O-PD") water = Atoms('H2O', positions=[(0.7601, 1.9270, 1), (1.9575, 1, 1), (1., 1., 1.)], cell=[100, 100, 100]) water.calc = dp print(water.get_potential_energy()) print(water.get_forces()) ## Run BFGS structure optimization from ase.optimize import BFGS dyn = BFGS(water) dyn.run(fmax=1e-6) print(water.get_positions()) ``` The variable named `head` of `DPCalculator` is not being used in `DeepPot` which directly leads to this error. ## Summary by CodeRabbit - **New Features** - Enhanced configurability of the DP class with an additional `head` parameter for initialization. - Updated default behavior in the calculate method to include stress calculations alongside energy and forces. - **Bug Fixes** - Clarified handling of stress property in the context of lattice relaxation, improving accuracy in calculations. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- deepmd/calculator.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/deepmd/calculator.py b/deepmd/calculator.py index 6f863ab09b..0fbc447aaa 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -45,6 +45,8 @@ class DP(Calculator): will infer this information from model, by default None neighbor_list : ase.neighborlist.NeighborList, optional The neighbor list object. If None, then build the native neighbor list. + head : Union[str, None], optional + a specific model branch choosing from pretrained model, by default None Examples -------- @@ -84,10 +86,15 @@ def __init__( label: str = "DP", type_dict: Optional[dict[str, int]] = None, neighbor_list=None, + head=None, **kwargs, ) -> None: Calculator.__init__(self, label=label, **kwargs) - self.dp = DeepPot(str(Path(model).resolve()), neighbor_list=neighbor_list) + self.dp = DeepPot( + str(Path(model).resolve()), + neighbor_list=neighbor_list, + head=head, + ) if type_dict: self.type_dict = type_dict else: