mae101.py

import numpy
import scipy.optimize
from matplotlib import pyplot as plt

def calculate_rmse(poly, x, y):
	y_poly = numpy.polyval(poly, x)
	mse = ((y - y_poly) ** 2).mean()
	return numpy.sqrt(mse)

def calculate_rmse_exp(popt, x, y):
	y_exp = exponential_fit(x, *popt)
	mse = ((y - y_exp) ** 2).mean()
	return numpy.sqrt(mse)

def exponential_fit(x, a, b, c):
	return a * numpy.exp(-b * x) + c

data = numpy.loadtxt("vietnam_gdp.csv", delimiter = ",")
x = data[:,0]
y = data[:,1]

polys = []
for i in range(3, 4):
	poly = numpy.polyfit(x, y, deg = i)
	print(f"{i} poly: ", poly)
	poly = numpy.poly1d(poly)
	polys.append(poly)


plt.title("Vietnam's GDP, 1989-2017 period")
plt.xlabel("Year")
plt.ylabel("GDP (in 10 billion USD)")

plt.plot(x, y, '.', label = "Original data", zorder = 10)

for poly in polys:
	range_x = numpy.linspace(min(x) - 3, max(x) + 3, num = 100)
	plt.plot(range_x, poly(range_x), label = f"{len(poly)}-degree polyfit (RMSE: {calculate_rmse(poly, x, y):.4f})")

popt, _ = scipy.optimize.curve_fit(exponential_fit, x, y, p0 = (1e-6, 1e-6, 1e-6), maxfev = 100000)
range_x = numpy.linspace(min(x) - 3, max(x) + 3, num = 100)
plt.plot(range_x, exponential_fit(range_x, *popt), label = f"Exponential fit, 100k calls (RMSE: {calculate_rmse_exp(popt, x, y):.4f})")
print("exp: ", popt)

maple_poly = numpy.poly1d([0.000671226382184633, -3.99601866044572, 7929.82423581614, -5245381.72886966])
plt.plot(range_x, maple_poly(range_x), label = f"{len(maple_poly)}-degree polyfit (Maple) (RMSE: {calculate_rmse(maple_poly, x, y):.4f})")

plt.legend(loc = "best")
plt.show()

vietnam_gdp.csv