1
#This file is part of QuTIP.
3
# QuTIP is free software: you can redistribute it and/or modify
4
# it under the terms of the GNU General Public License as published by
5
# the Free Software Foundation, either version 3 of the License, or
6
# (at your option) any later version.
8
# QuTIP is distributed in the hope that it will be useful,
9
# but WITHOUT ANY WARRANTY; without even the implied warranty of
10
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11
# GNU General Public License for more details.
13
# You should have received a copy of the GNU General Public License
14
# along with QuTIP. If not, see <http://www.gnu.org/licenses/>.
16
# Copyright (C) 2011, Paul D. Nation & Robert J. Johansson
18
###########################################################################
19
from multiprocessing import Pool,cpu_count
21
from scipy.integrate import *
22
import scipy.linalg as la
29
def mcsolve(Heff,psi0,tlist,ntraj,collapse_ops,expect_ops):
30
mc=MC_class(Heff,psi0,tlist,ntraj,collapse_ops,expect_ops)
32
if mc.num_collapse==0 and mc.num_expect==0:
34
elif mc.num_collapse==0 and mc.num_expect!=0:
36
elif mc.num_collapse!=0 and mc.num_expect==0:
37
return mc.collapse_times_out,mc.which_op_out
38
elif mc.num_collapse!=0 and mc.num_expect!=0:
43
######---Monte-Carlo class---######
45
def __init__(self,Heff,psi0,tlist,ntraj,collapse_ops,expect_ops):
48
self.num_times=len(tlist)
49
self.collapse_ops=collapse_ops
50
self.num_collapse=len(collapse_ops)
51
self.expect_ops=expect_ops
52
self.num_expect=len(expect_ops)
53
self.Hdata=-1.0j*Heff.data# extract matrix and multiply Heff by -1j so user doesn't have to.
54
self.psi_in=psi0.full() #need dense matrix as input to ODE solver.
55
self.psi_dims=psi0.dims
56
self.psi_shape=psi0.shape
57
if self.num_collapse==0 and self.num_expect==0:
58
self.psi_out=array([qobj() for k in xrange(self.num_times)])#preallocate array of qobjs
59
elif self.num_collapse==0 and self.num_expect!=0:#no collpase expectation values
61
self.isher=isherm(self.expect_ops)#checks if expectation operators are hermitian
62
for jj in xrange(self.num_expect):#expectation operators evaluated at initial conditions
64
self.expect_out.append(zeros(self.num_times))
66
self.expect_out.append(zeros(self.num_times,dtype=complex))
67
elif self.num_collapse!=0:
68
self.bar=Counter(self.ntraj)#create GUI element if avaliable
69
#extract matricies from collapse operators
70
self.norm_collapse_data=array([(op.dag()*op).data for op in self.collapse_ops])
71
self.collapse_ops_data=array([op.data for op in self.collapse_ops])
72
#preallocate #ntraj arrays for state vectors, collapse times, and which operator
73
self.collapse_times_out=zeros((self.ntraj),dtype=ndarray)
74
self.which_op_out=zeros((self.ntraj),dtype=ndarray)
75
if self.num_expect==0:# if no expectation operators, preallocate #ntraj arrays for state vectors
76
self.psi_out=zeros((self.ntraj),dtype=ndarray)
77
else: #preallocate array of lists for expectation values
78
self.isher=isherm(self.expect_ops)
79
self.expect_out=[[] for x in xrange(self.ntraj)]
80
#------------------------------------------------------------------------------------
81
def callback(self,results):
83
if self.num_expect==0:#output state-vector
84
self.psi_out[r]=array([qobj(psi,dims=self.psi_dims,shape=self.psi_shape) for psi in results[1]])
85
else:#output expectation values
86
self.expect_out[r]=results[1]
87
self.collapse_times_out[r]=results[2]
88
self.which_op_out[r]=results[3]
90
#########################
92
if self.num_collapse==0:
93
if self.ntraj!=1:#check if ntraj!=1 which is pointless for no collapse operators
94
print 'No collapse operators specified.\nRunning a single trajectory only.\n'
95
if self.num_expect==0:# return psi QOBJ at each requested time
96
self.psi_out=no_collapse_psi_out(self.Hdata,self.psi_in,self.times,self.num_times,self.psi_dims,self.psi_shape,self.psi_out)
97
else:# return expectation values of requested operators
98
self.expect_out=no_collapse_expect_out(self.Hdata,self.psi_in,self.times,self.expect_ops,self.num_expect,self.num_times,self.psi_dims,self.psi_shape,self.expect_out,self.isher)
99
elif self.num_collapse!=0:
100
args=(self.Hdata,self.psi_in,self.times,self.num_times,self.num_collapse,self.collapse_ops_data,self.norm_collapse_data,self.num_expect,self.expect_ops,self.isher)
101
pool=Pool(processes=cpu_count())
102
for nt in xrange(0,self.ntraj):
103
pool.apply_async(mc_alg_evolve,args=(nt,args),callback=self.callback)
106
self.bar.finish() #stop GUI if running
110
######---return psi at requested times for no collapse operators---######
111
def no_collapse_psi_out(Hdata,psi_in,tlist,num_times,psi_dims,psi_shape,psi_out):
112
def RHS(t,psi):#define RHS of ODE
113
return Hdata*psi #cannot use dot(a,b) since mat is matrix and not array.
115
ODE.set_integrator('zvode',method='adams',nsteps=500,atol=1e-6) #initialize ODE solver for RHS
116
ODE.set_initial_value(psi_in,tlist[0]) #set initial conditions
117
psi_out[0]=qobj(psi_in,dims=psi_dims,shape=psi_shape)
118
for k in xrange(1,num_times):
119
ODE.integrate(tlist[k],step=0) #integrate up to tlist[k]
121
psi_out[k]=qobj(ODE.y/la.norm(ODE.y),dims=psi_dims,shape=psi_shape)
123
raise ValueError('Error in ODE solver')
125
#------------------------------------------------------------------------
128
######---return expectation values at requested times for no collapse operators---######
129
def no_collapse_expect_out(Hdata,psi_in,tlist,expect_ops,num_expect,num_times,psi_dims,psi_shape,expect_out,isher):
130
######---Define RHS of ODE---##############
132
return Hdata*psi #cannot use dot(a,b) since mat is matrix and not array.
133
######-------------------------------------
135
ODE.set_integrator('zvode',method='adams',nsteps=500,atol=1e-6) #initialize ODE solver for RHS
136
ODE.set_initial_value(psi_in,tlist[0]) #set initial conditions
137
for jj in xrange(num_expect):
138
expect_out[jj][0]=mc_expect(expect_ops[jj],psi_in,isher[jj])
139
for k in xrange(1,num_times):
140
ODE.integrate(tlist[k],step=0) #integrate up to tlist[k]
142
state=ODE.y/la.norm(ODE.y)
143
for jj in xrange(num_expect):
144
expect_out[jj][k]=mc_expect(expect_ops[jj],state,isher[jj])
146
raise ValueError('Error in ODE solver')
147
return expect_out #return times and expectiation values
148
#------------------------------------------------------------------------
151
######---single-trajectory for monte-carlo---###########
152
def mc_alg_evolve(nt,args):
154
Monte-Carlo algorithm returning state-vector at times tlist[k] for a single trajectory
156
Hdata,psi_in,tlist,num_times,num_collapse,collapse_ops_data,norm_collapse_data,num_expect,expect_ops,isher=args
158
return Hdata*psi #cannot use dot(a,b) since mat is matrix and not array.
160
psi_out=array([zeros((len(psi_in),1),dtype=complex) for k in xrange(num_times)])#preallocate real array of qobjs
165
if i==1:#preallocate real array of zeros
166
expect_out.append(zeros(num_times))
167
else:#preallocate complex array of zeros
168
expect_out.append(zeros(num_times,dtype=complex))
169
for jj in xrange(num_expect):
170
expect_out[jj][0]=mc_expect(expect_ops[jj],psi_in,isher[jj])
171
collapse_times=[] #times at which collapse occurs
172
which_oper=[] # which operator did the collapse
174
rand_vals=random.random(2)#first rand is collapse norm, second is which operator
176
ODE.set_integrator('zvode',method='adams',nsteps=500,atol=1e-6) #initialize ODE solver for RHS
177
ODE.set_initial_value(psi_in,tlist[0]) #set initial conditions
178
for k in xrange(1,num_times):
179
last_t=ODE.t;last_y=ODE.y
181
while ODE.successful() and ODE.t<tlist[k]:
182
ODE.integrate(tlist[k],step=step_flag) #integrate up to tlist[k], one step at a time.
184
ODE.set_integrator('zvode',method='adams',nsteps=500,atol=1e-6,first_step=(tlist[k]-last_t))
185
ODE.set_initial_value(last_y,last_t)
188
psi_nrm=la.norm(ODE.y)
189
if psi_nrm<0.0:#safety check for norm<0
190
psi_nrm,ODE=norm_safety(ODE,tlist,psi_nrm,last_y,last_t)#find non-zero psi norm
191
if psi_nrm<=rand_vals[0]:#collpase has occured
192
collapse_times.append(ODE.t)
194
n_dp=array([real(dot(ODE.y.conj().T,op*ODE.y)[0,0]) for op in norm_collapse_data])
196
for j in xrange(num_collapse):
199
which_oper.append(j) #record which operator did collapse
200
state=collapse_ops_data[j]*ODE.y
201
psi_nrm=la.norm(state)
205
rand_vals=random.random(2)
206
#last_y=ODE.y;last_t=ODE.t
207
break #breaks out of for-loop
208
last_t=ODE.t;last_y=ODE.y
209
###--after while loop--####
210
psi_nrm=la.norm(ODE.y)
212
psi_out[k]=ODE.y/psi_nrm
215
for jj in xrange(num_expect):
216
expect_out[jj][k]=mc_expect(expect_ops[jj],state,isher[jj])
218
return nt,psi_out,array(collapse_times),array(which_oper)
220
return nt,expect_out,array(collapse_times),array(which_oper)
221
#------------------------------------------------------------------------------------------
225
######---if psi norm is less than zero---###########
226
def norm_safety(ODE,tlist,psi_nrm,last_y,last_t):
227
print 'wavefunction norm below zero, reducing step size.'
229
while ODE.successful() and la.norm(ODE.y)<0 and ntrys<=10:#reduce step-size by half and integrate again until norm>0 or more than 10 attempts
230
ODE.set_integrator('zvode',method='adams',nsteps=500,atol=1e-6,first_step=(ODE.t-last_t)/2.0)
231
ODE.set_initial_value(last_y,last_t)
232
ODE.integrate(tlist[k],step=1)
234
psi_nrm=la.norm(ODE.y)
235
if psi_nrm<0:# if norm still <0 return error
236
raise ValueError('State vector norm<0 after reducing ODE step-size 10 times.')
237
ODE.set_integrator('zvode',method='adams',nsteps=500,atol=1e-6)
239
#------------------------------------------------------------------------
242
def mc_expect(oper,state,isherm):
244
return real(dot(conj(state).T,oper.data*state))
246
return complex(dot(conj(state).T,oper.data*state))
b'\\ No newline at end of file'
added
removed
qutip/mcsolve.py
