TMVA_SOFIE_GNN.py

Go to the documentation of this file.

## \file
## \ingroup tutorial_ml
## \notebook -nodraw
##
## \macro_code
 
import ROOT
 
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.,
      "nodes": 10*np.random.rand(num_nodes, NODE_FEATURE_SIZE).astype(np.float32)-5.,
      "edges": 10*np.random.rand(num_edges, EDGE_FEATURE_SIZE).astype(np.float32)-5.,
      "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

encoder.Generate()
encoder.OutputGenerated()
 

core.Generate()
core.OutputGenerated()
 

decoder.Generate()
decoder.OutputGenerated()
 

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()