python + gnuplot でマンデルブロ集合（その１）

#!/usr/bin/env python
import numpy as np

def columns_save(data, filename='tmp.dat'):
    '''
    PyGnuplotのs関数は等高線を表示させるには力不足なので、別途空白行を出力できるsave関数を作成してデータ出力する。
    s関数のように data = [xs,ys] ではなく data= [[x1,y1],[x2,y2] ...]というようなデータ構造に変更。
    '''
    file = open(filename, 'w')
    columns = len(data)
    rows = len(data[0])
    for i in range(columns):
        if data[i] == []:
            file.write('\n')
            continue
        for j in range(rows):
            file.write(str(data[i][j]))
            file.write(' ')
        file.write('\n')
        if i % 1000 == 0 :
            file.flush()  # write once after every 1000 entries
    file.close()  # write the rest

import PyGnuplot as gp
gp.colss = columns_save

class Mandelbrot:
    #マンデルブロ集合を生成するクラス。
    scale = 2.0
    increments = 0.01
    def __init__(self):
        self._xbase = np.arange(-self.scale,self.scale,self.increments)
        self._ybase = np.arange(-self.scale,self.scale,self.increments)
        self._cbases = []
        self.mandelbrot_set = []

    def __gen_cbases(self):
        #マンデルブロ集合の土台となる複素領域を生成する関数。
        for x in self._xbase:
            self._cbases.append([])
            for y in self._ybase:
               self._cbases[-1].append(complex(x,y)) 
        return

    def mandelbrot(self,c):
        #マンデルブロ集合の力学系の計算をする関数。返り値は反復回数。反復回数の数に応じた等高線がマンデルブロ集合のあの絵になる。
        z = 0j
        for n in range(1,2000):
            z = z**2 + c
            if np.abs(z) > 2.0: return np.log(n)
        return 0

    def gen_mandelbrot(self):
        #マンデルブロ集合を生成する関数。
        self.__gen_cbases()
        for cs in self._cbases:
            for c in cs:
                self.mandelbrot_set.append([c.real, c.imag, self.mandelbrot(c)])
            self.mandelbrot_set.append([])
        return

if __name__ == "__main__":
    m = Mandelbrot()
    m.gen_mandelbrot()
    gp.colss(m.mandelbrot_set)
    gp.c("""
        #gnuplotの設定読み込み。
        set size square;
        set nosurface;
        set contour base;
        set view map;
        """)
    gp.c('splot "tmp.dat" using 1:2:3 with lines')