SciPy exercises

SciPy exercises#

import numpy as np
import matplotlib.pyplot as plt

Integration#

from scipy import integrate

Q1: integrating an analytic function#

Numerical integration methods work differently depending on whether you have the analytic function available (in which case you can evaluate it freely at any point you please) or if it is sampled for you.

Consider the function \(f(x) = e^{-x^2}\). We want to integrate this from \([-5, 5]\). The analytic integral is not easily obtained. Use integrate.quad to do the integration.

Q2: integrating a sampled function#

Consider now that you have data that represents a function sampled a N points, but you don’t know the analytic form of the function. Here, we create the sampling here for a Gaussian and we will do the same integral as in Q1.

N = 32
x = np.linspace(-5, 5, N)
f = np.exp(-x**2)
f

array([1.38879439e-11, 3.15061953e-10, 5.80457065e-09, 8.68481106e-08,
       1.05527775e-06, 1.04133225e-05, 8.34503173e-05, 5.43103745e-04,
       2.87047478e-03, 1.23208538e-02, 4.29481052e-02, 1.21580337e-01,
       2.79510942e-01, 5.21855680e-01, 7.91258065e-01, 9.74320895e-01,
       9.74320895e-01, 7.91258065e-01, 5.21855680e-01, 2.79510942e-01,
       1.21580337e-01, 4.29481052e-02, 1.23208538e-02, 2.87047478e-03,
       5.43103745e-04, 8.34503173e-05, 1.04133225e-05, 1.05527775e-06,
       8.68481106e-08, 5.80457065e-09, 3.15061953e-10, 1.38879439e-11])

Compute the integral of this sampled function using Simpson’s method (integrate.simps). Now, vary the number of sample points (try 64, 128, …) and see how the answer changes. Simpson’s method is 4-th order accurate, which means that the error should decrease by \(2^4\) when we double the number of sample points

Optional: Make a plot of the error (compared to the analytic integral from Q1) vs. N

Interpolation#

from scipy import interpolate

Q3: interpolation error#

There are a large number of different interpolation schemes available through scipy. Let’s test them out.

Create a python function, \(f(x)\), that is your true function. Now create \(N\) samples of it (either regularly spaced or irregularly spaced).

Try some of the different interolation routines. interpolate.interp1d takes a kind argument that let’s you choose the order of the interpolation. Measure the error in the method, by comparing the interpolated result with the actual function value. Also, try using cubic splines (look at CubicSpline)

Try plotting the resulting interpolant.

Root Finding#

Q4: scalar function roots#

Consider the function

\[q(x) = x^3 - 2x^2 - 11x + 12\]

This has 3 roots, but is known to cause problems for some root-finding methods (it exhibits basis of attraction: https://en.wikipedia.org/wiki/Newton’s_method#Basins_of_attraction – very closely spaced initial guesses leave to very different roots)

Use the SciPy optimize.brentq method to find the roots. You might need to play around with the intervals to find all 3 roots (try plotting the function to help)

from scipy import optimize

ODEs#

Q5: orbits#

We want to consider planetary orbits. To do this, we need to solve Newton’s second law together with Newton’s law of gravity. If we restrict ourselves to the x-y plane, then there are 4 quantities we need to solve for: \(x\), \(y\), \(v_x\), and \(v_y\). These evolve according to:

\[\begin{align*} \frac{dx}{dt} &= v_x \\ \frac{dy}{dt} &= v_y \\ \frac{dv_x}{dt} &= a_x = -\frac{GM_\star x}{r^3} \\ \frac{dv_y}{dt} &= a_y = -\frac{GM_\star y}{r^3} \end{align*}\]

To integrate these forward in time, we need an initial condition for each quantity. We’ll setup our system such that the Sun is at the origin (that will be one focus), and the planet begins at perihelion and orbits counterclockwise.

geometry

The distance of perihelion from the focus is:

\[r_p = a (1 - e)\]

where \(a\) is the semi-major axis and \(e\) is the eccentricity. The perihelion velocity is all in the \(y\) direction and is:

\[v_y = v_p = \sqrt{\frac{GM_\star}{a} \frac{1+e}{1-e}}\]

We’ll work in units of AU, years, and solar masses, in which case, \(GM_\star = 4\pi^2\) (for the Sun).

Your initial conditions should be:

\(x(t=0) = r_p\)
\(y(t=0) = 0\)
\(v_x(t=0) = 0\)
\(v_y(t=0) = v_p\)

Here’s a righthand side function for the ODEs:

def rhs(t, Y, GM=4*np.pi**2):
    """RHS for orbits, Y is the solution vector, containing
    x, y, v_x, and v_y"""

    x, y, vx, vy = Y
    f = np.zeros_like(Y)

    # dx/dt = vx
    f[0] = vx

    # dy/dt = vy
    f[1] = vy

    # d(vx)/dt = -GMx/r**3
    r = np.sqrt(x**2 + y**2)
    f[2] = -GM*x/r**3

    # d(vy)/dt = -GMy/r**3
    f[3] = -GM*y/r**3

    return f

Use the SciPy ODE integration methods to integrate an orbit and plot it

Q6: damped driven pendulum and chaos#

There are a large class of ODE integration methods available through the scipy.integrate.ode() function. Not all of them provide dense output – most will just give you the value at the end of the integration.

The explicit Runge-Kutta integrator will give you access to the solution at intermediate points and provides methods to interpolate to any value. You enable this via dense_output=True (see the example in our out-of-class notebook).

The damped driven pendulum obeys the following equations:

\[\dot{\theta} = \omega\]

\[\dot{\omega} = -q \omega - \sin \theta + b \cos \omega_d t\]

here, \(\theta\) is the angle of the pendulum from vertical and \(\omega\) is the angular velocity. \(q\) is a damping coefficient, \(b\) is a forcing amplitude, and \(\omega_d\) is a driving frequency.

Choose \(q = 0.5\) and \(\omega_d = 2/3\).

Integrate the system for different values of \(b\) (start with \(b = 0.9\) and increase by \(0.05\), and plot the results (\(\theta\) vs. \(t\)). Here’s a RHS function to get you started:

def rhs(t, Y, q, omega_d, b):
        """ damped driven pendulum system derivatives.  Here, Y = (theta, omega) are
        the solution variables. """
        f = np.zeros_like(Y)
        
        f[0] = Y[1]
        f[1] = -q*Y[1] - np.sin(Y[0]) + b*np.cos(omega_d*t)

        return f

Note that the pendulum can flip over, giving values of \(\theta\) outside of \([-\pi, \pi]\). The following function can be used to restrict it back to \([-\pi, \pi]\) for plotting.

def restrict_theta(theta):
    """ convert theta to be restricted to lie between -pi and pi"""
    tnew = theta + np.pi
    tnew += -2.0*np.pi*np.floor(tnew/(2.0*np.pi))
    tnew -= np.pi
    return tnew

Write a function that takes an initial angle, \(\theta_0\), and integrates the system and returns the solution.

Note, the righthand side function, rhs, takes additional arguments that you need to pass through the integrator. The preferred method to do this with the solve_ivp() interface appears to be to use functools.partial(), as:

from functools import partial

r = solve_ivp(partial(rhs, q=q, omega_d=omega_d, b=b), ...)

Some values of \(b\) will show very non-periodic behavior. To see chaos, integrate two different pendula that are the same except for \(\theta_0\), with only a small difference between then (like 60 degrees and 60.0001 degrees. You’ll see the solutions track for a while, but then diverge.