optimizer.py 18.4 KB
Newer Older
1
2
3
4
#!/usr/bin/env python
# optimizer.py

"""
Peters, Wouter's avatar
Peters, Wouter committed
5
6
.. module:: optimizer
.. moduleauthor:: Wouter Peters 
7
8
9
10
11
12
13
14
15
16
17
18

Revision History:
File created on 28 Jul 2010.

"""

import os
import sys
import logging
import datetime

identifier = 'Optimizer baseclass'
karolina's avatar
karolina committed
19
version = '0.0'
20
21
22
23
24
25
26
27
28
29
30
31

################### 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):
32
33
        self.Identifier = identifier
        self.Version = version
34

karolina's avatar
karolina committed
35
        logging.info('Optimizer object initialized: %s' % self.Identifier)
36
37
38


    def Initialize(self, dims):
karolina's avatar
karolina committed
39
40
41
42
        self.nlag = dims[0]
        self.nmembers = dims[1]
        self.nparams = dims[2]
        self.nobs = dims[3]
43
        self.create_matrices()
44

45
    def create_matrices(self):
46
47
48
49
        """ Create Matrix space needed in optimization routine """
        import numpy as np

        # mean state  [X]
karolina's avatar
karolina committed
50
        self.x = np.zeros((self.nlag * self.nparams,), float)
51
        # deviations from mean state  [X']
karolina's avatar
karolina committed
52
        self.X_prime = np.zeros((self.nlag * self.nparams, self.nmembers,), float)
53
        # mean state, transported to observation space [ H(X) ]
karolina's avatar
karolina committed
54
        self.Hx = np.zeros((self.nobs,), float)
55
        # deviations from mean state, transported to observation space [ H(X') ]
karolina's avatar
karolina committed
56
        self.HX_prime = np.zeros((self.nobs, self.nmembers), float)
57
        # observations
karolina's avatar
karolina committed
58
        self.obs = np.zeros((self.nobs,), float)
59
        # observation ids
karolina's avatar
karolina committed
60
        self.obs_ids = np.zeros((self.nobs,), float)
61
        # covariance of observations
karolina's avatar
karolina committed
62
        self.R = np.zeros((self.nobs, self.nobs,), float)
63
        # localization of obs
karolina's avatar
karolina committed
64
        self.may_localize = np.zeros(self.nobs, bool)
65
        # rejection of obs
karolina's avatar
karolina committed
66
        self.may_reject = np.zeros(self.nobs, bool)
67
        # flags of obs
karolina's avatar
karolina committed
68
        self.flags = np.zeros(self.nobs, int)
69
        # species type
karolina's avatar
karolina committed
70
        self.species = np.zeros(self.nobs, str)
71
        # species type
karolina's avatar
karolina committed
72
        self.sitecode = np.zeros(self.nobs, str)
73
74

        # species mask
karolina's avatar
karolina committed
75
        self.speciesmask = {}
76

77
        # Total covariance of fluxes and obs in units of obs [H P H^t + R]
karolina's avatar
karolina committed
78
        self.HPHR = np.zeros((self.nobs, self.nobs,), float)
79
        # Kalman Gain matrix
karolina's avatar
karolina committed
80
81
        self.KG = np.zeros((self.nlag * self.nparams, self.nobs,), float)

82
    def state_to_matrix(self, StateVector):
83
84
        import numpy as np

karolina's avatar
karolina committed
85
86
87
88
89
90
91
92
93
        allsites = []      # collect all obs for n=1,..,nlag
        allobs = []      # collect all obs for n=1,..,nlag
        allmdm = []      # collect all mdm 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
94

95
        for n in range(self.nlag):
karolina's avatar
karolina committed
96
97
98
99
            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]))
100

101
            if Samples != None:
karolina's avatar
karolina committed
102
                self.rejection_threshold = Samples.rejection_threshold
Peters, Wouter's avatar
Peters, Wouter committed
103

104
105
106
107
108
109
110
111
                allreject.extend(Samples.getvalues('may_reject'))
                alllocalize.extend(Samples.getvalues('may_localize'))
                allflags.extend(Samples.getvalues('flag'))
                allspecies.extend(Samples.getvalues('species'))
                allobs.extend(Samples.getvalues('obs'))
                allsites.extend(Samples.getvalues('code'))
                allmdm.extend(Samples.getvalues('mdm'))
                allids.extend(Samples.getvalues('id'))
112

113
                simulatedensemble = Samples.getvalues('simulated')
Peters, Wouter's avatar
Peters, Wouter committed
114

115
116
                if allsimulated == None :
                    allsimulated = np.array(simulatedensemble)
Peters, Wouter's avatar
Peters, Wouter committed
117
                else:
karolina's avatar
karolina committed
118
                    allsimulated = np.concatenate((allsimulated, np.array(simulatedensemble)), axis=0)
119

karolina's avatar
karolina committed
120
121
122
123
        self.obs[:] = np.array(allobs)
        self.obs_ids[:] = np.array(allids)
        self.HX_prime[:, :] = np.array(allsimulated)
        self.Hx[:] = self.HX_prime[:, 0]
124

karolina's avatar
karolina committed
125
126
127
128
129
        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
130

karolina's avatar
karolina committed
131
132
        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
133

karolina's avatar
karolina committed
134
135
        for i, mdm in enumerate(allmdm):
            self.R[i, i] = mdm ** 2
136

137
    def matrix_to_state(self, StateVector):
138
        for n in range(self.nlag):
karolina's avatar
karolina committed
139
140
141
            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]     
142

143
        StateVector.isOptimized = True
karolina's avatar
karolina committed
144
        logging.debug('Returning optimized data to the StateVector, setting "StateVector.isOptimized = True" ')
145

146
    def write_diagnostics(self, diagdir, timestart, StateVector, type='prior'):
147
148
149
150
151
152
153
154
155
156
        """
            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
157
158

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

163
164
165
        
        filename = os.path.join(diagdir, 'optimizer.%s.nc' % timestart.strftime('%Y%m%d'))
        #DaCycle.OutputFileList.append(filename)
166

167
168
169
        # Open or create file

        if type == 'prior':
karolina's avatar
karolina committed
170
171
            f = io.CT_CDF(filename, method='create')
            logging.debug('Creating new diagnostics file for optimizer (%s)' % filename)
172
        elif type == 'optimized':
karolina's avatar
karolina committed
173
174
            f = io.CT_CDF(filename, method='write')
            logging.debug('Opening existing diagnostics file for optimizer (%s)' % filename)
175
176

        # Add dimensions 
177

karolina's avatar
karolina committed
178
179
180
181
182
183
        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)
184

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

karolina's avatar
karolina committed
187
188
189
190
191
        savedict = io.std_savedict.copy() 
        savedict['name'] = "statevectormean_%s" % type
        savedict['long_name'] = "full_statevector_mean_%s" % type
        savedict['units'] = "unitless"
        savedict['dims'] = dimstate
karolina's avatar
karolina committed
192
        savedict['values'] = self.x.tolist()
karolina's avatar
karolina committed
193
194
195
196
197
198
199
200
        savedict['comment'] = 'Full %s state vector mean ' % type
        f.AddData(savedict)

        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
karolina's avatar
karolina committed
201
        savedict['values'] = self.X_prime.tolist()
karolina's avatar
karolina committed
202
203
204
205
206
207
208
209
        savedict['comment'] = 'Full state vector %s deviations as resulting from the optimizer' % type
        f.AddData(savedict)

        savedict = io.std_savedict.copy()
        savedict['name'] = "modelsamplesmean_%s" % type
        savedict['long_name'] = "modelsamplesforecastmean_%s" % type
        savedict['units'] = "mol mol-1"
        savedict['dims'] = dimobs
karolina's avatar
karolina committed
210
        savedict['values'] = self.Hx.tolist()
211
        savedict['comment'] = '%s mean mixing ratios based on %s state vector' % (type, type)
karolina's avatar
karolina committed
212
213
214
215
216
217
218
        f.AddData(savedict)

        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
karolina's avatar
karolina committed
219
        savedict['values'] = self.HX_prime.tolist()
220
        savedict['comment'] = '%s mixing ratio deviations based on %s state vector' % (type, type)
karolina's avatar
karolina committed
221
        f.AddData(savedict)
222

223
224
225
        # Continue with prior only data

        if type == 'prior':
226

karolina's avatar
karolina committed
227
228
229
230
231
            savedict = io.std_savedict.copy()
            savedict['name'] = "sitecode"
            savedict['long_name'] = "site code propagated from observation file"
            savedict['dtype'] = "char"
            savedict['dims'] = dimobs + dim200char
karolina's avatar
karolina committed
232
            savedict['values'] = self.sitecode
233
            savedict['missing_value'] = '!'
karolina's avatar
karolina committed
234
235
236
237
238
239
240
            f.AddData(savedict)

            savedict = io.std_savedict.copy()
            savedict['name'] = "observed"
            savedict['long_name'] = "observedvalues"
            savedict['units'] = "mol mol-1"
            savedict['dims'] = dimobs
karolina's avatar
karolina committed
241
            savedict['values'] = self.obs.tolist()
karolina's avatar
karolina committed
242
243
244
245
246
247
248
249
250
            savedict['comment'] = 'Observations used in optimization'
            f.AddData(savedict)

            savedict = io.std_savedict.copy()
            savedict['name'] = "obspack_num"
            savedict['dtype'] = "int64"
            savedict['long_name'] = "Unique_ObsPack_observation_number"
            savedict['units'] = ""
            savedict['dims'] = dimobs
karolina's avatar
karolina committed
251
            savedict['values'] = self.obs_ids.tolist()
karolina's avatar
karolina committed
252
253
254
255
256
257
258
259
            savedict['comment'] = 'Unique observation number across the entire ObsPack distribution'
            f.AddData(savedict)

            savedict = io.std_savedict.copy()
            savedict['name'] = "modeldatamismatchvariance"
            savedict['long_name'] = "modeldatamismatch variance"
            savedict['units'] = "[mol mol-1]^2"
            savedict['dims'] = dimobs + dimobs
karolina's avatar
karolina committed
260
            savedict['values'] = self.R.tolist()
karolina's avatar
karolina committed
261
262
            savedict['comment'] = 'Variance of mole fractions resulting from model-data mismatch'
            f.AddData(savedict)
263
264
265
266

        # Continue with posterior only data

        elif type == 'optimized':
karolina's avatar
karolina committed
267
            
karolina's avatar
karolina committed
268
269
270
271
272
            savedict = io.std_savedict.copy()
            savedict['name'] = "totalmolefractionvariance"
            savedict['long_name'] = "totalmolefractionvariance"
            savedict['units'] = "[mol mol-1]^2"
            savedict['dims'] = dimobs + dimobs
karolina's avatar
karolina committed
273
            savedict['values'] = self.HPHR.tolist()
karolina's avatar
karolina committed
274
275
276
277
278
279
280
281
            savedict['comment'] = 'Variance of mole fractions resulting from prior state and model-data mismatch'
            f.AddData(savedict)

            savedict = io.std_savedict.copy()
            savedict['name'] = "flag"
            savedict['long_name'] = "flag_for_obs_model"
            savedict['units'] = "None"
            savedict['dims'] = dimobs
karolina's avatar
karolina committed
282
            savedict['values'] = self.flags.tolist()
karolina's avatar
karolina committed
283
284
285
286
287
288
289
290
            savedict['comment'] = 'Flag (0/1/2/99) for observation value, 0 means okay, 1 means QC error, 2 means rejected, 99 means not sampled'
            f.AddData(savedict)

            savedict = io.std_savedict.copy()
            savedict['name'] = "kalmangainmatrix"
            savedict['long_name'] = "kalmangainmatrix"
            savedict['units'] = "unitless molefraction-1"
            savedict['dims'] = dimstate + dimobs
karolina's avatar
karolina committed
291
            savedict['values'] = self.KG.tolist()
karolina's avatar
karolina committed
292
            savedict['comment'] = 'Kalman gain matrix of all obs and state vector elements'
293
            #dummy                   = f.AddData(savedict)
294

karolina's avatar
karolina committed
295
296
        f.close()
        logging.debug('Diagnostics file closed')
297
        return filename
298

299
    def serial_minimum_least_squares(self):
300
301
        """ Make minimum least squares solution by looping over obs"""
        import numpy as np
302

303
304
        for n in range(self.nobs):

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

            if self.flags[n] != 0:
karolina's avatar
karolina committed
308
                logging.debug('Skipping observation (%s,%s) because of flag value %d' % (self.sitecode[n], self.obs_ids[n], self.flags[n]))
309
310
                continue

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

karolina's avatar
karolina committed
313
            res = self.obs[n] - self.Hx[n]
314
315

            if self.may_reject[n]:
karolina's avatar
karolina committed
316
                threshold = self.rejection_threshold * np.sqrt(self.R[n, n])
317
                if np.abs(res) > threshold:
karolina's avatar
karolina committed
318
                    logging.debug('Rejecting observation (%s,%s) because residual (%f) exceeds threshold (%f)' % (self.sitecode[n], self.obs_ids[n], res, threshold))
319
                    self.flags[n] = 2
320
                    continue
321

karolina's avatar
karolina committed
322
323
            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]
324

karolina's avatar
karolina committed
325
            self.KG[:, n] = PHt / self.HPHR[n, n]
326

327
            if self.may_localize[n]:
328
                self.localize(n)
karolina's avatar
karolina committed
329
                logging.debug('Localized observation %s' % self.obs_ids[n])
330
            else:
331
                logging.debug('Not allowed to localize observation %s' % self.obs_ids[n])
332

karolina's avatar
karolina committed
333
            alpha = np.double(1.0) / (np.double(1.0) + np.sqrt((self.R[n, n]) / self.HPHR[n, n]))
334

karolina's avatar
karolina committed
335
            self.x[:] = self.x + self.KG[:, n] * res
336
337

            for r in range(self.nmembers):
karolina's avatar
karolina committed
338
                self.X_prime[:, r] = self.X_prime[:, r] - alpha * self.KG[:, n] * (self.HX_prime[n, r])
339
340
341
342

#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 !!!!

karolina's avatar
karolina committed
343
344
345
346
347
            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, :]
348

karolina's avatar
karolina committed
349
350
351
352
353
            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, :]
354
355

            
356
    def bulk_minimum_least_squares(self):
357
358
359
360
361
362
        """ 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

karolina's avatar
karolina committed
363
364
365
366
        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)
367
368

        for n in range(self.nobs):
369
            self.localize(n)
370

karolina's avatar
karolina committed
371
        self.x[:] = self.x + np.dot(self.KG, self.obs - self.Hx)                             # xa = xp + K (y-Hx)
372
373
374
375
376

        # 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.

karolina's avatar
karolina committed
377
378
379
380
381
382
        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'
383
384
385
386
387


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

karolina's avatar
karolina committed
388
389
390
391
        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'
392

karolina's avatar
karolina committed
393
        logging.info('Minimum Least Squares solution was calculated, returning')
394

395
    def set_localization(self):
396
397
398
        """ determine which localization to use """
        self.localization = True
        self.localizetype = "None"
karolina's avatar
karolina committed
399
        logging.info("Current localization option is set to %s" % self.localizetype)
400

401
    def localize(self, n):
402
        """ localize the Kalman Gain matrix """
karolina's avatar
karolina committed
403
        pass
404
405
406
407
408
409
410
411
412

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



if __name__ == "__main__":

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

karolina's avatar
karolina committed
413
    from da.tools.initexit import StartLogger 
414
415
416
417
418
    from da.tools.initexit import CycleControl 
    from da.ct.obs import CtObservations
    import numpy as np

    opts = ['-v']
karolina's avatar
karolina committed
419
    args = {'rc':'../../da.rc', 'logfile':'da_initexit.log', 'jobrcfilename':'test.rc'}
420
421

    StartLogger()
karolina's avatar
karolina committed
422
    DaCycle = CycleControl(opts, args)
423
424
425
426
427

    DaCycle.Initialize()
    print DaCycle


karolina's avatar
karolina committed
428
    samples = CtObservations(DaCycle.DaSystem, datetime.datetime(2005, 3, 5))
429
430
    samples.add_observations()
    samples.add_simulations('/Users/peters/tmp/test_da/output/20050305/samples.000.nc')
431
432


433
    nobs = len(samples.datalist)
karolina's avatar
karolina committed
434
    dims = (int(DaCycle['time.nlag']),
435
                  int(DaCycle['da.optimizer.nmembers']),
436
                  int(DaCycle.DaSystem['nparameters']),
karolina's avatar
karolina committed
437
                  nobs,)
438

karolina's avatar
karolina committed
439
440
441
442
#   opt = CtOptimizer(dims)
#   opt.state_to_matrix(StateVector)
#   opt.MinimumLeastSquares()
#   opt.matrix_to_state(StateVector)