b_smoothlife.py

import argparse, math
import numpy as np
import matplotlib.pyplot as plt 
import matplotlib.animation as animation
from matplotlib.widgets import Button

class BesselJ:
    def __init__(self, radius, logRes):
        N = 2**logRes
        self.logRes = logRes
        self.N = N
        re = np.zeros((N,N))
        self.weight = 0.0
        for i in xrange(N):
            for j in xrange(N):
                ii = ((i + 0.5*N) % N) - 0.5*N
                jj = ((j + 0.5*N) % N) - 0.5*N

                r = np.sqrt(ii*ii + jj*jj) - radius
                v = 0.
                if logRes*r < 100:
                    v = 1.0 / (1.0 + np.exp(logRes * r))

                self.weight += v
                re[i,j] = v

        self.f = np.fft.fft2(re)

class SmoothLife:
    def __init__(self):
        self.logRes = 8
        self.color = True
        self.euler = True
        self.dt = 0.05
        self.speckleIntensity = 1.
        self.lutN = 256
        self.parameters(1)

    def parameters(self, which):
        ns = 1
        if which==0: # Soliton. With Fourier integration, c=93.75
            self.r0 = 7.0
            self.r1 = 3 * self.r0
            self.b1 = 0.278
            self.b2 = 0.365
            self.d1 = 0.267
            self.d2 = 0.445
            self.alphaN = 0.028
            self.alphaM = 0.147
            self.nSpeckles = 800 * 2**(self.logRes-8)
            self.euler = False
        elif which==1: # buds and cords (with Euler integration)
            self.r0 = 1.*20./3.
            self.r1 = 1.*20.
            self.b1 = 0.257
            self.b2 = 0.336
            self.d1 = 0.365
            self.d2 = 0.549
            self.alphaN = 0.028
            self.alphaM = 0.147
            self.nSpeckles = 100 #500 * 2**(self.logRes-8)
            self.euler = True
            ns = 6
        self.construct()

        for i in xrange(ns):
            self.addSpeckles()

    def sigma1(self, x, a, alpha):
        return 1.0 / (1.0 + np.exp(-4.0/alpha * (x - a)))

    def sigma2(self, x, a, b, alpha):
        return     self.sigma1(x, a, alpha) * \
            (1.0 - self.sigma1(x, b, alpha))

    def sigmaM(self, x, y, m, alpha):
        alive = self.sigma1(m, 0.5, alpha)
        return x * (1.0 - alive) + y * alive

    def lerp(self, a, b, t):
        return (1.0-t)*a + t*b

    def s2(self, n, m):
        return self.sigmaM(self.sigma2(n, self.b1, self.b2, self.alphaN),
                           self.sigma2(n, self.d1, self.d2, self.alphaN), m, self.alphaM)

    def construct(self):
        N = 2**self.logRes
        self.N = N

        self.field = np.zeros((N,N))
        self.Mf = np.zeros((N,N))
        self.Nf = np.zeros((N,N))
        
        self.grid = np.zeros((N,N,3))

        self.innerBessel = BesselJ(self.r0, self.logRes)
        self.outerBessel = BesselJ(self.r1, self.logRes)

        innerW = 1.0 / self.innerBessel.weight
        outerW = 1.0 / (self.outerBessel.weight - self.innerBessel.weight)

        self.outerBessel.f = outerW * (self.outerBessel.f - self.innerBessel.f)
        self.innerBessel.f *= innerW
        
        self.lut = np.zeros((1+self.lutN, 1+self.lutN))
        f = 1./self.lutN
        for n in xrange(1+self.lutN):
            for m in xrange(1+self.lutN):
                self.lut[n,m] = self.s2(f*n, f*m)

    def calcMfNf(self):
        ft = np.fft.fft2(self.field)
        self.Mf = np.fft.ifft2(ft * self.innerBessel.f).real
        self.Nf = np.fft.ifft2(ft * self.outerBessel.f).real

    def step(self):
        self.calcMfNf()
        N = (self.lutN * self.Nf).astype(int)
        M = (self.lutN * self.Mf).astype(int)
        if self.euler:
            self.field = np.clip(self.field + self.dt * (2.0*self.lut[N,M]-1.0), 0.0, 1.0)
        else:
            self.field = self.lut[N,M]

    def addSpeckles(self):
        N = self.N
        r0 = int(self.r0)
        for i in xrange(self.nSpeckles):
            u = int(np.random.rand() * (N-r0))
            v = int(np.random.rand() * (N-r0))
            for x in xrange(r0):
                for y in xrange(r0):
                    self.field[u+x, v+y] = self.speckleIntensity
    
    def buttonSpeckles(self, event): self.addSpeckles()
    def buttonParameters0(self, event): self.parameters(0)
    def buttonParameters1(self, event): self.parameters(1)
    def buttonParameters2(self, event): self.parameters(2)
    def buttonParameters3(self, event): self.parameters(3)
    def buttonColor(self, event):
        self.color = not self.color

    def setGrid(self):
        if self.color:
            self.grid[:,:,0] = self.field
            self.grid[:,:,1] = self.Mf
            self.grid[:,:,2] = self.Nf
        else:
            self.grid[:,:,0] = self.field
            self.grid[:,:,1] = self.field
            self.grid[:,:,2] = self.field

def update(frameNum, img, life):
    life.step()
    life.setGrid()
    img.set_data(life.grid)
    return img,

def main():
    parser = argparse.ArgumentParser(description="Runs Rafler's SmoothLife.")
    parser.add_argument('--log-res', dest='logres', required=False)
    parser.add_argument('--mov-file', dest='movfile', required=False)
    parser.add_argument('--interval', dest='interval', required=False)
    args = parser.parse_args()

    updateInterval = 20
    if args.interval:
        updateInterval = int(args.interval)

    life = SmoothLife()

    if args.logres:
        life.logRes = int(args.logres)
        assert life.logRes > 0 and life.logRes < 13

    life.setGrid()

    fig, ax = plt.subplots()
    img = ax.imshow(life.grid, interpolation='nearest', cmap='Greys')
    ani = animation.FuncAnimation(fig, update, fargs=(img, life, ),
                                  frames = 10,
                                  interval = updateInterval,
                                  save_count = 50)

    def makeButton(n, name, callback):
        axButton = plt.axes([0.83, 0.15+0.075*n, 0.165, 0.06])
        button = Button(axButton, name)
        button.on_clicked(callback)
        return axButton, button

    buttons = []
    buttons += [makeButton( 3, 'Solitons (Fourier)', life.buttonParameters0)]
    buttons += [makeButton( 2, 'Buds (Euler)',       life.buttonParameters1)]
    buttons += [makeButton( 1, 'Toggle Color',       life.buttonColor)]
    buttons += [makeButton( 0, 'Add Noise',          life.buttonSpeckles)]

    if args.movfile:
        ani.save(args.movfile, fps=30, extra_args=['-vcodec', 'libx264'])
 
    plt.show()

if __name__ == '__main__':
    main()