megnet.layers.graph.base module

A full GN block has the following computation steps

1. Compute updated edge attributes

2. Aggregate edge attributes per node

3. Compute updated node attributes

4. Aggregate edge attributes globally

5. Aggregate node attributes globally

6. Compute updated global attribute

[1] https://arxiv.org/pdf/1806.01261.pdf

class GraphNetworkLayer(*args, **kwargs)

Bases: keras.engine.base_layer.Layer

Implementation of a graph network layer. Current implementation is based on neural networks for each update function, and sum or mean for each aggregation function

Method:

call(inputs, mask=None): the logic of the layer, returns the final graph compute_output_shape(input_shape): compute static output shapes, returns list of tuple shapes build(input_shape): initialize the weights and biases for each function phi_e(inputs): update function for bonds and returns updated bond attribute e_p rho_e_v(e_p, inputs): aggregate updated bonds e_p to per atom attributes, b_e_p phi_v(b_e_p, inputs): update the atom attributes by the results from previous step b_e_p and all the inputs

returns v_p.

rho_e_u(e_p, inputs): aggregate bonds to global attribute rho_v_u(v_p, inputs): aggregate atom to global attributes get_config(): part of keras interface for serialization

Parameters

• activation (str) – Default: None. The activation function used for each sub-neural network. Examples include ‘relu’, ‘softmax’, ‘tanh’, ‘sigmoid’ and etc.

• use_bias (bool) – Default: True. Whether to use the bias term in the neural network.

• kernel_initializer (str) – Default: ‘glorot_uniform’. Initialization function for the layer kernel weights,

• bias_initializer (str) – Default: ‘zeros’

• activity_regularizer (str) – Default: None. The regularization function for the output

• kernel_constraint (str) – Default: None. Keras constraint for kernel values

• bias_constraint (str) – Default: None .Keras constraint for bias values

• **kwargs –

call(inputs: Sequence, mask=None) → Sequence

Core logic of graph network :param inputs: input tensors :type inputs: Sequence :param mask: mask tensor :type mask: tensor

Returns: output tensor

get_config() → Dict

Part of keras layer interface, where the signature is converted into a dict :returns: configurational dictionary

phi_e(inputs: Sequence) → tensorflow.python.framework.ops.Tensor

This is for updating the edge attributes ek’ = phi_e(ek, vrk, vsk, u)

Parameters

inputs (Sequence) – list or tuple for the graph inputs

Returns

updated edge/bond attributes

phi_u(b_e_p: tensorflow.python.framework.ops.Tensor, b_v_p: tensorflow.python.framework.ops.Tensor, inputs: Sequence) → tensorflow.python.framework.ops.Tensor

u’ = phi_u(bar e’, bar v’, u) :param b_e_p: edge/bond to global aggregated tensor :type b_e_p: tf.Tensor :param b_v_p: node/atom to global aggregated tensor :type b_v_p: tf.Tensor :param inputs: list or tuple for the graph inputs :type inputs: Sequence

Returns

updated globa/state attributes

phi_v(b_ei_p: tensorflow.python.framework.ops.Tensor, inputs: Sequence)

Step 3. Compute updated node attributes v_i’ = phi_v(bar e_i, vi, u)

Parameters

• b_ei_p (tf.Tensor) – edge-to-node aggregated tensor

• inputs (Sequence) – list or tuple for the graph inputs

Returns

updated node/atom attributes

rho_e_u(e_p: tensorflow.python.framework.ops.Tensor, inputs: Sequence) → tensorflow.python.framework.ops.Tensor

let V’ = {v’} i = 1:Nv let E’ = {(e_k’, rk, sk)} k = 1:Ne bar e’ = rho_e_u(E’)

Parameters

• e_p (tf.Tensor) – updated edge/bond attributes

• inputs (Sequence) – list or tuple for the graph inputs

Returns

edge/bond to global/state aggregated tensor

rho_e_v(e_p: tensorflow.python.framework.ops.Tensor, inputs: Sequence) → tensorflow.python.framework.ops.Tensor

This is for step 2, aggregate edge attributes per node Ei’ = {(ek’, rk, sk)} with rk =i, k=1:Ne

Parameters

• e_p (tf.Tensor) – the updated edge attributes

• inputs (Sequence) – list or tuple for the graph inputs

Returns

edge/bond to node/atom aggregated tensor

rho_v_u(v_p: tensorflow.python.framework.ops.Tensor, inputs: Sequence) → tensorflow.python.framework.ops.Tensor

bar v’ = rho_v_u(V’)

Parameters

• v_p (tf.Tensor) – updated atom/node attributes

• inputs (Sequence) – list or tuple for the graph inputs

Returns

atom/node to global/state aggregated tensor