How to plot a 4-parameter function in a grid of heatmaps in Python?

python

#1

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.

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:

Heatmap grid

This (more or less) has the form I am going for, but the heatmaps seem wrong. For example, looking at the bottom right plot (n: 9, p: 5), I ran this to check what it should have looked like:

n, p = 9, 5
for j in range(1,9):
    print [fun((i,j,n,p)) for i in range(1, 10)]

This gives:

[1, 0, 0, 0, 0, 0, 0, 0, 0]
[1, 1, 0, 0, 0, 0, 0, 0, 0]
[1, 1, 1, 1, 1, 1, 0, 0, 0]
[1, 1, 1, 1, 1, 1, 1, 1, 1]
[1, 1, 1, 1, 1, 1, 1, 1, 1]
[1, 1, 1, 1, 1, 1, 1, 1, 1]
[1, 1, 1, 1, 1, 1, 1, 1, 1]
[1, 1, 1, 1, 1, 1, 1, 1, 1]

Does anyone have any idea what’s going wrong and how I can fix it? :slight_smile: <3


#2

Whoops!! 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()

Corrected heatmap