{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Partly occupied Wannier Functions\n\nThis tutorial walks through building **partly occupied Wannier\nfunctions** with the :mod:`ase.dft.wannier` module\nand the [GPAW](https://wiki.fysik.dtu.dk/gpaw/) electronic structure code.\nFor more information on the details of the method and the implementation, see\n\n  | K. S. Thygesen, L. B. Hansen, and K. W. Jacobsen\n  | Partly occupied Wannier functions: Construction and applications\n  | [Phys. Rev. B 72, 125119 (2005)](https://doi.org/10.1103/PhysRevB.72.125119)_\n   :depth: 2\n   :local:\n\n\n## Benzene molecule\n\n### Step 1 \u2013 Ground-state calculation\n\nRun the script below to obtain the ground-state density and the\nKohn\u2013Sham (KS) orbitals. The result is stored in :file:`benzene.gpw`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\nfrom gpaw import GPAW, restart\n\nfrom ase import Atoms\nfrom ase.build import molecule\nfrom ase.dft.kpoints import monkhorst_pack\nfrom ase.dft.wannier import Wannier\n\natoms = molecule('C6H6')\natoms.center(vacuum=3.5)\n\ncalc = GPAW(mode='pw', xc='PBE', txt='benzene.txt', nbands=18)\natoms.calc = calc\natoms.get_potential_energy()\n\ncalc = calc.fixed_density(\n    txt='benzene-harris.txt',\n    nbands=40,\n    convergence={'bands': 35},\n)\natoms.get_potential_energy()\n\ncalc.write('benzene.gpw', mode='all')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 2 \u2013 Maximally localized WFs for the occupied subspace (15 WFs)\n\nThere are 15 occupied bands in the benzene molecule. We construct one\nWannier function per occupied band by setting ``nwannier = 15``.\nBy calling ``wan.localize()``, the code attempts to minimize the spread\nfunctional using a gradient-descent algorithm.\nThe resulting WFs are written to .cube files, which allows them\nto be inspected using e.g. VESTA.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "atoms, calc = restart('benzene.gpw', txt=None)\n\n# Make wannier functions of occupied space only\nwan = Wannier(nwannier=15, calc=calc)\nwan.localize()\nfor i in range(wan.nwannier):\n    wan.write_cube(i, f'benzene15_{i}.cube')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 3 \u2013 Adding three extra degrees of freedom (18 WFs)\n\nTo improve localization we augment the basis with three extra\nWannier functions - so-called *extra degrees of freedom*\n(``nwannier = 18``, ``fixedstates = 15``).\nThis will allow the Wannierization procedure to use the unoccupied states to\nminimize spread functional.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "atoms, calc = restart('benzene.gpw', txt=None)\n\n# Make wannier functions using (three) extra degrees of freedom.\nwan = Wannier(nwannier=18, calc=calc, fixedstates=15)\nwan.localize()\nwan.save('wan18.json')\nfor i in range(wan.nwannier):\n    wan.write_cube(i, f'benzene18_{i}.cube')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 4 \u2013 Spectral-weight analysis\n\nThe script below projects the WFs on the KS eigenstates. You should see\nthe 15 lowest bands perfectly reconstructed (weight \u2243 1.0) while higher\nbands are only partially represented.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "atoms, calc = restart('benzene.gpw', txt=None)\nwan = Wannier(nwannier=18, calc=calc, fixedstates=15, file='wan18.json')\n\nweight_n = np.sum(abs(wan.V_knw[0]) ** 2, 1)\nN = len(weight_n)\nF = wan.fixedstates_k[0]\nplt.figure(1, figsize=(12, 4))\nplt.bar(\n    range(1, N + 1),\n    weight_n,\n    width=0.65,\n    bottom=0,\n    color='k',\n    edgecolor='k',\n    linewidth=None,\n    align='center',\n    orientation='vertical',\n)\nplt.plot([F + 0.5, F + 0.5], [0, 1], 'k--')\nplt.axis(xmin=0.32, xmax=N + 1.33, ymin=0, ymax=1)\nplt.xlabel('Eigenstate')\nplt.ylabel('Projection of wannier functions')\nplt.savefig('spectral_weight.png')\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Polyacetylene chain (1-D periodic)\n\nWe now want to construct partially occupied Wannier functions\nto describe a polyacetylene chain.\n\n### Step 1 \u2013 Structure & ground-state calculation\n\nPolyacetylene is modelled as an infinite chain; we therefore enable\nperiodic boundary conditions along *x*.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "kpts = monkhorst_pack((13, 1, 1))\ncalc = GPAW(\n    mode='pw',\n    xc='PBE',\n    kpts=kpts,\n    nbands=12,\n    txt='poly.txt',\n    convergence={'bands': 9},\n    symmetry='off',\n)\n\nCC = 1.38\nCH = 1.094\na = 2.45\nx = a / 2.0\ny = np.sqrt(CC**2 - x**2)\natoms = Atoms(\n    'C2H2',\n    pbc=(True, False, False),\n    cell=(a, 8.0, 6.0),\n    calculator=calc,\n    positions=[[0, 0, 0], [x, y, 0], [x, y + CH, 0], [0, -CH, 0]],\n)\natoms.center()\natoms.get_potential_energy()\ncalc.write('poly.gpw', mode='all')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 2 \u2013 Wannierization\n\nWe repeat the localization procedure, keeping the five lowest\nbands fixed and adding one extra degree of freedom to aid localization.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nfrom gpaw import restart\n\nfrom ase.dft.wannier import Wannier\n\natoms, calc = restart('poly.gpw', txt=None)\n\n# Make wannier functions using (one) extra degree of freedom\nwan = Wannier(\n    nwannier=6,\n    calc=calc,\n    fixedenergy=1.5,\n    initialwannier='orbitals',\n    functional='var',\n)\nwan.localize()\nwan.save('poly.json')\nwan.translate_all_to_cell((2, 0, 0))\nfor i in range(wan.nwannier):\n    wan.write_cube(i, f'polyacetylene_{i}.cube')\n\n# Print Kohn-Sham bandstructure\nef = calc.get_fermi_level()\nwith open('KSbands.txt', 'w') as fd:\n    for k, kpt_c in enumerate(calc.get_ibz_k_points()):\n        for eps in calc.get_eigenvalues(kpt=k):\n            print(kpt_c[0], eps - ef, file=fd)\n\n# Print Wannier bandstructure\nwith open('WANbands.txt', 'w') as fd:\n    for k in np.linspace(-0.5, 0.5, 100):\n        ham = wan.get_hamiltonian_kpoint([k, 0, 0])\n        for eps in np.linalg.eigvalsh(ham).real:\n            print(k, eps - ef, file=fd)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 3 \u2013 High-resolution band structure\n\nUsing the Wannier Hamiltonian we can interpolate the band structure on a\nfine 100-point *k* mesh and compare it to the original DFT result.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig = plt.figure(dpi=80, figsize=(4.2, 6))\nfig.subplots_adjust(left=0.16, right=0.97, top=0.97, bottom=0.05)\n\n# Plot KS bands\nk, eps = np.loadtxt('KSbands.txt', unpack=True)\nplt.plot(k, eps, 'ro', label='DFT', ms=9)\n\n# Plot Wannier bands\nk, eps = np.loadtxt('WANbands.txt', unpack=True)\nplt.plot(k, eps, 'k.', label='Wannier')\n\nplt.plot([-0.5, 0.5], [1, 1], 'k:', label='_nolegend_')\nplt.text(-0.5, 1, 'fixedenergy', ha='left', va='bottom')\nplt.axis('tight')\nplt.xticks(\n    [-0.5, -0.25, 0, 0.25, 0.5],\n    [r'$X$', r'$\\Delta$', r'$\\Gamma$', r'$\\Delta$', r'$X$'],\n    size=16,\n)\nplt.ylabel(r'$E - E_F\\  \\rm{(eV)}$', size=16)\nplt.legend()\nplt.savefig('bands.png', dpi=80)\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Within the fixed-energy window\u2014that is, for energies below the fixed-energy\nline\u2014the Wannier-interpolated bands coincide perfectly with the DFT reference\n(red circles). Above this window the match is lost, because the degrees of\nfreedom deliberately mix several Kohn\u2013Sham states to achieve maximal\nreal-space localisation.\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.13.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}