dopyqo 0.0.1__py3-none-any.whl

dopyqo/__init__.py

@@ -0,0 +1,28 @@
+"""
+Dopyqo: Many-body analysis on top of Quantum ESPRESSO calculations
+"""
+
+from dopyqo.info import __version__, HOMEPAGE
+from dopyqo.helpers.vqe_helpers import *
+from dopyqo.helpers.printing import *
+from dopyqo.helpers.config import *
+from dopyqo.helpers.matrix_elements import *
+from dopyqo.helpers.atoms import elements_to_atomic_number
+from dopyqo.helpers.tcc_helpers import *
+from dopyqo.calc_matrix_elements import nuclear_repulsion_energy_ewald, iTj
+from dopyqo.eri_pair_densities import (
+    eri,
+    get_frozen_core_energy_pp,
+    get_frozen_core_energy_given_pp,
+    get_frozen_core_pot_and_energy_given_pp,
+    get_frozen_core_pot,
+)
+from dopyqo.calc_pseudo_pot import calc_pps
+from dopyqo.pseudopot import Pseudopot
+from dopyqo.units import *
+from dopyqo.wfc import Wfc
+from dopyqo.hamiltonian import Hamiltonian
+from dopyqo.wannier90 import read_u_mat
+from dopyqo.transform_matrices import transform_one_body_matrix, transform_two_body_matrix, to_density_matrix
+from dopyqo.wfc import runQE
+from dopyqo.scripts.main import run

dopyqo/calc_matrix_elements.py

@@ -0,0 +1,365 @@
+from warnings import warn
+import itertools
+import logging
+from scipy.special import erfc
+import numpy as np
+from numba import jit
+from dopyqo.colors import *
+
+
+def nuclear_repulsion_energy_ewald(
+    atom_positions: np.ndarray,
+    atomic_numbers: np.ndarray,
+    lattice_vectors: np.ndarray,
+    lattice_vectors_reciprocal: np.ndarray,
+    cell_volume: float,
+    gcut: float,
+    sigma: float | None = None,
+    rust_impl: bool = True,
+) -> np.ndarray:
+    r"""Calculate the nuclear repulsion energy E between all nuclei in a periodic lattice, i.e.,
+    E = 1/2 \sum_I \sum_J \sum_T' Z_I Z_J / |R_I-R_J-T|
+    where ' denotes that the term T=0 and I=J is omitted. Z_I are the charges of the nuclei
+    and R_I their positions. T are all possible lattice translation vectors.
+    This is equivalent to
+    E = 1/2 \sum_I Z_I \Phi_I(R_I)
+    where \Phi_I(r) is the electrostatic potential generated by all periodic nuclei EXCEPT
+    the nucleus Z_I at position R_I (its periodic images are still included).
+    To calculate \Phi_I(r) we define charge densities and solve the Poisson equation in
+    real- and reciprocal-space. The nuclei Z_I are modelled by delta distribution at positions R_I+T.
+    Additionally, Ewalds method is used, i.e., we add and substract Gaussian charge distributions
+    at all R_I+T, and subtract a homogenous background charge simulating the electron-charges.
+    A short-ranged term including all delta distributions minus all Gaussian charges (except the
+    term where T=0 and I=J) can be calculated in real-space. A long-ranged term including
+    all Gaussian charges at R_I+T minus the constant background charge can be calculated
+    in reciprocal-space. The term describing the subtraction of the Gaussian charge at R_I
+    is calculated in real-space and is the so-called self-energy term. The G -> 0 limit
+    in the reciprocal space calculation is also calculated separately and is called the
+    charged-system term. The value \sigma determining the splitting between the real-space
+    and reciprocal-space term is calculated as done in Quantum ESPRESSO (see PW/src/ewald.f90:
+    https://github.com/QEF/q-e/blob/de3035747f5d8f2ec9a67869827341ebb43f5b12/PW/src/ewald.f90)
+
+    We calculate E = E_S + E_L + E_self + E_charged with
+
+    Short-ranged term calculated in real-space
+    E_S    = 1/2 \sum_T' \sum_I \sum_J Z_I Z_J / |R_I - R_J - T| erfc(|R_I - R_J - T| \sqrt(\sigma))
+    where ' denotes that the term T=0 and I=J is omitted.
+
+    Long-ranged term calculated in reciprocal-space
+    E_L    = 1/2 4\pi/V \sum_{k \neq 0} \sum_I \sum_J Z_I Z_J / |k|^2 e^{i k . (R_I - R_J)} e^{- 1/(4 \sigma) |k|^2}
+
+    Self-energy originated in subtracting the Gaussian charge at position R_I from the periodic charge density
+    generated by all atoms (in real-space)
+    E_self = - \sum_I (Z_I^2) \sqrt(\sigma) / \sqrt(\pi)
+
+    Charged-system term from the G=0 term in the reciprocal-space sum
+    E_charged = - \pi / (2 V \sigma)  (\sum_I Z_I)^2
+
+    Notes:
+    - When pseudopotentials are used Z_I are the valence charges, i.e. the charge of the nucleus that is not
+      modelled with the pseudopotential.
+    - This function should calculate the same Ewald energy calculated in Quantum ESPRESSO, called
+      "ewald contribution" in the pw.x output file and "ewald" in the total_energy section in
+      the output-xml.
+    - Equivalent to Ewald calculation in PySCF of a cell in pyscf.pbc.get.cell.ewald function
+
+    References:
+    - http://metal.elte.hu/~groma/Anyagtudomany/kittel.pdf Appendix B
+    - https://journals.aps.org/prb/pdf/10.1103/PhysRevB.53.1814, equation (16)
+    - https://courses.physics.illinois.edu/phys466/sp2013/projects/2003/Team2/ewald_text.htm,
+    - S. W. de Leeuw et al. Proc. R. Soc. Lond. A 1980 373, 27-56, doi: 10.1098/rspa.1980.0135,  equation (1.5)
+    - M. P. Allen and D. J. Tildesley "Computer Simulation of Solids", equation (5.20)
+    - http://micro.stanford.edu/mediawiki/images/4/46/Ewald_notes.pdf
+    - https://homepages.uc.edu/~becktl/tlb_ewald.pdf
+    - Quantum ESPRESSO source code in q-e/PW/src/ewald.f90
+
+    Args:
+        atom_positions (np.ndarray): Array of the coordinates R_I of every atom described
+                                     by the pseudopotential. Shape (#atoms, 3)
+        atomic_numbers (np.ndarray):  1D array of atomic number of each atom. Shape (#atoms,)
+        lattice_vectors (np.ndarray): 2D array of lattice vectors. Each row is one lattice vector. Shape (3, 3)
+        lattice_vectors_reciprocal (np.ndarray): 2D array of reciprocal lattice vectors. Each row is one reciprocal lattice vector. Shape (3, 3)
+        cell_volume (float): Cell volume of the computationen real space cell.
+        gcut (float): Cutoff of the reciprocal space sum in E_L, i.e., k < gcutrho. In Quantum ESPRESSO the value of the G-cutoff of the density is used.
+        sigma (float | None, optional): Value for \sigma used in the Ewald summation to split real- and reciprocal-space sums. If None this is calculated to satisfy
+                                        \sum_I 2 Z_I \sqrt{\frac\sigma\pi} \mathrm{erfc}\left(\sqrt{\frac{\left(4G_\mathrm{cut}\right)^2}{4\sigma}}\right)  \leq 10^{-7}
+                                        Defaults to None.
+
+    Raises:
+        RuntimeError: If \sigma could not be calculated.
+
+    Returns:
+        float: Ewald energy
+    """
+
+    if rust_impl:
+        try:
+            import dopyqo_rs as calc_rs
+        except ImportError:
+            print(f"{ORANGE}Ewald warning: Could not import dopyqo_rs package. Falling back to python implementation.{RESET_COLOR}")
+            rust_impl = False  # Set such that variable accurately represents if rust implementation is used
+        else:  # No exception
+            logging.info("Using Rust implementation.")
+            return calc_rs.ewald(
+                atom_positions,
+                atomic_numbers,
+                lattice_vectors,
+                lattice_vectors_reciprocal,
+                cell_volume,
+                gcut,
+            )
+
+    if sigma is None:
+        sigma = 2.8
+        charge = np.sum(atomic_numbers)
+        gcutm = gcut**2
+        #
+        # choose sigma in order to have convergence in the sum over G
+        # upperbound is a safe upper bound for the error in the sum over G
+        #
+        while True:
+            if sigma <= 0.0:
+                raise RuntimeError("Optimal sigma for Ewald sum not found!")
+            upperbound = 2.0 * charge**2 * np.sqrt(sigma / np.pi) * erfc(np.sqrt(gcutm / 4.0 / sigma))
+            if upperbound > 1e-7:
+                sigma = sigma - 0.1
+            else:
+                break
+
+    logging.info("sigma %f", sigma)
+
+    e_short = 0.0  # real-space sum
+    converged_real = False
+    # n_max = 5
+    n_max_used = 0
+
+    alat = np.linalg.norm(lattice_vectors[0], ord=2)
+    t_vec_max_norm = 4.0 / np.sqrt(sigma) / alat
+    b_norms = np.linalg.norm(lattice_vectors_reciprocal, ord=2, axis=1)
+    n_max_x, n_max_y, n_max_z = (b_norms * t_vec_max_norm).astype(int) + 2
+
+    # Generate list or translation vectors ordered by their norm
+    # n_unordered = itertools.product(range(-n_max, n_max + 1), repeat=3)
+    logging.debug("Generating unordered translation vectors...")
+    t_vecs_unordered = [
+        [np.dot(lattice_vectors, np.array([nx, ny, nz])), [nx, ny, nz]]
+        for nx, ny, nz in itertools.product(
+            range(-n_max_x, n_max_x + 1),
+            range(-n_max_y, n_max_y + 1),
+            range(-n_max_z, n_max_z + 1),
+        )
+        # for nx, ny, nz in n_unordered
+    ]
+
+    t_vecs = sorted(t_vecs_unordered, key=lambda x: np.linalg.norm(x[0], ord=2))
+
+    for t_vec, (nx, ny, nz) in t_vecs:
+        for i, pos_i in enumerate(atom_positions):  # \sum_I
+            z_i = atomic_numbers[i]
+            for j, pos_j in enumerate(atom_positions):  # \sum_J
+                if np.all(t_vec == 0.0) and i == j:
+                    continue
+                z_j = atomic_numbers[j]
+                r = np.linalg.norm(pos_i - pos_j - t_vec, ord=2)
+                e_short += z_i * z_j * erfc(r * np.sqrt(sigma)) / r
+
+                # with this choice terms up to ZiZj*erfc(4) are counted (erfc(4)=1.5e-8)
+                if r >= 4 / np.sqrt(sigma):
+                    # print(f"Stopping real-sum at {nx}/{ny}/{nz} (sigma={sigma})")
+                    converged_real = True
+                    n_max_used = np.max([np.abs(nx), np.abs(ny), np.abs(nz), n_max_used])
+                    # break
+        # if converged_real:
+        #     break
+    if converged_real:
+        logging.info("Real-space sum converged! Largest n was %i", n_max_used)
+    else:
+        logging.warning("Real-space sum NOT converged!")
+
+    e_short *= 1 / 2
+
+    e_long = 0.0  # reciprocal-space sum
+
+    # Estimate size of Miller indices grid using the cutoff-energy and reciprocal lattice vectors
+    # A simple, but not accurate enough, estimate would be gcutrho/|b_i| for i \in {1, 2, 3}
+    # Here we take the shape of the reciprocal lattice into account:
+    #       max[ gcutrho/(|b_i| + \sum_{j \neq i} b_i . b_j ) ] for i \in {1, 2, 3}
+    # TODO: To use the same grid QE uses we have to read the Miller indices used for the density,
+    #       which is probably saved in charge-density.dat.
+    #       Alternatively check how pymatgen chooses the reciprocal grid in its Ewald calculation method
+    b1 = lattice_vectors_reciprocal[0]
+    b2 = lattice_vectors_reciprocal[1]
+    b3 = lattice_vectors_reciprocal[2]
+    mill_max = max(
+        [
+            round(abs(gcut / (np.linalg.norm(b1) + np.dot(b1, b2) + np.dot(b1, b3)))),
+            round(abs(gcut / (np.linalg.norm(b2) + np.dot(b2, b1) + np.dot(b2, b3)))),
+            round(abs(gcut / (np.linalg.norm(b3) + np.dot(b3, b1) + np.dot(b3, b2)))),
+        ]
+    )
+    x, y, z = np.meshgrid(
+        np.arange(-mill_max, mill_max + 1),
+        np.arange(-mill_max, mill_max + 1),
+        np.arange(-mill_max, mill_max + 1),
+        indexing="ij",
+    )
+    mill_rho = np.stack((x.ravel(), y.ravel(), z.ravel()), axis=1)
+    k_vecs = np.einsum("ij, kj -> ki", lattice_vectors_reciprocal.T, mill_rho)
+    norms = np.linalg.norm(k_vecs, axis=1)
+    k_vecs = k_vecs[norms <= gcut]
+    for k_vec in k_vecs:
+        k2 = np.linalg.norm(k_vec, ord=2) ** 2
+        if np.isclose(k2, 0.0):
+            continue
+        structure_factor = np.sum(atomic_numbers * np.exp(1.0j * np.dot(atom_positions, k_vec)))
+        e_long += np.exp(-k2 / (4 * sigma)) * np.abs(structure_factor) ** 2 / k2
+
+    e_long *= 4.0 * np.pi / 2.0 / cell_volume
+
+    logging.debug("Reciprocal-space sum calculated!")
+
+    # Term arising from substracting the Gaussian charge distribution at R_I
+    # in real-space
+    e_self = -np.sum(atomic_numbers**2) * np.sqrt(sigma / np.pi)
+
+    # The analytic limit for G → 0 in the reciprocal sum
+    # The G=0 term in the reciprocal term is only partly cancelled by the uniform
+    # background charge. The remaining term is the following
+    # Also see Martin, 2020, "Electronic Structure" eq. (F.6) at the text afterwards
+    # which unfortunately misses a derivation
+    e_charged = -np.pi / cell_volume / sigma * np.sum(atomic_numbers) ** 2 / 2
+    # e_charged = -np.sum(atomic_numbers) ** 2 / sigma / 2 * np.pi / cell_volume
+
+    e_tot = e_short + e_long + e_self + e_charged
+
+    logging.info(
+        "e_short: %f, e_long: %f, e_self: %f, e_charged: %f, e_tot: %f",
+        e_short,
+        e_long,
+        e_self,
+        e_charged,
+        e_tot,
+    )
+
+    return e_tot
+
+
+def check_symmetry_one_body_matrix(matrix: np.ndarray):
+    """Check if given matrix satisfies the symmetries of one-body matrices (hermitian)
+
+    Args:
+        matrix (np.ndarray): Matrix to check symmetries of
+
+    Returns:
+        bool: Bool specifying of symmetry is satisfied
+    """
+    matrix_hermitian = matrix.T.conj()
+    allclose_hermitian = np.allclose(matrix, matrix_hermitian)
+    return allclose_hermitian
+
+
+def check_symmetry_two_body_matrix(matrix: np.ndarray):
+    """Check if given matrix satisfies the symmetries of ERIs
+
+    Args:
+        matrix (np.ndarray): Matrix to check symmetries of
+
+    Returns:
+        tuple[bool, bool, bool]: Bools specifying if following symmetries are satisfied: swap symmetry, hermitian symmetry, hermitian+swap symmetry
+    """
+    matrix_swap = matrix.swapaxes(0, 1).swapaxes(2, 3)
+    matrix_hermitian = matrix.T.conj()
+    matrix_hemitian_swap = matrix.T.conj().swapaxes(0, 1).swapaxes(2, 3)
+
+    allclose_swap = np.allclose(matrix, matrix_swap)  # , rtol=1e-5, atol=1e-5)
+    allclose_hermitian = np.allclose(matrix, matrix_hermitian)  # , rtol=1e-5, atol=1e-5)
+    allclose_hermitian_swap = np.allclose(matrix, matrix_hemitian_swap)  # , rtol=1e-5, atol=1e-5)
+
+    allclose_ljki = None
+    allclose_ikjl = None
+    allclose_kilj = None
+    allclose_jlik = None
+    if np.allclose(matrix.real, matrix):
+        logging.info("Two-body matrix is real. Checking real-symmetries...")
+        allclose_hermitian = np.allclose(matrix, matrix.T)  # ijkl = lkji
+        allclose_swap = np.allclose(matrix, matrix.transpose(1, 0, 3, 2))  # ijkl = jilk
+        allclose_ljki = np.allclose(matrix, matrix.transpose(3, 1, 2, 0))  # ijkl = ljki
+        allclose_ikjl = np.allclose(matrix, matrix.transpose(0, 2, 1, 3))  # ijkl = ikjl
+        allclose_hermitian_swap = np.allclose(matrix, matrix.T.transpose(1, 0, 3, 2))  # ijkl = klij
+        allclose_kilj = np.allclose(matrix, matrix.transpose(2, 0, 3, 1))  # ijkl = kilj
+        allclose_jlik = np.allclose(matrix, matrix.transpose(1, 3, 0, 2))  # ijkl = jlik
+
+    return (
+        allclose_swap,
+        allclose_hermitian,
+        allclose_hermitian_swap,
+        allclose_ljki,
+        allclose_ikjl,
+        allclose_kilj,
+        allclose_jlik,
+    )
+
+
+def iTj(p: np.ndarray, c_ip: np.ndarray) -> np.ndarray:  # Calculates <i|T|j>
+    """Calculate kinetic energy matrix in Hartree units in the Kohn-Sham basis
+
+    Args:
+        p (np.ndarray): Array of momentum vectors, shape (#waves, 3)
+        c_ip (np.ndarray): Array of coefficients describing the Kohn-Sham orbitals
+            in the plane wave basis, shape (#bands, #waves)
+
+    Returns:
+        np.ndarray: Kinetic energy matrix
+    """
+    # Kinetic energy matrix in Hartree units
+    p_norm = p[:, 0] ** 2 + p[:, 1] ** 2 + p[:, 2] ** 2  # p^2
+    # 1/2 <i|p^2|j>
+    return 0.5 * np.einsum("ip, p, jp -> ij", c_ip.conjugate(), p_norm, c_ip)
+
+
+def iUj(
+    p: np.ndarray,
+    c_ip: np.ndarray,
+    atom_positions: np.ndarray,
+    atomic_numbers: np.ndarray,
+    cell_volume: float,
+) -> np.ndarray:
+    r"""Calculate nuclear interaction (U(r) = \sum_i Z_i / |R_I-r|) matrix in Hartree units in the Kohn-Sham basis
+
+    Args:
+        p (np.ndarray): Array of momentum vectors
+        c_ip (np.ndarray): Array of coefficients describing the Kohn-Sham orbitals in the plane wave basis
+        atom_positions (np.ndarray): 2D array of positions of each atom. Shape (#atoms, 3)
+        atomic_numbers (np.ndarray): 1D array of atomic number of each ato. Shape (#atoms,)m
+        cell_volume (float): Cell volume of the computationen real space cell.
+
+    Returns:
+        np.ndarray: Nuclear interaction matrix
+    """
+    # Nuclear interaction matrix in Hartree units
+    # Calculates <i|U|j>
+    # <p|U|q> = 4\pi / cell_volume \sum_I Z_I  exp(-i (p-q) . R_I) 1/(q-p)²
+    # <i|U|j> = \sum_{p,q,p!=q}  c_{p,i}^* c_{q,j} U_pq = 4\pi / cell_volume \sum_{p,q,p!=q} \sum_I Z_I  c_{p,i}^* c_{q,j} exp(-i (p-q) . R_I) 1/(q-p)²
+    q = p
+    Z_I = atomic_numbers  # shape (#atoms,)
+    R_I = atom_positions  # shape (#atoms, 3)
+
+    q = p
+    p_minus_q = p[:, None] - q[None]  # shape (#waves, #waves, 3)
+    # q[None] is of shape (1, #waves, 3) while p[:, None] is of shape (#waves, 1, 3)
+    # so p_minus_q[i,j] = p[i] - q[j]
+
+    p_minus_q = q[:, None] - p[None]  # shape (#waves, #waves, 3)
+    p_minus_q_norm = np.linalg.norm(p_minus_q, ord=2, axis=2)  # Has zeros on diagonal, shape (#waves, #waves)
+    p_minus_q_norm_filtered = p_minus_q_norm.copy() + np.eye(p_minus_q_norm.shape[0])  # Replace zeros on diagonal with ones
+
+    p_minus_q_dot_R = np.sum(p_minus_q[None] * R_I[:, None, None], axis=3)  # sum over 3D-coordinates, shape (#atoms, #waves, #waves)
+    # \sum_I 4\pi/V e^{-i (G-G').R_I}
+    prefactor = 4 * np.pi / cell_volume * np.sum(Z_I[:, None, None] * np.exp(-1j * p_minus_q_dot_R), axis=0)  # sum over atoms, shape (#waves, #waves)
+    one_over_p_minus_q_norm_squared = 1 / p_minus_q_norm_filtered**2
+    one_over_p_minus_q_norm_squared = one_over_p_minus_q_norm_squared.copy() - np.eye(one_over_p_minus_q_norm_squared.shape[0])
+
+    u_mat = prefactor * one_over_p_minus_q_norm_squared
+    assert (u_mat.diagonal() == 0.0).all()
+
+    return c_ip.conj() @ u_mat @ c_ip.T