fix for multi-dim input bug in Zero and Constant, tidy up docstrings, add tests

Author: bwengals

pymc3/gp/mean.py

@@ -3,10 +3,9 @@
 __all__ = ['Zero', 'Constant']
 
 class Mean(object):
-    """
+    R"""
     Base class for mean functions
     """
-    
     def __call__(self, X):
         R"""
         Evaluate the mean function.
@@ -22,49 +21,52 @@ def __add__(self, other):
 
     def __mul__(self, other):
         return Prod(self, other)
-        
+
 class Zero(Mean):
+    R"""
+    Zero mean function for Gaussian process.
+
+    """
     def __call__(self, X):
-        return tt.zeros(X.shape, dtype='float32')
-        
+        return tt.zeros(tt.stack([X.shape[0], ]), dtype='float32')
+
 class Constant(Mean):
-    """
+    R"""
     Constant mean function for Gaussian process.
-    
+
     Parameters
     ----------
     c : variable, array or integer
         Constant mean value
     """
-
     def __init__(self, c=0):
         Mean.__init__(self)
         self.c = c
 
     def __call__(self, X):
-        return tt.ones(X.shape) * self.c
+        return tt.ones(tt.stack([X.shape[0], ])) * self.c
+
 
 class Linear(Mean):
-    
+    R"""
+    Linear mean function for Gaussian process.
+
+    Parameters
+    ----------
+    coeffs : variables
+        Linear coefficients
+    intercept : variable, array or integer
+        Intercept for linear function (Defaults to zero)
+    """
     def __init__(self, coeffs, intercept=0):
-        """
-        Linear mean function for Gaussian process.
-        
-        Parameters
-        ----------
-        coeffs : variables
-            Linear coefficients
-        intercept : variable, array or integer
-            Intercept for linear function (Defaults to zero)
-        """
         Mean.__init__(self)
         self.b = intercept
         self.A = coeffs
-        
+
     def __call__(self, X):
         return tt.dot(X, self.A) + self.b
 
-        
+
 class Add(Mean):
     def __init__(self, first_mean, second_mean):
         Mean.__init__(self)
@@ -83,4 +85,4 @@ def __init__(self, first_mean, second_mean):
 
     def __call__(self, X):
         return tt.mul(self.m1(X), self.m2(X))
-        
+