binary carry accumulator

accum.py

# binary carry accumulator with O(log(N)) space requirement
# better numerical stability when adding lots of small numbers
class Accumulator:
    __slots__ = ('fn', 'data', 'index', 'init')
    def __init__(self, size=31, init=0., fn='add'):
        if fn == 'mean':
            fn=lambda x,wx,y,wy:(x*wx+y*wy)/(wx+wy)
        elif fn == 'add':
            fn=lambda x,wx,y,wy:x+y
        self.fn = fn
        self.data = [0.] * (size+1)
        self.index = 0
        self.init = init

    def reset(self):
        for i in range(len(self.data)):
            self.data[i] = self.init
        self.index = 0

    def add(self, value):
        self.data[0] = value
        nbits = len(self.data)-1
        new_index = self.index+1
        bits = (new_index ^ self.index) >> 1
        for i in range(nbits):
            if (bits>>i)&1:
                value = self.fn(self.data[i], 2<<i, value, 1<<i)
                self.data[i] = self.init
            else:
                self.data[i] = value
                break
        else:
            bitmask = (1<<nbits)-1
            self.data[nbits] = self.fn(self.data[-1], (self.index | bitmask)^bitmask, value, 1<<nbits)
        self.index = new_index

    def value(self):
        value = self.data[-1]
        nbits = len(self.data)-1
        bitmask = (1<<nbits)-1
        index = self.index
        weight = (index | bitmask)^bitmask
        bit = 1<<nbits
        for v in reversed(self.data[:-1]):
            bit >>=1
            if self.index & bit:
                value = self.fn(value, weight, v, bit)
                weight += bit
        return value