# 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[str]
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[str],
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