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
9import numpy as np
10import graph_nets as gn
11from graph_nets import utils_tf
12import sonnet as snt
13import time
14
15# defining graph properties
16num_nodes=5
17num_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')
20node_size=4
21edge_size=4
22global_size=1
23LATENT_SIZE = 100
24NUM_LAYERS = 4
25processing_steps = 5
26
27# 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):
29 return {
30 "globals": 10*np.random.rand(GLOBAL_FEATURE_SIZE).astype(np.float32)-5.,
31 "nodes": 10*np.random.rand(num_nodes, NODE_FEATURE_SIZE).astype(np.float32)-5.,
32 "edges": 10*np.random.rand(num_edges, EDGE_FEATURE_SIZE).astype(np.float32)-5.,
33 "senders": snd,
34 "receivers": rec
35 }
36
37# method to instantiate mlp model to be added in GNN
make_mlp_model():
39 return snt.Sequential([
40 snt.nets.MLP([LATENT_SIZE]*NUM_LAYERS, activate_final=True),
41 snt.LayerNorm(axis=-1, create_offset=True, create_scale=True)
42 ])
43
44# defining GraphIndependent class with MLP edge, node, and global models.
45class MLPGraphIndependent(snt.Module):
46 def __init__(self, name="MLPGraphIndependent"):
47 super(MLPGraphIndependent, self).__init__(name=name)
49 edge_model_fn = lambda: snt.nets.MLP([LATENT_SIZE]*NUM_LAYERS, activate_final=True),
50 node_model_fn = lambda: snt.nets.MLP([LATENT_SIZE]*NUM_LAYERS, activate_final=True),
51 global_model_fn = lambda: snt.nets.MLP([LATENT_SIZE]*NUM_LAYERS, activate_final=True))
52
53 def __call__(self, inputs):
54 return self._network(inputs)
55
56# defining Graph network class with MLP edge, node, and global models.
57class MLPGraphNetwork(snt.Module):
58 def __init__(self, name="MLPGraphNetwork"):
59 super(MLPGraphNetwork, self).__init__(name=name)
61 edge_model_fn=make_mlp_model,
62 node_model_fn=make_mlp_model,
63 global_model_fn=make_mlp_model)
64
65 def __call__(self, inputs):
66 return self._network(inputs)
67
68# defining a Encode-Process-Decode module for LHCb toy model
69class EncodeProcessDecode(snt.Module):
70
71 def __init__(self,
72 name="EncodeProcessDecode"):
73 super(EncodeProcessDecode, self).__init__(name=name)
78
79 def __call__(self, input_op, num_processing_steps):
80 latent = self._encoder(input_op)
81 latent0 = latent
82 output_ops = []
83 for _ in range(num_processing_steps):
84 core_input = utils_tf.concat([latent0, latent], axis=1)
85 latent = self._core(core_input)
86 decoded_op = self._decoder(latent)
87 output_ops.append(self._output_transform(decoded_op))
88 return output_ops
89
90
91# Instantiating EncodeProcessDecode Model
EncodeProcessDecode()
93
94# Initializing randomized input data
get_graph_data_dict(num_nodes,num_edges, node_size, edge_size, global_size)
96
97#input_graphs is a tuple representing the initial data
utils_tf.data_dicts_to_graphs_tuple([GraphData])
99
100# Initializing randomized input data for core
101# 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])
104
105#initialize graph data for decoder (input is LATENT_SIZE)
get_graph_data_dict(num_nodes,num_edges, LATENT_SIZE, LATENT_SIZE, LATENT_SIZE)
107
108# Make prediction of GNN
ep_model(input_graph_data, processing_steps)
110#print("---> Input:\n",input_graph_data)
111#print("\n\n------> Input core data:\n",input_core_graph_data)
112#print("\n\n---> Output:\n",output_gn)
113
114# Make SOFIE Model
ep_model._encoder._network, GraphData, filename = "gnn_encoder")
118
ep_model._core._network, CoreGraphData, filename = "gnn_core")
122
ep_model._decoder._network, DecodeGraphData, filename = "gnn_decoder")
126
ep_model._output_transform._network, DecodeGraphData, filename = "gnn_output_transform")
130
131# Compile now the generated C++ code from SOFIE
132gen_code = '''#pragma cling optimize(2)
133#include "gnn_encoder.hxx"
134#include "gnn_core.hxx"
135#include "gnn_decoder.hxx"
136#include "gnn_output_transform.hxx"'''
138
139#helper function to print SOFIE GNN data structure
PrintSofie(output, printShape = False):
144 if (printShape) :
145 print("SOFIE data ... shapes",n.shape,e.shape,g.shape)
146 print(" node data", n.reshape(n.size,))
147 print(" edge data", e.reshape(e.size,))
148 print(" global data",g.reshape(g.size,))
149
CopyData(input_data) :
151 output_data = ROOT.TMVA.Experimental.SOFIE.Copy(input_data)
152 return output_data
153
154# Build SOFIE GNN Model and run inference
162 def infer(self, graphData):
163 # copy the input data
164 input_data = CopyData(graphData)
165
166 # running inference on sofie
167 self.encoder_session.infer(input_data)
168 latent0 = CopyData(input_data)
169 latent = input_data
170 output_ops = []
171 for _ in range(processing_steps):
172 core_input = ROOT.TMVA.Experimental.SOFIE.Concatenate(latent0, latent, axis=1)
173 self.core_session.infer(core_input)
174 latent = CopyData(core_input)
175 self.decoder_session.infer(core_input)
176 self.output_transform_session.infer(core_input)
177 output = CopyData(core_input)
178 output_ops.append(output)
179
180 return output_ops
181
182# Test both GNN on some simulated events
GenerateData():
184 data = get_graph_data_dict(num_nodes,num_edges, node_size, edge_size, global_size)
185 return data
186
187numevts = 40
188dataSet = []
189for i in range(0,numevts):
GenerateData()
191 dataSet.append(data)
192
193# Run graph_nets model
194# First we convert input data to the required input format
195gnetData = []
196for i in range(0,numevts):
197 graphData = dataSet[i]
utils_tf.data_dicts_to_graphs_tuple([graphData])
199 gnetData.append(gnet_data_i)
200
201# Function to run the graph net
RunGNet(inputGraphData) :
203 output_gn = ep_model(inputGraphData, processing_steps)
204 return output_gn
205
time.time()
ROOT.TH1D("hG","Result from graphnet",20,1,0)
208for i in range(0,numevts):
RunGNet(gnetData[i])
globals.numpy()
211 hG.Fill(np.mean(g))
212
time.time()
214print("elapsed time for ",numevts,"events = ",end-start)
215
216# running SOFIE-GNN
217sofieData = []
218for i in range(0,numevts):
219 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'])))
225 sofieData.append(input_data)
226
227
time.time()
229print("time to convert data to SOFIE format",endSC-end)
230
ROOT.TH1D("hS","Result from SOFIE",20,1,0)
time.time()
SofieGNN()
234start = time.time()
235print("time to create SOFIE GNN class", start-start0)
236for i in range(0,numevts):
237 #print("inference event....",i)
238 out = gnn.infer(sofieData[i])
239 g = np.asarray(out[1].global_data)
240 hS.Fill(np.mean(g))
241
242end = time.time()
243print("elapsed time for ",numevts,"events = ",end-start)
244
ROOT.TCanvas()
246c0.Divide(1,2)
c0.cd(1)
248c1.Divide(2,1)
249c1.cd(1)
250hG.Draw()
251c1.cd(2)
252hS.Draw()
253
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)
257#compute differences between SOFIE and GNN
258for i in range(0,numevts):
gnn.infer(sofieData[i])
RunGNet(gnetData[i])
edges.numpy()
np.asarray(outSofie[1].edge_data)
263 if (i == 0) : print(edgesG.shape)
264 for j in range(0,edgesG.shape[0]) :
265 for k in range(0,edgesG.shape[1]) :
266 hDe.Fill(edgesG[j,k]-edgesS[j,k])
267
nodes.numpy()
np.asarray(outSofie[1].node_data)
270 for j in range(0,nodesG.shape[0]) :
271 for k in range(0,nodesG.shape[1]) :
272 hDn.Fill(nodesG[j,k]-nodesS[j,k])
273
globals.numpy()
np.asarray(outSofie[1].global_data)
276 for j in range(0,globG.shape[1]) :
277 hDg.Fill(globG[0,j]-globS[j])
278
279
c0.cd(2)
281c2.Divide(3,1)
282c2.cd(1)
283hDe.Draw()
284c2.cd(2)
285hDn.Draw()
286c2.cd(3)
287hDg.Draw()
288
289c0.Draw()
290
291
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)