recurrent neural network written in rpython

#!/usr/bin/python
## RPython Neural Network - 0.1b
## by Brett Hartshorn 2010
## goatman.py@gmail.com
## tested with Ubuntu Lucid, copy this file to your pypy root directory and run "python rAI.py"
## you can test just pypy compilation with "python rAI.py --pypy --subprocess"
## you need to install SDL headers, apt-get install libsdl-dev

'''
PyPY Tips

math.radians    not available
random.uniform  not available

incorrect:
    string.strip()  will not work without an arg
    list.sort()     not available

correct:
    string.strip(' \n')
    list = sortd( list )
    list.pop( index )       # no list.pop() with end as default

'''

COLUMNS = 24
LAYERS = 5
STEM = 8

import os, sys, time
#from random import *       # not valid in rpython?
import math             # math.radians is missing in pypy?

degToRad = math.pi / 180.0
def radians(x):
    """radians(x) -> converts angle x from degrees to radians
    """
    return x * degToRad

from pypy.rlib import streamio
from pypy.rlib import rpoll

# apt-get install libsdl-dev
from pypy.rlib.rsdl import RSDL, RSDL_helper
from pypy.rlib.rarithmetic import r_uint
from pypy.rpython.lltypesystem import lltype, rffi
from pypy.rlib.listsort import TimSort
from pypy.rlib.jit import hint, we_are_jitted, JitDriver, purefunction_promote

#random.randrange(0, up)        # not a replacement for random.uniform
#Choose a random item from range(start, stop[, step]).
#This fixes the problem with randint() which includes the
#endpoint; in Python this is usually not what you want.


if '--pypy' in sys.argv:        # random.random works in pypy, but random.uniform is missing?
    from pypy.rlib import rrandom
    RAND = rrandom.Random()
    RAND.init_by_array([1, 2, 3, 4])
    def random(): return RAND.random()
    def uniform(start,end): return ( RAND.random() * (end-start) ) - start
else:
    from random import *

def distance( v1, v2 ):
    dx = v1[0] - v2[0]
    dy = v1[1] - v2[1]
    dz = v1[2] - v2[2]
    t = dx*dx + dy*dy + dz*dz
    return math.sqrt( float(t) )


Njitdriver = JitDriver(
    greens = 'cur train times branes last_spike spikers temporal thresh triggers self abs_refactory abs_refactory_value'.split(),
    reds = 'i t a b c bias neuron'.split()
)
class RecurrentSpikingModel(object):
    '''
    Model runs in realtime and lossy - stores recent
    spikes in a list for realtime learning.
    Uses spike train, with time delay based on distance,        Spikes will trigger absolute refactory period.
    Membrane has simple linear falloff.

        notes:
            log(1.0) = 0.0
            log(1.5) = 0.40546510810816438
            log(2.0) = 0.69314718055994529
            log(2.7) = 0.99325177301028345
            log(3.0) = 1.0986122886681098

    '''
    def iterate( self ):
        now = float( time.time() )
        self._state_dirty = True; train = self._train
        brane = self._brane; rest = self._rest
        branes = self._branes; temporal = self._temporal
        abs_refactory = self._abs_refactory
        abs_refactory_value = self._abs_refactory_value
        fps = self._fps
        elapsed = now - self._lasttime
        clip = self._clip
        cur = self._lasttime
        last_spike = self._last_spike
        spikers = self._spikers
        thresh = self._thresh
        triggers = self._triggers

        ## times of incomming spikes ##
        times = train.keys()
        ## To do, if seziure, lower all connection strengths
        ##TimSort(times).sort()     # pypy note - list have no .sort(), and JIT dislikes builtin sorted(list)

        if not times:
            a = now - self._lasttime
            if a > 0.1: brane = rest
            else:
                b = a * 10  # 0.0-1.0
                brane += (rest - brane) * b
            branes.append( (brane,now) )

        i = 0
        t = a = b = c = bias = .0
        neuron = self
        ntrain = len(train)
        while train:        #the JIT does only properly support a while loop as the main dispatch loop
            Njitdriver.can_enter_jit(
                cur=cur, train=train, times=times, branes=branes, last_spike=last_spike, spikers=spikers,
                temporal=temporal, thresh=thresh, triggers=triggers, self=self,
                abs_refactory=abs_refactory, abs_refactory_value=abs_refactory_value,
                i=i, t=t, a=a, b=b, c=c,
                neuron=neuron, bias=bias
            )
            Njitdriver.jit_merge_point(
                cur=cur, train=train, times=times, branes=branes, last_spike=last_spike, spikers=spikers,
                temporal=temporal, thresh=thresh, triggers=triggers, self=self,
                abs_refactory=abs_refactory, abs_refactory_value=abs_refactory_value,
                i=i, t=t, a=a, b=b, c=c,
                neuron=neuron, bias=bias
            )
            ##bias, neuron = train.pop( t ) ## not pypy
            t = times[i]; i += 1
            bias,neuron = train[t]
            del train[t]


            if t <= cur:
                #print 'time error'
                pass
            else:
                a = t - cur
                cur += a
                #print 'delay', a
                if cur - last_spike < abs_refactory:
                    brane = abs_refactory_value
                    branes.append( (brane,cur) )
                else:
                    spikers.append( neuron )
                    if a > 0.1:
                        brane = rest + bias
                        branes.append( (brane,cur) )
                    else:
                        b = a * 10  # 0.0-1.0
                        brane += (rest - brane) * b
                        c = b * temporal
                        bias *= math.log( (temporal-c)+2.8 )
                        brane += bias
                        branes.append( (brane,cur) )

                        if brane > thresh:
                            triggers.append( neuron )
                            self.spike( cur )
                            brane = abs_refactory_value
                            #fps *= 0.25        # was this to play catch up?
                            break # this is efficient!

        #self._train = {}
        self._brane = brane
        #for lst in [branes, spikers, triggers]:    # not allowed in pypy
        while len(branes) > clip: branes.pop(0)
        while len(spikers) > clip: spikers.pop(0)
        while len(triggers) > clip: triggers.pop(0)

        self._lasttime = end = float(time.time())
        #print ntrain, end-now


    def detach( self ):
        self._inputs = {}       # time:neuron
        self._outputs = []  # children (neurons)
        self._train = {}
        self._branes = []
        self._spikers = []  # last spikers
        self._triggers = [] # last spikers to cause a spike

    def __init__( self, name='neuron', x=.0,y=.0,z=.0, column=0, layer=0, fps=12, thresh=100.0, dendrite_bias=35.0, dendrite_noise=1.0, temporal=1.0, distance_factor=0.1, red=.0, green=.0, blue=.0 ):
        self._name = name
        self._fps = fps     # active fps 10?
        self._thresh = thresh
        self._column = column
        self._layer = layer
        self._brane = self._rest = -65.0
        self._abs_refactory = 0.05
        self._abs_refactory_value = -200.0
        self._spike_value = 200.0
        self._temporal = temporal   # ranges 1.0 - inf
        self._distance_factor = distance_factor
        self._dendrite_bias = dendrite_bias
        self._dendrite_noise = dendrite_noise
        self._clip = 128
        self._learning = False
        self._learning_rate = 0.1

        ## the spike train of incomming spikes
        self._train = {}

        ## list of tuples, (brane value, abs time)
        ## use this list to render wave timeline ##
        self._branes = []

        self._spikers = []  # last spikers
        self._triggers = [] # last spikers to cause a spike
        self._lasttime = time.time()
        self._last_spike = 0
        self._inputs = {}       # time:neuron
        self._outputs = []  # children (neurons)

        self._color = [ red, green, blue ]
        #x = uniform( 0.2, 0.8 )
        #y = random()   #uniform( 0.2, 0.8 )
        #z = random()   #uniform( 0.2, 0.8 )
        self._pos = [ x,y,z ]
        self._spike_callback = None
        self._state_dirty = False
        self._active = False
        self._cached_distances = {}
        self._clipped_spikes = 0    # for debugging

        self._draw_spike = False
        #self._braneRect = self._spikeRect = None

    def randomize( self ):
        for n in self._inputs:
            bias = self._dendrite_bias + uniform( -self._dendrite_noise, self._dendrite_noise )
            self._inputs[ n ] = bias


    def mutate( self, v ):
        for n in self._spikers:
            self._inputs[ n ] += uniform(-v*2,v*2)
        for n in self._triggers:
            self._inputs[ n ] += uniform(-v,v)

        if not self._spikes or not self._triggers:
            for n in self._inputs:
                self._inputs[ n ] += uniform(-v,v)


    def reward( self, v ):
        if not self._spikers: print 'no spikers to reward'
        for n in self._spikers:
            bias = self._inputs[ n ]
            if abs(bias) < 100:
                if bias > 15:
                    self._inputs[ n ] += v
                else:
                    self._inputs[ n ] -= v


    def punish( self, v ):
        for n in self._inputs:
            self._inputs[n] *= 0.9

        for n in self._spikers:
            self._inputs[ n ] += uniform( -v, v )



    def attach_dendrite( self, neuron ):
        bias = self._dendrite_bias + uniform( -self._dendrite_noise, self._dendrite_noise )
        # add neuron to input list
        if neuron not in self._inputs:
            self._inputs[ neuron ] = bias
            #print 'attached neuron', neuron, bias
        if self not in neuron._outputs:
            neuron._outputs.append( self )
            neuron.update_distances()
        return bias

    def update_distances( self ):
        for child in self._outputs:
            dist = distance( self._pos, child._pos )
            self._cached_distances[ child ] = dist

    def spike( self, t ):
        #print 'spike', t
        self._draw_spike = True
        self._last_spike = t
        self._branes.append( (self._spike_value,t) )
        self._brane = self._abs_refactory_value
        self._branes.append( (self._brane,t) )
        if self._learning and self._triggers:
            #print 'learning'
            n = self._triggers[-1]
            bias = self._inputs[ n ]
            if bias > 0: self._inputs[n] += self._learning_rate
            else: self._inputs[n] -= self._learning_rate

        if self._spike_callback:
            self._spike_callback( t )

        for child in self._outputs:
            #print 'spike to', child
            bias = child._inputs[ self ]
            #dist = distance( self._pos, child._pos )
            dist = self._cached_distances[ child ]
            dist *= self._distance_factor
            child._train[ t+dist ] = (bias,self)

    def stop( self ):
        self._active = False
        print 'clipped spikes', self._clipped_spikes
    def start( self ):
        self._active = True
        self.iterate()

    def setup_draw( self, format, braneRect, groupRect, spikeRect, colors ):
        self._braneRect = braneRect
        self._groupRect = groupRect
        self._spikeRect = spikeRect
        self._sdl_colors = colors
        r,g,b = self._color
        self._group_color = RSDL.MapRGB(format, int(r*255), int(g*255), int(b*255))

    def draw( self, surf ):
        #if self._braneRect:
        fmt = surf.c_format
        b = int(self._brane * 6)
        if b > 255: b = 255
        elif b < 0: b = 0
        color = RSDL.MapRGB(fmt, 0, 0, b)
        RSDL.FillRect(surf, self._braneRect, color)
        RSDL.FillRect(surf, self._groupRect, self._group_color)
        if self._draw_spike:
            RSDL.FillRect(surf, self._spikeRect, self._sdl_colors['white'])
        else:
            RSDL.FillRect(surf, self._spikeRect, color )#self._sdl_colors['black'])
        self._draw_spike = False




Bjitdriver = JitDriver(
    reds=['loops'],
    greens='layers neurons pulse_layers'.split()
)

class Brain( object ):
    def loop( self ):
        self._active = True
        fmt = self.screen.c_format
        RSDL.FillRect(self.screen, lltype.nullptr(RSDL.Rect), self.ColorGrey)
        layers = self._layers
        neurons = self._neurons
        pulse_layers = self._pulse_layers
        screen = self.screen
        loops = 0
        now = start = float( time.time() )
        while self._active:
            #Bjitdriver.can_enter_jit( layers=layers, neurons=neurons, pulse_layers=pulse_layers, loops=loops )
            #Bjitdriver.jit_merge_point( layers=layers, neurons=neurons, pulse_layers=pulse_layers, loops=loops)
            now = float( time.time() )
            self._fps = loops / float(now-start)
            for i,lay in enumerate(self._layers):
                if self._pulse_layers[i] and False:
                    #print 'pulse layer: %s neurons: %s ' %(i, len(lay))
                    for n in lay:
                        if random()*random() > 0.8:
                            n.spike( now )
            for i,col in enumerate(self._columns):
                if self._pulse_columns[i]:
                    for n in col: n.spike(now)
            for n in self._neurons:
                n.iterate()
                n.draw(self.screen)
            #r,w,x = rpoll.select( [self._stdin], [], [], 1 )   # wait
            rl,wl,xl = rpoll.select( [0], [], [], 0.000001 )    # wait
            if rl:
                cmd = self._stdin.readline().strip('\n').strip(' ')
                self.do_command( cmd )
            loops += 1
            self._iterations = loops
            #print loops        # can not always print in mainloop, then select can never read from stdin
            RSDL.Flip(self.screen)
            #self._fps = float(time.time()) - now
        #return loops
        return 0

    def __init__(self):
        start = float(time.time())
        self._neurons = []
        self._columns = []
        self._layers = [ [] ] * LAYERS
        self._pulse_layers = [0] * LAYERS
        self._pulse_layers[ 0 ] = 1

        self._pulse_columns = [0] * COLUMNS
        self._pulse_columns[ 0 ] = 1
        self._pulse_columns[ 1 ] = 1
        self._pulse_columns[ 2 ] = 1
        self._pulse_columns[ 3 ] = 1


        inc = 360.0 / COLUMNS
        scale = float( LAYERS )
        expansion = 1.333
        linc = scale / LAYERS
        for column in range(COLUMNS):
            colNeurons = []
            self._columns.append( colNeurons )
            X = math.sin( radians(column*inc) )
            Y = math.cos( radians(column*inc) )
            expanding = STEM
            width = 1.0 / scale
            for layer in range(LAYERS):
                Z = layer * linc
                r = random() * random()
                g = 0.2
                b = 0.2
                for i in range(int(expanding)):
                    x = uniform( -width, width )
                    rr = random()*random()      # DJ's trick
                    y = uniform( -width*rr, width*rr ) + X
                    z = Z + Y
                    # create 50/50 exitatory/inhibitory
                    n = RecurrentSpikingModel(x=x, y=y, z=z, column=column, layer=layer, red=r, green=g, blue=b )
                    self._neurons.append( n )
                    colNeurons.append( n )
                    self._layers[ layer ].append( n )

                expanding *= expansion
                width *= expansion

        dendrites = 0
        interlayer = 0
        for lay in self._layers:
            for a in lay:
                for b in lay:
                    if a is not b and a._column == b._column:
                        a.attach_dendrite( b )
                        dendrites += 1
                        interlayer += 1

        intercol = 0
        for col in self._columns:
            for a in col:
                for b in col:
                    if a is not b and random()*random() > 0.75:
                        a.attach_dendrite( b )
                        intercol += 1
                        dendrites += 1

        intercore = 0
        core = self._layers[-1]
        for a in core:
            for b in core:
                if a is not b and random()*random() > 0.85:
                    a.attach_dendrite( b )
                    intercore += 1
                    dendrites += 1

        print 'brain creation time (seconds)', float(time.time())-start
        print 'neurons per column', len(self._columns[0])
        print 'inter-layer dendrites', interlayer
        print 'inter-column dendrites', intercol
        print 'inter-neocoretex dendrites', intercore
        print 'total dendrites', dendrites
        print 'total neurons', len(self._neurons)
        for i,lay in enumerate(self._layers):
            print 'layer: %s    neurons: %s' %(i,len(lay))
        for i,col in enumerate(self._columns):
            print 'column: %s   neurons: %s' %(i,len(col))



        self._stdin = streamio.fdopen_as_stream(0, 'r', 1)
        #self._stdout = streamio.fdopen_as_stream(1, 'w', 1)
        #self._stderr = streamio.fdopen_as_stream(2, 'w', 1)

        self._width = 640; self._height = 480
        assert RSDL.Init(RSDL.INIT_VIDEO) >= 0
        self.screen = RSDL.SetVideoMode(self._width, self._height, 32, 0)
        assert self.screen
        fmt = self.screen.c_format
        self.ColorWhite = white = RSDL.MapRGB(fmt, 255, 255, 255)
        self.ColorGrey = grey = RSDL.MapRGB(fmt, 128, 128, 128)
        self.ColorBlack = black = RSDL.MapRGB(fmt, 0, 0, 0)
        self.ColorBlue = blue = RSDL.MapRGB(fmt, 0, 0, 200)

        colors = {'white':white, 'grey':grey, 'black':black, 'blue':blue}

        x = 1; y = 1
        for i,n in enumerate(self._neurons):
            braneRect = RSDL_helper.mallocrect(x, y, 12, 12)
            groupRect = RSDL_helper.mallocrect(x, y, 12, 2)
            spikeRect = RSDL_helper.mallocrect(x+4, y+4, 4, 4)
            n.setup_draw( self.screen.c_format, braneRect, groupRect, spikeRect, colors )
            x += 13
            if x >= self._width-14:
                x = 1
                y += 13

    def do_command( self, cmd ):
        if cmd == 'spike-all':
            t = float(time.time())
            for n in self._neurons: n.spike(t)
        elif cmd == 'spike-one':
            t = float(time.time())
            self._neurons[0].spike(t)
        elif cmd == 'spike-column':
            t = float(time.time())
            for n in self._columns[0]:
                n.spike(t)
        elif cmd == 'info':
            info = self.info()
            #sys.stderr.write( info )
            #sys.stderr.flush()
            print info

    def info(self):
        r = ' "num-layers": %s,' %len(self._layers)
        r += ' "num-neurons": %s,' %len(self._neurons)
        r += ' "fps" : %s, ' %self._fps
        r += ' "iterations" : %s, ' %self._iterations
        return '<load_info> { %s }' %r






import subprocess, select, time
import gtk, glib

class App:
    def load_info( self, arg ): print arg

    def __init__(self):
        self._commands = cmds = []
        self.win = gtk.Window()
        self.win.connect('destroy', lambda w: gtk.main_quit())
        self.root = gtk.VBox(False,10); self.win.add( self.root )
        self.root.set_border_width(20)
        self.header = header = gtk.HBox()
        self.root.pack_start( header, expand=False )
        b = gtk.Button('spike all neurons')
        b.connect('clicked', lambda b,s: s._commands.append('spike-all'), self )
        self.header.pack_start( b, expand=False )

        b = gtk.Button('spike one neuron')
        b.connect('clicked', lambda b,s: s._commands.append('spike-one'), self )
        self.header.pack_start( b, expand=False )

        b = gtk.Button('spike column 1')
        b.connect('clicked', lambda b,s: s._commands.append('spike-column'), self )
        self.header.pack_start( b, expand=False )

        self.header.pack_start( gtk.SeparatorMenuItem() )

        b = gtk.Button('debug')
        b.connect('clicked', lambda b,s: s._commands.append('info'), self )
        self.header.pack_start( b, expand=False )

        da = gtk.DrawingArea()
        da.set_size_request( 640,480 )
        da.connect('realize', self.realize)
        self.root.pack_start( da )

        self._read = None
        glib.timeout_add( 33, self.loop )
        self.win.show_all()


    def realize(self, da ):
        print 'realize'
        xid = da.window.xid
        self._process = process = subprocess.Popen( 'python rAI.py --pypy --subprocess %s' %xid, stdin=subprocess.PIPE, stdout=subprocess.PIPE, bufsize=32, shell=True )
        self._write = write = process.stdin
        self._read = read = process.stdout
        print 'read', read
        print 'write', write

    def loop( self ):
        if self._read:
            rlist,wlist,xlist = select.select( [self._read], [], [], 0.001 )
            while self._commands:
                cmd = self._commands.pop()
                print 'sending cmd ->', cmd
                self._write.write( '%s\n'%cmd )
                self._write.flush()
            if rlist:
                a = self._read.readline().strip()
                if a:
                    print a
                    if a.startswith('<'):
                        func = a[ 1 : a.index('>') ]
                        arg = a[ a.index('>')+1 : ].strip()
                        func = getattr(self, func)
                        func( eval(arg) )
        return True


if '--subprocess' in sys.argv:
    os.putenv('SDL_WINDOWID', sys.argv[-1])
    def pypy_entry_point():
        def jitpolicy(*args):
            from pypy.jit.metainterp.policy import JitPolicy
            return JitPolicy()

        brain = Brain()
        brain.loop()
    if '--pypy' in sys.argv:
        from pypy.translator.interactive import Translation
        t = Translation( pypy_entry_point )
        ## NotImplementedError: --gcrootfinder=asmgcc requires standalone ##
        #t.config.translation.suggest(jit=True, jit_debug='steps', jit_backend='x86', gc='boehm')
        t.annotate()
        t.rtype()
        f = t.compile_c()
        f()
    else:
        pypy_entry_point()

else:
    a = App()
    gtk.main()
    print '-------------------exit toplevel-----------------'