Plotting a sphere in 3D space with the surface colour determined by another variable

I am trying to create a 3D plot in python to display a sphere where the surface colour indicates the value of a potential on the surface. The potential is dependent on the value of theta.

Right now my code is as follows

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import matplotlib as mp

fig, ax = plt.subplots(subplot_kw=dict(projection="3d"),
                        constrained_layout=1)

r = 1
theta = np.random.uniform(low=0.0, high=np.pi, size=(1,1000))
phi = np.random.uniform(low=0.0, high=2*np.pi, size=(1,1000))
    
x = r*np.cos(theta)*np.sin(phi)
y = r*np.sin(theta)*np.sin(phi)
z = r*np.cos(phi)
X, Y = np.meshgrid(x, y)
V = v_func(theta, r) # 4th value

norm = plt.Normalize(vmin=V.min().min(), vmax=V.max().max())
ax.plot_surface(X, Y, V, facecolors=plt.cm.plasma(norm(V)))
m = mp.cm.ScalarMappable(cmap=plt.cm.plasma, norm=norm)
plt.colorbar(m)
plt.show()

where v_func is

def v_func(theta, r):
    E_0 =1
    #e_0 =8.854e-12
    e_0 = 1
    a =1
    Q_0 =1
    return -E_0*r*np.cos(theta)+ (2*E_0*a**3/r**2)*np.cos(theta)+Q_0/(4*np.pi*e_0*r)

This code will just throw back an error of

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[59], line 16
     13 V = v_func(theta, r) # 4th value
     15 norm = plt.Normalize(vmin=V.min().min(), vmax=V.max().max())
---> 16 ax.plot_surface(X, Y, V, facecolors=plt.cm.plasma(norm(V)))
     17 m = mp.cm.ScalarMappable(cmap=plt.cm.plasma, norm=norm)
     18 plt.colorbar(m)

File ~\anaconda3\envs\muon\lib\site-packages\mpl_toolkits\mplot3d\axes3d.py:1640, in Axes3D.plot_surface(self, X, Y, Z, norm, vmin, vmax, lightsource, **kwargs)
   1637             polys.append(ps)
   1639             if fcolors is not None:
-> 1640                 colset.append(fcolors[rs][cs])
   1642 # In cases where there are NaNs in the data (possibly from masked
   1643 # arrays), artifacts can be introduced. Here check whether NaNs exist
   1644 # and remove the entries if so
   1645 if not isinstance(polys, np.ndarray) or np.isnan(polys).any():

IndexError: index 20 is out of bounds for axis 0 with size 1

And I have been unable to identify the source of this error and produce the required graph

Here is a way with PyVista:

import numpy as np
import pyvista as pv

def v_func(theta, r):
    E_0 =1
    #e_0 =8.854e-12
    e_0 = 1
    a =1
    Q_0 =1
    return -E_0*r*np.cos(theta)+ (2*E_0*a**3/r**2)*np.cos(theta)+Q_0/(4*np.pi*e_0*r)

mesh = pv.Sphere(theta_resolution=200, phi_resolution=200)
theta = np.arccos(2*mesh.points[:,2]) 
scalars = v_func(theta, 0.5)

pltr = pv.Plotter(window_size = [512, 512])
pltr.background_color = "#363940"
pltr.set_focus([0, 0, 0])
pltr.set_position((4, 0, 0))
mesh["dist"] = scalars
pltr.add_mesh(
    mesh, smooth_shading = True, cmap = "plasma", 
    show_scalar_bar = True, specular = 1
)
pltr.show()