Refactor quadrature code

[awav]

The current code state - spaghetti. Let's think about how we can make it more readable and extendable.

Contact: @awav, @st--

[gustavocmv]

Hello @awav !

I was about to post an Issue about quadrature code, but since you guys are thinking about refactoring it, I believe I can post it here.

There is a bug with the current code when you use it with tf.function and the inputs Fmu and Fvar are lists, below is an MWE:

import numpy as np
import tensorflow as tf

from gpflow.quadrature import ndiagquad
tf_ndiagquad = tf.function(ndiagquad)

# Generate Data
N = 10
D = 2
H = 11
mu = 4
var = 9
aux = np.ones((N, D))
Fmu = mu * aux
Fvar = var * aux

# Define Funcs
funcs = [
    lambda x1, x2: x1 + x2, 
    lambda x1, x2: x1**2 + x2**2
]

# Transform to Fmu, Fvar to lists
Fmu = tf.unstack(tf.reshape(Fmu, (-1, 1, D)), axis=-1)
Fvar = tf.unstack(tf.reshape(Fvar, (-1, 1, D)), axis=-1)


# This one works
EX, EX2 = ndiagquad(funcs, H, Fmu, Fvar) 

# This one fails
tf_EX, tf_EX2 = tf_ndiagquad(funcs, H, Fmu, Fvar)

# NameError: free variable 'Din' referenced before assignment in enclosing scope

I have made a quick-fix with the following function:

from gpflow.quadrature import mvhermgauss
def my_ndiagquad(funcs, H, Fmu, Fvar, logspace: bool = False, **Ys):
    if isinstance(Fmu, (tuple, list)):
        Din = len(Fmu)
        shape = tf.shape(Fmu[0])
        Fmu, Fvar = [tf.stack(f, axis=-1) for f in [Fmu, Fvar]] # both [N, 1, Din]
    else:
        Din = 1
        shape = tf.shape(Fmu)
        Fmu, Fvar = [tf.reshape(f, (-1, 1, 1)) for f in [Fmu, Fvar]] # both [N, 1, Din]
        
    xn, wn = mvhermgauss(H, Din)
    # xn: H**Din x Din, wn: H**Din
    
    gh_x = xn.reshape(1, -1, Din)  # [1, H]**Din x Din
    Xall = gh_x * tf.sqrt(2.0 * Fvar) + Fmu  # [N, H]**Din x Din
    Xs = [Xall[:, :, i] for i in range(Din)]  # [N, H]**Din  each
    
    gh_w = wn * np.pi**(-0.5 * Din)  # H**Din x 1

    for name, Y in Ys.items():
        Y = tf.reshape(Y, (-1, 1))
        Y = tf.tile(Y, [1, H**Din])  # broadcast Y to match X
        # without the tiling, some calls such as tf.where() (in bernoulli) fail
        Ys[name] = Y  # now [N, H]**Din

    def eval_func(f):
        feval = f(*Xs,
                  **Ys)  # f should be elementwise: return shape [N, H]**Din
        if logspace:
            log_gh_w = np.log(gh_w.reshape(1, -1))
            result = tf.reduce_logsumexp(feval + log_gh_w, axis=1)
        else:
            result = tf.linalg.matmul(feval, gh_w.reshape(-1, 1))
        return tf.reshape(result, shape)

    if isinstance(funcs, Iterable):
        return [eval_func(f) for f in funcs]

    return eval_func(funcs)

tf_my_nddiagquad = tf.function(my_nddiagquad)

I tested it with the MWE and it worked:

my_EX, my_EX2 = my_ndiagquad(funcs, H, Fmu, Fvar) 
tf_my_EX, tf_my_EX2 = tf_my_nddiagquad(funcs, H, Fmu, Fvar)

Hope you guys can use the code above to patch it.

Also, I have done some code to perform GH Quadrature before and I would be glad to contribute in the refactoring.

[awav]

@gustavocmv, Thanks! Could you make a PR with tests?