def test_set2set(self):
n_hidden = 10
set2set = Set2Set(n_hidden=n_hidden)
out = set2set([self.x, self.index])
shapes = set2set.compute_output_shape([self.x.shape, self.index.shape])
self.assertEqual(shapes[-1], n_hidden * 2)
model = Model(inputs=[self.x, self.index], outputs=out)
model.compile(loss='mse', optimizer='adam')
x = np.random.normal(size=(1, 5, 6))
index = np.array([[0, 0, 0, 1, 1]])
result = model.predict([x, index])
self.assertListEqual(list(result.shape), [1, 2, n_hidden * 2])
def __init__(self,
nfeat_edge,
nfeat_global,
nfeat_node=None,
nblocks=3,
lr=1e-3,
n1=64,
n2=32,
n3=16,
nvocal=95,
embedding_dim=16,
npass=3,
ntarget=1,
act=softplus2,
is_classification=False,
loss="mse",
l2_coef=None,
dropout=None,
graph_convertor=None,
optimizer_kwargs=None
):
int32 = 'int32'
if nfeat_node is None:
x1 = Input(shape=(None,), dtype=int32) # only z as feature
x1_ = Embedding(nvocal, embedding_dim)(x1)
else:
x1 = Input(shape=(None, nfeat_node))
x1_ = x1
x2 = Input(shape=(None, nfeat_edge))
x3 = Input(shape=(None, nfeat_global))
x4 = Input(shape=(None,), dtype=int32)
x5 = Input(shape=(None,), dtype=int32)
x6 = Input(shape=(None,), dtype=int32)
x7 = Input(shape=(None,), dtype=int32)
if l2_coef is not None:
reg = l2(l2_coef)
else:
reg = None
# two feedforward layers
def ff(x, n_hiddens=[n1, n2]):
out = x
for i in n_hiddens:
out = Dense(i, activation=act, kernel_regularizer=reg)(out)
return out
# a block corresponds to two feedforward layers + one MEGNetLayer layer
# Note the first block does not contain the feedforward layer since
# it will be explicitly added before the block
def one_block(a, b, c, has_ff=True):
if has_ff:
x1_ = ff(a)
x2_ = ff(b)
x3_ = ff(c)
else:
x1_ = a
x2_ = b
x3_ = c
out = MEGNetLayer(
[n1, n1, n2], [n1, n1, n2], [n1, n1, n2],
pool_method='mean', activation=act, kernel_regularizer=reg)(
[x1_, x2_, x3_, x4, x5, x6, x7])
x1_temp = out[0]
x2_temp = out[1]
x3_temp = out[2]
if dropout:
x1_temp = Dropout(dropout)(x1_temp)
x2_temp = Dropout(dropout)(x2_temp)
x3_temp = Dropout(dropout)(x3_temp)
return x1_temp, x2_temp, x3_temp
x1_ = ff(x1_)
x2_ = ff(x2)
x3_ = ff(x3)
for i in range(nblocks):
if i == 0:
has_ff = False
else:
has_ff = True
x1_1 = x1_
x2_1 = x2_
x3_1 = x3_
x1_1, x2_1, x3_1 = one_block(x1_1, x2_1, x3_1, has_ff)
# skip connection
x1_ = Add()([x1_, x1_1])
x2_ = Add()([x2_, x2_1])
x3_ = Add()([x3_, x3_1])
# set2set for both the atom and bond
node_vec = Set2Set(T=npass, n_hidden=n3, kernel_regularizer=reg)([x1_, x6])
edge_vec = Set2Set(T=npass, n_hidden=n3, kernel_regularizer=reg)([x2_, x7])
# concatenate atom, bond, and global
final_vec = Concatenate(axis=-1)([node_vec, edge_vec, x3_])
if dropout:
final_vec = Dropout(dropout)(final_vec)
# final dense layers
final_vec = Dense(n2, activation=act, kernel_regularizer=reg)(final_vec)
final_vec = Dense(n3, activation=act, kernel_regularizer=reg)(final_vec)
if is_classification:
final_act = 'sigmoid'
loss = 'binary_crossentropy'
else:
final_act = None
loss = loss
out = Dense(ntarget, activation=final_act)(final_vec)
model = Model(inputs=[x1, x2, x3, x4, x5, x6, x7], outputs=out)
opt_params = {'lr': lr}
if optimizer_kwargs is not None:
opt_params.update(optimizer_kwargs)
model.compile(Adam(**opt_params), loss)
if graph_convertor is None:
graph_convertor = CrystalGraph(cutoff=4, bond_convertor=GaussianDistance(np.linspace(0, 5, 100), 0.5))
super().__init__(model=model, graph_convertor=graph_convertor)
def make_megnet_model(nfeat_edge: int = None,
nfeat_global: int = None,
nfeat_node: int = None,
nblocks: int = 3,
n1: int = 64,
n2: int = 32,
n3: int = 16,
nvocal: int = 95,
embedding_dim: int = 16,
nbvocal: int = None,
bond_embedding_dim: int = None,
ngvocal: int = None,
global_embedding_dim: int = None,
npass: int = 3,
ntarget: int = 1,
act: Callable = softplus2,
is_classification: bool = False,
l2_coef: float = None,
dropout: float = None,
dropout_on_predict: bool = False) -> Model:
"""Make a MEGNet Model
Args:
nfeat_edge: (int) number of bond features
nfeat_global: (int) number of state features
nfeat_node: (int) number of atom features
nblocks: (int) number of MEGNetLayer blocks
n1: (int) number of hidden units in layer 1 in MEGNetLayer
n2: (int) number of hidden units in layer 2 in MEGNetLayer
n3: (int) number of hidden units in layer 3 in MEGNetLayer
nvocal: (int) number of total element
embedding_dim: (int) number of embedding dimension
nbvocal: (int) number of bond types if bond attributes are types
bond_embedding_dim: (int) number of bond embedding dimension
ngvocal: (int) number of global types if global attributes are types
global_embedding_dim: (int) number of global embedding dimension
npass: (int) number of recurrent steps in Set2Set layer
ntarget: (int) number of output targets
act: (object) activation function
l2_coef: (float or None) l2 regularization parameter
is_classification: (bool) whether it is a classification task
dropout: (float) dropout rate
dropout_on_predict (bool): Whether to use dropout during prediction and training
Returns:
(Model) Keras model, ready to run
"""
# Get the setting for the training kwarg of Dropout
dropout_training = True if dropout_on_predict else None
# Create the input blocks
int32 = 'int32'
if nfeat_node is None:
x1 = Input(shape=(None, ), dtype=int32) # only z as feature
x1_ = Embedding(nvocal, embedding_dim)(x1)
else:
x1 = Input(shape=(None, nfeat_node))
x1_ = x1
if nfeat_edge is None:
x2 = Input(shape=(None, ), dtype=int32)
x2_ = Embedding(nbvocal, bond_embedding_dim)(x2)
else:
x2 = Input(shape=(None, nfeat_edge))
x2_ = x2
if nfeat_global is None:
x3 = Input(shape=(None, ), dtype=int32)
x3_ = Embedding(ngvocal, global_embedding_dim)(x3)
else:
x3 = Input(shape=(None, nfeat_global))
x3_ = x3
x4 = Input(shape=(None, ), dtype=int32)
x5 = Input(shape=(None, ), dtype=int32)
x6 = Input(shape=(None, ), dtype=int32)
x7 = Input(shape=(None, ), dtype=int32)
if l2_coef is not None:
reg = l2(l2_coef)
else:
reg = None
# two feedforward layers
def ff(x, n_hiddens=[n1, n2]):
out = x
for i in n_hiddens:
out = Dense(i, activation=act, kernel_regularizer=reg)(out)
return out
# a block corresponds to two feedforward layers + one MEGNetLayer layer
# Note the first block does not contain the feedforward layer since
# it will be explicitly added before the block
def one_block(a, b, c, has_ff=True):
if has_ff:
x1_ = ff(a)
x2_ = ff(b)
x3_ = ff(c)
else:
x1_ = a
x2_ = b
x3_ = c
out = MEGNetLayer(
[n1, n1, n2], [n1, n1, n2], [n1, n1, n2],
pool_method='mean',
activation=act,
kernel_regularizer=reg)([x1_, x2_, x3_, x4, x5, x6, x7])
x1_temp = out[0]
x2_temp = out[1]
x3_temp = out[2]
if dropout:
x1_temp = Dropout(dropout)(x1_temp, training=dropout_training)
x2_temp = Dropout(dropout)(x2_temp, training=dropout_training)
x3_temp = Dropout(dropout)(x3_temp, training=dropout_training)
return x1_temp, x2_temp, x3_temp
x1_ = ff(x1_)
x2_ = ff(x2_)
x3_ = ff(x3_)
for i in range(nblocks):
if i == 0:
has_ff = False
else:
has_ff = True
x1_1 = x1_
x2_1 = x2_
x3_1 = x3_
x1_1, x2_1, x3_1 = one_block(x1_1, x2_1, x3_1, has_ff)
# skip connection
x1_ = Add()([x1_, x1_1])
x2_ = Add()([x2_, x2_1])
x3_ = Add()([x3_, x3_1])
# set2set for both the atom and bond
node_vec = Set2Set(T=npass, n_hidden=n3, kernel_regularizer=reg)([x1_, x6])
edge_vec = Set2Set(T=npass, n_hidden=n3, kernel_regularizer=reg)([x2_, x7])
# concatenate atom, bond, and global
final_vec = Concatenate(axis=-1)([node_vec, edge_vec, x3_])
if dropout:
final_vec = Dropout(dropout)(final_vec, training=dropout_training)
# final dense layers
final_vec = Dense(n2, activation=act, kernel_regularizer=reg)(final_vec)
final_vec = Dense(n3, activation=act, kernel_regularizer=reg)(final_vec)
if is_classification:
final_act = 'sigmoid'
else:
final_act = None
out = Dense(ntarget, activation=final_act)(final_vec)
model = Model(inputs=[x1, x2, x3, x4, x5, x6, x7], outputs=out)
return model
def make_megnet_model(
nfeat_edge: int = None,
nfeat_global: int = None,
nfeat_node: int = None,
nblocks: int = 3,
n1: int = 64,
n2: int = 32,
n3: int = 16,
nvocal: int = 95,
embedding_dim: int = 16,
nbvocal: int = None,
bond_embedding_dim: int = None,
ngvocal: int = None,
global_embedding_dim: int = None,
npass: int = 3,
ntarget: int = 1,
act: Callable = softplus2,
is_classification: bool = False,
l2_coef: float = None,
dropout: float = None,
dropout_on_predict: bool = False,
**kwargs,
) -> Model:
"""Make a MEGNet Model
Args:
nfeat_edge: (int) number of bond features
nfeat_global: (int) number of state features
nfeat_node: (int) number of atom features
nblocks: (int) number of MEGNetLayer blocks
n1: (int) number of hidden units in layer 1 in MEGNetLayer
n2: (int) number of hidden units in layer 2 in MEGNetLayer
n3: (int) number of hidden units in layer 3 in MEGNetLayer
nvocal: (int) number of total element
embedding_dim: (int) number of embedding dimension
nbvocal: (int) number of bond types if bond attributes are types
bond_embedding_dim: (int) number of bond embedding dimension
ngvocal: (int) number of global types if global attributes are types
global_embedding_dim: (int) number of global embedding dimension
npass: (int) number of recurrent steps in Set2Set layer
ntarget: (int) number of output targets
act: (object) activation function
l2_coef: (float or None) l2 regularization parameter
is_classification: (bool) whether it is a classification task
dropout: (float) dropout rate
dropout_on_predict (bool): Whether to use dropout during prediction and training
kwargs (dict): in the case where bond inputs are pure distances (not the expanded
distances nor integers for embedding, i.e., nfeat_edge=None and bond_embedding_dim=None),
kwargs can take additional inputs for expand the distance using Gaussian basis.
centers (np.ndarray): array for defining the Gaussian expansion centers
width (float): width for the Gaussian basis
Returns:
(Model) Keras model, ready to run
"""
# Get the setting for the training kwarg of Dropout
dropout_training = True if dropout_on_predict else None
# atom inputs
if nfeat_node is None:
# only z as feature
x1 = Input(shape=(None, ),
dtype=DataType.tf_int,
name="atom_int_input")
x1_ = Embedding(nvocal, embedding_dim, name="atom_embedding")(x1)
else:
x1 = Input(shape=(None, nfeat_node), name="atom_feature_input")
x1_ = x1
# bond inputs
if nfeat_edge is None:
if bond_embedding_dim is not None:
# bond attributes are integers for embedding
x2 = Input(shape=(None, ),
dtype=DataType.tf_int,
name="bond_int_input")
x2_ = Embedding(nbvocal, bond_embedding_dim,
name="bond_embedding")(x2)
else:
# the bond attributes are float distance
x2 = Input(shape=(None, ),
dtype=DataType.tf_float,
name="bond_float_input")
centers = kwargs.get("centers", None)
width = kwargs.get("width", None)
if centers is None and width is None:
raise ValueError(
"If the bond attributes are single float values, "
"we expect the value to be expanded before passing "
"to the models. Therefore, `centers` and `width` for "
"Gaussian basis expansion are needed")
x2_ = GaussianExpansion(centers=centers,
width=width)(x2) # type: ignore
else:
x2 = Input(shape=(None, nfeat_edge), name="bond_feature_input")
x2_ = x2
# state inputs
if nfeat_global is None:
if global_embedding_dim is not None:
# global state inputs are embedding integers
x3 = Input(shape=(None, ),
dtype=DataType.tf_int,
name="state_int_input")
x3_ = Embedding(ngvocal,
global_embedding_dim,
name="state_embedding")(x3)
else:
# take default vector of two zeros
x3 = Input(shape=(None, 2),
dtype=DataType.tf_float,
name="state_default_input")
x3_ = x3
else:
x3 = Input(shape=(None, nfeat_global), name="state_feature_input")
x3_ = x3
x4 = Input(shape=(None, ),
dtype=DataType.tf_int,
name="bond_index_1_input")
x5 = Input(shape=(None, ),
dtype=DataType.tf_int,
name="bond_index_2_input")
x6 = Input(shape=(None, ),
dtype=DataType.tf_int,
name="atom_graph_index_input")
x7 = Input(shape=(None, ),
dtype=DataType.tf_int,
name="bond_graph_index_input")
if l2_coef is not None:
reg = l2(l2_coef)
else:
reg = None
# two feedforward layers
def ff(x, n_hiddens=[n1, n2], name_prefix=None):
if name_prefix is None:
name_prefix = "FF"
out = x
for k, i in enumerate(n_hiddens):
out = Dense(i,
activation=act,
kernel_regularizer=reg,
name=f"{name_prefix}_{k}")(out)
return out
# a block corresponds to two feedforward layers + one MEGNetLayer layer
# Note the first block does not contain the feedforward layer since
# it will be explicitly added before the block
def one_block(a, b, c, has_ff=True, block_index=0):
if has_ff:
x1_ = ff(a, name_prefix=f"block_{block_index}_atom_ff")
x2_ = ff(b, name_prefix=f"block_{block_index}_bond_ff")
x3_ = ff(c, name_prefix=f"block_{block_index}_state_ff")
else:
x1_ = a
x2_ = b
x3_ = c
out = MEGNetLayer(
[n1, n1, n2],
[n1, n1, n2],
[n1, n1, n2],
pool_method="mean",
activation=act,
kernel_regularizer=reg,
name=f"megnet_{block_index}",
)([x1_, x2_, x3_, x4, x5, x6, x7])
x1_temp = out[0]
x2_temp = out[1]
x3_temp = out[2]
if dropout:
x1_temp = Dropout(dropout, name=f"dropout_atom_{block_index}")(
x1_temp, training=dropout_training)
x2_temp = Dropout(dropout, name=f"dropout_bond_{block_index}")(
x2_temp, training=dropout_training)
x3_temp = Dropout(dropout, name=f"dropout_state_{block_index}")(
x3_temp, training=dropout_training)
return x1_temp, x2_temp, x3_temp
x1_ = ff(x1_, name_prefix="preblock_atom")
x2_ = ff(x2_, name_prefix="preblock_bond")
x3_ = ff(x3_, name_prefix="preblock_state")
for i in range(nblocks):
if i == 0:
has_ff = False
else:
has_ff = True
x1_1 = x1_
x2_1 = x2_
x3_1 = x3_
x1_1, x2_1, x3_1 = one_block(x1_1, x2_1, x3_1, has_ff, block_index=i)
# skip connection
x1_ = Add(name=f"block_{i}_add_atom")([x1_, x1_1])
x2_ = Add(name=f"block_{i}_add_bond")([x2_, x2_1])
x3_ = Add(name=f"block_{i}_add_state")([x3_, x3_1])
# print(Set2Set(T=npass, n_hidden=n3, kernel_regularizer=reg, name='set2set_atom'
# ).compute_output_shape([i.shape for i in [x1_, x6]]))
# set2set for both the atom and bond
node_vec = Set2Set(T=npass,
n_hidden=n3,
kernel_regularizer=reg,
name="set2set_atom")([x1_, x6])
# print('Node vec', node_vec)
edge_vec = Set2Set(T=npass,
n_hidden=n3,
kernel_regularizer=reg,
name="set2set_bond")([x2_, x7])
# concatenate atom, bond, and global
final_vec = Concatenate(axis=-1)([node_vec, edge_vec, x3_])
if dropout:
final_vec = Dropout(dropout,
name="dropout_final")(final_vec,
training=dropout_training)
# final dense layers
final_vec = Dense(n2,
activation=act,
kernel_regularizer=reg,
name="readout_0")(final_vec)
final_vec = Dense(n3,
activation=act,
kernel_regularizer=reg,
name="readout_1")(final_vec)
if is_classification:
final_act = "sigmoid"
else:
final_act = None # type: ignore
out = Dense(ntarget, activation=final_act, name="readout_2")(final_vec)
model = Model(inputs=[x1, x2, x3, x4, x5, x6, x7], outputs=out)
return model
def set2set_model(n_feature,
n_connect,
n_global,
n_blocks=3,
lr=1e-3,
n1=64,
n2=32,
n3=16,
n_pass=3,
n_target=1,
act=softplus2,
dropout=None):
"""
construct a graph network model with explicit atom features
:param n_feature: (int) number of atom features
:param n_connect: (int) number of bond features
:param n_global: (int) number of state features
:param n_blocks: (int) number of MEGNet block
:param lr: (float) learning rate
:param n1: (int) number of hidden units in layer 1 in MEGNet
:param n2: (int) number of hidden units in layer 2 in MEGNet
:param n3: (int) number of hidden units in layer 3 in MEGNet
:param n_pass: (int) number of recurrent steps in Set2Set layer
:param n_target: (int) number of output targets
:param act: (object) activation function
:param dropout: (float) dropout rate
:return: keras model object
"""
int32 = 'int32'
x1 = Input(shape=(None, n_feature))
x2 = Input(shape=(None, n_connect))
x3 = Input(shape=(None, n_global))
x4 = Input(shape=(None, ), dtype=int32)
x5 = Input(shape=(None, ), dtype=int32)
x6 = Input(shape=(None, ), dtype=int32)
x7 = Input(shape=(None, ), dtype=int32)
# two feedforward layers
def ff(x, n_hiddens=[n1, n2]):
out = x
for i in n_hiddens:
out = Dense(i, activation=act)(out)
return out
# a block corresponds to two feedforward layers + one MEGNet layer
# Note the first block does not contain the feedforward layer since
# it will be explicitly added before the block
def one_block(a, b, c, has_ff=True):
if has_ff:
x1_ = ff(a)
x2_ = ff(b)
x3_ = ff(c)
else:
x1_ = a
x2_ = b
x3_ = c
out = MEGNet([n1, n1, n2], [n1, n1, n2], [n1, n1, n2],
pool_method='mean',
activation=act)([x1_, x2_, x3_, x4, x5, x6, x7])
x1_temp = out[0]
x2_temp = out[1]
x3_temp = out[2]
if dropout:
x1_temp = Dropout(dropout)(x1_temp)
x2_temp = Dropout(dropout)(x2_temp)
x3_temp = Dropout(dropout)(x3_temp)
return x1_temp, x2_temp, x3_temp
x1_ = ff(x1)
x2_ = ff(x2)
x3_ = ff(x3)
for i in range(n_blocks):
if i == 0:
has_ff = False
else:
has_ff = True
x1_1 = x1_
x2_1 = x2_
x3_1 = x3_
x1_1, x2_1, x3_1 = one_block(x1_1, x2_1, x3_1, has_ff)
# skip connection
x1_ = Add()([x1_, x1_1])
x2_ = Add()([x2_, x2_1])
x3_ = Add()([x3_, x3_1])
# set2set for both the atom and bond
node_vec = Set2Set(T=n_pass, n_hidden=n3)([x1_, x6])
edge_vec = Set2Set(T=n_pass, n_hidden=n3)([x2_, x7])
# concatenate atom, bond, and global
final_vec = Concatenate(axis=-1)([node_vec, edge_vec, x3_])
if dropout:
final_vec = Dropout(dropout)(final_vec)
# final dense layers
final_vec = Dense(n2, activation=act)(final_vec)
final_vec = Dense(n3, activation=act)(final_vec)
out = Dense(n_target)(final_vec)
model = Model(inputs=[x1, x2, x3, x4, x5, x6, x7], outputs=out)
model.compile(Adam(lr), mse_scale)
return model
def set2set_with_embedding_mp(n_connect,
n_global,
n_vocal=95,
embedding_dim=16,
n_blocks=3,
lr=1e-3,
n1=64,
n2=32,
n3=16,
n_pass=3,
n_target=1,
act=softplus2,
l2_coef=None,
is_classification=False):
"""
construct a graph network model with only Z as atom features
:param n_connect: (int) number of bond features
:param n_global: (int) number of state features
:param n_vocal: (int) number of vocabulary. Max Z number needs to be less than this.
since we have max Z number 94 in materials project, here the default n_vocal = 95 (the last index is 94)
:param embedding_dim: (int) embedding vector length
:param n_blocks: (int) number of MEGNet block
:param lr: (float) learning rate
:param n1: (int) number of hidden units in layer 1 in MEGNet
:param n2: (int) number of hidden units in layer 2 in MEGNet
:param n3: (int) number of hidden units in layer 3 in MEGNet
:param n_pass: (int) number of recurrent steps in Set2Set layer
:param n_target: (int) number of output targets
:param act: (object) activation function
:param l2_coef: (float) l2 regularization rate
:param is_classification: (bool) whether it is a classification problem
:return: keras model object
"""
int32 = 'int32'
x1 = Input(shape=(None, ), dtype=int32)
x2 = Input(shape=(None, n_connect))
x3 = Input(shape=(None, n_global))
x4 = Input(shape=(None, ), dtype=int32)
x5 = Input(shape=(None, ), dtype=int32)
x6 = Input(shape=(None, ), dtype=int32)
x7 = Input(shape=(None, ), dtype=int32)
if l2_coef is not None:
reg = l2(l2_coef)
else:
reg = None
def ff(x, n_hiddens=[n1, n2]):
out = x
for i in n_hiddens:
out = Dense(i, activation=act, kernel_regularizer=reg)(out)
return out
def one_block(a, b, c, has_ff=True):
if has_ff:
x1_ = ff(a)
x2_ = ff(b)
x3_ = ff(c)
else:
x1_ = a
x2_ = b
x3_ = c
out = MEGNet([n1, n1, n2], [n1, n1, n2], [n1, n1, n2],
pool_method='mean',
activation=act,
kernel_regularizer=reg)([x1_, x2_, x3_, x4, x5, x6, x7])
return out[0], out[1], out[2]
x1_ = Embedding(n_vocal, embedding_dim)(x1)
x1_ = ff(x1_)
x2_ = ff(x2)
x3_ = ff(x3)
for i in range(n_blocks):
if i == 0:
has_ff = False
else:
has_ff = True
x1_1 = x1_
x2_1 = x2_
x3_1 = x3_
x1_1, x2_1, x3_1 = one_block(x1_1, x2_1, x3_1, has_ff)
x1_ = Add()([x1_, x1_1])
x2_ = Add()([x2_, x2_1])
x3_ = Add()([x3_, x3_1])
node_vec = Set2Set(T=n_pass, n_hidden=n3)([x1_, x6])
edge_vec = Set2Set(T=n_pass, n_hidden=n3)([x2_, x7])
final_vec = Concatenate(axis=-1)([node_vec, edge_vec, x3_])
final_vec = Dense(n2, activation=act)(final_vec)
final_vec = Dense(n3, activation=act)(final_vec)
if is_classification:
final_act = 'sigmoid'
loss = 'binary_crossentropy'
else:
final_act = None
loss = mse_scale
out = Dense(n_target, activation=final_act)(final_vec)
model = Model(inputs=[x1, x2, x3, x4, x5, x6, x7], outputs=out)
model.compile(Adam(lr), loss)
return model
def megnet_model(n_connect,
n_global,
n_feature=None,
n_blocks=3,
lr=1e-3,
n1=64,
n2=32,
n3=16,
n_vocal=95,
embedding_dim=16,
n_pass=3,
n_target=1,
act=softplus2,
is_classification=False,
loss="mse",
l2_coef=None,
dropout=None,
graph_convertor=None,
distance_convertor=None):
"""
construct a graph network model with or without explicit atom features
if n_feature is specified then a general graph model is assumed, otherwise a crystal graph model with z number as
atom feature is assumed.
:param n_connect: (int) number of bond features
:param n_global: (int) number of state features
:param n_feature: (int) number of atom features
:param n_blocks: (int) number of MEGNet block
:param lr: (float) learning rate
:param n1: (int) number of hidden units in layer 1 in MEGNet
:param n2: (int) number of hidden units in layer 2 in MEGNet
:param n3: (int) number of hidden units in layer 3 in MEGNet
:param n_vocal: (int) number of total element
:param embedding_dim: (int) number of embedding dimension
:param n_pass: (int) number of recurrent steps in Set2Set layer
:param n_target: (int) number of output targets
:param act: (object) activation function
:param l2_coef: (float or None) l2 regularization parameter
:param is_classification: (bool) whether it is a classifiation task
:param loss: (object or str) loss function
:param dropout: (float) dropout rate
:param graph_convertor: (object) object that exposes a "convert" method for structure to graph conversion
:param distance_convertor: (object) object that exposes a "convert" method for distance to expanded vector conversion
:return: keras model object
"""
int32 = 'int32'
if n_feature is None:
x1 = Input(shape=(None, ), dtype=int32) # only z as feature
x1_ = Embedding(n_vocal, embedding_dim)(x1)
else:
x1 = Input(shape=(None, n_feature))
x1_ = x1
x2 = Input(shape=(None, n_connect))
x3 = Input(shape=(None, n_global))
x4 = Input(shape=(None, ), dtype=int32)
x5 = Input(shape=(None, ), dtype=int32)
x6 = Input(shape=(None, ), dtype=int32)
x7 = Input(shape=(None, ), dtype=int32)
if l2_coef is not None:
reg = l2(l2_coef)
else:
reg = None
# two feedforward layers
def ff(x, n_hiddens=[n1, n2]):
out = x
for i in n_hiddens:
out = Dense(i, activation=act, kernel_regularizer=reg)(out)
return out
# a block corresponds to two feedforward layers + one MEGNet layer
# Note the first block does not contain the feedforward layer since
# it will be explicitly added before the block
def one_block(a, b, c, has_ff=True):
if has_ff:
x1_ = ff(a)
x2_ = ff(b)
x3_ = ff(c)
else:
x1_ = a
x2_ = b
x3_ = c
out = MEGNet([n1, n1, n2], [n1, n1, n2], [n1, n1, n2],
pool_method='mean',
activation=act)([x1_, x2_, x3_, x4, x5, x6, x7])
x1_temp = out[0]
x2_temp = out[1]
x3_temp = out[2]
if dropout:
x1_temp = Dropout(dropout)(x1_temp)
x2_temp = Dropout(dropout)(x2_temp)
x3_temp = Dropout(dropout)(x3_temp)
return x1_temp, x2_temp, x3_temp
x1_ = ff(x1_)
x2_ = ff(x2)
x3_ = ff(x3)
for i in range(n_blocks):
if i == 0:
has_ff = False
else:
has_ff = True
x1_1 = x1_
x2_1 = x2_
x3_1 = x3_
x1_1, x2_1, x3_1 = one_block(x1_1, x2_1, x3_1, has_ff)
# skip connection
x1_ = Add()([x1_, x1_1])
x2_ = Add()([x2_, x2_1])
x3_ = Add()([x3_, x3_1])
# set2set for both the atom and bond
node_vec = Set2Set(T=n_pass, n_hidden=n3)([x1_, x6])
edge_vec = Set2Set(T=n_pass, n_hidden=n3)([x2_, x7])
# concatenate atom, bond, and global
final_vec = Concatenate(axis=-1)([node_vec, edge_vec, x3_])
if dropout:
final_vec = Dropout(dropout)(final_vec)
# final dense layers
final_vec = Dense(n2, activation=act)(final_vec)
final_vec = Dense(n3, activation=act)(final_vec)
if is_classification:
final_act = 'sigmoid'
loss = 'binary_crossentropy'
else:
final_act = None
loss = loss
out = Dense(n_target, activation=final_act)(final_vec)
model = Model(inputs=[x1, x2, x3, x4, x5, x6, x7],
outputs=out,
graph_convertor=graph_convertor,
distance_convertor=distance_convertor)
model.compile(Adam(lr), loss)
return model
def create_model():
n_atom_feature = 8
n_bond_feature = 1
n_global_feature = 0
Xavier_init = initializers.VarianceScaling(scale=1.0,
mode='fan_avg',
distribution='normal')
x1 = Input(shape=(None, n_atom_feature)) # atom feature placeholder
x2 = Input(shape=(None, n_bond_feature)) # bond feature placeholder
x3 = Input(shape=(None, n_global_feature)) # global feature placeholder
x4 = Input(shape=(None, ), dtype='int32') # bond index1 placeholder
x5 = Input(shape=(None, ), dtype='int32') # bond index2 placeholder
x6 = Input(shape=(None, ), dtype='int32') # atom_ind placeholder
x7 = Input(shape=(None, ), dtype='int32') # bond_ind placeholder
x8 = Input(shape=(None, ), dtype='int32'
) # Ga_index placeholder to gather Ga nodes to a single tensor
x9 = Input(shape=(None, ), dtype='int32'
) # As_index placeholder to gather As nodes to a single tensor
x10 = Input(
shape=(None, ), dtype='int32'
) # Ga_ind placeholder to sum of the atomic energies of Ga atoms for each single structure
x11 = Input(
shape=(None, ), dtype='int32'
) # As_ind placeholder to sum of the atomic energies of As atoms for each single structure
with kb.name_scope("embedding"):
embed_atom_fea = layers.Dense(16,
name="embedding_atom_fea",
activation='tanh',
kernel_initializer=Xavier_init)
embed_bond_fea = layers.Dense(16,
name="embedding_bond_fea",
activation='tanh',
kernel_initializer=Xavier_init)
GN_1 = MEGNetLayer([16, 16], [16, 16], [1],
pool_method='mean',
activation=softplus2)
GN_2 = MEGNetLayer([16, 16], [16, 16], [1],
pool_method='mean',
activation=softplus2)
x1_ = embed_atom_fea(x1)
x2_ = embed_bond_fea(x2)
out1 = GN_1([x1_, x2_, x3, x4, x5, x6, x7])
x1__ = layers.Add()([x1_, out1[0]])
x2__ = layers.Add()([x2_, out1[1]])
out2 = GN_2([x1__, x2__, out1[2], x4, x5, x6, x7])
x1___ = layers.Add()([x1__, out2[0]])
x2___ = layers.Add()([x2__, out2[1]])
Ga_idx = layers.Lambda(lambda x: tf.reshape(x, (-1, )),
name="Ga_idx")(x8)
As_idx = layers.Lambda(lambda x: tf.reshape(x, (-1, )),
name="As_idx")(x9)
Ga = layers.Lambda(lambda x: tf.gather(x, Ga_idx, axis=1),
name="Ga")(x1___)
As = layers.Lambda(lambda x: tf.gather(x, As_idx, axis=1),
name="As")(x1___)
Ga_grp = layers.Lambda(lambda x: tf.reshape(x, (-1, )),
name="Ga_grp")(x10)
As_grp = layers.Lambda(lambda x: tf.reshape(x, (-1, )),
name="As_grp")(x11)
#node = Set2Set(T=3, n_hidden=3)([x1___, x6])
edge = Set2Set(T=3, n_hidden=3)([x2___, x7])
edge = layers.Lambda(lambda x: kb.sum(x, axis=2, keepdims=True),
name="sum_edge")(edge)
zero = layers.Lambda(lambda x: tf.zeros_like(x),
name="zero_like_edge")(edge)
zero_edge = layers.Multiply(name="zero_edge")([edge, zero])
zero_glob = layers.Multiply(name="zero_glob")([out2[2], zero])
#final = layers.Concatenate(axis=-1)([node, edge, out2[2]])
with kb.name_scope("Ga"):
hidden_Ga1 = layers.Dense(10,
name="hidden_Ga1",
activation='tanh',
kernel_initializer=Xavier_init)
hidden_Ga2 = layers.Dense(10,
name="hidden_Ga2",
activation='tanh',
kernel_initializer=Xavier_init)
output_Ga = layers.Dense(1,
name="output_Ga",
activation=None,
kernel_initializer=Xavier_init)
E_Ga = hidden_Ga1(Ga)
E_Ga = hidden_Ga2(E_Ga)
E_Ga = output_Ga(E_Ga)
E_Ga = layers.Lambda(lambda x: tf.reshape(x, (-1, 1)),
name="reshape_E_Ga")(E_Ga)
sum_Ga = layers.Lambda(lambda x: tf.math.segment_sum(x, Ga_grp),
name="sum_Ga")(E_Ga)
#sum_Ga = layers.Lambda(lambda x: tf.reshape(x, glob.shape), name="reshape_sum_Ga")(sum_Ga)
sum_Ga = layers.Lambda(lambda x: tf.expand_dims(x, axis=0),
name="reshape_sum_Ga")(sum_Ga)
with kb.name_scope("As"):
hidden_As1 = layers.Dense(10,
name="hidden_As1",
activation='tanh',
kernel_initializer=Xavier_init)
hidden_As2 = layers.Dense(10,
name="hidden_As2",
activation='tanh',
kernel_initializer=Xavier_init)
output_As = layers.Dense(1,
name="output_As",
activation=None,
kernel_initializer=Xavier_init)
E_As = hidden_As1(As)
E_As = hidden_As2(E_As)
E_As = output_As(E_As)
E_As = layers.Lambda(lambda x: tf.reshape(x, (-1, 1)),
name="reshape_E_As")(E_As)
sum_As = layers.Lambda(lambda x: tf.math.segment_sum(x, As_grp),
name="sum_As")(E_As)
#sum_As = layers.Lambda(lambda x: tf.reshape(x, glob.shape), name="reshape_sum_As")(sum_As)
sum_As = layers.Lambda(lambda x: tf.expand_dims(x, axis=0),
name="reshape_sum_As")(sum_As)
total_E = layers.Add(name="total_E")([sum_Ga, sum_As])
final_E = layers.Add(name="final_E")([total_E, zero_edge, zero_glob])
# total_E = layers.Lambda(lambda x: tf.expand_dims(x, axis=0))(total_E)
return Model(inputs=[x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11],
outputs=final_E)
