diff --git a/deepmd/dpmodel/atomic_model/__init__.py b/deepmd/dpmodel/atomic_model/__init__.py index 3d90c738ae..4d882d5e4b 100644 --- a/deepmd/dpmodel/atomic_model/__init__.py +++ b/deepmd/dpmodel/atomic_model/__init__.py @@ -42,6 +42,9 @@ from .polar_atomic_model import ( DPPolarAtomicModel, ) +from .property_atomic_model import ( + DPPropertyAtomicModel, +) __all__ = [ "BaseAtomicModel", @@ -50,6 +53,7 @@ "DPDipoleAtomicModel", "DPEnergyAtomicModel", "DPPolarAtomicModel", + "DPPropertyAtomicModel", "DPZBLLinearEnergyAtomicModel", "LinearEnergyAtomicModel", "PairTabAtomicModel", diff --git a/deepmd/dpmodel/atomic_model/property_atomic_model.py b/deepmd/dpmodel/atomic_model/property_atomic_model.py index 6f69f8dfb6..e3c038e695 100644 --- a/deepmd/dpmodel/atomic_model/property_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/property_atomic_model.py @@ -1,4 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np + from deepmd.dpmodel.fitting.property_fitting import ( PropertyFittingNet, ) @@ -15,3 +17,25 @@ def __init__(self, descriptor, fitting, type_map, **kwargs): "fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel" ) super().__init__(descriptor, fitting, type_map, **kwargs) + + def apply_out_stat( + self, + ret: dict[str, np.ndarray], + atype: np.ndarray, + ): + """Apply the stat to each atomic output. + + In property fitting, each output will be multiplied by label std and then plus the label average value. + + Parameters + ---------- + ret + The returned dict by the forward_atomic method + atype + The atom types. nf x nloc. It is useless in property fitting. + + """ + out_bias, out_std = self._fetch_out_stat(self.bias_keys) + for kk in self.bias_keys: + ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][0] + return ret diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py index 8b903af00e..6d0aa3546f 100644 --- a/deepmd/dpmodel/fitting/property_fitting.py +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -41,10 +41,9 @@ class PropertyFittingNet(InvarFitting): this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable. intensive Whether the fitting property is intensive. - bias_method - The method of applying the bias to each atomic output, user can select 'normal' or 'no_bias'. - If 'normal' is used, the computed bias will be added to the atomic output. - If 'no_bias' is used, no bias will be added to the atomic output. + property_name: + The name of fitting property, which should be consistent with the property name in the dataset. + If the data file is named `humo.npy`, this parameter should be "humo". resnet_dt Time-step `dt` in the resnet construction: :math:`y = x + dt * \phi (Wx + b)` @@ -74,7 +73,7 @@ def __init__( rcond: Optional[float] = None, trainable: Union[bool, list[bool]] = True, intensive: bool = False, - bias_method: str = "normal", + property_name: str = "property", resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -89,9 +88,8 @@ def __init__( ) -> None: self.task_dim = task_dim self.intensive = intensive - self.bias_method = bias_method super().__init__( - var_name="property", + var_name=property_name, ntypes=ntypes, dim_descrpt=dim_descrpt, dim_out=task_dim, @@ -113,9 +111,9 @@ def __init__( @classmethod def deserialize(cls, data: dict) -> "PropertyFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version"), 3, 1) + check_version_compatibility(data.pop("@version"), 4, 1) data.pop("dim_out") - data.pop("var_name") + data["property_name"] = data.pop("var_name") data.pop("tot_ener_zero") data.pop("layer_name") data.pop("use_aparam_as_mask", None) @@ -131,6 +129,8 @@ def serialize(self) -> dict: **InvarFitting.serialize(self), "type": "property", "task_dim": self.task_dim, + "intensive": self.intensive, } + dd["@version"] = 4 return dd diff --git a/deepmd/dpmodel/model/property_model.py b/deepmd/dpmodel/model/property_model.py index 16fdedd36e..9bd07bd349 100644 --- a/deepmd/dpmodel/model/property_model.py +++ b/deepmd/dpmodel/model/property_model.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from deepmd.dpmodel.atomic_model.dp_atomic_model import ( - DPAtomicModel, +from deepmd.dpmodel.atomic_model import ( + DPPropertyAtomicModel, ) from deepmd.dpmodel.model.base_model import ( BaseModel, @@ -13,7 +13,7 @@ make_model, ) -DPPropertyModel_ = make_model(DPAtomicModel) +DPPropertyModel_ = make_model(DPPropertyAtomicModel) @BaseModel.register("property") diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index d9744246d7..5aeb84468d 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -779,9 +779,17 @@ def test_property( tuple[list[np.ndarray], list[int]] arrays with results and their shapes """ - data.add("property", dp.task_dim, atomic=False, must=True, high_prec=True) + var_name = dp.get_var_name() + assert isinstance(var_name, str) + data.add(var_name, dp.task_dim, atomic=False, must=True, high_prec=True) if has_atom_property: - data.add("atom_property", dp.task_dim, atomic=True, must=False, high_prec=True) + data.add( + f"atom_{var_name}", + dp.task_dim, + atomic=True, + must=False, + high_prec=True, + ) if dp.get_dim_fparam() > 0: data.add( @@ -832,12 +840,12 @@ def test_property( aproperty = ret[1] aproperty = aproperty.reshape([numb_test, natoms * dp.task_dim]) - diff_property = property - test_data["property"][:numb_test] + diff_property = property - test_data[var_name][:numb_test] mae_property = mae(diff_property) rmse_property = rmse(diff_property) if has_atom_property: - diff_aproperty = aproperty - test_data["atom_property"][:numb_test] + diff_aproperty = aproperty - test_data[f"atom_{var_name}"][:numb_test] mae_aproperty = mae(diff_aproperty) rmse_aproperty = rmse(diff_aproperty) @@ -854,7 +862,7 @@ def test_property( detail_path = Path(detail_file) for ii in range(numb_test): - test_out = test_data["property"][ii].reshape(-1, 1) + test_out = test_data[var_name][ii].reshape(-1, 1) pred_out = property[ii].reshape(-1, 1) frame_output = np.hstack((test_out, pred_out)) @@ -868,7 +876,7 @@ def test_property( if has_atom_property: for ii in range(numb_test): - test_out = test_data["atom_property"][ii].reshape(-1, 1) + test_out = test_data[f"atom_{var_name}"][ii].reshape(-1, 1) pred_out = aproperty[ii].reshape(-1, 1) frame_output = np.hstack((test_out, pred_out)) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 159f9bdf60..15e4a56280 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -70,8 +70,6 @@ class DeepEvalBackend(ABC): "dipole_derv_c_redu": "virial", "dos": "atom_dos", "dos_redu": "dos", - "property": "atom_property", - "property_redu": "property", "mask_mag": "mask_mag", "mask": "mask", # old models in v1 @@ -276,6 +274,10 @@ def get_has_spin(self) -> bool: """Check if the model has spin atom types.""" return False + def get_var_name(self) -> str: + """Get the name of the fitting property.""" + raise NotImplementedError + @abstractmethod def get_ntypes_spin(self) -> int: """Get the number of spin atom types of this model. Only used in old implement.""" diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index 389a0e8512..5944491cc0 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -37,25 +37,41 @@ class DeepProperty(DeepEval): Keyword arguments. """ - @property def output_def(self) -> ModelOutputDef: - """Get the output definition of this model.""" - return ModelOutputDef( + """ + Get the output definition of this model. + But in property_fitting, the output definition is not known until the model is loaded. + So we need to rewrite the output definition after the model is loaded. + See detail in change_output_def. + """ + pass + + def change_output_def(self) -> None: + """ + Change the output definition of this model. + In property_fitting, the output definition is known after the model is loaded. + We need to rewrite the output definition and related information. + """ + self.output_def = ModelOutputDef( FittingOutputDef( [ OutputVariableDef( - "property", - shape=[-1], + self.get_var_name(), + shape=[self.get_task_dim()], reducible=True, atomic=True, + intensive=self.get_intensive(), ), ] ) ) - - def change_output_def(self) -> None: - self.output_def["property"].shape = self.task_dim - self.output_def["property"].intensive = self.get_intensive() + self.deep_eval.output_def = self.output_def + self.deep_eval._OUTDEF_DP2BACKEND[self.get_var_name()] = ( + f"atom_{self.get_var_name()}" + ) + self.deep_eval._OUTDEF_DP2BACKEND[f"{self.get_var_name()}_redu"] = ( + self.get_var_name() + ) @property def task_dim(self) -> int: @@ -120,10 +136,12 @@ def eval( aparam=aparam, **kwargs, ) - atomic_property = results["property"].reshape( + atomic_property = results[self.get_var_name()].reshape( nframes, natoms, self.get_task_dim() ) - property = results["property_redu"].reshape(nframes, self.get_task_dim()) + property = results[f"{self.get_var_name()}_redu"].reshape( + nframes, self.get_task_dim() + ) if atomic: return ( @@ -141,5 +159,9 @@ def get_intensive(self) -> bool: """Get whether the property is intensive.""" return self.deep_eval.get_intensive() + def get_var_name(self) -> str: + """Get the name of the fitting property.""" + return self.deep_eval.get_var_name() + __all__ = ["DeepProperty"] diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 59b833d34c..facead838e 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -184,6 +184,15 @@ def get_dim_aparam(self) -> int: def get_intensive(self) -> bool: return self.dp.model["Default"].get_intensive() + def get_var_name(self) -> str: + """Get the name of the property.""" + if hasattr(self.dp.model["Default"], "get_var_name") and callable( + getattr(self.dp.model["Default"], "get_var_name") + ): + return self.dp.model["Default"].get_var_name() + else: + raise NotImplementedError + @property def model_type(self) -> type["DeepEvalWrapper"]: """The the evaluator of the model type.""" @@ -200,7 +209,7 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepGlobalPolar elif "wfc" in model_output_type: return DeepWFC - elif "property" in model_output_type: + elif self.get_var_name() in model_output_type: return DeepProperty else: raise RuntimeError("Unknown model type") diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index 07e394650a..9d42c81b45 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging +from typing import ( + Union, +) import torch import torch.nn.functional as F @@ -21,9 +24,13 @@ class PropertyLoss(TaskLoss): def __init__( self, task_dim, + var_name: str, loss_func: str = "smooth_mae", metric: list = ["mae"], beta: float = 1.00, + out_bias: Union[list, None] = None, + out_std: Union[list, None] = None, + intensive: bool = False, **kwargs, ) -> None: r"""Construct a layer to compute loss on property. @@ -32,18 +39,32 @@ def __init__( ---------- task_dim : float The output dimension of property fitting net. + var_name : str + The atomic property to fit, 'energy', 'dipole', and 'polar'. loss_func : str The loss function, such as "smooth_mae", "mae", "rmse". metric : list The metric such as mae, rmse which will be printed. - beta: + beta : float The 'beta' parameter in 'smooth_mae' loss. + out_bias : Union[list, None] + It is the average value of the label. The shape is nkeys * ntypes * task_dim. + In property fitting, nkeys = 1, so the shape is 1 * ntypes * task_dim. + out_std : Union[list, None] + It is the standard deviation of the label. The shape is nkeys * ntypes * task_dim. + In property fitting, nkeys = 1, so the shape is 1 * ntypes * task_dim. + intensive : bool + Whether the property is intensive. """ super().__init__() self.task_dim = task_dim self.loss_func = loss_func self.metric = metric self.beta = beta + self.out_bias = out_bias + self.out_std = out_std + self.intensive = intensive + self.var_name = var_name def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False): """Return loss on properties . @@ -69,34 +90,64 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False Other losses for display. """ model_pred = model(**input_dict) - assert label["property"].shape[-1] == self.task_dim - assert model_pred["property"].shape[-1] == self.task_dim + var_name = self.var_name + nbz = model_pred[var_name].shape[0] + assert model_pred[var_name].shape == (nbz, self.task_dim) + assert label[var_name].shape == (nbz, self.task_dim) + if not self.intensive: + model_pred[var_name] = model_pred[var_name] / natoms + label[var_name] = label[var_name] / natoms + + if self.out_std is None: + out_std = model.atomic_model.out_std[0][0] + else: + out_std = torch.tensor( + self.out_std, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) + if out_std.shape != (self.task_dim,): + raise ValueError( + f"Expected out_std to have shape ({self.task_dim},), but got {out_std.shape}" + ) + + if self.out_bias is None: + out_bias = model.atomic_model.out_bias[0][0] + else: + out_bias = torch.tensor( + self.out_bias, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) + if out_bias.shape != (self.task_dim,): + raise ValueError( + f"Expected out_bias to have shape ({self.task_dim},), but got {out_bias.shape}" + ) + loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] more_loss = {} # loss if self.loss_func == "smooth_mae": loss += F.smooth_l1_loss( - label["property"], - model_pred["property"], + (label[var_name] - out_bias) / out_std, + (model_pred[var_name] - out_bias) / out_std, reduction="sum", beta=self.beta, ) elif self.loss_func == "mae": loss += F.l1_loss( - label["property"], model_pred["property"], reduction="sum" + (label[var_name] - out_bias) / out_std, + (model_pred[var_name] - out_bias) / out_std, + reduction="sum", ) elif self.loss_func == "mse": loss += F.mse_loss( - label["property"], - model_pred["property"], + (label[var_name] - out_bias) / out_std, + (model_pred[var_name] - out_bias) / out_std, reduction="sum", ) elif self.loss_func == "rmse": loss += torch.sqrt( F.mse_loss( - label["property"], - model_pred["property"], + (label[var_name] - out_bias) / out_std, + (model_pred[var_name] - out_bias) / out_std, reduction="mean", ) ) @@ -106,28 +157,28 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False # more loss if "smooth_mae" in self.metric: more_loss["smooth_mae"] = F.smooth_l1_loss( - label["property"], - model_pred["property"], + label[var_name], + model_pred[var_name], reduction="mean", beta=self.beta, ).detach() if "mae" in self.metric: more_loss["mae"] = F.l1_loss( - label["property"], - model_pred["property"], + label[var_name], + model_pred[var_name], reduction="mean", ).detach() if "mse" in self.metric: more_loss["mse"] = F.mse_loss( - label["property"], - model_pred["property"], + label[var_name], + model_pred[var_name], reduction="mean", ).detach() if "rmse" in self.metric: more_loss["rmse"] = torch.sqrt( F.mse_loss( - label["property"], - model_pred["property"], + label[var_name], + model_pred[var_name], reduction="mean", ) ).detach() @@ -140,10 +191,10 @@ def label_requirement(self) -> list[DataRequirementItem]: label_requirement = [] label_requirement.append( DataRequirementItem( - "property", + self.var_name, ndof=self.task_dim, atomic=False, - must=False, + must=True, high_prec=True, ) ) diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index a64eca0fe9..c83e35dab3 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -125,6 +125,14 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.type_map + def get_compute_stats_distinguish_types(self) -> bool: + """Get whether the fitting net computes stats which are not distinguished between different types of atoms.""" + return True + + def get_intensive(self) -> bool: + """Whether the fitting property is intensive.""" + return False + def reinit_atom_exclude( self, exclude_types: list[int] = [], @@ -456,7 +464,6 @@ def change_out_bias( model_forward=self._get_forward_wrapper_func(), rcond=self.rcond, preset_bias=self.preset_out_bias, - atomic_output=self.atomic_output_def(), ) self._store_out_stat(delta_bias, out_std, add=True) elif bias_adjust_mode == "set-by-statistic": @@ -467,7 +474,8 @@ def change_out_bias( stat_file_path=stat_file_path, rcond=self.rcond, preset_bias=self.preset_out_bias, - atomic_output=self.atomic_output_def(), + stats_distinguish_types=self.get_compute_stats_distinguish_types(), + intensive=self.get_intensive(), ) self._store_out_stat(bias_out, std_out) else: diff --git a/deepmd/pt/model/atomic_model/property_atomic_model.py b/deepmd/pt/model/atomic_model/property_atomic_model.py index 1fdc72b2b6..3622c9f476 100644 --- a/deepmd/pt/model/atomic_model/property_atomic_model.py +++ b/deepmd/pt/model/atomic_model/property_atomic_model.py @@ -19,31 +19,31 @@ def __init__(self, descriptor, fitting, type_map, **kwargs): ) super().__init__(descriptor, fitting, type_map, **kwargs) + def get_compute_stats_distinguish_types(self) -> bool: + """Get whether the fitting net computes stats which are not distinguished between different types of atoms.""" + return False + + def get_intensive(self) -> bool: + """Whether the fitting property is intensive.""" + return self.fitting_net.get_intensive() + def apply_out_stat( self, ret: dict[str, torch.Tensor], atype: torch.Tensor, ): """Apply the stat to each atomic output. - This function defines how the bias is applied to the atomic output of the model. + In property fitting, each output will be multiplied by label std and then plus the label average value. Parameters ---------- ret The returned dict by the forward_atomic method atype - The atom types. nf x nloc + The atom types. nf x nloc. It is useless in property fitting. """ - if self.fitting_net.get_bias_method() == "normal": - out_bias, out_std = self._fetch_out_stat(self.bias_keys) - for kk in self.bias_keys: - # nf x nloc x odims, out_bias: ntypes x odims - ret[kk] = ret[kk] + out_bias[kk][atype] - return ret - elif self.fitting_net.get_bias_method() == "no_bias": - return ret - else: - raise NotImplementedError( - "Only 'normal' and 'no_bias' is supported for parameter 'bias_method'." - ) + out_bias, out_std = self._fetch_out_stat(self.bias_keys) + for kk in self.bias_keys: + ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][0] + return ret diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index 4581a2bc3e..7c50c75ff1 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/deepmd/pt/model/model/property_model.py @@ -37,8 +37,8 @@ def __init__( def translated_output_def(self): out_def_data = self.model_output_def().get_data() output_def = { - "atom_property": out_def_data["property"], - "property": out_def_data["property_redu"], + f"atom_{self.get_var_name()}": out_def_data[self.get_var_name()], + self.get_var_name(): out_def_data[f"{self.get_var_name()}_redu"], } if "mask" in out_def_data: output_def["mask"] = out_def_data["mask"] @@ -62,8 +62,8 @@ def forward( do_atomic_virial=do_atomic_virial, ) model_predict = {} - model_predict["atom_property"] = model_ret["property"] - model_predict["property"] = model_ret["property_redu"] + model_predict[f"atom_{self.get_var_name()}"] = model_ret[self.get_var_name()] + model_predict[self.get_var_name()] = model_ret[f"{self.get_var_name()}_redu"] if "mask" in model_ret: model_predict["mask"] = model_ret["mask"] return model_predict @@ -76,7 +76,12 @@ def get_task_dim(self) -> int: @torch.jit.export def get_intensive(self) -> bool: """Get whether the property is intensive.""" - return self.model_output_def()["property"].intensive + return self.model_output_def()[self.get_var_name()].intensive + + @torch.jit.export + def get_var_name(self) -> str: + """Get the name of the property.""" + return self.get_fitting_net().var_name @torch.jit.export def forward_lower( @@ -102,8 +107,8 @@ def forward_lower( extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) model_predict = {} - model_predict["atom_property"] = model_ret["property"] - model_predict["property"] = model_ret["property_redu"] + model_predict[f"atom_{self.get_var_name()}"] = model_ret[self.get_var_name()] + model_predict[self.get_var_name()] = model_ret[f"{self.get_var_name()}_redu"] if "mask" in model_ret: model_predict["mask"] = model_ret["mask"] return model_predict diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index dec0f1447b..c15e60fe04 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -43,17 +43,16 @@ class PropertyFittingNet(InvarFitting): dim_descrpt : int Embedding width per atom. task_dim : int - The dimension of outputs of fitting net. + The dimension of outputs of fitting net. + property_name: + The name of fitting property, which should be consistent with the property name in the dataset. + If the data file is named `humo.npy`, this parameter should be "humo". neuron : list[int] Number of neurons in each hidden layers of the fitting net. bias_atom_p : torch.Tensor, optional Average property per atom for each element. intensive : bool, optional Whether the fitting property is intensive. - bias_method : str, optional - The method of applying the bias to each atomic output, user can select 'normal' or 'no_bias'. - If 'normal' is used, the computed bias will be added to the atomic output. - If 'no_bias' is used, no bias will be added to the atomic output. resnet_dt : bool Using time-step in the ResNet construction. numb_fparam : int @@ -77,11 +76,11 @@ def __init__( self, ntypes: int, dim_descrpt: int, + property_name: str, task_dim: int = 1, neuron: list[int] = [128, 128, 128], bias_atom_p: Optional[torch.Tensor] = None, intensive: bool = False, - bias_method: str = "normal", resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -94,9 +93,8 @@ def __init__( ) -> None: self.task_dim = task_dim self.intensive = intensive - self.bias_method = bias_method super().__init__( - var_name="property", + var_name=property_name, ntypes=ntypes, dim_descrpt=dim_descrpt, dim_out=task_dim, @@ -113,9 +111,6 @@ def __init__( **kwargs, ) - def get_bias_method(self) -> str: - return self.bias_method - def output_def(self) -> FittingOutputDef: return FittingOutputDef( [ @@ -130,12 +125,16 @@ def output_def(self) -> FittingOutputDef: ] ) + def get_intensive(self) -> bool: + """Whether the fitting property is intensive.""" + return self.intensive + @classmethod def deserialize(cls, data: dict) -> "PropertyFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("dim_out") - data.pop("var_name") + data["property_name"] = data.pop("var_name") obj = super().deserialize(data) return obj @@ -146,7 +145,9 @@ def serialize(self) -> dict: **InvarFitting.serialize(self), "type": "property", "task_dim": self.task_dim, + "intensive": self.intensive, } + dd["@version"] = 4 return dd diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 8ca510492c..eca952d7f8 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -1240,7 +1240,11 @@ def get_loss(loss_params, start_lr, _ntypes, _model): return TensorLoss(**loss_params) elif loss_type == "property": task_dim = _model.get_task_dim() + var_name = _model.get_var_name() + intensive = _model.get_intensive() loss_params["task_dim"] = task_dim + loss_params["var_name"] = var_name + loss_params["intensive"] = intensive return PropertyLoss(**loss_params) else: loss_params["starter_learning_rate"] = start_lr diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 1c5e3f1c52..710d392ac3 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -12,9 +12,6 @@ import numpy as np import torch -from deepmd.dpmodel.output_def import ( - FittingOutputDef, -) from deepmd.pt.utils import ( AtomExcludeMask, ) @@ -27,6 +24,7 @@ to_torch_tensor, ) from deepmd.utils.out_stat import ( + compute_stats_do_not_distinguish_types, compute_stats_from_atomic, compute_stats_from_redu, ) @@ -136,11 +134,16 @@ def _post_process_stat( For global statistics, we do not have the std for each type of atoms, thus fake the output std by ones for all the types. + If the shape of out_std is already the same as out_bias, + we do not need to do anything. """ new_std = {} for kk, vv in out_bias.items(): - new_std[kk] = np.ones_like(vv) + if vv.shape == out_std[kk].shape: + new_std[kk] = out_std[kk] + else: + new_std[kk] = np.ones_like(vv) return out_bias, new_std @@ -242,7 +245,8 @@ def compute_output_stats( rcond: Optional[float] = None, preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None, model_forward: Optional[Callable[..., torch.Tensor]] = None, - atomic_output: Optional[FittingOutputDef] = None, + stats_distinguish_types: bool = True, + intensive: bool = False, ): """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -272,8 +276,10 @@ def compute_output_stats( If not None, the model will be utilized to generate the original energy prediction, which will be subtracted from the energy label of the data. The difference will then be used to calculate the delta complement energy bias for each type. - atomic_output : FittingOutputDef, optional - The output of atomic model. + stats_distinguish_types : bool, optional + Whether to distinguish different element types in the statistics. + intensive : bool, optional + Whether the fitting target is intensive. """ # try to restore the bias from stat file bias_atom_e, std_atom_e = _restore_from_file(stat_file_path, keys) @@ -362,7 +368,8 @@ def compute_output_stats( rcond, preset_bias, model_pred_g, - atomic_output, + stats_distinguish_types, + intensive, ) bias_atom_a, std_atom_a = compute_output_stats_atomic( sampled, @@ -405,7 +412,8 @@ def compute_output_stats_global( rcond: Optional[float] = None, preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None, model_pred: Optional[dict[str, np.ndarray]] = None, - atomic_output: Optional[FittingOutputDef] = None, + stats_distinguish_types: bool = True, + intensive: bool = False, ): """This function only handle stat computation from reduced global labels.""" # return directly if model predict is empty for global @@ -476,19 +484,22 @@ def compute_output_stats_global( std_atom_e = {} for kk in keys: if kk in stats_input: - if atomic_output is not None and atomic_output.get_data()[kk].intensive: - task_dim = stats_input[kk].shape[1] - assert merged_natoms[kk].shape == (nf[kk], ntypes) - stats_input[kk] = ( - merged_natoms[kk].sum(axis=1).reshape(-1, 1) * stats_input[kk] + if not stats_distinguish_types: + bias_atom_e[kk], std_atom_e[kk] = ( + compute_stats_do_not_distinguish_types( + stats_input[kk], + merged_natoms[kk], + assigned_bias=assigned_atom_ener[kk], + intensive=intensive, + ) + ) + else: + bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu( + stats_input[kk], + merged_natoms[kk], + assigned_bias=assigned_atom_ener[kk], + rcond=rcond, ) - assert stats_input[kk].shape == (nf[kk], task_dim) - bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu( - stats_input[kk], - merged_natoms[kk], - assigned_bias=assigned_atom_ener[kk], - rcond=rcond, - ) else: # this key does not have global labels, skip it. continue diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 9eac0e804d..50ef07b2af 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1580,7 +1580,7 @@ def fitting_property(): doc_seed = "Random seed for parameter initialization of the fitting net" doc_task_dim = "The dimension of outputs of fitting net" doc_intensive = "Whether the fitting property is intensive" - doc_bias_method = "The method of applying the bias to each atomic output, user can select 'normal' or 'no_bias'. If 'no_bias' is used, no bias will be added to the atomic output." + doc_property_name = "The names of fitting property, which should be consistent with the property name in the dataset." 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), @@ -1612,7 +1612,10 @@ def fitting_property(): Argument("task_dim", int, optional=True, default=1, doc=doc_task_dim), Argument("intensive", bool, optional=True, default=False, doc=doc_intensive), Argument( - "bias_method", str, optional=True, default="normal", doc=doc_bias_method + "property_name", + str, + optional=False, + doc=doc_property_name, ), ] diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index 4d0d788f8b..ecbd379e2d 100644 --- a/deepmd/utils/out_stat.py +++ b/deepmd/utils/out_stat.py @@ -130,3 +130,64 @@ def compute_stats_from_atomic( output[mask].std(axis=0) if output[mask].size > 0 else np.nan ) return output_bias, output_std + + +def compute_stats_do_not_distinguish_types( + output_redu: np.ndarray, + natoms: np.ndarray, + assigned_bias: Optional[np.ndarray] = None, + intensive: bool = False, +) -> tuple[np.ndarray, np.ndarray]: + """Compute element-independent statistics for property fitting. + + Computes mean and standard deviation of the output, treating all elements equally. + For extensive properties, the output is normalized by the total number of atoms + before computing statistics. + + Parameters + ---------- + output_redu + The reduced output value, shape is [nframes, *(odim0, odim1, ...)]. + natoms + The number of atoms for each atom, shape is [nframes, ntypes]. + Used for normalization of extensive properties and generating uniform bias. + assigned_bias + The assigned output bias, shape is [ntypes, *(odim0, odim1, ...)]. + Set to a tensor of shape (odim0, odim1, ...) filled with nan if the bias + of the type is not assigned. + intensive + Whether the output is intensive or extensive. + If False, the output will be normalized by the total number of atoms before computing statistics. + + Returns + ------- + np.ndarray + The computed output mean(fake bias), shape is [ntypes, *(odim0, odim1, ...)]. + The same bias is used for all atom types. + np.ndarray + The computed output standard deviation, shape is [ntypes, *(odim0, odim1, ...)]. + The same standard deviation is used for all atom types. + """ + natoms = np.array(natoms) # [nf, ntypes] + nf, ntypes = natoms.shape + output_redu = np.array(output_redu) + var_shape = list(output_redu.shape[1:]) + output_redu = output_redu.reshape(nf, -1) + if not intensive: + total_atoms = natoms.sum(axis=1) + output_redu = output_redu / total_atoms[:, np.newaxis] + # check shape + assert output_redu.ndim == 2 + assert natoms.ndim == 2 + assert output_redu.shape[0] == natoms.shape[0] # [nf,1] + + computed_output_bias = np.repeat( + np.mean(output_redu, axis=0)[np.newaxis, :], ntypes, axis=0 + ) + output_std = np.std(output_redu, axis=0) + + computed_output_bias = computed_output_bias.reshape([natoms.shape[1]] + var_shape) # noqa: RUF005 + output_std = output_std.reshape(var_shape) + output_std = np.tile(output_std, (computed_output_bias.shape[0], 1)) + + return computed_output_bias, output_std diff --git a/doc/model/index.rst b/doc/model/index.rst index c067ea4207..5e7ba32486 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -16,6 +16,7 @@ Model train-energy-spin train-fitting-tensor train-fitting-dos + train-fitting-property train-se-e2-a-tebd train-se-a-mask train-se-e3-tebd diff --git a/doc/model/train-fitting-property.md b/doc/model/train-fitting-property.md new file mode 100644 index 0000000000..be1b63bf6f --- /dev/null +++ b/doc/model/train-fitting-property.md @@ -0,0 +1,194 @@ +# Fit other properties {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} + +:::{note} +**Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +::: + +Here we present an API to DeepProperty model, which can be used to fit other properties like band gap, bulk modulus, critical temperature, etc. + +In this example, we will show you how to train a model to fit properties of `humo`, `lumo` and `band gap`. A complete training input script of the examples can be found in + +```bash +$deepmd_source_dir/examples/property/train +``` + +The training and validation data are also provided our examples. But note that **the data provided along with the examples are of limited amount, and should not be used to train a production model.** + +Similar to the `input.json` used in `ener` mode, training JSON is also divided into {ref}`model `, {ref}`learning_rate `, {ref}`loss ` and {ref}`training `. Most keywords remain the same as `ener` mode, and their meaning can be found [here](train-se-atten.md). To fit the `property`, one needs to modify {ref}`model[standard]/fitting_net ` and {ref}`loss `. + +## The fitting Network + +The {ref}`fitting_net ` section tells DP which fitting net to use. + +The JSON of `property` type should be provided like + +```json +"fitting_net" : { + "type": "property", + "intensive": true, + "property_name": "band_prop", + "task_dim": 3, + "neuron": [240,240,240], + "resnet_dt": true, + "seed": 1, +}, +``` + +- `type` specifies which type of fitting net should be used. It should be `property`. +- `intensive` indicates whether the fitting property is intensive. If `intensive` is `true`, the model output is the average of the property contribution of each atom. If `intensive` is `false`, the model output is the sum of the property contribution of each atom. +- `property_name` is the name of the property to be predicted. It should be consistent with the property name in the dataset. In each system, code will read `set.*/{property_name}.npy` file as prediction label if you use NumPy format data. +- `fitting_net/task_dim` is the dimension of model output. It should be consistent with the property dimension in the dataset, which means if the shape of data stored in `set.*/{property_name}.npy` is `batch size * 3`, `fitting_net/task_dim` should be set to 3. +- The rest arguments have the same meaning as they do in `ener` mode. + +## Loss + +DeepProperty supports trainings of the global system (one or more global labels are provided in a frame). For example, when fitting `property`, each frame will provide a `1 x task_dim` vector which gives the fitting properties. + +The loss section should be provided like + +```json +"loss" : { + "type": "property", + "metric": ["mae"], + "loss_func": "smooth_mae" +}, +``` + +- {ref}`type ` should be written as `property` as a distinction from `ener` mode. +- `metric`: The metric for display, which will be printed in `lcurve.out`. This list can include 'smooth_mae', 'mae', 'mse' and 'rmse'. +- `loss_func`: The loss function to minimize, you can use 'mae','smooth_mae', 'mse' and 'rmse'. + +## Training Data Preparation + +The label should be named `{property_name}.npy/raw`, `property_name` is defined by `fitting_net/property_name` in `input.json`. + +To prepare the data, you can use `dpdata` tools, for example: + +```py +import dpdata +import numpy as np +from dpdata.data_type import ( + Axis, + DataType, +) + +property_name = "band_prop" # fittng_net/property_name +task_dim = 3 # fitting_net/task_dim + +# register datatype +datatypes = [ + DataType( + property_name, + np.ndarray, + shape=(Axis.NFRAMES, task_dim), + required=False, + ), +] +datatypes.extend( + [ + DataType( + "energies", + np.ndarray, + shape=(Axis.NFRAMES, 1), + required=False, + ), + DataType( + "forces", + np.ndarray, + shape=(Axis.NFRAMES, Axis.NATOMS, 1), + required=False, + ), + ] +) + +for datatype in datatypes: + dpdata.System.register_data_type(datatype) + dpdata.LabeledSystem.register_data_type(datatype) + +ls = dpdata.MultiSystems() +frame = dpdata.System("POSCAR", fmt="vasp/poscar") +labelframe = dpdata.LabeledSystem() +labelframe.append(frame) +labelframe.data[property_name] = np.array([[-0.236, 0.056, 0.292]], dtype=np.float32) +ls.append(labelframe) +ls.to_deepmd_npy_mixed("deepmd") +``` + +## Train the Model + +The training command is the same as `ener` mode, i.e. + +::::{tab-set} + +:::{tab-item} PyTorch {{ pytorch_icon }} + +```bash +dp --pt train input.json +``` + +::: + +:::: + +The detailed loss can be found in `lcurve.out`: + +``` +# step mae_val mae_trn lr +# If there is no available reference data, rmse_*_{val,trn} will print nan + 1 2.72e-02 2.40e-02 2.0e-04 + 100 1.79e-02 1.34e-02 2.0e-04 + 200 1.45e-02 1.86e-02 2.0e-04 + 300 1.61e-02 4.90e-03 2.0e-04 + 400 2.04e-02 1.05e-02 2.0e-04 + 500 9.09e-03 1.85e-02 2.0e-04 + 600 1.01e-02 5.63e-03 2.0e-04 + 700 1.10e-02 1.76e-02 2.0e-04 + 800 1.14e-02 1.50e-02 2.0e-04 + 900 9.54e-03 2.70e-02 2.0e-04 + 1000 1.00e-02 2.73e-02 2.0e-04 +``` + +## Test the Model + +We can use `dp test` to infer the properties for given frames. + +::::{tab-set} + +:::{tab-item} PyTorch {{ pytorch_icon }} + +```bash + +dp --pt freeze -o frozen_model.pth + +dp --pt test -m frozen_model.pth -s ../data/data_0/ -d ${output_prefix} -n 100 +``` + +::: + +:::: + +if `dp test -d ${output_prefix}` is specified, the predicted properties for each frame are output in the working directory + +``` +${output_prefix}.property.out.0 ${output_prefix}.property.out.1 ${output_prefix}.property.out.2 ${output_prefix}.property.out.3 +``` + +for `*.property.out.*`, it contains matrix with shape of `(2, task_dim)`, + +``` +# ../data/data_0 - 0: data_property pred_property +-2.449000030755996704e-01 -2.315840660495154801e-01 +6.400000303983688354e-02 5.810663314446311983e-02 +3.088999986648559570e-01 2.917143316092784544e-01 +``` + +## Data Normalization + +When `fitting_net/type` is `ener`, the energy bias layer “$e_{bias}$” adds a constant bias to the atomic energy contribution according to the atomic number.i.e., +$$e_{bias} (Z_i) (MLP(D_i))= MLP(D_i) + e_{bias} (Z_i)$$ + +But when `fitting_net/type` is `property`. The property bias layer is used to normalize the property output of the model.i.e., +$$p_{bias} (MLP(D_i))= MLP(D_i) * std+ mean$$ + +1. `std`: The standard deviation of the property label +2. `mean`: The average value of the property label diff --git a/examples/property/data/data_0/set.000000/property.npy b/examples/property/data/data_0/set.000000/band_prop.npy similarity index 100% rename from examples/property/data/data_0/set.000000/property.npy rename to examples/property/data/data_0/set.000000/band_prop.npy diff --git a/examples/property/data/data_1/set.000000/property.npy b/examples/property/data/data_1/set.000000/band_prop.npy similarity index 100% rename from examples/property/data/data_1/set.000000/property.npy rename to examples/property/data/data_1/set.000000/band_prop.npy diff --git a/examples/property/data/data_2/set.000000/property.npy b/examples/property/data/data_2/set.000000/band_prop.npy similarity index 100% rename from examples/property/data/data_2/set.000000/property.npy rename to examples/property/data/data_2/set.000000/band_prop.npy diff --git a/examples/property/train/README.md b/examples/property/train/README.md new file mode 100644 index 0000000000..e4dc9ed704 --- /dev/null +++ b/examples/property/train/README.md @@ -0,0 +1,5 @@ +Some explanations of the parameters in `input.json`: + +1. `fitting_net/property_name` is the name of the property to be predicted. It should be consistent with the property name in the dataset. In each system, code will read `set.*/{property_name}.npy` file as prediction label if you use NumPy format data. +2. `fitting_net/task_dim` is the dimension of model output. It should be consistent with the property dimension in the dataset, which means if the shape of data stored in `set.*/{property_name}.npy` is `batch size * 3`, `fitting_net/task_dim` should be set to 3. +3. `fitting/intensive` indicates whether the fitting property is intensive. If `intensive` is `true`, the model output is the average of the property contribution of each atom. If `intensive` is `false`, the model output is the sum of the property contribution of each atom. diff --git a/examples/property/train/input_torch.json b/examples/property/train/input_torch.json index 33eaa28a07..1e6ce00048 100644 --- a/examples/property/train/input_torch.json +++ b/examples/property/train/input_torch.json @@ -33,6 +33,7 @@ "type": "property", "intensive": true, "task_dim": 3, + "property_name": "band_prop", "neuron": [ 240, 240, @@ -53,6 +54,11 @@ }, "loss": { "type": "property", + "metric": [ + "mae" + ], + "loss_func": "smooth_mae", + "beta": 1.0, "_comment": " that's all" }, "training": { diff --git a/source/tests/common/test_out_stat.py b/source/tests/common/test_out_stat.py index c175d7c643..0236c39f22 100644 --- a/source/tests/common/test_out_stat.py +++ b/source/tests/common/test_out_stat.py @@ -4,6 +4,7 @@ import numpy as np from deepmd.utils.out_stat import ( + compute_stats_do_not_distinguish_types, compute_stats_from_atomic, compute_stats_from_redu, ) @@ -89,6 +90,58 @@ def test_compute_stats_from_redu_with_assigned_bias(self) -> None: rtol=1e-7, ) + def test_compute_stats_do_not_distinguish_types_intensive(self) -> None: + """Test compute_stats_property function with intensive scenario.""" + bias, std = compute_stats_do_not_distinguish_types( + self.output_redu, self.natoms, intensive=True + ) + # Test shapes + assert bias.shape == (len(self.mean), self.output_redu.shape[1]) + assert std.shape == (len(self.mean), self.output_redu.shape[1]) + + # Test values + for fake_atom_bias in bias: + np.testing.assert_allclose( + fake_atom_bias, np.mean(self.output_redu, axis=0), rtol=1e-7 + ) + for fake_atom_std in std: + np.testing.assert_allclose( + fake_atom_std, np.std(self.output_redu, axis=0), rtol=1e-7 + ) + + def test_compute_stats_do_not_distinguish_types_extensive(self) -> None: + """Test compute_stats_property function with extensive scenario.""" + bias, std = compute_stats_do_not_distinguish_types( + self.output_redu, self.natoms + ) + # Test shapes + assert bias.shape == (len(self.mean), self.output_redu.shape[1]) + assert std.shape == (len(self.mean), self.output_redu.shape[1]) + + # Test values + for fake_atom_bias in bias: + np.testing.assert_allclose( + fake_atom_bias, + np.array( + [ + 6218.91610282, + 7183.82275736, + 4445.23155934, + 5748.23644722, + 5362.8519454, + ] + ), + rtol=1e-7, + ) + for fake_atom_std in std: + np.testing.assert_allclose( + fake_atom_std, + np.array( + [128.78691576, 36.53743668, 105.82372405, 96.43642486, 33.68885327] + ), + rtol=1e-7, + ) + def test_compute_stats_from_atomic(self) -> None: bias, std = compute_stats_from_atomic(self.output, self.atype) np.testing.assert_allclose(bias, self.mean) diff --git a/source/tests/consistent/fitting/test_property.py b/source/tests/consistent/fitting/test_property.py index 3abd672c88..4c359026c7 100644 --- a/source/tests/consistent/fitting/test_property.py +++ b/source/tests/consistent/fitting/test_property.py @@ -86,6 +86,7 @@ def data(self) -> dict: "seed": 20240217, "task_dim": task_dim, "intensive": intensive, + "property_name": "foo", } @property @@ -186,7 +187,7 @@ def eval_pt(self, pt_obj: Any) -> Any: aparam=torch.from_numpy(self.aparam).to(device=PT_DEVICE) if numb_aparam else None, - )["property"] + )[pt_obj.var_name] .detach() .cpu() .numpy() @@ -207,7 +208,7 @@ def eval_dp(self, dp_obj: Any) -> Any: self.atype.reshape(1, -1), fparam=self.fparam if numb_fparam else None, aparam=self.aparam if numb_aparam else None, - )["property"] + )[dp_obj.var_name] def eval_jax(self, jax_obj: Any) -> Any: ( @@ -225,7 +226,7 @@ def eval_jax(self, jax_obj: Any) -> Any: jnp.asarray(self.atype.reshape(1, -1)), fparam=jnp.asarray(self.fparam) if numb_fparam else None, aparam=jnp.asarray(self.aparam) if numb_aparam else None, - )["property"] + )[jax_obj.var_name] ) def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: @@ -244,7 +245,7 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: array_api_strict.asarray(self.atype.reshape(1, -1)), fparam=array_api_strict.asarray(self.fparam) if numb_fparam else None, aparam=array_api_strict.asarray(self.aparam) if numb_aparam else None, - )["property"] + )[array_api_strict_obj.var_name] ) def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: diff --git a/source/tests/consistent/model/test_property.py b/source/tests/consistent/model/test_property.py index 29786fb247..75aded98fd 100644 --- a/source/tests/consistent/model/test_property.py +++ b/source/tests/consistent/model/test_property.py @@ -56,6 +56,7 @@ def data(self) -> dict: "fitting_net": { "type": "property", "neuron": [4, 4, 4], + "property_name": "foo", "resnet_dt": True, "numb_fparam": 0, "precision": "float64", @@ -182,14 +183,15 @@ def eval_jax(self, jax_obj: Any) -> Any: def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: # shape not matched. ravel... + property_name = self.data["fitting_net"]["property_name"] if backend in {self.RefBackend.DP, self.RefBackend.JAX}: return ( - ret["property_redu"].ravel(), - ret["property"].ravel(), + ret[f"{property_name}_redu"].ravel(), + ret[property_name].ravel(), ) elif backend is self.RefBackend.PT: return ( - ret["property"].ravel(), - ret["atom_property"].ravel(), + ret[property_name].ravel(), + ret[f"atom_{property_name}"].ravel(), ) raise ValueError(f"Unknown backend: {backend}") diff --git a/source/tests/pd/model/test_permutation.py b/source/tests/pd/model/test_permutation.py index 4543348d3b..135c5ea819 100644 --- a/source/tests/pd/model/test_permutation.py +++ b/source/tests/pd/model/test_permutation.py @@ -331,10 +331,10 @@ }, "fitting_net": { "type": "property", + "property_name": "band_property", "task_dim": 3, "neuron": [24, 24, 24], "resnet_dt": True, - "bias_method": "normal", "intensive": True, "seed": 1, }, diff --git a/source/tests/pt/model/test_permutation.py b/source/tests/pt/model/test_permutation.py index 5c7b8db9a4..e4eb47a540 100644 --- a/source/tests/pt/model/test_permutation.py +++ b/source/tests/pt/model/test_permutation.py @@ -331,9 +331,9 @@ "fitting_net": { "type": "property", "task_dim": 3, + "property_name": "band_property", "neuron": [24, 24, 24], "resnet_dt": True, - "bias_method": "normal", "intensive": True, "seed": 1, }, diff --git a/source/tests/pt/model/test_property_fitting.py b/source/tests/pt/model/test_property_fitting.py index 305d1be951..6825924bc1 100644 --- a/source/tests/pt/model/test_property_fitting.py +++ b/source/tests/pt/model/test_property_fitting.py @@ -61,7 +61,7 @@ def test_consistency( self.atype_ext[:, : self.nloc], dtype=int, device=env.DEVICE ) - for nfp, nap, bias_atom_p, intensive, bias_method in itertools.product( + for nfp, nap, bias_atom_p, intensive in itertools.product( [0, 3], [0, 4], [ @@ -69,18 +69,17 @@ def test_consistency( np.array([[11, 12, 13, 4, 15], [16, 17, 18, 9, 20]]), ], [True, False], - ["normal", "no_bias"], ): ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, task_dim=5, + property_name="foo", numb_fparam=nfp, numb_aparam=nap, mixed_types=self.dd0.mixed_types(), bias_atom_p=bias_atom_p, intensive=intensive, - bias_method=bias_method, seed=GLOBAL_SEED, ).to(env.DEVICE) @@ -120,36 +119,35 @@ def test_consistency( aparam=to_numpy_array(iap), ) np.testing.assert_allclose( - to_numpy_array(ret0["property"]), - ret1["property"], + to_numpy_array(ret0[ft0.var_name]), + ret1[ft1.var_name], ) np.testing.assert_allclose( - to_numpy_array(ret0["property"]), - to_numpy_array(ret2["property"]), + to_numpy_array(ret0[ft0.var_name]), + to_numpy_array(ret2[ft2.var_name]), ) np.testing.assert_allclose( - to_numpy_array(ret0["property"]), - ret3["property"], + to_numpy_array(ret0[ft0.var_name]), + ret3[ft3.var_name], ) def test_jit( self, ) -> None: - for nfp, nap, intensive, bias_method in itertools.product( + for nfp, nap, intensive in itertools.product( [0, 3], [0, 4], [True, False], - ["normal", "no_bias"], ): ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, task_dim=5, + property_name="foo", numb_fparam=nfp, numb_aparam=nap, mixed_types=self.dd0.mixed_types(), intensive=intensive, - bias_method=bias_method, seed=GLOBAL_SEED, ).to(env.DEVICE) torch.jit.script(ft0) @@ -201,6 +199,7 @@ def test_trans(self) -> None: self.nt, self.dd0.dim_out, task_dim=11, + property_name="bar", numb_fparam=0, numb_aparam=0, mixed_types=self.dd0.mixed_types(), @@ -229,7 +228,7 @@ def test_trans(self) -> None: ) ret0 = ft0(rd0, atype, gr0, fparam=None, aparam=None) - res.append(ret0["property"]) + res.append(ret0[ft0.var_name]) np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1])) @@ -257,21 +256,20 @@ def test_rot(self) -> None: # use larger cell to rotate only coord and shift to the center of cell cell_rot = 10.0 * torch.eye(3, dtype=dtype, device=env.DEVICE) - for nfp, nap, intensive, bias_method in itertools.product( + for nfp, nap, intensive in itertools.product( [0, 3], [0, 4], [True, False], - ["normal", "no_bias"], ): ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, # dim_descrpt - task_dim=9, + task_dim=5, + property_name="bar", numb_fparam=nfp, numb_aparam=nap, mixed_types=self.dd0.mixed_types(), intensive=intensive, - bias_method=bias_method, seed=GLOBAL_SEED, ).to(env.DEVICE) if nfp > 0: @@ -312,7 +310,7 @@ def test_rot(self) -> None: ) ret0 = ft0(rd0, atype, gr0, fparam=ifp, aparam=iap) - res.append(ret0["property"]) + res.append(ret0[ft0.var_name]) np.testing.assert_allclose( to_numpy_array(res[1]), to_numpy_array(res[0]), @@ -324,6 +322,7 @@ def test_permu(self) -> None: self.nt, self.dd0.dim_out, task_dim=8, + property_name="abc", numb_fparam=0, numb_aparam=0, mixed_types=self.dd0.mixed_types(), @@ -353,7 +352,7 @@ def test_permu(self) -> None: ) ret0 = ft0(rd0, atype, gr0, fparam=None, aparam=None) - res.append(ret0["property"]) + res.append(ret0[ft0.var_name]) np.testing.assert_allclose( to_numpy_array(res[0][:, idx_perm]), @@ -372,6 +371,7 @@ def test_trans(self) -> None: self.nt, self.dd0.dim_out, task_dim=11, + property_name="foo", numb_fparam=0, numb_aparam=0, mixed_types=self.dd0.mixed_types(), @@ -400,7 +400,7 @@ def test_trans(self) -> None: ) ret0 = ft0(rd0, atype, gr0, fparam=None, aparam=None) - res.append(ret0["property"]) + res.append(ret0[ft0.var_name]) np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1])) @@ -422,6 +422,7 @@ def setUp(self) -> None: self.nt, self.dd0.dim_out, task_dim=3, + property_name="bar", numb_fparam=0, numb_aparam=0, mixed_types=self.dd0.mixed_types(), diff --git a/source/tests/pt/property/double/nopbc b/source/tests/pt/property/double/nopbc new file mode 100644 index 0000000000..e69de29bb2 diff --git a/source/tests/pt/property/double/set.000000/band_property.npy b/source/tests/pt/property/double/set.000000/band_property.npy new file mode 100644 index 0000000000..042c1a8b0d Binary files /dev/null and b/source/tests/pt/property/double/set.000000/band_property.npy differ diff --git a/source/tests/pt/property/double/set.000000/coord.npy b/source/tests/pt/property/double/set.000000/coord.npy new file mode 100644 index 0000000000..9c781a81f3 Binary files /dev/null and b/source/tests/pt/property/double/set.000000/coord.npy differ diff --git a/source/tests/pt/property/double/set.000000/real_atom_types.npy b/source/tests/pt/property/double/set.000000/real_atom_types.npy new file mode 100644 index 0000000000..3bfe0abd94 Binary files /dev/null and b/source/tests/pt/property/double/set.000000/real_atom_types.npy differ diff --git a/source/tests/pt/property/double/type.raw b/source/tests/pt/property/double/type.raw new file mode 100644 index 0000000000..d677b495ec --- /dev/null +++ b/source/tests/pt/property/double/type.raw @@ -0,0 +1,20 @@ +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 diff --git a/source/tests/pt/property/double/type_map.raw b/source/tests/pt/property/double/type_map.raw new file mode 100644 index 0000000000..c8a39f3a9e --- /dev/null +++ b/source/tests/pt/property/double/type_map.raw @@ -0,0 +1,4 @@ +H +C +N +O diff --git a/source/tests/pt/property/input.json b/source/tests/pt/property/input.json index 4e005f8277..44bc1e6005 100644 --- a/source/tests/pt/property/input.json +++ b/source/tests/pt/property/input.json @@ -27,6 +27,7 @@ "fitting_net": { "type": "property", "intensive": true, + "property_name": "band_property", "task_dim": 3, "neuron": [ 100, diff --git a/source/tests/pt/property/single/set.000000/property.npy b/source/tests/pt/property/single/set.000000/band_property.npy similarity index 100% rename from source/tests/pt/property/single/set.000000/property.npy rename to source/tests/pt/property/single/set.000000/band_property.npy diff --git a/source/tests/pt/test_dp_test.py b/source/tests/pt/test_dp_test.py index dbec472cc0..c2915c7ee7 100644 --- a/source/tests/pt/test_dp_test.py +++ b/source/tests/pt/test_dp_test.py @@ -183,7 +183,7 @@ def test_dp_test_1_frame(self) -> None: pred_property = np.loadtxt(self.detail_file + ".property.out.0")[:, 1] np.testing.assert_almost_equal( pred_property, - to_numpy_array(result["property"])[0], + to_numpy_array(result[model.get_var_name()])[0], ) def tearDown(self) -> None: diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index 1fbd01c39f..ad52c5db16 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -464,7 +464,7 @@ def setUp(self) -> None: property_input = str(Path(__file__).parent / "property/input.json") with open(property_input) as f: self.config_property = json.load(f) - prop_data_file = [str(Path(__file__).parent / "property/single")] + prop_data_file = [str(Path(__file__).parent / "property/double")] self.config_property["training"]["training_data"]["systems"] = prop_data_file self.config_property["training"]["validation_data"]["systems"] = prop_data_file self.config_property["model"]["descriptor"] = deepcopy(model_dpa1["descriptor"]) diff --git a/source/tests/universal/common/cases/model/model.py b/source/tests/universal/common/cases/model/model.py index cee69d9d6c..06ddd90970 100644 --- a/source/tests/universal/common/cases/model/model.py +++ b/source/tests/universal/common/cases/model/model.py @@ -165,7 +165,7 @@ def setUpClass(cls) -> None: cls.expected_dim_aparam = 0 cls.expected_sel_type = [0, 1] cls.expected_aparam_nall = False - cls.expected_model_output_type = ["property", "mask"] + cls.expected_model_output_type = ["band_prop", "mask"] cls.model_output_equivariant = [] cls.expected_sel = [46, 92] cls.expected_sel_mix = sum(cls.expected_sel) diff --git a/source/tests/universal/dpmodel/fitting/test_fitting.py b/source/tests/universal/dpmodel/fitting/test_fitting.py index db199c02a3..2fe0060003 100644 --- a/source/tests/universal/dpmodel/fitting/test_fitting.py +++ b/source/tests/universal/dpmodel/fitting/test_fitting.py @@ -208,6 +208,8 @@ def FittingParamProperty( "dim_descrpt": dim_descrpt, "mixed_types": mixed_types, "type_map": type_map, + "task_dim": 3, + "property_name": "band_prop", "exclude_types": exclude_types, "seed": GLOBAL_SEED, "precision": precision, diff --git a/source/tests/universal/dpmodel/loss/test_loss.py b/source/tests/universal/dpmodel/loss/test_loss.py index 6473c159da..79c67cdba4 100644 --- a/source/tests/universal/dpmodel/loss/test_loss.py +++ b/source/tests/universal/dpmodel/loss/test_loss.py @@ -189,11 +189,14 @@ def LossParamTensor( def LossParamProperty(): key_to_pref_map = { - "property": 1.0, + "foo": 1.0, } input_dict = { "key_to_pref_map": key_to_pref_map, - "task_dim": 2, + "var_name": "foo", + "out_bias": [0.1, 0.5, 1.2, -0.1, -10], + "out_std": [8, 10, 0.001, -0.2, -10], + "task_dim": 5, } return input_dict