popt.loop.ensemble

  1# External imports
  2import numpy as np
  3import sys
  4import warnings
  5
  6from copy import deepcopy
  7
  8
  9# Internal imports
 10from popt.misc_tools import optim_tools as ot
 11from pipt.misc_tools import analysis_tools as at
 12from ensemble.ensemble import Ensemble as PETEnsemble
 13from popt.loop.dist import GenOptDistribution
 14
 15
 16class Ensemble(PETEnsemble):
 17    """
 18    Class to store control states and evaluate objective functions.
 19
 20    Methods
 21    -------
 22    get_state()
 23        Returns control vector as ndarray
 24
 25    get_final_state(return_dict)
 26        Returns final control vector between [lb,ub]
 27
 28    get_cov()
 29        Returns the ensemble covariance matrix
 30
 31    function(x,*args)
 32        Objective function called during optimization
 33
 34    gradient(x,*args)
 35        Ensemble gradient
 36
 37    hessian(x,*args)
 38        Ensemble hessian
 39
 40    calc_ensemble_weights(self,x,*args):
 41        Calculate weights used in sequential monte carlo optimization
 42
 43    """
 44
 45    def __init__(self, keys_en, sim, obj_func):
 46        """
 47        Parameters
 48        ----------
 49        keys_en : dict
 50            Options for the ensemble class
 51
 52            - disable_tqdm: supress tqdm progress bar for clean output in the notebook
 53            - ne: number of perturbations used to compute the gradient
 54            - state: name of state variables passed to the .mako file
 55            - prior_<name>: the prior information the state variables, including mean, variance and variable limits
 56            - num_models: number of models (if robust optimization) (default 1)
 57            - transform: transform variables to [0,1] if true (default true)
 58
 59        sim : callable
 60            The forward simulator (e.g. flow)
 61
 62        obj_func : callable
 63            The objective function (e.g. npv)
 64        """
 65
 66        # Initialize PETEnsemble
 67        super(Ensemble, self).__init__(keys_en, sim)
 68        
 69        def __set__variable(var_name=None, defalut=None):
 70            if var_name in keys_en:
 71                return keys_en[var_name]
 72            else:
 73                return defalut
 74
 75        # Set number of models (default 1)
 76        self.num_models = __set__variable('num_models', 1)
 77
 78        # Set transform flag (defalult True)
 79        self.transform = __set__variable('transform', True)
 80
 81        # Number of samples to compute gradient
 82        self.num_samples = self.ne
 83
 84        # We need the limits to convert between [0, 1] and [lb, ub],
 85        # and we need the bounds as list of (min, max) pairs
 86        # Also set the state and covarianve equal to the values provided in the input.
 87        self.upper_bound = []
 88        self.lower_bound = []
 89        self.bounds = []
 90        self.cov = np.array([])
 91        for name in self.prior_info.keys():
 92            self.state[name] = np.asarray(self.prior_info[name]['mean'])
 93            num_state_var = len(self.state[name])
 94            value_cov = self.prior_info[name]['variance'] * np.ones((num_state_var,))
 95            if 'limits' in self.prior_info[name].keys():
 96                lb = self.prior_info[name]['limits'][0]
 97                ub = self.prior_info[name]['limits'][1]
 98                self.lower_bound.append(lb)
 99                self.upper_bound.append(ub)
100                if self.transform:
101                    value_cov = value_cov / (ub - lb)**2
102                    np.clip(value_cov, 0, 1, out=value_cov)
103                    self.bounds += num_state_var*[(0, 1)]
104                else:
105                    self.bounds += num_state_var*[(lb, ub)]
106            else:
107                self.bounds += num_state_var*[(None, None)]
108            self.cov = np.append(self.cov, value_cov)
109
110        self._scale_state()
111        self.cov = np.diag(self.cov)
112
113        # Set objective function (callable)
114        self.obj_func = obj_func
115
116        # Objective function values
117        self.state_func_values = None
118        self.ens_func_values = None
119
120        # Inflation factor used in SmcOpt
121        self.inflation_factor = None
122        self.survival_factor = None
123        self.particles = np.empty((self.cov.shape[0],0))
124        self.particle_values = np.empty((0))
125
126        # Initialize variables for bias correction
127        if 'bias_file' in self.sim.input_dict:  # use bias correction
128            self.bias_file = self.sim.input_dict['bias_file'].upper()  # mako file for simulations
129        else:
130            self.bias_file = None
131        self.bias_adaptive = None  # flag to adaptively update the bias correction (not implemented yet)
132        self.bias_factors = None  # this is J(x_j,m_j)/J(x_j,m)
133        self.bias_weights = np.ones(self.num_samples) / self.num_samples  # initialize with equal weights
134        self.bias_points = None  # this is the points used to estimate the bias correction
135
136        # Setup GenOpt
137        self.genopt = GenOptDistribution(self.get_state(), 
138                                         self.get_cov(), 
139                                         func=self.function, 
140                                         ne=self.num_samples)
141
142    def get_state(self):
143        """
144        Returns
145        -------
146        x : numpy.ndarray
147            Control vector as ndarray, shape (number of controls, number of perturbations)
148        """
149        x = ot.aug_optim_state(self.state, list(self.state.keys()))
150        return x
151
152    def get_cov(self):
153        """
154        Returns
155        -------
156        cov : numpy.ndarray
157            Covariance matrix, shape (number of controls, number of controls)
158        """
159
160        return self.cov
161
162    def get_bounds(self):
163        """
164        Returns
165        -------
166        bounds : list
167            (min, max) pairs for each element in x. None is used to specify no bound.
168        """
169
170        return self.bounds
171    
172    def get_final_state(self, return_dict=False):
173        """
174        Parameters
175        ----------
176        return_dict : bool
177            Retrun dictionary if true
178
179        Returns
180        -------
181        x : numpy.ndarray
182            Control vector as ndarray, shape (number of controls, number of perturbations)
183        """
184
185        self._invert_scale_state()
186        if return_dict:
187            x = self.state
188        else:
189            x = self.get_state()
190        return x
191
192    def function(self, x, *args, **kwargs):
193        """
194        This is the main function called during optimization.
195
196        Parameters
197        ----------
198        x : ndarray
199            Control vector, shape (number of controls, number of perturbations)
200
201        Returns
202        -------
203        obj_func_values : numpy.ndarray
204            Objective function values, shape (number of perturbations, )
205        """
206        self._aux_input()
207
208        if len(x.shape) == 1:
209            self.ne = self.num_models
210        else:
211            self.ne = x.shape[1]
212
213        self.state = ot.update_optim_state(x, self.state, list(self.state.keys()))  # go from nparray to dict
214        self._invert_scale_state()  # ensure that state is in [lb,ub]
215        run_success = self.calc_prediction()  # calculate flow data
216        self._scale_state()  # scale back to [0, 1]
217        if run_success:
218            func_values = self.obj_func(self.pred_data, input_dict=self.sim.input_dict,
219                                        true_order=self.sim.true_order, **kwargs)
220        else:
221            func_values = np.inf  # the simulations have crashed
222
223        if len(x.shape) == 1:
224            self.state_func_values = func_values
225        else:
226            self.ens_func_values = func_values
227        
228        return func_values
229
230    def gradient(self, x, *args, **kwargs):
231        r"""
232        Calculate the preconditioned gradient associated with ensemble, defined as:
233
234        .. math::
235            S \approx C_x \times G^T
236
237        where :math:`C_x` is the state covariance matrix, and :math:`G` is the standard
238        gradient. The ensemble sensitivity matrix is calculated as:
239
240        .. math::
241            S = X \times J^T /(N_e-1)
242
243        where :math:`X` and :math:`J` are ensemble matrices of :math:`x` (or control variables) and objective function
244        perturbed by their respective means. In practice (and in this method), :math:`S` is calculated by perturbing the
245        current control variable with Gaussian random numbers from :math:`N(0, C_x)` (giving :math:`X`), running
246        the generated ensemble (:math:`X`) through the simulator to give an ensemble of objective function values
247        (:math:`J`), and in the end calculate :math:`S`. Note that :math:`S` is an :math:`N_x \times 1` vector, where
248        :math:`N_x` is length of the control vector and the objective function is scalar.
249
250        Parameters
251        ----------
252        x : ndarray
253            Control vector, shape (number of controls, )
254
255        args : tuple
256            Covarice (:math:`C_x`), shape (number of controls, number of controls)
257
258        Returns
259        -------
260        gradient : numpy.ndarray
261                The gradient evaluated at x, shape (number of controls, )
262        """
263
264        # Set the ensemble state equal to the input control vector x
265        self.state = ot.update_optim_state(x, self.state, list(self.state.keys()))
266
267        # Set the covariance equal to the input
268        self.cov = args[0]
269
270        # If bias correction is used we need to temporarily store the initial state
271        initial_state = None
272        if self.bias_file is not None and self.bias_factors is None:  # first iteration
273            initial_state = deepcopy(self.state)  # store this to update current objective values
274
275        # Generate ensemble of states
276        self.ne = self.num_samples
277        nr = self._aux_input()
278        self.state = self._gen_state_ensemble()
279        
280        state_ens = at.aug_state(self.state, list(self.state.keys()))
281        self.function(state_ens, **kwargs)
282
283        # If bias correction is used we need to calculate the bias factors, J(u_j,m_j)/J(u_j,m)
284        if self.bias_file is not None:  # use bias corrections
285            self._bias_factors(self.ens_func_values, initial_state)
286
287        # Perturb state and function values with their mean
288        state_ens = at.aug_state(self.state, list(self.state.keys()))
289        pert_state = state_ens - np.dot(state_ens.mean(1)[:, None], np.ones((1, self.ne)))
290        if self.bias_file is not None:  # use bias corrections
291            self.ens_func_values *= self._bias_correction(self.state)
292            pert_obj_func = self.ens_func_values - np.mean(self.ens_func_values)
293        else:
294            pert_obj_func = self.ens_func_values - np.array(np.repeat(self.state_func_values, nr))
295
296        # Calculate the gradient
297        g_m = np.zeros(state_ens.shape[0])
298        for i in np.arange(self.ne):
299            g_m = g_m + pert_obj_func[i] * pert_state[:, i]
300
301        gradient = g_m / (self.ne - 1)
302
303        return gradient
304
305    def hessian(self, x=None, *args):
306        r"""
307        Calculate the hessian matrix associated with ensemble, defined as:
308
309        .. math::
310            H = J(XX^T - \Sigma)/ (N_e-1)
311
312        where :math:`X` and :math:`J` are ensemble matrices of :math:`x` (or control variables) and objective function
313        perturbed by their respective means.
314
315        Note: state and ens_func_values are assumed to already exist from computation of the gradient.
316        Save time by not running them again.
317
318        Parameters
319        ----------
320        x : ndarray
321            Control vector, shape (number of controls, number of perturbations)
322
323        Returns
324        -------
325        hessian: numpy.ndarray
326            The hessian evaluated at x, shape (number of controls, number of controls)
327
328        References
329        ----------
330        Zhang, Y., Stordal, A.S. & Lorentzen, R.J. A natural Hessian approximation for ensemble based optimization.
331        Comput Geosci 27, 355–364 (2023). https://doi.org/10.1007/s10596-022-10185-z
332        """
333
334        # Perturb state and function values with their mean
335        state_ens = at.aug_state(self.state, list(self.state.keys()))
336        pert_state = state_ens - np.dot(state_ens.mean(1)[:, None], np.ones((1, self.ne)))
337        nr = self._aux_input()
338        pert_obj_func = self.ens_func_values - np.array(np.repeat(self.state_func_values, nr))
339
340        # Calculate the gradient for mean and covariance matrix
341        g_c = np.zeros(self.cov.shape)
342        for i in np.arange(self.ne):
343            g_c = g_c + pert_obj_func[i] * (np.outer(pert_state[:, i], pert_state[:, i]) - self.cov)
344
345        hessian = g_c / (self.ne - 1)
346
347        return hessian
348    '''
349    def genopt_gradient(self, x, *args):
350        self.genopt.update_distribution(*args)
351        gradient = self.genopt.ensemble_gradient(func=self.function, 
352                                                 x=x, 
353                                                 ne=self.num_samples)
354        return gradient
355    
356    def genopt_mutation_gradient(self, x=None, *args, **kwargs):
357        return self.genopt.ensemble_mutation_gradient(return_ensembles=kwargs['return_ensembles'])
358    '''
359
360    def calc_ensemble_weights(self, x, *args, **kwargs):
361        r"""
362        Calculate weights used in sequential monte carlo optimization.
363
364        Parameters
365        ----------
366        x : ndarray
367            Control vector, shape (number of controls, )
368
369        args : tuple
370            Inflation factor, covariance (:math:`C_x`, shape (number of controls, number of controls)) and survival factor
371
372        Returns
373        -------
374        sens_matrix, best_ens, best_func : tuple
375                The weighted ensemble, the best ensemble member, and the best objective function value
376        """
377
378        # Set the ensemble state equal to the input control vector x
379        self.state = ot.update_optim_state(x, self.state, list(self.state.keys()))
380
381        # Set the inflation factor and covariance equal to the input
382        self.inflation_factor = args[0]
383        self.cov = args[1]
384        self.survival_factor = args[2]
385
386        # If bias correction is used we need to temporarily store the initial state
387        initial_state = None
388        if self.bias_file is not None and self.bias_factors is None:  # first iteration
389            initial_state = deepcopy(self.state)  # store this to update current objective values
390
391        # Generate ensemble of states
392        if self.particles.shape[1] == 0:
393            self.ne = self.num_samples
394        else:
395            self.ne = int(np.round(self.num_samples*self.survival_factor))
396        self._aux_input()
397        self.state = self._gen_state_ensemble()
398
399        self._invert_scale_state()  # ensure that state is in [lb,ub]
400        self.calc_prediction()  # calculate flow data
401        self._scale_state()  # scale back to [0, 1]
402
403        #self.ens_func_values = self.obj_func(self.pred_data, self.sim.input_dict, self.sim.true_order)
404        #self.ens_func_values = np.array(self.ens_func_values)
405        state_ens = at.aug_state(self.state, list(self.state.keys()))
406        self.function(state_ens, **kwargs)
407
408        self.particles = np.hstack((self.particles, state_ens))
409        self.particle_values = np.hstack((self.particle_values,self.ens_func_values))
410
411        # If bias correction is used we need to calculate the bias factors, J(u_j,m_j)/J(u_j,m)
412        if self.bias_file is not None:  # use bias corrections
413            self._bias_factors(self.ens_func_values, initial_state)
414
415        # Calculate the weights and ensemble sensitivity matrix
416        warnings.filterwarnings('ignore')  # suppress warnings
417        weights = np.zeros(self.num_samples)
418        for i in np.arange(self.num_samples):
419            weights[i] = np.exp(-(self.particle_values[i]-np.min(self.particle_values))*self.inflation_factor)
420
421        weights = weights + 0.000001
422        weights = weights/np.sum(weights)  # TODO: Sjekke at disse er riktig
423
424        sens_matrix = self.particles @ weights
425        index = np.argmin(self.particle_values)
426        best_ens = self.particles[:, index]
427        best_func = self.particle_values[index]
428        resample_index = np.random.choice(self.num_samples,int(np.round(self.num_samples-
429                                          self.num_samples*self.survival_factor)),replace=True,p=weights)
430        self.particles = self.particles[:, resample_index]
431        self.particle_values = self.particle_values[resample_index]
432        return sens_matrix, best_ens, best_func
433
434    def _gen_state_ensemble(self):
435        """
436        Generate ensemble with the current state (control variable) as the mean and using the covariance matrix
437        """
438
439        state_en = {}
440        cov_blocks = ot.corr2BlockDiagonal(self.state, self.cov)
441        for i, statename in enumerate(self.state.keys()):
442            mean = self.state[statename]
443            cov = cov_blocks[i]
444            temp_state_en = np.random.multivariate_normal(mean, cov, self.ne).transpose()
445            shifted_ensemble = np.array([mean]).T + temp_state_en - np.array([np.mean(temp_state_en, 1)]).T
446            if self.upper_bound and self.lower_bound:
447                if self.transform:
448                    np.clip(shifted_ensemble, 0, 1, out=shifted_ensemble)
449                else:
450                    np.clip(shifted_ensemble, self.lower_bound[i], self.upper_bound[i], out=shifted_ensemble)
451            state_en[statename] = shifted_ensemble
452
453        return state_en
454
455    def _aux_input(self):
456        """
457        Set the auxiliary input used for multiple geological realizations
458        """
459
460        nr = 1  # nr is the ratio of samples over models
461        if self.num_models > 1:
462            if np.remainder(self.num_samples, self.num_models) == 0:
463                nr = int(self.num_samples / self.num_models)
464                self.aux_input = list(np.repeat(np.arange(self.num_models), nr))
465            else:
466                print('num_samples must be a multiplum of num_models!')
467                sys.exit(0)
468        return nr
469
470    def _scale_state(self):
471        """
472        Transform the internal state from [lb, ub] to [0, 1]
473        """
474        if self.transform and (self.upper_bound and self.lower_bound):
475            for i, key in enumerate(self.state):
476                self.state[key] = (self.state[key] - self.lower_bound[i])/(self.upper_bound[i] - self.lower_bound[i])
477                np.clip(self.state[key], 0, 1, out=self.state[key])
478
479    def _invert_scale_state(self):
480        """
481        Transform the internal state from [0, 1] to [lb, ub]
482        """
483        if self.transform and (self.upper_bound and self.lower_bound):
484            for i, key in enumerate(self.state):
485                if self.transform:
486                    self.state[key] = self.lower_bound[i] + self.state[key]*(self.upper_bound[i] - self.lower_bound[i])
487                np.clip(self.state[key], self.lower_bound[i], self.upper_bound[i], out=self.state[key])
488
489    def _bias_correction(self, state):
490        """
491        Calculate bias correction. Currently, the bias correction is a constant independent of the state
492        """
493        if self.bias_factors is not None:
494            return np.sum(self.bias_weights * self.bias_factors)
495        else:
496            return 1
497
498    def _bias_factors(self, obj_func_values, initial_state):
499        """
500        Function for computing the bias factors
501        """
502
503        if self.bias_factors is None:  # first iteration
504            currentfile = self.sim.file
505            self.sim.file = self.bias_file
506            self.ne = self.num_samples
507            self.aux_input = list(np.arange(self.ne))
508            self.calc_prediction()
509            self.sim.file = currentfile
510            bias_func_values = self.obj_func(self.pred_data, self.sim.input_dict, self.sim.true_order)
511            bias_func_values = np.array(bias_func_values)
512            self.bias_factors = bias_func_values / obj_func_values
513            self.bias_points = deepcopy(self.state)
514            self.state_func_values *= self._bias_correction(initial_state)
515        elif self.bias_adaptive is not None and self.bias_adaptive > 0:  # update factors to account for new information
516            pass  # not implemented yet
517
518
519
520
521