import numpy as np from numba import njit import scipy.sparse as sp from . import hamiltonian as ham from . import iterators as it fctexpblock_debug = False # (parallel) worker function that receives assignments from a queue (-> queue) # calls core algorithm and writes the results in the shared arrays part_fct and fct_exp # p_max, pos_charge: parameters of the hamiltonian # v_tilde: fourier transform of the potential, calculated in advance # kin_coeff: coefficient of the kinetic contribution def fctExpWorker(mp_queue, p_max, pos_charge, spin_boundary_bits, v_tilde, kin_coeff, mp_part_fct, mp_fct_exp, mp_exponent_shift, shift_init, shift_thr, mu_grid, beta_grid, mp_lock, *args): local_part_fct = np.zeros((mu_grid.shape[0], beta_grid.size)) local_fct_exp = np.zeros((len(args), mu_grid.shape[0], beta_grid.size)) local_exponent_shift = np.ones((mu_grid.shape[0], beta_grid.size)) * shift_init while True: init_fstate, spin_dist, p_total, symmetry, dia_shift = mp_queue.get() if spin_dist is None: break spin_dist = np.array(spin_dist) fctExpBlock(init_fstate, spin_dist, p_total, dia_shift, spin_boundary_bits, v_tilde, kin_coeff, p_max, local_part_fct, local_fct_exp, local_exponent_shift, shift_thr, mu_grid, beta_grid, symmetry, *args) with mp_lock: for m_idx in range(mu_grid.shape[0]): for b_idx in range(beta_grid.size): global_shift = mp_exponent_shift[m_idx + b_idx * mu_grid.shape[0]] local_shift = local_exponent_shift[m_idx, b_idx] if global_shift > local_shift: local_part_fct[m_idx, b_idx] *= np.exp(local_shift - global_shift) local_fct_exp[:,m_idx, b_idx] *= np.exp(local_shift - global_shift) else: mp_part_fct[m_idx + b_idx * mu_grid.shape[0]] *= np.exp(global_shift - local_shift) for f_idx in range(len(args)): mp_fct_exp[m_idx + b_idx * mu_grid.shape[0] + f_idx * mu_grid.shape[0] * beta_grid.size ] *= np.exp(global_shift - local_shift) mp_exponent_shift[m_idx + b_idx * mu_grid.shape[0]] = local_shift mp_part_fct[m_idx + b_idx * mu_grid.shape[0]] += local_part_fct[m_idx, b_idx] for f_idx in range(len(args)): mp_fct_exp[m_idx + b_idx * mu_grid.shape[0] + f_idx * mu_grid.shape[0] * beta_grid.size ] += local_fct_exp[f_idx, m_idx, b_idx] # contains the core algorithm that calculates all eigenvalues of the block fixed # by a given spin occupation (-> spin_dist) # spin_block_size: dimension of this block # p_max, pos_charge: parameters of the hamiltonian # v_tilde: fourier transform of the potential, calculated in advance # kin_coeff: coefficient of the kinetic contribution def fctExpBlock(init_fstate, spin_dist, p_total, dia_shift, spin_boundary_bits, v_tilde, kin_coeff, p_max, part_fct, fct_exp, exponent_shift, shift_thr, mu_grid, beta_grid, symmetry, *fcts): row, col, data, block_size = cooIngredients(init_fstate, spin_boundary_bits, v_tilde, kin_coeff, p_max) coo_matrix = sp.coo_array((data, (row, col)), shape=(block_size, block_size)) ############ DEBUG SECTION START ####################### if fctexpblock_debug: csr_matrix = coo_matrix.tocsr() dia_max = csr_matrix[0,0] dia_min = csr_matrix[0,0] g_global = 0 g_min = 0 g_max = 0 for idx in range(1, block_size): entry = csr_matrix[idx, idx] g_local = - abs(entry) + np.linalg.norm( csr_matrix.data[csr_matrix.indptr[idx] : csr_matrix.indptr[idx+1]], 1) if g_local > g_global: g_global = g_local if entry > dia_max: dia_max = entry g_max = g_local if entry < dia_min: dia_min = entry g_min = g_local if block_size > 2: evs = sp.linalg.eigsh(csr_matrix, 2, which = "BE", return_eigenvectors = False) else: matrix = coo_matrix.toarray() evs = np.linalg.eigvalsh(matrix) print("p=%3d D=%6d low: %5.2f %5.2f high: %5.2f %5.2f g's: %4.2f %4.2f %4.2f" % ( p_total, block_size, evs[0], dia_min, dia_max, evs[-1], g_min, g_global, g_max)) return ######### DEBUG SECTION END ################ matrix = coo_matrix.toarray() evs = np.linalg.eigvalsh(matrix) + dia_shift occ = spin_dist.sum() updatePartFct(evs, occ, part_fct, fct_exp, exponent_shift, shift_thr, mu_grid, beta_grid, symmetry) for f_idx, fct in enumerate(fcts): updateFctExp(evs, occ, spin_dist, p_total, fct, f_idx, fct_exp, exponent_shift, shift_thr, mu_grid, beta_grid, symmetry) @njit def updatePartFct(evs, occ, part_fct, fct_exp, exponent_shift, shift_thr, mu_grid, beta_grid, symmetry): for m_idx in range(mu_grid.shape[0]): for b_idx in range(beta_grid.size): exponents = - beta_grid[b_idx] * (evs - mu_grid[m_idx, b_idx] * occ) shift = exponent_shift[m_idx, b_idx] if np.max(exponents) > shift + shift_thr: new_shift = np.max(exponents) part_fct[m_idx, b_idx] *= np.exp(shift - new_shift) fct_exp[:,m_idx, b_idx] *= np.exp(shift - new_shift) exponent_shift[m_idx, b_idx] = new_shift shift = new_shift part_fct[m_idx, b_idx] += symmetry * np.sum(np.exp(exponents - shift)) @njit def updateFctExp(evs, occ, spin_dist, p_total, fct, f_idx, fct_exp, exponent_shift, shift_thr, mu_grid, beta_grid, symmetry): for m_idx in range(mu_grid.shape[0]): for b_idx in range(beta_grid.size): exponents = - beta_grid[b_idx] * (evs - mu_grid[m_idx, b_idx] * occ) fct_exp[f_idx, m_idx, b_idx] += symmetry * np.sum(np.exp(exponents - exponent_shift[m_idx, b_idx]) * fct(occ, spin_dist, p_total, evs) ) @njit def cooIngredients(init_fstate, spin_boundary_bits, v_tilde, kin_coeff, p_max): fstate_to_ind = dict() fstate_to_ind[init_fstate] = 0 fstate_count = 1 entry_count = 1 max_size = 100000 row = np.zeros(max_size, dtype = np.int64) col = np.zeros(max_size, dtype = np.int64) data = np.zeros(max_size) stack_mimic = [init_fstate] while len(stack_mimic) > 0: fstate = stack_mimic.pop() ind = fstate_to_ind[fstate] dia = ham.exchangeInteraction(fstate, spin_boundary_bits, v_tilde, p_max, kin_coeff) if entry_count == max_size: row = np.hstack((row, np.zeros(max_size, dtype = np.int64))) col = np.hstack((col, np.zeros(max_size, dtype = np.int64))) data = np.hstack((data, np.zeros(max_size))) max_size *= 2 row[entry_count] = ind col[entry_count] = ind data[entry_count] = dia entry_count += 1 for neighbour in it.neighbourGen(fstate, spin_boundary_bits): if not(neighbour in fstate_to_ind): n_ind = len(fstate_to_ind) fstate_to_ind[neighbour] = fstate_count fstate_count += 1 stack_mimic.append(neighbour) else: n_ind = fstate_to_ind[neighbour] if fstate < neighbour: dyn_int = ham.dynamicInteraction(fstate, neighbour, spin_boundary_bits, v_tilde) if dyn_int == 0: continue if entry_count + 1 >= max_size: row = np.hstack((row, np.zeros(max_size, dtype = np.int64))) col = np.hstack((col, np.zeros(max_size, dtype = np.int64))) data = np.hstack((data, np.zeros(max_size))) max_size *= 2 row[entry_count] = n_ind col[entry_count] = ind data[entry_count] = dyn_int row[entry_count+1] = ind col[entry_count+1] = n_ind data[entry_count+1] = dyn_int entry_count += 2 return row, col, data, fstate_count