deepmd.tf.descriptor.hybrid

Module Contents

Classes

DescrptHybrid

Concate a list of descriptors to form a new descriptor.

class deepmd.tf.descriptor.hybrid.DescrptHybrid(list: List[deepmd.tf.descriptor.descriptor.Descriptor | Dict[str, Any]], multi_task: bool = False, ntypes: int | None = None, spin: deepmd.tf.utils.spin.Spin | None = None, **kwargs)

Bases: deepmd.tf.descriptor.descriptor.Descriptor

Concate a list of descriptors to form a new descriptor.

Parameters:

listlistList[Union[Descriptor, Dict[str, Any]]]

Build a descriptor from the concatenation of the list of descriptors. The descriptor can be either an object or a dictionary.

property explicit_ntypes: bool

Explicit ntypes with type embedding.

get_rcut() → float

Returns the cut-off radius.

get_ntypes() → int

Returns the number of atom types.

get_dim_out() → int

Returns the output dimension of this descriptor.

get_nlist() → Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, List[int], List[int]]

Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.

Returns:

nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_nlist_i(ii: int) → Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, List[int], List[int]]

Get the neighbor information of the ii-th descriptor.

Parameters:

iiint

The index of the descriptor

Returns:

nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, mixed_type: bool = False, real_natoms_vec: list | None = None, **kwargs) → None

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord

The coordinates. Can be generated by deepmd.tf.model.make_stat_input

data_box

The box. Can be generated by deepmd.tf.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.tf.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input

input_dict

Dictionary for additional 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.

real_natoms_vec

If mixed_type is True, it takes in the real natoms_vec for each frame.

**kwargs

Additional keyword arguments.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters:

stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr

The sum of radial statisitcs.

suma

The sum of relative coord statisitcs.

sumn

The sum of neighbor numbers.

sumr2

The sum of square of radial statisitcs.

suma2

The sum of square of relative coord statisitcs.

build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → deepmd.tf.env.tf.Tensor

Build the computational graph for the descriptor.

Parameters:

coord_

The coordinate of atoms

atype_

The type of atoms

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

box_tf.Tensor

The box of the system

mesh

For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.

input_dict

Dictionary for additional inputs

reuse

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

suffix

Name suffix to identify this descriptor

Returns:

descriptor

The output descriptor

prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) → Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor]

Compute force and virial.

Parameters:

atom_ener

The atomic energy

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

Returns:

force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

enable_compression(min_nbor_dist: float, graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, table_extrapolate: float = 5.0, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') → None

Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_distfloat

The nearest distance between atoms

graphtf.Graph

The graph of the model

graph_deftf.GraphDef

The graph_def of the model

table_extrapolatefloat, default: 5.

The scale of model extrapolation

table_stride_1float, default: 0.01

The uniform stride of the first table

table_stride_2float, default: 0.1

The uniform stride of the second table

check_frequencyint, default: -1

The overflow check frequency

suffixstr, optional

The suffix of the 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

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

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

suffixstr, optional

The suffix of the scope

get_tensor_names(suffix: str = '') → Tuple[str]

Get names of tensors.

Parameters:

suffixstr

The suffix of the scope

Returns:

Tuple[str]

Names of tensors

pass_tensors_from_frz_model(*tensors: deepmd.tf.env.tf.Tensor) → None

Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.

Parameters:

*tensorstf.Tensor

passed tensors

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

serialize(suffix: str = '') → dict

Serialize the model.

There is no suffix in a native DP model, but it is important for the TF backend.

Returns:

dict

The serialized data

suffixstr, optional

Name suffix to identify this descriptor

classmethod deserialize(data: dict, suffix: str = '') → DescrptHybrid

Deserialize the model.

There is no suffix in a native DP model, but it is important for the TF backend.

Parameters:

datadict

The serialized data

suffixstr, optional

Name suffix to identify this descriptor

Returns:

Descriptor

The deserialized descriptor