Need help making a 3D surface plot a 4D surface plot with color as separate dimension

As the title says. I have tried so many different ways of doing this. I have 4 vectors of length 48.

X: [ 25  25  25  25  25  25  50  50  50  50  50  50  75  75  75  75  75  75
 100 100 100 100 100 100 125 125 125 125 125 125 150 150 150 150 150 150
 175 175 175 175 175 175 200 200 200 200 200 200]
Y: [ 100  250  500 1000 1500 2000  100  250  500 1000 1500 2000  100  250
  500 1000 1500 2000  100  250  500 1000 1500 2000  100  250  500 1000
 1500 2000  100  250  500 1000 1500 2000  100  250  500 1000 1500 2000
  100  250  500 1000 1500 2000]
Z: [ 0.20900428  0.51286209  1.03853414  3.28220448  4.6407558   7.34891026
  0.2765902   0.7604821   1.76022537  5.10049512  8.61249235 12.96447849
  0.2623122   0.98286221  2.5040107   6.2533442  11.0721308  15.36910634
  0.32121766  0.97078288  2.66376145  7.51123161 12.98652091 20.21016505
  0.38653798  1.21371622  3.30200138  7.93705671 17.20774968 28.97923372
  0.46758823  1.23861806  3.72943289  8.38099084 19.04535632 32.7009341
  0.44258697  1.42894619  3.96008332 10.45831311 22.98130064 31.32277734
  0.4507597   1.7036628   4.69553339 10.92697349 25.68610439 45.02457106]
C: [38.96 39.48 40.34 41.04 41.08 41.06 39.76 40.62 40.88 41.06 41.04 41.2
 39.22 40.48 40.98 41.2  41.26 41.16 40.2  40.78 40.68 41.26 41.26 41.32
 39.96 40.56 40.86 41.26 41.26 41.52 40.36 40.6  41.22 41.26 41.78 41.7
 39.24 40.8  41.26 41.4  41.92 41.62 39.74 41.06 41.24 41.56 41.94 42.06]

This code snippet

X = overall_results.num_generations.values
Y = overall_results.population_size.values
Z = overall_results.avg_time.values
C = overall_results.avg_reward.values 
# note, C is not used as this is meant to be a fully working example
fig = plt.figure(figsize=(8,6))
ax = Axes3D(fig, auto_add_to_figure=False)
fig.add_axes(ax)
surf = ax.plot_trisurf(X, Y, Z, cmap=cm.jet, linewidth=.2)
ax.view_init(elev=5, azim=-140)
colorbar = fig.colorbar(surf, ax=ax, pad=0.1, shrink=.5, ticks=[5, 10, 15, 20, 25, 30], format="%d")
colorbar.ax.set_yticklabels(['<= 5', '10', '15', '20', '25', '>= 30'])
plt.title('GA Time Analysis by Population Size and Number of Generations')

plt.show()

produces this figure

The color is mapped to Z, and the various methods I have tried to incorporate C all throw errors. This also uses trisurf, and, the polycount is very low.

This code snippet

X = overall_results.num_generations.values
Y = overall_results.population_size.values
Z = overall_results.avg_time.values
C = overall_results.avg_reward.values

# Define a finer grid for interpolation
new_X = np.linspace(X.min(), X.max(), 100)
new_Y = np.linspace(Y.min(), Y.max(), 100)
new_X, new_Y = np.meshgrid(new_X, new_Y)

# Perform interpolation
new_Z = griddata((X, Y), Z, (new_X, new_Y), method='linear')

# Create the 3D plot
fig = plt.figure(figsize=(8, 8))
ax = Axes3D(fig, auto_add_to_figure=False)
fig.add_axes(ax)

surf = ax.plot_surface(new_X, new_Y, new_Z, cmap=cm.jet, antialiased=True)

ticks = np.linspace(Z.min(), Z.max(), 10)
#ticks = [5, 10, 15, 20, 25, 30, 35, 40, 45]

colorbar = fig.colorbar(surf, ax=ax, pad=0.1, shrink=0.35, ticks=ticks, format="%d")
#colorbar.ax.set_yticklabels(['<= 5', '10', '15', '20', '25', '>= 30'])

ax.view_init(elev=8, azim=-150)
plt.title('GA Time Analysis by Population Size and Number of Generations')
ax.set_xlabel('Number of Generations', labelpad=12, fontsize=14)
ax.set_ylabel('Population Size', labelpad=12, fontsize=14)
ax.zaxis.set_rotate_label(False)
ax.set_zlabel('Running Time (seconds)', rotation=90, fontsize=14)
plt.show()

produces this figure

A much better looking figure, and uses plot_surface instead of trisurf. But again, I'm not able to use C to set the color bar. I've looked at

Plot 3d surface with colormap as 4th dimension, function of x,y,z

How to make a 4d plot with matplotlib using arbitrary data

Color matplotlib plot_surface command with surface gradient

libraries used

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from scipy.interpolate import griddata
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from matplotlib.ticker import MaxNLocator
from matplotlib.colors import Normalize

Now, I'm able to create a scatter plot that does what I want, minus the surface, like so

X = overall_results.num_generations.values
Y = overall_results.population_size.values
Z = overall_results.avg_time.values
C = overall_results.avg_reward.values

cmap = plt.get_cmap('jet')  # You can choose any colormap you prefer
norm = Normalize(vmin=C.min(), vmax=C.max())

fig = plt.figure(figsize=(8,8))
ax = fig.add_subplot(111, projection='3d')
img = ax.scatter(X, Y, Z, c=C, s=100, cmap=cmap, norm=norm, alpha=1.0)
#plt.scatter(x, y, c=x, cmap=cmap, s=350, alpha=.7)
plt.xlabel('Average Reward', fontsize=14)
plt.ylabel('Running Time', fontsize=14)

cbar = fig.colorbar(img, pad=.1, shrink=.5)
cbar.set_label('Average Reward', fontsize=14, labelpad=10)
ax.view_init(elev=20, azim=-140)
plt.show()

Which produces this figure

But I'd like this effect as a surface.

Thanks to Gordon Stein and Caleb Vatral (go Vanderbilt!), this was the missing piece:

# Also color
new_C = griddata((X, Y), C, (new_X, new_Y), method='linear')

# Done to rescale it to show in color map, you likely need to change it
new_C = new_C - new_C.min()
new_C = cm.gist_rainbow(new_C / new_C.max())
surf = ax.plot_surface(new_X, new_Y, new_Z, facecolors=new_C, antialiased=True)

alternatively,

norm = (C - C.min()) / (C.max()-C.min())
new_C = griddata((X,Y), norm, (new_X, new_Y), method="linear")

colors = np.empty(new_X.shape, dtype=tuple)
for y in range(100):
    for x in range(100):
        colors[y, x] = cm.jet(new_C[x, y], )
surf = ax.plot_surface(new_X, new_Y, new_Z, facecolors=colors, antialiased=True)

and with some colorbar manipulation we now have