Littlewood Polynomial Heatmap

Question

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

Answer 1

As Oscar mentions, switching to python3 would probably make this easier, just because it's better designed for collection-like things that aren't lists. However, I don't think it's strictly necessary.

The first thing that I notice is that you're caching an awful lot of data into that temp array. The temp array itself, as far as I can tell, doesn't need to exist at all: you can just pack the numbers straight into x and y.

However, that would at best halve your memory usage, which isn't going to cut it here. Fundamentally, for the operations that you're doing, you should not need to store the x and y values either. Instead, you just want to tally up the histogram.

I suggest that you implement your own version of histogram2d. I'll grant that home-rolling a python function that works at even comparable speed to numpy is usually downright impossible. In this case, however, you just need to go faster than the combined numpy function and allocating and packing 12GB of data! That is likely perfectly tractable.

If in fact you can't get the performance that you want without letting numpy work its magic, consider batching up instead. Get the first batch of X thousand array elements, build a histogram, discard that batch and get the next X thousand, build another histogram, sum the histograms, and repeat. Looping against a generator and yield is probably the easiest way to go about doing this.

Because your program works fine with 2**16 in the first loop, and nothing beyond generating the histogram needs memory allocated in proportion to that loop size, that should solve the issue for essentially as big a number as you have the patience to run the program for.

Answer 2

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.