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 build (self,
input_d,
rot_mat,
natoms,
reuse = None,
suffix = '') :
start_index = 0
inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision)
rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]])
count = 0
outs_list = []
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice (inputs,
[ 0, start_index* self.dim_descrpt],
[-1, natoms[2+type_i]* self.dim_descrpt] )
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
rot_mat_i = tf.slice (rot_mat,
[ 0, start_index* 9],
[-1, natoms[2+type_i]* 9] )
rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3])
start_index += natoms[2+type_i]
if not type_i in self.sel_type :
continue
layer = inputs_i
for ii in range(0,len(self.n_neuron)) :
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
else :
layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
# (nframes x natoms) x 9
final_layer = one_layer(layer, 9, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, final_layer = True)
# (nframes x natoms) x 3 x 3
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 3, 3])
# (nframes x natoms) x 3 x 3
final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1])
# (nframes x natoms) x 3 x 3(coord)
final_layer = tf.matmul(final_layer, rot_mat_i)
# (nframes x natoms) x 3(coord) x 3(coord)
final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True)
# nframes x natoms x 3 x 3
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3])
# concat the results
outs_list.append(final_layer)
count += 1
outs = tf.concat(outs_list, axis = 1)
tf.summary.histogram('fitting_net_output', outs)
return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
def _filter(self,
inputs,
type_input,
natoms,
activation_fn=tf.nn.tanh,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
seed=None,
trainable=True):
# natom x (nei x 4)
shape = inputs.get_shape().as_list()
outputs_size = [1] + self.filter_neuron
outputs_size_2 = self.n_axis_neuron
with tf.variable_scope(name, reuse=reuse):
start_index = 0
xyz_scatter_total = []
for type_i in range(self.ntypes):
# cut-out inputs
# with natom x (nei_type_i x 4)
inputs_i = tf.slice(inputs, [0, start_index * 4],
[-1, self.sel_a[type_i] * 4])
start_index += self.sel_a[type_i]
shape_i = inputs_i.get_shape().as_list()
# with (natom x nei_type_i) x 4
inputs_reshape = tf.reshape(inputs_i, [-1, 4])
xyz_scatter = tf.reshape(
tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
if (type_input, type_i) not in self.exclude_types:
for ii in range(1, len(outputs_size)):
w = tf.get_variable(
'matrix_' + str(ii) + '_' + str(type_i),
[outputs_size[ii - 1], outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(
stddev=stddev / np.sqrt(outputs_size[ii] +
outputs_size[ii - 1]),
seed=seed),
trainable=trainable)
b = tf.get_variable(
'bias_' + str(ii) + '_' + str(type_i),
[1, outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(stddev=stddev,
mean=bavg,
seed=seed),
trainable=trainable)
if self.filter_resnet_dt:
idt = tf.get_variable(
'idt_' + str(ii) + '_' + str(type_i),
[1, outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(stddev=0.001,
mean=1.0,
seed=seed),
trainable=trainable)
if outputs_size[ii] == outputs_size[ii - 1]:
if self.filter_resnet_dt:
xyz_scatter += activation_fn(
tf.matmul(xyz_scatter, w) + b) * idt
else:
xyz_scatter += activation_fn(
tf.matmul(xyz_scatter, w) + b)
elif outputs_size[ii] == outputs_size[ii - 1] * 2:
if self.filter_resnet_dt:
xyz_scatter = tf.concat(
[xyz_scatter, xyz_scatter],
1) + activation_fn(
tf.matmul(xyz_scatter, w) + b) * idt
else:
xyz_scatter = tf.concat(
[xyz_scatter, xyz_scatter],
1) + activation_fn(
tf.matmul(xyz_scatter, w) + b)
else:
xyz_scatter = activation_fn(
tf.matmul(xyz_scatter, w) + b)
else:
w = tf.zeros((outputs_size[0], outputs_size[-1]),
dtype=global_tf_float_precision)
xyz_scatter = tf.matmul(xyz_scatter, w)
# natom x nei_type_i 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)
# 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[-1]])
return result, qmat
def build(self, input_d, rot_mat, natoms, reuse=None, suffix=''):
start_index = 0
inputs = tf.cast(
tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]),
self.fitting_precision)
rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]])
count = 0
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice(inputs, [0, start_index * self.dim_descrpt],
[-1, natoms[2 + type_i] * self.dim_descrpt])
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
rot_mat_i = tf.slice(rot_mat, [0, start_index * self.dim_rot_mat],
[-1, natoms[2 + type_i] * self.dim_rot_mat])
rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3])
start_index += natoms[2 + type_i]
if not type_i in self.sel_type:
continue
layer = inputs_i
for ii in range(0, len(self.n_neuron)):
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii - 1]:
layer += one_layer(
layer,
self.n_neuron[ii],
name='layer_' + str(ii) + '_type_' + str(type_i) +
suffix,
reuse=reuse,
seed=self.seed,
use_timestep=self.resnet_dt,
activation_fn=self.fitting_activation_fn,
precision=self.fitting_precision)
else:
layer = one_layer(layer,
self.n_neuron[ii],
name='layer_' + str(ii) + '_type_' +
str(type_i) + suffix,
reuse=reuse,
seed=self.seed,
activation_fn=self.fitting_activation_fn,
precision=self.fitting_precision)
if self.fit_diag:
bavg = np.zeros(self.dim_rot_mat_1)
# bavg[0] = self.avgeig[0]
# bavg[1] = self.avgeig[1]
# bavg[2] = self.avgeig[2]
# (nframes x natoms) x naxis
final_layer = one_layer(layer,
self.dim_rot_mat_1,
activation_fn=None,
name='final_layer_type_' +
str(type_i) + suffix,
reuse=reuse,
seed=self.seed,
bavg=bavg,
precision=self.fitting_precision)
# (nframes x natoms) x naxis
final_layer = tf.reshape(final_layer, [
tf.shape(inputs)[0] * natoms[2 + type_i],
self.dim_rot_mat_1
])
# (nframes x natoms) x naxis x naxis
final_layer = tf.matrix_diag(final_layer)
else:
bavg = np.zeros(self.dim_rot_mat_1 * self.dim_rot_mat_1)
# bavg[0*self.dim_rot_mat_1+0] = self.avgeig[0]
# bavg[1*self.dim_rot_mat_1+1] = self.avgeig[1]
# bavg[2*self.dim_rot_mat_1+2] = self.avgeig[2]
# (nframes x natoms) x (naxis x naxis)
final_layer = one_layer(
layer,
self.dim_rot_mat_1 * self.dim_rot_mat_1,
activation_fn=None,
name='final_layer_type_' + str(type_i) + suffix,
reuse=reuse,
seed=self.seed,
bavg=bavg,
precision=self.fitting_precision)
# (nframes x natoms) x naxis x naxis
final_layer = tf.reshape(final_layer, [
tf.shape(inputs)[0] * natoms[2 + type_i],
self.dim_rot_mat_1, self.dim_rot_mat_1
])
# (nframes x natoms) x naxis x naxis
final_layer = final_layer + tf.transpose(final_layer,
perm=[0, 2, 1])
# (nframes x natoms) x naxis x 3(coord)
final_layer = tf.matmul(final_layer, rot_mat_i)
# (nframes x natoms) x 3(coord) x 3(coord)
final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a=True)
# nframes x natoms x 3 x 3
final_layer = tf.reshape(
final_layer, [tf.shape(inputs)[0], natoms[2 + type_i], 3, 3])
# shift and scale
sel_type_idx = self.sel_type.index(type_i)
final_layer = final_layer * self.scale[sel_type_idx]
final_layer = final_layer + self.diag_shift[sel_type_idx] * tf.eye(
3,
batch_shape=[tf.shape(inputs)[0], natoms[2 + type_i]],
dtype=global_tf_float_precision)
# concat the results
if count == 0:
outs = final_layer
else:
outs = tf.concat([outs, final_layer], axis=1)
count += 1
return tf.cast(tf.reshape(outs, [-1]), global_tf_float_precision)
def _filter(self,
inputs,
type_input,
natoms,
type_embedding=None,
activation_fn=tf.nn.tanh,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
trainable=True):
nframes = tf.shape(tf.reshape(inputs,
[-1, natoms[0], self.ndescrpt]))[0]
# natom x (nei x 4)
shape = inputs.get_shape().as_list()
outputs_size = [1] + self.filter_neuron
outputs_size_2 = self.n_axis_neuron
all_excluded = all([(type_input, type_i) in self.exclude_types
for type_i in range(self.ntypes)])
if all_excluded:
# all types are excluded so result and qmat should be zeros
# we can safaly return a zero matrix...
# See also https://stackoverflow.com/a/34725458/9567349
# result: natom x outputs_size x outputs_size_2
# qmat: natom x outputs_size x 3
natom = tf.shape(inputs)[0]
result = tf.cast(
tf.fill((natom, outputs_size_2, outputs_size[-1]), 0.),
GLOBAL_TF_FLOAT_PRECISION)
qmat = tf.cast(tf.fill((natom, outputs_size[-1], 3), 0.),
GLOBAL_TF_FLOAT_PRECISION)
return result, qmat
with tf.variable_scope(name, reuse=reuse):
start_index = 0
type_i = 0
# natom x 4 x outputs_size
if type_embedding is None:
rets = []
for type_i in range(self.ntypes):
ret = self._filter_lower(type_i,
type_input,
start_index,
self.sel_a[type_i],
inputs,
nframes,
natoms,
type_embedding=type_embedding,
is_exclude=(type_input, type_i)
in self.exclude_types,
activation_fn=activation_fn,
stddev=stddev,
bavg=bavg,
trainable=trainable,
suffix="_" + str(type_i))
if (type_input, type_i) not in self.exclude_types:
# add zero is meaningless; skip
rets.append(ret)
start_index += self.sel_a[type_i]
# faster to use accumulate_n than multiple add
xyz_scatter_1 = tf.accumulate_n(rets)
else:
xyz_scatter_1 = self._filter_lower(
type_i,
type_input,
start_index,
np.cumsum(self.sel_a)[-1],
inputs,
nframes,
natoms,
type_embedding=type_embedding,
is_exclude=False,
activation_fn=activation_fn,
stddev=stddev,
bavg=bavg,
trainable=trainable)
# natom x nei x outputs_size
# xyz_scatter = tf.concat(xyz_scatter_total, axis=1)
# 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)
if self.original_sel is None:
# shape[1] = nnei * 4
nnei = shape[1] / 4
else:
nnei = tf.cast(
tf.Variable(np.sum(self.original_sel),
dtype=tf.int32,
trainable=False,
name="nnei"), self.filter_precision)
xyz_scatter_1 = xyz_scatter_1 / nnei
# 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_1 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[-1]])
return result, qmat
def build (self,
input_d : tf.Tensor,
rot_mat : tf.Tensor,
natoms : tf.Tensor,
reuse : bool = None,
suffix : str = '') :
"""
Build the computational graph for fitting net
Parameters
----------
input_d
The input descriptor
rot_mat
The rotation matrix from the descriptor.
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
reuse
The weights in the networks should be reused when get the variable.
suffix
Name suffix to identify this descriptor
Returns
-------
atomic_polar
The atomic polarizability
"""
start_index = 0
inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]])
rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]])
count = 0
outs_list = []
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice (inputs,
[ 0, start_index* self.dim_descrpt],
[-1, natoms[2+type_i]* self.dim_descrpt] )
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
rot_mat_i = tf.slice (rot_mat,
[ 0, start_index* self.dim_rot_mat],
[-1, natoms[2+type_i]* self.dim_rot_mat] )
rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3])
start_index += natoms[2+type_i]
if not type_i in self.sel_type :
continue
layer = inputs_i
for ii in range(0,len(self.n_neuron)) :
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec)
else :
layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
if self.fit_diag :
bavg = np.zeros(self.dim_rot_mat_1)
# bavg[0] = self.avgeig[0]
# bavg[1] = self.avgeig[1]
# bavg[2] = self.avgeig[2]
# (nframes x natoms) x naxis
final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
# (nframes x natoms) x naxis
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1])
# (nframes x natoms) x naxis x naxis
final_layer = tf.matrix_diag(final_layer)
else :
bavg = np.zeros(self.dim_rot_mat_1*self.dim_rot_mat_1)
# bavg[0*self.dim_rot_mat_1+0] = self.avgeig[0]
# bavg[1*self.dim_rot_mat_1+1] = self.avgeig[1]
# bavg[2*self.dim_rot_mat_1+2] = self.avgeig[2]
# (nframes x natoms) x (naxis x naxis)
final_layer = one_layer(layer, self.dim_rot_mat_1*self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
# (nframes x natoms) x naxis x naxis
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1, self.dim_rot_mat_1])
# (nframes x natoms) x naxis x naxis
final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1])
# (nframes x natoms) x naxis x 3(coord)
final_layer = tf.matmul(final_layer, rot_mat_i)
# (nframes x natoms) x 3(coord) x 3(coord)
final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True)
# nframes x natoms x 3 x 3
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3])
# shift and scale
sel_type_idx = self.sel_type.index(type_i)
final_layer = final_layer * self.scale[sel_type_idx]
final_layer = final_layer + self.constant_matrix[sel_type_idx] * tf.eye(3, batch_shape=[tf.shape(inputs)[0], natoms[2+type_i]], dtype = GLOBAL_TF_FLOAT_PRECISION)
# concat the results
outs_list.append(final_layer)
count += 1
outs = tf.concat(outs_list, axis = 1)
tf.summary.histogram('fitting_net_output', outs)
return tf.reshape(outs, [-1])
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
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 build(self, input_d, rot_mat, natoms, reuse=None, suffix=''):
start_index = 0
inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]])
rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]])
shape = inputs.get_shape().as_list()
count = 0
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice(inputs, [0, start_index * self.dim_descrpt],
[-1, natoms[2 + type_i] * self.dim_descrpt])
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
rot_mat_i = tf.slice(rot_mat, [0, start_index * self.dim_rot_mat],
[-1, natoms[2 + type_i] * self.dim_rot_mat])
rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3])
start_index += natoms[2 + type_i]
if not type_i in self.sel_type:
continue
layer = inputs_i
for ii in range(0, len(self.n_neuron)):
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii - 1]:
layer += one_layer(layer,
self.n_neuron[ii],
name='layer_' + str(ii) + '_type_' +
str(type_i) + suffix,
reuse=reuse,
seed=self.seed,
use_timestep=self.resnet_dt)
else:
layer = one_layer(layer,
self.n_neuron[ii],
name='layer_' + str(ii) + '_type_' +
str(type_i) + suffix,
reuse=reuse,
seed=self.seed)
# (nframes x natoms) x (naxis x naxis)
final_layer = one_layer(layer,
self.dim_rot_mat_1 * self.dim_rot_mat_1,
activation_fn=None,
name='final_layer_type_' + str(type_i) +
suffix,
reuse=reuse,
seed=self.seed)
# (nframes x natoms) x naxis x naxis
final_layer = tf.reshape(final_layer, [
tf.shape(inputs)[0] * natoms[2 + type_i], self.dim_rot_mat_1,
self.dim_rot_mat_1
])
# (nframes x natoms) x naxis x naxis
final_layer = final_layer + tf.transpose(final_layer,
perm=[0, 2, 1])
# (nframes x natoms) x naxis x 3(coord)
final_layer = tf.matmul(final_layer, rot_mat_i)
# (nframes x natoms) x 3(coord) x 3(coord)
final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a=True)
# nframes x natoms x 3 x 3
final_layer = tf.reshape(
final_layer, [tf.shape(inputs)[0], natoms[2 + type_i], 3, 3])
# concat the results
if count == 0:
outs = final_layer
else:
outs = tf.concat([outs, final_layer], axis=1)
count += 1
return tf.reshape(outs, [-1])
