{
  "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# Reduced Bloch mode expansion\n\nCalculation of the band diagram of a two-dimensional photonic crystal.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\n\nimport protis as pt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Reference results are taken from :cite:p:`Hussein2009`.\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": [
        "epsilon = lattice.ones() * 1\nhole = lattice.square((0.5, 0.5), 0.33)\nepsilon[hole] = 11.4\nmu = 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We define here the wavevector path:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def k_space_path(Nb):\n    K = np.linspace(0, np.pi / a, Nb)\n    bands = np.zeros((3 * Nb - 3, 2))\n    bands[:Nb, 0] = K\n    bands[Nb : 2 * Nb, 0] = K[-1]\n    bands[Nb : 2 * Nb - 1, 1] = K[1:]\n    bands[2 * Nb - 1 : 3 * Nb - 3, 0] = bands[2 * Nb - 1 : 3 * Nb - 3, 1] = np.flipud(\n        K\n    )[1:-1]\n    return bands, K"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Full model:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def full_model(bands, nh=100):\n    t0 = pt.tic()\n    sim = pt.Simulation(lattice, epsilon=epsilon, mu=mu, nh=nh)\n    BD = {}\n    for polarization in [\"TE\", \"TM\"]:\n        ev_band = []\n        for kx, ky in bands:\n            sim.k = kx, ky\n            sim.solve(polarization, vectors=False)\n            ev_norma = sim.eigenvalues * a / (2 * np.pi)\n            ev_band.append(ev_norma)\n        # append first value since this is the same point\n        ev_band.append(ev_band[0])\n        BD[polarization] = ev_band\n    BD[\"TM\"] = pt.backend.stack(BD[\"TM\"]).real\n    BD[\"TE\"] = pt.backend.stack(BD[\"TE\"]).real\n    t_full = pt.toc(t0, verbose=False)\n    return BD, t_full, sim"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Reduced Bloch mode expansion\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def rbme_model(bands, nh=100, Nmodel=2, N_RBME=8):\n    t0 = pt.tic()\n    sim = pt.Simulation(lattice, epsilon=epsilon, mu=mu, nh=nh)\n    q = pt.pi / a\n    if Nmodel == 2:\n        bands_RBME = [(0, 0), (q, 0), (q, q)]\n    elif Nmodel == 3:\n        bands_RBME = [(0, 0), (q / 2, 0), (q, 0), (q, q / 2), (q, q), (q / 2, q / 2)]\n    else:\n        raise ValueError\n    rbme = {\n        polarization: sim.get_rbme_matrix(N_RBME, bands_RBME, polarization)\n        for polarization in [\"TE\", \"TM\"]\n    }\n    BD_RBME = {}\n    for polarization in [\"TE\", \"TM\"]:\n        ev_band = []\n        for kx, ky in bands:\n            sim.k = kx, ky\n            sim.solve(polarization, vectors=False, rbme=rbme[polarization])\n            ev_norma = sim.eigenvalues * a / (2 * np.pi)\n            ev_band.append(ev_norma)\n        # append first value since this is the same point\n        ev_band.append(ev_band[0])\n        BD_RBME[polarization] = ev_band\n    BD_RBME[\"TM\"] = pt.backend.stack(BD_RBME[\"TM\"]).real\n    BD_RBME[\"TE\"] = pt.backend.stack(BD_RBME[\"TE\"]).real\n    t_rbme = pt.toc(t0, verbose=False)\n    return BD_RBME, t_rbme, sim\n\n\nbands, K = k_space_path(Nb=49)\nBD, t_full, sim_full = full_model(bands, nh=100)\nBD_RBME, t_rbme, sim_rbme = rbme_model(bands, nh=100, Nmodel=2, N_RBME=8)\nprint(f\"speedup = {t_full/t_rbme}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plot the bands:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def k_space_path_plot(Nb, K):\n    bands_plot = np.zeros(3 * Nb - 2)\n    bands_plot[:Nb] = K\n    bands_plot[Nb : 2 * Nb - 1] = K[-1] + K[1:]\n    bands_plot[2 * Nb - 1 : 3 * Nb - 2] = 2 * K[-1] + 2**0.5 * K[1:]\n    return bands_plot\n\n\nbands_plot = k_space_path_plot(49, K)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "TE polarization:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(3.2, 2.5))\nplotTE = plt.plot(bands_plot, BD[\"TE\"], c=\"#cf5268\", lw=1.5, alpha=0.5)\nplotTE_RBME = plt.plot(bands_plot, BD_RBME[\"TE\"], \"--\", c=\"#cf5268\")\nplt.ylim(0, 1.2)\nplt.xlim(0, bands_plot[-1])\nplt.xticks(\n    [0, K[-1], 2 * K[-1], bands_plot[-1]], [\"$\\Gamma$\", \"$X$\", \"$M$\", \"$\\Gamma$\"]\n)\nplt.axvline(K[-1], c=\"k\", lw=0.3)\nplt.axvline(2 * K[-1], c=\"k\", lw=0.3)\nplt.ylabel(r\"Frequency $\\omega a/2\\pi c$\")\nplt.legend([plotTE[0], plotTE_RBME[0]], [\"full\", \"2-point RBME\"], loc=(0.31, 0.02))\nplt.title(\"TE modes\", c=\"#cf5268\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "TM polarization:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(3.2, 2.5))\nplotTM = plt.plot(bands_plot, BD[\"TM\"], c=\"#4199b0\", lw=1.5, alpha=0.5)\nplotTM_RBME = plt.plot(bands_plot, BD_RBME[\"TM\"], \"--\", c=\"#4199b0\")\nplt.ylim(0, 1.2)\nplt.xlim(0, bands_plot[-1])\nplt.xticks(\n    [0, K[-1], 2 * K[-1], bands_plot[-1]], [\"$\\Gamma$\", \"$X$\", \"$M$\", \"$\\Gamma$\"]\n)\nplt.axvline(K[-1], c=\"k\", lw=0.3)\nplt.axvline(2 * K[-1], c=\"k\", lw=0.3)\nplt.ylabel(r\"Frequency $\\omega a/2\\pi c$\")\nplt.legend([plotTM[0], plotTM_RBME[0]], [\"full\", \"2-point RBME\"], loc=(0.31, 0.02))\nplt.title(\"TM modes\", c=\"#4199b0\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Performances with number of harmonics\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "bands, K = k_space_path(49)\nNH = np.arange(100, 350, 50)\nactual_nh = []\ns2 = []\ns3 = []\nfor nh in NH:\n    BD, t_full, sim_full = full_model(bands, nh=nh)\n    BD_RBME2, t_rbme2, sim_rbme2 = rbme_model(bands, nh=nh, Nmodel=2, N_RBME=8)\n    BD_RBME3, t_rbme3, sim_rbme3 = rbme_model(bands, nh=nh, Nmodel=3, N_RBME=8)\n    actual_nh.append(sim_full.nh)\n    s2.append(t_rbme2 / t_full)\n    s3.append(t_rbme3 / t_full)\n\n\nplt.figure()\nplt.plot(actual_nh, s2, \"o-\", c=\"#e98b34\", label=\"2-point RBME\")\nplt.plot(actual_nh, s3, \"s-\", c=\"#56c291\", label=\"3-point RBME\")\nplt.xlabel(r\"number of harmonics $n_h$\")\nplt.ylabel(r\"$s$\")\nplt.legend()\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Performances with number of k-space points\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = {}\nNH = [100, 200]\nnum_ks = np.arange(5, 100, 15)\nfor nh in NH:\n    actual_nh = []\n    nk = []\n    s2 = []\n    s3 = []\n    for Nb in num_ks:\n        bands, K = k_space_path(Nb)\n        nk.append(len(bands))\n        BD, t_full, sim_full = full_model(bands, nh=nh)\n        BD_RBME2, t_rbme2, sim_rbme2 = rbme_model(bands, nh=nh, Nmodel=2, N_RBME=8)\n        BD_RBME3, t_rbme3, sim_rbme3 = rbme_model(bands, nh=nh, Nmodel=3, N_RBME=8)\n        actual_nh.append(sim_full.nh)\n        s2.append(t_rbme2 / t_full)\n        s3.append(t_rbme3 / t_full)\n    result[nh] = dict(s2=s2, s3=s3, actual_nh=actual_nh)\n\n\nplt.figure()\nplt.plot(\n    nk,\n    result[NH[0]][\"s2\"],\n    \"o--\",\n    c=\"#e98b34\",\n    label=rf\"2-point RBME, $nh={result[NH[0]]['actual_nh'][0]}$\",\n)\nplt.plot(\n    nk,\n    result[NH[0]][\"s3\"],\n    \"s--\",\n    c=\"#56c291\",\n    label=rf\"3-point RBME, $nh={result[NH[0]]['actual_nh'][0]}$\",\n)\nplt.plot(\n    nk,\n    result[NH[1]][\"s2\"],\n    \"o-\",\n    c=\"#e98b34\",\n    label=rf\"2-point RBME, $nh={result[NH[1]]['actual_nh'][0]}$\",\n)\nplt.plot(\n    nk,\n    result[NH[1]][\"s3\"],\n    \"s-\",\n    c=\"#56c291\",\n    label=rf\"3-point RBME, $nh={result[NH[1]]['actual_nh'][0]}$\",\n)\nplt.xlabel(r\"number of k points $n_k$\")\nplt.ylabel(r\"$s$\")\nplt.legend()\nplt.yscale(\"log\")\nplt.tight_layout()"
      ]
    },
    {
      "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
}