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

"""
Population Monte Carlo sampling
"""

import sys
from random import sample as sampleWOR

import scipy as s
from scipy import random
from twisted.internet import defer

import params
from sample import Resampler


class Allocator(object):
    """
    Construct me with a population size I{N} and an integer number of
    allocation bins I{D}.

    @ivar W: A 1-D array containing the fraction of population members, i.e.,
      the population weight, currently allocated to each bin.

    @ivar R: A 1-D array containing the number of population members currently
      allocated to each bin.

    """
    def __init__(self, N, D):
        if N < 2*D:
            raise ValueError(
                "Population size must be at least twice the number of bins")
        self.N = N
        self.D = D
        self.rMin = max([2, int(round(0.01 * N))])
        self.rMax = N - self.rMin*(self.D-1)
        self.updateAllocations()
        self.II = None
    
    def subsetIndex(self, k):
        """
        Returns a subset index for population members that correspond to the
        bin indexed by I{k}.
        """
        I = sampleWOR(self.Is, self.R[k])
        self.Is = s.setdiff1d(self.Is, I)
        return I
    
    def assembler(self, resultList):
        """
        Call this method with a reference to an empty list that will hold
        results assembled from an allocation that I'll be (re)doing. If you
        want to re-use the same allocation as from a previous call to this
        method, set the I{lastAllocation} keyword C{True}.

        The method will allocate the population members into my I{D} bins,
        unless we're re-using the last allocation. For each bin, the method
        will yield the bin index, a 1-D array of indices for the subset of my
        population allocated to that bin, and a fresh deferred already fired
        with C{None}. You must add a callback to the deferred for each
        iteration that returns a tuple of 1-D arrays, each being the same
        length as the yielded index array.
        
        Each returned array for each subset will be assembled back into a
        single population-size array and placed into the supplied list, in
        order.
        """
        def gotResults(results, I):
            if not isinstance(results, tuple):
                results = (results,)
            if not resultList:
                for result in results:
                    if isinstance(result, list) and len(result) == len(I):
                        resultList.append(s.empty(self.N))
                    elif isinstance(result, s.ndarray):
                        resultList.append(s.empty(self.N))
                    else:
                        resultList.append(params.FlexArray(self.N))
            for k, result in enumerate(results):
                resultList[k][I] = result
        
        if resultList != []:
            raise ValueError(
                "You must supply an empty list to hold your assembled results")
        if self.II is None:
            lastAllocation = False
            self.II = [None] * self.D
        else:
            lastAllocation = True
        for k in s.flipud(s.argsort(self.R)):
            if lastAllocation:
                I = self.II[k]
            else:
                I = self.II[k] = self.subsetIndex(k)
            d = defer.succeed(None)
            yield k, I, d
            d.addCallback(gotResults, I)

    def updateAllocations(self, I=None):
        """
        Rescales my allocations based on the supplied 1-D array I{I} of
        resampling indices. Allocations with lots of repeated indices were
        highly successful and those with lots of missing indices were not.
        
        Highly successful allocations will have a large portion of the next
        allocation even if they last operated on only a few members of the
        population.

        If no index array is supplied, the allocations will be equalized.
        """
        if I is None:
            R = s.ones(self.D) / self.D
        elif self.II is not None:
            R = s.array(
                [sum([x in self.II[k] for x in I]) for k in xrange(self.D)])
            R = R.astype(float) / R.sum()
            self.II = None
        else:
            raise AttributeError(
                "You can only update with an index array after completion "+\
                "of an assembler run")

        # Turn the scaled array into an array of subsample sizes, with a
        # minimum size of two apiece.
        R = s.clip(s.round_(self.N*R), self.rMin, self.rMax)
        # Twiddle the biggest one as needed to keep sum = Number of members
        R[s.argmax(R)] += self.N - sum(R)
        # Replace the old allocation arrays and subset index
        self.W = R / R.sum()
        self.R = R.astype(int)
        self.Is = s.arange(self.N)


class PMC(object):
    """
    Population Monte Carlo with per Cappe et al.

    """
    def __init__(self, projectManager, socket=None):
        self.socket = socket
        self.pm = projectManager
        self.mm = projectManager.mgr
        self.V = projectManager.V
        self.queue = projectManager.queue
        self.resampler = Resampler(logWeights=True)            
        self.allocator = Allocator(projectManager.m, len(projectManager.V))

    def _oops(self, failure, **kw):
        text = "\nFailure"
        if kw:
            text += " with "
            for name, value in kw.iteritems():
                text += "%s=%s" % (name, value)
        print "%s\n%s" % (text, failure.getTraceback())
    
    def proposals(self, *args):
        """
        Call this method with a 1-D FlexArray of parameter containers and a
        jump deviation, and it will generate proposals from the existing
        parameters.

        Call it with just an integer number of proposals, and it will generate
        new ones directly from the priors.

        Returns a 1-D FlexArray of parameter containers embodying the
        proposals.
        """
        if len(args) == 2:
            new = False
            X, v = args
            N = len(X)
        else:
            new = True
            N = args[0]
        XP = params.FlexArray(N)
        for k in xrange(N):
            if new:
                proposalPC = self.pm.priors.new()
            else:
                proposalPC = self.pm.priors.proposal(X[k], v)
            XP[k] = proposalPC
        return XP
    
    def simulations(self, X):
        """
        Runs a log-volatility simulation for each parameter container supplied
        in the 1-D FlexArray I{X}. The simulations are run via the model
        manager, which constructs a 3-D array of the normal variates underlying
        the entire set of simulations.

        Returns a deferred that fires when all the log-volatility simulations
        are done.
        """
        def simulationDone(success):
            if success:
                print "+",
            else:
                print ".",

        dList = []
        for k, Xk in enumerate(X):
            d = self.mm.zSimulation(Xk, k)
            d.addCallback(simulationDone)
            d.addErrback(self._oops, k=k)
            dList.append(d)
        d = defer.DeferredList(dList)
        d.addCallback(lambda _: self.mm.nextIteration())
        d.addErrback(self._oops)
        return d
    
    def weights(self, XP, v):
        """
        Returns a deferred that fires with a 1-D array of log importance
        weights for the parameter proposals in the 1-D FlexArray I{XP}.
        
        The importance weights are to be combined with those from other calls
        to this method with different variance settings. The weights will be
        used to resample the supplied proposals so that the probability of any
        given proposal being included in the final result for this iteration is
        proportional to its likelihood and prior probability density, and
        inversely proportional to the probability density of the proposal
        offset.
        """
        def weight(paramContainer):
            if s.isfinite(paramContainer.Lx):
                L = paramContainer.Lx + paramContainer.Lp - paramContainer.Lj
                info = " ".join([
                    "%+12.2f" % (getattr(paramContainer, x),)
                    for x in ('Lx', 'Lp', 'Lj')])
            else:
                L = -s.inf
            if L < 0:
                print "%6.4f" % v
            else:
                print "%6.4f\t%+12.2f\t%s" % (v, L, info)
            return L
        
        # Evaluate the proposals (very time-consuming, done either locally in a
        # thread or in the cluster)
        dList = [self.mm.likelihood(XPk).addCallback(weight) for XPk in XP]
        return defer.gatherResults(dList).addCallback(lambda W: s.array(W))
    
    @defer.deferredGenerator
    def run(self, N_iter):
        """
        Does a PMC run with I{N_iter} iterations and the sequence of jump
        deviations in my I{V} attribute. The portion of the population members
        allocated to each jump variance will go up or down, depending on the
        performance for that setting.

        Returns a deferred that fires when the run is done. No output value is
        provided via the deferred, however. Instead, everything is written to
        my project manager's open NetCDF file.

        @param N_iter: The number of iterations to produce after burn-in.

        """
        X = None
        if self.socket:
            wfd = defer.waitForDeferred(self.mm.setRemoteMode(self.socket))
            yield wfd; wfd.getResult()
        N_members = self.pm.m
        print "\nRunning %d iterations with %d population members" \
              % (N_iter, N_members)
        # For each iteration...
        for i in xrange(N_iter):
            print "\n%03d\n%s" % (i+1, "-"*100)
            if X is None:
                # No existing members, generate proposals directly from the
                # priors and simulate log-volatilities for them
                print "Generating %d proposals from priors " % N_members +\
                      "and simulating log-volatilities for them"
                wfd = defer.waitForDeferred(
                    self.queue.call(self.proposals, N_members))
                yield wfd
                XP = wfd.getResult()
                wfd = defer.waitForDeferred(self.simulations(XP))
                yield wfd
                wfd.getResult()
            else:
                # Regular iteration, simulate log-volatilities for the last
                # iteration's members and use proposals generated from those
                # members
                print "\nSimulating log-volatilities for the previous "+\
                      "population and generating %d proposals " % N_members +\
                      "from the population:"
                resultList = []
                dList = [self.simulations(X)]
                for vIndex, I, d in self.allocator.assembler(resultList):
                    d.addCallback(
                        lambda _:
                        self.queue.call(self.proposals, X[I], self.V[vIndex]))
                    d.addErrback(self._oops, vIndex=vIndex)
                    dList.append(d)
                # Record the previous iteration's members while the nodes work on
                # generating the proposals and simulations for the current iteration...
                if i > 0:
                    self.pm.writeParams(X)
                wfd = defer.waitForDeferred(defer.DeferredList(dList))
                yield wfd
                wfd.getResult()
                XP = resultList[0]
            print "\n"
            # Evaluate the proposals
            dList, resultList = [], []
            for vIndex, I, d in self.allocator.assembler(resultList):
                d.addCallback(lambda _: self.weights(XP[I], self.V[vIndex]))
                d.addErrback(self._oops, vIndex=vIndex)
                dList.append(d)
            wfd = defer.waitForDeferred(defer.DeferredList(dList))
            yield wfd
            wfd.getResult()
            W = resultList[0]
            I0 = s.isfinite(W).nonzero()[0]
            if len(I0):
                # Resample everything together
                I1 = self.resampler(W[I0], N_members)
                X = XP[I0][I1]
                self.allocator.updateAllocations(I0[I1])
            else:
                # No valid parameters, start over
                X = None
                self.allocator.updateAllocations()
        wfd = defer.waitForDeferred(self.pm.done())
        yield wfd
        wfd.getResult()
        print "\nPMC Run done"