def numba_calculate_psi(vecs, dim, nz, omega, t):
psi = numba_zeros((dim, dim))
for k in range(0, dim):
partial = numba_zeros(dim)
for i in range(0, nz):
num = i_to_n(i, nz)
partial += np.exp(1j * omega * t * num) * vecs[k][i]
psi[k, :] = partial
return psi
def numba_assemble_k(hf, dim, k_dim, nz, nc, omega):
hf_max = (nc-1)/2
k = numba_zeros((k_dim, k_dim))
# 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 range(-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[n_to_i(n, nc)]
while row <= stop_row and col <= stop_col:
if n == 0:
block = current_component + np.identity(dim)*omega*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
def assemble_du(nz, nz_max, dim, npm, alphas, psi, vecsstar):
# Execute the sum defining dU, taking pre-computed factors into account
du = numba_zeros((npm, dim, dim))
for n2 in range(-nz_max, nz_max + 1):
for i1 in range(0, dim):
for i2 in range(0, dim):
product = numba_outer(psi[i1], vecsstar[i2, n_to_i(-n2, nz)])
for n1 in range(-nz_max, nz_max + 1):
idn = n_to_i(n1 - n2, 2 * nz)
for c in xrange(0, npm):
du[c] += alphas[c, idn, i1, i2] * product
return du
def numba_sum_components(vec, dim):
n = vec.shape[0]
result = numba_zeros(dim)
for i in range(n):
result += vec[i]
return result