Python: how to create a smoothed version of a 2D binned "color map"?

Question

I would like to create a version of this 2D binned "color map" with smoothed colors.

I am not even sure this would be the correct nomenclature for the plot, but, essentially, I want my figure to be color coded by the median values of a third variable for points that reside in each defined bin of my (X, Y) space.

Even though I am able to accomplish that to a certain degree (see example), I would like to find a way to create a version of the same plot with a smoothed color gradient. That would allow me to visualize the overall behavior of my distribution.

I tried ideas described here: Smoothing 2D map in python

and here: Python: binned_statistic_2d mean calculation ignoring NaNs in data

as well as links therein, but could not find a clear solution to the problem.

This is what I have so far:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from scipy.stats import binned_statistic_2d
import random
random.seed(999)

x = np.random.normal (0,10,5000)
y = np.random.normal (0,10,5000)
z = np.random.uniform(0,10,5000)

fig = plt.figure(figsize=(20, 20))
plt.rcParams.update({'font.size': 10})

ax = fig.add_subplot(3,3,1)

ax.set_axisbelow(True)
plt.grid(b=True, lw=0.5, zorder=-1)

x_bins = np.arange(-50., 50.5, 1.)
y_bins = np.arange(-50., 50.5, 1.)

cmap = plt.cm.get_cmap('jet_r',1000) #just a colormap

ret = binned_statistic_2d(x, y, z, statistic=np.median, bins=[x_bins, y_bins]) # Bin (X, Y) and create a map of the medians of "Colors"

plt.imshow(ret.statistic.T, origin='bottom', extent=(-50, 50, -50, 50), cmap=cmap) 

plt.xlim(-40,40)
plt.ylim(-40,40)
plt.xlabel("X", fontsize=15)
plt.ylabel("Y", fontsize=15)
ax.set_yticks([-40,-30,-20,-10,0,10,20,30,40])

bounds = np.arange(2.0, 20.0, 1.0)
plt.colorbar(ticks=bounds, label="Color", fraction=0.046, pad=0.04)

# save plots
plt.savefig("Whatever_name.png", bbox_inches='tight')

Which produces the following image (from random data):

Therefore, the simple question would be: how to smooth these colors?

Thanks in advance!

PS: sorry for excessive coding, but I believe a clear visualization is crucial for this particular problem.

Answer 1

Thanks to everyone who viewed this issue and tried to help!

I ended up being able to solve my own problem. In the end, it was all about image smoothing with Gaussian Kernel.

This link: Gaussian filtering a image with Nan in Python gave me the insight for the solution.

I, basically, implemented the exactly same code, but, in the end, mapped the previously known NaN pixels from the original 2D array to the resulting smoothed version. Unlike the solution from the link, my version does NOT fill NaN pixels with some value derived from the pixels around. Or, it does, but then I erase those again.

Here is the final figure produced for the example I provided:

Final code, for reference, for those who might need in the future:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from scipy.stats import binned_statistic_2d
import scipy.stats as st
import scipy.ndimage
import scipy as sp
import random
random.seed(999)

x = np.random.normal (0,10,5000)
y = np.random.normal (0,10,5000)
z = np.random.uniform(0,10,5000)

fig = plt.figure(figsize=(20, 20))
plt.rcParams.update({'font.size': 10})

ax = fig.add_subplot(3,3,1)

ax.set_axisbelow(True)
plt.grid(b=True, lw=0.5, zorder=-1)

x_bins = np.arange(-50., 50.5, 1.)
y_bins = np.arange(-50., 50.5, 1.)

cmap = plt.cm.get_cmap('jet_r',1000) #just a colormap

ret = binned_statistic_2d(x, y, z, statistic=np.median, bins=[x_bins, y_bins]) # Bin (X, Y) and create a map of the medians of "Colors"

sigma=1                    # standard deviation for Gaussian kernel
truncate=5.0               # truncate filter at this many sigmas

U = ret.statistic.T.copy()

V=U.copy()
V[np.isnan(U)]=0
VV=sp.ndimage.gaussian_filter(V,sigma=sigma)

W=0*U.copy()+1
W[np.isnan(U)]=0
WW=sp.ndimage.gaussian_filter(W,sigma=sigma)

np.seterr(divide='ignore', invalid='ignore')

Z=VV/WW

for i in range(len(Z)):
    for j in range(len(Z[0])):
        if np.isnan(U[i][j]):
             Z[i][j] = np.nan

plt.imshow(Z, origin='bottom', extent=(-50, 50, -50, 50), cmap=cmap) 

plt.xlim(-40,40)
plt.ylim(-40,40)
plt.xlabel("X", fontsize=15)
plt.ylabel("Y", fontsize=15)
ax.set_yticks([-40,-30,-20,-10,0,10,20,30,40])

bounds = np.arange(2.0, 20.0, 1.0)
plt.colorbar(ticks=bounds, label="Color", fraction=0.046, pad=0.04)

# save plots
plt.savefig("Whatever_name.png", bbox_inches='tight')