def test_pooling_LSTM(self):
node = tf.ragged.constant(self.n1, ragged_rank=1, inner_shape=(1, ))
edgeind = tf.ragged.constant(self.ei1,
ragged_rank=1,
inner_shape=(2, ))
edgefeat = tf.ragged.constant(self.e1,
ragged_rank=1,
inner_shape=(1, ))
node_indexing = 'sample'
partition_type = 'row_length'
tens_type = "values_partition"
n = ChangeTensorType(input_tensor_type="ragged",
output_tensor_type=tens_type)(node)
ed = ChangeTensorType(input_tensor_type="ragged",
output_tensor_type=tens_type)(edgefeat)
edi = ChangeTensorType(input_tensor_type="ragged",
output_tensor_type=tens_type)(edgeind)
edi = ChangeIndexing(input_tensor_type=tens_type,
to_indexing=node_indexing)([n, edi])
ns = GatherNodes(input_tensor_type=tens_type,
node_indexing=node_indexing,
partition_type=partition_type)([n, edi])
messages = Concatenate(axis=-1, input_tensor_type=tens_type)([ed, ns])
out = PoolingLocalEdgesLSTM(units=3,
input_tensor_type=tens_type,
node_indexing=node_indexing)(
[n, messages, edi])
self.assertTrue(np.all(np.array(out[0].shape) == np.array([23, 3])))
def test_gather_nodes(self):
node = tf.ragged.constant(self.n1, ragged_rank=1, inner_shape=(1, ))
edgeind = tf.ragged.constant(self.ei1,
ragged_rank=1,
inner_shape=(2, ))
gathered_nodes = GatherNodes(concat_nodes=False)([node, edgeind])
np_gather = np.array(self.n1[1])[np.array(self.ei1[1])]
test = np.sum(np.abs(np.array(gathered_nodes[1]) - np_gather)) < 1e-6
self.assertTrue(test)
def test_gather_nodes_concat(self):
node = tf.ragged.constant(self.n1, ragged_rank=1, inner_shape=(1, ))
edgeind = tf.ragged.constant(self.ei1,
ragged_rank=1,
inner_shape=(2, ))
gathered_nodes_concat = GatherNodes()([node, edgeind])
np_gather = np.reshape(
np.array(self.n1[1])[np.array(self.ei1[1])], (28, 2 * 1))
test = np.sum(
np.abs(np.array(gathered_nodes_concat[1]) - np_gather)) < 1e-6
self.assertTrue(test)
def __init__(self,
node_embed=None,
edge_embed=None,
env_embed=None,
pooling_method="mean",
use_bias=True,
activation=None,
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
kernel_initializer='glorot_uniform',
bias_initializer='zeros',
**kwargs):
"""Initialize layer."""
super(MEGnetBlock, self).__init__(**kwargs)
self.pooling_method = pooling_method
if node_embed is None:
node_embed = [16, 16, 16]
if env_embed is None:
env_embed = [16, 16, 16]
if edge_embed is None:
edge_embed = [16, 16, 16]
self.node_embed = node_embed
self.edge_embed = edge_embed
self.env_embed = env_embed
self.use_bias = use_bias
if activation is None and 'softplus2' in kgcnn_custom_act:
activation = 'softplus2'
elif activation is None:
activation = "selu"
kernel_args = {
"kernel_regularizer": kernel_regularizer,
"activity_regularizer": activity_regularizer,
"bias_regularizer": bias_regularizer,
"kernel_constraint": kernel_constraint,
"bias_constraint": bias_constraint,
"kernel_initializer": kernel_initializer,
"bias_initializer": bias_initializer,
"use_bias": use_bias
}
mlp_args = {
"input_tensor_type": self.input_tensor_type,
"ragged_validate": self.ragged_validate
}
mlp_args.update(kernel_args)
pool_args = {"pooling_method": self.pooling_method}
pool_args.update(self._all_kgcnn_info)
gather_args = self._all_kgcnn_info
# Node
self.lay_phi_n = Dense(units=self.node_embed[0],
activation=activation,
**mlp_args)
self.lay_phi_n_1 = Dense(units=self.node_embed[1],
activation=activation,
**mlp_args)
self.lay_phi_n_2 = Dense(units=self.node_embed[2],
activation='linear',
**mlp_args)
self.lay_esum = PoolingLocalEdges(**pool_args)
self.lay_gather_un = GatherState(**gather_args)
self.lay_conc_nu = Concatenate(
axis=-1, input_tensor_type=self.input_tensor_type)
# Edge
self.lay_phi_e = Dense(units=self.edge_embed[0],
activation=activation,
**mlp_args)
self.lay_phi_e_1 = Dense(units=self.edge_embed[1],
activation=activation,
**mlp_args)
self.lay_phi_e_2 = Dense(units=self.edge_embed[2],
activation='linear',
**mlp_args)
self.lay_gather_n = GatherNodes(**gather_args)
self.lay_gather_ue = GatherState(**gather_args)
self.lay_conc_enu = Concatenate(
axis=-1, input_tensor_type=self.input_tensor_type)
# Environment
self.lay_usum_e = PoolingGlobalEdges(**pool_args)
self.lay_usum_n = PoolingNodes(**pool_args)
self.lay_conc_u = Concatenate(axis=-1, input_tensor_type="tensor")
self.lay_phi_u = ks.layers.Dense(units=self.env_embed[0],
activation=activation,
**kernel_args)
self.lay_phi_u_1 = ks.layers.Dense(units=self.env_embed[1],
activation=activation,
**kernel_args)
self.lay_phi_u_2 = ks.layers.Dense(units=self.env_embed[2],
activation='linear',
**kernel_args)
def __init__(self,
node_embed=None,
edge_embed=None,
env_embed=None,
use_bias=True,
activation=None,
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
kernel_initializer='glorot_uniform',
bias_initializer='zeros',
pooling_method="mean",
is_sorted=False,
has_unconnected=True,
partition_type="row_length",
node_indexing='sample',
input_tensor_type="ragged",
ragged_validate=False,
**kwargs):
"""Initialize layer."""
super(MEGnetBlock, self).__init__(**kwargs)
self.pooling_method = pooling_method
self.is_sorted = is_sorted
self.has_unconnected = has_unconnected
self.partition_type = partition_type
self.node_indexing = node_indexing
self.input_tensor_type = input_tensor_type
self.ragged_validate = ragged_validate
self._tensor_input_type_implemented = [
"ragged", "values_partition", "disjoint", "tensor", "RaggedTensor"
]
self._supports_ragged_inputs = True
self._test_tensor_input = kgcnn_ops_static_test_tensor_input_type(
self.input_tensor_type, self._tensor_input_type_implemented,
self.node_indexing)
if node_embed is None:
node_embed = [16, 16, 16]
if env_embed is None:
env_embed = [16, 16, 16]
if edge_embed is None:
edge_embed = [16, 16, 16]
self.node_embed = node_embed
self.edge_embed = edge_embed
self.env_embed = env_embed
self.use_bias = use_bias
if activation is None and 'softplus2' in kgcnn_custom_act:
activation = 'softplus2'
elif activation is None:
activation = "selu"
self.megnet_activation = tf.keras.activations.get(activation)
self.megnet_kernel_initializer = tf.keras.initializers.get(
kernel_initializer)
self.megnet_bias_initializer = tf.keras.initializers.get(
bias_initializer)
self.megnet_kernel_regularizer = tf.keras.regularizers.get(
kernel_regularizer)
self.megnet_bias_regularizer = tf.keras.regularizers.get(
bias_regularizer)
self.megnet_activity_regularizer = tf.keras.regularizers.get(
activity_regularizer)
self.megnet_kernel_constraint = tf.keras.constraints.get(
kernel_constraint)
self.megnet_bias_constraint = tf.keras.constraints.get(bias_constraint)
kernel_args = {
"kernel_regularizer": kernel_regularizer,
"activity_regularizer": activity_regularizer,
"bias_regularizer": bias_regularizer,
"kernel_constraint": kernel_constraint,
"bias_constraint": bias_constraint,
"kernel_initializer": kernel_initializer,
"bias_initializer": bias_initializer
}
mlp_args = {
"input_tensor_type": self.input_tensor_type,
"ragged_validate": self.ragged_validate
}
mlp_args.update(kernel_args)
pool_args = {
"pooling_method": self.pooling_method,
"is_sorted": self.is_sorted,
"has_unconnected": self.has_unconnected,
"input_tensor_type": self.input_tensor_type,
"ragged_validate": self.ragged_validate,
"partition_type": self.partition_type,
"node_indexing": self.node_indexing
}
gather_args = {
"is_sorted": self.is_sorted,
"has_unconnected": self.has_unconnected,
"input_tensor_type": self.input_tensor_type,
"ragged_validate": self.ragged_validate,
"partition_type": self.partition_type,
"node_indexing": self.node_indexing
}
# Node
self.lay_phi_n = Dense(units=self.node_embed[0],
activation=activation,
use_bias=self.use_bias,
**mlp_args)
self.lay_phi_n_1 = Dense(units=self.node_embed[1],
activation=activation,
use_bias=self.use_bias,
**mlp_args)
self.lay_phi_n_2 = Dense(units=self.node_embed[2],
activation='linear',
use_bias=self.use_bias,
**mlp_args)
self.lay_esum = PoolingLocalEdges(**pool_args)
self.lay_gather_un = GatherState(**gather_args)
self.lay_conc_nu = Concatenate(
axis=-1, input_tensor_type=self.input_tensor_type)
# Edge
self.lay_phi_e = Dense(units=self.edge_embed[0],
activation=activation,
use_bias=self.use_bias,
**mlp_args)
self.lay_phi_e_1 = Dense(units=self.edge_embed[1],
activation=activation,
use_bias=self.use_bias,
**mlp_args)
self.lay_phi_e_2 = Dense(units=self.edge_embed[2],
activation='linear',
use_bias=self.use_bias,
**mlp_args)
self.lay_gather_n = GatherNodes(**gather_args)
self.lay_gather_ue = GatherState(**gather_args)
self.lay_conc_enu = Concatenate(
axis=-1, input_tensor_type=self.input_tensor_type)
# Environment
self.lay_usum_e = PoolingGlobalEdges(**pool_args)
self.lay_usum_n = PoolingNodes(**pool_args)
self.lay_conc_u = Concatenate(axis=-1, input_tensor_type="tensor")
self.lay_phi_u = ks.layers.Dense(units=self.env_embed[0],
activation=activation,
use_bias=self.use_bias,
**kernel_args)
self.lay_phi_u_1 = ks.layers.Dense(units=self.env_embed[1],
activation=activation,
use_bias=self.use_bias,
**kernel_args)
self.lay_phi_u_2 = ks.layers.Dense(units=self.env_embed[2],
activation='linear',
use_bias=self.use_bias,
**kernel_args)
def __init__(self,
node_embed=None,
edge_embed=None,
env_embed=None,
pooling_method="mean",
use_bias=True,
activation='kgcnn>softplus2',
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
kernel_initializer='glorot_uniform',
bias_initializer='zeros',
**kwargs):
"""Initialize layer."""
super(MEGnetBlock, self).__init__(**kwargs)
self.pooling_method = pooling_method
if node_embed is None:
node_embed = [16, 16, 16]
if env_embed is None:
env_embed = [16, 16, 16]
if edge_embed is None:
edge_embed = [16, 16, 16]
self.node_embed = node_embed
self.edge_embed = edge_embed
self.env_embed = env_embed
self.use_bias = use_bias
kernel_args = {
"kernel_regularizer": kernel_regularizer,
"activity_regularizer": activity_regularizer,
"bias_regularizer": bias_regularizer,
"kernel_constraint": kernel_constraint,
"bias_constraint": bias_constraint,
"kernel_initializer": kernel_initializer,
"bias_initializer": bias_initializer,
"use_bias": use_bias
}
# Node
self.lay_phi_n = Dense(units=self.node_embed[0],
activation=activation,
**kernel_args,
**self._kgcnn_info)
self.lay_phi_n_1 = Dense(units=self.node_embed[1],
activation=activation,
**kernel_args,
**self._kgcnn_info)
self.lay_phi_n_2 = Dense(units=self.node_embed[2],
activation='linear',
**kernel_args,
**self._kgcnn_info)
self.lay_esum = PoolingLocalEdges(pooling_method=self.pooling_method,
**self._kgcnn_info)
self.lay_gather_un = GatherState(**self._kgcnn_info)
self.lay_conc_nu = Concatenate(axis=-1, **self._kgcnn_info)
# Edge
self.lay_phi_e = Dense(units=self.edge_embed[0],
activation=activation,
**kernel_args,
**self._kgcnn_info)
self.lay_phi_e_1 = Dense(units=self.edge_embed[1],
activation=activation,
**kernel_args,
**self._kgcnn_info)
self.lay_phi_e_2 = Dense(units=self.edge_embed[2],
activation='linear',
**kernel_args,
**self._kgcnn_info)
self.lay_gather_n = GatherNodes(**self._kgcnn_info)
self.lay_gather_ue = GatherState(**self._kgcnn_info)
self.lay_conc_enu = Concatenate(axis=-1, **self._kgcnn_info)
# Environment
self.lay_usum_e = PoolingGlobalEdges(
pooling_method=self.pooling_method, **self._kgcnn_info)
self.lay_usum_n = PoolingNodes(pooling_method=self.pooling_method,
**self._kgcnn_info)
self.lay_conc_u = Concatenate(axis=-1, input_tensor_type="tensor")
self.lay_phi_u = ks.layers.Dense(units=self.env_embed[0],
activation=activation,
**kernel_args)
self.lay_phi_u_1 = ks.layers.Dense(units=self.env_embed[1],
activation=activation,
**kernel_args)
self.lay_phi_u_2 = ks.layers.Dense(units=self.env_embed[2],
activation='linear',
**kernel_args)
def make_dimnet_pp(
# Input
input_node_shape,
input_embedd: dict = None,
# Output
output_embedd: dict = None,
# Model specific parameter
emb_size = 128,
out_emb_size = 256,
int_emb_size = 64,
basis_emb_size =8,
num_blocks = 4,
num_spherical = 7,
num_radial= 6,
cutoff=5.0,
envelope_exponent=5,
num_before_skip=1,
num_after_skip=2,
num_dense_output=3,
num_targets=12,
activation="swish",
extensive=True,
output_init='zeros',
):
model_default = {'input_embedd': {'input_node_vocab': 95, 'input_node_embedd': 64, 'input_tensor_type': 'ragged'}
}
input_embedd = update_model_args(model_default['input_embedd'], input_embedd)
node_input, n, xyz_input, bond_index_input, angle_index_input, _ = generate_mol_graph_input(input_node_shape,
[None, 3],
[None, 2],
[None, 2],
**input_embedd)
x = xyz_input
edi = bond_index_input
adi = angle_index_input
# Calculate distances
d = NodeDistance()([x, edi])
rbf = BesselBasisLayer(num_radial=num_radial, cutoff=cutoff, envelope_exponent=envelope_exponent)(d)
# Calculate angles
a = EdgeAngle()([x, edi, adi])
sbf = SphericalBasisLayer(num_spherical=num_spherical, num_radial=num_radial, cutoff=cutoff,
envelope_exponent=envelope_exponent)([d, a, adi])
# Embedding block
rbf_emb = Dense(emb_size, use_bias=True, activation=activation, kernel_initializer="orthogonal")(rbf)
n_pairs = GatherNodes()([n, edi])
x = Concatenate(axis=-1)([n_pairs, rbf_emb])
x = Dense(emb_size, use_bias=True, activation=activation, kernel_initializer="orthogonal")(x)
ps = DimNetOutputBlock(emb_size, out_emb_size, num_dense_output, num_targets=num_targets,
output_kernel_initializer=output_init)([n, x, rbf, edi])
# Interaction blocks
add_xp = Add()
for i in range(num_blocks):
x = DimNetInteractionPPBlock(emb_size, int_emb_size, basis_emb_size, num_before_skip, num_after_skip)(
[x, rbf, sbf, adi])
p_update = DimNetOutputBlock(emb_size, out_emb_size, num_dense_output, num_targets=num_targets,
output_kernel_initializer=output_init)([n, x, rbf, edi])
ps = add_xp([ps, p_update])
if extensive:
main_output = PoolingNodes(pooling_method="sum")(ps)
else:
main_output = PoolingNodes(pooling_method="mean")(ps)
model = tf.keras.models.Model(inputs=[node_input, xyz_input, bond_index_input, angle_index_input],
outputs=main_output)
return model