Custom Mean Function Construction for GPFlow Regression

Question

I am trying to pass a custom mean function into GPflow 2.0. I have some (x,y,z) data with several observations for each x,y point. I wanted to pass the average z value for each (x,y) pair as the mean function. My code for the custom mean function is below where train_xy are the (x,y,z) tuples in my training data (x and y are inputs, and z is the output).

Below is my code setting up the mean function and passing it into a simple GPR model.

---

class CustomMean(gpflow.mean_functions.MeanFunction):
    """
    means = mean(grouped by unique tuples of input1 & input2)
    """

    def __init__(self, X):
        # MeanFunction.__init__(self)
        self.X = Parameter(X)

    def __call__(self, X):
        mean_arr = []
        X = np.unique(X, axis=0)
        for i in range(len(X)):
            y_list = train_xy[(train_xy['input1'] == X[i, 0]) & (train_xy['input2'] == X[i, 1])]['output']
            mean_arr.append(y_list.mean())
        mean_arr = np.asarray(mean_arr, dtype=np.float64)
        mean_arr[np.isnan(mean_arr)] = 0
        means = tf.convert_to_tensor(mean_arr)
        return means
    
    
k = gpflow.kernels.Linear()
mean_func = CustomMean(X_train.values.astype('float64'))
m = gpflow.models.GPR(data=(X_train, y_train), kernel=k, mean_function=mean_func)

pred_mean, pred_var = m.predict_f(X_test.values.astype('float64'))

---

I ran into problems when I called the predict_f function:

---

InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-140-c7b24d06c38e> in <module>()
----> 1 pred_mean, pred_var = m.predict_f(X_test.values.astype('float64'))

~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/gpflow/models/gpr.py in predict_f(self, Xnew, full_cov, full_output_cov)
    103             kmn, kmm + s, knn, err, full_cov=full_cov, white=False
    104         )  # [N, P], [N, P] or [P, N, N]
--> 105         f_mean = f_mean_zero + self.mean_function(Xnew)
    106         return f_mean, f_var

...

InvalidArgumentError: Incompatible shapes: [1250,2596] vs. [1104] [Op:AddV2]

---

And it looks like the error is tied to something called f_mean = f_mean_zero + self.mean_function(Xnew). When I look at the shape of this, the dimensions are very large. I have 2596 unique (x,y) pairs in my training data (out of 3750 total rows) and 1104 unique (x,y) pairs in test (out of 1250 total rows). f_mean_zero is (1250, 2596), whereas self.mean_function(Xnew) returns a one-dimensional tensor of means (as expected). I don't understand the dimension of f_mean_zero, particularly since (x,y) pairs in my test data are in my training data. It seems like a similar issue to this ticket but there is no solution. Thank you!

Answer 1

I have a couple of pointers that should help you solve your issue:

1. It is important to remember that GPflow uses TensorFlow for most of its computations. The inputs and outputs of many operations (like the __call__ of a mean function) are assumed to be tf.Tensors and not np.ndarrays as is the case in your snippet. Operations like len(), np.unique, etc. are not supported by tf.Tensors and will have to be correctly translated to TensorFlow ops. That being said, the task you want to achieve is not straightforward to implement with Tensor logic in TensorFlow. I believe you will have to use operations like tf.where.

2. The stack trace you provided shows mismatching shapes. The output of the mean function's call should have the same leading dimensions as the input (X in your case). Example: if X has shape [N, D], the output of the call should have shape [N, 1].

3. It is unclear to me why you are storing X as a parameter in the init of your mean function. Can you elaborate on that?

Hope this helps you on your way.

Answer 2

Have you considered subtracting the mean out of your data and modeling it with a zero mean GPR, then adding the mean back afterwards? That would achieve the same effect without having to create the custom class.