def process(self): #convert the npz into pytorch tensors and save them path = self.processed_dir for idx, raw_path in enumerate(tqdm(self.raw_paths)): g = load_graph(raw_path) Ro = g.Ro[0].T.astype(np.int64) Ri = g.Ri[0].T.astype(np.int64) i_out = Ro[Ro[:, 1].argsort()][:, 0] i_in = Ri[Ri[:, 1].argsort()][:, 0] x = g.X.astype(np.float32) edge_index = np.stack((i_out, i_in)) y = g.y.astype(np.float32) outdata = Data(x=torch.from_numpy(x), edge_index=torch.from_numpy(edge_index), y=torch.from_numpy(y)) if not self._directed and not outdata.is_undirected(): rows, cols = outdata.edge_index temp = torch.stack((cols, rows)) outdata.edge_index = torch.cat([outdata.edge_index, temp], dim=-1) outdata.y = torch.cat([outdata.y, outdata.y]) torch.save(outdata, osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
def generate_input_target(n_graphs, is_train=True): global _evt_id_tr_ global _sec_id_tr_ global _evt_id_te_ global _sec_id_te_ input_graphs = [] target_graphs = [] igraphs = 0 while igraphs < n_graphs: # determine while file to read if is_train: file_name = base_dir.format(evt_ids[_evt_id_tr_], _sec_id_tr_) _sec_id_tr_ += 1 if _sec_id_tr_ == n_sections: _evt_id_tr_ += 1 _sec_id_tr_ = 0 if _evt_id_tr_ >= n_max_evt_id_tr: _evt_id_tr_ = 0 else: file_name = base_dir.format(evt_ids[_evt_id_te_], _sec_id_te_) _sec_id_te_ += 1 if _sec_id_te_ == n_sections: _evt_id_te_ += 1 _sec_id_te_ = 0 if _evt_id_te_ >= n_events: _evt_id_te_ = n_max_evt_id_tr graph = hitsgraph_to_networkx_graph(load_graph(file_name)) input_graph, output_graph = graph_to_input_target(graph) input_graphs.append(input_graph) target_graphs.append(output_graph) igraphs += 1 return input_graphs, target_graphs
def process(self): #convert the npz into pytorch tensors and save them path = self.processed_dir bst = xgb.Booster() bst.load_model('../model_files/xgboost_denoiser_reference.model') bst.set_param({'tree_method': 'hist'}) #force inference on CPU for idx, raw_path in tqdm(enumerate(self.raw_paths), desc='event processed'): g = load_graph(raw_path, sparse=False) x = g.X.astype(np.float32)[:, 3:] pos = g.X.astype(np.float32)[:, :3] y = g.y.astype(np.float32) if (y.sum() != 0): weight_in_event = (y * (1 / y.sum()) + (1 - y) * (1 / (y.shape[0] - y.sum()))).astype( np.float32) else: weight_in_event = np.ones(y.shape[0]).astype(np.float32) dmatrix = xgb.DMatrix(g.X.astype(np.float32)) pred = bst.predict(dmatrix).astype(np.float32) del (dmatrix) outdata = Data(x=torch.from_numpy(x), pos=torch.from_numpy(pos), y=torch.from_numpy(y), weight_in_event=torch.from_numpy(weight_in_event), xgboost_score=torch.from_numpy(pred)) torch.save(outdata, osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
def draw_graph_result(G = None, graph_path = None): if G: X, Ri, Ro, y = G elif graph_path: X, Ri, Ro, y = load_graph(graph_path) else: assert False and "Nothing to draw" draw_single(X, Ri, Ro, y, c_fake = (0,0,0,0.0), xcord1=(0, 'x'), xcord2=(1, 'y'), ycord=(4, 'z'))
def plot_3d(filenames, n_section, n_files): def change_view(az=20): ax.view_init(elev=10, azim=az % 360) return ax fig = plt.figure() ax = fig.add_subplot(111, projection='3d') cmap = plt.get_cmap('bwr_r') theta_counter = np.pi / n_section # start 45 degree rotated for i in range(n_files): print('Plotting file: ' + filenames[i]) X, Ri, Ro, y = load_graph(filenames[i]) X[:, 1] = X[:, 1] * np.pi / n_section theta = (X[:, 1] + theta_counter) % (np.pi * 2) ax.scatter(1000 * X[:, 2], 1000 * X[:, 0] * np.cos(theta), 1000 * X[:, 0] * np.sin(theta), c='k', s=1) feats_o = X[np.where(Ri.T)[1]] feats_i = X[np.where(Ro.T)[1]] x_o = 1000 * feats_o[:, 0] * np.cos(feats_o[:, 1] + theta_counter) x_i = 1000 * feats_i[:, 0] * np.cos(feats_i[:, 1] + theta_counter) y_o = 1000 * feats_o[:, 0] * np.sin(feats_o[:, 1] + theta_counter) y_i = 1000 * feats_i[:, 0] * np.sin(feats_i[:, 1] + theta_counter) z_o = 1000 * feats_o[:, 2] z_i = 1000 * feats_i[:, 2] for j in range(y.shape[0]): seg_args = dict(c='C' + str(i), alpha=float(y[j])) ax.plot([z_o[j], z_i[j]], [x_o[j], x_i[j]], [y_o[j], y_i[j]], '-', **seg_args) if i % 2 == 1: theta_counter += 2 * np.pi / n_section theta_counter = theta_counter % (np.pi * 2) ax.set_xlabel('$Z [mm]$') ax.set_ylabel('$X [mm]$') ax.set_zlabel('$Y [mm]$') ax.grid(b=None) change_view(45) ax.dist = 8 plt.savefig(pdf_dir + 'Cartesian3D.pdf') #Make Gif """
def process(self): #convert the npz into pytorch tensors and save them path = self.processed_dir for idx, raw_path in enumerate(tqdm(self.raw_paths)): g = load_graph(raw_path) x = g.X.astype(np.float32) pos = g.X.astype(np.float32)[:, :3] y = g.y.astype(np.int_) outdata = DataG(x=torch.from_numpy(x), pos=torch.from_numpy(pos), y=torch.from_numpy(y)) torch.save(outdata, osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
def plot_cartesian(filenames, n_section, n_files): fig, ax = plt.subplots(figsize=(8, 8)) cmap = plt.get_cmap('bwr_r') theta_counter = np.pi / n_section # start rotated for i in range(n_files): print('Plotting file: ' + filenames[i * 2]) X, Ri, Ro, y = load_graph(filenames[i * 2]) #print('Zmin: %.2f, Zmax: %.2f' %(min(X[:,2]),max(X[:,2])) ) X[:, 1] = X[:, 1] * np.pi / n_section theta = (X[:, 1] + theta_counter) % (np.pi * 2) ax.scatter(1000 * X[:, 0] * np.cos(theta), 1000 * X[:, 0] * np.sin(theta), c='k') #ax1.scatter(1000*X[:,0]*np.cos(theta), 1000*X[:,2], c='k') feats_o = X[np.where(Ri.T)[1]] feats_i = X[np.where(Ro.T)[1]] x_o = 1000 * feats_o[:, 0] * np.cos(feats_o[:, 1] + theta_counter) x_i = 1000 * feats_i[:, 0] * np.cos(feats_i[:, 1] + theta_counter) y_o = 1000 * feats_o[:, 0] * np.sin(feats_o[:, 1] + theta_counter) y_i = 1000 * feats_i[:, 0] * np.sin(feats_i[:, 1] + theta_counter) for j in range(y.shape[0]): seg_args = dict(c='C' + str(i), alpha=float(y[j] * 2)) ax.plot([x_o[j], x_i[j]], [y_o[j], y_i[j]], '-', **seg_args) theta_counter += 2 * np.pi / n_section theta_counter = theta_counter % (np.pi * 2) ax.set_aspect('equal') ax.set_xlabel('x [mm]') ax.set_ylabel('y [mm]') #ax1.set_xlabel('$x [mm]$') #ax1.set_ylabel('$z [mm]$') #plt.tight_layout() plt.savefig(pdf_dir + 'Cartesian.pdf')
def get_tracks_of_one_sector(self, event_str, iSec, hits, truth): logging.debug("processing {} section".format(iSec)) file_name = os.path.join(self.hitgraph_dir, event_str+"_g{:03d}.npz".format(iSec)) id_name = file_name.replace('.npz', '_ID.npz') G = load_graph(file_name) with np.load(id_name) as f: hit_ids = f['ID'] n_hits = G.X.shape[0] batch_input = [torch.from_numpy(m[None]).float() for m in [G.X, G.Ri, G.Ro]] with torch.no_grad(): weights = self.model(batch_input).flatten().numpy() gnn_tracks = glue.get_tracks(G, weights, hit_ids, hits, truth) \ if self.method=='glue' else \ pathfinder.get_tracks(G, weights, hit_ids, self.weight_cutoff) true_tracks = glue.get_tracks(G, G.y, hit_ids, hits, truth) \ if self.method=='glue' else \ pathfinder.get_tracks(G, G.y, hit_ids, self.weight_cutoff) return gnn_tracks, true_tracks
def __getitem__(self, index): if self.is_concrete_files: return load_graph( os.path.join(self.input_dir, self.df.iloc[index][['filenames']].values[0])) return load_graph(self.filenames[index])
def __getitem__(self, index): return load_graph(self.filenames[index])
#!/usr/bin/env python3 from postprocess.glue import create_glue from postprocess.glue import n_det_layers import numpy as np import pandas as pd import torch import time if __name__ == "__main__": from datasets.graph import load_graph file_name = '/global/cscratch1/sd/xju/heptrkx/data/hitgraphs_001/event000001000_g000.npz' id_name = '/global/cscratch1/sd/xju/heptrkx/data/hitgraphs_002/event000001000_g000_ID.npz' G = load_graph(file_name) with np.load(id_name) as f: hit_ids = f['ID'] from score import load_model from score import load_config model_config_file = 'configs/segclf_small_new.yaml' model = load_model(load_config(model_config_file), reload_epoch=30).eval() batch_input = [ torch.from_numpy(m[None]).float() for m in [G.X, G.Ri, G.Ro] ] with torch.no_grad(): weights = model(batch_input).flatten().numpy() print("precision: ", np.sum(weights * G.y) / G.y.nonzero()[0].shape[0])
def plot_combined(filenames, n_section): print('Plotting file: ' + filenames[1] + ' to: ' + pdf_dir) X, Ri, Ro, y = load_graph(filenames[1]) feats_o = X[np.where(Ri.T)[1]] feats_i = X[np.where(Ro.T)[1]] feats_o[:, 0] = feats_o[:, 0] * 1000 feats_o[:, 1] = feats_o[:, 1] * np.pi / n_section feats_o[:, 2] = feats_o[:, 2] * 1000 feats_i[:, 0] = feats_i[:, 0] * 1000 feats_i[:, 1] = feats_i[:, 1] * np.pi / n_section feats_i[:, 2] = feats_i[:, 2] * 1000 print( 'Plotting: Initial Graph colored in Cylindrical and Cartesian coordinates!' ) fig, ax = plt.subplots(1, 3, figsize=(10, 3), sharey=False, tight_layout=True) cmap = plt.get_cmap('bwr_r') color = {0: 'red', 1: 'blue'} for j in range(y.shape[0]): seg_args = dict(c=color[int(y[j])], alpha=1.) ax[0].plot([feats_o[j, 1], feats_i[j, 1]], [feats_o[j, 0], feats_i[j, 0]], '-', **seg_args) ax[0].scatter((np.pi / 8) * X[:, 1], 1000 * X[:, 0], s=10, c='k') ax[0].set_ylabel('r [mm]') ax[0].set_xlabel(r'$\phi$' + '\n\n (a)') for j in range(y.shape[0]): seg_args = dict(c=color[int(y[j])], alpha=1.) ax[1].plot([feats_o[j, 2], feats_i[j, 2]], [feats_o[j, 0], feats_i[j, 0]], '-', **seg_args) ax[1].scatter(1000 * X[:, 2], 1000 * X[:, 0], s=10, c='k') ax[1].set_ylabel('r [mm]') ax[1].set_xlabel('z [mm]\n\n (b)') # Cartesian X[:, 1] = X[:, 1] * np.pi / n_section theta = X[:, 1] % (np.pi * 2) feats_o = X[np.where(Ri.T)[1]] feats_i = X[np.where(Ro.T)[1]] x_o = 1000 * feats_o[:, 0] * np.cos(feats_o[:, 1]) x_i = 1000 * feats_i[:, 0] * np.cos(feats_i[:, 1]) y_o = 1000 * feats_o[:, 0] * np.sin(feats_o[:, 1]) y_i = 1000 * feats_i[:, 0] * np.sin(feats_i[:, 1]) for j in range(y.shape[0]): seg_args = dict(c=color[int(y[j])], alpha=1.) ax[2].plot([x_o[j], x_i[j]], [y_o[j], y_i[j]], '-', **seg_args) ax[2].scatter(1000 * X[:, 0] * np.cos(theta), 1000 * X[:, 0] * np.sin(theta), s=10, c='k') ax[2].set_xlabel('x [mm]\n\n (c)') ax[2].set_ylabel('y [mm]') plt.savefig(pdf_dir + 'Initial_graph_colored_combined.pdf')
def plot_cylindrical(filenames, n_section): print('Plotting file: ' + filenames[1] + ' to: ' + pdf_dir) X, Ri, Ro, y = load_graph(filenames[1]) feats_o = X[np.where(Ri.T)[1]] feats_i = X[np.where(Ro.T)[1]] feats_o[:, 0] = feats_o[:, 0] * 1000 feats_o[:, 1] = feats_o[:, 1] * np.pi / n_section feats_o[:, 2] = feats_o[:, 2] * 1000 feats_i[:, 0] = feats_i[:, 0] * 1000 feats_i[:, 1] = feats_i[:, 1] * np.pi / n_section feats_i[:, 2] = feats_i[:, 2] * 1000 print('Plotting: After Preprocessing in Cylindrical coordinates!') fig, ax = plt.subplots(1, 2, figsize=(10, 5), sharey=True, tight_layout=True) cmap = plt.get_cmap('bwr_r') ax[0].scatter((np.pi / 8) * X[:, 1], 1000 * X[:, 0], c='k') for j in range(y.shape[0]): seg_args = dict(c='r', alpha=1.) ax[0].plot([feats_o[j, 1], feats_i[j, 1]], [feats_o[j, 0], feats_i[j, 0]], '-', **seg_args) ax[0].set_ylabel('r [mm]') ax[0].set_xlabel(r'$\phi $') ax[1].scatter(1000 * X[:, 2], 1000 * X[:, 0], c='k') for j in range(y.shape[0]): seg_args = dict(c='r', alpha=1.) ax[1].plot([feats_o[j, 2], feats_i[j, 2]], [feats_o[j, 0], feats_i[j, 0]], '-', **seg_args) ax[1].set_ylabel('r [mm]') ax[1].set_xlabel('z [mm]') plt.savefig(pdf_dir + 'Cylindrical_initial_graph.pdf') print('Plotting: Ground Truth in Cylindrical coordinates!') fig, ax = plt.subplots(1, 2, figsize=(10, 5), sharey=True, tight_layout=True) cmap = plt.get_cmap('bwr_r') ax[0].scatter((np.pi / 8) * X[:, 1], 1000 * X[:, 0], c='k') for j in range(y.shape[0]): seg_args = dict(c='r', alpha=float(y[j])) ax[0].plot([feats_o[j, 1], feats_i[j, 1]], [feats_o[j, 0], feats_i[j, 0]], '-', **seg_args) ax[0].set_ylabel('r [mm]') ax[0].set_xlabel(r'$\phi $') ax[1].scatter(1000 * X[:, 2], 1000 * X[:, 0], c='k') for j in range(y.shape[0]): seg_args = dict(c='r', alpha=float(y[j])) ax[1].plot([feats_o[j, 2], feats_i[j, 2]], [feats_o[j, 0], feats_i[j, 0]], '-', **seg_args) ax[1].set_ylabel('r [mm]') ax[1].set_xlabel('z [mm]') plt.savefig(pdf_dir + 'Cylindrical_truth_graphs.pdf') print('Plotting: Initial Graph colored in Cylindrical coordinates!') fig, ax = plt.subplots(1, 2, figsize=(10, 5), sharey=True, tight_layout=True) cmap = plt.get_cmap('bwr_r') color = {0: 'red', 1: 'blue'} ax[0].scatter((np.pi / 8) * X[:, 1], 1000 * X[:, 0], c='k') for j in range(y.shape[0]): seg_args = dict(c=color[int(y[j])], alpha=1.) ax[0].plot([feats_o[j, 1], feats_i[j, 1]], [feats_o[j, 0], feats_i[j, 0]], '-', **seg_args) ax[0].set_ylabel('r [mm]') ax[0].set_xlabel(r'$\phi $') ax[1].scatter(1000 * X[:, 2], 1000 * X[:, 0], c='k') for j in range(y.shape[0]): seg_args = dict(c=color[int(y[j])], alpha=1.) ax[1].plot([feats_o[j, 2], feats_i[j, 2]], [feats_o[j, 0], feats_i[j, 0]], '-', **seg_args) ax[1].set_ylabel('r [mm]') ax[1].set_xlabel('z [mm]') plt.savefig(pdf_dir + 'Cylindrical_initial_graph_colored.pdf')
n_files = 16 * 100 input_dir = os.path.expandvars(input_dir) filenames = sorted([ os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.startswith('event') and f.endswith('.npz') ]) filenames[:n_files] if n_files is not None else filenames n_bins = 20 r_combined = [] p_combined = [] z_combined = [] y_combined = [] for n_file in tqdm(range(n_files)): X, Ri, Ro, y = load_graph(filenames[n_file]) r_combined = np.append(r_combined, X[:, 0]) p_combined = np.append(p_combined, X[:, 1] * np.pi / n_section) z_combined = np.append(z_combined, X[:, 2]) y_combined = np.append(y_combined, y) print('Min r: %.2f, Max r: %.2f' % (np.min(r_combined), np.max(r_combined))) print('Min phi: %.2f, Max phi: %.2f' % (np.min(p_combined), np.max(p_combined))) print('Min z: %.2f, Max z: %.2f' % (np.min(z_combined), np.max(z_combined))) fig, axs = plt.subplots(1, 3, figsize=(10, 2),