def __init__(self, k, feature_dims, emb_dims, output_classes, input_dims=3,
dropout_prob=0.5):
super(DgcnnModel, self).__init__()
self.nng = KNNGraph(k)
self.conv = nn.ModuleList()
self.num_layers = len(feature_dims)
for i in range(self.num_layers):
self.conv.append(EdgeConv(
feature_dims[i - 1] if i > 0 else input_dims,
feature_dims[i],
batch_norm=True))
self.proj = nn.Linear(sum(feature_dims), emb_dims[0])
self.embs = nn.ModuleList()
self.bn_embs = nn.ModuleList()
self.dropouts = nn.ModuleList()
self.num_embs = len(emb_dims) - 1
for i in range(1, self.num_embs + 1):
self.embs.append(nn.Linear(
# * 2 because of concatenation of max- and mean-pooling
emb_dims[i - 1] if i > 1 else (emb_dims[i - 1] * 2),
emb_dims[i]))
# self.bn_embs.append(nn.BatchNorm1d(emb_dims[i]))
self.dropouts.append(nn.Dropout(dropout_prob))
self.proj_output = nn.Linear(emb_dims[-1], output_classes)
def __init__(self, k = 10, feature_dims = [64, 64, 128, 256], emb_dims = [512, 512, 256],
input_dims=3, output_dims=1024):
super(GNN, self).__init__()
self.nng = KNNGraph(k)
self.conv = nn.ModuleList()
self.num_layers = len(feature_dims)
for i in range(self.num_layers):
self.conv.append(EdgeConv(
feature_dims[i - 1] if i > 0 else input_dims,
feature_dims[i],
batch_norm=True))
self.proj = nn.Linear(sum(feature_dims), emb_dims[0])
def __init__(self, k, in_dim: int, emb_dims: list, out_dim: int):
super(DescripNet, self).__init__()
self.knng = KNNGraph(k)
self.conv = nn.ModuleList()
self.feat_nn = nn.Sequential(nn.Linear(emb_dims[-2], emb_dims[-1]),
nn.ReLU())
self.gate_nn = nn.Sequential(nn.Linear(emb_dims[-2], 1), nn.ReLU())
self.global_attention_pooling = GlobalAttentionPooling(
gate_nn=self.gate_nn, feat_nn=self.feat_nn)
self.last_layer = nn.Linear(emb_dims[-1], out_dim)
for i in range(len(emb_dims) - 1):
self.conv.append(
EdgeConv(emb_dims[i - 1] if i > 0 else in_dim,
emb_dims[i],
batch_norm=True))
def __getitem__(self, idx):
if idx == 0:
start = 0
start_hr = 0
else:
start = self.cumsum[idx-1]
start_hr = self.cumsum_highres[idx-1]
end = self.cumsum[idx]
end_hr = self.cumsum_highres[idx]
# print('start : ', start)
# print('end : ', end)
# print('start_hr : ', start_hr)
# print('end_hr : ', end_hr)
cell_xyz = self.file['CellXYLayer'][start:end]
energies = self.file['TotalEnergy'][start:end]
neu_energies = self.file['NeutralEnergy'][start:end]
cell_xyz_highres = np.concatenate([goes_to_dict[(l,x,y)] for l,x,y in cell_xyz if l < 3])
energies_highres = self.file['TotalEnergy_HighRes'][start_hr:end_hr]
neu_energies_highres = self.file['NeutralEnergy_HighRes'][start_hr:end_hr]
cell_xyz_val = np.array([ Get_XYZ_val(idx[0], idx[1], idx[2]) for idx in cell_xyz ])
cell_xyz_val = np.reshape( cell_xyz_val, (1, cell_xyz_val.shape[0], cell_xyz_val.shape[1]) )
cell_xyz_val = torch.FloatTensor(cell_xyz_val)
cell_xyz_val_hr = np.array([ Get_XYZ_val(idx[0], idx[1], idx[2], highres=True) for idx in cell_xyz_highres ])
cell_xyz_val_hr = np.reshape( cell_xyz_val_hr, (1, cell_xyz_val_hr.shape[0], cell_xyz_val_hr.shape[1]) )
cell_xyz_val_hr = torch.FloatTensor(cell_xyz_val_hr)
cell_layers = cell_xyz[:,0]
#b_factors = [scale_factors[l_i] for l_i in cell_layers]
b_factors = []
for l_i in cell_layers :
if(l_i < 3) :
b_factors.append( scale_factors[l_i] )
else :
b_factors.append( 0 )
graph = KNNGraph(graph_size)
g = graph(cell_xyz_val)
g = dgl.transform.remove_self_loop(g)
g.ndata['energy'] = torch.reshape(torch.FloatTensor(energies), ( torch.FloatTensor(energies).shape[0],1 ) )
g.ndata['broadcast'] = torch.tensor(b_factors)
g.ndata['parent_node'] = g.number_of_nodes() * torch.ones([ g.number_of_nodes() ], dtype=torch.int)
g.ndata['_ID'] = g.nodes()
g.ndata['cell_xyz'] = cell_xyz
g.ndata['neu_energy'] = neu_energies[:, None]
g_hr = graph(cell_xyz_val_hr)
g_hr = dgl.transform.remove_self_loop(g_hr)
g_hr.ndata['energy'] = torch.FloatTensor(energies_highres)
frac_hr = torch.FloatTensor(neu_energies_highres)/torch.FloatTensor(energies_highres)
frac_hr[ torch.isnan(frac_hr) ] = 0.
frac_hr[ torch.where(frac_hr < 0.) ] = 0.
g_hr.ndata['neu_frac'] = frac_hr[:, None]
g_hr.ndata['neu_energy'] = neu_energies_highres[:, None]
#g_hr.ndata['cell_xyz_highres'] = cell_xyz_highres
g_out_hr = graph(cell_xyz_val_hr)
g_out_hr = dgl.transform.remove_self_loop(g_out_hr)
g_out_hr.ndata['cell_xyz_highres'] = cell_xyz_highres
pi0_phi = self.file['Pi0_Phi'][idx:idx+1]
pi0_theta = self.file['Pi0_Theta'][idx:idx+1]
sample = {
'gr' : g,
'gr_hr' : g_hr,
'gr_out_hr' : g_out_hr,
'pi0_theta' : torch.FloatTensor( np.cos(pi0_theta) ),
'pi0_phi' : torch.FloatTensor( np.cos(pi0_phi) )
# 'cell_xyz' : cell_xyz,
# 'cell_xyz_highres' : cell_xyz_highres
}
return sample
def __getitem__(self, idx):
idx = idx + self.nstart
if idx == 0:
start = 0
start_hr = 0
else:
start = self.cumsum[idx - 1]
start_hr = self.cumsum_highres[idx - 1]
end = self.cumsum[idx]
end_hr = self.cumsum_highres[idx]
# print('start : ', start)
# print('end : ', end)
# print('start_hr : ', start_hr)
# print('end_hr : ', end_hr)
Trk_X_indx = self.file['Trk_X_indx'][idx:idx + 1]
Trk_Y_indx = self.file['Trk_Y_indx'][idx:idx + 1]
cell_xyz = self.file['CellXYLayer'][start:end]
energies = self.file['TotalEnergy'][start:end]
neu_energies = self.file['NeutralEnergy'][start:end]
cell_xyz_highres = np.concatenate(
[goes_to_dict[(l, x, y)] for l, x, y in cell_xyz if l < 3])
energies_highres = self.file['TotalEnergy_HighRes'][start_hr:end_hr]
neu_energies_highres = self.file['NeutralEnergy_HighRes'][
start_hr:end_hr]
cell_xyz_val = np.array(
[Get_XYZ_val(idx[0], idx[1], idx[2]) for idx in cell_xyz])
cell_xyz_val = np.reshape(
cell_xyz_val, (1, cell_xyz_val.shape[0], cell_xyz_val.shape[1]))
cell_xyz_val = torch.FloatTensor(cell_xyz_val)
cell_xyz_val_hr = np.array([
Get_XYZ_val(idx[0], idx[1], idx[2], highres=True)
for idx in cell_xyz_highres
])
cell_xyz_val_hr = np.reshape(
cell_xyz_val_hr,
(1, cell_xyz_val_hr.shape[0], cell_xyz_val_hr.shape[1]))
cell_xyz_val_hr = torch.FloatTensor(cell_xyz_val_hr)
cell_layers = cell_xyz[:, 0]
#b_factors = [scale_factors[l_i] for l_i in cell_layers]
b_factors = []
for l_i in cell_layers:
if (l_i < 3):
b_factors.append(scale_factors[l_i])
else:
b_factors.append(0)
graph = KNNGraph(graph_size)
g = graph(cell_xyz_val)
g = dgl.transform.remove_self_loop(g)
g.ndata['neu_energy'] = neu_energies[:, None]
g.ndata['cell_xyz'] = cell_xyz
broad_neu_energy = torch.repeat_interleave(torch.tensor(neu_energies),
torch.tensor(b_factors),
dim=0)
g_hr = graph(cell_xyz_val_hr)
g_hr = dgl.transform.remove_self_loop(g_hr)
g_hr.ndata['broad_neu_energy'] = torch.FloatTensor(
broad_neu_energy)[:, None]
frac_hr = torch.FloatTensor(neu_energies_highres) / torch.FloatTensor(
broad_neu_energy)
frac_hr[torch.isnan(frac_hr)] = 0.
frac_hr[torch.where(frac_hr < 0.)] = 0.
g_hr.ndata['neu_frac'] = frac_hr[:, None]
g_hr.ndata['neu_energy'] = neu_energies_highres[:, None]
g_hr.ndata['cell_xyz_highres'] = cell_xyz_highres
g_hr.ndata['broad_factor'] = torch.repeat_interleave(
torch.tensor(b_factors), torch.tensor(b_factors), dim=0)[:, None]
#print('Neu energy shape : ', cell_xyz[ np.where(cell_xyz[:,0] < 6) ].shape)
sample = {
'gr': g,
'gr_hr': g_hr,
'Trk_X_indx': Trk_X_indx,
'Trk_Y_indx': Trk_Y_indx
}
return sample
def __init__(self, k, feature_dims, emb_dims):
super(Model, self).__init__()
self.nng = KNNGraph(k)
self.conv = EdgeConv(feature_dims, emb_dims, False)