root/projects/AsynCluster/trunk/svpmc/params.py

# svpmc: Stochastic Volatility Inference via Population Monte Carlo
#
# Copyright (C) 2007-2008 by Edwin A. Suominen, http://www.eepatents.com
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
# 
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the file COPYING for more details.
# 
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc., 51
# Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA

"""
Model parameters
"""

import os.path
import scipy as s
from scipy import linalg
from scipy.stats import distributions

from twisted_goodies.pybywire.params import Parameterized


SHAPE_MISMATCH = ValueError(
    "shape mismatch: objects cannot be broadcast to a single shape")


class FlexArray(object):
    """
    I am an array of arbitrary objects, accessible in much the same manner as the
    elements of a SciPy array object. Initially, all my objects are C{None}.

    Call an instance of me with a hash array and you'll get a new instance with
    just the subset of objects referenced by the array. Call the instance with
    a single object ID and you'll get that object, just like you'd get a
    numeric scalar by addressing a single element of a SciPy array. In either
    case, you can supply a raveled list of the new values (objects, numeric
    scalars, whatever) to be used instead of the instance's current objects.

    You can add two like-dimensioned instances of me. Each element of the sum
    will be the sum of the respective objects of the added L{FlexArray}
    instances.

    You can call the L{call} method on an instance of me with a method name and
    zero or more arguments. The result will be a numeric array of the same
    shape as the instance, each element of which will have the numeric result
    of a call of the named method on each of the corresponding objects in the
    instance. Any argument that is an array (or L{FlexArray} instance) of the
    same shape as my instance will have each the value of each element
    corresponding to the called object's element extracted and used as the
    method argument instead of the original array argument.

    You can get and set an attribute of an instance's objects in array fashion
    by simply setting or getting the named attribute of the instance.

    B{Caveat}: Unlike the behavior of SciPy arrays, Setting elements of a
    subset or raveled version of an instance of me does I{not} set the
    corresponding element of the original instance.
    """
    @classmethod
    def concatenate(cls, FA_list):
        """
        Returns an instance of me that is a concatentation of the instances in
        the supplied list I{FA_list}. Analogous to SciPy's C{concatenate}
        function.
        """
        length = 0
        shape = None
        for FA in FA_list:
            length += len(FA)
            if shape is None:
                shape = FA.shape
            elif FA.shape[1:] != shape[1:]:
                raise ValueError(
                    "FlexArray dimensions must agree except for d_0")
        shape = list(shape)
        shape[0] = length
        newVersion = cls(*shape)
        k = 0
        for FA in FA_list:
            for FA_sub in FA:
                newVersion[k] = FA_sub
                k += 1
        return newVersion

    def __init__(self, *shape):
        self._shape = tuple([int(x) for x in shape])
        self._N = 1
        for dim in self._shape:
            self._N *= dim
        # Object list
        self._O = [None] * self._N
        # Index array
        self._I = s.arange(self._N).reshape(*self._shape)

    def __call__(self, I, valueList=None):
        if valueList is None:
            valueList = self._O
        if isinstance(I, int):
            return valueList[I]
        newVersion = self.__class__(*I.shape)
        for j, k in enumerate(I.ravel()):
            newVersion._O[j] = valueList[k]
        return newVersion

    def __len__(self):
        return self.shape[0]
    
    def __iter__(self):
        self._iterator = iter(self._I)
        return self

    def next(self):
        return self(self._iterator.next())
    
    def __getitem__(self, key):
        return self(self._I[key])

    def __setitem__(self, key, value):
        I = self._I[key]
        if isinstance(I, int):
            self._O[I] = value
        elif I.shape == s.shape(value):
            if not isinstance(value, (list, tuple)):
                value = value.ravel()
            for j, k in enumerate(I.ravel()):
                self._O[k] = value[j]
        else:
            raise SHAPE_MISMATCH

    def ravel(self):
        return self(self._I.ravel())

    def reshape(self, *shape):
        return self(self._I.reshape(*shape))

    def __add__(self, other):
        if other.shape != self.shape:
            raise SHAPE_MISMATCH
        newVersion = self(self._I)
        for k, thisObj in enumerate(newVersion._O):
            thisObj += other._O[k]
        return newVersion

    def __getattr__(self, name):
        """
        For each of my objects, get the specified numeric attribute.

        Returns the result of each attribute value in an array of the same
        shape as my object array.

        If any value is not a numeric scalar, the result will be another
        instance of me with the results (of whatever type) as its objects.
        """
        if name == 'shape':
            return self._shape
        x = []
        allNumeric = True
        for k in self._I.ravel():
            thisValue = getattr(self._O[k], name)
            if allNumeric and not isinstance(thisValue, (int, long, float)):
                allNumeric = False
            x.append(thisValue)
        if allNumeric:
            return s.array(x).reshape(*self._shape)
        return self(self._I, x)

    def __setattr__(self, name, value):
        """
        For each of my objects, sets the specified numeric attribute to the
        supplied value.

        If the value is an array (or L{FlexArray}) with the same shape as me
        (not a list or tuple), each object's attribute value will be set to the
        value of the corresponding array element.
        """
        if name.startswith('_'):
            object.__setattr__(self, name, value)
        else:
            if hasattr(value, 'shape') and value.shape == self.shape:
                values = value.ravel()
            else:
                values = [value] * self._N
            for k, thisValue in enumerate(values):
                setattr(self._O[k], name, thisValue)
    
    def call(self, methodName, *args, **kw):
        """
        For each of my objects, call the method specified by I{methodName} with
        any further args and keywords supplied.

        For any argument that is an array (or L{FlexArray}) with the same shape
        as me (not a list or tuple), only the corresponding array element will
        be supplied as the argument to each object's method call. You can force
        an array argument to be supplied to the call in its entirety (not
        element-by-element) even if its shape matches mine by including its
        position index (with 0 as the first argument after the method name) in
        a sequence defined via the keyword I{arrayArgs}. For example, you can
        force the first argument (after the method name) to be an array
        argument to all method calls with C{arrayArgs=(1,)}.
        
        Returns the result of each call in an array of the same shape as my
        object array. The results must all be numeric scalars.

        If the result of any call is not a numeric scalar, the result will be
        another instance of me with the results (of whatever type) as its
        objects.
        """
        indicesOfArrayArgs = [
            k for k, arg in enumerate(args)
            if hasattr(arg, 'shape') and \
                arg.shape == self.shape and \
                k not in kw.get('arrayArgs', [])]
        # This is a shallow list copy
        theseArgs = [arg for arg in args]
        x = []
        allNumeric = True
        for j in self._I.ravel():
            for k in indicesOfArrayArgs:
                theseArgs[k] = args[k].ravel()[j]
            thisValue = getattr(self._O[j], methodName)(*theseArgs)
            if allNumeric and not isinstance(thisValue, (int, long, float)):
                allNumeric = False
            x.append(thisValue)
        if allNumeric:
            return s.array(x).reshape(*self._shape)
        return self(self._I, x)


class ParameterContainer(Parameterized):
    """
    I am a container for arrays of model parameters.

    In addition to the parameter arrays, I house as key attributes:

        - A scalar I{Lx} with the log-likelihood of the model's data given my
          parameters and some externally supplied normal variates underlying
          sets of simulated log-volatilites.

        - A scalar I{Lp} with the prior log-probability density of the
          parameters.

        - A scalar I{Lj} with the log-probability density of the proposal jump
          resulting in my parameters.
    
    """
    keyAttrs = {'Lx':None, 'Lp':None, 'Lj':None}


class Prior(Parameterized):
    """
    You must define the name of an underlying dist object from the scipy.stats
    L{distributions} module via my parameter I{dname}.

    For all underlying dist objects, you must also specify I{loc} and I{scale} as
    additional parameters. Some of the objects also require I{a} and I{b}.

    The exception is that you can use 'constant' as a distribution name. Then
    the only pertinent parameter is the constant, unvarying I{loc}.

    """
    paramNames = ['dname', 'loc', 'scale', 'a', 'b']
    N_draw = 1000

    class Constant(object):
        def __init__(self, loc):
            self.loc = loc
        def pdf(self, x):
            return 1.0 * (x == self.loc)
        def rvs(self, N):
            return self.loc * s.ones(N)
        def ppf(self, *args):
            return 0.0

    class Jump(object):
        __slots__ = ('x', 'px', 'pJump')
        def __init__(self, x, px, pJump):
            self.x = x
            self.px = px
            self.pJump = pJump
    
    def _get_distObj(self):
        if 'dist' not in self.cache:
            if self.dname == 'constant':
                self.cache['dist'] = self.Constant(self.loc)
            elif not hasattr(distributions, self.dname):
                raise ImportError(
                    "No distribution '%s' in scipy.stats.distributions" \
                    % self.dname)
            else:
                distClass = getattr(distributions, self.dname)
                args = []
                for xName in ('a', 'b'):
                    xValue = getattr(self, xName, None)
                    if xValue is not None:
                        args.append(xValue)
                kw = {'loc':self.loc, 'scale':self.scale}
                self.cache['dist'] = distClass(*args, **kw)
        return self.cache['dist']
    distObj = property(_get_distObj)

    def _get_rdsObj(self):
        if 'rds' not in self.cache:
            width = s.diff(self.distObj.ppf([0.15, 0.85]))[0]
            self.cache['rds'] = distributions.norm(0, width)
        return self.cache['rds']
    rdsObj = property(_get_rdsObj)

    def _newDraw(self):
        rJumps = self.rdsObj.rvs(self.N_draw)
        pJumps = self.rdsObj.pdf(rJumps)
        return {'k':-1, 'N':self.N_draw, 'R':rJumps, 'P':pJumps}
    
    def _rJumpDraw(self):
        if 'draw' not in self.cache:
            draw = self.cache['draw'] = self._newDraw()
        else:
            draw = self.cache['draw']
            draw['k'] += 1
            if draw['k'] >= draw['N']:
                draw = self.cache['draw'] = self._newDraw()
        k = draw['k']
        return draw['R'][k], draw['P'][k]
    
    def pdf(self, x):
        """
        Returns probability density of the supplied parameter value I{x}.
        """
        if x.shape:
            return self.distObj.pdf(x)
        return float(self.distObj.pdf(x))
    
    def rvs(self):
        """
        Returns a parameter value drawn from my distribution.
        """
        return self.distObj.rvs(1)[0]

    def rJump(self, startValue, v, N_attempts=200):
        """
        Call this method with a I{startValue} and the jump deviation I{v} to
        use for the jump.

        Draws a jump value from a normal distribution with a standard deviation
        that is set to approximate the 70% probability width of my prior
        distribution, scaled by the jump deviation. (The 70% width is the +/- 1
        standard deviation width of a normal distribution.)
        
        Returns a tiny object with three attributes, conveniently accessible as
        part of a L{FlexArray}:

            - B{x}: the value after the jump, and

            - B{px}: the probability density of I{x}, given the prior
              distribution.

            - B{pj}: the probability density of the jump, given the jump
              deviation used.
        
        For my Gaussian-PDF deviate generator, the probability density of a
        jump is inversely proportional to the scaling to its standard deviation
        imparted by the jump deviation. Thus, the unit-normal PDF is divided by
        each possible jump deviation in the I{p} array of the result.
        """
        vInverse = 1.0 / v
        for k in xrange(N_attempts):
            rJump, pJump = self._rJumpDraw()
            rJump *= v
            pJump *= vInverse
            x = startValue + rJump
            px = self.pdf(x)
            if px > 0:
                break
        else:
            print "Oops"
            # Could't come up with a jump from this point, fall back to an
            # independent draw from the prior distribution
            x = self.rvs()
            px = self.pdf(x)
            pJump = 1.0  # No jump involved
        return self.Jump(x, px, pJump)


class PriorContainer(object):
    """
    I am a container for array-like instances of L{FlexArray} that contain
    L{Prior} objects for the array elements of my model's parameters.

    @ivar n: The number of samples.
    
    @ivar p: The number of time series.

    @ivar primaryParamNames: A list of the primary parameter names in sequence.

    @ivar derivedParamNames: A list of the derived parameter names in sequence.
    
    """
    typeChecking = True
    N_attempts = 200
    
    def __init__(self, p, n):
        self.p = p
        self.n = n
    
    def priorArray(self, paramName, *shape):
        array = getattr(self, paramName, None)
        if array is None:
            array = FlexArray(*shape)
            setattr(self, paramName, array)
        return array

    def covarMatrix(self, cs, cr):
        """
        Returns a covariance matrix from the vector or sequence of I{cs}
        standard deviations and I{cr} cross-correlations.

        The I{cs} argument is in the form s0, s1,..., sp. It must have C{p}
        elements.

        The I{cr} argument must have C{0.5*(p**2 + p)} elements and is in the
        form r00, r01, r02,..., r0p, r11, r12,..., r1p,..., rpp
        """
        i = 0
        p = len(cs)
        cv = s.zeros((p, p))
        for j in xrange(p):
            cv[j,j] = cs[j]**2
            for k in xrange(j+1, p):
                cv[j, k] = cv[k, j] = cr[i] * cs[j] * cs[k]
                i += 1
        return cv
    
    def derivedParams(self, paramContainer):
        """
        Sets derived parameters in the supplied I{paramContainer} based on its
        primary parameters, unless there is an error in generating the derived
        parameters.

        Returns C{True} if everything went OK, C{False} if not.
        """
        # Concurrent cross-correlations between volatility shocks and inverse
        # of concurrent cross-correlations between innovation shocks
        try:
            rz = linalg.cholesky(self.covarMatrix(
                paramContainer.fs, paramContainer.fr), lower=True)
            ri = linalg.inv(linalg.cholesky(self.covarMatrix(
                s.ones(self.p), paramContainer.cr), lower=True))
        except linalg.LinAlgError:
            return False
        paramContainer.rz = rz
        paramContainer.ri = ri
        # Scaling constants for multivariate normal pdf
        sc = s.empty((self.p, 2))
        sc[:,0] = 1.0 / s.sqrt(1.0 - paramContainer.g**2)
        sc[:,1] = s.log(s.sqrt(2*s.pi) / sc[:,0])
        paramContainer.sc = sc
        return True

    def new(self):
        """
        Returns a new parameter container with values intialized from a random
        draw of my priors.

        The method will generate derived parameters based on the primary
        parameter values, repeating as needed to get primary values that
        work. The derived values are included in the returned parameter
        container.
        """
        paramContainer = ParameterContainer(
            paramNames=self.primaryParamNames+self.derivedParamNames)
        for k in xrange(self.N_attempts):
            Lp = 0
            for name in self.primaryParamNames:
                priorFlexArray = getattr(self, name)
                paramArray = priorFlexArray.call('rvs')
                setattr(paramContainer, name, paramArray)
                Lp += s.log(priorFlexArray.call('pdf', paramArray)).sum()
            if self.derivedParams(paramContainer):
                break
        else:
            raise ValueError(
                "Couldn't generate a valid set of derived parameters")
        paramContainer.Lp = Lp
        paramContainer.Lj = 0 # There was no jump
        return paramContainer

    def proposal(self, paramContainer, v):
        """
        Returns a new instance of L{ParameterContainer} with a random-walk
        proposal based on the supplied I{paramContainer}, my priors, and the
        supplied jump deviation I{v}.

        The returned parameter container object will have new values of I{Lp}
        and I{Lj} corresponding to the proposal.

        The value of I{Lj} is computed for each parameter element from a normal
        jump distribution whose width is proportional to the jump deviation
        I{v} and the width of the parameter's prior distribution.

        The method will generate derived parameters based on the primary
        parameter values of the proposal, repeating as needed to get primary
        values that work. The derived values are included in the returned
        parameter container.
        """
        if self.typeChecking and \
               not isinstance(paramContainer, ParameterContainer):
            raise TypeError(
                "You must supply a ParameterContainer as the first "+\
                "argument, not a '%s'" % type(paramContainer))
        newVersion = ParameterContainer(
            paramNames=self.primaryParamNames + self.derivedParamNames)
        for k in xrange(self.N_attempts):
            Lp, Lj = 0, 0
            # Define a 1-D array holding the probability density for each jump,
            # joint across all parameters. Since the densities will be scaled
            # by the weight for each jump deviation, just start the array with
            # the weights.
            for name in self.primaryParamNames:
                priorsFA = getattr(self, name)
                startValues = getattr(paramContainer, name)
                jumpsFA = priorsFA.call(
                    'rJump', startValues, v, N_attempts=self.N_attempts)
                setattr(newVersion, name, jumpsFA.x)
                Lp += s.log(jumpsFA.px).sum()
                Lj += s.log(jumpsFA.pJump).sum()
            if self.derivedParams(newVersion):
                break
        else:
            raise ValueError(
                "Couldn't generate a valid set of derived parameters")
        newVersion.Lp = Lp
        newVersion.Lj = Lj
        return newVersion