In [1]:
#------------------------------------------------------------------------------------
import numpy as np
import pycuda.gpuarray as gpuarray
from pycuda.tools import make_default_context
import matplotlib as matplotlib
import pylab as plt
from mpl_toolkits.mplot3d import Axes3D
#-------------------------------------------------------------------------------------
from pywignercuda_path import SetPyWignerCUDA_Path
SetPyWignerCUDA_Path()
from GPU_Dirac3D import *
In [2]:
#import pyqtgraph as pg
In [3]:
%matplotlib inline
In [4]:
class DefinePotential_Dirac3D( GPU_Dirac3D ):
def __init__ (self):
#....................Defining the geometry.....................................
gridDIM_X = 256
gridDIM_Y = 256
gridDIM_Z = 256
X_amplitude =12
Y_amplitude =16
Z_amplitude =12
dt = 0.01
timeSteps = 300
skipFrames = 100
#...................Defining the kinematic-dynamical constants.................
self.mass = 1.
self.c = 1.
self.hBar = 1.
#...................Defining the potential and initial state parameters........
hBar = 1
self.px = 1.
self.py = 0.
self.pz = 0.
self.energy = self.c*np.sqrt( (self.mass*self.c)**2 + self.px**2 + self.py**2 + self.pz**2 )
#...................Defining the potential.....................................
self.Potential_0_String = '0.5*x'
self.Potential_1_String = '0.'
self.Potential_2_String = '0.'
self.Potential_3_String = '0.'
#..................Defining the output directory/file .........................
#self.fileName = '/home/rcabrera/DATA/Dirac3D/X.hdf5'
self.fileName = '/home/rcabrera/DATA/Dirac3D/X_SpinDown.hdf5'
#..............................................................................
amplitude = (X_amplitude,Y_amplitude,Z_amplitude)
gridDIM = (gridDIM_X,gridDIM_Y,gridDIM_Z)
self.Compute_Ehrenfest_P = True
GPU_Dirac3D.__init__( self,gridDIM, amplitude ,dt,timeSteps, skipFrames = skipFrames,frameSaveMode='Spinor')
def Set_Initial_Condition_SpinUp (self):
def gaussian(x,y,z):
return np.exp( - ( (x)**2 + y**2 + z**2)/2. )
self.Psi_init = self.Spinor_Particle_SpinUp( (self.px,self.py,self.pz) , gaussian )
norm = self.Norm( self.Psi_init )
#self.FilterElectrons(1)
self.Psi_init /= norm
def Set_Initial_Condition_SpinDown (self):
def gaussian(x,y,z):
return np.exp( - ( x**2 + y**2 + z**2)/2. )
self.Psi_init = self.Spinor_Particle_SpinDown( (self.px,self.py,self.pz) , gaussian )
norm = self.Norm( self.Psi_init )
self.Psi_init /= norm
In [5]:
instance = DefinePotential_Dirac3D()
In [6]:
instance = DefinePotential_Dirac3D()
#instance.Set_Initial_Condition_SpinUp()
instance.Set_Initial_Condition_SpinDown()
print ' '
print ' dx = ', instance.dX
Psi_end = instance.Run ()
In [7]:
np.sum( instance.Psi_end )
Out[7]:
In [8]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange , instance.X_average , 'g',label= '$x^1 $')
ax.plot( instance.timeRange , instance.Y_average , 'r',label= '$x^2 $')
ax.plot( instance.timeRange , instance.Z_average , 'b--',label= '$x^3 $')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 1.2), loc=1, prop={'size':22})
Out[8]:
In [9]:
if instance.Compute_Ehrenfest_P == True:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange , instance.Px_average , 'g',label= '$p^1 $')
ax.plot( instance.timeRange , instance.Py_average , 'r',label= '$p^2 $')
ax.plot( instance.timeRange , instance.Pz_average , 'b--',label= '$p^3 $')
ax.set_ylim(-1.5,1.5)
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 0.5), loc=1, prop={'size':22})
In [10]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange , np.gradient( instance.X_average , instance.dt) , 'g',
label= '$\\frac{dx^1}{dt} $')
ax.plot( instance.timeRange , instance.Alpha1_average ,'r--' ,label='$c \\alpha^1$')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 0.5), loc=1, prop={'size':22})
Out[10]:
In [11]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange , np.gradient( instance.Y_average , instance.dt) , 'g',
label= '$\\frac{dy^1}{dt} $')
ax.plot( instance.timeRange , instance.Alpha2_average ,'r--' ,label='$c \\alpha^2$')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 0.5), loc=1, prop={'size':22})
Out[11]:
In [12]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange , np.gradient( instance.X_average , instance.dt) , 'g',
label= '$\\frac{dx^3}{dt} $')
ax.plot( instance.timeRange , instance.Alpha1_average ,'r--' ,label='$c \\alpha^3$')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 0.5), loc=1, prop={'size':22})
Out[12]:
In [13]:
if instance.Compute_Ehrenfest_P == True:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange , np.gradient( instance.Px_average , instance.dt) , 'g',
label= '$\\frac{dp^1}{dt} $')
ax.set_ylim(-2,2)
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 0.5), loc=1, prop={'size':22})
In [14]:
if instance.Compute_Ehrenfest_P == True:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange , np.gradient( instance.Py_average , instance.dt) , 'g',
label= '$\\frac{dp^2}{dt} $')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 0.5), loc=1, prop={'size':22})
In [15]:
if instance.Compute_Ehrenfest_P == True:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange , np.gradient( instance.Pz_average , instance.dt) , 'g',
label= '$\\frac{dp^3}{dt} $')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 0.5), loc=1, prop={'size':22})
In [16]:
instance.Potential_0_average.shape
Out[16]:
In [17]:
if instance.Compute_Ehrenfest_P == True:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange , instance.K_Energy_average + instance.Potential_0_average , 'g',
label= '$ Energy$')
ax.set_ylim(0,2.5)
ax.set_xlabel(r'$t$',**axis_font)
ax.grid('on')
ax.legend(bbox_to_anchor=(1.05, 0.5), loc=1, prop={'size':22})
In [18]:
fig, ax = plt.subplots(figsize=(20, 10))
ax = fig.gca(projection='3d')
ax.scatter(
instance.X_average,
instance.Y_average,
instance.Z_average
)
ax.set_xlabel(r'$x$',**axis_font)
ax.set_ylabel(r'$y$',**axis_font)
ax.set_zlabel(r'$z$',**axis_font)
ax.set_xlim(-2 , instance.X_amplitude/10)
ax.set_ylim(-0.1, 1)
ax.set_zlim(-0.1, 0.2 )
Out[18]:
In [ ]: