来自3d坐标的Python热图

Question

I start from two linspaces and I meshgrid them. Then I calculate a function's values on grid. My function is called cpt_hcpv() . Then I would like to heatmap my data, with each point on the grid with its corresponding function value.

Code looks like

poro = np.linspace(min(poro), max(poro))
sw = np.linspace(min(sw), max(sw))

g = np.meshgrid(poro, sw)
points = zip(*(x.flat for x in g))

hcpv = []
for p in points:
    hcpv = hcpv + [cpt_hcpv(p[0], p[1], poro, sw)]

with

def cpt_hcpv(pCut, sCut, poro, sw):
    #find points belonging to calculation
    truncated = [(p, s) for p, s in zip(poro, sw) if p > pCut and s < sCut ]
    hcv = 0
    for k in truncated:
        hcv += p*(1-s)*0.5
    return hcv

Why I am not computing cpt_hcpv() directly on grid: because I have to deal with condition in comprehension truncated = [(p, s) for p, s in zip(poro, sw) if p > pCut and s < sCut ] so that I must iterate on the point in grid. I don't know how to iterate on a meshgrid.

So, I would like to heatmap from the 3d coordinates : in points I have x and y for the points and in hcpv I have the z parameters for each point, in same order.

From the examples I have found, there are pylab and matplotlib solutions to plot heatmap from meshgrid + values computed on the grid, with a method taking meshgrid as an argument.

Is there a way to plot heatmap from 3d coordinates ?

Answer 1

If you need to iterate over a meshgrid, try this:

g = np.meshgrid(poro, sw)

#Turn into 2x3x3 array
g_arr = np.array(g)                            

#Move first dimension to third: 3x3x2 array
g_arr = g_arr.swapaxes(0,1).swapaxes(1,2)

#Generate results by iterating over first and second dimension, and unpacking the third
hcpv = np.array([[cpt_hcpv(p, s, poro, sw) for p,s in r] for r in g_arr])

I don't know if matplotlib is going to have easy ploting for heatmaps from generic 3-d points. It would have to handle the generic case of scattered, out-of-order and missing points.

Answer 2

I came up with this solution using DrRobotNinja approach

g = np.meshgrid(poro_g, sw_g)
g_arr = np.array(g)                            
g_arr = g_arr.swapaxes(0,1).swapaxes(1,2)

#I compute z value on the grid, `g_arr`
hcpv = np.array([[cpt_hcpv(p, s, poro, sw) for p,s in r] for r in g_arr])

I superimpose heatmap and contours (levels)

#heatmap
im = plt.imshow(hcpv, origin='lower',extent=[min(poro)-EPS,max(poro)
                                    +EPS,min(sw)-EPS,max(sw)+EPS],aspect='auto')

#contours
levels = np.array([p10,p50,p90])
cset = plt.contour(hcpv,levels,linewidths=2,cmap=pylab.cm.hot, 
      origin='lower',extent=[min(poro),max(poro),min(sw),max(sw)],aspect='auto')
plt.clabel(cset,inline=True,fmt='%1.1f',fontsize=20)

and display

plt.show()