Littlewood Polynomial Heatmap The following Python program generates heat maps of the roots of Littlewood polynomials. It works fine with a small number of roots, however I tried to use 2^21 in the first loop and it ate up 12 gigabytes of RAM on my computer. How can I make the code less memory-intensive while still allowing reasonably fast performance? Ideally, I would like to use this code to generate very large images but I cannot do so with the program in its current state. Any help would be greatly appreciated. Thanks!

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from math import sqrt, sin, log10
def f(x): # function that maps [0,1] to itself for the color space so that values closer to 1 end up further from 1
    return sqrt(x)
def roots(n): # given n, returns a Littlewood polynomial with coefficients representing the binary representation of n
    roots = np.polynomial.polynomial.polyroots([int(i)*2-1 for i in bin(n)[2:]])
    return [[i.real,i.imag] for i in roots]
# get array of all possible roots and split values up into arrays x and y
temp=[]
for i in xrange(2**16):
    if i%100==0: print i
    for j in roots(i):
        temp.append(j)
x=[i[0] for i in temp]
y=[i[1] for i in temp]

# generate histogram of roots to make heatmap
gap=0.01
print len(np.arange(-2.5,2.5,gap))
normal_histogram = np.histogram2d(x,y,bins=np.arange(-2.5,2.5,gap),normed=True)

# generate image 
from PIL import Image
normal_histogram = normal_histogram[0]
size=len(normal_histogram)
im = Image.new("RGB",(size,size))
im2 = Image.new("RGB",(size,size))
#color = [int(256*k) for k in cm.hot(histogram[i][j])[0:3]]
for i in range(len(normal_histogram)):
    print i
    for j in range(len(normal_histogram)):
        temp = int(normal_histogram[i][j]*256)
        im.putpixel((i,j),tuple([int(256*k) for k in cm.hot(f(normal_histogram[i][j]))][0:3]))
        im2.putpixel((i,j),tuple([int(256*k) for k in cm.nipy_spectral(f(normal_histogram[i][j]))][0:3]))
im.save("test.png")
im2.save("test2.png")


Things I've tried


using a txt file to store the polynomial roots
using matplotlib's ability to make heatmaps to plot the data
Here are a few changes that make image generation near instant, and slightly speed up root generation, and lower memory overhead.

def roots(n): # given n, returns a Littlewood polynomial with coefficients representing the binary representation of n
    return np.polynomial.polynomial.polyroots([int(i)*2-1 for i in bin(n)[2:]])


The only change here is to just return the roots rather than messing with them, so we pass less data around.

x, y = [], []
for i in range(0,2**16):
    if i%10000==0:
        print(i)
    root = roots(i)
    x.extend(np.real(root))
    y.extend(np.imag(root))


Here I removed temp to halve memory use, and used vectorization where applicable for small speed boosts.
the major changes are in the image generation

hot = np.asarray(128*cm.hot(normal_histogram)[:,:,0:3], dtype=np.uint8)
im = Image.fromarray(hot, mode="RGB")
spectral = np.asarray(128*cm.nipy_spectral(normal_histogram)[:,:,0:3], dtype=np.uint8)
im2 = Image.fromarray(spectral, mode="RGB")
im.show()
im2.show()


here, rather than making a blank image and filling it bit by bit, we make an image from an array created using vectorization. This takes 3 seconds with gap = 0.001 instead of over 16 times that before. This is far from optimized, vectorizing roots would help a lot, but that is hard.