pipt.update_schemes.enrml

   1"""
   2EnRML type schemes
   3"""
   4# External imports
   5import pipt.misc_tools.analysis_tools as at
   6from geostat.decomp import Cholesky
   7from pipt.loop.ensemble import Ensemble
   8from pipt.update_schemes.update_methods_ns.subspace_update import subspace_update
   9from pipt.update_schemes.update_methods_ns.full_update import full_update
  10from pipt.update_schemes.update_methods_ns.approx_update import approx_update
  11import sys
  12import pkgutil
  13import inspect
  14import numpy as np
  15import copy as cp
  16from scipy.linalg import cholesky, solve
  17
  18import importlib.util
  19
  20# List all available packages in the namespace package
  21# Import those that are present
  22import pipt.update_schemes.update_methods_ns as ns_pkg
  23tot_ns_pkg = []
  24# extract all class methods from namespace
  25for finder, name, ispkg in pkgutil.walk_packages(ns_pkg.__path__):
  26    spec = finder.find_spec(name)
  27    _module = importlib.util.module_from_spec(spec)
  28    spec.loader.exec_module(_module)
  29    tot_ns_pkg.extend(inspect.getmembers(_module, inspect.isclass))
  30
  31# import standard libraries
  32
  33# Check and import (if present) from other namespace packages
  34if 'margIS_update' in [el[0] for el in tot_ns_pkg]:  # only compare package name
  35    from pipt.update_schemes.update_methods_ns.margIS_update import margIS_update
  36else:
  37    class margIS_update:
  38        pass
  39
  40# Internal imports
  41
  42
  43class lmenrmlMixIn(Ensemble):
  44    """
  45    This is an implementation of EnRML using Levenberg-Marquardt. The update scheme is selected by a MixIn with multiple
  46    update_methods_ns. This class must therefore facititate many different update schemes.
  47    """
  48
  49    def __init__(self, keys_da, keys_fwd, sim):
  50        """
  51        The class is initialized by passing the PIPT init. file upwards in the hierarchy to be read and parsed in
  52        `pipt.input_output.pipt_init.ReadInitFile`.
  53
  54        Parameters
  55        ----------
  56        init_file: str
  57            PIPT init. file containing info. to run the inversion algorithm
  58        """
  59        # Pass the init_file upwards in the hierarchy
  60        super().__init__(keys_da, keys_fwd, sim)
  61
  62        if self.restart is False:
  63            # Save prior state in separate variable
  64            self.prior_state = cp.deepcopy(self.state)
  65
  66            # Extract parameters like conv. tol. and damping param. from ITERATION keyword in DATAASSIM
  67            self._ext_iter_param()
  68
  69            # Within variables
  70            self.prev_data_misfit = None  # Data misfit at previous iteration
  71            if 'actnum' in self.keys_da.keys():
  72                try:
  73                    self.actnum = np.load(self.keys_da['actnum'])['actnum']
  74                except:
  75                    print('ACTNUM file cannot be loaded!')
  76            else:
  77                self.actnum = None
  78            # At the moment, the iterative loop is threated as an iterative smoother and thus we check if assim. indices
  79            # are given as in the Simultaneous loop.
  80            self.check_assimindex_simultaneous()
  81            # define the assimilation index
  82            self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]]
  83            # define the list of states
  84            self.list_states = list(self.state.keys())
  85            # define the list of datatypes
  86            self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(
  87                self.obs_data, self.assim_index)
  88            # Get the perturbed observations and observation scaling
  89            self.data_random_state = cp.deepcopy(np.random.get_state())
  90            self._ext_obs()
  91            # Get state scaling and svd of scaled prior
  92            self._ext_state()
  93            self.current_state = cp.deepcopy(self.state)
  94
  95    def calc_analysis(self):
  96        """
  97        Calculate the update step in LM-EnRML, which is just the Levenberg-Marquardt update algorithm with
  98        the sensitivity matrix approximated by the ensemble.
  99        """
 100
 101        # reformat predicted data
 102        _, self.aug_pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index,
 103                                                     self.list_datatypes)
 104
 105        if self.iteration == 1:  # first iteration
 106            data_misfit = at.calc_objectivefun(
 107                self.real_obs_data, self.aug_pred_data, self.cov_data)
 108
 109            # Store the (mean) data misfit (also for conv. check)
 110            self.data_misfit = np.mean(data_misfit)
 111            self.prior_data_misfit = np.mean(data_misfit)
 112            self.data_misfit_std = np.std(data_misfit)
 113
 114            if self.lam == 'auto':
 115                self.lam = (0.5 * self.prior_data_misfit)/self.aug_pred_data.shape[0]
 116
 117            self.logger.info(
 118                f'Prior run complete with data misfit: {self.prior_data_misfit:0.1f}. Lambda for initial analysis: {self.lam}')
 119
 120        if 'localanalysis' in self.keys_da:
 121            self.local_analysis_update()
 122        else:
 123            # Mean pred_data and perturbation matrix with scaling
 124            if len(self.scale_data.shape) == 1:
 125                self.pert_preddata = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1),
 126                                            np.ones((1, self.ne))) * np.dot(self.aug_pred_data, self.proj)
 127            else:
 128                self.pert_preddata = solve(
 129                    self.scale_data, np.dot(self.aug_pred_data, self.proj))
 130
 131            aug_state = at.aug_state(self.current_state, self.list_states)
 132            self.update()  # run ordinary analysis
 133            if hasattr(self, 'step'):
 134                aug_state_upd = aug_state + self.step
 135            if hasattr(self, 'w_step'):
 136                self.W = self.current_W + self.w_step
 137                aug_prior_state = at.aug_state(self.prior_state, self.list_states)
 138                aug_state_upd = np.dot(aug_prior_state, (np.eye(
 139                    self.ne) + self.W / np.sqrt(self.ne - 1)))
 140
 141            # Extract updated state variables from aug_update
 142            self.state = at.update_state(aug_state_upd, self.state, self.list_states)
 143            self.state = at.limits(self.state, self.prior_info)
 144
 145    def check_convergence(self):
 146        """
 147        Check if LM-EnRML have converged based on evaluation of change sizes of objective function, state and damping
 148        parameter. 
 149
 150        Returns
 151        -------
 152        conv: bool
 153            Logic variable telling if algorithm has converged
 154        why_stop: dict
 155            Dict. with keys corresponding to conv. criteria, with logical variable telling which of them that has been
 156            met
 157        """
 158
 159        _, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index,
 160                                            self.list_datatypes)
 161        # Initialize the initial success value
 162        success = False
 163
 164        # if inital conv. check, there are no prev_data_misfit
 165        self.prev_data_misfit = self.data_misfit
 166        self.prev_data_misfit_std = self.data_misfit_std
 167
 168        # Calc. std dev of data misfit (used to update lamda)
 169        # mat_obs = np.dot(obs_data_vector.reshape((len(obs_data_vector),1)), np.ones((1, self.ne))) # use the perturbed
 170        # data instead.
 171
 172        data_misfit = at.calc_objectivefun(self.real_obs_data, pred_data, self.cov_data)
 173
 174        self.data_misfit = np.mean(data_misfit)
 175        self.data_misfit_std = np.std(data_misfit)
 176
 177        # # Calc. mean data misfit for convergence check, using the updated state variable
 178        # self.data_misfit = np.dot((mean_preddata - obs_data_vector).T,
 179        #                      solve(cov_data, (mean_preddata - obs_data_vector)))
 180
 181        # Convergence check: Relative step size of data misfit or state change less than tolerance
 182        if abs(1 - (self.data_misfit / self.prev_data_misfit)) < self.data_misfit_tol \
 183                or self.lam >= self.lam_max:
 184            # Logical variables for conv. criteria
 185            why_stop = {'data_misfit_stop': 1 - (self.data_misfit / self.prev_data_misfit) < self.data_misfit_tol,
 186                        'data_misfit': self.data_misfit,
 187                        'prev_data_misfit': self.prev_data_misfit,
 188                        'lambda': self.lam,
 189                        'lambda_stop': self.lam >= self.lam_max}
 190
 191            if self.data_misfit >= self.prev_data_misfit:
 192                success = False
 193                self.logger.info(
 194                    f'Iterations have converged after {self.iteration} iterations. Objective function reduced '
 195                    f'from {self.prior_data_misfit:0.1f} to {self.prev_data_misfit:0.1f}')
 196            else:
 197                self.logger.info(
 198                    f'Iterations have converged after {self.iteration} iterations. Objective function reduced '
 199                    f'from {self.prior_data_misfit:0.1f} to {self.data_misfit:0.1f}')
 200            # Return conv = True, why_stop var.
 201            return True, success, why_stop
 202
 203        else:  # conv. not met
 204            # Logical variables for conv. criteria
 205            why_stop = {'data_misfit_stop': 1 - (self.data_misfit / self.prev_data_misfit) < self.data_misfit_tol,
 206                        'data_misfit': self.data_misfit,
 207                        'prev_data_misfit': self.prev_data_misfit,
 208                        'lambda': self.lam,
 209                        'lambda_stop': self.lam >= self.lam_max}
 210
 211            ###############################################
 212            ##### update Lambda step-size values ##########
 213            ###############################################
 214            # If reduction in mean data misfit, reduce damping param
 215            if self.data_misfit < self.prev_data_misfit and self.data_misfit_std < self.prev_data_misfit_std:
 216                # Reduce damping parameter (divide calculations for ANALYSISDEBUG purpose)
 217                if self.lam > self.lam_min:
 218                    self.lam = self.lam / self.gamma
 219                success = True
 220                self.current_state = cp.deepcopy(self.state)
 221                if hasattr(self, 'W'):
 222                    self.current_W = cp.deepcopy(self.W)
 223            elif self.data_misfit < self.prev_data_misfit and self.data_misfit_std >= self.prev_data_misfit_std:
 224                # accept itaration, but keep lam the same
 225                success = True
 226                self.current_state = cp.deepcopy(self.state)
 227                if hasattr(self, 'W'):
 228                    self.current_W = cp.deepcopy(self.W)
 229
 230            else:  # Reject iteration, and increase lam
 231                # Increase damping parameter (divide calculations for ANALYSISDEBUG purpose)
 232                self.lam = self.lam * self.gamma
 233                success = False
 234
 235            if success:
 236                self.logger.info(f'Successfull iteration number {self.iteration}! Objective function reduced from '
 237                                 f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}. New Lamba for next analysis: '
 238                                 f'{self.lam}')
 239                # self.prev_data_misfit = self.data_misfit
 240                # self.prev_data_misfit_std = self.data_misfit_std
 241            else:
 242                self.logger.info(f'Failed iteration number {self.iteration}! Objective function increased from '
 243                                 f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}. New Lamba for repeated analysis: '
 244                                 f'{self.lam}')
 245                # Reset the objective function after report
 246                self.data_misfit = self.prev_data_misfit
 247                self.data_misfit_std = self.prev_data_misfit_std
 248
 249            # Return conv = False, why_stop var.
 250            return False, success, why_stop
 251
 252    def _ext_iter_param(self):
 253        """
 254        Extract parameters needed in LM-EnRML from the ITERATION keyword given in the DATAASSIM part of PIPT init.
 255        file. These parameters include convergence tolerances and parameters for the damping parameter. Default
 256        values for these parameters have been given here, if they are not provided in ITERATION.
 257        """
 258
 259        # Predefine all the default values
 260        self.data_misfit_tol = 0.01
 261        self.step_tol = 0.01
 262        self.lam = 100
 263        self.lam_max = 1e10
 264        self.lam_min = 0.01
 265        self.gamma = 5
 266        self.trunc_energy = 0.95
 267        self.iteration = 0
 268
 269        # Loop over options in ITERATION and extract the parameters we want
 270        for i, opt in enumerate(list(zip(*self.keys_da['iteration']))[0]):
 271            if opt == 'data_misfit_tol':
 272                self.data_misfit_tol = self.keys_da['iteration'][i][1]
 273            if opt == 'step_tol':
 274                self.step_tol = self.keys_da['iteration'][i][1]
 275            if opt == 'lambda':
 276                self.lam = self.keys_da['iteration'][i][1]
 277            if opt == 'lambda_max':
 278                self.lam_max = self.keys_da['iteration'][i][1]
 279            if opt == 'lambda_min':
 280                self.lam_min = self.keys_da['iteration'][i][1]
 281            if opt == 'lambda_factor':
 282                self.gamma = self.keys_da['iteration'][i][1]
 283
 284        if 'energy' in self.keys_da:
 285            # initial energy (Remember to extract this)
 286            self.trunc_energy = self.keys_da['energy']
 287            if self.trunc_energy > 1:  # ensure that it is given as percentage
 288                self.trunc_energy /= 100.
 289
 290
 291class lmenrml_approx(lmenrmlMixIn, approx_update):
 292    pass
 293
 294
 295class lmenrml_full(lmenrmlMixIn, full_update):
 296    pass
 297
 298
 299class lmenrml_subspace(lmenrmlMixIn, subspace_update):
 300    pass
 301
 302
 303class gnenrmlMixIn(Ensemble):
 304    """
 305    This is an implementation of EnRML using the Gauss-Newton approach. The update scheme is selected by a MixIn with multiple
 306    update_methods_ns. This class must therefore facititate many different update schemes.
 307    """
 308
 309    def __init__(self, keys_da, keys_fwd, sim):
 310        """
 311        The class is initialized by passing the PIPT init. file upwards in the hierarchy to be read and parsed in
 312        `pipt.input_output.pipt_init.ReadInitFile`.
 313
 314        Parameters
 315        ----------
 316        init_file: str
 317            PIPT init. file containing info. to run the inversion algorithm
 318        """
 319        # Pass the init_file upwards in the hierarchy
 320        super().__init__(keys_da, keys_fwd, sim)
 321
 322        if self.restart is False:
 323            # Save prior state in separate variable
 324            self.prior_state = cp.deepcopy(self.state)
 325
 326            # extract and save state scaling
 327
 328            # Extract parameters like conv. tol. and damping param. from ITERATION keyword in DATAASSIM
 329            self._ext_iter_param()
 330
 331            # Within variables
 332            self.prev_data_misfit = None  # Data misfit at previous iteration
 333            if 'actnum' in self.keys_da.keys():
 334                try:
 335                    self.actnum = np.load(self.keys_da['actnum'])['actnum']
 336                except:
 337                    print('ACTNUM file cannot be loaded!')
 338            else:
 339                self.actnum = None
 340            # At the moment, the iterative loop is threated as an iterative smoother and thus we check if assim. indices
 341            # are given as in the Simultaneous loop.
 342            self.check_assimindex_simultaneous()
 343            # define the assimilation index
 344            self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]]
 345            # define the list of states
 346            self.list_states = list(self.state.keys())
 347            # define the list of datatypes
 348            self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(
 349                self.obs_data, self.assim_index)
 350            # Get the perturbed observations and observation scaling
 351            self._ext_obs()
 352            # Get state scaling and svd of scaled prior
 353            self._ext_state()
 354            self.current_state = cp.deepcopy(self.state)
 355            # ensure that the updates does not invoke the LM inflation of the Hessian.
 356            self.lam = 0
 357
 358    def _ext_obs(self):
 359
 360        self.obs_data_vector, _ = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index,
 361                                                       self.list_datatypes)
 362
 363        # Generate the data auto-covariance matrix
 364        if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes':
 365            if hasattr(self, 'cov_data'):  # cd matrix has been imported
 366                tmp_E = np.dot(cholesky(self.cov_data).T,
 367                               np.random.randn(self.cov_data.shape[0], self.ne))
 368            else:
 369                tmp_E = at.extract_tot_empirical_cov(
 370                    self.datavar, self.assim_index, self.list_datatypes, self.ne)
 371            # self.E = (tmp_E - tmp_E.mean(1)[:,np.newaxis])/np.sqrt(self.ne - 1)/
 372            if 'screendata' in self.keys_da and self.keys_da['screendata'] == 'yes':
 373                tmp_E = at.screen_data(tmp_E, self.aug_pred_data,
 374                                       self.obs_data_vector, self.iteration)
 375            self.E = tmp_E
 376            self.real_obs_data = self.obs_data_vector[:, np.newaxis] - tmp_E
 377
 378            self.cov_data = np.var(self.E, ddof=1,
 379                                   axis=1)  # calculate the variance, to be used for e.g. data misfit calc
 380            # self.cov_data = ((self.E * self.E)/(self.ne-1)).sum(axis=1) # calculate the variance, to be used for e.g. data misfit calc
 381            self.scale_data = np.sqrt(self.cov_data)
 382        else:
 383            if not hasattr(self, 'cov_data'):  # if cd is not loaded
 384                self.cov_data = at.gen_covdata(
 385                    self.datavar, self.assim_index, self.list_datatypes)
 386            # data screening
 387            if 'screendata' in self.keys_da and self.keys_da['screendata'] == 'yes':
 388                self.cov_data = at.screen_data(
 389                    self.cov_data, self.aug_pred_data, self.obs_data_vector, self.iteration)
 390
 391            init_en = Cholesky()  # Initialize GeoStat class for generating realizations
 392            self.real_obs_data, self.scale_data = init_en.gen_real(self.obs_data_vector, self.cov_data, self.ne,
 393                                                                   return_chol=True)
 394
 395    def _ext_state(self):
 396        # get vector of scaling
 397        self.state_scaling = at.calc_scaling(
 398            self.prior_state, self.list_states, self.prior_info)
 399
 400        delta_scaled_prior = self.state_scaling[:, None] * \
 401            np.dot(at.aug_state(self.prior_state, self.list_states), self.proj)
 402
 403        u_d, s_d, v_d = np.linalg.svd(delta_scaled_prior, full_matrices=False)
 404
 405        # remove the last singular value/vector. This is because numpy returns all ne values, while the last is actually
 406        # zero. This part is a good place to include eventual additional truncation.
 407        energy = 0
 408        trunc_index = len(s_d) - 1  # inititallize
 409        for c, elem in enumerate(s_d):
 410            energy += elem
 411            if energy / sum(s_d) >= self.trunc_energy:
 412                trunc_index = c  # take the index where all energy is preserved
 413                break
 414        u_d, s_d, v_d = u_d[:, :trunc_index +
 415                            1], s_d[:trunc_index + 1], v_d[:trunc_index + 1, :]
 416        self.Am = np.dot(u_d, np.eye(trunc_index+1) *
 417                         ((s_d**(-1))[:, None]))  # notation from paper
 418
 419    def calc_analysis(self):
 420        """
 421        Calculate the update step in LM-EnRML, which is just the Levenberg-Marquardt update algorithm with
 422        the sensitivity matrix approximated by the ensemble.
 423
 424        """
 425
 426        # reformat predicted data
 427        _, self.aug_pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index,
 428                                                     self.list_datatypes)
 429
 430        if self.iteration == 1:  # first iteration
 431            data_misfit = at.calc_objectivefun(
 432                self.real_obs_data, self.aug_pred_data, self.cov_data)
 433
 434            # Store the (mean) data misfit (also for conv. check)
 435            self.data_misfit = np.mean(data_misfit)
 436            self.prior_data_misfit = np.mean(data_misfit)
 437            self.data_misfit_std = np.std(data_misfit)
 438
 439            if self.gamma == 'auto':
 440                self.gamma = 0.1
 441
 442        # Mean pred_data and perturbation matrix with scaling
 443        if len(self.scale_data.shape) == 1:
 444            self.pert_preddata = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1),
 445                                        np.ones((1, self.ne))) * np.dot(self.aug_pred_data, self.proj)
 446        else:
 447            self.pert_preddata = solve(
 448                self.scale_data, np.dot(self.aug_pred_data, self.proj))
 449
 450        aug_state = at.aug_state(self.current_state, self.list_states)
 451
 452        self.update()  # run analysis
 453        if hasattr(self, 'step'):
 454            aug_state_upd = aug_state + self.gamma*self.step
 455        if hasattr(self, 'w_step'):
 456            self.W = self.current_W + self.gamma*self.w_step
 457            aug_prior_state = at.aug_state(self.prior_state, self.list_states)
 458            aug_state_upd = np.dot(aug_prior_state, (np.eye(
 459                self.ne) + self.W / np.sqrt(self.ne - 1)))
 460        if hasattr(self, 'sqrt_w_step'):  # if we do a sqrt update
 461            self.w = self.current_w + self.gamma*self.sqrt_w_step
 462            new_mean_state = self.mean_prior + np.dot(self.X, self.w)
 463            u, sigma, v = np.linalg.svd(self.C_w, full_matrices=True)
 464            sigma_inv_sqrt = np.diag([el_s ** (-1 / 2) for el_s in sigma])
 465            C_w_inv_sqrt = np.dot(np.dot(u, sigma_inv_sqrt), v.T)
 466            self.W = C_w_inv_sqrt * np.sqrt(self.ne - 1)
 467            aug_state_upd = np.tile(new_mean_state, (self.ne, 1)
 468                                    ).T + np.dot(self.X, self.W)
 469
 470        # Extract updated state variables from aug_update
 471        self.state = at.update_state(aug_state_upd, self.state, self.list_states)
 472        self.state = at.limits(self.state, self.prior_info)
 473
 474    def check_convergence(self):
 475        """
 476        Check if LM-EnRML have converged based on evaluation of change sizes of objective function, state and damping
 477        parameter.
 478
 479        Returns
 480        -------
 481        conv: bool
 482            Logic variable telling if algorithm has converged
 483        why_stop: dict
 484            Dict. with keys corresponding to conv. criteria, with logical variable telling which of them that has been
 485            met
 486        """
 487        # Prelude to calc. conv. check (everything done below is from calc_analysis)
 488        if hasattr(self, 'list_datatypes'):
 489            assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]]
 490            list_datatypes = self.list_datatypes
 491            cov_data = self.cov_data
 492            obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, assim_index,
 493                                                              list_datatypes)
 494            mean_preddata = np.mean(pred_data, 1)
 495        else:
 496            assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]]
 497            list_datatypes, _ = at.get_list_data_types(self.obs_data, assim_index)
 498            # cov_data = at.gen_covdata(self.datavar, assim_index, list_datatypes)
 499            obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, assim_index,
 500                                                              list_datatypes)
 501            # mean_preddata = np.mean(pred_data, 1)
 502
 503        # Initialize the initial success value
 504        success = False
 505
 506        # if inital conv. check, there are no prev_data_misfit
 507        if self.prev_data_misfit is None:
 508            self.data_misfit = np.mean(self.data_misfit)
 509            self.prev_data_misfit = self.data_misfit
 510            self.prev_data_misfit_std = self.data_misfit_std
 511            success = True
 512
 513        # update the last mismatch, only if this was a reduction of the misfit
 514        if self.data_misfit < self.prev_data_misfit:
 515            self.prev_data_misfit = self.data_misfit
 516            self.prev_data_misfit_std = self.data_misfit_std
 517            success = True
 518        # if there was no reduction of the misfit, retain the old "valid" data misfit.
 519
 520        # Calc. std dev of data misfit (used to update lamda)
 521        # mat_obs = np.dot(obs_data_vector.reshape((len(obs_data_vector),1)), np.ones((1, self.ne))) # use the perturbed
 522        # data instead.
 523        mat_obs = self.real_obs_data
 524        data_misfit = at.calc_objectivefun(mat_obs, pred_data, self.cov_data)
 525
 526        self.data_misfit = np.mean(data_misfit)
 527        self.data_misfit_std = np.std(data_misfit)
 528
 529        # # Calc. mean data misfit for convergence check, using the updated state variable
 530        # self.data_misfit = np.dot((mean_preddata - obs_data_vector).T,
 531        #                      solve(cov_data, (mean_preddata - obs_data_vector)))
 532
 533        # Convergence check: Relative step size of data misfit or state change less than tolerance
 534        if abs(1 - (self.data_misfit / self.prev_data_misfit)) < self.data_misfit_tol:
 535            # Logical variables for conv. criteria
 536            why_stop = {'data_misfit_stop': 1 - (self.data_misfit / self.prev_data_misfit) < self.data_misfit_tol,
 537                        'data_misfit': self.data_misfit,
 538                        'prev_data_misfit': self.prev_data_misfit,
 539                        'gamma': self.gamma,
 540                        }
 541
 542            if self.data_misfit >= self.prev_data_misfit:
 543                success = False
 544                self.logger.info(
 545                    f'Iterations have converged after {self.iteration} iterations. Objective function reduced '
 546                    f'from {self.prior_data_misfit:0.1f} to {self.prev_data_misfit:0.1f}')
 547            else:
 548                self.logger.info(
 549                    f'Iterations have converged after {self.iteration} iterations. Objective function reduced '
 550                    f'from {self.prior_data_misfit:0.1f} to {self.data_misfit:0.1f}')
 551            # Return conv = True, why_stop var.
 552            return True, success, why_stop
 553
 554        else:  # conv. not met
 555            # Logical variables for conv. criteria
 556            why_stop = {'data_misfit_stop': 1 - (self.data_misfit / self.prev_data_misfit) < self.data_misfit_tol,
 557                        'data_misfit': self.data_misfit,
 558                        'prev_data_misfit': self.prev_data_misfit,
 559                        'gamma': self.gamma}
 560
 561            ###############################################
 562            ##### update Lambda step-size values ##########
 563            ###############################################
 564            # If reduction in mean data misfit, reduce damping param
 565            if self.data_misfit < self.prev_data_misfit and self.data_misfit_std < self.prev_data_misfit_std:
 566                # Reduce damping parameter (divide calculations for ANALYSISDEBUG purpose)
 567                self.gamma = self.gamma + (self.gamma_max - self.gamma) * 2 ** (
 568                    -(self.iteration) / (self.gamma_factor - 1))
 569                success = True
 570                self.current_state = cp.deepcopy(self.state)
 571                if hasattr(self, 'W'):
 572                    self.current_W = cp.deepcopy(self.W)
 573
 574            elif self.data_misfit < self.prev_data_misfit and self.data_misfit_std >= self.prev_data_misfit_std:
 575                # accept itaration, but keep lam the same
 576                success = True
 577                self.current_state = cp.deepcopy(self.state)
 578                if hasattr(self, 'W'):
 579                    self.current_W = cp.deepcopy(self.W)
 580
 581            else:  # Reject iteration, and increase lam
 582                # Increase damping parameter (divide calculations for ANALYSISDEBUG purpose)
 583                err_str = f"Misfit increased. Set new start step length and try again. Final ojective function value is {self.data_misfit:0.1f}"
 584                self.logger.info(err_str)
 585                sys.exit(err_str)
 586                success = False
 587
 588            if success:
 589                self.logger.info(f'Successfull iteration number {self.iteration}! Objective function reduced from '
 590                                 f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}. New Gamma for next analysis: '
 591                                 f'{self.gamma}')
 592            else:
 593                self.logger.info(f'Failed iteration number {self.iteration}! Objective function increased from '
 594                                 f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}. New Gamma for repeated analysis: '
 595                                 f'{self.gamma}')
 596
 597            # Return conv = False, why_stop var.
 598            return False, success, why_stop
 599
 600    def _ext_iter_param(self):
 601        """
 602        Extract parameters needed in LM-EnRML from the ITERATION keyword given in the DATAASSIM part of PIPT init.
 603        file. These parameters include convergence tolerances and parameters for the damping parameter. Default
 604        values for these parameters have been given here, if they are not provided in ITERATION.
 605        """
 606
 607        # Predefine all the default values
 608        self.data_misfit_tol = 0.01
 609        self.step_tol = 0.01
 610        self.gamma = 0.2
 611        self.gamma_max = 0.5
 612        self.gamma_factor = 2.5
 613        self.trunc_energy = 0.95
 614        self.iteration = 0
 615
 616        # Loop over options in ITERATION and extract the parameters we want
 617        for i, opt in enumerate(list(zip(*self.keys_da['iteration']))[0]):
 618            if opt == 'data_misfit_tol':
 619                self.data_misfit_tol = self.keys_da['iteration'][i][1]
 620            if opt == 'step_tol':
 621                self.step_tol = self.keys_da['iteration'][i][1]
 622            if opt == 'gamma':
 623                self.gamma = self.keys_da['iteration'][i][1]
 624            if opt == 'gamma_max':
 625                self.gamma_max = self.keys_da['iteration'][i][1]
 626            if opt == 'gamma_factor':
 627                self.gamma_factor = self.keys_da['iteration'][i][1]
 628
 629        if 'energy' in self.keys_da:
 630            # initial energy (Remember to extract this)
 631            self.trunc_energy = self.keys_da['energy']
 632            if self.trunc_energy > 1:  # ensure that it is given as percentage
 633                self.trunc_energy /= 100.
 634
 635
 636class gnenrml_approx(gnenrmlMixIn, approx_update):
 637    pass
 638
 639
 640class gnenrml_full(gnenrmlMixIn, full_update):
 641    pass
 642
 643
 644class gnenrml_subspace(gnenrmlMixIn, subspace_update):
 645    pass
 646
 647
 648class gnenrml_margis(gnenrmlMixIn, margIS_update):
 649    '''
 650    The marg-IS scheme is currently not available in this version of PIPT. To utilize the scheme you have to import the
 651    *margIS_update* class from a standalone repository.
 652    '''
 653    pass
 654
 655
 656class co_lm_enrml(lmenrmlMixIn, approx_update):
 657    """
 658    This is the implementation of the approximative LM-EnRML algorithm as described in [1].
 659
 660    This algorithm is quite similar to the lm_enrml as provided above, and will therefore inherit most of its methods.
 661    We only change the calc_analysis part...
 662
 663    % Copyright (c) 2019-2022 NORCE, All Rights Reserved. 4DSEIS
 664    """
 665
 666    def __init__(self, keys_da):
 667        """
 668        The class is initialized by passing the PIPT init. file upwards in the hierarchy to be read and parsed in
 669        `pipt.input_output.pipt_init.ReadInitFile`.
 670
 671        Parameters
 672        ----------
 673        init_file: str
 674            PIPT init. file containing info. to run the inversion algorithm
 675
 676        References
 677        ----------
 678        [1] Chen Y. & Oliver D.S. 2013, Levenberg-Marquardt Forms of the Iterative Ensemble Smoother for Efficient
 679        History Matching and Uncertainty Quantification. Comput. Geosci., 17 (4): 689-703
 680        """
 681        # Call __init__ in parent class
 682        super().__init__(keys_da)
 683
 684    def calc_analysis(self):
 685        """
 686        Calculate the update step in approximate LM-EnRML code.
 687
 688        Parameters
 689        ----------
 690        iteration: int
 691            Iteration number
 692
 693        Returns
 694        -------
 695        success: bool
 696            True if data mismatch is decreasing, False if increasing
 697
 698        References
 699        ----------  
 700        [1] Chen Y. & Oliver D.S. 2013, Levenberg-Marquardt Forms of the Iterative Ensemble Smoother for Efficient
 701        History Matching and Uncertainty Quantification. Comput. Geosci., 17 (4): 689-703
 702        """
 703        # Get assimilation order as a list
 704        self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]]
 705
 706        # When handling large cases, it may be very costly to assemble the data covariance and localizaton matrix.
 707        # To alleviate this in the simultuaneus-iterative scheme we store these matrices, the list of states and
 708        # the list of data types after the first iteration.
 709
 710        if not hasattr(self, 'list_datatypes'):
 711            # Get list of data types to be assimilated and of the free states. Do this once, because listing keys from a
 712            # Python dictionary just when needed (in different places) may not yield the same list!
 713            self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(
 714                self.obs_data, self.assim_index)
 715            list_datatypes = self.list_datatypes
 716            self.list_states = list(self.state.keys())
 717            list_states = self.list_states
 718            list_act_datatypes = self.list_act_datatypes
 719
 720            # self.cov_data = np.load('CD.npz')['arr_0']
 721            # Generate the realizations of the observed data once
 722            # Augment observed and predicted data
 723            self.obs_data_vector, self.aug_pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index,
 724                                                                            self.list_datatypes)
 725            obs_data_vector = self.obs_data_vector
 726
 727            # Generate the data auto-covariance matrix
 728            if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes':
 729                if hasattr(self, 'cov_data'):  # cd matrix has been imported
 730                    tmp_E = np.dot(cholesky(self.cov_data).T,
 731                                   np.random.randn(self.cov_data.shape[0], self.ne))
 732                else:
 733                    tmp_E = at.extract_tot_empirical_cov(
 734                        self.datavar, self.assim_index, self.list_datatypes, self.ne)
 735                # self.E = (tmp_E - tmp_E.mean(1)[:,np.newaxis])/np.sqrt(self.ne - 1)/
 736                if 'screendata' in self.keys_da and self.keys_da['screendata'] == 'yes':
 737                    tmp_E = at.screen_data(tmp_E, self.aug_pred_data,
 738                                           obs_data_vector, self.iteration)
 739                self.E = tmp_E
 740                self.real_obs_data = obs_data_vector[:, np.newaxis] - tmp_E
 741
 742                self.cov_data = np.var(self.E, ddof=1,
 743                                       axis=1)  # calculate the variance, to be used for e.g. data misfit calc
 744                # self.cov_data = ((self.E * self.E)/(self.ne-1)).sum(axis=1) # calculate the variance, to be used for e.g. data misfit calc
 745                self.scale_data = np.sqrt(self.cov_data)
 746            else:
 747                if not hasattr(self, 'cov_data'):  # if cd is not loaded
 748                    self.cov_data = at.gen_covdata(
 749                        self.datavar, self.assim_index, self.list_datatypes)
 750                # data screening
 751                if 'screendata' in self.keys_da and self.keys_da['screendata'] == 'yes':
 752                    self.cov_data = at.screen_data(
 753                        self.cov_data, self.aug_pred_data, obs_data_vector, self.iteration)
 754
 755                init_en = Cholesky()  # Initialize GeoStat class for generating realizations
 756                self.real_obs_data, self.scale_data = init_en.gen_real(self.obs_data_vector, self.cov_data, self.ne,
 757                                                                       return_chol=True)
 758
 759            self.datavar = at.update_datavar(
 760                self.cov_data, self.datavar, self.assim_index, self.list_datatypes)
 761            self.current_state = cp.deepcopy(self.state)
 762
 763            # Calc. misfit for the initial iteration
 764            data_misfit = at.calc_objectivefun(
 765                self.real_obs_data, self.aug_pred_data, self.cov_data)
 766            # Store the (mean) data misfit (also for conv. check)
 767            self.data_misfit = np.mean(data_misfit)
 768            self.prior_data_misfit = np.mean(data_misfit)
 769            self.data_misfit_std = np.std(data_misfit)
 770
 771            if self.lam == 'auto':
 772                self.lam = 0.5 * self.prior_data_misfit
 773
 774        else:
 775            _, self.aug_pred_data = at.aug_obs_pred_data(
 776                self.obs_data, self.pred_data, assim_index, self.list_datatypes)
 777
 778        # Mean pred_data and perturbation matrix with scaling
 779        mean_preddata = np.mean(self.aug_pred_data, 1)
 780        if len(self.scale_data.shape) == 1:
 781            if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes':
 782                pert_preddata = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), np.ones((1, self.ne))) * (
 783                    self.aug_pred_data - np.dot(mean_preddata[:, None], np.ones((1, self.ne))))
 784            else:
 785                pert_preddata = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), np.ones((1, self.ne))) * (
 786                    self.aug_pred_data - np.dot(mean_preddata[:, None], np.ones((1, self.ne)))) / \
 787                    (np.sqrt(self.ne - 1))
 788        else:
 789            if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes':
 790                pert_preddata = solve(self.scale_data, self.aug_pred_data -
 791                                      np.dot(mean_preddata[:, None], np.ones((1, self.ne))))
 792            else:
 793                pert_preddata = solve(self.scale_data, self.aug_pred_data - np.dot(mean_preddata[:, None], np.ones((1, self.ne)))) / \
 794                    (np.sqrt(self.ne - 1))
 795        self.pert_preddata = pert_preddata
 796
 797        self.update()
 798        if hasattr(self, 'step'):
 799            aug_state_upd = aug_state + self.step
 800        if hasattr(self, 'w_step'):
 801            self.W = self.current_W - self.w_step
 802            aug_prior_state = at.aug_state(self.prior_state, self.list_states)
 803            aug_state_upd = np.dot(aug_prior_state, (np.eye(
 804                self.ne) + self.W / np.sqrt(self.ne - 1)))
 805
 806        # Extract updated state variables from aug_update
 807        self.state = at.update_state(aug_state_upd, self.state, self.list_states)
 808        self.state = at.limits(self.state, self.prior_info)
 809
 810
 811class gn_enrml(lmenrmlMixIn):
 812    """
 813    This is the implementation of the stochastig IES as  described in [1]. More information about the method is found in
 814    [2]. This implementation is the Gauss-Newton version.
 815
 816    This algorithm is quite similar to the `lm_enrml` as provided above, and will therefore inherit most of its methods.
 817    We only change the calc_analysis part...
 818    """
 819
 820    def __init__(self, keys_da):
 821        """
 822        The class is initialized by passing the PIPT init. file upwards in the hierarchy to be read and parsed in
 823        `pipt.input_output.pipt_init.ReadInitFile`.
 824
 825        Parameters
 826        ----------
 827        init_file: str
 828            PIPT init. file containing info. to run the inversion algorithm
 829
 830        References
 831        ----------
 832        [1] Raanes, P. N., Stordal, A. S., & Evensen, G. (2019). Revising the stochastic iterative ensemble smoother.
 833        Nonlinear Processes in Geophysics, 26(3), 325-338. https://doi.org/10.5194/npg-26-325-2019 <br>
 834        [2] Evensen, G., Raanes, P. N., Stordal, A. S., & Hove, J. (2019). Efficient Implementation of an Iterative
 835        Ensemble Smoother for Data Assimilation and Reservoir History Matching. Frontiers in Applied Mathematics and
 836        Statistics, 5(October), 114. https://doi.org/10.3389/fams.2019.00047
 837        """
 838        # Call __init__ in parent class
 839        super().__init__(keys_da)
 840
 841    def calc_analysis(self):
 842        """
 843        Changelog
 844        ---------
 845        - KF 25/2-20
 846        """
 847        # Get assimilation order as a list
 848        assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]]
 849
 850        # When handling large cases, it may be very costly to assemble the data covariance and localizaton matrix.
 851        # To alleviate this in the simultuaneus-iterative scheme we store these matrices, the list of states and
 852        # the list of data types after the first iteration.
 853
 854        if not hasattr(self, 'list_datatypes'):
 855            # Get list of data types to be assimilated and of the free states. Do this once, because listing keys from a
 856            # Python dictionary just when needed (in different places) may not yield the same list!
 857            self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(
 858                self.obs_data, assim_index)
 859            list_datatypes = self.list_datatypes
 860            self.list_states = list(self.state.keys())
 861            list_states = self.list_states
 862            list_act_datatypes = self.list_act_datatypes
 863
 864            # Generate the realizations of the observed data once
 865            # Augment observed and predicted data
 866            self.obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, assim_index,
 867                                                                   self.list_datatypes)
 868            obs_data_vector = self.obs_data_vector
 869
 870            if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes':
 871                if hasattr(self, 'cov_data'):  # cd matrix has been imported
 872                    tmp_E = np.dot(cholesky(self.cov_data).T, np.random.randn(
 873                        self.cov_data.shape[0], self.ne))
 874                else:
 875                    tmp_E = at.extract_tot_empirical_cov(
 876                        self.datavar, assim_index, self.list_datatypes, self.ne)
 877                # self.E = (tmp_E - tmp_E.mean(1)[:,np.newaxis])/np.sqrt(self.ne - 1)/
 878                self.real_obs_data = obs_data_vector[:, np.newaxis] - tmp_E
 879
 880                self.cov_data = np.var(tmp_E, ddof=1,
 881                                       axis=1)  # calculate the variance, to be used for e.g. data misfit calc
 882                # self.cov_data = ((self.E * self.E)/(self.ne-1)).sum(axis=1) # calculate the variance, to be used for e.g. data misfit calc
 883                self.scale_data = np.sqrt(self.cov_data)
 884            else:
 885                if not hasattr(self, 'cov_data'):  # if cd is not loaded
 886                    self.cov_data = at.gen_covdata(
 887                        self.datavar, assim_index, self.list_datatypes)
 888                # data screening
 889                if 'screendata' in self.keys_da and self.keys_da['screendata'] == 'yes':
 890                    self.cov_data = at.screen_data(
 891                        self.cov_data, pred_data, obs_data_vector, self.iteration)
 892
 893                init_en = Cholesky()  # Initialize GeoStat class for generating realizations
 894                self.real_obs_data, self.scale_data = init_en.gen_real(self.obs_data_vector, self.cov_data, self.ne,
 895                                                                       return_chol=True)
 896
 897            self.datavar = at.update_datavar(
 898                self.cov_data, self.datavar, assim_index, self.list_datatypes)
 899            cov_data = self.cov_data
 900            obs_data = self.real_obs_data
 901            #
 902            self.current_state = cp.deepcopy(self.state)
 903            #
 904            self.aug_prior = cp.deepcopy(at.aug_state(
 905                self.current_state, self.list_states))
 906            # self.mean_prior = aug_prior.mean(axis=1)
 907            # self.X = (aug_prior - np.dot(np.resize(self.mean_prior, (len(self.mean_prior), 1)),
 908            #                                                  np.ones((1, self.ne))))
 909            self.W = np.zeros((self.ne, self.ne))
 910
 911            self.proj = (np.eye(self.ne) - (1 / self.ne) *
 912                         np.ones((self.ne, self.ne))) / np.sqrt(self.ne - 1)
 913            self.E = np.dot(obs_data, self.proj)
 914
 915            # Calc. misfit for the initial iteration
 916            if len(cov_data.shape) == 1:
 917                tmp_data_misfit = np.diag(np.dot((pred_data - obs_data).T,
 918                                                 np.dot(np.expand_dims(self.cov_data ** (-1), axis=1),
 919                                                        np.ones((1, self.ne))) * (pred_data - obs_data)))
 920            else:
 921                tmp_data_misfit = np.diag(
 922                    np.dot((pred_data - obs_data).T, solve(self.cov_data, (pred_data - obs_data))))
 923            mean_data_misfit = np.mean(tmp_data_misfit)
 924            # mean_data_misfit = np.median(tmp_data_misfit)
 925            std_data_misfit = np.std(tmp_data_misfit)
 926
 927            # Store the (mean) data misfit (also for conv. check)
 928            self.data_misfit = mean_data_misfit
 929            self.prior_data_misfit = mean_data_misfit
 930            self.data_misfit_std = std_data_misfit
 931
 932        else:
 933            # for analysis debug...
 934            list_datatypes = self.list_datatypes
 935            list_act_datatypes = self.list_act_datatypes
 936            list_states = self.list_states
 937            cov_data = self.cov_data
 938            obs_data_vector = self.obs_data_vector
 939            _, pred_data = at.aug_obs_pred_data(
 940                self.obs_data, self.pred_data, assim_index, self.list_datatypes)
 941            obs_data = self.real_obs_data
 942
 943        if len(self.scale_data.shape) == 1:
 944            Y = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), np.ones((1, self.ne))) * \
 945                np.dot(pred_data, self.proj)
 946        else:
 947            Y = solve(self.scale_data, np.dot(pred_data, self.proj))
 948        omega = np.eye(self.ne) + np.dot(self.W, self.proj)
 949        LU = lu_factor(omega.T)
 950        S = lu_solve(LU, Y.T).T
 951        if len(self.scale_data.shape) == 1:
 952            scaled_misfit = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1),
 953                                   np.ones((1, self.ne))) * (obs_data - pred_data)
 954        else:
 955            scaled_misfit = solve(self.scale_data, (obs_data - pred_data))
 956
 957        u, s, v = np.linalg.svd(S, full_matrices=False)
 958        if self.trunc_energy < 1:
 959            ti = (np.cumsum(s) / sum(s)) <= self.trunc_energy
 960            u, s, v = u[:, ti].copy(), s[ti].copy(), v[ti, :].copy()
 961
 962        ps_inv = np.diag([el_s ** (-1) for el_s in s])
 963        if len(self.scale_data.shape) == 1:
 964            X = np.dot(ps_inv, np.dot(u.T, np.dot(np.expand_dims(self.scale_data ** (-1), axis=1),
 965                                                  np.ones((1, self.ne))) * self.E))
 966        else:
 967            X = np.dot(ps_inv, np.dot(u.T, solve(self.scale_data, self.E)))
 968        Lam, z = np.linalg.eig(np.dot(X, X.T))
 969
 970        X2 = np.dot(u, np.dot(ps_inv.T, z))
 971
 972        X3_m = np.dot(S.T, X2)
 973        # X3_old = np.dot(X2, np.linalg.solve(np.eye(len(Lam)) + np.diag(Lam), X2.T))
 974        step_m = np.dot(np.dot(X3_m, inv(np.eye(len(Lam)) + np.diag(Lam))),
 975                        np.dot(X3_m.T, self.W))
 976
 977        if 'localization' in self.keys_da:
 978            if self.keys_da['localization'][1][0] == 'autoadaloc':
 979                loc_step_d = np.dot(np.linalg.pinv(self.aug_prior), self.localization.auto_ada_loc(self.aug_prior,
 980                                                                                                   np.dot(np.dot(S.T, X2),
 981                                                                                                          np.dot(inv(
 982                                                                                                              np.eye(len(Lam)) + np.diag(Lam)),
 983                                                                                                       np.dot(X2.T, scaled_misfit))),
 984                                                                                                   self.list_states,
 985                                                                                                   **{'prior_info': self.prior_info}))
 986                self.step = self.lam * (self.W - (step_m + loc_step_d))
 987        else:
 988            step_d = np.dot(np.linalg.inv(omega).T,  np.dot(np.dot(Y.T, X2),
 989                                                            np.dot(inv(np.eye(len(Lam)) + np.diag(Lam)),
 990                                                                   np.dot(X2.T, scaled_misfit))))
 991            self.step = self.lam * (self.W - (step_m + step_d))
 992
 993        self.W -= self.step
 994
 995        aug_state_upd = np.dot(self.aug_prior, (np.eye(
 996            self.ne) + self.W / np.sqrt(self.ne - 1)))
 997
 998        # Extract updated state variables from aug_update
 999        self.state = at.update_state(aug_state_upd, self.state, self.list_states)
1000
1001        self.state = at.limits(self.state, self.prior_info)
1002
1003    def check_convergence(self):
1004        """
1005        Check if GN-EnRML have converged based on evaluation of change sizes of objective function, state and damping
1006        parameter. Very similar to original function, but exit if there is no reduction in obj. function.
1007
1008        Parameters
1009        ----------
1010        prev_data_misfit : float
1011            Data misfit calculated from the previous update.
1012
1013        step : float
1014            Step size.
1015
1016        lam : float
1017            LM damping parameter.
1018
1019        Returns
1020        -------
1021        conv : bool
1022            Logic variable indicating if the algorithm has converged.
1023
1024        status : bool
1025            Indicates whether the objective function has reduced.
1026
1027        why_stop : dict
1028            Dictionary with keys corresponding to convergence criteria, with logical variables indicating
1029            which of them has been met.
1030
1031        Changelog
1032        ---------
1033        - ST 3/6-16
1034        - ST 6/6-16: Added LM damping param. check
1035        - KF 16/11-20: Modified for GN-EnRML
1036        - KF 10/3-21: Output whether the method reduced the objective function
1037        """
1038        # Prelude to calc. conv. check (everything done below is from calc_analysis)
1039        if hasattr(self, 'list_datatypes'):
1040            assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]]
1041            list_datatypes = self.list_datatypes
1042            cov_data = self.cov_data
1043            obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, assim_index,
1044                                                              list_datatypes)
1045            mean_preddata = np.mean(pred_data, 1)
1046        else:
1047            assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]]
1048            list_datatypes, _ = at.get_list_data_types(self.obs_data, assim_index)
1049            # cov_data = at.gen_covdata(self.datavar, assim_index, list_datatypes)
1050            obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, assim_index,
1051                                                              list_datatypes)
1052            # mean_preddata = np.mean(pred_data, 1)
1053
1054        success = False
1055
1056        # if inital conv. check, there are no prev_data_misfit
1057        if self.prev_data_misfit is None:
1058            self.data_misfit = np.mean(self.data_misfit)
1059            self.prev_data_misfit = self.data_misfit
1060            self.prev_data_misfit_std = self.data_misfit_std
1061            success = True
1062        # update the last mismatch, only if this was a reduction of the misfit
1063        if self.data_misfit < self.prev_data_misfit:
1064            self.prev_data_misfit = self.data_misfit
1065            self.prev_data_misfit_std = self.data_misfit_std
1066            success = True
1067        # if there was no reduction of the misfit, retain the old "valid" data misfit.
1068
1069        # Calc. std dev of data misfit (used to update lamda)
1070        # mat_obs = np.dot(obs_data_vector.reshape((len(obs_data_vector), 1)), np.ones((1, self.ne)))  # use the perturbed
1071        # data instead.
1072        mat_obs = self.real_obs_data
1073        if len(cov_data.shape) == 1:
1074            data_misfit = np.diag(np.dot((pred_data - mat_obs).T,
1075                                         np.dot(np.expand_dims(self.cov_data ** (-1), axis=1),
1076                                                np.ones((1, self.ne))) * (pred_data - mat_obs)))
1077        else:
1078            data_misfit = np.diag(np.dot((pred_data - mat_obs).T,
1079                                  solve(self.cov_data, (pred_data - mat_obs))))
1080        self.data_misfit = np.mean(data_misfit)
1081        self.data_misfit_std = np.std(data_misfit)
1082
1083        # # Calc. mean data misfit for convergence check, using the updated state variable
1084        # self.data_misfit = np.dot((mean_preddata - obs_data_vector).T,
1085        #                      solve(cov_data, (mean_preddata - obs_data_vector)))
1086        # if self.data_misfit > self.prev_data_misfit:
1087        #    print(f'\n\nMisfit increased from {self.prev_data_misfit:.1f} to {self.data_misfit:.1f}. Exiting')
1088        #    self.logger.info(f'\n\nMisfit increased from {self.prev_data_misfit:.1f} to {self.data_misfit:.1f}. Exiting')
1089
1090        # Convergence check: Relative step size of data misfit or state change less than tolerance
1091        if abs(1 - (self.data_misfit / self.prev_data_misfit)) < self.data_misfit_tol \
1092                or np.any(abs(np.mean(self.step, 1)) < self.step_tol) \
1093                or self.lam >= self.lam_max:
1094            # or self.data_misfit > self.prev_data_misfit:
1095            # Logical variables for conv. criteria
1096            why_stop = {'data_misfit_stop': 1 - (self.data_misfit / self.prev_data_misfit) < self.data_misfit_tol,
1097                        'data_misfit': self.data_misfit,
1098                        'prev_data_misfit': self.prev_data_misfit,
1099                        'step_size_stop': np.any(abs(np.mean(self.step, 1)) < self.step_tol),
1100                        'step_size': self.step,
1101                        'lambda': self.lam,
1102                        'lambda_stop': self.lam >= self.lam_max}
1103
1104            if self.data_misfit >= self.prev_data_misfit:
1105                success = False
1106                self.logger.info(f'Iterations have converged after {self.iteration} iterations. Objective function reduced '
1107                                 f'from {self.prior_data_misfit:0.1f} to {self.prev_data_misfit:0.1f}')
1108            else:
1109                self.logger.info(f'Iterations have converged after {self.iteration} iterations. Objective function reduced '
1110                                 f'from {self.prior_data_misfit:0.1f} to {self.data_misfit:0.1f}')
1111
1112            # Return conv = True, why_stop var.
1113            return True, success, why_stop
1114
1115        else:  # conv. not met
1116            # Logical variables for conv. criteria
1117            why_stop = {'data_misfit_stop': 1 - (self.data_misfit / self.prev_data_misfit) < self.data_misfit_tol,
1118                        'data_misfit': self.data_misfit,
1119                        'prev_data_misfit': self.prev_data_misfit,
1120                        'step_size': self.step,
1121                        'step_size_stop': np.any(abs(np.mean(self.step, 1)) < self.step_tol),
1122                        'lambda': self.lam,
1123                        'lambda_stop': self.lam >= self.lam_max}
1124
1125            ###############################################
1126            ##### update Lambda step-size values ##########
1127            ###############################################
1128            if self.data_misfit < self.prev_data_misfit and self.data_misfit_std < self.prev_data_misfit_std:
1129                # If reduction in mean data misfit, increase step length
1130                self.lam = self.lam + (self.lam_max - self.lam) * \
1131                    2 ** (-(self.iteration) / (self.gamma - 1))
1132                success = True
1133                self.current_state = deepcopy(self.state)
1134            elif self.data_misfit < self.prev_data_misfit and self.data_misfit_std >= self.prev_data_misfit_std:
1135                # Accept itaration, but keep lam the same
1136                success = True
1137                self.current_state = deepcopy(self.state)
1138            else:  # Reject iteration, and decrease step length
1139                self.lam = self.lam / self.gamma
1140                success = False
1141
1142            if success:
1143                self.logger.info(f'Successfull iteration number {self.iteration}! Objective function reduced from '
1144                                 f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}. New Lamba for next analysis: '
1145                                 f'{self.lam}')
1146            else:
1147                self.logger.info(f'Failed iteration number {self.iteration}! Objective function increased from '
1148                                 f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}. New Lamba for repeated analysis: '
1149                                 f'{self.lam}')
1150                # Reset data misfit to prev_data_misfit (because the current state is neglected)
1151                self.data_misfit = self.prev_data_misfit
1152                self.data_misfit_std = self.prev_data_misfit_std
1153
1154            return False, success, why_stop