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.
`SeZMEnergyFittingNet`	SeZM energy fitting with GLU hidden layers.
`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: int, ntypes: int, attn_head: int = 8, prefactor: list[float] = [0.5, 0.5], activation_function: str = 'gelu', **kwargs: Any)[source]#

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

Base fitting provides the interfaces of fitting net.

feature_dim#

ntypes#

attn_head = 8#

prefactor#

lm_head#

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

forward(pair_weights: torch.Tensor, diff: torch.Tensor, nlist_mask: torch.Tensor, features: torch.Tensor, sw: torch.Tensor, masked_tokens: torch.Tensor | None = None) → dict[str, torch.Tensor][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, dim_case_embd: 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, default_fparam: list | None = None, **kwargs: Any)[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.
dim_case_embdint: Dimension of case specific embedding.
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.
default_fparam: list[float], optional: The default frame parameter. If set, when fparam.npy files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net.

embedding_width#

r_differentiable = True#

c_differentiable = True#

_net_out_dim() → int[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: collections.abc.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) → dict[str, torch.Tensor][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, dim_case_embd: 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, default_fparam: list | 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.
dim_case_embdint: Dimension of case specific embedding.
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.
default_fparam: list[float], optional: The default frame parameter. If set, when fparam.npy files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net.

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, dim_case_embd: 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, default_fparam: list | None = None, **kwargs: Any)[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.
dim_case_embdint: Dimension of case specific embedding.
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.
default_fparam: list[float], optional: The default frame parameter. If set, when fparam.npy files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net.

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: int, dim_descrpt: int, neuron: list[int], bias_atom_e: torch.Tensor | None = None, out_dim: int = 1, resnet_dt: bool = True, use_tebd: bool = True, return_energy: bool = False, **kwargs: Any)[source]#

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

Base fitting provides the interfaces of fitting net.

ntypes#

dim_descrpt#

use_tebd = True#

out_dim = 1#

filter_layers_dipole#

return_energy = False#

filter_layers#

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

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

abstractmethod deserialize() → EnergyFittingNetDirect[source]#

Deserialize the fitting.

Parameters:

datadict: The serialized data

Returns:

BF: The deserialized fitting

abstractmethod change_type_map(type_map: list[str], model_with_new_type_stat: Any | None = 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.

abstractmethod 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: Fitting, shared_level: int, model_prob: float = 1.0, protection: float = 0.01, resume: bool = 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.

save_to_file_fparam(stat_file_path: deepmd.utils.path.DPPath) → None[source]#

Save the statistics of fparam.

Parameters:

stat_file_pathDPPath: The path to save the statistics of fparam.

save_to_file_aparam(stat_file_path: deepmd.utils.path.DPPath) → None[source]#

Save the statistics of aparam.

Parameters:

stat_file_pathDPPath: The path to save the statistics of aparam.

restore_fparam_from_file(stat_file_path: deepmd.utils.path.DPPath) → None[source]#

Load the statistics of fparam.

Parameters:

stat_file_pathDPPath: The path to load the statistics of fparam.

restore_aparam_from_file(stat_file_path: deepmd.utils.path.DPPath) → None[source]#

Load the statistics of aparam.

Parameters:

stat_file_pathDPPath: The path to load the statistics of aparam.

compute_input_stats(merged: collections.abc.Callable[[], list[dict]] | list[dict], protection: float = 0.01, stat_file_path: deepmd.utils.path.DPPath | None = None) → None[source]#

Compute the input statistics (e.g. mean and stddev) for the fittings 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.

protectionfloat

Divided-by-zero protection

stat_file_pathOptional[DPPath]

The path to the stat file.

get_stats() → dict[str, list[deepmd.utils.env_mat_stat.StatItem]][source]#: Get the statistics of the fitting_net.

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, dim_case_embd: 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, default_fparam: list | None = None, **kwargs: Any)[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.
dim_case_embdint: Dimension of case specific embedding.
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.
default_fparam: list[float], optional: The default frame parameter. If set, when fparam.npy files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net.

embedding_width#

fit_diag = True#

scale = None#

shift_diag = True#

constant_matrix#

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

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

__getitem__(key: str) → Any[source]#

change_type_map(type_map: list[str], model_with_new_type_stat: Any | None = 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) → dict[str, torch.Tensor][source]#

exclude_types: list[int]#

class deepmd.pt.model.task.PropertyFittingNet(ntypes: int, dim_descrpt: int, property_name: str, task_dim: int = 1, neuron: list[int] = [128, 128, 128], bias_atom_p: torch.Tensor | None = None, intensive: bool = False, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, activation_function: str = 'tanh', precision: str = DEFAULT_PRECISION, mixed_types: bool = True, trainable: bool | list[bool] = True, seed: int | None = None, default_fparam: list | None = None, distinguish_types: bool = True, **kwargs: Any)[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.
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”.
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.
resnet_dtbool: Using time-step in the ResNet construction.
numb_fparamint: Number of frame parameters.
numb_aparamint: Number of atomic parameters.
dim_case_embdint: Dimension of case specific embedding.
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.
distinguish_typesbool: Whether to distinguish atom types when computing output statistics.

task_dim = 1#

intensive = False#

distinguish_types = True#

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

get_intensive() → bool[source]#: Whether the fitting property is intensive.

get_distinguish_types() → bool[source]#: Get whether to distinguish atom types when computing output statistics.

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.SeZMEnergyFittingNet(ntypes: int, dim_descrpt: int, neuron: list[int] | None = None, bias_atom_e: torch.Tensor | None = None, resnet_dt: bool = False, numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, case_film_embd: bool = False, activation_function: str = 'silu', bias_out: bool = False, precision: str = 'float32', mixed_types: bool = True, seed: int | list[int] | None = None, type_map: list[str] | None = None, default_fparam: list | None = None, **kwargs: Any)[source]#

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

SeZM energy fitting with GLU hidden layers.

This uses the same configuration keys as the standard energy fitting but replaces hidden MLP layers with GLU blocks.

bias_out = False#

case_film_embd#

_build_glu_fitting_layers() → None[source]#

_forward_common(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) → dict[str, torch.Tensor][source]#: Run the SeZM fitting path with optional case FiLM.

_forward_case_film(descriptor: torch.Tensor, atype: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None) → dict[str, torch.Tensor][source]#

Forward path for SeZM case FiLM.

Parameters:

descriptor: Descriptor tensor with shape (nf, nloc, dim_descrpt).
atype: Atom types with shape (nf, nloc).
fparam: Frame parameters with shape (nf, numb_fparam).
aparam: Atomic parameters with shape (nf, nloc, numb_aparam).

Returns:

dict[str, torch.Tensor]: Per-atom fitting outputs.

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.

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

Bases: deepmd.pt.model.task.Fitting

Base fitting provides the interfaces of fitting net.

feature_dim#

ntypes#

lm_head#

forward(features: torch.Tensor, masked_tokens: torch.Tensor | None = None) → torch.Tensor[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].