def forward(self, data):
"""
inputs data:
data[0] - dbscan data
data[1] - primary data
"""
# need to form graph, then pass through GNN
clusts = form_clusters_new(data[0])
# remove track-like particles
#types = get_cluster_label(data[0], clusts)
#selection = types > 1 # 0 or 1 are track-like
#clusts = clusts[selection]
# remove compton clusters
# if no cluster fits this condition, return
selection = filter_compton(clusts) # non-compton looking clusters
if not len(selection):
e = torch.tensor([], requires_grad=True)
if data[0].is_cuda:
e.cuda()
return e
clusts = clusts[selection]
# process group data
# data_grp = process_group_data(data[1], data[0])
# data_grp = data[1]
# form primary/secondary bipartite graph
primaries = assign_primaries(data[1], clusts, data[0])
batch = get_cluster_batch(data[0], clusts)
edge_index = primary_bipartite_incidence(batch, primaries, cuda=True)
# obtain vertex features
x = cluster_vtx_features(data[0], clusts, cuda=True)
# x = cluster_vtx_features_old(data[0], clusts, cuda=True)
#print("max input: ", torch.max(x.view(-1)))
#print("min input: ", torch.min(x.view(-1)))
# obtain edge features
e = cluster_edge_features(data[0], clusts, edge_index, cuda=True)
# go through layers
x = self.attn1(x, edge_index)
#print("max x: ", torch.max(x.view(-1)))
#print("min x: ", torch.min(x.view(-1)))
x = self.attn2(x, edge_index)
#print("max x: ", torch.max(x.view(-1)))
#print("min x: ", torch.min(x.view(-1)))
x = self.attn3(x, edge_index)
#print("max x: ", torch.max(x.view(-1)))
#print("min x: ", torch.min(x.view(-1)))
xbatch = torch.tensor(batch).cuda()
x, e, u = self.edge_predictor(x, edge_index, e, u=None, batch=xbatch)
print("max edge weight: ", torch.max(e.view(-1)))
print("min edge weight: ", torch.min(e.view(-1)))
return e
def find_shower_gnn(dbscan, groups, em_primaries, energy_data, types, model_name, model_checkpoint, gpu_ind=0, verbose=False):
"""
NOTE: THIS IS PROBABLY BROKEN; it was written right after the first pi0 workshop
dbscan: data parsed from "dbscan_label": ["parse_dbscan", "sparse3d_fivetypes"]
groups: data parsed from "group_label": ["parse_cluster3d_clean", "cluster3d_mcst", "sparse3d_fivetypes"]
em_primaries: data parsed from "em_primaries" : ["parse_em_primaries", "sparse3d_data", "particle_mcst"]
energy_data: data parsed from "input_data": ["parse_sparse3d_scn", "sparse3d_data"]
returns a list of length len(em_primaries) containing np arrays, each of which contains the indices corresponding to the voxels in the cone of the corresponding EM primary
"""
event_data = [torch.tensor(dbscan), torch.tensor(em_primaries)]
torch.cuda.set_device(0)
model_attn = DataParallel(BasicAttentionModel(model_name),
device_ids=[0],
dense=False)
model_attn.load_state_dict(torch.load(model_checkpoint, map_location='cuda:'+str(gpu_ind))['state_dict'])
model_attn.eval().cuda()
data_grp = process_group_data(torch.tensor(groups), torch.tensor(dbscan))
clusts = form_clusters_new(dbscan)
selection = filter_compton(clusts) # non-compton looking clusters
clusts = clusts[selection]
full_primaries = np.array(assign_primaries3(em_primaries, clusts, groups))
primaries = assign_primaries(torch.tensor(em_primaries), clusts, torch.tensor(groups))
batch = get_cluster_batch(dbscan, clusts)
edge_index = primary_bipartite_incidence(batch, primaries, cuda=True)
if len(edge_index) == 0: # no secondary clusters
selected_voxels = []
for p in full_primaries.astype(int):
if p == -1:
selected_voxels.append(np.array([]))
else:
selected_voxels.append(clusts[p])
return selected_voxels
n = len(clusts)
mask = np.array([(i not in primaries) for i in range(n)])
others = np.arange(n)[mask]
pred_labels = model_attn(event_data)
pred_nodes = assign_clusters(edge_index, pred_labels, primaries, others, n)
count = 0
selected_voxels = []
for i in range(len(full_primaries)):
p = full_primaries[i]
if p == -1:
selected_voxels.append(np.array([]))
else:
selected_clusts = clusts[np.where(pred_nodes == p)[0]]
selected_voxels.append(np.concatenate(selected_clusts))
return selected_voxels
def forward(self, data):
"""
input data:
data[0] - dbscan data
data[1] - primary data
output data:
dictionary with following keys:
edges : list of edge_index tensors used for edge prediction
edge_pred : list of torch tensors with edge prediction weights
matched : numpy array of group for each cluster (identified by primary index)
n_iter : number of iterations taken
each list is of length k, where k is the number of times the iterative network is applied
"""
# need to form graph, then pass through GNN
clusts = form_clusters_new(data[0])
# remove compton clusters
# if no cluster fits this condition, return
if self.remove_compton:
selection = filter_compton(
clusts, self.compton_thresh) # non-compton looking clusters
if not len(selection):
e = torch.tensor([], requires_grad=True)
if data[0].is_cuda:
e = e.cuda()
return e
clusts = clusts[selection]
#others = np.array([(i not in primaries) for i in range(n)])
batch = get_cluster_batch(data[0], clusts)
# get x batch
xbatch = torch.tensor(batch).cuda()
primaries = assign_primaries(data[1],
clusts,
data[0],
max_dist=self.pmd)
# keep track of who is matched. -1 is not matched
matched = np.repeat(-1, len(clusts))
matched[primaries] = primaries
# print(matched)
edges = []
edge_pred = []
counter = 0
found_match = True
while (-1 in matched) and (counter < self.maxiter) and found_match:
# continue until either:
# 1. everything is matched
# 2. we have exceeded the max number of iterations
# 3. we didn't find any matches
#print('iter ', counter)
counter = counter + 1
# get matched indices
assigned = np.where(matched > -1)[0]
# print(assigned)
others = np.where(matched == -1)[0]
edge_index = primary_bipartite_incidence(batch,
assigned,
cuda=True)
# check if there are any edges to predict
# also batch norm will fail on only 1 edge, so break if this is the case
if edge_index.shape[1] < 2:
counter -= 1
break
# obtain vertex features
x = cluster_vtx_features(data[0], clusts, cuda=True)
# obtain edge features
e = cluster_edge_features(data[0], clusts, edge_index, cuda=True)
# print(x.shape)
# print(torch.max(edge_index))
# print(torch.min(edge_index))
out = self.edge_predictor(x, edge_index, e, xbatch)
# predictions for this edge set.
edge_pred.append(out[0][0])
edges.append(edge_index)
#print(out[0][0].shape)
matched, found_match = self.assign_clusters(
edge_index, out[0][0][:, 1] - out[0][0][:, 0], others, matched,
self.thresh)
# print(edges)
# print(edge_pred)
#print('num iterations: ', counter)
matched = torch.tensor(matched)
counter = torch.tensor([counter])
if data[0].is_cuda:
matched = matched.cuda()
counter = counter.cuda()
return {
'edges': [edges],
'edge_pred': [edge_pred],
'matched': [matched],
'counter': [counter]
}
def forward(self, edge_pred, data0, data1, data2):
"""
edge_pred:
predicted edge weights from model forward
data:
data[0] - 5 types data
data[1] - groups data
data[2] - primary data
"""
data0 = data0[0]
data1 = data1[0]
data2 = data2[0]
# first decide what true edges should be
# need to form graph, then pass through GNN
# clusts = form_clusters(data0)
clusts = form_clusters_new(data0)
# remove track-like particles
# types = get_cluster_label(data0, clusts)
# selection = types > 1 # 0 or 1 are track-like
# clusts = clusts[selection]
# remove compton clusters
# if no cluster fits this condition, return
selection = filter_compton(clusts) # non-compton looking clusters
if not len(selection):
total_loss = self.lossfn(edge_pred, edge_pred)
return {'accuracy': 1., 'loss_seg': total_loss}
clusts = clusts[selection]
# process group data
# data_grp = process_group_data(data1, data0)
data_grp = data1
# form primary/secondary bipartite graph
primaries = assign_primaries(data2, clusts, data0)
batch = get_cluster_batch(data0, clusts)
edge_index = primary_bipartite_incidence(batch, primaries)
group = get_cluster_label(data_grp, clusts)
primaries_true = assign_primaries(data2,
clusts,
data1,
use_labels=True)
print("primaries (est): ", primaries)
print("primaries (true): ", primaries_true)
# determine true assignments
edge_assn = edge_assignment(edge_index, batch, group, cuda=True)
edge_assn = edge_assn.view(-1)
edge_pred = edge_pred.view(-1)
if self.balance:
# weight edges so that 0/1 labels appear equally often
ind0 = edge_assn == 0
ind1 = edge_assn == 1
# number in each class
n0 = torch.sum(ind0).float()
n1 = torch.sum(ind1).float()
print("n0 = ", n0, " n1 = ", n1)
# weights to balance classes
w0 = n1 / (n0 + n1)
w1 = n0 / (n0 + n1)
print("w0 = ", w0, " w1 = ", w1)
edge_assn[ind0] = w0 * edge_assn[ind0]
edge_assn[ind1] = w1 * edge_assn[ind1]
edge_pred = edge_pred.clone()
edge_pred[ind0] = w0 * edge_pred[ind0]
edge_pred[ind1] = w1 * edge_pred[ind1]
total_loss = self.lossfn(edge_pred, edge_assn)
# compute accuracy of assignment
# need to multiply by batch size to be accurate
total_acc = (np.max(batch) + 1) * torch.tensor(
secondary_matching_vox_efficiency(edge_index, edge_assn, edge_pred,
primaries, clusts, len(clusts)))
return {'accuracy': total_acc, 'loss_seg': total_loss}