TMVA_SOFIE_GNN_Parser.py

Go to the documentation of this file.

## \file
## \ingroup tutorial_ml
## \notebook -nodraw
##
## Tutorial showing how to parse a GNN from GraphNet and make a SOFIE model
## The tutorial also generate some  data which can serve as input for the tutorial TMVA_SOFIE_GNN_Application.C
##
## \macro_code
 
import ROOT
 
import numpy as np
import graph_nets as gn
from graph_nets import utils_tf
import sonnet as snt
#for getting time and memory
import time
import os
import psutil
 
# defining graph properties. Number of edges/modes are the maximum
num_max_nodes=100
num_max_edges=300
node_size=4
edge_size=4
global_size=1
LATENT_SIZE = 100
NUM_LAYERS = 4
processing_steps = 5
numevts = 100
 
verbose = False
 
#print the used memory in MB

    #get memory of current process
    pid = os.getpid()
    python_process = psutil.Process(pid)
    memoryUse = python_process.memory_info()[0]/(1024.*1024.)    #divide by 1024 * 1024 to get memory in MB
    print(s,"memory:",memoryUse,"(MB)")
 
 
# method for returning dictionary of graph data

   num_nodes = np.random.randint(num_max_nodes-2, size=1)[0] + 2
   num_edges = np.random.randint(num_max_edges-1, size=1)[0] + 1
   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": np.random.randint(num_nodes, size=num_edges, dtype=np.int32),
      "receivers": np.random.randint(num_nodes, size=num_edges, dtype=np.int32)
   }
 
# generate graph data with a fixed number of nodes/edges

   return {
      "globals": np.ones((GLOBAL_FEATURE_SIZE),dtype=np.float32),
      "nodes": np.ones((num_nodes, NODE_FEATURE_SIZE), dtype = np.float32),
      "edges": np.ones((num_edges, EDGE_FEATURE_SIZE), dtype = np.float32),
      "senders":  np.random.randint(num_nodes, size=num_edges, dtype=np.int32),
      "receivers": np.random.randint(num_nodes, size=num_edges, dtype=np.int32)
    }
 
 
 
 
# 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
 
printMemory("before instantiating")

printMemory("after instantiating")
 
# Initializing randomized input data with maximum number of nodes/edges

 
#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. This will initialize the GNN with weights
printMemory("before first eval")

printMemory("after first eval")
#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, the model will be made using a maximum number of nodes/edges which are inside GraphData
 

encoder.Generate()
encoder.OutputGenerated()
 

core.Generate()
core.OutputGenerated()
 

decoder.Generate()
decoder.OutputGenerated()
 

output_transform.Generate()
output_transform.OutputGenerated()
 
####################################################################################################################################
 
#generate data and save in a ROOT TTree
#
 


#need to store each element since annot store RTensor
 






 
tree.Branch("node_data", "std::vector<float>" , node_data)
tree.Branch("edge_data", "std::vector<float>" ,  edge_data)
tree.Branch("global_data", "std::vector<float>" ,  global_data)
tree.Branch("receivers", "std::vector<int>" ,  receivers)
tree.Branch("senders", "std::vector<int>" ,  senders)
 
 
print("\n\nSaving data in a ROOT File:")



dataset = []
for i in range(0,numevts):

    s_nodes = graphData['nodes'].size
    s_edges = graphData['edges'].size
    num_edges = graphData['edges'].shape[0]

    node_data.assign(tmp.begin(),tmp.end())
    tmp = ROOT.std.vector['float'](graphData['edges'].reshape((graphData['edges'].size)))
    edge_data.assign(tmp.begin(),tmp.end())
    tmp = ROOT.std.vector['float'](graphData['globals'].reshape((graphData['globals'].size)))
    global_data.assign(tmp.begin(),tmp.end())
    #make sure dtype of graphData['receivers'] and senders is int32
    tmp = ROOT.std.vector['int'](graphData['receivers'])
    receivers.assign(tmp.begin(),tmp.end())
    tmp = ROOT.std.vector['int'](graphData['senders'])
    senders.assign(tmp.begin(),tmp.end())
    if (i < 1 and verbose) :
      print("Nodes - shape:",int(node_data.size()/node_size),node_size,"data: ",node_data)
      print("Edges - shape:",num_edges, edge_size,"data: ", edge_data)
      print("Globals : ",global_data)
      print("Receivers : ",receivers)
      print("Senders   : ",senders)
#
#evaluate graph net on these events
#
    tree.Fill()

    dataset.append(tf_graph_data)
 
tree.Print()
 
#do a first evaluation
printMemory("before eval1")

printMemory("after eval1")
 

firstEvent = True
for tf_graph_data in dataset:
    output_gnn = ep_model(tf_graph_data, processing_steps)



    outgnn[0] = np.mean(output_nodes)
    outgnn[1] = np.mean(output_edges)
    outgnn[2] = np.mean(output_globals)
    h1.Fill(outgnn[0])
    h2.Fill(outgnn[1])
    h3.Fill(outgnn[2])
    if (firstEvent and verbose) :
      print("Output of first event")
      print("nodes data", output_gnn[-1].nodes.numpy())
      print("edge data", output_gnn[-1].edges.numpy())
      print("global data", output_gnn[-1].globals.numpy())
      firstEvent = False
 
 

 
print("time to evaluate events",end-start)
printMemory("after eval Nevts")
 

c1.Divide(1,3)
c1.cd(1)
h1.DrawCopy()
c1.cd(2)
h2.DrawCopy()
c1.cd(3)
h3.DrawCopy()
 
tree.Write()
h1.Write()
h2.Write()
h3.Write()
fileOut.Close()