User-Defined Function

Transcription

1 ENGR (Socolofsky) Week 11 Python scripts In the lecture this week, we are continuing to learn powerful things that can be done with userdefined functions. In several of the examples, we consider a quadratic equation in the form f(x) = ax 2 + bx + c (1) and these examples could be done in nearly the same way with any other function f(x). The set of scripts below contain a module of functions and a set of examples/tests contained within an if name is main : block. One thing you will notice in this script is that we import modules as needed throughout the code. This is useful to avoid importing large packages into memory before they are needed. This also isolates these complex packages from the rest of the code, where they are not needed. Whether all packages are imported at the top of the code or as needed is a matter of taste, and the examples below follows the approach of keeping numerical packages together with the code that performs the numerical analysis. A User-Defined Function The function quadratic eqn() is a user-defined function to compute the quadratic equation at a point x given the coefficients a, b, and c. This function can take x as an integer, real number, np.array object, or a complex number and returns the corresponding value of f(x). If we want to use a root-finding method to find the roots of the quadratic equation when the roots are complex, we need to adapt the quadratic eqn() function to return the real and imaginary parts of x and f(x) as separate, real, floating-point numbers. There are different ways to do this, and in the script below, we choose to return these as real numbers in a 2 1 np.array. For example, if the x-coordinate is i, then we would store this as x = np.array([7.2, -8.6]) The function quadratic complex() in the script below takes x in this format as input, converts it to a complex number using the built-in Python function complex(), passes this x value to quadratic eqn() to compute the value of f(x), and returns the possibly complex result as a 2 1 np.array object. B Passing Functions to Functions A key concept needed to make functions that solve mathematical problems more general is the ability to pass a user-defined function that contains the particular problem of interest to the function that performs the numerical solution. We demonstrate passing a function to a function using the function plot fun(), defined in the script below. This function takes a user-defined function in the form func(x, params) as input together with the lower and upper bounds on the x-axis for

2 plotting, an optional input that contains the params for func, and an optional figure number as input. We use tuples to pass multiple parameters through the variable params to the user-defined function func. Here, we demonstrate using quadratic eqn() as the function for func. plot fun() then computes an array of x and f(x) data and plots the data over the desired range in the given figure number. C Examples In the script below, we use scipy to solve several, typical mathematical problems in science and engineering. C.1 Integration As an example of integration, we use scipy.integrate.quad() to evaluate the integral b f(x)dx a using the quadratic equation in quadratic eqn() as an example for f(x). The function quadratic eq int() returns the analytical solution for the integral of a quadratic equation, and the script tests below compare the result returned by quad() to this analytical solution. C.2 Root-Finding There are multiple methods to find roots available within scipy, and here we demonstrate the function scipy.optimize.fsolve() for the problem f(x) = 0 (3) where we use the quadratic equation as an example of f(x). In the example in the script, we look for complex roots; hence, we use the function quadratic complex() with fsolve(). The function quadratic roots() computes the roots using the quadratic formula, and this analytical solution is compared with the roots returned by fsolve(). (2) C.3 Interpolation As an example of interpolation, we use some sparse data for the temperature and salinity as a function of depth in the ocean. There are interpolation methods in scipy and numpy, and we demonstrate scipy.interpolate.interp1d(). The data and the example are all contained within the examples at the bottom of the script, and this example does not use any of the user-defined function defined in the main body of the script module. C.4 Ordinary Differential Equations There are two main solver function for ordinary differential equations (ODEs) in scipy, and here we demonstrate scipy.integrate.odeint(). The general form of a one-dimensional ODE is dy dt = f(y, t) (4) 2

3 subject to the initial condition y(0) = y 0. In the example below, we consider f(y, t) = ky. To solve ODEs numerically, we need to write a function that returns f(y, t); in the sample script this is diff eq rhs(). We compare the results of odeint() to the analytical solution, given by y(t) = y 0 exp( kt) (5) The sample script that performs all of these examples and contains each of the required userdefined functions is given below in the script fun script.py 3

4 Listing for the script fun script.py 1 # fun_script.py 2 # 3 # This program demonstrates passing functions to functions and using SciPy 4 # 5 # S. Socolofsky 6 # ENGR # October import numpy as np # Fuction definitions def quadratic_eqn (x, a, b, c): 14 """ 15 Function to compute the quadratic equation Computes the quadratic equation, defined by: 18 f( x) = ax ^2 + bx + c 19 at a given value of x for given coefficients a, b, and c Parameters x : float or np. array 24 Value ( s) of the independent variable 25 a, b, c : float 26 Coefficients of the quadratic equation Returns f( x) : float or np. array 31 Value ( s) of the corresponding dependent variable """ 34 return a * x** 2 + b * x + c def quadratic_ complex (x, a, b, c): 37 """ 38 Compute complex values for the quadratic equation Compute the quadratic equation for values of x for which x and / or f( x) is 41 complex Parameters x : np. array 46 Value of the independent variable stored in an array containing the 47 real and imaginary parts separately, e. g., np. array ([ x. real, x. imag ]) 48 a, b, c : float 49 Coefficients of the quadratic equation 4

5 50 51 Returns x : np. array 54 Array containing the real and imaginary parts of f( x) Notes Uses the function quadratic_ eqn (x, a, b, c) to compute the quadratic 59 equation """ 62 fc = quadratic_eqn ( complex (x[0], x[1]), a, b, c) 63 return np. array ([ fc.real, fc. imag ]) def quadratic_eqn_int ( lim_a, lim_b, a, b, c): 67 """ 68 Function to compute the integral of the quadratic equation Computes the integral of the quadratic equation given by: 71 f( x) = ax ^2 + bx + c 72 with the limits of integration spanning ( lim_a, lim_b ) Parameters lim_a : float 77 Lower bound for the limits of integration 78 lim_b : float 79 Upper bound for the limits of integration 80 a, b, c : float 81 Coefficients of the quadratic equation Returns g( x) : float 86 The value of the definite integral of the quadratic equation with the 87 given limits of integration """ 90 fun = lambda x: a * x** 3 / 3. + b * x** 2 / 2. + c * x return fun ( lim_b ) - fun ( lim_a ) def quadratic_ roots (a, b, c): 96 """ 97 Compute the roots of the quadratic equation Computes the roots of the quadratic equation, defined by: 100 f( x) = ax ^2 + bx + c 101 using the analytical solution 102 xr = (-b +/ - sqrt (b^2-4ac )) / (2a) 5

6 103 where the solution is returned as a float ( complex as needed ) Parameters a, b, c : float 108 Coefficients of the quadratic equation Returns xr : np. array 113 Array containing the two roots of the quadratic equation, each as 114 an array containing the real and imaginary parts of the root """ 117 # Compute the discriminant 118 d = b**2-4. * a * c # Find the roots assuming a is not zero 121 if d >= 0.: 122 x0 = (-b - np. sqrt (d)) / (2 * a) 123 x1 = (-b + np. sqrt (d)) / (2 * a) 124 else : 125 x0 = complex (-b / (2 * a), -np. sqrt (-d) / (2 * a)) 126 x1 = complex (-b / (2 * a), np. sqrt (-d) / (2 * a)) # Return the roots as np. arrays 129 return (np. array ([x0.real, x0. imag ]), np. array ([x1.real, x1. imag ])) def diff_eq_rhs (y, t, k): 133 """ 134 Function containing the RHS of a differential equation Function to solve the ODE given by 137 dy/dt = -ky 138 using the scipy. integrate. odeint () function. This function contains 139 the right - hand - side ( RHS ) of this ODE Parameters y : np. array 144 Array of dependent variable value ( s) 145 t : float 146 Value of the independent variable 147 k : float 148 Die - off rate constant for the growth model dy/ dx = - ky in dimensions 149 of 1 / t Returns dy/dt : np. array 154 Array of value ( s) returning the RHS of the ODE to solve. This is the 155 local slope of the solution y( t). 6

7 """ 158 return - k * y def plot_fun (fun, x0, xf, params =(), fnum =1): 162 """ 163 Plot a function f( x) over the range ( x0, xf) Plot a function f( x) given by the name fun over the range of independent 166 variables between x0 and xf Parameters fun : function 171 Name of a Python function to compute the dependent variable. Should 172 be in the form fun (x, params ), where x can accept a np. array input 173 and the additional params are optional. 174 x0 : float 175 Minimum independent variable value in range to plot 176 xf : float 177 Maximum independent variable value in range to plot 178 params : tuple, default () 179 Optional parameters to pass to the function fun 180 fnum : int 181 Figure number to plot """ 184 # Create an array of independent variables 185 xp = np. linspace (x0, xf, num =250) # Compute the corresponding dependent variables 188 yp = fun (xp, * params ) # Create the desired plot 191 import matplotlib. pyplot as plt 192 plt. figure ( fnum ) 193 plt. gcf (). clear () 194 plt. plot (xp, yp, 'b-') 195 plt. xlabel ('x') 196 plt. ylabel ('f(x)') 197 plt. grid ( True ) 198 plt. show () # Testing procedures if name is ' main ': """ 206 Test procedures that execute if this script is run from the command prompt Examples using user - defined functions to perform typical simulation 7

8 209 exercises, including integration, root finding, interpolation, and 210 solving ordinary differential equations The examples below include the following : ) Compute the quadratic equation using a user - defined function ) Plot a user - defined function by passing a function and optional 216 arguments to another function that plots the results ) Integrate a quadratic equation over fixed limits of integration 218 a.) using scipy. integrate. quad () 219 b.) using the analytical solution ) Find the roots of the quadratic equation ( including imaginary 221 roots ): 222 a.) using scipy. optimize. fsolve () 223 b.) using the quadratic formula ) Interpolate multidimensional data using 225 a.) scipy. interpolate. interp 1d() for linear interpolation 226 b.) scipy. interpolate. interp 1d() for cubic interpolation ) Solve the differential equation dy/ dt = - ky; y( 0) = y0 228 a.) using scipy. integrate. odeint () 229 b.) using the analytical solution """ # Examples using the quadratic equation x = a, b, c = 1.0, 1.5, print (' Testing the quadratic equation function for the equation :') 238 print (' f(x) = ' + str (a) + 'x^2 + ' + str (b) + 'x + ' + str (c)) # Compute a single point 241 print ('\ nthe solution at x =', x, 'is:') 242 print (' ' + str ( quadratic_eqn (x, a, b, c ))) # Compute using a np. array of 5 points 245 x = np. linspace (-3., 3., num =5) 246 print ('\ nthe solution at x =', x, 'is:') 247 print (' ' + str ( quadratic_eqn (x, a, b, c ))) # Plot this quadratic equation over the range (- 3, 3) 250 fnum = print ('\ nplotting the quadratic equation over the range ') 252 print (' (' + str (x. min ()) + ', ' + str (x. max ()) + ') in Figure ' str ( fnum )) 254 plot_fun ( quadratic_eqn, -3, 3, (a, b, c), fnum ) # Example integrating a quadratic equation over the range (- 3, 3) from scipy. integrate import quad 258 lower_ limit = upper_ limit = scipy_ integral = quad ( quadratic_eqn, lower_limit, upper_limit, 261 args =(a, b, c ))[ 0] 8

9 262 print ('\ nscipy returns the integral of this quadratic equation with ') 263 print (' limits of integration spanning (' + str ( lower_ limit ) + ', ' str ( upper_limit ) + ') as:') 265 print (' ' + str ( scipy_ integral )) # Compare to the true value of the integral 268 true_ integral = quadratic_ eqn_ int ( lower_limit, upper_limit, *( a, b, c)) 269 print ('\ nthe true value of this integral is:') 270 print (' ' + str ( true_ integral )) # Example finding roots of a quadratic equation from scipy. optimize import fsolve 274 x0 = np. array ([ -1., -1.]) 275 sol = fsolve ( quadratic_ complex, x0, args =( a, b, c)) 276 print ('\ nscipy returns the roots of this quadratic equation near ') 277 print (' x0 = ' + str (x0[0]) + ' + (' + str (x0[1]) + ')i as:') 278 print (' ' + str ( sol [0]) + ' + (' + str ( sol [1]) + ')i') 279 x0 = np. array ([ -1., 1.]) 280 sol = fsolve ( quadratic_ complex, x0, args =( a, b, c)) 281 print (' And Scipy returns the roots of this quadratic equation near ') 282 print (' x0 = ' + str (x0[0]) + ' + (' + str (x0[1]) + ')i as:') 283 print (' ' + str ( sol [0]) + ' + (' + str ( sol [1]) + ')i') # Compare to the analytical solution for the roots of a quadratic equation 286 sol = quadratic_ roots (a, b, c) 287 print ('\ nthe true values for these roots are :') 288 print (' ' + str ( sol [0][0]) + ' + (' + str ( sol [0][1]) + ')i') 289 print (' ' + str ( sol [1][0]) + ' + (' + str ( sol [1][1]) + ')i') # Example using interpolation z = np. array ([0., 10., 50., 100., 200., 1000.]) 293 T = np. array ([23.4, 18.7, 17.9, 16.7, 10.5, 9.9]) 294 S = np. array ([32.5, 33.5, 33.7, 33.8, 34.5, 34.51]) 295 print ('\ nconsider the temperature ( deg C) and salinity ( psu ) data ') 296 print (' versus depth ( m) given by:') 297 print (' z :', z) 298 print (' T :', T) 299 print (' S :', S) # Create an interpolation function that has z as independent variable and 302 # T and S as dependent variables. 303 from scipy. interpolate import interp 1d 304 data = np. vstack ((T, S)) 305 interp_ linear = interp 1d(z, data ) 306 interp_ cubic = interp 1d(z, data, kind =' cubic ') # Compute points using interpolation 309 print ('\ nthe temperature and salinity at z = 75 m is:') 310 print (' Linear interpolation :', interp_ linear ( 75.)) 311 print (' Cubic interpolation : ', interp_ cubic ( 75.)) # Plot the profiles 314 zp = np. linspace (z. min (), z. max (), num =250) 9

10 315 yp_lin = interp_ linear ( zp) 316 yp_cub = interp_ cubic ( zp) 317 import matplotlib. pyplot as plt 318 fnum = print ('\ nplotting this temperature profile using linear and cubic ') 320 print (' interpolation in Figure ' + str ( fnum )) 321 plt. figure ( fnum ) 322 plt. gcf (). clear () 323 plt. plot (T, z, 'o') 324 plt. plot ( yp_lin [0,:], zp, 'b-') 325 plt. plot ( yp_cub [0,:], zp, 'g-') 326 plt. xlabel ('Temperature ( deg C)') 327 plt. ylabel ('Depth (m)') 328 plt. gca (). invert_yaxis () 329 plt. grid ( True ) 330 plt. legend (( ' Data ', ' Linear Interpolation ', ' Cubic Interpolation ')) 331 plt. show () # Example solving an ordinary differential equation ( ODE ) t = np. linspace (0., 10., num =5) 335 y0 = k = # Use scipy to solve the ODE dy/ dt = - ky subject to the initial condition 339 # y( 0) = from scipy. integrate import odeint 341 sol = odeint ( diff_eq_rhs, y0, t, args =( k,)) 342 print ('\ nthe Scipy solution to the differential equation :') 343 print (' dy/dt = -ky ') 344 print (' with initial condition y( 0) = ', y0, ' is:') 345 print (' t:', t) 346 print (' y:', sol [:,0]) # The analytical solution is y = y0 exp (- kt )... report this 349 print ('\ nthe analytical solution for this case is:') 350 print (' t:', t) 351 print (' y:', y0 * np.exp (-k * t)) 10

11 Code execution for fun script.py yields Testing the quadratic equation function for the equation: f(x) = 1.0x^ x The solution at x = 2.0 is: 9.0 The solution at x = [ ] is: [ ] Plotting the quadratic equation over the range (-3.0, 3.0) in Figure 1 Scipy returns the integral of this quadratic equation with limits of integration spanning (-3.0, 3.0) as: 30.0 The true value of this integral is: 30.0 Scipy returns the roots of this quadratic equation near x0 = (-1.0)i as: ( )i And Scipy returns the roots of this quadratic equation near x0 = (1.0)i as: ( )i The true values for these roots are: ( )i ( )i Consider the temperature (deg C) and salinity (psu) data versus depth (m) given by: z : [ ] T : [ ] S : [ ] The temperature and salinity at z = 75 m is: Linear interpolation: [ ] Cubic interpolation: [ ] 11

12 Plotting this temperature profile using linear and cubic interpolation in Figure 2 The Scipy solution to the differential equation: dy/dt = -ky with initial condition y(0) = 1.0 is: t: [ ] y: [ ] The analytical solution for this case is: t: [ ] y: [ ] 12

13 The two figures produced by this script are shown in Figures 1 and f(x) x Figure 1. Plot of the quadratic equation over the range ( 3, 3) using quadratic eqn() and plot fun() in the script fun script.py. 13

14 0 200 Data Linear Interpolation Cubic Interpolation Depth (m) Temperature (deg C) Figure 2. Plot of data for temperature as a function of depth in the ocean using the raw data, linear interpolation, and cubic interpolation based on scipy.interpolate.interp1d() as demonstrated in fun script.py. 14