diff --git a/deepmd/pt/model/model/linear_model.py b/deepmd/pt/model/model/linear_model.py index 326b413af1..0bd9ffba5b 100644 --- a/deepmd/pt/model/model/linear_model.py +++ b/deepmd/pt/model/model/linear_model.py @@ -10,22 +10,22 @@ from deepmd.model_format import ( FittingOutputDef, ) - +from deepmd.pt.utils.nlist import ( + build_multiple_neighbor_list, +) from .atomic_model import ( AtomicModel, ) +from .dp_atomic_model import ( + DPAtomicModel, +) from .model import ( BaseModel, ) from .pair_tab import ( PairTabModel, ) -from .dp_atomic_model import ( - DPAtomicModel, -) - -from deepmd.pt.utils.nlist import build_multiple_neighbor_list class LinearModel(BaseModel, AtomicModel): @@ -46,11 +46,13 @@ def __init__( self.models = models self.dp_model = models[0] self.zbl_model = models[1] - assert isinstance(self.zbl_model, PairTabModel) and isinstance(self.dp_model, DPAtomicModel), "The provided models are not in the correct order `DPAtomicModel` + `PairTabModel`." + assert ( + isinstance(self.zbl_model, PairTabModel) + and isinstance(self.dp_model, DPAtomicModel) + ), "The provided models are not in the correct order `DPAtomicModel` + `PairTabModel`." self.rcut = self.get_rcut() self.sel = self.get_sel() - def get_fitting_output_def(self) -> FittingOutputDef: """Get the output def of the fitting net.""" return ( @@ -62,18 +64,25 @@ def get_fitting_output_def(self) -> FittingOutputDef: def get_rcut(self) -> float: """Get the cut-off radius.""" return self.get_rcuts()[-1] - + def get_rcuts(self) -> float: """Get the cut-off radius for each individual models in ascending order.""" return sorted([model.get_rcut() for model in self.models]) def get_sel(self) -> int: - """Get the neighbor selection.""" + """Get the neighbor selection.""" return self.get_sels()[-1] - + def get_sels(self) -> int: - """Get the neighbor selection for each individual models in ascending order.""" - return sorted([sum(model.get_sel()) if isinstance(model.get_sel(), list) else model.get_sel() for model in self.models]) + """Get the neighbor selection for each individual models in ascending order.""" + return sorted( + [ + sum(model.get_sel()) + if isinstance(model.get_sel(), list) + else model.get_sel() + for model in self.models + ] + ) def distinguish_types(self) -> bool: """If distinguish different types by sorting.""" @@ -84,10 +93,9 @@ def forward_atomic( extended_coord, extended_atype, nlist, - ra: float, - rb: float, + ra: float, + rb: float, alpha: Optional[float] = 0.1, - mapping: Optional[torch.Tensor] = None, ) -> Dict[str, torch.Tensor]: """Return atomic prediction. @@ -96,7 +104,6 @@ def forward_atomic( the weight is calculated based on this paper: Appl. Phys. Lett. 114, 244101 (2019); https://doi.org/10.1063/1.5098061 - Parameters ---------- extended_coord @@ -107,11 +114,11 @@ def forward_atomic( neighbor list. nf x nloc x nsel mapping mapps the extended indices to local indices - ra : float - inclusive lower boundary of the range in which the ZBL potential and the deep potential are interpolated. - rb : float + ra : float + inclusive lower boundary of the range in which the ZBL potential and the deep potential are interpolated. + rb : float exclusive upper boundary of the range in which the ZBL potential and the deep potential are interpolated. - alpha : float + alpha : float a tunable scale of the distances between atoms. Returns @@ -119,38 +126,51 @@ def forward_atomic( result_dict the result dict, defined by the fitting net output def. """ - # the DPAtomicModel sel is always a List or Union[List, int]? - nlists = build_multiple_neighbor_list(extended_coord, nlist, self.get_rcuts(), self.get_sels()) + nlists = build_multiple_neighbor_list( + extended_coord, nlist, self.get_rcuts(), self.get_sels() + ) zbl_nlist = nlists[str(self.zbl_model.rcut) + "_" + str(self.zbl_model.sel)] - dp_nlist = nlists[str(self.dp_model.rcut) + "_" + str(self.dp_model.sel)] # need to handle sel dtype. + dp_nlist = nlists[ + str(self.dp_model.rcut) + "_" + str(self.dp_model.sel) + ] # need to handle sel dtype. zbl_nframe, zbl_nloc, zbl_nnei = zbl_nlist.shape dp_nframe, dp_nloc, dp_nnei = dp_nlist.shape - zbl_atype = extended_atype[:, :zbl_nloc] # nframe, nloc should all be the same, only difference is nnei based on rcut and nlist. + zbl_atype = extended_atype[ + :, :zbl_nloc + ] # nframe, nloc should all be the same, only difference is nnei based on rcut and nlist. dp_atype = extended_atype[:, :dp_nloc] - + # which rr should I use? this rr should be (nfrmaes, nloc, nnei) zbl_weight = self._compute_weight(rr, ra, rb) - - dp_energy = self.dp_model.forward_atomic(extended_coord, dp_atype, nlist)["energy"] - zbl_energy = self.zbl_model.forward_atomic(extended_coord, zbl_atype, nlist)["energy"] - fit_ret = zbl_weight * zbl_energy + (1 - zbl_weight) * dp_energy # (nframes, nloc) + + dp_energy = self.dp_model.forward_atomic(extended_coord, dp_atype, nlist)[ + "energy" + ] + zbl_energy = self.zbl_model.forward_atomic(extended_coord, zbl_atype, nlist)[ + "energy" + ] + fit_ret = ( + zbl_weight * zbl_energy + (1 - zbl_weight) * dp_energy + ) # (nframes, nloc) return fit_ret - def _compute_weight(self, rr: torch.Tensor, ra: float, rb: float, alpha: Optional[float] = 0.1) -> torch.Tensor: + def _compute_weight( + self, rr: torch.Tensor, ra: float, rb: float, alpha: Optional[float] = 0.1 + ) -> torch.Tensor: """ZBL weight. Parameters ---------- - rr : torch.Tensor + rr : torch.Tensor pairwise distance between atom i and atom j, (nframes, nloc, nnei). - ra : float - inclusive lower boundary of the range in which the ZBL potential and the deep potential are interpolated. - rb : float + ra : float + inclusive lower boundary of the range in which the ZBL potential and the deep potential are interpolated. + rb : float exclusive upper boundary of the range in which the ZBL potential and the deep potential are interpolated. - alpha : float + alpha : float a tunable scale of the distances between atoms. Returns @@ -158,14 +178,17 @@ def _compute_weight(self, rr: torch.Tensor, ra: float, rb: float, alpha: Optiona torch.Tensor the atomic ZBL weight for interpolation. (nframes, nloc) """ - assert rb > ra, "The upper boundary `rb` must be greater than the lower boundary `ra`." - - sigma = torch.sum(rr*torch.exp(-rr/alpha),dim=-1)/torch.sum(torch.exp(-rr/alpha),dim=-1) # (nframes, nloc) - u = (sigma - ra)/(rb - ra) + assert ( + rb > ra + ), "The upper boundary `rb` must be greater than the lower boundary `ra`." + + sigma = torch.sum(rr * torch.exp(-rr / alpha), dim=-1) / torch.sum( + torch.exp(-rr / alpha), dim=-1 + ) # (nframes, nloc) + u = (sigma - ra) / (rb - ra) if sigma < ra: return torch.ones_like(u) elif ra <= sigma < rb: return -6 * u**5 + 15 * u**4 - 10 * u**3 + 1 else: return torch.zeros_like(u) -