def __init__(self,
units,
use_edge_features=False,
activation='kgcnn>leaky_relu',
use_bias=True,
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(AttentionHeadGAT, self).__init__(**kwargs)
# graph args
self.use_edge_features = use_edge_features
# dense args
self.units = int(units)
kernel_args = {"use_bias": use_bias, "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}
self.lay_linear_trafo = Dense(units, activation="linear", **kernel_args, **self._kgcnn_info)
self.lay_alpha = Dense(1, activation=activation, **kernel_args, **self._kgcnn_info)
self.lay_gather_in = GatherNodesIngoing(**self._kgcnn_info)
self.lay_gather_out = GatherNodesOutgoing(**self._kgcnn_info)
self.lay_concat = Concatenate(axis=-1, **self._kgcnn_info)
self.lay_pool_attention = PoolingLocalEdgesAttention(**self._kgcnn_info)
self.lay_final_activ = Activation(activation=activation, **self._kgcnn_info)
def __init__(self, units,
cfconv_pool='segment_sum',
use_bias=True,
activation='kgcnn>shifted_softplus',
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(SchNetCFconv, self).__init__(**kwargs)
self.cfconv_pool = cfconv_pool
self.units = units
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}
# Layer
self.lay_dense1 = Dense(units=self.units, activation=activation, use_bias=self.use_bias, **kernel_args,
**self._kgcnn_info)
self.lay_dense2 = Dense(units=self.units, activation='linear', use_bias=self.use_bias, **kernel_args,
**self._kgcnn_info)
self.lay_sum = PoolingLocalEdges(pooling_method=cfconv_pool, **self._kgcnn_info)
self.gather_n = GatherNodesOutgoing(**self._kgcnn_info)
self.lay_mult = Multiply(**self._kgcnn_info)
def __init__(self,
units,
pooling_method='sum',
normalize_by_weights=False,
activation='kgcnn>leaky_relu',
use_bias=True,
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(GCN, self).__init__(**kwargs)
self.normalize_by_weights = normalize_by_weights
self.pooling_method = pooling_method
self.units = units
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}
pool_args = {"pooling_method": pooling_method, "normalize_by_weights": normalize_by_weights}
# Layers
self.lay_gather = GatherNodesOutgoing(**self._kgcnn_info)
self.lay_dense = Dense(units=self.units, activation='linear',
input_tensor_type=self.input_tensor_type, ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_pool = PoolingWeightedLocalEdges(**pool_args, **self._kgcnn_info)
self.lay_act = Activation(activation, ragged_validate=self.ragged_validate,
input_tensor_type=self.input_tensor_type)
def __init__(self,
units,
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(ResidualLayer, self).__init__(**kwargs)
if activation is None and 'swish' in kgcnn_custom_act:
activation = 'swish'
elif activation is None:
activation = "selu"
dense_args = {
"units": units,
"activation": activation,
"use_bias": use_bias,
"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
}
self.dense_1 = Dense(**dense_args)
self.dense_2 = Dense(**dense_args)
self.add_end = Add()
def __init__(self,
units=128,
cfconv_pool='sum',
use_bias=True,
activation='kgcnn>shifted_softplus',
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(SchNetInteraction, self).__init__(**kwargs)
self.cfconv_pool = cfconv_pool
self.use_bias = use_bias
self.units = units
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}
conv_args = {"units": self.units, "use_bias": use_bias, "activation": activation, "cfconv_pool": cfconv_pool}
# Layers
self.lay_cfconv = SchNetCFconv(**conv_args, **kernel_args, **self._kgcnn_info)
self.lay_dense1 = Dense(units=self.units, activation='linear', use_bias=False,
**self._kgcnn_info, **kernel_args)
self.lay_dense2 = Dense(units=self.units, activation=activation, use_bias=self.use_bias,
**self._kgcnn_info, **kernel_args)
self.lay_dense3 = Dense(units=self.units, activation='linear', use_bias=self.use_bias,
**self._kgcnn_info, **kernel_args)
self.lay_add = Add(**self._kgcnn_info)
def __init__(self,
emb_size,
out_emb_size,
num_dense,
num_targets=12,
use_bias=True,
output_kernel_initializer="zeros",
kernel_initializer='kgcnn>glorot_orthogonal',
bias_initializer='zeros',
activation='kgcnn>swish',
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
pooling_method="sum",
**kwargs):
"""Initialize layer."""
super(DimNetOutputBlock, self).__init__(**kwargs)
self.pooling_method = pooling_method
self.emb_size = emb_size
self.out_emb_size = out_emb_size
self.num_dense = num_dense
self.num_targets = num_targets
self.use_bias = use_bias
kernel_args = {
"kernel_regularizer": kernel_regularizer,
"activity_regularizer": activity_regularizer,
"kernel_constraint": kernel_constraint,
"bias_initializer": bias_initializer,
"bias_regularizer": bias_regularizer,
"bias_constraint": bias_constraint,
}
self.dense_rbf = Dense(emb_size,
use_bias=False,
kernel_initializer=kernel_initializer,
**kernel_args,
**self._kgcnn_info)
self.up_projection = Dense(out_emb_size,
use_bias=False,
kernel_initializer=kernel_initializer,
**kernel_args,
**self._kgcnn_info)
self.dense_mlp = MLP([out_emb_size] * num_dense,
activation=activation,
kernel_initializer=kernel_initializer,
use_bias=use_bias,
**kernel_args,
**self._kgcnn_info)
self.dimnet_mult = Multiply(**self._kgcnn_info)
self.pool = PoolingLocalEdges(pooling_method=self.pooling_method,
**self._kgcnn_info)
self.dense_final = Dense(num_targets,
use_bias=False,
kernel_initializer=output_kernel_initializer,
**kernel_args,
**self._kgcnn_info)
def __init__(self,
units,
use_edge_features=False,
activation=None,
use_bias=True,
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(AttentionHeadGAT, self).__init__(**kwargs)
# graph args
self.use_edge_features = use_edge_features
# dense args
self.units = int(units)
if activation is None and "leaky_relu" in kgcnn_custom_act:
activation = {"class_name": "leaky_relu", "config": {"alpha": 0.2}}
elif activation is None:
activation = "relu"
kernel_args = {
"use_bias": use_bias,
"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
}
dens_args = {
"ragged_validate": self.ragged_validate,
"input_tensor_type": self.input_tensor_type
}
dens_args.update(kernel_args)
gather_args = self._all_kgcnn_info
pooling_args = self._all_kgcnn_info
self.lay_linear_trafo = Dense(units, activation="linear", **dens_args)
self.lay_alpha = Dense(1, activation=activation, **dens_args)
self.lay_gather_in = GatherNodesIngoing(**gather_args)
self.lay_gather_out = GatherNodesOutgoing(**gather_args)
self.lay_concat = Concatenate(axis=-1,
input_tensor_type=self.input_tensor_type)
self.lay_pool_attention = PoolingLocalEdgesAttention(**pooling_args)
self.lay_final_activ = Activation(
activation=activation, input_tensor_type=self.input_tensor_type)
def __init__(self,
units,
depth=3,
pooling_method="sum",
activation='kgcnn>leaky_relu',
activation_context="elu",
use_bias=True,
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
kernel_initializer='glorot_uniform',
bias_initializer='zeros',
recurrent_activation='sigmoid',
recurrent_initializer='orthogonal',
recurrent_regularizer=None,
recurrent_constraint=None,
dropout=0.0,
recurrent_dropout=0.0,
reset_after=True,
**kwargs):
"""Initialize layer."""
super(AttentiveNodePooling, self).__init__(**kwargs)
self.pooling_method = pooling_method
self.depth = depth
# dense args
self.units = int(units)
kernel_args = {"use_bias": use_bias, "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}
gru_args = {"recurrent_activation": recurrent_activation,
"use_bias": use_bias, "kernel_initializer": kernel_initializer,
"recurrent_initializer": recurrent_initializer, "bias_initializer": bias_initializer,
"kernel_regularizer": kernel_regularizer, "recurrent_regularizer": recurrent_regularizer,
"bias_regularizer": bias_regularizer, "kernel_constraint": kernel_constraint,
"recurrent_constraint": recurrent_constraint, "bias_constraint": bias_constraint,
"dropout": dropout, "recurrent_dropout": recurrent_dropout, "reset_after": reset_after}
self.lay_linear_trafo = Dense(units, activation="linear", **kernel_args, **self._kgcnn_info)
self.lay_alpha = Dense(1, activation=activation, **kernel_args, **self._kgcnn_info)
self.lay_gather_s = GatherState(**self._kgcnn_info)
self.lay_concat = Concatenate(axis=-1, **self._kgcnn_info)
self.lay_pool_start = PoolingNodes(pooling_method=self.pooling_method, **self._kgcnn_info)
self.lay_pool_attention = PoolingNodesAttention(**self._kgcnn_info)
self.lay_final_activ = Activation(activation=activation_context, **self._kgcnn_info)
self.lay_gru = tf.keras.layers.GRUCell(units=units, activation="tanh", **gru_args)
def __init__(self,
units,
pooling_method='sum',
normalize_by_weights=False,
activation=None,
use_bias=True,
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(GCN, self).__init__(**kwargs)
self.normalize_by_weights = normalize_by_weights
self.pooling_method = pooling_method
self.units = units
if activation is None and 'leaky_relu' in kgcnn_custom_act:
activation = {"class_name": "leaky_relu", "config": {"alpha": 0.2}}
elif activation is None:
activation = "relu"
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}
gather_args = self._all_kgcnn_info
pool_args = {"pooling_method": pooling_method, "normalize_by_weights": normalize_by_weights}
pool_args.update(self._all_kgcnn_info)
# Layers
self.lay_gather = GatherNodesOutgoing(**gather_args)
self.lay_dense = Dense(units=self.units, activation='linear',
input_tensor_type=self.input_tensor_type, ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_pool = PoolingWeightedLocalEdges(**pool_args)
self.lay_act = Activation(activation, ragged_validate=self.ragged_validate,
input_tensor_type=self.input_tensor_type)
def __init__(self, units,
cfconv_pool='segment_sum',
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(SchNetCFconv, self).__init__(**kwargs)
self.cfconv_pool = cfconv_pool
self.units = units
self.use_bias = use_bias
if activation is None and 'shifted_softplus' in kgcnn_custom_act:
activation = 'shifted_softplus'
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}
pooling_args = {"pooling_method": cfconv_pool}
pooling_args.update(self._all_kgcnn_info)
# Layer
self.lay_dense1 = Dense(units=self.units, activation=activation, use_bias=self.use_bias,
input_tensor_type=self.input_tensor_type, ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_dense2 = Dense(units=self.units, activation='linear', use_bias=self.use_bias,
input_tensor_type=self.input_tensor_type, ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_sum = PoolingLocalEdges(**pooling_args)
self.gather_n = GatherNodesOutgoing(**self._all_kgcnn_info)
self.lay_mult = Multiply(input_tensor_type=self.input_tensor_type, ragged_validate=self.ragged_validate)
def __init__(self, units,
use_bias=True,
activation='kgcnn>swish',
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(ResidualLayer, self).__init__(**kwargs)
dense_args = {"units": units, "activation": activation, "use_bias": use_bias,
"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}
self.dense_1 = Dense(**dense_args, **self._kgcnn_info)
self.dense_2 = Dense(**dense_args, **self._kgcnn_info)
self.add_end = Add(**self._kgcnn_info)
class MEGnetBlock(GraphBaseLayer):
"""Megnet Block.
Args:
node_embed (list, optional): List of node embedding dimension. Defaults to [16,16,16].
edge_embed (list, optional): List of edge embedding dimension. Defaults to [16,16,16].
env_embed (list, optional): List of environment embedding dimension. Defaults to [16,16,16].
pooling_method (str): Pooling method information for layer. Default is 'mean'.
use_bias (bool, optional): Use bias. Defaults to True.
activation (str): Activation function. Default is 'softplus2' with fall-back 'selu'.
kernel_regularizer: Kernel regularization. Default is None.
bias_regularizer: Bias regularization. Default is None.
activity_regularizer: Activity regularization. Default is None.
kernel_constraint: Kernel constrains. Default is None.
bias_constraint: Bias constrains. Default is None.
kernel_initializer: Initializer for kernels. Default is 'glorot_uniform'.
bias_initializer: Initializer for bias. Default is 'zeros'.
"""
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 build(self, input_shape):
"""Build layer."""
super(MEGnetBlock, self).build(input_shape)
def call(self, inputs, **kwargs):
"""Forward pass.
Args:
inputs: [nodes, edges, edge_index, state]
- nodes (tf.ragged): Node embeddings of shape (batch, [N], F)
- edges (tf.ragged): Edge or message embeddings of shape (batch, [M], F)
- edge_index (tf.ragged): Edge indices of shape (batch, [M], 2)
- state (tf.tensor): State information for the graph, a single tensor of shape (batch, F)
Returns:
node_update: Updated node embeddings of shape (batch, [N], F)
"""
# Calculate edge Update
node_input, edge_input, edge_index_input, env_input = inputs
e_n = self.lay_gather_n([node_input, edge_index_input])
e_u = self.lay_gather_ue([env_input, edge_input])
ec = self.lay_conc_enu([e_n, edge_input, e_u])
ep = self.lay_phi_e(ec) # Learning of Update Functions
ep = self.lay_phi_e_1(ep) # Learning of Update Functions
ep = self.lay_phi_e_2(ep) # Learning of Update Functions
# Calculate Node update
vb = self.lay_esum([node_input, ep, edge_index_input
]) # Summing for each node connections
v_u = self.lay_gather_un([env_input, node_input])
vc = self.lay_conc_nu(
[vb, node_input,
v_u]) # Concatenate node features with new edge updates
vp = self.lay_phi_n(vc) # Learning of Update Functions
vp = self.lay_phi_n_1(vp) # Learning of Update Functions
vp = self.lay_phi_n_2(vp) # Learning of Update Functions
# Calculate environment update
es = self.lay_usum_e(ep)
vs = self.lay_usum_n(vp)
ub = self.lay_conc_u([es, vs, env_input])
up = self.lay_phi_u(ub)
up = self.lay_phi_u_1(up)
up = self.lay_phi_u_2(up) # Learning of Update Functions
return vp, ep, up
def get_config(self):
config = super(MEGnetBlock, self).get_config()
config.update({
"pooling_method": self.pooling_method,
"node_embed": self.node_embed,
"use_bias": self.use_bias,
"edge_embed": self.edge_embed,
"env_embed": self.env_embed
})
config_dense = self.lay_phi_n.get_config()
for x in [
"kernel_regularizer", "activity_regularizer",
"bias_regularizer", "kernel_constraint", "bias_constraint",
"kernel_initializer", "bias_initializer", "activation"
]:
config.update({x: config_dense[x]})
return config
def __init__(self,
units=128,
cfconv_pool='sum',
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(SchNetInteraction, self).__init__(**kwargs)
self.cfconv_pool = cfconv_pool
self.use_bias = use_bias
self.units = units
if activation is None and 'shifted_softplus' in kgcnn_custom_act:
activation = 'shifted_softplus'
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
}
conv_args = {
"units": self.units,
"use_bias": use_bias,
"activation": activation,
"cfconv_pool": cfconv_pool
}
conv_args.update(kernel_args)
conv_args.update(self._all_kgcnn_info)
# Layers
self.lay_cfconv = SchNetCFconv(**conv_args)
self.lay_dense1 = Dense(units=self.units,
activation='linear',
use_bias=False,
input_tensor_type=self.input_tensor_type,
ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_dense2 = Dense(units=self.units,
activation=activation,
use_bias=self.use_bias,
input_tensor_type=self.input_tensor_type,
ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_dense3 = Dense(units=self.units,
activation='linear',
use_bias=self.use_bias,
input_tensor_type=self.input_tensor_type,
ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_add = Add(input_tensor_type=self.input_tensor_type,
ragged_validate=self.ragged_validate)
def __init__(self,
emb_size,
int_emb_size,
basis_emb_size,
num_before_skip,
num_after_skip,
use_bias=True,
pooling_method="sum",
activation=None,
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
kernel_initializer='orthogonal',
bias_initializer='zeros',
**kwargs):
super(DimNetInteractionPPBlock, self).__init__(**kwargs)
self.use_bias = use_bias
self.pooling_method = pooling_method
self.emb_size = emb_size
self.int_emb_size = int_emb_size
self.basis_emb_size = basis_emb_size
self.num_before_skip = num_before_skip
self.num_after_skip = num_after_skip
if activation is None and 'swish' in kgcnn_custom_act:
activation = 'swish'
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
}
pool_args = {"pooling_method": pooling_method}
pool_args.update(self._all_kgcnn_info)
gather_args = self._all_kgcnn_info
# Transformations of Bessel and spherical basis representations
self.dense_rbf1 = Dense(basis_emb_size, use_bias=False, **kernel_args)
self.dense_rbf2 = Dense(emb_size, use_bias=False, **kernel_args)
self.dense_sbf1 = Dense(basis_emb_size, use_bias=False, **kernel_args)
self.dense_sbf2 = Dense(int_emb_size, use_bias=False, **kernel_args)
# Dense transformations of input messages
self.dense_ji = Dense(emb_size,
activation=activation,
use_bias=True,
**kernel_args)
self.dense_kj = Dense(emb_size,
activation=activation,
use_bias=True,
**kernel_args)
# Embedding projections for interaction triplets
self.down_projection = Dense(int_emb_size,
activation=activation,
use_bias=False,
**kernel_args)
self.up_projection = Dense(emb_size,
activation=activation,
use_bias=False,
**kernel_args)
# Residual layers before skip connection
self.layers_before_skip = []
for i in range(num_before_skip):
self.layers_before_skip.append(
ResidualLayer(emb_size,
activation=activation,
use_bias=True,
**kernel_args))
self.final_before_skip = Dense(emb_size,
activation=activation,
use_bias=True,
**kernel_args)
# Residual layers after skip connection
self.layers_after_skip = []
for i in range(num_after_skip):
self.layers_after_skip.append(
ResidualLayer(emb_size,
activation=activation,
use_bias=True,
**kernel_args))
self.lay_add1 = Add()
self.lay_add2 = Add()
self.lay_mult1 = Multiply()
self.lay_mult2 = Multiply()
self.lay_gather = GatherNodesOutgoing(**gather_args) # Are edges here
self.lay_pool = PoolingLocalEdges(**pool_args)
class DimNetInteractionPPBlock(GraphBaseLayer):
"""DimNetInteractionPPBlock as defined by DimNet.
Args:
emb_size: Embedding size used for the messages
int_emb_size (int): Embedding size used for interaction triplets
basis_emb_size: Embedding size used inside the basis transformation
num_before_skip: Number of residual layers in interaction block before skip connection
num_after_skip: Number of residual layers in interaction block before skip connection
use_bias (bool, optional): Use bias. Defaults to True.
pooling_method (str): Pooling method information for layer. Default is 'sum'.
activation (str): Activation function. Default is "swish".
kernel_regularizer: Kernel regularization. Default is None.
bias_regularizer: Bias regularization. Default is None.
activity_regularizer: Activity regularization. Default is None.
kernel_constraint: Kernel constrains. Default is None.
bias_constraint: Bias constrains. Default is None.
kernel_initializer: Initializer for kernels. Default is 'orthogonal'.
bias_initializer: Initializer for bias. Default is 'zeros'.
**kwargs:
"""
def __init__(self,
emb_size,
int_emb_size,
basis_emb_size,
num_before_skip,
num_after_skip,
use_bias=True,
pooling_method="sum",
activation=None,
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
kernel_initializer='orthogonal',
bias_initializer='zeros',
**kwargs):
super(DimNetInteractionPPBlock, self).__init__(**kwargs)
self.use_bias = use_bias
self.pooling_method = pooling_method
self.emb_size = emb_size
self.int_emb_size = int_emb_size
self.basis_emb_size = basis_emb_size
self.num_before_skip = num_before_skip
self.num_after_skip = num_after_skip
if activation is None and 'swish' in kgcnn_custom_act:
activation = 'swish'
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
}
pool_args = {"pooling_method": pooling_method}
pool_args.update(self._all_kgcnn_info)
gather_args = self._all_kgcnn_info
# Transformations of Bessel and spherical basis representations
self.dense_rbf1 = Dense(basis_emb_size, use_bias=False, **kernel_args)
self.dense_rbf2 = Dense(emb_size, use_bias=False, **kernel_args)
self.dense_sbf1 = Dense(basis_emb_size, use_bias=False, **kernel_args)
self.dense_sbf2 = Dense(int_emb_size, use_bias=False, **kernel_args)
# Dense transformations of input messages
self.dense_ji = Dense(emb_size,
activation=activation,
use_bias=True,
**kernel_args)
self.dense_kj = Dense(emb_size,
activation=activation,
use_bias=True,
**kernel_args)
# Embedding projections for interaction triplets
self.down_projection = Dense(int_emb_size,
activation=activation,
use_bias=False,
**kernel_args)
self.up_projection = Dense(emb_size,
activation=activation,
use_bias=False,
**kernel_args)
# Residual layers before skip connection
self.layers_before_skip = []
for i in range(num_before_skip):
self.layers_before_skip.append(
ResidualLayer(emb_size,
activation=activation,
use_bias=True,
**kernel_args))
self.final_before_skip = Dense(emb_size,
activation=activation,
use_bias=True,
**kernel_args)
# Residual layers after skip connection
self.layers_after_skip = []
for i in range(num_after_skip):
self.layers_after_skip.append(
ResidualLayer(emb_size,
activation=activation,
use_bias=True,
**kernel_args))
self.lay_add1 = Add()
self.lay_add2 = Add()
self.lay_mult1 = Multiply()
self.lay_mult2 = Multiply()
self.lay_gather = GatherNodesOutgoing(**gather_args) # Are edges here
self.lay_pool = PoolingLocalEdges(**pool_args)
def call(self, inputs, **kwargs):
"""Forward pass.
Args:
inputs: [edges, rbf, sbf, angle_index]
- edges (tf.ragged): Edge embeddings of shape (batch, [M], F)
- rbf (tf.ragged): Radial basis features of shape (batch, [M], F)
- sbf (tf.ragged): Spherical basis features of shape (batch, [K], F)
- edge_index (tf.ragged): Angle indices between two edges of shape (batch, [K], 2)
Returns:
tf.ragged: Updated edge embeddings.
"""
x, rbf, sbf, id_expand = inputs
# Initial transformation
x_ji = self.dense_ji(x)
x_kj = self.dense_kj(x)
# Transform via Bessel basis
rbf = self.dense_rbf1(rbf)
rbf = self.dense_rbf2(rbf)
x_kj = self.lay_mult1([x_kj, rbf])
# Down-project embeddings and generate interaction triplet embeddings
x_kj = self.down_projection(x_kj)
x_kj = self.lay_gather([x_kj, id_expand])
# Transform via 2D spherical basis
sbf = self.dense_sbf1(sbf)
sbf = self.dense_sbf2(sbf)
x_kj = self.lay_mult1([x_kj, sbf])
# Aggregate interactions and up-project embeddings
x_kj = self.lay_pool([rbf, x_kj, id_expand])
x_kj = self.up_projection(x_kj)
# Transformations before skip connection
x2 = self.lay_add1([x_ji, x_kj])
for layer in self.layers_before_skip:
x2 = layer(x2)
x2 = self.final_before_skip(x2)
# Skip connection
x = self.lay_add2([x, x2])
# Transformations after skip connection
for layer in self.layers_after_skip:
x = layer(x)
return x
def get_config(self):
config = super(DimNetInteractionPPBlock, self).get_config()
config.update({
"use_bias": self.use_bias,
"pooling_method": self.pooling_method,
"emb_size": self.emb_size,
"int_emb_size": self.int_emb_size,
"basis_emb_size": self.basis_emb_size,
"num_before_skip": self.num_before_skip,
"num_after_skip": self.num_after_skip
})
conf_dense = self.dense_ji.get_config()
for x in [
"kernel_regularizer", "activity_regularizer",
"bias_regularizer", "kernel_constraint", "bias_constraint",
"kernel_initializer", "bias_initializer", "activation"
]:
config.update({x: conf_dense[x]})
return config
class GCN(GraphBaseLayer):
r"""Graph convolution according to Kipf et al.
Computes graph conv as $\sigma(A_s*(WX+b))$ where $A_s$ is the precomputed and scaled adjacency matrix.
The scaled adjacency matrix is defined by $A_s = D^{-0.5} (A + I) D{^-0.5}$ with the degree matrix $D$.
In place of $A_s$, this layers uses edge features (that are the entries of $A_s$) and edge indices.
$A_s$ is considered pre-scaled, this is not done by this layer.
If no scaled edge features are available, you could consider use e.g. "segment_mean", or normalize_by_weights to
obtain a similar behaviour that is expected by a pre-scaled adjacency matrix input.
Edge features must be possible to broadcast to node features. Ideally they have shape (...,1).
Args:
units (int): Output dimension/ units of dense layer.
pooling_method (str): Pooling method for summing edges. Default is 'segment_sum'.
normalize_by_weights (bool): Normalize the pooled output by the sum of weights. Default is False.
In this case the edge features are considered weights of dimension (...,1) and are summed for each node.
activation (str): Activation. Default is {"class_name": "leaky_relu", "config": {"alpha": 0.2}},
with fall-back "relu".
use_bias (bool): Use bias. Default is True.
kernel_regularizer: Kernel regularization. Default is None.
bias_regularizer: Bias regularization. Default is None.
activity_regularizer: Activity regularization. Default is None.
kernel_constraint: Kernel constrains. Default is None.
bias_constraint: Bias constrains. Default is None.
kernel_initializer: Initializer for kernels. Default is 'glorot_uniform'.
bias_initializer: Initializer for bias. Default is 'zeros'.
"""
def __init__(self,
units,
pooling_method='sum',
normalize_by_weights=False,
activation=None,
use_bias=True,
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(GCN, self).__init__(**kwargs)
self.normalize_by_weights = normalize_by_weights
self.pooling_method = pooling_method
self.units = units
if activation is None and 'leaky_relu' in kgcnn_custom_act:
activation = {"class_name": "leaky_relu", "config": {"alpha": 0.2}}
elif activation is None:
activation = "relu"
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}
gather_args = self._all_kgcnn_info
pool_args = {"pooling_method": pooling_method, "normalize_by_weights": normalize_by_weights}
pool_args.update(self._all_kgcnn_info)
# Layers
self.lay_gather = GatherNodesOutgoing(**gather_args)
self.lay_dense = Dense(units=self.units, activation='linear',
input_tensor_type=self.input_tensor_type, ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_pool = PoolingWeightedLocalEdges(**pool_args)
self.lay_act = Activation(activation, ragged_validate=self.ragged_validate,
input_tensor_type=self.input_tensor_type)
def build(self, input_shape):
"""Build layer."""
super(GCN, self).build(input_shape)
def call(self, inputs, **kwargs):
"""Forward pass.
Args:
inputs: [nodes, edges, edge_index]
- nodes (tf.ragged): Node embeddings of shape (batch, [N], F)
- edges (tf.ragged): Edge or message embeddings of shape (batch, [M], F)
- edge_index (tf.ragged): Edge indices of shape (batch, [M], 2)
Returns:
embeddings: Node embeddings of shape (batch, [N], F)
"""
node, edges, edge_index = inputs
no = self.lay_dense(node)
no = self.lay_gather([no, edge_index])
nu = self.lay_pool([node, no, edge_index, edges]) # Summing for each node connection
out = self.lay_act(nu)
return out
def get_config(self):
"""Update config."""
config = super(GCN, self).get_config()
config.update({"normalize_by_weights": self.normalize_by_weights,
"pooling_method": self.pooling_method, "units": self.units})
conf_dense = self.lay_dense.get_config()
for x in ["kernel_regularizer", "activity_regularizer", "bias_regularizer", "kernel_constraint",
"bias_constraint", "kernel_initializer", "bias_initializer", "use_bias"]:
config.update({x: conf_dense[x]})
conf_act = self.lay_act.get_config()
config.update({"activation": conf_act["activation"]})
return config
class ResidualLayer(GraphBaseLayer):
"""Residual Layer as defined by DimNet.
Args:
units: Dimension of the kernel.
use_bias (bool, optional): Use bias. Defaults to True.
activation (str): Activation function. Default is "swish".
kernel_regularizer: Kernel regularization. Default is None.
bias_regularizer: Bias regularization. Default is None.
activity_regularizer: Activity regularization. Default is None.
kernel_constraint: Kernel constrains. Default is None.
bias_constraint: Bias constrains. Default is None.
kernel_initializer: Initializer for kernels. Default is 'glorot_uniform'.
bias_initializer: Initializer for bias. Default is 'zeros'.
**kwargs:
"""
def __init__(self,
units,
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(ResidualLayer, self).__init__(**kwargs)
if activation is None and 'swish' in kgcnn_custom_act:
activation = 'swish'
elif activation is None:
activation = "selu"
dense_args = {
"units": units,
"activation": activation,
"use_bias": use_bias,
"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
}
self.dense_1 = Dense(**dense_args)
self.dense_2 = Dense(**dense_args)
self.add_end = Add()
def build(self, input_shape):
"""Build layer."""
super(ResidualLayer, self).build(input_shape)
def call(self, inputs, **kwargs):
"""Forward pass
Args:
inputs (tf.ragged): Node or edge embedding of shape (batch, [N], F)
Returns:
embeddings: Node or edge embedding of shape (batch, [N], F)
"""
x = self.dense_1(inputs)
x = self.dense_2(x)
x = self.add_end([inputs, x])
return x
def get_config(self):
config = super(ResidualLayer, self).get_config()
conf_dense = self.dense_1.get_config()
for x in [
"kernel_regularizer", "activity_regularizer",
"bias_regularizer", "kernel_constraint", "bias_constraint",
"kernel_initializer", "bias_initializer", "activation",
"use_bias", "units"
]:
config.update({x: conf_dense[x]})
return config
def make_unet(
# Input
input_node_shape,
input_edge_shape,
input_embedd: dict = None,
# Output
output_embedd: dict = None,
output_mlp: dict = None,
# Model specific
hidden_dim=32,
depth=4,
k=0.3,
score_initializer='ones',
use_bias=True,
activation='relu',
is_sorted=False,
has_unconnected=True,
use_reconnect=True
):
"""
Make Graph U Net.
Args:
input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used.
input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used.
input_embedd (list): Dictionary of embedding parameters used if input shape is None. Default is
{'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100,
'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64,
'input_type': 'ragged'}
output_mlp (dict, optional): Parameter for MLP output classification/ regression. Defaults to
{"use_bias": [True, False], "output_dim": [25, 1],
"activation": ['relu', 'sigmoid']}
output_embedd (str): Dictionary of embedding parameters of the graph network. Default is
{"output_mode": 'graph', "output_type": 'padded'}
hidden_dim (int): Hidden node feature dimension 32,
depth (int): Depth of pooling steps. Default is 4.
k (float): Pooling ratio. Default is 0.3.
score_initializer (str): How to initialize score kernel. Default is 'ones'.
use_bias (bool): Use bias. Default is True.
activation (str): Activation function used. Default is 'relu'.
is_sorted (bool, optional): Edge edge_indices are sorted. Defaults to False.
has_unconnected (bool, optional): Has unconnected nodes. Defaults to True.
use_reconnect (bool): Reconnect nodes after pooling. I.e. adj_matrix=adj_matrix^2. Default is True.
Returns:
model (ks.models.Model): Unet model.
"""
# Default values update
model_default = {'input_embedd': {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100,
'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64,
'input_tensor_type': 'ragged'},
'output_embedd': {"output_mode": 'graph', "output_type": 'padded'},
'output_mlp': {"use_bias": [True, False], "units": [25, 1], "activation": ['relu', 'sigmoid']}
}
# Update model args
input_embedd = update_model_args(model_default['input_embedd'], input_embedd)
output_embedd = update_model_args(model_default['output_embedd'], output_embedd)
output_mlp = update_model_args(model_default['output_mlp'], output_mlp)
pooling_args = {"pooling_method": 'segment_mean', "is_sorted": is_sorted, "has_unconnected": has_unconnected}
# Make input embedding, if no feature dimension
node_input, n, edge_input, ed, edge_index_input, _, _ = generate_standard_graph_input(input_node_shape,
input_edge_shape, None,
**input_embedd)
tens_type = "values_partition"
node_indexing = "batch"
n = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(n)
ed = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(ed)
edi = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(edge_index_input)
edi = ChangeIndexing(input_tensor_type=tens_type, to_indexing=node_indexing)([n, edi]) # disjoint
output_mlp.update({"input_tensor_type": tens_type})
gather_args = {"input_tensor_type": tens_type, "node_indexing": node_indexing}
pooling_args.update({"input_tensor_type": tens_type, "node_indexing": node_indexing})
# Graph lists
n = Dense(hidden_dim, use_bias=use_bias, activation='linear', input_tensor_type=tens_type)(n)
in_graph = [n, ed, edi]
graph_list = [in_graph]
map_list = []
# U Down
i_graph = in_graph
for i in range(0, depth):
n, ed, edi = i_graph
# GCN layer
eu = GatherNodesOutgoing(**gather_args)([n, edi])
eu = Dense(hidden_dim, use_bias=use_bias, activation='linear', input_tensor_type=tens_type)(eu)
nu = PoolingLocalEdges(**pooling_args)([n, eu, edi]) # Summing for each node connection
n = Activation(activation=activation, input_tensor_type=tens_type)(nu)
if use_reconnect:
ed, edi = AdjacencyPower(n=2, node_indexing=node_indexing, input_tensor_type=tens_type)([n, ed, edi])
# Pooling
i_graph, i_map = PoolingTopK(k=k, kernel_initializer=score_initializer,
node_indexing=node_indexing, input_tensor_type=tens_type)([n, ed, edi])
graph_list.append(i_graph)
map_list.append(i_map)
# U Up
ui_graph = i_graph
for i in range(depth, 0, -1):
o_graph = graph_list[i - 1]
i_map = map_list[i - 1]
ui_graph = UnPoolingTopK(node_indexing=node_indexing, input_tensor_type=tens_type)(o_graph + i_map + ui_graph)
n, ed, edi = ui_graph
# skip connection
n = Add(input_tensor_type=tens_type)([n, o_graph[0]])
# GCN
eu = GatherNodesOutgoing(**gather_args)([n, edi])
eu = Dense(hidden_dim, use_bias=use_bias, activation='linear', input_tensor_type=tens_type)(eu)
nu = PoolingLocalEdges(**pooling_args)([n, eu, edi]) # Summing for each node connection
n = Activation(activation=activation, input_tensor_type=tens_type)(nu)
ui_graph = [n, ed, edi]
# Otuput
n = ui_graph[0]
if output_embedd["output_mode"] == 'graph':
out = PoolingNodes(**pooling_args)(n)
output_mlp.update({"input_tensor_type": "tensor"})
out = MLP(**output_mlp)(out)
main_output = ks.layers.Flatten()(out) # will be dense
else: # node embedding
out = MLP(**output_mlp)(n)
main_output = ChangeTensorType(input_tensor_type=tens_type, output_tensor_type="tensor")(out)
model = ks.models.Model(inputs=[node_input, edge_input, edge_index_input], outputs=main_output)
return model
class AttentiveHeadFP(GraphBaseLayer):
r"""Computes the attention head for Attentive FP model.
The attention coefficients are computed by $a_{ij} = \sigma_1( W_1 [h_i || h_j] )$.
The initial representation $h_i$ and $h_j$ must be calculated beforehand.
The attention is obtained by $\alpha_ij = softmax_j (a_{ij})$.
And finally pooled through for context $C_i = \sigma_2(\sum_j \alpha_{ij} W_2 h_j)$.
If graphs indices were in 'batch' mode, the layer's 'node_indexing' must be set to 'batch'.
Args:
units (int): Units for the linear trafo of node features before attention.
use_edge_features (bool): Append edge features to attention computation. Default is False.
activation (str): Activation. Default is {"class_name": "kgcnn>leaky_relu", "config": {"alpha": 0.2}}.
activation_context (str): Activation function for context. Default is "elu".
use_bias (bool): Use bias. Default is True.
kernel_regularizer: Kernel regularization. Default is None.
bias_regularizer: Bias regularization. Default is None.
activity_regularizer: Activity regularization. Default is None.
kernel_constraint: Kernel constrains. Default is None.
bias_constraint: Bias constrains. Default is None.
kernel_initializer: Initializer for kernels. Default is 'glorot_uniform'.
bias_initializer: Initializer for bias. Default is 'zeros'.
"""
def __init__(self,
units,
use_edge_features=False,
activation='kgcnn>leaky_relu',
activation_context="elu",
use_bias=True,
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(AttentiveHeadFP, self).__init__(**kwargs)
# graph args
self.use_edge_features = use_edge_features
# dense args
self.units = int(units)
kernel_args = {"use_bias": use_bias, "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}
self.lay_linear_trafo = Dense(units, activation="linear", **kernel_args, **self._kgcnn_info)
self.lay_alpha = Dense(1, activation=activation, **kernel_args, **self._kgcnn_info)
self.lay_gather_in = GatherNodesIngoing(**self._kgcnn_info)
self.lay_gather_out = GatherNodesOutgoing(**self._kgcnn_info)
self.lay_concat = Concatenate(axis=-1, **self._kgcnn_info)
self.lay_pool_attention = PoolingLocalEdgesAttention(**self._kgcnn_info)
self.lay_final_activ = Activation(activation=activation_context, **self._kgcnn_info)
if use_edge_features:
self.lay_fc1 = Dense(units, activation=activation, **kernel_args, **self._kgcnn_info)
self.lay_fc2 = Dense(units, activation=activation, **kernel_args, **self._kgcnn_info)
self.lay_concat_edge = Concatenate(axis=-1, **self._kgcnn_info)
def build(self, input_shape):
"""Build layer."""
super(AttentiveHeadFP, self).build(input_shape)
def call(self, inputs, **kwargs):
"""Forward pass.
Args:
inputs (list): of [node, edges, edge_indices]
- nodes: Node features of shape (batch, [N], F)
- edges: Edge or message features of shape (batch, [M], F)
- edge_index: Edge indices of shape (batch, [M], 2)
Returns:
features: Hidden tensor of pooled edge attentions for each node.
"""
node, edge, edge_index = inputs
if self.use_edge_features:
n_in = self.lay_gather_in([node, edge_index])
n_out = self.lay_gather_out([node, edge_index])
n_in = self.lay_fc1(n_in)
n_out = self.lay_concat_edge([n_out, edge])
n_out = self.self.lay_fc2(n_out)
else:
n_in = self.lay_gather_in([node, edge_index])
n_out = self.lay_gather_out([node, edge_index])
wn_out = self.lay_linear_trafo(n_out)
e_ij = self.lay_concat([n_in, n_out])
a_ij = self.lay_alpha(e_ij) # Should be dimension (batch,None,1)
n_i = self.lay_pool_attention([node, wn_out, a_ij, edge_index])
out = self.lay_final_activ(n_i)
return out
def get_config(self):
"""Update layer config."""
config = super(AttentiveHeadFP, self).get_config()
config.update({"use_edge_features": self.use_edge_features,
"units": self.units})
conf_sub = self.lay_alpha.get_config()
for x in ["kernel_regularizer", "activity_regularizer", "bias_regularizer", "kernel_constraint",
"bias_constraint", "kernel_initializer", "bias_initializer", "activation", "use_bias"]:
config.update({x: conf_sub[x]})
conf_context = self.lay_final_activ.get_config()
config.update({"activation_context": conf_context["activation"]})
return config
def __init__(self,
units,
activation=None,
use_bias=True,
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
kernel_initializer='glorot_uniform',
bias_initializer='zeros',
use_edge_features=False,
node_indexing="sample",
is_sorted=False,
has_unconnected=True,
partition_type="row_length",
input_tensor_type="ragged",
ragged_validate=False,
**kwargs):
"""Initialize layer."""
super(AttentionHeadGAT, self).__init__(**kwargs)
# graph args
self.is_sorted = is_sorted
self.use_edge_features = use_edge_features
self.has_unconnected = has_unconnected
self.node_indexing = node_indexing
self.partition_type = partition_type
self.input_tensor_type = input_tensor_type
self.ragged_validate = ragged_validate
self._supports_ragged_inputs = True
self._tensor_input_type_implemented = [
"ragged", "values_partition", "disjoint", "tensor", "RaggedTensor"
]
self._test_tensor_input = kgcnn_ops_static_test_tensor_input_type(
self.input_tensor_type, self._tensor_input_type_implemented,
self.node_indexing)
# dense args
self.units = int(units)
if activation is None and "leaky_relu" in kgcnn_custom_act:
activation = {"class_name": "leaky_relu", "config": {"alpha": 0.2}}
elif activation is None:
activation = "relu"
self.use_bias = use_bias
self.ath_activation = tf.keras.activations.get(activation)
self.ath_kernel_initializer = tf.keras.initializers.get(
kernel_initializer)
self.ath_bias_initializer = tf.keras.initializers.get(bias_initializer)
self.ath_kernel_regularizer = tf.keras.regularizers.get(
kernel_regularizer)
self.ath_bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
self.ath_activity_regularizer = tf.keras.regularizers.get(
activity_regularizer)
self.ath_kernel_constraint = tf.keras.constraints.get(
kernel_constraint)
self.ath_bias_constraint = tf.keras.constraints.get(bias_constraint)
kernel_args = {
"use_bias": use_bias,
"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
}
dens_args = {
"ragged_validate": self.ragged_validate,
"input_tensor_type": self.input_tensor_type
}
dens_args.update(kernel_args)
gather_args = {
"input_tensor_type": self.input_tensor_type,
"node_indexing": self.node_indexing
}
pooling_args = {
"node_indexing": node_indexing,
"partition_type": partition_type,
"has_unconnected": has_unconnected,
"is_sorted": is_sorted,
"ragged_validate": self.ragged_validate,
"input_tensor_type": self.input_tensor_type
}
self.lay_linear_trafo = Dense(units, activation="linear", **dens_args)
self.lay_alpha = Dense(1, activation=activation, **dens_args)
self.lay_gather_in = GatherNodesIngoing(**gather_args)
self.lay_gather_out = GatherNodesOutgoing(**gather_args)
self.lay_concat = Concatenate(axis=-1,
input_tensor_type=self.input_tensor_type)
self.lay_pool_attention = PoolingLocalEdgesAttention(**pooling_args)
self.lay_final_activ = Activation(
activation=activation, input_tensor_type=self.input_tensor_type)
def make_gat( # Input
input_node_shape,
input_edge_shape,
input_embedd: dict = None,
# Output
output_embedd: dict = None,
output_mlp: dict = None,
# Model specific parameter
depth=3,
attention_heads_num=5,
attention_heads_concat=False,
attention_args: dict = None):
"""
Generate Interaction network.
Args:
input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used.
input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used.
input_embedd (dict): Dictionary of embedding parameters used if input shape is None. Default is
{'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100,
'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64,
'input_type': 'ragged'}.
output_embedd (dict): Dictionary of embedding parameters of the graph network. Default is
{"output_mode": 'graph', "output_type": 'padded'}.
output_mlp (dict): Dictionary of arguments for final MLP regression or classifcation layer. Default is
{"use_bias": [True, True, False], "units": [25, 10, 1],
"activation": ['relu', 'relu', 'sigmoid']}.
depth (int): Number of convolution layers. Default is 3.
attention_heads_num (int): Number of attention heads. Default is 5.
attention_heads_concat (bool): Concat attention. Default is False.
attention_args (dict): Layer arguments for attention layer. Default is
{"units": 32, 'is_sorted': False, 'has_unconnected': True}
Returns:
model (tf.keras.model): Interaction model.
"""
# default values
model_default = {
'input_embedd': {
'input_node_vocab': 95,
'input_edge_vocab': 5,
'input_state_vocab': 100,
'input_node_embedd': 64,
'input_edge_embedd': 64,
'input_state_embedd': 64,
'input_tensor_type': 'ragged'
},
'output_embedd': {
"output_mode": 'graph',
"output_tensor_type": 'padded'
},
'output_mlp': {
"use_bias": [True, True, False],
"units": [25, 10, 1],
"activation": ['relu', 'relu', 'sigmoid']
},
'attention_args': {
"units": 32,
'is_sorted': False,
'has_unconnected': True
}
}
# Update default values
input_embedd = update_model_args(model_default['input_embedd'],
input_embedd)
output_embedd = update_model_args(model_default['output_embedd'],
output_embedd)
output_mlp = update_model_args(model_default['output_mlp'], output_mlp)
attention_args = update_model_args(model_default['attention_args'],
attention_args)
pooling_nodes_args = {}
# Make input embedding, if no feature dimension
node_input, n, edge_input, ed, edge_index_input, _, _ = generate_standard_graph_input(
input_node_shape, input_edge_shape, None, **input_embedd)
edi = edge_index_input
nk = Dense(units=attention_args["units"], activation="linear")(n)
for i in range(0, depth):
heads = [
AttentionHeadGAT(**attention_args)([nk, ed, edi])
for _ in range(attention_heads_num)
]
if attention_heads_concat:
nk = Concatenate(axis=-1)(heads)
else:
nk = Average()(heads)
n = nk
if output_embedd["output_mode"] == 'graph':
out = PoolingNodes(**pooling_nodes_args)(n)
output_mlp.update({"input_tensor_type": "tensor"})
out = MLP(**output_mlp)(out)
main_output = ks.layers.Flatten()(out) # will be dense
else: # node embedding
out = MLP(**output_mlp)(n)
main_output = ChangeTensorType(input_tensor_type="ragged",
output_tensor_type="tensor")(out)
model = tf.keras.models.Model(
inputs=[node_input, edge_input, edge_index_input], outputs=main_output)
return model
def make_megnet(
# Input
input_node_shape,
input_edge_shape,
input_state_shape,
input_embedd: dict = None,
# Output
output_embedd: dict = None, # Only graph possible for megnet
output_mlp: dict = None,
# Model specs
meg_block_args: dict = None,
node_ff_args: dict = None,
edge_ff_args: dict = None,
state_ff_args: dict = None,
set2set_args: dict = None,
nblocks: int = 3,
has_ff: bool = True,
dropout: float = None,
use_set2set: bool = True,
):
"""
Get Megnet model.
Args:
input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used.
input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used.
input_state_shape (list): Shape of state features. If shape is (,) embedding layer is used.
input_embedd (dict): Dictionary of embedding parameters used if input shape is None. Default is
{'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100,
'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64,
'input_type': 'ragged'}.
output_embedd (str): Dictionary of embedding parameters of the graph network. Default is
{"output_mode": 'graph', "output_type": 'padded'}
output_mlp (dict): Dictionary of MLP arguments for output regression or classifcation. Default is
{"use_bias": [True, True, True], "units": [32, 16, 1],
"activation": ['softplus2', 'softplus2', 'linear']}.
meg_block_args (dict): Dictionary of MegBlock arguments. Default is
{'node_embed': [64, 32, 32], 'edge_embed': [64, 32, 32],
'env_embed': [64, 32, 32], 'activation': 'softplus2', 'is_sorted': False,
'has_unconnected': True}.
node_ff_args (dict): Dictionary of Feed-Forward Layer arguments. Default is
{"units": [64, 32], "activation": "softplus2"}.
edge_ff_args (dict): Dictionary of Feed-Forward Layer arguments. Default is
{"units": [64, 32], "activation": "softplus2"}.
state_ff_args (dict): Dictionary of Feed-Forward Layer arguments. Default is
{"units": [64, 32], "activation": "softplus2"}.
set2set_args (dict): Dictionary of Set2Set Layer Arguments. Default is
{'channels': 16, 'T': 3, "pooling_method": "sum", "init_qstar": "0"}
nblocks (int): Number of block. Default is 3.
has_ff (bool): Use a Feed-Forward layer. Default is True.
dropout (float): Use dropout. Default is None.
use_set2set (bool): Use set2set. Default is True.
Returns:
model (tf.keras.models.Model): MEGnet model.
"""
# Default arguments if None
model_default = {'input_embedd': {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100,
'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64,
'input_tensor_type': 'ragged'},
'output_embedd': {"output_mode": 'graph', "output_tensor_type": 'padded'},
'output_mlp': {"use_bias": [True, True, True], "units": [32, 16, 1],
"activation": ['kgcnn>softplus2', 'kgcnn>softplus2', 'linear']},
'meg_block_args': {'node_embed': [64, 32, 32], 'edge_embed': [64, 32, 32],
'env_embed': [64, 32, 32], 'activation': 'kgcnn>softplus2', 'is_sorted': False,
'has_unconnected': True},
'set2set_args': {'channels': 16, 'T': 3, "pooling_method": "sum", "init_qstar": "0"},
'node_ff_args': {"units": [64, 32], "activation": "kgcnn>softplus2"},
'edge_ff_args': {"units": [64, 32], "activation": "kgcnn>softplus2"},
'state_ff_args': {"units": [64, 32], "activation": "kgcnn>softplus2"}
}
# Update default arguments
input_embedd = update_model_args(model_default['input_embedd'], input_embedd)
output_embedd = update_model_args(model_default['output_embedd'], output_embedd)
output_mlp = update_model_args(model_default['output_mlp'], output_mlp)
meg_block_args = update_model_args(model_default['meg_block_args'], meg_block_args)
set2set_args = update_model_args(model_default['set2set_args'], set2set_args)
node_ff_args = update_model_args(model_default['node_ff_args'], node_ff_args)
edge_ff_args = update_model_args(model_default['edge_ff_args'], edge_ff_args)
state_ff_args = update_model_args(model_default['state_ff_args'], state_ff_args)
state_ff_args.update({"input_tensor_type": "tensor"})
# Make input embedding, if no feature dimension
node_input, n, edge_input, ed, edge_index_input, env_input, uenv = generate_standard_graph_input(input_node_shape,
input_edge_shape,
input_state_shape,
**input_embedd)
edi = edge_index_input
# starting
vp = n
ep = ed
up = uenv
vp = MLP(**node_ff_args)(vp)
ep = MLP(**edge_ff_args)(ep)
up = MLP(**state_ff_args)(up)
vp2 = vp
ep2 = ep
up2 = up
for i in range(0, nblocks):
if has_ff and i > 0:
vp2 = MLP(**node_ff_args)(vp)
ep2 = MLP(**edge_ff_args)(ep)
up2 = MLP(**state_ff_args)(up)
# MEGnetBlock
vp2, ep2, up2 = MEGnetBlock(**meg_block_args)(
[vp2, ep2, edi, up2])
# skip connection
if dropout is not None:
vp2 = Dropout(dropout, name='dropout_atom_%d' % i)(vp2)
ep2 = Dropout(dropout, name='dropout_bond_%d' % i)(ep2)
up2 = Dropout(dropout, name='dropout_state_%d' % i)(up2)
vp = Add()([vp2, vp])
ep = Add()([ep2, ep])
up = Add(input_tensor_type="tensor")([up2, up])
if use_set2set:
vp = Dense(set2set_args["channels"], activation='linear')(vp) # to match units
ep = Dense(set2set_args["channels"], activation='linear')(ep) # to match units
vp = Set2Set(**set2set_args)(vp)
ep = Set2Set(**set2set_args)(ep)
else:
vp = PoolingNodes()(vp)
ep = PoolingGlobalEdges()(ep)
ep = ks.layers.Flatten()(ep)
vp = ks.layers.Flatten()(vp)
final_vec = ks.layers.Concatenate(axis=-1)([vp, ep, up])
if dropout is not None:
final_vec = ks.layers.Dropout(dropout, name='dropout_final')(final_vec)
# final dense layers
main_output = MLP(**output_mlp, input_tensor_type="tensor")(final_vec)
model = ks.models.Model(inputs=[node_input, edge_input, edge_index_input, env_input], outputs=main_output)
return model
class AttentiveNodePooling(GraphBaseLayer):
r"""Computes the attentive pooling for node embeddings.
Args:
units (int): Units for the linear trafo of node features before attention.
pooling_method(str): Initial pooling before iteration. Default is "sum".
depth (int): Number of iterations for graph embedding. Default is 3.
activation (str): Activation. Default is {"class_name": "kgcnn>leaky_relu", "config": {"alpha": 0.2}}.
activation_context (str): Activation function for context. Default is "elu".
use_bias (bool): Use bias. Default is True.
kernel_regularizer: Kernel regularization. Default is None.
bias_regularizer: Bias regularization. Default is None.
activity_regularizer: Activity regularization. Default is None.
kernel_constraint: Kernel constrains. Default is None.
bias_constraint: Bias constrains. Default is None.
kernel_initializer: Initializer for kernels. Default is 'glorot_uniform'.
bias_initializer: Initializer for bias. Default is 'zeros'.
"""
def __init__(self,
units,
depth=3,
pooling_method="sum",
activation='kgcnn>leaky_relu',
activation_context="elu",
use_bias=True,
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
kernel_initializer='glorot_uniform',
bias_initializer='zeros',
recurrent_activation='sigmoid',
recurrent_initializer='orthogonal',
recurrent_regularizer=None,
recurrent_constraint=None,
dropout=0.0,
recurrent_dropout=0.0,
reset_after=True,
**kwargs):
"""Initialize layer."""
super(AttentiveNodePooling, self).__init__(**kwargs)
self.pooling_method = pooling_method
self.depth = depth
# dense args
self.units = int(units)
kernel_args = {"use_bias": use_bias, "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}
gru_args = {"recurrent_activation": recurrent_activation,
"use_bias": use_bias, "kernel_initializer": kernel_initializer,
"recurrent_initializer": recurrent_initializer, "bias_initializer": bias_initializer,
"kernel_regularizer": kernel_regularizer, "recurrent_regularizer": recurrent_regularizer,
"bias_regularizer": bias_regularizer, "kernel_constraint": kernel_constraint,
"recurrent_constraint": recurrent_constraint, "bias_constraint": bias_constraint,
"dropout": dropout, "recurrent_dropout": recurrent_dropout, "reset_after": reset_after}
self.lay_linear_trafo = Dense(units, activation="linear", **kernel_args, **self._kgcnn_info)
self.lay_alpha = Dense(1, activation=activation, **kernel_args, **self._kgcnn_info)
self.lay_gather_s = GatherState(**self._kgcnn_info)
self.lay_concat = Concatenate(axis=-1, **self._kgcnn_info)
self.lay_pool_start = PoolingNodes(pooling_method=self.pooling_method, **self._kgcnn_info)
self.lay_pool_attention = PoolingNodesAttention(**self._kgcnn_info)
self.lay_final_activ = Activation(activation=activation_context, **self._kgcnn_info)
self.lay_gru = tf.keras.layers.GRUCell(units=units, activation="tanh", **gru_args)
def build(self, input_shape):
"""Build layer."""
super(AttentiveNodePooling, self).build(input_shape)
def call(self, inputs, **kwargs):
"""Forward pass.
Args:
inputs: nodes
- nodes: Node features of shape (batch, [N], F)
Returns:
features: Hidden tensor of pooled edge attentions for each node.
"""
node = inputs
h = self.lay_pool_start(node)
Wn = self.lay_linear_trafo(node)
for _ in range(self.depth):
hv = self.lay_gather_s(h, node)
ev = self.lay_concat([hv, node])
av = self.self.lay_alpha(ev)
cont = self.lay_pool_attention([Wn, av])
cont = self.lay_final_activ(cont)
h = self.lay_gru(cont, h, **kwargs)
out = h
return out
def get_config(self):
"""Update layer config."""
config = super(AttentiveNodePooling, self).get_config()
config.update({"units": self.units, "depth": self.depth, "pooling_method": self.pooling_method})
conf_sub = self.lay_alpha.get_config()
for x in ["kernel_regularizer", "activity_regularizer", "bias_regularizer", "kernel_constraint",
"bias_constraint", "kernel_initializer", "bias_initializer", "activation", "use_bias"]:
config.update({x: conf_sub[x]})
conf_context = self.lay_final_activ.get_config()
config.update({"activation_context": conf_context["activation"]})
conf_gru = self.lay_gru.get_config()
for x in ["recurrent_activation", "recurrent_initializer", "mrecurrent_regularizer", "recurrent_constraint",
"dropout", "recurrent_dropout", "reset_after"]:
config.update({x: conf_gru[x]})
return config
class DimNetOutputBlock(GraphBaseLayer):
"""DimNetOutputBlock.
Args:
emb_size (list): List of node embedding dimension.
out_emb_size (list): List of edge embedding dimension.
num_dense (list): List of environment embedding dimension.
num_targets (int): Number of output target dimension. Defaults to 12.
use_bias (bool, optional): Use bias. Defaults to True.
kernel_initializer: Initializer for kernels. Default is 'orthogonal'.
output_kernel_initializer: Initializer for last kernel. Default is 'zeros'.
bias_initializer: Initializer for bias. Default is 'zeros'.
activation (str): Activation function. Default is 'softplus2' with fall-back 'selu'.
kernel_regularizer: Kernel regularization. Default is None.
bias_regularizer: Bias regularization. Default is None.
activity_regularizer: Activity regularization. Default is None.
kernel_constraint: Kernel constrains. Default is None.
bias_constraint: Bias constrains. Default is None.
pooling_method (str): Pooling method information for layer. Default is 'mean'.
"""
def __init__(self,
emb_size,
out_emb_size,
num_dense,
num_targets=12,
use_bias=True,
output_kernel_initializer="zeros",
kernel_initializer='orthogonal',
bias_initializer='zeros',
activation=None,
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
pooling_method="sum",
**kwargs):
"""Initialize layer."""
super(DimNetOutputBlock, self).__init__(**kwargs)
self.pooling_method = pooling_method
self.emb_size = emb_size
self.out_emb_size = out_emb_size
self.num_dense = num_dense
self.num_targets = num_targets
self.use_bias = use_bias
if activation is None and 'swish' in kgcnn_custom_act:
activation = 'swish'
elif activation is None:
activation = "selu"
kernel_args = {
"kernel_regularizer": kernel_regularizer,
"activity_regularizer": activity_regularizer,
"kernel_constraint": kernel_constraint,
"bias_initializer": bias_initializer,
"bias_regularizer": bias_regularizer,
"bias_constraint": bias_constraint,
}
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)
self.dense_rbf = Dense(emb_size,
use_bias=False,
kernel_initializer=kernel_initializer,
**kernel_args)
self.up_projection = Dense(out_emb_size,
use_bias=False,
kernel_initializer=kernel_initializer,
**kernel_args)
self.dense_mlp = MLP([out_emb_size] * num_dense,
activation=activation,
kernel_initializer=kernel_initializer,
use_bias=use_bias,
**mlp_args)
self.dimnet_mult = Multiply(input_tensor_type=self.input_tensor_type)
self.pool = PoolingLocalEdges(**pool_args)
self.dense_final = Dense(num_targets,
use_bias=False,
kernel_initializer=output_kernel_initializer,
**kernel_args)
def build(self, input_shape):
"""Build layer."""
super(DimNetOutputBlock, self).build(input_shape)
def call(self, inputs, **kwargs):
"""Forward pass.
Args:
inputs: [nodes, edges, edge_index, state]
- nodes (tf.ragged): Node embeddings of shape (batch, [N], F)
- edges (tf.ragged): Edge or message embeddings of shape (batch, [M], F)
- rbf (tf.ragged): Edge distance basis of shape (batch, [M], F)
- edge_index (tf.ragged): Node indices of shape (batch, [M], 2)
Returns:
node_update (tf.ragged): Updated node embeddings.
"""
# Calculate edge Update
n_atoms, x, rbf, idnb_i = inputs
g = self.dense_rbf(rbf)
x = self.dimnet_mult([g, x])
x = self.pool([n_atoms, x, idnb_i])
x = self.up_projection(x)
x = self.dense_mlp(x)
x = self.dense_final(x)
return x
def get_config(self):
config = super(DimNetOutputBlock, self).get_config()
conf_mlp = self.dense_mlp.get_config()
for x in [
"kernel_regularizer", "activity_regularizer",
"bias_regularizer", "kernel_constraint", "bias_constraint",
"kernel_initializer", "bias_initializer", "activation"
]:
config.update({x: conf_mlp[x][0]})
conf_dense_output = self.dense_final.get_config()
config.update({
"output_kernel_initializer":
conf_dense_output["kernel_initializer"]
})
config.update({
"pooling_method": self.pooling_method,
"use_bias": self.use_bias
})
config.update({
"emb_size": self.emb_size,
"out_emb_size": self.out_emb_size,
"num_dense": self.num_dense,
"num_targets": self.num_targets
})
return config
def __init__(self,
units,
use_bias=True,
activation=None,
activity_regularizer=None,
kernel_regularizer=None,
bias_regularizer=None,
kernel_initializer='glorot_uniform',
bias_initializer='zeros',
kernel_constraint=None,
bias_constraint=None,
ragged_validate=False,
input_tensor_type="ragged",
**kwargs):
"""Initialize MLP as for dense."""
super(MLP, self).__init__(**kwargs)
self._supports_ragged_inputs = True
self.ragged_validate = ragged_validate
self.input_tensor_type = input_tensor_type
# Make to one element list
if isinstance(units, int):
units = [units]
if not isinstance(use_bias, list) and not isinstance(use_bias, tuple):
use_bias = [use_bias for _ in units]
if not isinstance(activation, list) and not isinstance(
activation, tuple):
activation = [activation for _ in units]
if not isinstance(kernel_regularizer, list) and not isinstance(
kernel_regularizer, tuple):
kernel_regularizer = [kernel_regularizer for _ in units]
if not isinstance(bias_regularizer, list) and not isinstance(
bias_regularizer, tuple):
bias_regularizer = [bias_regularizer for _ in units]
if not isinstance(activity_regularizer, list) and not isinstance(
activity_regularizer, tuple):
activity_regularizer = [activity_regularizer for _ in units]
if not isinstance(kernel_initializer, list) and not isinstance(
kernel_initializer, tuple):
kernel_initializer = [kernel_initializer for _ in units]
if not isinstance(bias_initializer, list) and not isinstance(
bias_initializer, tuple):
bias_initializer = [bias_initializer for _ in units]
if not isinstance(kernel_constraint, list) and not isinstance(
kernel_constraint, tuple):
kernel_constraint = [kernel_constraint for _ in units]
if not isinstance(bias_constraint, list) and not isinstance(
bias_constraint, tuple):
bias_constraint = [bias_constraint for _ in units]
for x in [
activation, kernel_regularizer, bias_regularizer,
activity_regularizer, kernel_initializer, bias_initializer,
kernel_constraint, bias_constraint, use_bias
]:
if len(x) != len(units):
raise ValueError("Error: Provide matching list of units",
units, "and", x, "or simply a single value.")
# Serialized props
self.mlp_units = list(units)
self.mlp_use_bias = list(use_bias)
self.mlp_activation = list(
[tf.keras.activations.get(x) for x in activation])
self.mlp_kernel_regularizer = list(
[tf.keras.regularizers.get(x) for x in kernel_regularizer])
self.mlp_bias_regularizer = list(
[tf.keras.regularizers.get(x) for x in bias_regularizer])
self.mlp_activity_regularizer = list(
[tf.keras.regularizers.get(x) for x in activity_regularizer])
self.mlp_kernel_initializer = list(
[tf.keras.initializers.get(x) for x in kernel_initializer])
self.mlp_bias_initializer = list(
[tf.keras.initializers.get(x) for x in bias_initializer])
self.mlp_kernel_constraint = list(
[tf.keras.constraints.get(x) for x in kernel_constraint])
self.mlp_bias_constraint = list(
[tf.keras.constraints.get(x) for x in bias_constraint])
self.mlp_dense_list = [
Dense(units=self.mlp_units[i],
use_bias=self.mlp_use_bias[i],
name=self.name + '_dense_' + str(i),
activation=self.mlp_activation[i],
activity_regularizer=self.mlp_activity_regularizer[i],
kernel_regularizer=self.mlp_kernel_regularizer[i],
bias_regularizer=self.mlp_bias_regularizer[i],
kernel_initializer=self.mlp_kernel_initializer[i],
bias_initializer=self.mlp_bias_initializer[i],
kernel_constraint=self.mlp_kernel_constraint[i],
bias_constraint=self.mlp_bias_constraint[i],
ragged_validate=self.ragged_validate,
input_tensor_type=self.input_tensor_type)
for i in range(len(self.mlp_units))
]
def __init__(self,
emb_size,
out_emb_size,
num_dense,
num_targets=12,
use_bias=True,
output_kernel_initializer="zeros",
kernel_initializer='orthogonal',
bias_initializer='zeros',
activation=None,
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
pooling_method="sum",
**kwargs):
"""Initialize layer."""
super(DimNetOutputBlock, self).__init__(**kwargs)
self.pooling_method = pooling_method
self.emb_size = emb_size
self.out_emb_size = out_emb_size
self.num_dense = num_dense
self.num_targets = num_targets
self.use_bias = use_bias
if activation is None and 'swish' in kgcnn_custom_act:
activation = 'swish'
elif activation is None:
activation = "selu"
kernel_args = {
"kernel_regularizer": kernel_regularizer,
"activity_regularizer": activity_regularizer,
"kernel_constraint": kernel_constraint,
"bias_initializer": bias_initializer,
"bias_regularizer": bias_regularizer,
"bias_constraint": bias_constraint,
}
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)
self.dense_rbf = Dense(emb_size,
use_bias=False,
kernel_initializer=kernel_initializer,
**kernel_args)
self.up_projection = Dense(out_emb_size,
use_bias=False,
kernel_initializer=kernel_initializer,
**kernel_args)
self.dense_mlp = MLP([out_emb_size] * num_dense,
activation=activation,
kernel_initializer=kernel_initializer,
use_bias=use_bias,
**mlp_args)
self.dimnet_mult = Multiply(input_tensor_type=self.input_tensor_type)
self.pool = PoolingLocalEdges(**pool_args)
self.dense_final = Dense(num_targets,
use_bias=False,
kernel_initializer=output_kernel_initializer,
**kernel_args)
def make_inorp( # Input
input_node_shape,
input_edge_shape,
input_state_shape,
input_embedd: dict = None,
# Output
output_embedd: dict = None,
output_mlp: dict = None,
# Model specific parameter
depth=3,
use_set2set: bool = False, # not in original paper
node_mlp_args: dict = None,
edge_mlp_args: dict = None,
set2set_args: dict = None,
pooling_args: dict = None):
"""
Generate Interaction network.
Args:
input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used.
input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used.
input_state_shape (list): Shape of state features. If shape is (,) embedding layer is used.
input_embedd (dict): Dictionary of embedding parameters used if input shape is None. Default is
{'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100,
'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64,
'input_type': 'ragged'}.
output_embedd (dict): Dictionary of embedding parameters of the graph network. Default is
{"output_mode": 'graph', "output_type": 'padded'}.
output_mlp (dict): Dictionary of arguments for final MLP regression or classifcation layer. Default is
{"use_bias": [True, True, False], "units": [25, 10, 1],
"activation": ['relu', 'relu', 'sigmoid']}.
depth (int): Number of convolution layers. Default is 3.
node_mlp_args (dict): Dictionary of arguments for MLP for node update. Default is
{"units": [100, 50], "use_bias": True, "activation": ['relu', "linear"]}
edge_mlp_args (dict): Dictionary of arguments for MLP for interaction update. Default is
{"units": [100, 100, 100, 100, 50],
"activation": ['relu', 'relu', 'relu', 'relu', "linear"]}
use_set2set (str): Use set2set pooling for graph embedding. Default is False.
set2set_args (dict): Dictionary of set2set layer arguments. Default is
{'channels': 32, 'T': 3, "pooling_method": "mean", "init_qstar": "mean"}.
pooling_args (dict): Dictionary for message pooling arguments. Default is
{'is_sorted': False, 'has_unconnected': True, 'pooling_method': "segment_mean"}
Returns:
model (tf.keras.model): Interaction model.
"""
# default values
model_default = {
'input_embedd': {
'input_node_vocab': 95,
'input_edge_vocab': 5,
'input_state_vocab': 100,
'input_node_embedd': 64,
'input_edge_embedd': 64,
'input_state_embedd': 64,
'input_tensor_type': 'ragged'
},
'output_embedd': {
"output_mode": 'graph',
"output_tensor_type": 'padded'
},
'output_mlp': {
"use_bias": [True, True, False],
"units": [25, 10, 1],
"activation": ['relu', 'relu', 'sigmoid']
},
'set2set_args': {
'channels': 32,
'T': 3,
"pooling_method": "mean",
"init_qstar": "mean"
},
'node_mlp_args': {
"units": [100, 50],
"use_bias": True,
"activation": ['relu', "linear"]
},
'edge_mlp_args': {
"units": [100, 100, 100, 100, 50],
"activation": ['relu', 'relu', 'relu', 'relu', "linear"]
},
'pooling_args': {
'is_sorted': False,
'has_unconnected': True,
'pooling_method': "segment_mean"
}
}
# Update default values
input_embedd = update_model_args(model_default['input_embedd'],
input_embedd)
output_embedd = update_model_args(model_default['output_embedd'],
output_embedd)
output_mlp = update_model_args(model_default['output_mlp'], output_mlp)
set2set_args = update_model_args(model_default['set2set_args'],
set2set_args)
node_mlp_args = update_model_args(model_default['node_mlp_args'],
node_mlp_args)
edge_mlp_args = update_model_args(model_default['edge_mlp_args'],
edge_mlp_args)
pooling_args = update_model_args(model_default['pooling_args'],
pooling_args)
gather_args = {"node_indexing": "sample"}
# Make input embedding, if no feature dimension
node_input, n, edge_input, ed, edge_index_input, env_input, uenv = generate_standard_graph_input(
input_node_shape, input_edge_shape, input_state_shape, **input_embedd)
# Preprocessing
edi = edge_index_input
ev = GatherState(**gather_args)([uenv, n])
# n-Layer Step
for i in range(0, depth):
# upd = GatherNodes()([n,edi])
eu1 = GatherNodesIngoing(**gather_args)([n, edi])
eu2 = GatherNodesOutgoing(**gather_args)([n, edi])
upd = Concatenate(axis=-1)([eu2, eu1])
eu = Concatenate(axis=-1)([upd, ed])
eu = MLP(**edge_mlp_args)(eu)
# Pool message
nu = PoolingLocalEdges(**pooling_args)(
[n, eu, edi]) # Summing for each node connection
# Add environment
nu = Concatenate(axis=-1)(
[n, nu, ev]) # Concatenate node features with new edge updates
n = MLP(**node_mlp_args)(nu)
if output_embedd["output_mode"] == 'graph':
if use_set2set:
# output
outss = Dense(set2set_args["channels"], activation="linear")(n)
out = Set2Set(**set2set_args)(outss)
else:
out = PoolingNodes(**pooling_args)(n)
output_mlp.update({"input_tensor_type": "tensor"})
main_output = MLP(**output_mlp)(out)
else: # Node labeling
out = n
main_output = MLP(**output_mlp)(out)
main_output = ChangeTensorType(
input_tensor_type="ragged",
output_tensor_type="tensor")(main_output)
# no ragged for distribution atm
model = ks.models.Model(
inputs=[node_input, edge_input, edge_index_input, env_input],
outputs=main_output)
return model
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)
class AttentionHeadGAT(GraphBaseLayer):
r"""Computes the attention head according to GAT.
The attention coefficients are computed by $a_{ij} = \sigma( W n_i || W n_j)$,
optionally by $a_{ij} = \sigma( W n_i || W n_j || e_{ij})$.
The attention is obtained by $\alpha_ij = softmax_j (a_{ij})$.
And the messages are pooled by $n_i = \sum_j \alpha_{ij} e_ij $.
And finally passed through an activation $h_i = \sigma(\sum_j \alpha_{ij} e_ij)$.
An edge is defined by index tuple (i,j) with i<-j connection.
If graphs indices were in 'batch' mode, the layer's 'node_indexing' must be set to 'batch'.
Args:
units (int): Units for the linear trafo of node features before attention.
use_edge_features (bool): Append edge features to attention computation. Default is False.
activation (str): Activation. Default is {"class_name": "leaky_relu", "config": {"alpha": 0.2}},
with fall-back "relu".
use_bias (bool): Use bias. Default is True.
kernel_regularizer: Kernel regularization. Default is None.
bias_regularizer: Bias regularization. Default is None.
activity_regularizer: Activity regularization. Default is None.
kernel_constraint: Kernel constrains. Default is None.
bias_constraint: Bias constrains. Default is None.
kernel_initializer: Initializer for kernels. Default is 'glorot_uniform'.
bias_initializer: Initializer for bias. Default is 'zeros'.
"""
def __init__(self,
units,
use_edge_features=False,
activation=None,
use_bias=True,
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(AttentionHeadGAT, self).__init__(**kwargs)
# graph args
self.use_edge_features = use_edge_features
# dense args
self.units = int(units)
if activation is None and "leaky_relu" in kgcnn_custom_act:
activation = {"class_name": "leaky_relu", "config": {"alpha": 0.2}}
elif activation is None:
activation = "relu"
kernel_args = {
"use_bias": use_bias,
"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
}
dens_args = {
"ragged_validate": self.ragged_validate,
"input_tensor_type": self.input_tensor_type
}
dens_args.update(kernel_args)
gather_args = self._all_kgcnn_info
pooling_args = self._all_kgcnn_info
self.lay_linear_trafo = Dense(units, activation="linear", **dens_args)
self.lay_alpha = Dense(1, activation=activation, **dens_args)
self.lay_gather_in = GatherNodesIngoing(**gather_args)
self.lay_gather_out = GatherNodesOutgoing(**gather_args)
self.lay_concat = Concatenate(axis=-1,
input_tensor_type=self.input_tensor_type)
self.lay_pool_attention = PoolingLocalEdgesAttention(**pooling_args)
self.lay_final_activ = Activation(
activation=activation, input_tensor_type=self.input_tensor_type)
def build(self, input_shape):
"""Build layer."""
super(AttentionHeadGAT, self).build(input_shape)
def call(self, inputs, **kwargs):
"""Forward pass.
Args:
inputs (list): of [node, edges, edge_indices]
- nodes: Node features of shape (batch, [N], F)
- edges: Edge or message features of shape (batch, [M], F)
- edge_index: Edge indices of shape (batch, [M], 2)
Returns:
features: Feature tensor of pooled edge attentions for each node.
"""
node, edge, edge_index = inputs
n_in = self.lay_gather_in([node, edge_index])
n_out = self.lay_gather_out([node, edge_index])
wn_in = self.lay_linear_trafo(n_in)
wn_out = self.lay_linear_trafo(n_out)
if self.use_edge_features:
e_ij = self.lay_concat([wn_in, wn_out, edge])
else:
e_ij = self.lay_concat([wn_in, wn_out])
a_ij = self.lay_alpha(e_ij) # Should be dimension (batch*None,1)
n_i = self.lay_pool_attention([node, wn_out, a_ij, edge_index])
out = self.lay_final_activ(n_i)
return out
def get_config(self):
"""Update layer config."""
config = super(AttentionHeadGAT, self).get_config()
config.update({
"use_edge_features": self.use_edge_features,
"units": self.units
})
conf_sub = self.lay_alpha.get_config()
for x in [
"kernel_regularizer", "activity_regularizer",
"bias_regularizer", "kernel_constraint", "bias_constraint",
"kernel_initializer", "bias_initializer", "activation",
"use_bias"
]:
config.update({x: conf_sub[x]})
return config
class SchNetInteraction(GraphBaseLayer):
"""
Schnet interaction block, which uses the continuous filter convolution from SchNetCFconv.
Args:
units (int): Dimension of node embedding. Default is 128.
cfconv_pool (str): Pooling method information for SchNetCFconv layer. Default is'segment_sum'.
use_bias (bool): Use bias in last layers. Default is True.
activation (str): Activation function. Default is 'shifted_softplus' with fall-back 'selu'.
kernel_regularizer: Kernel regularization. Default is None.
bias_regularizer: Bias regularization. Default is None.
activity_regularizer: Activity regularization. Default is None.
kernel_constraint: Kernel constrains. Default is None.
bias_constraint: Bias constrains. Default is None.
kernel_initializer: Initializer for kernels. Default is 'glorot_uniform'.
bias_initializer: Initializer for bias. Default is 'zeros'.
"""
def __init__(self,
units=128,
cfconv_pool='sum',
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(SchNetInteraction, self).__init__(**kwargs)
self.cfconv_pool = cfconv_pool
self.use_bias = use_bias
self.units = units
if activation is None and 'shifted_softplus' in kgcnn_custom_act:
activation = 'shifted_softplus'
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
}
conv_args = {
"units": self.units,
"use_bias": use_bias,
"activation": activation,
"cfconv_pool": cfconv_pool
}
conv_args.update(kernel_args)
conv_args.update(self._all_kgcnn_info)
# Layers
self.lay_cfconv = SchNetCFconv(**conv_args)
self.lay_dense1 = Dense(units=self.units,
activation='linear',
use_bias=False,
input_tensor_type=self.input_tensor_type,
ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_dense2 = Dense(units=self.units,
activation=activation,
use_bias=self.use_bias,
input_tensor_type=self.input_tensor_type,
ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_dense3 = Dense(units=self.units,
activation='linear',
use_bias=self.use_bias,
input_tensor_type=self.input_tensor_type,
ragged_validate=self.ragged_validate,
**kernel_args)
self.lay_add = Add(input_tensor_type=self.input_tensor_type,
ragged_validate=self.ragged_validate)
def build(self, input_shape):
"""Build layer."""
super(SchNetInteraction, self).build(input_shape)
def call(self, inputs, **kwargs):
"""Forward pass: Calculate node update.
Args:
inputs: [nodes, edges, edge_index]
- nodes: Node embeddings of shape (batch, [N], F)
- edges: Edge or message embeddings of shape (batch, [N], F)
- edge_index: Edge indices of shape (batch, [N], 2)
Returns:
node_update: Updated node embeddings.
"""
node, edge, indexlist = inputs
x = self.lay_dense1(node)
x = self.lay_cfconv([x, edge, indexlist])
x = self.lay_dense2(x)
x = self.lay_dense3(x)
out = self.lay_add([node, x])
return out
def get_config(self):
config = super(SchNetInteraction, self).get_config()
config.update({
"cfconv_pool": self.cfconv_pool,
"units": self.units,
"use_bias": self.use_bias
})
conf_dense = self.lay_dense2.get_config()
for x in [
"activation", "kernel_regularizer", "bias_regularizer",
"activity_regularizer", "kernel_constraint", "bias_constraint",
"kernel_initializer", "bias_initializer"
]:
config.update({x: conf_dense[x]})
return config
