diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index 18fd7a1ac1..33ff8c6600 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -151,6 +151,8 @@ jobs: - variant: "_cu11" cuda_version: "11" steps: + - name: Delete huge unnecessary tools folder + run: rm -rf /opt/hostedtoolcache - uses: actions/checkout@v4 - uses: actions/download-artifact@v4 with: diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index 15a20d889a..b024431d65 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -47,7 +47,7 @@ jobs: && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3 if: false # skip as we use nvidia image - run: python -m pip install -U "pip>=21.3.1,!=23.0.0" - - run: python -m pip install "tensorflow>=2.15.0rc0" "torch>=2.2.0" + - run: python -m pip install "tensorflow>=2.15.0rc0" "torch==2.3.1.*" - run: python -m pip install -v -e .[gpu,test,lmp,cu12,torch] mpi4py "ase @ https://gitlab.com/ase/ase/-/archive/8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f/ase-8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f.tar.gz" env: DP_VARIANT: cuda diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 8a40f8d238..8d8bbf90c4 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -104,6 +104,10 @@ def get_out_bias(self) -> np.ndarray: """Return the output bias of the atomic model.""" return self.fitting["bias_atom_e"] + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the atomic model needs sorted nlist when using `forward_lower`.""" + return self.descriptor.need_sorted_nlist_for_lower() + def forward_atomic( self, extended_coord: np.ndarray, diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 93a885f3ab..ef2c629aad 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -111,6 +111,10 @@ def _sort_rcuts_sels(self) -> Tuple[List[float], List[int]]: ) return [p[0] for p in zipped], [p[1] for p in zipped] + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the atomic model needs sorted nlist when using `forward_lower`.""" + return True + def forward_atomic( self, extended_coord, diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index 3e02a5d076..9daf38a82f 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -134,6 +134,10 @@ def mixed_types(self) -> bool: """ pass + @abstractmethod + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + @abstractmethod def fwd( self, diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 30ab58928b..3660c97742 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -177,6 +177,10 @@ def deserialize(cls, data) -> "PairTabAtomicModel": tab_model.tab_data = tab_model.tab.tab_data.reshape(ntypes, ntypes, nspline, 4) return tab_model + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the atomic model needs sorted nlist when using `forward_lower`.""" + return False + def forward_atomic( self, extended_coord, diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 96640d75c8..2219897661 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -140,6 +140,10 @@ def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None) for descrpt in self.descrpt_list: descrpt.compute_input_stats(merged, path) + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + return True + def call( self, coord_ext, diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index 940bd0cd27..b98de7fa1f 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -103,6 +103,10 @@ def compute_input_stats( """Update mean and stddev for descriptor elements.""" raise NotImplementedError + @abstractmethod + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + @abstractmethod def fwd( self, diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 531aa09f3a..2ba0bc9882 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -283,6 +283,10 @@ def reinit_exclude( self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + return False + def call( self, coord_ext, diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index 3128a28493..f27db3e975 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -232,6 +232,10 @@ def cal_g( gg = self.embeddings[(ll,)].call(ss) return gg + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + return False + def call( self, coord_ext, diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 68889ad331..43a91650c8 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -214,7 +214,12 @@ def call_lower( """ nframes, nall = extended_atype.shape[:2] extended_coord = extended_coord.reshape(nframes, -1, 3) - nlist = self.format_nlist(extended_coord, extended_atype, nlist) + nlist = self.format_nlist( + extended_coord, + extended_atype, + nlist, + extra_nlist_sort=self.need_sorted_nlist_for_lower(), + ) cc_ext, _, fp, ap, input_prec = self.input_type_cast( extended_coord, fparam=fparam, aparam=aparam ) @@ -309,6 +314,7 @@ def format_nlist( extended_coord: np.ndarray, extended_atype: np.ndarray, nlist: np.ndarray, + extra_nlist_sort: bool = False, ): """Format the neighbor list. @@ -334,6 +340,8 @@ def format_nlist( atomic type in extended region. nf x nall nlist neighbor list. nf x nloc x nsel + extra_nlist_sort + whether to forcibly sort the nlist. Returns ------- @@ -343,7 +351,12 @@ def format_nlist( """ n_nf, n_nloc, n_nnei = nlist.shape mixed_types = self.mixed_types() - ret = self._format_nlist(extended_coord, nlist, sum(self.get_sel())) + ret = self._format_nlist( + extended_coord, + nlist, + sum(self.get_sel()), + extra_nlist_sort=extra_nlist_sort, + ) if not mixed_types: ret = nlist_distinguish_types(ret, extended_atype, self.get_sel()) return ret @@ -353,6 +366,7 @@ def _format_nlist( extended_coord: np.ndarray, nlist: np.ndarray, nnei: int, + extra_nlist_sort: bool = False, ): n_nf, n_nloc, n_nnei = nlist.shape extended_coord = extended_coord.reshape([n_nf, -1, 3]) @@ -368,8 +382,9 @@ def _format_nlist( ], axis=-1, ) - elif n_nnei > nnei: - # make a copy before revise + + if n_nnei > nnei or extra_nlist_sort: + n_nf, n_nloc, n_nnei = nlist.shape m_real_nei = nlist >= 0 ret = np.where(m_real_nei, nlist, 0) coord0 = extended_coord[:, :n_nloc, :] @@ -382,9 +397,11 @@ def _format_nlist( ret = np.take_along_axis(ret, ret_mapping, axis=2) ret = np.where(rr > rcut, -1, ret) ret = ret[..., :nnei] - else: # n_nnei == nnei: - # copy anyway... + # not extra_nlist_sort and n_nnei <= nnei: + elif n_nnei == nnei: ret = nlist + else: + pass assert ret.shape[-1] == nnei return ret @@ -469,6 +486,10 @@ def mixed_types(self) -> bool: """ return self.atomic_model.mixed_types() + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the model needs sorted nlist when using `forward_lower`.""" + return self.atomic_model.need_sorted_nlist_for_lower() + def atomic_output_def(self) -> FittingOutputDef: """Get the output def of the atomic model.""" return self.atomic_model.atomic_output_def() diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 13b8f09a79..2205bd5b05 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -100,6 +100,10 @@ def mixed_types(self) -> bool: """ return self.descriptor.mixed_types() + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the atomic model needs sorted nlist when using `forward_lower`.""" + return self.descriptor.need_sorted_nlist_for_lower() + def serialize(self) -> dict: dd = BaseAtomicModel.serialize(self) dd.update( diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index f599399e66..2edf17a046 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -132,6 +132,10 @@ def _sort_rcuts_sels(self) -> Tuple[List[float], List[int]]: sorted_sels: List[int] = outer_sorted[:, 1].to(torch.int64).tolist() return sorted_rcuts, sorted_sels + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the atomic model needs sorted nlist when using `forward_lower`.""" + return True + def forward_atomic( self, extended_coord: torch.Tensor, diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index c20abf6a12..cff286e793 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -164,6 +164,10 @@ def mixed_types(self) -> bool: # to match DPA1 and DPA2. return True + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the atomic model needs sorted nlist when using `forward_lower`.""" + return False + def serialize(self) -> dict: dd = BaseAtomicModel.serialize(self) dd.update( diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py index 5aae848aa4..68c06efa52 100644 --- a/deepmd/pt/model/descriptor/descriptor.py +++ b/deepmd/pt/model/descriptor/descriptor.py @@ -159,6 +159,10 @@ def forward( """Calculate DescriptorBlock.""" pass + @abstractmethod + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor block needs sorted nlist when using `forward_lower`.""" + def make_default_type_embedding( ntypes, diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 21275317dc..afdaa8ea4b 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -150,6 +150,10 @@ def mixed_types(self) -> bool: """ return self.se_atten.mixed_types() + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + return self.se_atten.need_sorted_nlist_for_lower() + def share_params(self, base_class, shared_level, resume=False): """ Share the parameters of self to the base_class with shared_level during multitask training. diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index fb792a51e2..4976a45aa0 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -307,6 +307,10 @@ def mixed_types(self) -> bool: """ return True + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + return True + def share_params(self, base_class, shared_level, resume=False): """ Share the parameters of self to the base_class with shared_level during multitask training. diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index 204ca7589d..fb202cb3eb 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -140,6 +140,10 @@ def mixed_types(self): """ return any(descrpt.mixed_types() for descrpt in self.descrpt_list) + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + return True + def share_params(self, base_class, shared_level, resume=False): """ Share the parameters of self to the base_class with shared_level during multitask training. @@ -369,6 +373,10 @@ def mixed_types(self) -> bool: """ return all(descriptor.mixed_types() for descriptor in self.descriptor_list) + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor block needs sorted nlist when using `forward_lower`.""" + return True + @property def dim_out(self): """Returns the output dimension of this descriptor.""" diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index 16a38052b1..7a157def0d 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/deepmd/pt/model/descriptor/repformers.py @@ -205,6 +205,10 @@ def mixed_types(self) -> bool: """ return True + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor block needs sorted nlist when using `forward_lower`.""" + return False + @property def dim_out(self): """Returns the output dimension of this descriptor.""" diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index e17b7c5d54..ef32c42a10 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -132,6 +132,10 @@ def mixed_types(self): """ return self.sea.mixed_types() + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + return self.sea.need_sorted_nlist_for_lower() + def share_params(self, base_class, shared_level, resume=False): """ Share the parameters of self to the base_class with shared_level during multitask training. @@ -433,6 +437,10 @@ def mixed_types(self) -> bool: """ return False + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor block needs sorted nlist when using `forward_lower`.""" + return False + @property def dim_out(self): """Returns the output dimension of this descriptor.""" diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 051c66385c..ef4e32ddde 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -196,6 +196,10 @@ def mixed_types(self) -> bool: """ return True + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor block needs sorted nlist when using `forward_lower`.""" + return False + @property def dim_out(self): """Returns the output dimension of this descriptor.""" diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index ff922e0649..a06b08bf13 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -160,6 +160,10 @@ def mixed_types(self) -> bool: """ return False + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + return False + def share_params(self, base_class, shared_level, resume=False): """ Share the parameters of self to the base_class with shared_level during multitask training. diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index 45b120771b..2e6eeef641 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -76,6 +76,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) if self.get_fitting_net() is not None: model_predict = {} diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index e043700bee..510b14d31b 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -74,6 +74,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) if self.get_fitting_net() is not None: model_predict = {} diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index bbc82b8d77..d63327d1fd 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -86,6 +86,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) model_predict = {} diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 5217293623..35acef0f1b 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -78,6 +78,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) if self.get_fitting_net() is not None: model_predict = {} diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py index e3dcd389bb..6a6283b12e 100644 --- a/deepmd/pt/model/model/frozen.py +++ b/deepmd/pt/model/model/frozen.py @@ -102,6 +102,10 @@ def mixed_types(self) -> bool: """ return self.model.mixed_types() + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the model needs sorted nlist when using `forward_lower`.""" + return self.model.need_sorted_nlist_for_lower() + @torch.jit.export def forward( self, diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 595e5fba74..5eb9c4309d 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -216,6 +216,7 @@ def forward_common_lower( fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, do_atomic_virial: bool = False, + extra_nlist_sort: bool = False, ): """Return model prediction. Lower interface that takes extended atomic coordinates and types, nlist, and mapping @@ -238,6 +239,8 @@ def forward_common_lower( atomic parameter. nf x nloc x nda do_atomic_virial whether calculate atomic virial. + extra_nlist_sort + whether to forcibly sort the nlist. Returns ------- @@ -247,7 +250,9 @@ def forward_common_lower( """ nframes, nall = extended_atype.shape[:2] extended_coord = extended_coord.view(nframes, -1, 3) - nlist = self.format_nlist(extended_coord, extended_atype, nlist) + nlist = self.format_nlist( + extended_coord, extended_atype, nlist, extra_nlist_sort=extra_nlist_sort + ) cc_ext, _, fp, ap, input_prec = self.input_type_cast( extended_coord, fparam=fparam, aparam=aparam ) @@ -348,6 +353,7 @@ def format_nlist( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, + extra_nlist_sort: bool = False, ): """Format the neighbor list. @@ -373,6 +379,8 @@ def format_nlist( atomic type in extended region. nf x nall nlist neighbor list. nf x nloc x nsel + extra_nlist_sort + whether to forcibly sort the nlist. Returns ------- @@ -381,7 +389,12 @@ def format_nlist( """ mixed_types = self.mixed_types() - nlist = self._format_nlist(extended_coord, nlist, sum(self.get_sel())) + nlist = self._format_nlist( + extended_coord, + nlist, + sum(self.get_sel()), + extra_nlist_sort=extra_nlist_sort, + ) if not mixed_types: nlist = nlist_distinguish_types(nlist, extended_atype, self.get_sel()) return nlist @@ -391,6 +404,7 @@ def _format_nlist( extended_coord: torch.Tensor, nlist: torch.Tensor, nnei: int, + extra_nlist_sort: bool = False, ): n_nf, n_nloc, n_nnei = nlist.shape # nf x nall x 3 @@ -410,7 +424,9 @@ def _format_nlist( ], dim=-1, ) - elif n_nnei > nnei: + + if n_nnei > nnei or extra_nlist_sort: + n_nf, n_nloc, n_nnei = nlist.shape m_real_nei = nlist >= 0 nlist = torch.where(m_real_nei, nlist, 0) # nf x nloc x 3 @@ -427,7 +443,7 @@ def _format_nlist( nlist = torch.gather(nlist, 2, nlist_mapping) nlist = torch.where(rr > rcut, -1, nlist) nlist = nlist[..., :nnei] - else: # n_nnei == nnei: + else: # not extra_nlist_sort and n_nnei <= nnei: pass # great! assert nlist.shape[-1] == nnei return nlist @@ -533,4 +549,8 @@ def mixed_types(self) -> bool: """ return self.atomic_model.mixed_types() + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the model needs sorted nlist when using `forward_lower`.""" + return self.atomic_model.need_sorted_nlist_for_lower() + return CM diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 403058aa47..a7c97cdd39 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -68,6 +68,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) if self.get_fitting_net() is not None: model_predict = {} diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index df2f48e2e4..2fdef4fb8a 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -305,6 +305,10 @@ def has_spin(self) -> bool: """Returns whether it has spin input and output.""" return True + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the model needs sorted nlist when using `forward_lower`.""" + return self.backbone_model.need_sorted_nlist_for_lower() + def __getattr__(self, name): """Get attribute from the wrapped model.""" if ( @@ -418,6 +422,7 @@ def forward_common_lower( fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, do_atomic_virial: bool = False, + extra_nlist_sort: bool = False, ): nframes, nloc = nlist.shape[:2] ( @@ -436,6 +441,7 @@ def forward_common_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + extra_nlist_sort=extra_nlist_sort, ) model_output_type = self.backbone_model.model_output_type() if "mask" in model_output_type: @@ -546,6 +552,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + extra_nlist_sort=self.backbone_model.need_sorted_nlist_for_lower(), ) model_predict = {} model_predict["atom_energy"] = model_ret["energy"] diff --git a/deepmd/pt/model/model/transform_output.py b/deepmd/pt/model/model/transform_output.py index e8a99fb5fc..204d656482 100644 --- a/deepmd/pt/model/model/transform_output.py +++ b/deepmd/pt/model/model/transform_output.py @@ -33,15 +33,27 @@ def atomic_virial_corr( faked_grad = torch.ones_like(sumce0) lst = torch.jit.annotate(List[Optional[torch.Tensor]], [faked_grad]) extended_virial_corr0 = torch.autograd.grad( - [sumce0], [extended_coord], grad_outputs=lst, create_graph=False, retain_graph=True, + [sumce0], + [extended_coord], + grad_outputs=lst, + create_graph=False, + retain_graph=True, )[0] assert extended_virial_corr0 is not None extended_virial_corr1 = torch.autograd.grad( - [sumce1], [extended_coord], grad_outputs=lst, create_graph=False, retain_graph=True, + [sumce1], + [extended_coord], + grad_outputs=lst, + create_graph=False, + retain_graph=True, )[0] assert extended_virial_corr1 is not None extended_virial_corr2 = torch.autograd.grad( - [sumce2], [extended_coord], grad_outputs=lst, create_graph=False, retain_graph=True, + [sumce2], + [extended_coord], + grad_outputs=lst, + create_graph=False, + retain_graph=True, )[0] assert extended_virial_corr2 is not None extended_virial_corr = torch.concat( @@ -66,7 +78,11 @@ def task_deriv_one( faked_grad = torch.ones_like(energy) lst = torch.jit.annotate(List[Optional[torch.Tensor]], [faked_grad]) extended_force = torch.autograd.grad( - [energy], [extended_coord], grad_outputs=lst, create_graph=create_graph, retain_graph=True, + [energy], + [extended_coord], + grad_outputs=lst, + create_graph=create_graph, + retain_graph=True, )[0] assert extended_force is not None extended_force = -extended_force diff --git a/source/lmp/tests/run_mpi_pair_deepmd.py b/source/lmp/tests/run_mpi_pair_deepmd.py index b27774ce11..0c4291ab3a 100644 --- a/source/lmp/tests/run_mpi_pair_deepmd.py +++ b/source/lmp/tests/run_mpi_pair_deepmd.py @@ -54,8 +54,8 @@ ) lammps.pair_coeff("* *") lammps.run(0) -pe = lammps.eval("pe") if rank == 0: + pe = lammps.eval("pe") arr = [pe] np.savetxt(output, np.array(arr)) MPI.Finalize()