optimizer.py

#!/usr/bin/env python
# optimizer.py

"""
.. module:: optimizer
.. moduleauthor:: Wouter Peters 

Revision History:
File created on 28 Jul 2010.

"""

import os
import sys
import logging
import datetime

identifier = 'Optimizer baseclass'
version    = '0.0'

################### Begin Class Optimizer ###################

class Optimizer(object):
    """
        This creates an instance of an optimization object. It handles the minimum least squares optimization
        of the state vector given a set of sample objects. Two routines will be implemented: one where the optimization
        is sequential and one where it is the equivalent matrix solution. The choice can be made based on considerations of speed
        and efficiency.
    """

    def __init__(self):
        self.Identifier = self.getid()
        self.Version    = self.getversion()

        msg                 = 'Optimizer object initialized: %s'%self.Identifier ; logging.info(msg)

    def getid(self):
        return identifier

    def getversion(self):
        return version

    def Initialize(self, dims):

        self.nlag               = dims[0]
        self.nmembers           = dims[1]
        self.nparams            = dims[2]
        self.nobs               = dims[3]
        self.CreateMatrices()

        return None

    def CreateMatrices(self):
        """ Create Matrix space needed in optimization routine """
        import numpy as np

        # mean state  [X]
        self.x               = np.zeros( (self.nlag*self.nparams,), float)
        # deviations from mean state  [X']
        self.X_prime         = np.zeros( (self.nlag*self.nparams,self.nmembers,), float)
        # mean state, transported to observation space [ H(X) ]
        self.Hx              = np.zeros( (self.nobs,), float)
        # deviations from mean state, transported to observation space [ H(X') ]
        self.HX_prime        = np.zeros( (self.nobs,self.nmembers), float)
        # observations
        self.obs             =  np.zeros( (self.nobs,), float)
        # observation ids
        self.obs_ids         =  np.zeros( (self.nobs,), float)
        # covariance of observations
        self.R               =  np.zeros( (self.nobs,self.nobs,), float)
        # localization of obs
        self.may_localize       = np.zeros(self.nobs,bool)
        # rejection of obs
        self.may_reject         = np.zeros(self.nobs,bool)
        # flags of obs
        self.flags         = np.zeros(self.nobs,int)
        # species type
        self.species       = np.zeros(self.nobs,str)
        # species type
        self.sitecode      = np.zeros(self.nobs,str)

        # species mask
        self.speciesmask    = {}

        # Total covariance of fluxes and obs in units of obs [H P H^t + R]
        self.HPHR            =  np.zeros( (self.nobs,self.nobs,), float)
        # Kalman Gain matrix
        self.KG              =  np.zeros( (self.nlag*self.nparams,self.nobs,), float)

    def StateToMatrix(self,StateVector):
        import numpy as np
        try:
            import matplotlib.pyplot as plt
        except:
            pass

        allsites=[]      # collect all obs for n=1,..,nlag
        allobs=[]      # collect all obs for n=1,..,nlag
        allmdm=[]      # collect all mdm for n=1,..,nlag
        allsamples=[]  # collect all model samples for n=1,..,nlag
        allids=[]  # collect all model samples for n=1,..,nlag
        allreject=[]  # collect all model samples for n=1,..,nlag
        alllocalize=[]  # collect all model samples for n=1,..,nlag
        allflags=[]  # collect all model samples for n=1,..,nlag
        allspecies=[]  # collect all model samples for n=1,..,nlag
        allsimulated=None  # collect all members model samples for n=1,..,nlag

        for n in range(self.nlag):

            Samples                                             = StateVector.ObsToAssimmilate[n]
            members                                             = StateVector.EnsembleMembers[n]
            self.x[n*self.nparams:(n+1)*self.nparams]           = members[0].ParameterValues
            self.X_prime[n*self.nparams:(n+1)*self.nparams,:]   = np.transpose(np.array([m.ParameterValues for m in members]))

            if Samples != None:

                self.rejection_threshold                        = Samples.rejection_threshold

                allreject.extend(                                 Samples.Data.getvalues('may_reject') )
                alllocalize.extend(                               Samples.Data.getvalues('may_localize') )
                allflags.extend(                                  Samples.Data.getvalues('flag') )
                allspecies.extend(                                Samples.Data.getvalues('species') )
                allobs.extend(                                    Samples.Data.getvalues('obs') )
                allsites.extend(                                  Samples.Data.getvalues('code') )
                allmdm.extend(                                    Samples.Data.getvalues('mdm')  )
                allids.extend(                                    Samples.Data.getvalues('id')  )

                simulatedensemble =                               Samples.Data.getvalues('simulated')

                if allsimulated == None :
                    allsimulated = np.array(simulatedensemble)
                else:
                    allsimulated = np.concatenate((allsimulated,np.array(simulatedensemble)),axis=0)



        self.obs[:]                                             = np.array(allobs)
        self.obs_ids[:]                                         = np.array(allids)
        self.HX_prime[:,:]                                      = np.array(allsimulated)
        self.Hx[:]                                              = self.HX_prime[:,0]

        self.may_reject[:]                                      = np.array(allreject)
        self.may_localize[:]                                    = np.array(alllocalize)
        self.flags[:]                                           = np.array(allflags)
        self.species[:]                                         = np.array(allspecies)
        self.sitecode                                           = allsites

        self.X_prime                                            = self.X_prime - self.x[:,np.newaxis] # make into a deviation matrix
        self.HX_prime                                           = self.HX_prime - self.Hx[:,np.newaxis] # make a deviation matrix


        for i,mdm in enumerate(allmdm):

            self.R[i,i]                                             = mdm**2

        return None

    def MatrixToState(self,StateVector):
        import numpy as np

        for n in range(self.nlag):

            members                              = StateVector.EnsembleMembers[n]
            for m,mem in enumerate(members):
                members[m].ParameterValues[:]    = self.X_prime[n*self.nparams:(n+1)*self.nparams,m] + self.x[n*self.nparams:(n+1)*self.nparams]     

        StateVector.isOptimized = True
        msg = 'Returning optimized data to the StateVector, setting "StateVector.isOptimized = True" ' ; logging.debug(msg)
        return None

    def WriteDiagnostics(self,DaCycle, StateVector,type='prior'):
        """
            Open a NetCDF file and write diagnostic output from optimization process:

                - calculated residuals
                - model-data mismatches
                - HPH^T
                - prior ensemble of samples
                - posterior ensemble of samples
                - prior ensemble of fluxes
                - posterior ensemble of fluxes

            The type designation refers to the writing of prior or posterior data and is used in naming the variables"
        """
        import da.tools.io4 as io
        #import da.tools.io as io

        outdir          = DaCycle['dir.diagnostics']
        filename        = os.path.join(outdir,'optimizer.%s.nc'% DaCycle['time.start'].strftime('%Y%m%d') )
        DaCycle.OutputFileList += ( filename, )

        # Open or create file

        if type == 'prior':
            f           = io.CT_CDF(filename,method='create')
            msg         = 'Creating new diagnostics file for optimizer (%s)' % filename ; logging.debug(msg)
        elif type == 'optimized':
            f           = io.CT_CDF(filename,method='write')
            msg         = 'Opening existing diagnostics file for optimizer (%s)' % filename ; logging.debug(msg)

        # Add dimensions 

        dimparams       = f.AddParamsDim(self.nparams)
        dimmembers      = f.AddMembersDim(self.nmembers)
        dimlag          = f.AddLagDim(self.nlag, unlimited=False)
        dimobs          = f.AddObsDim(self.nobs)
        dimstate        = f.AddDim('nstate',self.nparams*self.nlag)
        dim200char      = f.AddDim('string_of200chars',200)

        # Add data, first the ones that are written both before and after the optimization

        data =  self.x

        savedict                = io.std_savedict.copy() 
        savedict['name']        = "statevectormean_%s" % type
        savedict['long_name']   = "full_statevector_mean_%s" % type
        savedict['units']       = "unitless"
        savedict['dims']        = dimstate
        savedict['values']      = data.tolist()
        savedict['comment']     = 'Full %s state vector mean '% type
        dummy                   = f.AddData(savedict)

        data =  self.X_prime

        savedict                = io.std_savedict.copy()
        savedict['name']        = "statevectordeviations_%s" % type
        savedict['long_name']   = "full_statevector_deviations_%s" % type
        savedict['units']       = "unitless"
        savedict['dims']        = dimstate+dimmembers
        savedict['values']      = data.tolist()
        savedict['comment']     = 'Full state vector %s deviations as resulting from the optimizer'% type
        dummy                   = f.AddData(savedict)

        data =  self.Hx

        savedict                = io.std_savedict.copy()
        savedict['name']        = "modelsamplesmean_%s"%type
        savedict['long_name']   = "modelsamplesforecastmean_%s" %type
        savedict['units']       = "mol mol-1"
        savedict['dims']        = dimobs
        savedict['values']      = data.tolist()
        savedict['comment']     = '%s mean mixing ratios based on %s state vector'% (type,type,)
        dummy                   = f.AddData(savedict)

        data =  self.HX_prime

        savedict                = io.std_savedict.copy()
        savedict['name']        = "modelsamplesdeviations_%s"% type
        savedict['long_name']   = "modelsamplesforecastdeviations_%s"%type
        savedict['units']       = "mol mol-1"
        savedict['dims']        = dimobs+dimmembers
        savedict['values']      = data.tolist()
        savedict['comment']     = '%s mixing ratio deviations based on %s state vector'% (type,type,)
        dummy                   = f.AddData(savedict)


        # Continue with prior only data

        if type == 'prior':

            data = self.sitecode

            savedict                = io.std_savedict.copy()
            savedict['name']        = "sitecode"
            savedict['long_name']   = "site code propagated from observation file"
            savedict['dtype']       = "char"
            savedict['dims']        = dimobs+dim200char
            savedict['values']      = data
            savedict['missing_value'] = '!'
            status                   = f.AddData(savedict)


            data =  self.obs

            savedict                = io.std_savedict.copy()
            savedict['name']        = "observed"
            savedict['long_name']   = "observedvalues"
            savedict['units']       = "mol mol-1"
            savedict['dims']        = dimobs
            savedict['values']      = data.tolist()
            savedict['comment']     = 'Observations used in optimization'
            dummy                   = f.AddData(savedict)

            data =  self.obs_ids

            savedict                = io.std_savedict.copy()
            savedict['name']        = "obspack_num"
            savedict['dtype']       = "int64"
            savedict['long_name']   = "Unique_ObsPack_observation_number"
            savedict['units']       = ""
            savedict['dims']        = dimobs
            savedict['values']      = data.tolist()
            savedict['comment']     = 'Unique observation number across the entire ObsPack distribution'
            dummy                   = f.AddData(savedict)


            data =  self.R

            savedict                = io.std_savedict.copy()
            savedict['name']        = "modeldatamismatchvariance"
            savedict['long_name']   = "modeldatamismatch variance"
            savedict['units']       = "[mol mol-1]^2"
            savedict['dims']        = dimobs+dimobs
            savedict['values']      = data.tolist()
            savedict['comment']     = 'Variance of mole fractions resulting from model-data mismatch'
            dummy                   = f.AddData(savedict)

        # Continue with posterior only data

        elif type == 'optimized':

            data =  self.HPHR

            savedict                = io.std_savedict.copy()
            savedict['name']        = "totalmolefractionvariance"
            savedict['long_name']   = "totalmolefractionvariance"
            savedict['units']       = "[mol mol-1]^2"
            savedict['dims']        = dimobs+dimobs
            savedict['values']      = data.tolist()
            savedict['comment']     = 'Variance of mole fractions resulting from prior state and model-data mismatch'
            dummy                   = f.AddData(savedict)

            data =  self.flags

            savedict                = io.std_savedict.copy()
            savedict['name']        = "flag"
            savedict['long_name']   = "flag_for_obs_model"
            savedict['units']       = "None"
            savedict['dims']        = dimobs
            savedict['values']      = data.tolist()
            savedict['comment']     = 'Flag (0/1/2/99) for observation value, 0 means okay, 1 means QC error, 2 means rejected, 99 means not sampled'
            dummy                   = f.AddData(savedict)

            data =  self.KG

            savedict                = io.std_savedict.copy()
            savedict['name']        = "kalmangainmatrix"
            savedict['long_name']   = "kalmangainmatrix"
            savedict['units']       = "unitless molefraction-1"
            savedict['dims']        = dimstate+dimobs
            savedict['values']      = data.tolist()
            savedict['comment']     = 'Kalman gain matrix of all obs and state vector elements'
            dummy                   = f.AddData(savedict)

        dummy                       = f.close()
        msg = 'Diagnostics file closed ' ; logging.debug(msg)

        return None

    def SerialMinimumLeastSquares(self):
        """ Make minimum least squares solution by looping over obs"""
        import numpy as np
        import numpy.linalg as la

        tvalue=1.97591

        for n in range(self.nobs):

            # Screen for flagged observations (for instance site not found, or no sample written from model)

            if self.flags[n] != 0:
                msg = 'Skipping observation %d because of flag value %d'%(n,self.flags[n]) ; logging.debug(msg)
                continue


            # Screen for outliers greather than 3x model-data mismatch, only apply if obs may be rejected

            res                         = self.obs[n]-self.Hx[n]

            if self.may_reject[n]:
                threshold =  self.rejection_threshold*np.sqrt(self.R[n,n])
                if np.abs(res) > threshold:
                    msg = 'Rejecting observation %d because residual (%f) exceeds threshold (%f)'%(n,res,threshold) ; logging.debug(msg)
                    self.flags[n] = 2
                    continue

            PHt                         = 1./(self.nmembers-1)*np.dot(self.X_prime,self.HX_prime[n,:])
            self.HPHR[n,n]              = 1./(self.nmembers-1)*(self.HX_prime[n,:]*self.HX_prime[n,:]).sum()+self.R[n,n]

            self.KG[:,n]                = PHt/self.HPHR[n,n]

            if self.may_localize[n]:
                dummy                   = self.Localize(n)
                msg = 'Localized observation %d'%(n,) ; logging.debug(msg)
            else:
                msg = 'Not allowed to Localize observation %d'%(n,) ; logging.debug(msg)

            alpha                       = np.double(1.0)/(np.double(1.0)+np.sqrt( (self.R[n,n])/self.HPHR[n,n] ) )

            self.x[:]                   = self.x + self.KG[:,n]*res

            for r in range(self.nmembers):
                self.X_prime[:,r]       = self.X_prime[:,r]-alpha*self.KG[:,n]*(self.HX_prime[n,r])

#WP !!!! Very important to first do all obervations from n=1 through the end, and only then update 1,...,n. The current observation
#WP      should always be updated last because it features in the loop of the adjustments !!!!

            for m in range(n+1,self.nobs):
                res                     = self.obs[n]-self.Hx[n]
                fac                     = 1.0/(self.nmembers-1)*(self.HX_prime[n,:]*self.HX_prime[m,:]).sum()/self.HPHR[n,n]
                self.Hx[m]              = self.Hx[m] + fac * res
                self.HX_prime[m,:]      = self.HX_prime[m,:] - alpha*fac*self.HX_prime[n,:]

            for m in range(1,n+1):
                res                     = self.obs[n]-self.Hx[n]
                fac                     = 1.0/(self.nmembers-1)*(self.HX_prime[n,:]*self.HX_prime[m,:]).sum()/self.HPHR[n,n]
                self.Hx[m]              = self.Hx[m] + fac * res
                self.HX_prime[m,:]      = self.HX_prime[m,:] - alpha*fac*self.HX_prime[n,:]

            
    def BulkMinimumLeastSquares(self):
        """ Make minimum least squares solution by solving matrix equations"""
        import numpy as np
        import numpy.linalg as la

        # Create full solution, first calculate the mean of the posterior analysis

        HPH                 = np.dot(self.HX_prime,np.transpose(self.HX_prime))/(self.nmembers-1)   # HPH = 1/N * HX' * (HX')^T
        self.HPHR[:,:]      = HPH+self.R                                                            # HPHR = HPH + R
        HPb                 = np.dot(self.X_prime,np.transpose(self.HX_prime))/(self.nmembers-1)    # HP = 1/N X' * (HX')^T
        self.KG[:,:]        = np.dot(HPb,la.inv(self.HPHR))                                         # K = HP/(HPH+R)

        for n in range(self.nobs):
            dummy           = self.Localize(n)

        self.x[:]           = self.x + np.dot(self.KG,self.obs-self.Hx)                             # xa = xp + K (y-Hx)

        # And next make the updated ensemble deviations. Note that we calculate P by using the full equation (10) at once, and 
        # not in a serial update fashion as described in Whitaker and Hamill. 
        # For the current problem with limited N_obs this is easier, or at least more straightforward to do.

        I                   = np.identity(self.nlag*self.nparams)
        sHPHR               = la.cholesky(self.HPHR)                                  # square root of HPH+R
        part1               = np.dot(HPb,np.transpose(la.inv(sHPHR)))                 # HP(sqrt(HPH+R))^-1
        part2               = la.inv(sHPHR+np.sqrt(self.R))                           # (sqrt(HPH+R)+sqrt(R))^-1
        Kw                  = np.dot(part1,part2)                                     # K~
        self.X_prime[:,:]   = np.dot(I,self.X_prime)-np.dot(Kw,self.HX_prime)         # HX' = I - K~ * HX'

        P_opt               = np.dot(self.X_prime,np.transpose(self.X_prime))/(self.nmembers-1)

        # Now do the adjustments of the modeled mixing ratios using the linearized ensemble. These are not strictly needed but can be used
        # for diagnosis.

        part3               = np.dot(HPH,np.transpose(la.inv(sHPHR)))                           # HPH(sqrt(HPH+R))^-1
        Kw                  = np.dot(part3,part2)                                               # K~
        self.Hx[:]          = self.Hx + np.dot(np.dot(HPH,la.inv(self.HPHR)),self.obs-self.Hx)  # Hx  = Hx+ HPH/HPH+R (y-Hx)
        self.HX_prime[:,:]  = self.HX_prime-np.dot(Kw,self.HX_prime)                            # HX' = HX'- K~ * HX'

        msg = 'Minimum Least Squares solution was calculated, returning' ; logging.info(msg)

        return None

    def SetLocalization(self):
        """ determine which localization to use """

        self.localization = True
        self.localizetype = "None"
    
        msg       = "Current localization option is set to %s"%self.localizetype  ; logging.info(msg)

    def Localize(self,n):
        """ localize the Kalman Gain matrix """
        import numpy as np
    
        if not self.localization: return 

        return

################### End Class Optimizer ###################



if __name__ == "__main__":

    sys.path.append('../../')

    import os
    import sys
    from da.tools.general import StartLogger 
    from da.tools.initexit import CycleControl 
    from da.ct.statevector import CtStateVector, PrepareState
    from da.ct.obs import CtObservations
    import numpy as np
    import datetime
    import da.tools.rc as rc

    opts = ['-v']
    args = {'rc':'../../da.rc','logfile':'da_initexit.log','jobrcfilename':'test.rc'}

    StartLogger()
    DaCycle = CycleControl(opts,args)

    DaCycle.Initialize()
    print DaCycle

    StateVector = PrepareState(DaCycle)

    samples = CtObservations(DaCycle.DaSystem,datetime.datetime(2005,3,5))
    dummy = samples.AddObs()
    dummy = samples.AddSimulations('/Users/peters/tmp/test_da/output/20050305/samples.000.nc')


    nobs      = len(samples.Data)
    dims      = ( int(DaCycle['time.nlag']),
                  int(DaCycle['da.optimizer.nmembers']),
                  int(DaCycle.DaSystem['nparameters']),
                  nobs,  )

    opt = CtOptimizer(dims)

    opt.StateToMatrix(StateVector)

    opt.MinimumLeastSquares()

    opt.MatrixToState(StateVector)