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 generate_fake_spectrum(unique_vals, dim, omega, nz):
vals = np.array([])
for i in xrange(0, nz):
offset = h.i_to_n(i, nz)
new = unique_vals + offset * omega * np.ones(dim)
vals = np.append(vals, new)
return vals
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 test_in_between(self):
self.assertEqual(h.i_to_n(37, 81), -3)
def test_middle(self):
self.assertEqual(h.i_to_n(40, 81), 0)
def test_end(self):
self.assertEqual(h.i_to_n(80, 81), 40)
def test_start(self):
self.assertEqual(h.i_to_n(0, 81), -40)