In [ ]:

    

import numpy as np
import h5py
import matplotlib.pyplot as plt 
from Loader3D import *
%matplotlib inline


    

In [ ]:

    

filepath="../../build/test.h5"
x0 = -6.0
x1 = 6.0
y0 = x0
y1 = x1
z0 = x0 
z1 = x1 
nx = 500
ny = 500
nz = 300
w1 = 0.1 
a1 = 2*np.pi/1000
l1 = 1.0 
dt = 0.001
nt = 1000
t0 = 0
t = np.linspace(0, dt*nt, nt, endpoint=False)
x = np.linspace(x0,x1,nx, endpoint=False)
y = np.linspace(y0,y1,ny, endpoint=False)
z = np.linspace(z0,z1,nz, endpoint=False)
x_mesh, y_mesh = np.meshgrid(x,y)


    

In [ ]:

    

class W_3D:
    """
    Implementation of a linear time-dependent disturbance 
    term. 
    """
    def __init__(self, a1, a2, a3, w1, w2, w3, ft):
        """
        ai: strength of the disturbance in dimension i
        wi: frequency of the disturbance in dimension i
        """
        self.a1 = a1 
        self.a2 = a2
        self.a3 = a3
        self.w1 = w1 
        self.w2 = w2 
        self.w3 = w3
        self.ft = ft
    
    def __call__(self, t, x=0, y=0, z=0):
        """
        Make the object callable so that 
        it can be used like a function, which
        it should in fac represent
        t: time
        """
        s1 = self.a1 * self.ft(self.w1* t) 
        s2 = self.a2 * self.ft(self.w2* t)
        s3 = self.a3 * self.ft(self.w3* t)
        return s1 + s2 + s3
    
    
class VHarm_3D:
    """
    Implementation of the 3D harmonic potential
    in atomic units. 
    """
    def __init__(self):
        pass
    
    def __call__(self, x, y, z, t):
        """
        Use objects of the class like 
        a function, since this what they
        are supposed to be represented. 
        x: x coordinate in a.u.
        y: y coordinate in a.u.
        z: z coordinate in a.u.
        t: time in a.u.
        """
        return 0.5 * (x**2 + y**2 + z **2)


    

In [ ]:

    

class U_Lin_Dist_3D:
    """
    Time evolution operator for a
    time-dependent disturbance term.
    """
    
    def __init__(self, E0, W1, dt):
        """
        E0: Eigenenergy for H_0 Term
        W1: Function object reprensenting t
        he time-dependent disturbance. It is important
        that [H(t),H(t')] = 0 for arbitrary points
        in time with t != t'. 
        dt: Chosen timestep
        """
        self.E0 = E0
        self.W1 = W1
        self.dt = dt
        
    def __call__(self, t):
        """
        Calling the timeevolution 
        operator.
        t: array/numpy linspace of the points in 
        time from [t0,t1]
        """
        w_val = self.W1(t)
        print(w_val)
        return np.exp(-1j * (self.E0 * (self.dt * t.size) + np.trapz(self.W1(t),dx=self.dt)))


    

In [ ]:

    

Psi0 = lambda x: np.pi**(-0.25) * np.exp(-x**2*0.5)
PsiD = lambda x,y,z: Psi0(x) * Psi0(y) * Psi0(0)

psi_an = PsiD(x_mesh, y_mesh, 0) + 0 * 1j
W = W_3D(a1, a1, a1, w1, w1, w1, np.sin)
V = VHarm_3D()
U = U_Lin_Dist_3D(1.5, W, dt)
psi_an *= U(t)


    

In [ ]:

    

loader = CLoader3D(nx, ny, nz, filepath)
psi_num = loader.get_complex_data("/real", "/imag")
psi_num = psi_num[:, :, int(nz/2)]
fig = plt.figure(figsize=(10,6))

ax = fig.gca(projection='3d')
print(np.max(np.abs(psi_an)**2)-np.max(np.abs(psi_num)**2))

surf = ax.plot_surface(x_mesh, y_mesh, psi_num.imag, cmap=cm.magma_r,linewidth=1,antialiased=False)
plt.title("Time evolution 3D Harm. Osc.")
ax.set_xlabel(r"$x \; (a.u.)$")
ax.set_ylabel(r"$y \; (a.u.)$")
ax.set_zlabel(r"$|\Psi|^2$")

fig.colorbar(surf, shrink=0.5, aspect=5)


plt.show()


    

In [ ]: