In [1]:
import numpy as np
import scipy.fftpack as fftpack
import pylab as plt
import matplotlib as matplotlib
import pycuda.gpuarray as gpuarray
#-------------------------------------------------------------------------------------
from pywignercuda_path import SetPyWignerCUDA_Path
SetPyWignerCUDA_Path()
from GPU_WignerDirac2D_4x4 import *
In [2]:
%matplotlib inline
In [3]:
class Klein(GPU_WignerDirac2D_4x4):
def __init__ (self):
#....................Defining the geometry.....................................
X_gridDIM = 512*2
P_gridDIM = 512
X_amplitude = 24
P_amplitude = 16
timeSteps = 1300
dt = 0.01 #dX/c
skipFrames = 100
#...................Defining the kinematic-dynamical constants.................
mass = 1.
c = 1.
#self.dt = dX/self.c
#...................Defining the potential and initial state parameters........
#V0 = 5.75*2
V0 = 10.
w = 1.
#.........................damping ........................................
self.gammaDamping = 0.0
epsilon = 0.01;
self.fp_Damping_String = ' p*p/sqrt( p*p + {epsilon} ) '.format( epsilon=epsilon )
self.dampingModel = 'CaldeiraLegget'
#............................................................................
self.D_Theta = 0.005
self.D_Lambda = 0.0
self.grossPitaevskiiCoefficient = 0.0
#self.pX = 9.5
self.Potential_0_String = '{0}*0.5*( 1. + tanh((x-5.)/{1}) )'.format(V0,w)
self.Potential_1_String = ' 0.'
self.Potential_2_String = ' 0.'
self.Potential_3_String = ' 0.'
#.............................................................................
GPU_WignerDirac2D_4x4.__init__(self,
X_gridDIM, P_gridDIM, X_amplitude, P_amplitude, mass, c, dt,
timeSteps,skipFrames,frameSaveMode='Density',antiParticleNorm = True, computeEnergy=True)
#.............................................................................
def Set_Initial_State (self) :
#..................Defining the output directory/file ........................
self.fileName = '/home/rcabrera/DATA/WignerDirac2D_4x4/Klein.hdf5'
self.W_init = np.empty([4,4,instance.P_gridDIM,instance.X_gridDIM],dtype = np.complex128)
init_x = 0
self.pX = 5
#
psiL1 = self.GaussianSpinor_ParticleUp( init_x , self.pX , 1, self.X - 0.5*self.Theta )
psiR1 = self.GaussianSpinor_ParticleUp( init_x , self.pX, 1, self.X + 0.5*self.Theta )
#
for i in range(4):
for j in range(4):
self.W_init[i,j][:,:] = psiL1[i]*psiR1[j].conj()
# To XP
self.Fourier_4X4_Theta_To_P(self.W_init)
#instance.FilterElectrons( self.W_init , 1)
norm = self.Wigner_4x4_Norm(self.W_init)
self.W_init *= 1./ norm
def TakabayashiAngle_CPU_(self,W):
WW = W.copy()
s = 2.*np.imag( WW[0,2] ) + 2.*np.imag( WW[1,3] );
c = np.real( WW[0,0] ) + np.real( WW[1,1] ) - np.real( WW[2,2] ) - np.real( WW[3,3] );
return np.arctan2( s , np.abs(c) );
def Set_Initial_State_Takabayashi (self) :
#..................Defining the output directory/file ........................
self.fileName = '/home/rcabrera/DATA/WignerDirac2D_4x4/Klein-Takabayashi0.0.hdf5'
self.W_init = np.empty([4,4,instance.P_gridDIM,instance.X_gridDIM],dtype = np.complex128)
init_x = 0
self.pX = 5
#
psiL_ = self.GaussianSpinor_ParticleUp( init_x , self.pX , 1, self.X - 0.5*self.Theta )
psiR_ = self.GaussianSpinor_ParticleUp( init_x , self.pX, 1, self.X + 0.5*self.Theta )
# Including Yvone-Takabayashi angle
beta = 0
psiL = psiL_.copy()
psiR = psiR_.copy()
psiL[0] = psiL_[0] *np.cos(beta/2.) + 1j*psiL_[2] *np.sin(beta/2.)
psiL[1] = psiL_[1] *np.cos(beta/2.) + 1j*psiL_[3] *np.sin(beta/2.)
psiL[2] = psiL_[2] *np.cos(beta/2.) + 1j*psiL_[0] *np.sin(beta/2.)
psiL[3] = psiL_[3] *np.cos(beta/2.) + 1j*psiL_[1] *np.sin(beta/2.)
psiR[0] = psiR_[0] *np.cos(beta/2.) + 1j*psiR_[2] *np.sin(beta/2.)
psiR[1] = psiR_[1] *np.cos(beta/2.) + 1j*psiR_[3] *np.sin(beta/2.)
psiR[2] = psiR_[2] *np.cos(beta/2.) + 1j*psiR_[0] *np.sin(beta/2.)
psiR[3] = psiR_[3] *np.cos(beta/2.) + 1j*psiR_[1] *np.sin(beta/2.)
#
for i in range(4):
for j in range(4):
self.W_init[i,j][:,:] = psiL[i]*psiR[j].conj()
# To XP
self.Fourier_4X4_Theta_To_P(self.W_init)
#instance.FilterElectrons( self.W_init , 1)
norm = self.Wigner_4x4_Norm(self.W_init)
self.W_init *= 1./ norm
betaArray = self.TakabayashiAngle_CPU(self.W_init)
plt.plot(
np.sort( np.abs(betaArray.flatten()) ) )
plt.imshow( fftpack.fftshift( betaArray ) , origin='lower',interpolation='nearest' )
plt.colorbar()
In [4]:
instance = Klein()
In [5]:
instance.Set_Initial_State_Takabayashi()
In [6]:
#instance.Set_Initial_State_Takabayashi()
#instance.Set_Initial_State()
instance.Run ()
In [7]:
def PlotWigner(W):
W0 = fftpack.fftshift( instance.Wigner_4X4__SpinTrace( W ).real)
x_min = -instance.X_amplitude
x_max = instance.X_amplitude - instance.dX
p_min = -instance.P_amplitude
p_max = instance.P_amplitude - instance.dP
global_max = 0.31 # Maximum value used to select the color range
global_min = -0.27 #
print 'min = ', np.min( W0 ), ' max = ', np.max( W0 )
print 'normalization = ', np.sum( W0 )*instance.dX*instance.dP
zero_position = abs( global_min) / (abs( global_max) + abs(global_min))
wigner_cdict = {'red' : ((0., 0., 0.),
(zero_position, 1., 1.),
(1., 1., 1.)),
'green' : ((0., 0., 0.),
(zero_position, 1., 1.),
(1., 0., 0.)),
'blue' : ((0., 1., 1.),
(zero_position, 1., 1.),
(1., 0., 0.)) }
wigner_cmap = matplotlib.colors.LinearSegmentedColormap('wigner_colormap', wigner_cdict, 256)
#wigner_cmap = plt.colors.LinearSegmentedColormap('wigner_colormap', wigner_cdict, 256)
fig, ax = plt.subplots(figsize=(20, 7))
cax = ax.imshow( W0 ,origin='lower',interpolation='nearest',\
extent=[x_min, x_max, p_min, p_max], vmin= global_min, vmax=global_max, cmap=wigner_cmap)
ax.set_xlabel('x')
ax.set_ylabel('p')
#ax.set_xlim((x_min,x_max))
#ax.set_ylim((-5 , p_max/3.5))
#ax.set_ylim((-16,16))
ax.set_aspect(1)
ax.grid('on')
In [8]:
def PlotMarginal_P(instance):
W_0 = fftpack.fftshift( instance.Wigner_4X4__SpinTrace(instance.W_init).real)
print ' norm = ', np.sum(W_0).real*instance.dX*instance.dP
fig, ax = plt.subplots(figsize=(10, 5))
prob_P = np.sum(W_0,axis=1)*instance.dX
ax.plot(instance.P_range, prob_P , label = 'init')
W_0 = fftpack.fftshift( instance.Wigner_4X4__SpinTrace(instance.W_end).real)
print ' norm = ', np.sum(W_0).real*instance.dX*instance.dP
prob_P = np.sum(W_0,axis=1)*instance.dX
ax.plot(instance.P_range, prob_P , label = 'final')
ax.set_xlim(-18,18)
ax.set_xlabel('p')
ax.set_ylabel('Prob')
ax.grid('on')
ax.legend(bbox_to_anchor=(0.75, 0.5), loc=2, prop={'size':22})
def PlotMarginal_X(instance):
W_0 = fftpack.fftshift( instance.Wigner_4X4__SpinTrace(instance.W_init).real)
fig, ax = plt.subplots(figsize=(10, 5))
prob_X = np.sum(W_0,axis=0)*instance.dP
ax.plot(instance.X_range, prob_X , label = 'init')
W_0 = fftpack.fftshift( instance.Wigner_4X4__SpinTrace(instance.W_end).real)
prob_X = np.sum(W_0,axis=0)*instance.dP
ax.plot(instance.X_range, prob_X , label = 'final')
ax.set_xlabel('x')
ax.set_ylabel('Prob')
ax.grid('on')
ax.legend(bbox_to_anchor=(0.75, 0.5), loc=2, prop={'size':22})
In [9]:
PlotWigner(10*instance.W_init)
In [10]:
PlotWigner( 10*instance.W_end )
print ' time = ', instance.timeRange[-1]
In [11]:
PlotMarginal_P( instance )
PlotMarginal_X( instance )
In [12]:
WendFilter = instance.W_end.copy()
instance.FilterElectrons( WendFilter , 1)
In [13]:
PlotWigner( 10*WendFilter )
In [14]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange[1:] , np.gradient( instance.X_Average.real , instance.dt) , 'g',
label= '$\\frac{dx^1}{dt} $')
ax.plot( instance.timeRange[1:] , instance.Alpha_1_Average.real ,'r--' ,label='$c \\alpha^1$')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 1), loc=1, prop={'size':22})
Out[14]:
In [15]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange[1:] , np.gradient( instance.P_Average.real , instance.dt) , 'g',
label= '$\\frac{dp^1}{dt} $')
ax.plot( instance.timeRange[1:] ,
-instance.D_1_Potential_0_Average.real - 2.*instance.mass*instance.gammaDamping*instance.Alpha_1_Average.real ,'r--' ,label='$-c e\, \\partial_1 A_0 $')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 1), loc=2, prop={'size':22})
Out[15]:
In [16]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange[1:] ,
2*np.gradient( instance.XP_Average.real , instance.dt) , 'g',
label= '$\\frac{d}{dt}( x^1 p_1 ) $')
ax.plot( instance.timeRange[1:] ,
-2*instance.X1_D_1_Potential_0_Average.real + 2*instance.c*instance.P1_Alpha_1_Average.real -4.*instance.mass*instance.gammaDamping*instance.X1_Alpha_1_Average,
'r--' ,label='$-2 x^1 \\partial_1 e A_0 + 2 c p_1 \\alpha^1$')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 1), loc=2, prop={'size':22})
Out[16]:
In [17]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange[1:] ,
np.gradient( instance.XX_Average.real , instance.dt) , 'g',
label= '$\\frac{d}{dt}( x^1 x_1 ) $')
ax.plot( instance.timeRange[1:] ,
2*instance.c*instance.X1_Alpha_1_Average.real,
'r--' ,label='$2 c x_1 \\alpha^1$')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 1), loc=2, prop={'size':22})
Out[17]:
In [18]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange[1:] ,
np.gradient( instance.PP_Average.real , instance.dt) , 'g',
label= '$\\frac{d}{dt}( x^1 x_1 ) $')
ax.plot( instance.timeRange[1:] ,
-2*instance.P1_D_1_Potential_0_Average.real - \
4.*instance.mass*instance.gammaDamping*instance.P1_Alpha_1_Average.real,
'r--' ,label='$- c p^1 \\partial_1 A_0 $')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 1), loc=2, prop={'size':22})
Out[18]:
In [19]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange[1:] ,
instance.antiParticle_population.real , 'g',
label= 'Antiparticle population')
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 1), loc=2, prop={'size':22})
Out[19]:
In [20]:
axis_font = {'size':'24'}
fig, ax = plt.subplots(figsize=(20, 7))
ax.plot( instance.timeRange[1:] ,
instance.Dirac_energy.real , 'g',
label= 'Energy')
ax.set_ylim(0, 8)
ax.set_xlabel(r'$t$',**axis_font)
ax.legend(bbox_to_anchor=(1.05, 1), loc=2, prop={'size':22})
Out[20]:
In [21]:
plt.imshow(
fftpack.fftshift( np.abs(instance.TakabayashiAngle_init ) ),
origin='lower',interpolation='nearest')
plt.colorbar()
Out[21]:
In [22]:
plt.imshow(
fftpack.fftshift( np.abs(instance.TakabayashiAngle_end ) ),
origin='lower',interpolation='nearest')
plt.colorbar()
Out[22]:
In [23]:
plt.plot(
np.sort( np.abs(instance.TakabayashiAngle_init.flatten() ) ) ,'-',
label='init')
plt.plot(
np.sort( np.abs(instance.TakabayashiAngle_end.flatten() ) ),'r--',
label='end')
plt.legend()
Out[23]:
In [23]: