-
Notifications
You must be signed in to change notification settings - Fork 0
/
operators.py
157 lines (129 loc) · 5.37 KB
/
operators.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
import numpy as np
from scipy.linalg import eig,eigh,eig_banded
from scipy.integrate import simps
# Operator and Derivative classes are used in the Finite Difference Eigen
# decomposition solver and the propogate solver
class Operator(object):
""" The most general matrix formulated operator class,
the other operator classes inherit from this one"""
def __init__(self,matrix,particle):
""" Constructor takes in a proper (2darray) matrix"""
assert matrix.shape[0]==matrix.shape[1] #Must be square
self.dx = particle.dx
self.matrix = matrix
def isHermitian(self):
"""This is an expensive O(n^2) operation"""
return self.matrix.H == self.matrix
def operate(self,vector):
"""Expensive O(n^2), performs regular matrix mult on a vector"""
return np.dot(self.matrix,vector)
def repOperate(self,vector,n):
"""Calls the operate method n times on vector"""
assert n>0
operatedVector = vector
for i in np.arange(n):
operatedVector = self.operate(vector)
return operatedVector
def expValue(self,vector):
""" as you would expect, \int_{-\inf}^{\inf} \psi^* A_{op} \psi dx"""
#return np.real(np.vdot(vector,self.operate(vector)))*self.dx
return simps(np.conj(vector)*self.operate(vector),dx=self.dx)
def uncertainty(self,vector):
""" as you would expect, \int_{-\inf}^{\inf} \psi^* A_{op} \psi dx"""
#expASqrd = np.real(np.vdot(vector,self.operate(self.operate(vector))))*self.dx
expASqrd = simps(np.conj(vector)*self.repOperate(vector,2),dx = self.dx)
return np.sqrt(expASqrd - self.expValue(vector)**2)
def computeEigenFunctions(self):
if self.isHermitian():
eigenRoutine = eigh
else:
eigenRoutine = eig
lambdas,basis = eigenRoutine(self.matrix,overwrite_a=True)
basis/=np.sqrt(particle.dx) # Normalization, see note
return lambdas, basis
class BandedOp(Operator): #Must also be hermitian
def __init__(self,bandedM,particle):
"""Takes in the banded matrix in form [D0,D-1,D-2,...],
initialization and multiplication is much faster than Operator"""
assert bandedM.shape[0]%2==1 #Must have odd M
self.dx = particle.dx
self.bandedMatrix = bandedM
def isHermitian(self):
return True
def computeEigenFunctions(self,jRange=None):
if jRange: sel='i'
else: sel ='a'
#sel = 'a'
lambdas,basis = eig_banded(self.bandedMatrix,lower=True, check_finite=False,
overwrite_a_band=True, select=sel, select_range = jRange)
basis/=np.sqrt(self.dx) #For normalization
return lambdas,basis
def operate(self,vector):
""" An O(n*m) operation where m is the number of bands and n
is the size of the vector being operated on"""
M = self.bandedMatrix.shape[0]
N = self.bandedMatrix.shape[1]
#print type(self.bandedMatrix[0][0]),type(self.bandedMatrix[0])
resultVector = self.bandedMatrix[0]*vector
for i in np.arange(1,M):
resultVector += self.shiftDown(np.conj(self.bandedMatrix[i])*vector,i)
resultVector += self.shiftDown(vector,-i)*self.bandedMatrix[i]
return resultVector
def shiftDown(self,vector,n):
"""Helper function that fills shifted space with zeros"""
size = vector.shape[0]
shifted = np.zeros_like(vector)
if (n>0):
shifted[n:] = vector[:size-n]
elif (n<0):
n*=-1
shifted[:size-n] = vector[n:]
elif n==0:
shifted = vector
return shifted
class XOP(BandedOp):
def __init__(self,particle):
self.dx = particle.dx
self.bandedMatrix = np.reshape(particle.X,(1,particle.N))
class MomentumOp(BandedOp):
"""Implements the Momentum operator efficiently"""
# I could implement the stencil table here, but there is no need
def __init__(self,particle):
self.dx = particle.dx
self.bandedMatrix = np.empty([2,particle.N],dtype=np.complex_)##
self.bandedMatrix[0] = np.zeros(particle.N)
self.bandedMatrix[1] = -1j*particle.hbar* np.ones(particle.N)/(2*particle.dx)
class HamiltonianOp(BandedOp):
STENCIL_TABLE = np.array([-2.0,1.0,
-5/2.,4/3.,-1/12.,
-49/18.,3/2.,-3/20.,1/90.,
-205/72.,8/5.,-1/5.,8/315.,-1/560.])
SIGMA_CUTOFF = 30
def __init__(self,particle):
self.dx = particle.dx
self.bandedMatrix = self.getDiagonals(particle)
#self.cutOffEnergys = (np.real(self.expValue(particle.psi))+ HamiltonianOp.SIGMA_CUTOFF*self.uncertainty(particle.psi),
#?? np.real(self.expValue(particle.psi)) - HamiltonianOp.SIGMA_CUTOFF*self.uncertainty(particle.psi)
self.ERanges = self.updateERanges(particle)
#print self.cutOffEnergy
#print self.expValue(particle.psi)
#I could make further optimization by overriding operate with specialized
# But it may not be necessary
def getDiagonals(self,P):
stNum = P.stencilNum
assert stNum>0, "problem with stencilnum not being positive"
baseIndex = stNum*(stNum+1)/2 -1 #Stencil table is triangular so we use triangular #'s
diagArray = np.empty([stNum+1,P.N],dtype=np.float )
rho = -P.hbar*P.hbar/(2*P.mass*P.dx*P.dx) #Add some latex equation to explain
diagArray[0] = P.Vx+rho*HamiltonianOp.STENCIL_TABLE[baseIndex]
for i in np.arange(1,stNum+1):
coeffi = rho*HamiltonianOp.STENCIL_TABLE[i+baseIndex]
diagArray[i] = coeffi*np.ones(P.N)
return diagArray
def computeEigenFunctions(self):
return super(HamiltonianOp,self).computeEigenFunctions(self.ERange)
def updateERanges(self,particle):
expEnergy = np.real(self.expValue(particle.psi))
width = max(2*HamiltonianOp.SIGMA_CUTOFF,HamiltonianOp.SIGMA_CUTOFF*self.uncertainty(particle.psi))
self.ERange = (0,particle.N/2)#(-np.inf,expEnergy+width)
print self.ERange