deepmd.tf.model.multi

Module Contents

Classes

MultiModel

Multi-task model.

class deepmd.tf.model.multi.MultiModel(descriptor: dict, fitting_net_dict: dict, fitting_type_dict: dict | None = None, type_embedding=None, type_map: List[str] | None = None, data_stat_nbatch: int = 10, data_stat_protect: float = 0.01, use_srtab: str | None = None, smin_alpha: float | None = None, sw_rmin: float | None = None, sw_rmax: float | None = None, **kwargs)

Bases: deepmd.tf.model.model.Model

Multi-task model.

Parameters:

descriptor

Descriptor

fitting_net_dict

Dictionary of fitting nets

fitting_type_dict

deprecated argument

type_embedding

Type embedding net

type_map

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

data_stat_nbatch

Number of frames used for data statistic

data_stat_protect

Protect parameter for atomic energy regression

use_srtab

The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.

smin_alpha

The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when use_srtab is provided.

sw_rmin

The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.

sw_rmin

The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.

model_type = 'multi_task'

get_rcut()

Get cutoff radius of the model.

get_ntypes()

Get the number of types.

get_type_map()

Get the type map.

data_stat(data)

Data staticis.

_compute_input_stat(all_stat, protection=0.01, mixed_type=False, fitting_key='')

_compute_output_stat(all_stat, mixed_type=False, fitting_key='')

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: str | None = None, suffix='', reuse=None)

Build the model.

Parameters:

coord_tf.Tensor

The coordinates of atoms

atype_tf.Tensor

The atom types of atoms

natomstf.Tensor

The number of atoms

boxtf.Tensor

The box vectors

meshtf.Tensor

The mesh vectors

input_dictdict

The input dict

frz_modelstr, optional

The path to the frozen model

ckpt_metastr, optional

The path prefix of the checkpoint and meta files

suffixstr, optional

The suffix of the scope

reusebool or tf.AUTO_REUSE, optional

Whether to reuse the variables

Returns:

dict

The output dict

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

Init the embedding net variables with the given frozen model.

Parameters:

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

enable_mixed_precision(mixed_prec: dict)

Enable mixed precision for the model.

Parameters:

mixed_precdict

The mixed precision config

get_numb_fparam() → dict

Get the number of frame parameters.

get_numb_aparam() → dict

Get the number of atomic parameters.

get_numb_dos() → dict

Get the number of gridpoints in energy space.

get_fitting() → dict

Get the fitting(s).

get_loss(loss: dict, lr: dict) → Dict[str, deepmd.tf.loss.loss.Loss]

Get the loss function(s).

classmethod update_sel(global_jdata: dict, local_jdata: dict)

Update the selection and perform neighbor statistics.

Parameters:

global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class