latest changes after KF14 update

interactive/inference/generate.py

@@ -3,11 +3,14 @@ import random
 import sys
 from torch.autograd import Variable
 import numpy as np
-
+import h5py
+import torch.nn as nn
+import matplotlib.pyplot as plt
 sys.path.append('/home/jovyan/pytorchjob')
 from models.generator import DCGAN_G
 from models.generatorFull import Hcal_ecalEMB
-    
+import interactive.physics.basics as B    
+import scipy.spatial.distance as dist
 
 
 def make_shower(eph, Etrue, nEvents): 
@@ -71,4 +74,145 @@ def make_shower(eph, Etrue, nEvents):
     imE = np.vstack(fE_list)
     imH = np.vstack(fH_list)
     
-    return im, imE, imH
\ No newline at end of file
+    return im, imE, imH
+
+
+
+
+def make_shower_rECAL(eph, Etrue, nEvents): 
+    
+    
+    ngf = 32
+    nz = 100
+    batch_size = 1000
+    device = torch.device("cuda")
+    ## LOAD REAL ECAL DATA
+    f50 = h5py.File('/eos/user/e/eneren/scratch/50GeV75k.hdf5', 'r')
+    #####
+
+    
+    ## LOAD HCAL GENERATOR
+    mGenH = Hcal_ecalEMB(ngf, 32, nz, emb_size=32).to(device)
+    mGenH = torch.nn.parallel.DataParallel(mGenH)
+    expH = 'wganHCAL50GeV_v2'
+    #expH = 'wganHCALv1'
+    
+    Tensor = torch.cuda.FloatTensor 
+    
+    
+    gen_checkpointHCAL = torch.load('/eos/user/e/eneren/experiments/' + expH + "_generator_"+ str(eph) + ".pt", map_location=torch.device('cuda'))
+    mGenH.load_state_dict(gen_checkpointHCAL['model_state_dict'])
+    mGenH.eval()
+        
+    fakeList = []
+    fE_list = []
+    fH_list = []
+    for b in range(batch_size, nEvents + 1 , batch_size):
+        
+        print ("Generating of showers: current batch {}  .......".format(b))
+        ##ECAL REAL
+        s50E = f50['ecal/layers'][b-batch_size:b]
+        ecal_shower = torch.from_numpy(s50E).float()
+        ecal_shower = ecal_shower.unsqueeze(1)
+        ecal_shower = ecal_shower.cuda()
+        E = torch.from_numpy(np.random.uniform(Etrue, Etrue, (batch_size,1))).float()
+        ####
+
+        ## HCAL generate
+        zH = Variable(Tensor(np.random.uniform(-1, 1, (batch_size, nz))))
+        hcal_shower = mGenH(zH, E, ecal_shower).detach()
+
+        ecal_shower = ecal_shower.squeeze(1)
+        comb = torch.cat((ecal_shower, hcal_shower),1)
+       
+        fakeList.append(comb.cpu().numpy())
+        fE_list.append(ecal_shower.cpu().numpy())
+        fH_list.append(hcal_shower.cpu().numpy())
+    
+    
+    im = np.vstack(fakeList)
+    imE = np.vstack(fE_list)
+    imH = np.vstack(fH_list)
+    
+    return im, imE, imH
+
+
+def fid_scan_rECAL(showers, eph_start, eph_end):
+    
+    ngf = 32
+    nz = 100
+    batch_size = 1000
+    device = torch.device("cuda")
+    ## LOAD REAL ECAL DATA
+    f50 = h5py.File('/eos/user/e/eneren/scratch/50GeV75k.hdf5', 'r')
+    #####
+
+    ## LOAD HCAL GENERATOR
+    mGenH = Hcal_ecalEMB(ngf, 32, nz, emb_size=32).to(device)
+    mGenH = nn.parallel.DataParallel(mGenH)
+    expH = 'wganHCAL50GeV_v2'
+    
+    Tensor = torch.cuda.FloatTensor 
+    eph_list = list(range(eph_start,eph_end,1))
+    
+    esum_mean = []
+    esum_down = []
+    esum_up = []
+    for eph in eph_list:
+        gen_checkpointHCAL = torch.load('/eos/user/e/eneren/experiments/' + expH + "_generator_"+ str(eph) + ".pt", map_location=torch.device('cuda'))
+        mGenH.load_state_dict(gen_checkpointHCAL['model_state_dict'])
+        mGenH.eval()
+        
+        print ("Epoch: " , eph)
+        jsd = []
+        for Etrue in [50]:
+
+            ##ECAL real
+            ##ECAL REAL
+            s50E = f50['ecal/layers'][:1000]
+            ecal_shower = torch.from_numpy(s50E).float()
+            ecal_shower = ecal_shower.unsqueeze(1)
+            ecal_shower = ecal_shower.cuda()
+            E = torch.from_numpy(np.random.uniform(Etrue, Etrue, (batch_size,1))).float()
+            ####
+            ####
+            
+            ## HCAL generate
+            zH = Variable(Tensor(np.random.uniform(-1, 1, (batch_size, nz))))
+            hcal_shower = mGenH(zH, E, ecal_shower).detach()
+
+            ecal_shower = ecal_shower.squeeze(1)
+            comb = torch.cat((ecal_shower, hcal_shower),1)
+
+            esumFake = B.getTotE(comb.cpu().numpy(), 30, 30, 78)
+            esumReal = showers[str(Etrue)+'F']
+
+            
+
+            jsd.append(B.jsdHist(esumReal, esumFake, 50, 0, 2500, eph, debug=False))
+        
+        #m = (jsd[0] + jsd[1] + jsd[2]) / 3.0
+        m = jsd[0]
+        esum_mean.append(m)
+        esum_down.append(m - min(jsd))
+        esum_up.append(max(jsd) - m)
+        
+    
+        print ("Epoch: ", jsd)
+        
+
+    # Plotting
+    # Energy sum with average and spread
+    plt.figure(figsize=(12,4), facecolor='none', dpi=400)
+    plt.scatter(eph_list, esum_mean, marker = 'x', color='red', label='Average Energy sum')
+    #plt.scatter(epoch_list, epoch_matrix_median[:,0], marker = '+', color='blue', label='Median Energy sum area difference')
+    plt.errorbar(eph_list, esum_mean, yerr=[esum_down, esum_up], color='black', label='Min-Max interval', linestyle='None')
+    #plt.errorbar(epoch_list, epoch_matrix_mean[:,0], yerr=epoch_matrix_std[:,0], color='green', label='Standard deviation', linestyle='None', capsize=3.0)
+
+
+    plt.legend(loc='upper right', fontsize=10)
+    plt.xlabel('epoch', fontsize=14)
+    plt.ylabel('JS difference', fontsize=16)
+    plt.ylim(0.1,1.0)
+    
+    plt.savefig('/eos/user/e/eneren/plots/sweep_min_max__'+str(eph_start)+'__'+str(eph_end)+'.png')
\ No newline at end of file

interactive/physics/basics.py

 import numpy as np
+import matplotlib.pyplot as plt
+import scipy.spatial.distance as dist
 
 def getTotE(data, xbins, ybins, layers):
     data = np.reshape(data,[-1, layers*xbins*ybins])
@@ -54,4 +56,34 @@ def getRadialDistribution(data, xbins, ybins, layers):
     phi_arr = np.asarray(phi_list)
     e_arr = np.asarray(e_list)
 
-    return r_arr, phi_arr, e_arr
\ No newline at end of file
+    return r_arr, phi_arr, e_arr
+
+
+def jsdHist(data_real, data_fake, nbins, minE, maxE, eph, debug=False):
+    
+    figSE = plt.figure(figsize=(6,6*0.77/0.67))
+    axSE = figSE.add_subplot(1,1,1)
+
+
+    pSEa = axSE.hist(data_real, bins=nbins, 
+            weights=np.ones_like(data_real)/(float(len(data_real))), 
+            histtype='step', color='black',
+            range=[minE, maxE])
+    pSEb = axSE.hist(data_fake, bins=nbins, 
+            weights=np.ones_like(data_fake)/(float(len(data_fake))),
+            histtype='step', color='red',
+             range=[minE, maxE])
+
+    frq1 = pSEa[0]
+    frq2 = pSEb[0]
+
+    JSD = dist.jensenshannon(frq1, frq2)    
+
+    if (debug):
+        plt.savefig('./jsd/esum_'+str(eph)+'.png')
+
+    plt.close()
+    
+    if len(frq1) != len(frq2):
+        print('ERROR JSD: Histogram bins are not matching!!')
+    return JSD
\ No newline at end of file

interactive/physics/plotting.py

@@ -11,7 +11,8 @@ plt.style.use('classic')
 mpl.rc('font', **font)
 
 
-def plot_esum(real, fake, nbins, minE, maxE, name):
+def plot_esum(real, fake, nbins, minE, maxE, ymax, name):
+    
     
     figSE = plt.figure(figsize=(6,6*0.77/0.67))
     axSE = figSE.add_subplot(1,1,1)
@@ -41,7 +42,7 @@ def plot_esum(real, fake, nbins, minE, maxE, name):
     figSE.patch.set_facecolor('white') 
         
     axSE.set_xlim([0, maxE])
-    axSE.set_ylim([0, 700])
+    axSE.set_ylim([0, ymax])
     axSE.set_xlabel("MeV", family='serif')
     
     plt.savefig('./esum'+str(name)+'.png')