statevector.py 25.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
"""CarbonTracker Data Assimilation Shell (CTDAS) Copyright (C) 2017 Wouter Peters. 
Users are recommended to contact the developers (wouter.peters@wur.nl) to receive
updates of the code. See also: http://www.carbontracker.eu. 

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, 
version 3. 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 GNU General Public License for more details. 

You should have received a copy of the GNU General Public License along with this 
program. If not, see <http://www.gnu.org/licenses/>."""
13
14
15
16
#!/usr/bin/env python
# ct_statevector_tools.py

"""
Peters, Wouter's avatar
Peters, Wouter committed
17
18
.. module:: statevector
.. moduleauthor:: Wouter Peters 
19
20
21
22

Revision History:
File created on 28 Jul 2010.

Peters, Wouter's avatar
Peters, Wouter committed
23
24
25
26
27
28
29
30
31
32
33
The module statevector implements the data structure and methods needed to work with state vectors (a set of unknown parameters to be optimized by a DA system) of different lengths, types, and configurations. Two baseclasses together form a generic framework:
    * :class:`~da.baseclasses.statevector.StateVector`
    * :class:`~da.baseclasses.statevector.EnsembleMember`

As usual, specific implementations of StateVector objects are done through inheritance form these baseclasses. An example of designing 
your own baseclass StateVector we refer to :ref:`tut_chapter5`.

.. autoclass:: da.baseclasses.statevector.StateVector 

.. autoclass:: da.baseclasses.statevector.EnsembleMember 

34
35
36
"""

import os
37
import sys
38
import logging
39
import numpy as np
40
from datetime import timedelta
41
import da.tools.io4 as io
42
43

identifier = 'Baseclass Statevector '
karolina's avatar
karolina committed
44
version = '0.0'
45
46
47
48
49

################### Begin Class EnsembleMember ###################

class EnsembleMember(object):
    """ 
Peters, Wouter's avatar
Peters, Wouter committed
50
        An ensemble member object consists of:
51
52
           * a member number
           * parameter values
Peters, Wouter's avatar
Peters, Wouter committed
53
54
55
56
           * an observation object to hold sampled values for this member

        Ensemble members are initialized by passing only an ensemble member number, all data is added by methods 
        from the :class:`~da.baseclasses.statevector.StateVector`. Ensemble member objects have almost no functionality 
57
        except to write their data to file using method :meth:`~da.baseclasses.statevector.EnsembleMember.write_to_file`
Peters, Wouter's avatar
Peters, Wouter committed
58
59

        .. automethod:: da.baseclasses.statevector.EnsembleMember.__init__ 
60
        .. automethod:: da.baseclasses.statevector.EnsembleMember.write_to_file 
61
        .. automethod:: da.baseclasses.statevector.EnsembleMember.AddCustomFields 
Peters, Wouter's avatar
Peters, Wouter committed
62

63
64
65
    """

    def __init__(self, membernumber):
Peters, Wouter's avatar
Peters, Wouter committed
66
67
68
69
70
71
        """
           :param memberno: integer ensemble number
           :rtype: None

           An EnsembleMember object is initialized with only a number, and holds two attributes as containter for later
           data:
72
                * param_values, will hold the actual values of the parameters for this data
Peters, Wouter's avatar
Peters, Wouter committed
73
74
75
                * ModelSample, will hold an :class:`~da.baseclasses.obs.Observation` object and the model samples resulting from this members' data

        """
karolina's avatar
karolina committed
76
        self.membernumber = membernumber   # the member number
77
        self.param_values = None           # Parameter values of this member
78
79
80
81
82

################### End Class EnsembleMember ###################

################### Begin Class StateVector ###################

karolina's avatar
karolina committed
83

84
class StateVector(object):
Peters, Wouter's avatar
Peters, Wouter committed
85
86
87
88
89
90
91
92
93
94
95
    """ 
    The StateVector object first of all contains the data structure of a statevector, defined by 3 attributes that define the 
    dimensions of the problem in parameter space:
        * nlag
        * nparameters
        * nmembers

    The fourth important dimension `nobs` is not related to the StateVector directly but is initialized to 0, and later on 
    modified to be used in other parts of the pipeline:
        * nobs

96
    These values are set as soon as the :meth:`~da.baseclasses.statevector.StateVector.setup` is called from the :ref:`pipeline`. 
Peters, Wouter's avatar
Peters, Wouter committed
97
98
99
100
101
102
103
    Additionally, the value of attribute `isOptimized` is set to `False` indicating that the StateVector holds a-priori values 
    and has not been modified by the :ref:`optimizer`.

    StateVector objects can be filled with data in two ways
        1. By reading the data from file
        2. By creating the data through a set of method calls

104
    Option (1) is invoked using method :meth:`~da.baseclasses.statevector.StateVector.read_from_file`. 
105
    Option (2) consists of a call to method :meth:`~da.baseclasses.statevector.StateVector.make_new_ensemble`
Peters, Wouter's avatar
Peters, Wouter committed
106

107
    Once the StateVector object has been filled with data, it is used in the pipeline and a few more methods are
Peters, Wouter's avatar
Peters, Wouter committed
108
    invoked from there:
109
        * :meth:`~da.baseclasses.statevector.StateVector.propagate`, to advance the StateVector from t=t to t=t+1
110
        * :meth:`~da.baseclasses.statevector.StateVector.write_to_file`, to write the StateVector to a NetCDF file for later use
Peters, Wouter's avatar
Peters, Wouter committed
111
112
113

    The methods are described below:

114
    .. automethod:: da.baseclasses.statevector.StateVector.setup 
115
116
    .. automethod:: da.baseclasses.statevector.StateVector.read_from_file
    .. automethod:: da.baseclasses.statevector.StateVector.write_to_file
117
118
119
    .. automethod:: da.baseclasses.statevector.StateVector.make_new_ensemble
    .. automethod:: da.baseclasses.statevector.StateVector.propagate
    .. automethod:: da.baseclasses.statevector.StateVector.write_members_to_file
120
121
122

    Finally, the StateVector can be mapped to a gridded array, or to a vector of TransCom regions, using:

123
124
125
126
    .. automethod:: da.baseclasses.statevector.StateVector.grid2vector
    .. automethod:: da.baseclasses.statevector.StateVector.vector2grid
    .. automethod:: da.baseclasses.statevector.StateVector.vector2tc
    .. automethod:: da.baseclasses.statevector.StateVector.state2tc
Peters, Wouter's avatar
Peters, Wouter committed
127
128

    """
129

130
    def __init__(self):
karolina's avatar
karolina committed
131
132
        self.ID = identifier
        self.version = version
133

134
        # The following code allows the object to be initialized with a dacycle object already present. Otherwise, it can
135
136
        # be added at a later moment.

karolina's avatar
karolina committed
137
        logging.info('Statevector object initialized: %s' % self.ID)
138

139
    def setup(self, dacycle):
Peters, Wouter's avatar
Peters, Wouter committed
140
        """
141
        setup the object by specifying the dimensions. 
142
143
144
145
146
147
        There are two major requirements for each statvector that you want to build:
        
            (1) is that the statevector can map itself onto a regular grid
            (2) is that the statevector can map itself (mean+covariance) onto TransCom regions

        An example is given below.
Peters, Wouter's avatar
Peters, Wouter committed
148
        """
149

150
151
152
        self.nlag = int(dacycle['time.nlag'])
        self.nmembers = int(dacycle['da.optimizer.nmembers'])
        self.nparams = int(dacycle.dasystem['nparameters'])
karolina's avatar
karolina committed
153
        self.nobs = 0
154
        
155
        self.obs_to_assimilate = ()  # empty containter to hold observations to assimilate later on
156
157
158
159

        # These list objects hold the data for each time step of lag in the system. Note that the ensembles for each time step consist 
        # of lists of EnsembleMember objects, we define member 0 as the mean of the distribution and n=1,...,nmembers as the spread.

brunner's avatar
brunner committed
160
        self.ensemble_members = list(range(self.nlag))
161
162

        for n in range(self.nlag):
163
            self.ensemble_members[n] = []
164

165

166
        self.gridmap = np.random.randint(low=1,high=self.nparams+1,size=(180,360,))
karolina's avatar
karolina committed
167
        self.griddict = {}
168
169
        for r in range(1, self.nparams+1):
            sel = np.nonzero(self.gridmap.flat == r) 
karolina's avatar
karolina committed
170
171
172
            if len(sel[0]) > 0: 
                self.griddict[r] = sel
        logging.debug("A dictionary to map grids to states and vice versa was created")
173

174
175
        # Create a mask for species/unknowns

176
        self.make_species_mask()
177

178
    def make_species_mask(self):
179
180
181
182
183
184
185
186
187
188
189
190

        """

        This method creates a dictionary with as key the name of a tracer, and as values an array of 0.0/1.0 values 
        specifying which StateVector elements are constrained by this tracer. This mask can be used in 
        the optimization to ensure that certain types of osbervations only update certain unknowns.

        An example would be that the tracer '14CO2' can be allowed to only map onto fossil fuel emissions in the state

        The form of the mask is:

        {'co2': np.ones(self.nparams), 'co2c14', np.zeros(self.nparams)  }
191

192
193
        so that 'co2' maps onto all parameters, and 'co2c14' on none at all. These arrays are used in the Class 
        optimizer when state updates are actually performed
194

195
196
        """
        self.speciesdict = {'co2': np.ones(self.nparams)}
karolina's avatar
karolina committed
197
        logging.debug("A species mask was created, only the following species are recognized in this system:")
198
        for k in self.speciesdict.keys(): 
karolina's avatar
karolina committed
199
            logging.debug("   ->    %s" % k)
200

201

202
    def make_new_ensemble(self, lag, covariancematrix=[None]):
Peters, Wouter's avatar
Peters, Wouter committed
203
204
205
206
207
208
209
210
211
212
213
214
        """ 
        :param lag: an integer indicating the time step in the lag order
        :param covariancematrix: a matrix to draw random values from
        :rtype: None
    
        Make a new ensemble, the attribute lag refers to the position in the state vector. 
        Note that lag=1 means an index of 0 in python, hence the notation lag-1 in the indexing below.
        The argument is thus referring to the lagged state vector as [1,2,3,4,5,..., nlag]

        The optional covariance object to be passed holds a matrix of dimensions [nparams, nparams] which is
        used to draw ensemblemembers from. If this argument is not passed it will ne substituted with an 
        identity matrix of the same dimensions.
215
216
217
218

        """    

        if covariancematrix == None: 
karolina's avatar
karolina committed
219
            covariancematrix = np.identity(self.nparams)
220
221
222

        # Make a cholesky decomposition of the covariance matrix

223

224
225
226
227
        try:
            _, s, _ = np.linalg.svd(covariancematrix)
        except:
            s = np.linalg.svd(covariancematrix, full_matrices=1, compute_uv=0) #Cartesius fix
karolina's avatar
karolina committed
228
229
        dof = np.sum(s) ** 2 / sum(s ** 2)
        C = np.linalg.cholesky(covariancematrix)
230

karolina's avatar
karolina committed
231
232
        logging.debug('Cholesky decomposition has succeeded ')
        logging.info('Appr. degrees of freedom in covariance matrix is %s' % (int(dof)))
233
234
235
236


        # Create mean values 

237
        newmean = np.ones(self.nparams, float) # standard value for a new time step is 1.0
238
239
240

        # If this is not the start of the filter, average previous two optimized steps into the mix

karolina's avatar
karolina committed
241
        if lag == self.nlag - 1 and self.nlag >= 3:
242
243
            newmean += self.ensemble_members[lag - 1][0].param_values + \
                                           self.ensemble_members[lag - 2][0].param_values 
244
            newmean = newmean / 3.0
245
246
247

        # Create the first ensemble member with a deviation of 0.0 and add to list

karolina's avatar
karolina committed
248
        newmember = EnsembleMember(0)
249
250
        newmember.param_values = newmean.flatten()  # no deviations
        self.ensemble_members[lag].append(newmember)
251

252
        # Create members 1:nmembers and add to ensemble_members list
253

karolina's avatar
karolina committed
254
255
        for member in range(1, self.nmembers):
            rands = np.random.randn(self.nparams)
256

karolina's avatar
karolina committed
257
            newmember = EnsembleMember(member)
258
259
            newmember.param_values = np.dot(C, rands) + newmean
            self.ensemble_members[lag].append(newmember)
260

karolina's avatar
karolina committed
261
        logging.debug('%d new ensemble members were added to the state vector # %d' % (self.nmembers, (lag + 1)))
262
263


264
    def propagate(self, dacycle):
265
        """ 
Peters, Wouter's avatar
Peters, Wouter committed
266
        :rtype: None
267

Peters, Wouter's avatar
Peters, Wouter committed
268
269
270
271
        Propagate the parameter values in the StateVector to the next cycle. This means a shift by one cycle 
        step for all states that will
        be optimized once more, and the creation of a new ensemble for the time step that just 
        comes in for the first time (step=nlag). 
272
273
274
        In the future, this routine can incorporate a formal propagation of the statevector.

        """
275
        
276
277
278
        # Remove State Vector n=1 by simply "popping" it from the list and appending a new empty list at the front. This empty list will
        # hold the new ensemble for the new cycle 

279
280
        self.ensemble_members.pop(0)
        self.ensemble_members.append([])
281

282
        # And now create a new time step of mean + members for n=nlag
283
284
        date = dacycle['time.start'] + timedelta(days=(self.nlag - 0.5) * int(dacycle['time.cycle']))
        cov = self.get_covariance(date, dacycle)
285
        self.make_new_ensemble(self.nlag - 1, cov)
286

karolina's avatar
karolina committed
287
        logging.info('The state vector has been propagated by one cycle')
288
289


290
    def write_to_file(self, filename, qual):
Peters, Wouter's avatar
Peters, Wouter committed
291
292
293
294
295
296
297
298
299
300
301
        """
        :param filename: the full filename for the output NetCDF file
        :rtype: None

        Write the StateVector information to a NetCDF file for later use. 
        In principle the output file will have only one two datasets inside 
        called:
            * `meanstate`, dimensions [nlag, nparamaters]
            * `ensemblestate`, dimensions [nlag,nmembers, nparameters]

        This NetCDF information can be read back into a StateVector object using 
302
        :meth:`~da.baseclasses.statevector.StateVector.read_from_file`
Peters, Wouter's avatar
Peters, Wouter committed
303
304

        """
305
        #import da.tools.io4 as io
306
        #import da.tools.io as io
307

308
        if qual == 'prior':
karolina's avatar
karolina committed
309
310
            f = io.CT_CDF(filename, method='create')
            logging.debug('Creating new StateVector output file (%s)' % filename)
311
            #qual = 'prior'
312
        else:
karolina's avatar
karolina committed
313
314
            f = io.CT_CDF(filename, method='write')
            logging.debug('Opening existing StateVector output file (%s)' % filename)
315
            #qual = 'opt'
316

317
318
319
        dimparams = f.add_params_dim(self.nparams)
        dimmembers = f.add_members_dim(self.nmembers)
        dimlag = f.add_lag_dim(self.nlag, unlimited=True)
320
321

        for n in range(self.nlag):
322
323
            members = self.ensemble_members[n]
            mean_state = members[0].param_values
karolina's avatar
karolina committed
324

325
            savedict = f.standard_var(varname='meanstate_%s' % qual)
karolina's avatar
karolina committed
326
            savedict['dims'] = dimlag + dimparams 
327
            savedict['values'] = mean_state
karolina's avatar
karolina committed
328
329
            savedict['count'] = n
            savedict['comment'] = 'this represents the mean of the ensemble'
330
            f.add_data(savedict)
karolina's avatar
karolina committed
331

332
333
334
            members = self.ensemble_members[n]
            devs = np.asarray([m.param_values.flatten() for m in members])
            data = devs - np.asarray(mean_state)
karolina's avatar
karolina committed
335

336
            savedict = f.standard_var(varname='ensemblestate_%s' % qual)
karolina's avatar
karolina committed
337
338
339
340
            savedict['dims'] = dimlag + dimmembers + dimparams 
            savedict['values'] = data
            savedict['count'] = n
            savedict['comment'] = 'this represents deviations from the mean of the ensemble'
341
            f.add_data(savedict)
karolina's avatar
karolina committed
342
343
344
345
        f.close()

        logging.info('Successfully wrote the State Vector to file (%s) ' % filename)

346
    def read_from_file(self, filename, qual='opt'):
Peters, Wouter's avatar
Peters, Wouter committed
347
348
        """ 
        :param filename: the full filename for the input NetCDF file
349
        :param qual: a string indicating whether to read the 'prior' or 'opt'(imized) StateVector from file
Peters, Wouter's avatar
Peters, Wouter committed
350
351
352
        :rtype: None

        Read the StateVector information from a NetCDF file and put in a StateVector object
353
        In principle the input file will have only one four datasets inside 
Peters, Wouter's avatar
Peters, Wouter committed
354
        called:
355
356
357
358
            * `meanstate_prior`, dimensions [nlag, nparamaters]
            * `ensemblestate_prior`, dimensions [nlag,nmembers, nparameters]
            * `meanstate_opt`, dimensions [nlag, nparamaters]
            * `ensemblestate_opt`, dimensions [nlag,nmembers, nparameters]
Peters, Wouter's avatar
Peters, Wouter committed
359
360

        This NetCDF information can be written to file using 
361
        :meth:`~da.baseclasses.statevector.StateVector.write_to_file`
Peters, Wouter's avatar
Peters, Wouter committed
362
363
364

        """

365
        #import da.tools.io as io
Peters, Wouter's avatar
Peters, Wouter committed
366
        f = io.ct_read(filename, 'read')
karolina's avatar
karolina committed
367
        meanstate = f.get_variable('statevectormean_' + qual)
368
        ensmembers = f.get_variable('statevectorensemble_' + qual)
karolina's avatar
karolina committed
369
        f.close()
370
371

        for n in range(self.nlag):
372
373
            if not self.ensemble_members[n] == []:
                self.ensemble_members[n] = []
karolina's avatar
karolina committed
374
                logging.warning('Existing ensemble for lag=%d was removed to make place for newly read data' % (n + 1))
375

376
            for m in range(self.nmembers):
karolina's avatar
karolina committed
377
                newmember = EnsembleMember(m)
378
379
                newmember.param_values = ensmembers[n, m, :].flatten() + meanstate[n]  # add the mean to the deviations to hold the full parameter values
                self.ensemble_members[n].append(newmember)
380

karolina's avatar
karolina committed
381
        logging.info('Successfully read the State Vector from file (%s) ' % filename)
382

383
    def write_members_to_file(self, lag, outdir,endswith='.nc'):
384
        """ 
385
           :param: lag: Which lag step of the filter to write, must lie in range [1,...,nlag]
386
387
           :param: outdir: Directory where to write files
           :param: endswith: Optional label to add to the filename, default is simply .nc
388
389
390
391
           :rtype: None

           Write ensemble member information to a NetCDF file for later use. The standard output filename is 
           *parameters.DDD.nc* where *DDD* is the number of the ensemble member. Standard output file location 
392
           is the `dir.input` of the dacycle object. In principle the output file will have only two datasets inside 
393
394
395
396
           called `parametervalues` which is of dimensions `nparameters` and `parametermap` which is of dimensions (180,360). 
           This dataset can be read and used by a :class:`~da.baseclasses.observationoperator.ObservationOperator` object. 

           .. note:: if more, or other information is needed to complete the sampling of the ObservationOperator you
397
                     can simply inherit from the StateVector baseclass and overwrite this write_members_to_file function.
398
399
400
401
402
403
404
405
406

        """

        # These import statements caused a crash in netCDF4 on MacOSX. No problems on Jet though. Solution was
        # to do the import already at the start of the module, not just in this method.
           
        #import da.tools.io as io
        #import da.tools.io4 as io

407
        members = self.ensemble_members[lag]
408

409
        for mem in members:
410
            filename = os.path.join(outdir, 'parameters.%03d%s' % (mem.membernumber, endswith))
karolina's avatar
karolina committed
411
            ncf = io.CT_CDF(filename, method='create')
412
413
            dimparams = ncf.add_params_dim(self.nparams)
            dimgrid = ncf.add_latlon_dim()
karolina's avatar
karolina committed
414

415
            data = mem.param_values
karolina's avatar
karolina committed
416
417
418
419
420
421
422
423

            savedict = io.std_savedict.copy()
            savedict['name'] = "parametervalues"
            savedict['long_name'] = "parameter_values_for_member_%d" % mem.membernumber
            savedict['units'] = "unitless"
            savedict['dims'] = dimparams 
            savedict['values'] = data
            savedict['comment'] = 'These are parameter values to use for member %d' % mem.membernumber
424
            ncf.add_data(savedict)
karolina's avatar
karolina committed
425

426
            griddata = self.vector2grid(vectordata=data)
karolina's avatar
karolina committed
427
428
429
430
431
432
433
434

            savedict = io.std_savedict.copy()
            savedict['name'] = "parametermap"
            savedict['long_name'] = "parametermap_for_member_%d" % mem.membernumber
            savedict['units'] = "unitless"
            savedict['dims'] = dimgrid 
            savedict['values'] = griddata.tolist()
            savedict['comment'] = 'These are gridded parameter values to use for member %d' % mem.membernumber
435
            ncf.add_data(savedict)
karolina's avatar
karolina committed
436
437
438
439
440

            ncf.close()

            logging.debug('Successfully wrote data from ensemble member %d to file (%s) ' % (mem.membernumber, filename))

441
    def grid2vector(self, griddata=None, method='avg'):
442
        """ 
443
            Map gridded data onto a vector of length (nparams,)
444

445
           :param griddata: a gridded dataset to use. This dataset is mapped onto a vector of length `nparams`
446
           :param method: a string that specifies the method to combine grid boxes in case reverse=True. Must be either ['avg','sum','minval']
447
           :rtype: ndarray: size (nparameters,)
448
449
450
451
452

           This method makes use of a dictionary that links every parameter number [1,...,nparams] to a series of gridindices. These 
           indices specify a location on a 360x180 array, stretched into a vector using `array.flat`. There are multiple ways of calling 
           this method::

453
               values       = self.grid2vector(griddata=mygriddeddata,method='minval') # 
454
455
                                                                    using the minimum value of all datapoints covered by that parameter index

456
               values       = self.grid2vector(griddata=mygriddeddata,method='avg') # 
457
458
                                                                    using the average value of all datapoints covered by that parameter index

459
               values       = self.grid2vector(griddata=mygriddeddata,method='sum') # 
460
461
                                                                    using the sum of values of all datapoints covered by that parameter index

462
           .. note:: This method uses a DaSystem object that must be initialized with a proper parameter map. See :class:`~da.baseclasses.dasystem` for details
463
464
465

        """

karolina's avatar
karolina committed
466
        methods = ['avg', 'sum', 'minval']
467
        if method not in methods:
karolina's avatar
karolina committed
468
            logging.error("To put data from a map into the statevector, please specify the method to use (%s)" % methods)
469
            raise ValueError
470

karolina's avatar
karolina committed
471
        result = np.zeros((self.nparams,), float)
brunner's avatar
brunner committed
472
        for k, v in self.griddict.items():
473
474
            #print k,k-1,result.shape, v
            if method == "avg": 
karolina's avatar
karolina committed
475
                result[k - 1] = griddata.take(v).mean()
476
            elif method == "sum" : 
karolina's avatar
karolina committed
477
                result[k - 1] = griddata.take(v).sum()
478
            elif method == "minval" : 
karolina's avatar
karolina committed
479
                result[k - 1] = griddata.take(v).min()
480
        return result # Note that the result is returned, but not yet placed in the member.param_values attrtibute!
481
482


483
    def vector2grid(self, vectordata=None):
484
485
        """ 
            Map vector elements to a map or vice cersa
486

487
488
489
490
491
492
493
           :param vectordata: a vector dataset to use in case `reverse = False`. This dataset is mapped onto a 1x1 grid and must be of length `nparams`
           :rtype: ndarray: an array of size (360,180,) 

           This method makes use of a dictionary that links every parameter number [1,...,nparams] to a series of gridindices. These 
           indices specify a location on a 360x180 array, stretched into a vector using `array.flat`. There are multiple ways of calling 
           this method::

494
               griddedarray = self.vector2grid(vectordata=param_values) # simply puts the param_values onto a (180,360,) array
495
496

           .. note:: This method uses a DaSystem object that must be initialzied with a proper parameter map. See :class:`~da.baseclasses.dasystem` for details
497

498
        """
karolina's avatar
karolina committed
499
        result = np.zeros(self.gridmap.shape, float)
brunner's avatar
brunner committed
500
        for k, v in self.griddict.items():
501
            #print k,v
karolina's avatar
karolina committed
502
            result.put(v, vectordata[k - 1])
503
        return result         
504

505
    def vector2tc(self, vectordata, cov=False):
506
507
508
509
510
511
512
513
514
515
        """ 
            project Vector onto TransCom regions 

           :param vectordata: a vector dataset to use, must be of length `nparams`
           :param cov: a Boolean to specify whether the input dataset is a vector (mean), or a matrix (covariance)
           :rtype: ndarray: an array of size (23,) (cov:F) or of size (23,23,) (cov:T)
        """

        M = self.tcmatrix
        if cov:
karolina's avatar
karolina committed
516
            return np.dot(np.transpose(M), np.dot(vectordata, M))
517
        else:
karolina's avatar
karolina committed
518
            return np.dot(vectordata.squeeze(), M)
519

520
    def state_to_grid(self, fluxvector=None, lag=1):
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
        """ 
            Transforms the StateVector information (mean + covariance) to a 1x1 degree grid.

            :param: fluxvector: a vector of length (nparams,) that holds the fluxes associated with each parameter in the StateVector
            :param: lag: the lag at which to evaluate the StateVector
            :rtype: a tuple of two arrays (gridmean,gridvariance) with dimensions (180,360,)

            If the attribute `fluxvector` is not passed, the function will return the mean parameter value and its variance on a 1x1 map.
            
            ..note:: Although we can return the variance information for each gridbox, the covariance information contained in the original ensemble is lost when mapping to 1x1 degree!

        """

        if fluxvector == None:
            fluxvector = np.ones(self.nparams)

537
538
        ensemble = self.ensemble_members[lag - 1]
        ensemblemean = ensemble[0].param_values
539
540

        # First transform the mean
541
        gridmean = self.vector2grid(vectordata=ensemblemean * fluxvector)
542
543

        # And now the covariance, first create covariance matrix (!), and then multiply
544
        deviations = np.array([mem.param_values * fluxvector - ensemblemean for mem in ensemble])
karolina's avatar
karolina committed
545
        ensemble = []
546
        for mem in deviations:
547
            ensemble.append(self.vector2grid(mem))
548

karolina's avatar
karolina committed
549
        return (gridmean, np.array(ensemble))
550

551
    def state2tc(self, fluxvector=None, lag=1):
552
553
554
555
556
557
558
559
        """ 
            Transforms the StateVector information (mean + covariance) to the TransCom regions.

            :param: fluxvector: a vector of length (nparams,) that holds the fluxes associated with each parameter in the StateVector
            :param: lag: the lag at which to evaluate the StateVector
            :rtype: a tuple of two arrays (mean,covariance) with dimensions ((23,), (23,23,) )

        """
560
561
        ensemble = self.ensemble_members[lag - 1]
        ensemblemean = ensemble[0].param_values
562
563
564

        # First transform the mean

565
        mean = self.vector2tc(vectordata=ensemble[0].param_values * fluxvector)
566
567
568

        # And now the covariance, first create covariance matrix (!), and then multiply

569
        deviations = np.array([mem.param_values * fluxvector - ensemblemean for mem in ensemble])
karolina's avatar
karolina committed
570
        covariance = np.dot(np.transpose(deviations), deviations) / (self.nmembers - 1)
571
        cov = self.vector2tc(covariance, cov=True)
572

karolina's avatar
karolina committed
573
        return (mean, cov)
574

575
576
577
    def get_covariance(self, date, cycleparams):
        pass
    
578
579
580
################### End Class StateVector ###################

if __name__ == "__main__":
581
    pass
582