Polynomial Regression in Julia

regression.jl

import Base.show
using Formatting


typealias AA AbstractArray


function xvec!(x, row)
    m = length(row)
    row[1] = 1
    for i in 1:(m-1)
        row[i+1] = row[i] * x  # row[i+1] = x^(i-1)
    end
    return row
end


xvec(x::Float64, order::Int) = xvec!(x, Array(Float64, order))


function generate_X{T}(samples::AA{T, 1}, order::Int)
    X = Array(Float64, size(samples, 1), order)
    for (i, x) in enumerate(samples)
        xvec!(x, slice(X, i, :))
    end
    return X
end


type Regressor
    coefficients::Array{Float64}
    order::Int
    Regressor(order) = new(zeros(Float64, order+1), order)
end


function Base.show(io::IO, r::Regressor)
    if all(r.coefficients .== 0)
        print(io, "0")
    end

    expr = FormatExpr("{:+e}")
    s = ""
    for (i, c) in enumerate(r.coefficients)
        if c == 0
            continue
        end

        s *= "$(format(expr, c))*x^$i"
    end

    print(io, s)
end


predict(r::Regressor, x::Float64) = dot(r.coefficients, xvec(x, r.order))


function predict{T}(r::Regressor, x::AA{T, 1})
    n = size(x, 1)
    y = Array(Float64, n)
    for i in 1:n
        y[i] = predict(r, x[i])
    end
    return y
end


function fit{T, U}(r::Regressor, x::AA{T, 1}, y::AA{U, 1})
    assert(size(x) == size(y))

    X = generate_X(x, r.order)
    r.coefficients = inv(X'*X)*X'*y  # least squares
    return r
end

runtests.jl

include("regression.jl")

r = Regressor(20)

x = -10:0.01:10
y = sin(x)
println("Getting ready for fitting")

fit(r, x, y)

y_true = y
y_pred = predict(r, x)
println("Regressor: $r")

using PyPlot

subplot(311)
ylim([-1.5, 1.5])
title("Ground truth")
plot(x, y_true)

subplot(312)
title("Prediction")
ylim([-1.5, 1.5])
plot(x, y_pred)

subplot(313)
title("Both")
ylim([-1.5, 1.5])
plot(x, y_true)
plot(x, y_pred)

show()