More advanced implementation

Let us extend the function above to also include the quadratic damping formulas. The \( F(t) \) function in (1) is required to be available as an array in the solver_linear_damping function, but now we will allow for more flexibility: the F argument may either be a Python function F(t) or a Python array. In addition, we add some checks that variables have correct values or are of correct type.

import numpy as np

def solver(I, V, m, b, s, F, t, damping='linear'):
    """
    Solve m*u'' + f(u') + s(u) = F for time points in t.
    u(0)=I and u'(0)=V,
    by a central finite difference method with time step dt.
    If damping is 'linear', f(u')=b*u, while if damping is
    'quadratic', we have f(u')=b*u'*abs(u').
    s(u) is a Python function, while F may be a function
    or an array (then F[i] corresponds to F at t[i]).
    """
    N = t.size - 1              # No of time intervals
    dt = t[1] - t[0]            # Time step
    u = np.zeros(N+1)           # Result array
    b = float(b); m = float(m)  # Avoid integer division

    # Convert F to array
    if callable(F):
        F = F(t)
    elif isinstance(F, (list,tuple,np.ndarray)):
        F = np.asarray(F)
    else:
        raise TypeError(
            'F must be function or array, not %s' % type(F))

    u[0] = I
    if damping == 'linear':
        u[1] = u[0] + dt*V + dt**2/(2*m)*(-b*V - s(u[0]) + F[0])
    elif damping == 'quadratic':
        u[1] = u[0] + dt*V + \ 
               dt**2/(2*m)*(-b*V*abs(V) - s(u[0]) + F[0])
    else:
        raise ValueError('Wrong value: damping="%s"' % damping)

    for n in range(1,N):
        if damping == 'linear':
            u[n+1] = (2*m*u[n] + (b*dt/2 - m)*u[n-1] +
                      dt**2*(F[n] - s(u[n])))/(m + b*dt/2)
        elif damping == 'quadratic':
            u[n+1] = (2*m*u[n] - m*u[n-1] + b*u[n]*abs(u[n] - u[n-1])
                      - dt**2*(s(u[n]) - F[n]))/\ 
                      (m + b*abs(u[n] - u[n-1]))
    return u, t

Dissection of the code

Two types of import. This time we replace from numpy import *, which imports over 500 variables and functions, by import numpy as np, which just imports one variable, the package numpy, here under the nickname np. All variables and functions in numpy must now be reached via the prefix np., as in np.zeros, and np.linspace, and np.pi (for \( \pi \)). The advantage of the prefix is that we clearly see where functionality comes from. The disadvantage is that mathematical formulas like \( \sin (\pi x) \) must be written np.sin(np.pi*x). We can always perform an explicit import of some names, like sin and pi, to write the formula as sin(pi*x):

import numpy as np
from numpy import sin, pi

def myfunction(x):
    return sin(pi*x)

Another disadvantage with the np. prefix is that names are no longer (almost) the same as in Matlab. Many will therefore prefer to do from numpy import * and skip the prefix.

Doc strings. The string, enclosed in triple double-quotes, right after the function definition, is a doc string used for documenting the function. Various tools can extract function definitions and doc strings to automatically produce nicely typeset manuals.

Avoiding integer division. At one line we explicitly convert b and m to float variables. This is not strictly necessary, but if we supply b=2 and m=4, a computation like b/m will give zero as result because both b and m are then integers and b/m implies integer division, not the mathematical division of real numbers (this is not a special feature of Python - all languages with a strong heritage from C invoke integer division if both operands in a division are integers). We need to make sure that at least one of the operands in a division is a real number (float) to ensure the intended mathematical operation. Looking at the formulas, there is never a problem with integer division in our implementation, because dt is computed from t, which has float elements (linspace makes float elements by default), and all divisions in our code involve dt as one of the operands. However, future edits may alter the way formulas are written, so to be on the safe side we ensure that real input parameters are float objects.

Flexible variable type. As mentioned, we allow F to be either a function or an array. In the former case, we convert F to an array such that the rest of the code can assume that F is indeed an array. It is easy to check the type of variables in Python. The test if callable(F) is true if the object F can be called as a function.

Checking correct variable type. To test that F is an array, we can use isinstance(F, np.ndarray) (ndarray is the name of the array type in numpy). In the code above we allow that F can also be a list or a tuple. Running F through the asarray function makes an array out of F in the cases where F is a list or tuple. The final else: clause takes care of the situation where F is neither a function, nor an object that can easily be converted to an array, and an error message is issued. More precisely, we raise a TypeError exception indicating that we have encountered a wrong type. The error message will contain the type of F as obtained from type(F). Instead of using the syntax isinstance(F, list) we may test type(F) == list or even type(F) in (list,tuple,np.ndarray).

We could simplify the if-else test involving F to just the two lines

if callable(F):
    F = F(t)

if we are sure that F is either a function or a numpy array. Should the user send in something else for F, Python will encounter a run-time error when trying to index F as in F[0] (in the statement computing u[1]).

To be absolutely safe, we should test that the other arguments are of right type as well. For example,

if not isinstance(I, (float,int)):
    raise TypeError('V must be float or int, not %s' % type(V))

and similar for V, m, b, while s must be tested by callable(s), t must be np.ndarray, and damping must be str.

Calling the function. Below is a simple example on how the solver function can be called to solve the differential equation problem \( mu'' + bu' + ku = A\sin\pi t \), \( u(0)=I \), \( u'(0)=0 \):

import numpy as np
from numpy import sin, pi  # for nice math

def F(t):
    # Sinusoidal bumpy road
    return A*sin(pi*t)

def s(u):
    return k*u

A = 0.25
k = 2
t = np.linspace(0, 20, 2001)
u, t = solver(I=0, V=0, m=2, b=0.05, s=s, F=F, t=t)

# Show u(t) as a curve plot
import matplotlib.pyplot as plt
plt.plot(t, u)
plt.show()

Local and global variables. We use in this example a linear spring function \( s(u) \),

def f(u):
    return k*u

Here, u is a local variable, which is accessible just inside in the function, while k is a global variable, which must be initialized outside the function prior to calling f.

class Spring:
    def __init__(self, k):
        self.k = k
    def __call__(self, u):
        return self.k*u

f = Spring(k)

• «
• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10
• 11
• »