From 1fa3e18cb7df7a9ec783938f548c91b350e949a1 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 15 Mar 2024 16:23:32 +0800 Subject: [PATCH 1/2] chore: move invar fitting --- deepmd/pt/model/task/ener.py | 180 +-------------------- deepmd/pt/model/task/invar_fitting.py | 215 ++++++++++++++++++++++++++ 2 files changed, 217 insertions(+), 178 deletions(-) create mode 100644 deepmd/pt/model/task/invar_fitting.py diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index b58b0c9b19..75bbe2be09 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -2,11 +2,9 @@ import copy import logging from typing import ( - Callable, List, Optional, Tuple, - Union, ) import numpy as np @@ -30,11 +28,8 @@ from deepmd.pt.utils.env import ( DEFAULT_PRECISION, ) -from deepmd.pt.utils.stat import ( - compute_output_stats, -) -from deepmd.utils.path import ( - DPPath, +from deepmd.pt.model.task.invar_fitting import( + InvarFitting ) from deepmd.utils.version import ( check_version_compatibility, @@ -46,177 +41,6 @@ log = logging.getLogger(__name__) -@GeneralFitting.register("invar") -@fitting_check_output -class InvarFitting(GeneralFitting): - """Construct a fitting net for energy. - - Parameters - ---------- - var_name : str - The atomic property to fit, 'energy', 'dipole', and 'polar'. - ntypes : int - Element count. - dim_descrpt : int - Embedding width per atom. - dim_out : int - The output dimension of the fitting net. - neuron : List[int] - Number of neurons in each hidden layers of the fitting net. - bias_atom_e : torch.Tensor, optional - Average enery 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. - rcond : float, optional - The condition number for the regression of atomic energy. - seed : int, optional - Random seed. - exclude_types: List[int] - Atomic contributions of the excluded atom types are set zero. - atom_ener: List[float], optional - Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. - - """ - - def __init__( - self, - var_name: str, - ntypes: int, - dim_descrpt: int, - dim_out: int, - neuron: List[int] = [128, 128, 128], - bias_atom_e: 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, - rcond: Optional[float] = None, - seed: Optional[int] = None, - exclude_types: List[int] = [], - atom_ener: Optional[List[float]] = None, - **kwargs, - ): - self.dim_out = dim_out - self.atom_ener = atom_ener - super().__init__( - var_name=var_name, - ntypes=ntypes, - dim_descrpt=dim_descrpt, - neuron=neuron, - bias_atom_e=bias_atom_e, - resnet_dt=resnet_dt, - numb_fparam=numb_fparam, - numb_aparam=numb_aparam, - activation_function=activation_function, - precision=precision, - mixed_types=mixed_types, - rcond=rcond, - seed=seed, - exclude_types=exclude_types, - remove_vaccum_contribution=None - if atom_ener is None or len([x for x in atom_ener if x is not None]) == 0 - else [x is not None for x in atom_ener], - **kwargs, - ) - - def _net_out_dim(self): - """Set the FittingNet output dim.""" - return self.dim_out - - def serialize(self) -> dict: - data = super().serialize() - data["type"] = "invar" - data["dim_out"] = self.dim_out - data["atom_ener"] = self.atom_ener - return data - - @classmethod - def deserialize(cls, data: dict) -> "GeneralFitting": - data = copy.deepcopy(data) - check_version_compatibility(data.pop("@version", 1), 1, 1) - return super().deserialize(data) - - 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. - - """ - bias_atom_e = compute_output_stats( - merged, self.ntypes, stat_file_path, self.rcond, self.atom_ener - ) - self.bias_atom_e.copy_( - torch.tensor(bias_atom_e, device=env.DEVICE).view( - [self.ntypes, self.dim_out] - ) - ) - - def output_def(self) -> FittingOutputDef: - return FittingOutputDef( - [ - OutputVariableDef( - self.var_name, - [self.dim_out], - reduciable=True, - r_differentiable=True, - c_differentiable=True, - ), - ] - ) - - def forward( - self, - descriptor: torch.Tensor, - atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): - """Based on embedding net output, alculate total energy. - - Args: - - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.dim_descrpt]. - - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. - - Returns - ------- - - `torch.Tensor`: Total energy with shape [nframes, natoms[0]]. - """ - return self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam) - - # make jit happy with torch 2.0.0 - exclude_types: List[int] - @Fitting.register("ener") class EnergyFittingNet(InvarFitting): diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py new file mode 100644 index 0000000000..40a36709f0 --- /dev/null +++ b/deepmd/pt/model/task/invar_fitting.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import copy +import logging +from typing import ( + Callable, + List, + Optional, + Tuple, + Union, +) + +import numpy as np +import torch + +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) + +from deepmd.pt.model.task.fitting import ( + GeneralFitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, +) +from deepmd.pt.utils.stat import ( + compute_output_stats, +) +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__) + + +@GeneralFitting.register("invar") +@fitting_check_output +class InvarFitting(GeneralFitting): + """Construct a fitting net for energy. + + Parameters + ---------- + var_name : str + The atomic property to fit, 'energy', 'dipole', and 'polar'. + ntypes : int + Element count. + dim_descrpt : int + Embedding width per atom. + dim_out : int + The output dimension of the fitting net. + neuron : List[int] + Number of neurons in each hidden layers of the fitting net. + bias_atom_e : torch.Tensor, optional + Average enery 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. + rcond : float, optional + The condition number for the regression of atomic energy. + seed : int, optional + Random seed. + exclude_types: List[int] + Atomic contributions of the excluded atom types are set zero. + atom_ener: List[float], optional + Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. + + """ + + def __init__( + self, + var_name: str, + ntypes: int, + dim_descrpt: int, + dim_out: int, + neuron: List[int] = [128, 128, 128], + bias_atom_e: 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, + rcond: Optional[float] = None, + seed: Optional[int] = None, + exclude_types: List[int] = [], + atom_ener: Optional[List[float]] = None, + **kwargs, + ): + self.dim_out = dim_out + self.atom_ener = atom_ener + super().__init__( + var_name=var_name, + ntypes=ntypes, + dim_descrpt=dim_descrpt, + neuron=neuron, + bias_atom_e=bias_atom_e, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + activation_function=activation_function, + precision=precision, + mixed_types=mixed_types, + rcond=rcond, + seed=seed, + exclude_types=exclude_types, + remove_vaccum_contribution=None + if atom_ener is None or len([x for x in atom_ener if x is not None]) == 0 + else [x is not None for x in atom_ener], + **kwargs, + ) + + def _net_out_dim(self): + """Set the FittingNet output dim.""" + return self.dim_out + + def serialize(self) -> dict: + data = super().serialize() + data["type"] = "invar" + data["dim_out"] = self.dim_out + data["atom_ener"] = self.atom_ener + return data + + @classmethod + def deserialize(cls, data: dict) -> "GeneralFitting": + data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) + return super().deserialize(data) + + 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. + + """ + bias_atom_e = compute_output_stats( + merged, self.ntypes, stat_file_path, self.rcond, self.atom_ener + ) + self.bias_atom_e.copy_( + torch.tensor(bias_atom_e, device=env.DEVICE).view( + [self.ntypes, self.dim_out] + ) + ) + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reduciable=True, + r_differentiable=True, + c_differentiable=True, + ), + ] + ) + + def forward( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + gr: Optional[torch.Tensor] = None, + g2: Optional[torch.Tensor] = None, + h2: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + ): + """Based on embedding net output, alculate total energy. + + Args: + - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.dim_descrpt]. + - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. + + Returns + ------- + - `torch.Tensor`: Total energy with shape [nframes, natoms[0]]. + """ + return self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam) + + # make jit happy with torch 2.0.0 + exclude_types: List[int] From 33e30ca832c805759baf9f9fe409ed7b196d13b5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 15 Mar 2024 08:24:39 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/ener.py | 7 +++---- deepmd/pt/model/task/invar_fitting.py | 3 --- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 75bbe2be09..12c0917dd2 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -22,15 +22,15 @@ 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.pt.model.task.invar_fitting import( - InvarFitting -) from deepmd.utils.version import ( check_version_compatibility, ) @@ -41,7 +41,6 @@ log = logging.getLogger(__name__) - @Fitting.register("ener") class EnergyFittingNet(InvarFitting): def __init__( diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index 40a36709f0..1699b440ac 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -5,11 +5,9 @@ Callable, List, Optional, - Tuple, Union, ) -import numpy as np import torch from deepmd.dpmodel import ( @@ -17,7 +15,6 @@ OutputVariableDef, fitting_check_output, ) - from deepmd.pt.model.task.fitting import ( GeneralFitting, )