deepmd.pt.model.task

deepmd.pt.model.task#

Submodules#

Attributes#

BaseFitting

Classes#

`DenoiseNet`	Base fitting provides the interfaces of fitting net.
`DipoleFittingNet`	Construct a dipole fitting net.
`DOSFittingNet`	Construct a fitting net for energy.
`EnergyFittingNet`	Construct a fitting net for energy.
`EnergyFittingNetDirect`	Base fitting provides the interfaces of fitting net.
`Fitting`	Base fitting provides the interfaces of fitting net.
`PolarFittingNet`	Construct a polar fitting net.
`PropertyFittingNet`	Fitting the rotationally invariant properties of task_dim of the system.
`TypePredictNet`	Base fitting provides the interfaces of fitting net.

Package Contents#

deepmd.pt.model.task.BaseFitting[source]#

class deepmd.pt.model.task.DenoiseNet(feature_dim, ntypes, attn_head=8, prefactor=[0.5, 0.5], activation_function='gelu', **kwargs)[source]#

Bases: deepmd.pt.model.task.fitting.Fitting

Base fitting provides the interfaces of fitting net.

feature_dim#

ntypes#

attn_head#

prefactor#

lm_head#

output_def()[source]#: Returns the output def of the fitting net.

forward(pair_weights, diff, nlist_mask, features, sw, masked_tokens: torch.Tensor | None = None)[source]#

Calculate the updated coord. Args: - coord: Input noisy coord with shape [nframes, nloc, 3]. - pair_weights: Input pair weights with shape [nframes, nloc, nnei, head]. - diff: Input pair relative coord list with shape [nframes, nloc, nnei, 3]. - nlist_mask: Input nlist mask with shape [nframes, nloc, nnei].

Returns:

denoised_coord: Denoised updated coord with shape [nframes, nloc, 3].

class deepmd.pt.model.task.DipoleFittingNet(ntypes: int, dim_descrpt: int, embedding_width: int, neuron: list[int] = [128, 128, 128], 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: float | None = None, seed: int | list[int] | None = None, exclude_types: list[int] = [], r_differentiable: bool = True, c_differentiable: bool = True, type_map: list[str] | None = None, **kwargs)[source]#

Bases: deepmd.pt.model.task.fitting.GeneralFitting

Construct a dipole fitting net.

Parameters:

ntypesint: Element count.
dim_descrptint: Embedding width per atom.
embedding_widthint: The dimension of rotation matrix, m1.
neuronlist[int]: Number of neurons in each hidden layers of the fitting net.
resnet_dtbool: Using time-step in the ResNet construction.
numb_fparamint: Number of frame parameters.
numb_aparamint: Number of atomic parameters.
activation_functionstr: Activation function.
precisionstr: Numerical precision.
mixed_typesbool: If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.
rcondfloat, optional: The condition number for the regression of atomic energy.
seedint, optional: Random seed.
r_differentiable: If the variable is differentiated with respect to coordinates of atoms. Only reducible variable are differentiable.
c_differentiable: If the variable is differentiated with respect to the cell tensor (pbc case). Only reducible variable are differentiable.
type_map: list[str], Optional: A list of strings. Give the name to each type of atoms.

embedding_width#

r_differentiable#

c_differentiable#

_net_out_dim()[source]#: Set the FittingNet output dim.

serialize() → dict[source]#: Serialize the fitting to dict.

classmethod deserialize(data: dict) → deepmd.pt.model.task.fitting.GeneralFitting[source]#

Deserialize the fitting.

Parameters:

datadict: The serialized data

Returns:

BF: The deserialized fitting

output_def() → deepmd.dpmodel.FittingOutputDef[source]#: Returns the output def of the fitting net.

compute_output_stats(merged: Callable[[], list[dict]] | list[dict], stat_file_path: deepmd.utils.path.DPPath | None = None) → None[source]#

Compute the output statistics (e.g. energy bias) for the fitting net from packed data.

Parameters:

mergedUnion[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_pathOptional[DPPath]

The path to the stat file.

forward(descriptor: torch.Tensor, atype: torch.Tensor, gr: torch.Tensor | None = None, g2: torch.Tensor | None = None, h2: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None)[source]#

exclude_types: list[int]#

class deepmd.pt.model.task.DOSFittingNet(ntypes: int, dim_descrpt: int, numb_dos: int = 300, neuron: list[int] = [128, 128, 128], resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, rcond: float | None = None, bias_dos: torch.Tensor | None = None, trainable: bool | list[bool] = True, seed: int | list[int] | None = None, activation_function: str = 'tanh', precision: str = DEFAULT_PRECISION, exclude_types: list[int] = [], mixed_types: bool = True, type_map: list[str] | None = None)[source]#

Bases: deepmd.pt.model.task.ener.InvarFitting

Construct a fitting net for energy.

Parameters:

var_namestr: The atomic property to fit, ‘energy’, ‘dipole’, and ‘polar’.
ntypesint: Element count.
dim_descrptint: Embedding width per atom.
dim_outint: The output dimension of the fitting net.
neuronlist[int]: Number of neurons in each hidden layers of the fitting net.
bias_atom_etorch.Tensor, optional: Average energy per atom for each element.
resnet_dtbool: Using time-step in the ResNet construction.
numb_fparamint: Number of frame parameters.
numb_aparamint: Number of atomic parameters.
activation_functionstr: Activation function.
precisionstr: Numerical precision.
mixed_typesbool: If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.
rcondfloat, optional: The condition number for the regression of atomic energy.
seedint, optional: Random seed.
exclude_types: list[int]: Atomic contributions of the excluded atom types are set zero.
atom_ener: list[Optional[torch.Tensor]], optional: Specifying atomic energy contribution in vacuum. The value is a list specifying the bias. the elements can be None or np.array of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] The set_davg_zero key in the descriptor should be set.
type_map: list[str], Optional: A list of strings. Give the name to each type of atoms.
use_aparam_as_mask: bool: If True, the aparam will not be used in fitting net for embedding.

output_def() → deepmd.dpmodel.FittingOutputDef[source]#: Returns the output def of the fitting net.

classmethod deserialize(data: dict) → DOSFittingNet[source]#

Deserialize the fitting.

Parameters:

datadict: The serialized data

Returns:

BF: The deserialized fitting

serialize() → dict[source]#: Serialize the fitting to dict.

exclude_types: list[int]#

class deepmd.pt.model.task.EnergyFittingNet(ntypes: int, dim_descrpt: int, neuron: list[int] = [128, 128, 128], bias_atom_e: torch.Tensor | None = 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: int | list[int] | None = None, type_map: list[str] | None = None, **kwargs)[source]#

Bases: deepmd.pt.model.task.invar_fitting.InvarFitting

Construct a fitting net for energy.

Parameters:

var_namestr: The atomic property to fit, ‘energy’, ‘dipole’, and ‘polar’.
ntypesint: Element count.
dim_descrptint: Embedding width per atom.
dim_outint: The output dimension of the fitting net.
neuronlist[int]: Number of neurons in each hidden layers of the fitting net.
bias_atom_etorch.Tensor, optional: Average energy per atom for each element.
resnet_dtbool: Using time-step in the ResNet construction.
numb_fparamint: Number of frame parameters.
numb_aparamint: Number of atomic parameters.
activation_functionstr: Activation function.
precisionstr: Numerical precision.
mixed_typesbool: If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.
rcondfloat, optional: The condition number for the regression of atomic energy.
seedint, optional: Random seed.
exclude_types: list[int]: Atomic contributions of the excluded atom types are set zero.
atom_ener: list[Optional[torch.Tensor]], optional: Specifying atomic energy contribution in vacuum. The value is a list specifying the bias. the elements can be None or np.array of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] The set_davg_zero key in the descriptor should be set.
type_map: list[str], Optional: A list of strings. Give the name to each type of atoms.
use_aparam_as_mask: bool: If True, the aparam will not be used in fitting net for embedding.

classmethod deserialize(data: dict) → deepmd.pt.model.task.fitting.GeneralFitting[source]#

Deserialize the fitting.

Parameters:

datadict: The serialized data

Returns:

BF: The deserialized fitting

serialize() → dict[source]#: Serialize the fitting to dict.

exclude_types: list[int]#

class deepmd.pt.model.task.EnergyFittingNetDirect(ntypes, dim_descrpt, neuron, bias_atom_e=None, out_dim=1, resnet_dt=True, use_tebd=True, return_energy=False, **kwargs)[source]#

Bases: deepmd.pt.model.task.fitting.Fitting

Base fitting provides the interfaces of fitting net.

ntypes#

dim_descrpt#

use_tebd#

out_dim#

filter_layers_dipole#

return_energy#

filter_layers#

output_def()[source]#: Returns the output def of the fitting net.

abstract serialize() → dict[source]#: Serialize the obj to dict.

abstract deserialize() → EnergyFittingNetDirect[source]#

Deserialize the fitting.

Parameters:

datadict: The serialized data

Returns:

BF: The deserialized fitting

abstract change_type_map(type_map: list[str], model_with_new_type_stat=None) → None[source]#: Change the type related params to new ones, according to type_map and the original one in the model. If there are new types in type_map, statistics will be updated accordingly to model_with_new_type_stat for these new types.

abstract get_type_map() → list[str][source]#: Get the name to each type of atoms.

forward(inputs: torch.Tensor, atype: torch.Tensor, gr: torch.Tensor | None = None, g2: torch.Tensor | None = None, h2: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None) → tuple[torch.Tensor, None][source]#

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]].

class deepmd.pt.model.task.Fitting[source]#

Bases: torch.nn.Module, deepmd.pt.model.task.base_fitting.BaseFitting

Base fitting provides the interfaces of fitting net.

share_params(base_class, shared_level, resume=False) → None[source]#: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes.

class deepmd.pt.model.task.PolarFittingNet(ntypes: int, dim_descrpt: int, embedding_width: int, neuron: list[int] = [128, 128, 128], 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: float | None = None, seed: int | list[int] | None = None, exclude_types: list[int] = [], fit_diag: bool = True, scale: list[float] | float | None = None, shift_diag: bool = True, type_map: list[str] | None = None, **kwargs)[source]#

Bases: deepmd.pt.model.task.fitting.GeneralFitting

Construct a polar fitting net.

Parameters:

ntypesint: Element count.
dim_descrptint: Embedding width per atom.
embedding_widthint: The dimension of rotation matrix, m1.
neuronlist[int]: Number of neurons in each hidden layers of the fitting net.
resnet_dtbool: Using time-step in the ResNet construction.
numb_fparamint: Number of frame parameters.
numb_aparamint: Number of atomic parameters.
activation_functionstr: Activation function.
precisionstr: Numerical precision.
mixed_typesbool: If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.
rcondfloat, optional: The condition number for the regression of atomic energy.
seedint, optional: Random seed.
fit_diagbool: Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix.
scalelist[float]: The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i]
shift_diagbool: Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale.
type_map: list[str], Optional: A list of strings. Give the name to each type of atoms.

embedding_width#

fit_diag#

scale#

shift_diag#

constant_matrix#

_net_out_dim()[source]#: Set the FittingNet output dim.

__setitem__(key, value) → None[source]#

__getitem__(key)[source]#

change_type_map(type_map: list[str], model_with_new_type_stat=None) → None[source]#: Change the type related params to new ones, according to type_map and the original one in the model. If there are new types in type_map, statistics will be updated accordingly to model_with_new_type_stat for these new types.

serialize() → dict[source]#: Serialize the fitting to dict.

classmethod deserialize(data: dict) → deepmd.pt.model.task.fitting.GeneralFitting[source]#

Deserialize the fitting.

Parameters:

datadict: The serialized data

Returns:

BF: The deserialized fitting

output_def() → deepmd.dpmodel.FittingOutputDef[source]#: Returns the output def of the fitting net.

forward(descriptor: torch.Tensor, atype: torch.Tensor, gr: torch.Tensor | None = None, g2: torch.Tensor | None = None, h2: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None)[source]#

exclude_types: list[int]#

class deepmd.pt.model.task.PropertyFittingNet(ntypes: int, dim_descrpt: int, task_dim: int = 1, neuron: list[int] = [128, 128, 128], bias_atom_p: torch.Tensor | None = None, intensive: bool = False, bias_method: str = 'normal', 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: int | None = None, **kwargs)[source]#

Bases: deepmd.pt.model.task.ener.InvarFitting

Fitting the rotationally invariant properties of task_dim of the system.

Parameters:

ntypesint: Element count.
dim_descrptint: Embedding width per atom.
task_dimint: The dimension of outputs of fitting net.
neuronlist[int]: Number of neurons in each hidden layers of the fitting net.
bias_atom_ptorch.Tensor, optional: Average property per atom for each element.
intensivebool, optional: Whether the fitting property is intensive.
bias_methodstr, 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_dtbool: Using time-step in the ResNet construction.
numb_fparamint: Number of frame parameters.
numb_aparamint: Number of atomic parameters.
activation_functionstr: Activation function.
precisionstr: Numerical precision.
mixed_typesbool: If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.
seedint, optional: Random seed.

task_dim#

intensive#

bias_method#

get_bias_method() → str[source]#

output_def() → deepmd.dpmodel.FittingOutputDef[source]#: Returns the output def of the fitting net.

classmethod deserialize(data: dict) → PropertyFittingNet[source]#

Deserialize the fitting.

Parameters:

datadict: The serialized data

Returns:

BF: The deserialized fitting

serialize() → dict[source]#: Serialize the fitting to dict.

exclude_types: list[int]#

class deepmd.pt.model.task.TypePredictNet(feature_dim, ntypes, activation_function='gelu', **kwargs)[source]#

Bases: deepmd.pt.model.task.Fitting

Base fitting provides the interfaces of fitting net.

feature_dim#

ntypes#

lm_head#

forward(features, masked_tokens: torch.Tensor | None = None)[source]#

Calculate the predicted logits. Args: - features: Input features with shape [nframes, nloc, feature_dim]. - masked_tokens: Input masked tokens with shape [nframes, nloc].

Returns:

logits: Predicted probs with shape [nframes, nloc, ntypes].