def calculate_psi(vecs, p):
# Given an array of eigenvectors vecs,
# sum over all frequency components in each,
# weighted by exp(- i omega t n), with n
# being the Fourier index of the component
psi = np.zeros([p.dim, p.dim], dtype='complex128')
for k in xrange(0, p.dim):
partial = np.zeros(p.dim, dtype='complex128')
for i in xrange(0, p.nz):
num = h.i_to_n(i, p.nz)
partial += np.exp(1j*p.omega*p.t*num)*vecs[k][i]
psi[k, :] = partial
return psi
def assemble_k(hf, p):
hf_max = (p.nc-1)/2
nz = p.nz
nc = p.nc
dim = p.dim
omega = p.omega
k = np.zeros([p.k_dim, p.k_dim], dtype='complex128')
# Assemble K by placing each component of Hf in turn, which
# for a fixed Fourier index lie on diagonals, with 0 on the
# main diagonal, positive numbers on the right and negative on the left
#
# The first row is therefore essentially Hf(0) Hf(1) ... Hf(hf_max) 0 0 0 ...
# The last row is then ... 0 0 0 Hf(-hf_max) ... Hf(0)
# Note that the main diagonal acquires a factor of omega*identity*(row/column number)
for n in xrange(-hf_max, hf_max+1):
start_row = max(0, n) # if n < 0, start at row 0
start_col = max(0, -n) # if n > 0, start at col 0
stop_row = min((nz-1)+n, nz-1)
stop_col = min((nz-1)-n, nz-1)
row = start_row
col = start_col
current_component = hf[h.n_to_i(n, nc)]
while row <= stop_row and col <= stop_col:
if n == 0:
block = current_component + np.identity(dim)*omega*h.i_to_n(row, nz)
bm.set_block_in_matrix(block, k, dim, nz, row, col)
else:
bm.set_block_in_matrix(current_component, k, dim, nz, row, col)
row += 1
col += 1
return k