heatmap.py

#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
from numexpr import evaluate
from math import pi

def testset():
    count = 450
    rcoords = np.random.rand(count, 3)
    rcoords[:,2] = rcoords[:,2] * 0.005 + 0.005
    return rcoords

def main():
    xd, yd = 1024, 1024
    img = np.zeros((yd,xd))
    samx, samy = np.meshgrid(np.linspace(0,1,xd), np.linspace(0,1,yd))
    for px, py, sz in testset():
        img = evaluate("img + (sz * 0.01) * (1.0 / ((samx - px) ** 2 + (samy - py) ** 2))")
    np.clip(img, 0, 1, img)
    mcmap = plt.get_cmap('hot')
    mcmap._init()
    mcmap._lut[:-3,-1] = np.clip(np.linspace(0,2,mcmap.N), 0, 1)
    plt.imsave('bleh.png', img, cmap = mcmap)

if __name__ == "__main__":
    main()