diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index eb95886598..1158d278e5 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -88,6 +88,10 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.type_map + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return False + def reinit_atom_exclude( self, exclude_types: list[int] = [], diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 2fa072cc78..8bae07dcad 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -233,6 +233,10 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.fitting.get_dim_aparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.fitting.has_default_fparam() + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py index fcaea43338..f49c148377 100644 --- a/deepmd/dpmodel/fitting/dipole_fitting.py +++ b/deepmd/dpmodel/fitting/dipole_fitting.py @@ -84,6 +84,9 @@ class DipoleFitting(GeneralFitting): Only reducible variable are differentiable. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -110,6 +113,7 @@ def __init__( c_differentiable: bool = True, type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, + default_fparam: Optional[list[float]] = None, ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") @@ -144,6 +148,7 @@ def __init__( exclude_types=exclude_types, type_map=type_map, seed=seed, + default_fparam=default_fparam, ) def _net_out_dim(self): @@ -161,7 +166,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) var_name = data.pop("var_name", None) assert var_name == "dipole" return super().deserialize(data) diff --git a/deepmd/dpmodel/fitting/dos_fitting.py b/deepmd/dpmodel/fitting/dos_fitting.py index 2f6df77eac..4bc34b8abf 100644 --- a/deepmd/dpmodel/fitting/dos_fitting.py +++ b/deepmd/dpmodel/fitting/dos_fitting.py @@ -46,6 +46,7 @@ def __init__( exclude_types: list[int] = [], type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, + default_fparam: Optional[list] = None, ) -> None: if bias_dos is not None: self.bias_dos = bias_dos @@ -70,12 +71,13 @@ def __init__( exclude_types=exclude_types, type_map=type_map, seed=seed, + default_fparam=default_fparam, ) @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data["numb_dos"] = data.pop("dim_out") data.pop("tot_ener_zero", None) data.pop("var_name", None) diff --git a/deepmd/dpmodel/fitting/ener_fitting.py b/deepmd/dpmodel/fitting/ener_fitting.py index 6435b6468f..794c074485 100644 --- a/deepmd/dpmodel/fitting/ener_fitting.py +++ b/deepmd/dpmodel/fitting/ener_fitting.py @@ -46,6 +46,7 @@ def __init__( exclude_types: list[int] = [], type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, + default_fparam: Optional[list] = None, ) -> None: super().__init__( var_name="energy", @@ -70,12 +71,13 @@ def __init__( exclude_types=exclude_types, type_map=type_map, seed=seed, + default_fparam=default_fparam, ) @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("var_name") data.pop("dim_out") return super().deserialize(data) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 651a2d0a96..868577a46f 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -94,6 +94,9 @@ class GeneralFitting(NativeOP, BaseFitting): A list of strings. Give the name to each type of atoms. seed: Optional[Union[int, list[int]]] Random seed for initializing the network parameters. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -120,6 +123,7 @@ def __init__( remove_vaccum_contribution: Optional[list[bool]] = None, type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, + default_fparam: Optional[list[float]] = None, ) -> None: self.var_name = var_name self.ntypes = ntypes @@ -129,6 +133,7 @@ def __init__( self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam self.dim_case_embd = dim_case_embd + self.default_fparam = default_fparam self.rcond = rcond self.tot_ener_zero = tot_ener_zero self.trainable = trainable @@ -177,6 +182,15 @@ def __init__( self.case_embd = np.zeros(self.dim_case_embd, dtype=self.prec) else: self.case_embd = None + + if self.default_fparam is not None: + if self.numb_fparam > 0: + assert len(self.default_fparam) == self.numb_fparam, ( + "default_fparam length mismatch!" + ) + self.default_fparam_tensor = np.array(self.default_fparam, dtype=self.prec) + else: + self.default_fparam_tensor = None # init networks in_dim = ( self.dim_descrpt @@ -217,6 +231,10 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.numb_aparam + def has_default_fparam(self) -> bool: + """Check if the fitting has default frame parameters.""" + return self.default_fparam is not None + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. @@ -274,6 +292,8 @@ def __setitem__(self, key, value) -> None: self.case_embd = value elif key in ["scale"]: self.scale = value + elif key in ["default_fparam_tensor"]: + self.default_fparam_tensor = value else: raise KeyError(key) @@ -292,6 +312,8 @@ def __getitem__(self, key): return self.case_embd elif key in ["scale"]: return self.scale + elif key in ["default_fparam_tensor"]: + return self.default_fparam_tensor else: raise KeyError(key) @@ -306,7 +328,7 @@ def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", - "@version": 3, + "@version": 4, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -315,6 +337,7 @@ def serialize(self) -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "rcond": self.rcond, "activation_function": self.activation_function, "precision": self.precision, @@ -403,6 +426,14 @@ def _call_common( xx_zeros = xp.zeros_like(xx) else: xx_zeros = None + + if self.numb_fparam > 0 and fparam is None: + # use default fparam + assert self.default_fparam_tensor is not None + fparam = xp.tile( + xp.reshape(self.default_fparam_tensor, (1, self.numb_fparam)), (nf, 1) + ) + # check fparam dim, concate to input descriptor if self.numb_fparam > 0: assert fparam is not None, "fparam should not be None" diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index b5d3a02d86..9e97eac22a 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -110,6 +110,9 @@ class InvarFitting(GeneralFitting): Atomic contributions of the excluded atom types are set zero. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ @@ -138,6 +141,7 @@ def __init__( exclude_types: list[int] = [], type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, + default_fparam: Optional[list[float]] = None, ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") @@ -173,6 +177,7 @@ def __init__( else [x is not None for x in atom_ener], type_map=type_map, seed=seed, + default_fparam=default_fparam, ) def serialize(self) -> dict: @@ -185,7 +190,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) return super().deserialize(data) def _net_out_dim(self): diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index bfc337a177..cc20e4c932 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -90,6 +90,9 @@ class PolarFitting(GeneralFitting): Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -117,6 +120,7 @@ def __init__( shift_diag: bool = True, type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, + default_fparam: Optional[list[float]] = None, ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") @@ -164,6 +168,7 @@ def __init__( exclude_types=exclude_types, type_map=type_map, seed=seed, + default_fparam=default_fparam, ) def _net_out_dim(self): @@ -189,7 +194,7 @@ def __getitem__(self, key): def serialize(self) -> dict: data = super().serialize() data["type"] = "polar" - data["@version"] = 4 + data["@version"] = 5 data["embedding_width"] = self.embedding_width data["fit_diag"] = self.fit_diag data["shift_diag"] = self.shift_diag @@ -200,7 +205,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) var_name = data.pop("var_name", None) assert var_name == "polar" return super().deserialize(data) diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py index 59b685d391..dbd415bde1 100644 --- a/deepmd/dpmodel/fitting/property_fitting.py +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -65,6 +65,9 @@ class PropertyFittingNet(InvarFitting): Atomic contributions of the excluded atom types are set zero. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -87,6 +90,7 @@ def __init__( mixed_types: bool = True, exclude_types: list[int] = [], type_map: Optional[list[str]] = None, + default_fparam: Optional[list] = None, # not used seed: Optional[int] = None, ) -> None: @@ -110,6 +114,7 @@ def __init__( mixed_types=mixed_types, exclude_types=exclude_types, type_map=type_map, + default_fparam=default_fparam, ) def output_def(self) -> FittingOutputDef: @@ -129,7 +134,7 @@ def output_def(self) -> FittingOutputDef: @classmethod def deserialize(cls, data: dict) -> "PropertyFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version"), 4, 1) + check_version_compatibility(data.pop("@version"), 5, 1) data.pop("dim_out") data["property_name"] = data.pop("var_name") data.pop("tot_ener_zero") @@ -149,6 +154,6 @@ def serialize(self) -> dict: "task_dim": self.task_dim, "intensive": self.intensive, } - dd["@version"] = 4 + dd["@version"] = 5 return dd diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index 4b30d97c29..241a93fc90 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -120,6 +120,10 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" return self.dp.get_dim_aparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.dp.has_default_fparam() + @property def model_type(self) -> type["DeepEvalWrapper"]: """The the evaluator of the model type.""" diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 7f07181087..7b7da0f2c6 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -564,6 +564,10 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.atomic_model.get_dim_aparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.atomic_model.has_default_fparam() + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index ed3ceda8c2..88b017b628 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -343,7 +343,11 @@ def test_ener( data.add("atom_ener", 1, atomic=True, must=True, high_prec=False) if dp.get_dim_fparam() > 0: data.add( - "fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False + "fparam", + dp.get_dim_fparam(), + atomic=False, + must=not dp.has_default_fparam(), + high_prec=False, ) if dp.get_dim_aparam() > 0: data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False) @@ -380,7 +384,7 @@ def test_ener( atype = test_data["type"][:numb_test].reshape([numb_test, -1]) else: atype = test_data["type"][0] - if dp.get_dim_fparam() > 0: + if dp.get_dim_fparam() > 0 and test_data["find_fparam"] != 0.0: fparam = test_data["fparam"][:numb_test] else: fparam = None diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 9ad2576697..5f29f08330 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -162,6 +162,10 @@ def get_type_map(self) -> list[str]: def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this DP.""" + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return False + @abstractmethod def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" @@ -432,6 +436,10 @@ def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this DP.""" return self.deep_eval.get_dim_fparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.deep_eval.has_default_fparam() + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" return self.deep_eval.get_dim_aparam() diff --git a/deepmd/jax/fitting/fitting.py b/deepmd/jax/fitting/fitting.py index d62681490c..e69bded640 100644 --- a/deepmd/jax/fitting/fitting.py +++ b/deepmd/jax/fitting/fitting.py @@ -35,6 +35,7 @@ def setattr_for_general_fitting(name: str, value: Any) -> Any: "fparam_inv_std", "aparam_avg", "aparam_inv_std", + "default_fparam_tensor", }: value = to_jax_array(value) if value is not None: diff --git a/deepmd/pd/model/task/ener.py b/deepmd/pd/model/task/ener.py index 789ef75066..738990b2d8 100644 --- a/deepmd/pd/model/task/ener.py +++ b/deepmd/pd/model/task/ener.py @@ -72,7 +72,7 @@ def __init__( @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = copy.deepcopy(data) - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("var_name") data.pop("dim_out") return super().deserialize(data) diff --git a/deepmd/pd/model/task/fitting.py b/deepmd/pd/model/task/fitting.py index 9dede6a897..953ec5bf0e 100644 --- a/deepmd/pd/model/task/fitting.py +++ b/deepmd/pd/model/task/fitting.py @@ -183,6 +183,10 @@ class GeneralFitting(Fitting): Number of frame parameters. numb_aparam : int Number of atomic parameters. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. + This parameter is not supported in PaddlePaddle. dim_case_embd : int Dimension of case specific embedding. activation_function : str @@ -233,6 +237,7 @@ def __init__( remove_vaccum_contribution: Optional[list[bool]] = None, type_map: Optional[list[str]] = None, use_aparam_as_mask: bool = False, + default_fparam: Optional[list[float]] = None, **kwargs, ) -> None: super().__init__() @@ -245,6 +250,7 @@ def __init__( self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam self.dim_case_embd = dim_case_embd + self.default_fparam = default_fparam self.activation_function = activation_function self.precision = precision self.prec = PRECISION_DICT[self.precision] @@ -372,7 +378,7 @@ def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", - "@version": 3, + "@version": 4, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -381,6 +387,7 @@ def serialize(self) -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "activation_function": self.activation_function, "precision": self.precision, "mixed_types": self.mixed_types, diff --git a/deepmd/pd/model/task/invar_fitting.py b/deepmd/pd/model/task/invar_fitting.py index b92c862dc8..176acdeb20 100644 --- a/deepmd/pd/model/task/invar_fitting.py +++ b/deepmd/pd/model/task/invar_fitting.py @@ -147,7 +147,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = copy.deepcopy(data) - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) return super().deserialize(data) def output_def(self) -> FittingOutputDef: diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index e2cd4b9e32..f3e52cdac0 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -218,6 +218,14 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" return self.dp.model["Default"].get_dim_aparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + try: + return self.dp.model["Default"].has_default_fparam() + except AttributeError: + # for compatibility with old models + return False + def get_intensive(self) -> bool: return self.dp.model["Default"].get_intensive() diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index a18834e40b..b8ba0a1981 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -135,6 +135,10 @@ def get_intensive(self) -> bool: """Whether the fitting property is intensive.""" return False + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return False + def reinit_atom_exclude( self, exclude_types: list[int] = [], diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 1a34eb986a..5b7d96560f 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -338,6 +338,10 @@ def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.fitting_net.get_dim_fparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.fitting_net.has_default_fparam() + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.fitting_net.get_dim_aparam() diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index e282ebd83a..53d32977b0 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -525,6 +525,11 @@ def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.atomic_model.get_dim_fparam() + @torch.jit.export + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.atomic_model.has_default_fparam() + @torch.jit.export def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index b2d8b598b9..b6a1477f7a 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -73,6 +73,9 @@ class DipoleFittingNet(GeneralFitting): Only reducible variable are differentiable. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -94,6 +97,7 @@ def __init__( r_differentiable: bool = True, c_differentiable: bool = True, type_map: Optional[list[str]] = None, + default_fparam: Optional[list] = None, **kwargs: Any, ) -> None: self.embedding_width = embedding_width @@ -115,6 +119,7 @@ def __init__( seed=seed, exclude_types=exclude_types, type_map=type_map, + default_fparam=default_fparam, **kwargs, ) @@ -133,7 +138,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("var_name", None) return super().deserialize(data) diff --git a/deepmd/pt/model/task/dos.py b/deepmd/pt/model/task/dos.py index 568ef81c92..afbed5f748 100644 --- a/deepmd/pt/model/task/dos.py +++ b/deepmd/pt/model/task/dos.py @@ -57,6 +57,7 @@ def __init__( exclude_types: list[int] = [], mixed_types: bool = True, type_map: Optional[list[str]] = None, + default_fparam: Optional[list] = None, ) -> None: if bias_dos is not None: self.bias_dos = bias_dos @@ -83,6 +84,7 @@ def __init__( exclude_types=exclude_types, trainable=trainable, type_map=type_map, + default_fparam=default_fparam, ) def output_def(self) -> FittingOutputDef: @@ -101,7 +103,7 @@ def output_def(self) -> FittingOutputDef: @classmethod def deserialize(cls, data: dict) -> "DOSFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("@class", None) data.pop("var_name", None) data.pop("tot_ener_zero", None) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 9027bfe288..af288bec10 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -57,6 +57,7 @@ def __init__( mixed_types: bool = True, seed: Optional[Union[int, list[int]]] = None, type_map: Optional[list[str]] = None, + default_fparam: Optional[list] = None, **kwargs: Any, ) -> None: super().__init__( @@ -75,13 +76,14 @@ def __init__( mixed_types=mixed_types, seed=seed, type_map=type_map, + default_fparam=default_fparam, **kwargs, ) @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("var_name") data.pop("dim_out") return super().deserialize(data) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 841c494e88..7ad72ba4b4 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -207,6 +207,9 @@ class GeneralFitting(Fitting): A list of strings. Give the name to each type of atoms. use_aparam_as_mask: bool If True, the aparam will not be used in fitting net for embedding. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -230,6 +233,7 @@ def __init__( remove_vaccum_contribution: Optional[list[bool]] = None, type_map: Optional[list[str]] = None, use_aparam_as_mask: bool = False, + default_fparam: Optional[list[float]] = None, **kwargs: Any, ) -> None: super().__init__() @@ -241,6 +245,7 @@ def __init__( self.resnet_dt = resnet_dt self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam + self.default_fparam = default_fparam self.dim_case_embd = dim_case_embd self.activation_function = activation_function self.precision = precision @@ -302,6 +307,20 @@ def __init__( else: self.case_embd = None + if self.default_fparam is not None: + if self.numb_fparam > 0: + assert len(self.default_fparam) == self.numb_fparam, ( + "default_fparam length mismatch!" + ) + self.register_buffer( + "default_fparam_tensor", + torch.tensor( + np.array(self.default_fparam), dtype=self.prec, device=device + ), + ) + else: + self.default_fparam_tensor = None + in_dim = ( self.dim_descrpt + self.numb_fparam @@ -371,7 +390,7 @@ def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", - "@version": 3, + "@version": 4, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -380,6 +399,7 @@ def serialize(self) -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "activation_function": self.activation_function, "precision": self.precision, "mixed_types": self.mixed_types, @@ -423,6 +443,10 @@ def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.numb_fparam + def has_default_fparam(self) -> bool: + """Check if the fitting has default frame parameters.""" + return self.default_fparam is not None + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.numb_aparam @@ -476,6 +500,8 @@ def __setitem__(self, key: str, value: torch.Tensor) -> None: self.case_embd = value elif key in ["scale"]: self.scale = value + elif key in ["default_fparam_tensor"]: + self.default_fparam_tensor = value else: raise KeyError(key) @@ -494,6 +520,8 @@ def __getitem__(self, key: str) -> torch.Tensor: return self.case_embd elif key in ["scale"]: return self.scale + elif key in ["default_fparam_tensor"]: + return self.default_fparam_tensor else: raise KeyError(key) @@ -520,6 +548,13 @@ def _forward_common( ) -> dict[str, torch.Tensor]: # cast the input to internal precsion xx = descriptor.to(self.prec) + nf, nloc, nd = xx.shape + + if self.numb_fparam > 0 and fparam is None: + # use default fparam + assert self.default_fparam_tensor is not None + fparam = torch.tile(self.default_fparam_tensor.unsqueeze(0), [nf, 1]) + fparam = fparam.to(self.prec) if fparam is not None else None aparam = aparam.to(self.prec) if aparam is not None else None @@ -532,7 +567,6 @@ def _forward_common( xx_zeros = torch.zeros_like(xx) else: xx_zeros = None - nf, nloc, nd = xx.shape net_dim_out = self._net_out_dim() if nd != self.dim_descrpt: diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index f7233352f8..4ec3407901 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -81,6 +81,9 @@ class InvarFitting(GeneralFitting): A list of strings. Give the name to each type of atoms. use_aparam_as_mask: bool If True, the aparam will not be used in fitting net for embedding. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -104,6 +107,7 @@ def __init__( atom_ener: Optional[list[Optional[torch.Tensor]]] = None, type_map: Optional[list[str]] = None, use_aparam_as_mask: bool = False, + default_fparam: Optional[list[float]] = None, **kwargs: Any, ) -> None: self.dim_out = dim_out @@ -129,6 +133,7 @@ def __init__( else [x is not None for x in atom_ener], type_map=type_map, use_aparam_as_mask=use_aparam_as_mask, + default_fparam=default_fparam, **kwargs, ) @@ -146,7 +151,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) return super().deserialize(data) def output_def(self) -> FittingOutputDef: diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index 5f8ad3c73a..bf63d9db4b 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -76,7 +76,9 @@ class PolarFittingNet(GeneralFitting): Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. - + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -99,6 +101,7 @@ def __init__( scale: Optional[Union[list[float], float]] = None, shift_diag: bool = True, type_map: Optional[list[str]] = None, + default_fparam: Optional[list] = None, **kwargs: Any, ) -> None: self.embedding_width = embedding_width @@ -140,6 +143,7 @@ def __init__( seed=seed, exclude_types=exclude_types, type_map=type_map, + default_fparam=default_fparam, **kwargs, ) @@ -196,7 +200,7 @@ def change_type_map( def serialize(self) -> dict: data = super().serialize() data["type"] = "polar" - data["@version"] = 4 + data["@version"] = 5 data["embedding_width"] = self.embedding_width data["fit_diag"] = self.fit_diag data["shift_diag"] = self.shift_diag @@ -207,7 +211,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data.pop("var_name", None) return super().deserialize(data) diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 23069ea0da..c2440b7de3 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -92,6 +92,7 @@ def __init__( mixed_types: bool = True, trainable: Union[bool, list[bool]] = True, seed: Optional[int] = None, + default_fparam: Optional[list] = None, **kwargs: Any, ) -> None: self.task_dim = task_dim @@ -112,6 +113,7 @@ def __init__( mixed_types=mixed_types, trainable=trainable, seed=seed, + default_fparam=default_fparam, **kwargs, ) @@ -136,7 +138,7 @@ def get_intensive(self) -> bool: @classmethod def deserialize(cls, data: dict) -> "PropertyFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data.pop("dim_out") data["property_name"] = data.pop("var_name") obj = super().deserialize(data) @@ -151,7 +153,7 @@ def serialize(self) -> dict: "task_dim": self.task_dim, "intensive": self.intensive, } - dd["@version"] = 4 + dd["@version"] = 5 return dd diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index ab98389426..52d2888081 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -1252,7 +1252,8 @@ def get_data( label_dict = {} for item_key in batch_data: if item_key in input_keys: - input_dict[item_key] = batch_data[item_key] + if item_key != "fparam" or batch_data["find_fparam"] != 0.0: + input_dict[item_key] = batch_data[item_key] else: if item_key not in ["sid", "fid"]: label_dict[item_key] = batch_data[item_key] @@ -1338,7 +1339,10 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: if _model.get_dim_fparam() > 0: fparam_requirement_items = [ DataRequirementItem( - "fparam", _model.get_dim_fparam(), atomic=False, must=True + "fparam", + _model.get_dim_fparam(), + atomic=False, + must=not _model.has_default_fparam(), ) ] additional_data_requirement += fparam_requirement_items diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index cbab88de44..7312d95a06 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -62,6 +62,14 @@ def make_stat_input( except StopIteration: iterator = iter(dataloaders[i]) stat_data = next(iterator) + if ( + "find_fparam" in stat_data + and "fparam" in stat_data + and stat_data["find_fparam"] == 0.0 + ): + # for model using default fparam + stat_data.pop("fparam") + stat_data.pop("find_fparam") for dd in stat_data: if stat_data[dd] is None: sys_stat[dd] = None diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index a081f38b17..2142e80f30 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -78,6 +78,9 @@ class DipoleFittingSeA(Fitting): different fitting nets for different atom types. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. trainable : list[bool], Optional If the weights of fitting net are trainable. Suppose that we have :math:`N_l` hidden layers in the fitting net, @@ -101,6 +104,7 @@ def __init__( uniform_seed: bool = False, mixed_types: bool = False, type_map: Optional[list[str]] = None, # to be compat with input + default_fparam: Optional[list[float]] = None, # to be compat with input trainable: Optional[list[bool]] = None, **kwargs, ) -> None: @@ -131,12 +135,15 @@ def __init__( self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam self.dim_case_embd = dim_case_embd + self.default_fparam = default_fparam if numb_fparam > 0: raise ValueError("numb_fparam is not supported in the dipole fitting") if numb_aparam > 0: raise ValueError("numb_aparam is not supported in the dipole fitting") if dim_case_embd > 0: raise ValueError("dim_case_embd is not supported in TensorFlow.") + if default_fparam is not None: + raise ValueError("default_fparam is not supported in TensorFlow.") self.fparam_avg = None self.fparam_std = None self.fparam_inv_std = None @@ -411,7 +418,7 @@ def serialize(self, suffix: str) -> dict: data = { "@class": "Fitting", "type": "dipole", - "@version": 3, + "@version": 4, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, "embedding_width": self.dim_rot_mat_1, @@ -422,6 +429,7 @@ def serialize(self, suffix: str) -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "activation_function": self.activation_function_name, "precision": self.fitting_precision.name, "exclude_types": [] @@ -468,7 +476,7 @@ def deserialize(cls, data: dict, suffix: str): The deserialized model """ data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) exclude_types = data.pop("exclude_types", []) if len(exclude_types) > 0: data["sel_type"] = [ diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index 96e9470692..7c90641153 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -101,6 +101,9 @@ class DOSFitting(Fitting): mixed_types : bool If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. """ @@ -125,6 +128,7 @@ def __init__( use_aparam_as_mask: bool = False, mixed_types: bool = False, type_map: Optional[list[str]] = None, # to be compat with input + default_fparam: Optional[list[float]] = None, # to be compat with input **kwargs, ) -> None: """Constructor.""" @@ -136,8 +140,11 @@ def __init__( self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam self.dim_case_embd = dim_case_embd + self.default_fparam = default_fparam if dim_case_embd > 0: raise ValueError("dim_case_embd is not supported in TensorFlow.") + if default_fparam is not None: + raise ValueError("default_fparam is not supported in TensorFlow.") self.numb_dos = numb_dos @@ -678,7 +685,7 @@ def deserialize(cls, data: dict, suffix: str = ""): The deserialized model """ data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data["numb_dos"] = data.pop("dim_out") fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( @@ -705,7 +712,7 @@ def serialize(self, suffix: str = "") -> dict: data = { "@class": "Fitting", "type": "dos", - "@version": 3, + "@version": 4, "var_name": "dos", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -716,6 +723,7 @@ def serialize(self, suffix: str = "") -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "rcond": self.rcond, "trainable": self.trainable, "activation_function": self.activation_function, diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 2458081a88..547c0eefb1 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -119,6 +119,8 @@ class EnerFitting(Fitting): Number of atomic parameter dim_case_embd Dimension of case specific embedding. + default_fparam + The default frame parameter. This parameter is not supported in TensorFlow. rcond The condition number for the regression of atomic energy. tot_ener_zero @@ -146,6 +148,9 @@ class EnerFitting(Fitting): mixed_types : bool If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. """ @@ -172,6 +177,7 @@ def __init__( spin: Optional[Spin] = None, mixed_types: bool = False, type_map: Optional[list[str]] = None, # to be compat with input + default_fparam: Optional[list[float]] = None, # to be compat with input **kwargs, ) -> None: """Constructor.""" @@ -196,6 +202,9 @@ def __init__( self.dim_case_embd = dim_case_embd if dim_case_embd > 0: raise ValueError("dim_case_embd is not supported in TensorFlow.") + self.default_fparam = default_fparam + if self.default_fparam is not None: + raise ValueError("default_fparam is not supported in TensorFlow.") self.n_neuron = neuron self.resnet_dt = resnet_dt self.rcond = rcond @@ -884,7 +893,7 @@ def deserialize(cls, data: dict, suffix: str = ""): The deserialized model """ data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( data["nets"], @@ -910,7 +919,7 @@ def serialize(self, suffix: str = "") -> dict: data = { "@class": "Fitting", "type": "ener", - "@version": 3, + "@version": 4, "var_name": "energy", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt + self.tebd_dim, @@ -921,6 +930,7 @@ def serialize(self, suffix: str = "") -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "rcond": self.rcond, "tot_ener_zero": self.tot_ener_zero, "trainable": self.trainable, diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index c44af58a5a..779cfbc8da 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -90,6 +90,9 @@ class PolarFittingSeA(Fitting): different fitting nets for different atom types. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. trainable : list[bool], Optional If the weights of fitting net are trainable. Suppose that we have :math:`N_l` hidden layers in the fitting net, @@ -117,6 +120,7 @@ def __init__( uniform_seed: bool = False, mixed_types: bool = False, type_map: Optional[list[str]] = None, # to be compat with input + default_fparam: Optional[list[float]] = None, # to be compat with input trainable: Optional[list[bool]] = None, **kwargs, ) -> None: @@ -175,12 +179,15 @@ def __init__( self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam self.dim_case_embd = dim_case_embd + self.default_fparam = default_fparam if numb_fparam > 0: raise ValueError("numb_fparam is not supported in the dipole fitting") if numb_aparam > 0: raise ValueError("numb_aparam is not supported in the dipole fitting") if dim_case_embd > 0: raise ValueError("dim_case_embd is not supported in TensorFlow.") + if default_fparam is not None: + raise ValueError("default_fparam is not supported in TensorFlow.") self.fparam_avg = None self.fparam_std = None self.fparam_inv_std = None @@ -629,7 +636,7 @@ def serialize(self, suffix: str) -> dict: data = { "@class": "Fitting", "type": "polar", - "@version": 4, + "@version": 5, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, "embedding_width": self.dim_rot_mat_1, @@ -640,6 +647,7 @@ def serialize(self, suffix: str) -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "activation_function": self.activation_function_name, "precision": self.fitting_precision.name, "exclude_types": [], @@ -687,7 +695,7 @@ def deserialize(cls, data: dict, suffix: str): """ data = data.copy() check_version_compatibility( - data.pop("@version", 1), 4, 1 + data.pop("@version", 1), 5, 1 ) # to allow PT version. fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 195b43dc8d..308d39b0a3 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1748,6 +1748,7 @@ def descrpt_variant_type_args(exclude_hybrid: bool = False) -> Variant: def fitting_ener() -> list[Argument]: doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -1775,6 +1776,13 @@ def fitting_ener() -> list[Argument]: return [ Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_only_pt_supported + doc_default_fparam, + ), Argument( "dim_case_embd", int, @@ -1832,6 +1840,7 @@ def fitting_ener() -> list[Argument]: def fitting_dos() -> list[Argument]: doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -1849,6 +1858,13 @@ def fitting_dos() -> list[Argument]: return [ Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_only_pt_supported + doc_default_fparam, + ), Argument( "dim_case_embd", int, @@ -1887,6 +1903,7 @@ def fitting_dos() -> list[Argument]: def fitting_property() -> list[Argument]: doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built" doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -1902,6 +1919,13 @@ def fitting_property() -> list[Argument]: return [ Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_only_pt_supported + doc_default_fparam, + ), Argument( "dim_case_embd", int, @@ -1949,6 +1973,7 @@ def fitting_property() -> list[Argument]: def fitting_polar() -> list[Argument]: doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -1978,6 +2003,13 @@ def fitting_polar() -> list[Argument]: default=0, doc=doc_only_pt_supported + doc_numb_aparam, ), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_only_pt_supported + doc_default_fparam, + ), Argument( "dim_case_embd", int, @@ -2027,6 +2059,7 @@ def fitting_polar() -> list[Argument]: def fitting_dipole() -> list[Argument]: doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -2049,6 +2082,13 @@ def fitting_dipole() -> list[Argument]: default=0, doc=doc_only_pt_supported + doc_numb_aparam, ), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_only_pt_supported + doc_default_fparam, + ), Argument( "dim_case_embd", int, diff --git a/source/tests/array_api_strict/fitting/fitting.py b/source/tests/array_api_strict/fitting/fitting.py index 323a49cfe8..c4a5674d2a 100644 --- a/source/tests/array_api_strict/fitting/fitting.py +++ b/source/tests/array_api_strict/fitting/fitting.py @@ -31,6 +31,7 @@ def setattr_for_general_fitting(name: str, value: Any) -> Any: "fparam_inv_std", "aparam_avg", "aparam_inv_std", + "default_fparam_tensor", }: value = to_array_api_strict_array(value) elif name == "emask": diff --git a/source/tests/consistent/fitting/test_ener.py b/source/tests/consistent/fitting/test_ener.py index f5a79acabe..ad70bd0bfa 100644 --- a/source/tests/consistent/fitting/test_ener.py +++ b/source/tests/consistent/fitting/test_ener.py @@ -70,7 +70,7 @@ (True, False), # resnet_dt ("float64", "float32", "bfloat16"), # precision (True, False), # mixed_types - (0, 1), # numb_fparam + ((0, None), (1, None), (1, [1.0])), # (numb_fparam, default_fparam) ((0, False), (1, False), (1, True)), # (numb_aparam, use_aparam_as_mask) ([], [-12345.6, None]), # atom_ener ) @@ -81,7 +81,7 @@ def data(self) -> dict: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -91,6 +91,7 @@ def data(self) -> dict: "precision": precision, "numb_fparam": numb_fparam, "numb_aparam": numb_aparam, + "default_fparam": default_fparam, "seed": 20240217, "atom_ener": atom_ener, "use_aparam_as_mask": use_aparam_as_mask, @@ -102,7 +103,7 @@ def skip_pt(self) -> bool: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -116,7 +117,7 @@ def skip_array_api_strict(self) -> bool: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -129,13 +130,25 @@ def skip_pd(self) -> bool: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param # Paddle do not support "bfloat16" in some kernels, # so skip this in CI test - return not INSTALLED_PD or precision == "bfloat16" + return not INSTALLED_PD or precision == "bfloat16" or default_fparam is not None + + @property + def skip_tf(self) -> bool: + ( + resnet_dt, + precision, + mixed_types, + (numb_fparam, default_fparam), + (numb_aparam, use_aparam_as_mask), + atom_ener, + ) = self.param + return not INSTALLED_TF or default_fparam is not None tf_class = EnerFittingTF dp_class = EnerFittingDP @@ -165,7 +178,7 @@ def additional_data(self) -> dict: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -180,7 +193,7 @@ def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -199,7 +212,7 @@ def eval_pt(self, pt_obj: Any) -> Any: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -209,7 +222,7 @@ def eval_pt(self, pt_obj: Any) -> Any: torch.from_numpy(self.atype.reshape(1, -1)).to(device=PT_DEVICE), fparam=( torch.from_numpy(self.fparam).to(device=PT_DEVICE) - if numb_fparam + if (numb_fparam and default_fparam is None) # test default_fparam else None ), aparam=( @@ -228,14 +241,14 @@ def eval_dp(self, dp_obj: Any) -> Any: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param return dp_obj( self.inputs, self.atype.reshape(1, -1), - fparam=self.fparam if numb_fparam else None, + fparam=self.fparam if (numb_fparam and default_fparam is None) else None, aparam=self.aparam if numb_aparam else None, )["energy"] @@ -244,7 +257,7 @@ def eval_jax(self, jax_obj: Any) -> Any: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -252,7 +265,9 @@ def eval_jax(self, jax_obj: Any) -> Any: jax_obj( jnp.asarray(self.inputs), jnp.asarray(self.atype.reshape(1, -1)), - fparam=jnp.asarray(self.fparam) if numb_fparam else None, + fparam=jnp.asarray(self.fparam) + if (numb_fparam and default_fparam is None) + else None, aparam=jnp.asarray(self.aparam) if numb_aparam else None, )["energy"] ) @@ -262,7 +277,7 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -270,7 +285,9 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: array_api_strict_obj( array_api_strict.asarray(self.inputs), array_api_strict.asarray(self.atype.reshape(1, -1)), - fparam=array_api_strict.asarray(self.fparam) if numb_fparam else None, + fparam=array_api_strict.asarray(self.fparam) + if (numb_fparam and default_fparam is None) + else None, aparam=array_api_strict.asarray(self.aparam) if numb_aparam else None, )["energy"] ) @@ -280,7 +297,7 @@ def eval_pd(self, pd_obj: Any) -> Any: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -317,7 +334,7 @@ def rtol(self) -> float: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -337,7 +354,7 @@ def atol(self) -> float: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param diff --git a/source/tests/universal/dpmodel/fitting/test_fitting.py b/source/tests/universal/dpmodel/fitting/test_fitting.py index 90b0668d20..29c5fcd4da 100644 --- a/source/tests/universal/dpmodel/fitting/test_fitting.py +++ b/source/tests/universal/dpmodel/fitting/test_fitting.py @@ -52,6 +52,7 @@ def FittingParamEnergy( "numb_fparam": numb_param, "numb_aparam": numb_param, "dim_case_embd": numb_param, + "default_fparam": [1.0] * numb_param if numb_param > 0 else None, } return input_dict