deepmd.pt.model.task.ener

Attributes

dtype
device
log

Classes

EnergyFittingNet	Construct a fitting net for energy.
EnergyFittingNetDirect	Base fitting provides the interfaces of fitting net.

Module Contents

deepmd.pt.model.task.ener.dtype

deepmd.pt.model.task.ener.device

deepmd.pt.model.task.ener.log

class deepmd.pt.model.task.ener.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)

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 enery 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 descrptor should be set.

type_map: List[str], Optional

A list of strings. Give the name to each type of atoms.

classmethod deserialize(data: dict) → deepmd.pt.model.task.fitting.GeneralFitting

Deserialize the fitting.

Parameters:

datadict

The serialized data

Returns:

BF

The deserialized fitting

serialize() → dict

Serialize the fitting to dict.

exclude_types: List[int]

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

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

Base fitting provides the interfaces of fitting net.

ntypes

dim_descrpt

use_tebd

out_dim

bias_atom_e

filter_layers_dipole = []

return_energy

filter_layers = []

output_def()

Returns the output def of the fitting net.

abstract serialize() → dict

Serialize the obj to dict.

abstract deserialize() → EnergyFittingNetDirect

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

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]

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]

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