Logo ROOT  
Reference Guide
 
Loading...
Searching...
No Matches
TMVA_SOFIE_GNN.py
Go to the documentation of this file.
1## \file
2## \ingroup tutorial_ml
3## \notebook -nodraw
4##
5## \macro_code
6
7import ROOT
8
9# Load system openblas library explicitly if available. This avoids pulling in
10# NumPys builtin openblas later, which will conflict with the system openblas.
11ROOT.gInterpreter.Load("libopenblaso.so")
12
13import numpy as np
14import graph_nets as gn
15from graph_nets import utils_tf
16import sonnet as snt
17import time
18
19# defining graph properties
20num_nodes=5
21num_edges=20
np.array([1,2,3,4,2,3,4,3,4,4,0,0,0,0,1,1,1,2,2,3], dtype='int32')
np.array([0,0,0,0,1,1,1,2,2,3,1,2,3,4,2,3,4,3,4,4], dtype='int32')
24node_size=4
25edge_size=4
26global_size=1
27LATENT_SIZE = 100
28NUM_LAYERS = 4
29processing_steps = 5
30
31# method for returning dictionary of graph data
get_graph_data_dict(num_nodes, num_edges, NODE_FEATURE_SIZE=2, EDGE_FEATURE_SIZE=2, GLOBAL_FEATURE_SIZE=1):
33 return {
34 "globals": 10*np.random.rand(GLOBAL_FEATURE_SIZE).astype(np.float32)-5.,
35 "nodes": 10*np.random.rand(num_nodes, NODE_FEATURE_SIZE).astype(np.float32)-5.,
36 "edges": 10*np.random.rand(num_edges, EDGE_FEATURE_SIZE).astype(np.float32)-5.,
37 "senders": snd,
38 "receivers": rec
39 }
40
41# method to instantiate mlp model to be added in GNN
make_mlp_model():
43 return snt.Sequential([
44 snt.nets.MLP([LATENT_SIZE]*NUM_LAYERS, activate_final=True),
45 snt.LayerNorm(axis=-1, create_offset=True, create_scale=True)
46 ])
47
48# defining GraphIndependent class with MLP edge, node, and global models.
49class MLPGraphIndependent(snt.Module):
50 def __init__(self, name="MLPGraphIndependent"):
51 super(MLPGraphIndependent, self).__init__(name=name)
53 edge_model_fn = lambda: snt.nets.MLP([LATENT_SIZE]*NUM_LAYERS, activate_final=True),
54 node_model_fn = lambda: snt.nets.MLP([LATENT_SIZE]*NUM_LAYERS, activate_final=True),
55 global_model_fn = lambda: snt.nets.MLP([LATENT_SIZE]*NUM_LAYERS, activate_final=True))
56
57 def __call__(self, inputs):
58 return self._network(inputs)
59
60# defining Graph network class with MLP edge, node, and global models.
61class MLPGraphNetwork(snt.Module):
62 def __init__(self, name="MLPGraphNetwork"):
63 super(MLPGraphNetwork, self).__init__(name=name)
65 edge_model_fn=make_mlp_model,
66 node_model_fn=make_mlp_model,
67 global_model_fn=make_mlp_model)
68
69 def __call__(self, inputs):
70 return self._network(inputs)
71
72# defining a Encode-Process-Decode module for LHCb toy model
73class EncodeProcessDecode(snt.Module):
74
75 def __init__(self,
76 name="EncodeProcessDecode"):
77 super(EncodeProcessDecode, self).__init__(name=name)
82
83 def __call__(self, input_op, num_processing_steps):
84 latent = self._encoder(input_op)
85 latent0 = latent
86 output_ops = []
87 for _ in range(num_processing_steps):
88 core_input = utils_tf.concat([latent0, latent], axis=1)
89 latent = self._core(core_input)
90 decoded_op = self._decoder(latent)
91 output_ops.append(self._output_transform(decoded_op))
92 return output_ops
93
94
95# Instantiating EncodeProcessDecode Model
EncodeProcessDecode()
97
98# Initializing randomized input data
get_graph_data_dict(num_nodes,num_edges, node_size, edge_size, global_size)
100
101#input_graphs is a tuple representing the initial data
utils_tf.data_dicts_to_graphs_tuple([GraphData])
103
104# Initializing randomized input data for core
105# note that the core network has as input a double number of features
get_graph_data_dict(num_nodes, num_edges, 2*LATENT_SIZE, 2*LATENT_SIZE, 2*LATENT_SIZE)
utils_tf.data_dicts_to_graphs_tuple([CoreGraphData])
108
109#initialize graph data for decoder (input is LATENT_SIZE)
get_graph_data_dict(num_nodes,num_edges, LATENT_SIZE, LATENT_SIZE, LATENT_SIZE)
111
112# Make prediction of GNN
ep_model(input_graph_data, processing_steps)
114#print("---> Input:\n",input_graph_data)
115#print("\n\n------> Input core data:\n",input_core_graph_data)
116#print("\n\n---> Output:\n",output_gn)
117
118# Make SOFIE Model
ep_model._encoder._network, GraphData, filename = "gnn_encoder")
122
ep_model._core._network, CoreGraphData, filename = "gnn_core")
126
ep_model._decoder._network, DecodeGraphData, filename = "gnn_decoder")
130
ep_model._output_transform._network, DecodeGraphData, filename = "gnn_output_transform")
134
135# Compile now the generated C++ code from SOFIE
136gen_code = '''#pragma cling optimize(2)
137#include "gnn_encoder.hxx"
138#include "gnn_core.hxx"
139#include "gnn_decoder.hxx"
140#include "gnn_output_transform.hxx"'''
142
143#helper function to print SOFIE GNN data structure
PrintSofie(output, printShape = False):
148 if (printShape) :
149 print("SOFIE data ... shapes",n.shape,e.shape,g.shape)
150 print(" node data", n.reshape(n.size,))
151 print(" edge data", e.reshape(e.size,))
152 print(" global data",g.reshape(g.size,))
153
CopyData(input_data) :
155 output_data = ROOT.TMVA.Experimental.SOFIE.Copy(input_data)
156 return output_data
157
158# Build SOFIE GNN Model and run inference
166 def infer(self, graphData):
167 # copy the input data
168 input_data = CopyData(graphData)
169
170 # running inference on sofie
171 self.encoder_session.infer(input_data)
172 latent0 = CopyData(input_data)
173 latent = input_data
174 output_ops = []
175 for _ in range(processing_steps):
176 core_input = ROOT.TMVA.Experimental.SOFIE.Concatenate(latent0, latent, axis=1)
177 self.core_session.infer(core_input)
178 latent = CopyData(core_input)
179 self.decoder_session.infer(core_input)
180 self.output_transform_session.infer(core_input)
181 output = CopyData(core_input)
182 output_ops.append(output)
183
184 return output_ops
185
186# Test both GNN on some simulated events
GenerateData():
188 data = get_graph_data_dict(num_nodes,num_edges, node_size, edge_size, global_size)
189 return data
190
191numevts = 40
192dataSet = []
193for i in range(0,numevts):
GenerateData()
195 dataSet.append(data)
196
197# Run graph_nets model
198# First we convert input data to the required input format
199gnetData = []
200for i in range(0,numevts):
201 graphData = dataSet[i]
utils_tf.data_dicts_to_graphs_tuple([graphData])
203 gnetData.append(gnet_data_i)
204
205# Function to run the graph net
RunGNet(inputGraphData) :
207 output_gn = ep_model(inputGraphData, processing_steps)
208 return output_gn
209
time.time()
ROOT.TH1D("hG","Result from graphnet",20,1,0)
212for i in range(0,numevts):
RunGNet(gnetData[i])
globals.numpy()
215 hG.Fill(np.mean(g))
216
time.time()
218print("elapsed time for ",numevts,"events = ",end-start)
219
220# running SOFIE-GNN
221sofieData = []
222for i in range(0,numevts):
223 graphData = dataSet[i]
ROOT.TMVA.Experimental.SOFIE.GNN_Data()
ROOT.TMVA.Experimental.AsRTensor(graphData['nodes'])
ROOT.TMVA.Experimental.AsRTensor(graphData['edges'])
ROOT.TMVA.Experimental.AsRTensor(graphData['globals'])
np.stack((graphData['receivers'],graphData['senders'])))
229 sofieData.append(input_data)
230
231
time.time()
233print("time to convert data to SOFIE format",endSC-end)
234
ROOT.TH1D("hS","Result from SOFIE",20,1,0)
time.time()
SofieGNN()
238start = time.time()
239print("time to create SOFIE GNN class", start-start0)
240for i in range(0,numevts):
241 #print("inference event....",i)
242 out = gnn.infer(sofieData[i])
243 g = np.asarray(out[1].global_data)
244 hS.Fill(np.mean(g))
245
246end = time.time()
247print("elapsed time for ",numevts,"events = ",end-start)
248
ROOT.TCanvas()
250c0.Divide(1,2)
c0.cd(1)
252c1.Divide(2,1)
253c1.cd(1)
254hG.Draw()
255c1.cd(2)
256hS.Draw()
257
ROOT.TH1D("hDe","Difference for edge data",40,1,0)
ROOT.TH1D("hDn","Difference for node data",40,1,0)
ROOT.TH1D("hDg","Difference for global data",40,1,0)
261#compute differences between SOFIE and GNN
262for i in range(0,numevts):
gnn.infer(sofieData[i])
RunGNet(gnetData[i])
edges.numpy()
np.asarray(outSofie[1].edge_data)
267 if (i == 0) : print(edgesG.shape)
268 for j in range(0,edgesG.shape[0]) :
269 for k in range(0,edgesG.shape[1]) :
270 hDe.Fill(edgesG[j,k]-edgesS[j,k])
271
nodes.numpy()
np.asarray(outSofie[1].node_data)
274 for j in range(0,nodesG.shape[0]) :
275 for k in range(0,nodesG.shape[1]) :
276 hDn.Fill(nodesG[j,k]-nodesS[j,k])
277
globals.numpy()
np.asarray(outSofie[1].global_data)
280 for j in range(0,globG.shape[1]) :
281 hDg.Fill(globG[0,j]-globS[j])
282
283
c0.cd(2)
285c2.Divide(3,1)
286c2.cd(1)
287hDe.Draw()
288c2.cd(2)
289hDn.Draw()
290c2.cd(3)
291hDg.Draw()
292
293c0.Draw()
294
295
ROOT::Detail::TRangeCast< T, true > TRangeDynCast
TRangeDynCast is an adapter class that allows the typed iteration through a TCollection.
__call__(self, input_op, num_processing_steps)
__init__(self, name="EncodeProcessDecode")
__init__(self, name="MLPGraphIndependent")
__init__(self, name="MLPGraphNetwork")
CopyData(input_data)
get_graph_data_dict(num_nodes, num_edges, NODE_FEATURE_SIZE=2, EDGE_FEATURE_SIZE=2, GLOBAL_FEATURE_SIZE=1)
RunGNet(inputGraphData)
PrintSofie(output, printShape=False)