TMVA_SOFIE_GNN.py

Go to the documentation of this file.

## \file
## \ingroup tutorial_ml
## \notebook -nodraw
##
## \macro_code
 
import ROOT
 
# Load system openblas library explicitly if available. This avoids pulling in
# NumPys builtin openblas later, which will conflict with the system openblas.
ROOT.gInterpreter.Load("libopenblaso.so")
 
import numpy as np
import graph_nets as gn
from graph_nets import utils_tf
import sonnet as snt
import time
 
# defining graph properties
num_nodes = 5
num_edges = 20


node_size = 4
edge_size = 4
global_size = 1
LATENT_SIZE = 100
NUM_LAYERS = 4
processing_steps = 5
 
 
# method for returning dictionary of graph data

    return {
        "globals": 10 * np.random.rand(GLOBAL_FEATURE_SIZE).astype(np.float32) - 5.0,
        "nodes": 10 * np.random.rand(num_nodes, NODE_FEATURE_SIZE).astype(np.float32) - 5.0,
        "edges": 10 * np.random.rand(num_edges, EDGE_FEATURE_SIZE).astype(np.float32) - 5.0,
        "senders": snd,
        "receivers": rec,
    }
 
 
# method to instantiate mlp model to be added in GNN

    return snt.Sequential(
        [
            snt.nets.MLP([LATENT_SIZE] * NUM_LAYERS, activate_final=True),
            snt.LayerNorm(axis=-1, create_offset=True, create_scale=True),
        ]
    )
 
 
# defining GraphIndependent class with MLP edge, node, and global models.
class MLPGraphIndependent(snt.Module):
    def __init__(self, name="MLPGraphIndependent"):
        super(MLPGraphIndependent, self).__init__(name=name)
        self._network = gn.modules.GraphIndependent(
            edge_model_fn=lambda: snt.nets.MLP([LATENT_SIZE] * NUM_LAYERS, activate_final=True),
            node_model_fn=lambda: snt.nets.MLP([LATENT_SIZE] * NUM_LAYERS, activate_final=True),
            global_model_fn=lambda: snt.nets.MLP([LATENT_SIZE] * NUM_LAYERS, activate_final=True),
        )
 
    def __call__(self, inputs):
        return self._network(inputs)
 
 
# defining Graph network class with MLP edge, node, and global models.
class MLPGraphNetwork(snt.Module):
    def __init__(self, name="MLPGraphNetwork"):
        super(MLPGraphNetwork, self).__init__(name=name)
        self._network = gn.modules.GraphNetwork(
            edge_model_fn=make_mlp_model, node_model_fn=make_mlp_model, global_model_fn=make_mlp_model
        )
 
    def __call__(self, inputs):
        return self._network(inputs)
 
 
# defining a Encode-Process-Decode module for LHCb toy model
class EncodeProcessDecode(snt.Module):
 
    def __init__(self, name="EncodeProcessDecode"):
        super(EncodeProcessDecode, self).__init__(name=name)
        self._encoder = MLPGraphIndependent()
        self._core = MLPGraphNetwork()
        self._decoder = MLPGraphIndependent()
        self._output_transform = MLPGraphIndependent()
 
    def __call__(self, input_op, num_processing_steps):
        latent = self._encoder(input_op)
        latent0 = latent
        output_ops = []
        for _ in range(num_processing_steps):
            core_input = utils_tf.concat([latent0, latent], axis=1)
            latent = self._core(core_input)
            decoded_op = self._decoder(latent)
            output_ops.append(self._output_transform(decoded_op))
        return output_ops
 
 
# Instantiating EncodeProcessDecode Model

 
# Initializing randomized input data

 
# input_graphs  is a tuple representing the initial data

 
# Initializing randomized input data for core
# note that the core network has as input a double number of features


 
# initialize graph data for decoder (input is LATENT_SIZE)

 
# Make prediction of GNN

# print("---> Input:\n",input_graph_data)
# print("\n\n------> Input core data:\n",input_core_graph_data)
# print("\n\n---> Output:\n",output_gn)
 
# Make SOFIE Model

    ep_model._encoder._network, GraphData, filename="gnn_encoder"
)
encoder.Generate()
encoder.OutputGenerated()
 

    ep_model._core._network, CoreGraphData, filename="gnn_core"
)
core.Generate()
core.OutputGenerated()
 

    ep_model._decoder._network, DecodeGraphData, filename="gnn_decoder"
)
decoder.Generate()
decoder.OutputGenerated()
 

    ep_model._output_transform._network, DecodeGraphData, filename="gnn_output_transform"
)
output_transform.Generate()
output_transform.OutputGenerated()
 
# Compile now the generated C++ code from SOFIE
gen_code = '''#pragma cling optimize(2)
#include "gnn_encoder.hxx"
#include "gnn_core.hxx"
#include "gnn_decoder.hxx"
#include "gnn_output_transform.hxx"'''
ROOT.gInterpreter.Declare(gen_code)
 
 
# helper function to print SOFIE GNN data structure

    n = np.asarray(output.node_data)
    e = np.asarray(output.edge_data)
    g = np.asarray(output.global_data)
    if printShape:
        print("SOFIE data ... shapes", n.shape, e.shape, g.shape)
    print(
        " node data",
        n.reshape(
            n.size,
        ),
    )
    print(
        " edge data",
        e.reshape(
            e.size,
        ),
    )
    print(
        " global data",
        g.reshape(
            g.size,
        ),
    )
 
 

    output_data = ROOT.TMVA.Experimental.SOFIE.Copy(input_data)
    return output_data
 
 
# Build  SOFIE GNN Model and run inference
class SofieGNN:
    def __init__(self):
        self.encoder_session = ROOT.TMVA_SOFIE_gnn_encoder.Session()
        self.core_session = ROOT.TMVA_SOFIE_gnn_core.Session()
        self.decoder_session = ROOT.TMVA_SOFIE_gnn_decoder.Session()
        self.output_transform_session = ROOT.TMVA_SOFIE_gnn_output_transform.Session()
 
    def infer(self, graphData):
        # copy the input data
        input_data = CopyData(graphData)
 
        # running inference on sofie
        self.encoder_session.infer(input_data)
        latent0 = CopyData(input_data)
        latent = input_data
        output_ops = []
        for _ in range(processing_steps):
            core_input = ROOT.TMVA.Experimental.SOFIE.Concatenate(latent0, latent, axis=1)
            self.core_session.infer(core_input)
            latent = CopyData(core_input)
            self.decoder_session.infer(core_input)
            self.output_transform_session.infer(core_input)
            output = CopyData(core_input)
            output_ops.append(output)
 
        return output_ops
 
 
# Test both GNN on some simulated events

    data = get_graph_data_dict(num_nodes, num_edges, node_size, edge_size, global_size)
    return data
 
 
numevts = 40
dataSet = []
for i in range(0, numevts):

    dataSet.append(data)
 
# Run graph_nets model
# First we convert input data to the required input format
gnetData = []
for i in range(0, numevts):
    graphData = dataSet[i]

    gnetData.append(gnet_data_i)
 
 
# Function to run the graph net

    output_gn = ep_model(inputGraphData, processing_steps)
    return output_gn
 
 


for i in range(0, numevts):


    hG.Fill(np.mean(g))
 

print("elapsed time for ", numevts, "events = ", end - start)
 
# running SOFIE-GNN
sofieData = []
for i in range(0, numevts):
    graphData = dataSet[i]





    sofieData.append(input_data)
 
 

print("time to convert data to SOFIE format", endSC - end)
 



start = time.time()
print("time to create SOFIE GNN class", start - start0)
for i in range(0, numevts):
    # print("inference event....",i)
    out = gnn.infer(sofieData[i])
    g = np.asarray(out[1].global_data)
    hS.Fill(np.mean(g))
 
end = time.time()
print("elapsed time for ", numevts, "events = ", end - start)
 

c0.Divide(1, 2)

c1.Divide(2, 1)
c1.cd(1)
hG.Draw()
c1.cd(2)
hS.Draw()
 



# compute differences between SOFIE and GNN
for i in range(0, numevts):




    if i == 0:
        print(edgesG.shape)
    for j in range(0, edgesG.shape[0]):
        for k in range(0, edgesG.shape[1]):
            hDe.Fill(edgesG[j, k] - edgesS[j, k])
 


    for j in range(0, nodesG.shape[0]):
        for k in range(0, nodesG.shape[1]):
            hDn.Fill(nodesG[j, k] - nodesS[j, k])
 


    for j in range(0, globG.shape[1]):
        hDg.Fill(globG[0, j] - globS[j])
 
 

c2.Divide(3, 1)
c2.cd(1)
hDe.Draw()
c2.cd(2)
hDn.Draw()
c2.cd(3)
hDg.Draw()
 
c0.Draw()