Search code examples
pythonmatplotlibkernel-density

3D data contour ploting using a kde


I have two Arrays of positional Data (X,Y) and a corresponding 1D Array of Integers (Z) that weighs the positional Data. So my Data set looks like that:

X = [ 507, 1100, 1105, 1080, 378, 398, 373]
Y = [1047,  838,  821,  838, 644, 644, 659]
Z = [ 300,   55,   15,   15,  55,  15,  15] 

I want to use that Data to create a KDE thats equivalent to a KDE that gets only X and Y as input but gets the X and Y values Z times. To apply that KDE to a np.mgrid to create a contourplot.

I already got it working by just iterating over the arrays in a FOR Loop and adding Z times X and Y, but that looks to me like a rather inelegant Solution and I hope you can help me to find a better way of doing this.


Solution

  • You could use the weights= parameter of scipy.stats.gaussian_kde:

    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import axes3d
    import numpy as np
    from scipy import stats
    
    X = [ 507, 1100, 1105, 1080, 378, 398, 373]
    Y = [1047,  838,  821,  838, 644, 644, 659]
    Z = [ 300,   55,   15,   15,  55,  15,  15]
    
    kernel = stats.gaussian_kde(np.array([X, Y]), weights=Z)
    
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")
    xs, ys = np.mgrid[0:1500:30j, 0:1500:30j]
    zs = kernel(np.array([xs.ravel(), ys.ravel()])).reshape(xs.shape)
    ax.plot_surface(xs, ys, zs, cmap="hot_r", lw=0.5, rstride=1, cstride=1, ec='k')
    plt.show()
    

    resulting plot