deepmd.dpmodel.fitting package

class deepmd.dpmodel.fitting.DipoleFitting(*args, **kwargs)

Bases: DipoleFitting

Fitting rotationally equivariant diploe of the system.

Parameters

var_name

The name of the output variable.

ntypes

The number of atom types.

dim_descrpt

The dimension of the input descriptor.

embedding_widthint

The dimension of rotation matrix, m1.

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.

activation_function

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

precision

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

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_types

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

exclude_types

Atomic contributions of the excluded atom types are set zero.

r_differentiable

If the variable is differentiated with respect to coordinates of atoms. Only reduciable variable are differentiable.

c_differentiable

If the variable is differentiated with respect to the cell tensor (pbc case). Only reduciable variable are differentiable.

Methods

call(descriptor, atype[, gr, g2, h2, ...])	Calculate the fitting.
compute_output_stats(merged)	Update the output bias for fitting net.
deserialize(data)	Deserialize the fitting.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
output_def()	Returns the output def of the fitting net.
register(key)	Register a descriptor plugin.
serialize()	Serialize the fitting to dict.

__call__
reinit_exclude

class deepmd.dpmodel.fitting.EnergyFittingNet(*args, **kwargs)

Bases: InvarFitting

Methods

call(descriptor, atype[, gr, g2, h2, ...])	Calculate the fitting.
compute_output_stats(merged)	Update the output bias for fitting net.
deserialize(data)	Deserialize the fitting.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
init_fitting_stat(result_dict)	Initialize the model bias by the statistics.
output_def()	Returns the output def of the fitting net.
register(key)	Register a descriptor plugin.
serialize()	Serialize the fitting to dict.

__call__
reinit_exclude

classmethod deserialize(data: dict) → GeneralFitting

Deserialize the fitting.

Parameters

datadict

The serialized data

Returns

BF

The deserialized fitting

serialize() → dict

Serialize the fitting to dict.

class deepmd.dpmodel.fitting.InvarFitting(*args, **kwargs)

Bases: InvarFitting

Fitting the energy (or a rotationally invariant porperty of dim_out) of the system. The force and the virial can also be trained.

Lets take the energy fitting task as an example. 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

var_name

The name of the output variable.

ntypes

The number of atom types.

dim_descrpt

The dimension of the input descriptor.

dim_out

The dimension of the output fit property.

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.

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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

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

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_types

If false, different atomic types uses different fitting net, otherwise different atom types share the same fitting net.

exclude_types: List[int]

Atomic contributions of the excluded atom types are set zero.

Methods

call(descriptor, atype[, gr, g2, h2, ...])	Calculate the fitting.
compute_output_stats(merged)	Update the output bias for fitting net.
deserialize(data)	Deserialize the fitting.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
init_fitting_stat(result_dict)	Initialize the model bias by the statistics.
output_def()	Returns the output def of the fitting net.
register(key)	Register a descriptor plugin.
serialize()	Serialize the fitting to dict.

__call__
reinit_exclude

class deepmd.dpmodel.fitting.PolarFitting(*args, **kwargs)

Bases: PolarFitting

Fitting rotationally equivariant polarizability of the system.

Parameters

var_name

The name of the output variable.

ntypes

The number of atom types.

dim_descrpt

The dimension of the input descriptor.

embedding_widthint

The dimension of rotation matrix, m1.

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.

activation_function

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

precision

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

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_types

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

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.

Methods

call(descriptor, atype[, gr, g2, h2, ...])	Calculate the fitting.
compute_output_stats(merged)	Update the output bias for fitting net.
deserialize(data)	Deserialize the fitting.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
output_def()	Returns the output def of the fitting net.
register(key)	Register a descriptor plugin.
serialize()	Serialize the fitting to dict.

__call__
reinit_exclude

deepmd.dpmodel.fitting.make_base_fitting(t_tensor, fwd_method_name: str = 'forward')

Make the base class for the fitting.

Parameters

t_tensor

The type of the tensor. used in the type hint.

fwd_method_name

Name of the forward method. For dpmodels, it should be “call”. For torch models, it should be “forward”.

Submodules

deepmd.dpmodel.fitting.base_fitting module

deepmd.dpmodel.fitting.dipole_fitting module

class deepmd.dpmodel.fitting.dipole_fitting.DipoleFitting(*args, **kwargs)

Bases: DipoleFitting

Fitting rotationally equivariant diploe of the system.

Parameters

var_name

The name of the output variable.

ntypes

The number of atom types.

dim_descrpt

The dimension of the input descriptor.

embedding_widthint

The dimension of rotation matrix, m1.

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.

activation_function

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

precision

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

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_types

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

exclude_types

Atomic contributions of the excluded atom types are set zero.

r_differentiable

If the variable is differentiated with respect to coordinates of atoms. Only reduciable variable are differentiable.

c_differentiable

If the variable is differentiated with respect to the cell tensor (pbc case). Only reduciable variable are differentiable.

Methods

call(descriptor, atype[, gr, g2, h2, ...])	Calculate the fitting.
compute_output_stats(merged)	Update the output bias for fitting net.
deserialize(data)	Deserialize the fitting.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
output_def()	Returns the output def of the fitting net.
register(key)	Register a descriptor plugin.
serialize()	Serialize the fitting to dict.

__call__
reinit_exclude

deepmd.dpmodel.fitting.ener_fitting module

class deepmd.dpmodel.fitting.ener_fitting.EnergyFittingNet(*args, **kwargs)

Bases: InvarFitting

Methods

call(descriptor, atype[, gr, g2, h2, ...])	Calculate the fitting.
compute_output_stats(merged)	Update the output bias for fitting net.
deserialize(data)	Deserialize the fitting.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
init_fitting_stat(result_dict)	Initialize the model bias by the statistics.
output_def()	Returns the output def of the fitting net.
register(key)	Register a descriptor plugin.
serialize()	Serialize the fitting to dict.

__call__
reinit_exclude

classmethod deserialize(data: dict) → GeneralFitting

Deserialize the fitting.

Parameters

datadict

The serialized data

Returns

BF

The deserialized fitting

serialize() → dict

Serialize the fitting to dict.

deepmd.dpmodel.fitting.general_fitting module

class deepmd.dpmodel.fitting.general_fitting.GeneralFitting(*args, **kwargs)

Bases: NativeOP, BF

General fitting class.

Parameters

var_name

The name of the output variable.

ntypes

The number of atom types.

dim_descrpt

The dimension of the input descriptor.

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.

activation_function

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

precision

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

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_types

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

exclude_types: List[int]

Atomic contributions of the excluded atom types are set zero.

remove_vaccum_contribution: List[bool], optional

Remove vaccum contribution before the bias is added. The list assigned each type. For mixed_types provide [True], otherwise it should be a list of the same length as ntypes signaling if or not removing the vaccum contribution for the atom types in the list.

Methods

__call__(*args, **kwargs)	Forward pass in NumPy implementation.
call(*args, **kwargs)	Forward pass in NumPy implementation.
compute_output_stats(merged)	Update the output bias for fitting net.
deserialize(data)	Deserialize the fitting.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
output_def()	Returns the output def of the fitting net.
register(key)	Register a descriptor plugin.
serialize()	Serialize the fitting to dict.

reinit_exclude

classmethod deserialize(data: dict) → GeneralFitting

Deserialize the fitting.

Parameters

datadict

The serialized data

Returns

BF

The deserialized fitting

get_dim_aparam() → int

Get the number (dimension) of atomic parameters of this atomic model.

get_dim_fparam() → int

Get the number (dimension) of frame parameters of this atomic model.

get_sel_type() → List[int]

Get the selected atom types of this model.

Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected.

reinit_exclude(exclude_types: List[int] = [])

serialize() → dict

Serialize the fitting to dict.

deepmd.dpmodel.fitting.invar_fitting module

class deepmd.dpmodel.fitting.invar_fitting.InvarFitting(*args, **kwargs)

Bases: InvarFitting

Fitting the energy (or a rotationally invariant porperty of dim_out) of the system. The force and the virial can also be trained.

Lets take the energy fitting task as an example. 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

var_name

The name of the output variable.

ntypes

The number of atom types.

dim_descrpt

The dimension of the input descriptor.

dim_out

The dimension of the output fit property.

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.

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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

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

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_types

If false, different atomic types uses different fitting net, otherwise different atom types share the same fitting net.

exclude_types: List[int]

Atomic contributions of the excluded atom types are set zero.

Methods

call(descriptor, atype[, gr, g2, h2, ...])	Calculate the fitting.
compute_output_stats(merged)	Update the output bias for fitting net.
deserialize(data)	Deserialize the fitting.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
init_fitting_stat(result_dict)	Initialize the model bias by the statistics.
output_def()	Returns the output def of the fitting net.
register(key)	Register a descriptor plugin.
serialize()	Serialize the fitting to dict.

__call__
reinit_exclude

deepmd.dpmodel.fitting.make_base_fitting module

deepmd.dpmodel.fitting.make_base_fitting.make_base_fitting(t_tensor, fwd_method_name: str = 'forward')

Make the base class for the fitting.

Parameters

t_tensor

The type of the tensor. used in the type hint.

fwd_method_name

Name of the forward method. For dpmodels, it should be “call”. For torch models, it should be “forward”.

deepmd.dpmodel.fitting.polarizability_fitting module

class deepmd.dpmodel.fitting.polarizability_fitting.PolarFitting(*args, **kwargs)

Bases: PolarFitting

Fitting rotationally equivariant polarizability of the system.

Parameters

var_name

The name of the output variable.

ntypes

The number of atom types.

dim_descrpt

The dimension of the input descriptor.

embedding_widthint

The dimension of rotation matrix, m1.

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.

activation_function

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

precision

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

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_types

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

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.

Methods

call(descriptor, atype[, gr, g2, h2, ...])	Calculate the fitting.
compute_output_stats(merged)	Update the output bias for fitting net.
deserialize(data)	Deserialize the fitting.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
output_def()	Returns the output def of the fitting net.
register(key)	Register a descriptor plugin.
serialize()	Serialize the fitting to dict.

__call__
reinit_exclude