diff --git a/deepmd/dpmodel/fitting/__init__.py b/deepmd/dpmodel/fitting/__init__.py index 866a710a3b..2a35a68d35 100644 --- a/deepmd/dpmodel/fitting/__init__.py +++ b/deepmd/dpmodel/fitting/__init__.py @@ -17,6 +17,9 @@ from .polarizability_fitting import ( PolarFitting, ) +from .property_fitting import ( + PropertyFittingNet, +) __all__ = [ "InvarFitting", @@ -25,4 +28,5 @@ "EnergyFittingNet", "PolarFitting", "DOSFittingNet", + "PropertyFittingNet", ] diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py new file mode 100644 index 0000000000..725108688d --- /dev/null +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import copy +from typing import ( + TYPE_CHECKING, + Any, + Callable, + List, + Optional, + Union, +) + +from deepmd.dpmodel.common import ( + DEFAULT_PRECISION, +) +from deepmd.dpmodel.fitting.invar_fitting import ( + InvarFitting, +) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.utils.path import ( + DPPath, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.fitting.general_fitting import ( + GeneralFitting, + ) + +from deepmd.utils.version import ( + check_version_compatibility, +) + + +@InvarFitting.register("property") +class PropertyFittingNet(InvarFitting): + def __init__( + self, + ntypes: int, + dim_descrpt: int, + task_dim: int = 1, + neuron: List[int] = [128, 128, 128], + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + rcond: Optional[float] = None, + tot_ener_zero: bool = False, + trainable: Optional[List[bool]] = None, + atom_ener: Optional[List[float]] = None, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + layer_name: Optional[List[Optional[str]]] = None, + use_aparam_as_mask: bool = False, + spin: Any = None, + mixed_types: bool = False, + exclude_types: List[int] = [], + # not used + seed: Optional[int] = None, + ): + self.task_dim = task_dim + super().__init__( + var_name="property", + ntypes=ntypes, + dim_descrpt=dim_descrpt, + dim_out=task_dim, + neuron=neuron, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + rcond=rcond, + tot_ener_zero=tot_ener_zero, + trainable=trainable, + atom_ener=atom_ener, + activation_function=activation_function, + precision=precision, + layer_name=layer_name, + use_aparam_as_mask=use_aparam_as_mask, + spin=spin, + mixed_types=mixed_types, + exclude_types=exclude_types, + ) + + @classmethod + def deserialize(cls, data: dict) -> "GeneralFitting": + data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("var_name") + data.pop("dim_out") + return super().deserialize(data) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return {**super().serialize(), "type": "property", "task_dim": self.task_dim} + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reduciable=True, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def compute_output_stats( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + stat_file_path: Optional[DPPath] = None, + ): + """ + Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The path to the stat file. + + """ + pass + + # make jit happy with torch 2.0.0 + exclude_types: List[int] diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index aae2082e13..5d1cc6ed43 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -73,6 +73,8 @@ 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", } diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py new file mode 100644 index 0000000000..ac88f0dd9c --- /dev/null +++ b/deepmd/infer/deep_property.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, + Dict, + List, + Optional, + Tuple, + Union, +) + +import numpy as np + +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + ModelOutputDef, + OutputVariableDef, +) + +from .deep_eval import ( + DeepEval, +) + + +class DeepProperty(DeepEval): + """Properties of structures. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + + @property + def output_def(self) -> ModelOutputDef: + """Get the output definition of this model.""" + return ModelOutputDef( + FittingOutputDef( + [ + OutputVariableDef( + "property", + shape=[-1], + reduciable=True, + atomic=True, + ), + ] + ) + ) + + @property + def numb_task(self) -> int: + """Get the number of task.""" + return self.get_numb_task() + + def eval( + self, + coords: np.ndarray, + cells: Optional[np.ndarray], + atom_types: Union[List[int], np.ndarray], + atomic: bool = False, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + mixed_type: bool = False, + **kwargs: Dict[str, Any], + ) -> Tuple[np.ndarray, ...]: + """Evaluate properties. If atomic is True, also return atomic property. + + Parameters + ---------- + coords : np.ndarray + The coordinates of the atoms, in shape (nframes, natoms, 3). + cells : np.ndarray + The cell vectors of the system, in shape (nframes, 9). If the system + is not periodic, set it to None. + atom_types : List[int] or np.ndarray + The types of the atoms. If mixed_type is False, the shape is (natoms,); + otherwise, the shape is (nframes, natoms). + atomic : bool, optional + Whether to return atomic property, by default False. + fparam : np.ndarray, optional + The frame parameters, by default None. + aparam : np.ndarray, optional + The atomic parameters, by default None. + mixed_type : bool, optional + Whether the atom_types is mixed type, by default False. + **kwargs : Dict[str, Any] + Keyword arguments. + + Returns + ------- + property + The properties of the system, in shape (nframes, num_tasks). + """ + ( + coords, + cells, + atom_types, + fparam, + aparam, + nframes, + natoms, + ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) + results = self.deep_eval.eval( + coords, + cells, + atom_types, + atomic, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + atomic_property = results["property"].reshape(nframes, natoms, -1) + property = results["property_redu"].reshape(nframes, -1) + + if atomic: + return ( + property, + atomic_property, + ) + else: + return (property,) + + def get_numb_task(self) -> int: + return self.deep_eval.get_numb_task() + + +__all__ = ["DeepProperty"] diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 1262a56310..22dcdf271c 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -35,6 +35,9 @@ from deepmd.infer.deep_pot import ( DeepPot, ) +from deepmd.infer.deep_property import ( + DeepProperty, +) from deepmd.infer.deep_wfc import ( DeepWFC, ) @@ -180,6 +183,8 @@ def model_type(self) -> "DeepEvalWrapper": return DeepGlobalPolar elif "wfc" in model_output_type: return DeepWFC + elif "property" in model_output_type: + return DeepProperty else: raise RuntimeError("Unknown model type") diff --git a/deepmd/pt/loss/__init__.py b/deepmd/pt/loss/__init__.py index 9c8bbc9a2a..0656db8dbe 100644 --- a/deepmd/pt/loss/__init__.py +++ b/deepmd/pt/loss/__init__.py @@ -11,6 +11,9 @@ from .loss import ( TaskLoss, ) +from .property import ( + PropertyLoss, +) from .tensor import ( TensorLoss, ) @@ -21,4 +24,5 @@ "EnergySpinLoss", "TensorLoss", "TaskLoss", + "PropertyLoss", ] diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py new file mode 100644 index 0000000000..a03dc82173 --- /dev/null +++ b/deepmd/pt/loss/property.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + List, +) + +import torch +import torch.nn.functional as F + +from deepmd.pt.loss.loss import ( + TaskLoss, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.utils.data import ( + DataRequirementItem, +) + +log = logging.getLogger(__name__) + + +class PropertyLoss(TaskLoss): + def __init__( + self, + task_dim, + loss_func: str = "smooth_mae", + metric: list = ["mae"], + **kwargs, + ): + r"""Construct a layer to compute loss on property. + + Parameters + ---------- + task_dim : float + The learning rate at the start of the training. + loss_func : str + The loss function, such as "smooth_mae", "mae", "rmse" + metric : list + The metric such as mae,rmse which will be printed. + **kwargs + Other keyword arguments. + """ + super().__init__() + self.loss_func = loss_func + self.metric = metric + self.task_dim = task_dim + self.mean = kwargs.get("mean", 0) + self.std = kwargs.get("std", 1) + self.beta = kwargs.get("beta", 1.00) + + def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): + """Return loss on properties . + + Parameters + ---------- + input_dict : dict[str, torch.Tensor] + Model inputs. + model : torch.nn.Module + Model to be used to output the predictions. + label : dict[str, torch.Tensor] + Labels. + natoms : int + The local atom number. + + Returns + ------- + model_pred: dict[str, torch.Tensor] + Model predictions. + loss: torch.Tensor + Loss for model to minimize. + more_loss: dict[str, torch.Tensor] + 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 + loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] + more_loss = {} + + label_mean = torch.tensor( + self.mean, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) + label_std = torch.tensor( + self.std, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) + + # loss + if self.loss_func == "smooth_mae": + loss += F.smooth_l1_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="sum", + beta=self.beta, + ) + elif self.loss_func == "mae": + loss += F.l1_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="sum", + ) + elif self.loss_func == "mse": + loss += F.mse_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="sum", + ) + elif self.loss_func == "rmse": + loss += torch.sqrt( + F.mse_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="mean", + ) + ) + else: + raise RuntimeError(f"Unknown loss function : {self.func}") + + # more loss + if "smooth_mae" in self.metric: + more_loss["smooth_mae"] = F.smooth_l1_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + beta=self.beta, + ).detach() + if "mae" in self.metric: + more_loss["mae"] = F.l1_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + ).detach() + if "mse" in self.metric: + more_loss["mse"] = F.mse_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + ).detach() + if "rmse" in self.metric: + more_loss["rmse"] = torch.sqrt( + F.mse_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + ) + ).detach() + + return model_pred, loss, more_loss + + @property + def label_requirement(self) -> List[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + label_requirement = [] + label_requirement.append( + DataRequirementItem( + "property", + ndof=self.task_dim, + atomic=False, + must=False, + high_prec=True, + ) + ) + return label_requirement diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index d7b3c4f4e2..1b29eebd9f 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -28,6 +28,9 @@ from deepmd.pt.model.task.polarizability import ( PolarFittingNet, ) +from deepmd.pt.model.task.property import ( + PropertyFittingNet, +) from .make_model import ( make_model, @@ -57,6 +60,9 @@ def __new__( from deepmd.pt.model.model.polar_model import ( PolarModel, ) + from deepmd.pt.model.model.property_model import ( + PropertyModel, + ) if atomic_model_ is not None: fitting = atomic_model_.fitting_net @@ -76,6 +82,8 @@ def __new__( cls = PolarModel elif isinstance(fitting, DOSFittingNet): cls = DOSModel + elif isinstance(fitting, PropertyFittingNet): + cls = PropertyModel # else: unknown fitting type, fall back to DPModel return super().__new__(cls) diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py new file mode 100644 index 0000000000..8ac62e6106 --- /dev/null +++ b/deepmd/pt/model/model/property_model.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Dict, + Optional, +) + +import torch + +from .dp_model import ( + DPModel, +) + + +class PropertyModel(DPModel): + model_type = "property" + + def __init__( + self, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + + def forward( + self, + coord, + atype, + box: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ) -> Dict[str, torch.Tensor]: + model_ret = self.forward_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_property"] = model_ret["property"] + natoms = model_predict["atom_property"].shape[1] + model_predict["property"] = model_ret["property_redu"] / natoms + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + @torch.jit.export + def forward_lower( + self, + extended_coord, + extended_atype, + nlist, + mapping: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ): + model_ret = self.forward_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_property"] = model_ret["property"] + natoms = model_predict["atom_property"].shape[1] + model_predict["property"] = model_ret["property_redu"] / natoms + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict diff --git a/deepmd/pt/model/task/__init__.py b/deepmd/pt/model/task/__init__.py index 9430ede766..9906d435cb 100644 --- a/deepmd/pt/model/task/__init__.py +++ b/deepmd/pt/model/task/__init__.py @@ -21,6 +21,9 @@ from .polarizability import ( PolarFittingNet, ) +from .property import ( + PropertyFittingNet, +) from .type_predict import ( TypePredictNet, ) @@ -35,4 +38,5 @@ "BaseFitting", "TypePredictNet", "PolarFittingNet", + "PropertyFittingNet", ] diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index afb1d73658..787ed8740e 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -192,7 +192,7 @@ def forward( fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, ): - """Based on embedding net output, alculate total energy. + """Based on embedding net output, calculate total energy. Args: - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.dim_descrpt]. diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py new file mode 100644 index 0000000000..b4c18ad378 --- /dev/null +++ b/deepmd/pt/model/task/property.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import copy +import logging +from typing import ( + Callable, + List, + Optional, + Union, +) + +import torch + +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.pt.model.task.fitting import ( + Fitting, + GeneralFitting, +) +from deepmd.pt.model.task.invar_fitting import ( + InvarFitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, +) +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE + +log = logging.getLogger(__name__) + + +@Fitting.register("property") +class PropertyFittingNet(InvarFitting): + """Construct a fitting net for energy. + + Parameters + ---------- + ntypes : int + Element count. + dim_descrpt : int + Embedding width per atom. + task_dim : int + The dimension of outputs of fitting net + 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. + resnet_dt : bool + Using time-step in the ResNet construction. + numb_fparam : int + Number of frame parameters. + numb_aparam : int + Number of atomic parameters. + activation_function : str + Activation function. + precision : str + Numerical precision. + mixed_types : bool + If true, use a uniform fitting net for all atom types, otherwise use + different fitting nets for different atom types. + seed : int, optional + Random seed. + + """ + + def __init__( + self, + ntypes: int, + dim_descrpt: int, + task_dim: int = 1, + neuron: List[int] = [128, 128, 128], + bias_atom_p: Optional[torch.Tensor] = None, + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + mixed_types: bool = True, + seed: Optional[int] = None, + **kwargs, + ): + self.task_dim = task_dim + super().__init__( + var_name="property", + ntypes=ntypes, + dim_descrpt=dim_descrpt, + dim_out=task_dim, + neuron=neuron, + bias_atom_e=bias_atom_p, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + activation_function=activation_function, + precision=precision, + mixed_types=mixed_types, + seed=seed, + **kwargs, + ) + + @classmethod + def deserialize(cls, data: dict) -> "GeneralFitting": + data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("dim_out") + data.pop("var_name") + return super().deserialize(data) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return {**super().serialize(), "type": "property", "task_dim": self.task_dim} + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reduciable=True, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def compute_output_stats( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + stat_file_path: Optional[DPPath] = None, + ): + """ + Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The path to the stat file. + + """ + pass + + # make jit happy with torch 2.0.0 + exclude_types: List[int] diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 1bea24d717..6ca718724b 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -27,6 +27,7 @@ DenoiseLoss, EnergySpinLoss, EnergyStdLoss, + PropertyLoss, TensorLoss, ) from deepmd.pt.model.model import ( @@ -297,6 +298,10 @@ def get_loss(loss_params, start_lr, _ntypes, _model): label_name = "polarizability" loss_params["label_name"] = label_name return TensorLoss(**loss_params) + elif loss_type == "property": + task_dim = _model.model_output_def()["property"].output_size + loss_params["task_dim"] = task_dim + return PropertyLoss(**loss_params) else: raise NotImplementedError diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 564039ccd0..42ab8b8a00 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1181,6 +1181,37 @@ def fitting_dipole(): ] +@fitting_args_plugin.register("property", doc=doc_only_pt_supported) +def fitting_property(): + 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.' + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' + doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." + doc_seed = "Random seed for parameter initialization of the fitting net" + doc_task_dim = "The dimension of outputs of fitting net" + return [ + Argument( + "neuron", + List[int], + optional=True, + default=[120, 120, 120], + alias=["n_neuron"], + doc=doc_neuron, + ), + Argument( + "activation_function", + str, + optional=True, + default="tanh", + doc=doc_activation_function, + ), + Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), + Argument("precision", str, optional=True, default="default", doc=doc_precision), + Argument("seed", [int, None], optional=True, doc=doc_seed), + Argument("task_dim", int, optional=True, default=1, doc=doc_task_dim), + ] + + # YWolfeee: Delete global polar mode, merge it into polar mode and use loss setting to support. def fitting_variant_type_args(): doc_descrpt_type = "The type of the fitting. See explanation below. \n\n\ @@ -1941,6 +1972,52 @@ def loss_tensor(): ] +@loss_args_plugin.register("property") +def loss_property(): + doc_loss_func = "The loss function, such as 'mae','smooth_mae'." + doc_metric = "The metric such as 'mae','smooth_mae' which will be printed." + doc_mean = "The averge value of target." + doc_std = "The standard deviation of the target." + doc_beta = "The 'beta' parameter in 'smooth_mae' loss." + return [ + Argument( + "loss_func", + str, + optional=True, + default="smooth_mae", + doc=doc_loss_func, + ), + Argument( + "metric", + list, + optional=True, + default=["mae"], + doc=doc_metric, + ), + Argument( + "mean", + [float, int, list], + optional=True, + default=0, + doc=doc_mean, + ), + Argument( + "std", + [float, int, list], + optional=True, + default=1, + doc=doc_std, + ), + Argument( + "beta", + [float, int], + optional=True, + default=1.00, + doc=doc_beta, + ), + ] + + def loss_variant_type_args(): doc_loss = "The type of the loss. When the fitting type is `ener`, the loss type should be set to `ener` or left unset. When the fitting type is `dipole` or `polar`, the loss type should be set to `tensor`." diff --git a/examples/property/data/3/nopbc b/examples/property/data/3/nopbc new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/property/data/3/set.000000/coord.npy b/examples/property/data/3/set.000000/coord.npy new file mode 100644 index 0000000000..324d8f63ba Binary files /dev/null and b/examples/property/data/3/set.000000/coord.npy differ diff --git a/examples/property/data/3/set.000000/property.npy b/examples/property/data/3/set.000000/property.npy new file mode 100644 index 0000000000..fcd3092954 Binary files /dev/null and b/examples/property/data/3/set.000000/property.npy differ diff --git a/examples/property/data/3/set.000000/real_atom_types.npy b/examples/property/data/3/set.000000/real_atom_types.npy new file mode 100644 index 0000000000..fb312ec6b8 Binary files /dev/null and b/examples/property/data/3/set.000000/real_atom_types.npy differ diff --git a/examples/property/data/3/type.raw b/examples/property/data/3/type.raw new file mode 100644 index 0000000000..bb0b1cf658 --- /dev/null +++ b/examples/property/data/3/type.raw @@ -0,0 +1,3 @@ +0 +0 +0 diff --git a/examples/property/data/3/type_map.raw b/examples/property/data/3/type_map.raw new file mode 100644 index 0000000000..e37cba22a9 --- /dev/null +++ b/examples/property/data/3/type_map.raw @@ -0,0 +1,118 @@ +H +He +Li +Be +B +C +N +O +F +Ne +Na +Mg +Al +Si +P +S +Cl +Ar +K +Ca +Sc +Ti +V +Cr +Mn +Fe +Co +Ni +Cu +Zn +Ga +Ge +As +Se +Br +Kr +Rb +Sr +Y +Zr +Nb +Mo +Tc +Ru +Rh +Pd +Ag +Cd +In +Sn +Sb +Te +I +Xe +Cs +Ba +La +Ce +Pr +Nd +Pm +Sm +Eu +Gd +Tb +Dy +Ho +Er +Tm +Yb +Lu +Hf +Ta +W +Re +Os +Ir +Pt +Au +Hg +Tl +Pb +Bi +Po +At +Rn +Fr +Ra +Ac +Th +Pa +U +Np +Pu +Am +Cm +Bk +Cf +Es +Fm +Md +No +Lr +Rf +Db +Sg +Bh +Hs +Mt +Ds +Rg +Cn +Nh +Fl +Mc +Lv +Ts +Og diff --git a/examples/property/data/5/nopbc b/examples/property/data/5/nopbc new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/property/data/5/set.000000/coord.npy b/examples/property/data/5/set.000000/coord.npy new file mode 100644 index 0000000000..4aa899142c Binary files /dev/null and b/examples/property/data/5/set.000000/coord.npy differ diff --git a/examples/property/data/5/set.000000/property.npy b/examples/property/data/5/set.000000/property.npy new file mode 100644 index 0000000000..4b9d4f9453 Binary files /dev/null and b/examples/property/data/5/set.000000/property.npy differ diff --git a/examples/property/data/5/set.000000/real_atom_types.npy b/examples/property/data/5/set.000000/real_atom_types.npy new file mode 100644 index 0000000000..497a9722c4 Binary files /dev/null and b/examples/property/data/5/set.000000/real_atom_types.npy differ diff --git a/examples/property/data/5/type.raw b/examples/property/data/5/type.raw new file mode 100644 index 0000000000..229972f292 --- /dev/null +++ b/examples/property/data/5/type.raw @@ -0,0 +1,5 @@ +0 +0 +0 +0 +0 diff --git a/examples/property/data/5/type_map.raw b/examples/property/data/5/type_map.raw new file mode 100644 index 0000000000..e37cba22a9 --- /dev/null +++ b/examples/property/data/5/type_map.raw @@ -0,0 +1,118 @@ +H +He +Li +Be +B +C +N +O +F +Ne +Na +Mg +Al +Si +P +S +Cl +Ar +K +Ca +Sc +Ti +V +Cr +Mn +Fe +Co +Ni +Cu +Zn +Ga +Ge +As +Se +Br +Kr +Rb +Sr +Y +Zr +Nb +Mo +Tc +Ru +Rh +Pd +Ag +Cd +In +Sn +Sb +Te +I +Xe +Cs +Ba +La +Ce +Pr +Nd +Pm +Sm +Eu +Gd +Tb +Dy +Ho +Er +Tm +Yb +Lu +Hf +Ta +W +Re +Os +Ir +Pt +Au +Hg +Tl +Pb +Bi +Po +At +Rn +Fr +Ra +Ac +Th +Pa +U +Np +Pu +Am +Cm +Bk +Cf +Es +Fm +Md +No +Lr +Rf +Db +Sg +Bh +Hs +Mt +Ds +Rg +Cn +Nh +Fl +Mc +Lv +Ts +Og diff --git a/examples/property/data/6/nopbc b/examples/property/data/6/nopbc new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/property/data/6/set.000000/coord.npy b/examples/property/data/6/set.000000/coord.npy new file mode 100644 index 0000000000..36a65100a6 Binary files /dev/null and b/examples/property/data/6/set.000000/coord.npy differ diff --git a/examples/property/data/6/set.000000/property.npy b/examples/property/data/6/set.000000/property.npy new file mode 100644 index 0000000000..66162f99db Binary files /dev/null and b/examples/property/data/6/set.000000/property.npy differ diff --git a/examples/property/data/6/set.000000/real_atom_types.npy b/examples/property/data/6/set.000000/real_atom_types.npy new file mode 100644 index 0000000000..0b0f7e0d55 Binary files /dev/null and b/examples/property/data/6/set.000000/real_atom_types.npy differ diff --git a/examples/property/data/6/type.raw b/examples/property/data/6/type.raw new file mode 100644 index 0000000000..f7eb44d66e --- /dev/null +++ b/examples/property/data/6/type.raw @@ -0,0 +1,6 @@ +0 +0 +0 +0 +0 +0 diff --git a/examples/property/data/6/type_map.raw b/examples/property/data/6/type_map.raw new file mode 100644 index 0000000000..e37cba22a9 --- /dev/null +++ b/examples/property/data/6/type_map.raw @@ -0,0 +1,118 @@ +H +He +Li +Be +B +C +N +O +F +Ne +Na +Mg +Al +Si +P +S +Cl +Ar +K +Ca +Sc +Ti +V +Cr +Mn +Fe +Co +Ni +Cu +Zn +Ga +Ge +As +Se +Br +Kr +Rb +Sr +Y +Zr +Nb +Mo +Tc +Ru +Rh +Pd +Ag +Cd +In +Sn +Sb +Te +I +Xe +Cs +Ba +La +Ce +Pr +Nd +Pm +Sm +Eu +Gd +Tb +Dy +Ho +Er +Tm +Yb +Lu +Hf +Ta +W +Re +Os +Ir +Pt +Au +Hg +Tl +Pb +Bi +Po +At +Rn +Fr +Ra +Ac +Th +Pa +U +Np +Pu +Am +Cm +Bk +Cf +Es +Fm +Md +No +Lr +Rf +Db +Sg +Bh +Hs +Mt +Ds +Rg +Cn +Nh +Fl +Mc +Lv +Ts +Og diff --git a/examples/property/train/input_torch.json b/examples/property/train/input_torch.json new file mode 100644 index 0000000000..bc5ca389e0 --- /dev/null +++ b/examples/property/train/input_torch.json @@ -0,0 +1,214 @@ +{ + "_comment": "that's all", + "model": { + "type_map": [ + "H", + "He", + "Li", + "Be", + "B", + "C", + "N", + "O", + "F", + "Ne", + "Na", + "Mg", + "Al", + "Si", + "P", + "S", + "Cl", + "Ar", + "K", + "Ca", + "Sc", + "Ti", + "V", + "Cr", + "Mn", + "Fe", + "Co", + "Ni", + "Cu", + "Zn", + "Ga", + "Ge", + "As", + "Se", + "Br", + "Kr", + "Rb", + "Sr", + "Y", + "Zr", + "Nb", + "Mo", + "Tc", + "Ru", + "Rh", + "Pd", + "Ag", + "Cd", + "In", + "Sn", + "Sb", + "Te", + "I", + "Xe", + "Cs", + "Ba", + "La", + "Ce", + "Pr", + "Nd", + "Pm", + "Sm", + "Eu", + "Gd", + "Tb", + "Dy", + "Ho", + "Er", + "Tm", + "Yb", + "Lu", + "Hf", + "Ta", + "W", + "Re", + "Os", + "Ir", + "Pt", + "Au", + "Hg", + "Tl", + "Pb", + "Bi", + "Po", + "At", + "Rn", + "Fr", + "Ra", + "Ac", + "Th", + "Pa", + "U", + "Np", + "Pu", + "Am", + "Cm", + "Bk", + "Cf", + "Es", + "Fm", + "Md", + "No", + "Lr", + "Rf", + "Db", + "Sg", + "Bh", + "Hs", + "Mt", + "Ds", + "Rg", + "Cn", + "Nh", + "Fl", + "Mc", + "Lv", + "Ts", + "Og" + ], + "descriptor": { + "type": "dpa2", + "tebd_dim": 8, + "repinit_rcut": 9.0, + "repinit_rcut_smth": 8.0, + "repinit_nsel": 120, + "repformer_rcut": 4.0, + "repformer_rcut_smth": 3.5, + "repformer_nsel": 40, + "repinit_neuron": [ + 25, + 50, + 100 + ], + "repinit_axis_neuron": 12, + "repinit_activation": "tanh", + "repformer_nlayers": 12, + "repformer_g1_dim": 128, + "repformer_g2_dim": 32, + "repformer_attn2_hidden": 32, + "repformer_attn2_nhead": 4, + "repformer_attn1_hidden": 128, + "repformer_attn1_nhead": 4, + "repformer_axis_dim": 4, + "repformer_update_h2": false, + "repformer_update_g1_has_conv": true, + "repformer_update_g1_has_grrg": true, + "repformer_update_g1_has_drrd": true, + "repformer_update_g1_has_attn": true, + "repformer_update_g2_has_g1g1": true, + "repformer_update_g2_has_attn": true, + "repformer_attn2_has_gate": true, + "repformer_add_type_ebd_to_seq": false + }, + "fitting_net": { + "type": "property", + "task_dim": 3, + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "seed": 1, + "_comment": " that's all" + }, + "_comment": " that's all" + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.0002, + "stop_lr": 3.51e-08, + "_comment": "that's all" + }, + "loss": { + "type": "property", + "loss_func": "smooth_mae", + "metric": [ + "mae" + ], + "beta": 1.00, + "_comment": " that's all" + }, + "training": { + "stat_file": "./dpa2", + "training_data": { + "systems": [ + "../data/3", + "../data/6" + ], + "batch_size": 2, + "_comment": "that's all" + }, + "validation_data": { + "systems": [ + "../data/5" + ], + "batch_size": 1, + "_comment": "that's all" + }, + "numb_steps": 1000000, + "warmup_steps": 0, + "gradient_max_norm": 5.0, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 2000, + "_comment": "that's all" + } +} diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index 91bb9c0174..9366e707fa 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -49,6 +49,7 @@ p_examples / "water" / "se_e2_a" / "input_torch.json", p_examples / "water" / "se_atten" / "input_torch.json", p_examples / "water" / "dpa2" / "input_torch.json", + p_examples / "property" / "train" / "input_torch.json", ) diff --git a/source/tests/pt/model/test_property_fitting.py b/source/tests/pt/model/test_property_fitting.py new file mode 100644 index 0000000000..a62f6ffa32 --- /dev/null +++ b/source/tests/pt/model/test_property_fitting.py @@ -0,0 +1,349 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +import os +import unittest + +import numpy as np +import torch +from scipy.stats import ( + special_ortho_group, +) + +from deepmd.dpmodel.fitting import PropertyFittingNet as DPProperFittingNet +from deepmd.infer.deep_property import ( + DeepProperty, +) +from deepmd.pt.model.descriptor.se_a import ( + DescrptSeA, +) +from deepmd.pt.model.model.property_model import ( + PropertyModel, +) +from deepmd.pt.model.task.property import ( + PropertyFittingNet, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, +) + +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION + + +class TestPropertyFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + self.rng = np.random.default_rng() + self.nf, self.nloc, _ = self.nlist.shape + self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + self.scale = self.rng.uniform(0, 1, self.nt).tolist() + + def test_consistency( + self, + ): + rd0, gr, _, _, _ = self.dd0( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + ) + atype = torch.tensor( + self.atype_ext[:, : self.nloc], dtype=int, device=env.DEVICE + ) + + for mixed_types, nfp, nap, fit_diag, scale in itertools.product( + [True, False], + [0, 3], + [0, 4], + [True, False], + [None, self.scale], + ): + ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, + task_dim=5, + numb_fparam=nfp, + numb_aparam=nap, + mixed_types=mixed_types, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + + ft1 = DPProperFittingNet.deserialize(ft0.serialize()) + ft2 = PropertyFittingNet.deserialize(ft0.serialize()) + ft3 = DPProperFittingNet.deserialize(ft1.serialize()) + + if nfp > 0: + ifp = torch.tensor( + self.rng.normal(size=(self.nf, nfp)), dtype=dtype, device=env.DEVICE + ) + else: + ifp = None + if nap > 0: + iap = torch.tensor( + self.rng.normal(size=(self.nf, self.nloc, nap)), + dtype=dtype, + device=env.DEVICE, + ) + else: + iap = None + + ret0 = ft0(rd0, atype, gr, fparam=ifp, aparam=iap) + ret1 = ft1( + rd0.detach().cpu().numpy(), + atype.detach().cpu().numpy(), + gr.detach().cpu().numpy(), + fparam=to_numpy_array(ifp), + aparam=to_numpy_array(iap), + ) + ret2 = ft2(rd0, atype, gr, fparam=ifp, aparam=iap) + ret3 = ft3( + rd0.detach().cpu().numpy(), + atype.detach().cpu().numpy(), + gr.detach().cpu().numpy(), + fparam=to_numpy_array(ifp), + aparam=to_numpy_array(iap), + ) + np.testing.assert_allclose( + to_numpy_array(ret0["property"]), + ret1["property"], + ) + np.testing.assert_allclose( + to_numpy_array(ret0["property"]), + to_numpy_array(ret2["property"]), + ) + np.testing.assert_allclose( + to_numpy_array(ret0["property"]), + ret3["property"], + ) + + def test_jit( + self, + ): + for mixed_types, nfp, nap, fit_diag in itertools.product( + [True, False], + [0, 3], + [0, 4], + [True, False], + ): + ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, + task_dim=5, + numb_fparam=nfp, + numb_aparam=nap, + mixed_types=mixed_types, + fit_diag=fit_diag, + ).to(env.DEVICE) + torch.jit.script(ft0) + + +class TestInvariance(unittest.TestCase): + def setUp(self) -> None: + self.natoms = 5 + self.rcut = 4 + self.rcut_smth = 0.5 + self.sel = [46, 92, 4] + self.nf = 1 + self.nt = 3 + self.rng = np.random.default_rng() + self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device=env.DEVICE) + self.shift = torch.tensor([4, 4, 4], dtype=dtype, device=env.DEVICE) + self.atype = torch.tensor([0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE) + self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + self.cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE) + self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3, device=env.DEVICE) + self.scale = self.rng.uniform(0, 1, self.nt).tolist() + + def test_rot(self): + atype = self.atype.reshape(1, 5) + rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype, device=env.DEVICE) + coord_rot = torch.matmul(self.coord, rmat) + + for mixed_types, nfp, nap, fit_diag, scale in itertools.product( + [True, False], + [0, 3], + [0, 4], + [True, False], + [None, self.scale], + ): + ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, # dim_descrpt + task_dim=9, + numb_fparam=nfp, + numb_aparam=nap, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + if nfp > 0: + ifp = torch.tensor( + self.rng.normal(size=(self.nf, nfp)), dtype=dtype, device=env.DEVICE + ) + else: + ifp = None + if nap > 0: + iap = torch.tensor( + self.rng.normal(size=(self.nf, self.natoms, nap)), + dtype=dtype, + device=env.DEVICE, + ) + else: + iap = None + + res = [] + for xyz in [self.coord, coord_rot]: + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + xyz + self.shift, atype, self.rcut, self.sel, mixed_types + ) + + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) + + ret0 = ft0(rd0, extended_atype, gr0, fparam=ifp, aparam=iap) + res.append(ret0["property"]) + np.testing.assert_allclose( + to_numpy_array(res[1]), + to_numpy_array(res[0]), + ) + + def test_permu(self): + coord = torch.matmul(self.coord, self.cell) + for fit_diag, scale in itertools.product([True, False], [None, self.scale]): + ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, + task_dim=8, + numb_fparam=0, + numb_aparam=0, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + res = [] + for idx_perm in [[0, 1, 2, 3, 4], [1, 0, 4, 3, 2]]: + atype = self.atype[idx_perm].reshape(1, 5) + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + coord[idx_perm], atype, self.rcut, self.sel, False + ) + + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) + + ret0 = ft0(rd0, extended_atype, gr0, fparam=None, aparam=None) + res.append(ret0["property"]) + + np.testing.assert_allclose( + to_numpy_array(res[0][:, idx_perm]), + to_numpy_array(res[1]), + ) + + def test_trans(self): + atype = self.atype.reshape(1, 5) + coord_s = torch.matmul( + torch.remainder( + torch.matmul(self.coord + self.shift, torch.linalg.inv(self.cell)), 1.0 + ), + self.cell, + ) + for fit_diag, scale in itertools.product([True, False], [None, self.scale]): + ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, + task_dim=11, + numb_fparam=0, + numb_aparam=0, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + res = [] + for xyz in [self.coord, coord_s]: + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + xyz, atype, self.rcut, self.sel, False + ) + + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) + + ret0 = ft0(rd0, extended_atype, gr0, fparam=0, aparam=0) + res.append(ret0["property"]) + + np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1])) + + +class TestPropertyModel(unittest.TestCase): + def setUp(self): + self.natoms = 5 + self.rcut = 4.0 + self.nt = 3 + self.rcut_smth = 0.5 + self.sel = [46, 92, 4] + self.nf = 1 + self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device="cpu") + cell = torch.rand([3, 3], dtype=dtype, device="cpu") + self.cell = (cell + cell.T) + 5.0 * torch.eye(3, device="cpu") + self.atype = torch.IntTensor([0, 0, 0, 1, 1], device="cpu") + self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + self.ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, + task_dim=3, + numb_fparam=0, + numb_aparam=0, + mixed_types=True, + ).to(env.DEVICE) + self.type_mapping = ["O", "H", "B"] + self.model = PropertyModel(self.dd0, self.ft0, self.type_mapping) + self.file_path = "model_output.pth" + + def test_deepproperty_infer(self): + atype = self.atype.view(self.nf, self.natoms) + coord = self.coord.reshape(1, 5, 3) + cell = self.cell.reshape(1, 9) + jit_md = torch.jit.script(self.model) + torch.jit.save(jit_md, self.file_path) + load_md = DeepProperty(self.file_path) + load_md.eval(coords=coord, atom_types=atype, cells=cell, atomic=True) + load_md.eval(coords=coord, atom_types=atype, cells=cell, atomic=False) + + def tearDown(self) -> None: + if os.path.exists(self.file_path): + os.remove(self.file_path) + + +if __name__ == "__main__": + unittest.main()