Skip to content
Snippets Groups Projects

Interface refactoring

Merged Interface refactoring
interface-refactoring into main
Merged Kostas Vilkelis requested to merge interface-refactoring into main
Compare
and
Show latest version
1 file
+ 152
− 9
Compare changes
  • Side-by-side
  • Inline
examples/testInterface.ipynb
+ 152
− 9
%% Cell type:code id: tags:
``` python
import numpy as np
import scipy
from scipy.integrate import quad, quad_vec
import matplotlib.pyplot as plt
from codes import utils
from functools import partial
```
%% Cell type:code id: tags:
``` python
def tb2kfunc(tb_model):
"""
Fourier transforms a real-space tight-binding model to a k-space function.
Parameters
----------
tb_model : dict
A dictionary with real-space vectors as keys and complex np.arrays as values.
Returns
-------
function
A function that takes a k-space vector and returns a complex np.array.
"""
def bloch_ham(k):
ham = 0
for vector in tb_model.keys():
ham += tb_model[vector] * np.exp(
-1j * np.dot(k, np.array(vector, dtype=float))
)
return ham
return bloch_ham
def kfunc2tbfunc(kfunc):
"""
Inverse Fourier transforms a k-space function to a real-space function.
Parameters
----------
kfunc : function
A function that takes a k-space vector and returns a np.array.
Returns
-------
function
A function that takes a real-space integer vector and returns a np.array.
"""
def tbfunc(vector):
def integrand(k):
return kfunc(k) * np.exp(1j * np.dot(k, np.array(vector, dtype=float)))
return quad_vec(integrand, -np.pi, np.pi)[0]
return tbfunc
```
%% Cell type:code id: tags:
``` python
def addTb(tb1, tb2):
"""
Add up two tight-binding models together.
Parameters:
-----------
tb1 : dict
Tight-binding model.
tb2 : dict
Tight-binding model.
Returns:
--------
dict
Sum of the two tight-binding models.
"""
return {k: tb1.get(k, 0) + tb2.get(k, 0) for k in frozenset(tb1) | frozenset(tb2)}
def generateRho(hkfunc, E_F):
"""
Generate a function that returns the density matrix at a given k-point.
Parameters
----------
hkfunc : function
Function that return Hamiltonian at a given k-point.
E_F : float
Fermi level
Returns
-------
rho : function
Density matrix at the same k-point as hk
"""
def rhofunc(k):
hk = hkfunc(k)
vals, vecs = np.linalg.eigh(hk)
unocc_vals = vals > E_F
occ_vecs = vecs
occ_vecs[:, unocc_vals] = 0
# Outter products between eigenvectors
return occ_vecs @ occ_vecs.T.conj()
return rhofunc
def fermiOnGrid(hkfunc, filling, Nk=500):
"""
Compute the Fermi energy on a grid of k-points.
Parameters
----------
hkfunc : function
Function that returns the Hamiltonian at a given k-point.
Nk : int
Number of k-points in the grid.
Returns
-------
E_F : float
Fermi energy
"""
ks = np.linspace(-np.pi, np.pi, Nk, endpoint=False)
hkarray = np.array([hkfunc(k) for k in ks])
vals = np.linalg.eigvalsh(hkarray)
norbs = vals.shape[-1]
vals_flat = np.sort(vals.flatten())
ne = len(vals_flat)
ifermi = int(round(ne * filling / norbs))
if ifermi >= ne:
return vals_flat[-1]
elif ifermi == 0:
return vals_flat[0]
else:
fermi = (vals_flat[ifermi - 1] + vals_flat[ifermi]) / 2
return fermi
def totalEnergy(rho, hk):
def integrand(k):
return np.real(np.trace(rho(k) @ hk(k)))
return quad(integrand, -np.pi, np.pi)[0]
def meanFieldTb(rhofunc, int_model):
"""
Compute the mean-field in real space.
Parameters
----------
rhofunc : function
Function that returns the density matrix at a given k-space vector.
int_model : dict
Interaction tb model.
Returns
-------
dict
Mean-field tb model.
"""
rhoTbfunc = kfunc2tbfunc(rhofunc)
# Compute direct interaction
intZero = tb2kfunc(int_model)(0)
direct = quad_vec(lambda k: np.diag(rhofunc(k)), -np.pi, np.pi)[0]
direct = np.diag(direct @ intZero) / (2*np.pi)
direct = {(0,) : direct}
exchange = {vec: -1*(int_model.get(vec, 0)*rhoTbfunc(vec)) / (2*np.pi) for vec in frozenset(int_model)}
return addTb(direct, exchange)
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 mf2rParams(mf_model):
"""
Convert a mean-field tight-binding model to a set of real parameters.
Parameters
----------
mf_model : dict
Mean-field tight-binding model.
Returns
-------
dict
Real parameters.
"""
return utils.complex_to_real(matrix_to_flat(mf_model[(0,)])) # placeholder for now
def rParams2mf(rParams, shape):
"""
Convert a set of real parameters to a mean-field tight-binding model.
Parameters
----------
r_params : dict
Real parameters.
shape : tuple
Shape of the mean-field tight-binding model.
Returns
-------
dict
Mean-field tight-binding model.
"""
return {(0,) : flat_to_matrix(utils.real_to_complex(rParams), shape)}
```
%% Cell type:code id: tags:
``` python
class Model:
def __init__(self, tb_model, int_model, filling):
self.tb_model = tb_model
self.int_model = int_model
self.filling = filling
def densityMatrix(self, mf_model):
self.hkfunc = tb2kfunc(addTb(self.tb_model, mf_model))
self.EF = fermiOnGrid(self.hkfunc, self.filling)
self.calculateEF()
return generateRho(self.hkfunc, self.EF)
def calculateEF(self):
self.EF = fermiOnGrid(self.hkfunc, self.filling)
def mfield(self, mf_model):
self.rhofunc = self.densityMatrix(mf_model)
return meanFieldTb(self.rhofunc, self.int_model)
return addTb(meanFieldTb(self.rhofunc, self.int_model), {(0,): - self.EF * np.eye(4)})
```
%% Cell type:code id: tags:
``` python
def cost(rParams, Model):
shape = Model.tb_model[(0,)].shape
mf_model = rParams2mf(rParams, shape) # limited placeholder
mf_new = Model.mfield(mf_model)
rParamsNew = mf2rParams(mf_new)
return rParamsNew - rParams
def solver(Model, x0, optimizer=scipy.optimize.anderson, optimizer_kwargs={}):
rParams = mf2rParams(x0)
f = partial(cost, Model=Model)
return rParams2mf(optimizer(f, rParams, **optimizer_kwargs), Model.tb_model[(0,)].shape)
result = rParams2mf(optimizer(f, rParams, **optimizer_kwargs), Model.tb_model[(0,)].shape)
Model.calculateEF()
return addTb(result, {(0,): - Model.EF * np.eye(4)})
```
%% Cell type:code id: tags:
``` python
filling = 1
filling = 2
hopp = np.kron(np.array([[0, 1], [0, 0]]), np.eye(2))
tb_model = {(0,): hopp + hopp.T.conj(), (1,): hopp, (-1,): hopp.T.conj()}
mf_model = {(0,) : np.zeros((4,4))} # placeholder for guess
guess = np.random.rand(4, 4) * np.exp(1j * 2 * np.pi * np.random.rand(4, 4))
guess += guess.T.conj()
guess /= 2
mf_model = {(0,) : guess}
int_model = {
(0,): 1 * np.kron(np.eye(2), np.ones((2, 2))),
(0,): 2 * np.kron(np.eye(2), np.ones((2, 2))),
}
tb = Model(tb_model, int_model, filling)
np.diag(solver(tb, mf_model)[(0,)])
# np.diag(solver(tb, mf_model)[(0,)])
plt.matshow(solver(tb, mf_model)[(0,)].real, cmap='RdBu_r')
plt.colorbar()
plt.show()
```
%% Output
array([0.25+0.j, 0.25+0.j, 0.25+0.j, 0.25+0.j])
%% Cell type:code id: tags:
``` python
# this is all wrong, because the mean-field is a diagonal constants term :(
def extractOnsite(U0):
int_model = {
(0,): U0 * np.kron(np.eye(2), np.ones((2, 2))),
}
tb = Model(tb_model, int_model, filling)
return np.abs(solver(tb, mf_model)[(0,)])[0, 0]
Us = np.linspace(0, 2, 10)
plt.plot(Us, [extractOnsite(U) for U in Us])
```
%% Output
[<matplotlib.lines.Line2D at 0x7f4a6520fe10>]
---------------------------------------------------------------------------
KeyboardInterrupt Traceback (most recent call last)
Cell In[48], line 10
7 return np.abs(solver(tb, mf_model)[(0,)])[0, 0]
9 Us = np.linspace(0, 2, 10)
---> 10 plt.plot(Us, [extractOnsite(U) for U in Us])
Cell In[48], line 10, in <listcomp>(.0)
7 return np.abs(solver(tb, mf_model)[(0,)])[0, 0]
9 Us = np.linspace(0, 2, 10)
---> 10 plt.plot(Us, [extractOnsite(U) for U in Us])
Cell In[48], line 7, in extractOnsite(U0)
3 int_model = {
4 (0,): U0 * np.kron(np.eye(2), np.ones((2, 2))),
5 }
6 tb = Model(tb_model, int_model, filling)
----> 7 return np.abs(solver(tb, mf_model)[(0,)])[0, 0]
Cell In[42], line 11, in solver(Model, x0, optimizer, optimizer_kwargs)
9 rParams = mf2rParams(x0)
10 f = partial(cost, Model=Model)
---> 11 result = rParams2mf(optimizer(f, rParams, **optimizer_kwargs), Model.tb_model[(0,)].shape)
12 Model.calculateEF()
13 return addTb(result, {(0,): - Model.EF * np.eye(4)})
File <string>:6, in anderson(F, xin, iter, alpha, w0, M, verbose, maxiter, f_tol, f_rtol, x_tol, x_rtol, tol_norm, line_search, callback, **kw)
File ~/micromamba/envs/python3/lib/python3.11/site-packages/scipy/optimize/_nonlin.py:214, in nonlin_solve(F, x0, jacobian, iter, verbose, maxiter, f_tol, f_rtol, x_tol, x_rtol, tol_norm, line_search, callback, full_output, raise_exception)
212 # Line search, or Newton step
213 if line_search:
--> 214 s, x, Fx, Fx_norm_new = _nonlin_line_search(func, x, Fx, dx,
215 line_search)
216 else:
217 s = 1.0
File ~/micromamba/envs/python3/lib/python3.11/site-packages/scipy/optimize/_nonlin.py:293, in _nonlin_line_search(func, x, Fx, dx, search_type, rdiff, smin)
290 s, phi1, phi0 = scalar_search_wolfe1(phi, derphi, tmp_phi[0],
291 xtol=1e-2, amin=smin)
292 elif search_type == 'armijo':
--> 293 s, phi1 = scalar_search_armijo(phi, tmp_phi[0], -tmp_phi[0],
294 amin=smin)
296 if s is None:
297 # XXX: No suitable step length found. Take the full Newton step,
298 # and hope for the best.
299 s = 1.0
File ~/micromamba/envs/python3/lib/python3.11/site-packages/scipy/optimize/_linesearch.py:692, in scalar_search_armijo(phi, phi0, derphi0, c1, alpha0, amin)
678 def scalar_search_armijo(phi, phi0, derphi0, c1=1e-4, alpha0=1, amin=0):
679 """Minimize over alpha, the function ``phi(alpha)``.
680
681 Uses the interpolation algorithm (Armijo backtracking) as suggested by
(...)
690
691 """
--> 692 phi_a0 = phi(alpha0)
693 if phi_a0 <= phi0 + c1*alpha0*derphi0:
694 return alpha0, phi_a0
File ~/micromamba/envs/python3/lib/python3.11/site-packages/scipy/optimize/_nonlin.py:277, in _nonlin_line_search.<locals>.phi(s, store)
275 return tmp_phi[0]
276 xt = x + s*dx
--> 277 v = func(xt)
278 p = _safe_norm(v)**2
279 if store:
File ~/micromamba/envs/python3/lib/python3.11/site-packages/scipy/optimize/_nonlin.py:168, in nonlin_solve.<locals>.func(z)
167 def func(z):
--> 168 return _as_inexact(F(_array_like(z, x0))).flatten()
Cell In[42], line 4, in cost(rParams, Model)
2 shape = Model.tb_model[(0,)].shape
3 mf_model = rParams2mf(rParams, shape) # limited placeholder
----> 4 mf_new = Model.mfield(mf_model)
5 rParamsNew = mf2rParams(mf_new)
6 return rParamsNew - rParams
Cell In[40], line 17, in Model.mfield(self, mf_model)
15 def mfield(self, mf_model):
16 self.rhofunc = self.densityMatrix(mf_model)
---> 17 return addTb(meanFieldTb(self.rhofunc, self.int_model), {(0,): - self.EF * np.eye(4)})
Cell In[3], line 107, in meanFieldTb(rhofunc, int_model)
105 direct = np.diag(direct @ intZero) / (2*np.pi)
106 direct = {(0,) : direct}
--> 107 exchange = {vec: -1*(int_model.get(vec, 0)*rhoTbfunc(vec)) / (2*np.pi) for vec in frozenset(int_model)}
108 return addTb(direct, exchange)
Cell In[3], line 107, in <dictcomp>(.0)
105 direct = np.diag(direct @ intZero) / (2*np.pi)
106 direct = {(0,) : direct}
--> 107 exchange = {vec: -1*(int_model.get(vec, 0)*rhoTbfunc(vec)) / (2*np.pi) for vec in frozenset(int_model)}
108 return addTb(direct, exchange)
Cell In[2], line 46, in kfunc2tbfunc.<locals>.tbfunc(vector)
43 def integrand(k):
44 return kfunc(k) * np.exp(1j * np.dot(k, np.array(vector, dtype=float)))
---> 46 return quad_vec(integrand, -np.pi, np.pi)[0]
File ~/micromamba/envs/python3/lib/python3.11/site-packages/scipy/integrate/_quad_vec.py:384, in quad_vec(f, a, b, epsabs, epsrel, norm, cache_size, limit, workers, points, quadrature, full_output, args)
381 err_sum += -neg_old_err
383 # Subdivide intervals
--> 384 for dint, derr, dround_err, subint, dneval in mapwrapper(_subdivide_interval, to_process):
385 neval += dneval
386 global_integral += dint
File ~/micromamba/envs/python3/lib/python3.11/site-packages/scipy/integrate/_quad_vec.py:441, in _subdivide_interval(args)
438 if getattr(_quadrature, 'cache_size', 0) > 0:
439 f = functools.lru_cache(_quadrature.cache_size)(f)
--> 441 s1, err1, round1 = _quadrature(a, c, f, norm_func)
442 dneval = _quadrature.num_eval
443 s2, err2, round2 = _quadrature(c, b, f, norm_func)
File ~/micromamba/envs/python3/lib/python3.11/site-packages/scipy/integrate/_quad_vec.py:597, in _quadrature_gk21(a, b, f, norm_func)
574 # 21-point weights
575 v = (0.011694638867371874278064396062192,
576 0.032558162307964727478818972459390,
577 0.054755896574351996031381300244580,
(...)
594 0.032558162307964727478818972459390,
595 0.011694638867371874278064396062192)
--> 597 return _quadrature_gk(a, b, f, norm_func, x, w, v)
File ~/micromamba/envs/python3/lib/python3.11/site-packages/scipy/integrate/_quad_vec.py:498, in _quadrature_gk(a, b, f, norm_func, x, w, v)
496 s_k_abs = 0.0
497 for i in range(len(x)):
--> 498 ff = f(c + h*x[i])
499 fv[i] = ff
501 vv = v[i]
Cell In[2], line 44, in kfunc2tbfunc.<locals>.tbfunc.<locals>.integrand(k)
43 def integrand(k):
---> 44 return kfunc(k) * np.exp(1j * np.dot(k, np.array(vector, dtype=float)))
Cell In[3], line 36, in generateRho.<locals>.rhofunc(k)
35 def rhofunc(k):
---> 36 hk = hkfunc(k)
37 vals, vecs = np.linalg.eigh(hk)
38 unocc_vals = vals > E_F
Cell In[2], line 20, in tb2kfunc.<locals>.bloch_ham(k)
17 ham = 0
18 for vector in tb_model.keys():
19 ham += tb_model[vector] * np.exp(
---> 20 -1j * np.dot(k, np.array(vector, dtype=float))
21 )
22 return ham
KeyboardInterrupt:
%% Cell type:code id: tags:
``` python
```