{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# This file is part of protis\n# License: GPLv3\n%matplotlib notebook"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Topological invariants\n\nCalculation of Berry phase and Chern number for a magneto-optical photonic crystal.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import time\n\nimport matplotlib.pyplot as plt\n\nimport protis as pt\n\npt.set_backend(\"autograd\")\n# pt.set_backend(\"torch\")\nbk = pt.backend\n\n\npi = bk.pi\nplt.ion()\nplt.close(\"all\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Reference results are taken from  :cite:p:`blancodepaz2020`.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "a = 1\nlattice = pt.Lattice([[a, 0], [0, a]], discretization=2**9)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define the permittivity\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "rod = lattice.circle(center=(0.5, 0.5), radius=0.11)\nepsilon = lattice.ones()\nepsilon[rod] = 15\nmu = lattice.ones()\none = lattice.ones()\nzero = lattice.zeros()\nmu = pt.block_z_anisotropic(one, zero, zero, one, one)\nmu[0, 0, rod] = mu[1, 1, rod] = 14\nmu[0, 1, rod] = 12.4j\nmu[1, 0, rod] = -12.4j"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We define here the wavevector path:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Gamma = (0, 0)\nX = (pi / a, 0)\nM = (pi / a, pi / a)\nsym_points = [Gamma, X, M, Gamma]\n\nNb = 31\nkpath = pt.init_bands(sym_points, Nb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Calculate the band diagram:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "polarization = \"TM\"\nnh = 200\n\n\ndef compute_bands(epsilon, mu):\n    sim = pt.Simulation(lattice, epsilon=epsilon, mu=mu, nh=nh)\n    ev_band = []\n    for kx, ky in kpath:\n        sim.k = kx, ky\n        sim.solve(polarization, vectors=False)\n        ev_norma = sim.eigenvalues * a / (2 * pi)\n        ev_band.append(ev_norma.real)\n    return ev_band\n\n\nev_band0 = compute_bands(epsilon, mu=1)\n\n\nev_band = compute_bands(epsilon, mu)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plot the bands (without magnetic field):\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "labels = [\"$\\Gamma$\", \"$X$\", \"$M$\", \"$\\Gamma$\"]\n\n\nplt.figure()\npt.plot_bands(sym_points, Nb, ev_band0, color=\"#8041b0\", xtickslabels=labels)\nplt.ylim(0, 1.2)\nplt.ylabel(r\"Frequency $\\omega a/2\\pi c$\")\nplt.title(\"without magnetic \ufb01eld\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plot the bands (with magnetic field):\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure()\npt.plot_bands(sym_points, Nb, ev_band, color=\"#8041b0\", xtickslabels=labels)\nplt.ylim(0, 1.2)\nplt.ylabel(r\"Frequency $\\omega a/2\\pi c$\")\nplt.title(\"with applied magnetic \ufb01eld\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compute modes in the first Brillouin zone:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "sim = pt.Simulation(lattice, epsilon=epsilon, mu=mu, nh=nh)\n\n\nt = -time.time()\n\nmethod = \"fourier\"\n\nn_eig = 3\nnk = 31\nkx = ky = bk.linspace(0, 2 * pi / a, nk)\nKx, Ky = bk.meshgrid(kx, ky)\n\n\nif method == \"fourier\":\n    eigenmodes = bk.empty((nk, nk, sim.nh, n_eig), dtype=bk.complex128)\nelse:\n    eigenmodes = bk.empty((nk, nk, *lattice.discretization, n_eig), dtype=bk.complex128)\n\n\nfor i in range(nk):\n    _mode = []\n    for j in range(nk):\n        k = kx[i], ky[j]\n        sim.k = k\n        eigs, modes = sim.solve(polarization)\n        if method == \"fourier\":\n            eigenmodes[i, j] = modes[:, :n_eig]\n        else:\n            eigenmodes[i, j] = sim.get_modes(range(n_eig))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compute Berry curvature and Chern number\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "for imode in range(n_eig):\n    mode = (\n        eigenmodes[:, :, :, imode]\n        if method == \"fourier\"\n        else eigenmodes[:, :, :, :, imode]\n    )\n    phi = sim.get_berry_curvature(kx, ky, mode, method=method)\n    C = sim.get_chern_number(kx, ky, phi)\n    print(f\"Mode {imode+1}: Chern number = {C}\")\n\n    plt.figure()\n    plt.pcolormesh(kx, ky, phi)\n    plt.axis(\"scaled\")\n    plt.xticks([0, pi / a, 2 * pi / a], [\"0\", r\"$\\pi/a$\", r\"$2\\pi/a$\"])\n    plt.yticks([0, pi / a, 2 * pi / a], [\"0\", r\"$\\pi/a$\", r\"$2\\pi/a$\"])\n    plt.title(f\"Berry curvature, mode {imode+1}, C={C:.0f}\")\n    plt.colorbar()\n    plt.tight_layout()\n    plt.pause(0.001)\n\n\nt += time.time()\nprint(f\"Elapsed time: {t:.1f}s\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Same but with reduced Bloch mode expansion\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "t = -time.time()\n\nU = (0, 2 * pi / a)\nV = (2 * pi / a, 0)\nW = (2 * pi / a, 2 * pi / a)\nY = (0, pi / a)\nZ = (2 * pi / a, pi / a)\nQ = (pi / a, 2 * pi / a)\nRBME_points = [Gamma, X, M, Y]  # ,U, V, W, Z, Q]\nN_RBME = 8\nR = sim.get_rbme_matrix(N_RBME, RBME_points, polarization)\n\nNred = N_RBME * len(RBME_points)\n\n\neigenmodes_rbme = bk.empty((nk, nk, Nred, n_eig), dtype=bk.complex128)\n\n\nfor i in range(nk):\n    _mode = []\n    for j in range(nk):\n        k = kx[i], ky[j]\n        sim.k = k\n        eigs, modes = sim.solve(polarization, rbme=R, reduced=True)\n        eigenmodes_rbme[i, j] = modes[:, :n_eig]\n\nfor imode in range(n_eig):\n    phi = sim.get_berry_curvature(\n        kx, ky, eigenmodes_rbme[:, :, :, imode], method=\"rbme\"\n    )\n    C = sim.get_chern_number(kx, ky, phi)\n    print(f\"Mode {imode+1}: Chern number = {C}\")\n\n    plt.figure()\n    plt.pcolormesh(kx, ky, phi)\n    plt.axis(\"scaled\")\n    plt.xticks([0, pi / a, 2 * pi / a], [\"0\", r\"$\\pi/a$\", r\"$2\\pi/a$\"])\n    plt.yticks([0, pi / a, 2 * pi / a], [\"0\", r\"$\\pi/a$\", r\"$2\\pi/a$\"])\n    plt.title(f\"Berry curvature, mode {imode+1}, C={C:.0f}\")\n    plt.colorbar()\n    plt.tight_layout()\n    plt.pause(0.001)\n\n\nt += time.time()\nprint(f\"Elapsed time RBME: {t:.1f}s\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import protis.utils.jupyter\n%protis_version_table"
      ]
    }
  ],
  "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.10.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}