def _get_wfs(self):
"""Calculate wavefunctions for all surfaces"""
# Get wavefunctions and sort them
wfs = self._get_wfs_raw()
wfs = [wfs[i] for i in self.idx]
# Fix phase and ordering
wfsold = self.wfs
S = get_overlaps(wfs, wfsold)
if self.switch:
# See if surfaces have switched places
this_permutation = np.abs(np.array(S.round(), dtype=int))
if world.rank == MASTER:
print '# DEBUG ######################'
print 'S:', S.reshape(-1)
print 'perm:', this_permutation.reshape(-1)
assert np.abs(np.linalg.det(this_permutation)) == 1
# Save the redordered index for later use
self.idx = np.dot(this_permutation,
self.idx)
self.data['idx'] = self.idx.copy()
self.permutation_matrix = np.dot(this_permutation,
self.permutation_matrix)
self.data['permutation_matrix'] = self.permutation_matrix.copy()
# Re-sort if necessary
this_idx = np.dot(this_permutation,
np.arange(len(self.idx), dtype=int))
wfs = [wfs[i] for i in this_idx]
S_permuted = np.dot(this_permutation, S)
S = get_overlaps(wfs, wfsold)
assert np.max(np.abs(S - S_permuted)) < 1e-3
# Change sign of wavefunctions if overlaps with itself is negative
for i in range(S.shape[0]):
if S[i, i] < 0.0:
wfs[i] *= -1
S = get_overlaps(wfs, wfsold)
self.data['S'] = S.copy()
return wfs
def set_position(self, x):
"""Set position"""
self.initialize()
ParticleOnPES.set_position(self, x)
wfsold = self.wfs
wfs = self._get_wfs()
# <psi_i|d/dt|psi_j>
ddt_wfs = []
for i in range(len(wfs)):
ddt_wf = (wfs[i] - wfsold[i]) / self.dt
ddt_wfs.append(ddt_wf)
V = get_overlaps(wfs, ddt_wfs)
self.data['V'] = V.copy()
# Energies
Epes = []
for i in self.idx:
pes = self.peslist[i]
Epes.append(pes.get_value(self.x))
self.data['Epes'] = Epes
# Hamiltonian
H = np.diag(Epes) - 1.0j * V
self.data['H'] = H.copy()
# Propagate expansion coefficients
cold = self.c.copy()
if not self.newint:
### Old integration
c = -1.0j / hbar * np.dot(H, cold) * self.dt + cold
else:
### New integration
# d d
# i -- c = H c => -- c = -i H c
# dt dt
#
# e_i, v_i is eigenvals/vectors of (-i H)
#
# then
#
# c(t) = k_1 v_1 exp(e_2 t) + k_2 v_2 exp(e_2 t)
#
# b.c: v k = c => k = v\c
if world.rank == MASTER:
print 'New integration'
e, v = np.linalg.eig(-1.0j * H / hbar)
k = np.linalg.solve(v, cold)
c = k[0] * v[:, 0] * np.exp(e[0] * self.dt) + \
k[1] * v[:, 1] * np.exp(e[1] * self.dt)
#
self.data['c'] = c.copy()
c /= np.linalg.norm(c)
self.data['cnorm'] = c.copy()
self.c = c.copy()
# Save wavefunctions for later
self.wfs = wfs
# Stuff done for logging purposes
self.data['nlist'] = self.nlist[:]
self.data['eigenvalues'] = \
[self.calc.get_eigenvalues()[n] for n in self.nlist]
self.data['potential_energy'] = self.calc.get_potential_energy()
