From f684be84a89fb1adddc4aa5f8f59e0c1f63f33e9 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Fri, 1 Mar 2024 02:17:11 -0500 Subject: [PATCH] feat: atom_ener in energy fitting (#3370) Also, fix the TF serialization issue (it tried to store a tensor instead of a NumPy array). --------- Signed-off-by: Jinzhe Zeng Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> --- deepmd/dpmodel/fitting/general_fitting.py | 33 +++++++++++++++++- deepmd/dpmodel/fitting/invar_fitting.py | 5 +-- deepmd/pt/model/task/ener.py | 3 ++ deepmd/pt/model/task/fitting.py | 36 +++++++++++++++++++- deepmd/tf/fit/ener.py | 2 +- source/tests/consistent/fitting/test_ener.py | 11 ++++++ 6 files changed, 85 insertions(+), 5 deletions(-) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 752a550849..5b4ca195b5 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -73,7 +73,10 @@ class GeneralFitting(NativeOP, BaseFitting): different fitting nets for different atom types. exclude_types: List[int] Atomic contributions of the excluded atom types are set zero. - + remove_vaccum_contribution: List[bool], optional + Remove vaccum contribution before the bias is added. The list assigned each + type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same + length as `ntypes` signaling if or not removing the vaccum contribution for the atom types in the list. """ def __init__( @@ -95,6 +98,7 @@ def __init__( spin: Any = None, mixed_types: bool = True, exclude_types: List[int] = [], + remove_vaccum_contribution: Optional[List[bool]] = None, ): self.var_name = var_name self.ntypes = ntypes @@ -119,6 +123,7 @@ def __init__( self.exclude_types = exclude_types if self.spin is not None: raise NotImplementedError("spin is not supported") + self.remove_vaccum_contribution = remove_vaccum_contribution self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) @@ -298,6 +303,14 @@ def _call_common( "which is not consistent with {self.dim_descrpt}." ) xx = descriptor + if self.remove_vaccum_contribution is not None: + # TODO: Idealy, the input for vaccum should be computed; + # we consider it as always zero for convenience. + # Needs a compute_input_stats for vaccum passed from the + # descriptor. + xx_zeros = np.zeros_like(xx) + else: + xx_zeros = None # check fparam dim, concate to input descriptor if self.numb_fparam > 0: assert fparam is not None, "fparam should not be None" @@ -312,6 +325,11 @@ def _call_common( [xx, fparam], axis=-1, ) + if xx_zeros is not None: + xx_zeros = np.concatenate( + [xx_zeros, fparam], + axis=-1, + ) # check aparam dim, concate to input descriptor if self.numb_aparam > 0: assert aparam is not None, "aparam should not be None" @@ -326,6 +344,11 @@ def _call_common( [xx, aparam], axis=-1, ) + if xx_zeros is not None: + xx_zeros = np.concatenate( + [xx_zeros, aparam], + axis=-1, + ) # calcualte the prediction if not self.mixed_types: @@ -335,11 +358,19 @@ def _call_common( (atype == type_i).reshape([nf, nloc, 1]), [1, 1, net_dim_out] ) atom_property = self.nets[(type_i,)](xx) + if self.remove_vaccum_contribution is not None and not ( + len(self.remove_vaccum_contribution) > type_i + and not self.remove_vaccum_contribution[type_i] + ): + assert xx_zeros is not None + atom_property -= self.nets[(type_i,)](xx_zeros) atom_property = atom_property + self.bias_atom_e[type_i] atom_property = atom_property * mask outs = outs + atom_property # Shape is [nframes, natoms[0], 1] else: outs = self.nets[()](xx) + self.bias_atom_e[atype] + if xx_zeros is not None: + outs -= self.nets[()](xx_zeros) # nf x nloc exclude_mask = self.emask.build_type_exclude_mask(atype) # nf x nloc x nod diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index 769dc45042..fd556ff074 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -136,8 +136,6 @@ def __init__( raise NotImplementedError("use_aparam_as_mask is not implemented") if layer_name is not None: raise NotImplementedError("layer_name is not implemented") - if atom_ener is not None and atom_ener != []: - raise NotImplementedError("atom_ener is not implemented") self.dim_out = dim_out self.atom_ener = atom_ener @@ -159,6 +157,9 @@ def __init__( spin=spin, mixed_types=mixed_types, exclude_types=exclude_types, + remove_vaccum_contribution=None + if atom_ener is None or len([x for x in atom_ener if x is not None]) == 0 + else [x is not None for x in atom_ener], ) def serialize(self) -> dict: diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 29ed5acaad..00bf049b97 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -126,6 +126,9 @@ def __init__( rcond=rcond, seed=seed, exclude_types=exclude_types, + remove_vaccum_contribution=None + if atom_ener is None or len([x for x in atom_ener if x is not None]) == 0 + else [x is not None for x in atom_ener], **kwargs, ) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 47535580db..c41d445f66 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -239,7 +239,10 @@ class GeneralFitting(Fitting): Random seed. exclude_types: List[int] Atomic contributions of the excluded atom types are set zero. - + remove_vaccum_contribution: List[bool], optional + Remove vaccum contribution before the bias is added. The list assigned each + type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same + length as `ntypes` signaling if or not removing the vaccum contribution for the atom types in the list. """ def __init__( @@ -258,6 +261,7 @@ def __init__( rcond: Optional[float] = None, seed: Optional[int] = None, exclude_types: List[int] = [], + remove_vaccum_contribution: Optional[List[bool]] = None, **kwargs, ): super().__init__() @@ -275,6 +279,7 @@ def __init__( self.rcond = rcond # order matters, should be place after the assignment of ntypes self.reinit_exclude(exclude_types) + self.remove_vaccum_contribution = remove_vaccum_contribution net_dim_out = self._net_out_dim() # init constants @@ -479,6 +484,14 @@ def _forward_common( aparam: Optional[torch.Tensor] = None, ): xx = descriptor + if self.remove_vaccum_contribution is not None: + # TODO: Idealy, the input for vaccum should be computed; + # we consider it as always zero for convenience. + # Needs a compute_input_stats for vaccum passed from the + # descriptor. + xx_zeros = torch.zeros_like(xx) + else: + xx_zeros = None nf, nloc, nd = xx.shape net_dim_out = self._net_out_dim() @@ -507,6 +520,11 @@ def _forward_common( [xx, fparam], dim=-1, ) + if xx_zeros is not None: + xx_zeros = torch.cat( + [xx_zeros, fparam], + dim=-1, + ) # check aparam dim, concate to input descriptor if self.numb_aparam > 0: assert aparam is not None, "aparam should not be None" @@ -526,6 +544,11 @@ def _forward_common( [xx, aparam], dim=-1, ) + if xx_zeros is not None: + xx_zeros = torch.cat( + [xx_zeros, aparam], + dim=-1, + ) outs = torch.zeros( (nf, nloc, net_dim_out), @@ -534,6 +557,7 @@ def _forward_common( ) # jit assertion if self.old_impl: assert self.filter_layers_old is not None + assert xx_zeros is None if self.mixed_types: atom_property = self.filter_layers_old[0](xx) + self.bias_atom_e[atype] outs = outs + atom_property # Shape is [nframes, natoms[0], 1] @@ -549,6 +573,8 @@ def _forward_common( atom_property = ( self.filter_layers.networks[0](xx) + self.bias_atom_e[atype] ) + if xx_zeros is not None: + atom_property -= self.filter_layers.networks[0](xx_zeros) outs = ( outs + atom_property ) # Shape is [nframes, natoms[0], net_dim_out] @@ -557,6 +583,14 @@ def _forward_common( mask = (atype == type_i).unsqueeze(-1) mask = torch.tile(mask, (1, 1, net_dim_out)) atom_property = ll(xx) + if xx_zeros is not None: + # must assert, otherwise jit is not happy + assert self.remove_vaccum_contribution is not None + if not ( + len(self.remove_vaccum_contribution) > type_i + and not self.remove_vaccum_contribution[type_i] + ): + atom_property -= ll(xx_zeros) atom_property = atom_property + self.bias_atom_e[type_i] atom_property = atom_property * mask outs = ( diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index a842df50bd..d605fbb0aa 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -1003,7 +1003,7 @@ def serialize(self, suffix: str = "") -> dict: "rcond": self.rcond, "tot_ener_zero": self.tot_ener_zero, "trainable": self.trainable, - "atom_ener": self.atom_ener, + "atom_ener": self.atom_ener_v, "activation_function": self.activation_function_name, "precision": self.fitting_precision.name, "layer_name": self.layer_name, diff --git a/source/tests/consistent/fitting/test_ener.py b/source/tests/consistent/fitting/test_ener.py index 994d967bc8..a22bcdb65f 100644 --- a/source/tests/consistent/fitting/test_ener.py +++ b/source/tests/consistent/fitting/test_ener.py @@ -43,6 +43,7 @@ ("float64", "float32"), # precision (True, False), # mixed_types (0, 1), # numb_fparam + ([], [-12345.6, None]), # atom_ener ) class TestEner(CommonTest, FittingTest, unittest.TestCase): @property @@ -52,6 +53,7 @@ def data(self) -> dict: precision, mixed_types, numb_fparam, + atom_ener, ) = self.param return { "neuron": [5, 5, 5], @@ -59,6 +61,7 @@ def data(self) -> dict: "precision": precision, "numb_fparam": numb_fparam, "seed": 20240217, + "atom_ener": atom_ener, } @property @@ -68,6 +71,7 @@ def skip_tf(self) -> bool: precision, mixed_types, numb_fparam, + atom_ener, ) = self.param # TODO: mixed_types return mixed_types or CommonTest.skip_pt @@ -79,6 +83,7 @@ def skip_pt(self) -> bool: precision, mixed_types, numb_fparam, + atom_ener, ) = self.param return CommonTest.skip_pt @@ -105,6 +110,7 @@ def addtional_data(self) -> dict: precision, mixed_types, numb_fparam, + atom_ener, ) = self.param return { "ntypes": self.ntypes, @@ -118,6 +124,7 @@ def build_tf(self, obj: Any, suffix: str) -> Tuple[list, dict]: precision, mixed_types, numb_fparam, + atom_ener, ) = self.param return self.build_tf_fitting( obj, @@ -134,6 +141,7 @@ def eval_pt(self, pt_obj: Any) -> Any: precision, mixed_types, numb_fparam, + atom_ener, ) = self.param return ( pt_obj( @@ -154,6 +162,7 @@ def eval_dp(self, dp_obj: Any) -> Any: precision, mixed_types, numb_fparam, + atom_ener, ) = self.param return dp_obj( self.inputs, @@ -175,6 +184,7 @@ def rtol(self) -> float: precision, mixed_types, numb_fparam, + atom_ener, ) = self.param if precision == "float64": return 1e-10 @@ -191,6 +201,7 @@ def atol(self) -> float: precision, mixed_types, numb_fparam, + atom_ener, ) = self.param if precision == "float64": return 1e-10