create solvers and interface modules

Co-authored-by: RonjaZijd <RonjaZijd@users.noreply.github.com>

codes/hf.py

 from scipy.ndimage import convolve
 import numpy as np
 import codes.utils as utils
-from scipy import optimize
 
 def density_matrix(vals, vecs, E_F, dim):
     """
@@ -158,145 +157,3 @@ def updated_matrices(mf_k, model):
         H_int=model.H_int,
         dim=model.dim) - E_F * np.eye(model.hamiltonians_0.shape[-1])
 
-def finite_system_solver(model, optimizer, 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.
-    """
-    mf = model.guess[()]
-    shape = mf.shape
-
-    def cost_function(mf):
-        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))
-
-    _ = optimizer(
-        cost_function,
-        utils.complex_to_real(utils.matrix_to_flat(mf)),
-        **optimizer_kwargs
-    )
-
-def rspace_solver(model, optimizer, 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)
-    mf = np.array([*model.guess.values()])
-    shape = mf.shape
-
-    def cost_function(mf):
-        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)
-        model.energy = total_energy(h=model.hamiltonians_0 + model.mf_k, rho=model.rho)
-        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))
-
-    _ = optimizer(
-        cost_function,
-        utils.complex_to_real(utils.matrix_to_flat(mf)),
-        **optimizer_kwargs
-    )
-
-
-def kspace_solver(model, optimizer, 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)
-    def cost_function(mf):
-        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)
-        model.energy = total_energy(h=model.hamiltonians_0 + model.mf_k, rho=model.rho)
-        delta_mf = model.mf_k - mf
-        return utils.complex_to_real(utils.matrix_to_flat(delta_mf))
-
-    _ = optimizer(
-        cost_function,
-        utils.complex_to_real(utils.matrix_to_flat(model.mf_k)),
-        **optimizer_kwargs
-    )
-
-def find_groundstate_ham(
-    model,
-    cutoff_Vk,
-    filling,
-    nk=10,
-    solver=kspace_solver,
-    optimizer=optimize.anderson,
-    optimizer_kwargs={'M':0, 'verbose': 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, optimizer_kwargs)
-    model.vectors=[*model.vectors, *model.tb_model.keys()]
-    assert np.allclose((model.mf_k - np.moveaxis(model.mf_k, -1, -2).conj())/2, 0, atol=1e-15)
-    if model.dim > 0:
-        return utils.hk2tb_model(model.hamiltonians_0 + model.mf_k, model.vectors, model.ks)
-    else:
-        return {() : model.hamiltonians_0 + model.mf_k}

codes/interface.py

+from scipy import optimize
+from . import utils, solvers
+import numpy as np
+
+def find_groundstate_ham(
+    model,
+    cutoff_Vk,
+    filling,
+    nk=10,
+    solver=solvers.kspace_solver,
+    cost_function=solvers.kspace_cost,
+    optimizer=optimize.anderson,
+    optimizer_kwargs={'M':0, 'verbose': 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())/2, 0, atol=1e-15)
+    if model.dim > 0:
+        return utils.hk2tb_model(model.hamiltonians_0 + model.mf_k, model.vectors, model.ks)
+    else:
+        return {() : model.hamiltonians_0 + model.mf_k}

codes/solvers.py

+import numpy as np
+from . import utils
+from .hf import updated_matrices
+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.
+    """
+    initial_mf = model.guess[()]
+    partial_cost = partial(cost_function, model=model)
+    optimize(initial_mf, partial_cost, optimizer, optimizer_kwargs)
+
+def real_space_cost(mf, model):
+    shape = mf.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()])
+    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)
+
+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_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)
\ No newline at end of file