

In [2]:

    
def load_wf(filepath):
    """
    Load a distinct energy levels 
    as a numpy array into memory
    """
    file = h5py.File(filepath)
    if "/numres" in file:
        data = np.array(file["/numres"])
        file.close()
        return data
    else:
        file.close()
        return -1


def load_timestep(filepath,nt,t0=0,):
    """
    Load a timestep from a result file
    and returns it as a complex numpy array.
    
    """
    file = h5py.File(filepath)
    rl = "/dset"+str(nt)+"real"
    im = "dset"+str(nt)+"img"
    if (rl in file) and (im in file):
        imag = np.array(file[rl])
        real = np.array(file[im])
        res = real + 1j * imag
        file.close()
        return res
    else:
        file.close()
        return -1
    
def get_coeff(filepath_stat, filepath_ev, t0, tmax, nt, toff,xmin, xmax, nx):
    """
    Returns the absolute square of the evolution coefficients c_n = <psi_n | psi_x >
    
    """
    # get the complex initial wavefunction
    psin = load_wf(filepath_stat)*(1+1j*0)
    cn = np.zeros(np.int32(nt/toff))
    c1 = np.zeros(np.int32(nt/toff))
    c2 = np.zeros(np.int32(nt/toff))
    c3 = np.zeros(np.int32(nt/toff))
    c4 = np.zeros(np.int32(nt/toff))
    
    dx = (xmax-xmin)/nx
    
    with progressbar.ProgressBar(max_value=int(nt)) as bar:
        n1 = 1 
        n2 = 2
        for i in range(0, int(nt),toff):
            index = np.int32(i/toff)
            psik = load_timestep(filepath_ev, i)
            bar.update(i)
            cn[index] = np.abs(np.trapz(np.conj(psik)*psin,dx=dx))**2
        #cn = np.abs(psin)**2
    return cn

fig = plt.figure(figsize=(14,10))
plt.subplot(221)
k4 = get_coeff("45_au_s2.h5","../../build/res.h5",0,5000,1e5,100,-45,45,1e5)
t = np.linspace(0, 5000,1e3)
plt.xlabel(r"time $(a.u.)$")
plt.ylabel(r"$|C_2(t)|^2$")
plt.plot(t,k4,label=r"$c_2(t)$")
plt.legend()
plt.subplot(222)
k3 = get_coeff("45_au_s3.h5","../../build/res.h5",0,5000,1e5,100,-45,45,1e5)
plt.xlabel(r"time $(a.u.)$")
plt.ylabel(r"$|C_n(t)|^2$")
plt.plot(t,k3,label=r"$c_3(t)$",c="green")
plt.legend()
plt.subplot(223)
k4 = get_coeff("45_au_s4.h5","../../build/res.h5",0,5000,1e5,100,-45,45,1e5)
plt.xlabel(r"time $(a.u.)$")
plt.ylabel(r"$|C_n(t)|^2$")
plt.plot(t,k4,label=r"$c_4(t)$",c="red")
plt.legend()

plt.subplot(224)
k4 = get_coeff("45_au_s5.h5","../../build/res.h5",0,5000,1e5,100,-45,45,1e5)
plt.xlabel(r"time $(a.u.)$")
plt.ylabel(r"$|C_n(t)|^2$")
plt.plot(t,k4,label=r"$c_5(t)$",c="purple")


plt.legend()
plt.show()









    



100% (100000 of 100000) |#######################################################| Elapsed Time: 0:00:06 Time: 0:00:06
100% (100000 of 100000) |#######################################################| Elapsed Time: 0:00:06 Time: 0:00:06
100% (100000 of 100000) |#######################################################| Elapsed Time: 0:00:06 Time: 0:00:06
100% (100000 of 100000) |#######################################################| Elapsed Time: 0:00:06 Time: 0:00:06