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()