14import graph_nets
as gn
15from graph_nets
import utils_tf
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')
get_graph_data_dict(num_nodes, num_edges, NODE_FEATURE_SIZE=2, EDGE_FEATURE_SIZE=2, GLOBAL_FEATURE_SIZE=1):
44 snt.nets.MLP([LATENT_SIZE]*NUM_LAYERS, activate_final=
True),
50 def __init__(self, name="MLPGraphIndependent"):
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))
65 edge_model_fn=make_mlp_model,
66 node_model_fn=make_mlp_model,
67 global_model_fn=make_mlp_model)
76 name="EncodeProcessDecode"):
83 def __call__(self, input_op, num_processing_steps):
87 for _
in range(num_processing_steps):
89 latent = self.
_core(core_input)
get_graph_data_dict(num_nodes,num_edges, node_size, edge_size, global_size)
utils_tf.data_dicts_to_graphs_tuple([GraphData])
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])
get_graph_data_dict(num_nodes,num_edges, LATENT_SIZE, LATENT_SIZE, LATENT_SIZE)
ep_model(input_graph_data, processing_steps)
ep_model._encoder._network, GraphData, filename = "gnn_encoder")
ep_model._core._network, CoreGraphData, filename = "gnn_core")
ep_model._decoder._network, DecodeGraphData, filename = "gnn_decoder")
ep_model._output_transform._network, DecodeGraphData, filename = "gnn_output_transform")
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"'''
PrintSofie(output, printShape = False):
168 input_data = CopyData(graphData)
172 latent0 = CopyData(input_data)
175 for _
in range(processing_steps):
178 latent = CopyData(core_input)
181 output = CopyData(core_input)
188 data = get_graph_data_dict(num_nodes,num_edges, node_size, edge_size, global_size)
193for i
in range(0,numevts):
200for i
in range(0,numevts):
201 graphData = dataSet[i]
utils_tf.data_dicts_to_graphs_tuple([graphData])
RunGNet(inputGraphData) :
207 output_gn = ep_model(inputGraphData, processing_steps)
ROOT.TH1D("hG","Result from graphnet",20,1,0)
212for i
in range(0,numevts):
218print(
"elapsed time for ",numevts,
"events = ",end-start)
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'])))
233print(
"time to convert data to SOFIE format",endSC-end)
ROOT.TH1D("hS","Result from SOFIE",20,1,0)
239print(
"time to create SOFIE GNN class", start-start0)
240for i
in range(0,numevts):
247print(
"elapsed time for ",numevts,
"events = ",end-start)
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)
262for i
in range(0,numevts):
np.asarray(outSofie[1].edge_data)
np.asarray(outSofie[1].node_data)
np.asarray(outSofie[1].global_data)
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")
get_graph_data_dict(num_nodes, num_edges, NODE_FEATURE_SIZE=2, EDGE_FEATURE_SIZE=2, GLOBAL_FEATURE_SIZE=1)
PrintSofie(output, printShape=False)