crank_seq-checkpoint.ipynb 5.35 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Crank-Nicholson scheme"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time dispersion term"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add time dispersion term as a new diffusion source $D_2$.\n",
    "* Add one dimension to tables;\n",
    "* Add one parameter to initialisation;\n",
    "* Border conditions: null conditions at $t_{min}$ and $t_{max}$;\n",
    "* Fill another scheme matrix;\n",
    "* Split step into two half steps: fill one dimension for each step inside the $z$ loop, separately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
32 33 34
   "metadata": {
    "collapsed": true
   },
35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.sparse\n",
    "import scipy.linalg as la\n",
    "\n",
    "class CrankNicolson:\n",
    "    \"\"\"A class that solves dE/dz = D*d2E/dr2 + V(r)*dE/dr + f(E)\"\"\"\n",
    "    \n",
    "    # Cylindrical grid\n",
    "    def set_grid(self, r_max, n_r, z_min, z_max, n_z):\n",
    "\n",
    "        self.r_max, self.n_r = r_max, n_r\n",
    "        self.z_min, self.z_max, self.n_z = z_min, z_max, n_z\n",
    "        self.r_pts, self.delta_r = np.linspace(0, r_max, n_r, retstep=True, endpoint=False)\n",
    "        self.z_pts, self.delta_z = np.linspace(z_min, z_max, n_z, retstep=True, endpoint=False)\n",
    "    \n",
    "    # Parameters of the scheme\n",
    "    def set_parameters(self, D, V, f):\n",
    "        \n",
    "        # V has to be vectorised\n",
    "        self.D, self.V, self.f = D, V, f\n",
    "    \n",
    "    \n",
    "    \n",
    "    # One solving step for r dependency\n",
    "    def solve(self, E_init):\n",
    "        \n",
    "        # Coefficient of matrices\n",
    "        sig = self.D * self.delta_z / 2. / self.delta_r**2\n",
    "        nu = -self.V * self.delta_z / 4. / self.delta_r # minus sign for V convention\n",
    "        \n",
    "        # Empty solution matrix\n",
    "        self.E_matrix = np.zeros([self.n_z, self.n_r], dtype=complex)\n",
    "        \n",
    "        # Sparse solver\n",
    "        A = self._fillA_sp(sig, nu, self.n_r)\n",
    "        B = self._fillB_sp(sig, nu, self.n_r)\n",
    "        \n",
    "        # Set boundary conditions \n",
    "        # Dirichlet at infinity\n",
    "        A[1,-1] = 1.0\n",
    "        A[2,-2] = 0.0\n",
    "        B[-1,-1] = 0.0\n",
    "        B[-1,-2] = 0.0\n",
    "        # Neumann at r=0\n",
    "        A[0,1] = -2*sig\n",
    "        B[0,1] = 2*sig\n",
    "        \n",
    "        # Propagate\n",
    "        E = E_init\n",
    "        for n in range(self.n_z):\n",
    "            self.E_matrix[n,:] = E\n",
    "            fE = f(E)\n",
    "            # Non linear term at origin\n",
    "            if n==0:\n",
    "                fE_old = fE\n",
    "            \n",
    "            # Non linear term with half sum\n",
    "            E = la.solve_banded((1,1),A, B.dot(E) + self.delta_z * (1.5 * fE - 0.5 * fE_old),\\\n",
    "                                        check_finite=False)\n",
    "            fE_old = fE\n",
    "    \n",
    "    \n",
    "    \n",
    "    # Deep copy of results for export\n",
    "    # np.savetxt ?\n",
    "    def get_E(self):\n",
    "        \n",
    "        return self.E_matrix.copy()\n",
    "    \n",
    "    # Diagonal packing for banded\n",
    "    def _fillA_sp(self, sig, nu, n):\n",
    "        \"\"\"Returns a tridiagonal matrix in compact form ab[1+i-j,j]=a[i,j]\"\"\"\n",
    "        \n",
    "        A = np.zeros([3,n], dtype=complex) # A has three diagonals and size n\n",
    "        # Superdiagonal\n",
    "        A[0,1:] = -(sig - nu(self.r_pts[:-1]))\n",
    "        # Diagonal\n",
    "        A[1] = 1+2*sig\n",
    "        # Subdiagonal\n",
    "        A[2,:-1] = -(sig + nu(self.r_pts[1:]))\n",
    "        \n",
    "        return A\n",
    "    \n",
    "    # Sparse tridiagonal storage\n",
    "    def _fillB_sp(self, sig, nu, n):\n",
    "        \"\"\"Returns a tridiagonal sparse matrix in csr-form\"\"\"\n",
    "        \n",
    "        _o = np.ones(n, dtype=complex)\n",
    "        supdiag = (sig - nu(self.r_pts[:-1]))\n",
    "        diag = (1-2*sig)*_o\n",
    "        subdiag = (sig + nu(self.r_pts[1:]))\n",
    "        \n",
    "        return scipy.sparse.diags([supdiag, diag, subdiag], [1,0,-1], (n,n), format=\"csr\")\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Non linear terms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute $\\mathcal{E} \\mapsto f(\\mathcal{E})$ to add non linear terms. Test one after another, with divergence if only Kerr effect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
151 152 153
   "metadata": {
    "collapsed": true
   },
154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173
   "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",
174
   "version": "3.6.8"
175 176 177 178 179
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}