Search code examples
python3d2dheatmap

Plot 4D data heatmap in Python


hey how can I plot a 2D heatmap in 3D? Now I create a python script to make a 2D Heatmap Plot with data from CSV (CSV format: x,y,z,v). For example:

First csv

0,000;-110,000;110,000;0,101

Second csv

0,000;-66,000;110,000;0,104

Third csv

0,000;-22,000;110,000;0,119

....

In this example, it is a heatmap in xz-plane and I create e.g. five more plots, so that I can insert six xz-plane Plots in a 3D room. In 4D heatmap plot with matplotlib there is a very nice example for doing it. But I don't know how to use it in my case.

import numpy as np
import os
import matplotlib.pyplot as plt
from scipy.interpolate import griddata


'Create a list for every parameter'
x = []
y = []
z = []
v = []

file_path = "path/."

'Insert data from csv into lists'
for root, dirs, files in os.walk(file_path, topdown=False):
   for name in files:
       if name[-4:] != '.csv': continue
       with open(os.path.join(root, name)) as data:
          data = np.genfromtxt((line.replace(',', '.') for line in data), delimiter=";")
          if data[1] == 22: 
            x.append(data[0])
            y.append(data[1])
            z.append(data[2])
            v.append(data[3])

'Create axis data'
xi = np.linspace(min(x), max(x), 1000)
zi = np.linspace(min(z), max(z), 1000)
vi = griddata((x, z), v, (xi[None,:], zi[:,None]), method='cubic')

'Create the contour plot'
CS = plt.contourf(xi, zi, vi, 20, cmap=plt.cm.rainbow)
plt.title("Heatmap xz-plane", y=1.05, 
          fontweight="bold")
plt.xlabel("length x in cm")
plt.xticks(np.arange(0, 201, step=40))
plt.ylabel("height z in cm")
plt.yticks(np.arange(110, 251, step=20))
cbar = plt.colorbar()
cbar.set_label("velocity v in m/s", labelpad=10)
plt.savefig('testplot.png', dpi=400)  
plt.show()

Satisfying the request of @keepAlive wishing to see the result of his untested answer... :

it actually works great :-)

enter image description here

enter image description here


Solution

  • Disclaimer: I am the author of the cited example, so I think that copying/pasting myself is not really a problem.

    Note that your dataset does not look (at least) 3-dimensional. But I will assume there is an unwilling selection bias.

    You first need to aggregate your "points" per level of altitude, which I assume is the third component of your vectors. They will be constitutive of your planes once gathered.

    # libraries
    from mpl_toolkits.mplot3d import Axes3D
    import matplotlib.pyplot as plt
    import scipy.interpolate as si
    from matplotlib import cm
    import collections as co  # <------------------
    import pandas as pd
    import numpy as np
    
    planes = co.defaultdict(list)
    
    for root, dirs, files in os.walk(file_path, topdown=False):
       # [...]
           # [...]
           # [...]
              # [...]
              # [...]
                level = data[2]  # <------ third component.
                planes[level].append(data)
    

    Now, at that stage, we have a list of arrays per level. Let's define our grids_maker function

    def grids_maker(arrays_list, colnames=list('xyzg')):
        # 0- The idea behind `list('xyzg')` is only to change the order
        #    of names, not the names as such. In case for example you
        #    want to use another component than the third to organize
        #    your planes.
        # 1- Instantiate a dataframe so as to minimize the modification
        #    of the function copied/pasted pasted from
        #    https://stackoverflow.com/a/54075350/4194079
        # 2- Pandas is also going to do some other jobs for us, such as
        #    stacking arrays, etc....
        df = pd.DataFrame(arrays_list, columns=colnames)
    
        # Make things more legible
        xy = df.loc[:, ['x', 'y']]
        x  = xy.x
        y  = xy.y
        z  = df.z
        g  = df.g
        reso_x = reso_y = 50
        interp = 'cubic' # or 'nearest' or 'linear'
    
        # Convert the 4d-space's dimensions into grids
        grid_x, grid_y = np.mgrid[
            x.min():x.max():1j*reso_x,
            y.min():y.max():1j*reso_y
        ]
    
        grid_z = si.griddata(
            xy, z.values,
            (grid_x, grid_y),
            method=interp
        )
    
        grid_g = si.griddata(
            xy, g.values,
            (grid_x, grid_y),
            method=interp
        )
    
        return {
            'x' : grid_x,
            'y' : grid_y,
            'z' : grid_z,
            'g' : grid_g,
        }
    

    Let's use grids_maker over our list of arrays and get the extrema of each z-level's 4th dimension.

    g_mins = []
    g_maxs = []
    lgrids = {}
    
    for level, arrays_list in planes.items():
        lgrids[level] = grids = grids_maker(arrays_list)
        g_mins.append(grids['g'].min())
        g_maxs.append(grids['g'].max())
    

    Let's create our (all-file unifying) color-scale and show the plot.

    # Create the 4th color-rendered dimension
    scam = plt.cm.ScalarMappable(
        norm=cm.colors.Normalize(min(g_mins), max(g_maxs)),
        cmap='jet' # see https://matplotlib.org/examples/color/colormaps_reference.html
    )
    fig = plt.figure()
    ax  = fig.gca(projection='3d')
    for grids in lgrids.values(): 
        scam.set_array([])   
        ax.plot_surface(
            grids['x'], grids['y'], grids['z'],
            facecolors  = scam.to_rgba(grids['g']),
            antialiased = True,
            rstride=1, cstride=1, alpha=None
        )
    plt.show()
    

    I would be glad to see the result.