{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Crank-Nicholson scheme" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Cylindrical Diffraction Term" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Initial class coming from TP 4.\n", "\n", "* Delete non parse solver.\n", "* Choose complex coefficient (data_type).\n", "* Choose boundary conditions: null gradient at $r=0$, null field at $r \\to +\\infty$.\n", "* Matrix dependency on $r$: varying coefficient along diagonal." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "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 du/dt + V du/dx = D du2/dx2 + f(u)\"\"\"\n", " \n", " def set_grid(self, x_min, x_max, n_x, t_min, t_max, n_t):\n", "\n", " self.x_min, self.x_max, self.n_x = x_min, x_max, n_x\n", " self.t_min, self.t_max, self.n_t = t_min, t_max, n_t\n", " self.x_pts, self.delta_x = np.linspace(x_min, x_max, n_x, retstep=True, endpoint=False)\n", " self.t_pts, self.delta_t = np.linspace(t_min, t_max, n_t, retstep=True, endpoint=False)\n", " \n", " def set_parameters(self, D, V, f):\n", " \n", " self.V, self.D, self.f = V, D, f\n", "\n", " def solve(self, u_init, sparse=True, boundary_conditions=('neumann','neumann')):\n", " \n", " sig = self.D * self.delta_t / 2. / self.delta_x**2\n", " nu = self.V * self.delta_t / 4. / self.delta_x\n", " \n", " # Figure the data type (real or complex)\n", " data_type = type(sig*nu*u_init[0])\n", " \n", " self.u_matrix = np.zeros([self.n_t, self.n_x], dtype=data_type)\n", "\n", " # Using sparse matrices and specialized tridiagonal solver speeds up the calculations\n", " if sparse:\n", " \n", " A = self._fillA_sp(sig, nu, self.n_x, data_type)\n", " B = self._fillB_sp(sig, nu, self.n_x, data_type)\n", " # Set boundary conditions\n", " for b in [0,1]:\n", " if boundary_conditions[b] == 'dirichlet':\n", " # u(x,t) = 0\n", " A[1,-b] = 1.0\n", " A[2*b,1-3*b] = 0.0\n", " B[-b,-b] = 0.0\n", " B[-b,1-3*b] = 0.0\n", " elif boundary_conditions[b] == 'neumann':\n", " # u'(x,t) = 0\n", " A[2*b,1-3*b] = -2*sig\n", " B[-b,1-3*b] = 2*sig\n", " \n", " # Propagate\n", " u = u_init\n", " for n in range(self.n_t):\n", " self.u_matrix[n,:] = u\n", " fu = f(u)\n", " if n==0: fu_old = fu\n", " u = la.solve_banded((1,1),A, B.dot(u) + self.delta_t * (1.5 * fu - 0.5 * fu_old),\\\n", " check_finite=False)\n", " fu_old = fu\n", "\n", " else:\n", " \n", " A = self._make_tridiag(sig, nu, self.n_x, data_type)\n", " B = self._make_tridiag(-sig, -nu, self.n_x, data_type)\n", "\n", " # Set boundary conditions\n", " for b in [0,1]:\n", " if boundary_conditions[b] == 'dirichlet':\n", " # u(x,t) = 0\n", " A[-b,-b] = 1.0\n", " A[-b,1-3*b] = 0.0\n", " B[-b,-b] = 0.0\n", " B[-b,1-3*b] = 0.0\n", " elif boundary_conditions[b] == 'neumann':\n", " # u'(x,t) = 0\n", " A[-b,1-3*b] = -2*sig\n", " B[-b,1-3*b] = 2*sig\n", "\n", " # Propagate\n", " u = u_init\n", " for n in range(self.n_t):\n", " self.u_matrix[n,:] = u\n", " fu = f(u)\n", " if n==0: fu_old = fu\n", " u = la.solve(A, B.dot(u) + self.delta_t * (1.5 * fu - 0.5 * fu_old))\n", " fu_old = fu\n", " \n", " def get_final_u(self):\n", " \n", " return self.u_matrix[-1,:].copy()\n", " \n", " def _make_tridiag(self, sig, nu, n, data_type):\n", " \n", " M = np.diagflat(np.full(n, (1+2*sig), dtype=data_type)) + \\\n", " np.diagflat(np.full(n-1, -(sig-nu), dtype=data_type), 1) + \\\n", " np.diagflat(np.full(n-1, -(sig+nu), dtype=data_type), -1)\n", "\n", " return M\n", " \n", " def _fillA_sp(self, sig, nu, n, data_type):\n", " \"\"\"Returns a tridiagonal matrix in compact form ab[1+i-j,j]=a[i,j]\"\"\"\n", " \n", " A = np.zeros([3,n], dtype=data_type) # A has three diagonals and size n\n", " A[0] = -(sig-nu) # superdiagonal\n", " A[1] = 1+2*sig # diagonal\n", " A[2] = -(sig+nu) # subdiagonal\n", " return A\n", "\n", " def _fillB_sp(self, sig, nu, n, data_type):\n", " \"\"\"Returns a tridiagonal sparse matrix in csr-form\"\"\"\n", " \n", " _o = np.ones(n, dtype=data_type)\n", " supdiag = (sig-nu)*_o[:-1]\n", " diag = (1-2*sig)*_o\n", " subdiag = (sig+nu)*_o[:-1]\n", " return scipy.sparse.diags([supdiag, diag, subdiag], [1,0,-1], (n,n), format=\"csr\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Adapted crank Nicholson class.\n", "\n", "* Prepare half-step method.\n", "* Prepare non linear term for one of the half-step.\n", "* Use $V$ advection to add first radial derivative.\n", "\n", "$V$ is a function of $r$\n", "$$\n", "V(r) = \\dfrac{i}{2 k_0} \\dfrac{1}{r}\n", "$$\n", "and $V(0) = 0$. **Warning**: do not use $V(0)=0$ for any calculation.\n", "\n", "**Warning**: Test scheme matrices separately." ] }, { "cell_type": "code", "execution_count": 15, "metadata": { "collapsed": true }, "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 = lambda x: - self.delta_z / 4. / self.delta_r * self.V(x) # 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 = self.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": [ "We define some constants for the whole problem, as well as some functions :\n", "\n", "- wavevector\n", "- potential\n", "- initialisation of E" ] }, { "cell_type": "code", "execution_count": 16, "metadata": { "collapsed": true }, "outputs": [], "source": [ "c = 299792458\n", "lambd = 700 * 10 ** -9\n", "\n", "k = c * 2 * np.pi / lambd\n", "\n", "diff_coeff = 1j/(2 * k)\n", "\n", "\n", "def potential(r):\n", " try:\n", " return diff_coeff / r\n", " except ZeroDivisionError:\n", " return 0\n", "\n", "\n", "def gaussian(r, mu=0, sigma=1, amplitude=1):\n", " return amplitude * np.exp(-(r - mu) ** 2 / (2 * sigma ** 2))\n", "\n", "\n", "def initial_enveloppe(r_pts):\n", " return gaussian(r_pts)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Instanciation of the CN Class" ] }, { "cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/usr/local/anaconda3-5.0.0/lib/python3.6/site-packages/ipykernel_launcher.py:11: RuntimeWarning: divide by zero encountered in true_divide\n", " # This is added back by InteractiveShellApp.init_path()\n", "/usr/local/anaconda3-5.0.0/lib/python3.6/site-packages/ipykernel_launcher.py:11: RuntimeWarning: invalid value encountered in true_divide\n", " # This is added back by InteractiveShellApp.init_path()\n", "/usr/local/anaconda3-5.0.0/lib/python3.6/site-packages/ipykernel_launcher.py:29: RuntimeWarning: invalid value encountered in multiply\n" ] } ], "source": [ "crank = CrankNicolson()\n", "\n", "crank.set_grid(r_max=5, n_r=50, z_min=0, z_max=10, n_z=100)\n", "crank.set_parameters(D=diff_coeff, V=potential, f=lambda x:0)\n", "\n", "crank.solve(initial_enveloppe(crank.r_pts))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Save the result for later analysis" ] }, { "cell_type": "code", "execution_count": 18, "metadata": { "collapsed": true }, "outputs": [], "source": [ "np.savetxt(\"CN_cylindric_complex_E.dat\", crank.E_matrix)\n", "np.savetxt(\"CN_cylindric_complex_r_pts.dat\", crank.r_pts)\n", "np.savetxt(\"CN_cylindric_complex_z_pts.dat\", crank.z_pts)" ] } ], "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.3" } }, "nbformat": 4, "nbformat_minor": 2 }