newton.py

newton.py

import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# scipy.integrate.solve_ivp
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html


# Newton equation
def f_newton(t, X, m, g, b):
    # X = [x, v_x, y, v_y]
    dXdt = np.array([
        X[1],              # dx/dt = v_x
        -(b/m) * X[1],     # dv_x/dt = -(b/m) v_x
        X[3],              # dy/dt = v_y
        -(b/m) * X[3] - g  # dv_y/dt = -(b/m) v_y - g
    ])
    return dXdt


def solve_newton(eq_params, X0, t_range, n_t):
    m, g, b = eq_params
    t_start, t_end = t_range

    # Solve ODE
    sol = solve_ivp(f_newton, (t_start, t_end), X0, args=(m, g, b), dense_output=True)
    print(sol.message)

    # Get dense output for plot
    t = np.linspace(t_start, t_end, n_t)
    Xt = sol.sol(t)
    assert Xt.shape == (4, n_t)

    # x(t), v_x(t), y(t), v_y(t)
    return Xt[0, :], Xt[1, :], Xt[2, :], Xt[3, :]


def plot(xt, yt):
    fig, ax = plt.subplots()
    ax.plot(xt, yt, marker='.')
    ax.set_xlim(left=0)
    ax.set_ylim(bottom=0)
    ax.set_xlabel(r'$x$')
    ax.set_ylabel(r'$y$')
    ax.set_aspect('equal')  # plot x and y in an equal scale
    # fig.show()
    fig.savefig("newton.pdf")


def main():
    # MKS units: m, kg, s
    g = 9.8  # [m/s^2]
    m = 0.1  # [kg]
    b = 0.1  # [kg/s]

    t_start = 0  # [s]
    t_end = 5.0  # [s]
    n_t = 101

    # initial state
    v0 = 100.0 / 3.6  # [m/s]
    theta = 45 * (np.pi / 180)  # radian

    # X0 = [x0, v0_x, y0, v0_y]
    X0 = np.array([0, v0 * np.cos(theta), 0, v0 * np.sin(theta)])

    # Solve Newton equation
    # x(t), y(t)
    xt, _, yt, _ = solve_newton(eq_params=(m, g, b), X0=X0, t_range=(t_start, t_end), n_t=n_t)

    # Plot
    plot(xt, yt)


if __name__ == '__main__':
    main()