How to plot a 2 dimensional density / heatmap in plotnine?

Question

I am trying to recreate the following graph in plotnine. It's asking me for more details but I don't want to distract from the example. I think it's pretty obvious what I'm trying to do. I have been given a function by a colleague. I'm not interested in rewriting the function. I want to take sm and use plotnine to plot it instead of matplotlib. I plot lots of dataframes with plotnine but I'm not sure how to use it in this case. I have tried on my own to figure it out and I keep getting lost. I hope that for someone more experienced I am overlooking something simple.

import matplotlib.pyplot as plt

def getSuccess(y,x):
    return((y*(-x))*.5+.5)

steps = 100
stepSize = 1/steps

sm = []

for y in range(steps*2+1):
    sm.append([getSuccess((y-steps)*stepSize,(x-steps)*stepSize) for x in range(steps*2+1)])

plt.imshow(sm)

plt.ylim(-1, 1)

plt.colorbar()
plt.yticks([0,steps,steps*2],[str(y) for y in [-1.0,0.0,1.0]])
plt.xticks([0,steps,steps*2],[str(x) for x in [-1.0,0.0,1.0]])
plt.show()

Answer 1

You could try geom_raster.

I have taken your synthetic data sm and converted to a dataframe as plotnine will need this.

import pandas as pd
import numpy as np
from plotnine import *

df = pd.DataFrame(sm).melt()
df.rename(columns={'variable':'x','value':'density'}, inplace=True)
df.insert(1,'y',df.index % 201)

p = (ggplot(df, aes('x','y'))
  + geom_raster(aes(fill='density'), interpolate=True)
  + labs(x=None,y=None)
  + scale_x_continuous(expand=(0,0), breaks=[0,100,200], labels=[-1,0,1])
  + scale_y_continuous(expand=(0,0), breaks=[0,100,200], labels=[-1,0,1])
  + theme_matplotlib()
  + theme(
      text = element_text(family="Calibri", size=9),
      legend_title = element_blank(),
      axis_ticks = element_blank(),
      legend_key_height = 29.6,
      legend_key_width = 6,
  )
)
p.save(filename='C:\\Users\\BRB\\geom_raster.png', height=10, width=10, units = 'cm', dpi=400)

This result is: