megnet.layers package

Module contents

class MEGNetLayer(units_v, units_e, units_u, pool_method='mean', activation=None, use_bias=True, kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None, **kwargs)[source]

Bases: megnet.layers.graph.base.GraphNetworkLayer

The MEGNet graph implementation as described in the paper

Chen, Chi; Ye, Weike Ye; Zuo, Yunxing; Zheng, Chen; Ong, Shyue Ping. Graph Networks as a Universal Machine Learning Framework for Molecules and Crystals, 2018, arXiv preprint. [arXiv:1812.05055](https://arxiv.org/abs/1812.05055)

Parameters
  • units_v (list of integers) – the hidden layer sizes for node update neural network

  • units_e (list of integers) – the hidden layer sizes for edge update neural network

  • units_u (list of integers) – the hidden layer sizes for state update neural network

  • pool_method (str) – ‘mean’ or ‘sum’, determines how information is gathered to nodes from neighboring edges

  • 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

call(inputs, mask=None)

the logic of the layer, returns the final graph

compute_output_shape(input_shape)[source]

compute static output shapes, returns list of tuple shapes

build(input_shape)[source]

initialize the weights and biases for each function

phi_e(inputs)[source]

update function for bonds and returns updated bond attribute e_p

rho_e_v(e_p, inputs)[source]

aggregate updated bonds e_p to per atom attributes, b_e_p

phi_v(b_e_p, inputs)[source]

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)[source]

aggregate bonds to global attribute

rho_v_u(v_p, inputs)[source]

aggregate atom to global attributes

get_config()[source]

part of keras interface for serialization

build(input_shapes)[source]

Creates the variables of the layer (optional, for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call.

This is typically used to create the weights of Layer subclasses.

Parameters

input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

If the layer has not been built, this method will call build on the layer. This assumes that the layer will later be used with inputs that match the input shape provided here.

Parameters

input_shape – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.

Returns

An input shape tuple.

get_config()[source]

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

phi_e(inputs)[source]

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, b_v_p, inputs)[source]

u’ = phi_u(ar e’, ar 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, inputs)[source]

Step 3. Compute updated node attributes v_i’ = phi_v(ar 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, inputs)[source]

let V’ = {v’} i = 1:Nv let E’ = {(e_k’, rk, sk)} k = 1:Ne ar 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, inputs)[source]

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, inputs)[source]

ar 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

class CrystalGraphLayer(activation='relu', use_bias=True, kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None, **kwargs)[source]

Bases: megnet.layers.graph.base.GraphNetworkLayer

The CGCNN graph implementation as described in the paper

Xie et al. PHYSICAL REVIEW LETTERS 120, 145301 (2018)

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

call(inputs, mask=None)

the logic of the layer, returns the final graph

compute_output_shape(input_shape)[source]

compute static output shapes, returns list of tuple shapes

build(input_shape)[source]

initialize the weights and biases for each function

phi_e(inputs)[source]

update function for bonds and returns updated bond attribute e_p

rho_e_v(e_p, inputs)[source]

aggregate updated bonds e_p to per atom attributes, b_e_p

phi_v(b_e_p, inputs)[source]

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)[source]

aggregate bonds to global attribute

rho_v_u(v_p, inputs)[source]

aggregate atom to global attributes

get_config()[source]

part of keras interface for serialization

build(input_shapes)[source]

Creates the variables of the layer (optional, for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call.

This is typically used to create the weights of Layer subclasses.

Parameters

input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

If the layer has not been built, this method will call build on the layer. This assumes that the layer will later be used with inputs that match the input shape provided here.

Parameters

input_shape – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.

Returns

An input shape tuple.

get_config()[source]

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

phi_e(inputs)[source]

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, b_v_p, inputs)[source]

u’ = phi_u(ar e’, ar 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, inputs)[source]

Step 3. Compute updated node attributes v_i’ = phi_v(ar 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, inputs)[source]

let V’ = {v’} i = 1:Nv let E’ = {(e_k’, rk, sk)} k = 1:Ne ar 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, inputs)[source]

Reduce edge attributes to node attribute, eqn 5 in the paper :param e_p: updated bond :param inputs: the whole input list

Returns: summed tensor

rho_v_u(v_p, inputs)[source]

ar 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

class InteractionLayer(activation=<function softplus2>, use_bias=True, kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None, **kwargs)[source]

Bases: megnet.layers.graph.base.GraphNetworkLayer

The Continuous filter InteractionLayer in Schnet

Schütt et al. SchNet: A continuous-filter convolutional neural network for modeling quantum interactions

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

call(inputs, mask=None)

the logic of the layer, returns the final graph

compute_output_shape(input_shape)[source]

compute static output shapes, returns list of tuple shapes

build(input_shape)[source]

initialize the weights and biases for each function

phi_e(inputs)[source]

update function for bonds and returns updated bond attribute e_p

rho_e_v(e_p, inputs)[source]

aggregate updated bonds e_p to per atom attributes, b_e_p

phi_v(b_e_p, inputs)[source]

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)[source]

aggregate bonds to global attribute

rho_v_u(v_p, inputs)[source]

aggregate atom to global attributes

get_config()[source]

part of keras interface for serialization

build(input_shapes)[source]

Creates the variables of the layer (optional, for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call.

This is typically used to create the weights of Layer subclasses.

Parameters

input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

If the layer has not been built, this method will call build on the layer. This assumes that the layer will later be used with inputs that match the input shape provided here.

Parameters

input_shape – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.

Returns

An input shape tuple.

get_config()[source]

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

phi_e(inputs)[source]

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, b_v_p, inputs)[source]

u’ = phi_u(ar e’, ar 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, inputs)[source]

Step 3. Compute updated node attributes v_i’ = phi_v(ar 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, inputs)[source]

let V’ = {v’} i = 1:Nv let E’ = {(e_k’, rk, sk)} k = 1:Ne ar 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, inputs)[source]

Reduce edge attributes to node attribute, eqn 5 in the paper :param e_p: updated bond :param inputs: the whole input list

Returns: summed tensor

rho_v_u(v_p, inputs)[source]

ar 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

class Set2Set(T=3, n_hidden=512, activation=None, activation_lstm='tanh', recurrent_activation='hard_sigmoid', kernel_initializer='glorot_uniform', recurrent_initializer='orthogonal', bias_initializer='zeros', use_bias=True, unit_forget_bias=True, kernel_regularizer=None, recurrent_regularizer=None, bias_regularizer=None, kernel_constraint=None, recurrent_constraint=None, bias_constraint=None, **kwargs)[source]

Bases: tensorflow.python.keras.engine.base_layer.Layer

For a set of vectors, the set2set neural network maps it to a single vector. The order invariance is acheived by a attention mechanism. See Vinyals, Oriol, Samy Bengio, and Manjunath Kudlur. “Order matters: Sequence to sequence for sets.” arXiv preprint arXiv:1511.06391 (2015).

Parameters
  • T – (int) recurrent step

  • n_hidden – (int) number of hidden units

  • activation – (str or object) activation function

  • activation_lstm – (str or object) activation function for lstm

  • recurrent_activation – (str or object) activation function for recurrent step

  • kernel_initializer – (str or object) initializer for kernel weights

  • recurrent_initializer – (str or object) initializer for recurrent weights

  • bias_initializer – (str or object) initializer for biases

  • use_bias – (bool) whether to use biases

  • unit_forget_bias – (bool) whether to use basis in forget gate

  • kernel_regularizer – (str or object) regularizer for kernel weights

  • recurrent_regularizer – (str or object) regularizer for recurrent weights

  • bias_regularizer – (str or object) regularizer for biases

  • kernel_constraint – (str or object) constraint for kernel weights

  • recurrent_constraint – (str or object) constraint for recurrent weights

  • bias_constraint – (str or object) constraint for biases

  • kwargs – other inputs for keras Layer class

build(input_shape)[source]

Creates the variables of the layer (optional, for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call.

This is typically used to create the weights of Layer subclasses.

Parameters

input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

call(inputs, mask=None)[source]

This is where the layer’s logic lives.

Parameters
  • inputs – Input tensor, or list/tuple of input tensors.

  • **kwargs – Additional keyword arguments.

Returns

A tensor or list/tuple of tensors.

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

If the layer has not been built, this method will call build on the layer. This assumes that the layer will later be used with inputs that match the input shape provided here.

Parameters

input_shape – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.

Returns

An input shape tuple.

get_config()[source]

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Returns

Python dictionary.

class LinearWithIndex(mode='mean', **kwargs)[source]

Bases: tensorflow.python.keras.engine.base_layer.Layer

Sum or average the node/edge attributes to get a structure-level vector

Parameters

mode – (str) ‘mean’ or ‘sum’

build(input_shape)[source]

Creates the variables of the layer (optional, for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call.

This is typically used to create the weights of Layer subclasses.

Parameters

input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

call(inputs, mask=None)[source]

This is where the layer’s logic lives.

Parameters
  • inputs – Input tensor, or list/tuple of input tensors.

  • **kwargs – Additional keyword arguments.

Returns

A tensor or list/tuple of tensors.

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

If the layer has not been built, this method will call build on the layer. This assumes that the layer will later be used with inputs that match the input shape provided here.

Parameters

input_shape – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.

Returns

An input shape tuple.

get_config()[source]

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Returns

Python dictionary.

keras_layer_deserialize(config, custom_objects=None)

Instantiates a layer from a config dictionary.

Parameters
  • config – dict of the form {‘class_name’: str, ‘config’: dict}

  • custom_objects – dict mapping class names (or function names) of custom (non-Keras) objects to class/functions

Returns

Layer instance (may be Model, Sequential, Network, Layer…)

mean_squared_error_with_scale(y_true, y_pred, scale=10000)[source]

Keras default log for tracking progress shows two decimal points, here we multiply the mse by a factor to fully show the loss in progress bar

Parameters
  • y_true – (tensor) training y

  • y_pred – (tensor) predicted y

  • scale – (int or float) factor to multiply with mse

Returns

scaled mse (float)

softplus2(x)[source]

out = log(exp(x)+1) - log(2) softplus function that is 0 at x=0, the implementation aims at avoiding overflow

Parameters

x – (Tensor) input tensor

Returns

(Tensor) output tensor