import SimEventViewer as SMV
import torch
import numpy as np
import matplotlib.pyplot as plt
%matplotlib tk
import pandas as pd
import os
import torch.nn as nn
import torch.optim as optim
from collections import OrderedDict
from tqdm import trange
from scipy.interpolate import interp1d

# bins_dir is the numpy file containing the bin edges used to fit to
bins_dir='/home/nr1315/Documents/Project/MachineLearning/test_bins.npy'

# coeffs_file is the h5 file containing the effective dose coefficients
coeffs_file='/home/nr1315/Documents/Project/effective_dose_coeffs.h5'

testbins=np.load(bins_dir)
coeffs=pd.read_hdf(coeffs_file)
seed=0

# Shuffles data order
def shuffle_inds(seed,batch):
    n_val=int(0.2*batch)
    torch.manual_seed(0)
    shuffled_indices=torch.randperm(batch)
    return shuffled_indices

# Prepare a specific simulated dataset for training
def SimTrainingInputs(file,avg,batch,shuffled_inds,seed):
    a=SMV.NewSimPlotter()
    path,name=os.path.split(file)
    a.add_data(file,name[:-3])
    df=a.data[name[:-3]]
    cubes=df.query('CapCube>0')['CapCube'].values
    z=(cubes/1e4).astype(int)
    x=((cubes-z*1e4)/100).astype(int)
    y=(cubes-z*1e4-x*100).astype(int)
    x-=1
    y-=1
    z-=1
    cubs=pd.DataFrame(np.array([x,y,z]).T,columns=['X','Y','Z'])
    ze=np.zeros([4,4,4])
    coords=np.argwhere(ze==0)
    counts=[]
    for coord in coords:
        counts.append(np.count_nonzero((cubs.values==coord).all(axis=1)))
    count=sum(counts)
    ncounts=np.array(counts)/sum(counts)*avg
    np.random.seed(seed)
    train_x=np.random.poisson(np.tile(ncounts,(batch,1)))
    train_x=64000*train_x/train_x.sum(axis=1)[:,np.newaxis]
    train_x=torch.from_numpy(train_x).to(torch.float)

    n_val=int(0.2*batch)
    if n_val>=1:
        val_x=train_x[shuffled_inds[-n_val:]]
        train_x=train_x[shuffled_inds[:-n_val]]
    else:
        val_x=None
    return train_x,val_x,count

# Prepare a specific real dataset for training
def RealTrainingInputs(file,avg,batch,shuffled_inds,seed):
    file=pd.read_hdf(file)
    file=file[file['Combine']]
    cubes=file.loc[:,['CubeX','CubeY','CubeZ']]
    ze=np.zeros([4,4,4])
    coords=np.argwhere(ze==0)
    counts=[]
    for coord in coords:
        counts.append(np.count_nonzero((cubes.values==coord).all(axis=1)))
    counts=np.array(counts)/2
    count=counts.sum()
    ncounts=counts/counts.sum()*avg
    np.random.seed(seed)
    train_x=np.random.poisson(np.tile(ncounts,(batch,1)))
    train_x=avg*train_x/train_x.sum(axis=1)[:,np.newaxis]
    train_x=torch.from_numpy(train_x).to(torch.float)


    n_val=int(0.2*batch)
    if n_val>=1:
        val_x=train_x[shuffled_inds[-n_val:]]
        train_x=train_x[shuffled_inds[:-n_val]]
    else:
        val_x=None
    return train_x,val_x,count

# Loads an expected energy bins file
def LoadTrainingTargets(file,avg,batch,shuffled_inds):
    train_y=np.nan_to_num(np.load(file))
    train_y=avg*train_y/train_y.sum()
    train_y=train_y[:,np.newaxis].repeat(batch,axis=1).T
    train_y=torch.from_numpy(train_y).to(torch.float)

    n_val=int(0.2*batch)
    if n_val>=1:
        val_y=train_y[shuffled_inds[-n_val:]]
        train_y=train_y[shuffled_inds[:-n_val]]
    else:
        val_y=None
    return train_y,val_y

# Prepares training data from a list of data files, associated labels, and whether each is a
# simulated or real dataset.
def prep_training_inputs(avg, batch, shuffled_inds,files,labels,real,seed=0):
    train_data={}
    val_data={}
    counts=[]
    for i in range(len(files)):
        if real[i]:
            train,val,count=RealTrainingInputs(files[i],avg,batch,shuffled_inds,seed)
        else:
            train,val,count=SimTrainingInputs(files[i],avg,batch,shuffled_inds,seed)
        train_data[labels[i]]=train
        val_data[labels[i]]=val
        counts.append(count)
    train=torch.cat(tuple(train_data.values()))
    try:
        val=torch.cat(tuple(val_data.values()))
    except TypeError:
        val=None
    return train,val,counts

# Prepares training targets from a list of datafiles
def prep_training_targets(avg,batch,shuffled_inds,files,labels):
    train_data={}
    val_data={}
    for i in range(len(files)):
        train_data[labels[i]],val_data[labels[i]]=LoadTrainingTargets(files[i],avg,batch,shuffled_inds)
    train=torch.cat(tuple(train_data.values()))
    try:
        val=torch.cat(tuple(val_data.values()))
    except TypeError:
        val=None
    return train, val

# Calculates effective dose from a bin distribution
def calc_dose(binvals,centres,coeffs):
    E=coeffs['Energy / MeV'].values
    AP=coeffs['AP'].values
    eff_dose=interp1d(E,AP)
    dose=((binvals/400)*eff_dose(centres)).sum(axis=1)
    return dose

xtrain=['/home/nr1315/Documents/Project/NPL_2017-06-13/data/nemenix2_npl_2017-06-13_cf-4774_1m5_0deg_sc-stand_1V0_v12.h5',
        '/home/nr1315/Documents/Project/NPL_2017-06-13/data/nemenix2_npl_2017-06-14_ambe-30-1679_1m5_+0deg_sc-stand_1V0_v67.h5',
        '/home/nr1315/Documents/Project/NPL_2017-06-13/data/nemenix2_npl_2017-06-14_amli_1m5_+0deg_sc-stand_1V0_v55.h5',
        '/home/nr1315/Documents/Project/Simulations/Data/Low_scatter_new/4x4normal/thermal.h5'
       ]
labels=['Cf',
        'AmBe',
        'AmLi',
        'Thermal'
       ]
real=[True,
      True,
      True,
      False
     ]

ytrain=['/home/nr1315/Documents/Project/MachineLearning/CfBinCounts.npy',
        '/home/nr1315/Documents/Project/MachineLearning/AmBeBinCounts.npy',
        '/home/nr1315/Documents/Project/MachineLearning/AmLiBinCounts.npy',
        '/home/nr1315/Documents/Project/MachineLearning/ThermBinCounts.npy'
        ]

batch=1250

torch.manual_seed(seed)
shuf=shuffle_inds(0,batch)
X_train,X_val,X_counts=prep_training_inputs(64000,batch,shuf,xtrain,labels,real,seed=seed)
Y_train,Y_val=prep_training_targets(64000,batch,shuf,ytrain,labels)

x_test=RealTrainingInputs('/home/nr1315/Documents/Project/NPL_2017-06-13/data/nemenix2_npl_2017-06-13_cf-4774_1m5_0deg_sc-stand_1V0_v12.h5',64000,1000,np.arange(0,1000,1),0)
x_test=torch.cat([x_test[0],x_test[1]])
y_test=LoadTrainingTargets('/home/nr1315/Documents/Project/MachineLearning/CfBinCounts.npy',64000,1000,np.arange(0,1000,1))
y_test=torch.cat([y_test[0],y_test[1]])

def iterate_minibatches(X, y, batchsize, shuffle=False):
    indices = np.arange(len(X))
    if shuffle:
        indices = np.random.permutation(indices)
    for start in range(0, len(indices), batchsize):
        ix = indices[start: start + batchsize]
        yield X[ix], y[ix]

def training_loop(n_epochs,model,optimizer,loss_fn,train_x,val_x,train_y,val_y,batchsize):
    train_loss_per_epoch=[]
    val_loss_per_epoch=[]
    weight_hist=[]
    for epoch in range(1,n_epochs+1):
        model.train(True)
        train_loss=0
        for batch_x,batch_y in iterate_minibatches(train_x,train_y,batchsize,shuffle=True):
            batch_pred_y=model(batch_x)
            batch_train_loss=loss_fn(batch_pred_y,batch_y)

            optimizer.zero_grad()
            batch_train_loss.backward()
            optimizer.step()

            train_loss+=batch_train_loss.data.numpy()
        train_loss_per_epoch.append(train_loss)



        if val_x is not None:
            val_loss=loss_fn(model(val_x),val_y)
            val_loss_per_epoch.append(val_loss)

        if epoch==1:
            if val_x is not None:
                print(f"Epoch {epoch}, Training loss {train_loss:.4f},"
                      f" Validation loss {val_loss.item():.4f}")
            else:
                print(f"Epoch {epoch}, Training loss {train_loss:.4f}")
        weight_hist.append(list(model.parameters()))

    return train_loss_per_epoch,val_loss_per_epoch,weight_hist

n_X_train=X_train/X_train.sum(axis=1)[:,np.newaxis]
n_X_val=X_val/X_val.sum(axis=1)[:,np.newaxis]
n_Y_train=Y_train/Y_train.sum(axis=1)[:,np.newaxis]
n_Y_val=Y_val/Y_val.sum(axis=1)[:,np.newaxis]

#model=nn.Sequential(OrderedDict([
#        ('hidden_linear1',nn.Linear(64,128)),
#        ('hidden_activation1',nn.Sigmoid()),
#        ('hidden_linear2',nn.Linear(128,64)),
#        ('hidden_activation2',nn.Sigmoid()),
#        ('hidden_linear3',nn.Linear(64,32)),
#        ('hidden_activation3',nn.Sigmoid()),
#        ('output_layer',nn.Linear(32,24))]))


model=nn.Sequential(nn.Linear(64,128),
                    nn.ReLU(),
                    nn.Linear(128,64),
                    nn.ReLU(),
                    nn.Linear(64,32),
                    nn.ReLU(),
                    nn.Linear(32,24),
                    )

optimizer=optim.Rprop(model.parameters(),lr=1e-2)
n_epochs=100

def weights(m):
    classname=m.__class__.__name__

    if classname=='Linear':
        n=m.in_features
        y=1.0/np.sqrt(n)
        m.weight.data.normal_(0.0,y)
        m.bias.data.fill_(4000)

model.apply(weights)

tloss,vloss,weights=training_loop(n_epochs=n_epochs,
                                  model=model,
                                  optimizer=optimizer,
                                  loss_fn=nn.MSELoss(),
                                  train_x=X_train,
                                  val_x=X_val,
                                  train_y=Y_train,
                                  val_y=Y_val,
                                  batchsize=256)


print()
if len(vloss)>0:
    print('Final training loss:  '+str(tloss[-1])+', final validation loss:  '+str(vloss[-1]))
else:
    print('Final training loss:  '+str(tloss[-1]))


if len(vloss)>0:
    fig,ax=plt.subplots(2,1,figsize=(10,20))
    ax[0].plot(np.arange(0,n_epochs,1),tloss,c='b',label='Training loss')
    ax[0].set_title('Training loss',fontsize=24)
    ax[1].plot(np.arange(0,n_epochs,1),vloss,c='r',label='Validation loss')
    ax[1].set_title('Validation loss',fontsize=24)
else:
    fig=plt.figure()
    ax=fig.add_subplot(111)
    ax.plot(np.arange(0,n_epochs,1),tloss,c='b',label='Training loss')
    ax.set_title('Training loss',fontsize=24)

Epoch 1, Training loss 388717416.5000, Validation loss 2317353.0000

Final training loss:  79519.21240234375, final validation loss:  tensor(5084.6782, grad_fn=<MseLossBackward>)

barcentres=testbins[:-1]+np.diff(testbins)/2
fig,ax=plt.subplots(4,2,figsize=(20,20))
fig.set_tight_layout(True)
title=['Cf','AmBe','AmLi','Thermal']
for i in range(4):
    ax[i,0].bar(barcentres,Y_val[250*i],edgecolor='xkcd:deep blue',width=np.diff(testbins))
    ax[i,0].set_title(title[i]+' training data',fontsize=22)
    ax[i,0].set_xscale('log')
    ax[i,1].bar(barcentres,model(X_val[250*i]).detach().numpy(),edgecolor='xkcd:deep blue',width=np.diff(testbins))
    ax[i,1].set_title(title[i]+' predicted data',fontsize=22)
    ax[i,1].set_xscale('log')
    ax[i,0].set_xlabel('Energy / MeV',fontsize=18)
    ax[i,1].set_xlabel('Energy / MeV',fontsize=18)
    ax[i,0].set_xlabel('Frequency',fontsize=18)
    ax[i,1].set_xlabel('Frequency',fontsize=18)

plt.show()

loss=nn.MSELoss()
l=loss(y_test,model(x_test)).item()

fig,ax=plt.subplots(1,2,figsize=(20,10))
ax[0].bar(barcentres,y_test[0],edgecolor='xkcd:deep blue',width=np.diff(testbins))
ax[0].set_title('Test data')
ax[1].bar(barcentres,model(x_test[0]).detach().numpy(),edgecolor='xkcd:deep blue',width=np.diff(testbins))
ax[1].set_title('Predicted data')
ax[0].set_xscale('log')
ax[1].set_xscale('log')

Now load in an IPNDV dataset, to compare to Cf & the neural net predictions

import DataProcess as DP
a=DP.Plotter()
folder='/home/nr1315/Documents/Project/IPNDVdata/data_BR1/10-09-2019/'
a.add_data(folder,'2019_09_10_bkg_0deg_0.95d_0.91h_afternoon.root','background')
a.add_data(folder,'2019_09_10_pin96_Bare_0deg_0.root','96_bare')

bkg=a.data['background']
bkg=bkg[bkg['Combine']]
bkgcubes=bkg.loc[:,['CubeX','CubeY','CubeZ']]

data=a.data['96_bare']
data=data[data['Combine']]
cubes=data.loc[:,['CubeX','CubeY','CubeZ']]

ze=np.zeros([4,4,4])
coords=np.argwhere(ze==0)
bkgcounts=[]
counts=[]
for coord in coords:
    counts.append(np.count_nonzero((cubes.values==coord).all(axis=1)))
    bkgcounts.append(np.count_nonzero((bkgcubes.values==coord).all(axis=1)))
bkgcounts=np.array(bkgcounts)/2
counts=np.array(counts)/2

norm=(a.data['background']['time'].max()-a.data['background']['time'].min())/(a.data['96_bare']['time'].max()-a.data['96_bare']['time'].min())
if norm>1:
    counts=counts - bkgcounts/norm
else:
    counts=norm*counts - bkgcounts

tt=np.zeros([4,4,4])
tt[coords[:,0],coords[:,1],coords[:,2]]=counts
tt=np.rot90(tt,2,axes=(0,1))
tcounts=tt.flatten()

Plotting Cf training inputs/target

fig,ax=plt.subplots(1,2,figsize=(20,10))

ax[0].bar(np.arange(0,64,1),X_val[0],edgecolor='xkcd:deep blue',width=1)
ax[0].set_xlabel('Cube ID',fontsize=18)
ax[0].set_ylabel('Frequency',fontsize=18)
ax[0].set_title('Cf training input',fontsize=22)

ax[1].bar(barcentres,Y_val[0],edgecolor='xkcd:deep blue',width=np.diff(testbins))
ax[1].set_xlabel('Energy / MeV',fontsize=18)
ax[1].set_ylabel('Frequency',fontsize=18)
ax[1].set_title('Cf training target',fontsize=22)
ax[1].set_xscale('log')

Plotting the IPNDV dataset prediction vs the Cf target used (i.e. comparing network prediction of IPNDV dataset with a known Cf fluence)

fig,ax=plt.subplots(1,2,figsize=(20,10))

ax[1].bar(barcentres,model(torch.from_numpy(counts).to(torch.float)).detach().numpy(),edgecolor='xkcd:deep blue',width=np.diff(testbins))
ax[1].set_xscale('log')
ax[1].set_xlabel('Energy / MeV',fontsize=18)
ax[1].set_ylabel('Frequency',fontsize=18)
ax[1].set_title('Fluence prediction, 96 bare (IPNDV)',fontsize=22)

ax[0].bar(barcentres,Y_val[0],edgecolor='xkcd:deep blue',width=np.diff(testbins))
ax[0].set_xscale('log')
ax[0].set_xlabel('Energy / MeV',fontsize=18)
ax[0].set_title('Cf fluence target',fontsize=22)
ax[0].set_ylabel('Frequency',fontsize=18)

Text(0, 0.5, 'Frequency')

tt=np.rot90(tt,2,axes=(0,1))

Plotting 3D cube view of background subtracted IPNDV data

import ndap
import matplotlib.colors as colors
import matplotlib.cm as cm
from mpl_toolkits.mplot3d import Axes3D

sc=(50.888-50)/50
p=ndap.NDArrayPlotter(tt)
p.set_alpha(0.05)

fig=plt.figure()
ax=fig.add_subplot(111,projection='3d')
fig.set_tight_layout(True)

cmap=cm.viridis
norm=colors.Normalize(vmin=0,vmax=np.max(tt))
p.colors=cmap(norm(tt))
alph=norm(tt)*0.95
p.alphas=alph
sm=plt.cm.ScalarMappable(cmap=cmap,norm=norm)
sm.set_array(tt)

p.render(azim=-56,elev=25,ax=ax,text=None,labels=True,space=0.5)
ax.quiver(-0.4,-0.4,-0.4,1,0,0,length=5,arrow_length_ratio=0.05,color='black')
ax.quiver(-0.4,-0.4,-0.4,0,1,0,length=5,arrow_length_ratio=0.05,color='black')
ax.quiver(-0.4,-0.4,-0.4,0,0,1,length=5,arrow_length_ratio=0.05,color='black')
cbar=plt.colorbar(sm,ax=ax)
cbar.set_label('Event count',rotation=270,fontsize=30,labelpad=30)
plt.show()

Comparing inputs and predictions for Cf and IPNDV data

d=X_val.detach().numpy()[0]
e=X_val.detach().numpy()[750]
ef=np.zeros([4,4,4])
ef[coords[:,0],coords[:,1],coords[:,2]]=e
ef=np.rot90(ef,2,axes=(0,1))

ecounts=ef.flatten()

fcounts=d+ecounts

dcounts=np.copy(counts)

Note this is trying to account for rotation of the detector

dcounts[:16]=dcounts[:16][::-1]
dcounts[16:32]=dcounts[16:32][::-1]
dcounts[32:48]=dcounts[32:48][::-1]
dcounts[48:]=dcounts[48:][::-1]

fig,ax=plt.subplots(2,2,figsize=(20,10))
fig.set_tight_layout(True)
ax[0,0].bar(np.arange(0,64,1),fcounts/fcounts.sum())
ax[0,0].set_title('Cf cube counts,normalized',fontsize=22)
ax[0,0].set_xlabel('Cube number',fontsize=18)
ax[0,1].bar(np.arange(0,64,1),tcounts/tcounts.sum())
ax[0,1].set_title('96 bare cube counts, normalized',fontsize=22)
ax[0,1].set_xlabel('Cube number',fontsize=18)

ax[1,1].bar(barcentres,model(torch.from_numpy(counts).to(torch.float)).detach().numpy(),edgecolor='xkcd:deep blue',width=np.diff(testbins))
ax[1,1].set_xscale('log')
ax[1,1].set_xlabel('Energy / MeV',fontsize=18)
ax[1,1].set_ylabel('Frequency',fontsize=18)
ax[1,1].set_title('Fluence prediction, 96 bare (IPNDV)',fontsize=22)

ax[1,0].bar(barcentres,Y_val[750]+Y_val[0],edgecolor='xkcd:deep blue',width=np.diff(testbins))
ax[1,0].set_xscale('log')
ax[1,0].set_xlabel('Energy / MeV',fontsize=18)
ax[1,0].set_title('Cf fluence target',fontsize=22)
ax[1,0].set_ylabel('Frequency',fontsize=18)

plt.show()

Beyond here is testing of dose calculation, unfinished

def calc_dose(binvals,centres,coeffs):
    E=coeffs['Energy / MeV'].values
    AP=coeffs['AP'].values
    eff_dose=interp1d(E,AP)
    dose=((binvals/400)*eff_dose(centres)).sum(axis=1)
    return dose

E=coeffs['Energy / MeV'].values
AP=coeffs['AP'].values
eff_dose=interp1d(E,AP)

effserrs=np.load('/home/nr1315/Documents/Project/Simulations/Data/Low_scatter_old/simeffs.npy')

effserrs[0]/=1e6

Y_val[0]

tensor([0.0000e+00, 7.7758e+02, 0.0000e+00, 1.5339e+03, 2.1336e+03, 1.5129e+03,
        2.4050e+03, 2.8502e+03, 3.5835e+03, 4.0967e+03, 4.8158e+03, 5.3545e+03,
        5.8893e+03, 6.2581e+03, 6.4400e+03, 5.8920e+03, 4.7713e+03, 3.3026e+03,
        1.6761e+03, 5.9166e+02, 1.1313e+02, 0.0000e+00, 2.2123e+00, 0.0000e+00])

xtrain

['/home/nr1315/Documents/Project/NPL_2017-06-13/data/nemenix2_npl_2017-06-13_cf-4774_1m5_0deg_sc-stand_1V0_v12.h5',
 '/home/nr1315/Documents/Project/NPL_2017-06-13/data/nemenix2_npl_2017-06-14_ambe-30-1679_1m5_+0deg_sc-stand_1V0_v67.h5',
 '/home/nr1315/Documents/Project/NPL_2017-06-13/data/nemenix2_npl_2017-06-14_amli_1m5_+0deg_sc-stand_1V0_v55.h5',
 '/home/nr1315/Documents/Project/Simulations/Data/Low_scatter_new/4x4normal/thermal.h5']

def calc_dose_new(binvals,centres,coeffs,eff):
    E=coeffs['Energy / MeV'].values
    AP=coeffs['AP'].values
    eff_dose=interp1d(E,AP)
    efficiency=interp1d(eff[0],eff[1])
    dose=((binvals.detach().numpy()/400))*eff_dose(centres)/efficiency(centres)
    return dose

times=[]
for file in xtrain[:3]:
    df=pd.read_hdf(file)
    time=df['time'].max()-df['time'].min()
    times.append(time)

efficiency=interp1d(effserrs[0],effserrs[1])

X_counts[:3]

[95078.0, 6386.0, 30535.0]

sctime=64000*np.array(times)/np.array(X_counts[:3])/1e8

calc_dose_new(Y_val[0]/sctime[0],barcentres,coeffs,effserrs)

array([ 0.        ,  0.51202556,  0.        ,  0.71689412,  1.1773983 ,
        1.04159186,  2.16550161,  3.47970837,  6.10912891,  9.55046197,
       15.63753709, 24.05948994, 38.3344697 , 58.01321766, 75.54448203,
       96.28895568, 89.83380318, 74.00285959, 46.41913814, 20.25659155,
        4.87955757,  0.        ,  0.14702784,  0.        ])

plt.bar(barcentres,Y_val[0]/sctime[0],edgecolor='xkcd:deep blue',width=np.diff(testbins))
plt.xscale('log')

X_val[0]

tensor([ 617.0620, 1067.6470, 1001.8536, 1323.8423,  544.2906, 1100.5436,
        1080.6062,  866.2794,  460.5536,  665.9086,  654.9431,  593.1372,
         257.1923,  328.9668,  439.6193,  380.8040, 1026.7753, 2423.3892,
        2194.1091, 1863.1486,  830.3920, 1653.8060, 1790.3771, 1633.8687,
         453.5755,  778.5549, 1151.3839,  821.4202,  455.5692,  435.6318,
         462.5473,  471.5191, 1112.5060, 1838.2268, 2118.3472, 2001.7134,
         939.0508, 1315.8673, 1655.7998, 1402.5950,  494.4471, 1030.7628,
         836.3733,  833.3826,  518.3720,  474.5097,  550.2718,  452.5786,
        1024.7815, 2227.0059, 2357.5957, 1797.3552,  783.5392,  917.1197,
        1592.9969, 1201.2274,  556.2530,  632.0151,  913.1322,  847.3388,
         382.7978,  434.6350,  481.4878,  446.5974])

Y_val[0]

tensor([0.0000e+00, 7.7758e+02, 0.0000e+00, 1.5339e+03, 2.1336e+03, 1.5129e+03,
        2.4050e+03, 2.8502e+03, 3.5835e+03, 4.0967e+03, 4.8158e+03, 5.3545e+03,
        5.8893e+03, 6.2581e+03, 6.4400e+03, 5.8920e+03, 4.7713e+03, 3.3026e+03,
        1.6761e+03, 5.9166e+02, 1.1313e+02, 0.0000e+00, 2.2123e+00, 0.0000e+00])