optimizer.py 18 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

Revision History:
File created on 28 Jul 2010.

"""

import os
import sys
import logging
import datetime
karolina's avatar
karolina committed
17
18
19
20
import numpy as np
import numpy.linalg as la

import da.tools.io4 as io
21
22

identifier = 'Optimizer baseclass'
karolina's avatar
karolina committed
23
version = '0.0'
24
25
26
27
28
29
30
31
32
33
34
35

################### 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):
36
37
        self.Identifier = identifier
        self.Version = version
38

karolina's avatar
karolina committed
39
        logging.info('Optimizer object initialized: %s' % self.Identifier)
40

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

48
    def create_matrices(self):
49
50
51
        """ Create Matrix space needed in optimization routine """

        # mean state  [X]
karolina's avatar
karolina committed
52
        self.x = np.zeros((self.nlag * self.nparams,), float)
53
        # deviations from mean state  [X']
karolina's avatar
karolina committed
54
        self.X_prime = np.zeros((self.nlag * self.nparams, self.nmembers,), float)
55
        # mean state, transported to observation space [ H(X) ]
karolina's avatar
karolina committed
56
        self.Hx = np.zeros((self.nobs,), float)
57
        # deviations from mean state, transported to observation space [ H(X') ]
karolina's avatar
karolina committed
58
        self.HX_prime = np.zeros((self.nobs, self.nmembers), float)
59
        # observations
karolina's avatar
karolina committed
60
        self.obs = np.zeros((self.nobs,), float)
61
        # observation ids
karolina's avatar
karolina committed
62
        self.obs_ids = np.zeros((self.nobs,), float)
63
        # covariance of observations
64
65
66
67
68
        # Total covariance of fluxes and obs in units of obs [H P H^t + R]
        if self.algorithm == 'Serial':
            self.R = np.zeros((self.nobs,), float)
            self.HPHR = np.zeros((self.nobs,), float)
        else:
karolina's avatar
karolina committed
69
            self.R = np.zeros((self.nobs, self.nobs,), float)
70
            self.HPHR = np.zeros((self.nobs, self.nobs,), float)
71
        # localization of obs
karolina's avatar
karolina committed
72
        self.may_localize = np.zeros(self.nobs, bool)
73
        # rejection of obs
karolina's avatar
karolina committed
74
        self.may_reject = np.zeros(self.nobs, bool)
75
        # flags of obs
karolina's avatar
karolina committed
76
        self.flags = np.zeros(self.nobs, int)
77
        # species type
karolina's avatar
karolina committed
78
        self.species = np.zeros(self.nobs, str)
79
        # species type
karolina's avatar
karolina committed
80
        self.sitecode = np.zeros(self.nobs, str)
81
82

        # species mask
karolina's avatar
karolina committed
83
        self.speciesmask = {}
84

85
        # Kalman Gain matrix
karolina's avatar
karolina committed
86
87
        self.KG = np.zeros((self.nlag * self.nparams, self.nobs,), float)

88
    def state_to_matrix(self, StateVector):
karolina's avatar
karolina committed
89
90
91
92
93
94
95
96
97
        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
98

99
        for n in range(self.nlag):
karolina's avatar
karolina committed
100
101
102
103
            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]))
104

105
            if Samples != None:        
karolina's avatar
karolina committed
106
                self.rejection_threshold = Samples.rejection_threshold
Peters, Wouter's avatar
Peters, Wouter committed
107

108
109
110
111
112
113
114
115
                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'))
116

117
                simulatedensemble = Samples.getvalues('simulated')
Peters, Wouter's avatar
Peters, Wouter committed
118

119
                if allsimulated == None :     
120
                    allsimulated = np.array(simulatedensemble)
Peters, Wouter's avatar
Peters, Wouter committed
121
                else:
karolina's avatar
karolina committed
122
                    allsimulated = np.concatenate((allsimulated, np.array(simulatedensemble)), axis=0)
123

karolina's avatar
karolina committed
124
125
126
127
        self.obs[:] = np.array(allobs)
        self.obs_ids[:] = np.array(allids)
        self.HX_prime[:, :] = np.array(allsimulated)
        self.Hx[:] = self.HX_prime[:, 0]
128

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

karolina's avatar
karolina committed
135
136
        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
137

138
139
140
141
142
        if self.algorithm == 'Serial':
            for i, mdm in enumerate(allmdm):
                self.R[i] = mdm ** 2
        else:
            for i, mdm in enumerate(allmdm):
karolina's avatar
karolina committed
143
                self.R[i, i] = mdm ** 2
144
                    
145
    def matrix_to_state(self, StateVector):
146
        for n in range(self.nlag):
karolina's avatar
karolina committed
147
148
149
            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]     
150

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

karolina's avatar
karolina committed
153
    def write_diagnostics(self, filename, type='prior'):
154
155
156
157
158
159
160
161
162
163
        """
            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
164
165

            The type designation refers to the writing of prior or posterior data and is used in naming the variables"
166
167
        """

168
169
170
        # Open or create file

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

        # Add dimensions 
178

179
180
181
182
183
184
        dimparams = f.add_params_dim(self.nparams)
        dimmembers = f.add_members_dim(self.nmembers)
        dimlag = f.add_lag_dim(self.nlag, unlimited=False)
        dimobs = f.add_obs_dim(self.nobs)
        dimstate = f.add_dim('nstate', self.nparams * self.nlag)
        dim200char = f.add_dim('string_of200chars', 200)
185

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

karolina's avatar
karolina committed
188
189
190
191
192
        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
193
        savedict['values'] = self.x.tolist()
karolina's avatar
karolina committed
194
        savedict['comment'] = 'Full %s state vector mean ' % type
195
        f.add_data(savedict)
karolina's avatar
karolina committed
196
197
198
199
200
201

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

        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
211
        savedict['values'] = self.Hx.tolist()
212
        savedict['comment'] = '%s mean mixing ratios based on %s state vector' % (type, type)
213
        f.add_data(savedict)
karolina's avatar
karolina committed
214
215
216
217
218
219

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

224
225
226
        # Continue with prior only data

        if type == 'prior':
227

karolina's avatar
karolina committed
228
229
230
231
232
            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
233
            savedict['values'] = self.sitecode
234
            savedict['missing_value'] = '!'
235
            f.add_data(savedict)
karolina's avatar
karolina committed
236
237
238
239
240
241

            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
242
            savedict['values'] = self.obs.tolist()
karolina's avatar
karolina committed
243
            savedict['comment'] = 'Observations used in optimization'
244
            f.add_data(savedict)
karolina's avatar
karolina committed
245
246
247
248
249
250
251

            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
252
            savedict['values'] = self.obs_ids.tolist()
karolina's avatar
karolina committed
253
            savedict['comment'] = 'Unique observation number across the entire ObsPack distribution'
254
            f.add_data(savedict)
karolina's avatar
karolina committed
255
256
257
258
259

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

        # Continue with posterior only data

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

            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
287
            savedict['values'] = self.flags.tolist()
karolina's avatar
karolina committed
288
            savedict['comment'] = 'Flag (0/1/2/99) for observation value, 0 means okay, 1 means QC error, 2 means rejected, 99 means not sampled'
289
            f.add_data(savedict)
karolina's avatar
karolina committed
290
291
292
293
294
295

            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
296
            savedict['values'] = self.KG.tolist()
karolina's avatar
karolina committed
297
            savedict['comment'] = 'Kalman gain matrix of all obs and state vector elements'
298
            #dummy                   = f.add_data(savedict)
299

karolina's avatar
karolina committed
300
301
        f.close()
        logging.debug('Diagnostics file closed')
302

303

304
    def serial_minimum_least_squares(self):
305
306
307
        """ Make minimum least squares solution by looping over obs"""
        for n in range(self.nobs):

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

            if self.flags[n] != 0:
311
                logging.debug('Skipping observation (%s,%i) because of flag value %d' % (self.sitecode[n], self.obs_ids[n], self.flags[n]))
312
313
                continue

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

karolina's avatar
karolina committed
316
            res = self.obs[n] - self.Hx[n]
317
318

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

karolina's avatar
karolina committed
325
            PHt = 1. / (self.nmembers - 1) * np.dot(self.X_prime, self.HX_prime[n, :])
326
            self.HPHR[n] = 1. / (self.nmembers - 1) * (self.HX_prime[n, :] * self.HX_prime[n, :]).sum() + self.R[n]
327

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

330
            if self.may_localize[n]:
karolina's avatar
karolina committed
331
                logging.debug('Trying to localize observation %s, %i' % (self.sitecode[n], self.obs_ids[n]))
332
                self.localize(n)
333
            else:
karolina's avatar
karolina committed
334
                logging.debug('Not allowed to localize observation %s, %i' % (self.sitecode[n], self.obs_ids[n]))
335

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

karolina's avatar
karolina committed
338
            self.x[:] = self.x + self.KG[:, n] * res
339
340

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

#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
346
347
            for m in range(n + 1, self.nobs):
                res = self.obs[n] - self.Hx[n]
348
                fac = 1.0 / (self.nmembers - 1) * (self.HX_prime[n, :] * self.HX_prime[m, :]).sum() / self.HPHR[n]
karolina's avatar
karolina committed
349
350
                self.Hx[m] = self.Hx[m] + fac * res
                self.HX_prime[m, :] = self.HX_prime[m, :] - alpha * fac * self.HX_prime[n, :]
351

karolina's avatar
karolina committed
352
353
            for m in range(1, n + 1):
                res = self.obs[n] - self.Hx[n]
354
                fac = 1.0 / (self.nmembers - 1) * (self.HX_prime[n, :] * self.HX_prime[m, :]).sum() / self.HPHR[n]
karolina's avatar
karolina committed
355
356
                self.Hx[m] = self.Hx[m] + fac * res
                self.HX_prime[m, :] = self.HX_prime[m, :] - alpha * fac * self.HX_prime[n, :]
357
358

            
359
    def bulk_minimum_least_squares(self):
360
        """ Make minimum least squares solution by solving matrix equations"""
karolina's avatar
karolina committed
361
        
362
363
364

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

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

        for n in range(self.nobs):
371
            self.localize(n)
372

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

        # 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
379
380
381
382
383
384
        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'
385
386
387
388
389


        # 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
390
391
392
393
        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'
394

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

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

403
    def localize(self, n):
404
        """ localize the Kalman Gain matrix """
karolina's avatar
karolina committed
405
        logging.debug('Not localized observation %s, %i' % (self.sitecode[n], self.obs_ids[n]))
406
407
408
409
    
    def set_algorithm(self):
        self.algorithm = 'Serial'
        logging.info("Current algorithm is set to %s in baseclass optimizer" % self.algorithm)
410
411
412
413
414
415

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



if __name__ == "__main__":
416
    pass