Merge branch 'interface-refactoring' into 'main'

create solvers and interface modules

See merge request qt/kwant-scf!3

codes/hf.py

 from scipy.ndimage import convolve
 import numpy as np
 import codes.utils as utils
-from functools import partial
-from scipy.optimize import anderson
 
-
-def mean_field_F(vals, vecs, E_F):
+def density_matrix(vals, vecs, E_F):
     """
     Returns the mean field F_ij(k) = <psi_i(k)|psi_j(k)> for all k-points and
     eigenvectors below the Fermi level.
@@ -21,8 +18,8 @@ def mean_field_F(vals, vecs, E_F):
 
     Returns
     -------
-    F : array_like
-        mean field F[kx, ky, ..., i, j] where i,j are cell indices.
+    rho : array_like
+        Density matrix rho=rho[kx, ky, ..., i, j] where i,j are cell indices.
     """
     norbs = vals.shape[-1]
     dim = len(vals.shape) - 1
@@ -37,19 +34,19 @@ def mean_field_F(vals, vecs, E_F):
         occ_vecs_flat = np.transpose(occ_vecs_flat, axes=[0, 2, 1])
 
         # inner products between eigenvectors
-        F_ij = np.einsum("kie,kje->kij", occ_vecs_flat, occ_vecs_flat.conj())
+        rho_ij = np.einsum("kie,kje->kij", occ_vecs_flat, occ_vecs_flat.conj())
         reshape_order = [nk for i in range(dim)]
         reshape_order.extend([norbs, norbs])
-        F = F_ij.reshape(*reshape_order)
+        rho = rho_ij.reshape(*reshape_order)
     else:
         unocc_vals = vals > E_F
         occ_vecs = vecs
         occ_vecs[:, unocc_vals] = 0
 
         # Outter products between eigenvectors
-        F = occ_vecs @ occ_vecs.T.conj()
+        rho = occ_vecs @ occ_vecs.T.conj()
 
-    return F
+    return rho
 
 
 def convolution(M1, M2):
@@ -82,46 +79,62 @@ def convolution(M1, M2):
     return V_output
 
 
-def compute_mf(vals, vecs, filling, H_int):
+def compute_mf(rho, H_int):
     """
     Compute mean-field correction at self-consistent loop.
 
     Parameters:
     -----------
-    vals : nd-array
-        Eigenvalues of current loop vals[k_x, ..., k_n, j].
-    vecs : nd_array
-        Eigenvectors of current loop vals[k_x, ..., k_n, i, j].
+    rho : nd_array
+        Density matrix.
     H_int : nd-array
         Interaction matrix.
-    filling: int
-        Number of electrons per cell.
 
     Returns:
     --------
     mf : nd-array
         Meanf-field correction with same format as `H_int`.
     """
-    dim = len(vals.shape) - 1
-
-    nk = vals.shape[0]
-
-    E_F = utils.get_fermi_energy(vals, filling)
-    F = mean_field_F(vals=vals, vecs=vecs, E_F=E_F)
+    
+    nk = rho.shape[0]
+    dim = len(rho.shape) - 2
 
     if dim > 0:
-        H0_int = H_int[*[0 for i in range(dim)]]
-        rho = np.diag(np.average(F, axis=tuple([i for i in range(dim)])))
-        exchange_mf = convolution(F, H_int) * nk ** (-dim)
-        direct_mf = np.diag(np.einsum("i,ij->j", rho, H0_int))
+        H0_int = H_int[*([0]*dim)]
+        local_density = np.diag(np.average(rho, axis=tuple([i for i in range(dim)])))
+        exchange_mf = convolution(rho, H_int) * nk ** (-dim)
+        direct_mf = np.diag(np.einsum("i,ij->j", local_density, H0_int))
+        dc_direct = local_density.T @ H0_int @ local_density
+        dc_exchange = np.einsum('kij, kji', exchange_mf, rho) * nk ** (-dim)
+        dc_energy = 0.5*(-dc_exchange + dc_direct) * np.eye(direct_mf.shape[-1])
     else:
-        rho = np.diag(F)
-        exchange_mf = F * H_int
-        direct_mf = np.diag(np.einsum("i,ij->j", rho, H_int))
-    return direct_mf - exchange_mf
+        local_density = np.diag(rho)
+        exchange_mf = rho * H_int
+        direct_mf = np.diag(np.einsum("i,ij->j", local_density, H_int))
+        dc_energy_direct = np.diag(np.einsum("ij, i, j->", H_int, local_density, local_density))
+        dc_energy_cross = np.diag(np.einsum("ij, ij, ji->", H_int, rho, rho))
+        dc_energy = 2 * dc_energy_direct - dc_energy_cross
+    return direct_mf - exchange_mf# - dc_energy * np.eye(direct_mf.shape[-1])
+
+def total_energy(h, rho):
+    """
+    Compute total energy.
 
+    Paramters:
+    ----------
+    h : nd-array
+        Hamiltonian.
+    rho : nd-array
+        Density matrix.
 
-def scf_loop(mf, H_int, filling, hamiltonians_0):
+    Returns:
+    --------
+    total_energy : float
+        System total energy computed as tr[h@rho].
+    """
+    return np.sum(np.trace(h @ rho, axis1=-1, axis2=-2)).real / np.prod(rho.shape[:-2])
+
+def updated_matrices(mf_k, model):
     """
     Self-consistent loop.
 
@@ -138,75 +151,16 @@ def scf_loop(mf, H_int, filling, hamiltonians_0):
 
     Returns:
     --------
-    diff : nd-array
-        Difference of mean-field matrix.
+    mf_new : nd-array
+        Updated mean-field solution.
     """
     # Generate the Hamiltonian
-    hamiltonians = hamiltonians_0 + mf
+    hamiltonians = model.hamiltonians_0 + mf_k
     vals, vecs = np.linalg.eigh(hamiltonians)
     vecs = np.linalg.qr(vecs)[0]
+    E_F = utils.get_fermi_energy(vals, model.filling)
+    rho = density_matrix(vals=vals, vecs=vecs, E_F=E_F)
+    return rho, compute_mf(
+        rho=rho,
+        H_int=model.H_int) - E_F * np.eye(model.hamiltonians_0.shape[-1])
 
-    mf_new = compute_mf(vals=vals, vecs=vecs, filling=filling, H_int=H_int)
-
-    diff = mf_new - mf
-
-    return diff
-
-
-def find_groundstate_ham(
-    tb_model,
-    int_model,
-    filling,
-    nk=10,
-    tol=1e-5,
-    guess=None,
-    mixing=0.01,
-    order=10,
-    verbose=False,
-    return_mf=False,
-):
-    """
-    Self-consistent loop to find groundstate Hamiltonian.
-
-    Parameters:
-    -----------
-    tb_model : dict
-        Tight-binding model. Must have the following structure:
-            - Keys are tuples for each hopping vector (in units of lattice vectors).
-            - Values are hopping matrices.
-    int_model : dict
-        Interaction matrix model. Must have same structure as `tb_model`
-    filling: int
-        Number of electrons per cell.
-    tol : float
-        Tolerance of meanf-field self-consistent loop.
-    guess : nd-array
-        Initial guess. Same format as `H_int`.
-    mixing : float
-        Regularization parameter in Anderson optimization. Default: 0.5.
-    order : int
-        Number of previous solutions to retain. Default: 1.
-    verbose : bool
-        Verbose of Anderson optimization. Default: False.
-    return_mf : bool
-        Returns mean-field result. Useful if wanted to reuse as guess in upcoming run.
-
-    Returns:
-    --------
-    scf_model : dict
-        Tight-binding model of Hartree-Fock solution.
-    """
-    hamiltonians_0, ks = utils.kgrid_hamiltonian(nk, tb_model, return_ks=True)
-    if guess is None:
-        guess = utils.generate_guess(nk, tb_model, int_model, scale=np.max(np.abs(np.array([*int_model.values()]))))
-    fun = partial(
-        scf_loop,
-        hamiltonians_0=hamiltonians_0,
-        H_int=utils.kgrid_hamiltonian(nk, int_model),
-        filling=filling,
-    )
-    mf = anderson(fun, guess, f_tol=tol, w0=mixing, M=order, verbose=verbose)
-    if return_mf:
-        return utils.hk2tb_model(hamiltonians_0 + mf, tb_model, int_model, ks), mf
-    else:
-        return utils.hk2tb_model(hamiltonians_0 + mf, tb_model, int_model, ks)

codes/interface.py

+from scipy import optimize
+from . import utils, solvers
+import numpy as np
+
+def find_groundstate_ham(
+    model,
+    filling,
+    nk=10,
+    cutoff_Vk=0,
+    solver=solvers.kspace_solver,
+    cost_function=solvers.kspace_cost,
+    optimizer=optimize.anderson,
+    optimizer_kwargs={},
+    return_mf=False,
+    return_kspace=False
+):
+    """
+    Self-consistent loop to find groundstate Hamiltonian.
+
+    Parameters:
+    -----------
+    tb_model : dict
+        Tight-binding model. Must have the following structure:
+            - Keys are tuples for each hopping vector (in units of lattice vectors).
+            - Values are hopping matrices.
+    filling: int
+        Number of electrons per cell.
+    guess : nd-array
+        Initial guess. Same format as `H_int`.
+    return_mf : bool
+        Returns mean-field result. Useful if wanted to reuse as guess in upcoming run.
+
+    Returns:
+    --------
+    scf_model : dict
+        Tight-binding model of Hartree-Fock solution.
+    """
+    model.nk=nk
+    model.filling=filling
+    if model.int_model is not None:
+        model.vectors=[*model.int_model.keys()]
+    else:
+        model.vectors = utils.generate_vectors(cutoff_Vk, model.dim)
+    if model.guess is None:
+        model.random_guess(model.vectors)
+    solver(model, optimizer, cost_function, optimizer_kwargs)
+    model.vectors=[*model.vectors, *model.tb_model.keys()]
+    assert np.allclose(model.mf_k - np.moveaxis(model.mf_k, -1, -2).conj(), 0, atol=1e-15)
+    vals, _ = np.linalg.eigh(model.hamiltonians_0 + model.mf_k)
+    EF = utils.get_fermi_energy(vals, filling)
+    model.mf_k -=  EF * np.eye(model.hamiltonians_0.shape[-1])
+    if return_kspace:
+        return model.hamiltonians_0 + model.mf_k
+    else:
+        if model.dim > 0:
+            scf_tb = utils.hk2tb_model(model.hamiltonians_0 + model.mf_k, model.vectors, model.ks)
+            if return_mf:
+                mf_tb = utils.hk2tb_model(model.mf_k, model.vectors, model.ks)
+                return scf_tb, mf_tb
+            else:
+                return scf_tb
+        else:
+            if return_mf:
+                return {() : model.hamiltonians_0 + model.mf_k}, {() : model.mf_k}
+            else:
+                return {() : model.hamiltonians_0 + model.mf_k}

codes/kwant_examples.py

@@ -27,7 +27,7 @@ def graphene_extended_hubbard():
         return V * np.ones((2, 2))
 
     syst_V = utils.build_interacting_syst(
-        syst=bulk_graphene,
+        builder=bulk_graphene,
         lattice = graphene,
         func_onsite = onsite_int,
         func_hop = nn_int,

codes/model.py

+from . import utils
+import numpy as np
+
+class Model:
+    """
+    A period tight-binding model class.
+
+    Attributes
+    ----------
+    tb_model : dict
+        Non-interacting tight-binding model.
+    int_model : dict
+        Interacting tight-binding model.
+    Vk : function
+        Interaction potential V(k). Used if `int_model = None`.
+    guess : dict
+        Initial guess for self-consistent calculation.
+    dim : int
+        Number of translationally invariant real-space dimensions.
+    ndof : int
+        Number of internal degrees of freedom (orbitals).
+    hamiltonians_0 : nd-array
+        Non-interacting amiltonian evaluated on a k-point grid.
+    H_int : nd-array
+        Interacting amiltonian evaluated on a k-point grid.
+    """
+
+    def __init__(self, tb_model, int_model=None, Vk=None, guess=None):
+        self.tb_model = tb_model
+        self.dim = len([*tb_model.keys()][0])
+        if self.dim > 0:
+            self.hk = utils.model2hk(tb_model=tb_model)
+        self.int_model = int_model
+        if self.int_model is not None:
+            self.int_model = int_model
+            if self.dim > 0:
+                self.Vk = utils.model2hk(tb_model=int_model)
+        else:
+            if self.dim > 0:
+                self.Vk = Vk
+        self.ndof = len([*tb_model.values()][0])
+        self.guess = guess
+        if self.dim == 0:
+            self.hamiltonians_0 = tb_model[()]
+            self.H_int = int_model[()]
+
+    def random_guess(self, vectors):
+        """
+        Generate random guess.
+
+        Parameters:
+        -----------
+        vectors : list of tuples
+            Hopping vectors for the mean-field corrections.
+        """
+        if self.int_model is None:
+            scale = 0.1
+        else:
+            scale = 0.1*(1+np.max(np.abs([*self.int_model.values()])))
+        self.guess = utils.generate_guess(
+            vectors=vectors,
+            ndof=self.ndof,
+            scale=scale
+        )
+
+    def kgrid_evaluation(self, nk):
+        """
+        Evaluates hamiltonians on a k-grid.
+
+        Parameters:
+        -----------
+        nk : int
+            Number of k-points along each direction.
+        """
+        self.hamiltonians_0, self.ks = utils.kgrid_hamiltonian(
+            nk=nk,
+            hk=self.hk,
+            dim=self.dim,
+            return_ks=True
+        )
+        self.H_int = utils.kgrid_hamiltonian(nk=nk, hk=self.Vk, dim=self.dim)
+        self.mf_k = utils.kgrid_hamiltonian(
+            nk=nk,
+            hk=utils.model2hk(self.guess),
+            dim=self.dim,
+        )

codes/solvers.py

+import numpy as np
+from . import utils
+from .hf import updated_matrices, total_energy
+from functools import partial
+
+def optimize(mf, cost_function, optimizer, optimizer_kwargs):
+    _ = optimizer(
+        cost_function,
+        mf,
+        **optimizer_kwargs
+    )
+
+def finite_system_cost(mf, model):
+    shape = mf.shape
+    mf = utils.flat_to_matrix(utils.real_to_complex(mf), shape)
+    model.rho, model.mf_k = updated_matrices(mf_k=mf, model=model)
+    delta_mf = model.mf_k - mf
+    return utils.complex_to_real(utils.matrix_to_flat(delta_mf))
+
+def finite_system_solver(model, optimizer, cost_function, optimizer_kwargs):
+    """
+    Real-space solver for finite systems.
+
+    Parameters:
+    -----------
+    model : model.Model
+        Physical model containting interacting and non-interacting Hamiltonian.
+    optimizer : function
+        Optimization function.
+    optimizer_kwargs : dict
+        Extra arguments passed to optimizer.
+    """
+    model.mf_k = model.guess[()]
+    initial_mf = utils.complex_to_real(utils.matrix_to_flat(model.mf_k))
+    partial_cost = partial(cost_function, model=model)
+    optimize(initial_mf, partial_cost, optimizer, optimizer_kwargs)
+
+def real_space_cost(mf, model, shape):
+    mf = utils.flat_to_matrix(utils.real_to_complex(mf), shape)
+    mf_dict = {}
+    for i, key in enumerate(model.guess.keys()):
+        mf_dict[key] = mf[i]
+    mf = utils.kgrid_hamiltonian(
+        nk=model.nk,
+        hk=utils.model2hk(mf_dict),
+        dim=model.dim,
+        hermitian=False
+    )
+    model.rho, model.mf_k = updated_matrices(mf_k=mf, model=model)
+    delta_mf = model.mf_k - mf
+    delta_mf = utils.hk2tb_model(delta_mf, model.vectors, model.ks)
+    delta_mf = np.array([*delta_mf.values()])
+    return utils.complex_to_real(utils.matrix_to_flat(delta_mf))
+
+def rspace_solver(model, optimizer, cost_function, optimizer_kwargs):
+    """
+    Real-space solver for infinite systems.
+
+    Parameters:
+    -----------
+    model : model.Model
+        Physical model containting interacting and non-interacting Hamiltonian.
+    optimizer : function
+        Optimization function.
+    optimizer_kwargs : dict
+        Extra arguments passed to optimizer.
+    """
+    model.kgrid_evaluation(nk=model.nk)
+    initial_mf = np.array([*model.guess.values()])
+    shape = initial_mf.shape
+    initial_mf = utils.complex_to_real(utils.matrix_to_flat(initial_mf))
+    partial_cost = partial(cost_function, model=model, shape=shape)
+    optimize(initial_mf, partial_cost, optimizer, optimizer_kwargs)
+
+def kspace_cost(mf, model):
+    mf = utils.flat_to_matrix(utils.real_to_complex(mf), model.mf_k.shape)
+    model.rho, model.mf_k = updated_matrices(mf_k=mf, model=model)
+    delta_mf = model.mf_k - mf
+    return utils.complex_to_real(utils.matrix_to_flat(delta_mf))
+
+def kspace_totalenergy_cost(mf, model):
+    mf = utils.flat_to_matrix(utils.real_to_complex(mf), model.mf_k.shape)
+    model.rho, model.mf_k = updated_matrices(mf_k=mf, model=model)
+    return total_energy(
+        model.hamiltonians_0 + model.mf_k,
+        model.rho,
+    )
+
+def kspace_solver(model, optimizer, cost_function, optimizer_kwargs):
+    """
+    k-space solver.
+
+    Parameters:
+    -----------
+    model : model.Model
+        Physical model containting interacting and non-interacting Hamiltonian.
+    optimizer : function
+        Optimization function.
+    optimizer_kwargs : dict
+        Extra arguments passed to optimizer.
+    """
+    model.kgrid_evaluation(nk=model.nk)
+    initial_mf = model.mf_k
+    initial_mf = utils.complex_to_real(utils.matrix_to_flat(initial_mf))
+    partial_cost = partial(cost_function, model=model)
+    optimize(initial_mf, partial_cost, optimizer, optimizer_kwargs)

codes/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
 
@@ -113,6 +113,11 @@ def builder2tb_model(builder, params={}, return_data=False):
                     .toarray()
                     .T.conj()
                 )
+            else:
+                # Hopping vector in the opposite direction
+                tb_model[tuple(-b_dom)] += coo_array(
+                    (data, (row, col)), shape=(norbs_tot, norbs_tot)
+                ).toarray().T.conj()
         else:
             tb_model[tuple(b_dom)] = coo_array(
                 (data, (row, col)), shape=(norbs_tot, norbs_tot)
@@ -123,6 +128,10 @@ def builder2tb_model(builder, params={}, return_data=False):
                     .toarray()
                     .T.conj()
                 )
+            else:
+                tb_model[tuple(-b_dom)] = coo_array(
+                    (data, (row, col)), shape=(norbs_tot, norbs_tot)
+                ).toarray().T.conj()
 
     if return_data:
         data = {}
@@ -133,32 +142,53 @@ def builder2tb_model(builder, params={}, return_data=False):
         return tb_model
 
 
-def dict2hk(tb_dict):
+def model2hk(tb_model):
     """
     Build Bloch Hamiltonian.
 
+    Paramters:
+    ----------
+    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.
+    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(tb_model))) > 0
+    ), "Zero-dimensional system. The Hamiltonian is simply tb_model[()]."
 
     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))
+        for vector in tb_model.keys():
+            ham += tb_model[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))
-                )
         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 +196,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 +212,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 +220,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,14 +260,12 @@ def build_interacting_syst(builder, lattice, func_onsite, func_hop, max_neighbor
     return int_builder
 
 
-def generate_guess(nk, tb_model, int_model, scale=0.1):
+def generate_guess(vectors, ndof, scale=0.1):
     """
-    nk : int
-        Number of k-points along each direction.
-    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.
 
@@ -253,22 +274,42 @@ def generate_guess(nk, tb_model, int_model, scale=0.1):
     guess : nd-array
         Guess evaluated on a k-point grid.
     """
-    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
+        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 generate_vectors(cutoff, dim):
+    """
+    Generates hopping vectors up to a cutoff.
 
-    return kgrid_hamiltonian(nk, guess) * scale
+    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 hk2tb_model(hk, tb_model, int_model, ks=None):
+
+def hk2tb_model(hk, hopping_vecs, ks=None):
     """
     Extract hopping matrices from Bloch Hamiltonian.
 
@@ -289,9 +330,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)
@@ -337,3 +375,50 @@ 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
+
+def complex_to_real(z):
+    """
+    Split real and imaginary parts of a complex array.
+
+    Parameters:
+    -----------
+    z : array
+    """
+    return np.concatenate((np.real(z), np.imag(z)))
+
+def real_to_complex(z):
+    """
+    Undo `complex_to_real`.
+    """
+    return z[:len(z)//2] + 1j * z[len(z)//2:]
\ No newline at end of file