Spherical coordinates


I am still relatively new to PyMC but I am trying to model GP regression 3 inputs to a value CZ.
The three inputs are ALPHA, PHI and RESOLUTION. Here RESOLUTION is a categorical variable with three possible options 0, 1, 2 so I used a coregion model for the covariance for that parameter. For the first two parameters they form a spherical domain so I used a WarpedInput model with warping the spherical coordinates to cartesian coordinates.

This model takes too long to sample for me and the time required shows too high. Is there any mistake that I’ve made or suggestions on how to improve that anyone can provide? Thank you

def warp_sphere(x, period=360):
    alpha = 2*np.pi/period*x[:,0] 
    phi = 2*np.pi/period*x[:,1]
    return at.stack((at.cos(phi)*at.sin(alpha), at.sin(phi)*at.sin(alpha), at.cos(alpha)), axis=1)

with pm.Model() as model:
    # Data
    x = pm.ConstantData('x', data[['ALPHA', 'PHI', 'RESOLUTION']])
    y = pm.ConstantData('y', data['CZ'])

    # Hyperparameter Priors
    sigma = pm.InverseGamma('sigma', sigma=np.array([1, 0.1, 0.05]), mu = np.array([1.2, 0.3, 0.1]), shape = (3,))
    ls = pm.Gamma('ls', sigma=0.8, mu=1, shape=3)
    eta = pm.Gamma('eta', sigma=20, mu=25)
    cov = eta**2 * pm.gp.cov.WarpedInput(3, cov_func=pm.gp.cov.ExpQuad(3, ls=ls), warp_func=warp_sphere, active_dims=[0,1]) + pm.gp.cov.Coregion(3, W = np.zeros((2,2)), kappa = sigma, active_dims=[2])

    mean = pm.gp.mean.Zero()

    # Marginal gp Likelihood
    gp = pm.gp.Marginal(mean_func = mean, cov_func = cov)
    y_marginal = gp.marginal_likelihood('y_marginal', X = x, y = y, noise = 1e-1)