deepmd.tf.fit.dipole

Module Contents

Classes

DipoleFittingSeA

Fit the atomic dipole with descriptor se_a.

class deepmd.tf.fit.dipole.DipoleFittingSeA(ntypes: int, dim_descrpt: int, embedding_width: int, neuron: List[int] = [120, 120, 120], resnet_dt: bool = True, sel_type: List[int] | None = None, seed: int | None = None, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, mixed_types: bool = False, **kwargs)

Bases: deepmd.tf.fit.fitting.Fitting

Fit the atomic dipole with descriptor se_a.

Parameters:

ntypes

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

dim_descrpt

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

embedding_width

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

neuronList[int]

Number of neurons in each hidden layer of the fitting net

resnet_dtbool

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

sel_typeList[int]

The atom types selected to have an atomic dipole prediction. If is None, all atoms are selected.

seedint

Random seed for initializing the network parameters.

activation_functionstr

The activation function in the embedding net. Supported options are “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.

precisionstr

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

mixed_typesbool

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

get_sel_type() → int

Get selected type.

get_out_size() → int

Get the output size. Should be 3.

_build_lower(start_index, natoms, inputs, rot_mat, suffix='', reuse=None)

build(input_d: deepmd.tf.env.tf.Tensor, rot_mat: 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

Build the computational graph for fitting net.

Parameters:

input_d

The input descriptor

rot_mat

The rotation matrix from the descriptor.

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

input_dict

Additional dict for inputs.

reuse

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

suffix

Name suffix to identify this descriptor

Returns:

dipole

The atomic dipole.

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

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

enable_mixed_precision(mixed_prec: dict | None = None) → None

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

Get the loss function.

Parameters:

lossdict

the loss dict

lrLearningRateExp

the learning rate

Returns:

Loss

the loss function

serialize(suffix: str) → dict

Serialize the model.

Returns:

dict

The serialized data

classmethod deserialize(data: dict, suffix: str)

Deserialize the model.

Parameters:

datadict

The serialized data

Returns:

Model

The deserialized model