newton_angles.py

newton_angles.py

import numpy as np
import matplotlib.pyplot as plt
from newton import solve_newton
from collections import namedtuple


def plot(results):
    fig, ax = plt.subplots()
    for r in results:
        ax.plot(r.xt, r.yt, marker='.', label=r.key)
    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
    ax.legend(fontsize='x-small')
    # fig.show()
    fig.savefig("newton_angles.png")


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]

    # Results will be stored into a list as namedtuple
    results = []
    Result = namedtuple('Result', ['key', 'xt', 'yt'])

    for theta_degree in range(5, 90, 5):
        theta = theta_degree * (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)

        # Store result
        results.append(Result(str(theta_degree), xt, yt))

    plot(results)


if __name__ == '__main__':
    main()