from numpy.testing import *
from numpy import *
from pymc.gp import Mean
from copy import copy

def quadfun(x, a, b, c):
    return (a * x ** 2 + b * x + c)

M = Mean(eval_fun = quadfun, a = 1., b = .5, c = 2.)

def constant(x, val):
    """docstring for parabolic_fun"""
    return zeros(x.shape[:-1],dtype=float) + val

N = Mean(constant, val = 0.)

x=arange(-1,1,.01)
y=vstack((x,x)).T

class test_mean(TestCase):
    def test(self):
        assert(M(x).shape == x.shape)
        assert(N(y).shape == (y.shape[0],))
