Skip to main content

No project description provided

Project description

An efficient tensorflow 2 implementation of the edge-convolution layer EdgeConv used in e.g. ParticleNet.

The structure of the layer is as described in ‘ParticleNet: Jet Tagging via Particle Clouds’


Install via:

pip install medgeconv

Use e.g. like this:

import medgeconv

nodes = medgeconv.DisjointEdgeConvBlock(
    units=[64, 64, 64],
)((nodes, coordinates))

Inputs to EdgeConv are 2 ragged tensors: nodes and coordinates

  • nodes, shape (batchsize, None, n_features)

    Node features of the graph. Secound dimension is the number of nodes, which can vary from graph to graph.

  • coordinates, shape (batchsize, None, n_coords)

    Features of each node used for calculating nearest neighbors.

Example: Input for a graph with 2 features per node, and all node features used as coordinates.

import tensorflow as tf

nodes = tf.ragged.constant([
   # graph 1: 2 nodes
   [[2., 4.],
    [2., 6.]],
   # graph 2: 4 nodes
   [[0., 2.],
    [3., 7.],
    [4., 0.],
    [1., 2.]],
], ragged_rank=1)

print(nodes.shape)  # output: (2, None, 2)

# using all node features as coordinates
coordinates = nodes


The full ParticleNet for n_features = n_coords = 2, and a dense layer with 2 neurons as the output can be built like this:

import tensorflow as tf
import medgeconv

inp = (
    tf.keras.Input((None, 2), ragged=True),
    tf.keras.Input((None, 2), ragged=True),
x = medgeconv.DisjointEdgeConvBlock(
    units=[64, 64, 64],

x = medgeconv.DisjointEdgeConvBlock(
    units=[128, 128, 128],

x = medgeconv.DisjointEdgeConvBlock(
    units=[256, 256, 256],

output = tf.keras.layers.Dense(2)(x)
model = tf.keras.Model(inp, output)

The last EdgeConv layer has pooling = True. This will attach a node-wise global average pooling layer in the end, producing normal not-ragged tensors again.

The model can then be used on ragged Tensors:

nodes = tf.RaggedTensor.from_tensor(tf.ones((3, 17, 2)))
model.predict((nodes, nodes))

Loading models

To load models, use the custom_objects:

import medgeconv

model = load_model(path, custom_objects=medgeconv.custom_objects)

knn_graph kernel

This package includes a cuda kernel for calculating the k nearest neighbors on a batch of graphs. It comes with a precompiled kernel for the version of tensorflow specified in requirements.txt.

To compile it locally, e.g. for a different version of tensorflow, go to medgeconv/tf_ops and adjust the bash script. Running it will download the specified tf dev docker image and produce the file medgeconv/tf_ops/python/ops/


Results using this model architecture in the context of particle physics were presented at the ICRC 2021 conference , as well as the VLVnT 2021 .

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

medgeconv-2.2.tar.gz (58.7 kB view hashes)

Uploaded source

Built Distribution

medgeconv-2.2-py3-none-any.whl (60.0 kB view hashes)

Uploaded py3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page