root/projects/AsynCluster/trunk/svpmc/model.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

"""
The stochastic volatility model and its manager
"""

from copy import copy

import scipy as s
from scipy import linalg, stats

from twisted.internet import defer
from twisted.spread import pb

from asyncluster.master import jobs
from twisted_goodies.pybywire.params import Parameterized
from twisted_goodies.pybywire import pack

import weave


INT_SCALE = 2000


class ZFetcher(pb.Referenceable):
    """
    I provide a node worker with access to the normal variates underlying the
    latest log-volatility simulations.

    @ivar z: The full-precision 3-D array of normal variates underlying the
      current set of valid log-volatility simulations.

    @ivar zPackedChunks: A list of strings containing a chunked, packed version
      of z.
    
    """
    chunkSize = 60000
    
    def __init__(self, m, p, n):
        self.m = m
        self._z = s.empty((m, p, n))
        self.restart()

    def _get_z(self):
        return self._z[self.successes,:,:]
    z = property(_get_z)
    
    def restart(self):
        self.failures = []
        self.successes = []
        self.zPackedChunks = None
    
    def update(self, member, z=None):
        """
        Call to update me with the results of a log-volatility simulation, with
        integer arguments specifying the I{member} of the population for the
        latest iteration of the overall PMC simulation.

        If the log-volatility simulation succeeded, also supply the 2-D array
        of its normal variates. Otherwise, the population member will be marked
        as a failure.

        Returns C{True} to indicate when this is the last update, C{False}
        until then.
        """
        if not isinstance(member, int) or member < 0 or member >= self.m:
            raise ValueError(
                "Invalid member ID '%s', must be integer between 0 and %d" \
                % (member, self.m-1))
        if z is None:
            self.failures.append(member)
        else:
            self.successes.append(member)
            self._z[member,:,:] = z
        if len(self.failures) + len(self.successes) == self.m:
            zInt = (INT_SCALE * self.z).astype('h')
            zPacked = pack.Packer(zInt)()
            k = 0
            N = len(zPacked)
            self.zPackedChunks = []
            while k < N:
                chunk = zPacked[k:k+self.chunkSize]
                self.zPackedChunks.append(chunk)
                k += self.chunkSize
            return True
        return False
    

class ModelManager(object):
    """
    I manage a L{Model} object and a collection of L{priors.Prior} objects.

    I handle calls to the model object in either local or remote mode. Calls in
    local mode are run through my project manager's thread queue.

    Instantiate me with the text of a specification for the project and model
    I'm going to manage.

    """
    remoteMode = False
    timeout = 80

    nodecode = """
    ###########################################################################
    from twisted_goodies.pybywire import pack
    
    MODEL = None
    CHUNKS = []

    def newModel(modelObj):
        global MODEL
        MODEL = modelObj
        MODEL.z = None

    def newZ(zPacked, k, N):
        global CHUNKS
        if k == 0:
            CHUNKS = []
            MODEL.z = None
        CHUNKS.append(zPacked)
        if k == N-1 and len(CHUNKS) == N:
            zInt = pack.Unpacker(''.join(CHUNKS))()
            MODEL.z = %f * zInt.astype('d')
    
    def hybridGibbs(paramContainer):
        if MODEL is None:
            raise RuntimeError('No model available')
        return pack.Packer(MODEL.hybridGibbs(paramContainer))()
    
    def likelihood(paramContainer):
        if MODEL is None:
            raise RuntimeError('No model available')
        if MODEL.z is None:
            raise RuntimeError(
                'Model not initialized with simulated log-volatility data')
        return MODEL.likelihood(paramContainer)

    ###########################################################################
    """ % (1.0/INT_SCALE,)
    
    def __init__(self, projectManager, y, **kw):
        kw['y'] = y
        kw['paramNames'] = projectManager.paramNames
        self.zFetcher = ZFetcher(
            projectManager.m, projectManager.p, projectManager.n)
        self.modelObj = Model(**kw)
        self.queue = projectManager.queue

    def _oops(self, failure):
        print "FAILURE:\n%s" % failure.getTraceback()

    def shutdown(self):
        """
        Call this when I'm done being used. Returns a deferred that fires when
        everything's shut down.
        """
        if hasattr(self, 'client'):
            d = self.client.shutdown()
        else:
            d = defer.succeed(None)
        return d

    def setLocalMode(self):
        """
        Sets me to run the models locally in the current interpreter.
        """
        self.remoteMode = False
    
    def setRemoteMode(self, socket):
        """
        Sets me to run the named model remotely on the asyncluster nodes via
        the supplied UNIX domain I{socket}. Returns a deferred that fires when
        remotely-run operations can commence.
        """
        if hasattr(self, 'client'):
            return defer.succeed(None)
        self.client = jobs.JobClient(socket, codeString=self.nodecode)
        d = self.client.startup()
        d.addCallback(
            lambda _: self.client.registerClasses(*Parameterized.registry))
        d.addCallback(lambda _: self.client.update('newModel', self.modelObj))
        d.addCallback(lambda _: setattr(self, 'remoteMode', True))
        d.addErrback(self._oops)
        return d

    @defer.deferredGenerator
    def nextIteration(self, remote=False, local=False):
        """
        """
        self.modelObj.z = self.zFetcher.z
        if (remote or self.remoteMode) and not local:
            chunks = self.zFetcher.zPackedChunks
            N = len(chunks)
            for k, chunk in enumerate(chunks):
                wfd = defer.waitForDeferred(
                    self.client.update('newZ', chunk, k, N, ephemeral=True))
                yield wfd
                wfd.getResult()
        self.zFetcher.restart()
    
    def zSimulation(self, paramContainer, member, remote=False, local=False):
        """
        Runs a new log-volatily simulation for the I{paramContainer}, saving
        its result in my I{z} array and updating my instance of L{ZFetcher}
        with the array when it has been completely updated.

        Updates the container in place with the last simulated log-volatility
        value.
        
        Returns a deferred that fires when the simulation is done, with C{True}
        if it succeeded or C{False} if not.
        """
        def update(z=None):
            simulationsDone = self.zFetcher.update(member, z)

        def unpackResult(packedResult):
            if packedResult:
                return pack.Unpacker(packedResult)()
        
        def simulationDone(result):
            if result is None:
                update()
                return False
            update(result)
            return True
        
        if (remote or self.remoteMode) and not local:
            d = self.client.run(
                'hybridGibbs', paramContainer, timeout=self.timeout)
            d.addCallback(unpackResult)
        else:
            d = self.queue.call(
                self.modelObj.hybridGibbs, paramContainer)
        return d.addCallbacks(simulationDone, self._oops)
    
    def likelihood(self, paramContainer, remote=False, local=False):
        """
        Computes the log-likelihood of the parameters in the supplied
        I{paramContainer} for my model.

        Updates the container in-place with the log-likelihood.

        Returns a deferred that fires with a reference to the paramContainer
        when the likelihood computation is done and it has been updated.
        """
        def gotLikelihood(result):
            if result is None:
                # An error from the remote likelihood method call is treated as
                # infinitely low likelihood
                result = -s.inf
            paramContainer.Lx = result
            return paramContainer

        # DEBUG
        if (remote or self.remoteMode) and not local:
            d = self.client.run(
                'likelihood', paramContainer, timeout=self.timeout)
        else:
            d = self.queue.call(
                self.modelObj.likelihood, paramContainer)
        return d.addCallbacks(gotLikelihood, self._oops)


class VAR(weave.Weaver):
    """
    Construct me with a 2-D array of time series values, one time series per
    row.
    
    Then you can call me with a C{p x 1} drift vector I{a} and a C{p x p}
    VAR(1) cross-correlation matrix I{b}, to get a C{p x n} vector of residuals
    from my observation data after application of the VAR(1) model::
    
        y[:,t] = a + b*y[:,t-1]

    @ivar y: Observations from I{n} intersecting dates of I{p} time series,
      C{p,n} array.
    
    """
    def __init__(self, y):
        self.y = y
    
    def reverse(self, a, b):
        v = s.empty_like(self.y)
        self.inline('v', 'y', 'a', 'b')
        return v

    def forward(self, v, a, b):
        y = s.empty_like(v)
        self.inline('y', 'v', 'a', 'b')
        return y


class Model(Parameterized, weave.Weaver):
    """
    I am a sophisticated econometric model wherein the residuals of a
    linear-drift, VAR(1) process follow a multivariate stochastic volatility
    process with correlated volatility (v) and return (w) shocks.

    @ivar y: A 2-D array containing my observation data.

    @ivar wiggle: The +/- range of the random-walk uniform proposal on z for
      each hybrid Gibbs proposal. Use +/-1.0 for balance between fast
      convergence and reasonable accept-reject efficiency in z proposal
      generator.

    @ivar Ni: Number of hybrid Gibbs iterations per log-volatility simulation.
    
    """
    keyAttrs = {'y':None, 'wiggle':1.0, 'Ni':50}

    #--- Properties -----------------------------------------------------------

    def _get_p(self):
        return self.y.shape[0]
    p = property(_get_p)
    
    def _get_n(self):
        return self.y.shape[1]
    n = property(_get_n)

    def _get_var(self):
        if 'var' not in self.cache:
            self.cache['var'] = VAR(self.y)
        return self.cache['var']
    var = property(_get_var)


    #--- Support --------------------------------------------------------------

    def setParams(self, paramContainer):
        """
        Sets up array variables for an inline call made by the caller,
        including innovations as residuals of my observations run through the
        reversed VAR(1).
        """
        self.setParamSequence(paramContainer.paramSequence())
        self.x = self.var.reverse(self.a, self.b)

    
    #--- The API --------------------------------------------------------------

    def hybridGibbs(self, paramContainer):
        """
        Call this method with a parameter object that defines a set of values
        for my parameters.

        A new 2-D array of log-volatilities will be drawn from a hybrid Gibbs
        simulation, with each Gibbs step driven by a short importance-sampling
        simulation of the IID variates, followed by cross-correlation to form
        volatility shocks, and VAR(1) to form the log-volatilities. The
        importance-sampling proposals are scaled based on the supplied
        fractional I{sigma}. The simulation is governed by the log-likelihood
        of my observation data given the supplied parameters and the simulated
        log-volatilities.

        If a L{linalg.LinAlgError} is raised due to an invalid correlation
        matrix parameter, the method returns with no result. Otherwise, it
        returns the IID normal variates underlying the simulated
        log-volatilities after the last simulation round.

        """
        self.setParams(paramContainer)
        # Create the normal variate array with a standard starting point for
        # every log-volatility simulation
        z = s.zeros_like(self.x)
        self.inline(
            'z', 'x', 'w', 'ni',
            'd', 'e', 'g', 'rz', 'ri', 'sc',
            w=self.wiggle, ni=self.Ni)
        return z

    def likelihood(self, paramContainer):
        """
        Call this method with a parameter object that defines a set of values
        for my parameters.

        If a L{linalg.LinAlgError} is raised due to an invalid correlation
        matrix parameter, the method returns with no result. Otherwise, it
        returns the log-likelihood of my observations given the parameters,
        from an integration of the joint probability of the observations and
        log-volatilities simulated from all the other parameters in the current
        PMC population.
        
        The returned likelihood does not consider the prior probability of the
        parameters. You have to account for that yourself, preferably before
        calling this method. If the prior probability of any parameter is zero,
        the posterior density for that parameter will also be zero no matter
        what, making any call to this method with it a waste of computing time.
        """
        self.setParams(paramContainer)
        return self.inline('z', 'x', 'd', 'e', 'g', 'rz', 'ri', 'sc')