Home > Software design >  Converting 1D distribution into matplotlib gradient
Converting 1D distribution into matplotlib gradient

Time:08-22

I have a 1D distribution (x values vs probability, shown below) and I would like to convert that to a 2D plot like the one shown below in which the color gradient is based on the values probabilities.

Currently, my code just plot in a qualitative manner because I am manually defining the array v1 and the color list. I tried my best to crack this and understand how to do it, but I failed. Does anyone have a suggestion?

def gradient_image(ax, extent, direction=0.3, cmap_range=(0, 1), **kwargs):
"""
Draw a gradient image based on a colormap.

Parameters
----------
ax : Axes
    The axes to draw on.
extent
    The extent of the image as (xmin, xmax, ymin, ymax).
    By default, this is in Axes coordinates but may be
    changed using the *transform* keyword argument.
direction : float
    The direction of the gradient. This is a number in
    range 0 (=vertical) to 1 (=horizontal).
cmap_range : float, float
    The fraction (cmin, cmax) of the colormap that should be
    used for the gradient, where the complete colormap is (0, 1).
**kwargs
    Other parameters are passed on to `.Axes.imshow()`.
    In particular useful is *cmap*.
"""
phi = direction * np.pi / 2
v = np.array([np.cos(phi), np.sin(phi)])
X = np.array([[v @ [1, 0], v @ [1, 1]],
              [v @ [0, 0], v @ [0, 1]]])
a, b = cmap_range
X = a   (b - a) / X.max() * X
im = ax.imshow(X, extent=extent, interpolation='bicubic',
               vmin=0, vmax=1, **kwargs)
return im

v1 = [0, 0.15, 0.5, 0.85, 1.0]                                # | Those two lines here
b = ["white","lightblue", "dodgerblue","lightblue", "white"]  # | were the best I could do
bl = list(zip(v1,b))
blue_grad=LinearSegmentedColormap.from_list('custom',bl, N=256)

xmin, xmax = xlim = 0, 4
ymin, ymax = ylim = -300, 300

fig, ax = plt.subplots()
ax.set(xlim=xlim, ylim=ylim, autoscale_on=False)

gradient_image(ax, direction=1, extent=(0 ,  2, -300, 300), cmap=blue_grad,  cmap_range=(0., 1), alpha=0.5)

enter image description here enter image description here

CodePudding user response:

Here is a minimal example with a gaussian distribution (code for generating the gaussian distribution was adapted from enter image description here

  • Related