#!/usr/bin/env python
"""
A python module with logging, doctest to implement simulation of
reversible birth-death-immigration processes.

Metadata:

  - Author:    W.S.Kendall
  - Copyright: opensource (c) W.S.Kendall 2008
  - Date:      2008-07-03
  - URL:       http://www.wilfridkendall.co.uk
  - Generator: 'E:/HOME/WSK/lib/python/generate.py' version 1.18.


Command-line options::

  usage: %prog [options] arguments
    --doctest:             perform doctests;
    --logging=LOGGING:     logging verbosity level.

Notes:
 - values for 'logging' option:
 'DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'
 (decreasing level of verbosity), with 'WARNING' as default.

"""
__RCSID__ = '$Revision: 1.2 $ $Date: 2008/07/03 20:42:02 $'


# \section{Definitions}
# We need the logging module regardless of whether the #logging# option
# is deployed, since #logging.info# calls are scattered through the text.
import logging, math, reversible


#\section{Exceptions}
class Error(Exception):
    """
    Generic exception to be raised by this module.
    """


# \section{Classes}
class BDI(reversible.Reversible):
    """
    This class implements a reversible birth-death-immigration process
    with rates as follows: immigration 'alpha', birth 'lambd', death
    'mu'::

     >>> chain = BDI(seed=2) ; print chain
     []
     >>> chain.equilibrium_draw()
     (0, 2)
     >>> chain.evolution(1) ; print ['%4.2f %d' % x for x in chain]
     ['0.00 2', '0.35 3', '0.38 4', '0.47 3', '0.52 2', '0.79 1', '1.00 1']
     >>> print ['%4.2f %d %d' % x for x in chain.reversed_statistics()]
     ['0.21 1 1', '0.26 1 2', '0.05 2 3', '0.09 3 4', '0.03 4 3', '0.35 3 2']
    """
    def __init__(self, alpha=1, lambd=1, mu=2, rng=None, seed=None):
        self.alpha, self.lambd, self.mu = alpha, lambd, mu
        super(BDI, self).__init__(rng=rng, seed=seed)

    def equilibrium_draw(self):
        """
        Draws from equilibrium of chain, result used to initialize
        chain. Note that state is tuple of time and place::

         >>> chain = BDI(seed=1)
         >>> N = 10000
         >>> result = {}
         >>> for t, n in [chain.equilibrium_draw() for i in range(N)]:
         ...   if result.has_key(n):
         ...      result[n] += 1
         ...   else:
         ...      result[n] = 1.0
         >>> res = sorted([(i, '%4.2f' % (f / N)) for i, f in result.items()])
         >>> [x[1] for x in res][:6]
         ['0.50', '0.26', '0.12', '0.06', '0.03', '0.02']
        """
        # First establish \(\rho=\lambda/\mu\) and \(\alpha\)
        # as floats.
        rho, alf = float(self.lambd) / self.mu, float(self.alpha)
        # We now compute initial probability as
        # \(\pi_0 = \left(1-\rho\right)^{\alpha/\lambda}\).
        prob = math.pow(1 - rho, alf/self.lambd)
        # Here is a convenient pre-computation based on
        # \[
        # \pi_n \;=\;
        # \frac{\lambda(n-1)+\alpha}{\mu n} \pi_{n-1}
        # \;=\;
        # \left(\rho + \frac{\alpha-\lambda}{\mu n}\right)\pi_{n-1}
        # \;=\;
        # \left(\rho + \frac{\alpha/\mu-\rho}{n}\right)\pi_{n-1}
        # \,.
        # \]
        alf_by_mu_minus_rho = (alf / self.mu) - rho
        # Finally the iteration \ldots.
        n, cond = 0, 1
        while self.rng.random() > prob / cond:
            n += 1
            cond -= prob
            prob *= (rho + alf_by_mu_minus_rho / n)
        self[:]=[(0, n)]
        return (0,n)

    def update(self):
        """
        Updates chain by single step, which is proposed to be '+1'
        with probability 'alpha + n lambd', '-1' with probability 'n'
        'mu', with time-increment being exponential of rate 'alpha + n'
        '(lambd + mu)'. The accumulated path is stored in 'self'.
        """
        t, n = self[-1]
        rate = self.alpha + n * (self.lambd + self.mu)

        t += self.rng.expovariate(rate)
        if rate * self.rng.random() < n * self.mu:
            n -= 1
        else:
            n += 1

        self.append((t, n))
        
    def evolution(self, T):
        """
        Repeat updates till time exceeds 'T', and then truncate at 'T'.
        """
        logging.info('Start evolution from trajectory\n%s' % self)
        while self[-1][0] < T:
            self.update()
            logging.info(' update to (%5.3f, %d)' % self[-1])
        self[-1] = (T, self[-2][1])

    def reversed_statistics(self):
        """
        Produce list of reversed transitions. Note that there must be
        at least one entry in 'self'::

         >>> chain = BDI(seed=1)
         >>> chain.reversed_statistics()
         Traceback (most recent call last):
          ...
         StopIteration
         >>> chain.equilibrium_draw()
         (0, 0)
         >>> print chain.reversed_statistics()
         []
        """
        state_list = super(BDI, self).reversed_statistics()
        return [(a[0]-b[0], a[1], b[1]) for a, b in state_list]
        

# \section{Functions}
def main(opt, arguments):
    """
    Make a single run.
    """
    logging.info('Entering main function')
    T = 8
    chain = BDI(seed=3)
    chain.equilibrium_draw()
    chain.evolution(T)
    print ['(%4.2f %d %d)' % x for x in chain.reversed_statistics()]


def _doctest():
    """
    Doctests.
    """
    logging.info('Commencing doctests')
    import doctest
    doctest.testmod()
    print __RCSID__.replace('$', '').strip()


# \section{Main script}
if __name__ == '__main__':
    import optionparse
    option, args = optionparse.parse(__doc__, version=__RCSID__)

    if option.logging is not None:
        logging.basicConfig(level=logging.__dict__[option.logging])
        logging.info('Logging level %s' % option.logging)

    if option.doctest:
        _doctest()


    # The action of the module as script should be concentrated
    # in the following function, to facilitate profiling.
    main(option, args)


    logging.info('Finishing')
    raise SystemExit


# \newpage \scriptsize
# \begin{multicols}{2}
# \section{History}
# \text{ }
#\\ $Log: bdi.py,v $
#\\ Revision 1.2  2008/07/03 20:42:02  wskendall
#\\ Correct equilibrium.
#\\
#\\ Revision 1.1  2008/07/03 19:12:15  wskendall
#\\ Initial revision
#\\
#\\ Revision 1.5  2008/07/03 12:35:28  wskendall
#\\ Display confidence interval in function #main#.
#\\
#\\ Revision 1.2  2008/07/03 12:22:32  wskendall
#\\ Added doctest, also statistical test of a longer run.
#\\
#\\ Revision 1.1  2008/07/03 11:36:12  wskendall
#\\ Initial revision
#\\
#
# \end{multicols}

'$Id: bdi.py,v 1.2 2008/07/03 20:42:02 wskendall Development wskendall $'
