fractal_markus_lyapunov.py

#!/usr/bin/env python
# Licensed under the Apache License, Version 2.0

import argparse
import multiprocessing
import os, signal, time
import numpy as np
import colorsys
import png

NUM_CPU=multiprocessing.cpu_count()

# Render size
SIZE = float(os.environ.get("RENDER_SIZE", 1))
WINSIZE = map(lambda x: int(x*SIZE), [ 100,  100])
CENTER  = map(lambda x: x/2, WINSIZE)


def compute_markus_lyapunov(param):
    window_size, offset, scale, seed, x0, max_iter, max_init, step_size, chunk = param
    # Return numpy array of image pixels

    results = np.zeros(step_size, dtype='i4')
    pos = 0

    while pos < step_size:
        step_pos = pos + chunk * step_size
        screen_coord = (step_pos / window_size[1], step_pos % window_size[1])
        c = np.complex128(complex(
            screen_coord[0] / scale[0] + offset[0],
            ((window_size[1] - screen_coord[1]) / scale[1] + offset[1])
        ))
        markus_func = lambda x: c.real if seed[idx % len(seed)] == "A" else c.imag

        # Init
        x = np.float128(x0)
        for idx in xrange(0, max_init):
            r = markus_func(idx)
            with np.errstate(over='ignore'):
                x = r * x * ( 1 - x )

        # Exponent
        total = np.float64(0)
        for idx in range(0, max_iter):
            r = markus_func(idx)
            with np.errstate(over='ignore'):
                x = r * x * ( 1 - x )
            v = abs(r - 2 * r * x)
            if v == 0:
                break
            total = total + np.log(v) / np.log(1.23)

        if total == 0 or total == float('Inf'):
            exponent = 0
        else:
            exponent = total / float(max_iter)
        results[pos] = exponent
        pos += 1
    return results


def usage():
    parser = argparse.ArgumentParser()
    parser.add_argument("scene")
    parser.add_argument("--action", choices=["compute", "render"])
    parser.add_argument("--size", type=float, default=1.)
    parser.add_argument("--length", type=int, default=8)
    parser.add_argument("--seed", type=str, default="AB")
    parser.add_argument("--center", type=str, default="2+2j")
    parser.add_argument("--radius", type=float, default=2.)
    parser.add_argument("--x0", type=float, default=0.5)
    return parser.parse_args()

# Multiprocessing abstraction
pool = multiprocessing.Pool(NUM_CPU, lambda : signal.signal(signal.SIGINT, signal.SIG_IGN))
#pool = None
def compute(method, params):
    if pool:
        params[-1] /= NUM_CPU
        params = map(lambda x: params + [x], range(NUM_CPU))
        # Compute
        res = pool.map(method, params)
        # Return flatten array
        return np.array(res).flatten()
    return method(params + [0])


def hsv(h, s, v):
    r, g, b = colorsys.hsv_to_rgb(h, s, v)
    return int(b * 0xff) | int((g * 0xff)) << 8 | int(r * 0xff) << 16


def color_factory(size, base_hue = 0.65):
    def color_scale(x):
        if x < 0:
            v = abs(x) / size
            hue = base_hue - .4 * v
            sat = 0.6 + 0.4 * v
        else:
            hue = base_hue
            sat = 0.6
        return colorsys.hsv_to_rgb(hue, sat, 0.7)
    return color_scale

class MarkusLyapunov:
    def __init__(self, args):
        self.window_size = args.window_size
        self.length = self.window_size[0] * self.window_size[1]
        self.seed = args.seed
        self.scene = args.scene
        self.x0 = args.x0
        self.max_iter = 100
        self.max_init = 50
        self.set_view(args.center, args.radius)
        self.color_vector = color_factory(22.)
        self.nparray = None

    def get_filename(self, frame, ext=''):
        return "%s-q%04d-%s-%04d%s" % (self.scene, self.window_size[0], self.seed, frame, ext)

    def set_view(self, center = None, radius = None):
        if center is not None:
            self.center = center
        if radius is not None:
            if radius == 0:
                raise RuntimeError("Radius can't be null")
            self.radius = radius
        plane_min = (self.center.real - self.radius, self.center.imag - self.radius)
        plane_max = (self.center.real + self.radius, self.center.imag + self.radius)
        # Coordinate conversion vector
        self.offset = (plane_min[0], plane_min[1])
        self.scale = (
            self.window_size[0] / float(plane_max[0] - plane_min[0]),
            self.window_size[1] / float(plane_max[1] - plane_min[1])
        )

    def render(self, frame):
        if self.nparray is None:
            self.nparray = np.load(self.get_filename(frame, '.npy'))
#        colorized = map(self.color_vector, self.nparray) #self.color_vector(self.nparray)
#        colorized = np.reshape(colorized, (self.window_size[0], self.window_size[1], 3))
        img = np.zeros((self.window_size[1], self.window_size[0], 3), dtype=np.uint8)
        #colorized = self.color_vector(self.nparray)
        for y in xrange(self.window_size[1]):
            for x in xrange(self.window_size[0]):
                r,g,b = self.color_vector(self.nparray[x+y*self.window_size[0]])
                img[y][x][0] = int(r * 255)
                img[y][x][1] = int(g * 255)
                img[y][x][2] = int(b * 255)
        #colorized = np.reshape(colorized, (self.window_size[1], self.window_size[0]))
        #print colorized
        #png.from_array(colorized, 'RGB', {'size': self.window_size, 'bitdepth': 8}).save(self.get_filename(frame, '.png'))
        #return
        png.from_array(img, 'RGB').save(self.get_filename(frame, '.png'))
        print "[+] Written %s" % self.get_filename(frame, '.png')
        return

        out_file = open(self.get_filename(frame, ".png"), "wb")
        w = png.Writer(size=self.window_size, planes=1) #bitdepth=16)
        #w.write_array(out_file, colorized)
        w.write_packed(out_file, colorized)
        out_file.close()
        print "[+] Written %s" % out_file.name

    def compute(self, frame):
        self.nparray = compute(compute_markus_lyapunov, [self.window_size, self.offset, self.scale, self.seed, self.x0, self.max_iter, self.max_init, self.length])
        np.save(self.get_filename(frame), self.nparray)
        return
        #print "min:", min(nparray), "max:", max(nparray), "mean:", nparray.mean()
        #self.color_vector = np.vectorize(color_factory(float(abs(min(nparray)))))
        self.blit(self.color_vector(nparray))
        print "%04d: %.2f sec: MarkusLyapunov(seed/center/radius = '%s' '%s' %s )" % (frame, time.time() - start_time, self.seed, self.center, self.radius)
        return nparray

class Test:
    def __init__(self, model, length):
        self.model = model
        self.length = length
        self.seeds = set()

    def get_seed(self, frame):
        seed = ''
        for i in xrange(int(frame/32.+2)):
            if frame & (1<<i) != 0:
                seed += 'B'
            else:
                seed += 'A'
        if 'A' not in seed or 'B' not in seed or seed in self.seeds:
            return False
        self.seeds.add(seed)
        return seed

    def update(self, frame):
        while True:
            seed = self.get_seed(frame)
            if seed:
                break
            frame += 1
        self.model.seed = seed

def main():
    args = usage()

    # image properties
    args.window_size = map(lambda x: int(x * args.size), [100, 100])
    args.center = np.complex128(complex(args.center))

    do_compute, do_render = True, True
    if args.action and args.action == "compute":
        do_render = False
    if args.action and args.action == "render":
        do_compute = False

    model = MarkusLyapunov(args)

    if args.scene == "test":
        scene = Test(model, args.length)

    for frame in xrange(args.length):
        scene.update(frame)

        if do_compute:
            start_time = time.time()
            model.compute(frame)
        if do_render:
            if not do_compute:
                start_time = time.time()
            model.render(frame)
        print "%04d: %.2f sec: MarkusLyapunov(seed/x0/max_init/center/radius = '%s' '%s' '%s' '%s' %s )" % (frame, time.time() - start_time, model.seed, model.x0, model.max_init, model.center, model.radius)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass
    pool.terminate()
    pool.join()
    del pool