Polynomial curve fitting example

fitting.py

#!/usr/bin/env python
#-*- coding: utf-8 -*-
#
# Polynomial curve fitting example

import sys
import math
import itertools
import numpy as np
import numpy.random as rand
import scipy.stats as stats
import matplotlib.pyplot as plt

# ground truth: y = sin(2πx), 0 <= x <= 1
xmin = 0; xmax = 1; ymin = -1.5; ymax = 1.5
ground_truth = lambda x: math.sin(x * math.pi * 2.0)

# M-th polynomial model regression
powervec = lambda x, M: np.matrix([ [x ** i] for i in range(M + 1) ])
poly = lambda w, x, M: np.dot(w.T, powervec(x, M)).item()

# regression by maximum-likelihood estimation; returns w
def poly_regression_mle(X, T, M):
    return poly_regression_mle_map("MLE", X, T, M)

# regression by maximum a posteriori estimation; returns w
def poly_regression_map(X, T, M, alpha):
    return poly_regression_mle_map("MAP", X, T, M, alpha)

# common function for MLE and MAP
def poly_regression_mle_map(est, X, T, M, alpha=None):
    A = np.matrix(np.zeros((M + 1, M + 1)))
    b = np.matrix(np.zeros((M + 1,     1)))
    for x, t in zip(X, T):
        x_power = powervec(x, M)
        A += np.dot(x_power, x_power.T)
        b += t * x_power
    if est == "MAP":
        A += alpha * np.identity(M + 1)
    return np.dot(A.I, b)

# regression by bayes estimation; returns distribution over Xaxis
def poly_regression_bayes(X, T, M, alpha, beta, Xaxis):
    vecT = np.matrix(T).T
    phi  = np.hstack(map(lambda x: powervec(x, M), X)).T
    S    = (alpha * np.identity(M+1) + beta * np.dot(phi.T, phi)).I
    m    = beta * np.dot(np.dot(S, phi.T), vecT)
    estTaxis = []
    for x in Xaxis:
        phi = powervec(x, M)
        mu  = np.dot(m.T, phi).item()
        var = 1.0 / beta + np.dot(np.dot(phi.T, S), phi).item()
        estTaxis.append((mu, var))
    return estTaxis

def main(N, TYPE, fixM=None, fixAlpha=None, fixBeta=None,
         resolution=0.01, dataError=0.2):
    # plot ground truth curve
    Xaxis = np.arange(xmin, xmax + resolution, resolution)
    Taxis = map(ground_truth, Xaxis)
    plt.plot(Xaxis, Taxis, "g-", label="ground truth")

    # generate training dataset X, T at random
    X = np.array([ rand.uniform() for _ in range(N) ])
    Y = map(ground_truth, X) # ideal output
    T = Y + np.array([ rand.normal(0, dataError) for y in Y ]) # + white noise
    plt.plot(X, T, "ko")

    # estimate and plot results
    if TYPE == "MLE":
        Ms = [fixM] if fixM is not None else [1, 3, 15]
        for M in Ms:
            w = poly_regression_mle(X, T, M)
            estTaxis = map(lambda x: poly(w, x, M), Xaxis)
            plt.plot(Xaxis, estTaxis, "--", label="M=%d" % M)
        plt.title("ML estimation")
    elif TYPE == "MAP":
        Ms = [fixM] if fixM is not None else [15]
        As = [fixAlpha] if fixAlpha is not None else [10e-15, 10e-5, 10e-1]
        for M, alpha in itertools.product(Ms, As):
            w = poly_regression_map(X, T, M, alpha)
            estTaxis = map(lambda x: poly(w, x, M), Xaxis)
            plt.plot(Xaxis, estTaxis, "--",
                     label="M=%d alpha=%.1e" % (M, alpha))
        plt.title("MAP estimation")
    elif TYPE == "BAYES":
        M = fixM if fixM is not None else 10
        alpha = fixAlpha if fixAlpha is not None else 10e-3
        beta  = fixBeta if fixBeta is not None else 1.0
        estTaxis = poly_regression_bayes(X, T, M, alpha, beta, Xaxis)
        Yaxis = np.arange(ymin, ymax, resolution)
        Xmesh, Ymesh = np.meshgrid(Xaxis, Yaxis)
        Zmesh = np.array([
            [ stats.norm.pdf(y[0], loc=estTaxis[i][0], scale=estTaxis[i][1])
              for i in range(Xmesh.shape[1]) ] for y in Ymesh ])
        plt.pcolor(Xmesh, Ymesh, Zmesh)
        plt.title("Bayes estimation")
    else:
        print "wrong type"
        quit()

    # show
    plt.xlim(xmin, xmax); plt.ylim(ymin, ymax)
    plt.legend()
    plt.xlabel("X"); plt.ylabel("T")
    plt.show()

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='polynomial curve fitting')
    parser.add_argument('-n', nargs=1, default=[10], type=int, metavar="N",
                        help='# of sampling points')
    parser.add_argument('-t', nargs=1, default=["MLE"], type=str,
                        metavar="type", choices=["MLE", "MAP", "BAYES"],
                        help='estimation type (MLE, MAP, or BAYES)')
    parser.add_argument('-m', nargs="?", default=None, type=int,
                        metavar="type", help='order of polynomial')    
    parser.add_argument('-a', nargs="?", default=None, type=float,
                        metavar="type", help='value of alpha')
    parser.add_argument('-b', nargs="?", default=None, type=float,
                        metavar="type", help='value of beta')
    args = parser.parse_args()
    main(args.n[0], args.t[0], args.m, args.a, args.b)