deepmd.tf.fit.ener

Module Contents

Classes

EnerFitting

Fitting the energy of the system. The force and the virial can also be trained.

Attributes

log

deepmd.tf.fit.ener.log[source]
class deepmd.tf.fit.ener.EnerFitting(ntypes: int, dim_descrpt: int, neuron: List[int] = [120, 120, 120], resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, rcond: float | None = None, tot_ener_zero: bool = False, trainable: List[bool] | None = None, seed: int | None = None, atom_ener: List[float] = [], activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, layer_name: List[str | None] | None = None, use_aparam_as_mask: bool = False, spin: deepmd.tf.utils.spin.Spin | None = None, mixed_types: bool = False, **kwargs)[source]

Bases: deepmd.tf.fit.fitting.Fitting

Fitting the energy of the system. The force and the virial can also be trained.

The potential energy \(E\) is a fitting network function of the descriptor \(\mathcal{D}\):

\[E(\mathcal{D}) = \mathcal{L}^{(n)} \circ \mathcal{L}^{(n-1)} \circ \cdots \circ \mathcal{L}^{(1)} \circ \mathcal{L}^{(0)}\]

The first \(n\) hidden layers \(\mathcal{L}^{(0)}, \cdots, \mathcal{L}^{(n-1)}\) are given by

\[\mathbf{y}=\mathcal{L}(\mathbf{x};\mathbf{w},\mathbf{b})= \boldsymbol{\phi}(\mathbf{x}^T\mathbf{w}+\mathbf{b})\]

where \(\mathbf{x} \in \mathbb{R}^{N_1}\) is the input vector and \(\mathbf{y} \in \mathbb{R}^{N_2}\) is the output vector. \(\mathbf{w} \in \mathbb{R}^{N_1 \times N_2}\) and \(\mathbf{b} \in \mathbb{R}^{N_2}\) are weights and biases, respectively, both of which are trainable if trainable[i] is True. \(\boldsymbol{\phi}\) is the activation function.

The output layer \(\mathcal{L}^{(n)}\) is given by

\[\mathbf{y}=\mathcal{L}^{(n)}(\mathbf{x};\mathbf{w},\mathbf{b})= \mathbf{x}^T\mathbf{w}+\mathbf{b}\]

where \(\mathbf{x} \in \mathbb{R}^{N_{n-1}}\) is the input vector and \(\mathbf{y} \in \mathbb{R}\) is the output scalar. \(\mathbf{w} \in \mathbb{R}^{N_{n-1}}\) and \(\mathbf{b} \in \mathbb{R}\) are weights and bias, respectively, both of which are trainable if trainable[n] is True.

Parameters:
ntypes

The ntypes of the descrptor \(\mathcal{D}\)

dim_descrpt

The dimension of the descrptor \(\mathcal{D}\)

neuron

Number of neurons \(N\) in each hidden layer of the fitting net

resnet_dt

Time-step dt in the resnet construction: \(y = x + dt * \phi (Wx + b)\)

numb_fparam

Number of frame parameter

numb_aparam

Number of atomic parameter

rcond

The condition number for the regression of atomic energy.

tot_ener_zero

Force the total energy to zero. Useful for the charge fitting.

trainable

If the weights of fitting net are trainable. Suppose that we have \(N_l\) hidden layers in the fitting net, this list is of length \(N_l + 1\), specifying if the hidden layers and the output layer are trainable.

seed

Random seed for initializing the network parameters.

atom_ener

Specifying atomic energy contribution in vacuum. The set_davg_zero key in the descrptor should be set.

activation_function

The activation function \(\boldsymbol{\phi}\) in the embedding net. Supported options are “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.

precision

The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

layer_namelist[Optional[str]], optional

The name of the each layer. If two layers, either in the same fitting or different fittings, have the same name, they will share the same neural network parameters.

use_aparam_as_mask: bool, optional

If True, the atomic parameters will be used as a mask that determines the atom is real/virtual. And the aparam will not be used as the atomic parameters for embedding.

mixed_typesbool

If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.

get_numb_fparam() int[source]

Get the number of frame parameters.

get_numb_aparam() int[source]

Get the number of atomic parameters.

compute_output_stats(all_stat: dict, mixed_type: bool = False) None[source]

Compute the ouput statistics.

Parameters:
all_stat

must have the following components: all_stat[‘energy’] of shape n_sys x n_batch x n_frame can be prepared by model.make_stat_input

mixed_type

Whether to perform the mixed_type mode. If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), in which frames in a system may have different natoms_vec(s), with the same nloc.

_compute_output_stats(all_stat, rcond=0.001, mixed_type=False)[source]
compute_input_stats(all_stat: dict, protection: float = 0.01) None[source]

Compute the input statistics.

Parameters:
all_stat

if numb_fparam > 0 must have all_stat[‘fparam’] if numb_aparam > 0 must have all_stat[‘aparam’] can be prepared by model.make_stat_input

protection

Divided-by-zero protection

_compute_std(sumv2, sumv, sumn)[source]
_build_lower(start_index, natoms, inputs, fparam=None, aparam=None, bias_atom_e=0.0, type_suffix='', suffix='', reuse=None)[source]
build(inputs: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, input_dict: dict | None = None, reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]

Build the computational graph for fitting net.

Parameters:
inputs

The input descriptor

input_dict

Additional dict for inputs. if numb_fparam > 0, should have input_dict[‘fparam’] if numb_aparam > 0, should have input_dict[‘aparam’]

natoms

The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns:
ener

The system energy

init_variables(graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, suffix: str = '') None[source]

Init the fitting net variables with the given dict.

Parameters:
graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

suffixstr

suffix to name scope

change_energy_bias(data, frozen_model, origin_type_map, full_type_map, bias_adjust_mode='change-by-statistic', ntest=10) None[source]
enable_mixed_precision(mixed_prec: dict | None = None) None[source]

Reveive the mixed precision setting.

Parameters:
mixed_prec

The mixed precision setting used in the embedding net

get_loss(loss: dict, lr) deepmd.tf.loss.loss.Loss[source]

Get the loss function.

Parameters:
lossdict

The loss function parameters.

lrLearningRateExp

The learning rate.

Returns:
Loss

The loss function.

classmethod deserialize(data: dict, suffix: str = '')[source]

Deserialize the model.

Parameters:
datadict

The serialized data

Returns:
Model

The deserialized model

serialize(suffix: str = '') dict[source]

Serialize the model.

Returns:
dict

The serialized data