The ability to color a 3D plot using a color function is geared more for functions of x and y only.
But quite often, the surface or pointwise 3D position of the plot is itself being specified as a height z which is a function of x and y. For the plot3d comand, that's pretty much the way it works (whether using an expression or a procedure).
So, of course, the very same rule that generates the z value for the position can be used in the color function. The `option remember` can be used to avoid repeating the z calculations.
An alternative is to produce an original plot (using plot3d and the GRID structure here -- it would need changes for other generated plot structures) and then simply replace its COLOR or SHADING section with the desired color specification.
In this way, any custom color function can be used, and in particular a scheme could be used which only depends on the precomputed z-positions. The precomputed z-values (z-positions) are right there and accessible in the GRID portion of the plot structure. And the maximal and minimal z-positions can be used so as to attain particular desired colors at the extremes of the plot (making the whole coloring scheme be relative to the plot's actual position -- not known in advance, in general -- which isn't really possible using the first example's color function method).
Here are some examples. I'd actually been looking at with this topic already, off and on, but Christopher2222's question spurred me to post earlier than planned. Ideally, this could all be wrapped into some nice procedure that allowed easy use (eg. special cases like named colors for top and bottom positions being simply supplied as optional arguments, etc).