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),