# Module #wskrand# defines a class #RNG# which inherits
# methods from the \emph{Python} 2.1 class #random.Random#.
# It is mildly more versatile, as it implements a #geometric#
# method to return Geometric random variates, and a (crude)
# #Poisson# method to return Poisson random variates.

# Here is the #RCS# identifier:
RCSID="$Id: wskrand.py,v 1.1 2009/04/10 15:48:07 wilfrid Exp wilfrid $"

# We import #log# from #math# as #_log#,
# as well as the base class #Random# from #random#, since the #geometric#
# method requires use of #log#.
from math import log as _log
from random import Random

# Class and other definitions
class RNG(Random):
    """
    \t$RCSfile: wskrand.py,v $
    \t$Revision: 1.1 $
    \t$Date: 2009/04/10 15:48:07 $
    The class RNG is a mildly more versatile elaboration on random.
    """
    # -------------------- Geometric distribution ----------------------\\
    # Note the #p# parameter should lie in the open interval \((0,1)\).
    # For reasons of speed we do not check for this. The method
    # #geometric# is now all in lower case, for compatibility with
    # other methods: this may cause trouble for old code based on
    # invocation of #wskrand# version #1.x#.
    def geometric(self,p): return int(self.expovariate(-_log(p)))
    
    # -------------------- Poisson distribution ------------------------\\
    # This is a crude implementation based on the definition of the
    # Poisson process in terms of exponentially distributed inter-incident
    # times. There are more efficient options for large mean values!
    def poisson(self,mean):
        # mean: should be positive!
        expovariate=self.expovariate
        count=-1
        while mean>0:
            mean=mean-expovariate(1)
            count=count+1
        return count
    
# Test section
if __name__ == "__main__":
    # Import useful functions from #math# module.
    from math import ceil as _ceil, floor as _floor
    from math import sqrt as _sqrt, exp as _exp
    
    # Factorial function borrowed from library example
    # in #doctest# module.
    def factorial(n):
        if not n >= 0:
            raise ValueError("n must be >= 0")
        if _floor(n) != n:
            raise ValueError("n must be exact integer")
        if n+1 == n: # e.g., 1e300
            raise OverflowError("n too large")
        result = 1
        factor = 2
        while factor <= n:
            try:
                result *= factor
            except OverflowError:
                result *= long(factor)
            factor += 1
        return result

    # Output of data for further analysis in \emph{Splus}.
    # Output of \emph{Splus} comment to identify \emph{Python} script.
    print '# Test output from $RCSfile: wskrand.py,v $'
    print '# '+RCSID
    print "\n".join(map(lambda l:'# '+l,RNG.__doc__.split("\n")))

    # Construct random number generator.
    rng=RNG()

    # Iterative loops to print \emph{Splus} script to perform
    # a number of rudimentary \(\chi^2\) tests.
    print '# Chi-square tests on output of RNG.poisson method'
    # for mu in [1.0,5.0,10.0,20.0]:
    for mu in [1.0,5.0,10.0,20.0,20.0,20.0]:
        # Compute sensible values for range of observations (#br#).
        r0=max(0,_floor(mu-2*_sqrt(mu)))
        r1=_ceil(mu+2*_sqrt(mu))
        br=range(r0+1,r1-1)
        br[0:0]=['-Inf']
        br.append('Inf')
        print 'br<-c'+str(tuple(br)).replace("'","")
        
        # Compute appropriate number of observations
        # needed to ensure that at least \(5\) observations
        # are expected to fall in each cell.
        e0=_exp(-mu)*mu**r0/factorial(r0)
        e1=_exp(-mu)*mu**r1/factorial(r1)
        n=2*_ceil(5/min(e0,e1))

        # Evaluate random number sample.
        print '# %d observations at mean %f' % (n,mu)
        x=map(lambda x,p=rng.poisson,m=mu:p(m),range(n))
        print 'x<-c'+str(tuple(x))
        
        # Compute \(\chi^2\) test.
        print 'o<-chisq.gof(x,dist="poisson",lambda=%f,cut.points=br)' % mu
        print "o \n rbind(o$expected,o$counts)"
        
    print '# Chi-square tests on output of RNG.geometric method'
    for sp in [0.9, 0.8, 0.5, 0.3, 0.1]:
        # Compute sensible range of observations (#br#).
        k=_floor(_log(0.01)/_log(sp))
        br=range(k)
        br[0:0]=['-Inf']
        br.append('Inf')
        print 'br<-c'+str(tuple(br)).replace("'","")
        
        # Compute appropriate number of observations
        # needed to ensure that at least \(5\) observations
        # are expected to fall in each cell.
        n=2*_ceil(5/((1-sp)*sp**k))

        # Evaluate random number sample.
        print '# %d observations at success probability %f' % (n,sp)
        x=map(lambda x,g=rng.geometric,p=sp:g(p),range(n))
        print 'x<-c'+str(tuple(x))
        
        # Compute and display \(\chi^2\) test.
        print 'o<-chisq.gof(x,dist="geometric",prob=%f,cut.points=br)' % (1-sp)
        print "o \n rbind(o$expected,o$counts)"
        
