deepmd.tf.fit.polar

deepmd.tf.fit.polar#

Classes#

`PolarFittingSeA`	Fit the atomic polarizability with descriptor se_a.
`GlobalPolarFittingSeA`	Fit the system polarizability with descriptor se_a.

Module Contents#

class deepmd.tf.fit.polar.PolarFittingSeA(ntypes: int, dim_descrpt: int, embedding_width: int, neuron: list[int] = [120, 120, 120], resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, sel_type: list[int] | None = None, fit_diag: bool = True, scale: list[float] | None = None, shift_diag: bool = True, seed: int | None = None, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, mixed_types: bool = False, type_map: list[str] | None = None, **kwargs)[source]#

Bases: deepmd.tf.fit.fitting.Fitting

Fit the atomic polarizability with descriptor se_a.

Parameters:

ntypes: The ntypes of the descriptor \(\mathcal{D}\)
dim_descrpt: The dimension of the descriptor \(\mathcal{D}\)
embedding_width: The rotation matrix dimension of the descriptor \(\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)
numb_fparam: Number of frame parameters
numb_aparam: Number of atomic parameters
sel_typelist[int]: The atom types selected to have an atomic polarizability prediction. If is None, all atoms are selected.
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]
diag_shiftlist[float]: The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale.
seedint: Random seed for initializing the network parameters.
activation_functionstr: The activation function in the embedding net. Supported options are “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.
precisionstr: The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.
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.
type_map: list[str], Optional: A list of strings. Give the name to each type of atoms.

ntypes[source]#

dim_descrpt[source]#

n_neuron[source]#

resnet_dt[source]#

sel_type[source]#

fit_diag[source]#

seed[source]#

uniform_seed[source]#

seed_shift[source]#

shift_diag[source]#

scale[source]#

activation_function_name[source]#

fitting_activation_fn[source]#

fitting_precision[source]#

sel_mask[source]#

constant_matrix[source]#

dim_rot_mat_1[source]#

dim_rot_mat[source]#

useBN = False[source]#

fitting_net_variables = None[source]#

mixed_prec = None[source]#

mixed_types[source]#

type_map[source]#

numb_fparam[source]#

numb_aparam[source]#

fparam_avg = None[source]#

fparam_std = None[source]#

fparam_inv_std = None[source]#

aparam_avg = None[source]#

aparam_std = None[source]#

aparam_inv_std = None[source]#

get_sel_type() → list[int][source]#: Get selected atom types.

get_out_size() → int[source]#: Get the output size. Should be 9.

compute_output_stats(all_stat) → None[source]#

Compute the output statistics.

Parameters:

all_stat: Dictionary of inputs. can be prepared by model.make_stat_input

_build_lower(start_index, natoms, inputs, rot_mat, suffix='', reuse=None)[source]#

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 = '')[source]#

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:

atomic_polar: The atomic polarizability

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

enable_mixed_precision(mixed_prec: dict | None = None) → None[source]#

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

serialize(suffix: str) → dict[source]#

Serialize the model.

Returns:

dict: The serialized data

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

Deserialize the model.

Parameters:

datadict: The serialized data

Returns:

Model: The deserialized model

class deepmd.tf.fit.polar.GlobalPolarFittingSeA(descrpt: deepmd.tf.env.tf.Tensor, neuron: list[int] = [120, 120, 120], resnet_dt: bool = True, sel_type: list[int] | None = None, fit_diag: bool = True, scale: list[float] | None = None, diag_shift: list[float] | None = None, seed: int | None = None, activation_function: str = 'tanh', precision: str = 'default')[source]#

Fit the system polarizability with descriptor se_a.

Parameters:

descrpttf.Tensor: The descriptor
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 polarizability prediction
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]
diag_shiftlist[float]: The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale.
seedint: Random seed for initializing the network parameters.
activation_functionstr: The activation function in the embedding net. Supported options are “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.
precisionstr: The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.

ntypes[source]#

dim_descrpt[source]#

polar_fitting[source]#

get_sel_type() → int[source]#: Get selected atom types.

get_out_size() → int[source]#: Get the output size. Should be 9.

build(input_d, rot_mat, natoms, input_dict: dict | None = None, reuse=None, suffix='') → deepmd.tf.env.tf.Tensor[source]#

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:

polar: The system polarizability

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

enable_mixed_precision(mixed_prec: dict | None = None) → None[source]#

Receive 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 dict
lrLearningRateExp: the learning rate

Returns:

Loss: the loss function

property input_requirement: list[deepmd.utils.data.DataRequirementItem][source]#: Return data requirements needed for the model input.

get_numb_fparam() → int[source]#: Get the number of frame parameters.

get_numb_aparam() → int[source]#: Get the number of atomic parameters.