### Exercise 5.18
### Fit a polynomial to data points
### OCL

import matplotlib.pyplot as plt
import numpy as np

# Below is a function for showing a scatter plot of y versus x
# together with the best fit to the data (in the least squares sense) 
# using a polynomial of degree deg. If deg is a list with more
# than one element, then the best fit to the data using a 
# polynomial of degree deg[0], the best fit using a polynomial
# of degree deg[1], etc, will all be plotted. 

def fit(x, y, deg):
    # We first make a scatter plot of y versus x
    plt.plot(x, y, 'bo') 
         
    # To plot a polynomial p(x), we first define some x-values
    # (say a<b<c<...) and then calculate p(x) in those values
    # so we obtain a "data set" (a,p(a)), (b,p(b)), (c,p(c)),..
    # which we can then plot as a curve. Since the curve we
    # plot will be a straight line between the x-values, the 
    # distance between the x-values should not be too large for
    # the curve to look smooth. Here, we simply choose 100 
    # equally spaced x-values in the range [min(x), max(x)].
    
    x1 = np.linspace(np.min(x), np.max(x), 100) 
    for d in deg:
        coef = np.polyfit(x, y, d)  # Find polynomial coefficients
        p = np.poly1d(coef)         # Create polynomial function
        y1 = p(x1)                  # Calculate polynomial in x1
        plt.plot(x1, y1, 'r-')      # Plot as a curve

# We want to create two plots: one for the density-vs-temperature
# relation of air and one for the density-vs-temperature relation
# of water. To avoid writing the same code twice, we define a list
# with the names of the two files we want to read from, and then
# we perform the same operations on both files and plot them side
# by side in the same figure, using the plt.subplot function in 
# the matplotlib.pyplot module. 

fnames = ['density_air.dat', 'density_water.dat']

for i in range(len(fnames)):
    plt.subplot(1,2,i+1)  # Splits the plot area in 1x2 subplots
    temp = []; dens = []
    infile = open(fnames[i], 'r')
    for line in infile:
        words = line.split(); 
        if len(words) >= 2 and words[0] != '#':
            temp.append(float(words[0]))
            dens.append(float(words[1]))
    infile.close()
    fit(temp, dens, deg=[1,2])
    plt.title(fnames[i])
    plt.xlabel('Temperature')
    plt.ylabel('Density')
plt.show()