pbchf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A simple PBC-SCF implementation\n",
    "## https://github.com/sunqm/pbchf\n",
    "\n",
    "Read this doc on nbviewer:\n",
    "https://nbviewer.jupyter.org/github/sunqm/pbchf/blob/master/pbchf.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "nbpresent": {
     "id": "dd885cde-5935-4bb1-8c37-138c6505b0f2"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pyscf.pbc import gto\n",
    "\n",
    "TEST = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "nbpresent": {
     "id": "f767ae35-a26a-4efa-a2e0-8de4be7dc858"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pyscf.pbc.gto.cell.Cell at 0x7f4c1e4ad1d0>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cell = gto.Cell()\n",
    "cell.a = np.eye(3) * 4.\n",
    "cell.atom = 'He 2 2 2'\n",
    "cell.basis = 'unc-sto-3g'\n",
    "cell.unit = 'Bohr'\n",
    "cell.build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gamma Point\n",
    "### Grids for numerical integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_grids(a, mesh):\n",
    "    grids_x = a[0,0] * np.arange(mesh[0]) / mesh[0]\n",
    "    grids_y = a[1,1] * np.arange(mesh[1]) / mesh[1]\n",
    "    grids_z = a[2,2] * np.arange(mesh[2]) / mesh[2]\n",
    "\n",
    "    grids = []\n",
    "    for x in grids_x:\n",
    "        for y in grids_y:\n",
    "            for z in grids_z:\n",
    "                grids.append((x, y, z))\n",
    "    return np.array(grids)\n",
    "\n",
    "def get_grids_ref(a, mesh):\n",
    "    grids_x = a[0,0] * np.arange(mesh[0]) / mesh[0]\n",
    "    grids_y = a[1,1] * np.arange(mesh[1]) / mesh[1]\n",
    "    grids_z = a[2,2] * np.arange(mesh[2]) / mesh[2]\n",
    "\n",
    "    grids = np.empty([mesh[0], mesh[1], mesh[2], 3])\n",
    "    grids[:,:,:,0] = grids_x[:,None,None]\n",
    "    grids[:,:,:,1] = grids_y[None,:,None]\n",
    "    grids[:,:,:,2] = grids_z[None,None,:]\n",
    "    return grids.reshape(-1, 3)\n",
    "\n",
    "if TEST:\n",
    "    mesh = [20,20,20]\n",
    "    print(abs(get_grids(cell.a, mesh) -\n",
    "              get_grids_ref(cell.a, mesh)).max())\n",
    "\n",
    "get_grids = get_grids_ref"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The value of PBC basis functions on real space grids\n",
    "\n",
    "\\begin{equation}\n",
    "\\phi(\\mathbf{r}) = \\frac{1}{\\sqrt{N}}\\sum_{\\mathbf{L}} \\mu(\\mathbf{r}-\\mathbf{L})\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_ao_values(cell, n_images, mesh):\n",
    "    cell = cell.copy()\n",
    "    grids = get_grids(cell.a, mesh)\n",
    "    Ls = get_lattice_Ls(cell, n_images)\n",
    "    atom_ref = cell.atom_coords()\n",
    "\n",
    "    ao = 0\n",
    "    for L in Ls:\n",
    "        cell.atom = [['He', atom_ref[0] + L]]\n",
    "        cell.build()\n",
    "        ao += cell.eval_gto('GTOval', grids)\n",
    "    return ao\n",
    "\n",
    "def get_lattice_Ls(cell, n_images):\n",
    "    n_images_x = n_images[0]\n",
    "    n_images_y = n_images[1]\n",
    "    n_images_z = n_images[2]\n",
    "    Ls = []\n",
    "    for ix in range(-n_images_x, n_images_x+1):\n",
    "        for iy in range(-n_images_y, n_images_y+1):\n",
    "            for iz in range(-n_images_z, n_images_z+1):\n",
    "                L = np.einsum('x,xy->y', (ix, iy, iz), cell.a)\n",
    "                Ls.append(L)\n",
    "    return Ls\n",
    "\n",
    "def get_ao_values_ref(cell, mesh=None):\n",
    "    if mesh is None:\n",
    "        mesh = cell.mesh\n",
    "    grids = get_grids(cell.a, mesh)\n",
    "    return cell.pbc_eval_gto('GTOval', grids)\n",
    "\n",
    "if TEST:    \n",
    "    mesh = [10,10,10]\n",
    "    ao_ref = get_ao_values_ref(cell, mesh)\n",
    "    ao = get_ao_values(cell, [0,0,0], mesh)\n",
    "    print(abs(ao-ao_ref).max())\n",
    "    ao = get_ao_values(cell, [1,1,1], mesh)\n",
    "    print(abs(ao-ao_ref).max())\n",
    "    ao = get_ao_values(cell, [2,2,2], mesh)\n",
    "    print(abs(ao-ao_ref).max())\n",
    "\n",
    "get_ao_values = get_ao_values_ref"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Overlap integrals\n",
    "The first method\n",
    "\\begin{equation}\n",
    "\\langle \\phi_{\\mu} | \\phi_{\\nu} \\rangle\n",
    "= \\frac{1}{N} \\int \\phi_\\mu(\\mathbf{r})^* \\phi_\\nu(\\mathbf{r}) d\\mathbf{r}\n",
    "= \\sum_{\\mathbf{r}\\in \\Omega} w \\phi_\\mu(\\mathbf{r})^* \\phi_\\nu(\\mathbf{r})\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_weight(cell):\n",
    "    return cell.vol / np.prod(cell.mesh)\n",
    "\n",
    "def get_ovlp1(cell):\n",
    "    w = get_weight(cell)\n",
    "    ao = get_ao_values(cell)\n",
    "    s = w * np.einsum('xi,xj->ij', ao.conj(), ao)\n",
    "    return s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The second method\n",
    "\\begin{align}\n",
    "\\langle \\phi_\\mu | \\phi_\\nu \\rangle\n",
    "&= \\frac{1}{N}\\sum_\\mathbf{LL'}\\int \\mu(\\mathbf{r}-\\mathbf{L})^* \\nu(\\mathbf{r} - \\mathbf{L}') d\\mathbf{r} \\\\\n",
    "&= \\frac{1}{N}\\sum_\\mathbf{LL'}\\int \\mu(\\mathbf{r})^* \\nu(\\mathbf{r} + \\mathbf{L} - \\mathbf{L}') d\\mathbf{r} \\\\\n",
    "&= \\frac{1}{N}\\sum_\\mathbf{L} \\sum_\\mathbf{L'}\\int \\mu(\\mathbf{r})^* \\nu(\\mathbf{r} + \\mathbf{L'}) d\\mathbf{r} \\\\\n",
    "&= \\int_{-\\infty}^{\\infty} \\mu(\\mathbf{r})^* \\nu(\\mathbf{r} + \\mathbf{L'}) d\\mathbf{r}\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_ovlp2(cell):\n",
    "    from pyscf import gto as mole_gto\n",
    "    cellL = cell.copy()\n",
    "    grids = get_grids(cell.a, cell.mesh)\n",
    "    Ls = get_lattice_Ls(cell, cell.nimgs)\n",
    "    atom_ref = cell.atom_coords()\n",
    "\n",
    "    s = 0\n",
    "    for L in Ls:\n",
    "        cellL.atom = [['He', atom_ref[0] + L]]\n",
    "        cellL.build()\n",
    "        s += mole_gto.intor_cross('int1e_ovlp', cell, cellL)\n",
    "    return s\n",
    "\n",
    "def get_ovlp_ref(cell):\n",
    "    return cell.pbc_intor('int1e_ovlp')\n",
    "get_ovlp = get_ovlp_ref\n",
    "\n",
    "if TEST:\n",
    "    print(abs(get_ovlp1(cell) - get_ovlp2(cell)).max())\n",
    "    print(abs(get_ovlp1(cell) - get_ovlp_ref(cell)).max())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The third method to compute overlap integrals:\n",
    " \n",
    "\\begin{equation}\n",
    "\\phi(\\mathbf{r}) \\rightarrow \\phi(\\mathbf{G})\n",
    "\\end{equation}\n",
    "where $\\mathbf{G}$ is the plane-wave vector (\"Gv\" in the code).\n",
    "Gv can be obtained by Fourier transform\n",
    "\\begin{equation}\n",
    "\\phi(\\mathbf{G}) = \\int_\\Omega e^{-i\\mathbf{G}\\cdot\\mathbf{r}} \\phi(\\mathbf{r}) d\\mathbf{r}\n",
    "\\end{equation}\n",
    "\n",
    "The discrete FT sample frequencys $G_x$ are\n",
    "\\begin{gather}\n",
    "G_x = \\frac{2\\pi n}{N L_x} \\\\\n",
    "\\begin{cases}\n",
    "n = 0,\\dots,N/2-1,\\, -N/2,\\dots,1 & N \\text{ is even} \\\\\n",
    "n = 0,\\dots,(N-1)/2,\\, -(N-1)/2,\\dots,1 & N \\text{ is odd}\n",
    "\\end{cases}\n",
    "\\end{gather}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_Gv(cell):\n",
    "    # Planewaves that corresponds to the mesh in real space\n",
    "    mesh = cell.mesh\n",
    "    Gx = np.fft.fftfreq(mesh[0], 1./mesh[0])\n",
    "    Gy = np.fft.fftfreq(mesh[1], 1./mesh[1])\n",
    "    Gz = np.fft.fftfreq(mesh[2], 1./mesh[2])\n",
    "\n",
    "    Gv = np.empty([mesh[0], mesh[1], mesh[2], 3])\n",
    "    Gv[:,:,:,0] = Gx[:,None,None]\n",
    "    Gv[:,:,:,1] = Gy[None,:,None]\n",
    "    Gv[:,:,:,2] = Gz[None,None,:]\n",
    "\n",
    "    b = cell.reciprocal_vectors()\n",
    "    Gv = np.dot(Gv, b)\n",
    "    return Gv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_aoG_values_ref(cell):\n",
    "    from pyscf.pbc import df\n",
    "    aoG = df.ft_ao.ft_ao(cell, cell.Gv)\n",
    "    return aoG\n",
    "\n",
    "from pyscf.pbc.tools import fft, ifft\n",
    "def get_aoG_values(cell):\n",
    "    aoR = get_ao_values(cell)\n",
    "    nao = aoR.shape[1]\n",
    "    aoG = []\n",
    "    for i in range(nao):\n",
    "        aoG.append(_______ * fft(aoR[:,i], cell.mesh))\n",
    "    aoG = np.array(aoG).T\n",
    "    return aoG\n",
    "\n",
    "if TEST:\n",
    "    print abs(get_aoG_values(cell) - get_aoG_values_ref(cell)).max()\n",
    "\n",
    "get_aoG_values = get_aoG_values_ref"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The overlap integrals can be numerically evaluated in reciprocal space\n",
    "\\begin{align}\n",
    "\\langle \\phi_\\mu | \\phi_\\nu \\rangle\n",
    "&= \\int \\phi_\\mu(\\mathbf{G})^* \\phi_nu(\\mathbf{G}) d\\mathbf{G} \\\\\n",
    "&= \\sum_{\\mathbf{G}} w_\\mathbf{G} \\phi_\\mu(\\mathbf{G})^* \\phi_nu(\\mathbf{G})\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "get_aoG_values = get_aoG_values_ref\n",
    "def get_ovlp3(cell):\n",
    "    aoR = get_ao_values(cell)\n",
    "    nao = aoR.shape[1]\n",
    "    aoG = get_aoG_values(cell)\n",
    "    w = 1./cell.vol\n",
    "    s = w * np.einsum('xi,xj->ij', aoG.conj(), aoG)\n",
    "    return s\n",
    "\n",
    "if TEST:\n",
    "    print(abs(get_ovlp3(cell) - get_ovlp_ref(cell)).max())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Kinetic integrals\n",
    "Method 1:\n",
    "\\begin{align}\n",
    "\\langle \\phi_\\mu | -\\frac{1}{2}\\nabla^2 | \\phi_\\nu \\rangle\n",
    "&= \\frac{1}{N} \\int (\\nabla \\phi_\\mu(\\mathbf{r})^*) (\\nabla\\phi_nu(\\mathbf{r})) d\\mathbf{r} \\\\\n",
    "&= \\sum_{\\mathbf{r}\\in \\Omega} w (\\nabla \\phi_\\mu(\\mathbf{r})^*) (\\nabla\\phi_nu(\\mathbf{r}))\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_kinetic1(cell):\n",
    "    from pyscf.pbc.dft.numint import eval_ao\n",
    "    grids = get_grids(cell.a, cell.mesh)\n",
    "    aoR, ao_dx, ao_dy, ao_dz = eval_ao(cell, grids, deriv=1)\n",
    "    w = get_weight(cell)\n",
    "\n",
    "    t = .5 * w * (np.einsum('xi,xj->ij', ao_dx.conj(), ao_dx) +\n",
    "                  np.einsum('xi,xj->ij', ao_dy.conj(), ao_dy) +\n",
    "                  np.einsum('xi,xj->ij', ao_dz.conj(), ao_dz))\n",
    "    return t"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Method 2: numerical integration in reciprocal space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_kinetic2(cell):\n",
    "    aoG = get_aoG_values(cell)\n",
    "    Gv = get_Gv(cell)\n",
    "    w = 1./cell.vol\n",
    "    t = .5 * w * ________\n",
    "    return t\n",
    "\n",
    "def get_kinetic_ref(cell):\n",
    "    return cell.pbc_intor('int1e_kin')\n",
    "get_kinetic = get_kinetic_ref\n",
    "\n",
    "if TEST:\n",
    "    print(abs(get_kinetic1(cell) - get_kinetic2(cell)).max())\n",
    "    print(abs(get_kinetic1(cell) - get_kinetic_ref(cell)).max())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Nuclear attraction integrals\n",
    "\\begin{align}\n",
    "\\langle \\phi_\\mu | V_{nuc} | \\phi_\\nu \\rangle\n",
    "&= \\frac{1}{N} \\sum_{LL'L''} \\int \\mu^*(\\mathbf{r}-\\mathbf{L}) \\nu(\\mathbf{r}-\\mathbf{L}') \\frac{Z_A}{|\\mathbf{r} - \\mathbf{R}_A - \\mathbf{L''}|} d\\mathbf{r} \\\\\n",
    "&= \\sum_{LL'} \\int \\mu^*(\\mathbf{r}-\\mathbf{L}) \\nu(\\mathbf{r}-\\mathbf{L}') \\frac{Z_A}{|\\mathbf{r} - \\mathbf{R}_A|} d\\mathbf{r} \\\\\n",
    "&=\\sum_{\\mathbf{G}\\neq 0} w_{\\mathbf{G}} \\rho_{\\mu\\nu}(-\\mathbf{G}) V_{nuc}(\\mathbf{G})\n",
    "\\end{align}\n",
    "\n",
    "\\begin{gather}\n",
    "\\rho_{\\mu\\nu}(\\mathbf{G}) = \\sum_{LL'}\\int e^{-i \\mathbf{G}\\cdot \\mathbf{r}} \\mu(\\mathbf{r}-\\mathbf{L})^* \\nu(\\mathbf{r}-\\mathbf{L'}) d\\mathbf{r}\n",
    "\\\\\n",
    "w_{\\mathbf{G}} = \\frac{1}{\\Omega}\n",
    "\\\\\n",
    "\\rho_{nuc}(\\mathbf{G}) = \\sum_A \\int e^{-i \\mathbf{G}\\cdot\\mathbf{r}} Z_A\\delta(\\mathbf{r} - \\mathbf{R}_A) d\\mathbf{r}\n",
    "= \\sum_A Z_A e^{-i \\mathbf{G}\\cdot\\mathbf{R}_A}\n",
    "\\\\\n",
    "V_{nuc}(\\mathbf{G}) = \\frac{4\\pi}{G^2} \\rho_{nuc}(\\mathbf{G})\n",
    "\\end{gather}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from pyscf.pbc.tools import fft, ifft\n",
    "def get_nuc(cell):\n",
    "    R = cell.atom_coords()\n",
    "    Gv = cell.get_Gv()\n",
    "    SI = np.exp(-1j*np.einsum('zs,xs->zx', R, Gv))\n",
    "    rho_nuc_G = -np.einsum('z,zx->x', cell.atom_charges(), SI)\n",
    "\n",
    "    aoR = get_ao_values(cell)\n",
    "    nao = aoR.shape[1]\n",
    "    pair_ao_real_space = np.einsum('xi,xj->xij', aoR.conj(), aoR)\n",
    "    pair_ao_G_space = ifft(pair_ao_real_space.reshape(-1,nao**2).T,\n",
    "                           cell.mesh).T.reshape(-1,nao,nao) * cell.vol\n",
    "\n",
    "    w = 1./cell.vol\n",
    "    G2 = np.einsum('xs,xs->x', Gv, Gv)\n",
    "    coulG = 4 * np.pi / G2\n",
    "    coulG[G2 == 0] = 0\n",
    "\n",
    "    v = w * np.einsum('x,x,xij->ij', rho_nuc_G, coulG, pair_ao_G_space)\n",
    "    return v\n",
    "\n",
    "def get_nuc_ref(cell):\n",
    "    from pyscf.pbc import df\n",
    "    return df.FFTDF(cell).get_nuc()\n",
    "\n",
    "if TEST:\n",
    "    print abs(get_nuc(cell) - get_nuc_ref(cell)).max()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_hcore(cell):\n",
    "    t = get_kinetic(cell)\n",
    "    v = get_nuc(cell)\n",
    "    return t + v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Electron repulsion integrals (ERI)\n",
    "\\begin{align}\n",
    "(ij|kl)\n",
    "&= \\int \\phi_i(\\mathbf{r})^* \\phi_j(\\mathbf{r}) \\frac{1}{|\\mathbf{r}-\\mathbf{r'}|} \\phi_k(\\mathbf{r'})^* \\phi_l(\\mathbf{r'}) d\\mathbf{r}d\\mathbf{r'} \\\\\n",
    "&=\\sum_{\\mathbf{G}\\neq 0} w_G \\frac{4\\pi}{G^2} \\rho_{ij}(\\mathbf{G}) \\rho_{kl}(-\\mathbf{G})\n",
    "\\end{align}\n",
    "where\n",
    "\\begin{equation}\n",
    "\\rho_{ij} (\\mathbf{G}) = \\int e^{-i \\mathbf{G}\\cdot\\mathbf{r}} \\phi_i^*(\\mathbf{r}) \\phi_j(\\mathbf{r}) d\\mathbf{r}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_eri(cell):\n",
    "    Gv = cell.get_Gv()\n",
    "\n",
    "    aoR = get_ao_values(cell)\n",
    "    nao = aoR.shape[1]\n",
    "\n",
    "    pair_ao_real_space = np.einsum('xi,xj->xij', aoR.conj(), aoR)\n",
    "    weight = cell.vol / np.prod(cell.mesh)\n",
    "    pair_ao_G_space = fft(pair_ao_real_space.reshape(-1,nao**2).T,\n",
    "                          cell.mesh).T.reshape(-1,nao,nao) * weight\n",
    "    pair_ao_G_inv_space = ifft(pair_ao_real_space.reshape(-1,nao**2).T,\n",
    "                               cell.mesh).T.reshape(-1,nao,nao) * cell.vol\n",
    "\n",
    "    w = 1./cell.vol\n",
    "    G2 = np.einsum('xs,xs->x', Gv, Gv)\n",
    "    coulG = 4 * np.pi / G2\n",
    "    coulG[G2 == 0] = 0\n",
    "\n",
    "    eri = ___________________________\n",
    "    return eri\n",
    "\n",
    "def get_eri_ref(cell):\n",
    "    from pyscf.pbc import df\n",
    "    nao = cell.nao_nr()\n",
    "    eri = df.FFTDF(cell).get_eri(compact=False).reshape(nao,nao,nao,nao)\n",
    "    return eri\n",
    "\n",
    "if TEST:\n",
    "    eri = get_eri(cell)\n",
    "    print(abs(eri - get_eri_ref(cell)).max())\n",
    "\n",
    "get_eri = get_eri_ref"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "get_eri = get_eri_ref\n",
    "def get_jk(eri, dm):\n",
    "    j = np.einsum('ijkl,ji->kl', eri, dm)\n",
    "    k = np.einsum('ijkl,jk->il', eri, dm)\n",
    "    return j, k\n",
    "\n",
    "def get_vhf(cell, dm, eri):\n",
    "    j, k = get_jk(eri, dm)\n",
    "    E_coul = .5 * np.einsum('ij,ji', j-.5*k, dm)\n",
    "    return j - .5 * k, E_coul\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Numerical integration for exchange-correlation functional\n",
    "The total XC energy is a functional of electron density $\\rho$\n",
    "\\begin{equation}\n",
    "E_{xc} = \\int e_{xc}[\\rho]\\, \\rho(\\mathbf{r}) d\\mathbf{r}\n",
    "\\end{equation}\n",
    "$e_{xc}$ is the XC energy per particle, which can be evaluated by the function pyscf.dft.libxc.eval_xc.\n",
    "\n",
    "In the DFT method, one also needs to calculate the XC potential $V_{xc}(\\mathbf{r})$:\n",
    "\\begin{equation}\n",
    "V_{xc}(\\mathbf{r}) = \\frac{\\delta E_{xc}}{\\delta \\rho} = e_{xc}[\\rho] + \\frac{\\delta e_{xc}}{\\delta \\rho} \\rho(\\mathbf{r})\n",
    "\\end{equation}\n",
    "\n",
    "and the matrix elements on top of that\n",
    "\\begin{align}\n",
    "\\langle \\phi_\\mu |V_{xc} |\\phi_\\nu\\rangle\n",
    "&= \\int \\phi_\\mu(\\mathbf{r})^* V_{xc}(\\mathbf{r}) \\phi_\\mu(\\mathbf{r}) d\\mathbf{r} \\\\\n",
    "&= \\sum_{\\mathbf{r}\\in\\Omega} w V_{xc}(\\mathbf{r}) \\rho_{\\mu\\nu}(\\mathbf{r})\n",
    "\\end{align}\n",
    "\\begin{equation}\n",
    "\\rho_{\\mu\\nu}(\\mathbf{r}) = \\phi_\\mu(\\mathbf{r})^* \\phi_\\nu(\\mathbf{r})\n",
    "\\end{equation}\n",
    "\n",
    "When calling pyscf.dft.libxc.eval_xc function, if deriv=1 is specified, $e_{xc}$ and $V_{xc}$ can be evaluated simultaneously. (Note \n",
    "$V_{xc}$ includes two terms as shown above, not just the derivative of $e_{xc}$.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_vxc(cell, dm, xc):\n",
    "    from pyscf.dft import libxc\n",
    "    aoR = get_ao_values(cell)\n",
    "    pair_ao = np.einsum('xi,xj->xij', aoR.conj(), aoR)\n",
    "    rho = np.einsum('xij,ji->x', pair_ao, dm)\n",
    "    e_xc, v_xc = libxc.eval_xc(xc, rho, deriv=1)[:2]\n",
    "    v_xc = v_xc[0]\n",
    "    w = get_weight(cell)\n",
    "    vxc = w * np.einsum('xij,x->ij', pair_ao, v_xc)\n",
    "    E_xc = w * np.einsum('x,x', rho, e_xc)\n",
    "    return vxc, E_xc\n",
    "\n",
    "def get_vxc_ref(cell, dm, xc):\n",
    "    from pyscf.pbc import dft\n",
    "    mf = dft.RKS(cell)\n",
    "    E_xc, vxc = mf._numint.nr_rks(cell, mf.grids, 'lda,vwn', dm)[1:]\n",
    "    return vxc, E_xc\n",
    "\n",
    "if TEST:\n",
    "    dm = cell.pbc_intor('int1e_ovlp')\n",
    "    vxc, E_xc = get_vxc(cell, dm, 'lda,vwn')\n",
    "    vxc_ref, E_xc_ref = get_vxc_ref(cell, dm, 'lda,vwn')\n",
    "    print(abs(E_xc_ref - E_xc).max())\n",
    "    print(abs(vxc_ref - vxc).max())\n",
    "\n",
    "get_vxc = get_vxc_ref"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import scipy.linalg\n",
    "def run_hf(cell):\n",
    "    nelectron = cell.nelectron\n",
    "    nocc = nelectron // 2\n",
    "\n",
    "    hcore = get_hcore(cell)\n",
    "    s = get_ovlp(cell)\n",
    "    eri = get_eri(cell)\n",
    "    dm = np.zeros_like(s)\n",
    "    vhf, E_coul = get_vhf(cell, dm, eri)\n",
    "\n",
    "    E_ewald = cell.ewald()\n",
    "    E = E_ewald\n",
    "    dE = 1e99\n",
    "\n",
    "    cycle = 0\n",
    "    while dE > 1e-4:\n",
    "        cycle += 1\n",
    "        Elast = E\n",
    "\n",
    "        fock = hcore + vhf\n",
    "        e, c = scipy.linalg.eigh(fock, s)\n",
    "        dm = np.einsum('pi,qi->pq', c[:,:nocc], c[:,:nocc].conj()) * 2\n",
    "        vhf, E_coul = get_vhf(cell, dm, eri)\n",
    "\n",
    "        E_elec = np.einsum('ij,ji', hcore, dm) + E_coul\n",
    "        E = E_ewald + E_elec\n",
    "        dE = abs(E - Elast)\n",
    "        print('SCF cycle', cycle, 'E(HF) =', E, 'dE =', dE)\n",
    "    return E, c\n",
    "\n",
    "if TEST:\n",
    "    E, c = run_hf(cell)\n",
    "    \n",
    "    # verify it by customizing pyscf SCF object\n",
    "    from pyscf.pbc import scf\n",
    "    mf = scf.RHF(cell)\n",
    "    mf.get_hcore = lambda *args: get_hcore(cell)\n",
    "    mf.get_ovlp = lambda *args: get_ovlp(cell)\n",
    "    mf._eri = get_eri(cell)\n",
    "    nao = cell.nao\n",
    "    mf.exxdiv = None  # Skip HF-X treatment\n",
    "    mf.kernel()\n",
    "    print(abs(E - mf.e_tot))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def ao2mo(cell, mos):\n",
    "    eri = get_eri(cell)\n",
    "    mo_i, mo_j, mo_k, mo_l = mos\n",
    "    eri_mo = np.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, mo_i.conj(), mo_j, mo_k.conj(), mo_l)\n",
    "    return eri_mo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## K-point sampling\n",
    "### PBC basis function with k-point label\n",
    "\\begin{equation}\n",
    "\\phi^{\\mathbf{k}}(\\mathbf{r}) = \\frac{1}{\\sqrt{N}}\\sum_{\\mathbf{L}} e^{i\\mathbf{k}\\cdot\\mathbf{L}} \\mu(\\mathbf{r}-\\mathbf{L})\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_ao_values(cell, kpt):\n",
    "    cell = cell.copy()\n",
    "    grids = get_grids(cell.a, cell.mesh)\n",
    "    Ls = get_lattice_Ls(cell, cell.nimgs)\n",
    "    atom_ref = cell.atom_coords()\n",
    "\n",
    "    ao = 0\n",
    "    for L in Ls:\n",
    "        cell.atom = [['He', atom_ref[0] + L]]\n",
    "        cell.build()\n",
    "        phase = np.exp(1j * np.dot(kpt, L))\n",
    "        ao += cell.eval_gto('GTOval', grids) * phase\n",
    "    return ao\n",
    "\n",
    "def get_ao_values_ref(cell, kpt):\n",
    "    grids = get_grids(cell.a, cell.mesh)\n",
    "    ao = cell.pbc_eval_gto('GTOval', grids, kpt=kpt)\n",
    "    return np.array(ao)\n",
    "\n",
    "if TEST:\n",
    "    kpt = np.random.rand(3)\n",
    "    print(abs(get_ao_values(cell, kpt)\n",
    "              - get_ao_values_ref(cell, kpt)).max())\n",
    "\n",
    "get_ao_values = get_ao_values_ref"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fourier transform with k point\n",
    "\n",
    "\\begin{equation}\n",
    "\\phi^\\mathbf{k}(\\tilde{\\mathbf{G}}) = \\int_\\Omega e^{-i\\tilde{\\mathbf{G}}\\cdot\\mathbf{r}} \\phi^\\mathbf{k}(\\mathbf{r}) d\\mathbf{r}\n",
    "\\end{equation}\n",
    "\n",
    "Due to the periodicity\n",
    "\\begin{equation}\n",
    "\\phi^\\mathbf{k}(\\mathbf{r}) = \\phi^\\mathbf{k}(\\mathbf{r} + \\mathbf{L})\n",
    "\\end{equation}\n",
    "the Fourier transform will be zero for most of the planewaves unless\n",
    "\\begin{equation}\n",
    "\\tilde{\\mathbf{G}} = \\mathbf{G} + \\mathbf{k}\n",
    "\\end{equation}\n",
    "where the \"Gv\" vectors $\\mathbf{G}$ are the discrete FT sample frequencies.  The DFT can be called as\n",
    "\n",
    "\\begin{equation}\n",
    "\\phi^\\mathbf{k}(\\mathbf{G} + \\mathbf{k}) = \\int_\\Omega e^{-i\\mathbf{G}\\cdot\\mathbf{r}} \\Big[ e^{-i\\mathbf{k}\\cdot\\mathbf{r}} \\phi^\\mathbf{k}(\\mathbf{r}) \\Big]d\\mathbf{r} \n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_aoG_values(cell, kpt):\n",
    "    aoR = get_ao_values(cell, kpt)\n",
    "    grids = cell.get_uniform_grids()\n",
    "    phase = ______\n",
    "    nao = cell.nao_nr()\n",
    "    aoG = []\n",
    "    for i in range(nao):\n",
    "        aoG.append(_____ * fft(np.einsum('x,x->x', phase, aoR[:,i]), cell.mesh))\n",
    "    aoG = np.array(aoG).T\n",
    "    return aoG\n",
    "\n",
    "def get_aoG_values_ref(cell, kpt):\n",
    "    from pyscf.pbc import df\n",
    "    return df.ft_ao.ft_ao(cell, cell.get_Gv(), kpt=kpt)\n",
    "\n",
    "if TEST:\n",
    "    kpt = np.random.rand(3)\n",
    "    print(abs(get_aoG_values(cell, kpt) -\n",
    "              get_aoG_values_ref(cell, kpt))).max()\n",
    "\n",
    "get_aoG_values = get_aoG_values_ref"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Overlap integrals with k point\n",
    "Method 1:\n",
    "\\begin{equation}\n",
    "\\langle \\phi^\\mathbf{k}_\\mu | \\phi^\\mathbf{k'}_\\nu\\rangle\n",
    "= \\frac{1}{N} \\int (\\phi^\\mathbf{k}_\\mu(\\mathbf{r}))^* \\phi^\\mathbf{k'}_\\nu(\\mathbf{r}) d\\mathbf{r}\n",
    "\\rightarrow \\sum_{\\mathbf{r}\\in \\Omega} w \\phi^\\mathbf{k}_\\mu(\\mathbf{r}))^* \\phi^\\mathbf{k'}_\\nu(\\mathbf{r})\n",
    "\\end{equation}\n",
    "\n",
    "Method 2:\n",
    "\\begin{align}\n",
    "\\langle \\phi^\\mathbf{k}_\\mu | \\phi^\\mathbf{k'}_\\nu \\rangle\n",
    "&= \\frac{1}{N}\\sum_\\mathbf{LL'}\\int [e^{i\\mathbf{k}\\cdot\\mathbf{L}}\\mu(\\mathbf{r}-\\mathbf{L})]^* e^{i\\mathbf{k'}\\cdot\\mathbf{L'}}\\nu(\\mathbf{r} - \\mathbf{L}') d\\mathbf{r} \\\\\n",
    "&= \\frac{1}{N}\\sum_\\mathbf{LL'}\\int [e^{i\\mathbf{k}\\cdot\\mathbf{L}}\\mu(\\mathbf{r})]^* e^{i\\mathbf{k'}\\cdot(\\mathbf{L'}-\\mathbf{L}) + i\\mathbf{k'}\\cdot\\mathbf{L}} \\nu(\\mathbf{r} + \\mathbf{L} - \\mathbf{L}') d\\mathbf{r} \\\\\n",
    "&= \\frac{1}{N}\\sum_\\mathbf{L} e^{i(\\mathbf{k'}-\\mathbf{k})\\cdot\\mathbf{L}}\n",
    "\\sum_\\mathbf{L'}\\int e^{i\\mathbf{k'}\\cdot\\mathbf{L'}}\\mu(\\mathbf{r})^* \\nu(\\mathbf{r} + \\mathbf{L'}) d\\mathbf{r} \\\\\n",
    "&=\n",
    "\\begin{cases}\n",
    "\\int_{-\\infty}^{\\infty} e^{i\\mathbf{k'}\\cdot\\mathbf{L'}}\\mu(\\mathbf{r})^* \\nu(\\mathbf{r} + \\mathbf{L'}) d\\mathbf{r} & \\mathbf{k} = \\mathbf{k'} \\\\\n",
    "0 & \\text{otherwise}\n",
    "\\end{cases}\n",
    "\\end{align}\n",
    "\n",
    "Why doesn't method 1 equal to method 2?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_ovlp1(cell, kpt1, kpt2):\n",
    "    w = get_weight(cell)\n",
    "    ao_k1 = get_ao_values(cell, kpt1)\n",
    "    ao_k2 = get_ao_values(cell, kpt2)\n",
    "    s = w * np.einsum('xi,xj->ij', ao_k1.conj(), ao_k2)\n",
    "    return np.array(s)\n",
    "\n",
    "if TEST:\n",
    "    kpt1 = np.random.rand(3)\n",
    "    kpt2 = np.random.rand(3)\n",
    "    print(get_ovlp1(cell, kpt1, kpt2).round(2))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_ovlp1(cell, kpts):\n",
    "    w = get_weight(cell)\n",
    "    s = []\n",
    "    for k, kpt in enumerate(kpts):\n",
    "        ao = get_ao_values(cell, kpt)\n",
    "        s.append(w * np.einsum('xi,xj->ij', ao.conj(), ao))\n",
    "    return np.array(s)\n",
    "\n",
    "def get_ovlp2(cell, kpts):\n",
    "    from pyscf import gto as mole_gto\n",
    "    cellL = cell.copy()\n",
    "    grids = get_grids(cell.a, cell.mesh)\n",
    "    Ls = get_lattice_Ls(cell, cell.nimgs)\n",
    "    atom_ref = cell.atom_coords()\n",
    "\n",
    "    s_k = []\n",
    "    for k, kpt in enumerate(kpts):\n",
    "        s = 0\n",
    "        for L in Ls:\n",
    "            cellL.atom = [['He', atom_ref[0] + L]]\n",
    "            cellL.build()\n",
    "            phase = np.exp(1j * np.dot(kpt, L))\n",
    "            s += phase * mole_gto.intor_cross('int1e_ovlp', cell, cellL)\n",
    "        s_k.append(s)\n",
    "    return np.array(s_k)\n",
    "\n",
    "def get_ovlp_ref(cell, kpts):\n",
    "    return cell.pbc_intor('int1e_ovlp', kpts=kpts)\n",
    "get_ovlp = get_ovlp_ref\n",
    "\n",
    "if TEST:\n",
    "    kpts = cell.make_kpts([2,1,1])\n",
    "    print(abs(get_ovlp1(cell, kpts) - get_ovlp_ref(cell, kpts)).max())\n",
    "    print(abs(get_ovlp2(cell, kpts) - get_ovlp_ref(cell, kpts)).max())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Kinetic integrals with k point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "get_aoG_values = get_aoG_values_ref\n",
    "def get_kinetic1(cell, kpts):\n",
    "    from pyscf.pbc.dft.numint import eval_ao\n",
    "    grids = get_grids(cell.a, cell.mesh)\n",
    "    w = get_weight(cell)\n",
    "    t_k = []\n",
    "    for k, kpt in enumerate(kpts):\n",
    "        aoR, ao_dx, ao_dy, ao_dz = eval_ao(cell, grids, deriv=1, kpt=kpt)\n",
    "        t = .5 * w * (np.einsum('xi,xj->ij', ao_dx.conj(), ao_dx) +\n",
    "                      np.einsum('xi,xj->ij', ao_dy.conj(), ao_dy) +\n",
    "                      np.einsum('xi,xj->ij', ao_dz.conj(), ao_dz))\n",
    "        t_k.append(t)\n",
    "    return np.array(t_k)\n",
    "\n",
    "def get_kinetic2(cell, kpts):\n",
    "    t = []\n",
    "    w = 1./cell.vol\n",
    "    for k, kpt in enumerate(kpts):\n",
    "        aoG = get_aoG_values(cell, kpt)\n",
    "        t.append(.5 * w * ______)\n",
    "    return np.array(t)\n",
    "\n",
    "def get_kinetic_ref(cell, kpts):\n",
    "    return cell.pbc_intor('int1e_kin', kpts=kpts)\n",
    "get_kinetic = get_kinetic_ref\n",
    "\n",
    "if TEST:\n",
    "    kpts = cell.make_kpts([2,1,1])\n",
    "    print(abs(get_kinetic1(cell, kpts) -\n",
    "              get_kinetic_ref(cell, kpts)).max())\n",
    "    print(abs(get_kinetic2(cell, kpts) -\n",
    "              get_kinetic_ref(cell, kpts)).max())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Nuclear attraction integrals with k point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_nuc(cell, kpts):\n",
    "    R = cell.atom_coords()\n",
    "    Gv = cell.get_Gv()\n",
    "    SI = np.exp(-1j*np.einsum('zs,xs->zx', R, Gv))\n",
    "    rho_nuc_G = -np.einsum('z,zx->x', cell.atom_charges(), SI)\n",
    "\n",
    "    nao = cell.nao_nr()\n",
    "    w = 1./cell.vol\n",
    "    G2 = np.einsum('xs,xs->x', Gv, Gv)\n",
    "    coulG = 4 * np.pi / G2\n",
    "    coulG[G2 == 0] = 0\n",
    "\n",
    "    nuc = []\n",
    "    for k, kpt in enumerate(kpts):\n",
    "        aoR = get_ao_values(cell, kpt)\n",
    "        pair_ao_real_space = np.einsum('xi,xj->xij', aoR.conj(), aoR)\n",
    "        pair_ao_G_space = ifft(pair_ao_real_space.reshape(-1,nao**2).T,\n",
    "                               cell.mesh).T.reshape(-1,nao,nao) * cell.vol\n",
    "\n",
    "        nuc.append(___________)\n",
    "    return np.array(nuc)\n",
    "\n",
    "def get_nuc_ref(cell, kpts):\n",
    "    from pyscf.pbc import df\n",
    "    return np.array(df.FFTDF(cell, kpts).get_nuc(kpts))\n",
    "\n",
    "if TEST:\n",
    "    kpts = cell.make_kpts([2,1,1])\n",
    "    print(abs(get_nuc_ref(cell, kpts)) - get_nuc(cell, kpts).max())\n",
    "    \n",
    "get_nuc = get_nuc_ref"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_hcore(cell, kpts):\n",
    "    t = get_kinetic(cell, kpts)\n",
    "    v = get_nuc(cell, kpts)\n",
    "    return t + v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coulomb integrals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_j(cell, dm, kpts1, kpts2):\n",
    "    ngrids = np.prod(cell.mesh)\n",
    "    rho = np.zeros(ngrids, dtype=np.complex)\n",
    "    for k, kpt in enumerate(kpts1):\n",
    "        aoR = get_ao_values(cell, kpt)\n",
    "        pair_ao_real_space = np.einsum('xi,xj->xij', aoR.conj(), aoR)\n",
    "        rho += np.einsum('xij,ji->x', pair_ao_real_space, dm[k])\n",
    "    rho *= 1./len(kpts1)\n",
    "\n",
    "    rhoG = fft(rho, cell.mesh) * cell.vol / np.prod(cell.mesh)\n",
    "\n",
    "    w = 1./cell.vol\n",
    "    nao = cell.nao_nr()\n",
    "    Gv = cell.get_Gv()\n",
    "    G2 = np.einsum('xs,xs->x', Gv, Gv)\n",
    "    coulG = 4 * np.pi / G2\n",
    "    coulG[G2 == 0] = 0\n",
    "\n",
    "    vj_kpts = []\n",
    "    for k, kpt in enumerate(kpts2):\n",
    "        vj = ________________\n",
    "        vj_kpts.append(vj)\n",
    "    vj_kpts = np.array(vj_kpts)\n",
    "    return vj_kpts\n",
    "\n",
    "def get_j_ref(cell, dm, kpts1, kpts2):\n",
    "    from pyscf.pbc import scf\n",
    "    return scf.KRHF(cell, kpts1).get_j(cell, dm, kpts=kpts1, kpts_band=kpts2)\n",
    "\n",
    "if TEST:\n",
    "    kpts = cell.make_kpts([2,1,1])\n",
    "    kpts2 = cell.make_kpts([1,3,1])\n",
    "    dm = cell.pbc_intor('int1e_ovlp', kpts=kpts)\n",
    "    print(abs(get_j(cell, dm, kpts, kpts2)\n",
    "              - get_j_ref(cell, dm, kpts, kpts2)).max()).max()\n",
    "    \n",
    "get_j = get_j_ref\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Integrals for XC functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_vxc(cell, dm, xc, kpts1, kpts2):\n",
    "    from pyscf.dft import libxc\n",
    "    ngrids = np.prod(cell.mesh)\n",
    "    rho = np.zeros(ngrids, dtype=np.complex)\n",
    "    for k, kpt in enumerate(kpts1):\n",
    "        aoR = get_ao_values(cell, kpt)\n",
    "        pair_ao = np.einsum('xi,xj->xij', aoR.conj(), aoR)\n",
    "        rho += np.einsum('xij,ji->x', pair_ao, dm[k])\n",
    "    rho *= 1./len(kpts1)\n",
    "    rho = rho.real\n",
    "\n",
    "    e_xc, v_xc = libxc.eval_xc(xc, rho, deriv=1)[:2]\n",
    "    v_xc = v_xc[0]\n",
    "    w = get_weight(cell)\n",
    "    E_xc = w * np.einsum('x,x', rho, e_xc)\n",
    "    vxc = []\n",
    "    for k, kpt in enumerate(kpts2):\n",
    "        vxc.append(___________)\n",
    "    return np.array(vxc), E_xc\n",
    "\n",
    "def get_vxc_ref(cell, dm, xc, kpts1, kpts2):\n",
    "    from pyscf.pbc import dft\n",
    "    mf = dft.KRKS(cell, kpts1)\n",
    "    E_xc, vxc = mf._numint.nr_rks(cell, mf.grids, 'lda,vwn', dm, kpts=kpts1,\n",
    "                                  kpts_band=kpts2)[1:]\n",
    "    return vxc, E_xc\n",
    "\n",
    "if TEST:\n",
    "    kpts2 = cell.make_kpts([1,3,1])\n",
    "    dm = cell.pbc_intor('int1e_ovlp', kpts=kpts)\n",
    "    vxc, E_xc = get_vxc(cell, dm, 'lda,vwn', kpts, kpts2)\n",
    "    vxc_ref, E_xc_ref = get_vxc_ref(cell, dm, 'lda,vwn', kpts, kpts2)\n",
    "    print(abs(vxc-vxc_ref).max())\n",
    "    print(abs(E_xc-E_xc_ref).max())\n",
    "\n",
    "get_vxc = get_vxc_ref"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Band structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "get_nuc = get_nuc_ref\n",
    "get_j = get_j_ref\n",
    "get_vxc = get_vxc_ref\n",
    "def get_bands(cell, mo_kpts, kpts, kpts_band):\n",
    "    nelectron = cell.nelectron\n",
    "    nocc = nelectron // 2\n",
    "    dm = [np.einsum('pi,qi->pq', c[:,:nocc], c[:,:nocc].conj()) * 2\n",
    "          for c in mo_kpts]\n",
    "\n",
    "    hcore = (np.array(get_kinetic(cell, kpts_band)) +\n",
    "             np.array(get_nuc(cell, kpts_band)))\n",
    "    vj = get_j(cell, dm, kpts, kpts_band)\n",
    "    vxc = get_vxc(cell, dm, 'lda,vwn', kpts, kpts_band)[0]\n",
    "    fock = hcore + vj + vxc\n",
    "    s = get_ovlp(cell, kpts_band)\n",
    "\n",
    "    bands = []\n",
    "    for k, kpt in enumerate(kpts_band):\n",
    "        e, c = scipy.linalg.eigh(fock[k], s[k])\n",
    "        bands.append(e)\n",
    "    return np.array(bands)\n",
    "\n",
    "def special_points(cell):\n",
    "    points = (('G', [0, 0, 0]),\n",
    "              ('A', [.5, .5, 0]),\n",
    "              ('C', [0, .5, .5]),\n",
    "              ('D', [.5, 0, .5]),\n",
    "              ('E', [.5, .5, .5]),\n",
    "              ('X', [0, .5, 0]),\n",
    "              ('Y', [0, 0, .5]),\n",
    "              ('Z', [.5, 0, 0]))\n",
    "    kpts_band = np.einsum('is,sx->ix',\n",
    "                          [x[1] for x in points],\n",
    "                          cell.reciprocal_vectors())\n",
    "    return points, kpts_band\n",
    "\n",
    "if TEST:\n",
    "    gamma = np.zeros((1,3))\n",
    "    from pyscf.pbc import dft\n",
    "    mf = dft.KRKS(cell, gamma).run()\n",
    "\n",
    "    points, kpts_band = special_points(cell)\n",
    "    bands = get_bands(cell, mf.mo_coeff, gamma, kpts_band)\n",
    "    \n",
    "    import matplotlib.pyplot as plt\n",
    "    fig = plt.figure()\n",
    "    af1 = fig.add_subplot(111)\n",
    "    nbands = cell.nao_nr()\n",
    "    kpath = np.arange(len(points))\n",
    "    for i in range(nbands):\n",
    "        af1.plot(kpath, bands[:,i], '-', color='k')\n",
    "    af1.set_xlim(0, 7)\n",
    "    af1.set_xticks(kpath)\n",
    "    af1.set_xticklabels([x[0] for x in points])\n",
    "    plt.show()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import scipy.linalg\n",
    "def run_hf(cell, kpts):\n",
    "    nelectron = cell.nelectron\n",
    "    nocc = nelectron // 2\n",
    "    nkpts = len(kpts)\n",
    "\n",
    "    hcore = get_hcore(cell, kpts)\n",
    "    s = get_ovlp(cell, kpts)\n",
    "    dm = np.zeros_like(s)\n",
    "    vxc, Exc = get_vxc(cell, dm, 'lda,vwn', kpts, kpts)\n",
    "    vj = get_j(cell, dm, kpts, kpts)\n",
    "\n",
    "    E_ewald = cell.ewald()\n",
    "    E = E_ewald\n",
    "    dE = 1e99\n",
    "\n",
    "    cycle = 0\n",
    "    while dE > 1e-4:\n",
    "        cycle += 1\n",
    "        Elast = E\n",
    "\n",
    "        fock = hcore + vj + vxc\n",
    "        e_k = []\n",
    "        c_k = []\n",
    "        for k in range(nkpts):\n",
    "            e, c = scipy.linalg.eigh(fock[k], s[k])\n",
    "            e_k.append(e)\n",
    "            c_k.append(c)\n",
    "        e_k = np.array(e_k)\n",
    "        c_k = np.array(c_k)\n",
    "        fermi = np.sort(e_k.ravel())[nkpts*nocc-1]\n",
    "        dm = __________\n",
    "        \n",
    "        vxc, Exc = get_vxc(cell, dm, 'lda,vwn', kpts, kpts)\n",
    "        vj = get_j(cell, dm, kpts, kpts)\n",
    "        \n",
    "        E_coul = .5 * np.einsum('kij,kji', vj, dm) / nkpts\n",
    "        E_elec = np.einsum('kij,kji', hcore, dm) / nkpts + E_coul + Exc\n",
    "        E = E_ewald + E_elec\n",
    "        E = E.real\n",
    "        dE = abs(E - Elast)\n",
    "        print('SCF cycle', cycle, 'E(HF) =', E, 'dE =', dE)\n",
    "    return E, c_k\n",
    "\n",
    "if TEST:\n",
    "    kpts = cell.make_kpts([2,2,2])\n",
    "    E, c_k = run_hf(cell, kpts)\n",
    "\n",
    "    from pyscf.pbc import dft\n",
    "    mf = dft.KRKS(cell, kpts)\n",
    "    mf.xc = 'lda,vwn'\n",
    "    mf.kernel()\n",
    "    print(abs(E - mf.e_tot))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ERIs with k points\n",
    "\\begin{align}\n",
    "(ij|kl)\n",
    "&=\\sum_{\\mathbf{G}\\neq 0} w_G \\frac{4\\pi}{|\\mathbf{G} + \\mathbf{k}_j - \\mathbf{k}_i|^2} \\rho_{ij}(\\mathbf{G} + \\mathbf{k}_j - \\mathbf{k}_i) \\rho_{kl}(-\\mathbf{G} - \\mathbf{k}_j + \\mathbf{k}_i)\n",
    "\\end{align}\n",
    "\n",
    "\\begin{equation}\n",
    "\\rho_{ij} (\\mathbf{G} + \\mathbf{k}) = \\int e^{-i (\\mathbf{G}+\\mathbf{k})\\cdot\\mathbf{r}} \\phi^{\\mathbf{k}_i}_i(\\mathbf{r})^* \\phi^{\\mathbf{k}_j}_j(\\mathbf{r}) d\\mathbf{r}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def ao2mo(cell, mos, kpts):\n",
    "    # momont conservation symmetry\n",
    "    dk = kpts[1] - kpts[0] + kpts[3] - kpts[2]\n",
    "    if abs(dk).max() > 1e-8:\n",
    "        raise RuntimeError('moment not conserved')\n",
    "\n",
    "    mo_i, mo_j, mo_k, mo_l = mos\n",
    "\n",
    "    mo_i = np.einsum('xp,pi->xi', get_ao_values(cell, kpts[0]), mo_i)\n",
    "    mo_j = np.einsum('xp,pi->xi', get_ao_values(cell, kpts[1]), mo_j)\n",
    "    mo_k = np.einsum('xp,pi->xi', get_ao_values(cell, kpts[2]), mo_k)\n",
    "    mo_l = np.einsum('xp,pi->xi', get_ao_values(cell, kpts[3]), mo_l)\n",
    "\n",
    "    ij_real_space = np.einsum('xi,xj->xij', mo_i.conj(), mo_j)\n",
    "    ij_G_space = fft(___________)\n",
    "    kl_G_space = ifft(____________)\n",
    "    \n",
    "    Gv = cell.get_Gv()\n",
    "    kGv = Gv + kpts[1] - kpts[0]\n",
    "    w = 1./cell.vol\n",
    "    G2 = np.einsum('xs,xs->x', kGv, kGv)\n",
    "    coulG = 4 * np.pi / G2\n",
    "    coulG[G2 == 0] = 0\n",
    "\n",
    "    eri_mo = w * np.einsum('x,xij,xkl->ijkl', coulG, ij_G_space, kl_G_space)\n",
    "    return eri_mo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using DF module to access ERIs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from pyscf.pbc import df\n",
    "kpts = cell.make_kpts([2,2,2])\n",
    "kpt_ijkl = [kpts[0], kpts[1], kpts[1], kpts[0]]\n",
    "\n",
    "eri_ao = df.FFTDF(cell, kpts).get_eri(kpt_ijkl)\n",
    "eri_mo = df.FFTDF(cell, kpts).get_eri(mos, kpt_ijkl)\n",
    "\n",
    "eri_ao = df.DF(cell, kpts).get_eri(kpt_ijkl)\n",
    "eri_mo = df.DF(cell, kpts).get_eri(mos, kpt_ijkl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test system for k-point to supercell gamma point transformation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if 0:\n",
    "    from pyscf.pbc import gto, scf, dft, tools\n",
    "    cell = gto.Cell()\n",
    "    cell.atom = '''\n",
    "    H 0.  0.  0.\n",
    "    H 0.5 0.3 0.4\n",
    "    '''\n",
    "\n",
    "    cell.basis = [[0, (1, 1)], [1, (.5, 1)]]\n",
    "    cell.a = np.eye(3) * 3.\n",
    "    cell.unit = 'B'\n",
    "    cell.build()\n",
    "\n",
    "    kmesh = [2, 1, 1]\n",
    "    kpts = cell.make_kpts(kmesh)\n",
    "    \n",
    "    kmf = scf.KRHF(cell, kpts).run()\n",
    "    mo_coeff_gamma = k2gamma(cell, kpts, kmf.mo_coeff)\n",
    "    \n",
    "    scell = tools.super_cell(cell, kmesh)\n",
    "    smf = scf.RHF(scell).run()\n",
    "    \n",
    "    #TODO: comparing smf.mo_coeff and mo_coeff_gamma\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda env:anaconda2]",
   "language": "python",
   "name": "conda-env-anaconda2-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  },
  "nbpresent": {
   "slides": {
    "226e5f85-1d19-495c-b78e-b18b4c27388f": {
     "id": "226e5f85-1d19-495c-b78e-b18b4c27388f",
     "prev": null,
     "regions": {}
    },
    "5c4b36dc-7fde-4edf-9057-86716d13e016": {
     "id": "5c4b36dc-7fde-4edf-9057-86716d13e016",
     "prev": "ee6a0647-6fce-4e7b-8e8a-662ca068a6de",
     "regions": {}
    },
    "ec873e20-3126-46a1-9b8c-00dd0ae84876": {
     "id": "ec873e20-3126-46a1-9b8c-00dd0ae84876",
     "prev": "5c4b36dc-7fde-4edf-9057-86716d13e016",
     "regions": {}
    },
    "ee6a0647-6fce-4e7b-8e8a-662ca068a6de": {
     "id": "ee6a0647-6fce-4e7b-8e8a-662ca068a6de",
     "prev": "226e5f85-1d19-495c-b78e-b18b4c27388f",
     "regions": {}
    }
   },
   "themes": {
    "default": "292796c3-c948-4e20-8f9a-9fbd60405712",
    "theme": {}
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}