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

# svpmc: Stochastic Volatility Inference via Population Monte Carlo
#
# Copyright (C) 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

"""
Management of an overall SVpmc run.
"""

import os.path, re
import scipy as s
from Scientific.IO import NetCDF
import asynqueue

import tseries, params, model


class NullQueue(object):
    """
    I act like an AsynQueue except that I run everything in the main
    thread. I'm mostly useful for debugging.
    """
    def __init__(self, *args):
        from twisted.internet import defer
        self.defer = defer

    def call(self, f, *args):
        result = f(*args)
        if isinstance(result, self.defer.Deferred):
            return result
        return self.defer.succeed(result)

    def shutdown(self):
        return self.defer.succeed(None)


class ProjectManager(object):
    """
    Instantiate me with the path of a file within a project directory that
    defines a project specification and the name of a NetCDF file to be created
    in the project directory. If no NetCDF file is specified, one will be
    created with the name 'svpmc.nc'.

    You can specify the PMC population size with the keyword I{m}. It defaults
    to 10,000.

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

    @ivar tables: A dict whose entries contain rows of tables parsed from my
      project specification.

    @ivar xcorrs: The number of cross-correlations between series, computed as
      C{(p**2 + p)/2}

    @ivar m: The number of population members in each population monte carlo
      iteration.

    @ivar paramNames: A list of the names of arrays in the parameter containers
      used in the project for both primary and secondary parameters.

    @ivar D: The number of jump deviations used by the D-kernel PMC
      inference engine that will be running the project.

    @ivar Df: The value of a given jump deviation as a fraction of the previous
      one (the first value is 1.0)

    @ivar priors: An instance of L{params.PriorContainer} constructed in
      accordance with my project specification.
    
    """
    re_dashes = re.compile("^\-{10,}$")
    re_xy = re.compile("([0-9]+),([0-9]+)")
    dimensions = {
        'n':        'sample',
        'p':        'series',               
        'xcorrs':   'cross_correlation',
        'm':        'member',    
        }
    def __init__(self, specFile, ncFile="svpmc.nc", m=1000, readOnly=False):
        self.m = m
        self.iteration = 0
        specDir = os.path.dirname(specFile)
        self.tables = self._parseSpec(specFile)
        # Set some simulation parameters
        for name, value, null in self.tables['variable']:
            if "." in value:
                value = float(value)
            else:
                value = int(value)
            setattr(self, name, value)
        V = [self.Ds]
        for k in xrange(self.D-1):
            V.append(self.Df * V[-1])
        self.V = s.array(V)
        # Observations
        tsData, seriesTitles = self._setupTimeSeries(specDir)
        self.p, self.n = tsData.shape
        # Parameters
        paramTitles, dimensions = self._setupParams()
        # Open the NetCDF file
        self.cdf = NetCDF.NetCDFFile(
            os.path.join(specDir, ncFile), ('w', 'r')[int(readOnly)])
        # If in read-only mode, we're done now
        if readOnly:
            return
        # Write mode, write some initial stuff to the NetCDF file...
        self._setupCDF(tsData, seriesTitles, paramTitles, dimensions)
        for k, title in enumerate(seriesTitles):
            obs = self.cdf.variables['observations']
            obs[k,:] = tsData[k,:].astype('f')
            setattr(obs, "series_%02d" % (k,), title)
        # ...and construct a queue and model manager for a simulation
        self.queue = asynqueue.ThreadQueue(1)
        #self.queue = NullQueue()
        self.mgr = model.ModelManager(
            self, tsData, wiggle=self.wiggle, Ni=self.Ni)
    
    def _parseSpec(self, filePath):
        tables = {}
        Nf = None
        linePrev = ""
        fh = open(filePath)
        for line in fh:
            line = line.strip()
            if line.startswith('#'):
                continue
            if not line:
                Nf = None
            elif Nf:
                fieldList = line.split()
                fieldList += [''] * (Nf - len(fieldList))
                fieldList[Nf-1] = " ".join(fieldList[Nf-1:])
                rows.append(fieldList[:Nf])
            elif self.re_dashes.match(line) and linePrev:
                fields = linePrev.split()
                Nf = len(fields)
                rows = tables[fields[0].lower()] = []
            else:
                linePrev = line
        fh.close()
        return tables
    
    def _setupTimeSeries(self, fileDir):
        """
        Given the specification defined in my I{tables} dict, loads a list of
        L{tseries.TimeSeries} objects from '.dat' files in the specified
        directory.

        Returns a 2-D array of the time series data, samples across columns and
        time series across rows, along with a list of the time series titles.
        """
        prevObj = None
        tsList, titles = [], []
        for series, transform, title in self.tables['series']:
            filePath = os.path.join(fileDir, series) + ".dat"
            tsObject = tseries.TimeSeries(title, filePath)
            for methodName in transform.split(','):
                getattr(tsObject, methodName)()
            if prevObj:
                tsObject = tsObject.intersect(prevObj)
            tsList.append(tsObject)
            titles.append(title)
            prevObj = tsObject
        return s.row_stack([ts() for ts in tsList]), titles

    def _setupParams(self):
        """
        Given the specification defined in my I{tables} dict, constructs an
        instance of L{params.PriorContainer} with L{params.FlexArray} objects
        loaded with a L{params.Prior} object for each array element of each of
        my parameters.

        Sets my I{priorContainer} attribute to the container and returns a list
        of the parameter titles and lists of dimensions compatible with my
        NetCDF specification.
        """
        def parseDims(dims):
            shape = []
            dimSequence = []
            for dim in dims.split(','):
                if dim not in self.dimensions:
                    raise ValueError(
                        "Invalid dimension specification '%s'" % dims)
                shape.append(getattr(self, dim))
                dimSequence.append(self.dimensions[dim])
            dimensions.append(dimSequence)
            return tuple(shape)

        def constructPriors(priorSpec, shape):
            index = priorSpec[0]
            kw = {'dname':priorSpec[1], 'V':self.V}
            for k, name in enumerate(('loc', 'scale', 'a', 'b')):
                stringValue = priorSpec[k+2]
                if stringValue != '':
                    kw[name] = float(stringValue)
            for j in xrange(shape[0]):
                if len(shape) == 1:
                    if index.isdigit() and j != int(index):
                        continue
                    pa[j] = params.Prior(**kw)
                else:
                    for k in xrange(shape[1]):
                        if index == "I" and j != k:
                            continue
                        if index != ":":
                            match = self.re_xy.match(index)
                            if match and match.group(0) != "%d,%d" % (j,k):
                                continue
                        pa[j,k] = params.Prior(**kw)
        
        self.primaryParamNames = []
        titles, dimensions = [], []
        self.paramInfo = {}
        xcorrs = self.xcorrs = int(0.5*(self.p**2 + self.p))
        self.priors = params.PriorContainer(self.p, self.n)
        for name, dims, title in self.tables['parameter']:
            shape = parseDims(dims)
            self.paramInfo[name] = shape, title
            pa = self.priors.priorArray(name, *shape)
            for priorSpec in self.tables[name]:
                constructPriors(priorSpec, shape)
            titles.append(title)
            self.primaryParamNames.append(name)
        self.priors.primaryParamNames = self.primaryParamNames
        self.derivedParamNames = []
        for name, sourceParams, description in self.tables['derivation']:
            self.derivedParamNames.append(name)
        self.priors.derivedParamNames = self.derivedParamNames
        self.paramNames = self.primaryParamNames + self.derivedParamNames
        return titles, dimensions

    def _setupCDF(self, tsData, sTitles, pTitles, pDims):
        """
        """
        cdf = self.cdf
        cdf.title = self.tables['project'][0][0]
        # Dimensions
        for attrName, name in self.dimensions.iteritems():
            cdf.createDimension(name, getattr(self, attrName))
        cdf.createDimension("iteration", None)
        # Time Series
        obs = cdf.createVariable(
            "observations", 'f', ('series', 'sample'))
        obs.long_name = "Observations, %d time series " % self.p +\
                        "with %d samples each" % self.n
        # Parameters
        for k, name in enumerate(self.primaryParamNames):
            pDim = ['iteration', 'member'] + pDims[k]
            param = cdf.createVariable(name, 'f', tuple(pDim))
            param.long_name = pTitles[k]
        # Log-Likelihood, log-density of priors and jumps
        L = cdf.createVariable("Lx", 'f', ('iteration', 'member'))
        L.long_name = "Log-Likelihood of Observations Given Parameters"
        L = cdf.createVariable("Lp", 'f', ('iteration', 'member'))
        L.long_name = "Log-Density of Parameters from Prior Distributions"
        L = cdf.createVariable("Lj", 'f', ('iteration', 'member'))
        L.long_name = "Log-Density of Proposal Jumps"
        # Done setting up, save what we've got thus far
        cdf.sync()

    def writeParams(self, parameters):
        """
        Write the supplied population of I{parameters}, housed in a
        L{param.FlexArray} instance containing L{param.ParameterContainer}
        objects, as the latest iteration of the open CDF record.
        """
        if len(parameters) != self.m:
            raise ValueError(
                "Each iteration record must have %d parameters" % self.m)
        # Write the parameters to their variables
        for name in self.primaryParamNames:
            var = self.cdf.variables[name]
            # Iterate over a 1-D FlexArray of the population's parameter arrays
            for member, paramArray in enumerate(getattr(parameters, name)):
                var[self.iteration, member, :] = paramArray.astype('f')
        # Write the log-volatility, prior and jump log-density values for each
        # population member
        for name in ('Lx', 'Lp', 'Lj'):
            var = self.cdf.variables[name]
            var[self.iteration, :] = getattr(parameters, name).astype('f')
        # Done writing this iteration of the CDF record
        self.cdf.sync()
        self.iteration += 1

    def done(self):
        """
        Call this when the project is done.
        """
        if hasattr(self, 'mgr'):
            d = self.mgr.shutdown()
            d.addCallback(lambda _: self.queue.shutdown())
            d.addCallback(lambda _: self.cdf.close())
            return d
        self.cdf.close()