Numba for Quantum Optimization

Another appreciation post for Numba. Couple of years ago, I had developed this package (pySKTB) to help with calculating Eigen spectrum of systems with co-dimensions greater than 1. This was done so that I could in theory do rudimentary testing of different symmetries and how they couple to the orbital and real-space symmetry. But one of the major problems I had was that the code was really clunky.

After dabbling with Numba, I learnt about the existence of couple of new tools like line_profiler which seems to be extremely useful for optimizing codes. It spits out time of run for each line in your code.

So I started with the biggest bottleneck, which was calculating , the matrix for connecting sites. The plan was to add all Numba functions in a separate module of optimized functions and call them when needed:

@nb.jit(nopython=False)
def get_gmat_jit(g_mat, all_iter, max_image,
                 n_orbitals, bond_mat, dist_mat_vec,
                 kpt_cart):
    for ind_1, (atom_1_i, orbit_1_i, element_1,
                orbit_1) in enumerate(all_iter):
        for ind_2, (atom_2_i, orbit_2_i, element_2,
                    orbit_2) in enumerate(all_iter):
            for image_ind in range(max_image):
                if bond_mat[image_ind, atom_1_i,
                            atom_2_i] == 0:
                    continue
                dist_vec = dist_mat_vec[image_ind,
                           atom_1_i, atom_2_i, :]
                phase = np.exp(2.*np.pi*1j
                        * np.dot(kpt_cart, dist_vec))
                g_mat[image_ind, ind_1, ind_2] = phase
    g_mat[int(max_image/2), :, :] += \
        np.eye(n_orbitals, dtype=complex)
    return g_mat

Almost order improvement, not bad for a couple of minutes of work. Next stop would be to parallelize the calculation with Dask.