initial code commit

Author: bmdillon

run_dir/extargs.py

@@ -0,0 +1,28 @@
+#!/bin/env python3
+logfile = "/path/to/logfile.txt"
+tr_dat_path = "/path/to/data.npy"
+tr_lab_path = "/path/to/labels.npy"
+nconstit = 50
+model_dim = 1000
+output_dim = 1000
+n_heads = 4
+dim_feedforward = 1000
+n_layers = 4
+n_head_layers = 2
+opt = "adam"
+sbratio = 1.0
+n_epochs = 500
+learning_rate = 0.00005
+batch_size = 128
+temperature = 0.10
+rot = True
+ptd = True
+ptcm = 0.1
+ptst = 0.1
+trs = True
+trsw = 1.0
+cf = True
+mask= False
+cmask = True
+expt = "experiment-name"
+

scripts/modules/fcn.py

@@ -0,0 +1,43 @@
+import os
+import sys
+import numpy as np
+import random
+import time
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# class for fully connected neural network
+class fully_connected_network( nn.Module ):
+    # define and intialize the structure of the neural network
+    def __init__( self, input_size, output_size, hidden_size, n_hidden, dropout_rate, opt, learning_rate ):
+        super( fully_connected_network, self ).__init__()
+        # define hyperparameters
+        self.input_size    = input_size
+        self.output_size   = output_size
+        self.hidden_size   = hidden_size
+        self.n_hidden      = n_hidden
+        self.opt = opt
+        self.learning_rate = learning_rate
+        self.dropout_rate = dropout_rate
+        # define layers
+        self.input_layer = nn.Linear( self.input_size, self.hidden_size )
+        self.hidden_layers = nn.ModuleList()
+        self.dropout = nn.Dropout( p=self.dropout_rate )
+        for i in range( self.n_hidden ):
+            self.hidden_layers.append( nn.Linear( self.hidden_size, self.hidden_size ) )    
+        self.output_layer = nn.Linear( self.hidden_size, self.output_size )
+        if self.opt == "adam":
+            self.optimizer = torch.optim.Adam( self.parameters(), lr=self.learning_rate )
+        if self.opt == "sgd":
+            self.optimizer = torch.optim.SGD( self.parameters(), lr=self.learning_rate )
+        
+    def forward( self, x ):
+        x = F.relu( self.input_layer( x ) )
+        for layer in self.hidden_layers:
+            x = F.relu(layer( x ))
+            x = self.dropout( x )
+        output = self.output_layer( x )
+        return output

scripts/modules/fcn_linear.py

@@ -0,0 +1,30 @@
+import os
+import sys
+import numpy as np
+import random
+import time
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# class for fully connected linear neural network
+class fully_connected_linear_network( nn.Module ):
+    # define and intialize the structure of the neural network
+    def __init__( self, input_size, output_size, opt, learning_rate ):
+        super( fully_connected_linear_network, self ).__init__()
+        # define hyperparameters
+        self.input_size    = input_size
+        self.output_size   = output_size
+        self.opt = opt
+        self.learning_rate = learning_rate
+        # define layers
+        self.layer = nn.Linear( self.input_size, self.output_size )
+        if self.opt == "adam":
+            self.optimizer = torch.optim.Adam( self.parameters(), lr=self.learning_rate )
+        if self.opt == "sgd":
+            self.optimizer = torch.optim.SGD( self.parameters(), lr=self.learning_rate )
+    def forward( self, x ):
+        output = self.layer( x )
+        return output

scripts/modules/jet_augs.py

@@ -0,0 +1,108 @@
+import os
+import sys
+import numpy as np
+import random
+import time
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def translate_jets( batch, width=1.0 ):
+    '''
+    Input: batch of jets, shape (batchsize, 3, n_constit)
+    dim 1 ordering: (pT, eta, phi)
+    Output: batch of eta-phi translated jets, same shape as input
+    '''
+    mask = (batch[:,0] > 0) # 1 for constituents with non-zero pT, 0 otherwise
+    ptp_eta  = np.ptp(batch[:,1,:], axis=-1, keepdims=True) # ptp = 'peak to peak' = max - min
+    ptp_phi  = np.ptp(batch[:,2,:], axis=-1, keepdims=True) # ptp = 'peak to peak' = max - min
+    low_eta  = -width*ptp_eta
+    high_eta = +width*ptp_eta
+    low_phi  = np.maximum(-width*ptp_phi, -np.pi-np.amin(batch[:,2,:], axis=1).reshape(ptp_phi.shape))
+    high_phi = np.minimum(+width*ptp_phi, +np.pi-np.amax(batch[:,2,:], axis=1).reshape(ptp_phi.shape))
+    shift_eta = mask*np.random.uniform(low=low_eta, high=high_eta, size=(batch.shape[0], 1))
+    shift_phi = mask*np.random.uniform(low=low_phi, high=high_phi, size=(batch.shape[0], 1))
+    shift = np.stack([np.zeros((batch.shape[0], batch.shape[2])), shift_eta, shift_phi], 1)
+    shifted_batch = batch+shift
+    return shifted_batch
+
+
+def rotate_jets( batch ):
+    '''
+    Input: batch of jets, shape (batchsize, 3, n_constit)
+    dim 1 ordering: (pT, eta, phi)
+    Output: batch of jets rotated independently in eta-phi, same shape as input
+    '''
+    rot_angle = np.random.rand(batch.shape[0])*2*np.pi
+    c = np.cos(rot_angle)
+    s = np.sin(rot_angle)
+    o = np.ones_like(rot_angle)
+    z = np.zeros_like(rot_angle)
+    rot_matrix = np.array([[o, z, z], [z, c, -s], [z, s, c]]) # (3, 3, batchsize)
+    return np.einsum('ijk,lji->ilk', batch, rot_matrix)
+
+def normalise_pts( batch ):
+    '''
+    Input: batch of jets, shape (batchsize, 3, n_constit)
+    dim 1 ordering: (pT, eta, phi)
+    Output: batch of pT-normalised jets, pT in each jet sums to 1, same shape as input
+    '''
+    batch_norm = batch.copy()
+    batch_norm[:,0,:] = np.nan_to_num(batch_norm[:,0,:]/np.sum(batch_norm[:,0,:], axis=1)[:, np.newaxis], posinf = 0.0, neginf = 0.0 )
+    return batch_norm
+
+def rescale_pts( batch ):
+    '''
+    Input: batch of jets, shape (batchsize, 3, n_constit)
+    dim 1 ordering: (pT, eta, phi)
+    Output: batch of pT-rescaled jets, each constituent pT is rescaled by 600, same shape as input
+    '''
+    batch_rscl = batch.copy()
+    batch_rscl[:,0,:] = np.nan_to_num(batch_rscl[:,0,:]/600, posinf = 0.0, neginf = 0.0 )
+    return batch_rscl
+
+def crop_jets( batch, nc ):
+    '''
+    Input: batch of jets, shape (batchsize, 3, n_constit)
+    dim 1 ordering: (pT, eta, phi)
+    Output: batch of cropped jets, each jet is cropped to nc constituents, shape (batchsize, 3, nc)
+    '''
+    batch_crop = batch.copy()
+    return batch_crop[:,:,0:nc]
+
+def distort_jets( batch, strength=0.1, pT_clip_min=0.1 ):
+    '''
+    Input: batch of jets, shape (batchsize, 3, n_constit)
+    dim 1 ordering: (pT, eta, phi)
+    Output: batch of jets with each constituents position shifted independently, shifts drawn from normal with mean 0, std strength/pT, same shape as input
+    '''
+    pT = batch[:,0]   # (batchsize, n_constit)
+    shift_eta = np.nan_to_num( strength * np.random.randn(batch.shape[0], batch.shape[2]) / pT.clip(min=pT_clip_min), posinf = 0.0, neginf = 0.0 )# * mask
+    shift_phi = np.nan_to_num( strength * np.random.randn(batch.shape[0], batch.shape[2]) / pT.clip(min=pT_clip_min), posinf = 0.0, neginf = 0.0 )# * mask
+    shift = np.stack( [ np.zeros( (batch.shape[0], batch.shape[2]) ), shift_eta, shift_phi ], 1)
+    return batch + shift
+
+def collinear_fill_jets( batch ):
+    '''
+    Input: batch of jets, shape (batchsize, 3, n_constit)
+    dim 1 ordering: (pT, eta, phi)
+    Output: batch of jets with collinear splittings, the function attempts to fill as many of the zero-padded args.nconstit
+    entries with collinear splittings of the constituents by splitting each constituent at most once, same shape as input
+    '''
+    batchb = batch.copy()
+    nc = batch.shape[2]
+    nzs = np.array( [ np.where( batch[:,0,:][i]>0.0)[0].shape[0] for i in range(len(batch)) ] )
+    for k in range(len(batch)):
+        nzs1 = np.max( [ nzs[k], int(nc/2) ] )
+        zs1 = int(nc-nzs1)
+        els = np.random.choice( np.linspace(0,nzs1-1,nzs1), size=zs1, replace=False )
+        rs = np.random.uniform( size=zs1 )
+        for j in range(zs1):
+            batchb[k,0,int(els[j])] = rs[j]*batch[k,0,int(els[j])]
+            batchb[k,0,int(nzs[k]+j)] = (1-rs[j])*batch[k,0,int(els[j])]
+            batchb[k,1,int(nzs[k]+j)] = batch[k,1,int(els[j])]
+            batchb[k,2,int(nzs[k]+j)] = batch[k,2,int(els[j])]
+    return batchb

scripts/modules/losses.py

@@ -0,0 +1,43 @@
+import os
+import sys
+import numpy as np
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def contrastive_loss( x_i, x_j, temperature ):
+    xdevice = x_i.get_device()
+    batch_size = x_i.shape[0]
+    z_i = F.normalize( x_i, dim=1 )
+    z_j = F.normalize( x_j, dim=1 )
+    z   = torch.cat( [z_i, z_j], dim=0 )
+    similarity_matrix = F.cosine_similarity( z.unsqueeze(1), z.unsqueeze(0), dim=2 )
+    sim_ij = torch.diag( similarity_matrix,  batch_size )
+    sim_ji = torch.diag( similarity_matrix, -batch_size )
+    positives = torch.cat( [sim_ij, sim_ji], dim=0 )
+    nominator = torch.exp( positives / temperature )
+    negatives_mask = ( ~torch.eye( 2*batch_size, 2*batch_size, dtype=bool ) ).float()
+    negatives_mask = negatives_mask.to( xdevice )
+    denominator = negatives_mask * torch.exp( similarity_matrix / temperature )
+    loss_partial = -torch.log( nominator / torch.sum( denominator, dim=1 ) )
+    loss = torch.sum( loss_partial )/( 2*batch_size )
+    return loss
+
+def align_loss(x, y, alpha=2):
+    xdevice = x.get_device()
+    reps_x = x.clone()
+    reps_y = y.clone()
+    reps_x = F.normalize(reps_x, dim=1).to(xdevice)
+    reps_y = F.normalize(reps_y, dim=1).to(xdevice)
+    loss_align = (reps_x-reps_y).norm(p=2, dim=1).pow(exponent=alpha).mean()
+    return loss_align
+
+def uniform_loss(x, t=2):
+    xdevice = x.get_device()
+    reps_x = x.clone()
+    reps_x = F.normalize(reps_x, dim=1).to(xdevice)
+    loss_uniform = torch.pdist(reps_x, p=2).pow(2).mul(-t).exp().mean().log()
+    return loss_uniform
+    

scripts/modules/perf_eval.py

@@ -0,0 +1,65 @@
+# import standard python modules
+import os
+import sys
+import numpy as np
+from sklearn import metrics
+
+# import torch modules
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# import simple FCN network
+from modules.fcn_linear import fully_connected_linear_network
+from modules.fcn import fully_connected_network
+
+# import preprocessing functions
+from sklearn.preprocessing import StandardScaler, MaxAbsScaler, RobustScaler
+
+def find_nearest( array, value ):
+    array = np.asarray( array )
+    idx = ( np.abs( array-value ) ).argmin()
+    return array[idx]
+
+def get_perf_stats( labels, measures ):
+    measures = np.nan_to_num( measures )
+    auc = metrics.roc_auc_score( labels, measures )
+    fpr,tpr,thresholds = metrics.roc_curve( labels, measures )
+    fpr2 = [ fpr[i] for i in range( len( fpr ) ) if tpr[i]>=0.5]
+    tpr2 = [ tpr[i] for i in range( len( tpr ) ) if tpr[i]>=0.5]
+    try:
+        imtafe = np.nan_to_num( 1 / fpr2[ list( tpr2 ).index( find_nearest( list( tpr2 ), 0.5 ) ) ] )
+    except:
+        imtafe = 1
+    return auc, imtafe
+
+def linear_classifier_test( linear_input_size, linear_batch_size, linear_n_epochs, linear_opt, linear_learning_rate, reps_tr_in, trlab_in, reps_te_in, telab_in ):
+    xdevice = torch.device( "cuda" if torch.cuda.is_available() else "cpu" )
+    fcn_linear = fully_connected_linear_network( linear_input_size, 1, linear_opt, linear_learning_rate )
+    fcn_linear.to( xdevice )
+    bce_loss = nn.BCELoss()
+    sigmoid = nn.Sigmoid()
+    losses = []
+    if linear_opt == "sgd":
+        scheduler = torch.optim.lr_scheduler.StepLR( fcn_linear.optimizer, 100, gamma=0.6, last_epoch=-1, verbose=False)
+    for epoch in range( linear_n_epochs ):
+        indices_list = torch.split( torch.randperm( reps_tr_in.shape[0] ), linear_batch_size )
+        losses_e = []
+        for i, indices in enumerate( indices_list ):
+            fcn_linear.optimizer.zero_grad()
+            x = reps_tr_in[indices,:]
+            l = trlab_in[indices]
+            x = torch.Tensor( x ).view( -1, linear_input_size ).to( xdevice )
+            l = torch.Tensor( l ).view( -1, 1 ).to( xdevice )
+            z = sigmoid( fcn_linear( x ) ).to( xdevice )
+            loss = bce_loss( z, l ).to( xdevice )
+            loss.backward()
+            fcn_linear.optimizer.step()
+            losses_e.append( loss.detach().cpu().numpy() )    
+        losses.append( np.mean( np.array( losses_e )  ) )
+        if linear_opt == "sgd":
+            scheduler.step()
+    out_dat = fcn_linear( torch.Tensor( reps_te_in ).view(-1, linear_input_size).to( xdevice ) ).detach().cpu().numpy()
+    out_lbs = telab_in
+    return out_dat, out_lbs, losses
+