Shooting method of solving ODE with boundary conditions

shooting.py

"""
Code for solving a second order differntial equation using RK2, I then modified
it to solve for an initial launch velocity given a final target height for a"""
Code for solving a second order differntial equation using RK2, I then modified
it to solve for an initial launch velocity given a final target height for a
projectile fired straight up.
"""
import numpy as np
import matplotlib.pyplot as plt

# This function solves a second order differential equation for
# d^2x/dt^2 = -g by recasting it as 2 first order differential equations
#   dx/dt = y, dy/dt = -g


# *** NEW ***:
#    Added g=9.8 to function definition
#    Changed return from -9.8 to -g
def func(conditions, t, g=9.8):
    y = conditions[0]
    vy = conditions[1]
    #
    return np.array([vy, -g])


# *** NEW ***:
#    Added **kwargs to rk2: Need to also add **kwargs to all calls to f
def rk2(f, x0, t, **kwargs):
    """
    Calculate solution to set differential equations

    dx/dt = f(x, t)

    at one or more times from initial conditions using the RK2 method.

    Note that ``x`` can be a single variable or multiple variables.

    Parameters
    ----------

    f : function
        Function with arguments x, t that returns the derivative(s) of x
        with respect to t at the time t.

    x0 : list or array of float
        Initial position(s) of the system.

    time : list or array of float
        Times at which the positions x should be calculated. Code assumes
        uniform spacing of points.

    Returns
    -------

    numpy array
        Positions calculated for each ``time``. Dimensions of the array are
        number of times  by number of variables in the initial position(s).
    """
    # One way to build up array of result xlist is to add new data as we
    # compute it. This sets that up.
    #
    # x is always the current position (i.e. the position at time)
    # xlist is the list of all positions (current and past).
    x = np.copy(x0)        # track the current value of
    try:
        # If the initial conditions are an array, use its length
        num_coords = len(np.asarray(x0))
    except TypeError as e:
        # Otherwise assume we have a single variable.
        num_coords = 1

    xlist = np.zeros([len(t), num_coords])
    lasttime = t[0]
    xlist[0, :] = x0
    for prev_time_step, time in enumerate(t[1:]):
        # This loop needs to compute h based on time and lasttime
        h = (time - lasttime)

        # Using the RK2 method, compute the new x based on derivative
        # at lastime.
        # *** NEW ***:
        #    Added **kwargs to call of f
        k1 = h * f(x, lasttime, **kwargs)
        # *** NEW ***:
        #    Added **kwargs to call of f
        k2 = h * f(x + 0.5 * k1, lasttime + 0.5 * h, **kwargs)
        x = x + k2

        # Update the appropriate element of xlist
        xlist[prev_time_step + 1, :] = x

        # Update the last time
        lasttime = time

    return xlist


# *** NEW ***:
#    Added **kwargs to make_solve: Need to also add **kwargs to all calls to rk2
#    Added yf to make_solve
def make_solve(times, yf, **kwargs):
    def solve(vy0):
        # Return the difference between the target height yf and the actual height
        # achieved for an initial velocity vy0.
        #
        # This version is kind of ugly (doesn't use **kwargs) so I have to hard
        # hard code some of the parameters here.

        yf = 0  # Desired final height
        conditions = np.array([0., vy0])  # Assuming initial height of 0.

        # *** NEW ***:
        #    Added **kwargs to call of rk2
        conditions_new = rk2(func, conditions, times, **kwargs)
        return conditions_new[-1, 0] - yf

    return solve


def secant_solver(function, initial_guesses,
                  tolerance=1e-16, max_iterations=100):
    x0, x1 = initial_guesses
    delta = abs(x1 - x0)
    n_iteration = 1

    while (delta > tolerance) and (n_iteration <= max_iterations):
        function_diff = function(x1) - function(x0)
        if abs(function_diff) < tolerance:  # This prevents a division by zero
            print('in underflow: ', abs(function_diff))
            x2 = x1
        else:
            x2 = x1 - function(x1) * (x1 - x0) / function_diff

        delta = abs(x2 - x1)
        x0, x1 = x1, x2  # This maps the x1 and x2 to x0 and x1 to recalculate

        n_iteration += 1

    return (n_iteration, x2, delta)


# Set up limits and step number
a = 0.
b = 0.3
N = 1000
h = (b - a) / N

# Initial condition
vy0 = 10.
y0 = 0
yf = 0

# *** NEW ***:
#    Define g
# Pick a planet
g = 9.8

print("Trying vy0 = ", vy0)
conditions = np.array([0., vy0])    # x(0) = 0, dxdt(0) = -10
t = np.arange(a, b, h)
# *** NEW ***:
#    Add g=g to rk2 call
x = rk2(func, conditions, t, g=g)
print("Located at x=", x[-1, 0], "at t=", b)
print("Desired x=", yf, "(Difference: ", abs(yf - x[-1, 0]), ")")

# Try an automated search and plot final result
guesses = (20, 30)
# *** NEW ***:
#    Add g=g to make_solve call
n_loop, root, delta = secant_solver(make_solve(t, g=g), guesses)
print("Shooting method finds vy0=", root, "gives yf =", yf)

conditions = np.array([0., root])
t = np.arange(a, b, h)
# *** NEW ***:
#    Add g=g to rk2 call
x = rk2(func, conditions, t, g=g)

# Make a 2D plot to show the position versus time
plt.plot(t, x[:, 0], c='g', marker='o')
plt.xlabel("t")
plt.ylabel("x(t)")
plt.title("Plot of trajectory")
plt.show()

projectile fired straight up.
"""
import numpy as np
import matplotlib.pyplot as plt

# This function solves a second order differential equation for
# d^2x/dt^2 = -g by recasting it as 2 first order differential equations
#   dx/dt = y, dy/dt = -g



def func(conditions, t):
    y = conditions[0]
    vy = conditions[1]
    #
    return np.array([vy, -9.8])


def rk2(f, x0, t):
    """
    Calculate solution to set differential equations

    dx/dt = f(x, t)

    at one or more times from initial conditions using the RK2 method.

    Note that ``x`` can be a single variable or multiple variables.

    Parameters
    ----------

    f : function
        Function with arguments x, t that returns the derivative(s) of x
        with respect to t at the time t.

    x0 : list or array of float
        Initial position(s) of the system.

    time : list or array of float
        Times at which the positions x should be calculated. Code assumes
        uniform spacing of points.

    Returns
    -------

    numpy array
        Positions calculated for each ``time``. Dimensions of the array are
        number of times  by number of variables in the initial position(s).
    """
    # One way to build up array of result xlist is to add new data as we
    # compute it. This sets that up.
    #
    # x is always the current position (i.e. the position at time)
    # xlist is the list of all positions (current and past).
    x = np.copy(x0)        # track the current value of
    try:
        # If the initial conditions are an array, use its length
        num_coords = len(np.asarray(x0))
    except TypeError as e:
        # Otherwise assume we have a single variable.
        num_coords = 1

    xlist = np.zeros([len(t), num_coords])
    lasttime = t[0]
    xlist[0, :] = x0
    for prev_time_step, time in enumerate(t[1:]):
        # This loop needs to compute h based on time and lasttime
        h = (time - lasttime)

        # Using the RK2 method, compute the new x based on derivative
        # at lastime.
        k1 = h * f(x, lasttime)
        k2 = h * f(x + 0.5 * k1, lasttime + 0.5 * h)
        x = x + k2

        # Update the appropriate element of xlist
        xlist[prev_time_step + 1, :] = x

        # Update the last time
        lasttime = time

    return xlist



def make_solve(times):
    def solve(vy0):
        # Return the difference between the target height yf and the actual height
        # achieved for an initial velocity vy0.
        #
        # This version is kind of ugly (doesn't use **kwargs) so I have to hard
        # hard code some of the parameters here.

        yf = 0  # Desired final height
        conditions = np.array([0., vy0])  # Assuming initial height of 0.
        conditions_new = rk2(func, conditions, times)
        return conditions_new[-1, 0] - yf

    return solve

def secant_solver(function, initial_guesses,
                  tolerance=1e-16, max_iterations=100):
    x0, x1 = initial_guesses
    delta = abs(x1 - x0)
    n_iteration = 1

    while (delta > tolerance) and (n_iteration <= max_iterations):
        function_diff = function(x1) - function(x0)
        if abs(function_diff) < tolerance:  # This prevents a division by zero
            print('in underflow: ', abs(function_diff))
            x2 = x1
        else:
            x2 = x1 - function(x1) * (x1 - x0) / function_diff

        delta = abs(x2 - x1)
        x0, x1 = x1, x2  # This maps the x1 and x2 to x0 and x1 to recalculate

        n_iteration += 1

    return (n_iteration, x2, delta)


# Set up limits and step number
a = 0.
b = 3.
N = 100
h = (b - a) / N

# Initial condition
vy0 = 14.553
y0 = 0
yf = 0

print("Trying vy0 = ", vy0)
conditions = np.array([0., vy0])    # x(0) = 0, dxdt(0) = -10
t = np.arange(a, b, h)
x = rk2(func, conditions, t)
print("Located at x=", x[-1, 0], "at t=", b)
print("Desired x=", yf, "(Difference: ", abs(yf - x[-1, 0]), ")")

# Try an automated search and plot final result
guesses = (20, 30)
n_loop, root, delta = secant_solver(make_solve(t), guesses)
print("Shooting method finds vy0=", root, "gives yf =", yf)

conditions = np.array([0., root])
t = np.arange(a, b, h)
x = rk2(func, conditions, t)

# Make a 2D plot to show the position versus time
plt.plot(t, x[:, 0], c='g', marker='o')
plt.xlabel("t")
plt.ylabel("x(t)")
plt.title("Plot of trajectory")
plt.show()