deepmie/.ipynb_checkpoints/my_tmm-checkpoint.py

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-11-01T07:40:08.911910Z",
     "start_time": "2018-11-01T07:40:08.818797Z"
    },
    "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",
    "import numpy as np\n",
    "import scipy as sp\n",
    "\n",
    "import sys\n",
    "EPSILON = sys.float_info.epsilon # typical floating-point calculation error\n",
    "\n",
    "def make_2x2_array(a, b, c, d, dtype=float):\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=dtype)\n",
    "    my_array[0,0] = a\n",
    "    my_array[0,1] = bggvG\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 * 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 * cos(th_i) - n_f * cos(th_f)) /\n",
    "                (n_i * cos(th_i) + n_f * cos(th_f)))\n",
    "    elif polarization == 'p':\n",
    "        return ((n_f * cos(th_i) - n_i * cos(th_f)) /\n",
    "                (n_f * cos(th_i) + n_i * cos(th_f)))\n",
    "    else:\n",
    "        raise ValueError(\"Polarization must be 's' or 'p'\")\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 * cos(th_i) / (n_i * cos(th_i) + n_f * cos(th_f))\n",
    "    elif polarization == 'p':\n",
    "        return 2 * n_i * cos(th_i) / (n_f * cos(th_i) + n_i * cos(th_f))\n",
    "    else:\n",
    "        raise ValueError(\"Polarization must be 's' or 'p'\")\n",
    "\n",
    "def R_from_r(r):\n",
    "    \"\"\"\n",
    "    Calculate reflected power R, starting with reflection amplitude r.\n",
    "    \"\"\"\n",
    "    return abs(r)**2\n",
    "\n",
    "def T_from_t(pol, t, n_i, n_f, th_i, th_f):\n",
    "    \"\"\"\n",
    "    Calculate transmitted power T, starting with transmission amplitude t.\n",
    "\n",
    "    n_i,n_f are refractive indices of incident and final medium.\n",
    "\n",
    "    th_i, th_f are (complex) propegation angles through incident & final medium\n",
    "    (in radians, where 0=normal). \"th\" stands for \"theta\".\n",
    "\n",
    "    In the case that n_i,n_f,th_i,th_f are real, formulas simplify to\n",
    "    T=|t|^2 * (n_f cos(th_f)) / (n_i cos(th_i)).\n",
    "\n",
    "    See manual for discussion of formulas\n",
    "    \"\"\"\n",
    "    if pol == 's':\n",
    "        return abs(t**2) * (((n_f*cos(th_f)).real) / (n_i*cos(th_i)).real)\n",
    "    elif pol == 'p':\n",
    "        return abs(t**2) * (((n_f*conj(cos(th_f))).real) /\n",
    "                                (n_i*conj(cos(th_i))).real)\n",
    "    else:\n",
    "        raise ValueError(\"Polarization must be 's' or 'p'\")\n",
    "\n",
    "def power_entering_from_r(pol, r, n_i, th_i):\n",
    "    \"\"\"\n",
    "    Calculate the power entering the first interface of the stack, starting with\n",
    "    reflection amplitude r. Normally this equals 1-R, but in the unusual case\n",
    "    that n_i is not real, it can be a bit different than 1-R. See manual.\n",
    "\n",
    "    n_i is refractive index of incident medium.\n",
    "\n",
    "    th_i is (complex) propegation angle through incident medium\n",
    "    (in radians, where 0=normal). \"th\" stands for \"theta\".\n",
    "    \"\"\"\n",
    "    if pol == 's':\n",
    "        return ((n_i*cos(th_i)*(1+conj(r))*(1-r)).real\n",
    "                     / (n_i*cos(th_i)).real)\n",
    "    elif pol == 'p':\n",
    "        return ((n_i*conj(cos(th_i))*(1+r)*(1-conj(r))).real\n",
    "                      / (n_i*conj(cos(th_i))).real)\n",
    "    else:\n",
    "        raise ValueError(\"Polarization must be 's' or 'p'\")\n",
    "\n",
    "def interface_R(polarization, n_i, n_f, th_i, th_f):\n",
    "    \"\"\"\n",
    "    Fraction of light intensity reflected at an interface.\n",
    "    \"\"\"\n",
    "    r = interface_r(polarization, n_i, n_f, th_i, th_f)\n",
    "    return R_from_r(r)\n",
    "\n",
    "def interface_T(polarization, n_i, n_f, th_i, th_f):\n",
    "    \"\"\"\n",
    "    Fraction of light intensity transmitted at an interface.\n",
    "    \"\"\"\n",
    "    t = interface_t(polarization, n_i, n_f, th_i, th_f)\n",
    "    return T_from_t(polarization, t, n_i, n_f, th_i, th_f)\n",
    "\n",
    "def coh_tmm(pol, n_list, d_list, th_0, 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",
    "    Outputs the following as a dictionary (see manual for details)\n",
    "\n",
    "    * r--reflection amplitude\n",
    "    * t--transmission amplitude\n",
    "    * R--reflected wave power (as fraction of incident)\n",
    "    * T--transmitted wave power (as fraction of incident)\n",
    "    * power_entering--Power entering the first layer, usually (but not always)\n",
    "      equal to 1-R (see manual).\n",
    "    * vw_list-- n'th element is [v_n,w_n], the forward- and backward-traveling\n",
    "      amplitudes, respectively, in the n'th medium just after interface with\n",
    "      (n-1)st medium.\n",
    "    * kz_list--normal component of complex angular wavenumber for\n",
    "      forward-traveling wave in each layer.\n",
    "    * th_list--(complex) propagation angle (in radians) in each layer\n",
    "    * pol, n_list, d_list, th_0, lam_vac--same as input\n",
    "\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",
    "    # Input tests\n",
    "    # if ((hasattr(lam_vac, 'size') and lam_vac.size > 1)\n",
    "    #       or (hasattr(th_0, 'size') and th_0.size > 1)):\n",
    "    #     raise ValueError('This function is not vectorized; you need to run one '\n",
    "    #                      'calculation at a time (1 wavelength, 1 angle, etc.)')\n",
    "    # if (n_list.ndim != 1) or (d_list.ndim != 1) or (n_list.size != d_list.size):\n",
    "    #     raise ValueError(\"Problem with n_list or d_list!\")\n",
    "    # assert d_list[0] == d_list[-1] == inf, 'd_list must start and end with inf!'\n",
    "    # assert abs((n_list[0]*np.sin(th_0)).imag) < 100*EPSILON, 'Error in n0 or th0!'\n",
    "    # assert is_forward_angle(n_list[0], th_0), 'Error in n0 or th0!'\n",
    "    num_layers = n_list.size\n",
    "\n",
    "    # th_list is a list with, for each layer, the angle that the light travels\n",
    "    # through the layer. Computed with Snell's law. Note that the \"angles\" may be\n",
    "    # complex!\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 * 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",
    "    olderr = sp.seterr(invalid='ignore')\n",
    "    delta = kz_list * d_list\n",
    "    sp.seterr(**olderr)\n",
    "\n",
    "    # For a very opaque layer, reset delta to avoid divide-by-0 and similar\n",
    "    # errors. The criterion imag(delta) > 35 corresponds to single-pass\n",
    "    # transmission < 1e-30 --- small enough that the exact value doesn't\n",
    "    # matter.\n",
    "    # for i in range(1, num_layers-1):\n",
    "    #     if delta[i].imag > 35:\n",
    "    #         delta[i] = delta[i].real + 35j\n",
    "    #         if 'opacity_warning' not in globals():\n",
    "    #             global opacity_warning\n",
    "    #             opacity_warning = True\n",
    "    #             print(\"Warning: Layers that are almost perfectly opaque \"\n",
    "    #                   \"are modified to be slightly transmissive, \"\n",
    "    #                   \"allowing 1 photon in 10^30 to pass through. It's \"\n",
    "    #                   \"for numerical stability. This warning will not \"\n",
    "    #                   \"be shown again.\")\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",
    "    # At the interface between the (n-1)st and nth material, let v_n be the\n",
    "    # amplitude of the wave on the nth side heading forwards (away from the\n",
    "    # boundary), and let w_n be the amplitude on the nth side heading backwards\n",
    "    # (towards the boundary). Then (v_n,w_n) = M_n (v_{n+1},w_{n+1}). M_n is\n",
    "    # M_list[n]. M_0 and M_{num_layers-1} are not defined.\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(exp(-1j*delta[i]), 0, 0, exp(1j*delta[i]),\n",
    "                           dtype=complex),\n",
    "            make_2x2_array(1, r_list[i,i+1], r_list[i,i+1], 1, dtype=complex))\n",
    "    Mtilde = make_2x2_array(1, 0, 0, 1, dtype=complex)\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,\n",
    "                                   dtype=complex)/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",
    "    # Net transmitted and reflected power, as a proportion of the incoming light\n",
    "    # power.\n",
    "    #R = R_from_r(r)\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
}