Commit 300858f1 authored by Mathieu RASSON's avatar Mathieu RASSON

Half step dispersion term

parent 3f549663
......@@ -23,7 +23,9 @@
"* 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."
"* Split step into two half steps: fill one dimension for each step inside the $z$ loop, separately.\n",
"\n",
"### Initial class"
]
},
{
......@@ -129,6 +131,158 @@
" "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Adapted class\n",
"\n",
"Term added with new table dimension:\n",
"$$\n",
" D_t \\dfrac{\\partial^2 \\mathcal{E}}{\\partial t^2}\n",
"$$\n",
"Matrix similar to TP4 without advection.\n",
"\n",
"3D matrix format:\n",
"$$\n",
" E(z_i, r_j, t_k)\n",
"$$\n",
"**Warning**: Non linear term treated with **space half step**. Initial conditions at $z=0$ in **2D**. Half step $\\Delta z/2$ in the scheme."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"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 = D1*d2E/dr2 + V(r)*dE/dr + D2*d2E/dt2 + f(E)\"\"\"\n",
" \n",
" # Cylindrical grid with rectangular time\n",
" def set_grid(self, r_max, n_r, z_min, z_max, n_z, t_min, t_max, n_t):\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.t_min, self.t_max, self.n_t = t_min, t_max, n_t\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",
" self.t_pts, self.delta_t = np.linspace(t_min, t_max, t_z, retstep=True, endpoint=False)\n",
" \n",
" # Parameters of the scheme\n",
" # D1 for diffraction, D2 for dispersion\n",
" def set_parameters(self, Dr, Dt, V, f):\n",
" \n",
" # V has to be vectorised\n",
" self.Dr, self.Dt, self.V, self.f = Dr, Dt, V, f\n",
" \n",
" \n",
" \n",
" # One solving step for r dependency and t dependency\n",
" def solve(self, E_init):\n",
" \n",
" # Coefficient of matrices (with half step)\n",
" sig_r = self.Dr * (self.delta_z/2.) / 2. / self.delta_r**2\n",
" sig_t = self.Dt * (self.delta_z/2.) / 2. / self.delta_t**2\n",
" nu = lambda x : -self.V(x) * (self.delta_z/2.) / 4. / self.delta_r # minus sign for V convention\n",
" \n",
" # Empty solution matrix (z,r,t)\n",
" self.E_matrix = np.zeros([self.n_z, self.n_r, self.n_t], dtype=complex)\n",
" \n",
" # Sparse solver\n",
" Ar = self._fillA_sp(sig_r, nu, self.n_r)\n",
" Br = self._fillB_sp(sig, nu, self.n_r)\n",
" At = self._fillA_sp(sig_t, np.vectorize(lambda x : 0), self.n_t) # null advection coefficient\n",
" Bt = self._fillB_sp(sig_t, np.vectorize(lambda x : 0), self.n_t) # null advection coefficient\n",
" \n",
" \n",
" # Set boundary conditions\n",
" # Spatial:\n",
" # Dirichlet at infinity\n",
" Ar[1,-1] = 1.0\n",
" Ar[2,-2] = 0.0\n",
" Br[-1,-1] = 0.0\n",
" Br[-1,-2] = 0.0\n",
" # Neumann at r=0\n",
" Ar[0,1] = -2*sig\n",
" Br[0,1] = 2*sig\n",
" \n",
" # Time\n",
" # Neumann at t=t_min, t_max\n",
" for b in [0,1]:\n",
" At[1,-b] = 1.0\n",
" At[2*b,1-3*b] = 0.0\n",
" Bt[-b,-b] = 0.0\n",
" Bt[-b,1-3*b] = 0.0\n",
" \n",
" # Propagate\n",
" E = E_init\n",
" for n in range(self.n_z):\n",
" # General step initialisation\n",
" self.E_matrix[n,:,:] = E\n",
" \n",
" # First half step\n",
" # r operator solver\n",
" fE = self.f(E)\n",
" # Non linear term at origin\n",
" if n==0:\n",
" fE_old = fE\n",
" # Non linear term with half sum\n",
" # r dependent solver for each t\n",
" for i in range(self.n_t):\n",
" E_intermediate[:,i] = la.solve_banded((1,1), Ar, Br.dot(E[:,i]) + (self.delta_z/2.) * \\\n",
" (1.5 * fE - 0.5 * fE_old), \\\n",
" check_finite=False)\n",
" # Non linear term update\n",
" fE_old = fE\n",
" \n",
" # Second half step\n",
" # t operator solver for each r\n",
" for i in range(self.n_r):\n",
" E[i] = la.solve_banded((1,1), At, Bt.dot(E_intermediate[i]), \\\n",
" check_finite=False)\n",
" \n",
" # E has been updated, put into E_matrix at next step\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": {},
......
This diff is collapsed.
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment