Interface refactoring

Find a design which is more intuitive and minimal.

codes/kwant_helper/utils.py

 import numpy as np
 import kwant
-from itertools import product
 from scipy.sparse import coo_array
+from itertools import product
 import inspect
 from copy import copy
 
@@ -28,7 +28,7 @@ def get_fermi_energy(vals, filling):
         return fermi
 
 
-def builder2tb_model(builder, params={}, return_data=False):
+def builder2h_0(builder, params={}, return_data=False):
     """
     Constructs a tight-binding model dictionary from a `kwant.Builder`.
 
@@ -43,7 +43,7 @@ def builder2tb_model(builder, params={}, return_data=False):
 
     Returns:
     --------
-    tb_model : dict
+    h_0 : dict
         Tight-binding model of non-interacting systems.
     data : dict
         Data with sites and number of orbitals. Only if `return_data=True`.
@@ -52,7 +52,7 @@ def builder2tb_model(builder, params={}, return_data=False):
     # Extract information from builder
     dims = len(builder.symmetry.periods)
     onsite_idx = tuple([0] * dims)
-    tb_model = {}
+    h_0 = {}
     sites_list = [*builder.sites()]
     norbs_list = [site[0].norbs for site in builder.sites()]
     positions_list = [site[0].pos for site in builder.sites()]
@@ -73,14 +73,14 @@ def builder2tb_model(builder, params={}, return_data=False):
                     _params[arg] = params[arg]
             val = val(site, **_params)
             data = val.flatten()
-        except:
+        except Exception:
             data = val.flatten()
-        if onsite_idx in tb_model:
-            tb_model[onsite_idx] += coo_array(
+        if onsite_idx in h_0:
+            h_0[onsite_idx] += coo_array(
                 (data, (row, col)), shape=(norbs_tot, norbs_tot)
             ).toarray()
         else:
-            tb_model[onsite_idx] = coo_array(
+            h_0[onsite_idx] = coo_array(
                 (data, (row, col)), shape=(norbs_tot, norbs_tot)
             ).toarray()
     # Hopping matrices
@@ -101,24 +101,37 @@ def builder2tb_model(builder, params={}, return_data=False):
                     _params[arg] = params[arg]
             val = val(a, b, **_params)
             data = val.flatten()
-        except:
+        except Exception:
             data = val.flatten()
-        if tuple(b_dom) in tb_model:
-            tb_model[tuple(b_dom)] += coo_array(
+        if tuple(b_dom) in h_0:
+            h_0[tuple(b_dom)] += coo_array(
                 (data, (row, col)), shape=(norbs_tot, norbs_tot)
             ).toarray()
             if np.linalg.norm(b_dom) == 0:
-                tb_model[tuple(b_dom)] += (
+                h_0[tuple(b_dom)] += (
+                    coo_array((data, (row, col)), shape=(norbs_tot, norbs_tot))
+                    .toarray()
+                    .T.conj()
+                )
+            else:
+                # Hopping vector in the opposite direction
+                h_0[tuple(-b_dom)] += (
                     coo_array((data, (row, col)), shape=(norbs_tot, norbs_tot))
                     .toarray()
                     .T.conj()
                 )
         else:
-            tb_model[tuple(b_dom)] = coo_array(
+            h_0[tuple(b_dom)] = coo_array(
                 (data, (row, col)), shape=(norbs_tot, norbs_tot)
             ).toarray()
             if np.linalg.norm(b_dom) == 0:
-                tb_model[tuple(b_dom)] += (
+                h_0[tuple(b_dom)] += (
+                    coo_array((data, (row, col)), shape=(norbs_tot, norbs_tot))
+                    .toarray()
+                    .T.conj()
+                )
+            else:
+                h_0[tuple(-b_dom)] = (
                     coo_array((data, (row, col)), shape=(norbs_tot, norbs_tot))
                     .toarray()
                     .T.conj()
@@ -128,37 +141,56 @@ def builder2tb_model(builder, params={}, return_data=False):
         data = {}
         data["norbs"] = norbs_list
         data["positions"] = positions_list
-        return tb_model, data
+        return h_0, data
     else:
-        return tb_model
+        return h_0
 
 
-def dict2hk(tb_dict):
+def model2hk(h_0):
     """
     Build Bloch Hamiltonian.
 
+    Paramters:
+    ----------
+    nk : int
+        Number of k-points along each direction.
+    h_0 : dictionary
+        Must have the following structure:
+            - Keys are tuples for each hopping vector (in units of lattice vectors).
+            - Values are hopping matrices.
+    return_ks : bool
+        Return k-points.
+
+    Returns:
+    --------
+    ham : nd.array
+        Hamiltonian evaluated on a k-point grid from k-points
+        along each direction evaluated from zero to 2*pi.
+        The indices are ordered as [k_1, ... , k_n, i, j], where
+        `k_m` corresponding to the k-point element along each
+        direction and `i` and `j` are the internal degrees of freedom.
+    ks : 1D-array
+        List of k-points over all directions. Only returned if `return_ks=True`.
+
     Returns:
     --------
     bloch_ham : function
         Evaluates the Hamiltonian at a given k-point.
     """
+    assert (
+        len(next(iter(h_0))) > 0
+    ), "Zero-dimensional system. The Hamiltonian is simply h_0[()]."
 
     def bloch_ham(k):
         ham = 0
-        for vector in tb_dict.keys():
-            ham += tb_dict[vector] * np.exp(
-                1j * np.dot(k, np.array(vector, dtype=float))
-            )
-            if np.linalg.norm(np.array(vector)) > 0:
-                ham += tb_dict[vector].T.conj() * np.exp(
-                    -1j * np.dot(k, np.array(vector))
-                )
+        for vector in h_0.keys():
+            ham += h_0[vector] * np.exp(-1j * np.dot(k, np.array(vector, dtype=float)))
         return ham
 
     return bloch_ham
 
 
-def kgrid_hamiltonian(nk, tb_model, return_ks=False):
+def kgrid_hamiltonian(nk, hk, dim, return_ks=False, hermitian=True):
     """
     Evaluates Hamiltonian on a k-point grid.
 
@@ -166,12 +198,10 @@ def kgrid_hamiltonian(nk, tb_model, return_ks=False):
     ----------
     nk : int
         Number of k-points along each direction.
-    tb_model : dictionary
-        Must have the following structure:
-            - Keys are tuples for each hopping vector (in units of lattice vectors).
-            - Values are hopping matrices.
+    hk : function
+        Calculates the Hamiltonian at a given k-point.
     return_ks : bool
-        Return k-points.
+        If `True`, returns k-points.
 
     Returns:
     --------
@@ -184,15 +214,6 @@ def kgrid_hamiltonian(nk, tb_model, return_ks=False):
     ks : 1D-array
         List of k-points over all directions. Only returned if `return_ks=True`.
     """
-    dim = len(next(iter(tb_model)))
-    if dim == 0:
-        if return_ks:
-            return syst[next(iter(tb_model))], None
-        else:
-            return syst[next(iter(tb_model))]
-    else:
-        hk = dict2hk(tb_model)
-
     ks = 2 * np.pi * np.linspace(0, 1, nk, endpoint=False)
 
     k_pts = np.tile(ks, dim).reshape(dim, nk)
@@ -201,6 +222,10 @@ def kgrid_hamiltonian(nk, tb_model, return_ks=False):
     for k in product(*k_pts):
         ham.append(hk(k))
     ham = np.array(ham)
+    if hermitian:
+        assert np.allclose(
+            ham, np.transpose(ham, (0, 2, 1)).conj()
+        ), "Tight-binding provided is non-Hermitian. Not supported yet"
     shape = (*[nk] * dim, ham.shape[-1], ham.shape[-1])
     if return_ks:
         return ham.reshape(*shape), ks
@@ -237,36 +262,56 @@ def build_interacting_syst(builder, lattice, func_onsite, func_hop, max_neighbor
     return int_builder
 
 
-def generate_guess(tb_model, int_model, scale=0.1):
+def generate_guess(vectors, ndof, scale=1):
     """
-    tb_model : dict
-        Tight-binding model of non-interacting systems.
-    int_model : dict
-        Tight-binding model for interacting Hamiltonian.
+    vectors : list
+        List of hopping vectors.
+    ndof : int
+        Number internal degrees of freedom (orbitals),
     scale : float
         The scale of the guess. Maximum absolute value of each element of the guess.
 
     Returns:
     --------
     guess : tb dictionary
-        TB guess.
+        Guess in the form of a tight-binding model.
     """
-    ndof = tb_model[next(iter(tb_model))].shape[-1]
     guess = {}
-    vectors = [*tb_model.keys(), *int_model.keys()]
     for vector in vectors:
-        amplitude = np.random.rand(ndof, ndof)
-        phase = 2 * np.pi * np.random.rand(ndof, ndof)
-        rand_hermitian = amplitude * np.exp(1j * phase)
-        if np.linalg.norm(np.array(vector)):
-            rand_hermitian += rand_hermitian.T.conj()
-            rand_hermitian /= 2
-        guess[vector] = rand_hermitian * scale
+        if vector not in guess.keys():
+            amplitude = scale * np.random.rand(ndof, ndof)
+            phase = 2 * np.pi * np.random.rand(ndof, ndof)
+            rand_hermitian = amplitude * np.exp(1j * phase)
+            if np.linalg.norm(np.array(vector)) == 0:
+                rand_hermitian += rand_hermitian.T.conj()
+                rand_hermitian /= 2
+                guess[vector] = rand_hermitian
+            else:
+                guess[vector] = rand_hermitian
+                guess[tuple(-np.array(vector))] = rand_hermitian.T.conj()
 
     return guess
 
 
-def hk2tb_model(hk, tb_model, int_model, ks=None):
+def generate_vectors(cutoff, dim):
+    """
+    Generates hopping vectors up to a cutoff.
+
+    Parameters:
+    -----------
+    cutoff : int
+        Maximum distance along each direction.
+    dim : int
+        Dimension of the vectors.
+
+    Returns:
+    --------
+    List of hopping vectors.
+    """
+    return [*product(*([[*range(-cutoff, cutoff + 1)]] * dim))]
+
+
+def hk2h_0(hk, hopping_vecs, ks=None):
     """
     Extract hopping matrices from Bloch Hamiltonian.
 
@@ -274,9 +319,9 @@ def hk2tb_model(hk, tb_model, int_model, ks=None):
     -----------
     hk : nd-array
         Bloch Hamiltonian matrix hk[k_x, ..., k_n, i, j]
-    tb_model : dict
+    h_0 : dict
         Tight-binding model of non-interacting systems.
-    int_model : dict
+    h_int : dict
         Tight-binding model for interacting Hamiltonian.
     ks : 1D-array
         Set of k-points. Repeated for all directions. If the system is finite, `ks=None`.
@@ -287,9 +332,6 @@ def hk2tb_model(hk, tb_model, int_model, ks=None):
         TIght-binding model of Hartree-Fock solution.
     """
     if ks is not None:
-        hopping_vecs = np.unique(
-            np.array([*tb_model.keys(), *int_model.keys()]), axis=0
-        )
         ndim = len(hk.shape) - 2
         dk = np.diff(ks)[0]
         nk = len(ks)
@@ -307,11 +349,11 @@ def hk2tb_model(hk, tb_model, int_model, ks=None):
             * (dk / (2 * np.pi)) ** ndim
         )
 
-        tb_model = {}
+        h_0 = {}
         for i, vector in enumerate(hopping_vecs):
-            tb_model[tuple(vector)] = hopps[i]
+            h_0[tuple(vector)] = hopps[i]
 
-        return tb_model
+        return h_0
     else:
         return {(): hk}
 
@@ -335,3 +377,35 @@ def calc_gap(vals, E_F):
     emax = np.max(vals[vals <= E_F])
     emin = np.min(vals[vals > E_F])
     return np.abs(emin - emax)
+
+
+def matrix_to_flat(matrix):
+    """
+    Flatten the upper triangle of a collection of matrices.
+
+    Parameters:
+    -----------
+    matrix : nd-array
+        Array with shape (..., n, n)
+    """
+    return matrix[..., *np.triu_indices(matrix.shape[-1])].flatten()
+
+
+def flat_to_matrix(flat, shape):
+    """
+    Undo `matrix_to_flat`.
+
+    Parameters:
+    -----------
+    flat : 1d-array
+        Output from `matrix_to_flat`.
+    shape : 1d-array
+        Shape of the input from `matrix_to_flat`.
+    """
+    matrix = np.zeros(shape, dtype=complex)
+    matrix[..., *np.triu_indices(shape[-1])] = flat.reshape(*shape[:-2], -1)
+    indices = np.arange(shape[-1])
+    diagonal = matrix[..., indices, indices]
+    matrix += np.moveaxis(matrix, -1, -2).conj()
+    matrix[..., indices, indices] -= diagonal
+    return matrix
\ No newline at end of file