How can I properly define a custom function to give as a ColorFunction option?

Question

So I thought that I knew what was going on. I defined my own color function like this:

myRainbow = Function[x, Blend[{Purple, Blue, Green, Yellow, Red}, x]];

Everything looks fine and dandy. So let's test it out!

nd[y_, a_, b_, n_] := 
  (Binomial[b - a + 1, y] * 
     Sum[((-1)^i)*Binomial[y, i]*((y - i)/(b - a + 1))^n, {i, 0, y}]);
ArrayPlot[Table[nd[3, 1, K, n], {K, 1, 15}, {n, 1, 15}], 
  ColorFunction -> myRainbow, DataReversed -> True]

Output:

Wow, it looks fantastic! It's doing exactly what I thought it would. Great, right? Well, let's try it on a DiscretePlot3D version of the same function:

DiscretePlot3D[nd[3, 1, k, n], {n, 1, 15}, {k, 1, 15}, 
  PlotRange -> {0, 1}, ExtentSize -> Full, ColorFunction -> myRainbow]

Output:

Whoa. Something's not right here. Suddenly Mathematica is deciding to use of the value of n in the color function instead of the output of my function nd. Clearly something hasn't been done right here, but I cannot figure out what. I'm assuming that I didn't define my color function correctly, because if, for example, I use the built-in color functions like "DarkRainbow," both plots behave as expected. So what needs to be changed in my color function definition to get the desired output? Also, if I turn off color function scaling, how do I control which values get mapped to each color?

The main problem here seems to be that I just don't know how to properly define my own custom color functions.

Answer 1

ArrayPlot calls the color function with only the matrix value, cf[z], whereas DiscretePlot3D calls it with the position and the value, cf[x, y, z]. Therefore, you need to select the third component:

rainbow[z_] := Blend[{Purple, Blue, Green, Yellow, Red}, z]
rainbow[_, _, z_] := rainbow[z]

DiscretePlot3D[
 nd[3, 1, k, n], {n, 1, 15}, {k, 1, 15},
 PlotRange -> {0, 1},
 ExtentSize -> Full,
 ColorFunction -> rainbow
 ]

We can look at the source code to see how this works for DarkRainbow:

GeneralUtilities`PrintDefinitions[DiscretePlot3D]]

If we look up how it selects the color function, we find that the relevant piece of code is located in a function called ConstructColorFunction. What this function does, given a color function cf, is essentially return

ColorData[cf, #3] &

if cf is a string, whereas if cf is a function it just returns

cf

So you see the difference – if you pass in a function you have to select to use the third argument yourself. If you pass in a string then there is a function used in DiscretePlot3D which does it automatically.

Answer 2

I can see why this is confusing. Plotting functions try to be too smart, and hide the true way in which they handle their colour functions.

Plotting functions typically pass multiple values to their colour function, e.g. Plot3D passes the x, y and z coordinates. You'll find a table in the ColorFunction documentation:

What is confusing is that if you pass the name of a built-in colour scheme, it applies it to the z value (vertical axis):

Plot3D[x^2 + y^2, {x, -1, 1}, {y, -1, 1}, ColorFunction -> "Rainbow"]

If you pass ColorData[something], it also uses it for the z value.

Plot3D[x^2 + y^2, {x, -1, 1}, {y, -1, 1}, 
 ColorFunction -> ColorData["Rainbow"]]

ColorData["Rainbow"] is a function, so this directly contradicts what the documentation said above! Plot3D should be trying to compute colours with expressions such as ColorData["Rainbow"][1,2,3], which don't even evaluate! In reality, it detects the head ColorData and treats it specially.

If we pass our own function (which may also use ColorData internally if we so wish), something different happens:

Plot3D[x^2 + y^2, {x, -1, 1}, {y, -1, 1}, 
 ColorFunction -> (ColorData["Rainbow"][#] &)]

It colours according to the x axis. Finally, this agrees with the documentation. Note that all Function objects will take any number of arguments, and will simply ignore those that you didn't use within the function body. Thus the y and z values (2nd and 3rd arguments) simply get dropped. Example:

# &[1, 2, 3]
(* 1 *)

If we want to colour according to the z axis, we must use the 3rd argument:

Plot3D[x^2 + y^2, {x, -1, 1}, {y, -1, 1}, 
 ColorFunction -> (ColorData["Rainbow"][#3] &)]

Thus, when using a custom colour function with some plotting function, you need to look up what values that plotting function will be passing to the colour function. DiscretePlot3D is missing from the above table, but you can check its documentation under Examples / Options / ColorFunction, and see that is uses the value values as Plot3D.

Therefore, you need

ColorFunction -> (myRainbow[#3]&)

Don't forget the parentheses around the pure function. ColorFunction -> myRainbow[#3]& is equivalent to (ColorFunction -> myRainbow[#3])& and will result in an error.

The typical way to work with colour functions is to define a function which converts a single number (single argument) to a colour, then build a pure function specific to your desired plotting function that chooses the proper axis. You can even have something like ColorFunction -> (myRainbow[Norm[{##}]]&), ColorFunctionScaling -> False to colour based on the distance from the origin.