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,
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,
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,
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 __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 __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',
cfconv_pool='segment_sum',
is_sorted=False,
has_unconnected=True,
partition_type="row_length",
node_indexing='sample',
input_tensor_type="ragged",
ragged_validate=False,
**kwargs):
"""Initialize Layer."""
super(SchNetCFconv, self).__init__(**kwargs)
self.cfconv_pool = cfconv_pool
self.is_sorted = is_sorted
self.partition_type = partition_type
self.has_unconnected = has_unconnected
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)
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"
self.cfc_activation = tf.keras.activations.get(activation)
self.cfc_kernel_initializer = tf.keras.initializers.get(
kernel_initializer)
self.cfc_bias_initializer = tf.keras.initializers.get(bias_initializer)
self.cfc_kernel_regularizer = tf.keras.regularizers.get(
kernel_regularizer)
self.cfc_bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
self.cfc_activity_regularizer = tf.keras.regularizers.get(
activity_regularizer)
self.cfc_kernel_constraint = tf.keras.constraints.get(
kernel_constraint)
self.cfc_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
}
# 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_method=self.cfconv_pool,
is_sorted=self.is_sorted,
has_unconnected=self.has_unconnected,
partition_type=self.partition_type,
node_indexing=self.node_indexing,
input_tensor_type=self.input_tensor_type,
ragged_validate=self.ragged_validate)
self.gather_n = GatherNodesOutgoing(
node_indexing=self.node_indexing,
partition_type=self.partition_type,
input_tensor_type=self.input_tensor_type,
is_sorted=self.is_sorted,
ragged_validate=self.ragged_validate,
has_unconnected=self.has_unconnected)
self.lay_mult = Multiply(input_tensor_type=self.input_tensor_type,
ragged_validate=self.ragged_validate)