Plot heat map from csv file using numpy and matplotlib

Question

There's a csv file with format:

x0, y0, v00
x0, y1, v01
...
x1, y0  v10
...

And what I want to do is to plot a heat map, in which at location (x, y) the value v is plotted with corresponding color. Below is my current implementation.

import random
import numpy as np
import matplotlib.pyplot as plt

def create_test_csv(file):
    random.seed(42)
    f = open(file, "w")
    for x in range(300):
        for y in range(600):
            value = random.randrange(255)
            f.write(str(x) + "," + str(y) + "," + str(value) + "\n")


def get_xyz_from_csv_file(csv_file_path):
    '''
    get x, y, z value from csv file
    csv file format: x0,y0,z0
    '''
    x = []
    y = []
    z = []
    map_value = {}

    for line in open(csv_file_path):
        list = line.split(",")
        temp_x = float(list[0])
        temp_y = float(list[1])
        temp_z = float(list[2])
        x.append(temp_x)
        y.append(temp_y)
        z.append(temp_z)
        map_value[(temp_x, temp_y)] = temp_z

    return x, y, map_value


def draw_heatmap(x, y, map_value):

    plt_x = np.asarray(list(set(x))) 
    plt_y = np.asarray(list(set(y)))  
    plt_z = np.zeros(shape = (len(plt_x), len(plt_y)))

    for i in range(len(plt_x)):
        for j in range(len(plt_y)):
            if map_value.has_key((plt_x.item(i), plt_y.item(j))): 
                plt_z[i][j] = map_value[(plt_x.item(i), plt_y.item(j))]    

    z_min = plt_z.min()
    z_max = plt_z.max()
    plt_z = np.transpose(plt_z)    

    plot_name = "demo"


    color_map = plt.cm.gist_heat #plt.cm.rainbow #plt.cm.hot #plt.cm.gist_heat
    plt.clf()
    plt.pcolor(plt_x, plt_y, plt_z, cmap=color_map, vmin=z_min, vmax=z_max)
    plt.axis([plt_x.min(), plt_x.max(), plt_y.min(), plt_y.max()])
    plt.title(plot_name)
    plt.colorbar().set_label(plot_name, rotation=270)   
    ax = plt.gca()
    ax.set_aspect('equal')
    figure = plt.gcf()
    plt.show()
    return figure    


if __name__ == "__main__":
    csv_file_name = "test.csv"
    create_test_csv(csv_file_name)
    x, y, map_value = get_xyz_from_csv_file(csv_file_name)
    draw_heatmap(x, y, map_value)

Function create_test_csv() created a test csv file. Function get_xyz_from_csv_file() create x, y coordinates list and a dict which key is tuple (x,y) and value is v. Function draw_heatmap() plot the heat map using list x, y and dict map_value.

It works but I would like to know if there is some more straightforward way to this, especially the transition from CSV to the matrix that created the heat map. It might worth to notice that in my real case the coordinate may not be integer.

Answer 1

Since you are already using numpy, you can use numpy's loadtxt function to read in all the data at once as numpy arrays from the start. This allows you to avoid having to worry about opening or closing files (this is done automatically), converting to numpy arrays, etc. Then it is a simple matter of converting the indexes to values.

You can also vectorize the test data creation, using numpy's meshgrid function to get a grid of corresponding X and Y coordinates.

You can make the plotting better, in my opinion at least, by using plt.subplots() to get a figure and axes object right at the beginning, then using those to do the plotting.

So here is how I would do it:

import numpy as np
import matplotlib.pyplot as plt


def create_test_csv(fname):
    np.random.seed(42)

    # Generate X and Y coordinates
    x = np.arange(300)
    y = np.arange(600)

    # Get corresponding X and Y coordinates
    xs, ys = np.meshgrid(x, y)

    # Get random values for each location
    zs = np.random.randint(0, 255, size=xs.size)

    # Convert 3 2D arrays to 1 2D array of columns
    data = np.vstack([xs.ravel(), ys.ravel(), zs.ravel()]).T

    # Save to file
    np.savetxt(fname, data, delimiter=',', fmt='%d')


def get_xyz_from_csv_file_np(csv_file_path):
    '''
    get a grid of values from a csv file
    csv file format: x0,y0,z0
    '''

    # Load the csv file into a single 2D array, 
    # then split the columns into individual variables.
    x, y, z = np.loadtxt(csv_file_path, delimiter=',', dtype=np.int).T

    # Create an empty 2D array of pixels and 
    # put all the values into the correct place
    plt_z = np.zeros((y.max()+1, x.max()+1))
    plt_z[y, x] = z

    return plt_z


def draw_heatmap(plt_z):
    # Generate y and x values from the dimension lengths
    plt_y = np.arange(plt_z.shape[0])
    plt_x = np.arange(plt_z.shape[1])

    # everything is the same from here on
    z_min = plt_z.min()
    z_max = plt_z.max() 

    plot_name = "demo"

    color_map = plt.cm.gist_heat #plt.cm.rainbow #plt.cm.hot #plt.cm.gist_heat
    fig, ax = plt.subplots()
    cax = ax.pcolor(plt_x, plt_y, plt_z, cmap=color_map, vmin=z_min, vmax=z_max)
    ax.set_xlim(plt_x.min(), plt_x.max())
    ax.set_ylim(plt_y.min(), plt_y.max())
    fig.colorbar(cax).set_label(plot_name, rotation=270) 
    ax.set_title(plot_name)  
    ax.set_aspect('equal')
    plt.show()
    return figure
    figure = plt.gcf()
    plt.show()
    return figure   


if __name__ == "__main__":
    fname = 'temp.csv'
    create_test_csv(fname)
    res = get_xyz_from_csv_file_np(fname)
    draw_heatmap(res)

Answer 2

One function at a time:

def create_test_csv(file):
    random.seed(42)
    with open(file, "w") as f:
        for x in range(300):
            for y in range(600):
                print('{:},{:},{:}'.format(str(x), str(y), str(random.randrange(255))), file=f)

Advantages:

• Using a "with" statement allows python to do the dirty work when it comes to manipulating files. It will close the file after writing it in this case.
• No need to assign a new variable "value" since you only use it once. Pass it directly to the write/print command. You do not need to have this value stored.

In the case of "get_xyz_from_csv_file" it seems you are getting the same thing twice with the list and dict representation. You can just build one of the two and jump to the other when you need to with a list comprehension statement but i will leave that for now.

def get_xyz_from_csv_file(csv_file_path):
    '''
    get x, y, z value from csv file
    csv file format: x0,y0,z0
    '''
    a_list = []
    with open(csv_file_path) as pic:
        for line in pic:
            a_list.append([(float(line.split(',')[0]), float(line.split(',')[1])), float(line.split(',')[2])])
        a_dict = {x: y for x, y in a_list}
    return a_dict 
# you can get the Xs & Ys like so:
# b = get_xyz_from_csv_file(csv_file_name)
# x = sorted([x[0] for x in list(b.keys())])
# y = sorted([y[1] for y in list(b.keys())])

As for the "draw_heatmap" function, i don't really have many things to add but i would generate the vectors X and Y from the dict returned before so that the if check is no longer needed.

def draw_heatmap(dict_value):

    plt_x = np.asarray(list(set(([x[0] for x in dict_value]))))
    plt_y = np.asarray(list(set(([y[1] for y in dict_value]))))
    plt_z = np.zeros(shape=(len(plt_x), len(plt_y)))

    for i_x in range(len((plt_x))):
        for i_y in range(len((plt_y))):
            plt_z[i_x][i_y] = dict_value[(plt_x.item(i_x), plt_y.item(i_y))]

    z_min = plt_z.min()
    z_max = plt_z.max()
    plt_z = np.transpose(plt_z)

    plot_name = "demo"


    color_map = plt.cm.gist_heat #plt.cm.rainbow #plt.cm.hot #plt.cm.gist_heat
    plt.clf()
    plt.pcolor(plt_x, plt_y, plt_z, cmap=color_map, vmin=z_min, vmax=z_max)
    plt.axis([plt_x.min(), plt_x.max(), plt_y.min(), plt_y.max()])
    plt.title(plot_name)
    plt.colorbar().set_label(plot_name, rotation=270)
    ax = plt.gca()
    ax.set_aspect('equal')
    figure = plt.gcf()
    plt.show()
    return figure

And finally,

if __name__ == "__main__":
    csv_file_name = "test.txt"
    create_test_csv(csv_file_name)
    the_values_dict = get_xyz_from_csv_file(csv_file_name)
    draw_heatmap(the_values_dict)