{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "from __future__ import division, print_function, absolute_import\n",
    "\n",
    "from numpy import cos, inf, zeros, array, exp, conj, nan, isnan, pi, sin\n",
    "\n",
    "cimport numpy as np\n",
    "cimport scipy as sp\n",
    "\n",
    "import sys\n",
    "EPSILON = sys.float_info.epsilon # typical floating-point calculation error\n",
    "\n",
    "def make_2x2_array(float complex a, float complex b, float complex c, float complex d):\n",
    "    \"\"\"\n",
    "    Makes a 2x2 numpy array of [[a,b],[c,d]]\n",
    "\n",
    "    Same as \"numpy.array([[a,b],[c,d]], dtype=float)\", but ten times faster\n",
    "    \"\"\"\n",
    "    my_array = np.empty((2,2), dtype=complex)\n",
    "    my_array[0,0] = a\n",
    "    my_array[0,1] = b\n",
    "    my_array[1,0] = c\n",
    "    my_array[1,1] = d\n",
    "    return my_array\n",
    "\n",
    "def is_forward_angle(n, theta):\n",
    "    \"\"\"\n",
    "    if a wave is traveling at angle theta from normal in a medium with index n,\n",
    "    calculate whether or not this is the forward-traveling wave (i.e., the one\n",
    "    going from front to back of the stack, like the incoming or outgoing waves,\n",
    "    but unlike the reflected wave). For real n & theta, the criterion is simply\n",
    "    -pi/2 < theta < pi/2, but for complex n & theta, it's more complicated.\n",
    "    See https://arxiv.org/abs/1603.02720 appendix D. If theta is the forward\n",
    "    angle, then (pi-theta) is the backward angle and vice-versa.\n",
    "    \"\"\"\n",
    "    assert n.real * n.imag >= 0, (\"For materials with gain, it's ambiguous which \"\n",
    "                                  \"beam is incoming vs outgoing. See \"\n",
    "                                  \"https://arxiv.org/abs/1603.02720 Appendix C.\\n\"\n",
    "                                  \"n: \" + str(n) + \"   angle: \" + str(theta))\n",
    "    ncostheta = n * np.cos(theta)\n",
    "    if abs(ncostheta.imag) > 100 * EPSILON:\n",
    "        # Either evanescent decay or lossy medium. Either way, the one that\n",
    "        # decays is the forward-moving wave\n",
    "        answer = (ncostheta.imag > 0)\n",
    "    else:\n",
    "        # Forward is the one with positive Poynting vector\n",
    "        # Poynting vector is Re[n cos(theta)] for s-polarization or\n",
    "        # Re[n cos(theta*)] for p-polarization, but it turns out they're consistent\n",
    "        # so I'll just assume s then check both below\n",
    "        answer = (ncostheta.real > 0)\n",
    "    # convert from numpy boolean to the normal Python boolean\n",
    "    answer = bool(answer)\n",
    "    # double-check the answer ... can't be too careful!\n",
    "#     error_string = (\"It's not clear which beam is incoming vs outgoing. Weird\"\n",
    "#                     \" index maybe?\\n\"\n",
    "#                     \"n: \" + str(n) + \"   angle: \" + str(theta))\n",
    "#     if answer is True:\n",
    "#         assert ncostheta.imag > -100 * EPSILON, error_string\n",
    "#         assert ncostheta.real > -100 * EPSILON, error_string\n",
    "#         assert (n * cos(theta.conjugate())).real > -100 * EPSILON, error_string\n",
    "#     else:\n",
    "#         assert ncostheta.imag < 100 * EPSILON, error_string\n",
    "#         assert ncostheta.real < 100 * EPSILON, error_string\n",
    "#         assert (n * cos(theta.conjugate())).real < 100 * EPSILON, error_string\n",
    "    return answer\n",
    "\n",
    "def snell(n_1, n_2, th_1):\n",
    "    \"\"\"\n",
    "    return angle theta in layer 2 with refractive index n_2, assuming\n",
    "    it has angle th_1 in layer with refractive index n_1. Use Snell's law. Note\n",
    "    that \"angles\" may be complex!!\n",
    "    \"\"\"\n",
    "    # Important that the arcsin here is scipy.arcsin, not numpy.arcsin! (They\n",
    "    # give different results e.g. for arcsin(2).)\n",
    "    th_2_guess = sp.arcsin(n_1*np.sin(th_1) / n_2)\n",
    "    if is_forward_angle(n_2, th_2_guess):\n",
    "        return th_2_guess\n",
    "    else:\n",
    "        return pi - th_2_guess\n",
    "\n",
    "def list_snell(n_list, th_0):\n",
    "    \"\"\"\n",
    "    return list of angle theta in each layer based on angle th_0 in layer 0,\n",
    "    using Snell's law. n_list is index of refraction of each layer. Note that\n",
    "    \"angles\" may be complex!!\n",
    "    \"\"\"\n",
    "    # Important that the arcsin here is scipy.arcsin, not numpy.arcsin! (They\n",
    "    # give different results e.g. for arcsin(2).)\n",
    "    angles = sp.arcsin(n_list[0]*np.sin(th_0) / n_list)\n",
    "    # The first and last entry need to be the forward angle (the intermediate\n",
    "    # layers don't matter, see https://arxiv.org/abs/1603.02720 Section 5)\n",
    "    if not is_forward_angle(n_list[0], angles[0]):\n",
    "        angles[0] = pi - angles[0]\n",
    "    if not is_forward_angle(n_list[-1], angles[-1]):\n",
    "        angles[-1] = pi - angles[-1]\n",
    "    return angles\n",
    "\n",
    "\n",
    "def interface_r(polarization, n_i, n_f, th_i, th_f):\n",
    "    \"\"\"\n",
    "    reflection amplitude (from Fresnel equations)\n",
    "\n",
    "    polarization is either \"s\" or \"p\" for polarization\n",
    "\n",
    "    n_i, n_f are (complex) refractive index for incident and final\n",
    "\n",
    "    th_i, th_f are (complex) propegation angle for incident and final\n",
    "    (in radians, where 0=normal). \"th\" stands for \"theta\".\n",
    "    \"\"\"\n",
    "    if polarization == 's':\n",
    "        return ((n_i * np.cos(th_i) - n_f * np.cos(th_f)) /\n",
    "                (n_i * np.cos(th_i) + n_f * np.cos(th_f)))\n",
    "    else:\n",
    "        return ((n_f * np.cos(th_i) - n_i * np.cos(th_f)) /\n",
    "                (n_f * np.cos(th_i) + n_i * np.cos(th_f)))\n",
    "\n",
    "def interface_t(polarization, n_i, n_f, th_i, th_f):\n",
    "    \"\"\"\n",
    "    transmission amplitude (frem Fresnel equations)\n",
    "\n",
    "    polarization is either \"s\" or \"p\" for polarization\n",
    "\n",
    "    n_i, n_f are (complex) refractive index for incident and final\n",
    "\n",
    "    th_i, th_f are (complex) propegation angle for incident and final\n",
    "    (in radians, where 0=normal). \"th\" stands for \"theta\".\n",
    "    \"\"\"\n",
    "    if polarization == 's':\n",
    "        return 2 * n_i * np.cos(th_i) / (n_i * np.cos(th_i) + n_f * np.cos(th_f))\n",
    "    else:\n",
    "        return 2 * n_i * np.cos(th_i) / (n_f * np.cos(th_i) + n_i * np.cos(th_f))\n",
    "\n",
    "\n",
    "\n",
    "def coh_tmm(pol, n_list, d_list, float th_0, float lam_vac):\n",
    "    \"\"\"\n",
    "    Main \"coherent transfer matrix method\" calc. Given parameters of a stack,\n",
    "    calculates everything you could ever want to know about how light\n",
    "    propagates in it. (If performance is an issue, you can delete some of the\n",
    "    calculations without affecting the rest.)\n",
    "\n",
    "    pol is light polarization, \"s\" or \"p\".\n",
    "\n",
    "    n_list is the list of refractive indices, in the order that the light would\n",
    "    pass through them. The 0'th element of the list should be the semi-infinite\n",
    "    medium from which the light enters, the last element should be the semi-\n",
    "    infinite medium to which the light exits (if any exits).\n",
    "\n",
    "    th_0 is the angle of incidence: 0 for normal, pi/2 for glancing.\n",
    "    Remember, for a dissipative incoming medium (n_list[0] is not real), th_0\n",
    "    should be complex so that n0 sin(th0) is real (intensity is constant as\n",
    "    a function of lateral position).\n",
    "\n",
    "    d_list is the list of layer thicknesses (front to back). Should correspond\n",
    "    one-to-one with elements of n_list. First and last elements should be \"inf\".\n",
    "\n",
    "    lam_vac is vacuum wavelength of the light.\n",
    "\n",
    "    * R--reflected wave power (as fraction of incident)\n",
    "    \"\"\"\n",
    "    # Convert lists to numpy arrays if they're not already.\n",
    "    n_list = array(n_list)\n",
    "    d_list = array(d_list, dtype=float)\n",
    "\n",
    "    cdef int num_layers\n",
    "    num_layers = n_list.size\n",
    "    \n",
    "    th_list = list_snell(n_list, th_0)\n",
    "\n",
    "    # kz is the z-component of (complex) angular wavevector for forward-moving\n",
    "    # wave. Positive imaginary part means decaying.\n",
    "    kz_list = 2 * np.pi * n_list * np.cos(th_list) / lam_vac\n",
    "\n",
    "    # delta is the total phase accrued by traveling through a given layer.\n",
    "    # Ignore warning about inf multiplication\n",
    "    delta = kz_list * d_list\n",
    "\n",
    "\n",
    "    # t_list[i,j] and r_list[i,j] are transmission and reflection amplitudes,\n",
    "    # respectively, coming from i, going to j. Only need to calculate this when\n",
    "    # j=i+1. (2D array is overkill but helps avoid confusion.)\n",
    "    t_list = zeros((num_layers, num_layers), dtype=complex)\n",
    "    r_list = zeros((num_layers, num_layers), dtype=complex)\n",
    "    for i in range(num_layers-1):\n",
    "        t_list[i,i+1] = interface_t(pol, n_list[i], n_list[i+1],\n",
    "                                    th_list[i], th_list[i+1])\n",
    "        r_list[i,i+1] = interface_r(pol, n_list[i], n_list[i+1],\n",
    "                                    th_list[i], th_list[i+1])\n",
    "        \n",
    "    # My M is a bit different than Sernelius's, but Mtilde is the same.\n",
    "    M_list = zeros((num_layers, 2, 2), dtype=complex)\n",
    "    for i in range(1, num_layers-1):\n",
    "        M_list[i] = (1/t_list[i,i+1]) * np.dot(\n",
    "            make_2x2_array(np.exp(-1j*delta[i]), 0, 0, np.exp(1j*delta[i]),dtype=complex),\n",
    "            make_2x2_array(1, r_list[i,i+1], r_list[i,i+1], 1))\n",
    "    Mtilde = make_2x2_array(1, 0, 0, 1)\n",
    "    for i in range(1, num_layers-1):\n",
    "        Mtilde = np.dot(Mtilde, M_list[i])\n",
    "    Mtilde = np.dot(make_2x2_array(1, r_list[0,1], r_list[0,1], 1)/t_list[0,1], Mtilde)\n",
    "\n",
    "    # Net complex transmission and reflection amplitudes\n",
    "    r = Mtilde[1,0]/Mtilde[0,0]\n",
    "#     t = 1/Mtilde[0,0]\n",
    "\n",
    "#     # vw_list[n] = [v_n, w_n]. v_0 and w_0 are undefined because the 0th medium\n",
    "#     # has no left interface.\n",
    "#     vw_list = zeros((num_layers, 2), dtype=complex)\n",
    "#     vw = array([[t],[0]])\n",
    "#     vw_list[-1,:] = np.transpose(vw)\n",
    "#     for i in range(num_layers-2, 0, -1):\n",
    "#         vw = np.dot(M_list[i], vw)\n",
    "#         vw_list[i,:] = np.transpose(vw)\n",
    "\n",
    "    return np.abs(r)**2\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.0"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}