10import graph_nets
as gn
11from 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):
40 snt.nets.MLP([LATENT_SIZE]*NUM_LAYERS, activate_final=
True),
46 def __init__(self, name="MLPGraphIndependent"):
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))
61 edge_model_fn=make_mlp_model,
62 node_model_fn=make_mlp_model,
63 global_model_fn=make_mlp_model)
72 name="EncodeProcessDecode"):
79 def __call__(self, input_op, num_processing_steps):
83 for _
in range(num_processing_steps):
85 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")
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"'''
PrintSofie(output, printShape = False):
164 input_data = CopyData(graphData)
168 latent0 = CopyData(input_data)
171 for _
in range(processing_steps):
174 latent = CopyData(core_input)
177 output = CopyData(core_input)
184 data = get_graph_data_dict(num_nodes,num_edges, node_size, edge_size, global_size)
189for i
in range(0,numevts):
196for i
in range(0,numevts):
197 graphData = dataSet[i]
utils_tf.data_dicts_to_graphs_tuple([graphData])
RunGNet(inputGraphData) :
203 output_gn = ep_model(inputGraphData, processing_steps)
ROOT.TH1D("hG","Result from graphnet",20,1,0)
208for i
in range(0,numevts):
214print(
"elapsed time for ",numevts,
"events = ",end-start)
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'])))
229print(
"time to convert data to SOFIE format",endSC-end)
ROOT.TH1D("hS","Result from SOFIE",20,1,0)
235print(
"time to create SOFIE GNN class", start-start0)
236for i
in range(0,numevts):
243print(
"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)
258for 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)