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_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 calculate_factors(dk, nz, nz_max, dim, npm, vals, vecs, vecsstar, omega, t):
# Factors in the sum for dU that only depend on dn=n1-n2, and therefore
# can be computed more efficiently outside the "full" loop
factors = np.empty([npm, 2*nz+1, dim, dim], dtype=np.complex128)
for dn in xrange(-nz_max*2, 2*nz_max+1):
idn = n_to_i(dn, 2*nz)
for i1 in xrange(0, dim):
for i2 in xrange(0, dim):
v1 = np.roll(vecsstar[i1], dn, axis=0) # not supported by numba!
for c in xrange(0, npm):
factors[c, idn, i1, i2] = (integral_factors(vals[i1], vals[i2], dn, omega, t) *
expectation_value(dk[c], v1, vecs[i2]))
return factors
def test_end(self):
self.assertEqual(h.n_to_i(40, 81), 80)
def test_start(self):
self.assertEqual(h.n_to_i(-40, 81), 0)
def test_too_small_a_lot(self):
self.assertEqual(h.n_to_i(-6 - 14, 7), h.n_to_i(1, 7))
def test_too_small_a_bit(self):
self.assertEqual(h.n_to_i(-6, 7), h.n_to_i(1, 7))
def test_too_big_a_lot(self):
self.assertEqual(h.n_to_i(5 + 7, 7), h.n_to_i(-2, 7))
def test_in_between(self):
self.assertEqual(h.n_to_i(-3, 81), 37)
def test_middle(self):
self.assertEqual(h.n_to_i(0, 81), 40)
