diff --git a/docs/source/algorithm.md b/docs/source/algorithm.md
new file mode 100644
index 0000000000000000000000000000000000000000..3845eab26bc00c7d7c8949eaf84d34596c07503f
--- /dev/null
+++ b/docs/source/algorithm.md
@@ -0,0 +1,43 @@
+# Algorithm overview
+
+## Self-consistency loop
+
+In order to calculate the mean-field interaction in {eq}`mf_infinite`, we require the ground-state density matrix $\rho_{mn}(R)$.
+However, the density matrix in {eq}`density` is a functional of the mean-field interaction $\hat{V}_{\text{MF}}$ itself.
+Therefore, we need to solve for both self-consistently.
+
+We define a single iteration of a self-consistency loop:
+
+$$
+\text{SCF}(\hat{V}_{\text{init, MF}}) \to \hat{V}_{\text{new, MF}},
+$$
+
+such that it performs the following operations given an initial mean-field interaction $\hat{V}_{\text{init, MF}}$:
+
+1. Calculate the total Hamiltonian $\hat{H}(R) = \hat{H_0}(R) + \hat{V}_{\text{init, MF}}(R)$ in real-space.
+2. Fourier transform the total Hamiltonian to the momentum space $\hat{H}(R) \to \hat{H}(k)$.
+3. Calculate the ground-state density matrix $\rho_{mn}(k)$ in momentum space.
+ 1. Diagonalize the Hamiltonian $\hat{H}(k)$ to obtain the eigenvalues and eigenvectors.
+ 2. Calculate the fermi level $\mu$ given the desired filling of the unit cell.
+ 3. Calculate the density matrix $\rho_{mn}(k)$ using the eigenvectors and the fermi level $\mu$ (currently we do not consider thermal effects so $\beta \to \infty$).
+4. Inverse Fourier transform the density matrix to real-space $\rho_{mn}(k) \to \rho_{mn}(R)$.
+5. Calculate the new mean-field interaction $\hat{V}_{\text{new, MF}}(R)$ via {eq}`mf_infinite`.
+
+## Self-consistency criterion
+
+To define the self-consistency condition, we first introduce an invertible function $f$ that uniquely maps $\hat{V}_{\text{MF}}$ to a real-valued vector which minimally parameterizes it:
+
+$$
+f : \hat{V}_{\text{MF}} \to f(\hat{V}_{\text{MF}}) \in \mathbb{R}^N.
+$$
+
+Currently, $f$ parameterizes the mean-field interaction by taking only the upper triangular elements of the matrix $V_{\text{MF}, nm}(R)$ (the lower triangular part is redundant due to the Hermiticity of the Hamiltonian) and splitting it into a real and imaginary parts to form a real-valued vector.
+
+With this function, we define the self-consistency criterion as a fixed-point problem:
+
+$$
+f(\text{SCF}(\hat{V}_{\text{MF}})) = f(\hat{V}_{\text{MF}}).
+$$
+
+To solve this fixed-point problem, we utilize a root-finding function `scipy.optimize.anderson` which uses the Anderson mixing method to find the fixed-point solution.
+However, our implementation also allows to use other custom fixed-point solvers by either providing it to `solver` or by re-defining the `solver` function.
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 6ae82d9f4039aea8f48fbc083edc0605b367ab57..7b401e6ff9586b062feab2a3054e0e3ea1701d2d 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -71,6 +71,8 @@ intersphinx_mapping = {
default_role = "autolink"
+latex_elements = {"extrapackages": r"\usepackage{braket}"}
+
# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]
diff --git a/docs/source/documentation/pymf.md b/docs/source/documentation/pymf.md
new file mode 100644
index 0000000000000000000000000000000000000000..7137c5a15d400e1c7371348cc4de6525e530636b
--- /dev/null
+++ b/docs/source/documentation/pymf.md
@@ -0,0 +1,76 @@
+# Package reference
+
+## Interactive problem definition
+
+To define the interactive problem, we use the following class:
+
+```{eval-rst}
+.. autoclass:: pymf.model.Model
+ :members: mfield
+```
+
+## Mean-field and density matrix
+
+```{eval-rst}
+.. automodule:: pymf.mf
+ :members: meanfield, construct_density_matrix, construct_density_matrix_kgrid, fermi_on_kgrid
+ :show-inheritance:
+```
+
+## Observables
+
+```{eval-rst}
+.. automodule:: pymf.observables
+ :members: expectation_value
+ :show-inheritance:
+```
+
+## Solvers
+
+```{eval-rst}
+.. automodule:: pymf.solvers
+ :members: solver, cost
+ :show-inheritance:
+```
+
+## Tight-binding dictionary
+
+### Manipulation
+
+```{eval-rst}
+.. automodule:: pymf.tb.tb
+ :members: add_tb, scale_tb
+ :show-inheritance:
+```
+
+### Brillouin zone transformations
+
+```{eval-rst}
+.. automodule:: pymf.tb.transforms
+ :members:
+ :show-inheritance:
+```
+
+### Parametrisation
+
+```{eval-rst}
+.. automodule:: pymf.params.rparams
+ :members:
+ :show-inheritance:
+```
+
+### Utility functions
+
+```{eval-rst}
+.. automodule:: pymf.tb.utils
+ :members:
+ :show-inheritance:
+```
+
+## `kwant` interface
+
+```{eval-rst}
+.. automodule:: pymf.kwant_helper.utils
+ :members:
+ :show-inheritance:
+```
diff --git a/docs/source/index.md b/docs/source/index.md
index 3df6663c32a93642efeb4d48c9c4e39058b0dcfc..9613885abede066d93b42f98cb5fd7978fe71505 100644
--- a/docs/source/index.md
+++ b/docs/source/index.md
@@ -18,7 +18,16 @@ kernelspec:
:maxdepth: 1
:caption: Tutorials
+```
+
+```{toctree}
+:hidden:
+:maxdepth: 1
+:caption: Documentation
+
mf_notes.md
+algorithm.md
+documentation/pymf.md
```
## What is pymf?
diff --git a/docs/source/mf_notes.md b/docs/source/mf_notes.md
index 6ff0d32d5cb7799306def2699bc89efd339340a2..33d49dd7a17619bf65e5a61ecb36e2eb9d1337fd 100644
--- a/docs/source/mf_notes.md
+++ b/docs/source/mf_notes.md
@@ -1,304 +1,75 @@
-# Self-consistent mean field algorithm
+# Theory overview
-## Mean-field approximation
+## Interacting problems
-The full hamiltonian is:
+In physics, one often encounters problems where a system of multiple particles interact with each other.
+In this package, we consider a general electronic system with density-density interparticle interaction:
-$$
-\hat{H} = \hat{H_0} + \hat{V} = \sum_{ij} h_{ij} c^\dagger_{i} c_{j} + \frac{1}{2} \sum_{ijkl} v_{ijkl} c_i^\dagger c_j^\dagger c_l c_k
-$$
-
-We assume the dominant part of the ground state wavefunction comes from $\hat{H}_0$. Let's assume $b_i$ operators diagonalize the unperturbed hamiltonian:
-
-$$
-c_i^\dagger = \sum_{k} U_{ik} b_k^\dagger,
-$$
-
-such that the unperturbed groundstate wavefunction is:
-
-$$
-| 0 \rangle = \Pi_{E_i \leq \mu } b_i^\dagger |\textrm{vac}\rangle.
-$$
-
-Based on this definition, we define the normal ordering operator $:ABC...:$ such that it fulfills:
-
-$$
-:ABC...: | 0 \rangle = 0
-$$
-
-which practically means it orders $b_i$ operators based on whether its above or below the Fermi level $\mu$.
-
-Under this definition of normal ordering, we define the Wick's expansion of the interaction term:
-
-$$
-\begin{multline}
-c_i^\dagger c_j^\dagger c_l c_k = :c_i^\dagger c_j^\dagger c_l c_k: \\+ \overline{c_i^\dagger c_j^\dagger} :c_l c_k: + \overline{c_i^\dagger c_k} :c_j^\dagger c_l: - \overline{c_i^\dagger c_l} :c_j^\dagger c_k: + \overline{c_l c_k} :c_i^\dagger c_j^\dagger: - \overline{c_j^\dagger c_k} :c_i^\dagger c_l: + \overline{c_j^\dagger c_l} :c_i^\dagger c_k: \\
-\overline{c_i^\dagger c_j^\dagger} \overline{c_l c_k} - \overline{c_i^\dagger c_l} \overline{c_j^\dagger c_k} + \overline{c_i^\dagger c_k} \overline{c_j^\dagger c_l}
-\end{multline}
-$$
+:::{math}
+:label: hamiltonian
+\hat{H} = \hat{H_0} + \hat{V} = \sum_{ij} h_{ij} c^\dagger_{i} c_{j} + \frac{1}{2} \sum_{ij} v_{ij} c_i^\dagger c_j^\dagger c_j c_i
+:::
-where the overline defines Wick's contraction:
+where $c_i^\dagger$ and $c_i$ are creation and annihilation operators respectively for fermion in state $i$.
+The first term $\hat{H_0}$ is the non-interacting Hamiltonian which by itself is straightforward to solve in a single-particle basis by direct diagonalizations made easy through packages such as [kwant](https://kwant-project.org/).
+The second term $\hat{V}$ is density-density interaction term between two particles, for example Coulomb interaction.
+In order to solve the interacting problem exactly, one needs to diagonalize the full Hamiltonian $\hat{H}$ in the many-particle basis which grows exponentially with the number of particles.
+Such a task is often infeasible for large systems and one needs to resort to approximations.
-$$
-\overline{AB} = AB - :AB:.
-$$
+## Mean-field approximaton
-The expectation value of the interaction with respect to the $| 0 \rangle$ is:
+The first-order perturbative approximation to the interacting Hamiltonian is the Hartree-Fock approximation also known as the mean-field approximation.
+The mean-field approximates the quartic term $\hat{V}$ in {eq}`hamiltonian` as a sum of bilinear terms weighted by the expectation values the remaining operators:
+:::{math}
+:label: mf_approx
+\hat{V} \approx \hat{V}_{\text{MF}} \equiv \sum_{ij} v_{ij} \left[
+\braket{c_i^\dagger c_i} c_j^\dagger c_j - \braket{c_i^\dagger c_j} c_j^\dagger c_i \right]
+:::
+we neglect the superconducting pairing and constant offset terms.
+The expectation value terms $\langle c_i^\dagger c_j \rangle$ are due to the ground-state density matrix and therefore act as an effective field acting on the system.
+The ground-state density matrix reads:
+:::{math}
+:label: density
+\rho_{ij} \equiv \braket{c_i^\dagger c_j } = \text{Tr}\left(e^{-\beta \left(\hat{H_0} + \hat{V}_{\text{MF}} - \mu \hat{N} \right)} c_i^\dagger c_j\right),
+:::
+where $\beta = 1/ (k_B T)$ is the inverse temperature, $\mu$ is the chemical potential, and $\hat{N} = \sum_i c_i^\dagger c_i$ is the number operator.
-$$
-\langle 0 | c_i^\dagger c_j^\dagger c_l c_k | 0 \rangle = \langle 0 | \overline{c_i^\dagger c_j^\dagger} \overline{c_l c_k} - \overline{c_i^\dagger c_l} \overline{c_j^\dagger c_k} + \overline{c_i^\dagger c_k} \overline{c_j^\dagger c_l} | 0 \rangle
-$$
+## Finite tight-binding grid
-where we can forget about all the normal ordered states since those give zero acting on the unperturbed groundstate. To evaluate this further, we utilize the mean-field approximation:
+To simplify the mean-field Hamiltonian, we assume a finite, normalised orthogonal tight-binding grid defined by the single-particle basis states:
$$
-A B \approx \langle A \rangle B + A \langle B \rangle - \langle A \rangle \langle B \rangle
+\ket{n} = c^\dagger_n\ket{\text{vac}}
$$
-onto the contractions such that we get:
+where $\ket{\text{vac}}$ is the vacuum state.
+We project our mean-field interaction in {eq}`mf_approx` onto the tight-binding grid:
-$$
-\langle c_i^\dagger c_j^\dagger c_l c_k \rangle \approx \langle c_i^\dagger c_j^\dagger \rangle \langle c_l c_k \rangle + \langle c_i^\dagger c_k \rangle \langle c_j^\dagger c_l \rangle - \langle c_i^\dagger c_l \rangle \langle c_j^\dagger c_k \rangle
-$$
+:::{math}
+:label: mf_finite
+V_{\text{MF}, nm} = \braket{n | \hat{V}_{\text{MF}} | m} = \sum_{i} \rho_{ii} v_{in} \delta_{nm} - \rho_{mn} v_{mn},
+:::
+where $\delta_{nm}$ is the Kronecker delta function.
-note $\langle A B \rangle \approx \langle A \rangle \langle B \rangle$ assuming mean-field.
+## Infinite tight-binding grid
-To consider excitations from the groundstate, we make use of the mean-field approximation defined above:
+In the limit of a translationally invariant system, the index $n$ that labels the basis states partitions into two independent variables: the unit cell internal degrees of freedom (spin, orbital, sublattice, etc.) and the position of the unit cell $R_n$:
$$
-\begin{multline}
-c_i^\dagger c_j^\dagger c_l c_k \approx \\
-\langle c_i^\dagger c_j^\dagger \rangle c_l c_k + \langle c_i^\dagger c_k \rangle c_j^\dagger c_l - \langle c_i^\dagger c_l \rangle c_j^\dagger c_k + \langle c_l c_k \rangle c_i^\dagger c_j^\dagger - \langle c_j^\dagger c_k \rangle c_i^\dagger c_l + \langle c_j^\dagger c_l \rangle c_i^\dagger c_k + \\
-\langle c_i^\dagger c_j^\dagger \rangle \langle c_l c_k \rangle + \langle c_i^\dagger c_k \rangle \langle c_j^\dagger c_l \rangle - \langle c_i^\dagger c_l \rangle \langle c_j^\dagger c_k \rangle
-\end{multline}
+n \to n, R_n.
$$
-Where we made use of the following operations:
+Because of the translationaly invariance, the physical properties of the system are independent of the absolute unit cell position $R_n$ and rather depend on the relative position between the two unit cells $R_{nm} = R_n - R_m$:
$$
-:c_i^\dagger c_j^\dagger c_l c_k: \approx 0
-$$
-
-$$
-\overline{c_i^\dagger c_k} \overline{c_j^\dagger c_l} \approx \langle \overline{c_i^\dagger c_k} \rangle \overline{c_j^\dagger c_l} + \overline{c_i^\dagger c_k} \langle \overline{c_j^\dagger c_l} \rangle - \langle \overline{c_i^\dagger c_k} \rangle \langle \overline{c_j^\dagger c_i} \rangle = \langle c_i^\dagger c_k \rangle \overline{c_j^\dagger c_l} + \overline{c_i^\dagger c_k} \langle c_j^\dagger c_l \rangle - \langle c_i^\dagger c_k \rangle \langle c_j^\dagger c_l \rangle
-$$
-
-$$
-\overline{c_i^\dagger c_k} :c_j^\dagger c_l: \approx \langle \overline{c_i^\dagger c_k} \rangle :c_j^\dagger c_l: + \overline{c_i^\dagger c_k} \langle :c_j^\dagger c_l: \rangle - \langle \overline{c_i^\dagger c_k} \rangle \langle :c_j^\dagger c_l: \rangle = \langle \overline{c_i^\dagger c_k} \rangle :c_j^\dagger c_l:
-$$
-
-
-$$
-\langle \overline{c_i^\dagger c_k} \rangle = \langle c_i^\dagger c_k - :c_i^\dagger c_k: \rangle = \langle c_i^\dagger c_k \rangle
-$$
-
-
-Without any superconducting terms, the form simplifies to:
-
-$$
-\begin{multline}
-c_i^\dagger c_j^\dagger c_l c_k \approx
-\langle c_i^\dagger c_k \rangle c_j^\dagger c_l - \langle c_i^\dagger c_l \rangle c_j^\dagger c_k - \langle c_j^\dagger c_k \rangle c_i^\dagger c_l + \langle c_j^\dagger c_l \rangle c_i^\dagger c_k + \\
-\langle c_i^\dagger c_k \rangle \langle c_j^\dagger c_l \rangle - \langle c_i^\dagger c_l \rangle \langle c_j^\dagger c_k \rangle
-\end{multline}
-$$
-
-## Finite size
-
-### Coulomb interaction
-
-We simplify the interaction term through the MF approximation to get:
-
-$$
-V = \frac{1}{2}\sum_{ijkl} v_{ijkl} c_i^{\dagger} c_j^{\dagger} c_l c_k
-\approx
-\frac12 \sum_{ijkl} v_{ijkl} \left[ \langle c_i^{\dagger} c_k \rangle c_j^{\dagger} c_l - \langle c_j^{\dagger} c_k \rangle c_i^{\dagger} c_l - \langle c_i^{\dagger} c_l \rangle c_j^{\dagger} c_k + \langle c_j^{\dagger} c_l \rangle c_i^{\dagger} c_k \right]
-$$
-(assuming no superconductivity)
-
-and an additional constant part:
-
-$$
-V_0 = \frac{1}{2} \sum_{ijkl} v_{ijkl} \left(\langle c_j^{\dagger} c_l \rangle \langle c_i^{\dagger} c_k \rangle - \langle c_j^{\dagger} c_k \rangle \langle c_i^{\dagger} c_l \rangle \right).
-$$
-
-The interaction reads:
-
-$$
-v_{ijkl} = \iint w_{i}^*(r) w_{j}^*(r') V(r, r') w_{k}(r) w_l(r') dr dr' = \\
-\iint V(|r - r'|) w_{i}^*(r)w_{k}(r) w_{j}^*(r') w_l(r') dr dr'
-$$
-
-whereas $w_i$ is a wannier function on site i (and corresponding dof). Whenever one interchanges $i \to j, k \to l$, the Coulomb term is preserved $v_{ijkl} = v_{jilk}$
-
-To make things more understandable, we are also going to explicitly split up position and spin indices: $i \to i \times \sigma$. In this notation, the Coulomb integral reads:
-
-$$
-v_{ijkl}^{\sigma_i \sigma_j \sigma_k \sigma_l} =
-\iint V(|r - r'|) w_{i\times\sigma_i}^{*} (r)w_{k \times \sigma_k}(r) w_{j \times \sigma_j}^{*}(r') w_{l\times \sigma_l}(r') dr dr' \delta_{\sigma_i \sigma_k} \delta_{\sigma_{j} \sigma_l}
-$$
-
-On a fine tight-binding model, we have:
-
-$$
-v_{ijkl}^{\sigma_i \sigma_j \sigma_k \sigma_l} = v_{ij} \delta_{ik} \delta_{jl} \delta_{\sigma_i \sigma_k} \delta_{\sigma_j \sigma_l}
-$$
-
-where $v_{ij} = V(r_i, r_j)$.
-
-We shall re-define $i$ index to absorb spin:
-
-$$
-\delta_{ik} \times \delta_{\sigma_{i} \sigma_{k}} \to \delta_{ik}
-$$
-
-in this notation the above reads:
-
-$$
-v_{ijkl} = v_{ij} \delta_{ik} \delta_{jl}
-$$
-
-The mean-field terms are:
-
+\rho_{mn} \to \rho_{mn}(R_{mn}).
$$
-\langle c_i^{\dagger} c_j\rangle = \langle \Psi_F|c_i^{\dagger} c_j | \Psi_F \rangle
-$$
-
-whereas $|\Psi_F \rangle = \Pi_{i=0}^{N_F} b_i^\dagger |0\rangle$. To make sense of things, we need to transform between $c_i$ basis (position + internal dof basis) into the $b_i$ basis (eigenfunction of a given mean-field Hamiltonian):
-
-$$
-c_i^\dagger = \sum_{k} U_{ik} b_k^\dagger
-$$
-
-where
-
-$$
-U_{ik} = \langle{i|\psi_k} \rangle.
-$$
-
-That gives us:
-
-
-$$
-c_i^{\dagger} c_j = \sum_{k, l} U_{ik} U_{lj}^* b_k^\dagger b_{l}
-$$
-
-and its expectation value gives us the mean-field ... field $F_{ij}$:
-
-$$
-F_{ij} = \langle c_i^{\dagger} c_j\rangle = \sum_{k, l} U_{ik} U_{lj}^* \langle \Psi_F| b_k^\dagger b_{l}| \Psi_F \rangle = \sum_{k} U_{ik} U_{kj}^{*}
-$$
-
-whereas I assumed the indices for wavefunctions $k,l$ are ordered in terms of increasing eigenvalue. We pop that into the definition of the mean-field correction $V$:
-
-
-$$
-\begin{multline}
-V_{nm} = \frac12 \sum_{ijkl} v_{ijkl} \langle n| \left[ \langle c_i^{\dagger} c_k \rangle c_j^{\dagger} c_l - \langle c_j^{\dagger} c_k \rangle c_i^{\dagger} c_l - \langle c_i^{\dagger} c_l \rangle c_j^{\dagger} c_k + \langle c_j^{\dagger} c_l \rangle c_i^{\dagger} c_k \right] |m\rangle = \\
- \frac12 \sum_{ijkl} v_{ijkl} \left[ +\delta_{jn}\delta_{lm} F_{ik} -\delta_{in}\delta_{lm} F_{jk} -\delta_{jn}\delta_{km} F_{il} + \delta_{in}\delta_{km} F_{jl} \right] = \\
-\frac12 \left[ \sum_{ik} v_{inkm} F_{ik} - \sum_{jk} v_{njkm} F_{jk} - \sum_{il} v_{inml} F_{il} + \sum_{jl} v_{njml} F_{jl} \right] = \\
--\sum_{ij} F_{ij} \left(v_{inmj} - v_{injm} \right)
-\end{multline}
-$$
-
-where I used the $v_{ijkl} = v_{jilk}$ symmetry from Coulomb.
-
-For a tight-binding grid, the mean-field correction simplifies to:
-
-$$
-\begin{multline}
-V_{nm} = - \sum_{ij} F_{ij} \left(v_{inmj} - v_{injm} \right) = \\
--\sum_{ij}F_{ij} v_{in} \delta_{im} \delta_{nj} + \sum_{ij}F_{ij} v_{in} \delta_{ij} \delta_{nm} = \\
--F_{mn} v_{mn} + \sum_{i} F_{ii} v_{in} \delta_{nm}
-\end{multline}
-$$
-
-the first term is the exchange interaction whereas the second one is the direct interaction.
-
-Similarly, the constant offset term reads:
-
-$$
-\begin{multline}
-V_0 = \frac{1}{2} \sum_{ijkl} v_{ijkl} \left(F_{jl} F_{ik} - F_{jk} F_{il} \right) = \\
-\frac{1}{2} \sum_{ijkl} v_{ij} \delta_{ik} \delta_{jl} \left(F_{jl} F_{ik} - F_{jk} F_{il}\right) \\
-= \frac{1}{2} \sum_{ij} v_{ij} \left(F_{ii} F_{jj} - F_{ji} F_{ij}\right)
-\end{multline}
-$$
-
-where we identify the first term as the exchange (mixes indices) and the right one as the direct (diagonal in indices).
-
-## Translational Invariance
-
-The above assumed a finite tight-binding model - all $nm$-indices contain the position of all atoms (among other dof). In this section tho we want to consider an infinite system with translational invariance.
-
-To begin with we deconstruct a general matrix $O_{nm}$ into the cell degrees of freedom ($nm$) and the position of the the cell itself ($ij$):
-
-$$
-O_{nm} \to O^{ij}_{nm}
-$$
-
-and we will Fourier transform the upper indices into k-space:
-
-$$
-O_{mn}(k) = \sum_{ij} O_{nm}^{ij} e^{-i k (R_i-R_j)}
-$$
-
-where I assumed $O$ (and thus all operators I will consider here) is local and thus diagonal in k-space.
-
-Now lets rewrite our main result in the previous section using our new notation:
-
-$$
-V_{nm}^{ij} =-F_{mn}^{ij} v_{mn}^{ij} + \sum_{r,p} F_{pp}^{rr} v_{pn}^{rj} \delta_{nm} \delta^{ij}
-$$
-
-Lets first consider the second (direct) term. Lets express the corresponding $F$ term in k-space:
-
-$$
-F_{pp}^{rr} = \int e^{i k (R_r-R_r)} F_{pp}(k) dk = \int F_{pp}(k) dk
-$$
-
-Notice that in the final expression, there is no $rr$ dependence and thus this term is cell-periodic. Therefore, we shall redefine it as cell electron density $\rho$:
-$$
-F_{pp}^0 = F_{pp}(R = 0) = \int F_{pp}(k) dk
-$$
-
-Now since $\rho$ has no $r$ dependence, we can proceed with the sum:
-
-$$
-\sum_{r} v_{pn}^{rj} = \int v_{pn}(k) e^{ik R_j} \sum_{r} e^{-i k R_r} dk = \int v_{pn}(k) e^{ik R_j} \delta_{k, 0} dk = v_{pn}(0)
-$$
-
-We are finally ready to Fourier transform the main result. Invoking convolution theorem and the results above gives us:
-
-$$
-V_{nm}(k) = \sum_{p} F_{pp}^0 v_{pn}(0) \delta_{nm} -F_{mn}(k) \circledast v_{mn}(k) = V_n^D - F_{mn}(k) \circledast v_{mn}(k)
-$$
-
-which does make sense. The first term (direct) is a potential term coming from the mean-field and the second term (exchange) is purely responsible for the hopping.
-
-The constant offset is:
-$$
-V_0 = \frac{1}{2} \sum_{r,s} \rho_r v_{rs}(0) \rho_s- \\ \frac{1}{2} tr\left[\int_{BZ} \left(F \circledast v\right)(k) F(k) dk \right]
-$$
-
-## Short-Range interactions
-
-In the case of short-range interactions, it is much more convenient to go back to real space to both store objects and perform the operations. In real space the mean-field part of the Hamiltonian reads:
-
-$$
-V_{nm}(\mathbf{R}) = V_n^D \delta(\mathbf{R}) - F_{mn}(\mathbf{R}) v_{mn}(\mathbf{R})
-$$
-
-(the first term might need some prefactor from Fourier transformation)
-where $\mathbf{R}$ is a sequence of integers representing real-space unit cell indices.
+That allows us to re-write the mean-field interaction in {eq}`mf_finite` as:
-### Proposed Algorithm
-Given an initial Hamiltonian $H_0 (R)$ and the interaction term $v(R)$ in real-space, the mean-field algorithm is the following:
+:::{math}
+:label: mf_infinite
+V_{\text{MF}, nm} (R) = \sum_{i} \rho_{ii} (0) v_{in} (0) \delta_{nm} \delta(R) - \rho_{mn}(R) v_{mn}(R),
+:::
-0. Start with a mean-field guess in real-space: $V(R)$.
-1. Fourier transform tight-binding model and the mean-field in real space to a given k-grid: $H_0(R) + V(R) \to H_0(k) + V(k)$
-2. Diagonalize and evaluate the density matrix: $H_0(k) + V(k) \to F(k)$
-3. Fourier transform the density matrix back to real-space: $F(k) \to F(R)$
-4. Evaluate the new mean-field Hamiltonian $V(R)$ according to the equation above.
-5. Evaluate self-consistency metric (could be based either on $V(R/k)$ or $F(R/k)$). Based on that, either stop or go back to 1 with a modified $V(R)$ starting guess.
+where now indices $i, n, m$ label the internal degrees of freedom of the unit cell and $R$ is the relative position between the two unit cells in terms of the lattice vectors.
diff --git a/pymf/__init__.py b/pymf/__init__.py
index a0f2b6e8598cdbc623b27aaeffdd5555fbcb8f96..bb870d26385e27f27b87fac4fbfd4b8f5b5b22fc 100644
--- a/pymf/__init__.py
+++ b/pymf/__init__.py
@@ -6,10 +6,30 @@ except ImportError:
__version__ = "unknown"
__version_tuple__ = (0, 0, "unknown", "unknown")
-from .mf import construct_density_matrix
+from .mf import (
+ construct_density_matrix,
+ meanfield,
+)
+from .solvers import solver
+from .model import Model
+from .observables import expectation_value
+from .tb.tb import add_tb, scale_tb
+from .tb.transforms import tb_to_kgrid, kgrid_to_tb
+from .tb.utils import generate_guess, calculate_fermi_energy
+
__all__ = [
+ "solver",
+ "Model",
+ "expectation_value",
+ "add_tb",
+ "scale_tb",
+ "generate_guess",
+ "calculate_fermi_energy",
"construct_density_matrix",
+ "meanfield",
+ "tb_to_kgrid",
+ "kgrid_to_tb",
"__version__",
"__version_tuple__",
]
diff --git a/pymf/kwant_helper/utils.py b/pymf/kwant_helper/utils.py
index debc0d99ac7ef73ef783418d17c1ebdaff186d78..650b636621b050aa592a29a9cf90e6d8f1fcf618 100644
--- a/pymf/kwant_helper/utils.py
+++ b/pymf/kwant_helper/utils.py
@@ -1,29 +1,36 @@
import inspect
from copy import copy
from itertools import product
+from typing import Callable
-import kwant
import numpy as np
from scipy.sparse import coo_array
+import kwant
+import kwant.lattice
+import kwant.builder
+from pymf.tb.tb import _tb_type
-def builder_to_tb(builder, params={}, return_data=False):
- """Construct a tight-binding model dictionary from a `kwant.Builder`.
+
+def builder_to_tb(
+ builder: kwant.builder.Builder, params: dict = {}, return_data: bool = False
+) -> _tb_type:
+ """Construct a tight-binding dictionary from a `kwant.builder.Builder` system.
Parameters
----------
- builder : `kwant.Builder`
- Either builder for non-interacting system or interacting Hamiltonian.
- params : dict
- Dictionary of parameters to evaluate the Hamiltonian.
- return_data : bool
+ builder :
+ system to convert to tight-binding dictionary.
+ params :
+ Dictionary of parameters to evaluate the builder on.
+ return_data :
Returns dictionary with sites and number of orbitals per site.
Returns
-------
- h_0 : dict
- Tight-binding model of non-interacting systems.
- data : dict
+ :
+ Tight-binding dictionary that corresponds to the builder.
+ :
Data with sites and number of orbitals. Only if `return_data=True`.
"""
builder = copy(builder)
@@ -124,28 +131,38 @@ def builder_to_tb(builder, params={}, return_data=False):
return h_0
-def build_interacting_syst(builder, lattice, func_onsite, func_hop, max_neighbor=1):
- """Construct an auxiliary `kwant` system to build Hamiltonian matrix.
+def build_interacting_syst(
+ builder: kwant.builder.Builder,
+ lattice: kwant.lattice.Polyatomic,
+ func_onsite: Callable,
+ func_hop: Callable,
+ max_neighbor: int = 1,
+) -> kwant.builder.Builder:
+ """
+ Construct an auxiliary `kwant` system that encodes the interactions.
Parameters
----------
- builder : `kwant.Builder`
- Non-interacting `kwant` system.
- func_onsite : function
- Onsite function.
- func_hop : function
- Hopping function.
- max_neighbor : int
- Maximal nearest-neighbor order.
+ builder :
+ Non-interacting `kwant.builder.Builder` system.
+ lattice :
+ Lattice of the system.
+ func_onsite :
+ Onsite interactions function.
+ func_hop :
+ Hopping/inter unit cell interactions function.
+ max_neighbor :
+ The maximal number of neighbouring unit cells (along a lattice vector)
+ connected by interaction. Interaction goes to zero after this distance.
Returns
-------
- int_builder : `kwant.Builder`
- Dummy `kwant.Builder` to compute interaction matrix.
+ :
+ Auxiliary `kwant.builder.Builder` that encodes the interactions of the system.
"""
- # lattice_info = list(builder.sites())[0][0]
- # lattice = kwant.lattice.general(lattice_info.prim_vecs, norbs=lattice_info.norbs)
- int_builder = kwant.Builder(kwant.TranslationalSymmetry(*builder.symmetry.periods))
+ int_builder = kwant.builder.Builder(
+ kwant.lattice.TranslationalSymmetry(*builder.symmetry.periods)
+ )
int_builder[builder.sites()] = func_onsite
for neighbors in range(max_neighbor):
int_builder[lattice.neighbors(neighbors + 1)] = func_hop
diff --git a/pymf/mf.py b/pymf/mf.py
index d04d8f3da1e14adc20c3cb088547d5fa99ab279a..fd2c33893bf919643adb67b29b4a1ca607003f0d 100644
--- a/pymf/mf.py
+++ b/pymf/mf.py
@@ -1,80 +1,97 @@
import numpy as np
-from scipy.fftpack import ifftn
+from typing import Tuple
-from pymf.tb.tb import add_tb
-from pymf.tb.transforms import ifftn_to_tb, tb_to_khamvector
+from pymf.tb.tb import add_tb, _tb_type
+from pymf.tb.transforms import tb_to_kgrid, kgrid_to_tb
-def construct_density_matrix_kgrid(kham, filling):
+def construct_density_matrix_kgrid(
+ kham: np.ndarray, filling: float
+) -> Tuple[np.ndarray, float]:
"""Calculate density matrix on a k-space grid.
Parameters
----------
- kham : npndarray
- Hamiltonian in k-space of shape (len(dim), norbs, norbs)
- filling : float
+ kham :
+ Hamiltonian from which to construct the density matrix.
+ The hamiltonian is sampled on a grid of k-points and has shape (nk, nk, ..., ndof, ndof),
+ where ndof is number of internal degrees of freedom.
+ filling :
Number of particles in a unit cell.
+ Used to determine the Fermi level.
Returns
-------
- np.ndarray, float
- Density matrix in k-space and Fermi energy.
+ :
+ Density matrix on a k-space grid with shape (nk, nk, ..., ndof, ndof) and Fermi energy.
"""
vals, vecs = np.linalg.eigh(kham)
- fermi = fermi_on_grid(vals, filling)
+ fermi = fermi_on_kgrid(vals, filling)
unocc_vals = vals > fermi
occ_vecs = vecs
np.moveaxis(occ_vecs, -1, -2)[unocc_vals, :] = 0
- rho_krid = occ_vecs @ np.moveaxis(occ_vecs, -1, -2).conj()
- return rho_krid, fermi
+ density_matrix_krid = occ_vecs @ np.moveaxis(occ_vecs, -1, -2).conj()
+ return density_matrix_krid, fermi
-def construct_density_matrix(h, filling, nk):
- """Compute the density matrix in real-space tight-binding format.
+def construct_density_matrix(
+ h: _tb_type, filling: float, nk: int
+) -> Tuple[_tb_type, float]:
+ """Compute the real-space density matrix tight-binding dictionary.
Parameters
----------
- h : dict
- Tight-binding model.
- filling : float
- Filling of the system.
- nk : int
- Number of k-points in the grid.
+ h :
+ Hamiltonian tight-binding dictionary from which to construct the density matrix.
+ filling :
+ Number of particles in a unit cell.
+ Used to determine the Fermi level.
+ nk :
+ Number of k-points in a grid to sample the Brillouin zone along each dimension.
+ If the system is 0-dimensional (finite), this parameter is ignored.
Returns
-------
- (dict, float)
- Density matrix in real-space tight-binding format and Fermi energy.
+ :
+ Density matrix tight-binding dictionary and Fermi energy.
"""
ndim = len(list(h)[0])
if ndim > 0:
- kham = tb_to_khamvector(h, nk=nk)
- rho_grid, fermi = construct_density_matrix_kgrid(kham, filling)
+ kham = tb_to_kgrid(h, nk=nk)
+ density_matrix_krid, fermi = construct_density_matrix_kgrid(kham, filling)
return (
- ifftn_to_tb(ifftn(rho_grid, axes=np.arange(ndim))),
+ kgrid_to_tb(density_matrix_krid),
fermi,
)
else:
- rho, fermi = construct_density_matrix_kgrid(h[()], filling)
- return {(): rho}, fermi
+ density_matrix, fermi = construct_density_matrix_kgrid(h[()], filling)
+ return {(): density_matrix}, fermi
-def meanfield(density_matrix_tb, h_int):
- """Compute the mean-field in k-space.
+def meanfield(density_matrix: _tb_type, h_int: _tb_type) -> _tb_type:
+ """Compute the mean-field correction from the density matrix.
Parameters
----------
- density_matrix : dict
- Density matrix in real-space tight-binding format.
- h_int : dict
- Interaction tb model.
-
+ density_matrix :
+ Density matrix tight-binding dictionary.
+ h_int :
+ Interaction hermitian Hamiltonian tight-binding dictionary.
Returns
-------
- dict
- Mean-field tb model.
+ :
+ Mean-field correction tight-binding dictionary.
+
+ Notes
+ -----
+
+ The interaction h_int must be of density-density type.
+ For example in 1D system with ndof internal degrees of freedom,
+ h_int[(2,)] = U * np.ones((ndof, ndof)) is a Coulomb repulsion interaction
+ with strength U between unit cells separated by 2 lattice vectors, where
+ the interaction is the same between all internal degrees of freedom.
"""
- n = len(list(density_matrix_tb)[0])
+ n = len(list(density_matrix)[0])
local_key = tuple(np.zeros((n,), dtype=int))
direct = {
@@ -82,7 +99,7 @@ def meanfield(density_matrix_tb, h_int):
np.array(
[
np.diag(
- np.einsum("pp,pn->n", density_matrix_tb[local_key], h_int[vec])
+ np.einsum("pp,pn->n", density_matrix[local_key], h_int[vec])
)
for vec in frozenset(h_int)
]
@@ -92,26 +109,26 @@ def meanfield(density_matrix_tb, h_int):
}
exchange = {
- vec: -1 * h_int.get(vec, 0) * density_matrix_tb[vec] # / (2 * np.pi)#**2
- for vec in frozenset(h_int)
+ vec: -1 * h_int.get(vec, 0) * density_matrix[vec] for vec in frozenset(h_int)
}
return add_tb(direct, exchange)
-def fermi_on_grid(vals, filling):
+def fermi_on_kgrid(vals: np.ndarray, filling: float) -> float:
"""Compute the Fermi energy on a grid of k-points.
Parameters
----------
- vals : ndarray
- Eigenvalues of the hamiltonian in k-space of shape (len(dim), norbs, norbs)
- filling : int
- Number of particles in a unit cell.
-
+ vals :
+ Eigenvalues of a hamiltonian sampled on a k-point grid with shape (nk, nk, ..., ndof, ndof),
+ where ndof is number of internal degrees of freedom.
+ filling :
+ Number of particles in a unit cell.
+ Used to determine the Fermi level.
Returns
-------
- fermi_energy : float
- Fermi energy
+ :
+ Fermi energy
"""
norbs = vals.shape[-1]
vals_flat = np.sort(vals.flatten())
diff --git a/pymf/model.py b/pymf/model.py
index 6a61b2d4cda86bc770691c23b970c238f6ca3928..00a7fd4153b0eede6498caae7223b0895583cf2f 100644
--- a/pymf/model.py
+++ b/pymf/model.py
@@ -4,7 +4,7 @@ from pymf.mf import (
construct_density_matrix,
meanfield,
)
-from pymf.tb.tb import add_tb
+from pymf.tb.tb import add_tb, _tb_type
def _check_hermiticity(h):
@@ -12,24 +12,52 @@ def _check_hermiticity(h):
op_vector = tuple(-1 * np.array(vector))
op_vector = tuple(-1 * np.array(vector))
if not np.allclose(h[vector], h[op_vector].conj().T):
- raise ValueError("Hamiltonian is not Hermitian.")
+ raise ValueError("Tight-binding dictionary must be hermitian.")
def _tb_type_check(tb):
for count, key in enumerate(tb):
if not isinstance(tb[key], np.ndarray):
- raise ValueError("Inputted dictionary values are not np.ndarray's")
+ raise ValueError(
+ "Values of the tight-binding dictionary must be numpy arrays"
+ )
shape = tb[key].shape
if count == 0:
size = shape[0]
if not len(shape) == 2:
- raise ValueError("Inputted dictionary values are not square matrices")
+ raise ValueError(
+ "Values of the tight-binding dictionary must be square matrices"
+ )
if not size == shape[0]:
- raise ValueError("Inputted dictionary elements shapes are not consistent")
+ raise ValueError(
+ "Values of the tight-binding dictionary must have consistent shape"
+ )
class Model:
- def __init__(self, h_0, h_int, filling):
+ """
+ Data class which defines the interacting tight-binding problem.
+
+ Parameters
+ ----------
+ h_0 :
+ Non-interacting hermitian Hamiltonian tight-binding dictionary.
+ h_int :
+ Interaction hermitian Hamiltonian tight-binding dictionary.
+ filling :
+ Number of particles in a unit cell.
+ Used to determine the Fermi level.
+
+ Notes
+ -----
+ The interaction h_int must be of density-density type.
+ For example in 1D system with ndof internal degrees of freedom,
+ h_int[(2,)] = U * np.ones((ndof, ndof)) is a Coulomb repulsion interaction
+ with strength U between unit cells separated by 2 lattice vectors, where
+ the interaction is the same between all internal degrees of freedom.
+ """
+
+ def __init__(self, h_0: _tb_type, h_int: _tb_type, filling: float) -> None:
_tb_type_check(h_0)
self.h_0 = h_0
_tb_type_check(h_int)
@@ -42,31 +70,32 @@ class Model:
_first_key = list(h_0)[0]
self._ndim = len(_first_key)
- self._size = h_0[_first_key].shape[0]
+ self._ndof = h_0[_first_key].shape[0]
self._local_key = tuple(np.zeros((self._ndim,), dtype=int))
_check_hermiticity(h_0)
_check_hermiticity(h_int)
- def mfield(self, mf_tb, nk=200): # method or standalone?
- """Compute single mean field iteration.
+ def mfield(self, mf: _tb_type, nk: int = 20) -> _tb_type:
+ """Computes a new mean-field correction from a given one.
Parameters
----------
- mf_tb : dict
- Mean-field tight-binding model.
- nk : int
- Number of k-points in the grid.
+ mf :
+ Initial mean-field correction tight-binding dictionary.
+ nk :
+ Number of k-points in a grid to sample the Brillouin zone along each dimension.
+ If the system is 0-dimensional (finite), this parameter is ignored.
Returns
-------
- dict
- New mean-field tight-binding model.
+ :
+ new mean-field correction tight-binding dictionary.
"""
rho, fermi_energy = construct_density_matrix(
- add_tb(self.h_0, mf_tb), self.filling, nk
+ add_tb(self.h_0, mf), self.filling, nk
)
return add_tb(
meanfield(rho, self.h_int),
- {self._local_key: -fermi_energy * np.eye(self._size)},
+ {self._local_key: -fermi_energy * np.eye(self._ndof)},
)
diff --git a/pymf/observables.py b/pymf/observables.py
index bf770349e11798bbe4d86dae5f3f076cf86043f1..bd47485a9299786551bd22455600c31bb89a6d11 100644
--- a/pymf/observables.py
+++ b/pymf/observables.py
@@ -1,19 +1,21 @@
import numpy as np
+from pymf.tb.tb import _tb_type
-def expectation_value(density_matrix, observable):
+
+def expectation_value(density_matrix: _tb_type, observable: _tb_type) -> complex:
"""Compute the expectation value of an observable with respect to a density matrix.
Parameters
----------
- density_matrix : dict
- Density matrix in tight-binding format.
- observable : dict
- Observable in tight-binding format.
+ density_matrix :
+ Density matrix tight-binding dictionary.
+ observable :
+ Observable tight-binding dictionary.
Returns
-------
- complex
+ :
Expectation value.
"""
return np.sum(
diff --git a/pymf/params/param_transforms.py b/pymf/params/param_transforms.py
index 98b68e2485597b9479ab30755e76b37312e15968..2095e345634d7588dc6eed6a90689fc7f7e6d756 100644
--- a/pymf/params/param_transforms.py
+++ b/pymf/params/param_transforms.py
@@ -1,18 +1,20 @@
import numpy as np
+from pymf.tb.tb import _tb_type
-def tb_to_flat(tb):
- """Convert a hermitian tight-binding dictionary to flat complex matrix.
+
+def tb_to_flat(tb: _tb_type) -> np.ndarray:
+ """Parametrise a hermitian tight-binding dictionary by a flat complex vector.
Parameters
----------
- tb : dict with nd-array elements
+ tb :
Hermitian tigh-binding dictionary
Returns
-------
- flat : complex 1d numpy array
- Flattened tight-binding dictionary
+ :
+ 1D complex array that parametrises the tight-binding dictionary.
"""
if len(list(tb)[0]) == 0:
matrix = np.array(list(tb.values()))
@@ -23,34 +25,39 @@ def tb_to_flat(tb):
return sorted_vals[:N].flatten()
-def flat_to_tb(flat, shape, tb_keys):
+def flat_to_tb(
+ tb_param_complex: np.ndarray,
+ ndof: int,
+ tb_keys: list[tuple[None] | tuple[int, ...]],
+) -> _tb_type:
"""Reverse operation to `tb_to_flat`.
It takes a flat complex 1d array and return the tight-binding dictionary.
Parameters
----------
- flat : dict with nd-array elements
- Hermitian tigh-binding dictionary
- shape : tuple
- shape of the tb elements
- tb_keys : iterable
- original tb key elements
+ tb_param_complex :
+ 1d complex array that parametrises the tb model.
+ ndof :
+ Number internal degrees of freedom within the unit cell.
+ tb_keys :
+ List of keys of the tight-binding dictionary.
Returns
-------
- tb : dict
+ tb :
tight-binding dictionary
"""
+ shape = (len(tb_keys), ndof, ndof)
if len(tb_keys[0]) == 0:
matrix = np.zeros((shape[-1], shape[-2]), dtype=complex)
- matrix[np.triu_indices(shape[-1])] = flat
+ matrix[np.triu_indices(shape[-1])] = tb_param_complex
matrix += matrix.conj().T
matrix[np.diag_indices(shape[-1])] /= 2
return {(): matrix}
matrix = np.zeros(shape, dtype=complex)
N = len(tb_keys) // 2 + 1
- matrix[:N] = flat.reshape(N, *shape[1:])
+ matrix[:N] = tb_param_complex.reshape(N, *shape[1:])
matrix[N:] = np.moveaxis(matrix[-(N + 1) :: -1], -1, -2).conj()
tb_keys = np.array(list(tb_keys))
@@ -59,17 +66,22 @@ def flat_to_tb(flat, shape, tb_keys):
return tb
-def complex_to_real(z):
- """Split real and imaginary parts of a complex array.
+def complex_to_real(z: np.ndarray) -> np.ndarray:
+ """Split and concatenate real and imaginary parts of a complex array.
Parameters
----------
- z : array
+ z :
Complex array.
+
+ Returns
+ -------
+ :
+ Real array that concatenates the real and imaginary parts of the input array.
"""
return np.concatenate((np.real(z), np.imag(z)))
-def real_to_complex(z):
+def real_to_complex(z: np.ndarray) -> np.ndarray:
"""Undo `complex_to_real`."""
return z[: len(z) // 2] + 1j * z[len(z) // 2 :]
diff --git a/pymf/params/rparams.py b/pymf/params/rparams.py
index bb78ed086ef08366f3d6c91f90ea0f91d838793b..dcb0479ec413f675732d12497b7a0c9a2331d221 100644
--- a/pymf/params/rparams.py
+++ b/pymf/params/rparams.py
@@ -1,44 +1,48 @@
+import numpy as np
+
from pymf.params.param_transforms import (
complex_to_real,
flat_to_tb,
real_to_complex,
tb_to_flat,
)
+from pymf.tb.tb import _tb_type
-def tb_to_rparams(tb):
- """Convert a mean-field tight-binding model to a set of real parameters.
+def tb_to_rparams(tb: _tb_type) -> np.ndarray:
+ """Parametrise a hermitian tight-binding dictionary by a real vector.
Parameters
----------
- tb : dict
- Mean-field tight-binding model.
+ tb :
+ tight-binding dictionary.
Returns
-------
- dict
- Real parameters.
+ :
+ 1D real vector that parametrises the tight-binding dictionary.
"""
- return complex_to_real(tb_to_flat(tb)) # placeholder for now
+ return complex_to_real(tb_to_flat(tb))
-def rparams_to_tb(r_params, key_list, size):
- """Extract mean-field tight-binding model from a set of real parameters.
+def rparams_to_tb(
+ tb_params: np.ndarray, tb_keys: list[tuple[None] | tuple[int, ...]], ndof: int
+) -> _tb_type:
+ """Extract a hermitian tight-binding dictionary from a real vector parametrisation.
Parameters
----------
- r_params : dict
- Real parameters.
- key_list : list
- List of the keys of the mean-field tight-binding model, meaning all the
- hoppings.
- size : tuple
- Shape of the mean-field tight-binding model.
+ tb_params :
+ 1D real array that parametrises the tight-binding dictionary.
+ tb_keys :
+ List of keys of the tight-binding dictionary.
+ ndof :
+ Number internal degrees of freedom within the unit cell.
Returns
-------
- dict
- Mean-field tight-binding model.
+ :
+ Tight-biding dictionary.
"""
- flat_matrix = real_to_complex(r_params)
- return flat_to_tb(flat_matrix, (len(key_list), size, size), key_list)
+ flat_matrix = real_to_complex(tb_params)
+ return flat_to_tb(flat_matrix, ndof, tb_keys)
diff --git a/pymf/solvers.py b/pymf/solvers.py
index 9099f39749878dfac9374f19291593d6262e9814..520941232e403673ae99b998c21ed896f4c7efa2 100644
--- a/pymf/solvers.py
+++ b/pymf/solvers.py
@@ -1,65 +1,76 @@
from functools import partial
-
import numpy as np
import scipy
+from typing import Optional, Callable
from pymf.params.rparams import rparams_to_tb, tb_to_rparams
-from pymf.tb.tb import add_tb
+from pymf.tb.tb import add_tb, _tb_type
+from pymf.model import Model
from pymf.tb.utils import calculate_fermi_energy
-def cost(mf_param, Model, nk=100):
- """Define the cost function for fixed point iteration.
+def cost(mf_param: np.ndarray, model: Model, nk: int = 20) -> np.ndarray:
+ """Defines the cost function for root solver.
- The cost function is the difference between the input mean-field real space
- parametrisation and a new mean-field.
+ The cost function is the difference between the computed and inputted mean-field.
Parameters
----------
- mf_param : numpy.array
- The mean-field real space parametrisation.
- Model : Model
- The model object.
- nk : int, optional
- The number of k-points to use in the grid. The default is 100.
+ mf_param :
+ 1D real array that parametrises the mean-field correction.
+ Model :
+ Interacting tight-binding problem definition.
+ nk :
+ Number of k-points in a grid to sample the Brillouin zone along each dimension.
+ If the system is 0-dimensional (finite), this parameter is ignored.
+
+ Returns
+ -------
+ :
+ 1D real array that is the difference between the computed and inputted mean-field
+ parametrisations
"""
- shape = Model._size
- mf_tb = rparams_to_tb(mf_param, list(Model.h_int), shape)
- mf_tb_new = Model.mfield(mf_tb, nk=nk)
- mf_params_new = tb_to_rparams(mf_tb_new)
+ shape = model._ndof
+ mf = rparams_to_tb(mf_param, list(model.h_int), shape)
+ mf_new = model.mfield(mf, nk=nk)
+ mf_params_new = tb_to_rparams(mf_new)
return mf_params_new - mf_param
def solver(
- Model, mf_guess, nk=100, optimizer=scipy.optimize.anderson, optimizer_kwargs={}
-):
- """Solve the mean-field self-consistent equation.
+ model: Model,
+ mf_guess: np.ndarray,
+ nk: int = 20,
+ optimizer: Optional[Callable] = scipy.optimize.anderson,
+ optimizer_kwargs: Optional[dict[str, str]] = {"M": 0},
+) -> _tb_type:
+ """Solve for the mean-field correction through self-consistent root finding.
Parameters
----------
- Model : Model
- The model object.
- mf_guess : numpy.array
- The initial guess for the mean-field tight-binding model.
- nk : int, optional
- The number of k-points to use in the grid. The default is 100. In the
- 0-dimensional case, this parameter is ignored.
- optimizer : scipy.optimize, optional
- The optimizer to use to solve for fixed-points. The default is
- scipy.optimize.anderson.
- optimizer_kwargs : dict, optional
- The keyword arguments to pass to the optimizer. The default is {}.
+ model :
+ Interacting tight-binding problem definition.
+ mf_guess :
+ The initial guess for the mean-field correction in the tight-binding dictionary format.
+ nk :
+ Number of k-points in a grid to sample the Brillouin zone along each dimension.
+ If the system is 0-dimensional (finite), this parameter is ignored.
+ optimizer :
+ The solver used to solve the fixed point iteration.
+ Default uses `scipy.optimize.anderson`.
+ optimizer_kwargs :
+ The keyword arguments to pass to the optimizer.
Returns
-------
- result : numpy.array
- The mean-field tight-binding model.
+ :
+ Mean-field correction solution in the tight-binding dictionary format.
"""
- shape = Model._size
+ shape = model._ndof
mf_params = tb_to_rparams(mf_guess)
- f = partial(cost, Model=Model, nk=nk)
+ f = partial(cost, model=model, nk=nk)
result = rparams_to_tb(
- optimizer(f, mf_params, **optimizer_kwargs), list(Model.h_int), shape
+ optimizer(f, mf_params, **optimizer_kwargs), list(model.h_int), shape
)
- fermi = calculate_fermi_energy(add_tb(Model.h_0, result), Model.filling, nk=nk)
- return add_tb(result, {Model._local_key: -fermi * np.eye(Model._size)})
+ fermi = calculate_fermi_energy(add_tb(model.h_0, result), model.filling, nk=nk)
+ return add_tb(result, {model._local_key: -fermi * np.eye(model._ndof)})
diff --git a/pymf/tb/tb.py b/pymf/tb/tb.py
index a059323b98b434dd41b3a8b413442fa21c95b88d..445a99b483292f885368842b52a240c9e86f1ab3 100644
--- a/pymf/tb/tb.py
+++ b/pymf/tb/tb.py
@@ -1,43 +1,45 @@
import numpy as np
+_tb_type = dict[tuple[()] | tuple[int, ...], np.ndarray]
-def add_tb(tb1, tb2):
- """Add up two tight-binding models together.
+
+def add_tb(tb1: _tb_type, tb2: _tb_type) -> _tb_type:
+ """Add up two tight-binding dictionaries together.
Parameters
----------
- tb1 : dict
- Tight-binding model.
- tb2 : dict
- Tight-binding model.
+ tb1 :
+ Tight-binding dictionary.
+ tb2 :
+ Tight-binding dictionary.
Returns
-------
- dict
- Sum of the two tight-binding models.
+ :
+ Sum of the two tight-binding dictionaries.
"""
return {k: tb1.get(k, 0) + tb2.get(k, 0) for k in frozenset(tb1) | frozenset(tb2)}
-def scale_tb(tb, scale):
- """Scale a tight-binding model.
+def scale_tb(tb: _tb_type, scale: float) -> _tb_type:
+ """Scale a tight-binding dictionary by a constant.
Parameters
----------
- tb : dict
- Tight-binding model.
- scale : float
- The scaling factor.
+ tb :
+ Tight-binding dictionary.
+ scale :
+ Constant to scale the tight-binding dictionary by.
Returns
-------
- dict
- Scaled tight-binding model.
+ :
+ Scaled tight-binding dictionary.
"""
return {k: tb.get(k, 0) * scale for k in frozenset(tb)}
-def compare_dicts(dict1, dict2, atol=1e-10):
+def compare_dicts(dict1: dict, dict2: dict, atol: float = 1e-10) -> None:
"""Compare two dictionaries."""
for key in frozenset(dict1) | frozenset(dict2):
assert np.allclose(dict1[key], dict2[key], atol=atol)
diff --git a/pymf/tb/transforms.py b/pymf/tb/transforms.py
index 4befee2240393754d58623f0095745e3de6d7efb..6d7a8715bdc01b64f18710ac89066989d1820f2c 100644
--- a/pymf/tb/transforms.py
+++ b/pymf/tb/transforms.py
@@ -1,29 +1,30 @@
import itertools
import numpy as np
+from scipy.fftpack import ifftn
+from pymf.tb.tb import _tb_type
-def tb_to_khamvector(tb, nk, ks=None):
- """Real-space tight-binding model to hamiltonian on k-space grid.
+
+def tb_to_kgrid(tb: _tb_type, nk: int) -> np.ndarray:
+ """Evaluate a tight-binding dictionary on a k-space grid.
Parameters
----------
- tb : dict
- A dictionary with real-space vectors as keys and complex np.arrays as values.
- nk : int
- Number of k-points along each direction.
- ks : 1D-array
- Set of k-points. Repeated for all directions.
+ tb :
+ Tight-binding dictionary to evaluate on a k-space grid.
+ nk :
+ Number of k-points in a grid to sample the Brillouin zone along each dimension.
+ If the system is 0-dimensional (finite), this parameter is ignored.
Returns
-------
- ndarray
- Hamiltonian evaluated on a k-point grid.
-
+ :
+ Tight-binding dictionary evaluated on a k-space grid.
+ Has shape (nk, nk, ..., ndof, ndof), where ndof is number of internal degrees of freedom.
"""
ndim = len(list(tb)[0])
- if ks is None:
- ks = np.linspace(-np.pi, np.pi, nk, endpoint=False)
- ks = np.concatenate((ks[nk // 2 :], ks[: nk // 2]), axis=0) # shift for ifft
+ ks = np.linspace(-np.pi, np.pi, nk, endpoint=False)
+ ks = np.concatenate((ks[nk // 2 :], ks[: nk // 2]), axis=0) # shift for ifft
kgrid = np.meshgrid(*([ks] * ndim), indexing="ij")
num_keys = len(list(tb.keys()))
@@ -37,66 +38,41 @@ def tb_to_khamvector(tb, nk, ks=None):
return np.sum(tb_array * k_dependency, axis=0)
-def ifftn_to_tb(ifft_array):
- """Convert an array from ifftn to a tight-binding model format.
+def kgrid_to_tb(kgrid_array: np.ndarray) -> _tb_type:
+ """
+ Convert a k-space grid array to a tight-binding dictionary.
Parameters
----------
- ifft_array : ndarray
- An array obtained from ifftn.
-
+ kgrid_array :
+ K-space grid array to convert to a tight-binding dictionary.
+ The array should be of shape (nk, nk, ..., ndof, ndof),
+ where ndof is number of internal degrees of freedom.
Returns
-------
- dict
- A dictionary with real-space vectors as keys and complex np.arrays as values.
+ :
+ Tight-binding dictionary.
"""
- size = ifft_array.shape[:-2]
-
- keys = [np.arange(-size[0] // 2 + 1, size[0] // 2) for i in range(len(size))]
- keys = itertools.product(*keys)
- return {tuple(k): ifft_array[tuple(k)] for k in keys}
+ ndim = len(kgrid_array.shape) - 2
+ return ifftn_to_tb(ifftn(kgrid_array, axes=np.arange(ndim)))
-def kham_to_tb(kham, hopping_vecs, ks=None):
- """Extract hopping matrices from Bloch Hamiltonian.
+def ifftn_to_tb(ifft_array: np.ndarray) -> _tb_type:
+ """
+ Convert the result of `scipy.fft.ifftn` to a tight-binding dictionary.
Parameters
----------
- kham : nd-array
- Bloch Hamiltonian matrix kham[k_x, ..., k_n, i, j]
- hopping_vecs : list
- List of hopping vectors, will be the keys to the tb.
- ks : 1D-array
- Set of k-points. Repeated for all directions. If the system is finite,
- ks=None`.
-
+ ifft_array :
+ Result of `scipy.fft.ifftn` to convert to a tight-binding dictionary.
+ The input to `scipy.fft.ifftn` should be from `tb_to_khamvector`.
Returns
-------
- scf_model : dict
- Tight-binding model of Hartree-Fock solution.
+ :
+ Tight-binding dictionary.
"""
- if ks is not None:
- ndim = len(kham.shape) - 2
- dk = np.diff(ks)[0]
- nk = len(ks)
- k_pts = np.tile(ks, ndim).reshape(ndim, nk)
- k_grid = np.array(np.meshgrid(*k_pts))
- k_grid = k_grid.reshape(k_grid.shape[0], np.prod(k_grid.shape[1:]))
- kham = kham.reshape(np.prod(kham.shape[:ndim]), *kham.shape[-2:])
-
- hopps = (
- np.einsum(
- "ij,jkl->ikl",
- np.exp(1j * np.einsum("ij,jk->ik", hopping_vecs, k_grid)),
- kham,
- )
- * (dk / (2 * np.pi)) ** ndim
- )
-
- h_0 = {}
- for i, vector in enumerate(hopping_vecs):
- h_0[tuple(vector)] = hopps[i]
+ size = ifft_array.shape[:-2]
- return h_0
- else:
- return {(): kham}
+ keys = [np.arange(-size[0] // 2 + 1, size[0] // 2) for i in range(len(size))]
+ keys = itertools.product(*keys)
+ return {tuple(k): ifft_array[tuple(k)] for k in keys}
diff --git a/pymf/tb/utils.py b/pymf/tb/utils.py
index cc20753f704bf160e6b0eec752ae1ae50adebc46..123b559f3cd80683cb1f28143aa86f21df085303 100644
--- a/pymf/tb/utils.py
+++ b/pymf/tb/utils.py
@@ -1,30 +1,31 @@
from itertools import product
-
import numpy as np
-from pymf.mf import fermi_on_grid
-from pymf.tb.transforms import tb_to_khamvector
+from pymf.tb.tb import _tb_type
+from pymf.mf import fermi_on_kgrid
+from pymf.tb.transforms import tb_to_kgrid
-def generate_guess(vectors, ndof, scale=1):
- """Generate guess for a tight-binding model.
+def generate_guess(
+ tb_keys: list[tuple[None] | tuple[int, ...]], ndof: int, scale: float = 1
+) -> _tb_type:
+ """Generate hermitian guess tight-binding dictionary.
Parameters
----------
- 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.
-
+ tb_keys :
+ List of hopping vectors (tight-binding dictionary keys) the guess contains.
+ ndof :
+ Number internal degrees of freedom within the unit cell.
+ scale :
+ Scale of the random guess.
Returns
-------
- guess : tb dictionary
- Guess in the form of a tight-binding model.
+ :
+ Hermitian guess tight-binding dictionary.
"""
guess = {}
- for vector in vectors:
+ for vector in tb_keys:
if vector not in guess.keys():
amplitude = scale * np.random.rand(ndof, ndof)
phase = 2 * np.pi * np.random.rand(ndof, ndof)
@@ -40,25 +41,44 @@ def generate_guess(vectors, ndof, scale=1):
return guess
-def generate_vectors(cutoff, dim):
- """Generate hopping vectors up to a cutoff.
+def generate_tb_keys(cutoff: int, dim: int) -> list[tuple[None] | tuple[int, ...]]:
+ """Generate tight-binding dictionary keys up to a cutoff.
Parameters
----------
- cutoff : int
- Maximum distance along each direction.
- dim : int
- Dimension of the vectors.
+ cutoff :
+ Maximum distance along each dimension to generate tight-bindign dictionary keys for.
+ dim :
+ Dimension of the tight-binding dictionary.
Returns
-------
- List of hopping vectors.
+ :
+ List of generated tight-binding dictionary keys up to a cutoff.
"""
return [*product(*([[*range(-cutoff, cutoff + 1)]] * dim))]
-def calculate_fermi_energy(tb, filling, nk=100):
- """Calculate the Fermi energy for a given filling."""
- kham = tb_to_khamvector(tb, nk, ks=None)
+def calculate_fermi_energy(tb: _tb_type, filling: float, nk: int = 100):
+ """
+ Calculate the Fermi energy of a given tight-binding dictionary.
+
+ Parameters
+ ----------
+ tb :
+ Tight-binding dictionary.
+ filling :
+ Number of particles in a unit cell.
+ Used to determine the Fermi level.
+ nk :
+ Number of k-points in a grid to sample the Brillouin zone along each dimension.
+ If the system is 0-dimensional (finite), this parameter is ignored.
+
+ Returns
+ -------
+ :
+ Fermi energy.
+ """
+ kham = tb_to_kgrid(tb, nk)
vals = np.linalg.eigvalsh(kham)
- return fermi_on_grid(vals, filling)
+ return fermi_on_kgrid(vals, filling)
diff --git a/pymf/tests/test_graphene.py b/pymf/tests/test_graphene.py
index 8657f7d36f16c9bea7c4cea804c61f832b8f27d8..90e1c1aecd2fdc22dcba7824da0bca036b5d82ba 100644
--- a/pymf/tests/test_graphene.py
+++ b/pymf/tests/test_graphene.py
@@ -3,11 +3,13 @@ import numpy as np
import pytest
from pymf.kwant_helper import kwant_examples, utils
-from pymf.model import Model
-from pymf.solvers import solver
-from pymf.tb.tb import add_tb
-from pymf.tb.transforms import tb_to_khamvector
-from pymf.tb.utils import generate_guess
+from pymf import (
+ Model,
+ solver,
+ tb_to_kgrid,
+ generate_guess,
+ add_tb,
+)
def compute_gap(tb, fermi_energy=0, nk=100):
@@ -27,7 +29,7 @@ def compute_gap(tb, fermi_energy=0, nk=100):
gap : float
Indirect gap.
"""
- kham = tb_to_khamvector(tb, nk, ks=None)
+ kham = tb_to_kgrid(tb, nk)
vals = np.linalg.eigvalsh(kham)
emax = np.max(vals[vals <= fermi_energy])
diff --git a/pymf/tests/test_hat.py b/pymf/tests/test_hat.py
index 154daf071b9b4a22914602d87106abe5b9a329c1..6f8fedc27db004ff02d6ad22ae70d76716c8914a 100644
--- a/pymf/tests/test_hat.py
+++ b/pymf/tests/test_hat.py
@@ -1,17 +1,21 @@
# %%
import numpy as np
-from pymf.solvers import solver
-from pymf.tb import utils
-from pymf.model import Model
-from pymf.tb.tb import add_tb, scale_tb
-from pymf import mf
-from pymf import observables
import pytest
+from pymf import (
+ Model,
+ solver,
+ generate_guess,
+ scale_tb,
+ add_tb,
+ expectation_value,
+ construct_density_matrix,
+)
+
# %%
def total_energy(ham_tb, rho_tb):
- return np.real(observables.expectation_value(rho_tb, ham_tb))
+ return np.real(expectation_value(rho_tb, ham_tb))
# %%
@@ -35,10 +39,10 @@ for ndim in np.arange(4):
# %%
@np.vectorize
-def mf_rescaled(alpha, mf0, h0):
- hamiltonian = add_tb(h0, scale_tb(mf0, alpha))
- rho, _ = mf.construct_density_matrix(hamiltonian, filling=filling, nk=nk)
- hamiltonian = add_tb(h0, scale_tb(mf0, np.sign(alpha)))
+def mf_rescaled(alpha, mf0):
+ hamiltonian = add_tb(h_0, scale_tb(mf0, alpha))
+ rho, _ = construct_density_matrix(hamiltonian, filling=filling, nk=nk)
+ hamiltonian = add_tb(h_0, scale_tb(mf0, np.sign(alpha)))
return total_energy(hamiltonian, rho)
@@ -46,7 +50,7 @@ def mf_rescaled(alpha, mf0, h0):
def test_mexican_hat(seed):
np.random.seed(seed)
for h0, h_int in zip(h0s, h_ints):
- guess = utils.generate_guess(frozenset(h_int), ndof)
+ guess = generate_guess(frozenset(h_int), ndof)
_model = Model(h0, h_int, filling=filling)
mf_sol_groundstate = solver(
_model, mf_guess=guess, nk=nk, optimizer_kwargs={"M": 0}
diff --git a/pymf/tests/test_hubbard.py b/pymf/tests/test_hubbard.py
index f06103409b066f16f97174563058a8791c922245..9007d995909c750cb007d7673913b7651201d786 100644
--- a/pymf/tests/test_hubbard.py
+++ b/pymf/tests/test_hubbard.py
@@ -2,11 +2,13 @@
import numpy as np
import pytest
-from pymf.model import Model
-from pymf.solvers import solver
-from pymf.tb import utils
-from pymf.tb.tb import add_tb
from pymf.tests.test_graphene import compute_gap
+from pymf import (
+ Model,
+ solver,
+ generate_guess,
+ add_tb,
+)
repeat_number = 10
@@ -33,7 +35,7 @@ def gap_relation_hubbard(Us, nk, nk_dense, tol=1e-3):
h_int = {
(0,): U * np.kron(np.ones((2, 2)), np.eye(2)),
}
- guess = utils.generate_guess(frozenset(h_int), len(list(h_0.values())[0]))
+ guess = generate_guess(frozenset(h_int), len(list(h_0.values())[0]))
full_model = Model(h_0, h_int, filling=2)
mf_sol = solver(full_model, guess, nk=nk)
_gap = compute_gap(add_tb(h_0, mf_sol), fermi_energy=0, nk=nk_dense)
diff --git a/pymf/tests/test_params.py b/pymf/tests/test_params.py
index 1b636b310188fd5db9f6bdcbed8cc25ce9b6b5c4..640b2b68b3ecff30dc776fc6f8ba05b752d82878 100644
--- a/pymf/tests/test_params.py
+++ b/pymf/tests/test_params.py
@@ -3,7 +3,7 @@ import pytest
import numpy as np
from pymf.params.rparams import rparams_to_tb, tb_to_rparams
from pymf.tb.tb import compare_dicts
-from pymf.tb.utils import generate_guess
+from pymf import generate_guess
repeat_number = 10
diff --git a/pymf/tests/test_tb.py b/pymf/tests/test_tb.py
index 1b0a667b6b7d2698838ff60f60cfd4647d2d0b8e..fecb06d6dca555ae3bb7912f574b139517053f45 100644
--- a/pymf/tests/test_tb.py
+++ b/pymf/tests/test_tb.py
@@ -6,7 +6,7 @@ import pytest
from scipy.fftpack import ifftn
from pymf.tb.tb import compare_dicts
-from pymf.tb.transforms import ifftn_to_tb, tb_to_khamvector
+from pymf.tb.transforms import ifftn_to_tb, tb_to_kgrid
repeat_number = 10
@@ -24,6 +24,6 @@ def test_fourier(seed):
keys = [np.arange(-max_order + 1, max_order) for i in range(ndim)]
keys = it.product(*keys)
h_0 = {key: (np.random.rand(matrix_size, matrix_size) - 0.5) * 2 for key in keys}
- kham = tb_to_khamvector(h_0, nk=nk)
+ kham = tb_to_kgrid(h_0, nk=nk)
tb_new = ifftn_to_tb(ifftn(kham, axes=np.arange(ndim)))
compare_dicts(h_0, tb_new)
diff --git a/pymf/tests/test_zero_hint.py b/pymf/tests/test_zero_hint.py
index db5b872ba4706df9fddb6e01e45776a744fa962b..70deb57497fd3ca6a92c40ab3c41b1a550e925f1 100644
--- a/pymf/tests/test_zero_hint.py
+++ b/pymf/tests/test_zero_hint.py
@@ -2,10 +2,9 @@
import numpy as np
import pytest
-from pymf.model import Model
-from pymf.solvers import solver
from pymf.tb import utils
-from pymf.tb.tb import add_tb, compare_dicts
+from pymf.tb.tb import compare_dicts
+from pymf import Model, solver, generate_guess, add_tb, calculate_fermi_energy
# %%
repeat_number = 10
@@ -21,16 +20,16 @@ def test_zero_hint(seed):
dim = np.random.randint(0, 3)
ndof = np.random.randint(2, 10)
filling = np.random.randint(1, ndof)
- random_hopping_vecs = utils.generate_vectors(cutoff, dim)
+ random_hopping_vecs = utils.generate_tb_keys(cutoff, dim)
zero_key = tuple([0] * dim)
- h_0_random = utils.generate_guess(random_hopping_vecs, ndof, scale=1)
- h_int_only_phases = utils.generate_guess(random_hopping_vecs, ndof, scale=0)
- guess = utils.generate_guess(random_hopping_vecs, ndof, scale=1)
+ h_0_random = generate_guess(random_hopping_vecs, ndof, scale=1)
+ h_int_only_phases = generate_guess(random_hopping_vecs, ndof, scale=0)
+ guess = generate_guess(random_hopping_vecs, ndof, scale=1)
model = Model(h_0_random, h_int_only_phases, filling=filling)
mf_sol = solver(model, guess, nk=40, optimizer_kwargs={"M": 0, "f_tol": 1e-10})
- h_fermi = utils.calculate_fermi_energy(mf_sol, filling=filling, nk=20)
+ h_fermi = calculate_fermi_energy(mf_sol, filling=filling, nk=20)
mf_sol[zero_key] -= h_fermi * np.eye(mf_sol[zero_key].shape[0])
compare_dicts(add_tb(mf_sol, h_0_random), h_0_random, atol=1e-10)
diff --git a/pyproject.toml b/pyproject.toml
index aba62d50bf13f427a4789e7e7c22f58803f9ae1b..f910fb94f9202d8e255d315537702302cd53e303 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -47,5 +47,5 @@ include = [
]
[tool.codespell]
-skip = "*.ipynb,"
-ignore-words-list = "multline,"
+skip = "*.ipynb"
+ignore-words-list = "multline, ket, bra, braket"