Plot heat map from csv file using numpy and matplotlib 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.
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)