{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Python for Scientists, 2nd Edition: Code Samples" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The book contains many valid Python code snippets. Many are trivial, but most of those whose length is at least five lines are available here." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Chapter 2" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §2.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "a = 3\n", "b = 2.7\n", "c = 12 + 5j\n", "s = \"Hello World!\"\n", "L = [a, b, c, s, 77.77]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §2.5" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "#%%writefile fib.py\n", "# File: fib.py Fibonacci numbers\n", "\n", "\"\"\"Module fib: Fibonacci numbers\n", " Contains one function fib(n).\n", "\"\"\"\n", "\n", "def fib(n):\n", " \"\"\" Returns n'th Fibonacci number. \"\"\"\n", " a, b = 0, 1\n", " for i in range(n):\n", " a, b = b, a+b\n", " return a\n", "\n", "##########################################################\n", "\n", "if __name__ == \"__main__\":\n", " for i in range(1001):\n", " print \"fib(\", i, \") = \", fib(i)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §2.5.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "#%%writefile gcd.py\n", "# File gcd.py Implementing the GCD Euclidean algorithm.\n", "\n", "\"\"\" Module gcd: contains two implementations of the Euclid\n", " GCD algorithm, gcdr and gcd.\n", "\"\"\" \n", "\n", "def gcdr(a, b):\n", " \"\"\" Euclidean algorithm, recursive vers., returns GCD. \"\"\"\n", " if b == 0:\n", " return a\n", " else:\n", " return gcdr(b, a%b)\n", "\n", "def gcd(a, b):\n", " \"\"\" Euclidean algorithm, non-recursive vers., returns GCD. \"\"\"\n", " while b:\n", " a, b = b, a%b\n", " return a\n", "\n", "##########################################################\n", "if __name__ == \"__main__\":\n", " import fib\n", "\n", " for i in range(963):\n", " print i, ' ', gcdr(fib.fib(i), fib.fib(i+1))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Chapter 3" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §3.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "a = 4; b = 5.5; c = 1.5+2j; d = 'a'\n", "e = 6.0*a - b*b + \\\n", " c**(a+b+c)\n", "f = 6.0*a - b*b + c**(\n", " a+b+c)\n", "a, b, c, d, e, f" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §3.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "p =3.14\n", "p\n", "q = p\n", "p = 'pi'\n", "p\n", "q" ] }, { "cell_type": "markdown", "metadata": { "collapsed": true }, "source": [ "### §3.5.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "[1, 4.0, 'a']\n", "u = [1, 4.0, 'a']\n", "v = [3.14, 2.78, u, 42]\n", "v\n", "len(v) \n", "len? # or help(len) \n", "v*2\n", "v + u\n", "v.append('foo')\n", "v" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §3.5.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "u = [0, 1, 2, 3, 4, 5, 6, 7]\n", "su = u[2 : 4]\n", "su\n", "su[0] = 17\n", "su\n", "u\n", "u[2 : 4] = [3.14,' a']\n", "u" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §3.5.4" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "a = 4\n", "b = a\n", "b = 'foo'\n", "a\n", "b\n", "u = [0, 1, 4, 9, 16]\n", "v = u\n", "v[2] = 'foo'\n", "v\n", "u" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "u = [0, 1, 4, 9, 16]\n", "v = u[ : ]\n", "v[2] = 'foo'\n", "v\n", "u" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §3.5.7" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "empty={} \n", "parms = {'alpha':1.3, 'beta':2.74} \n", "#parms = dict(alpha=1.3,beta=2.74) \n", "parms['gamma']=0.999\n", "parms" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §3.7.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "y = 107\n", "for x in range(2, y):\n", " if y%x == 0:\n", " print y, \" has a factor \", x \n", " break\n", "else:\n", " print y, \"is prime.\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §3.8.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def add_one(x):\n", " \"\"\" Takes x and returns x + 1. \"\"\"\n", " x = x + 1\n", " return x\n", "\n", "x=23\n", "#add_one? # or help(add_one)\n", "add_one(0.456)\n", "x" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def add_x_and_y(x, y):\n", " \"\"\" Add x and y and return them and their sum. \"\"\"\n", " z = x + y\n", " return x, y, z\n", "\n", "a, b, c = add_x_and_y(1,0.456)\n", "a, b, c\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "L = [0, 1, 2]\n", "print id(L)\n", "def add_with_side_effects(M):\n", " \"\"\" Increment first element of list. \"\"\"\n", " M[0] += 1\n", "\n", "add_with_side_effects(L)\n", "print L\n", "id(L)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "L = [0, 1, 2]\n", "id(L)\n", "def add_without_side_effects(M):\n", " \"\"\" Increment first element of list. \"\"\"\n", " MC = M[ : ]\n", " MC[0] +=1 \n", " return MC\n", "\n", "L = add_without_side_effects(L)\n", "L\n", "id(L)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §3.8.4" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def average(*args):\n", " \"\"\" Return mean of a non-empty tuple of numbers. \"\"\"\n", " print args\n", " sum = 0.0\n", " for x in args:\n", " sum += x\n", " return sum/len(args)\n", "\n", "print average(1,2,3,4)\n", "print average(1,2,3,4,5)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §3.9" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "#%%writefile frac.py\n", "# File: frac.py\n", "\n", "import gcd\n", "\n", "class Frac:\n", " \"\"\" Fractional class. A Frac is a pair of integers num, den\n", " (with den!=0) whose GCD is 1.\n", " \"\"\"\n", "\n", " def __init__(self, n, d):\n", " \"\"\" Construct a Frac from integers n and d.\n", " Needs error message if d = 0!\n", " \"\"\"\n", " hcf=gcd.gcd(n, d)\n", " self.num, self.den = n/hcf, d/hcf\n", "\n", " def __str__(self):\n", " \"\"\" Generate a string representation of a Frac. \"\"\"\n", " return \"%d/%d\" % (self.num,self.den)\n", "\n", " def __mul__(self, another):\n", " \"\"\" Multiply two Fracs to produce a Frac. \"\"\"\n", " return Frac(self.num*another.num, self.den*another.den)\n", "\n", " def __add__(self,another):\n", " \"\"\" Add two Fracs to produce a Frac. \"\"\"\n", " return Frac(self.num*another.den + self.den*another.num,\n", " self.den*another.den)\n", "\n", " def to_real(self):\n", " \"\"\" Return floating point value of Frac. \"\"\"\n", " return float(self.num)/float(self.den)\n", "\n", "if __name__ == \"__main__\":\n", " a=Frac(3,7)\n", " b=Frac(24,56)\n", " print \"a.num= \", a.num, \", b.den= \", b.den\n", " print a\n", " print b\n", " print \"floating point value of a is \", a.to_real()\n", " print \"product= \",a*b,\", sum= \",a + b\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §3.11" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "def sieve_v1(n):\n", " \"\"\"\n", " Use Sieve of Eratosthenes to compute list of primes <= n.\n", " Version 1\n", " \"\"\"\n", " primes = range(2, n+1)\n", " for p in primes:\n", " if p*p > n:\n", " break\n", " product = 2*p\n", " while product <= n:\n", " if product in primes:\n", " primes.remove(product)\n", " product += p\n", " return len(primes), primes" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def sieve_v2(n):\n", " \"\"\" \n", " Sieve of Eratosthenes to compute list of primes <= n.\n", " Version 2.\n", " \"\"\"\n", " sieve = [True]*(n+1)\n", " for i in xrange(3, n+1, 2):\n", " if i*i > n:\n", " break\n", " if sieve[i]:\n", " sieve[i*i : : 2*i] = [False]*((n - i*i) // (2*i) + 1)\n", " answer = [2] + [i for i in xrange(3, n+1, 2) if sieve[i]]\n", " return len(answer), answer" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Chapter 4" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.1.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "xl = np.linspace(0, 0.5, 5, endpoint=False)\n", "xm = np.linspace(0.5, 0.6, 10,endpoint=False)\n", "xr = np.linspace(0.6, 1.0, 5)\n", "xs = np.hstack((xl, xm, xr))\n", "xs" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.1.5" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def m1(x):\n", " return np.exp(2*x)\n", "def m2(x):\n", " return 1.0\n", "def m3(x):\n", " return np.exp(1.0-x)\n", " \n", "x = np.linspace(-10, 10, 21) \n", "conditions=[ x >= 0, x >= 1, x < 0 ]\n", "functions=[ m2, m3, m1 ]\n", "m = np.piecewise(x, conditions, functions)\n", "m" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.2.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "x=np.array([[0, 1, 2, 3], [10, 11, 12, 13], [20, 21, 22, 23]]) \n", "r=np.array([2, 3, 4, 5])\n", "c=np.array([[5], [6], [7]])\n", "print \"2*x = \\n\", 2*x\n", "print \"r*x = \\n\", r*x \n", "print \"x*r = \\n\", x*r\n", "print \"c*x = \\n\", c*x" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.2.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "xv=np.linspace(-1, 1, 5)\n", "yv=np.linspace(0, 1, 3)\n", "[xa, ya] = np.meshgrid(xv, yv)\n", "print \"xa = \\n\", xa\n", "print \"ya = \\n\", ya\n", "print \"xa*ya = \\n\", xa*ya" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "[xo, yo]=np.ogrid[-1: 1: 5j, 0: 1: 3j] \n", "print \"xo = \\n\", xo\n", "print \"yo = \\n\", yo\n", "print \"xo.shape = \", xo.shape, \" yo.shape = \", yo.shape\n", "print \"xo*yo = \\n\", xo*yo" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.4.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "import numpy as np\n", "quarter = np.array([1, 2, 3, 4],dtype=int)\n", "results=np.array([37.4, 47.3, 73.4, 99])\n", "outfile=open(\"q4.txt\",\"w\")\n", "outfile.write(\"The results for the first four quarters\\n\\n\")\n", "for q,r in zip(quarter, results):\n", " outfile.write(\"For quarter %d the result is %5.1f\\n\" % (q,r))\n", "outfile.close()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "infile=open(\"q4.txt\", \"r\")\n", "lquarter = []\n", "lresult = []\n", "temp = infile.readline()\n", "temp = infile.readline()\n", "for line in infile:\n", " words = line.split()\n", " lquarter.append(int(words[2]))\n", " lresult.append(float(words[6]))\n", "infile.close()\n", "import numpy as np\n", "aquarter = np.array(lquarter, dtype=int)\n", "aresult = np.array(lresult)\n", "aquarter, aresult" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.4.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "len = 21\n", "x = np.linspace(0, 2*np.pi, len)\n", "c = np.cos(x)\n", "s = np.sin(x)\n", "t = np.tan(x)\n", "arr = np.empty((4, len), dtype=float)\n", "arr[0, : ] = x\n", "arr[1, : ] = c\n", "arr[2, : ] = s\n", "arr[3, : ] = t\n", "np.savetxt('x.txt', x)\n", "np.savetxt('xcst.txt', (x, c, s, t))\n", "np.savetxt('xarr.txt', arr)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.4.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "np.savez('test.npz',x = x, c = c, s = s, t = t)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "temp=np.load('test.npz')\n", "temp.files\n", "xc = temp['x']\n", "cc = temp['c']\n", "sc = temp['s']\n", "tc = temp['t']" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.7.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "\n", "roots = [0, 1, 1, 2]\n", "coeffs = np.poly(roots)\n", "print \"coeffs = \\n\", coeffs\n", "print \"np.roots(coeffs) = \\n\", np.roots(coeffs)\n", "x = np.linspace(0,0.5*np.pi,7)\n", "y = np.sin(x)\n", "c = np.polyfit(x, y, 3)\n", "print \"c = \\n\", c\n", "y1 = np.polyval(c, x)\n", "print \"y = \\n\", y, \"\\n y1 = \\n\", y1, \"\\n y1-y = \\n\", y1-y" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.8.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import numpy.linalg as npl\n", "\n", "a=np.array([[4, 2, 0], [9, 3, 7], [1, 2, 1]])\n", "print \"a = \", a\n", "print \"det(a) = \", npl.det(a)\n", "b=npl.inv(a)\n", "print \"b = \", b, \"\\n b*a = \", np.dot(b,a)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import numpy.linalg as npl\n", "\n", "a = np.array([[-2, -4, 2], [-2, 1, 2], [4, 2, 5]])\n", "evals, evecs = npl.eig(a)\n", "eval1 = evals[0]\n", "evec1 = evecs[:,0]\n", "eval1, evec1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "npl.norm(evec1), np.dot(evec1, evec1)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "np.dot(a, evec1) - eval1*evec1" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.8.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import numpy.linalg as npl\n", "a=np.array([[3,2,1],[5,5,5],[1,4,6]])\n", "b=np.array([[5,1],[5,0],[-3,-7.0/2]])\n", "x=npl.solve(a,b)\n", "x" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "np.dot(a,x)-b" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §4.9.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import scipy.optimize as sco\n", "def fun(x):\n", " return np.cosh(x)/np.sinh(x)-x\n", "\n", "roots = sco.fsolve(fun, 1.0)\n", "root = roots[0]\n", "print \"root is %15.12f and value is %e\" % (root, fun(root))" ] }, { "cell_type": "markdown", "metadata": { "collapsed": true }, "source": [ "## Chapter 5" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §5.2.3" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "If you are comparing the two notebook options, you should reset the kernel before switching." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false, "scrolled": false }, "outputs": [], "source": [ "#%%writefile foo.py \n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "\n", "# Use the next command in Ipython terminal mode\n", "#plt.ion()\n", "# Use the next command in Ipython non-inline notebook mode\n", "#matplotlib\n", "# Use the next command in Ipython inline notebook mode\n", "%matplotlib notebook\n", "\n", "x = np.linspace(-np.pi, np.pi, 101)\n", "y = np.sin(x) + np.sin(3*x)/3.0\n", " \n", "plt.plot(x, y)\n", "\n", "plt.xlabel('x')\n", "plt.ylabel('y')\n", "plt.title('A simple plot')\n", "\n", "#plt.show()\n", "plt.savefig('foo.pdf')" ] }, { "cell_type": "markdown", "metadata": { "collapsed": true }, "source": [ "### §5.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "\n", "x = np.linspace(-np.pi, np.pi, 101)\n", "y = np.sin(x) + np.sin(3*x)/3.0\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "ax.plot(x, y)\n", "ax.set_xlabel('x')\n", "ax.set_ylabel('y')\n", "ax.set_title('A simple plot')\n", "plt.plot(x, y)\n", "\n", "fig.savefig('foo.pdf')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §5.4" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "\n", "x=np.linspace(-np.pi,np.pi,101)\n", "f=np.ones_like(x)\n", "f[x<0] = -1\n", "y1 = (4/np.pi)*(np.sin(x) + np.sin(3*x)/3.0)\n", "y2 = y1 + (4/np.pi)*(np.sin(5*x)/5.0 + np.sin(7*x)/7.0)\n", "y3 = y2 + (4/np.pi)*(np.sin(9*x)/9.0 + np.sin(11*x)/11.0)\n", "\n", "plt.ion()\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "ax.plot(x,f,'b-',lw=3,label='f(x)')\n", "ax.plot(x,y1,'c--',lw=2,label='two terms')\n", "ax.plot(x,y2,'r-.',lw=2,label='four terms')\n", "ax.plot(x, y3,'b:',lw=2,label='six terms')\n", "ax.legend(loc='best')\n", "ax.set_xlabel('x',style='italic')\n", "ax.set_ylabel('partial sums',style='italic')\n", "fig.suptitle('Partial sums for Fourier series of f(x)',\n", " size=16,weight='bold') " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §5.5" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt \n", "%matplotlib notebook\n", "\n", "theta = np.linspace(0,2*np.pi,201) \n", "r1 = np.abs(np.cos(5.0*theta) - 1.5*np.sin(3.0*theta)) \n", "r2 = theta/np.pi\n", "r3 = 2.25*np.ones_like(theta)\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111,projection='polar')\n", "ax.plot(theta, r1,label='trig')\n", "ax.plot(5*theta, r2,label='spiral')\n", "ax.plot(theta, r3,label='circle')\n", "ax.legend(loc='best') " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §5.6" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import numpy.random as npr \n", "x = np.linspace(0,4,21) \n", "y = np.exp(-x) \n", "xe = 0.08*npr.randn(len(x)) \n", "ye = 0.1*npr.randn(len(y))\n", "\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "ax.errorbar(x, y, fmt='bo', lw=2, xerr=xe, yerr=ye, ecolor='r', elinewidth=1)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §5.7" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np \n", "x = np.linspace(0,2,101) \n", "y = (x-1)**3+1\n", "\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "ax.plot(x,y)\n", "ax.annotate('point of inflection at x=1', xy=(1,1), xytext=(0.8,0.5), \n", " arrowprops=dict(facecolor='black',width=1, shrink=0.05))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §5.9" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false, "scrolled": true }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "[X,Y] = np.mgrid[-2.5:2.5:51j, -3:3:61j]\n", "Z=X**2-Y**2\n", "curves = ax.contour(X, Y, Z, 12,colors='k')\n", "ax.clabel(curves)\n", "fig.suptitle(r'The level contours of $z=x^2-y^2$', fontsize=20)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "[X,Y] = np.mgrid[-2.5:2.5:51j, -3:3:61j]\n", "Z=X**2-Y**2\n", "im = ax.contourf(X, Y, Z, 12)\n", "fig.colorbar(im, orientation='vertical')\n", "\n", "ax.set_title(r'${\\rm The\\ level\\ contours\\ of\\:} z=x^2-y^2$', fontsize=20)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "\n", "from matplotlib import rc\n", "rc('font',family='serif')\n", "rc('text',usetex = True)\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "[X,Y] = np.mgrid[-2.5:2.5:51j, -3:3:61j]\n", "Z=X**2-Y**2\n", "im = ax.imshow(Z)\n", "fig.colorbar(im, orientation='vertical')\n", "\n", "ax.set_title(r'The level contours of $z=x^2-y^2$', fontsize=20)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §5.10.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false, "scrolled": false }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook \n", "\n", "x = np.linspace(0, 2*np.pi, 301)\n", "y = np.cos(x)\n", "z = np.sin(x)\n", "\n", "plt.ion()\n", "fig1 = plt.figure()\n", "ax1 = fig1.add_subplot(111) \n", "ax1.plot(x,y) # First figure \n", "fig2 = plt.figure()\n", "ax2 = fig2.add_subplot(111) \n", "ax2.plot(x,z) # Second figure" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §5.10.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "\n", "x = np.linspace(0, 5, 101)\n", "y1 = 1.0/(x + 1.0)\n", "y2 = np.exp(-x)\n", "y3 = np.exp(-0.1*x**2)\n", "y4 = np.exp(-5*x**2)\n", " \n", "fig = plt.figure()\n", "ax1 = fig.add_subplot(2, 2, 1)\n", "ax1.plot(x, y1)\n", "ax1.set_xlabel('x') \n", "ax1.set_ylabel('y1')\n", "ax2 = fig.add_subplot(222)\n", "ax2.plot(x, y2)\n", "ax2.set_xlabel('x')\n", "ax2.set_ylabel('y2')\n", "ax3 = fig.add_subplot(223)\n", "ax3.plot(x, y3)\n", "ax3.set_xlabel('x')\n", "ax3.set_ylabel('y3')\n", "ax4 = fig.add_subplot(224)\n", "ax4.plot(x, y4)\n", "ax4.set_xlabel('x')\n", "ax4.set_ylabel('y4')\n", "fig.suptitle('Various decay functions')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §5.11" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np \n", "from time import time\n", "\n", "# Set the parameters\n", "max_iter = 256 # maximum number of iterations\n", "nx, ny = 1024, 1024 # x- and y-image resolutions\n", "x_lo, x_hi = -2.0, 1.0 # x bounds in complex plane\n", "y_lo, y_hi = -1.5, 1.5 # y bounds in complex plane\n", "\n", "start_time = time()\n", "\n", "# Construct the two dimensional arrays\n", "ix, iy = np.mgrid[0:nx, 0:ny] \n", "x, y = np.mgrid[x_lo:x_hi:1j*nx, y_lo:y_hi:1j*ny]\n", "c = x + 1j*y\n", "esc_parm = np.zeros((ny, nx, 3),dtype='uint8') # holds pixel rgb data\n", "\n", "# Flattened arrays\n", "nxny = nx*ny\n", "ix_f = np.reshape(ix,nxny)\n", "iy_f = np.reshape(iy,nxny)\n", "c_f = np.reshape(c,nxny)\n", "z_f=c_f.copy() # the iterated variable\n", "\n", "for iter in xrange(max_iter): # do the iterations\n", " if not len(z_f): # all points have escaped\n", " break\n", " # rgb values for this choice of iter\n", " n = iter + 1\n", " r, g, b = n % 4 * 64, n % 8 * 32, n % 16 * 16\n", " \n", " # Mandelbrot evolution\n", " z_f *= z_f\n", " z_f += c_f\n", " escape=np.abs(z_f) > 2.0 # points which are escaping\n", " # Set the rgb pixel value for the escaping points\n", " esc_parm[iy_f[escape], ix_f[escape], :] = r, g, b\n", " escape = -escape # points not escaping\n", " # Remove batch of newly escaped points from flattened arrays\n", " ix_f = ix_f[escape]\n", " iy_f = iy_f[escape]\n", " c_f = c_f[escape]\n", " z_f = z_f[escape]\n", "\n", "print \"Time taken = \", time() - start_time\n", "\n", "from PIL import Image\n", "\n", "picture = Image.fromarray(esc_parm)\n", "picture.show()\n", "picture.save(\"mandelbrot.jpg\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Chapter 6" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 6.5.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt \n", "%matplotlib notebook\n", "from matplotlib.widgets import Slider, Button, RadioButtons\n", "\n", "def solwave(x, t, c):\n", " \"\"\" Solitary wave solution of the K deV equation.\"\"\" \n", " return c/(2*np.cosh(np.sqrt(c)*(x-c*t)/2)**2)\n", "\n", "fig, ax = plt.subplots()\n", "plt.subplots_adjust(left=0.15, bottom=0.30)\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"u\")\n", "x = np.linspace(-5.0, 20.0, 1001)\n", "t0 = 5.0\n", "c0 = 1.0\n", "line, = plt.plot(x, solwave(x, t0, c0), lw=2, color='blue')\n", "plt.axis([-5, 20, 0, 2])\n", "\n", "axcolor = 'lightgoldenrodyellow'\n", "axtime = plt.axes([0.20, 0.15, 0.65, 0.03], axisbg=axcolor)\n", "axvely = plt.axes([0.20, 0.1, 0.65, 0.03], axisbg=axcolor)\n", "\n", "stime = Slider(axtime, 'Time', 0.0, 20.0, valinit=t0)\n", "svely = Slider(axvely, 'Vely', 0.1, 3.0, valinit=c0)\n", "\n", "def update(val):\n", " time = stime.val\n", " vely = svely.val\n", " line.set_ydata(solwave(x, time, vely))\n", " fig.canvas.draw_idle()\n", "\n", "svely.on_changed(update)\n", "stime.on_changed(update)\n", " \n", "resetax = plt.axes([0.75, 0.025, 0.1, 0.04])\n", "button = Button(resetax, 'Reset', color=axcolor, \n", " hovercolor='0.975')\n", "\n", "def reset(event):\n", " svely.reset()\n", " stime.reset()\n", " \n", "button.on_clicked(reset)\n", "\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §6.5.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt \n", "%matplotlib inline\n", "from matplotlib import animation\n", "from JSAnimation import IPython_display\n", "\n", "def solwave(t, x, c=1):\n", " \"\"\" Solitary wave solution of the K deV equation.\"\"\" \n", " return c/(2*np.cosh(np.sqrt(c)*(x-c*t)/2)**2)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Initialization\n", "fig = plt.figure()\n", "ax = plt.axes(xlim=(-5, 20), ylim=(0, 0.6))\n", "line, = ax.plot([], [], lw=2)\n", "\n", "t=np.linspace(-10,25,91)\n", "x = np.linspace(-5, 20.0, 101)\n", "\n", "def init():\n", " line.set_data([], [])\n", " return line,\n", "\n", "def animate(i):\n", " y = solwave(t[i], x)\n", " line.set_data(x, y)\n", " return line,\n", "\n", "animation.FuncAnimation(fig, animate, init_func=init, \n", " frames=90, interval=30, blit=True)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §6.5.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "\n", "# Fix speed\n", "c = 1.0\n", "\n", "def solwave(t, x):\n", " \"\"\" Solitary wave solution of the K deV equation.\"\"\" \n", " return c/(2*np.cosh(np.sqrt(c)*(x-c*t)/2)**2)\n", " \n", "import matplotlib.pyplot as plt \n", "\n", "def plot_solwave(t, x):\n", " fig = plt.figure()\n", " ax = fig.add_subplot(111)\n", " ax.plot(x, solwave(t, x))\n", " ax.set_ylim(0, 0.6*c)\n", " ax.text(-4, 0.55*c, \"t = \" + str(t))\n", " ax.text(-4, 0.50*c, \"c = \" + str(c))\n", "\n", "x = np.linspace(-5, 20.0, 101)\n", "t = np.linspace(-10, 25, 701)\n", "\n", "for i in range(len(t)):\n", " file_name='_temp%05d.png' % i\n", " plot_solwave(t[i], x)\n", " plt.savefig(file_name)\n", " plt.clf()\n", "\n", "import os\n", "os.system(\"rm _movie.mpg\")\n", "os.system(\"/opt/local/bin/ffmpeg -r 25 \" +\n", " \" -i _temp%05d.png -b:v 1800 _movie.mpg\")\n", "os.system(\"rm _temp*.png\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §6.7.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false, "scrolled": false }, "outputs": [], "source": [ "import numpy as np\n", " \n", "theta = np.linspace(0, 2*np.pi, 401)\n", "a = 0.3 # specific but arbitrary choice of the parameters\n", "m = 11\n", "n = 9\n", "x = (1 + a*np.cos(n*theta))*np.cos(m*theta)\n", "y = (1 + a*np.cos(n*theta))*np.sin(m*theta)\n", "z = a*np.sin(n*theta)\n", "\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "from mpl_toolkits.mplot3d import Axes3D\n", "\n", "plt.ion() \n", "fig = plt.figure()\n", "ax = Axes3D(fig)\n", "ax.plot(x, y, z, 'g', linewidth=4)\n", "ax.set_zlim3d(-1.0, 1.0)\n", "ax.set_xlabel('x')\n", "ax.set_ylabel('y')\n", "ax.set_zlabel('z')\n", "ax.set_title('A spiral as a parametric curve', weight='bold', size=16)\n", "#ax.elev, ax.azim = 60, -120\n", "fig.savefig('torus.pdf')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §6.7.2" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "This snippet should be saved as a file. You then need to invoke the Mayavi application." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile torus.py \n", "\n", "import numpy as np\n", "\n", "theta = np.linspace(0, 2*np.pi, 401)\n", "a = 0.3 # specific but arbitrary choice of the parameters\n", "m = 11\n", "n = 9\n", "x = (1 + a*np.cos(n*theta))*np.cos(m*theta)\n", "y = (1 + a*np.cos(n*theta))*np.sin(m*theta)\n", "z = a*np.sin(n*theta)\n", "\n", "from mayavi import mlab \n", "mlab.plot3d(x, y, z, np.sin(n*theta),\n", " tube_radius=0.025, colormap='spectral')\n", "mlab.axes(line_width=2, nb_labels=5)\n", "mlab.title('A spiral wrapped around a torus',size=1.2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §6.8.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "\n", "xx, yy = np.mgrid[-2:2:81j, -3:3:91j]\n", "zz = np.exp(-2*xx**2 - yy**2)*np.cos(2*xx)*np.cos(3*yy) \n", "\n", "import matplotlib.pyplot as plt\n", "from mpl_toolkits.mplot3d import Axes3D\n", "%matplotlib notebook\n", "\n", "plt.ion()\n", "\n", "fig = plt.figure()\n", "ax = Axes3D(fig)\n", "ax.plot_surface(xx, yy, zz, rstride=4, cstride=3, color='c',alpha=0.9)\n", "ax.contour(xx, yy, zz, zdir='x',offset=-3.0, colors='black')\n", "ax.contour(xx, yy, zz, zdir='y', offset=4.0, colors='blue')\n", "ax.contour(xx, yy, zz, zdir='z', offset=-2.0)\n", "ax.set_xlim3d(-3.0, 2.0)\n", "ax.set_ylim3d(-3.0, 4.0)\n", "ax.set_zlim3d(-2.0, 1.0)\n", "ax.set_xlabel('x')\n", "ax.set_ylabel('y')\n", "ax.set_zlabel('z')\n", "\n", "fig.savefig('surf1.pdf')\n" ] }, { "cell_type": "markdown", "metadata": { "collapsed": true }, "source": [ "### §6.8.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile surf2.py\n", "import numpy as np\n", " \n", "xx, yy = np.mgrid[-2:2:81j, -3:3:91j]\n", "zz = np.exp(-2*xx**2 - yy**2)*np.cos(2*xx)*np.cos(3*yy)\n", "\n", "from mayavi import mlab\n", "fig = mlab.figure()\n", "s = mlab.surf(xx, yy, zz, representation='surface')\n", "ax = mlab.axes(line_width=2, nb_labels=5)\n", "#mlab.title('Simple surface plot',size=0.4)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §6.9.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "\n", "[u, v] = np.mgrid[-2:2:51j, -2:2:61j] \n", "x, y, z = u*(1 - u**2/3 + v**2), v*(1 - v**2/3 + u**2), u**2 - v**2\n", "\n", "import matplotlib.pyplot as plt\n", "from mpl_toolkits.mplot3d import Axes3D\n", "plt.ion()\n", "\n", "fig = plt.figure()\n", "ax = Axes3D(fig)\n", "ax.plot_surface(x.T, y.T, z.T, rstride=2, cstride=2, color='r',\n", " alpha=0.2, linewidth=0.5)\n", "ax.elev, ax.azim = 50, -80\n", "ax.set_xlabel('x')\n", "ax.set_ylabel('y')\n", "ax.set_zlabel('z')\n", "#ax.set_title(’A parametric surface plot’, # weight=’bold’,size=18)\n", "fig.savefig('surf3.pdf')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 6.9.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile surf4.py\n", "\n", "import numpy as np\n", "[u, v] = np.mgrid[-2:2:51j, -2:2:61j]\n", "x, y, z = u*(1 - u**2/3 + v**2), v*(1 - v**2/3 + u**2), u**2 - v**2\n", "\n", "from mayavi import mlab\n", "fig = mlab.figure()\n", "s = mlab.mesh(x, y, z, representation='surface',\n", " line_width=0.5, opacity=0.5)\n", "ax = mlab.axes(line_width=2, nb_labels=5)" ] }, { "cell_type": "markdown", "metadata": { "collapsed": true }, "source": [ "### §6.10" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile Julia.py\n", "\n", "import numpy as np\n", "\n", "# Set up initial grid\n", "x, y = np.ogrid[-1.5:0.5:1000j, -1.0:1.0:1000j]\n", "z= x + 1j*y\n", "julia = np.zeros(z.shape)\n", "c = -0.7 - 0.4j\n", "\n", "# Build the Julia set\n", "for it in range(1,101): \n", " z = z**2 + c\n", " escape = z*z.conj()>4 \n", " julia += (1/float(it))*escape\n", " \n", "from mayavi import mlab \n", "mlab.figure(size = (800,600)) \n", "mlab.surf(julia, colormap='gist_ncar', warp_scale='auto', vmax=1.5)\n", "mlab.view(15,30,500,[-0.5,-0.5,2.0])\n", "mlab.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Chapter 7" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §7.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import sympy as sy\n", "sy.init_printing()\n", "x, y = sy.symbols(\"x y\")\n", "D = (x + y)*sy.exp(x)*sy.cos(y)\n", "D" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "f, g = sy.symbols(\"f g\", cls=sy.Function)\n", "f(x), f(x, y), g(x)\n", "i, j = sy.symbols(\"i j\", integer=True)\n", "u, v = sy.symbols(\"u v\", real=True)\n", "i.is_integer, j*j, (j*j).is_integer, u.is_real" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §7.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "D_s = sy.sympify('(x+y)*exp(x)*cos(y)')\n", "cosdiff = sy.sympify(\"cos(x)*cos(y) + sin(x)*sin(y)\")\n", "D_s, cosdiff" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "func = sy.lambdify((x, y), cosdiff, \"numpy\") \n", "xn = np.linspace(0, np.pi, 13)\n", "func(xn, 0.0) # should return np.cos(xn)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §7.4" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "M = sy.Matrix([[1,x],[y,1]])\n", "V = sy.Matrix([[u],[v]])\n", "M, V, M*V " ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "M[0,1]= u\n", "M" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §7.8" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from sympy.solvers import dsolve\n", "dsolve?" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "ode1 = f(x).diff(x, 2) + 4*f(x)\n", "sol1 = dsolve(ode1, f(x))\n", "ode1, sol1 " ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "C1, C2 = sy.symbols(\"C1, C2\")\n", "fun = sol1.rhs\n", "fund = fun.diff(x)\n", "fun.subs(x, 0), fund.subs(x, 0)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "psol = sol1.subs([(C2, 2), (C1, 0)])\n", "psol" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ " ode2 = sy.sin(f(x)) + (x*sy.cos(f(x)) + f(x))*f(x).diff(x)\n", " ode2, dsolve(ode2, f(x))" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "ode3 = x*f(x).diff(x) + f(x) - sy.log(x)*f(x)**2\n", "ode3, dsolve(ode3)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ " ode5 = f(x).diff(x, 2) + (f(x).diff(x))**2/f(x) + f(x).diff(x)/x\n", " ode5, dsolve(ode5)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ " ode6 = x*(f(x).diff(x, 2)) + 2*(f(x).diff(x)) + x*f(x)\n", " ode6, dsolve(ode6)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "ode7 = [f(x).diff(x) - 2*f(x) - g(x), g(x).diff(x) -f(x) - 2*g(x)]\n", "ode7, dsolve(ode7)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "ode8 = f(x).diff(x) - (x + f(x))**2\n", "ode8, dsolve(ode8)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "ode9 = g(x).diff(x) -1 - (g(x))**2\n", "dsolve(ode9)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §7.9" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%matplotlib notebook\n", "import sympy.plotting as syp" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "s1 = syp.plot(sy.sin(x), x, x-x**3/6, x-x**3/6 + x**5/120, (x, -4, 4),\n", " title='sin(x) and its first three Taylor approximants')" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false, "scrolled": true }, "outputs": [], "source": [ "xc = sy.cos(u) + sy.cos(7*u)/2 + sy.sin(17*u)/3\n", "yc = sy.sin(u) + sy.sin(7*u)/2 + sy.cos(17*u)/3\n", "fig = syp.plot_parametric(xc, yc, (u, 0, 2*sy.pi))\n", "fig.save(\"sympy2.pdf\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "x, y = sy.symbols(\"x y\", real=True)\n", "z = x + sy.I*y\n", "w = sy.cos(z**2).expand(complex=True)\n", "wa = sy.Abs(w).expand(complex=True)\n", "syp.plot_implicit(wa**2- 1)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false, "scrolled": true }, "outputs": [], "source": [ "fig = syp.plot3d_parametric_surface((3 + sy.sin(u) + sy.cos(v))*sy.cos(2*u),\n", " (3 + sy.sin(u) + sy.cos(v))*sy.sin(2*u),\n", " 2*sy.cos(u) + sy.sin(v),\n", " (u, 0, 2*sy.pi), (v, 0, 2*sy.pi))\n", "fig.save(\"foo.pdf\")" ] }, { "cell_type": "markdown", "metadata": { "collapsed": true }, "source": [ "## Chapter 8" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §8.3.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ " import numpy as np\n", " import matplotlib.pyplot as plt \n", " from scipy.integrate import odeint\n", " %matplotlib notebook\n", "\n", " def rhs(Y, t, omega):\n", " y, ydot = Y\n", " return ydot, -omega**2*y\n", "\n", " t_arr = np.linspace(0, 2*np.pi, 101)\n", " y_init = [1, 0]\n", " omega = 2.0\n", " y_arr=odeint(rhs, y_init, t_arr, args=(omega,))\n", " y, ydot = y_arr[:, 0], y_arr[:, 1]\n", "\n", " fig = plt.figure()\n", " ax1 = fig.add_subplot(121)\n", " ax1.plot(t_arr, y, t_arr, ydot)\n", " ax1.set_xlabel('t')\n", " ax1.set_ylabel('y and ydot')\n", " ax2 = fig.add_subplot(122)\n", " ax2.plot(y, ydot)\n", " ax2.set_xlabel('y')\n", " ax2.set_ylabel('ydot')\n", " plt.suptitle(\"Solution curve when omega = %5g\" % omega)\n", " fig.tight_layout()\n", " fig.subplots_adjust(top=0.90)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The next snippet must be run in terminal mode!" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile traj.py\n", "\n", "import numpy as np\n", "import matplotlib.pyplot as plt \n", "from scipy.integrate import odeint\n", "\n", "def rhs(Y, t, omega):\n", " y,ydot=Y\n", " return ydot, -omega**2*y\n", "\n", "t_arr = np.linspace(0, 2*np.pi, 101)\n", "y_init = [1, 0]\n", "omega = 2.0\n", "\n", "fig = plt.figure()\n", "y, ydot = np.mgrid[-3:3:21j, -6:6:21j]\n", "u, v = rhs(np.array([y, ydot]), 0.0, omega)\n", "mag = np.hypot(u, v)\n", "mag[mag==0] = 1.0\n", "plt.quiver(y, ydot, u/mag, v/mag, color='red')\n", "\n", "# Enable drawing of arbitrary number of trajectories\n", "print \"\\n\\n\\nUse mouse to select each starting point\"\n", "print \"Timeout after 30 seconds\"\n", "choice = [(0, 0)]\n", "while len(choice) > 0:\n", " y01 = np.array([choice[0][0], choice[0][1]])\n", " y = odeint(rhs, y01, t_arr, args=(omega,))\n", " plt.plot(y[:, 0], y[:, 1], lw=2)\n", " choice = plt.ginput()\n", "print \"Timed out!\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §8.3.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "from scipy.integrate import odeint\n", "\n", "def rhs(y, t, mu):\n", " return [ y[1], mu*(1-y[0]**2)*y[1] - y[0]]\n", "\n", "def jac(y, t, mu):\n", " return [ [0, 1], [-2*mu*y[0]*y[1]-1, mu*(1-y[0]**2)] ]\n", "\n", "mu = 1.0\n", "t_final = 15.0 if mu < 10 else 4.0*mu\n", "n_points = 1001 if mu < 10 else 1001*mu\n", "t = np.linspace(0,t_final,n_points)\n", "y0 = np.array([2.0, 0.0])\n", "y, info = odeint(rhs, y0, t, args=(mu,), Dfun=jac, full_output=True)\n", "\n", "print \" mu = %g, number of Jacobian calls is %d\" % \\\n", " (mu, info['nje'][-1])\n", "\n", "import matplotlib.pyplot as plt \n", "%matplotlib notebook\n", "plt.plot(y[:,0], y[:,1])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §8.3.4" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "from scipy.integrate import odeint\n", "\n", "def rhs(u, t, beta, rho, sigma):\n", " x, y, z = u\n", " return [sigma*(y-x), rho*x-y-x*z, x*y-beta*z]\n", "\n", "sigma = 10.0\n", "beta = 8.0/3.0\n", "rho1 = 29.0\n", "rho2 = 28.8\n", "\n", "u01 = [1.0, 1.0, 1.0]\n", "u02 = [1.0, 1.0, 1.0]\n", "\n", "t = np.linspace(0.0, 50.0, 10001)\n", "u1 = odeint(rhs, u01, t, args=(beta,rho1,sigma))\n", "u2 = odeint(rhs, u02, t, args=(beta,rho2,sigma))\n", "\n", "x1, y1, z1 = u1[:, 0], u1[:, 1], u1[:, 2]\n", "x2, y2, z2 = u2[:, 0], u2[:, 1], u2[:, 2]\n", "\n", "import matplotlib.pyplot as plt\n", "from mpl_toolkits.mplot3d import Axes3D\n", "%matplotlib notebook\n", "\n", "plt.ion()\n", "fig = plt.figure()\n", "ax=Axes3D(fig)\n", "ax.plot(x1, y1, z1, 'b-')\n", "ax.plot(x2, y2, z2, 'r:')\n", "ax.set_xlabel('x')\n", "ax.set_ylabel('y')\n", "ax.set_zlabel('z')\n", "ax.set_title('Lorenz equations with rho = %g, %g' % (rho1,rho2))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §8.4.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import scikits.bvp1lg.colnew as colnew\n", "\n", "degrees = [2]\n", "boundary_points = np.array([0,np.pi])\n", "tol = 1.0e-8*np.ones_like(boundary_points)\n", "\n", "def fsub(x, Z):\n", " \"\"\"The equations\"\"\"\n", " u, du = Z \n", " return np.array([-u])\n", "\n", "def gsub(Z):\n", " \"\"\"The boundary conditions\"\"\"\n", " u, du = Z\n", " return np.array([u[0], du[1]-1.0])\n", "\n", "solution = colnew.solve(boundary_points, degrees, fsub, gsub,\n", " is_linear=True, tolerances=tol,\n", " vectorized=True, maximum_mesh_size=300)\n", "x = solution.mesh\n", "u_exact = -np.sin(x)\n", "\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "plt.ion()\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "ax.plot(x,solution(x)[:,0],'b.')\n", "ax.plot(x,u_exact,'g-')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §8.4.4" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import scikits.bvp1lg.colnew as colnew\n", "\n", "degrees = [2, 1, 1]\n", "boundary_points = np.array([0.0, 0.0, 1.0, 1.0])\n", "tol = 1.0e-5*np.ones_like(boundary_points)\n", "\n", "def fsub(x, Z):\n", " \"\"\"The equations\"\"\"\n", " u, du, v, w = Z\n", " ddu = -u*v\n", " dv = np.zeros_like(x)\n", " dw = u*u\n", " return np.array([ddu, dv, dw])\n", "\n", "def gsub(Z):\n", " \"\"\"The boundary conditions\"\"\"\n", " u, du, v, w = Z\n", " return np.array([u[0], w[1], u[2], w[3]-1.0])\n", "\n", "\n", "guess_lambda = 100.0\n", "def guess(x):\n", " u = x*(1.0/3.0 - x)*(2.0/3.0 - x)*(1.0 - x)\n", " du = 2.0*(1.0 - 2.0*x)*(1.0 - 9.0*x + 9.0*x*x)/9.0\n", " v = guess_lambda*np.ones_like(x)\n", " w = u*u\n", " Z_guess = np.array([u, du, v, w])\n", " f_guess = fsub(x, Z_guess)\n", " return Z_guess, f_guess\n", "\n", "\n", "solution=colnew.solve(boundary_points, degrees,fsub,gsub,\n", " is_linear=False, initial_guess=guess,\n", " tolerances=tol, vectorized=True,\n", " maximum_mesh_size=300)\n", "x= solution.mesh\n", "u_exact = np.sqrt(2)*np.sin(3.0*np.pi*x)\n", "\n", "\n", "# plot solution\n", "\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "\n", "plt.ion()\n", "plt.plot(x, solution(x)[:, 0], 'b.', x, u_exact, 'g-')\n", "print \"Third eigenvalue is %16.10e .\" % solution(x)[0, 2]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §8.4.5" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import scikits.bvp1lg.colnew as colnew\n", "\n", "degrees = [2, 1, 1]\n", "boundary_points = np.array([0.0, 0.0, 1.0, 1.0])\n", "tol = 1.0e-8*np.ones_like(boundary_points)\n", "\n", "def fsub(x, Z):\n", " \"\"\"The equations\"\"\"\n", " u, du, v, w = Z\n", " ddu = -v*np.exp(u)\n", " dv = np.zeros_like(x)\n", " dw = u*u\n", " return np.array([ddu, dv, dw])\n", "\n", "\n", "def gsub(Z):\n", " \"\"\"The boundary conditions\"\"\"\n", " u, du, v, w = Z\n", " return np.array([u[0], w[1], u[2], w[3]-gamma])\n", "\n", "def guess(x):\n", " u = 0.5*x*(1.0 - x)\n", " du = 0.5*(1.0 - 2.0*x)\n", " v = np.zeros_like(x)\n", " w = u*u\n", " Z_guess = np.array([u, du, v, w])\n", " f_guess = fsub(x, Z_guess)\n", " return Z_guess, f_guess\n", "\n", "solution = guess\n", "gaml = []\n", "laml = []\n", "\n", "for gamma in np.linspace(0.01, 5.01, 1001):\n", " solution = colnew.solve(boundary_points, degrees, fsub, gsub,\n", " is_linear=False, initial_guess=solution,\n", " tolerances=tol, vectorized=True, \n", " maximum_mesh_size=300)\n", " x = solution.mesh\n", " lam = solution(x)[:, 2]\n", " gaml.append(gamma)\n", " laml.append(np.max(lam))\n", "\n", "# plot solution\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "\n", "plt.ion()\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "ax.plot(gaml, laml)\n", "ax.set_xlabel(r'$\\gamma$', size=20)\n", "ax.set_ylabel(r'$\\lambda$', size=20)\n", "ax.grid(b=True)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §8.5.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false, "scrolled": true }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "from pydelay import dde23\n", "\n", "t_final = 50\n", "delay = 2.0\n", "x_initial = 0.1\n", "\n", "equations={'x' : 'x*(1.0-x(t-tau))'}\n", "parameters={'tau' : delay}\n", "dde=dde23(eqns=equations, params=parameters)\n", "dde.set_sim_params(tfinal=t_final, dtmax=1.0,\n", " AbsTol=1.0e-6, RelTol=1.0e-3)\n", "histfunc={'x': lambda t: x_initial} \n", "dde.hist_from_funcs(histfunc, 101)\n", "dde.run()\n", "\n", "t_vis = 0.1*t_final\n", "sol = dde.sample(tstart=t_vis+delay, tfinal=t_final, dt=0.1)\n", "t = sol['t']\n", "x = sol['x']\n", "sold = dde.sample(tstart=t_vis, tfinal=t_final-delay, dt=0.1)\n", "xd = sold['x']\n", "\n", "plt.ion()\n", "fig = plt.figure()\n", "ax1 = fig.add_subplot(121)\n", "ax1.plot(t, x)\n", "ax1.set_xlabel('t')\n", "ax1.set_ylabel('x(t)')\n", "\n", "ax2 = fig.add_subplot(122)\n", "ax2.plot(x, xd)\n", "ax2.set_xlabel('tx')\n", "ax2.set_ylabel('x(t-tau)')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §8.5.4" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false, "scrolled": true }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "from pydelay import dde23\n", "\n", "t_final = 500\n", "a = 2.0\n", "b = 1.0\n", "m = 7.0\n", "delay = 2.0\n", "x_initial = 0.5\n", "\n", "equations={'x' : 'a*x(t-tau)/(1.0+pow(x(t-tau), m))- b*x'}\n", "parameters={'a': a, 'b' : b, 'm': m, 'tau' : delay}\n", "dde=dde23(eqns=equations, params=parameters)\n", "dde.set_sim_params(tfinal=t_final, dtmax=1.0,\n", " AbsTol=1.0e-6, RelTol=1.0e-3)\n", "histfunc={'x': lambda t: x_initial} \n", "dde.hist_from_funcs(histfunc, 101)\n", "dde.run()\n", "\n", "t_vis = 0.95*t_final\n", "sol = dde.sample(tstart=t_vis+delay, tfinal=t_final, dt=0.1)\n", "t = sol['t']\n", "x = sol['x']\n", "sold = dde.sample(tstart=t_vis, tfinal=t_final-delay, dt=0.1)\n", "xd = sold['x']\n", "\n", "plt.ion()\n", "fig = plt.figure()\n", "ax1 = fig.add_subplot(121)\n", "ax1.plot(t, x)\n", "ax1.set_xlabel('t')\n", "ax1.set_ylabel('x(t)')\n", "\n", "ax2 = fig.add_subplot(122)\n", "ax2.plot(x, xd)\n", "ax2.set_xlabel('tx')\n", "ax2.set_ylabel('x(t-tau)')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §8.6.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import numpy.random as npr\n", "\n", "T = 1\n", "N = 500\n", "t, dt = np.linspace(0, T, N+1, retstep=True)\n", "dW = npr.normal(0.0, np.sqrt(dt), N+1)\n", "dW[0] = 0.0\n", "W = np.cumsum(dW)\n", "\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "plt.ion()\n", "\n", "plt.plot(t,W)\n", "plt.xlabel('t')\n", "plt.ylabel('W(t)')\n", "#plt.title('Sample Wiener Process',weight='bold',size=16)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §8.6.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import numpy.random as npr\n", "\n", "T = 1\n", "N = 1000\n", "M = 5000\n", "t, dt = np.linspace(0, T, N+1, retstep=True)\n", "dW = npr.normal(0.0, np.sqrt(dt), (M,N+1))\n", "dW[ : ,0] = 0.0\n", "W = np.cumsum(dW, axis=1)\n", "U = np.exp(- 0.5*W)\n", "Umean = np.mean(U, axis=0)\n", "Uexact = np.exp(t/8)\n", "\n", "import matplotlib.pyplot as plt \n", "%matplotlib notebook\n", "plt.ion()\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "\n", "ax.plot(t, Umean, 'b-', label=\"mean of %d paths\" % M)\n", "ax.plot(t, Uexact, 'r-', label=\"exact \" + r'$\\langle U\\rangle$')\n", "for i in range(5):\n", " ax.plot(t,U[i, : ], '--')\n", " \n", "ax.set_xlabel('t')\n", "ax.set_ylabel('U')\n", "ax.legend(loc='best')\n", "\n", "maxerr = np.max(np.abs(Umean-Uexact))\n", "print \"With %d paths and %d intervals the max error is %g\" % \\\n", " (M, N, maxerr)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 8.6.4" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import numpy.random as npr\n", "\n", "# Set up grid\n", "T=1.0\n", "N=1000\n", "t, dt = np.linspace(0, T, N+1, retstep=True)\n", "\n", "# Get Brownian motion\n", "dW = npr.normal(0.0, np.sqrt(dt), N+1)\n", "dW[0] = 0.0\n", "W = np.cumsum(dW)\n", "\n", "# Equation parameters and functions\n", "lamda = 2.0\n", "mu = 1.0\n", "Xzero = 1.0\n", "def a(X): return lamda*X\n", "def b(X): return mu*X\n", "def bd(X): return mu*np.ones_like(X)\n", "\n", "# Analytic solution\n", "Xanal = Xzero*np.exp((lamda-0.5*mu*mu)*t+mu*W)\n", "\n", "# Milstein solution\n", "Xmil = np.empty_like(t)\n", "Xmil[0] = Xzero\n", "for n in range(N):\n", " Xmil[n+1] = Xmil[n] + dt*a(Xmil[n]) + dW[n+1]*b(Xmil[n]) + 0.5*(\n", " b(Xmil[n])*bd(Xmil[n])*(dW[n+1]**2-dt))\n", "\n", "# Plot solution\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "plt.ion()\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111)\n", "\n", "ax.plot(t, Xanal, 'b-', label='analytic')\n", "ax.plot(t, Xmil, 'g-.', label='Milstein')\n", "ax.legend(loc='best')\n", "ax.set_xlabel('t')\n", "ax.set_ylabel('X(t)')\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import numpy.random as npr\n", "\n", "# Problem definition\n", "M = 1000 # Number of paths sampled\n", "P = 6 # Number of discretizations \n", "T = 1 # Endpoint of time interval \n", "N = 2**12 # Finest grid size\n", "dt = 1.0*T/N\n", " \n", "# Problem parameters\n", "lamda = 2.0 \n", "mu = 1.0\n", "Xzero = 1.0\n", "\n", "def a(X): return lamda*X\n", "def b(X): return mu*X\n", "def bd(X): return mu*np.ones_like(X)\n", "\n", "# Build the Brownian paths.\n", "dW = npr.normal(0.0, np.sqrt(dt), (M,N+1))\n", "dW[ : , 0] = 0.0\n", "W = np.cumsum(dW, axis=1)\n", "\n", "# Build the exact solutions at the ends of the paths\n", "ones = np.ones(M)\n", "Xexact = Xzero*np.exp((lamda-0.5*mu*mu)*ones+mu*W[ : , -1])\n", "Xemerr = np.empty((M,P))\n", "Xmilerr = np.empty((M,P))\n", "\n", "# Loop over refinements\n", "for p in range(P):\n", " R = 2**p\n", " L = N/R # must be an integer!\n", " Dt = R*dt\n", " Xem = Xzero*ones\n", " Xmil = Xzero*ones\n", " Wc = W[ : , : :R]\n", " for j in range(L): # integration\n", " deltaW = Wc[ : , j+1]-Wc[ : , j]\n", " Xem += Dt*a(Xem) + deltaW*b(Xem)\n", " Xmil += Dt*a(Xmil) + deltaW*b(Xmil)+ \\\n", " 0.5*b(Xmil)*bd(Xmil)*(deltaW**2-Dt)\n", " Xemerr[ : ,p] = np.abs(Xem-Xexact)\n", " Xmilerr[ : ,p] = np.abs(Xmil-Xexact)\n", "\n", "Dtvals = dt*np.array([2**p for p in range(P)])\n", "lDtvals = np.log10(Dtvals)\n", "Xemerrmean = np.mean(Xemerr,axis=0)\n", "\n", "# Do some plotting\n", "import matplotlib.pyplot as plt\n", "%matplotlib notebook\n", "plt.ion()\n", "\n", "fig = plt.figure()\n", "ax = fig.add_subplot(111\n", " )\n", "ax.plot(lDtvals, np.log10(Xemerrmean), 'bo')\n", "ax.plot(lDtvals, np.log10(Xemerrmean), 'b:', label='EM actual')\n", "ax.plot(lDtvals, 0.5*np.log10(Dtvals), 'b-.', label='EM theoretical')\n", "Xmilerrmean = np.mean(Xmilerr, axis=0)\n", "ax.plot(lDtvals, np.log10(Xmilerrmean), 'bo')\n", "ax.plot(lDtvals, np.log10(Xmilerrmean),' b--', label='Mil actual')\n", "ax.plot(lDtvals, np.log10(Dtvals), 'b-', label='Mil theoretical')\n", "ax.legend(loc='best')\n", "ax.set_xlabel(r'$\\log_{10}\\Delta t$', size=16)\n", "ax.set_ylabel(r'$\\log_{10}\\left(\\langle|X_n-X(\\tau)|\\rangle\\right)$', size=16)\n", " \n", "emslope = ((np.log10(Xemerrmean[-1])-np.log10(Xemerrmean[0])) /\n", " (lDtvals[-1]-lDtvals[0]))\n", "print 'Empirical EM slope is %g' % emslope\n", "milslope = ((np.log10(Xmilerrmean[-1])- np.log10(Xmilerrmean[0])) / \n", " (lDtvals[-1]-lDtvals[0]))\n", "print 'Empirical MIL slope is %g' % milslope" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Chapter 9" ] }, { "cell_type": "markdown", "metadata": { "collapsed": true }, "source": [ "### §9.4" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "from scipy.fftpack import diff\n", "\n", "def fd(u):\n", " \"\"\" Return 2*dx*finite-difference x-derivative of u. \"\"\"\n", " ud = np.empty_like(u)\n", " ud[1:-1]=u[2: ]-u[ :-2]\n", " ud[0]=u[1]-u[-1]\n", " ud[-1]=u[0]-u[-2]\n", " return ud\n", "\n", "for N in [4,8,16,32,64,128,256]:\n", " dx=2.0*np.pi/N\n", " x=np.linspace(0,2.0*np.pi,N,endpoint=False)\n", " u=np.exp(np.sin(x))\n", " du_ex=np.cos(x)*u\n", " du_sp=diff(u)\n", " du_fd=fd(u)/(2.0*dx)\n", " err_sp=np.max(np.abs(du_sp-du_ex))\n", " err_fd=np.max(np.abs(du_fd-du_ex))\n", " print \"N=%3d, err_sp=%.1e err_fd=%.1e\" % (N, err_sp,err_fd)\n", " " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §9.5" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "from scipy.fftpack import diff\n", "from scipy.integrate import odeint\n", "import matplotlib.pyplot as plt \n", "from mpl_toolkits.mplot3d import Axes3D\n", "%matplotlib notebook\n", "\n", "def u_exact(t,x):\n", " \"\"\" Exact solution. \"\"\"\n", " return np.exp(np.sin(x-2.0*np.pi*t))\n", "\n", "def rhs(u,t):\n", " \"\"\" Return rhs. \"\"\"\n", " return -2.0*np.pi*diff(u)\n", "\n", "N=32\n", "x=np.linspace(0,2.0*np.pi,N,endpoint=False)\n", "u0 =u_exact(0.0,x)\n", "t_initial=0.0\n", "t_final=64.0*np.pi\n", "t=np.linspace(t_initial,t_final,101)\n", "sol=odeint(rhs,u0,t,mxstep=5000)\n", "\n", "plt.ion()\n", "fig=plt.figure()\n", "ax=Axes3D(fig)\n", "t_gr,x_gr=np.meshgrid(x,t)\n", "ax.plot_surface(t_gr,x_gr,sol,cstride=2,rstride=8,alpha=0.1)\n", "ax.elev,ax.azim=47,-137\n", "ax.set_xlabel('x')\n", "ax.set_ylabel('t')\n", "ax.set_zlabel('u')\n", "\n", "u_ex=u_exact(t[-1],x)\n", "err=np.abs(np.max(sol[-1,: ]-u_ex))\n", "print \"With %d Fourier modes the final error = %g\" % (N,err)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §9.6" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "\n", "def u(x): return 1.0/(1.0+25.0*x**2)\n", "\n", "N=20\n", "x_grid=np.linspace(-1.0,1.0,N+1)\n", "u_grid=u(x_grid)\n", "z=np.polyfit(x_grid,u_grid,N)\n", "p=np.poly1d(z)\n", "x_fine=np.linspace(-1.0,1.0,5*N+1)\n", "u_fine=p(x_fine)\n", "u_fine" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Introduction to f2py" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Run the following line in the notebook, after changing the address to that of your copy of _f2py_, so that you can access _f2py_ from your notebook in a compact manner." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "%alias f2py '/Users/john/Library/Enthought/Canopy_64bit/User/bin/f2py'" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §9.7.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile norm3.f90\n", "! File:norm3.f90 A simple subroutine in f90\n", "\n", "subroutine norm(u,v,w,s)\n", "real(8), intent(in) :: u,v,w\n", "real(8), intent(out) :: s\n", "s=sqrt(u*u+v*v+w*w)\n", "end subroutine norm" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "!gfortran norm3.f90" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "f2py -c norm3.f90 --f90flags=-O3 -m normv3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "import normv3\n", "normv3?" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "normv3.norm(3, 4, 5)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile norm3.f\n", "C FILE NORM3.F A SIMPLE SUBROUTINE IN F77\n", "\n", " SUBROUTINE NORM(U,V,W,S)\n", " REAL*8 U,V,W,S\n", "Cf2py intent(in) U,V,W\n", "Cf2py intent(out) S\n", " S=SQRT(U*U+V*V+W*W)\n", " END" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "!rm normv3.so" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "f2py -c norm3.f -m normv3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import normv3\n", "normv3?" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "normv3.norm(3,4,5)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "rm normv3.so" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §9.7.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile normn.f\n", "C FILE NORMN.F EXAMPLE OF N-DIMENSIONAL NORM \n", "\n", " SUBROUTINE NORM(U,S,N)\n", " INTEGER N\n", " REAL*8 U(N)\n", " REAL*8 S\n", "Cf2py intent(in) N\n", "Cf2py intent(in) U\n", "Cf2py depend(U) N\n", "Cf2py intent(out) S\n", " REAL*8 SUM\n", " INTEGER I\n", " SUM=0\n", " DO 100 I=1,N\n", " 100 SUM=SUM + U(I)*U(I)\n", " S=SQRT(SUM)\n", " END" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "f2py -c normn.f --f77flags=-O3 -m normn" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import normn\n", "normn.norm([3,4,5,6,7])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §9.7.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile modmat.f\n", "C FILE MODMAT.F MODIFYING A MULTIDIMENSIONAL ARRAY\n", "\n", " SUBROUTINE MMAT(A,B,C,R,M,N)\n", " INTEGER M,N\n", " REAL*8 A(M,N),B(M),C(N),R(M,N)\n", "Cf2py intent(in) M,N\n", "Cf2py intent(in) A,B,C\n", "Cf2py depend(A) M,N\n", "Cf2py intent(out) R\n", " INTEGER I,J\n", " DO 10 I=1,M\n", " DO 10 J=1,N\n", " 10 R(I,J)=A(I,J)\n", " DO 20 I=1,M\n", " 20 R(I,1)=R(I,1)+B(I)\n", " DO 30 J=1,N\n", " 30 R(1,J)=R(1,J)+C(J)\n", " END\n", " " ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "f2py -c modmat.f --f77flags=-O3 -m modmat" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import modmat\n", "\n", "a=np.array([[1,2,3],[4,5,6]],dtype='float',order='F')\n", "b=np.array([10,20],dtype='float')\n", "c=np.array([7,11,13],dtype='float')\n", "r=modmat.mmat(a,b,c)\n", "r" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §9.8" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The code from Appendix B" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile cheb.f\n", "\n", " SUBROUTINE CHEBPTS(X, N)\n", "C CREATE N+1 CHEBYSHEV DATA POINTS IN X\n", " IMPLICIT REAL*8 (A-H,O-Z)\n", " DIMENSION X(0:N)\n", "Cf2py intent(in) N\n", "Cf2py intent(out) X\n", "Cf2py depend(N) X\n", " PI = 4.D0*ATAN(1.D0)\n", " DO 10 I=0,N \n", " 10 X(I) = -COS(PI*I/N)\n", " RETURN\n", " END\n", "\n", " SUBROUTINE FFT (A,B,IS,N,ID)\n", "C-- +------------------------------------------------------------------\n", "C-- | A CALL TO FFT REPLACES THE COMPLEX DATA VALUES A(J) + i B(J), \n", "C-- | J=0,1,... ,N-1 WITH THEIR TRANSFORM\n", "C-- |\n", "C-- | 2 i ID PI K J / N \n", "C-- | SUM (A(K) + iB(K))e , J=0,1,.. .,N-1\n", "C-- | K=0..N-1\n", "C-- |\n", "C-- | INPUT AND OUTPUT PARAMETERS \n", "C-- | A ARRAY A (0: *), REAL PART OF DATA/TRANSFORM \n", "C-- | B ARRAY B (0: *), IMAGINARY PART OF DATA/TRANSFORM\n", "C-- | INPUT PARAMETERS \n", "C-- | IS SPACING BETWEEN CONSECUTIVE ELEMENTS IN A AND B\n", "C-- | (USE IS=+1 FOR ELEMENTS STORED CONSECUTIVELY) \n", "C-- | N NUMBER OF DATA VALUES, MUST BE A POWER OF TWO\n", "C-- | ID USE +1 OR -1 TO SPECIFY DIRECTION OF TRANSFORM\n", "C-- +------------------------------------------------------------------\n", " IMPLICIT REAL*8 (A-H,O-Z) \n", " DIMENSION A(0:*),B(0:*)\n", "Cf2py intent(in, out) A, B\n", "Cf2py intent(in) IS, ID\n", "Cf2py integer intent(hide), depend(A) :: N\n", " J=0 \n", "C--- APPLY PERMUTATION MATRIX ----\n", " DO 20 I=0,(N-2)*IS,IS \n", " IF (I.LT.J) THEN\n", " TR = A(J) \n", " A(J) = A(I) \n", " A(I) = TR\n", " TI = B(J)\n", " B(J) = B(I)\n", " B(I) = TI \n", " ENDIF\n", " K = IS*N/2\n", " 10 IF (K.LE.J) THEN\n", " J = J-K \n", " K = K/2\n", " GOTO 10 \n", " ENDIF \n", " 20 J = J+K \n", "C--- PERFORM THE LOG2 N MATRIX-VECTOR MULTIPLICATIONS ---\n", " S = 0.0D0 \n", " C = -1.0D0 \n", " L = IS\n", " 30 LH=L \n", " L = L+L\n", " UR = 1.0D0 \n", " UI = 0.0D0\n", " DO 50 J=0,LH-IS,IS \n", " DO 40 I=J,(N-1)*IS,L\n", " IP = I+LH \n", " TR = A(IP)*UR-B(IP)*UI \n", " TI = A(IP)*UI+B(IP)*UR \n", " A(IP) = A(I)-TR \n", " B(IP) = B(I)-TI \n", " A(I) = A(I)+TR\n", " 40 B(I) = B(I)+TI \n", " TI = UR*S+UI*C \n", " UR = UR*C-UI*S\n", " 50 UI=TI \n", " S = SQRT (0.5D0*(1.0D0-C))*ID\n", " C = SQRT (0.5D0*(1.0D0+C))\n", " IF (L.LT.N*IS) GOTO 30\n", " RETURN \n", " END\n", "\n", " SUBROUTINE FCT (A,X,N,B) \n", "C-- +---------------------------------------------------------------\n", "C-- | A CALL TO FCT PLACES IN B(0:N) THE COSINE TRANSFORM OF THE \n", "C-- | VALUES IN A(0:N)\n", "C-- | \n", "C-- | B(J) = SUM C(K)*A(K)*COS(PI*K*J/N) , J=0,1,...,N, K=0..N\n", "C-- |\n", "C-- | WHERE C(K) = 1.0 FOR K=0,N, C(K) =2.0 FOR K=1,2,...,N-1\n", "C__ |\n", "C-- | INPUT PARAMETERS: \n", "C-- | A A(0:N) ARRAY WITH INPUT DATA \n", "C-- | X X(0:N) ARRAY WITH CHEBYSHEV GRID POINT LOCATIONS\n", "C-- | X(J) = -COS(PI*J/N) , J=0,1,...,N \n", "C-- | N SIZE OF TRANSFORM - MUST BE A POWER OF TWO\n", "C-- | OUTPUT PARAMETER \n", "C-- | B B(0:N) ARRAY RECEIVING TRANSFORM COEFFICIENTS\n", "C-- | (MAY BE IDENTICAL WITH THE ARRAY A)\n", "C-- |\n", "C-- +---------------------------------------------------------------\n", " IMPLICIT REAL*8 (A-H,O-Z) \n", " DIMENSION A(0:*),X(0:*),B(0:*)\n", "Cf2py intent(in) A, X\n", "Cf2py intent(out) B\n", "Cf2py integer intent(hide), depend(A) :: N \n", " N2 = N/2 \n", " A0 = A(N2-1)+A(N2+1) \n", " A9 = A(1) \n", " DO 10 I=2,N2-2,2\n", " A0 = A0+A9+A(N+1-I) \n", " A1 = A( I+1)-A9 \n", " A2 = A(N+1-I)-A(N-1-I) \n", " A3 = A(I)+A(N-I) \n", " A4 = A(I)-A(N-I) \n", " A5 = A1-A2 \n", " A6 = A1+A2 \n", " A1 = X(N2-I) \n", " A2 = X(I) \n", " A7 = A1*A4+A2*A6 \n", " A8 = A1*A6-A2*A4 \n", " A9 = A(I+1) \n", " B(I ) = A3+A7 \n", " B(N-I) = A3-A7 \n", " B(I+1 ) = A8+A5\n", " 10 B(N+1-I) = A8-A5\n", " B(1) = A(0)-A(N) \n", " B(0) = A(0)+A(N)\n", " B(N2 ) = 2.D0*A(N2)\n", " B(N2+1) = 2.D0*(A9-A(N2+1)) \n", " CALL FFT(B(0),B(1),2,N2,1)\n", " A0 = 2.D0*A0\n", " B(N) = B(0)-A0\n", " B(0) = B(0)+A0 \n", " DO 20 I=1,N2-1\n", " A1 = 0.5 D0 *(B(I)+B(N-I)) \n", " A2 = 0.25D0/X(N2+I)*(B(I)-B(N-I)) \n", " B(I ) = A1+A2\n", " 20 B(N-I) = A1-A2 \n", " RETURN\n", " END\n", "\n", " \n", " SUBROUTINE FROMCHEB (A,X,N,B)\n", " IMPLICIT REAL*8 (A-H,O-Z)\n", " DIMENSION A(0:N),X(0:N),B(0:N)\n", "Cf2py intent(in) A, X\n", "Cf2py intent(out) B\n", "Cf2py integer intent(hide), depend(A) :: N\n", " B(0) = A(0) \n", " A1 = 0.5D0\n", " DO 10 I=1,N-1 \n", " A1 = -A1\n", " 10 B(I) = A1*A(I)\n", " B(N) = A(N) \n", " CALL FCT(B,X,N,B)\n", " RETURN \n", " END\n", "\n", "\n", " SUBROUTINE TOCHEB (A,X,N,B)\n", " IMPLICIT REAL*8 (A-H,O-Z)\n", " DIMENSION A(0:N),X(0:N),B(0:N) \n", "Cf2py intent(in) A, X\n", "Cf2py intent(out) B\n", "Cf2py integer intent(hide), depend(A) :: N\n", " CALL FCT(A,X,N,B)\n", " B1 = 0.5D0/N\n", " B(0) = B(0)*B1\n", " B(N) = B(N)*B1 \n", " B1 = 2.D0*B1\n", " DO 10 I=1,N-1 \n", " B1 = -B1\n", " 10 B(I) = B(I)*B1 \n", " RETURN\n", " END\n", "\n", " \n", " SUBROUTINE DIFFCHEB (A,N,B)\n", " IMPLICIT REAL*8 (A-H,O-Z)\n", " DIMENSION A(0:N),B(0:N) \n", "Cf2py intent(in) A\n", "Cf2py intent(out) B\n", "Cf2py integer intent(hide), depend(A) :: N\n", " A1 = A(N) \n", " A2 = A(N-1) \n", " B(N) = 0.D0\n", " B(N-1) = 2.D0*N*A1 \n", " A1 = A2 \n", " A2 = A(N-2) \n", " B(N-2) = 2.D0*(N-1)*A1 \n", " DO 10 I=N-2,2,-1\n", " A1 = A2 \n", " A2 = A(I-1)\n", " 10 B(I-1) = B(I+1)+2.D0*I*A1\n", " B(0) = 0.5D0*B(2)+A2 \n", " RETURN\n", " END\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "f2py -c cheb.f --f77flags=-O3 -m cheb" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import cheb\n", "print cheb.__doc__\n", "\n", "for N in [8,16,32,64]:\n", " x=cheb.chebpts(N)\n", " ex=np.exp(x)\n", " u=np.sin(ex)\n", " dudx=ex*np.cos(ex)\n", " \n", " uc=cheb.tocheb(u,x)\n", " duc=cheb.diffcheb(uc)\n", " du=cheb.fromcheb(duc,x)\n", " err=np.max(np.abs(du-dudx))\n", " print \"With N = %d error is %e\" % (N,err)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §9.9.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "from scipy.integrate import odeint\n", "import matplotlib.pyplot as plt \n", "from mpl_toolkits.mplot3d import Axes3D\n", "import cheb\n", "%matplotlib notebook\n", "\n", "c=1.0\n", "mu=0.1\n", "\n", "N=64\n", "x=cheb.chebpts(N)\n", "tau1=N**2\n", "tau2=N**2\n", "t_initial=-2.0\n", "t_final=2.0\n", "\n", "def u_exact(t,x):\n", " \"\"\" Exact kink solution of Burgers' equation. \"\"\"\n", " return c*(1.0+np.tanh(c*(c*t-x)/(2*mu)))\n", "\n", "def mu_ux_exact(t,x):\n", " \"\"\" Return mu*du/dx for exact solution. \"\"\"\n", " arg=np.tanh(c*(c*t-x)/(2*mu))\n", " return 0.5*c*c*(arg*arg-1.0)\n", "\n", "def f(x):\n", " \"\"\" Return initial data. \"\"\"\n", " return u_exact(t_initial,x)\n", "\n", "def g1(t):\n", " \"\"\" Return function needed at left boundary. \"\"\"\n", " return (u_exact(t,-1.0)**2-mu_ux_exact(t,-1.0))\n", "\n", "def g2(t):\n", " \"\"\" Return function needed at right boundary. \"\"\"\n", " return mu_ux_exact(t,1.0)\n", "\n", "def rhs(u,t):\n", " \"\"\" Return du/dt. \"\"\"\n", " u_cheb=cheb.tocheb(u,x)\n", " ux_cheb=cheb.diffcheb(u_cheb)\n", " uxx_cheb=cheb.diffcheb(ux_cheb)\n", " ux=cheb.fromcheb(ux_cheb,x)\n", " uxx=cheb.fromcheb(uxx_cheb,x)\n", " dudt=-u*ux+mu*uxx\n", " dudt[0]-=tau1*(u[0]**2-mu*ux[0]-g1(t))\n", " dudt[-1]-=tau2*(mu*ux[-1]-g2(t))\n", " return dudt\n", "\n", "t=np.linspace(t_initial,t_final,81)\n", "u_initial=f(x)\n", "sol=odeint(rhs,u_initial,t,rtol=1.0e-12,atol=1.0e-12,mxstep=5000)\n", "xg,tg=np.meshgrid(x,t)\n", "ueg=u_exact(tg,xg)\n", "err=sol-ueg\n", "print \"With %d points error is %e\" % (N,np.max(np.abs(err)))\n", "\n", "plt.ion()\n", "fig=plt.figure()\n", "ax=Axes3D(fig)\n", "ax.plot_surface(xg,tg,sol,rstride=1,cstride=2,alpha=0.1)\n", "ax.set_xlabel('x',style='italic')\n", "ax.set_ylabel('t',style='italic')\n", "ax.set_zlabel('u',style='italic')\n" ] }, { "cell_type": "markdown", "metadata": { "collapsed": true }, "source": [ "## Chapter 10" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §10.2.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt \n", "%matplotlib notebook\n", "\n", "N=16\n", "omega=2.0/3.0\n", "max_its=8\n", "\n", "x=np.linspace(0,1,N+1)\n", "s1=np.sin(np.pi*x)\n", "s13=np.sin(13*np.pi*x)\n", "e_old=s1+s13/3.0\n", "e=np.zeros_like(e_old)\n", "\n", "plt.ion()\n", "plt.plot(e_old)\n", "for it in range(max_its):\n", " e[1:-1]=(1.0-omega)*e_old[1:-1]+\\\n", " 0.5*omega*(e_old[0:-2]+e_old[2: ])\n", " plt.plot(e)\n", " e_old=e.copy()\n", "plt.xlabel('The 17 grid points')\n", "plt.ylabel('The first %d iterates of the error' % N)" ] }, { "cell_type": "markdown", "metadata": { "collapsed": true }, "source": [ "### §10.4.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile util.py\n", "# File: util.py: useful 2D utilities\n", "\n", "import numpy as np\n", "\n", "def l2_norm(a,h):\n", " \"\"\" Returns L2-norm of a-values with equal spacing h. \"\"\"\n", " return h*np.sqrt(np.sum(a**2))\n", "\n", "def prolong_lin(a):\n", " \"\"\" Return linear prolongation of 2-dimensional array a. \"\"\"\n", " pshape=(2*np.shape(a)[0]-1, 2*np.shape(a)[1] -1)\n", " p=np.empty(pshape,float)\n", " p[0: :2,0: :2]=a[0: ,0: ]\n", " p[1:-1:2,0: :2]=0.5*(a[0:-1,0: ]+a[1: ,0: ])\n", " p[0: :2,1:-1:2]=0.5*(a[0: ,0:-1]+a[0: ,1: ])\n", " p[1:-1:2,1:-1:2]=0.25*(a[0:-1,0:-1]+a[1: ,0:-1]+\n", " a[0:-1,1: ]+a[1: ,1:])\n", " return p\n", "\n", "def restrict_hw(a):\n", " \"\"\" Return half-weighted restriction of 2-dimensional array a.\"\"\"\n", " rshape=(np.shape(a)[0]/2+1, np.shape(a)[1]/2+1)\n", " r=np.empty(rshape,float)\n", " r[1:-1,1:-1]=0.5*a[2:-1:2,2:-1:2]+ \\\n", " 0.125*(a[2:-1:2,1:-2:2]+a[2:-1:2,3: :2]+\n", " a[1:-2:2,2:-1:2]+a[3: :2,2:-1:2])\n", " r[0,0: ]=a[0,0: :2]\n", " r[-1,0: ]=a[-1,0: :2]\n", " r[0: ,0]=a[0: :2,0]\n", " r[0: ,-1]=a[0: :2,-1]\n", " return r\n", "\n", "#---------------------------------------------------------\n", "if __name__=='__main__':\n", " a=np.linspace(1,81,81)\n", " b=a.reshape(9,9)\n", " c=restrict_hw(b)\n", " d=prolong_lin(c)\n", " print \" original grid:\\n\",b\n", " print \"\\n with spacing 1 its norm is %g\" % l2_norm(b,1)\n", " print \"\\n restricted grid:\\n\",c\n", " print \"\\n prolonged restricted grid:\\n\", d\n", " " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §10.4.2" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile smooth.py\n", "# File: smooth.py: problem-dependent 2D utilities.\n", "\n", "import numpy as np\n", "\n", "def get_lhs(u,h2):\n", " \"\"\" Return discretized operator L(u). h2=h**2 for spacing h. \"\"\"\n", " w=np.zeros_like(u)\n", " w[1:-1,1:-1]=(u[0:-2,1:-1]+u[2: ,1:-1]+\n", " u[1:-1,0:-2]+u[1:-1,2: ]-\n", " 4*u[1:-1,1:-1])/h2+u[1:-1,1:-1]**2\n", " \n", " return w\n", "\n", "\n", "def gs_rb_step(v,f,h2):\n", " \"\"\" Carry out single Gauss-Seidel iteration step on v.\n", " f is source term, h2 is square of grid spacing. \n", " \"\"\"\n", " u=v.copy()\n", " res=np.empty_like(v)\n", " \n", " res[1:-1:2,1:-1:2]=(u[0:-2:2,1:-1:2]+u[2: :2,1:-1:2]+\n", " u[1:-1:2,0:-2:2]+u[1:-1:2,2: :2]-\n", " 4*u[1:-1:2,1:-1:2])/h2 +\\\n", " u[1:-1:2,1:-1:2]**2-f[1:-1:2,1:-1:2]\n", " u[1:-1:2,1:-1:2]-=res[1:-1:2,1:-1:2]/(\n", " -4.0/h2+2*u[1:-1:2,1:-1:2])\n", " \n", " res[2:-2:2,2:-2:2]=(u[1:-3:2,2:-2:2]+u[3:-1:2,2:-2:2]+\n", " u[2:-2:2,1:-3:2]+u[2:-2:2,3:-1:2]-\n", " 4*u[2:-2:2,2:-2:2])/h2 +\\\n", " u[2:-2:2,2:-2:2]**2-f[2:-2:2,2:-2:2]\n", " u[2:-2:2,2:-2:2]-=res[2:-2:2,2:-2:2]/(\n", " -4.0/h2+2*u[2:-2:2,2:-2:2])\n", " \n", "\n", " res[2:-2:2,1:-1:2]=(u[1:-3:2,1:-1:2]+u[3:-1:2,1:-1:2]+\n", " u[2:-2:2,0:-2:2]+u[2:-2:2,2: :2]-\n", " 4*u[2:-2:2,1:-1:2])/h2 +\\\n", " u[2:-2:2,1:-1:2]**2-f[2:-2:2,1:-1:2]\n", " u[2:-2:2,1:-1:2]-=res[2:-2:2,1:-1:2]/(\n", " -4.0/h2+2*u[2:-2:2,1:-1:2])\n", " \n", "\n", " res[1:-1:2,2:-2:2]=(u[0:-2:2,2:-2:2]+u[2: :2,2:-2:2]+\n", " u[1:-1:2,1:-3:2]+u[1:-1:2,3:-1:2]-\n", " 4*u[1:-1:2,2:-2:2])/h2 +\\\n", " u[1:-1:2,2:-2:2]**2-f[1:-1:2,2:-2:2]\n", " u[1:-1:2,2:-2:2]-=res[1:-1:2,2:-2:2]/(\n", " -4.0/h2+2*u[1:-1:2,2:-2:2])\n", " \n", " return u\n", "\n", "def solve(rhs):\n", " \"\"\" Return exact solution on the coarsest 3x3 grid.\"\"\"\n", " h=0.5\n", " u=np.zeros_like(rhs)\n", " fac=2.0/h**2\n", " dis=np.sqrt(fac**2+rhs[1,1])\n", " u[1,1]=-rhs[1,1]/(fac+dis)\n", " return u\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### §10.4.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%writefile grid.py\n", "# File: grid.py: linked grid structures and associated algorithms.\n", "\n", "import numpy as np\n", "from util import l2_norm, restrict_hw, prolong_lin\n", "from smooth import gs_rb_step, get_lhs, solve\n", "\n", "class Grid:\n", " \"\"\"\n", " A Grid contains the structures and algorithms for a \n", " given level, together with a pointer to a coarser grid.\n", " \"\"\"\n", " def __init__(self,name,steplength,u,f,coarser=None):\n", " self.name=name\n", " self.co=coarser # pointer to a coarser grid\n", " self.h=steplength # step length h\n", " self.h2=steplength**2 # h**2\n", " self.u=u # improved variable array\n", " self.f=f # right hand side array\n", " \n", " def __str__(self):\n", " \"\"\" Generate an information string about this level. \"\"\"\n", " sme='Grid at %s with steplength = %0.4g\\n' % (\n", " self.name,self.h)\n", " if self.co:\n", " sco='Coarser grid with name %s\\n' % self.co.name\n", " else:\n", " sco='No coarser grid\\n'\n", " return sme+sco\n", " \n", " def smooth(self,nu):\n", " \"\"\" \n", " Carry out Newton-Raphson/Gauss-Seidel red-black\n", " iteration u --> u, nu times.\n", " \"\"\"\n", " print 'Relax in %s for %d times' % (self.name,nu)\n", " v=self.u.copy()\n", " for i in range(nu):\n", " v=gs_rb_step(v,self.f,self.h2)\n", " self.u=v\n", " \n", " def fas_v_cycle(self,nu1,nu2):\n", " \"\"\" Recursive implementation of (nu1,nu2) FAS V-cycle. \"\"\"\n", " print 'FAS-V-cycle called for grid at %s\\n' % self.name\n", " # Initial smoothing\n", " self.smooth(nu1)\n", " if self.co:\n", " # There is a coarser grid\n", " self.co.u=restrict_hw(self.u)\n", " # Get residual\n", " res=self.f-get_lhs(self.u,self.h2)\n", " # Get coarser f\n", " self.co.f=restrict_hw(res)+get_lhs(self.co.u,self.co.h2)\n", " oldc=self.co.u\n", " # Get new coarse solution\n", " newc=self.co.fas_v_cycle(nu1,nu2)\n", " # Correct current u\n", " self.u+=prolong_lin(newc-oldc)\n", " self.smooth(nu2)\n", " return self.u\n", " \n", " def fmg_fas_v_cycle(self,nu0,nu1,nu2):\n", " \"\"\" Recursive implementation of FMG-FAS_V-cycle. \"\"\"\n", " print 'FMG-FAS-V-cycle called for grid at %s\\n' % self.name\n", " if not self.co:\n", " # Coarsest grid\n", " self.u=solve(self.f)\n", " else:\n", " # Restrict f\n", " self.co.f=restrict_hw(self.f)\n", " # Use recursion to get coarser u\n", " self.co.u=self.co.fmg_fas_v_cycle(nu0,nu1,nu2)\n", " # Prolong to current u\n", " self.u=prolong_lin(self.co.u)\n", " for it in range(nu0):\n", " self.u=self.fas_v_cycle(nu1,nu2)\n", " return self.u\n", "\n", " \n", " " ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false, "scrolled": true }, "outputs": [], "source": [ "# File rungrid.py: runs the multigrid programme.\n", "\n", "import numpy as np\n", "from grid import Grid\n", "from smooth import get_lhs\n", "\n", "n_grids=5\n", "size_init=2**n_grids+1\n", "size=size_init-1\n", "h=1.0/size\n", "foo=[] # Container list for grids\n", "\n", "for k in range(n_grids): \n", " # Set up list of grids\n", " u=np.zeros((size+1,size+1),float)\n", " f=np.zeros_like(u)\n", " name='Level '+str(n_grids-1-k)\n", " temp=Grid(name,h,u,f)\n", " foo.append(temp)\n", " size/=2\n", " h*=2\n", "\n", "for k in range(1,n_grids):\n", " # Set up coarser Grid links\n", " foo[k-1].co=foo[k]\n", "\n", "# Check that the initial construction works\n", "for k in range(n_grids):\n", " print foo[k]\n", "\n", "# Set up data for the Numerical Recipes problem\n", "u_init=np.zeros((size_init,size_init))\n", "f_init=np.zeros_like(u_init)\n", "f_init[size_init/2,size_init/2]=2.0\n", "foo[0].u=u_init\n", "foo[0].f=f_init\n", "\n", "foo[0].fmg_fas_v_cycle(1,1,1)\n", "\n", "# As a check, get lhs for the final grid\n", "lhs=get_lhs(foo[0].u,foo[0].h2)\n", "print \"max abs lhs = \", np.max(np.abs(lhs))\n", "\n", "import matplotlib.pyplot as plt \n", "from mpl_toolkits.mplot3d import Axes3D\n", "%matplotlib notebook\n", "\n", "plt.ion()\n", "fig1=plt.figure()\n", "ax1=Axes3D(fig1)\n", "\n", "xx,yy=np.mgrid[0:1:1j*size_init,0:1:1j*size_init]\n", "ax1.plot_surface(xx,yy,1000*foo[0].u,rstride=1,cstride=1,alpha=0.2)\n", "ax1.set_xlabel('x',style='italic')\n", "ax1.set_ylabel('y',style='italic')\n", "ax1.set_zlabel('1000*u',style='italic')\n", "\n", "fig2=plt.figure()\n", "ax2=Axes3D(fig2)\n", "ax2.plot_surface(xx,yy,lhs,rstride=1,cstride=1,alpha=0.2)\n", "ax1.set_xlabel('x',style='italic')\n", "ax2.set_ylabel('y',style='italic')\n", "ax2.set_zlabel('lhs',style='italic')\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [] } ], "metadata": { "anaconda-cloud": {}, "kernelspec": { "display_name": "Python 2", "language": "python", "name": "python2" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 2 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython2", "version": "2.7.11" } }, "nbformat": 4, "nbformat_minor": 0 }