Plotting a function with 4 arguments in a grid of heatmaps

Question

I am using plotnine in python 2 (but would be happy for a solution using matplotlib or any other plotting package). I have a function (slightly simplified below) with 4 arguments. I want to plot a grid of heatmaps, with "super axes" varying two of the parameters, and each heatmap varying the other two. A bit like this:

I came up with this code:

from plotnine import ggplot, aes, facet_grid, geom_tile
import pandas as pd
import itertools

def fun((i, j, n, p)):
    if n > j:
        return 1 if (p**(3*n)) > ((p+i)**(3*(n-j))) else 0
    return -1

ilist, jlist, nlist, plist = range(1,10), range(1,9), range(8,10), range(4,6)
rows = itertools.product(ilist, jlist, nlist, plist)

df = pd.DataFrame(list(rows))
df.columns = ['i','j','n','p']
df['fun'] = df.apply(fun, axis=1)

(ggplot(df, aes('factor(i)', 'factor(j)')) +
 geom_tile(aes(fill='fun')) +
facet_grid('n ~ p', labeller='label_both'))

This produces the following:

This (more or less) has the form I am going for, but the heatmaps seem wrong. (EDIT: I discovered the error was in the definition of fun, not in the plotting - see my answer).

score 0 · Answer 1 · answered Jun 01 '18 at 14:58

OK... it turned out not to be a problem with the plotting, but with the function. I don't fully understand what went wrong, but I think some of the computations were running into "huge integer" problems. When I ran it to make the comparison in log space it started working.

I'll leave this here as it might help others who are trying to make similar grids of heatmaps, and maybe someone will answer with a better way of doing it.

Here's my annotated code to make the heatmap itself from a dataframe with columns ['i', 'j', 'n', 'p']:

# The first two lines are defining the subplots, i.e. the heatmaps
# themselves. I am passing the dataframe, and specifying the names of the
# columns to use as axes. I say 'factor(i)' etc to treat the i column as
# discrete, not continuous (in my case it's integers).
gg = (ggplot(df, aes('factor(i)', 'factor(j)')) +
    geom_tile(aes(fill='fun')) +
# The last bit is to call facet_grid which applies the above code
# in a grid. The parameter 'n ~ p' specifies that I want the grid to be
# over the columns 'n' and 'p' from the dataframe. The labeller
# parameter is what makes the labels at the edges (see top and right in
# image below) show both the column name and the value).
    facet_grid('n ~ p', labeller='label_both'))

Here's the full corrected code with the fixed result:

from plotnine import ggplot, aes, facet_grid, geom_tile
import pandas as pd
import numpy as np
import itertools

def fun((i, j, n, p)):
    if n > j:
        return 1 if np.log10(p) * (3*n) > np.log10(p+i) * (3*(n-j)) else 0

    return -1

ilist, jlist, nlist, plist = range(1,10), range(1,9), range(8,10), range(4,6)
rows = itertools.product(ilist, jlist, nlist, plist)
df = pd.DataFrame(list(rows))
df.columns = ['i','j','n','p']
df['fun'] = df.apply(fun, axis=1)

gg = (ggplot(df, aes('factor(i)', 'factor(j)')) +
    geom_tile(aes(fill='fun')) +
    facet_grid('n ~ p', labeller='label_both'))

gg.draw()

Plotting a function with 4 arguments in a grid of heatmaps

1 Answers1