.. !split .. _basics:func:branching: Functions and branching ======================= .. _basics:func: Functions --------- .. index:: function (in Python) Since :math:`s(t)` is a function in mathematics, it is convenient to have it as a function in Python too: .. code-block:: python def s(t): return v0*t + 0.5*a*t**2 v0 = 0.2 a = 4 value = s(3) # Call the function Note that * functions start with the keyword ``def`` * statements belonging to the function must be indented * function input is represented by arguments (separated by comma if more than one) * function output is returned to the calling code .. index:: global variables In this program, ``v0`` and ``a`` are *global variables*, which must be defined before calling the ``s(t)`` function, otherwise ``v0`` and ``a`` becomes undefined variables in the expression inside the function. Instead of having ``v0`` and ``a`` as global variables, we may let them be function arguments: .. code-block:: python def s(t, v0, a): return v0*t + 0.5*a*t**2 value = s(3, 0.2, 4) # Call the function # More readable call value = s(t=3, v0=0.2, a=4) .. index:: keyword arguments .. index:: default arguments Function arguments may also be declared to have a default value, allowing us to drop the argument in the call if the default value is appropriate. Such arguments are called *keyword arguments* or *default arguments*. Here is an example: .. code-block:: python def s(t, v0=1, a=1): return v0*t + 0.5*a*t**2 value = s(3, 0.2, 4) # specify new v0 and a value = s(3) # rely on v0=1 and a=1 value = s(3, a=2) # rely on v0=1 value = s(3, v0=2) # rely on a=1 value = s(t=3, v0=2, a=2) # specify everything value = s(a=2, t=3, v0=2) # any sequence allowed .. note:: * Arguments without the argument name are called *positional arguments*. * Positional arguments must always be listed before the keyword arguments in the function and in any call. * The sequence of the keyword arguments can be arbitrary. * If argument names are used for all arguments (as in the last line above) the sequence of the arguments is not important. .. index:: vectorization .. index:: vectorized functions .. admonition:: Vectorized functions Applying the function ``s(t, v0, a)`` to an array ``t`` can be done in two ways: .. code-block:: python # Scalar code: work with one element at a time for i in range(len(t)): s_values[i] = s(t_values[i], v0, a) # Vectorized code: apply s to the entire array s_values = s(t_values, v0, a) For the last line to work, the function ``s`` must contain statements that work correctly when ``t`` is an array argument. The advantage of vectorized code is that we avoid a loop in Python. Instead, we carry out mathematical operations on entire arrays, e.g., ``v0*t`` and ``a*t**2``. Technically, such (binary) operations are executed as loops in very fast (compiled) C code. .. index:: multiple return values A function can return more than one value, say :math:`s(t)` and the velocity :math:`s'(t)=v_0 + at`: .. code-block:: python def movement(t, v0, a): s = v0*t + 0.5*a*t**2 v = v0 + a*t return s, v s_value, v_value = movement(t=0.2, v0=2, a=4) When :math:`n` values are returned, we list :math:`n` variables on the left-hand side in the call. .. admonition:: Python functions return only one object Even when we return several values, as in ``return s, v``, actually only one object is returned. The ``s`` and ``v`` values are packed together in a *tuple* object (which is very similar to a list). .. code-block:: python >>> def f(x): ... return x+1, x+2, x+3 ... >>> r = f(3) # Store all three return values in one object r >>> print r (4, 5, 6) >>> type(r) # What type of object is r? >>> print r[1] 5 Tuples are constant lists, so you can index them as lists, but you cannot change the contents (``append`` or ``del`` is illegal). .. _basics:formula:piecewise: A more general mathematical formula ----------------------------------- The formula :ref:`(1) ` arises from the basic differential equations in kinematics: .. _Eq:basics:vdef: .. math:: \tag{2} v = \frac{ds}{dt},\quad s(0)=s_0, .. _Eq:basics:adef: .. math:: \tag{3} a = \frac{dv}{dt},\quad v(0)=v_0\thinspace . Given any acceleration :math:`a(t)`, we can solve for :math:`s(t)` through integration. First, we integrate to find :math:`v(t)`: .. math:: \int_0^t a(t)dt = \int_0^t \frac{dv}{dt} dt, which gives .. math:: v(t) = v_0 + \int_0^t a(t)dt\thinspace . Then we integrate again over :math:`[0,t]` to find :math:`s(t)`: .. _Eq:basics:s:at: .. math:: \tag{4} s(t) = s_0 + v_0t + \int_0^t\left( \int_0^t a(t)dt \right) dt\thinspace . Suppose we have some constant acceleration :math:`a_0` in :math:`[0,t_1]` and no acceleration thereafter. We find .. _Eq:basics:seq2: .. math:: \tag{5} s(t) = \left\lbrace\begin{array}{ll} s_0 + v_0 t + \frac{1}{2}a_0 t^2,& t\leq t_1\\ s_0 + v_0t_1 + \frac{1}{2}a_0 t_1^2 + a_0t_1(t-t_1),& t> t_1 \end{array}\right. To implement a function like :ref:`(5) `, we need to branch into one type of code (formula) if :math:`t\leq t_1` and another type of code (formula) if :math:`t>t_1`. This is called *branching* and the *if test* is the primary construction to use. .. _basics:if: If tests -------- .. index:: branching .. index:: if test .. index:: boolean expression An if test has the structure .. code-block:: python if condition: else: Here, ``condition`` is a boolean expression that evaluates to ``True`` or ``False``. For the piecewisely defined function :ref:`(5) ` we would use an if test in the implementation: .. code-block:: python if t <= t1: s = v0*t + 0.5*a0*t**2 else: s = v0*t + 0.5*a0*t1**2 + a0*t1*(t-t1) The ``else`` part can be omitted when not needed. Several branches are also possible: .. code-block:: python if condition1: elif condition2: elif condition3: else: A Python function implementing the mathematical function :ref:`(5) ` reads .. code-block:: python def s_func(t, v0, a0, t1): if t <= t1: s = v0*t + 0.5*a0*t**2 else: s = v0*t + 0.5*a0*t1**2 + a0*t1*(t-t1) return s To plot this :math:`s(t)`, we need to compute points on the curve. Trying to call ``s_func`` with an array ``t`` as argument will not work because of the if test. In general, functions with if tests will not automatically work with arrays because a test like ``if t <= t1`` evaluates to an if applied to a boolean array (``t <= t1`` becomes a boolean array, not just a boolean). One solution is to compute the function values one by one in a loop such that the ``s`` function is always called with a scalar value for ``t``. Appropriate code is .. code-block:: python n = 201 # No of t values for plotting t1 = 1.5 t = np.linspace(0, 2, n+1) s = np.zeros(n+1) for i in range(len(t)): s[i] = s_func(t=t[i], v0=0.2, a0=20, t1=t1) Relevant statements for plotting are now .. code-block:: python plt.plot(t, s, 'b-') plt.plot([t1, t1], [0, s_func(t=t1, v0=0.2, a0=20, t1=t1)], 'r--') plt.xlabel('t') plt.ylabel('s') plt.savefig('myplot.png') plt.show() .. index:: vectorizing if tests .. admonition:: Vectorization of functions with if tests To vectorize the computation of the array of :math:`s(t)` values, i.e., avoid a loop where ``s_func`` is called for each ``t[i]`` element, we must completely rewrite the if test. There are two methods: the ``numpy.where`` function and array indexing. Using the ``where`` function ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ A vectorized if-else test can be coded as .. code-block:: python s = np.where(condition, s1, s2) Here, ``condition`` is an array of boolean values, and ``s[i] = s1[i]`` if ``condition[i]`` is ``True``, and otherwise ``s[i] = s2[i]``. Our example then becomes .. code-block:: python s = np.where(t <= t1, v0*t + 0.5*a0*t**2, v0*t + 0.5*a0*t1**2 + a0*t1*(t-t1)) Note that ``t <= t1`` with array ``t`` and scalar ``t1`` results in a boolean array ``b`` where ``b[i] = t[i] <= t1``. Using array indexing ~~~~~~~~~~~~~~~~~~~~ It is possible to index a subset of indices in an array ``s`` using a boolean array ``b``: ``s[b]``. This construction picks out all the elements ``s[i]`` where ``b[i]`` is ``True``. On the right-hand side we can then assign some array expression ``expr`` of the same length as ``s[b]``: .. code-block:: python s[b] = (expr)[b] Our example can utilize this technique with ``b`` as ``t <= t1`` and ``t > t1``: .. code-block:: python s = np.zeros_like(t) # Make s as zeros, same size & type as t s[t <= t1] = (v0*t + 0.5*a0*t**2)[t <= t1] s[t > t1] = (v0*t + 0.5*a0*t1**2 + a0*t1*(t-t1))[t > t1] Array view versus array copying ------------------------------- Arrays are usually large and most array libraries have carefully designed functionality for avoiding unnecessary copying of large of amounts of data. If you do a simple assignment, .. code-block:: python a = np.linspace(1, 5, 5) b = a ``b`` becomes a just *view* of ``a``: the variables ``b`` and ``a`` point to the same data. In other languages one may say that ``b`` is a pointer or reference to the array. This means that changing ``a`` changes ``b`` and vice versa: .. code-block:: python array([ 1., 2., 3., 4., 5.]) >>> b[0] = 5 # changes a[0] to 5 >>> a array([ 5., 2., 3., 4., 5.]) >>> a[1] = 9 # changes b[1] to 9 >>> b array([ 5., 9., 3., 4., 5.]) Similarly, ``b[1:-1]`` gives a view to a subset of ``a`` and no copy of data. If we want to change ``b`` without affecting ``a``, a copy of the array is required: .. code-block:: python >>> c = a.copy() # copy all elements to new array c >>> c[0] = 6 # a is not changed >>> a array([ 1., 2., 3., 4., 5.]) >>> c array([ 6., 2., 3., 4., 5.]) >>> b array([ 5., 2., 3., 4., 5.]) Note that ``b`` remains unchanged by the change in ``c`` since the ``b`` and ``c`` variables now refer to different data. Copying of the elements of a sub-array is also done by the ``copy()`` method: ``b = a[1:-1].copy()``. .. admonition:: Understand the difference between view and copy Mathematical errors and/or inefficient code may easily arise from confusion between array view and array copying. Make sure you understand the implication of any variable assignment to arrays! .. code-block:: python a = np.array([-1, 4.5, 8, 9]) b = a[1:-2] # view: changes in a are reflected in b b = a[1:-2].copy() # copy: changes in a are not reflected in b Linear systems -------------- .. index:: linear systems .. index:: sparse matrix .. index:: tridigonal matrix Solving linear systems of the form :math:`Ax=b` with a matrix :math:`A` and vectors :math:`x` and :math:`b` is a frequently encountered task. A sample :math:`2\times 2` system can be coded as .. code-block:: python import numpy as np A = np.array([[1., 2], [4, 2]]) b = np.array([1., -2]) x = np.linalg.solve(A, b) Many scientific computing problems involves large matrices with special *sparse* structures. Significant memory and computing time can be saved by utilizing *sparse matrix* data structures and associated algorithms. Python has sparse matrix package `scipy.sparse `__ that can be used to construct various types of sparse matrices. The particular function ``scipy.sparse.spdiags`` can construct a sparse matrix from a few vectors representing the diagonals in the matrix (all the diagonals must have the same length as the dimension of their sparse matrix, consequently some elements of the diagonals are not used: the first :math:`k` elements are not used of the :math:`k` super-diagonal, and the last :math:`k` elements are not used of the :math:`-k` sub-diagonal). Here is an example of constructing a *tridiagonal* matrix: .. code-block:: python >>> import numpy as np >>> N = 6 >>> diagonals = np.zeros((3, N)) # 3 diagonals diagonals[0,:] = np.linspace(-1, -N, N) diagonals[1,:] = -2 diagonals[2,:] = np.linspace(1, N, N) >>> import scipy.sparse >>> A = scipy.sparse.spdiags(diagonals, [-1,0,1], N, N, format='csc') >>> A.toarray() # look at corresponding dense matrix [[-2. 2. 0. 0. 0. 0.] [-1. -2. 3. 0. 0. 0.] [ 0. -2. -2. 4. 0. 0.] [ 0. 0. -3. -2. 5. 0.] [ 0. 0. 0. -4. -2. 6.] [ 0. 0. 0. 0. -5. -2.]] Solving a linear system with a tridiagonal coefficient matrix requires a special function with a special algorithm to take advantage of the matrix' structure. The next example chooses a solution ``x``, computes the corresponding right-hand side :math:`b=Ax` using the sparse matrix vector product associated with the sparse matrix :math:`A`, and then solves :math:`Ax=b`: .. code-block:: python >>> x = np.linspace(-1, 1, N) # choose solution >>> b = A.dot(x) # sparse matrix vector product >>> import scipy.sparse.linalg >>> x = scipy.sparse.linalg.spsolve(A, b) >>> print x [-1. -0.6 -0.2 0.2 0.6 1. ] Although we can easily check that ``x`` is correctly computed, we can do another check by converting the linear system to its dense counterpart: .. code-block:: python >>> A_d = A.toarray() # corresponding dense matrix >>> b = np.dot(A_d, x) # standard matrix vector product >>> x = np.linalg.solve(A_d, b) # standard Ax=b algorithm >>> print x [-1. -0.6 -0.2 0.2 0.6 1. ]