In [1]:
    
import matplotlib.pyplot as plt
import os
import pandas as pd
import numpy as np
import quantumpropagator as qp
from quantumpropagator import readWholeH5toDict, abs2
#from __future__ import print_function
from ipywidgets import interact#, interactive, fixed, interact_manual
import ipywidgets as widgets
def heatThis(h5File,gams,thes,vmaxV,state,save=None):
    save = save or False
    from mpl_toolkits.axes_grid1 import make_axes_locatable
    wf   = qp.retrieve_hdf5_data(h5File,'WF')[:,:,state]
    time = qp.retrieve_hdf5_data(h5File,'Time')[0]
    fig = plt.figure(figsize=(10, 10), dpi= 80, facecolor='w', edgecolor='k')
    plt.title('Time = {:10.5f} fs'.format(time))
    plt.xlabel('Gamma')
    plt.ylabel('Theta')
    
    # this is to get a nice colorbar on the side
    ax = plt.gca()
    aaa = np.rad2deg
    ext = [aaa(gams[0]),aaa(gams[-1]),aaa(thes[0])*2,aaa(thes[-1])*2]
    im = ax.imshow(qp.abs2(wf), extent=ext, cmap='hot', vmax=vmaxV)
    #im = ax.imshow(qp.abs2(wf), cmap='PuBu_r', vmax=0.4)
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    plt.colorbar(im, cax=cax)
    
    if save:
        fig.savefig(os.path.splitext(h5File)[0] + '.png')
        plt.close()
    
def f(frame,state):
    return heatThis(os.path.join(a,filesList[frame]),gams,thes,vmaxV,state)
    
In [2]:
    
subfolders = sorted([ dir for dir in os.listdir('.') if os.path.isdir(dir) and dir != '.ipynb_checkpoints'])
print(''.join(['{} -> {}\n'.format(a,b) for a,b in enumerate(subfolders)]))
    
    
In [3]:
    
a = subfolders[-1]
    
In [4]:
    
filesList = [ fn for fn in sorted(os.listdir(a)) if fn[:8] == 'Gaussian' and fn[-3:] == '.h5']
outh5 = os.path.join(a,'allInput.h5')
dictio = readWholeH5toDict(outh5)
gams,thes = dictio['gams'],dictio['thes']
lastV = len(filesList)-1
# dictio.keys()
zeroWF = qp.retrieve_hdf5_data(os.path.join(a,filesList[0]),'WF')
vmaxV = abs2(zeroWF).max()
gamL,theL,nstates = (qp.retrieve_hdf5_data(os.path.join(a,filesList[0]),'WF')).shape
gamsT = np.rad2deg(gams)
thesT = np.rad2deg(thes)*2
    
In [5]:
    
interact(f, frame = widgets.IntSlider(min=0,max=lastV,step=1,value=lastV),state = widgets.IntSlider(min=0,max=nstates-1,step=1,value=0));
createimages = False
if createimages:
    for filn in filesList:
        filna = os.path.join(a,filn)
        heatThis(filna, gams, thes, vmaxV, save=True)
    
    
 
 
In [ ]:
    
    
In [6]:
    
def sliceGammas(h5File, gamma, exa):
    wf  = qp.retrieve_hdf5_data(h5File,'WF')
    fig = plt.figure(figsize=(11, 6), dpi= 80, facecolor='w', edgecolor='k')
    ys = wf[gamma,:]*exa
    time = qp.retrieve_hdf5_data(h5File,'Time')[0]
    plt.title('Time = {:10.5f} fs --- gammaL = {:8.3f}'.format(time,gamsT[gamma]))
    plt.ylim(-0.5,0.5)
    plt.plot(thesT, np.real(ys), linewidth=1,ls='--')
    plt.plot(thesT, np.imag(ys), linewidth=1,ls='--')
    plt.plot(thesT, abs2(ys), linewidth=3,ls='-')
def sliceThetas(h5File, theta, exa):
    wf  = qp.retrieve_hdf5_data(h5File,'WF')
    fig = plt.figure(figsize=(11, 6), dpi= 80, facecolor='w', edgecolor='k')
    ys = wf[:,theta]*exa
    time = qp.retrieve_hdf5_data(h5File,'Time')[0]
    plt.title('Time = {:10.5f} fs --- thetaL = {:8.3f}'.format(time,thesT[theta]))
    plt.ylim(-0.5,0.5)
    plt.plot(gamsT, np.real(ys), linewidth=1,ls='--')
    plt.plot(gamsT, np.imag(ys), linewidth=1,ls='-.')
    plt.plot(gamsT, abs2(ys), linewidth=3,ls='-')    
    
def fgam(file_number,gamma,exa):
    return sliceGammas(os.path.join(a, filesList[file_number]), gamma, exa)
def fthe(file_number,theta,exa):
    return sliceThetas(os.path.join(a, filesList[file_number]), theta, exa)
    
In [7]:
    
interact(fgam, file_number = widgets.IntSlider(min=0,max=lastV,step=1,value=lastV), gamma = widgets.IntSlider(min=0,max=gamL-1,step=1,value=10), exa = widgets.IntSlider(min=1,max=10,step=1,value=1));
    
    
 
 
In [8]:
    
interact(fthe, file_number = widgets.IntSlider(min=0,max=lastV,step=1,value=lastV), theta = widgets.IntSlider(min=0,max=theL-1,step=1,value=16), exa = widgets.IntSlider(min=1,max=10,step=1,value=1));
    
    
 
 
In [9]:
    
outfn = os.path.join(a,'output')
outfnP = os.path.join(a,'outputPopul')
    
In [10]:
    
data = pd.read_csv(outfn, delim_whitespace=True, header=None);
dataP = pd.read_csv(outfnP, delim_whitespace=True, header=None);
data.columns = ['steps','fs','Norm Deviation','Kinetic','Potential','Total','Total deviation','Xpulse','Ypulse','Zpulse']
    
In [11]:
    
result = pd.concat([data, dataP], axis=1)
result
    
    Out[11]:
In [12]:
    
data.plot(title = 'Norm Deviation', x='fs', y = 'Norm Deviation', figsize=(15,4));
    
    
In [13]:
    
data.plot(title = 'Kinetic', x='fs', y = 'Kinetic', figsize=(15,4));
    
    
In [14]:
    
data.plot(title = 'Potential', x='fs', y = 'Potential', figsize=(15,4));
    
    
In [15]:
    
data['Kinetic_Moved'] = data['Kinetic'] + data['Potential'][0]
data.plot(title = 'Comparison Potential Total Kinetic', x=['fs'] ,y=['Kinetic_Moved','Potential','Total'], figsize=(15,5));
    
    
In [16]:
    
fig = plt.figure(figsize=(9,5))
ax1 = fig.add_subplot(111)
ax2 = ax1.twinx()
ax1.set_ylabel('Population')
ax2.set_ylabel('Pulse')
popul = np.arange(nstates)+1
result.plot(title = 'Population and Pulse', ax = ax1,  x=['fs'] ,y=popul, linewidth=0.8)
result.plot(title = 'Population and Pulse', ax = ax2,  x=['fs'] ,y=['Xpulse','Ypulse','Zpulse'], linewidth=0.5,ls='--');
    
    
In [17]:
    
def expected(h5file):
    wf  = qp.retrieve_hdf5_data(h5file,'WF')
    _,_,nstates = wf.shape
    for i in range(nstates):
        wfState = wf[:,:,i]
        popu = np.linalg.norm(wfState) # this is norm, the real popu is this squared
        if popu == 0:
            print('Expected Values State {}:\nPopul: {:10.3f} \nTheta: - \nGamma: -'.format(i, popu**2, theA, gamA))
        else:
            wfA = abs2(wfState/popu)
            gamAvg = np.sum(wfA,axis=1)
            theAvg = np.sum(wfA,axis=0)
            gamA = sum([ gamsT[i] * gamAvg[i] for i in range(gamAvg.size) ])
            theA = sum([ thesT[i] * theAvg[i] for i in range(theAvg.size) ])
            print('Expected Values State {}:\nPopul: {:10.3f} \nTheta: {:10.3f} \nGamma: {:10.3f}'.format(i, popu, theA, gamA))
def fexp(x):
    return expected(os.path.join(a,filesList[x]))
    
In [18]:
    
interact(fexp, x = widgets.IntSlider(min=0,max=lastV,step=1,value=lastV));
    
    
 
 
In [19]:
    
# filesList
    
In [20]:
    
%matplotlib notebook
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from matplotlib import cm
def plotWF(wf):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    
    #something I do not like here, as X and Y are not in the order I like...
    X,Y = np.meshgrid(thesT,gamsT)
    ax.plot_wireframe(X, Y, wf.real, alpha=0.3)
    ax.plot_wireframe(X, Y, wf.imag, alpha=0.2, color='orange')
    print(wf.shape)
    
In [21]:
    
def plotdiffStates(state,frame):
    wf0  = qp.retrieve_hdf5_data(os.path.join(a,filesList[frame]),'WF')
    _,_,nstates = wf0.shape
    return plotWF(wf0[:,:,state])
interact(plotdiffStates, state = widgets.IntSlider(min=0,max=nstates-1,step=1,value=0),frame = widgets.IntSlider(min=0,max=lastV,step=1,value=0));