def get_graph(d, i):
    """ get all graphs from a directory
          return [("event<id1>", graph1, size), ...]
    """
    files = os.listdir(d)
    f = files[i]
    return load_graph(d+f)
def get_graphs(d, idxs):
    """ get all graphs from a directory
          return [("event<id1>", graph1, size), ...]
    """
    files = np.array(os.listdir(d))
    files = files[idxs]
    return [load_graph(d+f) for f in files]
Exemplo n.º 3
0
    "{0}/{1}".format(model_dir, model) for model in models if job_name in model
]

# load in test graph paths
graph_dir = "/tigress/jdezoort/IN_samples_large/IN_{0}_{1}/".format(
    prep, pt_cut)
graph_files = []
if (phi_reflect):
    phi_graph_dir = "/tigress/jdezoort/IN_samples_large/IN_{0}_{1}_phi_reflect/".format(
        prep, pt_cut)
    graph_files += os.listdir(phi_graph_dir)[0:800]

graph_files += os.listdir(graph_dir)
print("len(graph_files)={0}".format(len(graph_files)))
test_graphs = np.array([(int(f.split('_')[0].split('t00000')[1]),
                         load_graph(graph_dir + f))
                        for f in graph_files[train_size:]])

# prepare test graphs
size = len(test_graphs)
print("size={0}".format(size))
test_O = [
    Variable(torch.FloatTensor(test_graphs[i][1].X)) for i in range(size)
]
test_Rs = [
    Variable(torch.FloatTensor(test_graphs[i][1].Ro)) for i in range(size)
]
test_Rr = [
    Variable(torch.FloatTensor(test_graphs[i][1].Ri)) for i in range(size)
]
test_Ra = [
def get_graphs(d):
    """ get all graphs from a directory
          return [("event<id1>", graph1, size), ...]
    """
    files = os.listdir(d)
    return [load_graph(d + f) for f in files]
graph_dir = '/scratch/gpfs/jdezoort/hitgraphs_2/{}_{}/'.format(
    method, pt_cut_str[pt_cut])

truth = {}
purities, efficiencies = [], []
sizes, nodes, edges = [], [], []
for i, evtid in enumerate(evt_ids):
    #if (int(evtid.split("00000")[1].split(".")[0]) > 1010): continue

    print("evtid", evtid)
    #if (i == N_avg): break
    print('...', evtid)

    # load in graph
    graph_path = graph_dir + evtid + '_g000.npz'
    graph = load_graph(graph_path)
    X, Ra = graph.X, graph.Ra
    Ri, Ro = graph.Ri, graph.Ro
    y = graph.y
    size = sys.getsizeof(X) + sys.getsizeof(Ra)
    size += sys.getsizeof(Ri) + sys.getsizeof(Ro) + sys.getsizeof(y)

    for j, Ri_row in enumerate(Ri):
        Ri_row = Ri_row[y > 0.5]
        if (np.sum(Ri_row) > 1):
            print("ERREREREREREROR!!!!")

    print("graph.X: {}, graph.Ra: {}, graph.Ri: {}, graph.y: {}".format(
        graph.X.shape, graph.Ra.shape, graph.Ri.shape, graph.y.shape))

    n_edges, n_nodes = Ri.shape[0], Ri.shape[1]