Source code for deepmd.descriptor.se_a_ebd

# SPDX-License-Identifier: LGPL-3.0-or-later
from typing import (
    List,
    Optional,
)

import numpy as np

from deepmd.common import (
    add_data_requirement,
)
from deepmd.env import (
    GLOBAL_TF_FLOAT_PRECISION,
    op_module,
    tf,
)
from deepmd.utils.network import (
    embedding_net,
    one_layer,
)

from .descriptor import (
    Descriptor,
)
from .se_a import (
    DescrptSeA,
)


[docs]@Descriptor.register("se_a_tpe") @Descriptor.register("se_a_ebd") class DescrptSeAEbd(DescrptSeA): r"""DeepPot-SE descriptor with type embedding approach. Parameters ---------- rcut The cut-off radius rcut_smth From where the environment matrix should be smoothed sel : list[int] sel[i] specifies the maxmum number of type i atoms in the cut-off radius neuron : list[int] Number of neurons in each hidden layers of the embedding net axis_neuron Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix) resnet_dt Time-step `dt` in the resnet construction: y = x + dt * \phi (Wx + b) trainable If the weights of embedding net are trainable. seed Random seed for initializing the network parameters. type_one_side Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets type_nchanl Number of channels for type representation type_nlayer Number of hidden layers for the type embedding net (skip connected). numb_aparam Number of atomic parameters. If >0 it will be embedded with atom types. set_davg_zero Set the shift of embedding net input to zero. activation_function The activation function in the embedding net. Supported options are {0} precision The precision of the embedding net parameters. Supported options are {1} exclude_types : List[List[int]] The excluded pairs of types which have no interaction with each other. For example, `[[0, 1]]` means no interaction between type 0 and type 1. """ def __init__( self, rcut: float, rcut_smth: float, sel: List[int], neuron: List[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: Optional[int] = None, type_one_side: bool = True, type_nchanl: int = 2, type_nlayer: int = 1, numb_aparam: int = 0, set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", exclude_types: List[List[int]] = [], **kwargs, ) -> None: """Constructor.""" DescrptSeA.__init__( self, rcut, rcut_smth, sel, neuron=neuron, axis_neuron=axis_neuron, resnet_dt=resnet_dt, trainable=trainable, seed=seed, type_one_side=type_one_side, set_davg_zero=set_davg_zero, activation_function=activation_function, precision=precision, ) self.type_nchanl = type_nchanl self.type_nlayer = type_nlayer self.type_one_side = type_one_side self.numb_aparam = numb_aparam if self.numb_aparam > 0: add_data_requirement("aparam", 3, atomic=True, must=True, high_prec=False)
[docs] def build( self, coord_: tf.Tensor, atype_: tf.Tensor, natoms: tf.Tensor, box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = "", ) -> 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 """ nei_type = [] for ii in range(self.ntypes): nei_type.append(ii * np.ones(self.sel_a[ii], dtype=int)) nei_type = np.concatenate(nei_type) self.nei_type = tf.get_variable( "t_nei_type", [self.nnei], dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(nei_type), ) self.dout = DescrptSeA.build( self, coord_, atype_, natoms, box_, mesh, input_dict, suffix=suffix, reuse=reuse, ) tf.summary.histogram("embedding_net_output", self.dout) return self.dout
def _type_embed(self, atype, ndim=1, reuse=None, suffix="", trainable=True): ebd_type = tf.cast(atype, self.filter_precision) ebd_type = ebd_type / float(self.ntypes) ebd_type = tf.reshape(ebd_type, [-1, ndim]) for ii in range(self.type_nlayer): name = "type_embed_layer_" + str(ii) ebd_type = one_layer( ebd_type, self.type_nchanl, activation_fn=self.filter_activation_fn, precision=self.filter_precision, name=name, reuse=reuse, seed=self.seed + ii, trainable=trainable, ) name = "type_embed_layer_" + str(self.type_nlayer) ebd_type = one_layer( ebd_type, self.type_nchanl, activation_fn=None, precision=self.filter_precision, name=name, reuse=reuse, seed=self.seed + ii, trainable=trainable, ) ebd_type = tf.reshape(ebd_type, [tf.shape(atype)[0], self.type_nchanl]) return ebd_type def _embedding_net( self, inputs, natoms, filter_neuron, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name="linear", reuse=None, seed=None, trainable=True, ): """inputs: nf x na x (nei x 4) outputs: nf x na x nei x output_size. """ # natom x (nei x 4) inputs = tf.reshape(inputs, [-1, self.ndescrpt]) shape = inputs.get_shape().as_list() outputs_size = [1, *filter_neuron] with tf.variable_scope(name, reuse=reuse): xyz_scatter_total = [] # with natom x (nei x 4) inputs_i = inputs shape_i = inputs_i.get_shape().as_list() # with (natom x nei) x 4 inputs_reshape = tf.reshape(inputs_i, [-1, 4]) # with (natom x nei) x 1 xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1]) # with (natom x nei) x out_size xyz_scatter = embedding_net( xyz_scatter, self.filter_neuron, self.filter_precision, activation_fn=activation_fn, resnet_dt=self.filter_resnet_dt, stddev=stddev, bavg=bavg, seed=seed, trainable=trainable, ) # natom x nei x out_size xyz_scatter = tf.reshape( xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1]) ) xyz_scatter_total.append(xyz_scatter) # natom x nei x outputs_size xyz_scatter = tf.concat(xyz_scatter_total, axis=1) # nf x natom x nei x outputs_size xyz_scatter = tf.reshape( xyz_scatter, [tf.shape(inputs)[0], natoms[0], self.nnei, outputs_size[-1]] ) return xyz_scatter def _type_embedding_net_two_sides( self, mat_g, atype, natoms, name="", reuse=None, seed=None, trainable=True ): outputs_size = self.filter_neuron[-1] nframes = tf.shape(mat_g)[0] # (nf x natom x nei) x (outputs_size x chnl x chnl) mat_g = tf.reshape(mat_g, [nframes * natoms[0] * self.nnei, outputs_size]) mat_g = one_layer( mat_g, outputs_size * self.type_nchanl * self.type_nchanl, activation_fn=None, precision=self.filter_precision, name=name + "_amplify", reuse=reuse, seed=self.seed, trainable=trainable, ) # nf x natom x nei x outputs_size x chnl x chnl mat_g = tf.reshape( mat_g, [ nframes, natoms[0], self.nnei, outputs_size, self.type_nchanl, self.type_nchanl, ], ) # nf x natom x outputs_size x chnl x nei x chnl mat_g = tf.transpose(mat_g, perm=[0, 1, 3, 4, 2, 5]) # nf x natom x outputs_size x chnl x (nei x chnl) mat_g = tf.reshape( mat_g, [ nframes, natoms[0], outputs_size, self.type_nchanl, self.nnei * self.type_nchanl, ], ) # nei x nchnl ebd_nei_type = self._type_embed( self.nei_type, reuse=reuse, trainable=True, suffix="" ) # (nei x nchnl) ebd_nei_type = tf.reshape(ebd_nei_type, [self.nnei * self.type_nchanl]) # (nframes x natom) x nchnl ebd_atm_type = self._type_embed(atype, reuse=True, trainable=True, suffix="") # (nframes x natom x nchnl) ebd_atm_type = tf.reshape( ebd_atm_type, [nframes * natoms[0] * self.type_nchanl] ) # nf x natom x outputs_size x chnl x (nei x chnl) mat_g = tf.multiply(mat_g, ebd_nei_type) # nf x natom x outputs_size x chnl x nei x chnl mat_g = tf.reshape( mat_g, [ nframes, natoms[0], outputs_size, self.type_nchanl, self.nnei, self.type_nchanl, ], ) # nf x natom x outputs_size x chnl x nei mat_g = tf.reduce_mean(mat_g, axis=5) # outputs_size x nei x nf x natom x chnl mat_g = tf.transpose(mat_g, perm=[2, 4, 0, 1, 3]) # outputs_size x nei x (nf x natom x chnl) mat_g = tf.reshape( mat_g, [outputs_size, self.nnei, nframes * natoms[0] * self.type_nchanl] ) # outputs_size x nei x (nf x natom x chnl) mat_g = tf.multiply(mat_g, ebd_atm_type) # outputs_size x nei x nf x natom x chnl mat_g = tf.reshape( mat_g, [outputs_size, self.nnei, nframes, natoms[0], self.type_nchanl] ) # outputs_size x nei x nf x natom mat_g = tf.reduce_mean(mat_g, axis=4) # nf x natom x nei x outputs_size mat_g = tf.transpose(mat_g, perm=[2, 3, 1, 0]) # (nf x natom) x nei x outputs_size mat_g = tf.reshape(mat_g, [nframes * natoms[0], self.nnei, outputs_size]) return mat_g def _type_embedding_net_one_side( self, mat_g, atype, natoms, name="", reuse=None, seed=None, trainable=True ): outputs_size = self.filter_neuron[-1] nframes = tf.shape(mat_g)[0] # (nf x natom x nei) x (outputs_size x chnl x chnl) mat_g = tf.reshape(mat_g, [nframes * natoms[0] * self.nnei, outputs_size]) mat_g = one_layer( mat_g, outputs_size * self.type_nchanl, activation_fn=None, precision=self.filter_precision, name=name + "_amplify", reuse=reuse, seed=self.seed, trainable=trainable, ) # nf x natom x nei x outputs_size x chnl mat_g = tf.reshape( mat_g, [nframes, natoms[0], self.nnei, outputs_size, self.type_nchanl] ) # nf x natom x outputs_size x nei x chnl mat_g = tf.transpose(mat_g, perm=[0, 1, 3, 2, 4]) # nf x natom x outputs_size x (nei x chnl) mat_g = tf.reshape( mat_g, [nframes, natoms[0], outputs_size, self.nnei * self.type_nchanl] ) # nei x nchnl ebd_nei_type = self._type_embed( self.nei_type, reuse=reuse, trainable=True, suffix="" ) # (nei x nchnl) ebd_nei_type = tf.reshape(ebd_nei_type, [self.nnei * self.type_nchanl]) # nf x natom x outputs_size x (nei x chnl) mat_g = tf.multiply(mat_g, ebd_nei_type) # nf x natom x outputs_size x nei x chnl mat_g = tf.reshape( mat_g, [nframes, natoms[0], outputs_size, self.nnei, self.type_nchanl] ) # nf x natom x outputs_size x nei mat_g = tf.reduce_mean(mat_g, axis=4) # nf x natom x nei x outputs_size mat_g = tf.transpose(mat_g, perm=[0, 1, 3, 2]) # (nf x natom) x nei x outputs_size mat_g = tf.reshape(mat_g, [nframes * natoms[0], self.nnei, outputs_size]) return mat_g def _type_embedding_net_one_side_aparam( self, mat_g, atype, natoms, aparam, name="", reuse=None, seed=None, trainable=True, ): outputs_size = self.filter_neuron[-1] nframes = tf.shape(mat_g)[0] # (nf x natom x nei) x (outputs_size x chnl x chnl) mat_g = tf.reshape(mat_g, [nframes * natoms[0] * self.nnei, outputs_size]) mat_g = one_layer( mat_g, outputs_size * self.type_nchanl, activation_fn=None, precision=self.filter_precision, name=name + "_amplify", reuse=reuse, seed=self.seed, trainable=trainable, ) # nf x natom x nei x outputs_size x chnl mat_g = tf.reshape( mat_g, [nframes, natoms[0], self.nnei, outputs_size, self.type_nchanl] ) # outputs_size x nf x natom x nei x chnl mat_g = tf.transpose(mat_g, perm=[3, 0, 1, 2, 4]) # outputs_size x (nf x natom x nei x chnl) mat_g = tf.reshape( mat_g, [outputs_size, nframes * natoms[0] * self.nnei * self.type_nchanl] ) # nf x natom x nnei embed_type = tf.tile( tf.reshape(self.nei_type, [1, self.nnei]), [nframes * natoms[0], 1] ) # (nf x natom x nnei) x 1 embed_type = tf.reshape(embed_type, [nframes * natoms[0] * self.nnei, 1]) # nf x (natom x naparam) aparam = tf.reshape(aparam, [nframes, -1]) # nf x natom x nnei x naparam embed_aparam = op_module.map_aparam( aparam, self.nlist, natoms, n_a_sel=self.nnei_a, n_r_sel=self.nnei_r ) # (nf x natom x nnei) x naparam embed_aparam = tf.reshape( embed_aparam, [nframes * natoms[0] * self.nnei, self.numb_aparam] ) # (nf x natom x nnei) x (naparam+1) embed_input = tf.concat((embed_type, embed_aparam), axis=1) # (nf x natom x nnei) x nchnl ebd_nei_type = self._type_embed( embed_input, ndim=self.numb_aparam + 1, reuse=reuse, trainable=True, suffix="", ) # (nf x natom x nei x nchnl) ebd_nei_type = tf.reshape( ebd_nei_type, [nframes * natoms[0] * self.nnei * self.type_nchanl] ) # outputs_size x (nf x natom x nei x chnl) mat_g = tf.multiply(mat_g, ebd_nei_type) # outputs_size x nf x natom x nei x chnl mat_g = tf.reshape( mat_g, [outputs_size, nframes, natoms[0], self.nnei, self.type_nchanl] ) # outputs_size x nf x natom x nei mat_g = tf.reduce_mean(mat_g, axis=4) # nf x natom x nei x outputs_size mat_g = tf.transpose(mat_g, perm=[1, 2, 3, 0]) # (nf x natom) x nei x outputs_size mat_g = tf.reshape(mat_g, [nframes * natoms[0], self.nnei, outputs_size]) return mat_g def _pass_filter( self, inputs, atype, natoms, input_dict, reuse=None, suffix="", trainable=True ): # nf x na x ndescrpt # nf x na x (nnei x 4) inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]) layer, qmat = self._ebd_filter( tf.cast(inputs, self.filter_precision), atype, natoms, input_dict, name="filter_type_all" + suffix, reuse=reuse, seed=self.seed, trainable=trainable, activation_fn=self.filter_activation_fn, ) output = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0], self.get_dim_out()]) output_qmat = tf.reshape( qmat, [tf.shape(inputs)[0], natoms[0], self.get_dim_rot_mat_1() * 3] ) return output, output_qmat def _ebd_filter( self, inputs, atype, natoms, input_dict, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name="linear", reuse=None, seed=None, trainable=True, ): outputs_size = self.filter_neuron[-1] outputs_size_2 = self.n_axis_neuron # nf x natom x (nei x 4) nframes = tf.shape(inputs)[0] shape = tf.reshape(inputs, [-1, self.ndescrpt]).get_shape().as_list() # nf x natom x nei x outputs_size mat_g = self._embedding_net( inputs, natoms, self.filter_neuron, activation_fn=activation_fn, stddev=stddev, bavg=bavg, name=name, reuse=reuse, seed=seed, trainable=trainable, ) # nf x natom x nei x outputs_size mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size]) # (nf x natom) x nei x outputs_size if self.type_one_side: if self.numb_aparam > 0: aparam = input_dict["aparam"] xyz_scatter = self._type_embedding_net_one_side_aparam( mat_g, atype, natoms, aparam, name=name, reuse=reuse, seed=seed, trainable=trainable, ) else: xyz_scatter = self._type_embedding_net_one_side( mat_g, atype, natoms, name=name, reuse=reuse, seed=seed, trainable=trainable, ) else: xyz_scatter = self._type_embedding_net_two_sides( mat_g, atype, natoms, name=name, reuse=reuse, seed=seed, trainable=trainable, ) # natom x nei x 4 inputs_reshape = tf.reshape(inputs, [-1, shape[1] // 4, 4]) # natom x 4 x outputs_size xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a=True) xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape[1]) # natom x 4 x outputs_size_2 xyz_scatter_2 = tf.slice(xyz_scatter_1, [0, 0, 0], [-1, -1, outputs_size_2]) # # natom x 3 x outputs_size_2 # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1]) # natom x 3 x outputs_size_1 qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1]) # natom x outputs_size_2 x 3 qmat = tf.transpose(qmat, perm=[0, 2, 1]) # natom x outputs_size x outputs_size_2 result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True) # natom x (outputs_size x outputs_size_2) result = tf.reshape(result, [-1, outputs_size_2 * outputs_size]) return result, qmat