From 9e3a8216a6639f41e3430773aa06f82ff390f797 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 9 Jul 2025 13:05:24 +0200 Subject: [PATCH 01/94] update to TrustRegion --- popt/loop/optimize.py | 74 ++++++------- popt/update_schemes/linesearch.py | 17 ++- popt/update_schemes/trust_region.py | 160 +++++++++++++--------------- 3 files changed, 128 insertions(+), 123 deletions(-) diff --git a/popt/loop/optimize.py b/popt/loop/optimize.py index 43f99c6..6fa80a4 100644 --- a/popt/loop/optimize.py +++ b/popt/loop/optimize.py @@ -166,48 +166,48 @@ def run_loop(self): self.save() # Check if max iterations was reached - if self.iteration > self.max_iter: + if self.iteration >= self.max_iter: self.optimize_result['message'] = 'Iterations stopped due to max iterations reached!' else: if not isinstance(self.msg, str): self.msg = '' self.optimize_result['message'] = self.msg - # Logging some info to screen - logger.info(' Optimization converged in %d iterations ', self.iteration-1) - logger.info(' Optimization converged with final obj_func = %.4f', - np.mean(self.optimize_result['fun'])) - logger.info(' Total number of function evaluations = %d', self.optimize_result['nfev']) - logger.info(' Total number of jacobi evaluations = %d', self.optimize_result['njev']) - if self.start_time is not None: - logger.info(' Total elapsed time = %.2f minutes', (time.perf_counter()-self.start_time)/60) - logger.info(' ============================================') - - # Test for convergence of outer epf loop - epf_not_converged = False - if self.epf: - if self.epf_iteration > self.epf['max_epf_iter']: # max epf_iterations set to 10 - logger.info(f' -----> EPF-EnOpt: maximum epf iterations reached') # print epf info - break - p = np.abs(previous_state-self.mean_state) / (np.abs(previous_state) + 1.0e-9) - conv_crit = self.epf['conv_crit'] - if np.any(p > conv_crit): - epf_not_converged = True - previous_state = self.mean_state - self.epf['r'] *= self.epf['r_factor'] # increase penalty factor - self.obj_func_tol *= self.epf['tol_factor'] # decrease tolerance - self.obj_func_values = self.fun(self.mean_state, **self.epf) - self.iteration = 0 - self.epf_iteration += 1 - optimize_result = ot.get_optimize_result(self) - ot.save_optimize_results(optimize_result) - self.nfev += 1 - self.iteration = +1 - r = self.epf['r'] - logger.info(f' -----> EPF-EnOpt: {self.epf_iteration}, {r} (outer iteration, penalty factor)') # print epf info - else: - logger.info(f' -----> EPF-EnOpt: converged, no variables changed more than {conv_crit*100} %') # print epf info - final_obj_no_penalty = str(round(float(self.fun(self.mean_state)),4)) - logger.info(f' -----> EPF-EnOpt: objective value without penalty = {final_obj_no_penalty}') # print epf info + # Logging some info to screen + logger.info(' Optimization converged in %d iterations ', self.iteration-1) + logger.info(' Optimization converged with final obj_func = %.4f', + np.mean(self.optimize_result['fun'])) + logger.info(' Total number of function evaluations = %d', self.optimize_result['nfev']) + logger.info(' Total number of jacobi evaluations = %d', self.optimize_result['njev']) + if self.start_time is not None: + logger.info(' Total elapsed time = %.2f minutes', (time.perf_counter()-self.start_time)/60) + logger.info(' ============================================') + + # Test for convergence of outer epf loop + epf_not_converged = False + if self.epf: + if self.epf_iteration > self.epf['max_epf_iter']: # max epf_iterations set to 10 + logger.info(f' -----> EPF-EnOpt: maximum epf iterations reached') # print epf info + break + p = np.abs(previous_state-self.mean_state) / (np.abs(previous_state) + 1.0e-9) + conv_crit = self.epf['conv_crit'] + if np.any(p > conv_crit): + epf_not_converged = True + previous_state = self.mean_state + self.epf['r'] *= self.epf['r_factor'] # increase penalty factor + self.obj_func_tol *= self.epf['tol_factor'] # decrease tolerance + self.obj_func_values = self.fun(self.mean_state, **self.epf) + self.iteration = 0 + self.epf_iteration += 1 + optimize_result = ot.get_optimize_result(self) + ot.save_optimize_results(optimize_result) + self.nfev += 1 + self.iteration = +1 + r = self.epf['r'] + logger.info(f' -----> EPF-EnOpt: {self.epf_iteration}, {r} (outer iteration, penalty factor)') # print epf info + else: + logger.info(f' -----> EPF-EnOpt: converged, no variables changed more than {conv_crit*100} %') # print epf info + final_obj_no_penalty = str(round(float(self.fun(self.mean_state)),4)) + logger.info(f' -----> EPF-EnOpt: objective value without penalty = {final_obj_no_penalty}') # print epf info def save(self): """ diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index b8fe92f..956cd9c 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -13,6 +13,7 @@ # Internal imports from popt.misc_tools import optim_tools as ot from popt.loop.optimize import Optimize +from popt.update_schemes import optimizers def LineSearch(fun, x, jac, method='GD', hess=None, args=(), bounds=None, callback=None, **options): ''' @@ -203,7 +204,7 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c self.saveit = options.get('saveit', True) # Check method - valid_methods = ['GD', 'BFGS', 'Newton'] + valid_methods = ['GD', 'BFGS', 'Newton', 'Adam'] if not self.method in valid_methods: raise ValueError(f"'{self.method}' is not a valid method. Valid methods are: {valid_methods}") @@ -373,6 +374,20 @@ def calc_update(self, iter_resamp=0): pk = - np.matmul(self.Hk_inv, self.jk) if self.method == 'Newton': pk = - np.matmul(la.inv(self.Hk), self.jk) + if self.method == 'Adam': + if self.iteration == 1: + pk = - self.jk + else: + optimizer = optimizers.Adam(1) + pk = - optimizer.apply_update(np.zeros_like(self.xk), self.jk, iter=self.iteration-1)[1] + optimizer.restore_parameters() + + # remove components that point out of the hybercube given by [lb,ub] + lb = np.array(self.bounds)[:, 0] + ub = np.array(self.bounds)[:, 1] + for i in range(self.xk.size): + if (self.xk[i] <= lb[i] and pk[i] < 0) or (self.xk[i] >= ub[i] and pk[i] > 0): + pk[i] = 0 # Set step_size step_size = self._set_step_size(pk) diff --git a/popt/update_schemes/trust_region.py b/popt/update_schemes/trust_region.py index 5272d75..9ed6e5e 100644 --- a/popt/update_schemes/trust_region.py +++ b/popt/update_schemes/trust_region.py @@ -11,8 +11,12 @@ from popt.misc_tools import optim_tools as ot from popt.loop.optimize import Optimize +# Impors from scipy +from scipy.optimize._trustregion_ncg import CGSteihaugSubproblem +from scipy.optimize._trustregion_exact import IterativeSubproblem -def TrustRegion(fun, x, jac, hess, args=(), bounds=None, callback=None, **options): + +def TrustRegion(fun, x, jac, hess, method='iterative', args=(), bounds=None, callback=None, **options): ''' Trust region optimization algorithm. @@ -29,6 +33,10 @@ def TrustRegion(fun, x, jac, hess, args=(), bounds=None, callback=None, **option hess : callable Hessian of objective function. The calling signature is `hess(x, *args)`. + + method : str, optional + Method to use for solving the trust-region subproblem. Options are 'iterative' or 'CG-Steihaug'. + Default is 'iterative'. args : tuple, optional Extra arguments passed to the objective function and its derivatives (Jacobian, Hessian). @@ -110,12 +118,12 @@ def TrustRegion(fun, x, jac, hess, args=(), bounds=None, callback=None, **option - nfev: number of function evaluations - njev: number of jacobian evaluations ''' - tr_obj = TrustRegionClass(fun, x, jac, hess, args, bounds, callback, **options) + tr_obj = TrustRegionClass(fun, x, jac, hess, method, args, bounds, callback, **options) return tr_obj.optimize_result class TrustRegionClass(Optimize): - def __init__(self, fun, x, jac, hess, args=(), bounds=None, callback=None, **options): + def __init__(self, fun, x, jac, hess, method='iterative', args=(), bounds=None, callback=None, **options): # Initialize the parent class super().__init__(**options) @@ -125,6 +133,7 @@ def __init__(self, fun, x, jac, hess, args=(), bounds=None, callback=None, **opt self.xk = x self.jacobian = jac self.hessian = hess + self.method = method self.args = args self.bounds = bounds self.options = options @@ -144,12 +153,18 @@ def __init__(self, fun, x, jac, hess, args=(), bounds=None, callback=None, **opt self.resample = options.get('resample', 3) self.saveit = options.get('saveit', True) self.rho_tol = options.get('rho_tol', 1e-6) - self.eta1 = options.get('eta1', 0.001) - self.eta2 = options.get('eta2', 0.1) - self.gam1 = options.get('gam1', 0.7) - self.gam2 = options.get('gam2', 1.5) + self.eta1 = options.get('eta1', 0.1) # reduce raduis if rho < 10% + self.eta2 = options.get('eta2', 0.5) # increase radius if rho > 50% + self.gam1 = options.get('gam1', 0.5) # reduce by 50% + self.gam2 = options.get('gam2', 1.5) # increase by 50% self.rho = 0.0 + # Check if method is valid + if self.method not in ['iterative', 'CG-Steihaug']: + self.method = 'iterative' + raise ValueError(f'Method {self.method} is not valid!. Method is set to "iterative"') + + if not self.restart: self.start_time = time.perf_counter() @@ -242,17 +257,66 @@ def _log(self, msg): if self.logger is not None: self.logger.info(msg) + def solve_subproblem(self, g, B, delta): + """ + Solve the trust region subproblem using the iterative method. + (A big thanks to copilot for the help with this implementation) + + Parameters: + g (numpy.ndarray): Gradient vector at the current point. + B (numpy.ndarray): Hessian matrix at the current point. + delta (float): Trust region radius. + + Returns: + pk (numpy.ndarray): Step direction. + pk_hits_boundary (bool): True if the step hits the boundary of the trust region. + """ + + # Define quadratic model + quad = lambda p: self.fk + np.dot(g,p) + np.dot(p,np.dot(B,p))/2 + + + if self.method == 'iterative': + subproblem = IterativeSubproblem( + x=self.xk, + fun=quad, + jac=lambda _: g, + hess=lambda _: B, + ) + pk, pk_hits_boundary = subproblem.solve(tr_radius=delta) + + elif self.method == 'CG-Steihaug': + subproblem = CGSteihaugSubproblem( + x=self.xk, + fun=quad, + jac=lambda _: g, + hess=lambda _: B, + ) + pk, pk_hits_boundary = subproblem.solve(trust_radius=delta) + + else: + raise ValueError(f"Method {self.method} is not valid!") + + return pk, pk_hits_boundary + + def calc_update(self, iter_resamp=0): # Initialize variables for this step success = True # Solve subproblem - self._log('Solving trust region subproblem using the CG-Steihaug method') - sk = self.solve_sub_problem_CG_Steihaug(self.jk, self.Hk, self.trust_radius) + self._log('Solving trust region subproblem') + sk, hits_boundary = self.solve_subproblem(self.jk, self.Hk, self.trust_radius) + + # truncate sk to respect bounds + if self.bounds is not None: + lb = np.array(self.bounds)[:, 0] + ub = np.array(self.bounds)[:, 1] + sk = np.clip(sk, lb - self.xk, ub - self.xk) # Calculate the actual function value - xk_new = ot.clip_state(self.xk + sk, self.bounds) + xk_new = self.xk + sk fun_new = self._fun(xk_new) # Calculate rho @@ -283,7 +347,7 @@ def calc_update(self, iter_resamp=0): # update the trust region radius delta_old = self.trust_radius - if self.rho >= self.eta2: + if (self.rho >= self.eta2) and hits_boundary: delta_new = min(self.gam2*delta_old, self.trust_radius_max) elif self.eta1 <= self.rho < self.eta2: delta_new = delta_old @@ -328,80 +392,6 @@ def calc_update(self, iter_resamp=0): return success - - def solve_sub_problem_CG_Steihaug(self, g, B, delta): - """ - Solve the trust region subproblem using Steihaug's Conjugate Gradient method. - (A big thanks to copilot for the help with this implementation) - - Parameters: - g (numpy.ndarray): Gradient vector at the current point. - B (numpy.ndarray): Hessian matrix at the current point. - delta (float): Trust region radius. - tol (float): Tolerance for convergence. - max_iter (int): Maximum number of iterations. - - Returns: - p (numpy.ndarray): Solution vector. - """ - z = np.zeros_like(g) - r = g - d = -g - - # Set same default tolerance as scipy - tol = min(0.5, la.norm(g)**2)*la.norm(g) - - if la.norm(g) <= tol: - return z - - # make quadratic model - mc = lambda s: self.fk + np.dot(g,s) + np.dot(s,np.dot(B,s))/2 - - while True: - dBd = np.dot(d, np.dot(B,d)) - - if dBd <= 0: - # Solve the quadratic equation: (p + tau*d)**2 = delta**2 - tau_lo, tau_hi = self.get_tau_at_delta(z, d, delta) - p_lo = z + tau_lo*d - p_hi = z + tau_hi*d - - if mc(p_lo) < mc(p_hi): - return p_lo - else: - return p_hi - - alpha = np.dot(r,r)/dBd - z_new = z + alpha*d - - if la.norm(z_new) >= delta: - # Solve the quadratic equation: (p + tau*d)**2 = delta**2, for tau > 0 - _ , tau = self.get_tau_at_delta(z, d, delta) - return z + tau * d - - r_new = r + alpha*np.dot(B,d) - - if la.norm(r_new) < tol: - return z_new - - beta = np.dot(r_new,r_new)/np.dot(r,r) - d = -r_new + beta*d - r = r_new - z = z_new - - - def get_tau_at_delta(self, p, d, delta): - """ - Solve the quadratic equation: (p + tau*d)**2 = delta**2, for tau > 0 - """ - a = np.dot(d,d) - b = 2*np.dot(p,d) - c = np.dot(p,p) - delta**2 - tau_lo = -b/(2*a) - np.sqrt(b**2 - 4*a*c)/(2*a) - tau_hi = -b/(2*a) + np.sqrt(b**2 - 4*a*c)/(2*a) - return tau_lo, tau_hi - - From 400c79a6c571d06bca3ca9a57f6b2fb8953c2abc Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 16 Jul 2025 14:13:39 +0200 Subject: [PATCH 02/94] some design changes to TrustRegion --- popt/update_schemes/trust_region.py | 99 +++++++++++++++++++---------- 1 file changed, 67 insertions(+), 32 deletions(-) diff --git a/popt/update_schemes/trust_region.py b/popt/update_schemes/trust_region.py index 9ed6e5e..cd616f9 100644 --- a/popt/update_schemes/trust_region.py +++ b/popt/update_schemes/trust_region.py @@ -44,9 +44,10 @@ def TrustRegion(fun, x, jac, hess, method='iterative', args=(), bounds=None, cal bounds : sequence, optional Bounds for variables. Each element of the sequence must be a tuple of two scalars, representing the lower and upper bounds for that variable. Use None for one of the bounds if there are no bounds. + Bounds are handle by clipping the state to the bounds before evaluating the objective function and its derivatives. callback: callable, optional - A callable called after each successful iteration. The class instance of LineSearch + A callable called after each successful iteration. The class instance is passed as the only argument to the callback function: callback(self) **options : keyword arguments, optional @@ -66,6 +67,9 @@ def TrustRegion(fun, x, jac, hess, method='iterative', args=(), bounds=None, cal Minimum trust-region radius. Optimization is terminated if trust_radius = trust_radius_min. Default is trust_radius/100. + trust_radius_cuts: int + Number of allowed trust-region radius reductions if a step is not successful. Default is 4. + rho_tol: float Tolerance for rho (ratio of actual to predicted reduction). Default is 1e-6. @@ -81,6 +85,13 @@ def TrustRegion(fun, x, jac, hess, method='iterative', args=(), bounds=None, cal eta2 = 0.1 \n gam1 = 0.7 \n gam2 = 1.5 \n + + saveit: bool + If True, save the optimization results to a file. Default is True. + + convergence_criteria: callable + A callable that takes the current optimization object as an argument and returns True if the optimization should stop. + It can be used to implement custom convergence criteria. Default is None. save_folder: str Name of folder to save the results to. Defaul is ./ (the current directory). @@ -94,9 +105,9 @@ def TrustRegion(fun, x, jac, hess, method='iterative', args=(), bounds=None, cal hess0: ndarray Hessian value of the initial control. - resample: int - Number of jacobian re-computations allowed if a line search fails. Default is 4. - (useful if jacobian is stochastic) + resample: bool + If True, resample the Jacobian and Hessian if a step is not successful. Default is False. + (Only makes sense if the Jacobian and Hessian are stochastic). savedata: list[str] Further specification of which class variables to save to the result files. @@ -144,13 +155,21 @@ def __init__(self, fun, x, jac, hess, method='iterative', args=(), bounds=None, else: self.callback = None + # Custom convergence criteria (callable) + convergence_criteria = options.get('convergence_criteria', None) + if callable(convergence_criteria): + self.convergence_criteria = self.convergence_criteria + else: + self.convergence_criteria = None + # Set options for trust-region radius - self.trust_radius = options.get('trust_radius', 1.0) - self.trust_radius_max = options.get('trust_radius_max', 10*self.trust_radius) - self.trust_radius_min = options.get('trust_radius_min', self.trust_radius/100) + self.trust_radius = options.get('trust_radius', 1.0) + self.trust_radius_max = options.get('trust_radius_max', 10*self.trust_radius) + self.trust_radius_min = options.get('trust_radius_min', self.trust_radius/100) + self.trust_radius_cuts = options.get('trust_radius_cuts', 4) # Set other options - self.resample = options.get('resample', 3) + self.resample = options.get('resample', False) self.saveit = options.get('saveit', True) self.rho_tol = options.get('rho_tol', 1e-6) self.eta1 = options.get('eta1', 0.1) # reduce raduis if rho < 10% @@ -222,15 +241,17 @@ def _hess(self, x): return h def update_results(self): - res = {'fun': self.fk, - 'x': self.xk, - 'jac': self.jk, - 'hess': self.Hk, - 'nfev': self.nfev, - 'njev': self.njev, - 'nit': self.iteration, - 'trust_radius': self.trust_radius, - 'save_folder': self.options.get('save_folder', './')} + res = { + 'fun': self.fk, + 'x': self.xk, + 'jac': self.jk, + 'hess': self.Hk, + 'nfev': self.nfev, + 'njev': self.njev, + 'nit': self.iteration, + 'trust_radius': self.trust_radius, + 'save_folder': self.options.get('save_folder', './') + } for a, arg in enumerate(self.args): res[f'args[{a}]'] = arg @@ -300,7 +321,7 @@ def solve_subproblem(self, g, B, delta): return pk, pk_hits_boundary - def calc_update(self, iter_resamp=0): + def calc_update(self, inner_iter=0): # Initialize variables for this step success = True @@ -321,7 +342,7 @@ def calc_update(self, iter_resamp=0): # Calculate rho actual_reduction = self.fk - fun_new - predicted_reduction = - np.dot(self.jk, sk) - 0.5*np.dot(sk, np.dot(self.Hk, sk)) + predicted_reduction = - np.dot(self.jk, sk) - np.dot(sk, np.dot(self.Hk, sk))/2 self.rho = actual_reduction/predicted_reduction if self.rho > self.rho_tol: @@ -355,12 +376,19 @@ def calc_update(self, iter_resamp=0): delta_new = self.gam1*delta_old # Log new trust-radius - self.trust_radius = delta_new + self.trust_radius = np.clip(delta_new, self.trust_radius_min, self.trust_radius_max) if not (delta_old == delta_new): self._log(f'Trust-radius updated: {delta_old:<10.4e} --> {delta_new:<10.4e}') + # Check for custom convergence + if callable(self.convergence_criteria): + if self.convergence_criteria(self): + self._log('Custom convergence criteria met. Stopping optimization.') + success = False + return success + # check for convergence - if (self.trust_radius < self.trust_radius_min) or (self.iteration==self.max_iter): + if self.iteration==self.max_iter: success = False else: # Calculate the jacobian and hessian @@ -371,21 +399,28 @@ def calc_update(self, iter_resamp=0): self.iteration += 1 else: - if iter_resamp < self.resample: - - iter_resamp += 1 + if inner_iter < self.trust_radius_cuts: + + # Log the failure + self._log(f'Step not successful: rho < {self.rho_tol:<10.4e}') + + # Reduce trust region radius to 75% of current value + self._log('Reducing trust-radius by 75%') + self.trust_radius = 0.25*self.trust_radius - # Calculate the jacobian and hessian - self._log('Resampling gradient and hessian') - self.jk = self._jac(self.xk) - self.Hk = self._hess(self.xk) + if self.trust_radius < self.trust_radius_min: + self._log(f'Trust radius {self.trust_radius} is below minimum {self.trust_radius_min}. Stopping optimization.') + success = False + return success - # Reduce trust region radius to 50% of current value - self._log('Reducing trust-radius by 50%') - self.trust_radius = 0.5*self.trust_radius + # Check for resampling of Jac and Hess + if self.resample: + self._log('Resampling gradient and hessian') + self.jk = self._jac(self.xk) + self.Hk = self._hess(self.xk) # Recursivly call function - success = self.calc_update(iter_resamp=iter_resamp) + success = self.calc_update(inner_iter=inner_iter+1) else: success = False From 98211a043dd0469c9c03eb89921693ef46ffdac0 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 5 Aug 2025 13:57:56 +0200 Subject: [PATCH 03/94] decoupled GenOpt from Ensemble --- popt/update_schemes/linesearch.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index 956cd9c..7bf1081 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -204,7 +204,7 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c self.saveit = options.get('saveit', True) # Check method - valid_methods = ['GD', 'BFGS', 'Newton', 'Adam'] + valid_methods = ['GD', 'BFGS', 'Newton'] if not self.method in valid_methods: raise ValueError(f"'{self.method}' is not a valid method. Valid methods are: {valid_methods}") @@ -374,13 +374,6 @@ def calc_update(self, iter_resamp=0): pk = - np.matmul(self.Hk_inv, self.jk) if self.method == 'Newton': pk = - np.matmul(la.inv(self.Hk), self.jk) - if self.method == 'Adam': - if self.iteration == 1: - pk = - self.jk - else: - optimizer = optimizers.Adam(1) - pk = - optimizer.apply_update(np.zeros_like(self.xk), self.jk, iter=self.iteration-1)[1] - optimizer.restore_parameters() # remove components that point out of the hybercube given by [lb,ub] lb = np.array(self.bounds)[:, 0] From 6152b7f6e0822fc91c79393c76a25e1589d27d5d Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 5 Aug 2025 14:00:39 +0200 Subject: [PATCH 04/94] decoupled GenOpt from Ensemble --- popt/loop/ensemble.py | 7 -- popt/loop/{base.py => ensemble_base.py} | 108 +++++++++++++++--------- popt/loop/generalized_ensemble.py | 31 ++++--- 3 files changed, 82 insertions(+), 64 deletions(-) rename popt/loop/{base.py => ensemble_base.py} (57%) diff --git a/popt/loop/ensemble.py b/popt/loop/ensemble.py index 77ff998..507ad8c 100644 --- a/popt/loop/ensemble.py +++ b/popt/loop/ensemble.py @@ -10,7 +10,6 @@ from popt.misc_tools import optim_tools as ot from pipt.misc_tools import analysis_tools as at from ensemble.ensemble import Ensemble as PETEnsemble -from popt.loop.extensions import GenOptExtension class Ensemble(PETEnsemble): @@ -137,12 +136,6 @@ def __set__variable(var_name=None, defalut=None): self.bias_weights = np.ones(self.num_samples) / self.num_samples # initialize with equal weights self.bias_points = None # this is the points used to estimate the bias correction - # Setup GenOpt - self.genopt = GenOptExtension(self.get_state(), - self.get_cov(), - func=self.function, - ne=self.num_samples) - def get_state(self): """ Returns diff --git a/popt/loop/base.py b/popt/loop/ensemble_base.py similarity index 57% rename from popt/loop/base.py rename to popt/loop/ensemble_base.py index 33d4497..2500fc2 100644 --- a/popt/loop/base.py +++ b/popt/loop/ensemble_base.py @@ -9,63 +9,81 @@ from popt.misc_tools import optim_tools as ot from pipt.misc_tools import analysis_tools as at from ensemble.ensemble import Ensemble as PETEnsemble +from simulator.simple_models import noSimulation -class EnsembleOptimizationBase(PETEnsemble): +class EnsembleOptimizationBaseClass(PETEnsemble): ''' Base class for the popt ensemble ''' - def __init__(self, kwargs_ens, sim, obj_func): + def __init__(self, options, simulator, objective): ''' Parameters ---------- - kwargs_ens : dict + options : dict Options for the ensemble class - sim : callable - The forward simulator (e.g. flow) + simulator : callable + The forward simulator (e.g. flow). If None, no simulation is performed. - obj_func : callable + objective : callable The objective function (e.g. npv) ''' + if simulator is None: + sim = noSimulation() + else: + sim = simulator # Initialize PETEnsemble - super().__init__(kwargs_ens, sim) - - self.save_prediction = kwargs_ens.get('save_prediction', None) - self.num_models = kwargs_ens.get('num_models', 1) - self.transform = kwargs_ens.get('transform', False) - self.num_samples = self.ne + super().__init__(options, sim) - # Get bounds and varaince - self.upper_bound = [] - self.lower_bound = [] + # Unpack some options + self.save_prediction = options.get('save_prediction', None) + self.num_models = options.get('num_models', 1) + self.transform = options.get('transform', False) + self.num_samples = self.ne + + # Define some variables + self.lb = [] + self.ub = [] self.bounds = [] self.cov = np.array([]) - for name in self.prior_info.keys(): - self.state[name] = np.asarray(self.prior_info[name]['mean']) - num_state_var = len(self.state[name]) - value_cov = self.prior_info[name]['variance'] * np.ones((num_state_var,)) - if 'limits' in self.prior_info[name].keys(): - lb = self.prior_info[name]['limits'][0] - ub = self.prior_info[name]['limits'][1] - self.lower_bound.append(lb) - self.upper_bound.append(ub) + + # Get bounds and varaince, and initialize state + for key in self.prior_info.keys(): + variable = self.prior_info[key] + + # mean + self.state[key] = np.asarray(variable['mean']) + + # Covariance + dim = self.state[key].size + cov = variable['variance']*np.ones(dim) + + if 'limits' in variable.keys(): + lb, ub = variable['limits'] + self.lb(lb) + self.ub(ub) + + # transform cov to [0, 1] if transform is True if self.transform: - value_cov = value_cov / (ub - lb)**2 - np.clip(value_cov, 0, 1, out=value_cov) - self.bounds += num_state_var*[(0, 1)] + cov = np.clip(cov/(ub - lb)**2, 0, 1, out=cov) + self.bounds += dim*[(0, 1)] else: - self.bounds += num_state_var*[(lb, ub)] - self.cov = np.append(self.cov, value_cov) + self.bounds += dim*[(lb, ub)] else: - self.bounds += num_state_var*[(None, None)] + self.bounds += dim*[(None, None)] + + # Add to covariance + self.cov = np.append(self.cov, cov) - - self._scale_state() + # Make cov full covariance matrix self.cov = np.diag(self.cov) + # Scale the state to [0, 1] if transform is True + self._scale_state() + # Set objective function (callable) - self.obj_func = obj_func + self.obj_func = objective # Objective function values self.state_func_values = None @@ -78,8 +96,13 @@ def get_state(self): x : numpy.ndarray Control vector as ndarray, shape (number of controls, number of perturbations) """ - x = ot.aug_optim_state(self.state, list(self.state.keys())) - return x + return ot.aug_optim_state(self.state, list(self.state.keys())) + + def vec_to_state(self, x): + """ + Converts a control vector to the internal state representation. + """ + return ot.update_optim_state(x, self.state, list(self.state.keys())) def get_bounds(self): """ @@ -112,7 +135,10 @@ def function(self, x, *args): else: self.ne = x.shape[1] - self.state = ot.update_optim_state(x, self.state, list(self.state.keys())) # go from nparray to dict + # convert x to state + self.state = self.vec_to_state(x) # go from nparray to dict + + # run the simulation self._invert_scale_state() # ensure that state is in [lb,ub] run_success = self.calc_prediction(save_prediction=self.save_prediction) # calculate flow data self._scale_state() # scale back to [0, 1] @@ -147,17 +173,17 @@ def _scale_state(self): """ Transform the internal state from [lb, ub] to [0, 1] """ - if self.transform and (self.upper_bound and self.lower_bound): + if self.transform and (self.lb and self.ub): for i, key in enumerate(self.state): - self.state[key] = (self.state[key] - self.lower_bound[i])/(self.upper_bound[i] - self.lower_bound[i]) + self.state[key] = (self.state[key] - self.lb[i])/(self.ub[i] - self.lb[i]) np.clip(self.state[key], 0, 1, out=self.state[key]) def _invert_scale_state(self): """ Transform the internal state from [0, 1] to [lb, ub] """ - if self.transform and (self.upper_bound and self.lower_bound): + if self.transform and (self.lb and self.ub): for i, key in enumerate(self.state): if self.transform: - self.state[key] = self.lower_bound[i] + self.state[key]*(self.upper_bound[i] - self.lower_bound[i]) - np.clip(self.state[key], self.lower_bound[i], self.upper_bound[i], out=self.state[key]) \ No newline at end of file + self.state[key] = self.lb[i] + self.state[key]*(self.ub[i] - self.lb[i]) + np.clip(self.state[key], self.lb[i], self.ub[i], out=self.state[key]) \ No newline at end of file diff --git a/popt/loop/generalized_ensemble.py b/popt/loop/generalized_ensemble.py index cacbe99..81809c2 100644 --- a/popt/loop/generalized_ensemble.py +++ b/popt/loop/generalized_ensemble.py @@ -10,33 +10,32 @@ # Internal imports from popt.misc_tools import optim_tools as ot from pipt.misc_tools import analysis_tools as at -from popt.loop.base import EnsembleOptimizationBase +from popt.loop.ensemble_base import EnsembleOptimizationBaseClass -class GeneralizedEnsemble(EnsembleOptimizationBase): +class GeneralizedEnsemble(EnsembleOptimizationBaseClass): - def __init__(self, kwargs_ens, sim, obj_func): + def __init__(self, options, simulator, objective): ''' Parameters ---------- - kwargs_ens : dict + options : dict Options for the ensemble class - sim : callable - The forward simulator (e.g. flow) + simulator : callable + The forward simulator (e.g. flow). If None, no simulation is performed. - obj_func : callable + objective : callable The objective function (e.g. npv) ''' - super().__init__(kwargs_ens, sim, obj_func) - - self.dim = self.get_state().size + super().__init__(options, simulator, objective) # construct corr matrix std = np.sqrt(np.diag(self.cov)) self.corr = self.cov/np.outer(std, std) + self.dim = std # choose marginal - marginal = kwargs_ens.get('marginal', 'Beta') + marginal = options.get('marginal', 'BetaMC') if marginal in ['Beta', 'BetaMC', 'Logistic', 'TruncGaussian', 'Gaussian']: @@ -45,7 +44,7 @@ def __init__(self, kwargs_ens, sim, obj_func): if marginal == 'Beta': self.margs = Beta() - self.theta = kwargs_ens.get('theta', np.array([[20.0, 20.0] for _ in range(self.dim)])) + self.theta = options.get('theta', np.array([[20.0, 20.0] for _ in range(self.dim)])) self.eps = self.var2eps() self.grad_scale = 1/(2*self.eps) self.hess_scale = 1/(4*self.eps**2) @@ -56,20 +55,20 @@ def __init__(self, kwargs_ens, sim, obj_func): var = np.diag(self.cov) self.margs = BetaMC(lb, ub, 0.1*np.sqrt(var[0])) default_theta = np.array([var_to_concentration(state[i], var[i], lb[i], ub[i]) for i in range(self.dim)]) - self.theta = kwargs_ens.get('theta', default_theta) + self.theta = options.get('theta', default_theta) elif marginal == 'Logistic': self.margs = Logistic() - self.theta = kwargs_ens.get('theta', self.margs.var_to_scale(np.diag(self.cov))) + self.theta = options.get('theta', self.margs.var_to_scale(np.diag(self.cov))) elif marginal == 'TruncGaussian': lb, ub = np.array(self.bounds).T self.margs = TruncGaussian(lb,ub) - self.theta = kwargs_ens.get('theta', np.sqrt(np.diag(self.cov))) + self.theta = options.get('theta', np.sqrt(np.diag(self.cov))) elif marginal == 'Gaussian': self.margs = Gaussian() - self.theta = kwargs_ens.get('theta', np.sqrt(np.diag(self.cov))) + self.theta = options.get('theta', np.sqrt(np.diag(self.cov))) def get_theta(self): return self.theta From 8ff8d5f20eab6cd1a672cd3073d65e8e7be9aaa6 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 7 Aug 2025 10:19:10 +0200 Subject: [PATCH 05/94] Cleaned up code duplication and renamed some stuff --- popt/loop/ensemble_base.py | 104 ++++++---- .../{ensemble.py => ensemble_gaussian.py} | 188 +----------------- ...ed_ensemble.py => ensemble_generalized.py} | 0 popt/loop/extensions.py | 1 + 4 files changed, 71 insertions(+), 222 deletions(-) rename popt/loop/{ensemble.py => ensemble_gaussian.py} (70%) rename popt/loop/{generalized_ensemble.py => ensemble_generalized.py} (100%) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index 2500fc2..951a2be 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -8,10 +8,10 @@ # Internal imports from popt.misc_tools import optim_tools as ot from pipt.misc_tools import analysis_tools as at -from ensemble.ensemble import Ensemble as PETEnsemble +from ensemble.ensemble import Ensemble as SupEnsemble from simulator.simple_models import noSimulation -class EnsembleOptimizationBaseClass(PETEnsemble): +class EnsembleOptimizationBaseClass(SupEnsemble): ''' Base class for the popt ensemble ''' @@ -33,7 +33,7 @@ def __init__(self, options, simulator, objective): else: sim = simulator - # Initialize PETEnsemble + # Initialize the PET Ensemble super().__init__(options, sim) # Unpack some options @@ -41,32 +41,44 @@ def __init__(self, options, simulator, objective): self.num_models = options.get('num_models', 1) self.transform = options.get('transform', False) self.num_samples = self.ne - - # Define some variables + + # Set objective function (callable) + self.obj_func = objective + self.state_func_values = None + self.ens_func_values = None + + # Initialize prior + self._initialize_state_info() # Initialize cov, bounds, and state + self._scale_state() # Scale self.state to [0, 1] if transform is True + + def _initialize_state_info(self): + ''' + Initialize covariance and bounds based on prior information. + ''' + self.cov = np.array([]) self.lb = [] self.ub = [] self.bounds = [] - self.cov = np.array([]) - - # Get bounds and varaince, and initialize state + for key in self.prior_info.keys(): variable = self.prior_info[key] - + # mean self.state[key] = np.asarray(variable['mean']) # Covariance dim = self.state[key].size - cov = variable['variance']*np.ones(dim) - + var = variable['variance']*np.ones(dim) + if 'limits' in variable.keys(): lb, ub = variable['limits'] - self.lb(lb) - self.ub(ub) - - # transform cov to [0, 1] if transform is True + self.lb.append(lb) + self.ub.append(ub) + + # transform var to [0, 1] if transform is True if self.transform: - cov = np.clip(cov/(ub - lb)**2, 0, 1, out=cov) + var = var/(ub - lb)**2 + var = np.clip(var, 0, 1, out=var) self.bounds += dim*[(0, 1)] else: self.bounds += dim*[(lb, ub)] @@ -74,20 +86,11 @@ def __init__(self, options, simulator, objective): self.bounds += dim*[(None, None)] # Add to covariance - self.cov = np.append(self.cov, cov) - + self.cov = np.append(self.cov, var) + self.dim = self.cov.shape[0] + # Make cov full covariance matrix self.cov = np.diag(self.cov) - - # Scale the state to [0, 1] if transform is True - self._scale_state() - - # Set objective function (callable) - self.obj_func = objective - - # Objective function values - self.state_func_values = None - self.ens_func_values = None def get_state(self): """ @@ -98,6 +101,15 @@ def get_state(self): """ return ot.aug_optim_state(self.state, list(self.state.keys())) + def get_cov(self): + """ + Returns + ------- + cov : numpy.ndarray + Covariance matrix, shape (number of controls, number of controls) + """ + return self.cov + def vec_to_state(self, x): """ Converts a control vector to the internal state representation. @@ -114,7 +126,7 @@ def get_bounds(self): return self.bounds - def function(self, x, *args): + def function(self, x, *args, **kwargs): """ This is the main function called during optimization. @@ -130,29 +142,41 @@ def function(self, x, *args): """ self._aux_input() - if len(x.shape) == 1: - self.ne = self.num_models - else: - self.ne = x.shape[1] + # check for ensmble + if len(x.shape) == 1: self.ne = self.num_models + else: self.ne = x.shape[1] - # convert x to state - self.state = self.vec_to_state(x) # go from nparray to dict + # convert x (nparray) to state (dict) + self.state = self.vec_to_state(x) # run the simulation self._invert_scale_state() # ensure that state is in [lb,ub] + self._set_multilevel_state(self.state, x) # set multilevel state if applicable run_success = self.calc_prediction(save_prediction=self.save_prediction) # calculate flow data + self._set_multilevel_state(self.state, x) # For some reason this has to be done again after calc_prediction self._scale_state() # scale back to [0, 1] + + # Evaluate the objective function if run_success: - func_values = self.obj_func(self.pred_data, self.sim.input_dict, self.sim.true_order) + func_values = self.obj_func( + self.pred_data, + input_dict=self.sim.input_dict, + true_order=self.sim.true_order, + **kwargs + ) else: func_values = np.inf # the simulations have crashed - if len(x.shape) == 1: - self.state_func_values = func_values - else: - self.ens_func_values = func_values + if len(x.shape) == 1: self.state_func_values = func_values + else: self.ens_func_values = func_values return func_values + + def _set_multilevel_state(self, state, x): + if 'multilevel' in self.keys_en.keys() and len(x.shape) > 1: + en_size = ot.get_list_element(self.keys_en['multilevel'], 'en_size') + self.state = ot.toggle_ml_state(self.state, en_size) + def _aux_input(self): """ diff --git a/popt/loop/ensemble.py b/popt/loop/ensemble_gaussian.py similarity index 70% rename from popt/loop/ensemble.py rename to popt/loop/ensemble_gaussian.py index 62f7389..f4bf832 100644 --- a/popt/loop/ensemble.py +++ b/popt/loop/ensemble_gaussian.py @@ -5,14 +5,13 @@ from copy import deepcopy - # Internal imports from popt.misc_tools import optim_tools as ot from pipt.misc_tools import analysis_tools as at -from ensemble.ensemble import Ensemble as PETEnsemble +from popt.loop.ensemble_base import EnsembleOptimizationBaseClass -class Ensemble(PETEnsemble): +class GaussianEnsemble(EnsembleOptimizationBaseClass): """ Class to store control states and evaluate objective functions. @@ -41,7 +40,7 @@ class Ensemble(PETEnsemble): """ - def __init__(self, keys_en, sim, obj_func): + def __init__(self, options, simulator, objective): """ Parameters ---------- @@ -63,57 +62,7 @@ def __init__(self, keys_en, sim, obj_func): """ # Initialize PETEnsemble - super(Ensemble, self).__init__(keys_en, sim) - - def __set__variable(var_name=None, defalut=None): - if var_name in keys_en: - return keys_en[var_name] - else: - return defalut - - # Set number of models (default 1) - self.num_models = __set__variable('num_models', 1) - - # Set transform flag (defalult True) - self.transform = __set__variable('transform', True) - - # Number of samples to compute gradient - self.num_samples = self.ne - - # Save pred data? - self.save_prediction = __set__variable('save_prediction', None) - - # We need the limits to convert between [0, 1] and [lb, ub], - # and we need the bounds as list of (min, max) pairs - # Also set the state and covarianve equal to the values provided in the input. - self.upper_bound = [] - self.lower_bound = [] - self.bounds = [] - self.cov = np.array([]) - for name in self.prior_info.keys(): - self.state[name] = np.asarray(self.prior_info[name]['mean']) - num_state_var = len(self.state[name]) - value_cov = self.prior_info[name]['variance'] * np.ones((num_state_var,)) - if 'limits' in self.prior_info[name].keys(): - lb = self.prior_info[name]['limits'][0] - ub = self.prior_info[name]['limits'][1] - self.lower_bound.append(lb) - self.upper_bound.append(ub) - if self.transform: - value_cov = value_cov / (ub - lb)**2 - np.clip(value_cov, 0, 1, out=value_cov) - self.bounds += num_state_var*[(0, 1)] - else: - self.bounds += num_state_var*[(lb, ub)] - else: - self.bounds += num_state_var*[(None, None)] - self.cov = np.append(self.cov, value_cov) - - self._scale_state() - self.cov = np.diag(self.cov) - - # Set objective function (callable) - self.obj_func = obj_func + super().__init__(options, simulator, objective) # Objective function values self.state_func_values = None @@ -135,36 +84,6 @@ def __set__variable(var_name=None, defalut=None): self.bias_factors = None # this is J(x_j,m_j)/J(x_j,m) self.bias_weights = np.ones(self.num_samples) / self.num_samples # initialize with equal weights self.bias_points = None # this is the points used to estimate the bias correction - - def get_state(self): - """ - Returns - ------- - x : numpy.ndarray - Control vector as ndarray, shape (number of controls, number of perturbations) - """ - x = ot.aug_optim_state(self.state, list(self.state.keys())) - return x - - def get_cov(self): - """ - Returns - ------- - cov : numpy.ndarray - Covariance matrix, shape (number of controls, number of controls) - """ - - return self.cov - - def get_bounds(self): - """ - Returns - ------- - bounds : list - (min, max) pairs for each element in x. None is used to specify no bound. - """ - - return self.bounds def get_final_state(self, return_dict=False): """ @@ -186,56 +105,6 @@ def get_final_state(self, return_dict=False): x = self.get_state() return x - def function(self, x, *args, **kwargs): - """ - This is the main function called during optimization. - - Parameters - ---------- - x : ndarray - Control vector, shape (number of controls, number of perturbations) - - Returns - ------- - obj_func_values : numpy.ndarray - Objective function values, shape (number of perturbations, ) - """ - self._aux_input() - - if len(x.shape) == 1: - self.ne = self.num_models - else: - self.ne = x.shape[1] - - self.state = ot.update_optim_state(x, self.state, list(self.state.keys())) # go from nparray to dict - self._invert_scale_state() # ensure that state is in [lb,ub] - - # Here we need to account for the possibility of having a multilevel ensemble and make a list of levels - if 'multilevel' in self.keys_en.keys() and len(x.shape) > 1: - en_size = ot.get_list_element(self.keys_en['multilevel'], 'en_size') - self.state = ot.toggle_ml_state(self.state, en_size) - - run_success = self.calc_prediction() # calculate flow data - - # Here we need to account for the possibility of having a multilevel ensemble and remove list of levels - if 'multilevel' in self.keys_en.keys() and len(x.shape) > 1: - en_size = ot.get_list_element(self.keys_en['multilevel'], 'en_size') - self.state = ot.toggle_ml_state(self.state, en_size) - - self._scale_state() # scale back to [0, 1] - if run_success: - func_values = self.obj_func(self.pred_data, input_dict=self.sim.input_dict, - true_order=self.sim.true_order, **kwargs) - else: - func_values = np.inf # the simulations have crashed - - if len(x.shape) == 1: - self.state_func_values = func_values - else: - self.ens_func_values = func_values - - return func_values - def gradient(self, x, *args, **kwargs): r""" Calculate the preconditioned gradient associated with ensemble, defined as: @@ -393,17 +262,6 @@ def hessian(self, x=None, *args): hessian = level_hessian[0] return hessian - ''' - def genopt_gradient(self, x, *args): - self.genopt.update_distribution(*args) - gradient = self.genopt.ensemble_gradient(func=self.function, - x=x, - ne=self.num_samples) - return gradient - - def genopt_mutation_gradient(self, x=None, *args, **kwargs): - return self.genopt.ensemble_mutation_gradient(return_ensembles=kwargs['return_ensembles']) - ''' def calc_ensemble_weights(self, x, *args, **kwargs): r""" @@ -527,49 +385,15 @@ def _gen_state_ensemble(self): cov = cov_blocks[i] temp_state_en = np.random.multivariate_normal(mean, cov, self.ne).transpose() shifted_ensemble = np.array([mean]).T + temp_state_en - np.array([np.mean(temp_state_en, 1)]).T - if self.upper_bound and self.lower_bound: + if self.lb and self.ub: if self.transform: np.clip(shifted_ensemble, 0, 1, out=shifted_ensemble) else: - np.clip(shifted_ensemble, self.lower_bound[i], self.upper_bound[i], out=shifted_ensemble) + np.clip(shifted_ensemble, self.lb[i], self.ub[i], out=shifted_ensemble) state_en[statename] = shifted_ensemble return state_en - def _aux_input(self): - """ - Set the auxiliary input used for multiple geological realizations - """ - - nr = 1 # nr is the ratio of samples over models - if self.num_models > 1: - if np.remainder(self.num_samples, self.num_models) == 0: - nr = int(self.num_samples / self.num_models) - self.aux_input = list(np.repeat(np.arange(self.num_models), nr)) - else: - print('num_samples must be a multiplum of num_models!') - sys.exit(0) - return nr - - def _scale_state(self): - """ - Transform the internal state from [lb, ub] to [0, 1] - """ - if self.transform and (self.upper_bound and self.lower_bound): - for i, key in enumerate(self.state): - self.state[key] = (self.state[key] - self.lower_bound[i])/(self.upper_bound[i] - self.lower_bound[i]) - np.clip(self.state[key], 0, 1, out=self.state[key]) - - def _invert_scale_state(self): - """ - Transform the internal state from [0, 1] to [lb, ub] - """ - if self.transform and (self.upper_bound and self.lower_bound): - for i, key in enumerate(self.state): - if self.transform: - self.state[key] = self.lower_bound[i] + self.state[key]*(self.upper_bound[i] - self.lower_bound[i]) - np.clip(self.state[key], self.lower_bound[i], self.upper_bound[i], out=self.state[key]) - def _bias_correction(self, state): """ Calculate bias correction. Currently, the bias correction is a constant independent of the state diff --git a/popt/loop/generalized_ensemble.py b/popt/loop/ensemble_generalized.py similarity index 100% rename from popt/loop/generalized_ensemble.py rename to popt/loop/ensemble_generalized.py diff --git a/popt/loop/extensions.py b/popt/loop/extensions.py index 5bcba7d..2dc7536 100644 --- a/popt/loop/extensions.py +++ b/popt/loop/extensions.py @@ -6,6 +6,7 @@ # Internal imports from popt.misc_tools import optim_tools as ot +# NB! THIS FILE IS NOT USED ANYMORE __all__ = ['GenOptExtension'] From 235948d0e4b81002c2fc8f92a21692f43f38b46e Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 7 Aug 2025 13:33:14 +0200 Subject: [PATCH 06/94] comments --- popt/loop/__init__.py | 2 +- popt/loop/ensemble_base.py | 2 ++ popt/loop/ensemble_gaussian.py | 1 + popt/loop/ensemble_generalized.py | 4 +++- 4 files changed, 7 insertions(+), 2 deletions(-) diff --git a/popt/loop/__init__.py b/popt/loop/__init__.py index ead116b..5ded178 100644 --- a/popt/loop/__init__.py +++ b/popt/loop/__init__.py @@ -1 +1 @@ -"""Main loop for running optimization.""" +"""Main loop for running optimization.""" \ No newline at end of file diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index 951a2be..e5363a2 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -11,6 +11,8 @@ from ensemble.ensemble import Ensemble as SupEnsemble from simulator.simple_models import noSimulation +__all__ = ['EnsembleOptimizationBaseClass'] + class EnsembleOptimizationBaseClass(SupEnsemble): ''' Base class for the popt ensemble diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index f4bf832..2d334ca 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -10,6 +10,7 @@ from pipt.misc_tools import analysis_tools as at from popt.loop.ensemble_base import EnsembleOptimizationBaseClass +__all__ = ['GaussianEnsemble'] class GaussianEnsemble(EnsembleOptimizationBaseClass): """ diff --git a/popt/loop/ensemble_generalized.py b/popt/loop/ensemble_generalized.py index 81809c2..c786627 100644 --- a/popt/loop/ensemble_generalized.py +++ b/popt/loop/ensemble_generalized.py @@ -12,6 +12,8 @@ from pipt.misc_tools import analysis_tools as at from popt.loop.ensemble_base import EnsembleOptimizationBaseClass +__all__ = ['GeneralizedEnsemble'] + class GeneralizedEnsemble(EnsembleOptimizationBaseClass): def __init__(self, options, simulator, objective): @@ -32,7 +34,7 @@ def __init__(self, options, simulator, objective): # construct corr matrix std = np.sqrt(np.diag(self.cov)) self.corr = self.cov/np.outer(std, std) - self.dim = std + self.dim = std.size # choose marginal marginal = options.get('marginal', 'BetaMC') From 1d0988f0dfce6d324fa675a712b104101552c579 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 20 Aug 2025 13:01:42 +0200 Subject: [PATCH 07/94] improved and cleaned up LineSearch --- popt/update_schemes/line_search_step.py | 295 +++++++++++++ popt/update_schemes/linesearch.py | 549 +++++++++++------------- 2 files changed, 534 insertions(+), 310 deletions(-) create mode 100644 popt/update_schemes/line_search_step.py diff --git a/popt/update_schemes/line_search_step.py b/popt/update_schemes/line_search_step.py new file mode 100644 index 0000000..5207940 --- /dev/null +++ b/popt/update_schemes/line_search_step.py @@ -0,0 +1,295 @@ +# This is a an implementation of the Line Search Algorithm (Alg. 3.5) in Numerical Optimization from Nocedal 2006. + +import numpy as np +from functools import cache +from scipy.optimize._linesearch import _quadmin, _cubicmin + + +def line_search(step_size, xk, pk, fun, jac, fk=None, jk=None, **kwargs): + ''' + Line search algorithm to find step size alpha that satisfies the Wolfe conditions. + + Parameters + ---------- + step_size : float + Initial step size to start the line search. + + xk : ndarray + Current point in the optimization process. + + pk : ndarray + Search direction. + + fun : callable + Objective function + + jac : callable + Gradient of the objective function + + fk : float, optional + Function value at xk. If None, it will be computed. + + jk : ndarray, optional + Gradient at xk. If None, it will be computed. + + **kwargs : dict + Additional parameters for the line search, such as: + - amax : float, maximum step size (default: 1000) + - maxiter : int, maximum number of iterations (default: 10) + - c1 : float, sufficient decrease condition (default: 1e-4) + - c2 : float, curvature condition (default: 0.9) + + Returns + ------- + alpha : float + Step size that satisfies the Wolfe conditions. + + fval : float + Function value at the new point xk + step_size*pk. + + jval : ndarray + Gradient at the new point xk + step_size*pk. + + nfev : int + Number of function evaluations. + + njev : int + Number of gradient evaluations. + ''' + + global ls_nfev + global ls_njev + ls_nfev = 0 + ls_njev = 0 + + # Unpack some kwargs + amax = kwargs.get('amax', 1000) + maxiter = kwargs.get('maxiter', 10) + c1 = kwargs.get('c1', 1e-4) + c2 = kwargs.get('c2', 0.9) + + # assertions + assert step_size <= amax, "Initial step size must be less than or equal to amax." + + # Define phi and derivative of phi + @cache + def phi(alpha): + global ls_nfev + if (alpha == 0): + if (fk is None): + phi.fun_val = fun(xk) + ls_nfev += 1 + else: + phi.fun_val = fk + else: + phi.fun_val = fun(xk + alpha*pk) + ls_nfev += 1 + return phi.fun_val + + @cache + def dphi(alpha): + global ls_njev + if (alpha == 0): + if (jk is None): + dphi.jac_val = jac(xk) + ls_njev += 1 + else: + dphi.jac_val = jk + else: + dphi.jac_val = jac(xk + alpha*pk) + ls_njev += 1 + return np.dot(dphi.jac_val, pk) + + # Define initial values of phi and dphi + phi_0 = phi(0) + dphi_0 = dphi(0) + + # Start loop + a = [0, step_size] + for i in range(1, maxiter+1): + # Evaluate phi(ai) + phi_i = phi(a[i]) + + # Check for sufficient decrease + if (phi_i > phi_0 + c1*a[i]*dphi_0) or (phi_i >= phi(a[i-1]) and i>0): + # Call zoom function + step_size = zoom(a[i-1], a[i], phi, dphi, phi_0, dphi_0, maxiter+1-i, c1, c2) + return step_size, phi.fun_val, dphi.jac_val, ls_nfev, ls_njev + + # Evaluate dphi(ai) + dphi_i = dphi(a[i]) + + # Check curvature condition + if abs(dphi_i) <= -c2*dphi_0: + step_size = a[i] + return step_size, phi.fun_val, dphi.jac_val, ls_nfev, ls_njev + + # Check for posetive derivative + if dphi_i >= 0: + # Call zoom function + step_size = zoom(a[i], a[i-1], phi, dphi, phi_0, dphi_0, maxiter+1-i, c1, c2) + return step_size, phi.fun_val, dphi.jac_val, ls_nfev, ls_njev + + # Increase ai + a.append(min(2*a[i], amax)) + + # If we reached this point, the line search failed + return None, None, None, ls_nfev, ls_njev + + +def zoom(alo, ahi, f, df, f0, df0, maxiter, c1, c2): + '''Zoom function for line search algorithm. (This is the same as for scipy)''' + + phi_lo = f(alo) + phi_hi = f(ahi) + dphi_lo = df(alo) + + for j in range(maxiter): + + tol_cubic = 0.2*(ahi-alo) + tol_quad = 0.1*(ahi-alo) + + if (j > 0): + # cubic interpolation for alo, phi(alo), dphi(alo) and ahi, phi(ahi) + aj = _cubicmin(alo, phi_lo, dphi_lo, ahi, phi_hi, aold, phi_old) + if (j == 0) or (aj is None) or (aj < alo + tol_cubic) or (aj > ahi - tol_cubic): + # quadratic interpolation for alo, phi(alo), dphi(alo) and ahi, phi(ahi) + aj = _quadmin(alo, phi_lo, dphi_lo, ahi, phi_hi) + + # Ensure aj is within bounds + if (aj is None) or (aj < alo + tol_quad) or (aj > ahi - tol_quad): + aj = alo + 0.5*(ahi - alo) + + # Evaluate phi(aj) + phi_j = f(aj) + + # Check for sufficient decrease + if (phi_j > f0 + c1*aj*df0) or (phi_j >= phi_lo): + # store old values + aold = ahi + phi_old = phi_hi + # update ahi + ahi = aj + phi_hi = phi_j + else: + # check curvature condition + dphi_j = df(aj) + if abs(dphi_j) <= -c2*df0: + return aj + + if dphi_j*(ahi-alo) >= 0: + # store old values + aold = ahi + phi_old = phi_hi + # update alo + ahi = alo + phi_hi = phi_lo + else: + # store old values + aold = alo + phi_old = phi_lo + + alo = aj + phi_lo = phi_j + dphi_lo = dphi_j + + # If we reached this point, the line search failed + return None + + + + +def line_search_backtracking(step_size, xk, pk, fun, jac, fk=None, jk=None, **kwargs): + ''' + Backtracking line search algorithm to find step size alpha that satisfies the Wolfe conditions. + + Parameters + ---------- + step_size : float + Initial step size to start the line search. + + xk : ndarray + Current point in the optimization process. + + pk : ndarray + Search direction. + + fun : callable + Objective function + + jac : callable + Gradient of the objective function + + fk : float, optional + Function value at xk. If None, it will be computed. + + jk : ndarray, optional + Gradient at xk. If None, it will be computed. + + **kwargs : dict + Additional parameters for the line search, such as: + - rho : float, backtracking factor (default: 0.5) + - maxiter : int, maximum number of iterations (default: 10) + - c1 : float, sufficient decrease condition (default: 1e-4) + - c2 : float, curvature condition (default: 0.9) + + Returns + ------- + alpha : float + Step size that satisfies the Wolfe conditions. + + fval : float + Function value at the new point xk + step_size*pk. + + jval : ndarray + Gradient at the new point xk + step_size*pk. + + nfev : int + Number of function evaluations. + + njev : int + Number of gradient evaluations. + ''' + + global ls_nfev + global ls_njev + ls_nfev = 0 + ls_njev = 0 + + # Unpack some kwargs + rho = kwargs.get('rho', 0.5) + maxiter = kwargs.get('maxiter', 10) + c1 = kwargs.get('c1', 1e-4) + + # Define phi and derivative of phi + @cache + def phi(alpha): + global ls_nfev + if (alpha == 0): + if (fk is None): + fun_val = fun(xk) + ls_nfev += 1 + else: + fun_val = fk + else: + fun_val = fun(xk + alpha*pk) + ls_nfev += 1 + return fun_val + + + # run the backtracking line search loop + for i in range(maxiter): + # Evaluate phi(alpha) + phi_i = phi(step_size) + + # Check for sufficient decrease + if (phi_i <= phi(0) + c1*step_size*np.dot(jk, pk)): + # Evaluate jac at new point + jac_new = jac(xk + step_size*pk) + return step_size, phi_i, jac_new, ls_nfev, ls_njev + + # Reduce step size + step_size *= rho + + # If we reached this point, the line search failed + return None, None, None, ls_nfev, ls_njev \ No newline at end of file diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index 7bf1081..df4a0a8 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -13,7 +13,13 @@ # Internal imports from popt.misc_tools import optim_tools as ot from popt.loop.optimize import Optimize -from popt.update_schemes import optimizers +from popt.update_schemes.line_search_step import line_search, line_search_backtracking + +# some symbols for logger +subk = '\u2096' +jac_inf_symbol = f'‖jac(x{subk})‖\u221E' +fun_xk_symbol = f'fun(x{subk})' + def LineSearch(fun, x, jac, method='GD', hess=None, args=(), bounds=None, callback=None, **options): ''' @@ -48,72 +54,83 @@ def LineSearch(fun, x, jac, method='GD', hess=None, args=(), bounds=None, callba A callable called after each successful iteration. The class instance of LineSearch is passed as the only argument to the callback function: callback(self) - **options: keyword arguments, optional + **options: + keyword arguments, optional LineSearch Options (**options) ------------------------------ - maxiter: int + - maxiter: int, Maximum number of iterations. Default is 20. - step_size: float - Step-size for optimizer. Default is 0.25/inf-norm(jac(x0)). - - step_size_maxiter: int + - lsmaxiter: int, Maximum number of iterations for the line search. Default is 10. + + - step_size: float, + Step-size for optimizer. Default is 0.25/inf-norm(jac(x0)). - step_size_max: float - Maximum step-size. Default is 1e5 + - step_size_max: float, + Maximum step-size. Default is 1e5. If bounds are specified, + the maximum step-size is set to the maximum step-size allowed by the bounds. - step_size_adapt: int + - step_size_adapt: int, Set method for choosing initial step-size for each iteration. If 0, step_size value is used. If 1, Equation (3.6) from "Numercal Optimization" [1] is used. If 2, the equation above Equation (3.6) is used. Default is 0. - c1: float + - c1: float, Tolerance parameter for the Armijo condition. Default is 1e-4. - c2: float + - c2: float, Tolerance parameter for the Curvature condition. Default is 0.9. - xtol: float + - xtol: float, Optimization stop whenever |dx|>> from popt.update_schemes.linesearch import LineSearch >>> x0 = np.random.uniform(-3, 3, 2) >>> kwargs = {'maxiter': 100, - 'line_search_maxiter': 10, + 'lsmaxiter': 10, 'step_size_adapt': 1, 'saveit': False} >>> res = LineSearch(fun=rosen, x=x0, jac=rosen_der, method='BFGS', **kwargs) @@ -182,20 +199,27 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c self.callback = callback else: self.callback = None + + # Custom convergence criteria (callable) + convergence_criteria = options.get('convergence_criteria', None) + if callable(convergence_criteria): + self.convergence_criteria = self.convergence_criteria + else: + self.convergence_criteria = None # Set options for step-size self.step_size = options.get('step_size', None) self.step_size_max = options.get('step_size_max', 1e5) self.step_size_adapt = options.get('step_size_adapt', 0) - # Set options for line-search method - self.line_search_kwargs = { + # Set options for line-search + self.lskwargs = { 'c1': options.get('c1', 1e-4), 'c2': options.get('c2', 0.9), + 'rho': options.get('rho', 0.5), 'amax': self.step_size_max, - 'xtol': options.get('xtol', 1e-8), - 'maxiter': options.get('line_search_maxiter', 10), - 'method' : options.get('line_search_method', 1) + 'maxiter': options.get('lsmaxiter', 10), + 'method' : options.get('lsmethod', 1) } # Set other options @@ -203,6 +227,11 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c self.resample = options.get('resample', 0) self.saveit = options.get('saveit', True) + # set tolerance for convergence + self.xtol = options.get('xtol', 1e-8) # tolerance for control vector + self.ftol = options.get('ftol', 1e-4) # relative tolerance for function value + self.gtol = options.get('gtol', 1e-5) # tolerance for inf-norm of jacobian + # Check method valid_methods = ['GD', 'BFGS', 'Newton'] if not self.method in valid_methods: @@ -241,12 +270,12 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c self.optimize_result = self.update_results() if self.saveit: ot.save_optimize_results(self.optimize_result) - if self.logger is not None: self.logger.info(f' ====== Running optimization - Line search ({method}) ======') - self.logger.info('Specified options\n'+pprint.pformat(OptimizeResult(self.options))) - self.logger.info(f' {"iter.":<10} {"fun":<15} {"step-size":<15} {"|grad|":<15}') - self.logger.info(f' {self.iteration:<10} {self.fk:<15.4e} {0.0:<15.4e} {la.norm(self.jk):<15.4e}') + self.logger.info('\nSPECIFIED OPTIONS:\n'+pprint.pformat(OptimizeResult(self.options))) + self.logger.info('') + self.logger.info(f' {"iter.":<10} {fun_xk_symbol:<15} {jac_inf_symbol:<15} {"step-size":<15}') + self.logger.info(f' {self.iteration:<10} {self.fk:<15.4e} {la.norm(self.jk, np.inf):<15.4e} {0:<15.4e}') self.logger.info('') self.run_loop() @@ -268,6 +297,11 @@ def _jac(self, x): g = self.jacobian(x) else: g = self.jacobian(x, *self.args) + + # project gradient onto the feasible set + if self.bounds is not None: + g = - self._project_pk(-g, x) + return g def _hess(self, x): @@ -281,78 +315,13 @@ def _hess(self, x): h = self.hessian(x, *self.args) return make_matrix_psd(h) - def update_results(self): - - res = {'fun': self.fk, - 'x': self.xk, - 'jac': self.jk, - 'hess': self.Hk, - 'hess_inv': self.Hk_inv, - 'nfev': self.nfev, - 'njev': self.njev, - 'nit': self.iteration, - 'step-size': self.step_size, - 'method': self.method, - 'save_folder': self.options.get('save_folder', './')} - - for a, arg in enumerate(self.args): - res[f'args[{a}]'] = arg - - if 'savedata' in self.options: - # Make sure "SAVEDATA" gives a list - if isinstance(self.options['savedata'], list): - savedata = self.options['savedata'] - else: - savedata = [self.options['savedata']] - - # Loop over variables to store in save list - for save_typ in savedata: - if save_typ in locals(): - res[save_typ] = eval('{}'.format(save_typ)) - elif hasattr(self, save_typ): - res[save_typ] = eval(' self.{}'.format(save_typ)) - else: - print(f'Cannot save {save_typ}!\n\n') - - return OptimizeResult(res) - def _set_step_size(self, pk): - ''' Sets the step-size ''' - - # If first iteration - if self.iteration == 1: - if self.step_size is None: - self.step_size = 0.25/la.norm(pk, np.inf) - alpha = self.step_size - else: - alpha = self.step_size - else: - if np.dot(pk, self.jk) != 0 and self.step_size_adapt != 0: - if self.step_size_adapt == 1: - alpha = 2*(self.fk - self.f_old)/np.dot(pk, self.jk) - if self.step_size_adapt == 2: - slope_old = np.dot(self.p_old, self.j_old) - slope_new = np.dot(pk, self.jk) - alpha = self.step_size*slope_old/slope_new - else: - alpha = self.step_size - - if alpha < 0: - alpha = abs(alpha) - - #if self.method in ['BFGS', 'Newton']: - # From "Numerical Optimization" - # alpha = min(1, 1.01*alpha) - - return min(alpha, self.step_size_max) - - def calc_update(self, iter_resamp=0): # Initialize variables for this step success = False - # If in resampling mode, compute jacobian + # If in resampling mode, compute jacobian # Else, jacobian from in __init__ or from latest line_search is used if self.jk is None: self.jk = self._jac(self.xk) @@ -375,41 +344,50 @@ def calc_update(self, iter_resamp=0): if self.method == 'Newton': pk = - np.matmul(la.inv(self.Hk), self.jk) - # remove components that point out of the hybercube given by [lb,ub] - lb = np.array(self.bounds)[:, 0] - ub = np.array(self.bounds)[:, 1] - for i in range(self.xk.size): - if (self.xk[i] <= lb[i] and pk[i] < 0) or (self.xk[i] >= ub[i] and pk[i] > 0): - pk[i] = 0 + # porject search direction onto the feasible set + if self.bounds is not None: + pk = self._project_pk(pk, self.xk) # Set step_size - step_size = self._set_step_size(pk) - - # Set maximum step-size if self.bounds is not None: - mean_bound_range = np.mean([b[1]-b[0] for b in self.bounds]) - step_size_max = mean_bound_range/np.linalg.norm(pk) - self.line_search_kwargs['amax'] = step_size_max + self.step_size_max = self._set_max_step_size(pk, self.xk) + self.lskwargs['amax'] = self.step_size_max + step_size = self._set_step_size(pk, self.step_size_max) # Perform line-search self.logger.info('Performing line search...') - ls_res = line_search( - fun=self._fun, - jac=self._jac, - xk=self.xk, - pk=pk, - ak=step_size, - fk=self.fk, - gk=self.jk, - logger=self.logger, - **self.line_search_kwargs - ) - step_size, f_new, f_old, j_new, self.msg = ls_res + if self.lskwargs['method'] == 0: + ls_res = line_search_backtracking( + step_size=step_size, + xk=self.xk, + pk=pk, + fun=self._fun, + jac=self._jac, + fk=self.fk, + jk=self.jk, + **self.lskwargs + ) + else: + ls_res = line_search( + step_size=step_size, + xk=self.xk, + pk=pk, + fun=self._fun, + jac=self._jac, + fk=self.fk, + jk=self.jk, + **self.lskwargs + ) + step_size, f_new, j_new, _, _ = ls_res if not (step_size is None): - + + # Save old values x_old = self.xk j_old = self.jk + f_old = self.fk + + # Update control x_new = ot.clip_state(x_old + step_size*pk, self.bounds) # Update state @@ -421,6 +399,7 @@ def calc_update(self, iter_resamp=0): self.j_old = j_old self.f_old = f_old self.p_old = pk + sk = x_new - x_old # Call the callback function if callable(self.callback): @@ -428,14 +407,11 @@ def calc_update(self, iter_resamp=0): # Update BFGS if self.method == 'BFGS': - sk = x_new - x_old - yk = j_new - j_old - rho = 1/np.dot(yk,sk) - id_mat = np.eye(sk.size) + yk = j_new - j_old + if self.iteration == 1: + self.Hk_inv = np.dot(yk,sk)/np.dot(yk,yk) * np.eye(sk.size) - matrix1 = (id_mat - rho*np.outer(sk, yk)) - matrix2 = (id_mat - rho*np.outer(yk, sk)) - self.Hk_inv = matrix1@self.Hk_inv@matrix2 + rho*np.outer(sk, sk) + self.Hk_inv = bfgs_update(self.Hk_inv, sk, yk) # Update status success = True @@ -448,9 +424,36 @@ def calc_update(self, iter_resamp=0): # Write logging info if self.logger is not None: self.logger.info('') - self.logger.info(f' {"iter.":<10} {"fun":<15} {"step-size":<15} {"|grad|":<15}') - self.logger.info(f' {self.iteration:<10} {self.fk:<15.4e} {step_size:<15.4e} {la.norm(self.jk):<15.4e}') + self.logger.info(f' {"iter.":<10} {fun_xk_symbol:<15} {jac_inf_symbol:<15} {"step-size":<15}') + self.logger.info(f' {self.iteration:<10} {self.fk:<15.4e} {la.norm(self.jk, np.inf):<15.4e} {step_size:<15.4e}') self.logger.info('') + + # Check for convergence + if (la.norm(sk, np.inf) < self.xtol): + self.msg = 'Convergence criteria met: |dx| < xtol' + self.logger.info(self.msg) + success = False + return success + if (np.abs(self.fk - f_old) < self.ftol * np.abs(f_old)): + self.msg = 'Convergence criteria met: |f(x+dx) - f(x)| < ftol * |f(x)|' + self.logger.info(self.msg) + success = False + return success + if (la.norm(self.jk, np.inf) < self.gtol): + self.msg = f'Convergence criteria met: {jac_inf_symbol} < gtol' + self.logger.info(self.msg) + success = False + return success + + # Check for custom convergence + if callable(self.convergence_criteria): + if self.convergence_criteria(self): + self.logger.info('Custom convergence criteria met. Stopping optimization.') + success = False + return success + + if self.step_size_adapt == 2: + self.step_size = step_size # Update iteration self.iteration += 1 @@ -469,195 +472,122 @@ def calc_update(self, iter_resamp=0): success = False return success + + def update_results(self): + + res = {'fun': self.fk, + 'x': self.xk, + 'jac': self.jk, + 'hess': self.Hk, + 'hess_inv': self.Hk_inv, + 'nfev': self.nfev, + 'njev': self.njev, + 'nit': self.iteration, + 'step-size': self.step_size, + 'method': self.method, + 'save_folder': self.options.get('save_folder', './')} + + for a, arg in enumerate(self.args): + res[f'args[{a}]'] = arg + if 'savedata' in self.options: + # Make sure "SAVEDATA" gives a list + if isinstance(self.options['savedata'], list): + savedata = self.options['savedata'] + else: + savedata = [self.options['savedata']] -def line_search(fun, jac, xk, pk, ak, fk=None, gk=None, c1=0.0001, c2=0.9, maxiter=10, **kwargs): - ''' - Performs a single line search step - ''' - line_search_step = LineSearchStepBase( - fun, - jac, - xk, - pk, - ak, - fk, - gk, - c1, - c2, - maxiter, - **kwargs - ) - return line_search_step() - -class LineSearchStepBase: - - def __init__(self, fun, jac, xk, pk, ak, fk=None, gk=None, c1=0.0001, c2=0.9, maxiter=10, **kwargs): - self.fun = fun - self.jac = jac - self.xk = xk - self.pk = pk - self.ak = ak - self.fk = fk - self.gk = gk - self.c1 = c1 - self.c2 = c2 - self.maxiter = maxiter - self.msg = '' - - # kwargs - self.amax = kwargs.get('amax', 1e5) - self.amin = kwargs.get('amin', 0.0) - self.xtol = kwargs.get('xtol', 1e-8) - self.method = kwargs.get('method', 1) - self.logger = kwargs.get('logger', None) - - # If c2 is None, the curvature condition is not used - if self.c2 is None: - self.c2 = np.inf - self.method = 0 - - # Check for initial values - if self.fk is None: - self.phi0 = self.phi(0, eval=False) - else: - self.phi0 = self.fk + # Loop over variables to store in save list + for save_typ in savedata: + if save_typ in locals(): + res[save_typ] = eval('{}'.format(save_typ)) + elif hasattr(self, save_typ): + res[save_typ] = eval(' self.{}'.format(save_typ)) + else: + print(f'Cannot save {save_typ}!\n\n') + + return OptimizeResult(res) + + def _set_step_size(self, pk, amax): + ''' Sets the step-size ''' + + # If first iteration + if (self.iteration == 1): + if (self.step_size is None): + self.step_size = 0.25/la.norm(pk, np.inf) + alpha = self.step_size + else: + alpha = self.step_size - if self.gk is None: - self.dphi0 = self.dphi(0, eval=False) else: - self.dphi0 = np.dot(self.pk, self.gk) + if (self.step_size_adapt == 1) and (np.dot(pk, self.jk) != 0): + alpha = 2*(self.fk - self.f_old)/np.dot(pk, self.jk) + elif (self.step_size_adapt == 2) and (np.dot(pk, self.jk) == 0): + slope_old = np.dot(self.p_old, self.j_old) + slope_new = np.dot(pk, self.jk) + alpha = self.step_size*slope_old/slope_new + else: + alpha = self.step_size + if alpha < 0: + alpha = abs(alpha) - def __call__(self): + if alpha >= amax: + alpha = 0.75*amax - if self.method == 0: - step_size, fnew = self._line_search_alpha_cut(step_size=self.ak) - - if self.method == 1: - step_size, fnew = self._line_search_alpha_interpol(step_size=self.ak) - - if self.method == 2: - dcsrch = DCSRCH( - self.phi, - self.dphi, - self.c1, - self.c2, - self.xtol, - self.amin, - self.amax - ) - dcsrch_res = dcsrch( - self.ak, - phi0=self.phi0, - derphi0=self.dphi0, - maxiter=self.maxiter - ) - step_size, fnew, _, self.msg = dcsrch_res - self.msg = str(self.msg) - - if step_size is None: - if self.msg is None: - self.msg = 'Line search did not find a solution' - return None, None, None, None, self.msg - elif la.norm(step_size*self.pk) <= self.xtol: - self.msg = f'|dx| < {self.xtol}' - return None, None, None, None, self.msg - else: - step_size = min(step_size, self.amax) - self.msg = 'Line search was successful' - return step_size, fnew, self.phi0, self.jac_val, self.msg + return alpha + + def _project_pk(self, pk, xk): + ''' Projects the jacobian onto the feasible set defined by bounds ''' + lb = np.array(self.bounds)[:, 0] + ub = np.array(self.bounds)[:, 1] + for i, pk_val in enumerate(pk): + if (xk[i] <= lb[i] and pk_val < 0) or (xk[i] >= ub[i] and pk_val > 0): + pk[i] = 0 + return pk + + def _set_max_step_size(self, pk, xk): + lb = np.array(self.bounds)[:, 0] + ub = np.array(self.bounds)[:, 1] + amax = [] + for i, pk_val in enumerate(pk): + if pk_val < 0: + amax.append((lb[i] - xk[i])/pk_val) + elif pk_val > 0: + amax.append((ub[i] - xk[i])/pk_val) + else: + amax.append(np.inf) + amax = min(amax) + return amax - def _line_search_alpha_cut(self, step_size): - ak = step_size - for i in range(self.maxiter): - phi_new = self.phi(ak) - # Check Armijo Condition - if phi_new < self.phi0 + self.c1*ak*self.dphi0: - dphi_new = self.dphi(ak) +def bfgs_update(Hk, sk, yk): + """ + Perform the BFGS update of the inverse Hessian approximation. - # Curvature condition - if abs(dphi_new) <= abs(self.c2*self.dphi0): - return ak, phi_new - - ak = ak/2 - - return None, None - - def _line_search_alpha_interpol(self, step_size): - ak = step_size + Parameters: + - Hk: np.ndarray, current inverse Hessian approximation (n x n) + - sk: np.ndarray, step vector (x_{k+1} - x_k), shape (n,) + - yk: np.ndarray, gradient difference (grad_{k+1} - grad_k), shape (n,) - # Some lists - alpha = [0.0] - phi = [self.phi0] - dphi = [self.dphi0] + Returns: + - Hk_new: np.ndarray, updated inverse Hessian approximation + """ + sk = sk.reshape(-1, 1) + yk = yk.reshape(-1, 1) + rho = 1.0 / (yk.T @ sk) - for i in range(1, self.maxiter+1): - - # Append lists - alpha.append(ak) - phi.append(self.phi(ak)) - dphi.append(self.dphi(ak)) - - # Check Armijo Condition - if phi[i] > self.phi0 + self.c1*alpha[i]*self.dphi0 or (phi[i] >= phi[i-1] and i>1): - step_size_new, phi_new = self._zoom(alpha[i-1], alpha[i], phi[i-1], phi[i], dphi[i-1]) - return step_size_new, phi_new - - if abs(dphi[i]) < - self.c2*self.dphi0: - return alpha[i], phi[i] - - # Check Curvature condition - if dphi[i] >= 0: - step_size_new, phi_new = self._zoom(alpha[i], alpha[i-1], phi[i], phi[i-1], dphi[i]) - return step_size_new, phi_new - - if alpha[i] >= self.amax: - return None, None - else: - ak = ak*2 - - return None, None + if rho <= 0: + raise ValueError("Non-positive curvature detected. BFGS update skipped.") + I = np.eye(Hk.shape[0]) + Vk = I - rho * sk @ yk.T + Hk_new = Vk @ Hk @ Vk.T + rho * sk @ sk.T - def log(self, msg): - if self.logger is None: - print(msg) - else: - self.logger.info(msg) - - @cache - def phi(self, a, eval=True): - if eval: - self.log(' Evaluating Armijo Condition') - return self.fun(self.xk + a*self.pk) - - @cache - def dphi(self, a, eval=True): - if eval: - self.log(' Evaluating Curvature Condition') - jval = self.jac(self.xk + a*self.pk) - self.jac_val = jval - return np.dot(self.pk, jval) - - def _zoom(self, a_lo, a_hi, phi_lo, phi_hi, dphi_low): - alpha_new, phi_new, _ = _zoom(a_lo=a_lo, - a_hi=a_hi, - phi_lo=phi_lo, - phi_hi=phi_hi, - derphi_lo=dphi_low, - phi=self.phi, - derphi=self.dphi, - phi0=self.phi0, - derphi0=self.dphi0, - c1=self.c1, - c2=self.c2, - extra_condition=lambda *args: True) - return alpha_new, phi_new + return Hk_new def get_near_psd(A): @@ -688,7 +618,6 @@ def make_matrix_psd(A, maxiter=100): return None - From af046a0b0d4212ca1f85d1f3b48a2287530fe384 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 21 Aug 2025 13:12:54 +0200 Subject: [PATCH 08/94] added Newton-CG to LineSearch --- popt/update_schemes/linesearch.py | 147 ++++++++++++++++------------ popt/update_schemes/trust_region.py | 76 +++++++------- 2 files changed, 122 insertions(+), 101 deletions(-) diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index df4a0a8..e5137f8 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -17,8 +17,10 @@ # some symbols for logger subk = '\u2096' +sup2 = '\u00b2' jac_inf_symbol = f'‖jac(x{subk})‖\u221E' fun_xk_symbol = f'fun(x{subk})' +nabla_symbol = "\u2207" def LineSearch(fun, x, jac, method='GD', hess=None, args=(), bounds=None, callback=None, **options): @@ -39,7 +41,7 @@ def LineSearch(fun, x, jac, method='GD', hess=None, args=(), bounds=None, callba method: str Which optimization method to use. Default is 'GD' for 'Gradient Descent'. Other options are 'BFGS' for the 'Broyden–Fletcher–Goldfarb–Shanno' method, - and 'Newton' for the Newton method. + and 'Newton-CG'. hess: callable, optional Hessian function, hess(x, *args). Default is None. @@ -233,14 +235,12 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c self.gtol = options.get('gtol', 1e-5) # tolerance for inf-norm of jacobian # Check method - valid_methods = ['GD', 'BFGS', 'Newton'] + valid_methods = ['GD', 'BFGS', 'Newton-CG'] if not self.method in valid_methods: raise ValueError(f"'{self.method}' is not a valid method. Valid methods are: {valid_methods}") - # Make sure hessian is callable if mehtod='Newton' - if (self.method == 'Newton') and (not callable(self.hessian)): - warnings.warn('Newton’s method requires a hessian, method changed to BFGS') - self.method = 'BFGS' + if (self.method == 'Newton-CG') and (self.hessian is None): + print(f'Warning: No hessian function provided. Finite difference approximation is used: {nabla_symbol}{sup2}f(x{subk})d ≈ ({nabla_symbol}f(x{subk}+hd)-{nabla_symbol}f(x{subk}))/h') # Calculate objective function of startpoint if not self.restart: @@ -267,7 +267,7 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c self.p_old = None # Initial results - self.optimize_result = self.update_results() + self.optimize_result = self.get_intermediate_results() if self.saveit: ot.save_optimize_results(self.optimize_result) if self.logger is not None: @@ -313,7 +313,7 @@ def _hess(self, x): h = self.hessian(x) else: h = self.hessian(x, *self.args) - return make_matrix_psd(h) + return h def calc_update(self, iter_resamp=0): @@ -341,8 +341,8 @@ def calc_update(self, iter_resamp=0): pk = - self.jk if self.method == 'BFGS': pk = - np.matmul(self.Hk_inv, self.jk) - if self.method == 'Newton': - pk = - np.matmul(la.inv(self.Hk), self.jk) + if self.method == 'Newton-CG': + pk = newton_cg(self.jk, Hk=self.Hk, xk=self.xk, jac=self._jac, eps=1e-4) # porject search direction onto the feasible set if self.bounds is not None: @@ -408,16 +408,14 @@ def calc_update(self, iter_resamp=0): # Update BFGS if self.method == 'BFGS': yk = j_new - j_old - if self.iteration == 1: - self.Hk_inv = np.dot(yk,sk)/np.dot(yk,yk) * np.eye(sk.size) - + if self.iteration == 1: self.Hk_inv = np.dot(yk,sk)/np.dot(yk,yk) * np.eye(sk.size) self.Hk_inv = bfgs_update(self.Hk_inv, sk, yk) # Update status success = True # Save Results - self.optimize_result = self.update_results() + self.optimize_result = self.get_intermediate_results() if self.saveit: ot.save_optimize_results(self.optimize_result) @@ -473,23 +471,19 @@ def calc_update(self, iter_resamp=0): return success - - def update_results(self): - - res = {'fun': self.fk, - 'x': self.xk, - 'jac': self.jk, - 'hess': self.Hk, - 'hess_inv': self.Hk_inv, - 'nfev': self.nfev, - 'njev': self.njev, - 'nit': self.iteration, - 'step-size': self.step_size, - 'method': self.method, - 'save_folder': self.options.get('save_folder', './')} - - for a, arg in enumerate(self.args): - res[f'args[{a}]'] = arg + def get_intermediate_results(self): + + # Define default results + results = { + 'fun': self.fk, + 'x': self.xk, + 'jac': self.jk, + 'nfev': self.nfev, + 'njev': self.njev, + 'nit': self.iteration, + 'method': self.method, + 'save_folder': self.options.get('save_folder', './') + } if 'savedata' in self.options: # Make sure "SAVEDATA" gives a list @@ -498,16 +492,20 @@ def update_results(self): else: savedata = [self.options['savedata']] + if 'args' in savedata: + for a, arg in enumerate(self.args): + results[f'args[{a}]'] = arg + # Loop over variables to store in save list - for save_typ in savedata: - if save_typ in locals(): - res[save_typ] = eval('{}'.format(save_typ)) - elif hasattr(self, save_typ): - res[save_typ] = eval(' self.{}'.format(save_typ)) + for variable in savedata: + if variable in locals(): + results[variable] = eval('{}'.format(variable)) + elif hasattr(self, variable): + results[variable] = eval('self.{}'.format(variable)) else: - print(f'Cannot save {save_typ}!\n\n') + print(f'Cannot save {variable}!\n\n') - return OptimizeResult(res) + return OptimizeResult(results) def _set_step_size(self, pk, amax): ''' Sets the step-size ''' @@ -589,33 +587,56 @@ def bfgs_update(Hk, sk, yk): return Hk_new +def newton_cg(gk, Hk=None, maxiter=None, **kwargs): + print('\nRunning Newton-CG subroutine...') -def get_near_psd(A): - eigval, eigvec = np.linalg.eig((A + A.T)/2) - eigval[eigval < 0] = 1.0 - return eigvec.dot(np.diag(eigval)).dot(eigvec.T) + if Hk is None: + jac = kwargs.get('jac') + eps = kwargs.get('eps', 1e-4) + xk = kwargs.get('xk') + + # define a finite difference approximation of the Hessian times a vector + def Hessd(d): + return (jac(xk + eps*d) - gk)/eps + + if maxiter is None: + maxiter = 20*gk.size # Same dfault as in scipy + + tol = min(0.5, np.sqrt(la.norm(gk)))*la.norm(gk) + z = 0 + r = gk + d = -r + + for j in range(maxiter): + print('iteration: ', j) + if Hk is None: + Hd = Hessd(d) + else: + Hd = np.matmul(Hk, d) + + dTHd = np.dot(d, Hd) + + if dTHd <= 0: + print('Negative curvature detected, terminating subroutine') + print('\n') + if j == 0: + return -gk + else: + return z + + rold = r + a = np.dot(r,r)/dTHd + z = z + a*d + r = r + a*Hd + + if la.norm(r) < tol: + print('Subroutine converged') + print('\n') + return z + + b = np.dot(r, r)/np.dot(rold, rold) + d = -r + b*d -def make_matrix_psd(A, maxiter=100): - # Set beta to Frobenius norm of A - beta = np.linalg.norm(A, 'fro') - - # Initialize tau - if np.min(np.diag(A)) > 0: - tau = 0 - else: - tau = beta/2 - - for _ in range(maxiter): - try: - M = A + tau*np.eye(A.shape[0]) - # Attempt Cholesky - np.linalg.cholesky(A + tau*np.eye(A.shape[0])) - return M - except np.linalg.LinAlgError: - # Set new tau - tau = max(2*tau, beta/2) - - return None diff --git a/popt/update_schemes/trust_region.py b/popt/update_schemes/trust_region.py index cd616f9..b23f6b6 100644 --- a/popt/update_schemes/trust_region.py +++ b/popt/update_schemes/trust_region.py @@ -239,44 +239,6 @@ def _hess(self, x): else: h = self.hessian(x, *self.args) return h - - def update_results(self): - res = { - 'fun': self.fk, - 'x': self.xk, - 'jac': self.jk, - 'hess': self.Hk, - 'nfev': self.nfev, - 'njev': self.njev, - 'nit': self.iteration, - 'trust_radius': self.trust_radius, - 'save_folder': self.options.get('save_folder', './') - } - - for a, arg in enumerate(self.args): - res[f'args[{a}]'] = arg - - if 'savedata' in self.options: - # Make sure "SAVEDATA" gives a list - if isinstance(self.options['savedata'], list): - savedata = self.options['savedata'] - else: - savedata = [self.options['savedata']] - - # Loop over variables to store in save list - for save_typ in savedata: - if save_typ in locals(): - res[save_typ] = eval('{}'.format(save_typ)) - elif hasattr(self, save_typ): - res[save_typ] = eval(' self.{}'.format(save_typ)) - else: - print(f'Cannot save {save_typ}!\n\n') - - return OptimizeResult(res) - - def _log(self, msg): - if self.logger is not None: - self.logger.info(msg) def solve_subproblem(self, g, B, delta): """ @@ -426,6 +388,44 @@ def calc_update(self, inner_iter=0): success = False return success + + def update_results(self): + res = { + 'fun': self.fk, + 'x': self.xk, + 'jac': self.jk, + 'hess': self.Hk, + 'nfev': self.nfev, + 'njev': self.njev, + 'nit': self.iteration, + 'trust_radius': self.trust_radius, + 'save_folder': self.options.get('save_folder', './') + } + + for a, arg in enumerate(self.args): + res[f'args[{a}]'] = arg + + if 'savedata' in self.options: + # Make sure "SAVEDATA" gives a list + if isinstance(self.options['savedata'], list): + savedata = self.options['savedata'] + else: + savedata = [self.options['savedata']] + + # Loop over variables to store in save list + for save_typ in savedata: + if save_typ in locals(): + res[save_typ] = eval('{}'.format(save_typ)) + elif hasattr(self, save_typ): + res[save_typ] = eval(' self.{}'.format(save_typ)) + else: + print(f'Cannot save {save_typ}!\n\n') + + return OptimizeResult(res) + + def _log(self, msg): + if self.logger is not None: + self.logger.info(msg) From ab6543ea2aa73d3544e7b2ec2019ad88dc4c13e8 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 1 Sep 2025 09:51:56 +0200 Subject: [PATCH 09/94] improved logging for LineSearch --- popt/update_schemes/line_search_step.py | 12 ++++++++++++ popt/update_schemes/linesearch.py | 11 ++++++----- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/popt/update_schemes/line_search_step.py b/popt/update_schemes/line_search_step.py index 5207940..3f4af38 100644 --- a/popt/update_schemes/line_search_step.py +++ b/popt/update_schemes/line_search_step.py @@ -68,6 +68,12 @@ def line_search(step_size, xk, pk, fun, jac, fk=None, jk=None, **kwargs): c1 = kwargs.get('c1', 1e-4) c2 = kwargs.get('c2', 0.9) + # check for logger in kwargs + global logger + logger = kwargs.get('logger', None) + if logger is None: + logger = print + # assertions assert step_size <= amax, "Initial step size must be less than or equal to amax." @@ -82,6 +88,7 @@ def phi(alpha): else: phi.fun_val = fk else: + logger(' Evaluating Armijo condition') phi.fun_val = fun(xk + alpha*pk) ls_nfev += 1 return phi.fun_val @@ -96,6 +103,7 @@ def dphi(alpha): else: dphi.jac_val = jk else: + logger(' Evaluating curvature condition') dphi.jac_val = jac(xk + alpha*pk) ls_njev += 1 return np.dot(dphi.jac_val, pk) @@ -107,6 +115,8 @@ def dphi(alpha): # Start loop a = [0, step_size] for i in range(1, maxiter+1): + logger(f'Line search iteration: {i-1}') + # Evaluate phi(ai) phi_i = phi(a[i]) @@ -134,6 +144,7 @@ def dphi(alpha): a.append(min(2*a[i], amax)) # If we reached this point, the line search failed + logger('Line search failed to find a suitable step size \n') return None, None, None, ls_nfev, ls_njev @@ -145,6 +156,7 @@ def zoom(alo, ahi, f, df, f0, df0, maxiter, c1, c2): dphi_lo = df(alo) for j in range(maxiter): + logger(f'Line search iteration: {j+1}') tol_cubic = 0.2*(ahi-alo) tol_quad = 0.1*(ahi-alo) diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index e5137f8..e723f78 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -221,7 +221,8 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c 'rho': options.get('rho', 0.5), 'amax': self.step_size_max, 'maxiter': options.get('lsmaxiter', 10), - 'method' : options.get('lsmethod', 1) + 'method' : options.get('lsmethod', 1), + 'logger' : self.logger.info } # Set other options @@ -355,7 +356,7 @@ def calc_update(self, iter_resamp=0): step_size = self._set_step_size(pk, self.step_size_max) # Perform line-search - self.logger.info('Performing line search...') + self.logger.info('Performing line search.............') if self.lskwargs['method'] == 0: ls_res = line_search_backtracking( step_size=step_size, @@ -556,9 +557,9 @@ def _set_max_step_size(self, pk, xk): elif pk_val > 0: amax.append((ub[i] - xk[i])/pk_val) else: - amax.append(np.inf) - amax = min(amax) - return amax + continue + + return max(amax) From 237b80a1f0886c72ccdbc78531ebf95479014623 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 3 Sep 2025 10:57:11 +0200 Subject: [PATCH 10/94] dummy message for github check --- popt/update_schemes/linesearch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index e723f78..20f8fc1 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -23,6 +23,7 @@ nabla_symbol = "\u2207" + def LineSearch(fun, x, jac, method='GD', hess=None, args=(), bounds=None, callback=None, **options): ''' A Line Search Optimizer. From e29c159cc034ee7559285e2c1a8d951bf4ee3252 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 3 Sep 2025 11:06:51 +0200 Subject: [PATCH 11/94] empty commit --- popt/update_schemes/linesearch.py | 1 - 1 file changed, 1 deletion(-) diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index 20f8fc1..e723f78 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -23,7 +23,6 @@ nabla_symbol = "\u2207" - def LineSearch(fun, x, jac, method='GD', hess=None, args=(), bounds=None, callback=None, **options): ''' A Line Search Optimizer. From 32abdce95f93868f239c510562a8d94263862b33 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 3 Sep 2025 13:28:27 +0200 Subject: [PATCH 12/94] udpate BFGS to skip update if negetive curvature --- popt/update_schemes/linesearch.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index e723f78..57a2a09 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -5,17 +5,14 @@ import warnings from numpy import linalg as la -from functools import cache from scipy.optimize import OptimizeResult -from scipy.optimize._dcsrch import DCSRCH -from scipy.optimize._linesearch import _zoom # Internal imports from popt.misc_tools import optim_tools as ot from popt.loop.optimize import Optimize from popt.update_schemes.line_search_step import line_search, line_search_backtracking -# some symbols for logger +# Some symbols for logger subk = '\u2096' sup2 = '\u00b2' jac_inf_symbol = f'‖jac(x{subk})‖\u221E' @@ -580,7 +577,8 @@ def bfgs_update(Hk, sk, yk): rho = 1.0 / (yk.T @ sk) if rho <= 0: - raise ValueError("Non-positive curvature detected. BFGS update skipped.") + print('Non-positive curvature detected. BFGS update skipped....') + return Hk I = np.eye(Hk.shape[0]) Vk = I - rho * sk @ yk.T From 4ae14f76339094b190cd06f20b0cbbbb02fe4f3b Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 8 Sep 2025 09:52:14 +0200 Subject: [PATCH 13/94] made input more elegant and cleaned up popt structure --- ensemble/ensemble.py | 349 ++++++------------ input_output/organize.py | 8 + input_output/read_config.py | 11 +- pipt/loop/assimilation.py | 34 +- pipt/update_schemes/enrml.py | 82 ++-- pipt/update_schemes/esmda.py | 38 +- popt/cost_functions/ecalc_npv.py | 4 +- popt/cost_functions/ecalc_pareto_npv.py | 4 +- popt/cost_functions/npv.py | 4 +- popt/cost_functions/ren_npv.py | 4 +- popt/update_schemes/enopt.py | 2 +- popt/update_schemes/genopt.py | 4 +- popt/update_schemes/linesearch.py | 84 +---- popt/update_schemes/smcopt.py | 2 +- popt/update_schemes/subroutines/__init__.py | 3 + popt/update_schemes/{ => subroutines}/cma.py | 2 + .../{ => subroutines}/optimizers.py | 2 + .../subroutines.py} | 96 ++++- 18 files changed, 294 insertions(+), 439 deletions(-) create mode 100644 popt/update_schemes/subroutines/__init__.py rename popt/update_schemes/{ => subroutines}/cma.py (99%) rename popt/update_schemes/{ => subroutines}/optimizers.py (99%) rename popt/update_schemes/{line_search_step.py => subroutines/subroutines.py} (79%) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 304b240..d4dc46a 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -80,7 +80,7 @@ def __init__(self, keys_en, sim, redund_sim=None): # If it is a restart run, we do not need to initialize anything, only load the self info. that exists in the # pickle save file. If it is not a restart run, we initialize everything below. - if 'restart' in self.keys_en and self.keys_en['restart'] == 'yes': + if ('restart' in self.keys_en) and (self.keys_en['restart'] == 'yes'): # Initiate a restart run self.logger.info('\033[92m--- Restart run initiated! ---\033[92m') # Check if the pickle save file exists in folder @@ -117,7 +117,7 @@ def __init__(self, keys_en, sim, redund_sim=None): self.disable_tqdm = False # extract information that is given for the prior model - self._ext_prior_info() + self.prior_info = self._extract_prior_info() # Calculate initial ensemble if IMPORTSTATICVAR has not been given in init. file. # Prior info. on state variables must be given by PRIOR_ keyword. @@ -172,243 +172,116 @@ def _ext_ml_info(self): self.error_comp_scheme = self.keys_en['multilevel'][i][2] self.ML_corr_done = False - def _ext_prior_info(self): - """ + def _extract_prior_info(self) -> dict: + ''' Extract prior information on STATE from keyword(s) PRIOR_. - """ - # Parse prior info on each state entered in STATE. - # Store names given in STATE - if not isinstance(self.keys_en['state'], list): # Single string - state_names = [self.keys_en['state']] - else: # List - state_names = self.keys_en['state'] - - # Check if PRIOR_ exists for each entry in STATE + ''' + + # Get state names as list + state_names = self.keys_en['state'] + if not isinstance(state_names, list): state_names = [state_names] + + # Check if PRIOR_ exists for each entry in state for name in state_names: - assert 'prior_' + name in self.keys_en, \ + assert f'prior_{name}' in self.keys_en, \ 'PRIOR_{0} is missing! This keyword is needed to make initial ensemble for {0} entered in ' \ 'STATE'.format(name.upper()) + + # define dict to store prior information in + prior_info = {name: None for name in state_names} - # Init. prior info variable - self.prior_info = {keys: None for keys in state_names} - - # Loop over each prior keyword and make an initial. ensemble for each state in STATE, - # which is subsequently stored in the state dictionary. If 3D grid dimensions are inputted, information for - # each layer must be inputted, else the single information will be copied to all layers. - grid_dim = np.array([0, 0]) + # loop over state priors for name in state_names: - # initiallize an empty dictionary inside the dictionary. - self.prior_info[name] = {} - # List the option names inputted in prior keyword - # opt_list = list(zip(*self.keys_da['prior_'+name])) - mean = None - self.prior_info[name]['mean'] = mean - vario = [None] - self.prior_info[name]['vario'] = vario - aniso = [None] - self.prior_info[name]['aniso'] = aniso - angle = [None] - self.prior_info[name]['angle'] = angle - corr_length = [None] - self.prior_info[name]['corr_length'] = corr_length - self.prior_info[name]['nx'] = self.prior_info[name]['ny'] = self.prior_info[name]['nz'] = None - - # Extract info. from the prior keyword - for i, opt in enumerate(list(zip(*self.keys_en['prior_' + name]))[0]): - if opt == 'vario': # Variogram model - if not isinstance(self.keys_en['prior_' + name][i][1], list): - vario = [self.keys_en['prior_' + name][i][1]] - else: - vario = self.keys_en['prior_' + name][i][1] - elif opt == 'mean': # Mean - mean = self.keys_en['prior_' + name][i][1] - elif opt == 'var': # Variance - if not isinstance(self.keys_en['prior_' + name][i][1], list): - variance = [self.keys_en['prior_' + name][i][1]] - else: - variance = self.keys_en['prior_' + name][i][1] - elif opt == 'aniso': # Anisotropy factor - if not isinstance(self.keys_en['prior_' + name][i][1], list): - aniso = [self.keys_en['prior_' + name][i][1]] - else: - aniso = self.keys_en['prior_' + name][i][1] - elif opt == 'angle': # Anisotropy angle - if not isinstance(self.keys_en['prior_' + name][i][1], list): - angle = [self.keys_en['prior_' + name][i][1]] - else: - angle = self.keys_en['prior_' + name][i][1] - elif opt == 'range': # Correlation length - if not isinstance(self.keys_en['prior_' + name][i][1], list): - corr_length = [self.keys_en['prior_' + name][i][1]] + prior = self.keys_en[f'prior_{name}'] + + # Check if is a list (old way) + if isinstance(prior, list): + # list of lists - old way of inputting prior information + prior_dict = {} + for i, opt in enumerate(list(zip(*prior))[0]): + if opt == 'limits': + prior_dict[opt] = prior[i][1:] else: - corr_length = self.keys_en['prior_' + name][i][1] - elif opt == 'grid': # Grid dimensions - grid_dim = self.keys_en['prior_' + name][i][1] - elif opt == 'limits': # Truncation values - limits = self.keys_en['prior_' + name][i][1:] - elif opt == 'active': # Number of active cells (single number) - active = self.keys_en['prior_' + name][i][1] - - # Check if mean needs to be loaded, or if loaded - if type(mean) is str: - assert mean.endswith('.npz'), 'File name does not end with \'.npz\'!' - load_file = np.load(mean) + prior_dict[opt] = prior[i][1] + prior = prior_dict + else: + assert isinstance(prior, dict), 'PRIOR_{0} must be a dictionary or list of lists!'.format(name.upper()) + + + # load mean if in file + if isinstance(prior['mean'], str): + assert prior['mean'].endswith('.npz'), 'File name does not end with \'.npz\'!' + load_file = np.load(prior['mean']) assert len(load_file.files) == 1, \ 'More than one variable located in {0}. Only the mean vector can be stored in the .npz file!' \ - .format(mean) - mean = load_file[load_file.files[0]] + .format(prior['mean']) + prior['mean'] = load_file[load_file.files[0]] else: # Single number inputted, make it a list if not already - if not isinstance(mean, list): - mean = [mean] - - # Check if limits exists - try: - limits - except NameError: - limits = None - - # check if active exists - try: - active - except NameError: - active = None - - # Extract x- and y-dim - nx = int(grid_dim[0]) - ny = int(grid_dim[1]) - - # Check if 3D grid inputted. If so, we check if info. has been given on all layers. In the case it has - # not been given, we just copy the info. given. - if len(grid_dim) == 3 and grid_dim[2] > 1: # 3D - nz = int(grid_dim[2]) - - # Check mean when values have been inputted directly (not when mean has been loaded) - if isinstance(mean, list) and len(mean) < nz: - # Check if it is more than one entry and give error - assert len(mean) == 1, \ - 'Information from MEAN has been given for {0} layers, whereas {1} is needed!' \ - .format(len(mean), nz) - - # Only 1 entry; copy this to all layers - print( - '\033[1;33mSingle entry for MEAN will be copied to all {0} layers\033[1;m'.format(nz)) - self.prior_info[name]['mean'] = mean * nz - - else: - self.prior_info[name]['mean'] = mean - - # Check variogram model - if len(vario) < nz: - # Check if it is more than one entry and give error - assert len(vario) == 1, \ - 'Information from VARIO has been given for {0} layers, whereas {1} is needed!' \ - .format(len(vario), nz) - - # Only 1 entry; copy this to all layers - print( - '\033[1;33mSingle entry for VARIO will be copied to all {0} layers\033[1;m'.format(nz)) - self.prior_info[name]['vario'] = vario * nz - - else: - self.prior_info[name]['vario'] = vario - - # Variance - if len(variance) < nz: - # Check if it is more than one entry and give error - assert len(variance) == 1, \ - 'Information from VAR has been given for {0} layers, whereas {1} is needed!' \ - .format(len(variance), nz) - - # Only 1 entry; copy this to all layers - print( - '\033[1;33mSingle entry for VAR will be copied to all {0} layers\033[1;m'.format(nz)) - self.prior_info[name]['variance'] = variance * nz - - else: - self.prior_info[name]['variance'] = variance - - # Aniso factor - if len(aniso) < nz: - # Check if it is more than one entry and give error - assert len(aniso) == 1, \ - 'Information from ANISO has been given for {0} layers, whereas {1} is needed!' \ - .format(len(aniso), nz) - - # Only 1 entry; copy this to all layers - print( - '\033[1;33mSingle entry for ANISO will be copied to all {0} layers\033[1;m'.format(nz)) - self.prior_info[name]['aniso'] = aniso * nz - - else: - self.prior_info[name]['aniso'] = aniso - - # Aniso factor - if len(angle) < nz: - # Check if it is more than one entry and give error - assert len(angle) == 1, \ - 'Information from ANGLE has been given for {0} layers, whereas {1} is needed!' \ - .format(len(angle), nz) - - # Only 1 entry; copy this to all layers - print( - '\033[1;33mSingle entry for ANGLE will be copied to all {0} layers\033[1;m'.format(nz)) - self.prior_info[name]['angle'] = angle * nz + if not isinstance(prior['mean'], list): + prior['mean'] = [prior['mean']] + + # loop over keys in prior + for key in prior.keys(): + # ensure that entry is a list + if (not isinstance(prior[key], list)) and (key != 'mean'): + prior[key] = [prior[key]] + + # change the name of some keys + prior['variance'] = prior.pop('var', None) + prior['corr_length'] = prior.pop('range', None) + + # process grid + if 'grid' in prior: + grid_dim = prior['grid'] + + # check if 3D-grid + if (len(grid_dim) == 3) and (grid_dim[2] > 1): + nz = int(grid_dim[2]) + prior['nz'] = nz + prior['nx'] = int(grid_dim[0]) + prior['ny'] = int(grid_dim[1]) + + + # Check mean when values have been inputted directly (not when mean has been loaded) + mean = prior['mean'] + if isinstance(mean, list) and len(mean) < nz: + # Check if it is more than one entry and give error + assert len(mean) == 1, \ + 'Information from MEAN has been given for {0} layers, whereas {1} is needed!' \ + .format(len(mean), nz) + + # Only 1 entry; copy this to all layers + print( + '\033[1;33mSingle entry for MEAN will be copied to all {0} layers\033[1;m'.format(nz)) + prior['mean'] = mean * nz + + #check if info. has been given on all layers. In the case it has not been given, we just copy the info. given. + for key in ['vario', 'variance', 'aniso', 'angle', 'corr_length']: + if key in prior.keys(): + val = prior[key] + if len(val) < nz: + # Check if it is more than one entry and give error + assert len(val) == 1, \ + 'Information from {0} has been given for {1} layers, whereas {2} is needed!' \ + .format(key.upper(), len(val), nz) + + # Only 1 entry; copy this to all layers + print( + '\033[1;33mSingle entry for {0} will be copied to all {1} layers\033[1;m'.format(key.upper(), nz)) + prior[key] = val * nz else: - self.prior_info[name]['angle'] = angle + prior['nx'] = int(grid_dim[0]) + prior['ny'] = int(grid_dim[1]) + prior['nz'] = 1 - # Corr. length - if len(corr_length) < nz: - # Check if it is more than one entry and give error - assert len(corr_length) == 1, \ - 'Information from RANGE has been given for {0} layers, whereas {1} is needed!' \ - .format(len(corr_length), nz) + prior.pop('grid', None) - # Only 1 entry; copy this to all layers - print( - '\033[1;33mSingle entry for RANGE will be copied to all {0} layers\033[1;m'.format(nz)) - self.prior_info[name]['corr_length'] = corr_length * nz - - else: - self.prior_info[name]['corr_length'] = corr_length - - # Limits, if exists - if limits is not None: - self.prior_info[name]['limits'] = limits - - # if isinstance(limits[0], list) and len(limits) < nz or \ - # not isinstance(limits[0], list) and len(limits) < 2 * nz: - # # Check if it is more than one entry and give error - # assert (isinstance(limits[0], list) and len(limits) == 1), \ - # 'Information from LIMITS has been given for {0} layers, whereas {1} is needed!' \ - # .format(len(limits), nz) - # assert (not isinstance(limits[0], list) and len(limits) == 2), \ - # 'Information from LIMITS has been given for {0} layers, whereas {1} is needed!' \ - # .format(len(limits) / 2, nz) - # - # # Only 1 entry; copy this to all layers - # print( - # '\033[1;33mSingle entry for RANGE will be copied to all {0} layers\033[1;m'.format(nz)) - # self.prior_info[name]['limits'] = [limits] * nz - - else: # 2D grid only, or optimization case - nz = 1 - self.prior_info[name]['mean'] = mean - self.prior_info[name]['vario'] = vario - self.prior_info[name]['variance'] = variance - self.prior_info[name]['aniso'] = aniso - self.prior_info[name]['angle'] = angle - self.prior_info[name]['corr_length'] = corr_length - if limits is not None: - self.prior_info[name]['limits'] = limits - if active is not None: - self.prior_info[name]['active'] = active - - self.prior_info[name]['nx'] = nx - self.prior_info[name]['ny'] = ny - self.prior_info[name]['nz'] = nz - - # Loop over keys and input + # add prior to prior_info + prior_info[name] = prior + + return prior_info + def gen_init_ensemble(self): """ @@ -427,21 +300,21 @@ def gen_init_ensemble(self): ind_end = 0 # Extract info. - nz = self.prior_info[name]['nz'] - mean = self.prior_info[name]['mean'] - nx = self.prior_info[name]['nx'] - ny = self.prior_info[name]['ny'] + nx = self.prior_info[name].get('nx', 0) + ny = self.prior_info[name].get('ny', 0) + nz = self.prior_info[name].get('nz', 0) + mean = self.prior_info[name].get('mean', None) + if nx == ny == 0: # assume ensemble will be generated elsewhere if dimensions are zero break - variance = self.prior_info[name]['variance'] - corr_length = self.prior_info[name]['corr_length'] - aniso = self.prior_info[name]['aniso'] - vario = self.prior_info[name]['vario'] - angle = self.prior_info[name]['angle'] - if 'limits' in self.prior_info[name]: - limits = self.prior_info[name]['limits'] - else: - limits = None + + variance = self.prior_info[name].get('variance', None) + corr_length = self.prior_info[name].get('corr_length', None) + aniso = self.prior_info[name].get('aniso', None) + vario = self.prior_info[name].get('vario', None) + angle = self.prior_info[name].get('angle', None) + limits= self.prior_info[name].get('limits',None) + # Loop over nz to make layers of 2D priors for i in range(self.prior_info[name]['nz']): diff --git a/input_output/organize.py b/input_output/organize.py index 3accf02..8e500e3 100644 --- a/input_output/organize.py +++ b/input_output/organize.py @@ -3,6 +3,7 @@ from copy import deepcopy import csv import datetime as dt +import pandas as pd class Organize_input(): @@ -109,6 +110,13 @@ def _org_report(self): pred_prim.extend(csv_data) self.keys_fwd['reportpoint'] = pred_prim + elif isinstance(self.keys_fwd['reportpoint'], dict): + self.keys_fwd['reportpoint'] = pd.date_range(**self.keys_fwd['reportpoint']).to_pydatetime().tolist() + + else: + pass + + # Check if assimindex is given as a csv file. If so, we read and make a potential 2D list (if sequential). if 'assimindex' in self.keys_pr: if isinstance(self.keys_pr['assimindex'], str) and self.keys_pr['assimindex'].endswith('.csv'): diff --git a/input_output/read_config.py b/input_output/read_config.py index 986e755..ba00da2 100644 --- a/input_output/read_config.py +++ b/input_output/read_config.py @@ -51,6 +51,12 @@ def ndarray_constructor(loader, node): y = yaml.load(fid, Loader=FullLoader) # Check for dataassim and fwdsim + if 'ensemble' in y.keys(): + keys_en = y['ensemble'] + check_mand_keywords_en(keys_en) + else: + keys_en = None + if 'optim' in y.keys(): keys_pr = y['optim'] check_mand_keywords_opt(keys_pr) @@ -59,16 +65,17 @@ def ndarray_constructor(loader, node): check_mand_keywords_da(keys_pr) else: raise KeyError + if 'fwdsim' in y.keys(): keys_fwd = y['fwdsim'] else: raise KeyError # Organize keywords - org = Organize_input(keys_pr, keys_fwd) + org = Organize_input(keys_pr, keys_fwd, keys_en) org.organize() - return org.get_keys_pr(), org.get_keys_fwd() + return org.get_keys_pr(), org.get_keys_fwd(), org.get_keys_en() def convert_txt_to_toml(init_file): diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index 0ccd6dd..f5caa72 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -301,32 +301,20 @@ def _ext_max_iter(self): - ST 7/6-16 """ if 'iteration' in self.ensemble.keys_da: - # Make sure ITERATION is a list - if not isinstance(self.ensemble.keys_da['iteration'][0], list): - iter_opts = [self.ensemble.keys_da['iteration']] - else: - iter_opts = self.ensemble.keys_da['iteration'] - + iter_opts = dict(self.ensemble.keys_da['iteration']) # Check if 'max_iter' has been given; if not, give error (mandatory in ITERATION) - assert 'max_iter' in list( - zip(*iter_opts))[0], 'MAX_ITER has not been given in ITERATION!' - - # Extract max. iter - max_iter = [item[1] for item in iter_opts if item[0] == 'max_iter'][0] + try: + max_iter = iter_opts['max_iter'] + except KeyError: + raise AssertionError('MAX_ITER has not been given in ITERATION') elif 'mda' in self.ensemble.keys_da: - # Make sure ITERATION is a list - if not isinstance(self.ensemble.keys_da['mda'][0], list): - iter_opts = [self.ensemble.keys_da['mda']] - else: - iter_opts = self.ensemble.keys_da['mda'] - - # Check if 'max_iter' has been given; if not, give error (mandatory in ITERATION) - assert 'tot_assim_steps' in list( - zip(*iter_opts))[0], 'TOT_ASSIM_STEPS has not been given in MDA!' - - # Extract max. iter - max_iter = [item[1] for item in iter_opts if item[0] == 'tot_assim_steps'][0] + iter_opts = dict(self.ensemble.keys_da['mda']) + # Check if 'tot_assim_steps' has been given; if not, raise error (mandatory in MDA) + try: + max_iter = iter_opts['tot_assim_steps'] + except KeyError: + raise AssertionError('TOT_ASSIM_STEPS has not been given in MDA!') else: max_iter = 1 diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index 632564b..741b700 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -250,36 +250,22 @@ def _ext_iter_param(self): file. These parameters include convergence tolerances and parameters for the damping parameter. Default values for these parameters have been given here, if they are not provided in ITERATION. """ - - # Predefine all the default values - self.data_misfit_tol = 0.01 - self.step_tol = 0.01 - self.lam = 100 - self.lam_max = 1e10 - self.lam_min = 0.01 - self.gamma = 5 - self.trunc_energy = 0.95 + try: + options = dict(self.keys_da['iteration']) + except: + options = dict([self.keys_da['iteration']]) + + # unpack options + self.data_misfit_tol = options.get('data_misfit_tol', 0.01) + self.trunc_energy = options.get('energy', 0.95) + self.step_tol = options.get('step_tol', 0.01) + self.lam = options.get('lambda', 100) + self.lam_max = options.get('lambda_max', 1e10) + self.lam_min = options.get('lambda_min', 0.01) + self.gamma = options.get('lambda_factor', 5) self.iteration = 0 - # Loop over options in ITERATION and extract the parameters we want - for i, opt in enumerate(list(zip(*self.keys_da['iteration']))[0]): - if opt == 'data_misfit_tol': - self.data_misfit_tol = self.keys_da['iteration'][i][1] - if opt == 'step_tol': - self.step_tol = self.keys_da['iteration'][i][1] - if opt == 'lambda': - self.lam = self.keys_da['iteration'][i][1] - if opt == 'lambda_max': - self.lam_max = self.keys_da['iteration'][i][1] - if opt == 'lambda_min': - self.lam_min = self.keys_da['iteration'][i][1] - if opt == 'lambda_factor': - self.gamma = self.keys_da['iteration'][i][1] - - if 'energy' in self.keys_da: - # initial energy (Remember to extract this) - self.trunc_energy = self.keys_da['energy'] - if self.trunc_energy > 1: # ensure that it is given as percentage + if self.trunc_energy > 1: # ensure that it is given as percentage self.trunc_energy /= 100. @@ -593,33 +579,19 @@ def _ext_iter_param(self): file. These parameters include convergence tolerances and parameters for the damping parameter. Default values for these parameters have been given here, if they are not provided in ITERATION. """ - - # Predefine all the default values - self.data_misfit_tol = 0.01 - self.step_tol = 0.01 - self.gamma = 0.2 - self.gamma_max = 0.5 - self.gamma_factor = 2.5 - self.trunc_energy = 0.95 - self.iteration = 0 - - # Loop over options in ITERATION and extract the parameters we want - for i, opt in enumerate(list(zip(*self.keys_da['iteration']))[0]): - if opt == 'data_misfit_tol': - self.data_misfit_tol = self.keys_da['iteration'][i][1] - if opt == 'step_tol': - self.step_tol = self.keys_da['iteration'][i][1] - if opt == 'gamma': - self.gamma = self.keys_da['iteration'][i][1] - if opt == 'gamma_max': - self.gamma_max = self.keys_da['iteration'][i][1] - if opt == 'gamma_factor': - self.gamma_factor = self.keys_da['iteration'][i][1] - - if 'energy' in self.keys_da: - # initial energy (Remember to extract this) - self.trunc_energy = self.keys_da['energy'] - if self.trunc_energy > 1: # ensure that it is given as percentage + try: + options = dict(self.keys_da['iteration']) + except: + options = dict([self.keys_da['iteration']]) + + self.data_misfit_tol = options.get('data_misfit_tol', 0.01) + self.trunc_energy = options.get('energy', 0.95) + self.step_tol = options.get('step_tol', 0.01) + self.gamma = options.get('gamma', 0.2) + self.gamma_max = options.get('gamma_max', 0.5) + self.gamma_factor = options.get('gamma_factor', 2.5) + + if self.trunc_energy > 1: # ensure that it is given as percentage self.trunc_energy /= 100. diff --git a/pipt/update_schemes/esmda.py b/pipt/update_schemes/esmda.py index b848096..7cefd16 100644 --- a/pipt/update_schemes/esmda.py +++ b/pipt/update_schemes/esmda.py @@ -31,7 +31,7 @@ def __init__(self, keys_da, keys_en, sim): Parameters ---------- - keys_da['mda'] : list + keys_da['mda'] : dict - tot_assim_steps: total number of iterations in MDA, e.g., 3 - inflation_param: covariance inflation factors, e.g., [2, 4, 4] @@ -222,17 +222,16 @@ def _ext_inflation_param(self): alpha: list Data covariance inflation factor """ - # Make sure MDA is a list - if not isinstance(self.keys_da['mda'][0], list): - mda_opts = [self.keys_da['mda']] - else: - mda_opts = self.keys_da['mda'] + try: + mda_opts = dict(self.keys_da['mda']) + except: + mda_opts = dict([self.keys_da['mda']]) # Check if INFLATION_PARAM has been provided, and if so, extract the value(s). If not, we set alpha to the # default value equal to the tot. no. assim. steps - if 'inflation_param' in list(zip(*mda_opts))[0]: + if 'inflation_param' in mda_opts: # Extract value - alpha_tmp = [item[1] for item in mda_opts if item[0] == 'inflation_param'][0] + alpha_tmp = mda_opts['inflation_param'] # If one value is given, we copy it to all assim. steps. If multiple values are given, we check the # number of parameters corresponds to tot. no. assim. steps @@ -279,22 +278,17 @@ def _ext_assim_steps(self): - ST 7/6-16 - ST 1/3-17: Changed to output list of assim. steps instead of just tot. assim. steps """ - # Make sure MDA is a list - if not isinstance(self.keys_da['mda'][0], list): - mda_opts = [self.keys_da['mda']] - else: - mda_opts = self.keys_da['mda'] + try: + mda_opts = dict(self.keys_da['mda']) + except: + mda_opts = dict([self.keys_da['mda']]) + # Check if 'max_iter' has been given; if not, give error (mandatory in ITERATION) - assert 'tot_assim_steps' in list( - zip(*mda_opts))[0], 'TOT_ASSIM_STEPS has not been given in MDA!' - - # Extract max. iter - tot_no_assim = int([item[1] - for item in mda_opts if item[0] == 'tot_assim_steps'][0]) - - # Make a list of assim. steps - assim_steps = list(range(tot_no_assim)) + try: + assim_steps = list(range(int(mda_opts['tot_assim_steps']))) + except KeyError: + raise AssertionError('TOT_ASSIM_STEPS has not been given in MDA!') # If it is a restart run, we remove simulations already done if self.restart is True: diff --git a/popt/cost_functions/ecalc_npv.py b/popt/cost_functions/ecalc_npv.py index d13d8df..420e6ca 100644 --- a/popt/cost_functions/ecalc_npv.py +++ b/popt/cost_functions/ecalc_npv.py @@ -44,9 +44,7 @@ def ecalc_npv(pred_data, **kwargs): report = kwargs.get('true_order', []) # Economic values - npv_const = {} - for name, value in keys_opt['npv_const']: - npv_const[name] = value + npv_const = dict(keys_opt['npv_const']) # Collect production data Qop = [] diff --git a/popt/cost_functions/ecalc_pareto_npv.py b/popt/cost_functions/ecalc_pareto_npv.py index 2847dd9..9375f9f 100644 --- a/popt/cost_functions/ecalc_pareto_npv.py +++ b/popt/cost_functions/ecalc_pareto_npv.py @@ -46,9 +46,7 @@ def ecalc_pareto_npv(pred_data, kwargs): report = kwargs.get('true_order', []) # Economic values - npv_const = {} - for name, value in keys_opt['npv_const']: - npv_const[name] = value + npv_const = dict(keys_opt['npv_const']) # Collect production data Qop = [] diff --git a/popt/cost_functions/npv.py b/popt/cost_functions/npv.py index dfb3f6f..bb18c8c 100644 --- a/popt/cost_functions/npv.py +++ b/popt/cost_functions/npv.py @@ -38,9 +38,7 @@ def npv(pred_data, **kwargs): report = kwargs.get('true_order', []) # Economic values - npv_const = {} - for name, value in keys_opt['npv_const']: - npv_const[name] = value + npv_const = dict(keys_opt['npv_const']) values = [] for i in np.arange(1, len(pred_data)): diff --git a/popt/cost_functions/ren_npv.py b/popt/cost_functions/ren_npv.py index e76c567..0035f4a 100644 --- a/popt/cost_functions/ren_npv.py +++ b/popt/cost_functions/ren_npv.py @@ -32,9 +32,7 @@ def ren_npv(pred_data, kwargs): report = kwargs.get('true_order', []) # Economic values - npv_const = {} - for name, value in keys_opt['npv_const']: - npv_const[name] = value + npv_const = dict(keys_opt['npv_const']) # Loop over timesteps values = [] diff --git a/popt/update_schemes/enopt.py b/popt/update_schemes/enopt.py index a397ad9..ae9ed13 100644 --- a/popt/update_schemes/enopt.py +++ b/popt/update_schemes/enopt.py @@ -8,7 +8,7 @@ # Internal imports from popt.misc_tools import optim_tools as ot from popt.loop.optimize import Optimize -import popt.update_schemes.optimizers as opt +import popt.update_schemes.subroutines.optimizers as opt class EnOpt(Optimize): diff --git a/popt/update_schemes/genopt.py b/popt/update_schemes/genopt.py index 2316f6d..408baa7 100644 --- a/popt/update_schemes/genopt.py +++ b/popt/update_schemes/genopt.py @@ -7,8 +7,8 @@ # Internal imports from popt.misc_tools import optim_tools as ot from popt.loop.optimize import Optimize -import popt.update_schemes.optimizers as opt -from popt.update_schemes.cma import CMA +import popt.update_schemes.subroutines.optimizers as opt +from popt.update_schemes.subroutines.cma import CMA class GenOpt(Optimize): diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index 57a2a09..e82b305 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -10,7 +10,7 @@ # Internal imports from popt.misc_tools import optim_tools as ot from popt.loop.optimize import Optimize -from popt.update_schemes.line_search_step import line_search, line_search_backtracking +from popt.update_schemes.subroutines import line_search, line_search_backtracking, bfgs_update, newton_cg # Some symbols for logger subk = '\u2096' @@ -37,8 +37,8 @@ def LineSearch(fun, x, jac, method='GD', hess=None, args=(), bounds=None, callba method: str Which optimization method to use. Default is 'GD' for 'Gradient Descent'. - Other options are 'BFGS' for the 'Broyden–Fletcher–Goldfarb–Shanno' method, - and 'Newton-CG'. + Other options are 'BFGS' for the 'Broyden-Fletcher-Goldfarb-Shanno' method, + and 'Newton-CG' for the Newton-conjugate gradient method. hess: callable, optional Hessian function, hess(x, *args). Default is None. @@ -340,7 +340,7 @@ def calc_update(self, iter_resamp=0): if self.method == 'BFGS': pk = - np.matmul(self.Hk_inv, self.jk) if self.method == 'Newton-CG': - pk = newton_cg(self.jk, Hk=self.Hk, xk=self.xk, jac=self._jac, eps=1e-4) + pk = newton_cg(self.jk, Hk=self.Hk, xk=self.xk, jac=self._jac, logger=self.logger.info) # porject search direction onto the feasible set if self.bounds is not None: @@ -560,82 +560,6 @@ def _set_max_step_size(self, pk, xk): -def bfgs_update(Hk, sk, yk): - """ - Perform the BFGS update of the inverse Hessian approximation. - - Parameters: - - Hk: np.ndarray, current inverse Hessian approximation (n x n) - - sk: np.ndarray, step vector (x_{k+1} - x_k), shape (n,) - - yk: np.ndarray, gradient difference (grad_{k+1} - grad_k), shape (n,) - - Returns: - - Hk_new: np.ndarray, updated inverse Hessian approximation - """ - sk = sk.reshape(-1, 1) - yk = yk.reshape(-1, 1) - rho = 1.0 / (yk.T @ sk) - - if rho <= 0: - print('Non-positive curvature detected. BFGS update skipped....') - return Hk - - I = np.eye(Hk.shape[0]) - Vk = I - rho * sk @ yk.T - Hk_new = Vk @ Hk @ Vk.T + rho * sk @ sk.T - - return Hk_new - -def newton_cg(gk, Hk=None, maxiter=None, **kwargs): - print('\nRunning Newton-CG subroutine...') - - if Hk is None: - jac = kwargs.get('jac') - eps = kwargs.get('eps', 1e-4) - xk = kwargs.get('xk') - - # define a finite difference approximation of the Hessian times a vector - def Hessd(d): - return (jac(xk + eps*d) - gk)/eps - - if maxiter is None: - maxiter = 20*gk.size # Same dfault as in scipy - - tol = min(0.5, np.sqrt(la.norm(gk)))*la.norm(gk) - z = 0 - r = gk - d = -r - - for j in range(maxiter): - print('iteration: ', j) - if Hk is None: - Hd = Hessd(d) - else: - Hd = np.matmul(Hk, d) - - dTHd = np.dot(d, Hd) - - if dTHd <= 0: - print('Negative curvature detected, terminating subroutine') - print('\n') - if j == 0: - return -gk - else: - return z - - rold = r - a = np.dot(r,r)/dTHd - z = z + a*d - r = r + a*Hd - - if la.norm(r) < tol: - print('Subroutine converged') - print('\n') - return z - - b = np.dot(r, r)/np.dot(rold, rold) - d = -r + b*d - diff --git a/popt/update_schemes/smcopt.py b/popt/update_schemes/smcopt.py index f67503b..f0944fe 100644 --- a/popt/update_schemes/smcopt.py +++ b/popt/update_schemes/smcopt.py @@ -6,7 +6,7 @@ # Internal imports from popt.loop.optimize import Optimize -import popt.update_schemes.optimizers as opt +import popt.update_schemes.subroutines.optimizers as opt from popt.misc_tools import optim_tools as ot diff --git a/popt/update_schemes/subroutines/__init__.py b/popt/update_schemes/subroutines/__init__.py new file mode 100644 index 0000000..eab7b84 --- /dev/null +++ b/popt/update_schemes/subroutines/__init__.py @@ -0,0 +1,3 @@ +from .subroutines import * +from .cma import * +from .optimizers import * \ No newline at end of file diff --git a/popt/update_schemes/cma.py b/popt/update_schemes/subroutines/cma.py similarity index 99% rename from popt/update_schemes/cma.py rename to popt/update_schemes/subroutines/cma.py index 0a99874..bee93e5 100644 --- a/popt/update_schemes/cma.py +++ b/popt/update_schemes/subroutines/cma.py @@ -2,6 +2,8 @@ import numpy as np from popt.misc_tools import optim_tools as ot +__all__ = ['CMA'] + class CMA: def __init__(self, ne, dim, alpha_mu=None, n_mu=None, alpha_1=None, alpha_c=None, corr_update=False, equal_weights=True): diff --git a/popt/update_schemes/optimizers.py b/popt/update_schemes/subroutines/optimizers.py similarity index 99% rename from popt/update_schemes/optimizers.py rename to popt/update_schemes/subroutines/optimizers.py index 83514cd..a211093 100644 --- a/popt/update_schemes/optimizers.py +++ b/popt/update_schemes/subroutines/optimizers.py @@ -1,6 +1,8 @@ """Gradient acceleration.""" import numpy as np +__all__ = ['GradientAscent', 'Adam', 'AdaMax', 'Steihaug', ] + class GradientAscent: r""" diff --git a/popt/update_schemes/line_search_step.py b/popt/update_schemes/subroutines/subroutines.py similarity index 79% rename from popt/update_schemes/line_search_step.py rename to popt/update_schemes/subroutines/subroutines.py index 3f4af38..5e00c28 100644 --- a/popt/update_schemes/line_search_step.py +++ b/popt/update_schemes/subroutines/subroutines.py @@ -1,9 +1,16 @@ -# This is a an implementation of the Line Search Algorithm (Alg. 3.5) in Numerical Optimization from Nocedal 2006. - import numpy as np +import numpy.linalg as la from functools import cache from scipy.optimize._linesearch import _quadmin, _cubicmin +__all__ = [ + 'line_search', + 'zoom', + 'line_search_backtracking', + 'bfgs_update', + 'newton_cg' +] + def line_search(step_size, xk, pk, fun, jac, fk=None, jk=None, **kwargs): ''' @@ -304,4 +311,87 @@ def phi(alpha): step_size *= rho # If we reached this point, the line search failed - return None, None, None, ls_nfev, ls_njev \ No newline at end of file + return None, None, None, ls_nfev, ls_njev + + +def bfgs_update(Hk, sk, yk): + """ + Perform the BFGS update of the inverse Hessian approximation. + + Parameters: + - Hk: np.ndarray, current inverse Hessian approximation (n x n) + - sk: np.ndarray, step vector (x_{k+1} - x_k), shape (n,) + - yk: np.ndarray, gradient difference (grad_{k+1} - grad_k), shape (n,) + + Returns: + - Hk_new: np.ndarray, updated inverse Hessian approximation + """ + sk = sk.reshape(-1, 1) + yk = yk.reshape(-1, 1) + rho = 1.0 / (yk.T @ sk) + + if rho <= 0: + print('Non-positive curvature detected. BFGS update skipped....') + return Hk + + I = np.eye(Hk.shape[0]) + Vk = I - rho * sk @ yk.T + Hk_new = Vk @ Hk @ Vk.T + rho * sk @ sk.T + + return Hk_new + +def newton_cg(gk, Hk=None, maxiter=None, **kwargs): + + # Check for logger + logger = kwargs.get('logger', None) + if logger is None: + logger = print + + logger('Running Newton-CG subroutine..........') + + if Hk is None: + jac = kwargs.get('jac') + eps = kwargs.get('eps', 1e-4) + xk = kwargs.get('xk') + + # define a finite difference approximation of the Hessian times a vector + def Hessd(d): + return (jac(xk + eps*d) - gk)/eps + + if maxiter is None: + maxiter = 20*gk.size # Same dfault as in scipy + + tol = min(0.5, np.sqrt(la.norm(gk)))*la.norm(gk) + z = 0 + r = gk + d = -r + + for j in range(maxiter): + logger(f'iteration: {j}') + if Hk is None: + Hd = Hessd(d) + else: + Hd = np.matmul(Hk, d) + + dTHd = np.dot(d, Hd) + + if dTHd <= 0: + logger('Negative curvature detected, terminating subroutine') + logger('') + if j == 0: + return -gk + else: + return z + + rold = r + a = np.dot(r,r)/dTHd + z = z + a*d + r = r + a*Hd + + if la.norm(r) < tol: + logger('Subroutine converged') + logger('') + return z + + b = np.dot(r, r)/np.dot(rold, rold) + d = -r + b*d \ No newline at end of file From 3cdd192a8f84eadb71ece6026fcfdf87a5025a4e Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 8 Sep 2025 11:18:09 +0200 Subject: [PATCH 14/94] changed @cache to @lru_cache --- popt/update_schemes/subroutines/subroutines.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/popt/update_schemes/subroutines/subroutines.py b/popt/update_schemes/subroutines/subroutines.py index 5e00c28..ddd1105 100644 --- a/popt/update_schemes/subroutines/subroutines.py +++ b/popt/update_schemes/subroutines/subroutines.py @@ -1,6 +1,6 @@ import numpy as np import numpy.linalg as la -from functools import cache +from functools import lru_cache from scipy.optimize._linesearch import _quadmin, _cubicmin __all__ = [ @@ -85,7 +85,7 @@ def line_search(step_size, xk, pk, fun, jac, fk=None, jk=None, **kwargs): assert step_size <= amax, "Initial step size must be less than or equal to amax." # Define phi and derivative of phi - @cache + @lru_cache(maxsize=None) def phi(alpha): global ls_nfev if (alpha == 0): @@ -100,7 +100,7 @@ def phi(alpha): ls_nfev += 1 return phi.fun_val - @cache + @lru_cache(maxsize=None) def dphi(alpha): global ls_njev if (alpha == 0): @@ -281,7 +281,7 @@ def line_search_backtracking(step_size, xk, pk, fun, jac, fk=None, jk=None, **kw c1 = kwargs.get('c1', 1e-4) # Define phi and derivative of phi - @cache + @lru_cache(maxsize=None) def phi(alpha): global ls_nfev if (alpha == 0): From 1061cbfaa39df9a3e50ac230e8db33cfd54e880b Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 11 Sep 2025 09:47:47 +0200 Subject: [PATCH 15/94] branch commit --- ensemble/__init__.py | 2 +- ensemble/ensemble.py | 210 +++++++++++--------------- input_output/read_config.py | 2 +- pipt/loop/ensemble.py | 34 +++-- pipt/misc_tools/cov_regularization.py | 4 +- pipt/update_schemes/enrml.py | 9 +- 6 files changed, 123 insertions(+), 138 deletions(-) diff --git a/ensemble/__init__.py b/ensemble/__init__.py index 2145053..c8b7821 100644 --- a/ensemble/__init__.py +++ b/ensemble/__init__.py @@ -1 +1 @@ -"""Multiple realisations management.""" +"""Multiple realisations management.""" \ No newline at end of file diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index d4dc46a..aad7902 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -18,6 +18,8 @@ # Internal imports import pipt.misc_tools.analysis_tools as at +import pipt.misc_tools.extract_tools as extract + from geostat.decomp import Cholesky # Making realizations from pipt.misc_tools import cov_regularization from pipt.misc_tools import wavelet_tools as wt @@ -25,6 +27,7 @@ from misc.system_tools.environ_var import OpenBlasSingleThread # Single threaded OpenBLAS runs + class Ensemble: """ Class for organizing misc. variables and simulator for an ensemble-based inversion run. Here, the forecast step @@ -56,11 +59,13 @@ def __init__(self, keys_en, sim, redund_sim=None): self.aux_input = None # Setup logger - logging.basicConfig(level=logging.INFO, - filename='pet_logger.log', - filemode='w', - format='%(asctime)s : %(levelname)s : %(name)s : %(message)s', - datefmt='%Y-%m-%d %H:%M:%S') + logging.basicConfig( + level=logging.INFO, + filename='pet_logger.log', + filemode='w', + format='%(asctime)s : %(levelname)s : %(name)s : %(message)s', + datefmt='%Y-%m-%d %H:%M:%S' + ) self.logger = logging.getLogger('PET') # Check if folder contains any En_ files, and remove them! @@ -117,7 +122,7 @@ def __init__(self, keys_en, sim, redund_sim=None): self.disable_tqdm = False # extract information that is given for the prior model - self.prior_info = self._extract_prior_info() + self.prior_info = extract.extract_prior_info(self.keys_en) # Calculate initial ensemble if IMPORTSTATICVAR has not been given in init. file. # Prior info. on state variables must be given by PRIOR_ keyword. @@ -143,7 +148,12 @@ def __init__(self, keys_en, sim, redund_sim=None): print('\033[1;33mInput states have different ensemble size\033[1;m') sys.exit(1) self.ne = min(tmp_ne) - self._ext_ml_info() + + # extract multi-level info (if needed) + if 'multilevel' in self.keys_en: + ml_info = extract.extract_multilevel_info(self.keys_en) + self.multilevel, self.tot_level, self.ml_ne, self.ML_error_corr, self.error_comp_scheme, self.ML_corr_done = ml_info + #self._ext_ml_info() def _ext_ml_info(self): ''' @@ -172,117 +182,7 @@ def _ext_ml_info(self): self.error_comp_scheme = self.keys_en['multilevel'][i][2] self.ML_corr_done = False - def _extract_prior_info(self) -> dict: - ''' - Extract prior information on STATE from keyword(s) PRIOR_. - ''' - - # Get state names as list - state_names = self.keys_en['state'] - if not isinstance(state_names, list): state_names = [state_names] - - # Check if PRIOR_ exists for each entry in state - for name in state_names: - assert f'prior_{name}' in self.keys_en, \ - 'PRIOR_{0} is missing! This keyword is needed to make initial ensemble for {0} entered in ' \ - 'STATE'.format(name.upper()) - - # define dict to store prior information in - prior_info = {name: None for name in state_names} - - # loop over state priors - for name in state_names: - prior = self.keys_en[f'prior_{name}'] - - # Check if is a list (old way) - if isinstance(prior, list): - # list of lists - old way of inputting prior information - prior_dict = {} - for i, opt in enumerate(list(zip(*prior))[0]): - if opt == 'limits': - prior_dict[opt] = prior[i][1:] - else: - prior_dict[opt] = prior[i][1] - prior = prior_dict - else: - assert isinstance(prior, dict), 'PRIOR_{0} must be a dictionary or list of lists!'.format(name.upper()) - - - # load mean if in file - if isinstance(prior['mean'], str): - assert prior['mean'].endswith('.npz'), 'File name does not end with \'.npz\'!' - load_file = np.load(prior['mean']) - assert len(load_file.files) == 1, \ - 'More than one variable located in {0}. Only the mean vector can be stored in the .npz file!' \ - .format(prior['mean']) - prior['mean'] = load_file[load_file.files[0]] - else: # Single number inputted, make it a list if not already - if not isinstance(prior['mean'], list): - prior['mean'] = [prior['mean']] - - # loop over keys in prior - for key in prior.keys(): - # ensure that entry is a list - if (not isinstance(prior[key], list)) and (key != 'mean'): - prior[key] = [prior[key]] - - # change the name of some keys - prior['variance'] = prior.pop('var', None) - prior['corr_length'] = prior.pop('range', None) - - # process grid - if 'grid' in prior: - grid_dim = prior['grid'] - - # check if 3D-grid - if (len(grid_dim) == 3) and (grid_dim[2] > 1): - nz = int(grid_dim[2]) - prior['nz'] = nz - prior['nx'] = int(grid_dim[0]) - prior['ny'] = int(grid_dim[1]) - - - # Check mean when values have been inputted directly (not when mean has been loaded) - mean = prior['mean'] - if isinstance(mean, list) and len(mean) < nz: - # Check if it is more than one entry and give error - assert len(mean) == 1, \ - 'Information from MEAN has been given for {0} layers, whereas {1} is needed!' \ - .format(len(mean), nz) - - # Only 1 entry; copy this to all layers - print( - '\033[1;33mSingle entry for MEAN will be copied to all {0} layers\033[1;m'.format(nz)) - prior['mean'] = mean * nz - - #check if info. has been given on all layers. In the case it has not been given, we just copy the info. given. - for key in ['vario', 'variance', 'aniso', 'angle', 'corr_length']: - if key in prior.keys(): - val = prior[key] - if len(val) < nz: - # Check if it is more than one entry and give error - assert len(val) == 1, \ - 'Information from {0} has been given for {1} layers, whereas {2} is needed!' \ - .format(key.upper(), len(val), nz) - - # Only 1 entry; copy this to all layers - print( - '\033[1;33mSingle entry for {0} will be copied to all {1} layers\033[1;m'.format(key.upper(), nz)) - prior[key] = val * nz - - else: - prior['nx'] = int(grid_dim[0]) - prior['ny'] = int(grid_dim[1]) - prior['nz'] = 1 - - prior.pop('grid', None) - - # add prior to prior_info - prior_info[name] = prior - - return prior_info - - + def gen_init_ensemble(self): """ Generate the initial ensemble of (joint) state vectors using the GeoStat class in the "geostat" package. @@ -353,6 +253,80 @@ def gen_init_ensemble(self): # Save the ensemble for later inspection np.savez('prior.npz', **self.state) + def generate_state_ensemble(self): + # Initialize GeoStat + generator = Cholesky() + + # Initialize state and cov + enX = {} + covX = {} + + # Loop over statenames in prior_info + for name in self.prior_info.keys(): + # Init. indices to pick out correct mean vector for each layer + ind_end = 0 + + # Extract info. + nx = self.prior_info[name].get('nx', 0) + ny = self.prior_info[name].get('ny', 0) + nz = self.prior_info[name].get('nz', 0) + mean = self.prior_info[name].get('mean', None) + + if nx == ny == 0: # assume ensemble will be generated elsewhere if dimensions are zero + break + + variance = self.prior_info[name].get('variance', None) + corr_length = self.prior_info[name].get('corr_length', None) + aniso = self.prior_info[name].get('aniso', None) + vario = self.prior_info[name].get('vario', None) + angle = self.prior_info[name].get('angle', None) + limits= self.prior_info[name].get('limits',None) + + # Loop over nz to make layers of 2D priors + for i in range(self.prior_info[name]['nz']): + # If mean is scalar, no covariance matrix is needed + + if type(self.prior_info[name]['mean']).__module__ == 'numpy': + # Generate covariance matrix + cov = generator.gen_cov2d( + nx, + ny, + variance[i], + corr_length[i], + aniso[i], + angle[i], + vario[i] + ) + else: + cov = np.array(variance[i]) + + # Pick out the mean vector for the current layer + ind_start = ind_end + ind_end = int((i + 1) * (len(mean) / nz)) + mean_layer = mean[ind_start:ind_end] + + # Generate realizations. If LIMITS have been entered, they must be taken account for here + if limits is None: + real = generator.gen_real(mean_layer, cov, self.ne) + else: + real = generator.gen_real(mean_layer, cov, self.ne, limits) + + # Stack realizations for each layer + if i == 0: + real_out = real + else: + real_out = np.vstack((real_out, real)) + + # Fill in dicts + enX[name] = real_out + covX[name]= cov + + idX_state = {key: enX[key].shape for key in enX} + enX = np.vstack([enX[key] for key in enX]) + + return enX, idX_state, covX + + def get_list_assim_steps(self): """ Returns list of assimilation steps. Useful in a 'loop'-script. diff --git a/input_output/read_config.py b/input_output/read_config.py index ba00da2..842d0dc 100644 --- a/input_output/read_config.py +++ b/input_output/read_config.py @@ -61,7 +61,7 @@ def ndarray_constructor(loader, node): keys_pr = y['optim'] check_mand_keywords_opt(keys_pr) elif 'dataassim' in y.keys(): - keys_pr = y['datasssim'] + keys_pr = y['dataassim'] check_mand_keywords_da(keys_pr) else: raise KeyError diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index 5a2aaed..30376c1 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -17,8 +17,9 @@ from ensemble.ensemble import Ensemble as PETEnsemble import misc.read_input_csv as rcsv from pipt.misc_tools import wavelet_tools as wt -from pipt.misc_tools import cov_regularization +from pipt.misc_tools.cov_regularization import localization, _calc_distance import pipt.misc_tools.analysis_tools as at +import pipt.misc_tools.extract_tools as extract class Ensemble(PETEnsemble): @@ -113,15 +114,24 @@ def __init__(self, keys_da, keys_en, sim): # Initialize localization if 'localization' in self.keys_da: - self.localization = cov_regularization.localization(self.keys_da['localization'], - self.keys_da['truedataindex'], - self.keys_da['datatype'], - self.keys_da['staticvar'], - self.ne) + + if isinstance(self.keys_da['localization'], dict): + # Make 2D list of Dict (this should only be temporary) + loc_info = [[key, value] for key, value in self.keys_da['localization'].items()] + self.keys_da['localization'] = loc_info + + self.localization = localization( + self.keys_da['localization'], + self.keys_da['truedataindex'], + self.keys_da['datatype'], + self.keys_da['staticvar'], + self.ne + ) # Initialize local analysis if 'localanalysis' in self.keys_da: - self.local_analysis = at.init_local_analysis( - init=self.keys_da['localanalysis'], state=self.state.keys()) + self.local_analysis = extract.extract_local_analysis_info(self.keys_da['localanalysis'], self.state.keys()) + #self.local_analysis = at.init_local_analysis( + # init=self.keys_da['localanalysis'], state=self.state.keys()) self.pred_data = [{k: np.zeros((1, self.ne), dtype='float32') for k in self.keys_da['datatype']} for _ in self.obs_data] @@ -769,7 +779,7 @@ def local_analysis_update(self): self.list_datatypes = [elem for elem in self.list_datatypes if elem in self.local_analysis['update_mask'][state]] self.list_states = [deepcopy(state)] - self._ext_state() # scaling for this state + self._ext_scaling() # scaling for this state if 'localization' in self.keys_da: self.localization.loc_info['field'] = self.state_scaling.shape del self.cov_data @@ -799,7 +809,7 @@ def local_analysis_update(self): elem in self.local_analysis['update_mask'][state][state_indx]] if len(self.list_datatypes): self.list_states = [deepcopy(state)] - self._ext_state() # scaling for this state + self._ext_scaling() # scaling for this state if 'localization' in self.keys_da: self.localization.loc_info['field'] = self.state_scaling.shape del self.cov_data @@ -826,7 +836,7 @@ def local_analysis_update(self): for state in self.local_analysis['cell_parameter']: self.list_states = [deepcopy(state)] - self._ext_state() # scaling for this state + self._ext_scaling() # scaling for this state orig_state_scaling = deepcopy(self.state_scaling) param_position = self.local_analysis['parameter_position'][state] field_size = param_position.shape @@ -863,7 +873,7 @@ def local_analysis_update(self): if 'localization' in self.keys_da: self.localization.loc_info['field'] = ( len(self.cell_index),) - self.localization.loc_info['distance'] = cov_regularization._calc_distance( + self.localization.loc_info['distance'] = _calc_distance( self.local_analysis['data_position'], self.local_analysis['unique'], current_data_list, self.assim_index, diff --git a/pipt/misc_tools/cov_regularization.py b/pipt/misc_tools/cov_regularization.py index 82b48c3..f4d981a 100644 --- a/pipt/misc_tools/cov_regularization.py +++ b/pipt/misc_tools/cov_regularization.py @@ -37,6 +37,7 @@ from shutil import rmtree from scipy import sparse from scipy.spatial import distance +from typing import Union # internal import import pipt.misc_tools.analysis_tools as at @@ -47,13 +48,12 @@ class localization(): # TODO: Check field dimensions, should always ensure that we can provide i ,j ,k (x, y, z) ### - def __init__(self, parsed_info, assimIndex, data_typ, free_parameter, ne): + def __init__(self, parsed_info: list, assimIndex, data_typ, free_parameter, ne): """ Format the parsed info from the input file, and generate the unique localization masks """ # if the next element is a .p file (pickle), assume that this has been correctly formated and can be automatically # imported. NB: it is important that we use the pickle format since we have a dictionary containing dictionaries - # to make this as robust as possible, we always try to load the file try: if parsed_info[1][0].upper() == 'AUTOADALOC': diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index a570217..1b7101f 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -13,7 +13,7 @@ import inspect import numpy as np import copy as cp -from scipy.linalg import cholesky, solve +from scipy.linalg import cholesky, solve, inv, lu_solve, lu_factor import importlib.util @@ -38,6 +38,7 @@ class margIS_update: pass # Internal imports +from pipt.misc_tools.analysis_tools import aug_state class lmenrmlMixIn(Ensemble): @@ -720,7 +721,7 @@ def calc_analysis(self): else: _, self.aug_pred_data = at.aug_obs_pred_data( - self.obs_data, self.pred_data, assim_index, self.list_datatypes) + self.obs_data, self.pred_data, self.assim_index, self.list_datatypes) # Mean pred_data and perturbation matrix with scaling mean_preddata = np.mean(self.aug_pred_data, 1) @@ -1055,11 +1056,11 @@ def check_convergence(self): self.lam = self.lam + (self.lam_max - self.lam) * \ 2 ** (-(self.iteration) / (self.gamma - 1)) success = True - self.current_state = deepcopy(self.state) + self.current_state = cp.deepcopy(self.state) elif self.data_misfit < self.prev_data_misfit and self.data_misfit_std >= self.prev_data_misfit_std: # Accept itaration, but keep lam the same success = True - self.current_state = deepcopy(self.state) + self.current_state = cp.deepcopy(self.state) else: # Reject iteration, and decrease step length self.lam = self.lam / self.gamma success = False From 37796f4a1c47e2e259957fd186f2a03a75c46d42 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 11 Sep 2025 09:48:28 +0200 Subject: [PATCH 16/94] branch commit --- ensemble/ensemble.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index aad7902..82830df 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -128,9 +128,10 @@ def __init__(self, keys_en, sim, redund_sim=None): # Prior info. on state variables must be given by PRIOR_ keyword. if 'importstaticvar' not in self.keys_en: self.ne = int(self.keys_en['ne']) + self.enX, self.idX, self.cov_prior = self.generate_state_ensemble() # Output = self.state, self.cov_prior - self.gen_init_ensemble() + #self.gen_init_ensemble() else: # State variable imported as a Numpy save file From 89f5f38be5f5d563505243b03db2907634c5b987 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 11 Sep 2025 09:57:05 +0200 Subject: [PATCH 17/94] re-added stuff that got deleted --- ensemble/ensemble.py | 119 +++---------------------------------------- 1 file changed, 6 insertions(+), 113 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index d4dc46a..3898b82 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -18,6 +18,7 @@ # Internal imports import pipt.misc_tools.analysis_tools as at +import pipt.misc_tools.extract_tools as extract from geostat.decomp import Cholesky # Making realizations from pipt.misc_tools import cov_regularization from pipt.misc_tools import wavelet_tools as wt @@ -117,7 +118,7 @@ def __init__(self, keys_en, sim, redund_sim=None): self.disable_tqdm = False # extract information that is given for the prior model - self.prior_info = self._extract_prior_info() + self.prior_info = extract.extract_prior_info() # Calculate initial ensemble if IMPORTSTATICVAR has not been given in init. file. # Prior info. on state variables must be given by PRIOR_ keyword. @@ -143,7 +144,9 @@ def __init__(self, keys_en, sim, redund_sim=None): print('\033[1;33mInput states have different ensemble size\033[1;m') sys.exit(1) self.ne = min(tmp_ne) - self._ext_ml_info() + + self.multilevel, self.tot_level, self.ml_ne, self.ML_error_corr, self.error_comp_scheme, self.ML_corr_done = extract.extract_multilevel_info() + #self._ext_ml_info() def _ext_ml_info(self): ''' @@ -172,117 +175,7 @@ def _ext_ml_info(self): self.error_comp_scheme = self.keys_en['multilevel'][i][2] self.ML_corr_done = False - def _extract_prior_info(self) -> dict: - ''' - Extract prior information on STATE from keyword(s) PRIOR_. - ''' - - # Get state names as list - state_names = self.keys_en['state'] - if not isinstance(state_names, list): state_names = [state_names] - - # Check if PRIOR_ exists for each entry in state - for name in state_names: - assert f'prior_{name}' in self.keys_en, \ - 'PRIOR_{0} is missing! This keyword is needed to make initial ensemble for {0} entered in ' \ - 'STATE'.format(name.upper()) - - # define dict to store prior information in - prior_info = {name: None for name in state_names} - - # loop over state priors - for name in state_names: - prior = self.keys_en[f'prior_{name}'] - - # Check if is a list (old way) - if isinstance(prior, list): - # list of lists - old way of inputting prior information - prior_dict = {} - for i, opt in enumerate(list(zip(*prior))[0]): - if opt == 'limits': - prior_dict[opt] = prior[i][1:] - else: - prior_dict[opt] = prior[i][1] - prior = prior_dict - else: - assert isinstance(prior, dict), 'PRIOR_{0} must be a dictionary or list of lists!'.format(name.upper()) - - - # load mean if in file - if isinstance(prior['mean'], str): - assert prior['mean'].endswith('.npz'), 'File name does not end with \'.npz\'!' - load_file = np.load(prior['mean']) - assert len(load_file.files) == 1, \ - 'More than one variable located in {0}. Only the mean vector can be stored in the .npz file!' \ - .format(prior['mean']) - prior['mean'] = load_file[load_file.files[0]] - else: # Single number inputted, make it a list if not already - if not isinstance(prior['mean'], list): - prior['mean'] = [prior['mean']] - - # loop over keys in prior - for key in prior.keys(): - # ensure that entry is a list - if (not isinstance(prior[key], list)) and (key != 'mean'): - prior[key] = [prior[key]] - - # change the name of some keys - prior['variance'] = prior.pop('var', None) - prior['corr_length'] = prior.pop('range', None) - - # process grid - if 'grid' in prior: - grid_dim = prior['grid'] - - # check if 3D-grid - if (len(grid_dim) == 3) and (grid_dim[2] > 1): - nz = int(grid_dim[2]) - prior['nz'] = nz - prior['nx'] = int(grid_dim[0]) - prior['ny'] = int(grid_dim[1]) - - - # Check mean when values have been inputted directly (not when mean has been loaded) - mean = prior['mean'] - if isinstance(mean, list) and len(mean) < nz: - # Check if it is more than one entry and give error - assert len(mean) == 1, \ - 'Information from MEAN has been given for {0} layers, whereas {1} is needed!' \ - .format(len(mean), nz) - - # Only 1 entry; copy this to all layers - print( - '\033[1;33mSingle entry for MEAN will be copied to all {0} layers\033[1;m'.format(nz)) - prior['mean'] = mean * nz - - #check if info. has been given on all layers. In the case it has not been given, we just copy the info. given. - for key in ['vario', 'variance', 'aniso', 'angle', 'corr_length']: - if key in prior.keys(): - val = prior[key] - if len(val) < nz: - # Check if it is more than one entry and give error - assert len(val) == 1, \ - 'Information from {0} has been given for {1} layers, whereas {2} is needed!' \ - .format(key.upper(), len(val), nz) - - # Only 1 entry; copy this to all layers - print( - '\033[1;33mSingle entry for {0} will be copied to all {1} layers\033[1;m'.format(key.upper(), nz)) - prior[key] = val * nz - - else: - prior['nx'] = int(grid_dim[0]) - prior['ny'] = int(grid_dim[1]) - prior['nz'] = 1 - - prior.pop('grid', None) - - # add prior to prior_info - prior_info[name] = prior - - return prior_info - - + def gen_init_ensemble(self): """ Generate the initial ensemble of (joint) state vectors using the GeoStat class in the "geostat" package. From 16ac7480ac6714c24388d43e9ce3b662d6b5aa1b Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 11 Sep 2025 10:10:10 +0200 Subject: [PATCH 18/94] fixed what got lost --- ensemble/ensemble.py | 82 ++++---------------------------------------- 1 file changed, 7 insertions(+), 75 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 383f208..4c0e724 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -18,6 +18,7 @@ # Internal imports import pipt.misc_tools.analysis_tools as at +import pipt.misc_tools.extract_tools as extract from geostat.decomp import Cholesky # Making realizations from pipt.misc_tools import cov_regularization from pipt.misc_tools import wavelet_tools as wt @@ -120,7 +121,7 @@ def __init__(self, keys_en, sim, redund_sim=None): self.disable_tqdm = False # extract information that is given for the prior model - self.prior_info = self._extract_prior_info() + self.prior_info = extract.extract_prior_info(self.keys_en) # Calculate initial ensemble if IMPORTSTATICVAR has not been given in init. file. # Prior info. on state variables must be given by PRIOR_ keyword. @@ -146,7 +147,11 @@ def __init__(self, keys_en, sim, redund_sim=None): print('\033[1;33mInput states have different ensemble size\033[1;m') sys.exit(1) self.ne = min(tmp_ne) - self._ext_ml_info() + gi + if 'multilevel' in self.keys_en: + ml_info = extract.extract_multilevel_info(self.keys_en) + self.multilevel, self.tot_level, self.ml_ne, self.ML_error_corr, self.error_comp_scheme, self.ML_corr_done = ml_info + #self._ext_ml_info() def _ext_ml_info(self): ''' @@ -246,79 +251,6 @@ def gen_init_ensemble(self): # Save the ensemble for later inspection np.savez('prior.npz', **self.state) - def generate_state_ensemble(self): - # Initialize GeoStat - generator = Cholesky() - - # Initialize state and cov - enX = {} - covX = {} - - # Loop over statenames in prior_info - for name in self.prior_info.keys(): - # Init. indices to pick out correct mean vector for each layer - ind_end = 0 - - # Extract info. - nx = self.prior_info[name].get('nx', 0) - ny = self.prior_info[name].get('ny', 0) - nz = self.prior_info[name].get('nz', 0) - mean = self.prior_info[name].get('mean', None) - - if nx == ny == 0: # assume ensemble will be generated elsewhere if dimensions are zero - break - - variance = self.prior_info[name].get('variance', None) - corr_length = self.prior_info[name].get('corr_length', None) - aniso = self.prior_info[name].get('aniso', None) - vario = self.prior_info[name].get('vario', None) - angle = self.prior_info[name].get('angle', None) - limits= self.prior_info[name].get('limits',None) - - # Loop over nz to make layers of 2D priors - for i in range(self.prior_info[name]['nz']): - # If mean is scalar, no covariance matrix is needed - - if type(self.prior_info[name]['mean']).__module__ == 'numpy': - # Generate covariance matrix - cov = generator.gen_cov2d( - nx, - ny, - variance[i], - corr_length[i], - aniso[i], - angle[i], - vario[i] - ) - else: - cov = np.array(variance[i]) - - # Pick out the mean vector for the current layer - ind_start = ind_end - ind_end = int((i + 1) * (len(mean) / nz)) - mean_layer = mean[ind_start:ind_end] - - # Generate realizations. If LIMITS have been entered, they must be taken account for here - if limits is None: - real = generator.gen_real(mean_layer, cov, self.ne) - else: - real = generator.gen_real(mean_layer, cov, self.ne, limits) - - # Stack realizations for each layer - if i == 0: - real_out = real - else: - real_out = np.vstack((real_out, real)) - - # Fill in dicts - enX[name] = real_out - covX[name]= cov - - idX_state = {key: enX[key].shape for key in enX} - enX = np.vstack([enX[key] for key in enX]) - - return enX, idX_state, covX - def get_list_assim_steps(self): """ From 390c26bf9be2a6f03f62a7cbbb9fae800b8257bd Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 11 Sep 2025 10:10:25 +0200 Subject: [PATCH 19/94] fixed what got lost --- ensemble/ensemble.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 4c0e724..3613e5a 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -147,7 +147,7 @@ def __init__(self, keys_en, sim, redund_sim=None): print('\033[1;33mInput states have different ensemble size\033[1;m') sys.exit(1) self.ne = min(tmp_ne) - gi + if 'multilevel' in self.keys_en: ml_info = extract.extract_multilevel_info(self.keys_en) self.multilevel, self.tot_level, self.ml_ne, self.ML_error_corr, self.error_comp_scheme, self.ML_corr_done = ml_info From 42ad4abc4603869d349cca83de9e9d4abc54129c Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 11 Sep 2025 10:54:51 +0200 Subject: [PATCH 20/94] fixed confusing naming of input keys --- input_output/read_config.py | 4 ++-- pipt/loop/ensemble.py | 2 +- pipt/pipt_init.py | 4 ++-- pipt/update_schemes/enkf.py | 4 ++-- pipt/update_schemes/enrml.py | 8 ++++---- pipt/update_schemes/es.py | 4 ++-- 6 files changed, 13 insertions(+), 13 deletions(-) diff --git a/input_output/read_config.py b/input_output/read_config.py index 842d0dc..0d9d10d 100644 --- a/input_output/read_config.py +++ b/input_output/read_config.py @@ -55,7 +55,7 @@ def ndarray_constructor(loader, node): keys_en = y['ensemble'] check_mand_keywords_en(keys_en) else: - keys_en = None + keys_en = {} if 'optim' in y.keys(): keys_pr = y['optim'] @@ -109,7 +109,7 @@ def read_toml(init_file): keys_en = t['ensemble'] check_mand_keywords_en(keys_en) else: - keys_en = None + keys_en = {} if 'optim' in t.keys(): keys_pr = t['optim'] check_mand_keywords_opt(keys_pr) diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index 30376c1..e2b3482 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -63,7 +63,7 @@ def __init__(self, keys_da, keys_en, sim): # do the initiallization of the PETensemble - super(Ensemble, self).__init__(keys_en, sim) + super(Ensemble, self).__init__(keys_da|keys_en, sim) # set logger self.logger = logging.getLogger('PET.PIPT') diff --git a/pipt/pipt_init.py b/pipt/pipt_init.py index afeab6d..da4ff3d 100644 --- a/pipt/pipt_init.py +++ b/pipt/pipt_init.py @@ -5,7 +5,7 @@ from importlib import import_module -def init_da(da_input, fwd_input, sim): +def init_da(da_input, en_input, sim): "initialize the ensemble object based on the DA inputs" assert len( @@ -15,4 +15,4 @@ def init_da(da_input, fwd_input, sim): da_input['daalg'][0]), f'{da_input["daalg"][1]}_{da_input["analysis"]}') # Init. update scheme class, and get an object of that class - return da_import(da_input, fwd_input, sim) + return da_import(da_input, en_input, sim) diff --git a/pipt/update_schemes/enkf.py b/pipt/update_schemes/enkf.py index 4d72638..fa77b9c 100644 --- a/pipt/update_schemes/enkf.py +++ b/pipt/update_schemes/enkf.py @@ -23,13 +23,13 @@ class enkfMixIn(Ensemble): ordering of data. If only one-step EnKF is to be done, use `es` instead. """ - def __init__(self, keys_da, keys_fwd, sim): + def __init__(self, keys_da, keys_en, sim): """ The class is initialized by passing the PIPT init. file upwards in the hierarchy to be read and parsed in `pipt.input_output.pipt_init.ReadInitFile`. """ # Pass the init_file upwards in the hierarchy - super().__init__(keys_da, keys_fwd, sim) + super().__init__(keys_da, keys_en, sim) self.prev_data_misfit = None diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index 1b7101f..7b52267 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -47,13 +47,13 @@ class lmenrmlMixIn(Ensemble): update_methods_ns. This class must therefore facititate many different update schemes. """ - def __init__(self, keys_da, keys_fwd, sim): + def __init__(self, keys_da, keys_en, sim): """ The class is initialized by passing the PIPT init. file upwards in the hierarchy to be read and parsed in `pipt.input_output.pipt_init.ReadInitFile`. """ # Pass the init_file upwards in the hierarchy - super().__init__(keys_da, keys_fwd, sim) + super().__init__(keys_da, keys_en, sim) if self.restart is False: # Save prior state in separate variable @@ -288,13 +288,13 @@ class gnenrmlMixIn(Ensemble): update_methods_ns. This class must therefore facititate many different update schemes. """ - def __init__(self, keys_da, keys_fwd, sim): + def __init__(self, keys_da, keys_en, sim): """ The class is initialized by passing the PIPT init. file upwards in the hierarchy to be read and parsed in `pipt.input_output.pipt_init.ReadInitFile`. """ # Pass the init_file upwards in the hierarchy - super().__init__(keys_da, keys_fwd, sim) + super().__init__(keys_da, keys_en, sim) if self.restart is False: # Save prior state in separate variable diff --git a/pipt/update_schemes/es.py b/pipt/update_schemes/es.py index c633eb4..c3a2bf5 100644 --- a/pipt/update_schemes/es.py +++ b/pipt/update_schemes/es.py @@ -21,13 +21,13 @@ class esMixIn(): structure and `enkf` is inherited to get `calc_analysis`, so we do not have to implement it again. """ - def __init__(self, keys_da, keys_fwd, sim): + def __init__(self, keys_da, keys_en, sim): """ The class is initialized by passing the PIPT init. file upwards in the hierarchy to be read and parsed in `pipt.input_output.pipt_init.ReadInitFile`. """ # Pass init. file to Simultaneous parent class (Python searches parent classes from left to right). - super().__init__(keys_da, keys_fwd, sim) + super().__init__(keys_da, keys_en, sim) if self.restart is False: # At the moment, the iterative loop is threated as an iterative smoother an thus we check if assim. indices From 16acee2960f307d1dfabeb2e975473d8ff2dd4bf Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 12 Sep 2025 11:21:09 +0200 Subject: [PATCH 21/94] more input stuff --- pipt/loop/assimilation.py | 28 ++-- pipt/loop/ensemble.py | 53 ++------ pipt/misc_tools/cov_regularization.py | 123 +++++++++++++++--- pipt/misc_tools/extract_tools.py | 121 +++++++++++++---- .../update_methods_ns/approx_update.py | 6 +- 5 files changed, 234 insertions(+), 97 deletions(-) diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index f5caa72..d9b7e79 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -16,7 +16,7 @@ from importlib import import_module # Internal imports -from pipt.misc_tools import qaqc_tools +from pipt.misc_tools.qaqc_tools import QAQC from pipt.loop.ensemble import Ensemble from misc.system_tools.environ_var import OpenBlasSingleThread from pipt.misc_tools import analysis_tools as at @@ -83,15 +83,20 @@ def run(self): success_iter = True # Initiallize progressbar - pbar_out = tqdm(total=self.max_iter, - desc='Iterations (Obj. func. val: )', position=0) + pbar_out = tqdm(total=self.max_iter, desc='Iterations (Obj. func. val: )', position=0) # Check if we want to perform a Quality Assurance of the forecast qaqc = None - if 'qa' in self.ensemble.sim.input_dict or 'qc' in self.ensemble.keys_da: - qaqc = qaqc_tools.QAQC({**self.ensemble.keys_da, **self.ensemble.sim.input_dict}, - self.ensemble.obs_data, self.ensemble.datavar, self.ensemble.logger, - self.ensemble.prior_info, self.ensemble.sim, self.ensemble.prior_state) + if ('qa' in self.ensemble.sim.input_dict) or ('qc' in self.ensemble.keys_da): + qaqc = QAQC( + self.ensemble.keys_da|self.ensemble.sim.input_dict, + self.ensemble.obs_data, + self.ensemble.datavar, + self.ensemble.logger, + self.ensemble.prior_info, + self.ensemble.sim, + self.ensemble.prior_state + ) # Run a while loop until max. iterations or convergence is reached while self.ensemble.iteration < self.max_iter and conv is False: @@ -107,20 +112,17 @@ def run(self): if 'qa' in self.ensemble.keys_da: # Check if we want to perform a Quality Assurance of the forecast # set updated prediction, state and lam - qaqc.set(self.ensemble.pred_data, - self.ensemble.state, self.ensemble.lam) + qaqc.set(self.ensemble.pred_data, self.ensemble.state, self.ensemble.lam) # Level 1,2 all data, and subspace qaqc.calc_mahalanobis((1, 'time', 2, 'time', 1, None, 2, None)) qaqc.calc_coverage() # Compute data coverage - qaqc.calc_kg({'plot_all_kg': True, 'only_log': False, - 'num_store': 5}) # Compute kalman gain + qaqc.calc_kg({'plot_all_kg': True, 'only_log': False, 'num_store': 5}) # Compute kalman gain success_iter = True # always store prior forcast, unless specifically told not to if 'nosave' not in self.ensemble.keys_da: - np.savez('prior_forecast.npz', ** - {'pred_data': self.ensemble.pred_data}) + np.savez('prior_forecast.npz', pred_data=self.ensemble.pred_data) # For the remaining iterations we start by applying the analysis and finish by running the forecast else: diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index e2b3482..5a44ba3 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -48,6 +48,7 @@ def __init__(self, keys_da, keys_en, sim): - assimindex: index for the data that will be used for assimilation - datatype: list with the name of the datatypes - staticvar: name of the static variables + - dynamicvar: name of the dynamic variables - datavar: data variance, e.g., provided as a .csv file keys_en : dict @@ -57,6 +58,9 @@ def __init__(self, keys_da, keys_en, sim): - state: name of state variables passed to the .mako file - prior_: the prior information the state variables, including mean, variance and variable limits + NB: If keys_en is empty dict, it is assumed that the prior info is contained in keys_da. + The merged dict keys_da|keys_en is what is sent to the parent class. + sim : callable The forward simulator (e.g. flow) """ @@ -90,7 +94,7 @@ def __init__(self, keys_da, keys_en, sim): # Prepare sparse representation if 'compress' in self.keys_da: - self._org_sparse_representation() + self.sparse_info = extract.organize_sparse_representation(self.keys_da['compress']) self._org_obs_data() self._org_data_var() @@ -100,12 +104,13 @@ def __init__(self, keys_da, keys_en, sim): np.ones((self.ne, self.ne))) / np.sqrt(self.ne - 1) # If we have dynamic state variables, we allocate keys for them in 'state'. Since we do not know the size - # of the arrays of the dynamic variables, we only allocate an NE list to be filled in later (in + # of the arrays of the dynamic variables, we only allocate an NE list to be filled in later (in # calc_forecast) if 'dynamicvar' in self.keys_da: - dyn_var = self.keys_da['dynamicvar'] if isinstance(self.keys_da['dynamicvar'], list) else \ - [self.keys_da['dynamicvar']] - for name in dyn_var: + dyn_vars = self.keys_da['dynamicvar'] + if not isinstance(dyn_vars, list): + dyn_vars = [dyn_vars] + for name in dyn_vars: self.state[name] = [None] * self.ne # Option to store the dictionaries containing observed data and data variance @@ -114,12 +119,6 @@ def __init__(self, keys_da, keys_en, sim): # Initialize localization if 'localization' in self.keys_da: - - if isinstance(self.keys_da['localization'], dict): - # Make 2D list of Dict (this should only be temporary) - loc_info = [[key, value] for key, value in self.keys_da['localization'].items()] - self.keys_da['localization'] = loc_info - self.localization = localization( self.keys_da['localization'], self.keys_da['truedataindex'], @@ -127,11 +126,10 @@ def __init__(self, keys_da, keys_en, sim): self.keys_da['staticvar'], self.ne ) + # Initialize local analysis if 'localanalysis' in self.keys_da: self.local_analysis = extract.extract_local_analysis_info(self.keys_da['localanalysis'], self.state.keys()) - #self.local_analysis = at.init_local_analysis( - # init=self.keys_da['localanalysis'], state=self.state.keys()) self.pred_data = [{k: np.zeros((1, self.ne), dtype='float32') for k in self.keys_da['datatype']} for _ in self.obs_data] @@ -502,35 +500,6 @@ def _org_data_var(self): self.datavar[i][datatype[j]] = est_noise # override the given value vintage = vintage + 1 - def _org_sparse_representation(self): - """ - Function for reading input to wavelet sparse representation of data. - """ - self.sparse_info = {} - parsed_info = self.keys_da['compress'] - dim = [int(elem) for elem in parsed_info[0][1]] - # flip to align with flow / eclipse - self.sparse_info['dim'] = [dim[2], dim[1], dim[0]] - self.sparse_info['mask'] = [] - for vint in range(1, len(parsed_info[1])): - if not os.path.exists(parsed_info[1][vint]): - mask = np.ones(self.sparse_info['dim'], dtype=bool) - np.savez(f'mask_{vint-1}.npz', mask=mask) - else: - mask = np.load(parsed_info[1][vint])['mask'] - self.sparse_info['mask'].append(mask.flatten()) - self.sparse_info['level'] = parsed_info[2][1] - self.sparse_info['wname'] = parsed_info[3][1] - self.sparse_info['colored_noise'] = True if parsed_info[4][1] == 'yes' else False - self.sparse_info['threshold_rule'] = parsed_info[5][1] - self.sparse_info['th_mult'] = parsed_info[6][1] - self.sparse_info['use_hard_th'] = True if parsed_info[7][1] == 'yes' else False - self.sparse_info['keep_ca'] = True if parsed_info[8][1] == 'yes' else False - self.sparse_info['inactive_value'] = parsed_info[9][1] - self.sparse_info['use_ensemble'] = True if parsed_info[10][1] == 'yes' else None - self.sparse_info['order'] = parsed_info[11][1] - self.sparse_info['min_noise'] = parsed_info[12][1] - def _ext_obs(self): self.obs_data_vector, _ = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, self.list_datatypes) diff --git a/pipt/misc_tools/cov_regularization.py b/pipt/misc_tools/cov_regularization.py index f4d981a..e2b786a 100644 --- a/pipt/misc_tools/cov_regularization.py +++ b/pipt/misc_tools/cov_regularization.py @@ -41,6 +41,7 @@ # internal import import pipt.misc_tools.analysis_tools as at +from pipt.misc_tools.extract_tools import list_to_dict class localization(): @@ -48,10 +49,114 @@ class localization(): # TODO: Check field dimensions, should always ensure that we can provide i ,j ,k (x, y, z) ### - def __init__(self, parsed_info: list, assimIndex, data_typ, free_parameter, ne): + def __init__(self, parsed_info: Union[dict,list], assimIndex: list, data_typ: list, free_parameter: list, ne: int): """ Format the parsed info from the input file, and generate the unique localization masks """ + # Make parsed_info to a dict + if isinstance(parsed_info, list): + parsed_info = list_to_dict(parsed_info) + assert isinstance(parsed_info, dict) + + # Initialize + init_local = {} + + # Assert field keyword in parsed_info + assert 'field' in parsed_info + init_local['field'] = [int(elem) for elem in parsed_info['field']] + + # Check for ACTNUM + init_local['actnum'] = None + if 'actnum' in parsed_info: + file = parsed_info['actnum'] + assert file.endswith('.npz') # this must be a .npz file!! + init_local['actnum'] = np.load(file) + + # Check for threshold + if 'threshold' in parsed_info: + init_local['threshold'] = parsed_info['threshold'] + + # Check localization method/type + try: + if 'autoadaloc' in parsed_info: + init_local = {'autoadaloc': True, 'nstd': parsed_info['autoadaloc']} + if 'type' in parsed_info: + init_local['type'] = parsed_info['type'] + elif 'localanalysis' in parsed_info: + init_local = {'localanalysis': True} + if 'type' in parsed_info: + init_local['type'] = parsed_info['type'] + if 'range' in parsed_info: + init_local['range'] = float(parsed_info['range']) + else: + # Load from pickle file + picklefile = None + for key, val in parsed_info.items(): + if (str(val).endswith('.p')) or (str(val).endswith('.pkl')): + picklefile = key + break + init_local = pickle.load(open(parsed_info[picklefile], 'rb')) + + except: + # no file could be loaded, initiallize the outer dictionary + init_local = {} + for time in assimIndex: + for datum in data_typ: + for parameter in free_parameter: + init_local[(datum, time, parameter)] = { + 'taper_func': None, + 'position': None, + 'anisotropi': None, + 'range': None + } + # If you expect a key with a CSV filename, find it: + csv_key = next((k for k in parsed_info if str(k).endswith('.csv')), None) + if csv_key: + with open(csv_key) as csv_file: + reader = csv.reader(csv_file) + info = [elem for elem in reader] + info = [item for sublist in info for item in sublist] + # Else find the key-string that contains the info + else: + for key, val in parsed_info.items(): + if len(key.split(',')) > 1: + info = key.split(',') + break + else: + info = [] + + for elem in info: + # If a predefined mask is to be imported the localization keyword must be + # [import filename.npz] + # where filename is the name of the .npz file to be uploaded. + tmp_info = elem.split() + + # format the data and time elements + if len(tmp_info) == 11: # data has only one name + name = (tmp_info[8].lower(), float(tmp_info[9]), tmp_info[10].lower()) + else: + name = (tmp_info[8].lower() + ' ' + tmp_info[9].lower(), + float(tmp_info[10]), tmp_info[11].lower()) + + # assert if the data to be localized actually exists + if name in init_local.keys(): + + # input the correct info into the localization dictionary + init_local[name]['taper_func'] = tmp_info[0] + if tmp_info[0] == 'import': + # if a predefined mask is to be imported, the name is the following element. + init_local[name]['file'] = tmp_info[1] + else: + # the position can span over multiple cells, e.g., 55:100. Hence keep this input as a string + init_local[name]['position'] = [ + [int(float(tmp_info[1])), int(float(tmp_info[2])), int(float(tmp_info[3]))]] + init_local[name]['range'] = [int(tmp_info[4]), int( + tmp_info[5])] # the range is always an integer + init_local[name]['anisotropi'] = [ + float(tmp_info[6]), float(tmp_info[7])] + + + ''' # if the next element is a .p file (pickle), assume that this has been correctly formated and can be automatically # imported. NB: it is important that we use the pickle format since we have a dictionary containing dictionaries # to make this as robust as possible, we always try to load the file @@ -126,21 +231,7 @@ def __init__(self, parsed_info: list, assimIndex, data_typ, free_parameter, ne): tmp_info[5])] # the range is always an integer init_local[name]['anisotropi'] = [ float(tmp_info[6]), float(tmp_info[7])] - - # fist element of the parsed info is field size - assert parsed_info[0][0].upper() == 'FIELD' - init_local['field'] = [int(elem) for elem in parsed_info[0][1]] - - # check if final parsed info is the actnum - try: - if parsed_info[2][0].upper() == 'ACTNUM': - assert parsed_info[2][1].endswith('.npz') # this must be a .npz file!! - tmp_file = np.load(parsed_info[2][1]) - init_local['actnum'] = tmp_file['actnum'] - else: - init_local['actnum'] = None - except: - init_local['actnum'] = None + ''' # generate the unique localization masks. Recall that the parameters: "taper_type", "anisotropi", and "range" # gives a unique mask. diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index bdea52e..b592c12 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -2,12 +2,15 @@ __all__ = [ 'extract_prior_info', 'extract_multilevel_info', - 'extract_local_analysis_info' + 'extract_local_analysis_info', + 'organize_sparse_representation' + 'list_to_dict' ] # Imports import numpy as np import pickle +import os from scipy.spatial import cKDTree from typing import Union @@ -35,14 +38,7 @@ def extract_prior_info(keys: dict) -> dict: # Check if is a list (old way) if isinstance(prior, list): - # list of lists - old way of inputting prior information - prior_dict = {} - for i, opt in enumerate(list(zip(*prior))[0]): - if opt == 'limits': - prior_dict[opt] = prior[i][1:] - else: - prior_dict[opt] = prior[i][1] - prior = prior_dict + prior = list_to_dict(prior) else: assert isinstance(prior, dict), f'PRIOR_{name.upper()} must be a dictionary or list of lists!' @@ -125,22 +121,15 @@ def extract_prior_info(keys: dict) -> dict: return prior_info -def extract_multilevel_info(keys: dict) -> dict: +def extract_multilevel_info(keys: Union[dict, list]) -> dict: ''' Extract the info needed for ML simulations. Note if the ML keyword is not in keys_en we initialize such that we only have one level -- the high fidelity one ''' - try: - ml_info = dict(keys['multilevel']) - except: - # In this case it is a list which is converted into a dict - ml_info = {} - for line in keys['multilevel']: - if len(line) > 2: - ml_info[line[0]] = line[1:] - else: - ml_info[line[0]] = line[1] - + if isinstance(keys, list): + ml_info = list_to_dict(keys) + assert isinstance(ml_info, dict) + # Set levels levels = int(ml_info['levels']) ml_info['levels'] = [elem for elem in range(levels)] @@ -167,7 +156,8 @@ def extract_multilevel_info(keys: dict) -> dict: def extract_local_analysis_info(keys: Union[dict, list], state: list) -> dict: # Check if keys are list, and make it a dict if not if isinstance(keys, list): - keys = dict(keys) + keys = list_to_dict(keys) + assert isinstance(keys, dict) # Initialize local dict local = { @@ -231,5 +221,90 @@ def extract_local_analysis_info(keys: Union[dict, list], state: list) -> dict: [data_ind[count] for count, val in enumerate(in_region) if val]) return local + - +def organize_sparse_representation(info: Union[dict,list]) -> dict: + """ + Function for reading input to wavelet sparse representation of data. + + This function takes a dictionary (or a list convertible to a dictionary) describing + the configuration for wavelet sparse representation, standardizes boolean options + (interpreting 'yes'/'no' as True/False), loads or creates mask files, and collects + all relevant parameters into a new dictionary suitable for downstream processing. + + Parameters + ---------- + info : dict or list + Input configuration for sparse representation. If a list, it will be converted + to a dictionary. Expected keys include: + - 'dim': list of 3 ints, the dimensions of the data grid. + - 'mask': list of filenames for mask arrays. + - 'level', 'wname', 'threshold_rule', 'th_mult', 'order', 'min_noise', + 'colored_noise', 'use_hard_th', 'keep_ca', 'inactive_value', 'use_ensemble'. + + Returns + ------- + sparse : dict + Dictionary containing the processed sparse representation configuration, + with masks loaded or created, dimensions flipped for compatibility, and + all options standardized. + """ + # Ensure a dict + if isinstance(info, list): + info = list_to_dict(info) + assert isinstance(info, dict) + + # Redefine all 'yes' and 'no' values to bool + for key, val in info.items(): + if val == 'yes': info[key] == True + if val == 'no': info[key] == False + + # Intial dict + sparse = {} + + # Flip dim to align with flow/eclipse + dim = [int(x) for x in info['dim']] + sparse['dim'] = [dim[2], dim[1], dim[0]] + + # Read mask_files + sparse['mask'] = [] + for idx, filename in enumerate(info['mask'], start=1): + if not os.path.exists(filename): + mask = np.ones(sparse['dim'], dtype=bool) + np.savez(f'mask_{idx}.npz', mask=mask) + else: + mask = np.load(filename)['mask'] + sparse['mask'].append(mask.flatten()) + + # Read rest of keywords + sparse['level'] = info['level'] + sparse['wname'] = info['wname'] + sparse['threshold_rule'] = info['threshold_rule'] + sparse['th_mult'] = info['th_mult'] + sparse['order'] = info['order'] + sparse['min_noise'] = info['min_noise'] + sparse['colored_noise'] = info.get('colored_noise', False) + sparse['use_hard_th'] = info.get('use_hard_th', False) + sparse['keep_ca'] = info.get('keep_ca', False) + sparse['inactive_value'] = info['inactive_value'] + sparse['use_ensemble'] = info.get('use_ensemble', None) + + return sparse + + +def list_to_dict(info_list: list) -> dict: + assert isinstance(info_list) + # Initialize and loop over entries + info_dict = {} + for entry in info_list: + if not isinstance(entry, list): + entry = [entry] + # Fill in values + if len(entry) == 1: + info_dict[str(entry[0])] = None + elif len(entry) == 2: + info_dict[str(entry[0])] = entry[1] + else: + info_dict[str(entry[0])] = entry[1:] + + return info_dict \ No newline at end of file diff --git a/pipt/update_schemes/update_methods_ns/approx_update.py b/pipt/update_schemes/update_methods_ns/approx_update.py index 475e87b..1ec2194 100644 --- a/pipt/update_schemes/update_methods_ns/approx_update.py +++ b/pipt/update_schemes/update_methods_ns/approx_update.py @@ -50,9 +50,9 @@ def update(self): data_size = [[self.obs_data[int(time)][data].size if self.obs_data[int(time)][data] is not None else 0 for data in self.list_datatypes] for time in self.assim_index[1]] - f = self.keys_da['localization'] + #f = self.keys_da['localization'] - if f[1][0] == 'autoadaloc': + if 'autoadaloc' in self.localization.loc_info: # Mean state and perturbation matrix mean_state = np.mean(aug_state, 1) @@ -101,7 +101,7 @@ def update(self): except: self.step = (weight*(np.dot(pert_state, X))).dot(scaled_delta_data) - elif sum(['dist_loc' in el for el in f]) >= 1: + elif ('dist_loc' in self.keys_da['localization'].keys()) or ('dist_loc' in self.keys_da['localization'].values()): local_mask = self.localization.localize(self.list_datatypes, [self.keys_da['truedataindex'][int(elem)] for elem in self.assim_index[1]], self.list_states, self.ne, self.prior_info, data_size) From b10f6217dbedffb8cca8695da523874bef3ff76b Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 12 Sep 2025 13:10:13 +0200 Subject: [PATCH 22/94] added extract_maxiter to extract_tools --- pipt/loop/assimilation.py | 45 ++------------------------------ pipt/misc_tools/extract_tools.py | 27 ++++++++++++++++++- pipt/update_schemes/enrml.py | 15 +++++------ popt/loop/optimize.py | 30 ++++++--------------- 4 files changed, 43 insertions(+), 74 deletions(-) diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index d9b7e79..6803d27 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -20,6 +20,7 @@ from pipt.loop.ensemble import Ensemble from misc.system_tools.environ_var import OpenBlasSingleThread from pipt.misc_tools import analysis_tools as at +import pipt.misc_tools.extract_tools as extract class Assimilate: @@ -50,7 +51,7 @@ def __init__(self, ensemble: Ensemble): if hasattr(ensemble, 'max_iter'): self.max_iter = self.ensemble.max_iter else: - self.max_iter = self._ext_max_iter() + self.max_iter = extract.extract_maxiter(self.ensemble.keys_da) # Within variables self.why_stop = None # Output of why iter. loop stopped @@ -281,48 +282,6 @@ def remove_outliers(self): self.ensemble.pred_data[i][el][:, index] = deepcopy( self.ensemble.pred_data[i][el][:, new_index]) - def _ext_max_iter(self): - """ - Extract max iterations from ITERATION keyword in DATAASSIM part (mandatory keyword for iteration loops). - - Parameters - ---------- - keys_da : dict - A dictionary containing all keywords from DATAASSIM part. - - - 'iteration' : object - Information for iterative methods. - - Returns - ------- - max_iter : int - The maximum number of iterations allowed before abort. - - Changelog - --------- - - ST 7/6-16 - """ - if 'iteration' in self.ensemble.keys_da: - iter_opts = dict(self.ensemble.keys_da['iteration']) - # Check if 'max_iter' has been given; if not, give error (mandatory in ITERATION) - try: - max_iter = iter_opts['max_iter'] - except KeyError: - raise AssertionError('MAX_ITER has not been given in ITERATION') - - elif 'mda' in self.ensemble.keys_da: - iter_opts = dict(self.ensemble.keys_da['mda']) - # Check if 'tot_assim_steps' has been given; if not, raise error (mandatory in MDA) - try: - max_iter = iter_opts['tot_assim_steps'] - except KeyError: - raise AssertionError('TOT_ASSIM_STEPS has not been given in MDA!') - - else: - max_iter = 1 - # Return max. iter - return max_iter - def _save_iteration_information(self): """ More general method for saving all relevant information from a analysis/forecast step. Note that this is diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index b592c12..4504b18 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -3,7 +3,8 @@ 'extract_prior_info', 'extract_multilevel_info', 'extract_local_analysis_info', - 'organize_sparse_representation' + 'extract_maxiter', + 'organize_sparse_representation', 'list_to_dict' ] @@ -290,6 +291,30 @@ def organize_sparse_representation(info: Union[dict,list]) -> dict: sparse['use_ensemble'] = info.get('use_ensemble', None) return sparse + + +def extract_maxiter(keys: dict) -> dict: + + if 'iteration' in keys: + if isinstance(keys['iteration'], list): + keys['iteration'] = list_to_dict(keys['iteration']) + try: + max_iter = keys['iteration']['max_iter'] + except KeyError: + raise AssertionError('MAX_ITER has not been given in ITERATION') + + elif 'mda' in keys: + if isinstance(keys['mda'], list): + keys['mda'] = list_to_dict(keys['mda']) + try: + max_iter = keys['mda']['max_iter'] + except KeyError: + raise AssertionError('MAX_ITER has not been given in MDA') + + else: + max_iter = 1 + + return max_iter def list_to_dict(info_list: list) -> dict: diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index 7b52267..8eb3dea 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -3,6 +3,7 @@ """ # External imports import pipt.misc_tools.analysis_tools as at +import pipt.misc_tools.extract_tools as extract from geostat.decomp import Cholesky from pipt.loop.ensemble import Ensemble from pipt.update_schemes.update_methods_ns.subspace_update import subspace_update @@ -251,10 +252,9 @@ def _ext_iter_param(self): file. These parameters include convergence tolerances and parameters for the damping parameter. Default values for these parameters have been given here, if they are not provided in ITERATION. """ - try: - options = dict(self.keys_da['iteration']) - except: - options = dict([self.keys_da['iteration']]) + options = self.keys_da['iteration'] + if isinstance(options, list): + options = extract.list_to_dict(options) # unpack options self.data_misfit_tol = options.get('data_misfit_tol', 0.01) @@ -580,10 +580,9 @@ def _ext_iter_param(self): file. These parameters include convergence tolerances and parameters for the damping parameter. Default values for these parameters have been given here, if they are not provided in ITERATION. """ - try: - options = dict(self.keys_da['iteration']) - except: - options = dict([self.keys_da['iteration']]) + options = self.keys_da['iteration'] + if isinstance(options, list): + options = extract.list_to_dict(options) self.data_misfit_tol = options.get('data_misfit_tol', 0.01) self.trunc_energy = options.get('energy', 0.95) diff --git a/popt/loop/optimize.py b/popt/loop/optimize.py index 73c6cf3..35ccaa4 100644 --- a/popt/loop/optimize.py +++ b/popt/loop/optimize.py @@ -72,13 +72,6 @@ def __init__(self, **options): options : dict Optimization options """ - - def __set__variable(var_name=None, defalut=None): - if var_name in options: - return options[var_name] - else: - return defalut - # Set the logger self.logger = logger @@ -96,16 +89,16 @@ def __set__variable(var_name=None, defalut=None): self.rnd = None # Max number of iterations - self.max_iter = __set__variable('maxiter', 20) + self.max_iter = options.get('maxiter', 20) # Restart flag - self.restart = __set__variable('restart', False) + self.restart = options.get('restart', False) # Save restart information flag - self.restartsave = __set__variable('restartsave', False) + self.restartsave = options.get('restartsave', False) # Optimze with external penalty function for constraints, provide r_0 as input - self.epf = __set__variable('epf', {}) + self.epf = options.get('epf', {}) self.epf_iteration = 0 # Initialize variables (set in subclasses) @@ -128,23 +121,16 @@ def run_loop(self): # If it is a restart run, we load the self info that exists in the pickle save file. if self.restart: - - # Check if the pickle save file exists in folder - assert (self.pickle_restart_file in [f for f in os.listdir('.') if os.path.isfile(f)]), \ - 'The restart file "{0}" does not exist in folder. Cannot restart!'.format(self.pickle_restart_file) - - # Load restart file - self.load() - + try: + self.load() + except (FileNotFoundError, pickle.UnpicklingError) as e: + raise RuntimeError(f"Failed to load restart file '{self.pickle_restart_file}': {e}") # Set the random generator to be the saved value np.random.set_state(self.rnd) - else: - # delete potential restart files to avoid any problems if self.pickle_restart_file in [f for f in os.listdir('.') if os.path.isfile(f)]: os.remove(self.pickle_restart_file) - self.iteration += 1 # Check if external penalty function (epf) for handling constraints should be used From ff64901a0689fbfc6f114321b46643b3ca3b6327 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 15 Sep 2025 10:34:55 +0200 Subject: [PATCH 23/94] assertion fix --- pipt/misc_tools/extract_tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index 4504b18..fe2d546 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -318,7 +318,7 @@ def extract_maxiter(keys: dict) -> dict: def list_to_dict(info_list: list) -> dict: - assert isinstance(info_list) + assert isinstance(info_list, list) # Initialize and loop over entries info_dict = {} for entry in info_list: From b2e3b273de05282054855c4b72021099e516e732 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 15 Sep 2025 13:21:24 +0200 Subject: [PATCH 24/94] fixed input for sim options --- simulator/eclipse.py | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/simulator/eclipse.py b/simulator/eclipse.py index 38ef1ae..c2256a1 100644 --- a/simulator/eclipse.py +++ b/simulator/eclipse.py @@ -19,6 +19,7 @@ # Internal imports from misc.system_tools.environ_var import EclipseRunEnvironment from pipt.misc_tools.analysis_tools import store_ensemble_sim_information +from pipt.misc_tools.extract_tools import list_to_dict class eclipse: @@ -111,24 +112,18 @@ def _extInfoInputDict(self): # In the ecl framework, all reference to the filename should be uppercase self.file = self.input_dict['runfile'].upper() + + # Extract sim options + if isinstance(self.input_dict['simoptions'], list): + self.input_dict['simoptions'] = list_to_dict(self.input_dict['simoptions']) + + simoptions = self.input_dict['simoptions'] self.options = {} - self.options['sim_path'] = '' - self.options['sim_flag'] = '' - self.options['mpi'] = '' - self.options['parsing-strictness'] = '' - # Loop over options in SIMOPTIONS and extract the parameters we want - if 'simoptions' in self.input_dict: - if type(self.input_dict['simoptions'][0]) == str: - self.input_dict['simoptions'] = [self.input_dict['simoptions']] - for i, opt in enumerate(list(zip(*self.input_dict['simoptions']))[0]): - if opt == 'sim_path': - self.options['sim_path'] = self.input_dict['simoptions'][i][1] - if opt == 'sim_flag': - self.options['sim_flag'] = self.input_dict['simoptions'][i][1] - if opt == 'mpi': - self.options['mpi'] = self.input_dict['simoptions'][i][1] - if opt == 'parsing-strictness': - self.options['parsing-strictness'] = self.input_dict['simoptions'][i][1] + self.options['sim_path'] = simoptions.get('sim_path', '') + self.options['sim_flag'] = simoptions.get('sim_flag', '') + self.options['mpi'] = simoptions.get('mpi', '') + self.options['parsing-strictness'] = simoptions.get('parsing-strictness', '') + if 'sim_limit' in self.input_dict: self.options['sim_limit'] = self.input_dict['sim_limit'] From 2b22df6d3c6e2b1efa0d6b8c93848038f911231d Mon Sep 17 00:00:00 2001 From: "Rolf J. Lorentzen" Date: Fri, 19 Sep 2025 11:11:13 +0200 Subject: [PATCH 25/94] Update extract_tools.py --- pipt/misc_tools/extract_tools.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index fe2d546..46cbfe0 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -49,7 +49,7 @@ def extract_prior_info(keys: dict) -> dict: assert prior['mean'].endswith('.npz'), 'File name does not end with \'.npz\'!' mean_file = np.load(prior['mean']) assert len(mean_file.files) == 1, \ - f'More than one variable located in {prior['mean']}. Only the mean vector can be stored in the .npz file!' + f"More than one variable located in {prior['mean']}. Only the mean vector can be stored in the .npz file!" prior['mean'] = mean_file[mean_file.files[0]] else: # Single number inputted, make it a list if not already if not isinstance(prior['mean'], list): @@ -332,4 +332,4 @@ def list_to_dict(info_list: list) -> dict: else: info_dict[str(entry[0])] = entry[1:] - return info_dict \ No newline at end of file + return info_dict From da70d84045de1a1f92aafef6bb4d5c67a634b6c1 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 19 Sep 2025 15:18:52 +0200 Subject: [PATCH 26/94] added a general file reader --- input_output/read_config.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/input_output/read_config.py b/input_output/read_config.py index 0d9d10d..ccc6e22 100644 --- a/input_output/read_config.py +++ b/input_output/read_config.py @@ -9,6 +9,18 @@ import numpy as np +def read(filename: str): + ''' Read configuration file. Supported formats are toml, .yaml, .pipt and .popt.''' + if filename.endswith('.pipt') or filename.endswith('.popt'): + return read_txt(filename) + elif filename.endswith('.yaml'): + return read_yaml(filename) + elif filename.endswith('.toml'): + return read_toml(filename) + else: + raise ValueError('File format not supported. Supported formats are toml, .yaml, .pipt, .popt') + + def convert_txt_to_yaml(init_file): # Read .pipt or .popt file pr, fwd = read_txt(init_file) @@ -46,30 +58,31 @@ def ndarray_constructor(loader, node): # Add constructor to yaml with tag !ndarray yaml.add_constructor('!ndarray', ndarray_constructor) - # Read + # Read yaml file with open(init_file, 'rb') as fid: y = yaml.load(fid, Loader=FullLoader) - # Check for dataassim and fwdsim + # Check for ensemble if 'ensemble' in y.keys(): keys_en = y['ensemble'] check_mand_keywords_en(keys_en) else: keys_en = {} - if 'optim' in y.keys(): - keys_pr = y['optim'] - check_mand_keywords_opt(keys_pr) - elif 'dataassim' in y.keys(): + # Check for dataassim + if 'dataassim' in y.keys(): keys_pr = y['dataassim'] check_mand_keywords_da(keys_pr) + elif 'optim' in y.keys(): + keys_pr = y['optim'] + check_mand_keywords_opt(keys_pr) else: - raise KeyError + keys_pr = {} if 'fwdsim' in y.keys(): keys_fwd = y['fwdsim'] else: - raise KeyError + keys_fwd = {} # Organize keywords org = Organize_input(keys_pr, keys_fwd, keys_en) From b1b566f42b056acb3a29f6699e3c7bfec6712cea Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 19 Sep 2025 15:41:19 +0200 Subject: [PATCH 27/94] fixed an import error --- ensemble/ensemble.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 1208eff..a9bc9cb 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -385,7 +385,7 @@ def calc_prediction(self, input_state=None, save_prediction=None): self.sim.extract_data(member_i) en_pred.append(deepcopy(self.sim.pred_data)) if self.sim.saveinfo is not None: # Try to save information - store_ensemble_sim_information(self.sim.saveinfo, member_i) + at.store_ensemble_sim_information(self.sim.saveinfo, member_i) else: en_pred.append(False) self.sim.remove_folder(member_i) From 0d6b6877c7c340aa79ddba118d4e05a9301374c4 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 22 Sep 2025 10:26:43 +0200 Subject: [PATCH 28/94] fixed a bug --- popt/update_schemes/linesearch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index e82b305..bf5040d 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -519,7 +519,7 @@ def _set_step_size(self, pk, amax): else: if (self.step_size_adapt == 1) and (np.dot(pk, self.jk) != 0): alpha = 2*(self.fk - self.f_old)/np.dot(pk, self.jk) - elif (self.step_size_adapt == 2) and (np.dot(pk, self.jk) == 0): + elif (self.step_size_adapt == 2) and (np.dot(pk, self.jk) != 0): slope_old = np.dot(self.p_old, self.j_old) slope_new = np.dot(pk, self.jk) alpha = self.step_size*slope_old/slope_new From c0e008e138dd859f1c5bda0baf32c1bc74d318a9 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 23 Sep 2025 13:47:43 +0200 Subject: [PATCH 29/94] Started to introduce ensemble matrix (and code simplification) --- ensemble/ensemble.py | 249 ++++++++++-------- pipt/loop/assimilation.py | 13 +- pipt/loop/ensemble.py | 6 +- pipt/misc_tools/analysis_tools.py | 61 +++++ pipt/misc_tools/cov_regularization.py | 3 +- pipt/update_schemes/enrml.py | 56 ++-- .../update_methods_ns/approx_update.py | 219 ++++++++------- 7 files changed, 377 insertions(+), 230 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index a9bc9cb..a18bf62 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -51,7 +51,12 @@ def __init__(self, keys_en, sim, redund_sim=None): self.keys_en = keys_en self.sim = sim self.sim.redund_sim = redund_sim + + # Initialize some attributes self.pred_data = None + self.enX_temp = None + self.enX = None + self.idX = {} # Auxilliary input to the simulator - can be used e.g., # to allow for different models when optimizing. @@ -137,17 +142,19 @@ def __init__(self, keys_en, sim, redund_sim=None): # We assume that the user has saved the state dict. as **state (effectively saved all keys in state # individually). - self.state = {key: val for key, val in tmp_load.items()} - - # Find the number of ensemble members from state variable - tmp_ne = [] - for tmp_state in self.state.keys(): - tmp_ne.extend([self.state[tmp_state].shape[1]]) - if max(tmp_ne) != min(tmp_ne): - print('\033[1;33mInput states have different ensemble size\033[1;m') - sys.exit(1) - self.ne = min(tmp_ne) + for key in self.keys_en['staticvar']: + if self.enX is None: + self.enX = tmp_load[key] + self.ne = self.enX.shape[1] + else: + assert self.ne == tmp_load[key].shape[1], 'Ensemble size of imported state variables do not match!' + self.enX = np.vstack((self.enX, tmp_load[key])) + + # fill in indices + self.idX[key] = (self.enX.shape[0] - tmp_load[key].shape[0], self.enX.shape[0]) + self.list_states = list(self.keys_en['staticvar']) + if 'multilevel' in self.keys_en: ml_info = extract.extract_multilevel_info(self.keys_en) self.multilevel, self.tot_level, self.ml_ne, self.ML_error_corr, self.error_comp_scheme, self.ML_corr_done = ml_info @@ -277,7 +284,7 @@ def get_list_assim_steps(self): # Return tot. assim. steps return list_assim - def calc_prediction(self, input_state=None, save_prediction=None): + def calc_prediction(self, enX=None, save_prediction=None): """ Method for making predictions using the state variable. Will output the simulator response for all report steps and all data values provided to the simulator. @@ -297,106 +304,73 @@ def calc_prediction(self, input_state=None, save_prediction=None): """ - if isinstance(self.state,list) and hasattr(self, 'multilevel'): # assume multilevel is used if state is a list - success = self.calc_ml_prediction(input_state) + if isinstance(self.enX,list) and hasattr(self, 'multilevel'): # assume multilevel is used if state is a list + success = self.calc_ml_prediction(enX) else: - # Number of parallel runs - if 'parallel' in self.sim.input_dict: - no_tot_run = int(self.sim.input_dict['parallel']) + + # Use input state if given + if enX is None: + use_input_ensemble = False + enX = self.enX + self.enX = None # free memory else: - no_tot_run = 1 + use_input_ensemble = True + + # Number of parallel runs + nparallel = int(self.sim.input_dict.get('parallel', 1)) self.pred_data = [] - # for level in self.multilevel['level']: # - # Setup forward simulator and redundant simulator at the correct fidelity + # Run setup function for redund simulator if self.sim.redund_sim is not None: - self.sim.redund_sim.setup_fwd_run() - self.sim.setup_fwd_run(redund_sim=self.sim.redund_sim) + if hasattr(self.sim.redund_sim, 'setup_fwd_run'): + self.sim.redund_sim.setup_fwd_run() + + # Run setup function for simulator + if hasattr(self.sim, 'setup_fwd_run'): + self.sim.setup_fwd_run(redund_sim=self.sim.redund_sim) + + # Convert ensemble matrix to list of dictionaries + enX = at.ensmeble_matrix_to_list(enX, self.idX) - # Ensure that we put all the states in a list - list_state = [deepcopy({}) for _ in range(self.ne)] - for i in range(self.ne): - if input_state is None: - for key in self.state.keys(): - if self.state[key].ndim == 1: - list_state[i][key] = deepcopy(self.state[key]) - elif self.state[key].ndim == 2: - list_state[i][key] = deepcopy(self.state[key][:, i]) - # elif self.state[key].ndim == 3: - # list_state[i][key] = deepcopy(self.state[key][level,:, i]) - else: - for key in self.state.keys(): - if input_state[key].ndim == 1: - list_state[i][key] = deepcopy(input_state[key]) - elif input_state[key].ndim == 2: - list_state[i][key] = deepcopy(input_state[key][:, i]) - # elif input_state[key].ndim == 3: - # list_state[i][key] = deepcopy(input_state[key][:,:, i]) - if self.aux_input is not None: # several models are used - list_state[i]['aux_input'] = self.aux_input[i] - - # Index list of ensemble members - list_member_index = list(range(self.ne)) - - if no_tot_run==1: # if not in parallel we use regular loop - en_pred = [self.sim.run_fwd_sim(state, member_index) for state, member_index in - tqdm(zip(list_state, list_member_index), total=len(list_state))] - elif self.sim.input_dict.get('hpc', False): # Run prediction in parallel on hpc - batch_size = no_tot_run # If more than 500 ensemble members, we limit the runs to batches of 500 - # Split the ensemble into batches of 500 - if batch_size >= 1000: - self.logger.info(f'Cannot run batch size of {no_tot_run}. Set to 1000') - batch_size = 1000 + if not (self.aux_input is None): + for n in range(self.ne): + enX[n]['aux_input'] = self.aux_input[n] + + ###################################################################################################################### + # No parralelization + if nparallel==1: en_pred = [] - batch_en = [np.arange(start, start + batch_size) for start in - np.arange(0, self.ne - batch_size, batch_size)] - if len(batch_en): # if self.ne is less than batch_size - batch_en.append(np.arange(batch_en[-1][-1]+1, self.ne)) - else: - batch_en.append(np.arange(0, self.ne)) - for n_e in batch_en: - _ = [self.sim.run_fwd_sim(state, member_index, nosim=True) for state, member_index in - zip([list_state[curr_n] for curr_n in n_e], [list_member_index[curr_n] for curr_n in n_e])] - # Run call_sim on the hpc - if self.sim.options['mpiarray']: - job_id = self.sim.SLURM_ARRAY_HPC_run( - n_e, - venv=os.path.join(os.path.dirname(sys.executable), 'activate'), - filename=self.sim.file, - **self.sim.options - ) - else: - job_id=self.sim.SLURM_HPC_run( - n_e, - venv=os.path.join(os.path.dirname(sys.executable),'activate'), - filename=self.sim.file, - **self.sim.options - ) - - # Wait for the simulations to finish - if job_id: - sim_status = self.sim.wait_for_jobs(job_id) - else: - print("Job submission failed. Exiting.") - sim_status = [False]*len(n_e) - # Extract the results. Need a local counter to check the results in the correct order - for c_member, member_i in enumerate([list_member_index[curr_n] for curr_n in n_e]): - if sim_status[c_member]: - self.sim.extract_data(member_i) - en_pred.append(deepcopy(self.sim.pred_data)) - if self.sim.saveinfo is not None: # Try to save information - at.store_ensemble_sim_information(self.sim.saveinfo, member_i) - else: - en_pred.append(False) - self.sim.remove_folder(member_i) - else: # Run prediction in parallel using p_map - en_pred = p_map(self.sim.run_fwd_sim, list_state, - list_member_index, num_cpus=no_tot_run, disable=self.disable_tqdm) + for member_index, state in tqdm(enumerate(enX), total=self.ne, desc="Running simulations"): + en_pred.append(self.sim.run_fwd_sim(state, member_index)) + + # Parallelization on HPC using SLURM + elif self.sim.input_dict.get('hpc', False): # Run prediction in parallel on hpc + en_pred = self.run_on_HPC(enX, batch_size=nparallel) + + # Parallelization on local machine using p_map + else: + en_pred = p_map( + self.sim.run_fwd_sim, + enX, + list(range(self.ne)), + num_cpus=nparallel, + disable=self.disable_tqdm + ) + ###################################################################################################################### + + # Convert state enemble back to matrix form + enX = at.ensemble_list_to_matrix(enX, self.idX) + + # restore state ensemble if it was not inputted + if not use_input_ensemble: + self.enX = enX + enX = None # free memory + # List successful runs and crashes - list_crash = [indx for indx, el in enumerate(en_pred) if el is False] - list_success = [indx for indx, el in enumerate(en_pred) if el is not False] success = True - + list_success = [indx for indx, el in enumerate(en_pred) if el is not False] + list_crash = [indx for indx, el in enumerate(en_pred) if el is False] + # Dump all information and print error if all runs have crashed if not list_success: self.save() @@ -414,22 +388,21 @@ def calc_prediction(self, input_state=None, save_prediction=None): # Replace crashed runs with (random) successful runs. If there are more crashed runs than successful once, # we draw with replacement. if len(list_crash) < len(list_success): - copy_member = np.random.choice( - list_success, size=len(list_crash), replace=False) + copy_member = np.random.choice(list_success, size=len(list_crash), replace=False) else: - copy_member = np.random.choice( - list_success, size=len(list_crash), replace=True) + copy_member = np.random.choice(list_success, size=len(list_crash), replace=True) # Insert the replaced runs in prediction list for indx, el in enumerate(copy_member): - print(f'\033[92m--- Ensemble member {list_crash[indx]} failed, has been replaced by ensemble member ' - f'{el}! ---\033[92m') - self.logger.info(f'\033[92m--- Ensemble member {list_crash[indx]} failed, has been replaced by ' - f'ensemble member {el}! ---\033[92m') + msg = ( + f"\033[92m--- Ensemble member {list_crash[indx]} failed, " + f"has been replaced by ensemble member {el}! ---\033[92m" + ) + print(msg) + self.logger.info(msg) for key in self.state.keys(): if self.state[key].ndim > 1: - self.state[key][:, list_crash[indx]] = deepcopy( - self.state[key][:, el]) + self.state[key][:, list_crash[indx]] = deepcopy(self.state[key][:, el]) en_pred[list_crash[indx]] = deepcopy(en_pred[el]) # Convert ensemble specific result into pred_data, and filter for NONE data @@ -445,6 +418,58 @@ def calc_prediction(self, input_state=None, save_prediction=None): np.savez(f'{save_prediction}.npz', **{'pred_data': self.pred_data}) return success + + def run_on_HPC(self, enX, batch_size=None, **kwargs): + list_member_index = list(range(self.ne)) + + # Split the ensemble into batches of 500 + if batch_size >= 1000: + self.logger.info(f'Cannot run batch size of {batch_size}. Set to 1000') + batch_size = 1000 + en_pred = [] + batch_en = [np.arange(start, start + batch_size) for start in + np.arange(0, self.ne - batch_size, batch_size)] + if len(batch_en): # if self.ne is less than batch_size + batch_en.append(np.arange(batch_en[-1][-1]+1, self.ne)) + else: + batch_en.append(np.arange(0, self.ne)) + for n_e in batch_en: + _ = [self.sim.run_fwd_sim(state, member_index, nosim=True) for state, member_index in + zip([enX[curr_n] for curr_n in n_e], [list_member_index[curr_n] for curr_n in n_e])] + # Run call_sim on the hpc + if self.sim.options['mpiarray']: + job_id = self.sim.SLURM_ARRAY_HPC_run( + n_e, + venv=os.path.join(os.path.dirname(sys.executable), 'activate'), + filename=self.sim.file, + **self.sim.options + ) + else: + job_id=self.sim.SLURM_HPC_run( + n_e, + venv=os.path.join(os.path.dirname(sys.executable),'activate'), + filename=self.sim.file, + **self.sim.options + ) + + # Wait for the simulations to finish + if job_id: + sim_status = self.sim.wait_for_jobs(job_id) + else: + print("Job submission failed. Exiting.") + sim_status = [False]*len(n_e) + # Extract the results. Need a local counter to check the results in the correct order + for c_member, member_i in enumerate([list_member_index[curr_n] for curr_n in n_e]): + if sim_status[c_member]: + self.sim.extract_data(member_i) + en_pred.append(deepcopy(self.sim.pred_data)) + if self.sim.saveinfo is not None: # Try to save information + at.store_ensemble_sim_information(self.sim.saveinfo, member_i) + else: + en_pred.append(False) + self.sim.remove_folder(member_i) + + return en_pred def save(self): """ diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index 6803d27..e1c634e 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -200,11 +200,11 @@ def run(self): # always store posterior forcast and state, unless specifically told not to if 'nosave' not in self.ensemble.keys_da: try: # first try to save as npz file - np.savez('posterior_state_estimate.npz', **self.ensemble.state) + np.savez('posterior_state_estimate.npz', **self.ensemble.enX) np.savez('posterior_forecast.npz', **{'pred_data': self.ensemble.pred_data}) except: # If this fails, store as pickle with open('posterior_state_estimate.p', 'wb') as file: - pickle.dump(self.ensemble.state, file) + pickle.dump(self.ensemble.enX, file) with open('posterior_forecast.p', 'wb') as file: pickle.dump(self.ensemble.pred_data, file) @@ -351,6 +351,10 @@ def _save_analysis_debug(self): else: analysisdebug = [self.ensemble.keys_da['analysisdebug']] + if 'state' in analysisdebug: + analysisdebug.remove('state') + analysisdebug.append('enX') + # Loop over variables to store in save list for save_typ in analysisdebug: if hasattr(self, save_typ): @@ -440,7 +444,10 @@ def calc_forecast(self): l_prim = [int(assim_ind[1])] # Run forecast. Predicted data solved in self.ensemble.pred_data - self.ensemble.calc_prediction() + if self.ensemble.enX_temp is None: + self.ensemble.calc_prediction() + else: + self.ensemble.calc_prediction(enX=self.ensemble.enX_temp) # Filter pred. data needed at current assimilation step. This essentially means deleting pred. data not # contained in the assim. indices for current assim. step or does not have obs. data at this index diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index 5a44ba3..2b98624 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -126,7 +126,7 @@ def __init__(self, keys_da, keys_en, sim): self.keys_da['staticvar'], self.ne ) - + # Initialize local analysis if 'localanalysis' in self.keys_da: self.local_analysis = extract.extract_local_analysis_info(self.keys_da['localanalysis'], self.state.keys()) @@ -538,8 +538,8 @@ def _ext_obs(self): def _ext_scaling(self): # get vector of scaling self.state_scaling = at.calc_scaling( - self.prior_state, self.list_states, self.prior_info) - + self.prior_enX, self.idX.keys(), self.prior_info) + self.Am = None def save_temp_state_assim(self, ind_save): diff --git a/pipt/misc_tools/analysis_tools.py b/pipt/misc_tools/analysis_tools.py index 59f748f..ebd6e59 100644 --- a/pipt/misc_tools/analysis_tools.py +++ b/pipt/misc_tools/analysis_tools.py @@ -1181,6 +1181,67 @@ def compute_x(pert_preddata, cov_data, keys_da, alfa=None): return X +def ensmeble_matrix_to_list(matrix: np.ndarray, indecies: dict) -> list[dict]: + ''' + Convert an ensemble matrix to a list of dictionaries. + + Parameters + ---------- + matrix : np.ndarray + Ensemble matrix where each column represents an ensemble member. + indecies : dict + Dictionary with keys as variable names and values as tuples indicating the start and end row indices + for each variable in the ensemble matrix. + + Returns + ------- + ensemble_list : list of dict + ''' + ne = matrix.shape[1] + ensemble_list = [] + + for n in range(ne): + member = {} + for key, (start, end) in indecies.items(): + if matrix[start:end].ndim == 2: + member[key] = matrix[start:end, n] + else: + member[key] = matrix[start:end] + ensemble_list.append(member) + + return ensemble_list + +def ensemble_list_to_matrix(ensemble_list: list[dict], indecies: dict) -> np.ndarray: + ''' + Convert a list of dictionaries to an ensemble matrix. + + Parameters + ---------- + ensemble_list : list of dict + List where each dictionary represents an ensemble member with variable names as keys. + indecies : dict + Dictionary with keys as variable names and values as tuples indicating the start and end row indices + for each variable in the ensemble matrix. + + Returns + ------- + matrix : np.ndarray + Ensemble matrix where each column represents an ensemble member. + ''' + ne = len(ensemble_list) + nx = sum(end - start for start, end in indecies.values()) + matrix = np.zeros((nx, ne)) + + for n, member in enumerate(ensemble_list): + for key, (start, end) in indecies.items(): + if member[key].ndim == 2: + matrix[start:end, n] = member[key][:,n] + else: + matrix[start:end, n] = member[key] + + return matrix + + def aug_state(state, list_state, cell_index=None): """ Augment the state variables to an array. diff --git a/pipt/misc_tools/cov_regularization.py b/pipt/misc_tools/cov_regularization.py index e2b786a..31d762a 100644 --- a/pipt/misc_tools/cov_regularization.py +++ b/pipt/misc_tools/cov_regularization.py @@ -79,7 +79,7 @@ def __init__(self, parsed_info: Union[dict,list], assimIndex: list, data_typ: li # Check localization method/type try: if 'autoadaloc' in parsed_info: - init_local = {'autoadaloc': True, 'nstd': parsed_info['autoadaloc']} + init_local.update({'autoadaloc': True, 'nstd': parsed_info['autoadaloc']}) if 'type' in parsed_info: init_local['type'] = parsed_info['type'] elif 'localanalysis' in parsed_info: @@ -312,6 +312,7 @@ def __init__(self, parsed_info: Union[dict,list], assimIndex: list, data_typ: li field_size=init_local['field'], ne=ne ) + self.loc_info = init_local def localize(self, curr_data, curr_time, curr_param, ne, prior_info, data_size): diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index 8eb3dea..c01df81 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -58,7 +58,8 @@ def __init__(self, keys_da, keys_en, sim): if self.restart is False: # Save prior state in separate variable - self.prior_state = cp.deepcopy(self.state) + #self.prior_state = cp.deepcopy(self.state) + self.prior_enX = cp.deepcopy(self.enX) # not sure if this is wise! # Extract parameters like conv. tol. and damping param. from ITERATION keyword in DATAASSIM self._ext_iter_param() @@ -77,8 +78,7 @@ def __init__(self, keys_da, keys_en, sim): self.check_assimindex_simultaneous() # define the assimilation index self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] - # define the list of states - self.list_states = list(self.state.keys()) + # define the list of datatypes self.list_datatypes, self.list_act_datatypes = at.get_list_data_types( self.obs_data, self.assim_index) @@ -87,7 +87,8 @@ def __init__(self, keys_da, keys_en, sim): self._ext_obs() # Get state scaling and svd of scaled prior self._ext_scaling() - self.current_state = cp.deepcopy(self.state) + + def calc_analysis(self): """ @@ -119,25 +120,24 @@ def calc_analysis(self): else: # Mean pred_data and perturbation matrix with scaling if len(self.scale_data.shape) == 1: - self.pert_preddata = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), - np.ones((1, self.ne))) * np.dot(self.aug_pred_data, self.proj) + self.pert_preddata = (self.scale_data ** -1)[:, None] * np.dot(self.aug_pred_data, self.proj) else: - self.pert_preddata = solve( - self.scale_data, np.dot(self.aug_pred_data, self.proj)) + self.pert_preddata = solve(self.scale_data, np.dot(self.aug_pred_data, self.proj)) - aug_state = at.aug_state(self.current_state, self.list_states) - self.update() # run ordinary analysis + # Calculate update to get the step (found in update_methods_ns) + self.update() + + # Update the state ensemble and weights if hasattr(self, 'step'): - aug_state_upd = aug_state + self.step + self.enX_temp = self.enX + self.step if hasattr(self, 'w_step'): self.W = self.current_W + self.w_step - aug_prior_state = at.aug_state(self.prior_state, self.list_states) - aug_state_upd = np.dot(aug_prior_state, (np.eye( - self.ne) + self.W / np.sqrt(self.ne - 1))) + self.enX_temp = np.dot(self.prior_enX, (np.eye(self.ne) + self.W/np.sqrt(self.ne - 1))) + # Extract updated state variables from aug_update - self.state = at.update_state(aug_state_upd, self.state, self.list_states) - self.state = at.limits(self.state, self.prior_info) + #self.state = at.update_state(aug_state_upd, self.state, self.list_states) + #self.state = at.limits(self.state, self.prior_info) def check_convergence(self): """ @@ -214,13 +214,25 @@ def check_convergence(self): if self.lam > self.lam_min: self.lam = self.lam / self.gamma success = True - self.current_state = cp.deepcopy(self.state) + + # Update state ensemble + self.enX = cp.deepcopy(self.enX_temp) + self.enX_temp = None + + # Update ensemble weights if hasattr(self, 'W'): self.current_W = cp.deepcopy(self.W) + + elif self.data_misfit < self.prev_data_misfit and self.data_misfit_std >= self.prev_data_misfit_std: # accept itaration, but keep lam the same success = True - self.current_state = cp.deepcopy(self.state) + + # Update state ensemble + self.enX = cp.deepcopy(self.enX_temp) + self.enX_temp = None + + # Update ensemble weights if hasattr(self, 'W'): self.current_W = cp.deepcopy(self.W) @@ -298,7 +310,8 @@ def __init__(self, keys_da, keys_en, sim): if self.restart is False: # Save prior state in separate variable - self.prior_state = cp.deepcopy(self.state) + #self.prior_state = cp.deepcopy(self.state) + self.prior_enX = cp.deepcopy(self.enX) # not sure if this is wise! # extract and save state scaling @@ -319,8 +332,7 @@ def __init__(self, keys_da, keys_en, sim): self.check_assimindex_simultaneous() # define the assimilation index self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] - # define the list of states - self.list_states = list(self.state.keys()) + # define the list of datatypes self.list_datatypes, self.list_act_datatypes = at.get_list_data_types( self.obs_data, self.assim_index) @@ -328,7 +340,7 @@ def __init__(self, keys_da, keys_en, sim): self._ext_obs() # Get state scaling and svd of scaled prior self._ext_scaling() - self.current_state = cp.deepcopy(self.state) + # ensure that the updates does not invoke the LM inflation of the Hessian. self.lam = 0 diff --git a/pipt/update_schemes/update_methods_ns/approx_update.py b/pipt/update_schemes/update_methods_ns/approx_update.py index 1ec2194..2728b1e 100644 --- a/pipt/update_schemes/update_methods_ns/approx_update.py +++ b/pipt/update_schemes/update_methods_ns/approx_update.py @@ -19,111 +19,46 @@ class approx_update(): def update(self): # calc the svd of the scaled data pertubation matrix u_d, s_d, v_d = np.linalg.svd(self.pert_preddata, full_matrices=False) - aug_state = at.aug_state(self.current_state, self.list_states, self.cell_index) + #aug_state = at.aug_state(self.current_state, self.list_states, self.cell_index) # remove the last singular value/vector. This is because numpy returns all ne values, while the last is actually # zero. This part is a good place to include eventual additional truncation. if self.trunc_energy < 1: ti = (np.cumsum(s_d) / sum(s_d)) <= self.trunc_energy u_d, s_d, v_d = u_d[:, ti].copy(), s_d[ti].copy(), v_d[ti, :].copy() + + # Check for localization methods if 'localization' in self.keys_da: + if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - if len(self.scale_data.shape) == 1: - E_hat = np.dot(np.expand_dims(self.scale_data ** (-1), - axis=1), np.ones((1, self.ne))) * self.E - x_0 = np.dot(np.diag(s_d[:] ** (-1)), np.dot(u_d[:, :].T, E_hat)) - Lam, z = np.linalg.eig(np.dot(x_0, x_0.T)) + + # Scale data matrix + if self.scale_data.ndim == 1: + E_hat = (self.scale_data ** -1)[:, None] * self.E else: E_hat = solve(self.scale_data, self.E) - x_0 = np.dot(np.diag(s_d[:] ** (-1)), np.dot(u_d[:, :].T, E_hat)) - Lam, z = np.linalg.eig(np.dot(x_0, x_0.T)) - X = np.dot(np.dot(v_d.T, z), solve((self.lam + 1) * np.diag(Lam) + np.eye(len(Lam)), - np.dot(u_d[:, :], np.dot(np.diag(s_d[:] ** (-1)).T, z)).T)) + x_0 = np.diag(1/s_d) @ u_d.T @ E_hat + Lam, z = np.linalg.eig(x_0 @ x_0.T) + X = (v_d.T @ z) @ solve( (self.lam + 1)*np.diag(Lam) + np.eye(len(Lam)), (u_d.T @ (np.diag(1/s_d) @ z)).T ) else: - X = np.dot(np.dot(v_d.T, np.diag(s_d)), - solve(((self.lam + 1) * np.eye(len(s_d)) + np.diag(s_d ** 2)), u_d.T)) - - # we must perform localization - # store the size of all data - data_size = [[self.obs_data[int(time)][data].size if self.obs_data[int(time)][data] is not None else 0 - for data in self.list_datatypes] for time in self.assim_index[1]] - - #f = self.keys_da['localization'] + X = v_d.T @ np.diag(s_d) @ solve( (self.lam + 1)*np.eye(len(s_d)) + np.diag(s_d**2), u_d.T) + + # Check for adaptive localization if 'autoadaloc' in self.localization.loc_info: + self.step = self._update_with_auto_adaptive_localization(X) - # Mean state and perturbation matrix - mean_state = np.mean(aug_state, 1) - if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - pert_state = (self.state_scaling**(-1))[:, None] * (aug_state - np.dot(np.resize(mean_state, (len(mean_state), 1)), - np.ones((1, self.ne)))) - else: - pert_state = (self.state_scaling**(-1) - )[:, None] * np.dot(aug_state, self.proj) - if len(self.scale_data.shape) == 1: - scaled_delta_data = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), - np.ones((1, pert_state.shape[1]))) * ( - self.real_obs_data - self.aug_pred_data) - else: - scaled_delta_data = solve( - self.scale_data, (self.real_obs_data - self.aug_pred_data)) - - self.step = self.localization.auto_ada_loc(self.state_scaling[:, None] * pert_state, np.dot(X, scaled_delta_data), - self.list_states, - **{'prior_info': self.prior_info}) - elif 'localanalysis' in self.localization.loc_info and self.localization.loc_info['localanalysis']: - if 'distance' in self.localization.loc_info: - weight = _calc_loc(self.localization.loc_info['range'], self.localization.loc_info['distance'], - self.prior_info[self.list_states[0]], self.localization.loc_info['type'], self.ne) - else: - # if no distance, do full update - weight = np.ones((aug_state.shape[0], X.shape[1])) - mean_state = np.mean(aug_state, 1) - if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - pert_state = (aug_state - np.dot(np.resize(mean_state, (len(mean_state), 1)), - np.ones((1, self.ne)))) - else: - pert_state = (aug_state - np.dot(np.resize(mean_state, (len(mean_state), 1)), - np.ones((1, self.ne)))) / (np.sqrt(self.ne - 1)) - - if len(self.scale_data.shape) == 1: - scaled_delta_data = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), - np.ones((1, pert_state.shape[1]))) * ( - self.real_obs_data - self.aug_pred_data) - else: - scaled_delta_data = solve( - self.scale_data, (self.real_obs_data - self.aug_pred_data)) - try: - self.step = weight.multiply( - np.dot(pert_state, X)).dot(scaled_delta_data) - except: - self.step = (weight*(np.dot(pert_state, X))).dot(scaled_delta_data) + # Check for local analysis + elif ('localanalysis' in self.localization.loc_info) and (self.localization.loc_info['localanalysis']): + self.step = self._update_with_local_analysis(X) + # Check for distance based localization elif ('dist_loc' in self.keys_da['localization'].keys()) or ('dist_loc' in self.keys_da['localization'].values()): - local_mask = self.localization.localize(self.list_datatypes, [self.keys_da['truedataindex'][int(elem)] - for elem in self.assim_index[1]], - self.list_states, self.ne, self.prior_info, data_size) - mean_state = np.mean(aug_state, 1) - if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - pert_state = (aug_state - np.dot(np.resize(mean_state, (len(mean_state), 1)), - np.ones((1, self.ne)))) - else: - pert_state = (aug_state - np.dot(np.resize(mean_state, (len(mean_state), 1)), - np.ones((1, self.ne)))) / (np.sqrt(self.ne - 1)) - - if len(self.scale_data.shape) == 1: - scaled_delta_data = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), - np.ones((1, pert_state.shape[1]))) * ( - self.real_obs_data - self.aug_pred_data) - else: - scaled_delta_data = solve( - self.scale_data, (self.real_obs_data - self.aug_pred_data)) - - self.step = local_mask.multiply( - np.dot(pert_state, X)).dot(scaled_delta_data) + self.step = self._update_with_distance_based_localization(X) + # Else do parallel update else: act_data_list = {} count = 0 @@ -164,13 +99,13 @@ def update(self): else: # Mean state and perturbation matrix - mean_state = np.mean(aug_state, 1) + mean_state = np.mean(self.enX, 1) if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - pert_state = (self.state_scaling**(-1))[:, None] * (aug_state - np.dot(np.resize(mean_state, (len(mean_state), 1)), + pert_state = (self.state_scaling**(-1))[:, None] * (self.enX - np.dot(np.resize(mean_state, (len(mean_state), 1)), np.ones((1, self.ne)))) else: pert_state = (self.state_scaling**(-1) - )[:, None] * np.dot(aug_state, self.proj) + )[:, None] * np.dot(self.enX, self.proj) if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': if len(self.scale_data.shape) == 1: E_hat = np.dot(np.expand_dims(self.scale_data ** (-1), @@ -202,3 +137,109 @@ def update(self): x_2 = solve(((self.lam + 1) * np.eye(len(s_d)) + np.diag(s_d ** 2)), x_1) x_3 = np.dot(np.dot(v_d.T, np.diag(s_d)), x_2) self.step = np.dot(self.state_scaling[:, None] * pert_state, x_3) + + + + def _update_with_auto_adaptive_localization(self, X): + + # Center ensemble matrix + if ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes'): + pert_state = self.enX - np.mean(self.enX, 1)[:,None] + else: + pert_state = np.dot(self.enX, self.proj) + + # Scale centered ensemble matrix + pert_state = pert_state * (self.state_scaling**(-1))[:, None] + + # Calculate difference between observations and predictions + if len(self.scale_data.shape) == 1: + scaled_delta_data = (self.scale_data ** (-1))[:, None] * (self.real_obs_data - self.aug_pred_data) + else: + scaled_delta_data = solve(self.scale_data, (self.real_obs_data - self.aug_pred_data)) + + # Compute the update step with auto-adaptive localization + step = self.localization.auto_ada_loc( + pert_state = self.state_scaling[:, None]*pert_state, + proj_pred_data = np.dot(X, scaled_delta_data), + curr_param = self.list_states, + prior_info = self.prior_info + ) + + return step + + + def _update_with_local_analysis(self, X): + + # Calculate weights + if 'distance' in self.localization.loc_info: + weight = _calc_loc( + max_dist = self.localization.loc_info['range'], + distance = self.localization.loc_info['distance'], + prior_info = self.prior_info[self.list_states[0]], + loc_type = self.localization.loc_info['type'], + ne = self.ne + ) + else: # if no distance, do full update + weight = np.ones((self.enX.shape[0], X.shape[1])) + + # Center ensemble matrix + mean_state = np.mean(self.enX, axis=1, keepdims=True) + pert_state = self.enX - mean_state + + if not ('emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes'): + pert_state /= np.sqrt(self.ne - 1) + + # Calculate difference between observations and predictions + if self.scale_data.ndim == 1: + scaled_delta_data = (self.scale_data ** -1)[:, None] * (self.real_obs_data - self.aug_pred_data) + else: + scaled_delta_data = solve(self.scale_data, self.real_obs_data - self.aug_pred_data) + + # Compute the update step with local analysis + try: + step = weight.multiply(np.dot(pert_state, X)).dot(scaled_delta_data) + except: + step = (weight*(np.dot(pert_state, X))).dot(scaled_delta_data) + + return step + + + + def _update_with_distance_based_localization(self, X): + + # Get data size + data_size = [[self.obs_data[int(time)][data].size if self.obs_data[int(time)][data] is not None else 0 + for data in self.list_datatypes] for time in self.assim_index[1]] + + # Setup localization + local_mask = self.localization.localize( + self.list_datatypes, + [self.keys_da['truedataindex'][int(elem)] for elem in self.assim_index[1]], + self.list_states, + self.ne, + self.prior_info, + data_size + ) + + # Center ensemble matrix + mean_state = np.mean(self.enX, axis=1, keepdims=True) + pert_state = self.enX - mean_state + if not ('emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes'): + pert_state /= np.sqrt(self.ne - 1) + + # Calculate difference between observations and predictions + if self.scale_data.ndim == 1: + scaled_delta_data = (self.scale_data ** -1)[:, None] * (self.real_obs_data - self.aug_pred_data) + else: + scaled_delta_data = solve(self.scale_data, self.real_obs_data - self.aug_pred_data) + + # Compute the update step with distance-based localization + step = local_mask.multiply(np.dot(pert_state, X)).dot(scaled_delta_data) + + return step + + def _update_with_loclization(self): + pass + + def _update_without_localization(self): + pass From 37b843a5e9472dfec2e21fccec8841ad223bc745 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 24 Sep 2025 12:33:59 +0200 Subject: [PATCH 30/94] replaced self.state with self.enX (matrix) for EnRML and ESMDA --- ensemble/ensemble.py | 128 ++------- pipt/loop/assimilation.py | 40 +-- pipt/loop/ensemble.py | 26 +- pipt/misc_tools/analysis_tools.py | 68 +---- pipt/misc_tools/ensemble_tools.py | 261 ++++++++++++++++++ pipt/misc_tools/extract_tools.py | 3 +- pipt/update_schemes/enrml.py | 8 +- pipt/update_schemes/esmda.py | 28 +- .../update_methods_ns/approx_update.py | 52 ++-- 9 files changed, 363 insertions(+), 251 deletions(-) create mode 100644 pipt/misc_tools/ensemble_tools.py diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index a18bf62..c29da52 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -19,7 +19,7 @@ # Internal imports import pipt.misc_tools.analysis_tools as at import pipt.misc_tools.extract_tools as extract -from geostat.decomp import Cholesky # Making realizations +import pipt.misc_tools.ensemble_tools as entools from pipt.misc_tools import cov_regularization from pipt.misc_tools import wavelet_tools as wt from misc import read_input_csv as rcsv @@ -34,7 +34,7 @@ class Ensemble: implemented here. """ - def __init__(self, keys_en, sim, redund_sim=None): + def __init__(self, keys_en: dict, sim, redund_sim=None): """ Class extends the ReadInitFile class. First the PIPT init. file is passed to the parent class for reading and parsing. Rest of the initialization uses the keywords parsed in ReadInitFile (parent) class to set up observed, @@ -133,8 +133,12 @@ def __init__(self, keys_en, sim, redund_sim=None): if 'importstaticvar' not in self.keys_en: self.ne = int(self.keys_en['ne']) - # Output = self.state, self.cov_prior - self.gen_init_ensemble() + # Generate prior ensemble + self.enX, self.idX, self.cov_prior = entools.generate_prior_ensemble( + prior_info = self.prior_info, + size = self.ne, + save = self.keys_en.get('save_prior', True) + ) else: # State variable imported as a Numpy save file @@ -158,107 +162,8 @@ def __init__(self, keys_en, sim, redund_sim=None): if 'multilevel' in self.keys_en: ml_info = extract.extract_multilevel_info(self.keys_en) self.multilevel, self.tot_level, self.ml_ne, self.ML_error_corr, self.error_comp_scheme, self.ML_corr_done = ml_info - #self._ext_ml_info() - - def _ext_ml_info(self): - ''' - Extract the info needed for ML simulations. Note if the ML keyword is not in keys_en we initialize - such that we only have one level -- the high fidelity one - ''' - - if 'multilevel' in self.keys_en: - # parse - self.multilevel = {} - self.ML_error_corr = 'none' - for i, opt in enumerate(list(zip(*self.keys_en['multilevel']))[0]): - if opt == 'levels': - self.multilevel['levels'] = [elem for elem in range( - int(self.keys_en['multilevel'][i][1]))] - self.tot_level = int(self.keys_en['multilevel'][i][1]) - if opt == 'en_size': - self.multilevel['ne'] = [range(int(el)) - for el in self.keys_en['multilevel'][i][1]] - self.ml_ne = [int(el) for el in self.keys_en['multilevel'][i][1]] - if opt == 'ml_error_corr': - # options for ML_error_corr are: bias_corr, deterministic, stochastic, telescopic - self.ML_error_corr = self.keys_en['multilevel'][i][1] - if not self.ML_error_corr == 'none': - # options for error_comp_scheme are: once, ens, sep - self.error_comp_scheme = self.keys_en['multilevel'][i][2] - self.ML_corr_done = False - - - def gen_init_ensemble(self): - """ - Generate the initial ensemble of (joint) state vectors using the GeoStat class in the "geostat" package. - TODO: Merge this function with the perturbation function _gen_state_ensemble in popt. - """ - # Initialize GeoStat class - init_en = Cholesky() - - # (Re)initialize state variable as dictionary - self.state = {} - self.cov_prior = {} - - for name in self.prior_info: - # Init. indices to pick out correct mean vector for each layer - ind_end = 0 - - # Extract info. - nx = self.prior_info[name].get('nx', 0) - ny = self.prior_info[name].get('ny', 0) - nz = self.prior_info[name].get('nz', 0) - mean = self.prior_info[name].get('mean', None) - - if nx == ny == 0: # assume ensemble will be generated elsewhere if dimensions are zero - break - - variance = self.prior_info[name].get('variance', None) - corr_length = self.prior_info[name].get('corr_length', None) - aniso = self.prior_info[name].get('aniso', None) - vario = self.prior_info[name].get('vario', None) - angle = self.prior_info[name].get('angle', None) - limits= self.prior_info[name].get('limits',None) - - # Loop over nz to make layers of 2D priors - for i in range(self.prior_info[name]['nz']): - # If mean is scalar, no covariance matrix is needed - if type(self.prior_info[name]['mean']).__module__ == 'numpy': - # Generate covariance matrix - cov = init_en.gen_cov2d( - nx, ny, variance[i], corr_length[i], aniso[i], angle[i], vario[i]) - else: - cov = np.array(variance[i]) - - # Pick out the mean vector for the current layer - ind_start = ind_end - ind_end = int((i + 1) * (len(mean) / nz)) - mean_layer = mean[ind_start:ind_end] - - # Generate realizations. If LIMITS have been entered, they must be taken account for here - if limits is None: - real = init_en.gen_real(mean_layer, cov, self.ne) - else: - real = init_en.gen_real(mean_layer, cov, self.ne, { - 'upper': limits[i][1], 'lower': limits[i][0]}) - - # Stack realizations for each layer - if i == 0: - real_out = real - else: - real_out = np.vstack((real_out, real)) - - # Store realizations in dictionary with name given in STATICVAR - self.state[name] = real_out - - # Store the covariance matrix - self.cov_prior[name] = cov - # Save the ensemble for later inspection - np.savez('prior.npz', **self.state) - - def get_list_assim_steps(self): """ Returns list of assimilation steps. Useful in a 'loop'-script. @@ -330,7 +235,7 @@ def calc_prediction(self, enX=None, save_prediction=None): self.sim.setup_fwd_run(redund_sim=self.sim.redund_sim) # Convert ensemble matrix to list of dictionaries - enX = at.ensmeble_matrix_to_list(enX, self.idX) + enX = entools.matrix_to_list(enX, self.idX) if not (self.aux_input is None): for n in range(self.ne): @@ -359,7 +264,7 @@ def calc_prediction(self, enX=None, save_prediction=None): ###################################################################################################################### # Convert state enemble back to matrix form - enX = at.ensemble_list_to_matrix(enX, self.idX) + enX = entools.list_to_matrix(enX, self.idX) # restore state ensemble if it was not inputted if not use_input_ensemble: @@ -393,17 +298,16 @@ def calc_prediction(self, enX=None, save_prediction=None): copy_member = np.random.choice(list_success, size=len(list_crash), replace=True) # Insert the replaced runs in prediction list - for indx, el in enumerate(copy_member): + for index, element in enumerate(copy_member): msg = ( - f"\033[92m--- Ensemble member {list_crash[indx]} failed, " - f"has been replaced by ensemble member {el}! ---\033[92m" + f"\033[92m--- Ensemble member {list_crash[index]} failed, " + f"has been replaced by ensemble member {element}! ---\033[92m" ) print(msg) self.logger.info(msg) - for key in self.state.keys(): - if self.state[key].ndim > 1: - self.state[key][:, list_crash[indx]] = deepcopy(self.state[key][:, el]) - en_pred[list_crash[indx]] = deepcopy(en_pred[el]) + if enX.shape[1] > 1: + enX[:, list_crash[index]] = deepcopy(self.enX[:, element]) + en_pred[list_crash[index]] = deepcopy(en_pred[element]) # Convert ensemble specific result into pred_data, and filter for NONE data self.pred_data.extend([{typ: np.concatenate(tuple((el[ind][typ][:, np.newaxis]) for el in en_pred), axis=1) diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index e1c634e..40158f1 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -20,7 +20,9 @@ from pipt.loop.ensemble import Ensemble from misc.system_tools.environ_var import OpenBlasSingleThread from pipt.misc_tools import analysis_tools as at + import pipt.misc_tools.extract_tools as extract +import pipt.misc_tools.ensemble_tools as entools class Assimilate: @@ -96,7 +98,7 @@ def run(self): self.ensemble.logger, self.ensemble.prior_info, self.ensemble.sim, - self.ensemble.prior_state + entools.matrix_to_dict(self.ensemble.prior_enX, self.ensemble.idX) ) # Run a while loop until max. iterations or convergence is reached @@ -113,7 +115,12 @@ def run(self): if 'qa' in self.ensemble.keys_da: # Check if we want to perform a Quality Assurance of the forecast # set updated prediction, state and lam - qaqc.set(self.ensemble.pred_data, self.ensemble.state, self.ensemble.lam) + qaqc.set( + self.ensemble.pred_data, + entools.matrix_to_dict(self.ensemble.enX, self.ensemble.idX), + self.ensemble.lam + ) + # Level 1,2 all data, and subspace qaqc.calc_mahalanobis((1, 'time', 2, 'time', 1, None, 2, None)) qaqc.calc_coverage() # Compute data coverage @@ -166,13 +173,19 @@ def run(self): self._save_analysis_debug() if 'qc' in self.ensemble.keys_da: # Check if we want to perform a Quality Control of the updated state # set updated prediction, state and lam - qaqc.set(self.ensemble.pred_data, - self.ensemble.state, self.ensemble.lam) + qaqc.set( + self.ensemble.pred_data, + entools.matrix_to_dict(self.ensemble.enX, self.ensemble.idX), + self.ensemble.lam + ) qaqc.calc_da_stat() # Compute statistics for updated parameters if 'qa' in self.ensemble.keys_da: # Check if we want to perform a Quality Assurance of the forecast # set updated prediction, state and lam - qaqc.set(self.ensemble.pred_data, - self.ensemble.state, self.ensemble.lam) + qaqc.set( + self.ensemble.pred_data, + entools.matrix_to_dict(self.ensemble.enX, self.ensemble.idX), + self.ensemble.lam + ) qaqc.calc_mahalanobis( (1, 'time', 2, 'time', 1, None, 2, None)) # Level 1,2 all data, and subspace # qaqc.calc_coverage() # Compute data coverage @@ -266,9 +279,11 @@ def remove_outliers(self): new_index = np.random.choice(members) # replace state - for el in self.ensemble.state.keys(): - self.ensemble.state[el][:, index] = deepcopy( - self.ensemble.state[el][:, new_index]) + if self.ensemble.enX_temp is not None: + self.ensemble.enX[:, index] = deepcopy(self.ensemble.enX[:, new_index]) + else: + self.ensemble.enX_temp[:, index] = deepcopy(self.ensemble.enX_temp[:, new_index]) + # replace the failed forecast for i, data_ind in enumerate(self.ensemble.pred_data): @@ -465,13 +480,6 @@ def calc_forecast(self): if 'post_process_forecast' in self.ensemble.keys_da and self.ensemble.keys_da['post_process_forecast'] == 'yes': self.post_process_forecast() - # If we have dynamic variables, and we are in the first assimilation step, we must convert lists to (2D) - # numpy arrays - if 'dynamicvar' in self.ensemble.keys_da and assim_step == 0: - for dyn_state in self.ensemble.keys_da['dynamicvar']: - self.ensemble.state[dyn_state] = np.array( - self.ensemble.state[dyn_state]).T - # Extra option debug if 'saveforecast' in self.ensemble.sim.input_dict: with open('sim_results.p', 'wb') as f: diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index 2b98624..d4e494b 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -18,8 +18,11 @@ import misc.read_input_csv as rcsv from pipt.misc_tools import wavelet_tools as wt from pipt.misc_tools.cov_regularization import localization, _calc_distance + +# Import internal tools import pipt.misc_tools.analysis_tools as at import pipt.misc_tools.extract_tools as extract +import pipt.misc_tools.ensemble_tools as entools class Ensemble(PETEnsemble): @@ -99,19 +102,8 @@ def __init__(self, keys_da, keys_en, sim): self._org_obs_data() self._org_data_var() - # define projection for centring and scaling - self.proj = (np.eye(self.ne) - (1 / self.ne) * - np.ones((self.ne, self.ne))) / np.sqrt(self.ne - 1) - - # If we have dynamic state variables, we allocate keys for them in 'state'. Since we do not know the size - # of the arrays of the dynamic variables, we only allocate an NE list to be filled in later (in - # calc_forecast) - if 'dynamicvar' in self.keys_da: - dyn_vars = self.keys_da['dynamicvar'] - if not isinstance(dyn_vars, list): - dyn_vars = [dyn_vars] - for name in dyn_vars: - self.state[name] = [None] * self.ne + # Define projection operator for centring and scaling ensemble matrix + self.proj = (np.eye(self.ne) - np.ones((self.ne, self.ne))/self.ne) / np.sqrt(self.ne - 1) # Option to store the dictionaries containing observed data and data variance if 'obsvarsave' in self.keys_da and self.keys_da['obsvarsave'] == 'yes': @@ -129,7 +121,7 @@ def __init__(self, keys_da, keys_en, sim): # Initialize local analysis if 'localanalysis' in self.keys_da: - self.local_analysis = extract.extract_local_analysis_info(self.keys_da['localanalysis'], self.state.keys()) + self.local_analysis = extract.extract_local_analysis_info(self.keys_da['localanalysis'], self.idX.keys()) self.pred_data = [{k: np.zeros((1, self.ne), dtype='float32') for k in self.keys_da['datatype']} for _ in self.obs_data] @@ -558,7 +550,7 @@ def save_temp_state_assim(self, ind_save): self.temp_state = [None]*(len(self.get_list_assim_steps()) + 1) # Save the state - self.temp_state[ind_save] = deepcopy(self.state) + self.temp_state[ind_save] = entools.matrix_to_dict(self.enX) np.savez('temp_state_assim', self.temp_state) def save_temp_state_iter(self, ind_save, max_iter): @@ -579,7 +571,7 @@ def save_temp_state_iter(self, ind_save, max_iter): self.temp_state = [None] * (int(max_iter) + 1) # +1 due to init. ensemble # Save state - self.temp_state[ind_save] = deepcopy(self.state) + self.temp_state[ind_save] = entools.matrix_to_dict(self.enX) np.savez('temp_state_iter', self.temp_state) def save_temp_state_mda(self, ind_save): @@ -602,7 +594,7 @@ def save_temp_state_mda(self, ind_save): self.temp_state = [None] * (int(self.tot_assim) + 1) # Save state - self.temp_state[ind_save] = deepcopy(self.state) + self.temp_state[ind_save] = entools.matrix_to_dict(self.enX) np.savez('temp_state_mda', self.temp_state) def save_temp_state_ml(self, ind_save): diff --git a/pipt/misc_tools/analysis_tools.py b/pipt/misc_tools/analysis_tools.py index ebd6e59..97b8fe4 100644 --- a/pipt/misc_tools/analysis_tools.py +++ b/pipt/misc_tools/analysis_tools.py @@ -6,6 +6,13 @@ implementing, leave it in that class. """ +__all__ = [ + 'parallel_upd', + 'calc_autocov', + 'calc_crosscov', + 'calc_objectivefun' +] + # External imports import numpy as np # Numerical tools from scipy import linalg # Linear algebra tools @@ -1181,67 +1188,6 @@ def compute_x(pert_preddata, cov_data, keys_da, alfa=None): return X -def ensmeble_matrix_to_list(matrix: np.ndarray, indecies: dict) -> list[dict]: - ''' - Convert an ensemble matrix to a list of dictionaries. - - Parameters - ---------- - matrix : np.ndarray - Ensemble matrix where each column represents an ensemble member. - indecies : dict - Dictionary with keys as variable names and values as tuples indicating the start and end row indices - for each variable in the ensemble matrix. - - Returns - ------- - ensemble_list : list of dict - ''' - ne = matrix.shape[1] - ensemble_list = [] - - for n in range(ne): - member = {} - for key, (start, end) in indecies.items(): - if matrix[start:end].ndim == 2: - member[key] = matrix[start:end, n] - else: - member[key] = matrix[start:end] - ensemble_list.append(member) - - return ensemble_list - -def ensemble_list_to_matrix(ensemble_list: list[dict], indecies: dict) -> np.ndarray: - ''' - Convert a list of dictionaries to an ensemble matrix. - - Parameters - ---------- - ensemble_list : list of dict - List where each dictionary represents an ensemble member with variable names as keys. - indecies : dict - Dictionary with keys as variable names and values as tuples indicating the start and end row indices - for each variable in the ensemble matrix. - - Returns - ------- - matrix : np.ndarray - Ensemble matrix where each column represents an ensemble member. - ''' - ne = len(ensemble_list) - nx = sum(end - start for start, end in indecies.values()) - matrix = np.zeros((nx, ne)) - - for n, member in enumerate(ensemble_list): - for key, (start, end) in indecies.items(): - if member[key].ndim == 2: - matrix[start:end, n] = member[key][:,n] - else: - matrix[start:end, n] = member[key] - - return matrix - - def aug_state(state, list_state, cell_index=None): """ Augment the state variables to an array. diff --git a/pipt/misc_tools/ensemble_tools.py b/pipt/misc_tools/ensemble_tools.py new file mode 100644 index 0000000..2f88b81 --- /dev/null +++ b/pipt/misc_tools/ensemble_tools.py @@ -0,0 +1,261 @@ +# This module contains functions and tools for ensembles + +__all__ = [ + 'matrix_to_dict', + 'matrix_to_list', + 'list_to_matrix', + 'generate_prior_ensemble', + 'clip_matrix' +] + +# Imports +import numpy as np + +# Internal imports +from geostat.decomp import Cholesky + + +def matrix_to_dict(matrix: np.ndarray, indecies: dict[tuple]) -> dict: + ''' + Convert an ensemble matrix to a dictionary of arrays. + + Parameters + ---------- + matrix : np.ndarray + Ensemble matrix where each column represents an ensemble member. + indecies : dict + Dictionary with keys as variable names and values as tuples indicating the start and end row indices + for each variable in the ensemble matrix. + + Returns + ------- + ensemble_dict : dict + Dictionary with keys as variable names and values as arrays of shape (nx, ne). + ''' + ensemble_dict = {} + for key, (start, end) in indecies.items(): + ensemble_dict[key] = matrix[start:end] + + return ensemble_dict + + +def matrix_to_list(matrix: np.ndarray, indecies: dict[tuple]) -> list[dict]: + ''' + Convert an ensemble matrix to a list of dictionaries. + + Parameters + ---------- + matrix : np.ndarray + Ensemble matrix where each column represents an ensemble member. + indecies : dict + Dictionary with keys as variable names and values as tuples indicating the start and end row indices + for each variable in the ensemble matrix. + + Returns + ------- + ensemble_list : list of dict + ''' + ne = matrix.shape[1] + ensemble_list = [] + + for n in range(ne): + member = matrix_to_dict(matrix[:,n], indecies) + ensemble_list.append(member) + + return ensemble_list + + +def list_to_matrix(ensemble_list: list[dict], indecies: dict[tuple]) -> np.ndarray: + ''' + Convert a list of dictionaries to an ensemble matrix. + + Parameters + ---------- + ensemble_list : list of dict + List where each dictionary represents an ensemble member with variable names as keys. + indecies : dict + Dictionary with keys as variable names and values as tuples indicating the start and end row indices + for each variable in the ensemble matrix. + + Returns + ------- + matrix : np.ndarray + Ensemble matrix where each column represents an ensemble member. + ''' + ne = len(ensemble_list) + nx = sum(end - start for start, end in indecies.values()) + matrix = np.zeros((nx, ne)) + + for n, member in enumerate(ensemble_list): + for key, (start, end) in indecies.items(): + if member[key].ndim == 2: + matrix[start:end, n] = member[key][:,n] + else: + matrix[start:end, n] = member[key] + + return matrix + + +def generate_prior_ensemble(prior_info: dict, size: int, save: bool = True) -> tuple[np.ndarray, dict, dict]: + ''' + Generate a prior ensemble based on provided prior information. + + Parameters + ---------- + prior_info : dict + Dictionary containing prior information for each state variable. + + size : int + Size of ensemble. + + save : bool, optional + Whether to save the generated ensemble to a file. Default is True. + + Returns + ------- + enX : np.ndarray + The generated ensemble matrix, shape: (nx, ne). + + idX : dict + Dictionary with keys as variable names and values as tuples indicating the start and end row indices + for each variable in the ensemble matrix. + + cov_prior : dict + Dictionary containing the covariance matrices for each state variable. + ''' + + # Initialize sampler + generator = Cholesky() + + # Initialize variables + enX = None + idX = {} + cov_prior = {} + + # Loop over all state variables + for name, info in prior_info.items(): + + # Extract info + nx = info.get('nx', 0) + ny = info.get('ny', 0) + nz = info.get('nz', 0) + mean = info.get('mean', None) + + # if no dimensions are given, nothing is generated for this variable + if nx == ny == 0: + break + + # Extract more options + variance = info.get('variance', None) + corr_length = info.get('corr_length', None) + aniso = info.get('aniso', None) + vario = info.get('vario', None) + angle = info.get('angle', None) + limits= info.get('limits',None) + + # Loop over nz to make layers of 2D priors + index_stop = 0 + for idz in range(nz): + # If mean is scalar, no covariance matrix is needed + if isinstance(mean, (list, np.ndarray)) and len(mean) > 1: + # Generate covariance matrix + cov = generator.gen_cov2d( + x_size = nx, + y_size = ny, + variance = variance[idz], + var_range = corr_length[idz], + aspect = aniso[idz], + angle = angle[idz], + var_type = vario[idz] + ) + else: + cov = np.array(variance[idz]) + + # Pick out the mean vector for the current layer + index_start = index_stop + index_stop = int((idz + 1) * (len(mean)/nz)) + mean_layer = mean[index_start:index_stop] + + # Generate realizations. If LIMITS have been entered, they must be taken account for here + if limits is None: + real = generator.gen_real(mean_layer, cov, size) + else: + real = generator.gen_real(mean_layer, cov, size, limits[idz]) + + # Stack realizations for each layer + if idz == 0: + real_out = real + else: + real_out = np.vstack((real_out, real)) + + # Fill in the ensemble matrix and indecies + if enX is None: + idX[name] = (0, real_out.shape[0]) + enX = real_out + else: + idX[name] = (enX.shape[0], enX.shape[0] + real_out.shape[0]) + enX = np.vstack((enX, real_out)) + + # Store the covariance matrix + cov_prior[name] = cov + + # Save prior ensemble + if save: + np.savez( + 'prior_ensemble.npz', + **{name: enX[idX[name][0]:idX[name][1]] for name in idX.keys()} + ) + + return enX, idX, cov_prior + + +def clip_matrix(matrix: np.ndarray, limits: dict|tuple|list, indecies: dict|None = None) -> np.ndarray: + ''' + Clip the values in an ensemble matrix based on provided limits. + + Parameters + ---------- + matrix : np.ndarray + Ensemble matrix where each column represents an ensemble member. + + limits : dict, tuple, or list + If tuple, it should be (lower_bound, upper_bound) applied to all variables. + If dict, it should have variable names as keys and (lower_bound, upper_bound) as values. + If list, it should contain (lower_bound, upper_bound) tuples for each variable in the order of indecies. + + indecies : dict, optional + Dictionary with keys as variable names and values as tuples indicating the start and end row indices + for each variable in the ensemble matrix. Required if limits is a dict or list. Default is None. + + Returns + ------- + matrix : np.ndarray + ''' + if isinstance(limits, tuple): + lb, ub = limits + if not (lb is None and ub is None): + matrix = np.clip(matrix, lb, ub) + + elif isinstance(limits, dict) and isinstance(limits, dict): + if indecies is None: + raise ValueError("When limits is a dictionary, indecies must also be provided.") + + for key, (start, end) in indecies.items(): + if key in limits: + lb, ub = limits[key] + if not (lb is None and ub is None): + matrix[start:end] = np.clip(matrix[start:end], lb, ub) + + elif isinstance(limits, list): + if indecies is None: + raise ValueError("When limits is a list, indecies must also be provided.") + + if len(limits) != len(indecies): + raise ValueError("Length of limits list must match number of variables in indecies.") + + for (key, (start, end)), (lb, ub) in zip(indecies.items(), limits): + if not (lb is None and ub is None): + matrix[start:end] = np.clip(matrix[start:end], lb, ub) + + return matrix + diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index 46cbfe0..186c342 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -1,4 +1,5 @@ -# This module inlcudes fucntions for extracting information from input dicts +# This module inlcudes functions for extracting information from input dicts + __all__ = [ 'extract_prior_info', 'extract_multilevel_info', diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index c01df81..8e0742d 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -4,6 +4,8 @@ # External imports import pipt.misc_tools.analysis_tools as at import pipt.misc_tools.extract_tools as extract +import pipt.misc_tools.ensemble_tools as entools + from geostat.decomp import Cholesky from pipt.loop.ensemble import Ensemble from pipt.update_schemes.update_methods_ns.subspace_update import subspace_update @@ -135,9 +137,9 @@ def calc_analysis(self): self.enX_temp = np.dot(self.prior_enX, (np.eye(self.ne) + self.W/np.sqrt(self.ne - 1))) - # Extract updated state variables from aug_update - #self.state = at.update_state(aug_state_upd, self.state, self.list_states) - #self.state = at.limits(self.state, self.prior_info) + # Ensure limits are respected + limits = {key: self.prior_info[key].get('limits', (None, None)) for key in self.idX.keys()} + self.enX_temp = entools.clip_matrix(self.enX_temp, limits, self.idX) def check_convergence(self): """ diff --git a/pipt/update_schemes/esmda.py b/pipt/update_schemes/esmda.py index 2e5ad59..0ab6fdc 100644 --- a/pipt/update_schemes/esmda.py +++ b/pipt/update_schemes/esmda.py @@ -11,6 +11,7 @@ # Internal imports from pipt.loop.ensemble import Ensemble import pipt.misc_tools.analysis_tools as at +import pipt.misc_tools.ensemble_tools as entools # import update schemes from pipt.update_schemes.update_methods_ns.approx_update import approx_update @@ -45,8 +46,8 @@ def __init__(self, keys_da, keys_en, sim): self.prev_data_misfit = None if self.restart is False: - self.prior_state = deepcopy(self.state) - self.list_states = list(self.state.keys()) + self.prior_enX = deepcopy(self.enX) + self.list_states = list(self.idX.keys()) # At the moment, the iterative loop is threated as an iterative smoother an thus we check if assim. indices # are given as in the Simultaneous loop. self.check_assimindex_simultaneous() @@ -71,7 +72,7 @@ def __init__(self, keys_da, keys_en, sim): self.real_obs_data_conv = deepcopy(self.real_obs_data) # Get state scaling and svd of scaled prior self._ext_scaling() - self.current_state = deepcopy(self.state) + # Extract the inflation parameter from MDA keyword self.alpha = self._ext_inflation_param() @@ -147,20 +148,19 @@ def calc_analysis(self): self.pert_preddata = scilinalg.solve( self.scale_data, np.dot(self.aug_pred_data, self.proj)) - aug_state = at.aug_state(self.current_state, self.list_states) - self.update() + + # Update the state ensemble and weights if hasattr(self, 'step'): - aug_state_upd = aug_state + self.step + self.enX_temp = self.enX + self.step if hasattr(self, 'w_step'): self.W = self.current_W + self.w_step - aug_prior_state = at.aug_state(self.prior_state, self.list_states) - aug_state_upd = np.dot(aug_prior_state, (np.eye( - self.ne) + self.W / np.sqrt(self.ne - 1))) + self.enX_temp = np.dot(self.prior_enX, (np.eye(self.ne) + self.W/np.sqrt(self.ne - 1))) + - # Extract updated state variables from aug_update - self.state = at.update_state(aug_state_upd, self.state, self.list_states) - self.state = at.limits(self.state, self.prior_info) + # Ensure limits are respected + limits = {key: self.prior_info[key].get('limits', (None, None)) for key in self.idX.keys()} + self.enX_temp = entools.clip_matrix(self.enX_temp, limits, self.idX) def check_convergence(self): """ @@ -200,7 +200,9 @@ def check_convergence(self): self.logger.info( f'MDA iteration number {self.iteration}! Objective function increased from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') # Return conv = False, why_stop var. - self.current_state = deepcopy(self.state) + # Update state ensemble + self.enX = deepcopy(self.enX_temp) + self.enX_temp = None if hasattr(self, 'W'): self.current_W = deepcopy(self.W) diff --git a/pipt/update_schemes/update_methods_ns/approx_update.py b/pipt/update_schemes/update_methods_ns/approx_update.py index 2728b1e..fe74113 100644 --- a/pipt/update_schemes/update_methods_ns/approx_update.py +++ b/pipt/update_schemes/update_methods_ns/approx_update.py @@ -31,10 +31,10 @@ def update(self): if 'localization' in self.keys_da: if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - + # Scale data matrix - if self.scale_data.ndim == 1: - E_hat = (self.scale_data ** -1)[:, None] * self.E + if len(self.scale_data.shape) == 1: + E_hat = (1/self.scale_data)[:, None] * self.E else: E_hat = solve(self.scale_data, self.E) @@ -98,42 +98,40 @@ def update(self): self.step = at.aug_state(self.step, self.list_states) else: - # Mean state and perturbation matrix - mean_state = np.mean(self.enX, 1) + # Centered ensemble matrix if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - pert_state = (self.state_scaling**(-1))[:, None] * (self.enX - np.dot(np.resize(mean_state, (len(mean_state), 1)), - np.ones((1, self.ne)))) + pert_state = (self.state_scaling**(-1))[:, None] * (self.enX - np.mean(self.enX, axis=1, keepdims=True)) else: - pert_state = (self.state_scaling**(-1) - )[:, None] * np.dot(self.enX, self.proj) + pert_state = (self.state_scaling**(-1))[:, None] * np.dot(self.enX, self.proj) + if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': + + # Scale data matrix if len(self.scale_data.shape) == 1: - E_hat = np.dot(np.expand_dims(self.scale_data ** (-1), - axis=1), np.ones((1, self.ne))) * self.E - x_0 = np.dot(np.diag(s_d[:] ** (-1)), np.dot(u_d[:, :].T, E_hat)) - Lam, z = np.linalg.eig(np.dot(x_0, x_0.T)) - x_1 = np.dot(np.dot(u_d[:, :], np.dot(np.diag(s_d[:] ** (-1)).T, z)).T, - np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), np.ones((1, self.ne))) * - (self.real_obs_data - self.aug_pred_data)) + E_hat = (1/self.scale_data)[:, None] * self.E else: E_hat = solve(self.scale_data, self.E) - x_0 = np.dot(np.diag(s_d[:] ** (-1)), np.dot(u_d[:, :].T, E_hat)) - Lam, z = np.linalg.eig(np.dot(x_0, x_0.T)) - x_1 = np.dot(np.dot(u_d[:, :], np.dot(np.diag(s_d[:] ** (-1)).T, z)).T, - solve(self.scale_data, (self.real_obs_data - self.aug_pred_data))) + x_0 = np.diag(s_d ** -1) @ u_d.T @ E_hat + Lam, z = np.linalg.eig(x_0 @ x_0.T) + + if len(self.scale_data.shape) == 1: + delta_data = (1/self.scale_data)[:, None] * (self.real_obs_data - self.aug_pred_data) + else: + delta_data = solve(self.scale_data, self.real_obs_data - self.aug_pred_data) + + x_1 = (u_d @ (np.diag(s_d ** -1).T @ z)).T @ delta_data x_2 = solve((self.lam + 1) * np.diag(Lam) + np.eye(len(Lam)), x_1) x_3 = np.dot(np.dot(v_d.T, z), x_2) - delta_1 = np.dot(self.state_scaling[:, None] * pert_state, x_3) - self.step = delta_1 + self.step = np.dot(self.state_scaling[:, None] * pert_state, x_3) + else: # Compute the approximate update (follow notation in paper) if len(self.scale_data.shape) == 1: - x_1 = np.dot(u_d.T, np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), np.ones((1, self.ne))) * - (self.real_obs_data - self.aug_pred_data)) + x_1 = np.dot(u_d.T, (1/self.scale_data)[:, None] * (self.real_obs_data - self.aug_pred_data)) else: - x_1 = np.dot(u_d.T, solve(self.scale_data, - (self.real_obs_data - self.aug_pred_data))) + x_1 = np.dot(u_d.T, solve(self.scale_data, self.real_obs_data - self.aug_pred_data)) + x_2 = solve(((self.lam + 1) * np.eye(len(s_d)) + np.diag(s_d ** 2)), x_1) x_3 = np.dot(np.dot(v_d.T, np.diag(s_d)), x_2) self.step = np.dot(self.state_scaling[:, None] * pert_state, x_3) @@ -241,5 +239,3 @@ def _update_with_distance_based_localization(self, X): def _update_with_loclization(self): pass - def _update_without_localization(self): - pass From 1e1b19a8bf5a10484c64645e60fb9355ea6717fb Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 24 Sep 2025 12:42:22 +0200 Subject: [PATCH 31/94] removed temp_save functions (they are not in use anymore) --- pipt/loop/assimilation.py | 3 --- pipt/loop/ensemble.py | 8 ++++---- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index 40158f1..099a384 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -166,9 +166,6 @@ def run(self): # self._save_iteration_information() if self.ensemble.iteration > 0: - # Temporary save state if options in TEMPSAVE have been given and the option is not 'no' - if 'tempsave' in self.ensemble.keys_da and self.ensemble.keys_da['tempsave'] != 'no': - self._save_during_iteration(self.ensemble.keys_da['tempsave']) if 'analysisdebug' in self.ensemble.keys_da: self._save_analysis_debug() if 'qc' in self.ensemble.keys_da: # Check if we want to perform a Quality Control of the updated state diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index d4e494b..4bd4bbb 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -550,7 +550,7 @@ def save_temp_state_assim(self, ind_save): self.temp_state = [None]*(len(self.get_list_assim_steps()) + 1) # Save the state - self.temp_state[ind_save] = entools.matrix_to_dict(self.enX) + self.temp_state[ind_save] = deepcopy(entools.matrix_to_dict(self.enX)) np.savez('temp_state_assim', self.temp_state) def save_temp_state_iter(self, ind_save, max_iter): @@ -571,7 +571,7 @@ def save_temp_state_iter(self, ind_save, max_iter): self.temp_state = [None] * (int(max_iter) + 1) # +1 due to init. ensemble # Save state - self.temp_state[ind_save] = entools.matrix_to_dict(self.enX) + self.temp_state[ind_save] = deepcopy(entools.matrix_to_dict(self.enX)) np.savez('temp_state_iter', self.temp_state) def save_temp_state_mda(self, ind_save): @@ -594,7 +594,7 @@ def save_temp_state_mda(self, ind_save): self.temp_state = [None] * (int(self.tot_assim) + 1) # Save state - self.temp_state[ind_save] = entools.matrix_to_dict(self.enX) + self.temp_state[ind_save] = deepcopy(entools.matrix_to_dict(self.enX)) np.savez('temp_state_mda', self.temp_state) def save_temp_state_ml(self, ind_save): @@ -617,7 +617,7 @@ def save_temp_state_ml(self, ind_save): self.temp_state = [None] * (int(self.tot_assim) + 1) # Save state - self.temp_state[ind_save] = deepcopy(self.state) + self.temp_state[ind_save] = deepcopy(entools.matrix_to_dict(self.enX)) np.savez('temp_state_ml', self.temp_state) def compress(self, data=None, vintage=0, aug_coeff=None): From 52e41abca136ac7e5fef5a4e595b82e463060e17 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 5 Nov 2025 09:16:23 +0100 Subject: [PATCH 32/94] Clean up and redsign update for EnRML and ESMDA --- pipt/loop/assimilation.py | 29 --- pipt/loop/ensemble.py | 85 -------- pipt/misc_tools/analysis_tools.py | 11 + pipt/update_schemes/enrml.py | 15 +- pipt/update_schemes/esmda.py | 15 +- .../update_methods_ns/approx_update.py | 191 ++++++++++-------- .../update_methods_ns/full_update.py | 90 +++++---- 7 files changed, 185 insertions(+), 251 deletions(-) diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index 099a384..dfe2660 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -318,35 +318,6 @@ def _save_iteration_information(self): # Note: the function must be named main, and we pass the full current instance of the object. iter_info_func.main(self) - def _save_during_iteration(self, tempsave): - """ - Save during an iteration. How often is determined by the `TEMPSAVE` keyword; confer the manual for all the - different options. - - Parameters - ---------- - tempsave : list - Info. from the TEMPSAVE keyword - """ - self.ensemble.logger.info( - 'The TEMPSAVE feature is no longer supported. Please you debug_analyses, or iterinfo.') - # Save at specific points - # if isinstance(tempsave, list): - # # Save at regular intervals - # if tempsave[0] == 'each' or tempsave[0] == 'every' and self.ensemble.iteration % tempsave[1] == 0: - # self.ensemble.save_temp_state_iter(self.ensemble.iteration + 1, self.max_iter) - # - # # Save at points given by input - # elif tempsave[0] == 'list' or tempsave[0] == 'at': - # # Check if one or more save points have been given, and save if we are at that point - # savepoint = tempsave[1] if isinstance(tempsave[1], list) else [tempsave[1]] - # if self.ensemble.iteration in savepoint: - # self.ensemble.save_temp_state_iter(self.ensemble.iteration + 1, self.max_iter) - # - # # Save at all assimilation steps - # elif tempsave == 'yes' or tempsave == 'all': - # self.ensemble.save_temp_state_iter(self.ensemble.iteration + 1, self.max_iter) - def _save_analysis_debug(self): """ Moved Old analysis debug here to retain consistency. diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index 4bd4bbb..b5ada98 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -534,91 +534,6 @@ def _ext_scaling(self): self.Am = None - def save_temp_state_assim(self, ind_save): - """ - Method to save the state variable during the assimilation. It is stored in a list with length = tot. no. - assim. steps + 1 (for the init. ensemble). The list of temporary states are also stored as a .npz file. - - Parameters - ---------- - ind_save : int - Assim. step to save (0 = prior) - """ - # Init. temp. save - if ind_save == 0: - # +1 due to init. ensemble - self.temp_state = [None]*(len(self.get_list_assim_steps()) + 1) - - # Save the state - self.temp_state[ind_save] = deepcopy(entools.matrix_to_dict(self.enX)) - np.savez('temp_state_assim', self.temp_state) - - def save_temp_state_iter(self, ind_save, max_iter): - """ - Save a snapshot of state at current iteration. It is stored in a list with length equal to max. iteration - length + 1 (due to prior state being 0). The list of temporary states are also stored as a .npz file. - - !!! warning - Max. iterations must be defined before invoking this method. - - Parameters - ---------- - ind_save : int - Iteration step to save (0 = prior) - """ - # Initial save - if ind_save == 0: - self.temp_state = [None] * (int(max_iter) + 1) # +1 due to init. ensemble - - # Save state - self.temp_state[ind_save] = deepcopy(entools.matrix_to_dict(self.enX)) - np.savez('temp_state_iter', self.temp_state) - - def save_temp_state_mda(self, ind_save): - """ - Save a snapshot of the state during a MDA loop. The temporary state will be stored as a list with length - equal to the tot. no. of assimilations + 1 (init. ensemble saved in 0 entry). The list of temporary states - are also stored as a .npz file. - - !!! warning - Tot. no. of assimilations must be defined before invoking this method. - - Parameter - --------- - ind_save : int - Assim. step to save (0 = prior) - """ - # Initial save - if ind_save == 0: - # +1 due to init. ensemble - self.temp_state = [None] * (int(self.tot_assim) + 1) - - # Save state - self.temp_state[ind_save] = deepcopy(entools.matrix_to_dict(self.enX)) - np.savez('temp_state_mda', self.temp_state) - - def save_temp_state_ml(self, ind_save): - """ - Save a snapshot of the state during a ML loop. The temporary state will be stored as a list with length - equal to the tot. no. of assimilations + 1 (init. ensemble saved in 0 entry). The list of temporary states - are also stored as a .npz file. - - !!! warning - Tot. no. of assimilations must be defined before invoking this method. - - Parameters - ---------- - ind_save : int - Assim. step to save (0 = prior) - """ - # Initial save - if ind_save == 0: - # +1 due to init. ensemble - self.temp_state = [None] * (int(self.tot_assim) + 1) - - # Save state - self.temp_state[ind_save] = deepcopy(entools.matrix_to_dict(self.enX)) - np.savez('temp_state_ml', self.temp_state) def compress(self, data=None, vintage=0, aug_coeff=None): """ diff --git a/pipt/misc_tools/analysis_tools.py b/pipt/misc_tools/analysis_tools.py index 97b8fe4..37ed43b 100644 --- a/pipt/misc_tools/analysis_tools.py +++ b/pipt/misc_tools/analysis_tools.py @@ -1568,3 +1568,14 @@ def init_local_analysis(init, state): [data_ind[count] for count, val in enumerate(in_region) if val]) return local + + +def get_obs_size(obs_data, time_index, datatypes): + """Return a 2D list of sizes for each observation array.""" + return [ + [ + obs_data[int(time)][data].size if obs_data[int(time)][data] is not None else 0 + for data in datatypes + ] + for time in time_index + ] \ No newline at end of file diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index 8e0742d..64c3182 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -120,14 +120,13 @@ def calc_analysis(self): if 'localanalysis' in self.keys_da: self.local_analysis_update() else: - # Mean pred_data and perturbation matrix with scaling - if len(self.scale_data.shape) == 1: - self.pert_preddata = (self.scale_data ** -1)[:, None] * np.dot(self.aug_pred_data, self.proj) - else: - self.pert_preddata = solve(self.scale_data, np.dot(self.aug_pred_data, self.proj)) - - # Calculate update to get the step (found in update_methods_ns) - self.update() + # Perform the update + self.update( + enX = self.enX, + enY = self.aug_pred_data, + enE = self.real_obs_data, + prior = self.prior_enX + ) # Update the state ensemble and weights if hasattr(self, 'step'): diff --git a/pipt/update_schemes/esmda.py b/pipt/update_schemes/esmda.py index 0ab6fdc..b9a2494 100644 --- a/pipt/update_schemes/esmda.py +++ b/pipt/update_schemes/esmda.py @@ -141,14 +141,13 @@ def calc_analysis(self): if 'localanalysis' in self.keys_da: self.local_analysis_update() else: - if len(self.scale_data.shape) == 1: - self.pert_preddata = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), - np.ones((1, self.ne))) * np.dot(self.aug_pred_data, self.proj) - else: - self.pert_preddata = scilinalg.solve( - self.scale_data, np.dot(self.aug_pred_data, self.proj)) - - self.update() + # Perform the update + self.update( + enX = self.enX, + enY = self.aug_pred_data, + enE = self.real_obs_data, + prior = self.prior_enX + ) # Update the state ensemble and weights if hasattr(self, 'step'): diff --git a/pipt/update_schemes/update_methods_ns/approx_update.py b/pipt/update_schemes/update_methods_ns/approx_update.py index fe74113..d607862 100644 --- a/pipt/update_schemes/update_methods_ns/approx_update.py +++ b/pipt/update_schemes/update_methods_ns/approx_update.py @@ -16,13 +16,28 @@ class approx_update(): https://doi.org/10.1007/s10596-013-9351-5". Note that for a EnKF or ES update, or for update within GN scheme, lambda = 0. """ - def update(self): - # calc the svd of the scaled data pertubation matrix - u_d, s_d, v_d = np.linalg.svd(self.pert_preddata, full_matrices=False) - #aug_state = at.aug_state(self.current_state, self.list_states, self.cell_index) + def update(self, enX, enY, enE, **kwargs): + ''' + Perform the approximate LM update. + + Parameters: + ---------- + enX : np.ndarray + State ensemble matrix (nx, ne) + + enY : np.ndarray + Predicted data ensemble matrix (nd, ne) + + enE : np.ndarray + Ensemble of perturbed observations (nd, ne) + ''' + + # Scale and center the ensemble matrecies + enYcentered = self.scale(np.dot(enY, self.proj), self.scale_data) + + # Perform truncated SVD + u_d, s_d, v_d = np.linalg.svd(enYcentered, full_matrices=False) - # remove the last singular value/vector. This is because numpy returns all ne values, while the last is actually - # zero. This part is a good place to include eventual additional truncation. if self.trunc_energy < 1: ti = (np.cumsum(s_d) / sum(s_d)) <= self.trunc_energy u_d, s_d, v_d = u_d[:, ti].copy(), s_d[ti].copy(), v_d[ti, :].copy() @@ -30,33 +45,94 @@ def update(self): # Check for localization methods if 'localization' in self.keys_da: + # Calculate the localization projection matrix if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - - # Scale data matrix - if len(self.scale_data.shape) == 1: - E_hat = (1/self.scale_data)[:, None] * self.E - else: - E_hat = solve(self.scale_data, self.E) - - x_0 = np.diag(1/s_d) @ u_d.T @ E_hat + enEcentered = self.scale(np.dot(enE, self.proj), self.scale_data) + x_0 = np.diag(1/s_d) @ u_d.T @ enEcentered Lam, z = np.linalg.eig(x_0 @ x_0.T) X = (v_d.T @ z) @ solve( (self.lam + 1)*np.diag(Lam) + np.eye(len(Lam)), (u_d.T @ (np.diag(1/s_d) @ z)).T ) - else: X = v_d.T @ np.diag(s_d) @ solve( (self.lam + 1)*np.eye(len(s_d)) + np.diag(s_d**2), u_d.T) # Check for adaptive localization if 'autoadaloc' in self.localization.loc_info: - self.step = self._update_with_auto_adaptive_localization(X) + + # Scale and center the state ensemble matrix + if ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes'): + enXcentered = self.scale(self.enX - np.mean(self.enX, 1)[:,None], self.state_scaling) + else: + enXcentered = self.scale(np.dot(enX, self.proj), self.state_scaling) + + # Calculate and scale difference between observations and predictions + scaled_delta_data = self.scale(enE - enY, self.scale_data) + + # Compute the update step with auto-adaptive localization + self.step = self.localization.auto_ada_loc( + pert_state = self.state_scaling[:, None]*enXcentered, + proj_pred_data = np.dot(X, scaled_delta_data), + curr_param = self.list_states, + prior_info = self.prior_info + ) + # Check for local analysis elif ('localanalysis' in self.localization.loc_info) and (self.localization.loc_info['localanalysis']): - self.step = self._update_with_local_analysis(X) + + # Calculate weights + if 'distance' in self.localization.loc_info: + weight = _calc_loc( + max_dist = self.localization.loc_info['range'], + distance = self.localization.loc_info['distance'], + prior_info = self.prior_info[self.list_states[0]], + loc_type = self.localization.loc_info['type'], + ne = self.ne + ) + else: # if no distance, do full update + weight = np.ones((enX.shape[0], X.shape[1])) + + # Center ensemble matrix + enXcentered = enX - np.mean(self.enX, axis=1, keepdims=True) + + if (not ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes')): + enXcentered /= np.sqrt(self.ne - 1) + + # Calculate and scale difference between observations and predictions + scaled_delta_data = self.scale(enE - enY, self.scale_data) + + # Compute the update step with local analysis + try: + self.step = weight.multiply(np.dot(enXcentered, X)).dot(scaled_delta_data) + except: + self.step = (weight*(np.dot(enXcentered, X))).dot(scaled_delta_data) + # Check for distance based localization elif ('dist_loc' in self.keys_da['localization'].keys()) or ('dist_loc' in self.keys_da['localization'].values()): - self.step = self._update_with_distance_based_localization(X) + + # Setup localization mask + mask = self.localization.localize( + self.list_datatypes, + [self.keys_da['truedataindex'][int(elem)] for elem in self.assim_index[1]], + self.list_states, + self.ne, + self.prior_info, + at.get_obs_size(self.obs_data, self.assim_index[1], self.list_datatypes) + ) + + # Center ensemble matrix + enXcentered = enX - np.mean(self.enX, axis=1, keepdims=True) + + if not ('emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes'): + enXcentered /= np.sqrt(self.ne - 1) + + # Calculate and scale difference between observations and predictions + scaled_delta_data = self.scale(enE - enY, self.scale_data) + + # Compute the update step with distance-based localization + self.step = mask.multiply(np.dot(enXcentered, X)).dot(scaled_delta_data) + + # Else do parallel update else: @@ -138,68 +214,6 @@ def update(self): - def _update_with_auto_adaptive_localization(self, X): - - # Center ensemble matrix - if ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes'): - pert_state = self.enX - np.mean(self.enX, 1)[:,None] - else: - pert_state = np.dot(self.enX, self.proj) - - # Scale centered ensemble matrix - pert_state = pert_state * (self.state_scaling**(-1))[:, None] - - # Calculate difference between observations and predictions - if len(self.scale_data.shape) == 1: - scaled_delta_data = (self.scale_data ** (-1))[:, None] * (self.real_obs_data - self.aug_pred_data) - else: - scaled_delta_data = solve(self.scale_data, (self.real_obs_data - self.aug_pred_data)) - - # Compute the update step with auto-adaptive localization - step = self.localization.auto_ada_loc( - pert_state = self.state_scaling[:, None]*pert_state, - proj_pred_data = np.dot(X, scaled_delta_data), - curr_param = self.list_states, - prior_info = self.prior_info - ) - - return step - - - def _update_with_local_analysis(self, X): - - # Calculate weights - if 'distance' in self.localization.loc_info: - weight = _calc_loc( - max_dist = self.localization.loc_info['range'], - distance = self.localization.loc_info['distance'], - prior_info = self.prior_info[self.list_states[0]], - loc_type = self.localization.loc_info['type'], - ne = self.ne - ) - else: # if no distance, do full update - weight = np.ones((self.enX.shape[0], X.shape[1])) - - # Center ensemble matrix - mean_state = np.mean(self.enX, axis=1, keepdims=True) - pert_state = self.enX - mean_state - - if not ('emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes'): - pert_state /= np.sqrt(self.ne - 1) - - # Calculate difference between observations and predictions - if self.scale_data.ndim == 1: - scaled_delta_data = (self.scale_data ** -1)[:, None] * (self.real_obs_data - self.aug_pred_data) - else: - scaled_delta_data = solve(self.scale_data, self.real_obs_data - self.aug_pred_data) - - # Compute the update step with local analysis - try: - step = weight.multiply(np.dot(pert_state, X)).dot(scaled_delta_data) - except: - step = (weight*(np.dot(pert_state, X))).dot(scaled_delta_data) - - return step @@ -236,6 +250,19 @@ def _update_with_distance_based_localization(self, X): return step - def _update_with_loclization(self): - pass + def scale(self, data, scaling): + """ + Scale the data perturbations by the data error standard deviation. + Args: + data (np.ndarray): data perturbations + scaling (np.ndarray): data error standard deviation + + Returns: + np.ndarray: scaled data perturbations + """ + + if len(scaling.shape) == 1: + return (scaling ** (-1))[:, None] * data + else: + return solve(scaling, data) diff --git a/pipt/update_schemes/update_methods_ns/full_update.py b/pipt/update_schemes/update_methods_ns/full_update.py index b0cd2e1..709096f 100644 --- a/pipt/update_schemes/update_methods_ns/full_update.py +++ b/pipt/update_schemes/update_methods_ns/full_update.py @@ -18,14 +18,62 @@ class full_update(): no localization is implemented for this method yet. """ + def update(self, enX, enY, enE, **kwargs): + + # Get prior ensemble if provided + priorX = kwargs.get('prior', self.prior_enX) + + if self.Am is None: + self.ext_Am() # do this only once + + # Scale and center the ensemble matrecies + enYcentered = self.scale(np.dot(enY, self.proj), self.scale_data) + enXcentered = self.scale(np.dot(enX, self.proj), self.state_scaling) + + # Perform tuncated SVD + u_d, s_d, v_d = np.linalg.svd(enYcentered, full_matrices=False) + if self.trunc_energy < 1: + ti = (np.cumsum(s_d) / sum(s_d)) <= self.trunc_energy + u_d, s_d, v_d = u_d[:, ti].copy(), s_d[ti].copy(), v_d[ti, :].copy() + + # Compute the update step + x_1 = np.dot(u_d.T, self.scale(enE - enY, self.scale_data)) + x_2 = solve(((self.lam + 1) * np.eye(len(s_d)) + np.diag(s_d ** 2)), x_1) + x_3 = np.dot(np.dot(v_d.T, np.diag(s_d)), x_2) + delta_m1 = np.dot((self.state_scaling[:, None]*enXcentered), x_3) + + x_4 = np.dot(self.Am.T, (self.state_scaling**(-1))[:, None]*(enX - priorX)) + x_5 = np.dot(self.Am, x_4) + x_6 = np.dot(enXcentered.T, x_5) + x_7 = np.dot(v_d.T, solve(((self.lam + 1) * np.eye(len(s_d)) + np.diag(s_d ** 2)), np.dot(v_d, x_6))) + delta_m2 = -np.dot((self.state_scaling[:, None]*enXcentered), x_7) + + self.step = delta_m1 + delta_m2 + + + def scale(self, data, scaling): + """ + Scale the data perturbations by the data error standard deviation. + + Args: + data (np.ndarray): data perturbations + scaling (np.ndarray): data error standard deviation + + Returns: + np.ndarray: scaled data perturbations + """ + + if len(scaling.shape) == 1: + return (scaling ** (-1))[:, None] * data + else: + return solve(scaling, data) + def ext_Am(self, *args, **kwargs): """ The class is initialized by calculating the required Am matrix. """ - delta_scaled_prior = self.state_scaling[:, None] * \ - np.dot(at.aug_state(self.prior_state, self.list_states), self.proj) - + delta_scaled_prior = self.state_scaling[:, None] * np.dot(self.prior_enX, self.proj) u_d, s_d, v_d = np.linalg.svd(delta_scaled_prior, full_matrices=False) # remove the last singular value/vector. This is because numpy returns all ne values, while the last is actually @@ -41,39 +89,3 @@ def ext_Am(self, *args, **kwargs): 1], s_d[:trunc_index + 1], v_d[:trunc_index + 1, :] self.Am = np.dot(u_d, np.eye(trunc_index + 1) * ((s_d ** (-1))[:, None])) # notation from paper - - - def update(self): - - if self.Am is None: - self.ext_Am() # do this only once - - aug_state = at.aug_state(self.current_state, self.list_states) - aug_prior_state = at.aug_state(self.prior_state, self.list_states) - - delta_state = (self.state_scaling**(-1))[:, None]*np.dot(aug_state, self.proj) - - u_d, s_d, v_d = np.linalg.svd(self.pert_preddata, full_matrices=False) - if self.trunc_energy < 1: - ti = (np.cumsum(s_d) / sum(s_d)) <= self.trunc_energy - u_d, s_d, v_d = u_d[:, ti].copy(), s_d[ti].copy(), v_d[ti, :].copy() - - if len(self.scale_data.shape) == 1: - x_1 = np.dot(u_d.T, np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), np.ones((1, self.ne))) * - (self.real_obs_data - self.aug_pred_data)) - else: - x_1 = np.dot(u_d.T, solve(self.scale_data, - (self.real_obs_data - self.aug_pred_data))) - x_2 = solve(((self.lam + 1) * np.eye(len(s_d)) + np.diag(s_d ** 2)), x_1) - x_3 = np.dot(np.dot(v_d.T, np.diag(s_d)), x_2) - delta_m1 = np.dot((self.state_scaling[:, None]*delta_state), x_3) - - x_4 = np.dot(self.Am.T, (self.state_scaling**(-1)) - [:, None]*(aug_state - aug_prior_state)) - x_5 = np.dot(self.Am, x_4) - x_6 = np.dot(delta_state.T, x_5) - x_7 = np.dot(v_d.T, solve( - ((self.lam + 1) * np.eye(len(s_d)) + np.diag(s_d ** 2)), np.dot(v_d, x_6))) - delta_m2 = -np.dot((self.state_scaling[:, None]*delta_state), x_7) - - self.step = delta_m1 + delta_m2 From f4d2907dacac351b1b1310b8e163a4209f47f432 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 5 Nov 2025 09:28:34 +0100 Subject: [PATCH 33/94] Remove old function from analysis_tools.py Remove old extract function for local analysis. The new function lies in extract tools. --- pipt/misc_tools/analysis_tools.py | 90 ------------------------------- 1 file changed, 90 deletions(-) diff --git a/pipt/misc_tools/analysis_tools.py b/pipt/misc_tools/analysis_tools.py index 37ed43b..f22078b 100644 --- a/pipt/misc_tools/analysis_tools.py +++ b/pipt/misc_tools/analysis_tools.py @@ -1480,96 +1480,6 @@ def subsample_state(index, aug_state, pert_state): return new_state -def init_local_analysis(init, state): - """Initialize local analysis. - - Initialize the local analysis by reading the input variables, defining the parameter classes and search ranges. Build - the map of data/parameter positions. - - Args - ---- - init : dictionary containing the parsed information form the input file. - state : list of states that will be updated - - Returns - ------- - local : dictionary of initialized values. - """ - - local = {} - local['cell_parameter'] = [] - local['region_parameter'] = [] - local['vector_region_parameter'] = [] - local['unique'] = True - - for i, opt in enumerate(list(zip(*init))[0]): - if opt.lower() == 'region_parameter': # define scalar parameters valid in a region - local['region_parameter'] = [ - elem for elem in init[i][1].split(' ') if elem in state] - if opt.lower() == 'vector_region_parameter': # Sometimes it useful to define the same parameter for multiple - # regions as a vector. - local['vector_region_parameter'] = [ - elem for elem in init[i][1].split(' ') if elem in state] - if opt.lower() == 'cell_parameter': # define cell specific vector parameters - local['cell_parameter'] = [ - elem for elem in init[i][1].split(' ') if elem in state] - if opt.lower() == 'search_range': - local['search_range'] = int(init[i][1]) - if opt.lower() == 'column_update': - local['column_update'] = [elem for elem in init[i][1].split(',')] - if opt.lower() == 'parameter_position_file': # assume pickled format - with open(init[i][1], 'rb') as file: - local['parameter_position'] = pickle.load(file) - if opt.lower() == 'data_position_file': # assume pickled format - with open(init[i][1], 'rb') as file: - local['data_position'] = pickle.load(file) - if opt.lower() == 'update_mask_file': - with open(init[i][1], 'rb') as file: - local['update_mask'] = pickle.load(file) - - if 'update_mask' in local: - return local - else: - assert 'parameter_position' in local, 'A pickle file containing the binary map of the parameters is MANDATORY' - assert 'data_position' in local, 'A pickle file containing the position of the data is MANDATORY' - - data_name = [elem for elem in local['data_position'].keys()] - if type(local['data_position'][data_name[0]][0]) == list: # assim index has spesific position - local['unique'] = False - data_pos = [elem for data in data_name for assim_elem in local['data_position'][data] - for elem in assim_elem] - data_ind = [f'{data}_{assim_indx}' for data in data_name for assim_indx, assim_elem in enumerate(local['data_position'][data]) - for _ in assim_elem] - else: - data_pos = [elem for data in data_name for elem in local['data_position'][data]] - # store the name for easy index - data_ind = [data for data in data_name for _ in local['data_position'][data]] - kde_search = cKDTree(data=data_pos) - - local['update_mask'] = {} - for param in local['cell_parameter']: # find data in a distance from the parameter - field_size = local['parameter_position'][param].shape - local['update_mask'][param] = [[[[] for _ in range(field_size[2])] for _ in range(field_size[1])] for _ - in range(field_size[0])] - for k in range(field_size[0]): - for j in range(field_size[1]): - new_iter = [elem for elem, val in enumerate( - local['parameter_position'][param][k, j, :]) if val] - if len(new_iter): - for i in new_iter: - local['update_mask'][param][k][j][i] = set( - [data_ind[elem] for elem in kde_search.query_ball_point(x=(k, j, i), - r=local['search_range'], workers=-1)]) - - # see if data is inside the region. Note parameter_position is boolean map - for param in local['region_parameter']: - in_region = [local['parameter_position'][param][elem] for elem in data_pos] - local['update_mask'][param] = set( - [data_ind[count] for count, val in enumerate(in_region) if val]) - - return local - - def get_obs_size(obs_data, time_index, datatypes): """Return a 2D list of sizes for each observation array.""" return [ From 162414ace5b22949a6356a7057ceb80665038ac8 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 5 Nov 2025 10:22:05 +0100 Subject: [PATCH 34/94] Redefines state related variables - self.stateX: Current state vector - self.stateF: Function value(s) of current state - self.bounds: Bounds for each variable in stateX - self.varX: Variance for state vector - self.covX: Covariance matrix for state vector - self.enX: Ensemble of state vectors (nx, ne) - self.enF: Ensemble of function values (ne, ) --- popt/loop/ensemble_base.py | 73 ++++++++++++++++++-------------------- 1 file changed, 34 insertions(+), 39 deletions(-) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index e5363a2..c7c5b44 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -46,53 +46,48 @@ def __init__(self, options, simulator, objective): # Set objective function (callable) self.obj_func = objective + + # Initialize state-related variables self.state_func_values = None self.ens_func_values = None - # Initialize prior - self._initialize_state_info() # Initialize cov, bounds, and state - self._scale_state() # Scale self.state to [0, 1] if transform is True - - def _initialize_state_info(self): - ''' - Initialize covariance and bounds based on prior information. - ''' - self.cov = np.array([]) - self.lb = [] - self.ub = [] - self.bounds = [] + self.stateX = np.array([]) # Current state vector + self.stateF = None # Function value(s) of current state + self.bounds = [] # Bounds for each variable in stateX + self.varX = np.array([]) # Variance for state vector + self.covX = None # Covariance matrix for state vector + self.enX = None # Ensemble of state vectors (nx, ne) + self.enF = None # Ensemble of function values (ne, ) + # Intialize state information for key in self.prior_info.keys(): - variable = self.prior_info[key] - - # mean - self.state[key] = np.asarray(variable['mean']) - - # Covariance - dim = self.state[key].size - var = variable['variance']*np.ones(dim) - - if 'limits' in variable.keys(): - lb, ub = variable['limits'] - self.lb.append(lb) - self.ub.append(ub) - - # transform var to [0, 1] if transform is True - if self.transform: - var = var/(ub - lb)**2 - var = np.clip(var, 0, 1, out=var) - self.bounds += dim*[(0, 1)] - else: - self.bounds += dim*[(lb, ub)] + + # Extract prior information for this variable + mean = np.asarray(self.prior_info[key]['mean']) + var = self.prior_info[key]['variance']*np.ones(mean.size) + lb, ub = self.prior_info[key].get('limits', (None, None)) + + # Fill in state vector and index information + self.stateX = np.append(self.stateX, mean) + self.idX[key] = (self.stateX.size - mean.size, self.stateX.size) + + # Set bounds and transform variance if applicable + if self.transform and (lb is not None) and (ub is not None): + var = var/(ub - lb)**2 + var = np.clip(var, 0, 1, out=var) + self.bounds += mean.size*[(0, 1)] else: - self.bounds += dim*[(None, None)] + self.bounds.append((lb, ub)) - # Add to covariance - self.cov = np.append(self.cov, var) - self.dim = self.cov.shape[0] + # Fill in variance vector + self.varX = np.append(self.varX, var) - # Make cov full covariance matrix - self.cov = np.diag(self.cov) + self.covX = np.diag(self.varX) # Covariance matrix + self.dimX = self.stateX.size # Dimension of state vector + + # Scale state if applicable + self._scale_state() # Scale self.state to [0, 1] if transform is True + def get_state(self): """ From 419584416f7e1c08a21d1bf17a427aa4defab826 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 5 Nov 2025 10:32:56 +0100 Subject: [PATCH 35/94] Redefine the state scaler functions. A redefinition of the state scaler functions is necessary to accommodate changes in the state vector and ensemble structure. --- popt/loop/ensemble_base.py | 49 +++++++++++++++++++++++++++----------- 1 file changed, 35 insertions(+), 14 deletions(-) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index c7c5b44..a3fc4c4 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -87,7 +87,7 @@ def __init__(self, options, simulator, objective): # Scale state if applicable self._scale_state() # Scale self.state to [0, 1] if transform is True - + def get_state(self): """ @@ -190,21 +190,42 @@ def _aux_input(self): sys.exit(0) return nr - def _scale_state(self): + def scale_state(self, x): """ Transform the internal state from [lb, ub] to [0, 1] - """ - if self.transform and (self.lb and self.ub): - for i, key in enumerate(self.state): - self.state[key] = (self.state[key] - self.lb[i])/(self.ub[i] - self.lb[i]) - np.clip(self.state[key], 0, 1, out=self.state[key]) - def _invert_scale_state(self): + Parameters + ---------- + x : array_like + The input state + + Returns + ------- + x : array_like + The scaled state + """ + if self.bounds: + lb = np.array(self.bounds)[:, 0] + ub = np.array(self.bounds)[:, 1] + x = (x - lb) / (ub - lb) + return x + + def invert_scale_state(self, u): """ Transform the internal state from [0, 1] to [lb, ub] - """ - if self.transform and (self.lb and self.ub): - for i, key in enumerate(self.state): - if self.transform: - self.state[key] = self.lb[i] + self.state[key]*(self.ub[i] - self.lb[i]) - np.clip(self.state[key], self.lb[i], self.ub[i], out=self.state[key]) \ No newline at end of file + + Parameters + ---------- + u : array_like + The scaled state + + Returns + ------- + x : array_like + The unscaled state + """ + if self.bounds: + lb = np.array(self.bounds)[:, 0] + ub = np.array(self.bounds)[:, 1] + u = lb + u * (ub - lb) + return u \ No newline at end of file From 042e90aa2eab4f462e5848a74d45bb1914b68d5c Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 5 Nov 2025 10:51:56 +0100 Subject: [PATCH 36/94] Correct logical bug for state scalers --- popt/loop/ensemble_base.py | 53 ++++++++++++++++++++++++++++---------- 1 file changed, 39 insertions(+), 14 deletions(-) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index a3fc4c4..5a64d15 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -58,6 +58,8 @@ def __init__(self, options, simulator, objective): self.covX = None # Covariance matrix for state vector self.enX = None # Ensemble of state vectors (nx, ne) self.enF = None # Ensemble of function values (ne, ) + self.lb = np.array([]) # Lower bounds for state vector + self.ub = np.array([]) # Upper bounds for state vector # Intialize state information for key in self.prior_info.keys(): @@ -79,14 +81,21 @@ def __init__(self, options, simulator, objective): else: self.bounds.append((lb, ub)) + # Fill in lb and ub vectors + self.lb = np.append(self.lb, lb*np.ones(mean.size)) + self.ub = np.append(self.ub, ub*np.ones(mean.size)) + # Fill in variance vector self.varX = np.append(self.varX, var) - + self.covX = np.diag(self.varX) # Covariance matrix self.dimX = self.stateX.size # Dimension of state vector - + # Scale state if applicable - self._scale_state() # Scale self.state to [0, 1] if transform is True + self.stateX = self.scale_state(self.stateX) + + print(self.stateX) + print(self.state) def get_state(self): @@ -204,12 +213,20 @@ def scale_state(self, x): x : array_like The scaled state """ - if self.bounds: - lb = np.array(self.bounds)[:, 0] - ub = np.array(self.bounds)[:, 1] - x = (x - lb) / (ub - lb) - return x - + x = np.asarray(x) + scaled_x = np.zeros_like(x) + + if self.transform is False: + return x + + for i in range(len(x)): + if (self.lb[i] is not None) and (self.ub[i] is not None): + scaled_x[i] = (x[i] - self.lb[i]) / (self.ub[i] - self.lb[i]) + else: + scaled_x[i] = x[i] # No scaling if bounds are None + + return scaled_x + def invert_scale_state(self, u): """ Transform the internal state from [0, 1] to [lb, ub] @@ -224,8 +241,16 @@ def invert_scale_state(self, u): x : array_like The unscaled state """ - if self.bounds: - lb = np.array(self.bounds)[:, 0] - ub = np.array(self.bounds)[:, 1] - u = lb + u * (ub - lb) - return u \ No newline at end of file + u = np.asarray(u) + x = np.zeros_like(u) + + if self.transform is False: + return u + + for i in range(len(u)): + if (self.lb[i] is not None) and (self.ub[i] is not None): + x[i] = self.lb[i] + u[i] * (self.ub[i] - self.lb[i]) + else: + x[i] = u[i] # No scaling if bounds are None + + return x \ No newline at end of file From ff151a4e0ee67e7d248a887119db0a2473e56e34 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 5 Nov 2025 11:12:31 +0100 Subject: [PATCH 37/94] Update get_state() and get_cov() --- popt/loop/ensemble_base.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index 5a64d15..f9ab1e1 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -51,13 +51,13 @@ def __init__(self, options, simulator, objective): self.state_func_values = None self.ens_func_values = None - self.stateX = np.array([]) # Current state vector + self.stateX = np.array([]) # Current state vector, (nx,) self.stateF = None # Function value(s) of current state self.bounds = [] # Bounds for each variable in stateX self.varX = np.array([]) # Variance for state vector self.covX = None # Covariance matrix for state vector - self.enX = None # Ensemble of state vectors (nx, ne) - self.enF = None # Ensemble of function values (ne, ) + self.enX = None # Ensemble of state vectors ,(nx, ne) + self.enF = None # Ensemble of function values, (ne, ) self.lb = np.array([]) # Lower bounds for state vector self.ub = np.array([]) # Upper bounds for state vector @@ -94,9 +94,6 @@ def __init__(self, options, simulator, objective): # Scale state if applicable self.stateX = self.scale_state(self.stateX) - print(self.stateX) - print(self.state) - def get_state(self): """ @@ -105,7 +102,7 @@ def get_state(self): x : numpy.ndarray Control vector as ndarray, shape (number of controls, number of perturbations) """ - return ot.aug_optim_state(self.state, list(self.state.keys())) + return self.stateX def get_cov(self): """ @@ -114,7 +111,7 @@ def get_cov(self): cov : numpy.ndarray Covariance matrix, shape (number of controls, number of controls) """ - return self.cov + return self.covX def vec_to_state(self, x): """ From 3fe6b5f908abba86c2e6ca1934b57978361c0dc5 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 5 Nov 2025 12:59:25 +0100 Subject: [PATCH 38/94] Rewrite self.function to accommodate state changes --- popt/loop/ensemble_base.py | 32 +++++++++++++++++++------------- popt/loop/ensemble_gaussian.py | 4 ---- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index f9ab1e1..8d88f58 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -47,10 +47,7 @@ def __init__(self, options, simulator, objective): # Set objective function (callable) self.obj_func = objective - # Initialize state-related variables - self.state_func_values = None - self.ens_func_values = None - + # Initialize state-related attributes self.stateX = np.array([]) # Current state vector, (nx,) self.stateF = None # Function value(s) of current state self.bounds = [] # Bounds for each variable in stateX @@ -146,18 +143,25 @@ def function(self, x, *args, **kwargs): self._aux_input() # check for ensmble - if len(x.shape) == 1: self.ne = self.num_models + if len(x.shape) == 1: + x = x[:,np.newaxis] + self.ne = self.num_models else: self.ne = x.shape[1] + # Run simulation + x = self.invert_scale_state(x) + run_success = self.calc_prediction(enX=x, save_prediction=self.save_prediction) + x = self.scale_state(x).flatten() + # convert x (nparray) to state (dict) - self.state = self.vec_to_state(x) + #self.state = self.vec_to_state(x) # run the simulation - self._invert_scale_state() # ensure that state is in [lb,ub] - self._set_multilevel_state(self.state, x) # set multilevel state if applicable - run_success = self.calc_prediction(save_prediction=self.save_prediction) # calculate flow data - self._set_multilevel_state(self.state, x) # For some reason this has to be done again after calc_prediction - self._scale_state() # scale back to [0, 1] + #self._invert_scale_state() # ensure that state is in [lb,ub] + #self._set_multilevel_state(self.state, x) # set multilevel state if applicable + #run_success = self.calc_prediction(save_prediction=self.save_prediction) # calculate flow data + #self._set_multilevel_state(self.state, x) # For some reason this has to be done again after calc_prediction + #self._scale_state() # scale back to [0, 1] # Evaluate the objective function if run_success: @@ -170,8 +174,10 @@ def function(self, x, *args, **kwargs): else: func_values = np.inf # the simulations have crashed - if len(x.shape) == 1: self.state_func_values = func_values - else: self.ens_func_values = func_values + if len(x.shape) == 1: + self.stateF = func_values + else: + self.enF = func_values return func_values diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index 2d334ca..ef5e4a9 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -65,10 +65,6 @@ def __init__(self, options, simulator, objective): # Initialize PETEnsemble super().__init__(options, simulator, objective) - # Objective function values - self.state_func_values = None - self.ens_func_values = None - # Inflation factor used in SmcOpt self.inflation_factor = None self.survival_factor = None From 450b3251ab1094ca09068265287139f676ce9b55 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 6 Nov 2025 13:37:30 +0100 Subject: [PATCH 39/94] Rewrite ensemble.gradient to state changes Rewrite the ensemble-based gradient function to accommodate recent changes to state and ensemble representations. Also improved readability of the function. --- popt/loop/ensemble_base.py | 6 +- popt/loop/ensemble_gaussian.py | 142 ++++++++++++++------------------- 2 files changed, 64 insertions(+), 84 deletions(-) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index 8d88f58..5f9e721 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -55,8 +55,8 @@ def __init__(self, options, simulator, objective): self.covX = None # Covariance matrix for state vector self.enX = None # Ensemble of state vectors ,(nx, ne) self.enF = None # Ensemble of function values, (ne, ) - self.lb = np.array([]) # Lower bounds for state vector - self.ub = np.array([]) # Upper bounds for state vector + self.lb = np.array([]) # Lower bounds for state vector, (nx,) + self.ub = np.array([]) # Upper bounds for state vector, (nx,) # Intialize state information for key in self.prior_info.keys(): @@ -151,7 +151,7 @@ def function(self, x, *args, **kwargs): # Run simulation x = self.invert_scale_state(x) run_success = self.calc_prediction(enX=x, save_prediction=self.save_prediction) - x = self.scale_state(x).flatten() + x = self.scale_state(x).squeeze() # convert x (nparray) to state (dict) #self.state = self.vec_to_state(x) diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index ef5e4a9..1e90544 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -101,102 +101,82 @@ def get_final_state(self, return_dict=False): else: x = self.get_state() return x - + def gradient(self, x, *args, **kwargs): - r""" - Calculate the preconditioned gradient associated with ensemble, defined as: - - $$ S \approx C_x \times G^T $$ - - where $C_x$ is the state covariance matrix, and $G$ is the standard - gradient. The ensemble sensitivity matrix is calculated as: - - $$ S = X \times J^T /(N_e-1) $$ - - where $X$ and $J$ are ensemble matrices of $x$ (or control variables) and objective function - perturbed by their respective means. In practice (and in this method), $S$ is calculated by perturbing the - current control variable with Gaussian random numbers from $N(0, C_x)$ (giving $X$), running - the generated ensemble ($X$) through the simulator to give an ensemble of objective function values - ($J$), and in the end calculate $S$. Note that $S$ is an $N_x \times 1$ vector, where - $N_x$ is length of the control vector and the objective function is scalar. - - Note: In the case of multi-fidelity optimization, it is possible to specify 0 members for some of the levels - in order to skip these levels. In that case, cov_wgt should have the same length as the number of levels - that is acutally used. + ''' + Ensemble-based Gradient (EnOpt) Parameters ---------- x : ndarray Control vector, shape (number of controls, ) - + args : tuple Covarice ($C_x$), shape (number of controls, number of controls) - + Returns ------- - gradient : numpy.ndarray - The gradient evaluated at x, shape (number of controls, ) - """ + gradient : ndarray + Ensemble gradient, shape (number of controls, ) + ''' + # Update state vector + self.stateX = x - # Set the ensemble state equal to the input control vector x - self.state = ot.update_optim_state(x, self.state, list(self.state.keys())) + # Set covariance equal to the input + self.covX = args[0] - # Set the covariance equal to the input - self.cov = args[0] - - # If bias correction is used we need to temporarily store the initial state - initial_state = None - if self.bias_file is not None and self.bias_factors is None: # first iteration - initial_state = deepcopy(self.state) # store this to update current objective values - - # Generate ensemble of states + # Generate state ensemble self.ne = self.num_samples - nr = self._aux_input() - self.state = self._gen_state_ensemble() - - state_ens = at.aug_state(self.state, list(self.state.keys())) - self.function(state_ens, **kwargs) - - # If bias correction is used we need to calculate the bias factors, J(u_j,m_j)/J(u_j,m) - if self.bias_file is not None: # use bias corrections - self._bias_factors(self.ens_func_values, initial_state) - - # Perturb state and function values with their mean - state_ens = at.aug_state(self.state, list(self.state.keys())) - pert_state = state_ens - np.dot(state_ens.mean(1)[:, None], np.ones((1, self.ne))) - - if not isinstance(self.ens_func_values,list): - self.ens_func_values = [self.ens_func_values] - start_index = 0 - level_gradient = [] - gradient = np.zeros(state_ens.shape[0]) - L = len(self.ens_func_values) - for l in range(L): - - if self.bias_file is not None: # use bias corrections - self.ens_func_values[l] *= self._bias_correction(self.state) - pert_obj_func = self.ens_func_values[l] - np.mean(self.ens_func_values[l]) - else: - pert_obj_func = self.ens_func_values[l] - np.array(np.repeat(self.state_func_values, nr)) - - # Calculate the gradient - ml_ne = self.ens_func_values[l].size - g_m = np.zeros(state_ens.shape[0]) - for i in np.arange(ml_ne): - g_m = g_m + pert_obj_func[i] * pert_state[:, start_index + i] - - start_index += ml_ne - level_gradient.append(g_m / (ml_ne - 1)) - - if 'multilevel' in self.keys_en.keys(): - cov_wgt = ot.get_list_element(self.keys_en['multilevel'], 'cov_wgt') - for l in range(L): - gradient += level_gradient[l]*cov_wgt[l] - gradient /= self.ne + nr = self._aux_input() + self.enX = np.random.multivariate_normal(self.stateX, self.covX, self.ne).T + + # Shift ensemble to have correct mean + self.enX = self.enX - self.enX.mean(axis=1, keepdims=True) + self.stateX[:,None] + + # Truncate to bounds + if (self.lb is not None) and (self.ub is not None): + self.enX = np.clip(self.enX, self.lb[:, None], self.ub[:, None]) + + # Evaluate objective function for ensemble + self.enF = self.function(self.enX, *args, **kwargs) + + # Make function ensemble to a list (for Multilevel) + if not isinstance(self.enF, list): + self.enF = [self.enF] + + # Define some variables for gradient calculation + index = 0 + nlevels = len(self.enF) + grad_ml = np.zeros((nlevels, self.dimX)) + + # Loop over levels (only one level if not multilevel) + for id_level in range(nlevels): + dF = self.enF[id_level] - np.repeat(self.stateF, nr) + ne = self.enF[id_level].shape[0] + + # Calculate ensemble gradient for level + g = np.zeros(self.dimX) + for n in range(ne): + g = g + dF[n] * (self.enX[:, index+n] - self.stateX) + + grad_ml[id_level] = g/ne + index += ne + + if 'multilevel' in self.keys_en: + weight = ot.get_list_element(self.keys_en['multilevel'], 'cov_wgt') + weight = np.array(weight) + if not np.sum(weight) == 1.0: + weight = weight / np.sum(weight) + grad = np.dot(grad_ml, weight) else: - gradient = level_gradient[0] + grad = grad_ml[0] + + # Check if natural or averaged gradient (default is natural) + if not self.keys_en.get('natural_gradient', True): + cov_inv = np.linalg.inv(self.covX) + grad = np.matmul(cov_inv, grad) - return gradient + return grad def hessian(self, x=None, *args): r""" From 2f6521aa16bda4cb6c4e7c818bef81a900bd0156 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 6 Nov 2025 14:26:13 +0100 Subject: [PATCH 40/94] Rewrite ensemble.hessian Rewrite ensemble Hessian function to accommodate changes in gradient. --- popt/loop/ensemble_gaussian.py | 94 ++++++++++++++++++---------------- 1 file changed, 50 insertions(+), 44 deletions(-) diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index 1e90544..c7c6365 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -104,7 +104,7 @@ def get_final_state(self, return_dict=False): def gradient(self, x, *args, **kwargs): ''' - Ensemble-based Gradient (EnOpt) + Ensemble-based Gradient (EnOpt). Parameters ---------- @@ -178,66 +178,72 @@ def gradient(self, x, *args, **kwargs): return grad - def hessian(self, x=None, *args): - r""" - Calculate the hessian matrix associated with ensemble, defined as: - - $$ H = J(XX^T - \Sigma)/ (N_e-1) $$ - - where $X$ and $J$ are ensemble matrices of $x$ (or control variables) and objective function - perturbed by their respective means. - - !!! note - state and ens_func_values are assumed to already exist from computation of the gradient. - Save time by not running them again. + def hessian(self, x=None, *args, **kwargs): + ''' + Ensemble-based Hessian. Parameters ---------- x : ndarray - Control vector, shape (number of controls, number of perturbations) + Control vector, shape (number of controls, ). If None, use the last x used in gradient. + If x is not None and it does not match the last x used in gradient, recompute the gradient first. + args : tuple + Additional arguments passed to function + Returns ------- - hessian: numpy.ndarray - The hessian evaluated at x, shape (number of controls, number of controls) - + hessian : ndarray + Ensemble hessian, shape (number of controls, number of controls) + References ---------- Zhang, Y., Stordal, A.S. & Lorentzen, R.J. A natural Hessian approximation for ensemble based optimization. Comput Geosci 27, 355–364 (2023). https://doi.org/10.1007/s10596-022-10185-z - """ + ''' + # Check if self.gradient has been called with this x + if (not np.array_equal(x, self.stateX)) and (x is not None): + self.gradient(x, *args, **kwargs) - # Perturb state and function values with their mean - state_ens = at.aug_state(self.state, list(self.state.keys())) - pert_state = state_ens - np.dot(state_ens.mean(1)[:, None], np.ones((1, self.ne))) nr = self._aux_input() - if not isinstance(self.ens_func_values,list): - self.ens_func_values = [self.ens_func_values] - start_index = 0 - level_hessian = [] - L = len(self.ens_func_values) - hessian = np.zeros(self.cov.shape) - for l in range(L): - pert_obj_func = self.ens_func_values[l] - np.array(np.repeat(self.state_func_values, nr)) - ml_ne = self.ens_func_values[l].size - - # Calculate the gradient for mean and covariance matrix - g_c = np.zeros(self.cov.shape) - for i in np.arange(ml_ne): - g_c = g_c + pert_obj_func[i] * (np.outer(pert_state[:, start_index + i], pert_state[:, start_index + i]) - self.cov) + # Make function ensemble to a list (for Multilevel) + if not isinstance(self.enF, list): + self.enF = [self.enF] + + # Define some variables for gradient calculation + index = 0 + nlevels = len(self.enF) + hess_ml = np.zeros((nlevels, self.dimX, self.dimX)) - start_index += ml_ne - level_hessian.append(g_c / (ml_ne - 1)) + # Loop over levels (only one level if not multilevel) + for id_level in range(nlevels): + dF = self.enF[id_level] - np.repeat(self.stateF, nr) + ne = self.enF[id_level].shape[0] - if 'multilevel' in self.keys_en.keys(): - cov_wgt = ot.get_list_element(self.keys_en['multilevel'], 'cov_wgt') - for l in range(L): - hessian += level_hessian[l]*cov_wgt[l] - hessian /= self.ne + # Calculate ensemble Hessian for level + h = np.zeros((self.dimX, self.dimX)) + for n in range(ne): + dx = (self.enX[:, index+n] - self.stateX) + h = h + dF[n] * (np.outer(dx, dx) - self.covX) + + hess_ml[id_level] = h/ne + index += ne + + if 'multilevel' in self.keys_en: + weight = ot.get_list_element(self.keys_en['multilevel'], 'cov_wgt') + weight = np.array(weight) + if not np.sum(weight) == 1.0: + weight = weight / np.sum(weight) + hessian = np.sum([h*w for h, w in zip(hess_ml, weight)], axis=0) else: - hessian = level_hessian[0] - + hessian = hess_ml[0] + + # Check if natural or averaged Hessian (default is natural) + if not self.keys_en.get('natural_gradient', True): + cov_inv = np.linalg.inv(self.covX) + hessian = cov_inv @ hessian @ cov_inv + return hessian def calc_ensemble_weights(self, x, *args, **kwargs): From 78f9fb523ce49643b7205470727bdba6a5963f5a Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 6 Nov 2025 14:58:52 +0100 Subject: [PATCH 41/94] Add save_stateX function --- popt/loop/ensemble_base.py | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index 5f9e721..1e33a67 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -1,5 +1,6 @@ # External imports import numpy as np +import pandas as pd import sys import warnings @@ -10,6 +11,7 @@ from pipt.misc_tools import analysis_tools as at from ensemble.ensemble import Ensemble as SupEnsemble from simulator.simple_models import noSimulation +from pipt.misc_tools.ensemble_tools import matrix_to_dict __all__ = ['EnsembleOptimizationBaseClass'] @@ -256,4 +258,31 @@ def invert_scale_state(self, u): else: x[i] = u[i] # No scaling if bounds are None - return x \ No newline at end of file + return x + + def save_stateX(self, path='./', filetype='npz'): + ''' + Save the state vector. + + Parameters + ---------- + path : str + Path to save the state vector. Default is current directory. + + filetype : str + File type to save the state vector. Options are 'csv', 'npz' or 'npy'. Default is 'npz'. + ''' + if self.transform: + stateX = self.invert_scale_state(self.stateX) + else: + stateX = self.stateX + + if filetype == 'csv': + state_dict = matrix_to_dict(stateX, self.idX) + state_df = pd.DataFrame(data=state_dict) + state_df.to_csv(path + 'stateX.csv', index=False) + elif filetype == 'npz': + state_dict = matrix_to_dict(stateX, self.idX) + np.savez_compressed(path + 'stateX.npz', **state_dict) + elif filetype == 'npy': + np.save(path + 'stateX.npy', stateX) \ No newline at end of file From 1d3f7744d02407d8727299327c1d9e54ef537508 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 6 Nov 2025 14:59:47 +0100 Subject: [PATCH 42/94] remove save_final_state function --- popt/loop/ensemble_gaussian.py | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index c7c6365..0e7fa94 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -82,26 +82,6 @@ def __init__(self, options, simulator, objective): self.bias_weights = np.ones(self.num_samples) / self.num_samples # initialize with equal weights self.bias_points = None # this is the points used to estimate the bias correction - def get_final_state(self, return_dict=False): - """ - Parameters - ---------- - return_dict : bool - Retrun dictionary if true - - Returns - ------- - x : numpy.ndarray - Control vector as ndarray, shape (number of controls, number of perturbations) - """ - - self._invert_scale_state() - if return_dict: - x = self.state - else: - x = self.get_state() - return x - def gradient(self, x, *args, **kwargs): ''' Ensemble-based Gradient (EnOpt). From 1753335e3b185801cf1e6c4f1ef5e8f735b283e1 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 6 Nov 2025 15:15:48 +0100 Subject: [PATCH 43/94] Remove vec_to_state function --- popt/loop/ensemble_base.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index 1e33a67..8aec913 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -111,12 +111,6 @@ def get_cov(self): Covariance matrix, shape (number of controls, number of controls) """ return self.covX - - def vec_to_state(self, x): - """ - Converts a control vector to the internal state representation. - """ - return ot.update_optim_state(x, self.state, list(self.state.keys())) def get_bounds(self): """ From f8efda8e6e8935129507c72132054578c7b1f9c5 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 7 Nov 2025 09:11:39 +0100 Subject: [PATCH 44/94] Change order of some functions in ensemble --- popt/loop/ensemble_base.py | 99 +++++++++++++++++++------------------- 1 file changed, 49 insertions(+), 50 deletions(-) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index 8aec913..0e4ce98 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -52,14 +52,14 @@ def __init__(self, options, simulator, objective): # Initialize state-related attributes self.stateX = np.array([]) # Current state vector, (nx,) self.stateF = None # Function value(s) of current state - self.bounds = [] # Bounds for each variable in stateX + self.bounds = [] # Bounds (untransformed) for each variable in stateX self.varX = np.array([]) # Variance for state vector self.covX = None # Covariance matrix for state vector self.enX = None # Ensemble of state vectors ,(nx, ne) self.enF = None # Ensemble of function values, (ne, ) - self.lb = np.array([]) # Lower bounds for state vector, (nx,) - self.ub = np.array([]) # Upper bounds for state vector, (nx,) - + self.lb = np.array([]) # Lower bounds (transformed) for state vector, (nx,) + self.ub = np.array([]) # Upper bounds (transformed) for state vector, (nx,) + # Intialize state information for key in self.prior_info.keys(): @@ -93,35 +93,6 @@ def __init__(self, options, simulator, objective): # Scale state if applicable self.stateX = self.scale_state(self.stateX) - - def get_state(self): - """ - Returns - ------- - x : numpy.ndarray - Control vector as ndarray, shape (number of controls, number of perturbations) - """ - return self.stateX - - def get_cov(self): - """ - Returns - ------- - cov : numpy.ndarray - Covariance matrix, shape (number of controls, number of controls) - """ - return self.covX - - def get_bounds(self): - """ - Returns - ------- - bounds : list - (min, max) pairs for each element in x. None is used to specify no bound. - """ - - return self.bounds - def function(self, x, *args, **kwargs): """ This is the main function called during optimization. @@ -176,27 +147,34 @@ def function(self, x, *args, **kwargs): self.enF = func_values return func_values - - def _set_multilevel_state(self, state, x): - if 'multilevel' in self.keys_en.keys() and len(x.shape) > 1: - en_size = ot.get_list_element(self.keys_en['multilevel'], 'en_size') - self.state = ot.toggle_ml_state(self.state, en_size) + def get_state(self): + """ + Returns + ------- + x : numpy.ndarray + Control vector as ndarray, shape (number of controls, number of perturbations) + """ + return self.stateX + + def get_cov(self): + """ + Returns + ------- + cov : numpy.ndarray + Covariance matrix, shape (number of controls, number of controls) + """ + return self.covX - def _aux_input(self): + def get_bounds(self): """ - Set the auxiliary input used for multiple geological realizations + Returns + ------- + bounds : list + (min, max) pairs for each element in x. None is used to specify no bound. """ - nr = 1 # nr is the ratio of samples over models - if self.num_models > 1: - if np.remainder(self.num_samples, self.num_models) == 0: - nr = int(self.num_samples / self.num_models) - self.aux_input = list(np.repeat(np.arange(self.num_models), nr)) - else: - print('num_samples must be a multiplum of num_models!') - sys.exit(0) - return nr + return self.bounds def scale_state(self, x): """ @@ -279,4 +257,25 @@ def save_stateX(self, path='./', filetype='npz'): state_dict = matrix_to_dict(stateX, self.idX) np.savez_compressed(path + 'stateX.npz', **state_dict) elif filetype == 'npy': - np.save(path + 'stateX.npy', stateX) \ No newline at end of file + np.save(path + 'stateX.npy', stateX) + + def _set_multilevel_state(self, state, x): + if 'multilevel' in self.keys_en.keys() and len(x.shape) > 1: + en_size = ot.get_list_element(self.keys_en['multilevel'], 'en_size') + self.state = ot.toggle_ml_state(self.state, en_size) + + + def _aux_input(self): + """ + Set the auxiliary input used for multiple geological realizations + """ + + nr = 1 # nr is the ratio of samples over models + if self.num_models > 1: + if np.remainder(self.num_samples, self.num_models) == 0: + nr = int(self.num_samples / self.num_models) + self.aux_input = list(np.repeat(np.arange(self.num_models), nr)) + else: + print('num_samples must be a multiplum of num_models!') + sys.exit(0) + return nr \ No newline at end of file From 4dd5582db80e81d5464334c95cbe495b5b2cd13a Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 7 Nov 2025 09:21:07 +0100 Subject: [PATCH 45/94] Remove initialization of bias variables --- popt/loop/ensemble_gaussian.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index 0e7fa94..11550d8 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -71,16 +71,6 @@ def __init__(self, options, simulator, objective): self.particles = [] # list in case of multilevel self.particle_values = [] # list in case of multilevel self.resample_index = None - - # Initialize variables for bias correction - if 'bias_file' in self.sim.input_dict: # use bias correction - self.bias_file = self.sim.input_dict['bias_file'].upper() # mako file for simulations - else: - self.bias_file = None - self.bias_adaptive = None # flag to adaptively update the bias correction (not implemented yet) - self.bias_factors = None # this is J(x_j,m_j)/J(x_j,m) - self.bias_weights = np.ones(self.num_samples) / self.num_samples # initialize with equal weights - self.bias_points = None # this is the points used to estimate the bias correction def gradient(self, x, *args, **kwargs): ''' From f0a47b61ea581d3854c6b520d15dacd8334d3ed9 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 7 Nov 2025 10:54:24 +0100 Subject: [PATCH 46/94] Rewrite calc_ensemble_weights to accomodate changes --- popt/loop/ensemble_gaussian.py | 62 +++++++++++++++++----------------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index 11550d8..0c185a4 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -219,6 +219,7 @@ def hessian(self, x=None, *args, **kwargs): def calc_ensemble_weights(self, x, *args, **kwargs): r""" Calculate weights used in sequential monte carlo optimization. + Updated version that accommodates new base class changes. Parameters ---------- @@ -233,54 +234,53 @@ def calc_ensemble_weights(self, x, *args, **kwargs): sens_matrix, best_ens, best_func : tuple The weighted ensemble, the best ensemble member, and the best objective function value """ + # Update state vector using new base class method + self.stateX = x - # Set the ensemble state equal to the input control vector x - self.state = ot.update_optim_state(x, self.state, list(self.state.keys())) - - # Set the inflation factor and covariance equal to the input + # Set the inflation factor, covariance and survival factor equal to the input self.inflation_factor = args[0] - self.cov = args[1] + self.covX = args[1] self.survival_factor = args[2] - # If bias correction is used we need to temporarily store the initial state - initial_state = None - if self.bias_file is not None and self.bias_factors is None: # first iteration - initial_state = deepcopy(self.state) # store this to update current objective values - # Generate ensemble of states if self.resample_index is None: self.ne = self.num_samples else: self.ne = int(np.round(self.num_samples*self.survival_factor)) - self._aux_input() - self.state = self._gen_state_ensemble() + + nr = self._aux_input() + + # Generate state ensemble + self.enX = np.random.multivariate_normal(self.stateX, self.covX, self.ne).T + + # Truncate to bounds + if (self.lb is not None) and (self.ub is not None): + self.enX = np.clip(self.enX, self.lb[:, None], self.ub[:, None]) - state_ens = at.aug_state(self.state, list(self.state.keys())) - self.function(state_ens, **kwargs) + # Evaluate objective function for ensemble + self.enF = self.function(self.enX, **kwargs) - if not isinstance(self.ens_func_values, list): - self.ens_func_values = [self.ens_func_values] - L = len(self.ens_func_values) + if not isinstance(self.enF, list): + self.enF = [self.enF] + + L = len(self.enF) if self.resample_index is None: self.resample_index = [None]*L - # If bias correction is used we need to calculate the bias factors, J(u_j,m_j)/J(u_j,m) - if self.bias_file is not None: # use bias corrections - self._bias_factors(self.ens_func_values, initial_state) - warnings.filterwarnings('ignore') # suppress warnings start_index = 0 level_sens = [] - sens_matrix = np.zeros(state_ens.shape[0]) + sens_matrix = np.zeros(self.enX.shape[0]) best_ens = 0 best_func = 0 ml_ne_new_total = 0 + if 'multilevel' in self.keys_en.keys(): en_size = ot.get_list_element(self.keys_en['multilevel'], 'en_size') else: en_size = [self.num_samples] + for l in range(L): - ml_ne = en_size[l] if L > 1 and l == L-1: ml_ne_new = int(np.round(self.num_samples*self.survival_factor)) - ml_ne_new_total @@ -290,29 +290,29 @@ def calc_ensemble_weights(self, x, *args, **kwargs): ml_ne_surv = ml_ne - ml_ne_new # surviving samples if self.resample_index[l] is None: - self.particles.append(deepcopy(state_ens[:, start_index:start_index + ml_ne])) - self.particle_values.append(deepcopy(self.ens_func_values[l])) + self.particles.append(deepcopy(self.enX[:, start_index:start_index + ml_ne])) + self.particle_values.append(deepcopy(self.enF[l])) else: self.particles[l][:, :ml_ne_surv] = self.particles[l][:, self.resample_index[l]] - self.particles[l][:, ml_ne_surv:] = deepcopy(state_ens[:, start_index:start_index + ml_ne_new]) + self.particles[l][:, ml_ne_surv:] = deepcopy(self.enX[:, start_index:start_index + ml_ne_new]) self.particle_values[l][:ml_ne_surv] = self.particle_values[l][self.resample_index[l]] - self.particle_values[l][ml_ne_surv:] = deepcopy(self.ens_func_values[l]) + self.particle_values[l][ml_ne_surv:] = deepcopy(self.enF[l]) # Calculate the weights and ensemble sensitivity matrix weights = np.zeros(ml_ne) - for i in np.arange(ml_ne): + for i in range(ml_ne): weights[i] = np.exp(np.clip(-(self.particle_values[l][i] - np.min( self.particle_values[l])) * self.inflation_factor, None, 10)) - weights = weights + 0.000001 - weights = weights/np.sum(weights) # TODO: Sjekke at disse er riktig + weights = weights + 1e-6 # Add small regularization + weights = weights/np.sum(weights) level_sens.append(self.particles[l] @ weights) if l == L-1: # keep the best from the finest level index = np.argmin(self.particle_values[l]) best_ens = self.particles[l][:, index] best_func = self.particle_values[l][index] - self.resample_index[l] = np.random.choice(ml_ne,ml_ne_surv,replace=True,p=weights) + self.resample_index[l] = np.random.choice(ml_ne, ml_ne_surv, replace=True, p=weights) start_index += ml_ne_new From 160da50c4f647c468efef339043365992b006f3f Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 7 Nov 2025 10:55:27 +0100 Subject: [PATCH 47/94] Remove unused bias functions --- popt/loop/ensemble_gaussian.py | 29 ----------------------------- 1 file changed, 29 deletions(-) diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index 0c185a4..7e4ee26 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -347,35 +347,6 @@ def _gen_state_ensemble(self): return state_en - def _bias_correction(self, state): - """ - Calculate bias correction. Currently, the bias correction is a constant independent of the state - """ - if self.bias_factors is not None: - return np.sum(self.bias_weights * self.bias_factors) - else: - return 1 - - def _bias_factors(self, obj_func_values, initial_state): - """ - Function for computing the bias factors - """ - - if self.bias_factors is None: # first iteration - currentfile = self.sim.file - self.sim.file = self.bias_file - self.ne = self.num_samples - self.aux_input = list(np.arange(self.ne)) - self.calc_prediction() - self.sim.file = currentfile - bias_func_values = self.obj_func(self.pred_data, self.sim.input_dict, self.sim.true_order) - bias_func_values = np.array(bias_func_values) - self.bias_factors = bias_func_values / obj_func_values - self.bias_points = deepcopy(self.state) - self.state_func_values *= self._bias_correction(initial_state) - elif self.bias_adaptive is not None and self.bias_adaptive > 0: # update factors to account for new information - pass # not implemented yet - From a8e7b5acd554e58d333e8462ea18ec9b007f7b87 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 7 Nov 2025 10:57:33 +0100 Subject: [PATCH 48/94] Remove unused function --- popt/loop/ensemble_gaussian.py | 23 +---------------------- 1 file changed, 1 insertion(+), 22 deletions(-) diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index 7e4ee26..5976865 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -82,7 +82,7 @@ def gradient(self, x, *args, **kwargs): Control vector, shape (number of controls, ) args : tuple - Covarice ($C_x$), shape (number of controls, number of controls) + Covarice matrix, shape (number of controls, number of controls) Returns ------- @@ -326,27 +326,6 @@ def calc_ensemble_weights(self, x, *args, **kwargs): return sens_matrix, best_ens, best_func - def _gen_state_ensemble(self): - """ - Generate ensemble with the current state (control variable) as the mean and using the covariance matrix - """ - - state_en = {} - cov_blocks = ot.corr2BlockDiagonal(self.state, self.cov) - for i, statename in enumerate(self.state.keys()): - mean = self.state[statename] - cov = cov_blocks[i] - temp_state_en = np.random.multivariate_normal(mean, cov, self.ne).transpose() - shifted_ensemble = np.array([mean]).T + temp_state_en - np.array([np.mean(temp_state_en, 1)]).T - if self.lb and self.ub: - if self.transform: - np.clip(shifted_ensemble, 0, 1, out=shifted_ensemble) - else: - np.clip(shifted_ensemble, self.lb[i], self.ub[i], out=shifted_ensemble) - state_en[statename] = shifted_ensemble - - return state_en - From 7556c4d77b02b8702166378d32a9aa9f30fed4e9 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 7 Nov 2025 10:59:05 +0100 Subject: [PATCH 49/94] Remove unused imports --- ensemble/ensemble.py | 3 --- popt/loop/ensemble_base.py | 4 ---- popt/loop/ensemble_gaussian.py | 2 -- 3 files changed, 9 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index c29da52..60fc445 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -20,9 +20,6 @@ import pipt.misc_tools.analysis_tools as at import pipt.misc_tools.extract_tools as extract import pipt.misc_tools.ensemble_tools as entools -from pipt.misc_tools import cov_regularization -from pipt.misc_tools import wavelet_tools as wt -from misc import read_input_csv as rcsv from misc.system_tools.environ_var import OpenBlasSingleThread # Single threaded OpenBLAS runs diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index 0e4ce98..32f98bb 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -2,13 +2,9 @@ import numpy as np import pandas as pd import sys -import warnings - -from copy import deepcopy # Internal imports from popt.misc_tools import optim_tools as ot -from pipt.misc_tools import analysis_tools as at from ensemble.ensemble import Ensemble as SupEnsemble from simulator.simple_models import noSimulation from pipt.misc_tools.ensemble_tools import matrix_to_dict diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index 5976865..f59446d 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -1,13 +1,11 @@ # External imports import numpy as np -import sys import warnings from copy import deepcopy # Internal imports from popt.misc_tools import optim_tools as ot -from pipt.misc_tools import analysis_tools as at from popt.loop.ensemble_base import EnsembleOptimizationBaseClass __all__ = ['GaussianEnsemble'] From 582228e821ae043f707cb01508464e26835a362f Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 7 Nov 2025 11:01:30 +0100 Subject: [PATCH 50/94] Update docstring --- popt/loop/ensemble_gaussian.py | 27 ++++++++------------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index f59446d..51bd9a6 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -16,25 +16,13 @@ class GaussianEnsemble(EnsembleOptimizationBaseClass): Methods ------- - get_state() - Returns control vector as ndarray - - get_final_state(return_dict) - Returns final control vector between [lb,ub] - - get_cov() - Returns the ensemble covariance matrix - - function(x,*args) - Objective function called during optimization - - gradient(x,*args) + gradient(x, *args, **kwargs) Ensemble gradient - - hessian(x,*args) + + hessian(x, *args, **kwargs) Ensemble hessian - calc_ensemble_weights(self,x,*args): + calc_ensemble_weights(self,x, *args, **kwargs): Calculate weights used in sequential monte carlo optimization """ @@ -43,7 +31,7 @@ def __init__(self, options, simulator, objective): """ Parameters ---------- - keys_en : dict + options : dict Options for the ensemble class - disable_tqdm: supress tqdm progress bar for clean output in the notebook @@ -52,11 +40,12 @@ def __init__(self, options, simulator, objective): - prior_: the prior information the state variables, including mean, variance and variable limits - num_models: number of models (if robust optimization) (default 1) - transform: transform variables to [0,1] if true (default true) + - natural_gradient: use natural gradient if true (default false) - sim : callable + simulator : callable The forward simulator (e.g. flow) - obj_func : callable + objective : callable The objective function (e.g. npv) """ From 0196097a5944eca5e6c0cbb20e93eda394f5c063 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 17 Nov 2025 08:24:29 +0100 Subject: [PATCH 51/94] Update docstring --- popt/loop/ensemble_gaussian.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index 51bd9a6..c94c201 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -12,7 +12,7 @@ class GaussianEnsemble(EnsembleOptimizationBaseClass): """ - Class to store control states and evaluate objective functions. + Gaussian Ensemble class for ensemble-based optimization. Methods ------- @@ -24,7 +24,6 @@ class GaussianEnsemble(EnsembleOptimizationBaseClass): calc_ensemble_weights(self,x, *args, **kwargs): Calculate weights used in sequential monte carlo optimization - """ def __init__(self, options, simulator, objective): From 08fce6b3be7608c6e611fb949c5052c3ce074309 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 19 Nov 2025 14:15:07 +0100 Subject: [PATCH 52/94] Update GenOpt to state changes --- popt/loop/ensemble_generalized.py | 46 +++++++++++++++---------------- 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/popt/loop/ensemble_generalized.py b/popt/loop/ensemble_generalized.py index c786627..431a8b2 100644 --- a/popt/loop/ensemble_generalized.py +++ b/popt/loop/ensemble_generalized.py @@ -32,8 +32,8 @@ def __init__(self, options, simulator, objective): super().__init__(options, simulator, objective) # construct corr matrix - std = np.sqrt(np.diag(self.cov)) - self.corr = self.cov/np.outer(std, std) + std = np.sqrt(np.diag(self.covX)) + self.corr = self.covX/np.outer(std, std) self.dim = std.size # choose marginal @@ -61,16 +61,16 @@ def __init__(self, options, simulator, objective): elif marginal == 'Logistic': self.margs = Logistic() - self.theta = options.get('theta', self.margs.var_to_scale(np.diag(self.cov))) + self.theta = options.get('theta', self.margs.var_to_scale(np.diag(self.covX))) elif marginal == 'TruncGaussian': lb, ub = np.array(self.bounds).T self.margs = TruncGaussian(lb,ub) - self.theta = options.get('theta', np.sqrt(np.diag(self.cov))) + self.theta = options.get('theta', np.sqrt(np.diag(self.covX))) elif marginal == 'Gaussian': self.margs = Gaussian() - self.theta = options.get('theta', np.sqrt(np.diag(self.cov))) + self.theta = options.get('theta', np.sqrt(np.diag(self.covX))) def get_theta(self): return self.theta @@ -90,15 +90,15 @@ def sample(self, size=None): def gradient(self, x, *args, **kwargs): - # Set the ensemble state equal to the input control vector x - self.state = ot.update_optim_state(x, self.state, list(self.state.keys())) + # Update state vector + self.stateX = x if args: self.theta, self.corr = args self.enZ = kwargs.get('enZ', None) self.enX = kwargs.get('enX', None) - self.enJ = kwargs.get('enJ', None) + self.enF = kwargs.get('enF', None) ne = self.num_samples nr = self._aux_input() @@ -109,15 +109,15 @@ def gradient(self, x, *args, **kwargs): self.enX, self.enZ = self.sample(size=ne) # Evaluate - if self.enJ is None: - self.enJ = self.function(self._trafo_ensemble(x).T) + if self.enF is None: + self.enF = self.function(self._trafo_ensemble(x).T) self.avg_hess = np.zeros((dim,dim)) self.avg_grad = np.zeros(dim) H = np.linalg.inv(self.corr)-np.eye(dim) O = np.ones((dim,dim))-np.eye(dim) - enJ = self.enJ - np.array(np.repeat(self.state_func_values, nr)) + enF = self.enF - np.repeat(self.stateF, nr) for n in range(self.ne): @@ -138,8 +138,8 @@ def gradient(self, x, *args, **kwargs): # calc grad and hess grad_log_p = G + D hess_log_p = np.diag(K)+M - self.avg_grad += enJ[n]*grad_log_p - self.avg_hess += enJ[n]*(np.outer(grad_log_p, grad_log_p) + hess_log_p) + self.avg_grad += enF[n]*grad_log_p + self.avg_hess += enF[n]*(np.outer(grad_log_p, grad_log_p) + hess_log_p) self.avg_grad = -self.avg_grad*self.grad_scale/ne self.avg_hess = self.avg_hess*self.hess_scale/ne @@ -148,8 +148,8 @@ def gradient(self, x, *args, **kwargs): def hessian(self, x, *args, **kwargs): - # Set the ensemble state equal to the input control vector x - self.state = ot.update_optim_state(x, self.state, list(self.state.keys())) + # Update state vector + self.stateX = x if kwargs.get('sample', False): self.gradient(x, *args, **kwargs) @@ -165,7 +165,7 @@ def mutation_gradient(self, x, *args, **kwargs): self.enZ = kwargs.get('enZ', None) self.enX = kwargs.get('enX', None) - self.enJ = kwargs.get('enJ', None) + self.enF = kwargs.get('enF', None) ne = self.num_samples nr = self._aux_input() @@ -176,10 +176,10 @@ def mutation_gradient(self, x, *args, **kwargs): self.enX, self.enZ = self.sample(size=ne) # Evaluate - if self.enJ is None: - self.enJ = self.function(self._trafo_ensemble(x).T) + if self.enF is None: + self.enF = self.function(self._trafo_ensemble(x).T) - enJ = self.enJ - np.array(np.repeat(self.state_func_values, nr)) + enF = self.enF - np.repeat(self.stateF, nr) self.nat_grad = np.zeros(dim) self.nat_hess = np.zeros(dim) @@ -189,8 +189,8 @@ def mutation_gradient(self, x, *args, **kwargs): dm_log_p = self.margs.grad_theta_log_pdf(X, self.theta, mean=x) hm_log_p = self.margs.hess_theta_log_pdf(X, self.theta, mean=x) - self.nat_grad += enJ[n]*dm_log_p - self.nat_hess += enJ[n]*(hm_log_p + dm_log_p**2) + self.nat_grad += enF[n]*dm_log_p + self.nat_hess += enF[n]*(hm_log_p + dm_log_p**2) # Fisher self.nat_grad = self.nat_grad/ne @@ -199,8 +199,8 @@ def mutation_gradient(self, x, *args, **kwargs): def mutation_hessian(self, x, *args, **kwargs): - # Set the ensemble state equal to the input control vector x - self.state = ot.update_optim_state(x, self.state, list(self.state.keys())) + # Update state vector + self.stateX = x if kwargs.get('sample', False): self.gradient(x, *args, **kwargs) From b709fd08004037eeda0cfed2020a2f089bbb059c Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 20 Nov 2025 10:49:02 +0100 Subject: [PATCH 53/94] Improve readability of code --- pipt/loop/assimilation.py | 3 +- pipt/loop/ensemble.py | 66 +++++++- pipt/misc_tools/analysis_tools.py | 49 +++++- pipt/update_schemes/enrml.py | 93 ++++++----- .../update_methods_ns/approx_update.py | 157 ++++++++---------- 5 files changed, 233 insertions(+), 135 deletions(-) diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index dfe2660..39acefb 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -102,8 +102,9 @@ def run(self): ) # Run a while loop until max. iterations or convergence is reached - while self.ensemble.iteration < self.max_iter and conv is False: + while (self.ensemble.iteration < self.max_iter) and (conv is False): # Add a check to see if this is the prior model + if self.ensemble.iteration == 0: # Calc forecast for prior model # Inset 0 as input to forecast all data diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index b5ada98..842bf2c 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -493,8 +493,69 @@ def _org_data_var(self): vintage = vintage + 1 def _ext_obs(self): - self.obs_data_vector, _ = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - self.list_datatypes) + #self.obs_data_vector, _ = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, + # self.list_datatypes) + + self.vecObs, _ = at.aug_obs_pred_data( + self.obs_data, + self.pred_data, + self.assim_index, + self.list_datatypes + ) + + # Generate ensemble of perturbed observed data + if ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes'): + + if hasattr(self, 'cov_data'): # cd matrix has been imported + # enObs: samples from N(0,Cd) + enObs = cholesky(self.cov_data).T @ np.random.randn(self.cov_data.shape[0], self.ne) + else: + enObs = at.extract_tot_empirical_cov( + self.datavar, + self.assim_index, + self.list_datatypes, + self.ne + ) + + # Screen data if required + if ('screendata' in self.keys_da) and (self.keys_da['screendata'] == 'yes'): + enObs = at.screen_data( + enObs, + self.enPred, + self.vecObs, + self.iteration + ) + + # Center the ensemble of perturbed observed data + self.enObs = self.vecObs[:, np.newaxis] + enObs + self.cov_data = np.var(self.enObs, ddof=1, axis=1) + self.scale_data = np.sqrt(self.cov_data) + + else: + if not hasattr(self, 'cov_data'): # if cd is not loaded + self.cov_data = at.gen_covdata( + datavar = self.datavar, + assim_index = self.assim_index, + list_data = self.list_datatypes, + ) + # data screening + if ('screendata' in self.keys_da) and (self.keys_da['screendata'] == 'yes'): + self.cov_data = at.screen_data( + data = self.cov_data, + aug_pred_data = self.enPred, + obs_data_vector = self.vecObs, + iteration = self.iteration + ) + + generator = Cholesky() # Initialize GeoStat class for generating realizations + self.enObs, self.scale_data = generator.gen_real( + mean = self.vecObs, + var = self.cov_data, + number = self.ne, + return_chol = True + ) + + '''' # Generate the data auto-covariance matrix if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': if hasattr(self, 'cov_data'): # cd matrix has been imported @@ -526,6 +587,7 @@ def _ext_obs(self): init_en = Cholesky() # Initialize GeoStat class for generating realizations self.real_obs_data, self.scale_data = init_en.gen_real(self.obs_data_vector, self.cov_data, self.ne, return_chol=True) + ''' def _ext_scaling(self): # get vector of scaling diff --git a/pipt/misc_tools/analysis_tools.py b/pipt/misc_tools/analysis_tools.py index f22078b..14504b0 100644 --- a/pipt/misc_tools/analysis_tools.py +++ b/pipt/misc_tools/analysis_tools.py @@ -946,6 +946,7 @@ def aug_obs_pred_data(obs_data, pred_data, assim_index, list_data): tot_pred = tuple(pred_data[el][dat] for el in l_prim if pred_data[el] is not None for dat in list_data if obs_data[el][dat] is not None) + if len(tot_pred): # if this is done during the initiallization tot_pred contains nothing pred = np.concatenate(tot_pred) else: @@ -1488,4 +1489,50 @@ def get_obs_size(obs_data, time_index, datatypes): for data in datatypes ] for time in time_index - ] \ No newline at end of file + ] + +def truncSVD(matrix, r=None, energy=None, full_matrices=False): + ''' + Perform truncated SVD on input matrix. + + Parameters + ---------- + matrix : ndarray, shape (m, n) + Input matrix to perform SVD on. + + r : int, optional + Rank to truncate the SVD to. If None, energy must be specified. + + energy : float, optional + Percentage of energy to retain in the truncated SVD. If None, r must be specified. + + full_matrices : bool, optional + Whether to compute full or reduced SVD. Default is False. + + Returns + ------- + U : ndarray, shape (m, r) + Left singular vectors. + + S : ndarray, shape (r,) + Singular values. + + VT : ndarray, shape (r, n) + Right singular vectors transposed. + ''' + # Perform SVD on input matrix + U, S, VT = np.linalg.svd(matrix, full_matrices=full_matrices) + + # If not specified rank, energy must be given + if r is None: + if energy is not None: + # If energy is less than 100 we truncate the SVD matrices + if energy < 100: + r = (np.cumsum(S) / sum(S)) * 100 <= energy + U, S, VT = U[:,r], S[r], VT[r,:] + else: + raise ValueError("Either rank 'r' or 'energy' must be specified for truncSVD.") + else: + U, S, VT = U[:,:r], S[:r], VT[:r,:] + + return U, S, VT \ No newline at end of file diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index 64c3182..6eec85b 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -60,14 +60,31 @@ def __init__(self, keys_da, keys_en, sim): if self.restart is False: # Save prior state in separate variable - #self.prior_state = cp.deepcopy(self.state) self.prior_enX = cp.deepcopy(self.enX) # not sure if this is wise! - # Extract parameters like conv. tol. and damping param. from ITERATION keyword in DATAASSIM - self._ext_iter_param() - - # Within variables + # Set parameters needed for LM-EnRML + options = self.keys_da['iteration'] + if isinstance(options, list): + options = extract.list_to_dict(options) + + self.data_misfit_tol = options.get('data_misfit_tol', 0.01) + self.trunc_energy = options.get('energy', 0.95) + self.step_tol = options.get('step_tol', 0.01) + self.lam = options.get('lambda', 100) + self.lam_max = options.get('lambda_max', 1e10) + self.lam_min = options.get('lambda_min', 0.01) + self.gamma = options.get('lambda_factor', 5) + self.iteration = 0 + + # Ensure that it is given as percentage + if self.trunc_energy > 1: + self.trunc_energy /= 100. + + # Initalize some variables self.prev_data_misfit = None # Data misfit at previous iteration + self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] + + # Load ACTNUM if given if 'actnum' in self.keys_da.keys(): try: self.actnum = np.load(self.keys_da['actnum'])['actnum'] @@ -75,15 +92,12 @@ def __init__(self, keys_da, keys_en, sim): print('ACTNUM file cannot be loaded!') else: self.actnum = None + # At the moment, the iterative loop is threated as an iterative smoother and thus we check if assim. indices # are given as in the Simultaneous loop. self.check_assimindex_simultaneous() - # define the assimilation index - self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] - # define the list of datatypes - self.list_datatypes, self.list_act_datatypes = at.get_list_data_types( - self.obs_data, self.assim_index) + self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(self.obs_data, self.assim_index) # Get the perturbed observations and observation scaling self.data_random_state = cp.deepcopy(np.random.get_state()) self._ext_obs() @@ -97,14 +111,22 @@ def calc_analysis(self): Calculate the update step in LM-EnRML, which is just the Levenberg-Marquardt update algorithm with the sensitivity matrix approximated by the ensemble. """ - - # reformat predicted data - _, self.aug_pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - self.list_datatypes) + # Get Ensemble of predicted data + _, self.enPred = at.aug_obs_pred_data( + self.obs_data, + self.pred_data, + self.assim_index, + self.list_datatypes + ) if self.iteration == 1: # first iteration + + # Calculate the prior data misfit data_misfit = at.calc_objectivefun( - self.real_obs_data, self.aug_pred_data, self.cov_data) + pert_obs=self.enObs, + pred_data=self.enPred, + Cd=self.cov_data + ) # Store the (mean) data misfit (also for conv. check) self.data_misfit = np.mean(data_misfit) @@ -112,7 +134,7 @@ def calc_analysis(self): self.data_misfit_std = np.std(data_misfit) if self.lam == 'auto': - self.lam = (0.5 * self.prior_data_misfit)/self.aug_pred_data.shape[0] + self.lam = (0.5 * self.prior_data_misfit)/self.enPred.shape[0] self.logger.info( f'Prior run complete with data misfit: {self.prior_data_misfit:0.1f}. Lambda for initial analysis: {self.lam}') @@ -123,8 +145,8 @@ def calc_analysis(self): # Perform the update self.update( enX = self.enX, - enY = self.aug_pred_data, - enE = self.real_obs_data, + enY = self.enPred, + enE = self.enObs, prior = self.prior_enX ) @@ -154,8 +176,14 @@ def check_convergence(self): met """ - _, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - self.list_datatypes) + # Get Ensemble of predicted data + _, enPred = at.aug_obs_pred_data( + self.obs_data, + self.pred_data, + self.assim_index, + self.list_datatypes + ) + # Initialize the initial success value success = False @@ -167,7 +195,7 @@ def check_convergence(self): # mat_obs = np.dot(obs_data_vector.reshape((len(obs_data_vector),1)), np.ones((1, self.ne))) # use the perturbed # data instead. - data_misfit = at.calc_objectivefun(self.real_obs_data, pred_data, self.cov_data) + data_misfit = at.calc_objectivefun(self.enObs, enPred, self.cov_data) self.data_misfit = np.mean(data_misfit) self.data_misfit_std = np.std(data_misfit) @@ -259,29 +287,6 @@ def check_convergence(self): # Return conv = False, why_stop var. return False, success, why_stop - def _ext_iter_param(self): - """ - Extract parameters needed in LM-EnRML from the ITERATION keyword given in the DATAASSIM part of PIPT init. - file. These parameters include convergence tolerances and parameters for the damping parameter. Default - values for these parameters have been given here, if they are not provided in ITERATION. - """ - options = self.keys_da['iteration'] - if isinstance(options, list): - options = extract.list_to_dict(options) - - # unpack options - self.data_misfit_tol = options.get('data_misfit_tol', 0.01) - self.trunc_energy = options.get('energy', 0.95) - self.step_tol = options.get('step_tol', 0.01) - self.lam = options.get('lambda', 100) - self.lam_max = options.get('lambda_max', 1e10) - self.lam_min = options.get('lambda_min', 0.01) - self.gamma = options.get('lambda_factor', 5) - self.iteration = 0 - - if self.trunc_energy > 1: # ensure that it is given as percentage - self.trunc_energy /= 100. - class lmenrml_approx(lmenrmlMixIn, approx_update): pass diff --git a/pipt/update_schemes/update_methods_ns/approx_update.py b/pipt/update_schemes/update_methods_ns/approx_update.py index d607862..4ba20e9 100644 --- a/pipt/update_schemes/update_methods_ns/approx_update.py +++ b/pipt/update_schemes/update_methods_ns/approx_update.py @@ -36,75 +36,86 @@ def update(self, enX, enY, enE, **kwargs): enYcentered = self.scale(np.dot(enY, self.proj), self.scale_data) # Perform truncated SVD - u_d, s_d, v_d = np.linalg.svd(enYcentered, full_matrices=False) - - if self.trunc_energy < 1: - ti = (np.cumsum(s_d) / sum(s_d)) <= self.trunc_energy - u_d, s_d, v_d = u_d[:, ti].copy(), s_d[ti].copy(), v_d[ti, :].copy() + #u_d, s_d, v_d = at.truncSVD(enYcentered, energy=self.trunc_energy) + Ud, Sd, VTd = at.truncSVD(enYcentered, energy=self.trunc_energy) # Check for localization methods if 'localization' in self.keys_da: + loc_info = self.localization.loc_info # Calculate the localization projection matrix if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': + # Scale and center the data ensemble matrix enEcentered = self.scale(np.dot(enE, self.proj), self.scale_data) - x_0 = np.diag(1/s_d) @ u_d.T @ enEcentered - Lam, z = np.linalg.eig(x_0 @ x_0.T) - X = (v_d.T @ z) @ solve( (self.lam + 1)*np.diag(Lam) + np.eye(len(Lam)), (u_d.T @ (np.diag(1/s_d) @ z)).T ) + + # Calculate intermediate matrix + Sinv = np.diag(1/Sd) + X0 = Sinv @ Ud.T @ enEcentered + + # Eigen decomposition of X0 X0^T + eigval, eigvec = np.linalg.eig(X0 @ X0.T) + reg_term = (self.lam + 1) * np.diag(eigval) + np.eye(len(eigval)) + X = (VTd.T @ eigvec) @ solve(reg_term, (Ud.T @ (Sinv @ eigvec)).T) + + + #x_0 = np.diag(1/s_d) @ u_d.T @ enEcentered + #Lam, z = np.linalg.eig(x_0 @ x_0.T) + #X = (v_d.T @ z) @ solve( (self.lam + 1)*np.diag(Lam) + np.eye(len(Lam)), (u_d.T @ (np.diag(1/s_d) @ z)).T ) else: - X = v_d.T @ np.diag(s_d) @ solve( (self.lam + 1)*np.eye(len(s_d)) + np.diag(s_d**2), u_d.T) + reg_term = (self.lam + 1)*np.eye(Sd.size) + np.diag(Sd**2) + X = VTd.T @ np.diag(Sd) @ solve(reg_term, Ud.T) # Check for adaptive localization - if 'autoadaloc' in self.localization.loc_info: + if 'autoadaloc' in loc_info: - # Scale and center the state ensemble matrix + # Scale and center the state ensemble matrix, enX if ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes'): - enXcentered = self.scale(self.enX - np.mean(self.enX, 1)[:,None], self.state_scaling) + enXcentered = self.scale(enX - np.mean(enX, 1)[:,None], self.state_scaling) else: enXcentered = self.scale(np.dot(enX, self.proj), self.state_scaling) - # Calculate and scale difference between observations and predictions - scaled_delta_data = self.scale(enE - enY, self.scale_data) + # Calculate and scale difference between observations and predictions (residuals) + enRes = self.scale(enE - enY, self.scale_data) # Compute the update step with auto-adaptive localization self.step = self.localization.auto_ada_loc( pert_state = self.state_scaling[:, None]*enXcentered, - proj_pred_data = np.dot(X, scaled_delta_data), + proj_pred_data = np.dot(X, enRes), curr_param = self.list_states, prior_info = self.prior_info ) # Check for local analysis - elif ('localanalysis' in self.localization.loc_info) and (self.localization.loc_info['localanalysis']): + elif ('localanalysis' in loc_info) and (loc_info['localanalysis']): # Calculate weights - if 'distance' in self.localization.loc_info: + if 'distance' in loc_info: weight = _calc_loc( - max_dist = self.localization.loc_info['range'], - distance = self.localization.loc_info['distance'], + max_dist = loc_info['range'], + distance = loc_info['distance'], prior_info = self.prior_info[self.list_states[0]], - loc_type = self.localization.loc_info['type'], + loc_type = loc_info['type'], ne = self.ne ) else: # if no distance, do full update weight = np.ones((enX.shape[0], X.shape[1])) # Center ensemble matrix - enXcentered = enX - np.mean(self.enX, axis=1, keepdims=True) + enXcentered = enX - np.mean(enX, axis=1, keepdims=True) if (not ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes')): enXcentered /= np.sqrt(self.ne - 1) - # Calculate and scale difference between observations and predictions - scaled_delta_data = self.scale(enE - enY, self.scale_data) + # Calculate and scale difference between observations and predictions (residuals) + enRes = self.scale(enE - enY, self.scale_data) # Compute the update step with local analysis try: - self.step = weight.multiply(np.dot(enXcentered, X)).dot(scaled_delta_data) + self.step = weight.multiply(np.dot(enXcentered, X)).dot(enRes) except: - self.step = (weight*(np.dot(enXcentered, X))).dot(scaled_delta_data) + self.step = (weight*(np.dot(enXcentered, X))).dot(enRes) # Check for distance based localization @@ -126,15 +137,15 @@ def update(self, enX, enY, enE, **kwargs): if not ('emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes'): enXcentered /= np.sqrt(self.ne - 1) - # Calculate and scale difference between observations and predictions - scaled_delta_data = self.scale(enE - enY, self.scale_data) + # Calculate and scale difference between observations and predictions (residuals) + enRes = self.scale(enE - enY, self.scale_data) # Compute the update step with distance-based localization - self.step = mask.multiply(np.dot(enXcentered, X)).dot(scaled_delta_data) + self.step = mask.multiply(np.dot(enXcentered, X)).dot(enRes) - # Else do parallel update + # Else do parallel update (NOT UPDATED, TO NEW DEFINITIONS OF ENSEMBLE MATRICES) else: act_data_list = {} count = 0 @@ -174,34 +185,43 @@ def update(self, enX, enY, enE, **kwargs): self.step = at.aug_state(self.step, self.list_states) else: - # Centered ensemble matrix - if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - pert_state = (self.state_scaling**(-1))[:, None] * (self.enX - np.mean(self.enX, axis=1, keepdims=True)) - else: - pert_state = (self.state_scaling**(-1))[:, None] * np.dot(self.enX, self.proj) - if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': + if ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes'): - # Scale data matrix - if len(self.scale_data.shape) == 1: - E_hat = (1/self.scale_data)[:, None] * self.E - else: - E_hat = solve(self.scale_data, self.E) + # Scale and center the ensemble matrecies: enX and enE + enXcentered = self.scale(enX - np.mean(enX, 1)[:,None], self.state_scaling) + enEcentered = self.scale(enE - np.mean(enE, 1)[:,None], self.scale_data) - x_0 = np.diag(s_d ** -1) @ u_d.T @ E_hat - Lam, z = np.linalg.eig(x_0 @ x_0.T) + Sinv = np.diag(1/Sd) + X0 = Sinv @ Ud.T @ enEcentered + eigval, eigvec = np.linalg.eig(X0 @ X0.T) - if len(self.scale_data.shape) == 1: - delta_data = (1/self.scale_data)[:, None] * (self.real_obs_data - self.aug_pred_data) - else: - delta_data = solve(self.scale_data, self.real_obs_data - self.aug_pred_data) + # Calculate and scale difference between observations and predictions (residuals) + enRes = self.scale(enE - enY, self.scale_data) - x_1 = (u_d @ (np.diag(s_d ** -1).T @ z)).T @ delta_data - x_2 = solve((self.lam + 1) * np.diag(Lam) + np.eye(len(Lam)), x_1) - x_3 = np.dot(np.dot(v_d.T, z), x_2) - self.step = np.dot(self.state_scaling[:, None] * pert_state, x_3) + # Compute the update step + X1 = (Ud @ Sinv @ eigvec).T @ enRes + X2 = solve((self.lam + 1) * np.diag(eigval) + np.eye(len(eigval)), X1) + X3 = np.dot(VTd.T, eigvec) @ X2 + self.step = np.dot(self.state_scaling[:, None]*enXcentered, X3) + + + #x_0 = np.diag(s_d ** -1) @ u_d.T @ E_hat + #Lam, z = np.linalg.eig(x_0 @ x_0.T) + + #if len(self.scale_data.shape) == 1: + # delta_data = (1/self.scale_data)[:, None] * (self.real_obs_data - self.aug_pred_data) + #else: + # delta_data = solve(self.scale_data, self.real_obs_data - self.aug_pred_data) + + #x_1 = (u_d @ (np.diag(s_d ** -1).T @ z)).T @ delta_data + #x_2 = solve((self.lam + 1) * np.diag(Lam) + np.eye(len(Lam)), x_1) + #x_3 = np.dot(np.dot(v_d.T, z), x_2) + #self.step = np.dot(self.state_scaling[:, None] * pert_state, x_3) else: + enXcentered = self.scale(np.dot(enX, self.proj), self.state_scaling) + # Compute the approximate update (follow notation in paper) if len(self.scale_data.shape) == 1: x_1 = np.dot(u_d.T, (1/self.scale_data)[:, None] * (self.real_obs_data - self.aug_pred_data)) @@ -213,43 +233,6 @@ def update(self, enX, enY, enE, **kwargs): self.step = np.dot(self.state_scaling[:, None] * pert_state, x_3) - - - - - def _update_with_distance_based_localization(self, X): - - # Get data size - data_size = [[self.obs_data[int(time)][data].size if self.obs_data[int(time)][data] is not None else 0 - for data in self.list_datatypes] for time in self.assim_index[1]] - - # Setup localization - local_mask = self.localization.localize( - self.list_datatypes, - [self.keys_da['truedataindex'][int(elem)] for elem in self.assim_index[1]], - self.list_states, - self.ne, - self.prior_info, - data_size - ) - - # Center ensemble matrix - mean_state = np.mean(self.enX, axis=1, keepdims=True) - pert_state = self.enX - mean_state - if not ('emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes'): - pert_state /= np.sqrt(self.ne - 1) - - # Calculate difference between observations and predictions - if self.scale_data.ndim == 1: - scaled_delta_data = (self.scale_data ** -1)[:, None] * (self.real_obs_data - self.aug_pred_data) - else: - scaled_delta_data = solve(self.scale_data, self.real_obs_data - self.aug_pred_data) - - # Compute the update step with distance-based localization - step = local_mask.multiply(np.dot(pert_state, X)).dot(scaled_delta_data) - - return step - def scale(self, data, scaling): """ Scale the data perturbations by the data error standard deviation. From 8eb2a5c5308a559807e1e1c0f4fcf01ee8f7305f Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 21 Nov 2025 08:53:32 +0100 Subject: [PATCH 54/94] Improve readability --- pipt/loop/ensemble.py | 8 ++-- pipt/update_schemes/enrml.py | 3 +- .../update_methods_ns/approx_update.py | 37 ++++--------------- 3 files changed, 13 insertions(+), 35 deletions(-) diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index 842bf2c..a62ad1d 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -493,16 +493,14 @@ def _org_data_var(self): vintage = vintage + 1 def _ext_obs(self): - #self.obs_data_vector, _ = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - # self.list_datatypes) - + self.vecObs, _ = at.aug_obs_pred_data( self.obs_data, self.pred_data, self.assim_index, self.list_datatypes ) - + # Generate ensemble of perturbed observed data if ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes'): @@ -527,7 +525,7 @@ def _ext_obs(self): ) # Center the ensemble of perturbed observed data - self.enObs = self.vecObs[:, np.newaxis] + enObs + self.enObs = self.vecObs[:, np.newaxis] - enObs self.cov_data = np.var(self.enObs, ddof=1, axis=1) self.scale_data = np.sqrt(self.cov_data) diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index 6eec85b..cf3fdf8 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -82,7 +82,6 @@ def __init__(self, keys_da, keys_en, sim): # Initalize some variables self.prev_data_misfit = None # Data misfit at previous iteration - self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] # Load ACTNUM if given if 'actnum' in self.keys_da.keys(): @@ -96,6 +95,8 @@ def __init__(self, keys_da, keys_en, sim): # At the moment, the iterative loop is threated as an iterative smoother and thus we check if assim. indices # are given as in the Simultaneous loop. self.check_assimindex_simultaneous() + self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] + # define the list of datatypes self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(self.obs_data, self.assim_index) # Get the perturbed observations and observation scaling diff --git a/pipt/update_schemes/update_methods_ns/approx_update.py b/pipt/update_schemes/update_methods_ns/approx_update.py index 4ba20e9..1947af1 100644 --- a/pipt/update_schemes/update_methods_ns/approx_update.py +++ b/pipt/update_schemes/update_methods_ns/approx_update.py @@ -36,7 +36,6 @@ def update(self, enX, enY, enE, **kwargs): enYcentered = self.scale(np.dot(enY, self.proj), self.scale_data) # Perform truncated SVD - #u_d, s_d, v_d = at.truncSVD(enYcentered, energy=self.trunc_energy) Ud, Sd, VTd = at.truncSVD(enYcentered, energy=self.trunc_energy) # Check for localization methods @@ -57,10 +56,6 @@ def update(self, enX, enY, enE, **kwargs): reg_term = (self.lam + 1) * np.diag(eigval) + np.eye(len(eigval)) X = (VTd.T @ eigvec) @ solve(reg_term, (Ud.T @ (Sinv @ eigvec)).T) - - #x_0 = np.diag(1/s_d) @ u_d.T @ enEcentered - #Lam, z = np.linalg.eig(x_0 @ x_0.T) - #X = (v_d.T @ z) @ solve( (self.lam + 1)*np.diag(Lam) + np.eye(len(Lam)), (u_d.T @ (np.diag(1/s_d) @ z)).T ) else: reg_term = (self.lam + 1)*np.eye(Sd.size) + np.diag(Sd**2) X = VTd.T @ np.diag(Sd) @ solve(reg_term, Ud.T) @@ -85,6 +80,7 @@ def update(self, enX, enY, enE, **kwargs): curr_param = self.list_states, prior_info = self.prior_info ) + print(self.step) # Check for local analysis @@ -204,33 +200,16 @@ def update(self, enX, enY, enE, **kwargs): X2 = solve((self.lam + 1) * np.diag(eigval) + np.eye(len(eigval)), X1) X3 = np.dot(VTd.T, eigvec) @ X2 self.step = np.dot(self.state_scaling[:, None]*enXcentered, X3) - - - #x_0 = np.diag(s_d ** -1) @ u_d.T @ E_hat - #Lam, z = np.linalg.eig(x_0 @ x_0.T) - - #if len(self.scale_data.shape) == 1: - # delta_data = (1/self.scale_data)[:, None] * (self.real_obs_data - self.aug_pred_data) - #else: - # delta_data = solve(self.scale_data, self.real_obs_data - self.aug_pred_data) - - #x_1 = (u_d @ (np.diag(s_d ** -1).T @ z)).T @ delta_data - #x_2 = solve((self.lam + 1) * np.diag(Lam) + np.eye(len(Lam)), x_1) - #x_3 = np.dot(np.dot(v_d.T, z), x_2) - #self.step = np.dot(self.state_scaling[:, None] * pert_state, x_3) else: enXcentered = self.scale(np.dot(enX, self.proj), self.state_scaling) - - # Compute the approximate update (follow notation in paper) - if len(self.scale_data.shape) == 1: - x_1 = np.dot(u_d.T, (1/self.scale_data)[:, None] * (self.real_obs_data - self.aug_pred_data)) - else: - x_1 = np.dot(u_d.T, solve(self.scale_data, self.real_obs_data - self.aug_pred_data)) - - x_2 = solve(((self.lam + 1) * np.eye(len(s_d)) + np.diag(s_d ** 2)), x_1) - x_3 = np.dot(np.dot(v_d.T, np.diag(s_d)), x_2) - self.step = np.dot(self.state_scaling[:, None] * pert_state, x_3) + enRes = self.scale(enE - enY, self.scale_data) + + # Compute the update step + X1 = Ud.T @ enRes + X2 = solve((self.lam + 1)*np.eye(Sd.size) + np.diag(Sd**2), X1) + X3 = VTd.T @ np.diag(Sd) @ X2 + self.step = np.dot(self.state_scaling[:, None] * enXcentered, X3) def scale(self, data, scaling): From e59eb02511b7d55ac26f528386ec60955aff5ad5 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 21 Nov 2025 12:48:14 +0100 Subject: [PATCH 55/94] Improve readability --- pipt/loop/ensemble.py | 58 +++++-------------- pipt/misc_tools/analysis_tools.py | 11 ++-- pipt/misc_tools/extract_tools.py | 3 + pipt/update_schemes/enrml.py | 4 +- .../update_methods_ns/approx_update.py | 51 +++++++++------- 5 files changed, 55 insertions(+), 72 deletions(-) diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index a62ad1d..0d9e01b 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -492,9 +492,13 @@ def _org_data_var(self): self.datavar[i][datatype[j]] = est_noise # override the given value vintage = vintage + 1 - def _ext_obs(self): - - self.vecObs, _ = at.aug_obs_pred_data( + + def set_observations(self): + ''' + Generate the perturbed observed data ensemble + ''' + # Make observed data vector + vecObs, _ = at.aug_obs_pred_data( self.obs_data, self.pred_data, self.assim_index, @@ -520,13 +524,13 @@ def _ext_obs(self): enObs = at.screen_data( enObs, self.enPred, - self.vecObs, + vecObs, self.iteration ) # Center the ensemble of perturbed observed data - self.enObs = self.vecObs[:, np.newaxis] - enObs - self.cov_data = np.var(self.enObs, ddof=1, axis=1) + enObs = vecObs[:, np.newaxis] - enObs + self.cov_data = np.var(enObs, ddof=1, axis=1) self.scale_data = np.sqrt(self.cov_data) else: @@ -541,51 +545,19 @@ def _ext_obs(self): self.cov_data = at.screen_data( data = self.cov_data, aug_pred_data = self.enPred, - obs_data_vector = self.vecObs, + obs_data_vector = vecObs, iteration = self.iteration ) generator = Cholesky() # Initialize GeoStat class for generating realizations - self.enObs, self.scale_data = generator.gen_real( - mean = self.vecObs, + enObs, self.scale_data = generator.gen_real( + mean = vecObs, var = self.cov_data, number = self.ne, return_chol = True ) - - '''' - # Generate the data auto-covariance matrix - if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - if hasattr(self, 'cov_data'): # cd matrix has been imported - tmp_E = np.dot(cholesky(self.cov_data).T, - np.random.randn(self.cov_data.shape[0], self.ne)) - else: - tmp_E = at.extract_tot_empirical_cov( - self.datavar, self.assim_index, self.list_datatypes, self.ne) - # self.E = (tmp_E - tmp_E.mean(1)[:,np.newaxis])/np.sqrt(self.ne - 1)/ - if 'screendata' in self.keys_da and self.keys_da['screendata'] == 'yes': - tmp_E = at.screen_data(tmp_E, self.aug_pred_data, - self.obs_data_vector, self.iteration) - self.E = tmp_E - self.real_obs_data = self.obs_data_vector[:, np.newaxis] - tmp_E - - self.cov_data = np.var(self.E, ddof=1, - axis=1) # calculate the variance, to be used for e.g. data misfit calc - # 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 - self.scale_data = np.sqrt(self.cov_data) - else: - if not hasattr(self, 'cov_data'): # if cd is not loaded - self.cov_data = at.gen_covdata( - self.datavar, self.assim_index, self.list_datatypes) - # data screening - if 'screendata' in self.keys_da and self.keys_da['screendata'] == 'yes': - self.cov_data = at.screen_data( - self.cov_data, self.aug_pred_data, self.obs_data_vector, self.iteration) - - init_en = Cholesky() # Initialize GeoStat class for generating realizations - self.real_obs_data, self.scale_data = init_en.gen_real(self.obs_data_vector, self.cov_data, self.ne, - return_chol=True) - ''' + + return vecObs, enObs def _ext_scaling(self): # get vector of scaling diff --git a/pipt/misc_tools/analysis_tools.py b/pipt/misc_tools/analysis_tools.py index 14504b0..4726dc4 100644 --- a/pipt/misc_tools/analysis_tools.py +++ b/pipt/misc_tools/analysis_tools.py @@ -1527,12 +1527,11 @@ def truncSVD(matrix, r=None, energy=None, full_matrices=False): if r is None: if energy is not None: # If energy is less than 100 we truncate the SVD matrices - if energy < 100: - r = (np.cumsum(S) / sum(S)) * 100 <= energy - U, S, VT = U[:,r], S[r], VT[r,:] + if energy < 1: + r = np.sum((np.cumsum(S) / sum(S)) <= energy) + else: + r = np.sum((np.cumsum(S) / sum(S)) <= energy/100) else: raise ValueError("Either rank 'r' or 'energy' must be specified for truncSVD.") - else: - U, S, VT = U[:,:r], S[:r], VT[:r,:] - return U, S, VT \ No newline at end of file + return U[:,:r], S[:r], VT[:r,:] \ No newline at end of file diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index 186c342..032057b 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -17,6 +17,9 @@ from scipy.spatial import cKDTree from typing import Union +# Internal imports +import pipt.misc_tools.analysis_tools as at + def extract_prior_info(keys: dict) -> dict: ''' diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index cf3fdf8..2bed3b8 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -99,9 +99,11 @@ def __init__(self, keys_da, keys_en, sim): # define the list of datatypes self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(self.obs_data, self.assim_index) + # Get the perturbed observations and observation scaling self.data_random_state = cp.deepcopy(np.random.get_state()) - self._ext_obs() + self.vecObs, self.enObs = self.set_observations() + # Get state scaling and svd of scaled prior self._ext_scaling() diff --git a/pipt/update_schemes/update_methods_ns/approx_update.py b/pipt/update_schemes/update_methods_ns/approx_update.py index 1947af1..35a4f30 100644 --- a/pipt/update_schemes/update_methods_ns/approx_update.py +++ b/pipt/update_schemes/update_methods_ns/approx_update.py @@ -5,7 +5,10 @@ import copy as cp from scipy.linalg import solve, solve_banded, cholesky, lu_solve, lu_factor, inv import pickle + +import pipt.misc_tools.ensemble_tools as entools import pipt.misc_tools.analysis_tools as at + from pipt.misc_tools.cov_regularization import _calc_loc @@ -80,7 +83,6 @@ def update(self, enX, enY, enE, **kwargs): curr_param = self.list_states, prior_info = self.prior_info ) - print(self.step) # Check for local analysis @@ -128,7 +130,7 @@ def update(self, enX, enY, enE, **kwargs): ) # Center ensemble matrix - enXcentered = enX - np.mean(self.enX, axis=1, keepdims=True) + enXcentered = enX - np.mean(enX, axis=1, keepdims=True) if not ('emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes'): enXcentered /= np.sqrt(self.ne - 1) @@ -141,21 +143,19 @@ def update(self, enX, enY, enE, **kwargs): - # Else do parallel update (NOT UPDATED, TO NEW DEFINITIONS OF ENSEMBLE MATRICES) + # Else do parallel update (NOT TESTED AFTER UPDATES) else: act_data_list = {} count = 0 for i in self.assim_index[1]: - for el in self.list_datatypes: + for el in list(self.idX.keys()): if self.real_obs_data[int(i)][el] is not None: - act_data_list[( - el, float(self.keys_da['truedataindex'][int(i)]))] = count + act_data_list[(el, float(self.keys_da['truedataindex'][int(i)]))] = count count += 1 - well = [w for w in - set([el[0] for el in self.localization.loc_info.keys() if type(el) == tuple])] - times = [t for t in set( - [el[1] for el in self.localization.loc_info.keys() if type(el) == tuple])] + well = [w for w in set([el[0] for el in loc_info.keys() if type(el) == tuple])] + times = [t for t in set([el[1] for el in loc_info.keys() if type(el) == tuple])] + tot_dat_index = {} for uniq_well in well: tmp_index = [] @@ -163,22 +163,29 @@ def update(self, enX, enY, enE, **kwargs): if (uniq_well, t) in act_data_list: tmp_index.append(act_data_list[(uniq_well, t)]) tot_dat_index[uniq_well] = tmp_index - if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': + + if ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes'): emp_cov = True else: emp_cov = False - self.step = at.parallel_upd(self.list_states, self.prior_info, self.current_state, X, - self.localization.loc_info, self.real_obs_data, self.aug_pred_data, - int(self.keys_fwd['parallel']), - actnum=self.localization.loc_info['actnum'], - field_dim=self.localization.loc_info['field'], - act_data_list=tot_dat_index, - scale_data=self.scale_data, - num_states=len( - [el for el in self.list_states]), - emp_d_cov=emp_cov) - self.step = at.aug_state(self.step, self.list_states) + self.step = at.parallel_upd( + list(self.idX.keys()), + self.prior_info, + entools.matrix_to_dict(enX, self.idX), + X, + loc_info, + enE, + enY, + int(self.keys_fwd['parallel']), + actnum=loc_info['actnum'], + field_dim=loc_info['field'], + act_data_list=tot_dat_index, + scale_data=self.scale_data, + num_states=len([el for el in list(self.idX.keys())]), + emp_d_cov=emp_cov + ) + self.step = at.aug_state(self.step, list(self.idX.keys())) else: From 918b1690f1560a085c73d8e99b60f82adf0351b9 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 25 Nov 2025 08:51:37 +0100 Subject: [PATCH 56/94] Improve readability of subspace_update --- pipt/misc_tools/analysis_tools.py | 6 +++ .../update_methods_ns/subspace_update.py | 51 ++++++++++++++++--- 2 files changed, 51 insertions(+), 6 deletions(-) diff --git a/pipt/misc_tools/analysis_tools.py b/pipt/misc_tools/analysis_tools.py index 4726dc4..6587860 100644 --- a/pipt/misc_tools/analysis_tools.py +++ b/pipt/misc_tools/analysis_tools.py @@ -1533,5 +1533,11 @@ def truncSVD(matrix, r=None, energy=None, full_matrices=False): r = np.sum((np.cumsum(S) / sum(S)) <= energy/100) else: raise ValueError("Either rank 'r' or 'energy' must be specified for truncSVD.") + + if r == 0: + r = 1 # Ensure at least one singular value is retained + if r > len(S): + print("Warning: Specified rank exceeds number of singular values. Using maximum available rank.") + r = len(S) return U[:,:r], S[:r], VT[:r,:] \ No newline at end of file diff --git a/pipt/update_schemes/update_methods_ns/subspace_update.py b/pipt/update_schemes/update_methods_ns/subspace_update.py index 0f3280b..616e2b2 100644 --- a/pipt/update_schemes/update_methods_ns/subspace_update.py +++ b/pipt/update_schemes/update_methods_ns/subspace_update.py @@ -18,17 +18,38 @@ class subspace_update(): Frontiers in Applied Mathematics and Statistics, 5(October), 114. https://doi.org/10.3389/fams.2019.00047 """ - def update(self): + def update(self, enX, enY, enE, **kwargs): + if self.iteration == 1: # method requires some initiallization self.current_W = np.zeros((self.ne, self.ne)) - self.E = np.dot(self.real_obs_data, self.proj) - Y = np.dot(self.aug_pred_data, self.proj) - # Y = self.pert_preddata + self.E = np.dot(enE, self.proj) + + # Center ensemble matrices + Y = np.dot(enY, self.proj) omega = np.eye(self.ne) + np.dot(self.current_W, self.proj) - LU = lu_factor(omega.T) - S = lu_solve(LU, Y.T).T + S = lu_solve(lu_factor(omega.T), Y.T).T + + # Compute scaled misfit (residual between predicted and observed data) + enRes = self.scale(enY - enE, self.scale_data) + + # Truncate SVD of S + Us, Ss, VsT = at.truncSVD(S, energy=self.trunc_energy) + Sinv = np.diag(1/Ss) + + # Compute update step + X = Sinv @ Us.T @ self.scale(self.E, self.scale_data) + eigval, eigvec = np.linalg.eig(X @ X.T) + X2 = Us @ Sinv.T @ eigvec + X3 = S.T @ X2 + + lam_term = np.eye(len(eigval)) + (1+self.lam) * np.diag(eigval) + deltaM = X3 @ solve(lam_term, X3.T @ self.current_W) + deltaD = X3 @ solve(lam_term, X2.T @ enRes) + self.w_step = -self.current_W/(1 + self.lam) - (deltaD - deltaM)/(1 + self.lam) + + '''' # scaled_misfit = (self.aug_pred_data - self.real_obs_data) if len(self.scale_data.shape) == 1: scaled_misfit = (self.scale_data ** (-1) @@ -73,3 +94,21 @@ def update(self): # solve((np.eye(len(Lam)) + (self.lam+1)*np.diag(Lam)), # np.dot(X2.T, scaled_misfit)))) self.w_step = -self.current_W/(1+self.lam) - (step_d - step_m/(1+self.lam)) + ''' + + def scale(self, data, scaling): + """ + Scale the data perturbations by the data error standard deviation. + + Args: + data (np.ndarray): data perturbations + scaling (np.ndarray): data error standard deviation + + Returns: + np.ndarray: scaled data perturbations + """ + + if len(scaling.shape) == 1: + return (scaling ** (-1))[:, None] * data + else: + return solve(scaling, data) From 07339f28bbebab9bf79aabc4c3213ae0a331b830 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 25 Nov 2025 10:02:32 +0100 Subject: [PATCH 57/94] Rewrite ESMDA --- pipt/update_schemes/esmda.py | 79 ++++++++++++------- .../update_methods_ns/subspace_update.py | 52 +----------- 2 files changed, 50 insertions(+), 81 deletions(-) diff --git a/pipt/update_schemes/esmda.py b/pipt/update_schemes/esmda.py index b9a2494..c794308 100644 --- a/pipt/update_schemes/esmda.py +++ b/pipt/update_schemes/esmda.py @@ -48,17 +48,18 @@ def __init__(self, keys_da, keys_en, sim): if self.restart is False: self.prior_enX = deepcopy(self.enX) self.list_states = list(self.idX.keys()) + # At the moment, the iterative loop is threated as an iterative smoother an thus we check if assim. indices # are given as in the Simultaneous loop. self.check_assimindex_simultaneous() self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] - self.list_datatypes, self.list_act_datatypes = at.get_list_data_types( - self.obs_data, self.assim_index) + self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(self.obs_data, self.assim_index) # Extract no. assimilation steps from MDA keyword in DATAASSIM part of init. file and set this equal to # the number of iterations pluss one. Need one additional because the iter=0 is the prior run. self.max_iter = len(self._ext_assim_steps())+1 self.iteration = 0 + self.lam = 0 # set LM lamda to zero as we are doing one full update. if 'energy' in self.keys_da: # initial energy (Remember to extract this) @@ -67,9 +68,11 @@ def __init__(self, keys_da, keys_en, sim): self.trunc_energy /= 100. else: self.trunc_energy = 0.98 + # Get the perturbed observations and observation scaling - self._ext_obs() - self.real_obs_data_conv = deepcopy(self.real_obs_data) + self.vecObs, self.enObs = self.set_observations() + self.enObs_conv = deepcopy(self.enObs) + # Get state scaling and svd of scaled prior self._ext_scaling() @@ -103,15 +106,25 @@ def calc_analysis(self): where $N_a$ being the total number of assimilation steps. """ - # Get assimilation order as a list - # reformat predicted data - _, self.aug_pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - self.list_datatypes) + # Get Ensemble of predicted data + _, self.enPred = at.aug_obs_pred_data( + self.obs_data, + self.pred_data, + self.assim_index, + self.list_datatypes + ) + + # Initialize GeoStat class for generating realizations + generator = Cholesky() - init_en = Cholesky() # Initialize GeoStat class for generating realizations if self.iteration == 1: # first iteration + + # Calculate the prior data misfit data_misfit = at.calc_objectivefun( - self.real_obs_data_conv, self.aug_pred_data, self.cov_data) + pert_obs=self.enObs, + pred_data=self.enPred, + Cd=self.cov_data + ) # Store the (mean) data misfit (also for conv. check) self.prior_data_misfit = np.mean(data_misfit) @@ -122,21 +135,24 @@ def calc_analysis(self): self.logger.info( f'Prior run complete with data misfit: {self.prior_data_misfit:0.1f}.') self.data_random_state = deepcopy(np.random.get_state()) - self.real_obs_data, self.scale_data = init_en.gen_real(self.obs_data_vector, - self.alpha[self.iteration-1] * - self.cov_data, self.ne, - return_chol=True) - self.E = np.dot(self.real_obs_data, self.proj) + + self.enObs, self.scale_data = generator.gen_real( + self.vecObs, + self.alpha[self.iteration - 1] * self.cov_data, + self.ne, + return_chol=True + ) + self.E = np.dot(self.enObs, self.proj) + else: self.data_random_state = deepcopy(np.random.get_state()) - self.obs_data_vector, _ = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - self.list_datatypes) - self.real_obs_data, self.scale_data = init_en.gen_real(self.obs_data_vector, - self.alpha[self.iteration - - 1] * self.cov_data, - self.ne, - return_chol=True) - self.E = np.dot(self.real_obs_data, self.proj) + self.enObs, self.scale_data = generator.gen_real( + self.vecObs, + self.alpha[self.iteration - 1] * self.cov_data, + self.ne, + return_chol=True + ) + self.E = np.dot(self.enObs, self.proj) if 'localanalysis' in self.keys_da: self.local_analysis_update() @@ -144,8 +160,8 @@ def calc_analysis(self): # Perform the update self.update( enX = self.enX, - enY = self.aug_pred_data, - enE = self.real_obs_data, + enY = self.enPred, + enE = self.enObs, prior = self.prior_enX ) @@ -178,12 +194,15 @@ def check_convergence(self): self.prev_data_misfit = self.data_misfit self.prev_data_misfit_std = self.data_misfit_std - # Prelude to calc. conv. check (everything done below is from calc_analysis) - obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - self.list_datatypes) + # Get Ensemble of predicted data + _, enPred = at.aug_obs_pred_data( + self.obs_data, + self.pred_data, + self.assim_index, + self.list_datatypes + ) - data_misfit = at.calc_objectivefun( - self.real_obs_data_conv, pred_data, self.cov_data) + data_misfit = at.calc_objectivefun(self.enObs_conv, enPred, self.cov_data) self.data_misfit = np.mean(data_misfit) self.data_misfit_std = np.std(data_misfit) diff --git a/pipt/update_schemes/update_methods_ns/subspace_update.py b/pipt/update_schemes/update_methods_ns/subspace_update.py index 616e2b2..54ccff3 100644 --- a/pipt/update_schemes/update_methods_ns/subspace_update.py +++ b/pipt/update_schemes/update_methods_ns/subspace_update.py @@ -1,10 +1,7 @@ """Stochastic iterative ensemble smoother (IES, i.e. EnRML) with *subspace* implementation.""" import numpy as np -from copy import deepcopy -import copy as cp -from scipy.linalg import solve, solve_banded, cholesky, lu_solve, lu_factor, inv -import pickle +from scipy.linalg import solve, lu_solve, lu_factor import pipt.misc_tools.analysis_tools as at @@ -49,53 +46,6 @@ def update(self, enX, enY, enE, **kwargs): self.w_step = -self.current_W/(1 + self.lam) - (deltaD - deltaM)/(1 + self.lam) - '''' - # scaled_misfit = (self.aug_pred_data - self.real_obs_data) - if len(self.scale_data.shape) == 1: - scaled_misfit = (self.scale_data ** (-1) - )[:, None] * (self.aug_pred_data - self.real_obs_data) - else: - scaled_misfit = solve( - self.scale_data, (self.aug_pred_data - self.real_obs_data)) - - u, s, v = np.linalg.svd(S, full_matrices=False) - if self.trunc_energy < 1: - ti = (np.cumsum(s) / sum(s)) <= self.trunc_energy - if sum(ti) == 0: - # the first singular value contains more than the prescibed trucation energy. - ti[0] = True - u, s, v = u[:, ti].copy(), s[ti].copy(), v[ti, :].copy() - - ps_inv = np.diag([el_s ** (-1) for el_s in s]) - # if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': - X = np.dot(ps_inv, np.dot(u.T, self.E)) - if len(self.scale_data.shape) == 1: - X = np.dot(ps_inv, np.dot(u.T, (self.scale_data ** (-1))[:, None]*self.E)) - else: - X = np.dot(ps_inv, np.dot(u.T, solve(self.scale_data, self.E))) - Lam, z = np.linalg.eig(np.dot(X, X.T)) - # else: - # X = np.dot(np.dot(ps_inv, np.dot(u.T, np.diag(self.cov_data))),np.dot(u,ps_inv)) - # Lam, z = np.linalg.eig(X) - # Lam = s**2 - # z = np.eye(len(s)) - - X2 = np.dot(u, np.dot(ps_inv.T, z)) - X3 = np.dot(S.T, X2) - - # X3_old = np.dot(X2, np.linalg.solve(np.eye(len(Lam)) + np.diag(Lam), X2.T)) - step_m = np.dot(X3, solve(np.eye(len(Lam)) + (1+self.lam) * - np.diag(Lam), np.dot(X3.T, self.current_W))) - - step_d = np.dot(X3, solve(np.eye(len(Lam)) + (1+self.lam) * - np.diag(Lam), np.dot(X2.T, scaled_misfit))) - - # step_d = np.dot(np.linalg.inv(omega).T, np.dot(np.dot(Y.T, X2), - # solve((np.eye(len(Lam)) + (self.lam+1)*np.diag(Lam)), - # np.dot(X2.T, scaled_misfit)))) - self.w_step = -self.current_W/(1+self.lam) - (step_d - step_m/(1+self.lam)) - ''' - def scale(self, data, scaling): """ Scale the data perturbations by the data error standard deviation. From f52ca5ce03446bef7e7a79e03df48cecd2874dba Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 25 Nov 2025 10:29:28 +0100 Subject: [PATCH 58/94] Rewrite EnKF for readability --- pipt/update_schemes/enkf.py | 116 +++++++++++++++++++++++------------- 1 file changed, 73 insertions(+), 43 deletions(-) diff --git a/pipt/update_schemes/enkf.py b/pipt/update_schemes/enkf.py index fa77b9c..090a3df 100644 --- a/pipt/update_schemes/enkf.py +++ b/pipt/update_schemes/enkf.py @@ -11,6 +11,7 @@ from pipt.loop.ensemble import Ensemble # Misc. tools used in analysis schemes from pipt.misc_tools import analysis_tools as at +import pipt.misc_tools.ensemble_tools as entools from pipt.update_schemes.update_methods_ns.approx_update import approx_update from pipt.update_schemes.update_methods_ns.full_update import full_update @@ -34,8 +35,9 @@ def __init__(self, keys_da, keys_en, sim): self.prev_data_misfit = None if self.restart is False: - self.prior_state = deepcopy(self.state) - self.list_states = list(self.state.keys()) + self.prior_enX = deepcopy(self.enX) + self.list_states = list(self.idX.keys()) + # At the moment, the iterative loop is threated as an iterative smoother an thus we check if assim. indices # are given as in the Simultaneous loop. self.check_assimindex_sequential() @@ -45,6 +47,7 @@ def __init__(self, keys_da, keys_en, sim): self.max_iter = len(self.keys_da['assimindex'])+1 self.iteration = 0 self.lam = 0 # set LM lamda to zero as we are doing one full update. + if 'energy' in self.keys_da: # initial energy (Remember to extract this) self.trunc_energy = self.keys_da['energy'] @@ -52,10 +55,12 @@ def __init__(self, keys_da, keys_en, sim): self.trunc_energy /= 100. else: self.trunc_energy = 0.98 - self.current_state = deepcopy(self.state) self.state_scaling = at.calc_scaling( - self.prior_state, self.list_states, self.prior_info) + self.prior_enX, + self.list_states, + self.prior_info + ) def calc_analysis(self): """ @@ -74,16 +79,27 @@ def calc_analysis(self): self.datavar, assim_index, list_datatypes) else: self.full_cov_data = self.cov_data - obs_data_vector, pred_data = at.aug_obs_pred_data( - self.obs_data, self.pred_data, assim_index, list_datatypes) + + #obs_data_vector, pred_data = at.aug_obs_pred_data( + # self.obs_data, self.pred_data, assim_index, list_datatypes) + + vecObs, enPred = at.aug_obs_pred_data( + self.obs_data, + self.pred_data, + assim_index, + list_datatypes + ) + # Generate realizations of the observed data - init_en = Cholesky() # Initialize GeoStat class for generating realizations - self.full_real_obs_data = init_en.gen_real( - obs_data_vector, self.full_cov_data, self.ne) + generator = Cholesky() # Initialize GeoStat class for generating realizations + self.enObs = generator.gen_real( + vecObs, + self.full_cov_data, + self.ne + ) # Calc. misfit for the initial iteration - data_misfit = at.calc_objectivefun( - self.full_real_obs_data, pred_data, self.full_cov_data) + data_misfit = at.calc_objectivefun(self.enObs, enPred, self.full_cov_data) # Store the (mean) data misfit (also for conv. check) self.data_misfit = np.mean(data_misfit) @@ -95,8 +111,7 @@ def calc_analysis(self): # Get assimilation order as a list # must subtract one to be inline - self.assim_index = [self.keys_da['obsname'], - self.keys_da['assimindex'][self.iteration-1]] + self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][self.iteration-1]] # Get list of data types to be assimilated and of the free states. Do this once, because listing keys from a # Python dictionary just when needed (in different places) may not yield the same list! @@ -104,57 +119,69 @@ def calc_analysis(self): self.obs_data, self.assim_index) # Augment observed and predicted data - self.obs_data_vector, self.aug_pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - self.list_datatypes) + self.vecObs, self.enPred = at.aug_obs_pred_data( + self.obs_data, + self.pred_data, + self.assim_index, + self.list_datatypes + ) + self.cov_data = at.gen_covdata( - self.datavar, self.assim_index, self.list_datatypes) + self.datavar, + self.assim_index, + self.list_datatypes + ) - init_en = Cholesky() # Initialize GeoStat class for generating realizations + generator = Cholesky() # Initialize GeoStat class for generating realizations self.data_random_state = deepcopy(np.random.get_state()) - self.real_obs_data, self.scale_data = init_en.gen_real(self.obs_data_vector, self.cov_data, self.ne, - return_chol=True) - - self.E = np.dot(self.real_obs_data, self.proj) + self.enObs, self.scale_data = generator.gen_real( + self.vecObs, + self.cov_data, + self.ne, + return_chol=True + ) + self.E = np.dot(self.enObs, self.proj) if 'localanalysis' in self.keys_da: self.local_analysis_update() else: - # Mean pred_data and perturbation matrix with scaling - if len(self.scale_data.shape) == 1: - self.pert_preddata = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), - np.ones((1, self.ne))) * np.dot(self.aug_pred_data, self.proj) - else: - self.pert_preddata = solve( - self.scale_data, np.dot(self.aug_pred_data, self.proj)) - - aug_state = at.aug_state(self.current_state, self.list_states) - self.update() + self.update( + enX = self.enX, + enY = self.enPred, + enE = self.enObs, + prior = self.prior_enX + ) + # Update the state ensemble and weights if hasattr(self, 'step'): - aug_state_upd = aug_state + self.step + self.enX_temp = self.enX + self.step if hasattr(self, 'w_step'): self.W = self.current_W + self.w_step - aug_prior_state = at.aug_state(self.prior_state, self.list_states) - aug_state_upd = np.dot(aug_prior_state, (np.eye( - self.ne) + self.W / np.sqrt(self.ne - 1))) - # Extract updated state variables from aug_update - self.state = at.update_state(aug_state_upd, self.state, self.list_states) - self.state = at.limits(self.state, self.prior_info) + self.enX_temp = np.dot(self.prior_enX, (np.eye(self.ne) + self.W/np.sqrt(self.ne - 1))) + + # Ensure limits are respected + limits = {key: self.prior_info[key].get('limits', (None, None)) for key in self.idX.keys()} + self.enX_temp = entools.clip_matrix(self.enX_temp, limits, self.idX) def check_convergence(self): """ Calculate the "convergence" of the method. Important to """ self.prev_data_misfit = self.prior_data_misfit + # only calulate for the final (posterior) estimate if self.iteration == len(self.keys_da['assimindex']): assim_index = [self.keys_da['obsname'], list( np.concatenate(self.keys_da['assimindex']))] list_datatypes = self.list_datatypes - obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, assim_index, - list_datatypes) - data_misfit = at.calc_objectivefun( - self.full_real_obs_data, pred_data, self.full_cov_data) + _, enPred = at.aug_obs_pred_data( + self.obs_data, + self.pred_data, + assim_index, + list_datatypes + ) + + data_misfit = at.calc_objectivefun(self.enObs, enPred, self.full_cov_data) self.data_misfit = np.mean(data_misfit) self.data_misfit_std = np.std(data_misfit) @@ -166,7 +193,10 @@ def check_convergence(self): 'data_misfit': self.data_misfit, 'prev_data_misfit': self.prev_data_misfit} - self.current_state = deepcopy(self.state) + # Update state ensemble + self.enX = deepcopy(self.enX_temp) + self.enX_temp = None + if self.data_misfit == self.prev_data_misfit: self.logger.info( f'EnKF update {self.iteration} complete!') From 78edbf3b724a1c0c6e8616a3c406035aa2305cfd Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 25 Nov 2025 10:44:45 +0100 Subject: [PATCH 59/94] Use truncSVD function --- pipt/update_schemes/es.py | 3 ++- pipt/update_schemes/update_methods_ns/full_update.py | 5 +---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/pipt/update_schemes/es.py b/pipt/update_schemes/es.py index c3a2bf5..b292cbe 100644 --- a/pipt/update_schemes/es.py +++ b/pipt/update_schemes/es.py @@ -67,7 +67,8 @@ def check_convergence(self): if self.data_misfit == self.prev_data_misfit: self.logger.info( f'ES update {self.iteration} complete!') - self.current_state = deepcopy(self.state) + self.enX = deepcopy(self.enX_temp) + self.enX_temp = None else: if self.data_misfit < self.prior_data_misfit: self.logger.info( diff --git a/pipt/update_schemes/update_methods_ns/full_update.py b/pipt/update_schemes/update_methods_ns/full_update.py index 709096f..406081c 100644 --- a/pipt/update_schemes/update_methods_ns/full_update.py +++ b/pipt/update_schemes/update_methods_ns/full_update.py @@ -31,10 +31,7 @@ def update(self, enX, enY, enE, **kwargs): enXcentered = self.scale(np.dot(enX, self.proj), self.state_scaling) # Perform tuncated SVD - u_d, s_d, v_d = np.linalg.svd(enYcentered, full_matrices=False) - if self.trunc_energy < 1: - ti = (np.cumsum(s_d) / sum(s_d)) <= self.trunc_energy - u_d, s_d, v_d = u_d[:, ti].copy(), s_d[ti].copy(), v_d[ti, :].copy() + u_d, s_d, v_d = at.truncSVD(enYcentered, energy=self.trunc_energy) # Compute the update step x_1 = np.dot(u_d.T, self.scale(enE - enY, self.scale_data)) From b0770096ed76d95238eca5650f4c7575df363e3b Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 8 Dec 2025 11:00:50 +0100 Subject: [PATCH 60/94] Rewrite local analysis for region params --- pipt/loop/ensemble.py | 57 +++++++++++++++++++++++++++++++------------ 1 file changed, 41 insertions(+), 16 deletions(-) diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index 0d9e01b..f8a68d4 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -675,41 +675,64 @@ def local_analysis_update(self): Function for updates that can be used by all algorithms. Do this once to avoid duplicate code for local analysis. ''' + # Copy original info to restore after local updates orig_list_data = deepcopy(self.list_datatypes) orig_list_state = deepcopy(self.list_states) orig_cd = deepcopy(self.cov_data) orig_real_obs_data = deepcopy(self.real_obs_data) orig_data_vector = deepcopy(self.obs_data_vector) + # loop over the states that we want to update. Assume that the state and data combinations have been # determined by the initialization. # TODO: augment parameters with identical mask. + + # REGION PARAMETERS + ############################################################################################################ for state in self.local_analysis['region_parameter']: - self.list_datatypes = [elem for elem in self.list_datatypes if - elem in self.local_analysis['update_mask'][state]] + self.list_datatypes = [ + elem for elem in self.list_datatypes if + elem in self.local_analysis['update_mask'][state] + ] self.list_states = [deepcopy(state)] + self._ext_scaling() # scaling for this state if 'localization' in self.keys_da: self.localization.loc_info['field'] = self.state_scaling.shape del self.cov_data + # reset the random state for consistency np.random.set_state(self.data_random_state) - self._ext_obs() # get the data that's in the list of data. - _, self.aug_pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - self.list_datatypes) - # Mean pred_data and perturbation matrix with scaling - if len(self.scale_data.shape) == 1: - self.pert_preddata = np.dot(np.expand_dims(self.scale_data ** (-1), axis=1), - np.ones((1, self.ne))) * np.dot(self.aug_pred_data, self.proj) - else: - self.pert_preddata = solve( - self.scale_data, np.dot(self.aug_pred_data, self.proj)) + self.vecObs, self.enObs = self.set_observations() + _, self.enPred = at.aug_obs_pred_data( + self.obs_data, + self.pred_data, + self.assim_index, + self.list_datatypes + ) + + # Get state ensemble for list_states + enX = [] + idX = {} + for idx in self.list_states: + start, end = self.idX[idx] + tempX = self.enX[start:end, :] + enX.append(tempX) + idX[idx] = (enX.shape[0] - tempX.shape[0], enX.shape[0]) + + # Compute the analysis update + self.update( + enX = np.vstack(enX), + enY = self.enPred, + enE = self.enObs, + ) - aug_state = at.aug_state(self.current_state, self.list_states) - self.update() + # Update the state if hasattr(self, 'step'): - aug_state_upd = aug_state + self.step - self.state = at.update_state(aug_state_upd, self.state, self.list_states) + self.enX_temp = self.enX + self.step + ############################################################################################################ + # VECTOR REGION PARAMETERS + ############################################################################################################ for state in self.local_analysis['vector_region_parameter']: current_list_datatypes = deepcopy(self.list_datatypes) for state_indx in range(self.state[state].shape[0]): # loop over the elements in the region @@ -741,6 +764,8 @@ def local_analysis_update(self): self.state[state][state_indx,:] = aug_state_upd self.list_datatypes = deepcopy(current_list_datatypes) + ############################################################################################################ + for state in self.local_analysis['cell_parameter']: self.list_states = [deepcopy(state)] From d13d0f29097fb56e54335380380ada66f00313f6 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 9 Dec 2025 08:58:59 +0100 Subject: [PATCH 61/94] Update the progressbar --- ensemble/ensemble.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 60fc445..2e9de83 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -22,7 +22,15 @@ import pipt.misc_tools.ensemble_tools as entools from misc.system_tools.environ_var import OpenBlasSingleThread # Single threaded OpenBLAS runs - +# Settings +################################################################################################ +progbar_settings = { + 'desc': 'Progress', + 'ncols': 100, + 'colour': '#305069', + 'bar_format': '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]' +} +################################################################################################ class Ensemble: """ @@ -242,7 +250,8 @@ def calc_prediction(self, enX=None, save_prediction=None): # No parralelization if nparallel==1: en_pred = [] - for member_index, state in tqdm(enumerate(enX), total=self.ne, desc="Running simulations"): + pbar = tqdm(enumerate(enX), total=self.ne, **progbar_settings) + for member_index, state in pbar: en_pred.append(self.sim.run_fwd_sim(state, member_index)) # Parallelization on HPC using SLURM @@ -256,7 +265,8 @@ def calc_prediction(self, enX=None, save_prediction=None): enX, list(range(self.ne)), num_cpus=nparallel, - disable=self.disable_tqdm + disable=self.disable_tqdm, + **progbar_settings ) ###################################################################################################################### From ded9473b68643b63f25f1832acbb85fb5644a469 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 9 Dec 2025 09:20:20 +0100 Subject: [PATCH 62/94] Change style of logger --- ensemble/ensemble.py | 2 +- popt/loop/optimize.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 2e9de83..bb46885 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -25,7 +25,7 @@ # Settings ################################################################################################ progbar_settings = { - 'desc': 'Progress', + #'desc': ' Progress', 'ncols': 100, 'colour': '#305069', 'bar_format': '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]' diff --git a/popt/loop/optimize.py b/popt/loop/optimize.py index 35ccaa4..72c8529 100644 --- a/popt/loop/optimize.py +++ b/popt/loop/optimize.py @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) # set log level file_handler = logging.FileHandler('popt.log') # define file handler and set formatter -formatter = logging.Formatter('%(asctime)s : %(levelname)s : %(name)s : %(message)s') +formatter = logging.Formatter('%(asctime)s : %(message)s') file_handler.setFormatter(formatter) logger.addHandler(file_handler) # add file handler to logger console_handler = logging.StreamHandler() From 0b4b9fd31ea18673d7851b6f3d4a86723e9c26d2 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 9 Dec 2025 10:58:32 +0100 Subject: [PATCH 63/94] Add alternative way of defining initial controls for optimization --- ensemble/ensemble.py | 6 +- input_output/read_config.py | 2 +- pipt/misc_tools/extract_tools.py | 158 ++++++++++++++++++++++++++++++- 3 files changed, 163 insertions(+), 3 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index bb46885..02cf470 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -131,8 +131,12 @@ def __init__(self, keys_en: dict, sim, redund_sim=None): self.disable_tqdm = False # extract information that is given for the prior model - self.prior_info = extract.extract_prior_info(self.keys_en) + if 'state' in self.keys_en: + self.prior_info = extract.extract_prior_info(self.keys_en) + elif 'controls' in self.keys_en: + self.prior_info = extract.extract_initial_controls(self.keys_en) + # Calculate initial ensemble if IMPORTSTATICVAR has not been given in init. file. # Prior info. on state variables must be given by PRIOR_ keyword. if 'importstaticvar' not in self.keys_en: diff --git a/input_output/read_config.py b/input_output/read_config.py index ccc6e22..c6af002 100644 --- a/input_output/read_config.py +++ b/input_output/read_config.py @@ -392,7 +392,7 @@ def check_mand_keywords_en(keys_en): # Mandatory keywords in ENSEMBLE assert 'ne' in keys_en, 'NE not in ENSEMBLE!' - assert 'state' in keys_en, 'STATE not in ENSEMBLE!' + assert ('state' in keys_en) or ('controls' in keys_en), 'STATE or CONTROLS not in ENSEMBLE!' if 'importstaticvar' not in keys_en: assert filter(list(keys_en.keys()), 'prior_*') != [], 'No PRIOR_ in DATAASSIM' diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index 032057b..c5965f3 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -10,7 +10,8 @@ ] # Imports -import numpy as np +import numpy as np +import pandas as pd import pickle import os @@ -126,6 +127,161 @@ def extract_prior_info(keys: dict) -> dict: return prior_info +def extract_initial_controls(keys: dict) -> dict: + """ + Extract and process control variable information from configuration dictionary. + + This function parses control variable specifications from the input configuration, + handling various formats for initial values, bounds, and variance. + It supports loading data from files (.npy, .npz, .csv). + + Parameters + ---------- + keys : dict + Configuration dictionary containing a 'controls' key. Each control variable + should be a nested dictionary with the name of the control variable as the key. + The dictionary for each control variable should contain the following possible keys: + + - 'initial' or 'mean' : Initial value or mean of control variable + Can be scalar, list, numpy array, or filename (.npy, .npz, .csv). + If .npz or .csv, the variable name should match the control variable name. + Multiple variables can be specified in the same file. + + - 'limits' : tuple or list, optional + (lower_bound, upper_bound) for the control variable + + - 'var' or 'variance' : float, list, or array, optional + Variance of the control variable + + - 'std' : float, list, array, or str, optional + Standard deviation. If string ending with '%', interpreted as percentage + of the bound range (requires 'limits' to be specified). Only if 'var'/'variance' + is not provided. + + Returns + ------- + control_info : dict + Dictionary with control variable names as keys. Each value is a dict containing: + + - 'mean' : numpy.ndarray + Initial/mean values for the control variable + - 'limits' : list + [lower_bound, upper_bound], or [None, None] if not specified + - 'variance' : float, numpy.ndarray, or None + Variance of the control variable (if provided) + + Raises + ------ + AssertionError + If neither 'initial' nor 'mean' is provided for a control variable + If attempting to use percentage-based 'std' without specifying 'limits' + If loading from file fails (e.g., variable name not found in file) + + Examples + -------- + >>> keys = { + ... 'controls': { + ... 'pressure': { + ... 'initial': 100.0, + ... 'limits': [50.0, 150.0], + ... 'std': '10%' + ... }, + ... 'rate': { + ... 'mean': [10, 20, 30], + ... 'variance': 2.5 + ... } + ... } + ... } + >>> control_info = extract_initial_controls(keys) + >>> control_info['pressure']['mean'] + array([100.]) + >>> control_info['pressure']['variance'] + 100.0 # (10% of range [50, 150])^2 + """ + control_info = {} + + # Loop over names + for name in keys['controls'].keys(): + info = keys['controls'][name] + + # Assert that initial or mean is there + assert ('initial' in info) or ('mean' in info), f'INITIAL or MEAN missing in CONTROLS for {name}!' + + # Rename to mean if initial is there + if 'initial' in info: + info['mean'] = info.pop('initial', None) + + # Mean + ############################################################################################################ + if isinstance(info['mean'], str): + # Check if NPZ file + if info['mean'].endswith('.npz'): + file = np.load(info['mean'], allow_pickle=True) + if not (name in file.files): + # Assume only one variable in file + msg = f'Variable {name} not in {info["mean"]} and more than one variable located in the file!' + assert len(file.files) == 1, msg + info['mean'] = file[file.files[0]] + else: + info['mean'] = file[name] + + # Check for NPY file + elif info['mean'].endswith('.npy'): + info['mean'] = np.load(info['mean']) + + # Check for CSV file + elif info['mean'].endswith('.csv'): + df = pd.read_csv(info['mean']) + assert name in df.columns, f'Column {name} not in {info["mean"]}!' + info['mean'] = df[name].to_numpy() + + elif isinstance(info['mean'], (int, float)): + info['mean'] = np.array([info['mean']]) + else: + info['mean'] = np.asarray(info['mean']) + ############################################################################################################ + + # Limits + info['limits'] = info.get('limits', [None, None]) + + # Clip mean to limits if limits are given + if info['limits'][0] is not None: + info['mean'] = np.maximum(info['mean'], info['limits'][0]) + if info['limits'][1] is not None: + info['mean'] = np.minimum(info['mean'], info['limits'][1]) + + + # Check for var VAR or STD + ############################################################################################################ + if ('var' in info) or ('variance' in info): + if 'var' in info: + info['variance'] = info.pop('var', None) + + elif 'std' in info: + std = info.pop('std', None) + + # Standard deviation can be given as percentage of bound range + if isinstance(std, str) and (info['limits'][0] is not None) and (info['limits'][1] is not None): + if std.endswith('%'): + std, _ = std.split('%') + std = float(std)/100.0 * (info['limits'][1] - info['limits'][0]) + else: + raise AssertionError(f'If STD for {name} does not end with %') + + info['variance'] = np.square(std) + ############################################################################################################ + + # Add control_info + control_info[name] = info + + return control_info + + + + + + + def extract_multilevel_info(keys: Union[dict, list]) -> dict: ''' Extract the info needed for ML simulations. Note if the ML keyword is not in keys_en we initialize From ada42352b4aadf8cdfe56c3ce84d55763cf74ada Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 9 Dec 2025 11:22:29 +0100 Subject: [PATCH 64/94] Update style of progressbar --- ensemble/ensemble.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 02cf470..69f6a94 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -23,14 +23,15 @@ from misc.system_tools.environ_var import OpenBlasSingleThread # Single threaded OpenBLAS runs # Settings -################################################################################################ +####################################################################################################### progbar_settings = { - #'desc': ' Progress', 'ncols': 100, - 'colour': '#305069', - 'bar_format': '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]' + 'colour': "#285475", + 'bar_format': '{percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]', + 'ascii': '-◼', # Custom bar characters for a sleeker look + 'unit': 'member', } -################################################################################################ +####################################################################################################### class Ensemble: """ From afcecd8ee7f33c0b6a0450d481b70a5ae41a7257 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 12 Dec 2025 10:29:18 +0100 Subject: [PATCH 65/94] Update logger for ESMDA and ENRML --- ensemble/ensemble.py | 12 +++++---- pipt/loop/assimilation.py | 14 +++++----- pipt/loop/ensemble.py | 5 ++-- pipt/misc_tools/extract_tools.py | 1 + pipt/update_schemes/enrml.py | 44 +++++++++++++++++++++++--------- pipt/update_schemes/esmda.py | 38 +++++++++++++++++++++------ 6 files changed, 80 insertions(+), 34 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 69f6a94..9a16bd2 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -25,7 +25,7 @@ # Settings ####################################################################################################### progbar_settings = { - 'ncols': 100, + 'ncols': 110, 'colour': "#285475", 'bar_format': '{percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]', 'ascii': '-◼', # Custom bar characters for a sleeker look @@ -71,10 +71,12 @@ def __init__(self, keys_en: dict, sim, redund_sim=None): # Setup logger logging.basicConfig( level=logging.INFO, - filename='pet_logger.log', - filemode='w', - format='%(asctime)s : %(levelname)s : %(name)s : %(message)s', - datefmt='%Y-%m-%d %H:%M:%S' + format='%(asctime)s : %(levelname)s : %(message)s', + datefmt='%Y-%m-%d %H:%M:%S', + handlers=[ + logging.FileHandler('pet_logger.log', mode='w'), + logging.StreamHandler() + ] ) self.logger = logging.getLogger('PET') diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index 39acefb..ee38e97 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -86,7 +86,7 @@ def run(self): success_iter = True # Initiallize progressbar - pbar_out = tqdm(total=self.max_iter, desc='Iterations (Obj. func. val: )', position=0) + #pbar_out = tqdm(total=self.max_iter, desc='Iterations (Obj. func. val: )', position=0) # Check if we want to perform a Quality Assurance of the forecast qaqc = None @@ -193,16 +193,16 @@ def run(self): if self.ensemble.iteration >= 0 and success_iter is True: if self.ensemble.iteration == 0: self.ensemble.iteration += 1 - pbar_out.update(1) + #pbar_out.update(1) # pbar_out.set_description(f'Iterations (Obj. func. val:{self.data_misfit:.1f})') # self.prior_data_misfit = self.data_misfit # self.pbar_out.refresh() else: self.ensemble.iteration += 1 - pbar_out.update(1) - pbar_out.set_description( - f'Iterations (Obj. func. val:{self.ensemble.data_misfit:.1f}' - f' Reduced: {100 * (1 - (self.ensemble.data_misfit / self.ensemble.prev_data_misfit)):.0f} %)') + #pbar_out.update(1) + #pbar_out.set_description( + # f'Iterations (Obj. func. val:{self.ensemble.data_misfit:.1f}' + # f' Reduced: {100 * (1 - (self.ensemble.data_misfit / self.ensemble.prev_data_misfit)):.0f} %)') # self.pbar_out.refresh() if 'restartsave' in self.ensemble.keys_da and self.ensemble.keys_da['restartsave'] == 'yes': @@ -235,7 +235,7 @@ def run(self): with open('why_iter_loop_stopped.p', 'wb') as f: pickle.dump(why, f, protocol=4) # pbar.close() - pbar_out.close() + #pbar_out.close() if self.ensemble.prev_data_misfit is not None: out_str = 'Convergence was met.' if self.ensemble.prior_data_misfit > self.ensemble.data_misfit: diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index f8a68d4..83cfd40 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -76,8 +76,9 @@ def __init__(self, keys_da, keys_en, sim): self.logger = logging.getLogger('PET.PIPT') # write initial information - self.logger.info(f'Starting a {keys_da["daalg"][0]} run with the {keys_da["daalg"][1]} algorithm applying the ' - f'{keys_da["analysis"]} update scheme with {keys_da["energy"]} Energy.') + self.logger.info('') + self.logger.info(f' =========== Running Data Assimilation - {keys_da["daalg"][0].upper()} ===========') + self.logger.info('') # Internalize PIPT dictionary if not hasattr(self, 'keys_da'): diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index c5965f3..06865de 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -2,6 +2,7 @@ __all__ = [ 'extract_prior_info', + 'extract_initial_controls', 'extract_multilevel_info', 'extract_local_analysis_info', 'extract_maxiter', diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index 2bed3b8..6227757 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -139,8 +139,8 @@ def calc_analysis(self): if self.lam == 'auto': self.lam = (0.5 * self.prior_data_misfit)/self.enPred.shape[0] - self.logger.info( - f'Prior run complete with data misfit: {self.prior_data_misfit:0.1f}. Lambda for initial analysis: {self.lam}') + # Log initial data misfit + self.log_update(success=True, prior_run=True) if 'localanalysis' in self.keys_da: self.local_analysis_update() @@ -219,13 +219,16 @@ def check_convergence(self): if self.data_misfit >= self.prev_data_misfit: success = False + self.logger.info('') self.logger.info( f'Iterations have converged after {self.iteration} iterations. Objective function reduced ' f'from {self.prior_data_misfit:0.1f} to {self.prev_data_misfit:0.1f}') else: + self.logger.info('') self.logger.info( f'Iterations have converged after {self.iteration} iterations. Objective function reduced ' f'from {self.prior_data_misfit:0.1f} to {self.data_misfit:0.1f}') + # Return conv = True, why_stop var. return True, success, why_stop @@ -273,16 +276,10 @@ def check_convergence(self): self.lam = self.lam * self.gamma success = False - if success: - self.logger.info(f'Successfull iteration number {self.iteration}! Objective function reduced from ' - f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}. New Lamba for next analysis: ' - f'{self.lam}') - # self.prev_data_misfit = self.data_misfit - # self.prev_data_misfit_std = self.data_misfit_std - else: - self.logger.info(f'Failed iteration number {self.iteration}! Objective function increased from ' - f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}. New Lamba for repeated analysis: ' - f'{self.lam}') + # Log update results + self.log_update(success=success) + + if not success: # Reset the objective function after report self.data_misfit = self.prev_data_misfit self.data_misfit_std = self.prev_data_misfit_std @@ -290,6 +287,29 @@ def check_convergence(self): # Return conv = False, why_stop var. return False, success, why_stop + def log_update(self, success, prior_run=False): + ''' + Log the update results in a formatted table. + ''' + def _log_table(iteration, status, misfit, change_pct, lambda_val, change_label="Reduction"): + """Helper method to log iteration results in a formatted table.""" + self.logger.info('') + self.logger.info(f' {"Iteration":<11}| {"Status":<11}| {"Data Misfit":<16}| {f"{change_label} (%)":<15}| {"λ":<10}') + self.logger.info(f' {"—"*11}|{"—"*12}|{"—"*17}|{"—"*16}|{"—"*10}') + self.logger.info(f' {iteration:<11}| {status:<11}| {misfit:<16.2f}| {change_pct:<15.2f}| {lambda_val:<10.2f}') + self.logger.info('') + + if prior_run: + _log_table(0, "Success", self.data_misfit, 0.0, self.lam) + elif success: + reduction = 100 * (1 - self.data_misfit / self.prev_data_misfit) + _log_table(self.iteration, "Success", self.data_misfit, reduction, self.lam) + else: + increase = 100 * (1 - self.prev_data_misfit / self.data_misfit) + _log_table(self.iteration, "Failed", self.data_misfit, increase, self.lam, "Increase") + + + class lmenrml_approx(lmenrmlMixIn, approx_update): pass diff --git a/pipt/update_schemes/esmda.py b/pipt/update_schemes/esmda.py index c794308..142df35 100644 --- a/pipt/update_schemes/esmda.py +++ b/pipt/update_schemes/esmda.py @@ -132,8 +132,8 @@ def calc_analysis(self): self.data_misfit = np.mean(data_misfit) self.data_misfit_std = np.std(data_misfit) - self.logger.info( - f'Prior run complete with data misfit: {self.prior_data_misfit:0.1f}.') + # Log initial data misfit + self.log_update(prior_run=True) self.data_random_state = deepcopy(np.random.get_state()) self.enObs, self.scale_data = generator.gen_real( @@ -211,12 +211,10 @@ def check_convergence(self): 'data_misfit': self.data_misfit, 'prev_data_misfit': self.prev_data_misfit} - if self.data_misfit < self.prev_data_misfit: - self.logger.info( - f'MDA iteration number {self.iteration}! Objective function reduced from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') - else: - self.logger.info( - f'MDA iteration number {self.iteration}! Objective function increased from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') + # Log update results + success = self.data_misfit < self.prev_data_misfit + self.log_update(success=success) + # Return conv = False, why_stop var. # Update state ensemble self.enX = deepcopy(self.enX_temp) @@ -226,6 +224,30 @@ def check_convergence(self): return False, True, why_stop + def log_update(self, success=None, prior_run=False): + ''' + Log the update results in a formatted table. + ''' + def _log_table(iteration, status, misfit, change_pct, change_label="Reduction"): + """Helper method to log iteration results in a formatted table.""" + self.logger.info('') + self.logger.info(f' {"Iteration":<11}| {"Status":<11}| {"Data Misfit":<16}| {f"{change_label} (%)":<15}') + self.logger.info(f' {"—"*11}|{"—"*12}|{"—"*17}|{"—"*16}') + self.logger.info(f' {iteration:<11}| {status:<11}| {misfit:<16.2f}| {change_pct:<15.2f}') + self.logger.info('') + + if prior_run: + iteration_str = f'0/{self.max_iter}' + _log_table(iteration_str, "Success", self.data_misfit, 0.0) + elif success: + iteration_str = f'{self.iteration}/{self.max_iter}' + reduction = 100 * (1 - self.data_misfit / self.prev_data_misfit) + _log_table(iteration_str, "Success", self.data_misfit, reduction) + else: + iteration_str = f'{self.iteration}/{self.max_iter}' + increase = 100 * (1 - self.prev_data_misfit / self.data_misfit) + _log_table(iteration_str, "Failed", self.data_misfit, increase, "Increase") + def _ext_inflation_param(self): r""" Extract the data covariance inflation parameter from the MDA keyword in DATAASSIM part. Also, we check that From b864d197e12e5ee4d773bae3f9c9e85afdcab88f Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 12 Dec 2025 12:27:37 +0100 Subject: [PATCH 66/94] Add savefolder option to Assimilation --- pipt/loop/assimilation.py | 22 +++++++++++++++------- pipt/misc_tools/analysis_tools.py | 5 +++-- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index ee38e97..e18d447 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -48,6 +48,12 @@ def __init__(self, ensemble: Ensemble): # Internalize ensemble and simulator class instances self.ensemble = ensemble + # Save folder + if 'nosave' not in self.ensemble.keys_da: + self.save_folder = self.ensemble.keys_da.get('savefolder', 'SaveOutputs') + if not os.path.exists(self.save_folder): + os.makedirs(self.save_folder) + if self.ensemble.restart is False: # Default max. iter if not defined in the ensemble if hasattr(ensemble, 'max_iter'): @@ -131,7 +137,7 @@ def run(self): # always store prior forcast, unless specifically told not to if 'nosave' not in self.ensemble.keys_da: - np.savez('prior_forecast.npz', pred_data=self.ensemble.pred_data) + np.savez(f'{self.save_folder}/prior_forecast.npz', pred_data=self.ensemble.pred_data) # For the remaining iterations we start by applying the analysis and finish by running the forecast else: @@ -211,12 +217,12 @@ def run(self): # always store posterior forcast and state, unless specifically told not to if 'nosave' not in self.ensemble.keys_da: try: # first try to save as npz file - np.savez('posterior_state_estimate.npz', **self.ensemble.enX) - np.savez('posterior_forecast.npz', **{'pred_data': self.ensemble.pred_data}) + np.savez(f'{self.save_folder}/posterior_state_estimate.npz', **self.ensemble.enX) + np.savez(f'{self.save_folder}/posterior_forecast.npz', **{'pred_data': self.ensemble.pred_data}) except: # If this fails, store as pickle - with open('posterior_state_estimate.p', 'wb') as file: + with open(f'{self.save_folder}/posterior_state_estimate.p', 'wb') as file: pickle.dump(self.ensemble.enX, file) - with open('posterior_forecast.p', 'wb') as file: + with open(f'{self.save_folder}/posterior_forecast.p', 'wb') as file: pickle.dump(self.ensemble.pred_data, file) # If none of the convergence criteria were met, max. iteration was the reason iterations stopped. @@ -232,7 +238,7 @@ def run(self): why = self.why_stop if why is not None: why['conv_string'] = reason - with open('why_iter_loop_stopped.p', 'wb') as f: + with open(f'{self.save_folder}/why_iter_loop_stopped.p', 'wb') as f: pickle.dump(why, f, protocol=4) # pbar.close() #pbar_out.close() @@ -349,6 +355,8 @@ def _save_analysis_debug(self): else: print(f'Cannot save {save_typ}, because it is a local variable!\n\n') + save_dict['savefolder'] = self.save_folder + # Save the variables at.save_analysisdebug(self.ensemble.iteration, **save_dict) @@ -451,7 +459,7 @@ def calc_forecast(self): # Extra option debug if 'saveforecast' in self.ensemble.sim.input_dict: - with open('sim_results.p', 'wb') as f: + with open(f'{self.save_folder}/sim_results.p', 'wb') as f: pickle.dump(self.ensemble.pred_data, f) def post_process_forecast(self): diff --git a/pipt/misc_tools/analysis_tools.py b/pipt/misc_tools/analysis_tools.py index 6587860..b48f290 100644 --- a/pipt/misc_tools/analysis_tools.py +++ b/pipt/misc_tools/analysis_tools.py @@ -661,10 +661,11 @@ def save_analysisdebug(ind_save, **kwargs): is passed to np.savez (kwargs) the variable will be stored with their original name. """ # Save input variables + folder = kwargs.pop('savefolder') try: - np.savez('debug_analysis_step_{0}'.format(str(ind_save)), **kwargs) + np.savez(f'{folder}/debug_analysis_step_{ind_save}', **kwargs) except: # if npz save fails dump to a pickle file - with open(f'debug_analysis_step_{ind_save}.p', 'wb') as file: + with open(f'{folder}/debug_analysis_step_{ind_save}.p', 'wb') as file: pickle.dump(kwargs, file) From cf63b5687795fc665ce0fcab884b64e62baf03cd Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 12 Dec 2025 15:03:08 +0100 Subject: [PATCH 67/94] Fix logging bug --- ensemble/ensemble.py | 11 ----------- pipt/loop/ensemble.py | 13 +++++++++++-- popt/loop/optimize.py | 24 +++++++++++------------- simulator/eclipse.py | 6 +++--- 4 files changed, 25 insertions(+), 29 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 9a16bd2..dc9146d 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -68,17 +68,6 @@ def __init__(self, keys_en: dict, sim, redund_sim=None): # to allow for different models when optimizing. self.aux_input = None - # Setup logger - logging.basicConfig( - level=logging.INFO, - format='%(asctime)s : %(levelname)s : %(message)s', - datefmt='%Y-%m-%d %H:%M:%S', - handlers=[ - logging.FileHandler('pet_logger.log', mode='w'), - logging.StreamHandler() - ] - ) - self.logger = logging.getLogger('PET') # Check if folder contains any En_ files, and remove them! for folder in glob('En_*'): diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index 83cfd40..9231980 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -72,8 +72,17 @@ def __init__(self, keys_da, keys_en, sim): # do the initiallization of the PETensemble super(Ensemble, self).__init__(keys_da|keys_en, sim) - # set logger - self.logger = logging.getLogger('PET.PIPT') + # Setup logger + logging.basicConfig( + level=logging.INFO, + format='%(asctime)s : %(levelname)s : %(message)s', + datefmt='%Y-%m-%d %H:%M:%S', + handlers=[ + logging.FileHandler('assim.log', mode='w'), + logging.StreamHandler() + ] + ) + self.logger = logging.getLogger(__name__) # write initial information self.logger.info('') diff --git a/popt/loop/optimize.py b/popt/loop/optimize.py index 72c8529..4d8c6a2 100644 --- a/popt/loop/optimize.py +++ b/popt/loop/optimize.py @@ -8,17 +8,6 @@ # Internal imports import popt.misc_tools.optim_tools as ot -# Gets or creates a logger -logger = logging.getLogger(__name__) -logger.setLevel(logging.DEBUG) # set log level -file_handler = logging.FileHandler('popt.log') # define file handler and set formatter -formatter = logging.Formatter('%(asctime)s : %(message)s') -file_handler.setFormatter(formatter) -logger.addHandler(file_handler) # add file handler to logger -console_handler = logging.StreamHandler() -console_handler.setFormatter(formatter) -logger.addHandler(console_handler) - class Optimize: """ @@ -72,8 +61,17 @@ def __init__(self, **options): options : dict Optimization options """ - # Set the logger - self.logger = logger + # Setup logger + logging.basicConfig( + level=logging.INFO, + format='%(asctime)s : %(levelname)s : %(message)s', + datefmt='%Y-%m-%d %H:%M:%S', + handlers=[ + logging.FileHandler('optim.log', mode='w'), + logging.StreamHandler() + ] + ) + self.logger = logging.getLogger(__name__) # Save name for (potential) pickle dump/load self.pickle_restart_file = 'popt_restart_dump' diff --git a/simulator/eclipse.py b/simulator/eclipse.py index c591d63..0d2b85c 100644 --- a/simulator/eclipse.py +++ b/simulator/eclipse.py @@ -112,12 +112,12 @@ def _extInfoInputDict(self): # In the ecl framework, all reference to the filename should be uppercase self.file = self.input_dict['runfile'].upper() - + # Extract sim options - if isinstance(self.input_dict['simoptions'], list): + if isinstance(self.input_dict.get('simoptions', None), list): self.input_dict['simoptions'] = list_to_dict(self.input_dict['simoptions']) - simoptions = self.input_dict['simoptions'] + simoptions = self.input_dict.get('simoptions', {}) self.options = {} self.options['sim_path'] = simoptions.get('sim_path', '') From 00673cd62e1fca1a70e1c319bcb546822f9740bf Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 12 Dec 2025 15:45:25 +0100 Subject: [PATCH 68/94] Fix logging bug --- popt/loop/optimize.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/popt/loop/optimize.py b/popt/loop/optimize.py index 4d8c6a2..206d8e3 100644 --- a/popt/loop/optimize.py +++ b/popt/loop/optimize.py @@ -136,7 +136,7 @@ def run_loop(self): previous_state = None if self.epf: previous_state = self.mean_state - logger.info(f' -----> EPF-EnOpt: {self.epf_iteration}, {self.epf["r"]} (outer iteration, penalty factor)') # print epf info + self.logger.info(f' -----> EPF-EnOpt: {self.epf_iteration}, {self.epf["r"]} (outer iteration, penalty factor)') # print epf info while epf_not_converged: # outer loop using epf @@ -154,26 +154,28 @@ def run_loop(self): # Check if max iterations was reached if self.iteration >= self.max_iter: - self.optimize_result['message'] = 'Iterations stopped due to max iterations reached!' + self.msg = 'Optimization stopped due to maximum iterations reached!' + self.optimize_result['message'] = self.msg else: if not isinstance(self.msg, str): self.msg = '' self.optimize_result['message'] = self.msg # Logging some info to screen - logger.info(' Optimization converged in %d iterations ', self.iteration-1) - logger.info(' Optimization converged with final obj_func = %.4f', + self.logger.info(' ============================================') + self.logger.info(' Optimization converged in %d iterations ', self.iteration-1) + self.logger.info(' Optimization converged with final obj_func = %.4f', np.mean(self.optimize_result['fun'])) - logger.info(' Total number of function evaluations = %d', self.optimize_result['nfev']) - logger.info(' Total number of jacobi evaluations = %d', self.optimize_result['njev']) + self.logger.info(' Total number of function evaluations = %d', self.optimize_result['nfev']) + self.logger.info(' Total number of jacobi evaluations = %d', self.optimize_result['njev']) if self.start_time is not None: - logger.info(' Total elapsed time = %.2f minutes', (time.perf_counter()-self.start_time)/60) - logger.info(' ============================================') + self.logger.info(' Total elapsed time = %.2f minutes', (time.perf_counter()-self.start_time)/60) + self.logger.info(' ============================================') # Test for convergence of outer epf loop epf_not_converged = False if self.epf: if self.epf_iteration > self.epf['max_epf_iter']: # max epf_iterations set to 10 - logger.info(f' -----> EPF-EnOpt: maximum epf iterations reached') # print epf info + self.logger.info(f' -----> EPF-EnOpt: maximum epf iterations reached') # print epf info break p = np.abs(previous_state-self.mean_state) / (np.abs(previous_state) + 1.0e-9) conv_crit = self.epf['conv_crit'] @@ -190,12 +192,11 @@ def run_loop(self): self.nfev += 1 self.iteration = +1 r = self.epf['r'] - logger.info(f' -----> EPF-EnOpt: {self.epf_iteration}, {r} (outer iteration, penalty factor)') # print epf info + self.logger.info(f' -----> EPF-EnOpt: {self.epf_iteration}, {r} (outer iteration, penalty factor)') # print epf info else: - logger.info(f' -----> EPF-EnOpt: converged, no variables changed more than {conv_crit*100} %') # print epf info + self.logger.info(f' -----> EPF-EnOpt: converged, no variables changed more than {conv_crit*100} %') # print epf info final_obj_no_penalty = str(round(float(self.fun(self.mean_state)),4)) - logger.info(f' -----> EPF-EnOpt: objective value without penalty = {final_obj_no_penalty}') # print epf info - + self.logger.info(f' -----> EPF-EnOpt: objective value without penalty = {final_obj_no_penalty}') # print epf info def save(self): """ We use pickle to dump all the information we have in 'self'. Can be used, e.g., if some error has occurred. From 9b4bfbb2f9dd38ca030d850e6ba8718d657bd3bf Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 5 Jan 2026 10:46:27 +0100 Subject: [PATCH 69/94] Add logger.py --- ensemble/logger.py | 73 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 ensemble/logger.py diff --git a/ensemble/logger.py b/ensemble/logger.py new file mode 100644 index 0000000..8e37ee1 --- /dev/null +++ b/ensemble/logger.py @@ -0,0 +1,73 @@ +import logging + +class PetLogger: + ''' + A custom logger that logs messages and key-value pairs in a formatted table. + + Parameters: + filename (str): The name of the log file. Defaults to 'PET.log'. + ''' + def __init__(self, filename=None): + + self.filename = filename if filename else 'PET.log' + self.ns = 10 # Number of spaces for table formatting + + # Configurate logging + logging.basicConfig( + level=logging.INFO, + format='%(asctime)s : %(levelname)s : %(message)s', + datefmt='%Y-%m-%d %H:%M:%S', + handlers=[ + logging.FileHandler(self.filename, mode='w'), + logging.StreamHandler() + ] + ) + self._logger = logging.getLogger(__name__) + + + + def __call__(self, *args, **kwargs): + ''' + Log messages or key-value pairs in a formatted table. + + Parameters: + *args: Positional arguments to log as a single message. + **kwargs: Keyword arguments to log in a formatted table. + ''' + + if args: + # Log message from args + self._logger.info('') + msg = ' ' + ' '.join(str(arg) for arg in args) + self._logger.info(msg) + + if kwargs: + # Make strings for table logging + header_parts = [] + values_parts = [] + for key, value in kwargs.items(): + + # Make sure the ns is large enough + try: + if (len(key) > self.ns) or (len(f'{value:.2e}') > self.ns): + self.ns = max(len(key), len(f'{value:.2e}')) + 2 + except: + if len(key) > self.ns: + self.ns = len(key) + 2 + + header_parts.append(f'{key:^{self.ns}}') + try: + if isinstance(value, int) or isinstance(value, str): + values_parts.append(f'{value:^{self.ns}}') + else: + values_parts.append(f'{value:^{self.ns}.2e}') + except: + values_parts.append(f'{"":^{self.ns}}') + + # Log table + self._logger.info('') + self._logger.info(' ' + '─' * (len(kwargs) * self.ns + (len(kwargs) - 1) * 3)) + self._logger.info(' │' + ' │ '.join(header_parts) + '│') + self._logger.info(' │' + '─│─'.join(['─' * self.ns for _ in kwargs.keys()]) + '│') + self._logger.info(' │' + ' │ '.join(values_parts) + '│') + self._logger.info(' ' + '─' * (len(kwargs) * self.ns + (len(kwargs) - 1) * 3)) \ No newline at end of file From dd995335e62bce06e5f7ab5e1fc03394a69d56a1 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 5 Jan 2026 13:41:20 +0100 Subject: [PATCH 70/94] Use PetLogger --- ensemble/ensemble.py | 14 +++--- ensemble/logger.py | 42 ++++++++++------- pipt/loop/assimilation.py | 4 +- pipt/loop/ensemble.py | 19 ++------ pipt/update_schemes/enrml.py | 39 ++++++++-------- pipt/update_schemes/esmda.py | 34 +++++++------- popt/loop/optimize.py | 41 +++++++---------- popt/update_schemes/linesearch.py | 45 +++++++++++-------- .../update_schemes/subroutines/subroutines.py | 1 + 9 files changed, 117 insertions(+), 122 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index dc9146d..f29a683 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -88,13 +88,13 @@ def __init__(self, keys_en: dict, sim, redund_sim=None): # pickle save file. If it is not a restart run, we initialize everything below. if ('restart' in self.keys_en) and (self.keys_en['restart'] == 'yes'): # Initiate a restart run - self.logger.info('\033[92m--- Restart run initiated! ---\033[92m') + self.logger('\033[92m--- Restart run initiated! ---\033[92m') # Check if the pickle save file exists in folder try: assert (self.pickle_restart_file in [ f for f in os.listdir('.') if os.path.isfile(f)]) except AssertionError as err: - self.logger.exception('The restart file "{0}" does not exist in folder. Cannot restart!'.format( + self.logger('The restart file "{0}" does not exist in folder. Cannot restart!'.format( self.pickle_restart_file)) raise err @@ -286,7 +286,7 @@ def calc_prediction(self, enX=None, save_prediction=None): if len(list_crash) > 1: print( '\n\033[1;31mERROR: All started simulations has failed! We dump all information and exit!\033[1;m') - self.logger.info( + self.logger( '\n\033[1;31mERROR: All started simulations has failed! We dump all information and exit!\033[1;m') sys.exit(1) return success @@ -307,7 +307,7 @@ def calc_prediction(self, enX=None, save_prediction=None): f"has been replaced by ensemble member {element}! ---\033[92m" ) print(msg) - self.logger.info(msg) + self.logger(msg) if enX.shape[1] > 1: enX[:, list_crash[index]] = deepcopy(self.enX[:, element]) en_pred[list_crash[index]] = deepcopy(en_pred[element]) @@ -331,7 +331,7 @@ def run_on_HPC(self, enX, batch_size=None, **kwargs): # Split the ensemble into batches of 500 if batch_size >= 1000: - self.logger.info(f'Cannot run batch size of {batch_size}. Set to 1000') + self.logger(f'Cannot run batch size of {batch_size}. Set to 1000') batch_size = 1000 en_pred = [] batch_en = [np.arange(start, start + batch_size) for start in @@ -464,7 +464,7 @@ def calc_ml_prediction(self, input_state=None): if len(list_crash) > 1: print( '\n\033[1;31mERROR: All started simulations has failed! We dump all information and exit!\033[1;m') - self.logger.info( + self.logger( '\n\033[1;31mERROR: All started simulations has failed! We dump all information and exit!\033[1;m') sys.exit(1) return success @@ -484,7 +484,7 @@ def calc_ml_prediction(self, input_state=None): for indx, el in enumerate(copy_member): print(f'\033[92m--- Ensemble member {list_crash[indx]} failed, has been replaced by ensemble member ' f'{el}! ---\033[92m') - self.logger.info(f'\033[92m--- Ensemble member {list_crash[indx]} failed, has been replaced by ' + self.logger(f'\033[92m--- Ensemble member {list_crash[indx]} failed, has been replaced by ' f'ensemble member {el}! ---\033[92m') for key in self.state[level].keys(): self.state[level][key][:, list_crash[indx]] = deepcopy( diff --git a/ensemble/logger.py b/ensemble/logger.py index 8e37ee1..614a049 100644 --- a/ensemble/logger.py +++ b/ensemble/logger.py @@ -10,13 +10,13 @@ class PetLogger: def __init__(self, filename=None): self.filename = filename if filename else 'PET.log' - self.ns = 10 # Number of spaces for table formatting + self.ns = 12 # Number of spaces for table formatting # Configurate logging logging.basicConfig( level=logging.INFO, - format='%(asctime)s : %(levelname)s : %(message)s', - datefmt='%Y-%m-%d %H:%M:%S', + format='%(asctime)s : %(message)s', + datefmt='%Y-%m-%d│%H:%M:%S', handlers=[ logging.FileHandler(self.filename, mode='w'), logging.StreamHandler() @@ -25,7 +25,6 @@ def __init__(self, filename=None): self._logger = logging.getLogger(__name__) - def __call__(self, *args, **kwargs): ''' Log messages or key-value pairs in a formatted table. @@ -34,33 +33,24 @@ def __call__(self, *args, **kwargs): *args: Positional arguments to log as a single message. **kwargs: Keyword arguments to log in a formatted table. ''' - + if args: # Log message from args - self._logger.info('') msg = ' ' + ' '.join(str(arg) for arg in args) self._logger.info(msg) if kwargs: # Make strings for table logging + self._set_ns(**kwargs) header_parts = [] values_parts = [] for key, value in kwargs.items(): - - # Make sure the ns is large enough - try: - if (len(key) > self.ns) or (len(f'{value:.2e}') > self.ns): - self.ns = max(len(key), len(f'{value:.2e}')) + 2 - except: - if len(key) > self.ns: - self.ns = len(key) + 2 - header_parts.append(f'{key:^{self.ns}}') try: if isinstance(value, int) or isinstance(value, str): values_parts.append(f'{value:^{self.ns}}') else: - values_parts.append(f'{value:^{self.ns}.2e}') + values_parts.append(f'{value:^{self.ns}.3e}') except: values_parts.append(f'{"":^{self.ns}}') @@ -70,4 +60,22 @@ def __call__(self, *args, **kwargs): self._logger.info(' │' + ' │ '.join(header_parts) + '│') self._logger.info(' │' + '─│─'.join(['─' * self.ns for _ in kwargs.keys()]) + '│') self._logger.info(' │' + ' │ '.join(values_parts) + '│') - self._logger.info(' ' + '─' * (len(kwargs) * self.ns + (len(kwargs) - 1) * 3)) \ No newline at end of file + self._logger.info(' ' + '─' * (len(kwargs) * self.ns + (len(kwargs) - 1) * 3)) + + def info(self, *args, **kwargs): + self._logger.info(*args, **kwargs) + + def _set_ns(self, **kwargs): + ''' + Adjust the number of spaces for table formatting based on the length of keys and values. + + Parameters: + **kwargs: Keyword arguments to consider for adjusting the space width. + ''' + for key, value in kwargs.items(): + try: + if (len(key) > self.ns) or (len(f'{value:.3e}') > self.ns): + self.ns = max(len(key), len(f'{value:.3e}')) + 2 + except: + if len(key) > self.ns: + self.ns = len(key) + 2 \ No newline at end of file diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index e18d447..b982ae4 100644 --- a/pipt/loop/assimilation.py +++ b/pipt/loop/assimilation.py @@ -247,8 +247,8 @@ def run(self): if self.ensemble.prior_data_misfit > self.ensemble.data_misfit: out_str += f' Obj. function reduced from {self.ensemble.prior_data_misfit:0.1f} ' \ f'to {self.ensemble.data_misfit:0.1f}' - tqdm.write(out_str) - self.ensemble.logger.info(out_str) + #tqdm.write(out_str) + self.ensemble.logger(out_str) def remove_outliers(self): diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index 9231980..d5cd22b 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -1,7 +1,7 @@ """Descriptive description.""" # External import -import logging +#import logging import os.path import numpy @@ -15,6 +15,7 @@ # Internal import from ensemble.ensemble import Ensemble as PETEnsemble +from ensemble.logger import PetLogger import misc.read_input_csv as rcsv from pipt.misc_tools import wavelet_tools as wt from pipt.misc_tools.cov_regularization import localization, _calc_distance @@ -73,21 +74,9 @@ def __init__(self, keys_da, keys_en, sim): super(Ensemble, self).__init__(keys_da|keys_en, sim) # Setup logger - logging.basicConfig( - level=logging.INFO, - format='%(asctime)s : %(levelname)s : %(message)s', - datefmt='%Y-%m-%d %H:%M:%S', - handlers=[ - logging.FileHandler('assim.log', mode='w'), - logging.StreamHandler() - ] - ) - self.logger = logging.getLogger(__name__) + self.logger = PetLogger(filename='assim.log') + self.logger(f'=========== Running Data Assimilation - {keys_da["daalg"][0].upper()} ===========') - # write initial information - self.logger.info('') - self.logger.info(f' =========== Running Data Assimilation - {keys_da["daalg"][0].upper()} ===========') - self.logger.info('') # Internalize PIPT dictionary if not hasattr(self, 'keys_da'): diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index 6227757..1c149de 100644 --- a/pipt/update_schemes/enrml.py +++ b/pipt/update_schemes/enrml.py @@ -219,15 +219,15 @@ def check_convergence(self): if self.data_misfit >= self.prev_data_misfit: success = False - self.logger.info('') - self.logger.info( + self.logger( f'Iterations have converged after {self.iteration} iterations. Objective function reduced ' - f'from {self.prior_data_misfit:0.1f} to {self.prev_data_misfit:0.1f}') + f'from {self.prior_data_misfit:0.1f} to {self.prev_data_misfit:0.1f}' + ) else: - self.logger.info('') self.logger.info( f'Iterations have converged after {self.iteration} iterations. Objective function reduced ' - f'from {self.prior_data_misfit:0.1f} to {self.data_misfit:0.1f}') + f'from {self.prior_data_misfit:0.1f} to {self.data_misfit:0.1f}' + ) # Return conv = True, why_stop var. return True, success, why_stop @@ -291,22 +291,21 @@ def log_update(self, success, prior_run=False): ''' Log the update results in a formatted table. ''' - def _log_table(iteration, status, misfit, change_pct, lambda_val, change_label="Reduction"): - """Helper method to log iteration results in a formatted table.""" - self.logger.info('') - self.logger.info(f' {"Iteration":<11}| {"Status":<11}| {"Data Misfit":<16}| {f"{change_label} (%)":<15}| {"λ":<10}') - self.logger.info(f' {"—"*11}|{"—"*12}|{"—"*17}|{"—"*16}|{"—"*10}') - self.logger.info(f' {iteration:<11}| {status:<11}| {misfit:<16.2f}| {change_pct:<15.2f}| {lambda_val:<10.2f}') - self.logger.info('') - - if prior_run: - _log_table(0, "Success", self.data_misfit, 0.0, self.lam) - elif success: - reduction = 100 * (1 - self.data_misfit / self.prev_data_misfit) - _log_table(self.iteration, "Success", self.data_misfit, reduction, self.lam) + log_data = { + "Iteration": f'{0 if prior_run else self.iteration}', + "Status": "Success" if (prior_run or success) else "Failed", + "Data Misfit": self.data_misfit, + "λ": self.lam + } + if not prior_run: + if success: + log_data["Reduction (%)"] = 100 * (1 - self.data_misfit / self.prev_data_misfit) + else: + log_data["Increase (%)"] = 100 * (self.data_misfit / self.prev_data_misfit - 1) else: - increase = 100 * (1 - self.prev_data_misfit / self.data_misfit) - _log_table(self.iteration, "Failed", self.data_misfit, increase, self.lam, "Increase") + log_data["Reduction (%)"] = 'N/A' + + self.logger(**log_data) diff --git a/pipt/update_schemes/esmda.py b/pipt/update_schemes/esmda.py index 142df35..0a5b110 100644 --- a/pipt/update_schemes/esmda.py +++ b/pipt/update_schemes/esmda.py @@ -228,25 +228,23 @@ def log_update(self, success=None, prior_run=False): ''' Log the update results in a formatted table. ''' - def _log_table(iteration, status, misfit, change_pct, change_label="Reduction"): - """Helper method to log iteration results in a formatted table.""" - self.logger.info('') - self.logger.info(f' {"Iteration":<11}| {"Status":<11}| {"Data Misfit":<16}| {f"{change_label} (%)":<15}') - self.logger.info(f' {"—"*11}|{"—"*12}|{"—"*17}|{"—"*16}') - self.logger.info(f' {iteration:<11}| {status:<11}| {misfit:<16.2f}| {change_pct:<15.2f}') - self.logger.info('') - - if prior_run: - iteration_str = f'0/{self.max_iter}' - _log_table(iteration_str, "Success", self.data_misfit, 0.0) - elif success: - iteration_str = f'{self.iteration}/{self.max_iter}' - reduction = 100 * (1 - self.data_misfit / self.prev_data_misfit) - _log_table(iteration_str, "Success", self.data_misfit, reduction) + iteration_str = f'{0 if prior_run else self.iteration}/{self.max_iter}' + + log_data = { + "Iteration": iteration_str, + "Status": "Success" if (prior_run or success) else "Failed", + "Data Misfit": self.data_misfit + } + + if not prior_run: + if success: + log_data["Reduction (%)"] = 100 * (1 - self.data_misfit / self.prev_data_misfit) + else: + log_data["Increase (%)"] = 100 * (self.data_misfit / self.prev_data_misfit - 1) else: - iteration_str = f'{self.iteration}/{self.max_iter}' - increase = 100 * (1 - self.prev_data_misfit / self.data_misfit) - _log_table(iteration_str, "Failed", self.data_misfit, increase, "Increase") + log_data["Reduction (%)"] = 'N/A' + + self.logger(**log_data) def _ext_inflation_param(self): r""" diff --git a/popt/loop/optimize.py b/popt/loop/optimize.py index 206d8e3..756acfa 100644 --- a/popt/loop/optimize.py +++ b/popt/loop/optimize.py @@ -1,12 +1,12 @@ # External imports import os import numpy as np -import logging import time import pickle # Internal imports import popt.misc_tools.optim_tools as ot +from ensemble.logger import PetLogger class Optimize: @@ -62,16 +62,7 @@ def __init__(self, **options): Optimization options """ # Setup logger - logging.basicConfig( - level=logging.INFO, - format='%(asctime)s : %(levelname)s : %(message)s', - datefmt='%Y-%m-%d %H:%M:%S', - handlers=[ - logging.FileHandler('optim.log', mode='w'), - logging.StreamHandler() - ] - ) - self.logger = logging.getLogger(__name__) + self.logger = PetLogger('optim.log') # Save name for (potential) pickle dump/load self.pickle_restart_file = 'popt_restart_dump' @@ -136,7 +127,9 @@ def run_loop(self): previous_state = None if self.epf: previous_state = self.mean_state - self.logger.info(f' -----> EPF-EnOpt: {self.epf_iteration}, {self.epf["r"]} (outer iteration, penalty factor)') # print epf info + self.logger( + f'─────> EPF-EnOpt: {self.epf_iteration}, {self.epf["r"]} (outer iteration, penalty factor)' + ) # print epf info while epf_not_converged: # outer loop using epf @@ -161,21 +154,21 @@ def run_loop(self): self.optimize_result['message'] = self.msg # Logging some info to screen - self.logger.info(' ============================================') - self.logger.info(' Optimization converged in %d iterations ', self.iteration-1) - self.logger.info(' Optimization converged with final obj_func = %.4f', - np.mean(self.optimize_result['fun'])) - self.logger.info(' Total number of function evaluations = %d', self.optimize_result['nfev']) - self.logger.info(' Total number of jacobi evaluations = %d', self.optimize_result['njev']) + self.logger('') + self.logger('============================================') + self.logger(f'Optimization converged in {self.iteration-1} iterations ') + self.logger(f'Optimization converged with final obj_func = {np.mean(self.optimize_result["fun"]):.4f}') + self.logger(f'Total number of function evaluations = {self.optimize_result["nfev"]}') + self.logger(f'Total number of jacobi evaluations = {self.optimize_result["njev"]}') if self.start_time is not None: - self.logger.info(' Total elapsed time = %.2f minutes', (time.perf_counter()-self.start_time)/60) - self.logger.info(' ============================================') + self.logger(f'Total elapsed time = {(time.perf_counter()-self.start_time)/60:.2f} minutes') + self.logger('============================================') # Test for convergence of outer epf loop epf_not_converged = False if self.epf: if self.epf_iteration > self.epf['max_epf_iter']: # max epf_iterations set to 10 - self.logger.info(f' -----> EPF-EnOpt: maximum epf iterations reached') # print epf info + self.logger(f'─────> EPF-EnOpt: maximum epf iterations reached') # print epf info break p = np.abs(previous_state-self.mean_state) / (np.abs(previous_state) + 1.0e-9) conv_crit = self.epf['conv_crit'] @@ -192,11 +185,11 @@ def run_loop(self): self.nfev += 1 self.iteration = +1 r = self.epf['r'] - self.logger.info(f' -----> EPF-EnOpt: {self.epf_iteration}, {r} (outer iteration, penalty factor)') # print epf info + self.logger(f'─────> EPF-EnOpt: {self.epf_iteration}, {r} (outer iteration, penalty factor)') # print epf info else: - self.logger.info(f' -----> EPF-EnOpt: converged, no variables changed more than {conv_crit*100} %') # print epf info + self.logger(f'─────> EPF-EnOpt: converged, no variables changed more than {conv_crit*100} %') # print epf info final_obj_no_penalty = str(round(float(self.fun(self.mean_state)),4)) - self.logger.info(f' -----> EPF-EnOpt: objective value without penalty = {final_obj_no_penalty}') # print epf info + self.logger(f'─────> EPF-EnOpt: objective value without penalty = {final_obj_no_penalty}') # print epf info def save(self): """ We use pickle to dump all the information we have in 'self'. Can be used, e.g., if some error has occurred. diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index bf5040d..3c944bf 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -191,7 +191,7 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c self.hessian = hess self.args = args self.bounds = bounds - self.options = options + self.options = options # Check for Callback function if callable(callback): @@ -199,6 +199,9 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c else: self.callback = None + # Remove 'datatype' form options if present (This is a temporary bugfix) + self.options.pop('datatype', None) + # Custom convergence criteria (callable) convergence_criteria = options.get('convergence_criteria', None) if callable(convergence_criteria): @@ -219,7 +222,7 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c 'amax': self.step_size_max, 'maxiter': options.get('lsmaxiter', 10), 'method' : options.get('lsmethod', 1), - 'logger' : self.logger.info + 'logger' : self.logger } # Set other options @@ -269,12 +272,14 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c if self.saveit: ot.save_optimize_results(self.optimize_result) if self.logger is not None: - self.logger.info(f' ====== Running optimization - Line search ({method}) ======') - self.logger.info('\nSPECIFIED OPTIONS:\n'+pprint.pformat(OptimizeResult(self.options))) - self.logger.info('') - self.logger.info(f' {"iter.":<10} {fun_xk_symbol:<15} {jac_inf_symbol:<15} {"step-size":<15}') - self.logger.info(f' {self.iteration:<10} {self.fk:<15.4e} {la.norm(self.jk, np.inf):<15.4e} {0:<15.4e}') - self.logger.info('') + self.logger(f'========== Running optimization - Line search ({method}) ==========') + self.logger(f'\n \nSPECIFIED OPTIONS:\n{pprint.pformat(OptimizeResult(self.options))}\n') + self.logger(**{ + 'iter.': 0, + fun_xk_symbol: self.fk, + jac_inf_symbol: la.norm(self.jk, np.inf), + 'step-size': self.step_size + }) self.run_loop() @@ -340,7 +345,7 @@ def calc_update(self, iter_resamp=0): if self.method == 'BFGS': pk = - np.matmul(self.Hk_inv, self.jk) if self.method == 'Newton-CG': - pk = newton_cg(self.jk, Hk=self.Hk, xk=self.xk, jac=self._jac, logger=self.logger.info) + pk = newton_cg(self.jk, Hk=self.Hk, xk=self.xk, jac=self._jac, logger=self.logger) # porject search direction onto the feasible set if self.bounds is not None: @@ -353,7 +358,7 @@ def calc_update(self, iter_resamp=0): step_size = self._set_step_size(pk, self.step_size_max) # Perform line-search - self.logger.info('Performing line search.............') + self.logger('Performing line search.............') if self.lskwargs['method'] == 0: ls_res = line_search_backtracking( step_size=step_size, @@ -419,32 +424,34 @@ def calc_update(self, iter_resamp=0): # Write logging info if self.logger is not None: - self.logger.info('') - self.logger.info(f' {"iter.":<10} {fun_xk_symbol:<15} {jac_inf_symbol:<15} {"step-size":<15}') - self.logger.info(f' {self.iteration:<10} {self.fk:<15.4e} {la.norm(self.jk, np.inf):<15.4e} {step_size:<15.4e}') - self.logger.info('') + self.logger(**{ + 'iter.': self.iteration, + fun_xk_symbol: self.fk, + jac_inf_symbol: la.norm(self.jk, np.inf), + 'step-size': step_size + }) # Check for convergence if (la.norm(sk, np.inf) < self.xtol): self.msg = 'Convergence criteria met: |dx| < xtol' - self.logger.info(self.msg) + self.logger(self.msg) success = False return success if (np.abs(self.fk - f_old) < self.ftol * np.abs(f_old)): self.msg = 'Convergence criteria met: |f(x+dx) - f(x)| < ftol * |f(x)|' - self.logger.info(self.msg) + self.logger(self.msg) success = False return success if (la.norm(self.jk, np.inf) < self.gtol): self.msg = f'Convergence criteria met: {jac_inf_symbol} < gtol' - self.logger.info(self.msg) + self.logger(self.msg) success = False return success # Check for custom convergence if callable(self.convergence_criteria): if self.convergence_criteria(self): - self.logger.info('Custom convergence criteria met. Stopping optimization.') + self.logger('Custom convergence criteria met. Stopping optimization.') success = False return success @@ -457,7 +464,7 @@ def calc_update(self, iter_resamp=0): else: if iter_resamp < self.resample: - self.logger.info('Resampling Gradient') + self.logger('Resampling Gradient') iter_resamp += 1 self.jk = None diff --git a/popt/update_schemes/subroutines/subroutines.py b/popt/update_schemes/subroutines/subroutines.py index ddd1105..bd69274 100644 --- a/popt/update_schemes/subroutines/subroutines.py +++ b/popt/update_schemes/subroutines/subroutines.py @@ -347,6 +347,7 @@ def newton_cg(gk, Hk=None, maxiter=None, **kwargs): if logger is None: logger = print + logger('') logger('Running Newton-CG subroutine..........') if Hk is None: From 3f9abf17f498290d32f46efae03281b9a1e298e6 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 5 Jan 2026 13:56:10 +0100 Subject: [PATCH 71/94] Update string --- popt/update_schemes/linesearch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index 3c944bf..f9a417f 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -273,7 +273,7 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c ot.save_optimize_results(self.optimize_result) if self.logger is not None: self.logger(f'========== Running optimization - Line search ({method}) ==========') - self.logger(f'\n \nSPECIFIED OPTIONS:\n{pprint.pformat(OptimizeResult(self.options))}\n') + self.logger(f'\n \nUSER-SPECIFIED OPTIONS:\n{pprint.pformat(OptimizeResult(self.options))}\n') self.logger(**{ 'iter.': 0, fun_xk_symbol: self.fk, From 9078c9f1ea790a1e003bb8d333455a0ce4a17542 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 6 Jan 2026 08:32:31 +0100 Subject: [PATCH 72/94] Add convergence message --- popt/loop/optimize.py | 1 + 1 file changed, 1 insertion(+) diff --git a/popt/loop/optimize.py b/popt/loop/optimize.py index 756acfa..87cd3b8 100644 --- a/popt/loop/optimize.py +++ b/popt/loop/optimize.py @@ -156,6 +156,7 @@ def run_loop(self): # Logging some info to screen self.logger('') self.logger('============================================') + self.logger(self.msg) self.logger(f'Optimization converged in {self.iteration-1} iterations ') self.logger(f'Optimization converged with final obj_func = {np.mean(self.optimize_result["fun"]):.4f}') self.logger(f'Total number of function evaluations = {self.optimize_result["nfev"]}') From a83a0e28ea2f57767f94ab2d7f4b492b13fbc8f2 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 8 Jan 2026 13:10:58 +0100 Subject: [PATCH 73/94] Remove lines --- ensemble/ensemble.py | 2 +- popt/cost_functions/npv.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index f29a683..7ecfe76 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -265,7 +265,7 @@ def calc_prediction(self, enX=None, save_prediction=None): **progbar_settings ) ###################################################################################################################### - + # Convert state enemble back to matrix form enX = entools.list_to_matrix(enX, self.idX) diff --git a/popt/cost_functions/npv.py b/popt/cost_functions/npv.py index bb18c8c..e7e6384 100644 --- a/popt/cost_functions/npv.py +++ b/popt/cost_functions/npv.py @@ -42,7 +42,6 @@ def npv(pred_data, **kwargs): values = [] for i in np.arange(1, len(pred_data)): - Qop = np.squeeze(pred_data[i]['fopt']) - np.squeeze(pred_data[i - 1]['fopt']) Qgp = np.squeeze(pred_data[i]['fgpt']) - np.squeeze(pred_data[i - 1]['fgpt']) Qwp = np.squeeze(pred_data[i]['fwpt']) - np.squeeze(pred_data[i - 1]['fwpt']) From 91cf5d1f20410c9afee3710c6b23e4c20b77f342 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 12 Jan 2026 10:42:51 +0100 Subject: [PATCH 74/94] Add cosmetic changes to PetLogger --- ensemble/logger.py | 38 +++++++++++++++++++++++++++----------- 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/ensemble/logger.py b/ensemble/logger.py index 614a049..0f23ebb 100644 --- a/ensemble/logger.py +++ b/ensemble/logger.py @@ -32,6 +32,20 @@ def __call__(self, *args, **kwargs): Parameters: *args: Positional arguments to log as a single message. **kwargs: Keyword arguments to log in a formatted table. + + Example: + __call__('This is a log message.') ----> + 2024-06-01│12:00:00 : This is a log message. + + __call__(iteration=1, fun=0.5, step-size=0.1) ----> + + 2024-06-01│12:00:00 : + 2024-06-01│12:00:00 : ┌────────────┬────────────┬────────────┐ + 2024-06-01│12:00:00 : │ iteration │ fun │ step-size │ + 2024-06-01│12:00:00 : ├────────────┼────────────┼────────────┤ + 2024-06-01│12:00:00 : │ 1 │ 5.000e-01 │ 1.000e-01 │ + 2024-06-01│12:00:00 : └────────────┴────────────┴────────────┘ + 2024-06-01│12:00:00 : ''' if args: @@ -42,25 +56,27 @@ def __call__(self, *args, **kwargs): if kwargs: # Make strings for table logging self._set_ns(**kwargs) - header_parts = [] - values_parts = [] + header = [] + values = [] for key, value in kwargs.items(): - header_parts.append(f'{key:^{self.ns}}') + header.append(f'{key:^{self.ns}}') try: if isinstance(value, int) or isinstance(value, str): - values_parts.append(f'{value:^{self.ns}}') + values.append(f'{value:^{self.ns}}') else: - values_parts.append(f'{value:^{self.ns}.3e}') + values.append(f'{value:^{self.ns}.3e}') except: - values_parts.append(f'{"":^{self.ns}}') + values.append(f'{"":^{self.ns}}') # Log table + seperator = ['─' * self.ns for _ in kwargs.keys()] + self._logger.info('') + self._logger.info(' ┌' + '┬'.join(seperator) + '┐') + self._logger.info(' │' + '│'.join(header) + '│') + self._logger.info(' ├' + '┼'.join(seperator) + '┤') + self._logger.info(' │' + '│'.join(values) + '│') + self._logger.info(' └' + '┴'.join(seperator) + '┘') self._logger.info('') - self._logger.info(' ' + '─' * (len(kwargs) * self.ns + (len(kwargs) - 1) * 3)) - self._logger.info(' │' + ' │ '.join(header_parts) + '│') - self._logger.info(' │' + '─│─'.join(['─' * self.ns for _ in kwargs.keys()]) + '│') - self._logger.info(' │' + ' │ '.join(values_parts) + '│') - self._logger.info(' ' + '─' * (len(kwargs) * self.ns + (len(kwargs) - 1) * 3)) def info(self, *args, **kwargs): self._logger.info(*args, **kwargs) From 8a263fd615502db0eaf44668b6fcf6b4b979460e Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 12 Jan 2026 10:45:14 +0100 Subject: [PATCH 75/94] Update dosctring for PetLogger --- ensemble/logger.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/ensemble/logger.py b/ensemble/logger.py index 0f23ebb..50bf5ac 100644 --- a/ensemble/logger.py +++ b/ensemble/logger.py @@ -34,18 +34,17 @@ def __call__(self, *args, **kwargs): **kwargs: Keyword arguments to log in a formatted table. Example: - __call__('This is a log message.') ----> + >>> __call__('This is a log message.') 2024-06-01│12:00:00 : This is a log message. - - __call__(iteration=1, fun=0.5, step-size=0.1) ----> - - 2024-06-01│12:00:00 : - 2024-06-01│12:00:00 : ┌────────────┬────────────┬────────────┐ - 2024-06-01│12:00:00 : │ iteration │ fun │ step-size │ - 2024-06-01│12:00:00 : ├────────────┼────────────┼────────────┤ - 2024-06-01│12:00:00 : │ 1 │ 5.000e-01 │ 1.000e-01 │ - 2024-06-01│12:00:00 : └────────────┴────────────┴────────────┘ - 2024-06-01│12:00:00 : + >>> + >>> __call__(iteration=1, fun=0.5, step_size=0.1) + 2024-06-01│12:00:00 : + 2024-06-01│12:00:00 : ┌────────────┬────────────┬────────────┐ + 2024-06-01│12:00:00 : │ iteration │ fun │ step_size │ + 2024-06-01│12:00:00 : ├────────────┼────────────┼────────────┤ + 2024-06-01│12:00:00 : │ 1 │ 5.000e-01 │ 1.000e-01 │ + 2024-06-01│12:00:00 : └────────────┴────────────┴────────────┘ + 2024-06-01│12:00:00 : ''' if args: From f34a809ab6ddb8cda1d260d1b7e3fcf6636b66f7 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Mon, 12 Jan 2026 11:14:40 +0100 Subject: [PATCH 76/94] Improve logging for LineSearch --- popt/update_schemes/linesearch.py | 1 - .../update_schemes/subroutines/subroutines.py | 30 +++++++++++++++++-- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index f9a417f..ba560a9 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -358,7 +358,6 @@ def calc_update(self, iter_resamp=0): step_size = self._set_step_size(pk, self.step_size_max) # Perform line-search - self.logger('Performing line search.............') if self.lskwargs['method'] == 0: ls_res = line_search_backtracking( step_size=step_size, diff --git a/popt/update_schemes/subroutines/subroutines.py b/popt/update_schemes/subroutines/subroutines.py index bd69274..3700be2 100644 --- a/popt/update_schemes/subroutines/subroutines.py +++ b/popt/update_schemes/subroutines/subroutines.py @@ -81,6 +81,9 @@ def line_search(step_size, xk, pk, fun, jac, fk=None, jk=None, **kwargs): if logger is None: logger = print + logger('Performing line search..........') + logger('──────────────────────────────────────────────────') + # assertions assert step_size <= amax, "Initial step size must be less than or equal to amax." @@ -131,6 +134,7 @@ def dphi(alpha): if (phi_i > phi_0 + c1*a[i]*dphi_0) or (phi_i >= phi(a[i-1]) and i>0): # Call zoom function step_size = zoom(a[i-1], a[i], phi, dphi, phi_0, dphi_0, maxiter+1-i, c1, c2) + logger('──────────────────────────────────────────────────') return step_size, phi.fun_val, dphi.jac_val, ls_nfev, ls_njev # Evaluate dphi(ai) @@ -139,19 +143,23 @@ def dphi(alpha): # Check curvature condition if abs(dphi_i) <= -c2*dphi_0: step_size = a[i] + logger('──────────────────────────────────────────────────') return step_size, phi.fun_val, dphi.jac_val, ls_nfev, ls_njev # Check for posetive derivative if dphi_i >= 0: # Call zoom function - step_size = zoom(a[i], a[i-1], phi, dphi, phi_0, dphi_0, maxiter+1-i, c1, c2) + step_size = zoom(a[i], a[i-1], phi, dphi, phi_0, dphi_0, maxiter+1-i, c1, c2) + logger('──────────────────────────────────────────────────') return step_size, phi.fun_val, dphi.jac_val, ls_nfev, ls_njev # Increase ai a.append(min(2*a[i], amax)) + logger(f' Step-size: {a[i]:.3e} ──> {a[i+1]:.3e}') # If we reached this point, the line search failed - logger('Line search failed to find a suitable step size \n') + logger('Line search failed to find a suitable step size') + logger('──────────────────────────────────────────────────') return None, None, None, ls_nfev, ls_njev @@ -179,6 +187,9 @@ def zoom(alo, ahi, f, df, f0, df0, maxiter, c1, c2): if (aj is None) or (aj < alo + tol_quad) or (aj > ahi - tol_quad): aj = alo + 0.5*(ahi - alo) + + logger(f' New step-size ──> {aj:.3e}') + # Evaluate phi(aj) phi_j = f(aj) @@ -213,6 +224,8 @@ def zoom(alo, ahi, f, df, f0, df0, maxiter, c1, c2): dphi_lo = dphi_j # If we reached this point, the line search failed + logger('Line search failed to find a suitable step size') + logger('──────────────────────────────────────────────────') return None @@ -280,6 +293,15 @@ def line_search_backtracking(step_size, xk, pk, fun, jac, fk=None, jk=None, **kw maxiter = kwargs.get('maxiter', 10) c1 = kwargs.get('c1', 1e-4) + # check for logger in kwargs + global logger + logger = kwargs.get('logger', None) + if logger is None: + logger = print + + logger('Performing backtracking line search..........') + logger('──────────────────────────────────────────────────') + # Define phi and derivative of phi @lru_cache(maxsize=None) def phi(alpha): @@ -298,6 +320,7 @@ def phi(alpha): # run the backtracking line search loop for i in range(maxiter): + logger(f'iteration: {i}') # Evaluate phi(alpha) phi_i = phi(step_size) @@ -305,12 +328,15 @@ def phi(alpha): if (phi_i <= phi(0) + c1*step_size*np.dot(jk, pk)): # Evaluate jac at new point jac_new = jac(xk + step_size*pk) + logger('──────────────────────────────────────────────────') return step_size, phi_i, jac_new, ls_nfev, ls_njev # Reduce step size step_size *= rho # If we reached this point, the line search failed + logger('Backtracking failed to find a suitable step size') + logger('──────────────────────────────────────────────────') return None, None, None, ls_nfev, ls_njev From aaa9786c353e5fb50db190299f79e753b6ad0f16 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 13 Jan 2026 08:18:14 +0100 Subject: [PATCH 77/94] Update docstring for Petlogger --- ensemble/logger.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ensemble/logger.py b/ensemble/logger.py index 50bf5ac..8076d30 100644 --- a/ensemble/logger.py +++ b/ensemble/logger.py @@ -34,10 +34,11 @@ def __call__(self, *args, **kwargs): **kwargs: Keyword arguments to log in a formatted table. Example: - >>> __call__('This is a log message.') + >>> logger = PetLogger() + >>> logger('This is a log message.') 2024-06-01│12:00:00 : This is a log message. >>> - >>> __call__(iteration=1, fun=0.5, step_size=0.1) + >>> logger(iteration=1, fun=0.5, step_size=0.1) 2024-06-01│12:00:00 : 2024-06-01│12:00:00 : ┌────────────┬────────────┬────────────┐ 2024-06-01│12:00:00 : │ iteration │ fun │ step_size │ From f1f739612cc650f54e0fd760ceb633ce2a88d7d9 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 13 Jan 2026 08:43:50 +0100 Subject: [PATCH 78/94] Update PetLogger --- ensemble/logger.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ensemble/logger.py b/ensemble/logger.py index 8076d30..4943ac9 100644 --- a/ensemble/logger.py +++ b/ensemble/logger.py @@ -63,6 +63,8 @@ def __call__(self, *args, **kwargs): try: if isinstance(value, int) or isinstance(value, str): values.append(f'{value:^{self.ns}}') + elif '%' in key: + values.append(f'{value:^{self.ns}.2f}%') else: values.append(f'{value:^{self.ns}.3e}') except: From 3d021a7b1e704bc38b5e36e8b0ef680f54d7b3d9 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 13 Jan 2026 09:26:00 +0100 Subject: [PATCH 79/94] Update calc_ml_prediction --- ensemble/ensemble.py | 67 +- ensemble/logger.py | 2 +- pipt/update_schemes/multilevel.py | 1194 +++++++++++++++++ .../update_methods_ns/hybrid_udpate.py | 65 + 4 files changed, 1295 insertions(+), 33 deletions(-) create mode 100644 pipt/update_schemes/multilevel.py create mode 100644 pipt/update_schemes/update_methods_ns/hybrid_udpate.py diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 7ecfe76..74dc154 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -405,7 +405,7 @@ def load(self): # Save in 'self' self.__dict__.update(tmp_load) - def calc_ml_prediction(self, input_state=None): + def calc_ml_prediction(self, enX=None): """ Function for running the simulator over several levels. We assume that it is sufficient to provide the level integer to the setup of the forward run. This will initiate the correct simulator fidelity. @@ -413,7 +413,7 @@ def calc_ml_prediction(self, input_state=None): Parameters ---------- - input_state: + enX: If simulation is run stand-alone one can input any state. """ @@ -421,36 +421,37 @@ def calc_ml_prediction(self, input_state=None): ml_pred_data = [] for level in tqdm(self.multilevel['levels'], desc='Fidelity level', position=1): + # Setup forward simulator and redundant simulator at the correct fidelity if self.sim.redund_sim is not None: - self.sim.redund_sim.setup_fwd_run(level=level) - self.sim.setup_fwd_run(level=level) + if hasattr(self.sim.redund_sim, 'setup_fwd_run'): + self.sim.redund_sim.setup_fwd_run(level=level) + + # Run setup function for simulator + if hasattr(self.sim, 'setup_fwd_run'): + self.sim.setup_fwd_run(level=level) + ml_ne = self.multilevel['ne'][level] if ml_ne: - # Ensure that we put all the states in a list - list_state = [deepcopy({}) for _ in ml_ne] - for i in ml_ne: - if input_state is None: - for key in self.state[level].keys(): - if self.state[level][key].ndim == 1: - list_state[i][key] = deepcopy(self.state[level][key]) - elif self.state[level][key].ndim == 2: - list_state[i][key] = deepcopy(self.state[level][key][:, i]) - else: - for key in self.state.keys(): - if input_state[level][key].ndim == 1: - list_state[i][key] = deepcopy(input_state[level][key]) - elif input_state[level][key].ndim == 2: - list_state[i][key] = deepcopy(input_state[level][key][:, i]) - if self.aux_input is not None: # several models are used - list_state[i]['aux_input'] = self.aux_input[i] + + level_enX = entools.matrix_to_list(enX[level], self.idX) + for n in range(ml_ne): + if self.aux_input is not None: + level_enX[n]['aux_input'] = self.aux_input[n] + # Index list of ensemble members list_member_index = list(ml_ne) # Run prediction in parallel using p_map - en_pred = p_map(self.sim.run_fwd_sim, list_state, - list_member_index, num_cpus=no_tot_run, disable=self.disable_tqdm) + en_pred = p_map( + self.sim.run_fwd_sim, + level_enX, + list_member_index, + num_cpus=no_tot_run, + disable=self.disable_tqdm, + **progbar_settings, + ) # List successful runs and crashes list_crash = [indx for indx, el in enumerate(en_pred) if el is False] @@ -481,15 +482,17 @@ def calc_ml_prediction(self, input_state=None): list_success, size=len(list_crash), replace=True) # Insert the replaced runs in prediction list - for indx, el in enumerate(copy_member): - print(f'\033[92m--- Ensemble member {list_crash[indx]} failed, has been replaced by ensemble member ' - f'{el}! ---\033[92m') - self.logger(f'\033[92m--- Ensemble member {list_crash[indx]} failed, has been replaced by ' - f'ensemble member {el}! ---\033[92m') - for key in self.state[level].keys(): - self.state[level][key][:, list_crash[indx]] = deepcopy( - self.state[level][key][:, el]) - en_pred[list_crash[indx]] = deepcopy(en_pred[el]) + for index, element in enumerate(copy_member): + msg = ( + f"\033[92m--- Ensemble member {list_crash[index]} failed, " + f"has been replaced by ensemble member {element}! ---\033[92m" + ) + print(msg) + self.logger(msg) + if enX[level].shape[1] > 1: + enX[level][:, list_crash[index]] = deepcopy(enX[level][:, element]) + + en_pred[list_crash[index]] = deepcopy(en_pred[element]) # Convert ensemble specific result into pred_data, and filter for NONE data ml_pred_data.append([{typ: np.concatenate(tuple((el[ind][typ][:, np.newaxis]) for el in en_pred), axis=1) diff --git a/ensemble/logger.py b/ensemble/logger.py index 4943ac9..57010bc 100644 --- a/ensemble/logger.py +++ b/ensemble/logger.py @@ -64,7 +64,7 @@ def __call__(self, *args, **kwargs): if isinstance(value, int) or isinstance(value, str): values.append(f'{value:^{self.ns}}') elif '%' in key: - values.append(f'{value:^{self.ns}.2f}%') + values.append(f'{value:^{self.ns}.1f}') else: values.append(f'{value:^{self.ns}.3e}') except: diff --git a/pipt/update_schemes/multilevel.py b/pipt/update_schemes/multilevel.py new file mode 100644 index 0000000..4b4b610 --- /dev/null +++ b/pipt/update_schemes/multilevel.py @@ -0,0 +1,1194 @@ +''' +Here we place the classes that are required to run the multilevel schemes developed in the 4DSeis project. All methods +inherit the ensemble class, hence the main loop is inherited. These classes will consider the analysis step. +''' + +# local imports. Note, it is assumed that PET is installed and available in the path. +from pipt.loop.ensemble import Ensemble +from pipt.update_schemes.esmda import esmda_approx +from pipt.update_schemes.esmda import esmdaMixIn +from pipt.misc_tools import analysis_tools as at +from geostat.decomp import Cholesky +from misc import ecl + +from pipt.update_schemes.update_methods_ns.hybrid_update import hybrid_update +# system imports +import numpy as np +from scipy.sparse import coo_matrix +from scipy import linalg +import time +import shutil +import pickle +from scipy.linalg import solve # For linear system solvers +from scipy.stats import multivariate_normal +from scipy import sparse +from copy import deepcopy +import random +import os +import sys +from scipy.stats import ortho_group +from shutil import copyfile +import math + + +class multilevel(Ensemble): + """ + Inititallize the multilevel class. Similar for all ML schemes, hence make one class for all. + """ + def __init__(self, keys_da,keys_fwd,sim): + super().__init__(keys_da, keys_fwd, sim) + self._ext_ml_feat() + #self.ML_state = [{} for _ in range(self.tot_level)] + #self.ML_state[0] = deepcopy(self.state) + self.data_size = self.ext_data_size() + self.list_states = list(self.state.keys()) + self.init_ml_prior() + self.prior_state = deepcopy(self.state) + self._init_sim() + self.iteration = 0 + self.lam = 0 # set LM lamda to zero as we are doing one full update. + if 'energy' in self.keys_da: + self.trunc_energy = self.keys_da['energy'] # initial energy (Remember to extract this) + if self.trunc_energy > 1: # ensure that it is given as percentage + self.trunc_energy /= 100. + else: + self.trunc_energy = 0.98 + + self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] + # define the list of states + # define the list of datatypes + self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(self.obs_data, self.assim_index) + + self.current_state = deepcopy(self.state) + self.cov_wgt = self.ext_cov_mat_wgt() + self.cov_data = at.gen_covdata(self.datavar, self.assim_index, self.list_datatypes) + self.obs_data_vector, _ = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, + self.list_datatypes) + + def _ext_ml_feat(self): + """ + Extract ML specific info from input. + """ + # Make sure ML is a list + if not isinstance(self.keys_da['multilevel'][0], list): + ml_opts = [self.keys_da['multilevel']] + else: + ml_opts = self.keys_da['multilevel'] + + # set default + self.ML_Nested = False + + # Check if 'levels' has been given; if not, give error (mandatory in MULTILEVEL) + assert 'levels' in list(zip(*ml_opts))[0], 'LEVELS has not been given in MULTILEVEL!' + # Check if ensemble size has been given; if not, give error (mandatory in MULTILEVEL) + assert 'en_size' in list(zip(*ml_opts))[0], 'En_Size has not been given in MULTILEVEL!' + # Check if the Hybrid Weights are provided. If not, give error + assert 'cov_wgt' in list(zip(*ml_opts))[0], 'COV_WGT has not been given in MLDA!' + + for i, opt in enumerate(list(zip(*self.keys_da['multilevel']))[0]): + if opt == 'levels': + self.tot_level = int(self.keys_da['multilevel'][i][1]) + if opt == 'nested_states': + if self.keys_da['multilevel'][i][1] == 'true': + self.ML_Nested = True + if opt == 'en_size': + self.ml_ne = [int(el) for el in self.keys_da['multilevel'][i][1]] + if opt == 'ml_error_corr': + #options for ML_error_corr are: bias_corr, deterministic, stochastic, telescopic + self.ML_error_corr = self.keys_da['multilevel'][i][1] + if not self.ML_error_corr=='none': + #options for error_comp_scheme are: once, ens, sep + self.error_comp_scheme = self.keys_da['multilevel'][i][2] + if opt == 'cov_wgt': + try: + cov_mat_wgt = [float(elem) for elem in [item for item in self.keys_da['multilevel'][i][1]]] + except: + cov_mat_wgt = [float(item) for item in self.keys_da['multilevel'][i][1]] + Sum = 0 + for i in range(len(cov_mat_wgt)): + Sum += cov_mat_wgt[i] + for i in range(len(cov_mat_wgt)): + cov_mat_wgt[i] /= Sum + self.cov_wgt = cov_mat_wgt + # Check that we have specified a size for all levels: + assert len(self.ml_ne) == self.tot_level, 'The Ensemble Size must be specified for all levels!' + + def _init_sim(self): + """ + Ensure that the simulator is initiallized to handle ML forward simulation. + """ + self.sim.multilevel = [l for l in range(self.tot_level)] + # self.sim.mlne = + + self.sim.rawmap = [None] * self.tot_level + self.sim.ecl_coarse = [None] * self.tot_level + self.sim.well_cells = [None] * self.tot_level + + def ext_cov_mat_wgt(self): + # Make sure MULTILEVEL is a list + if not isinstance(self.keys_da['multilevel'][0], list): + mda_opts = [self.keys_da['multilevel']] + else: + mda_opts = self.keys_da['multilevel'] + + # Check if 'max_iter' has been given; if not, give error (mandatory in ITERATION) + assert 'cov_wgt' in list(zip(*mda_opts))[0], 'COV_WGT has not been given in MLDA!' + + # Extract max. iter + try: + cov_mat_wgt = [float(elem) for elem in [item[1] for item in mda_opts if item[0] == 'cov_wgt'][0]] + except: + cov_mat_wgt = [float(item[1]) for item in mda_opts if item[0]=='cov_wgt'] + Sum=0 + for i in range(len(cov_mat_wgt)): + Sum+=cov_mat_wgt[i] + for i in range(len(cov_mat_wgt)): + cov_mat_wgt[i]/=Sum + # Return max. iter + return cov_mat_wgt + + def init_ml_prior(self): + ''' + This function changes the structure of prior from independent + ensembles to a nested structure. + ''' + if self.ML_Nested: + ''' + for i in range(self.tot_level-1,0,-1): + TMP = self.ML_state[i][self.keys_da['staticvar']].shape + for j in range(i): + self.ML_state[j][self.keys_da['staticvar']][0:TMP[0], 0:TMP[1]] = \ + self.ML_state[i][self.keys_da['staticvar']] + ''' + #for el in self.state.keys(): + # self.state[el] = np.repeat(self.state[el][np.newaxis,:,:], self.tot_level,axis=0) + + self.state = [deepcopy(self.state) for _ in range(self.tot_level)] + for l in range(self.tot_level): + for el in self.state[0].keys(): + self.state[l][el] = self.state[l][el][:,:self.ml_ne[l]] + else: + # initiallize the state as an empty list of dictionaries with length equal self.tot_level + self.ml_state = [{} for _ in range(self.tot_level)] + # distribute the initial ensemble of states to the levels according to the given ensemble size. + start = 0 # intiallize + for l in range(self.tot_level): + stop = start + self.ml_ne[l] + for el in self.state.keys(): + self.ml_state[l][el] = self.state[el][:,start:stop] + start = stop + + del self.state + self.state = deepcopy(self.ml_state) + del self.ml_state + + def ext_data_size(self): + + # Make sure MULTILEVEL is a list + if not isinstance(self.keys_da['multilevel'][0], list): + mda_opts = [self.keys_da['multilevel']] + else: + mda_opts = self.keys_da['multilevel'] + + # Check if 'data_size' has been given + if not 'data_size' in list(zip(*mda_opts))[0]: # DATA_SIZE has not been given in MDA! + return None + + # Extract data_size + try: + data_size = [int(elem) for elem in [item[1] for item in mda_opts if item[0] == 'data_size'][0]] + except: + data_size = [int(item[1]) for item in mda_opts if item[0]=='data_size'] + + # Return data_size + return data_size + + +class mlhs_full(multilevel): + # sp_mfda_sim orig inherit this + ''' + Multilevel Hybrid Ensemble Smoother + ''' + + def __init__(self, keys_da,keys_fwd,sim): + """ + Standard initiallization + """ + super().__init__(keys_da,keys_fwd,sim) + + self.obs_data_BU = [] + for item in self.obs_data: + self.obs_data_BU.append(dict(item)) + + self.check_assimindex_simultaneous() + # define the assimilation index + + + self.check_fault() + + self.max_iter = 2 # No iterations + + def calc_analysis(self): + ''' + This class has been written based on the enkf class. It is designed for simultaneous assimilation of seismic + data with multilevel spatial data. + Using this calc_analysis tool we generate a seperate level-based covariance matrix for + every level and update them seperate from each other + This scheme updates each level based on the covariance and cross-covariance matrix + which are generated based on all the levels which are up/down-scaled to that specific level. + In short Modified Kristian's Idea + ''' + + # As the with statement in the ensemble code limits our access to the data and python does not seem to support + # pointers, I re-initialize some part of the code so that we'll access some essential data for the assimilaiton + #self.gen_ness_data() + self.obs_data = self.obs_data_BU + + self.B = [None] * self.tot_level + #self.B_gen(len(self.assim_index[1])) + + self.Dns_mat = [None] * self.tot_level + #self.Dns_mat_gen(len(self.assim_index[1])) + + #self.treat_modeling_error(0) + + # Generate the data auto-covariance matrix + #cov_data = self.gen_covdata(self.datavar, self.assim_index, self.list_datatypes) + + tot_pred = [] + self.Temp_State=[None]*self.tot_level + for level in range(self.tot_level): + obs_data_vector, pred = at.aug_obs_pred_data(self.obs_data, [time_dat[level] for time_dat in self.pred_data] + , self.assim_index, self.list_datatypes) # get some data + #pred = self.Dns_mat[level] * pred + tot_pred.append(pred) + + if not self.ML_error_corr == 'none': + if self.error_comp_scheme=='ens': + if self.ML_error_corr =='bias_corr': + L_mean = np.mean(tot_pred[-1], axis=1) + for l in range(self.tot_level-1): + tot_pred[l] += (L_mean - np.mean(tot_pred[l], axis=1))[:,np.newaxis] + + w_auto = self.cov_wgt + level_data_misfit = [None] * self.tot_level + if self.iteration == 1: # first iteration + misfit_data = 0 + for l in range(self.tot_level): + level_data_misfit[l] = at.calc_objectivefun(np.tile(obs_data_vector[:,np.newaxis],(1,self.ml_ne[l])), + tot_pred[l],self.cov_data) + misfit_data += w_auto[l] * np.mean(level_data_misfit[l]) + self.data_misfit = misfit_data + self.prior_data_misfit = misfit_data + self.prev_data_misfit = misfit_data + # Store the (mean) data misfit (also for conv. check) + if self.lam == 'auto': + self.lam = (0.5 * self.data_misfit)/len(self.obs_data_vector) + + self.logger.info(f'Prior run complete with data misfit: {self.prior_data_misfit:0.1f}. Lambda for initial analysis: {self.lam}') + + # Augment all the joint state variables (originally a dictionary) + aug_state = [at.aug_state(self.state[elem], self.list_states) for elem in range(self.tot_level)] + # concantenate all the elements + tot_aug_state = np.concatenate(aug_state, axis=1) + + # Mean state + mean_state = np.mean(tot_aug_state, 1) + + pert_state = [(aug_state[l] - np.dot(np.resize(mean_state, (len(mean_state), 1)), + np.ones((1, self.ml_ne[l])))) for l in + range(self.tot_level)] + + mean_preddata = [np.mean(tot_pred[elem], 1) for elem in range(self.tot_level)] + + tot_mean_preddata = sum([w_auto[elem] * np.mean(tot_pred[elem], 1) for elem in range(self.tot_level)]) + tot_mean_preddata /= sum([w_auto[elem] for elem in range(self.tot_level)]) + + # calculate the GMA covariances + pert_preddata = [(tot_pred[l] - np.dot(np.resize(mean_preddata[l], (len(mean_preddata[l]), 1)), + np.ones((1, self.ml_ne[l])))) for l in + range(self.tot_level)] + + self.update(pert_preddata,pert_state,mean_preddata,tot_mean_preddata,w_auto, tot_pred, aug_state) + + self.ML_state=self.Temp_State + + def gen_ness_data(self): + for i in range(self.tot_level): + self.ne=1 + self.sim.flow.level=i + self.level=i + assim_step=0 + assim_ind = [self.keys_da['obsname'], self.keys_da['assimindex'][assim_step]] + true_order = [self.keys_da['obsname'], self.keys_da['truedataindex']] + self.state=self.ML_state[i] + self.sim.setup_fwd_run(self.state, assim_ind, true_order) + os.mkdir(f'Test{i}') + folder=f'Test{i}'+os.sep + if self.Treat_Fault: + copyfile(f'IF/FL_{int(self.data_size[i])}.faults', 'IF/FL.faults') + self.sim.flow.run_fwd_sim(0, folder, wait_for_proc=True) + self.ecl_case = ecl.EclipseCase(f'Test{i}' + os.sep + self.sim.flow.file + + '.DATA') + tmp = self.ecl_case.cell_data('PORO') + self.sim.rawmap[self.level] = tmp + time.sleep(5) + for i in range(self.tot_level): + shutil.rmtree(f'Test{i}') + + def check_fault(self): + """ + Checks if there is a statement for generating a fault in the input file and if so + generates a synthetic fault based on the given input in there. + """ + # Make sure MULTILEVEL is a list + if not isinstance(self.keys_da['multilevel'][0], list): + fault_opts = [self.keys_da['multilevel']] + else: + fault_opts = self.keys_da['multilevel'] + + for i, opt in enumerate(list(zip(*self.keys_da['multilevel']))[0]): + if opt == 'generate_fault': + #options for ML_error_corr are: bias_corr, deterministic, stochastic, telescopic + fault_type=self.keys_da['multilevel'][i][1] + fault_dim=[float(item) for item in self.keys_da['multilevel'][i][2]] + if fault_type=='oblique': + self.generate_oblique_fault(fault_dim) + elif fault_type=='horizontal': + self.generate_horizontal_fault(fault_dim) + + self.Treat_Fault = False + for i, opt in enumerate(list(zip(*self.keys_da['multilevel']))[0]): + if opt=='treat_ml_fault': + self.Treat_Fault=True + + def generate_oblique_fault(self,fault_dim): + Dims=[int(self.prior_info[self.keys_da['staticvar']]['nx']), \ + int(self.prior_info[self.keys_da['staticvar']]['ny'])] + Margin=[int(np.floor(Dims[0]/10)),int(np.floor(Dims[1]/10))] + Temp_mat=np.zeros((Dims[0]-2*Margin[0],Dims[1]-2*Margin[1])) + for i in range(Temp_mat.shape[0]): + for j in range(Temp_mat.shape[1]): + if abs(i-j)<=np.floor(fault_dim[0]/2): + Temp_mat[i,Temp_mat.shape[1]-j-1]=fault_dim[1] + Temp_mat1=np.zeros((Dims[0],Dims[1])) + for i in range(Temp_mat.shape[0]): + for j in range(Temp_mat.shape[1]): + Temp_mat1[Margin[0]+i,Margin[1]+j]=Temp_mat[i,j] + Temp_mat = np.reshape(Temp_mat1, (np.product(Temp_mat1.shape), 1)) + for j in range(Temp_mat.shape[0]): + if Temp_mat[j, 0] != 0: + for l in range(self.tot_level): + for i in range(self.ml_ne[l]): + self.ML_state[l][self.keys_da['staticvar']][j,i]=Temp_mat[j] + + def generate_horizontal_fault(self,fault_dim): + Dims=[int(self.prior_info[self.keys_da['staticvar']]['nx']), \ + int(self.prior_info[self.keys_da['staticvar']]['ny'])] + Margin=[int(np.floor(Dims[0]/10)),int(np.floor(Dims[1]/10))] + Temp_mat=np.zeros((Dims[0]-2*Margin[0],Dims[1]-2*Margin[1])) + for i in range(int(fault_dim[0])): + for j in range(Temp_mat.shape[1]): + Temp_mat[int(Temp_mat.shape[0]/2)+i-int(fault_dim[0]/2),j]=fault_dim[1] + Temp_mat1=np.zeros((Dims[0],Dims[1])) + for i in range(Temp_mat.shape[0]): + for j in range(Temp_mat.shape[1]): + Temp_mat1[Margin[0]+i,Margin[1]+j]=Temp_mat[i,j] + Temp_mat = np.reshape(Temp_mat1, (np.product(Temp_mat1.shape), 1)) + for j in range(Temp_mat.shape[0]): + if Temp_mat[j, 0] != 0: + for l in range(self.tot_level): + for i in range(self.ml_ne[l]): + self.ML_state[l][self.keys_da['staticvar']][j,i]=Temp_mat[j] + + def B_gen(self,Multiplier): + for kk in range(self.tot_level): + self.level=kk + Ecl_coarse=self.sim.flow.ecl_coarse[self.level] + try: + Ecl_coarse = np.array(Ecl_coarse) + Ecl_coarse -= 1 + Rawmap_mask=self.sim.rawmap[self.level].mask + Rawmap_mask = Rawmap_mask[0, :, :] + ######### Notice !!!! + nx=self.prior_info[self.keys_da['staticvar']]['nx'] + ny=self.prior_info[self.keys_da['staticvar']]['ny'] + Shape=(nx,ny) + ######### + rows = np.zeros(Shape).flatten() + cols = np.zeros(Shape).flatten() + data = np.zeros(Shape).flatten() + mark = np.zeros(Shape).flatten() + Counter = 0 + for unit in range(Ecl_coarse.shape[0]): + I = None + J = None + Data = 1 / ((Ecl_coarse[unit, 3] - Ecl_coarse[unit, 2] + 1) * ( + Ecl_coarse[unit, 1] - Ecl_coarse[unit, 0] + 1)) + for i in range(Ecl_coarse[unit, 2], Ecl_coarse[unit, 3] + 1): + for j in range(Ecl_coarse[unit, 0], Ecl_coarse[unit, 1] + 1): + if Rawmap_mask[i, j] == False: + I = i + J = j + break + for i in range(Ecl_coarse[unit, 2], Ecl_coarse[unit, 3] + 1): + for j in range(Ecl_coarse[unit, 0], Ecl_coarse[unit, 1] + 1): + rows[Counter] = i * Shape[1] + j + cols[Counter] = I * Shape[1] + J + data[Counter] = Data + mark[i * Shape[1] + j] = 1 + Counter += 1 + + for i in range(Shape[0] * Shape[1]): + if mark[i] == 0: + rows[Counter] = i + cols[Counter] = i + data[Counter] = 1 + mark[i] = 1 + Counter += 1 + + rows = rows.reshape((Shape[0] * Shape[1], 1)) + cols = cols.reshape((Shape[0] * Shape[1], 1)) + data = data.reshape((Shape[0] * Shape[1], 1)) + COO = np.block([rows, cols, data]) + COO = COO[COO[:, 1].argsort(kind='mergesort')] + + Counter = 0 + for i in range(Shape[0] * Shape[1] - 1): + if COO[i, 1] != COO[i + 1, 1]: + Counter += 1 + + Counter = 0 + CXX = np.zeros(COO.shape) + CXX[0, 1] = Counter + CXX[0, 0] = COO[0, 0] + CXX[0, 2] = COO[0, 2] + for i in range(1, Shape[0] * Shape[1]): + if COO[i, 1] != COO[i - 1, 1]: + Counter += 1 + CXX[i, 1] = Counter + CXX[i, 0] = COO[i, 0] + CXX[i, 2] = COO[i, 2] + + S1 = self.data_size[self.level] + S2= CXX.shape[0] + Final_mat=np.zeros((CXX.shape[0]*Multiplier,CXX.shape[1])) + for i in range(Multiplier): + Final_mat[i*S2:(i+1)*S2,2]=CXX[:,2] + Final_mat[i*S2:(i+1)*S2,0]=CXX[:,0]+S2*i + Final_mat[i*S2:(i+1)*S2,1]=CXX[:,1]+S1*i + + rows = Final_mat[:, 1] + cols = Final_mat[:, 0] + data = Final_mat[:, 2] + + S2=np.product(Shape)*Multiplier + S1=self.data_size[self.level]*Multiplier + self.B[self.level]=coo_matrix((data,(rows,cols)),shape=(S1,S2)) + except: + COO=np.zeros((self.data_size[self.level]*Multiplier,3)) + S1=COO.shape[0] + for i in range(S1): + COO[i,0]=i + COO[i,1]=i + COO[i,2]=1 + self.B[self.level]=coo_matrix((COO[:,2],(COO[:,0],COO[:,1])),shape=(S1,S1)) + + def Dns_mat_gen(self, Multiplier): + for kk in range(self.tot_level): + self.level = kk + Ecl_coarse = self.sim.flow.ecl_coarse[self.level] + try: + Ecl_coarse = np.array(Ecl_coarse) + Ecl_coarse -= 1 + Rawmap_mask = self.sim.rawmap[self.level].mask + Rawmap_mask = Rawmap_mask[0, :, :] + ######### Notice !!!! + nx = self.prior_info[self.keys_da['staticvar']]['nx'] + ny = self.prior_info[self.keys_da['staticvar']]['ny'] + Shape = (nx, ny) + ######### + rows = np.zeros(Shape).flatten() + cols = np.zeros(Shape).flatten() + data = np.zeros(Shape).flatten() + mark = np.zeros(Shape).flatten() + Counter = 0 + for unit in range(Ecl_coarse.shape[0]): + I = None + J = None + for i in range(Ecl_coarse[unit, 2], Ecl_coarse[unit, 3] + 1): + for j in range(Ecl_coarse[unit, 0], Ecl_coarse[unit, 1] + 1): + if Rawmap_mask[i, j] == False: + I = i + J = j + break + for i in range(Ecl_coarse[unit, 2], Ecl_coarse[unit, 3] + 1): + for j in range(Ecl_coarse[unit, 0], Ecl_coarse[unit, 1] + 1): + rows[Counter] = i * Shape[1] + j + cols[Counter] = I * Shape[1] + J + data[Counter] = 1 + mark[i * Shape[1] + j] = 1 + Counter += 1 + + for i in range(Shape[0] * Shape[1]): + if mark[i] == 0: + rows[Counter] = i + cols[Counter] = i + data[Counter] = 1 + mark[i] = 1 + Counter += 1 + + rows = rows.reshape((Shape[0] * Shape[1], 1)) + cols = cols.reshape((Shape[0] * Shape[1], 1)) + data = data.reshape((Shape[0] * Shape[1], 1)) + COO = np.block([rows, cols, data]) + COO = COO[COO[:, 1].argsort(kind='mergesort')] + + Counter = 0 + for i in range(Shape[0] * Shape[1] - 1): + if COO[i, 1] != COO[i + 1, 1]: + Counter += 1 + + Counter = 0 + CXX = np.zeros(COO.shape) + CXX[0, 1] = Counter + CXX[0, 0] = COO[0, 0] + CXX[0, 2] = COO[0, 2] + for i in range(1, Shape[0] * Shape[1]): + if COO[i, 1] != COO[i - 1, 1]: + Counter += 1 + CXX[i, 1] = Counter + CXX[i, 0] = COO[i, 0] + CXX[i, 2] = COO[i, 2] + + S1 = self.data_size[self.level] + S2 = CXX.shape[0] + Final_mat = np.zeros((CXX.shape[0] * Multiplier, CXX.shape[1])) + for i in range(Multiplier): + Final_mat[i * S2:(i + 1) * S2, 2] = CXX[:, 2] + Final_mat[i * S2:(i + 1) * S2, 0] = CXX[:, 0] + S2 * i + Final_mat[i * S2:(i + 1) * S2, 1] = CXX[:, 1] + S1 * i + + rows = Final_mat[:, 0] + cols = Final_mat[:, 1] + data = Final_mat[:, 2] + + S2 = np.product(Shape) * Multiplier + S1 = self.data_size[self.level] * Multiplier + self.Dns_mat[self.level] = coo_matrix((data, (rows, cols)), shape=(S2, S1)) + except: + COO = np.zeros((self.data_size[self.level] * Multiplier, 3)) + S1 = COO.shape[0] + for i in range(S1): + COO[i, 0] = i + COO[i, 1] = i + COO[i, 2] = 1 + self.Dns_mat[self.level] = coo_matrix((COO[:, 2], (COO[:, 0], COO[:, 1])), shape=(S1, S1)) + + + def update(self,pert_preddata,pert_state,mean_preddata,tot_mean_preddata,w_auto, tot_pred, aug_state): + + #level_pert_preddata = [self.B[level] * pert_preddata[l] for l in range(self.tot_level)] + level_pert_preddata = [pert_preddata[l] for l in range(self.tot_level)] + #level_mean_preddata = [self.B[level] * mean_preddata[l] for l in range(self.tot_level)] + level_mean_preddata = [mean_preddata[l] for l in range(self.tot_level)] + #level_tot_mean_preddata = self.B[level] * tot_mean_preddata + level_tot_mean_preddata = tot_mean_preddata + + cov_auto = sum([w_auto[l] * at.calc_autocov(level_pert_preddata[l]) for l in range(self.tot_level)]) + \ + sum([w_auto[l] * np.outer((level_mean_preddata[l] - level_tot_mean_preddata), + (level_mean_preddata[l] - level_tot_mean_preddata)) for l in + range(self.tot_level)]) + cov_auto /= sum([w_auto[l] for l in range(self.tot_level)]) + + cov_cross = sum([w_auto[l] * at.calc_crosscov(pert_state[l], level_pert_preddata[l]) + for l in range(self.tot_level)]) + cov_cross /= sum([w_auto[l] for l in range(self.tot_level)]) + + #joint_data_cov = self.B[level] * self.cov_data * self.B[level].transpose() + joint_data_cov = self.cov_data + + kalman_gain_param = self.calc_kalmangain(cov_cross, cov_auto, + joint_data_cov) # global cov_cross and cov_auto + + for level in range(self.tot_level): + obs_data = self.efficient_real_gen(self.obs_data_vector, self.cov_data, self.ml_ne[level], \ + level) + + #level_tot_pred = self.B[level] * tot_pred[level] + level_tot_pred = tot_pred[level] + aug_state_upd = at.calc_kalman_filter_eq(aug_state[level], kalman_gain_param, obs_data, + level_tot_pred) # update levelwise + + self.Temp_State[level] = at.update_state(aug_state_upd, self.state[level], self.list_states) + + def efficient_real_gen(self, mean, var, number, level,original_size=False, limits=None, return_chol=False): + """ + This function is added to prevent additional computational cost if var is diagonal + MN 04/20 + """ + if not original_size: + var = np.array(var) #to enable var.shape + parsize = len(mean) + if parsize == 1 or len(var.shape) == 1: + l = np.sqrt(var) + # real = mean + L*np.random.randn(1, number) + else: + # Check if the covariance matrix is diagonal (only entries in the main diagonal). If so, we can use + # numpy.sqrt for efficiency + if 4==2: #np.count_nonzero(var - np.diagonal(var)) == 0: + l = np.sqrt(var) # only variance (diagonal) term + l=np.reshape(l,(l.size,1)) + else: + # Cholesky decomposition + l = linalg.cholesky(var) # cov. matrix has off-diag. terms + #Mean=deepcopy(mean) + Mean=np.reshape(mean,(mean.size,1)) + #Mean=self.B[level]*Mean + # Gen. realizations + # if len(var.shape) == 1: + # real = np.dot(Mean, np.ones((1, number))) + np.expand_dims((self.B[level]*l).flatten(), axis=1)*np.random.randn( + # np.size(Mean), number) + # else: + # real = np.tile(Mean, (1, number)) + np.dot(self.B[level]*l.T, np.random.randn(np.size(mean), + # number)) + if len(var.shape) == 1: + real = np.dot(Mean, np.ones((1, number))) + np.expand_dims((l).flatten(), axis=1)*np.random.randn( + np.size(Mean), number) + else: + real = np.tile(Mean, (1, number)) + np.dot(l.T, np.random.randn(np.size(mean), + number)) + + # Truncate values that are outside limits + # TODO: Make better truncation rules, or switch truncation on/off + if limits is not None: + # Truncate + real[real > limits['upper']] = limits['upper'] + real[real < limits['lower']] = limits['lower'] + + if return_chol: + return real, l + else: + return real + else: + var = np.array(var) # to enable var.shape + parsize = len(mean) + if parsize == 1 or len(var.shape) == 1: + l = np.sqrt(var) + # real = mean + L*np.random.randn(1, number) + else: + # Check if the covariance matrix is diagonal (only entries in the main diagonal). If so, we can use + # numpy.sqrt for efficiency + if 4 == 2: # np.count_nonzero(var - np.diagonal(var)) == 0: + l = np.sqrt(var) # only variance (diagonal) term + l = np.reshape(l, (l.size, 1)) + else: + # Cholesky decomposition + l = linalg.cholesky(var) # cov. matrix has off-diag. terms + # Mean=deepcopy(mean) + Mean = np.reshape(mean, (mean.size, 1)) + # Gen. realizations + if len(var.shape) == 1: + real = np.dot(Mean, np.ones((1, number))) + np.expand_dims((l).flatten(), + axis=1) * np.random.randn( + np.size(Mean), number) + else: + real = np.tile(Mean, (1, number)) + np.dot(l.T, np.random.randn(np.size(mean), + number)) + + # Truncate values that are outside limits + # TODO: Make better truncation rules, or switch truncation on/off + if limits is not None: + # Truncate + real[real > limits['upper']] = limits['upper'] + real[real < limits['lower']] = limits['lower'] + + if return_chol: + return real, l + else: + return real + def calc_kalmangain(self, cov_cross, cov_auto, cov_data, opt=None): + """ + Calculate the Kalman gain + Using mainly two options: linear soultion and pseudo inverse of the matrix + MN 04/2020 + """ + if opt is None: + calc_opt = 'lu' + + # Add data and predicted data auto-covariance matrices + if len(cov_data.shape)==1: + cov_data = np.diag(cov_data) + c_auto = cov_auto + cov_data + + if calc_opt == 'lu': + try: + kg = linalg.solve(c_auto.T, cov_cross.T) + kalman_gain = kg.T + except: + #Margin=10**5 + #kalman_gain = cov_cross * self.calc_pinv(c_auto, Margin=Margin) + #kalman_gain = cov_cross * self.calc_pinv(c_auto) + kalman_gain = cov_cross * np.linalg.pinv(c_auto) + #kalman_gain = cov_cross * np.linalg.pinv(c_auto, rcond=10**(-15)) + + elif calc_opt == 'chol': + # Cholesky decomp (upper triangular matrix) + u = linalg.cho_factor(c_auto.T, check_finite=False) + + # Solve linear system with cholesky square-root + kalman_gain = linalg.cho_solve(u, cov_cross.T, check_finite=False) + + # Return Kalman gain + return kalman_gain + + def check_convergence(self): + """ + Check if LM-EnRML have converged based on evaluation of change sizes of objective function, state and damping + parameter. + + Returns + ------- + conv: bool + Logic variable telling if algorithm has converged + why_stop: dict + Dict. with keys corresponding to conv. criteria, with logical variable telling which of them that has been + met + """ + success = False # init as false + + if hasattr(self, 'list_datatypes'): + assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] + list_datatypes = self.list_datatypes +# cov_data = self.gen_covdata(self.datavar, assim_index, list_datatypes) + pred_data = [None] * self.tot_level + level_mean_preddata = [None] * self.tot_level + for l in range(self.tot_level): + obs_data_vector, pred_data[l] = at.aug_obs_pred_data(self.obs_data, + [time_dat[l] for time_dat in self.pred_data], + assim_index, list_datatypes) + level_mean_preddata[l] = np.mean(pred_data[l], 1) + else: + assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] + list_datatypes, _ = at.get_list_data_types(self.obs_data, assim_index) + # cov_data = at.gen_covdata(self.datavar, assim_index, list_datatypes) + #cov_data = self.gen_covdata(self.datavar, assim_index, list_datatypes) + pred_data = [None] * self.tot_level + level_mean_preddata = [None] * self.tot_level + for l in range(self.tot_level): + obs_data_vector, pred_data[l] = at.aug_obs_pred_data(self.obs_data, + [time_dat[l] for time_dat in self.pred_data], + assim_index, list_datatypes) + level_mean_preddata[l] = np.mean(pred_data[l], 1) + + # self.prev_data_misfit_std = self.data_misfit_std + # if there was no reduction of the misfit, retain the old "valid" data misfit. + + # Calc. std dev of data misfit (used to update lamda) + # mat_obs = np.dot(obs_data_vector.reshape((len(obs_data_vector),1)), np.ones((1, self.ne))) # use the perturbed + # data instead. + # mat_obs = self.real_obs_data + level_data_misfit = [None] * self.tot_level + #list_states = list(self.state.keys()) + #cov_prior = at.block_diag_cov(self.cov_prior, list_states) + #ML_prior_state = [at.aug_state(self.ML_prior_state[elem], list_states) for elem in range(self.tot_level)] + #ML_state = [at.aug_state(self.state[elem], list_states) for elem in range(self.tot_level)] + # level_state_misfit = [None] * self.tot_level + # if len(self.cov_data.shape) == 1: + for l in range(self.tot_level): + + level_data_misfit[l] = at.calc_objectivefun(np.tile(obs_data_vector[:,np.newaxis],(1,self.ml_ne[l])), + pred_data[l],self.cov_data) + +# obs_data = self.Dns_mat[l] * self.obs_reals[l] + ##### This part is not done correctly as we do not need it now!!! ###### + # level_data_misfit[l] = np.diag(np.dot((pred_data[l] - obs_data).T * self.Dns_mat[l].transpose() * + # self.B[0].transpose(), + # np.dot(np.expand_dims(self.cov_data ** (-1), axis=1), + # np.ones((1, self.ne))) * self.B[0] * self.Dns_mat[l] * ( + # pred_data[l] - obs_data))) + #level_state_misfit[l] = np.diag(np.dot((ML_state[l] - ML_prior_state[l]).T, solve( + # cov_prior, (ML_state[l] - ML_prior_state[l])))) + # else: + # for l in range(self.tot_level): + # obs_data = self.Dns_mat[l]*self.obs_reals[l] + # obs_data = self.obs_reals[l] + # ''' + # level_data_misfit[l] = np.diag(np.dot((pred_data [l]- obs_data).T*self.Dns_mat[l].transpose()* + # self.B[0].transpose(),solve(self.B[0]*cov_data*self.B[0].transpose(), + # self.B[0]*self.Dns_mat[l]*(pred_data[l] - obs_data)))) + # level_state_misfit[l]=np.diag(np.dot((ML_state[l]-ML_prior_state[l]).T,solve( + # cov_prior,(ML_state[l]-ML_prior_state[l])))) + # ''' + # level_data_misfit[l] = np.diag(np.dot((pred_data[l] - obs_data).T * self.Dns_mat[l].transpose(), + # solve(self.cov_data, self.Dns_mat[l] * (pred_data[l] - obs_data)))) + + misfit_data = 0 +# misfit_state = 0 + w_auto = self.cov_wgt + for l in range(self.tot_level): + misfit_data += w_auto[l] * np.mean(level_data_misfit[l]) + # misfit_state+=w_auto[l]*np.mean(level_state_misfit[l]) + + self.data_misfit = misfit_data + # self.data_misfit_std = np.std(con_misfit) + + # # Calc. mean data misfit for convergence check, using the updated state variable + # self.data_misfit = np.dot((mean_preddata - obs_data_vector).T, + # solve(cov_data, (mean_preddata - obs_data_vector))) + + # Convergence check: Relative step size of data misfit or state change less than tolerance + why_stop = {} # todo: populate + + # update the last mismatch, only if this was a reduction of the misfit + if self.data_misfit < self.prev_data_misfit: + success = True + + + if success: + self.logger.info(f'ML Hybrid Smoother update complete! Objective function reduced from ' + f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') + # self.prev_data_misfit = self.data_misfit + # self.prev_data_misfit_std = self.data_misfit_std + else: + self.logger.info(f'ML Hybrid Smoother update complete! Objective function increased from ' + f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') + + # Return conv = False, why_stop var. + return False, True, why_stop + +class smlses_s(multilevel,esmda_approx): + """ + The Sequential multilevel ensemble smoother with the "straightforward" flavour as descibed in Nezhadali, M., + Bhakta, T., Fossum, K., & Mannseth, T. (2023). Sequential multilevel assimilation of inverted seismic data. + Computational Geosciences, 27(2), 265–287. https://doi.org/10.1007/s10596-023-10191-9 + + Since the update schemes are basically a esmda update we inherit the esmda_approx method. Hence, we only have to + care about handling the multi-level features. + """ + + def __init__(self,keys_da, keys_fwd, sim): + super().__init__(keys_da, keys_fwd, sim) + + self.current_state = [self.current_state[0]] + self.state = [self.state[0]] + + # Overwrite the method for extracting ml_information. Here, we should only get the first level + def _ext_ml_info(self, grab_level=0): + ''' + Extract the info needed for ML simulations. Grab the first level info + ''' + + if 'multilevel' in self.keys_en: + # parse + self.multilevel = {} + for i, opt in enumerate(list(zip(*self.keys_en['multilevel']))[0]): + if opt == 'levels': + self.multilevel['levels'] = [elem for elem in range( + int(self.keys_en['multilevel'][i][1]))] + if opt == 'en_size': + self.multilevel['ne'] = [range(int(el)) + for el in self.keys_en['multilevel'][i][1]] + try: + self.multilevel['levels'] = [self.multilevel['levels'][grab_level]] + except IndexError: # When converged, we need to set the level to the final one + self.multilevel['levels'] = [self.multilevel['levels'][-1]] + #self.multilevel['ne'] = [self.multilevel['ne'][grab_level]] + def calc_analysis(self): + # Some preamble for multilevel + # Do this. + # flatten the level element of the predicted data + tmp = [] + for elem in self.pred_data: + tmp += elem + self.pred_data = tmp + + self.current_state = self.current_state[self.multilevel['levels'][0]] + self.state = self.state[self.multilevel['levels'][0]] + # call the inherited version via super() + super().calc_analysis() + + # Afterwork + self._ext_ml_info(grab_level=self.iteration) + + # Grab the prior for the next mda step. Draw the top scoring values. + self._update_ensemble() + + def _update_ensemble(self): + # Prelude to calc. conv. check (everything done below is from calc_analysis) + obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, + self.list_datatypes) + + data_misfit = at.calc_objectivefun( + self.real_obs_data_conv, pred_data, self.cov_data) + + # sort the data_misfit after the percentile score + sort_ind = np.argsort(data_misfit)[self.multilevel['ne'][self.multilevel['levels'][0]]] + + # initialize self.current_state and self.state as empty lists with lenght equal to self.multilevel['levels'][0] + tmp_current_state = [[] for _ in range(self.multilevel['levels'][0]+1)] + tmp_state = [[] for _ in range(self.multilevel['levels'][0]+1)] + + tmp_current_state[self.multilevel['levels'][0]] = {el:self.current_state[el][:,sort_ind] for el in self.current_state.keys()} + tmp_state[self.multilevel['levels'][0]] = {el:self.state[el][:,sort_ind] for el in self.state.keys()} + + + #reduce the size of these ensembles as well + self.real_obs_data_conv = self.real_obs_data_conv[:,sort_ind] + self.real_obs_data = self.real_obs_data[:,sort_ind] + + # set the current state and state to the new values + self.current_state = tmp_current_state + self.state = tmp_state + + # update self.ne to be inline with new ensemble size + self.ne = len(self.multilevel['ne'][self.multilevel['levels'][0]]) + + # and update the projection to be inline with new ensemble size + self.proj = (np.eye(self.ne) - (1 / self.ne) * + np.ones((self.ne, self.ne))) / np.sqrt(self.ne - 1) + + def check_convergence(self): + """ + Check convergence for the smlses-s method + """ + + self.prev_data_misfit = self.data_misfit + self.prev_data_misfit_std = self.data_misfit_std + + # extract pred_data for the current level + level_pred_data = [el[0] for el in self.pred_data] + + # Prelude to calc. conv. check (everything done below is from calc_analysis) + obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, level_pred_data, self.assim_index, + self.list_datatypes) + + data_misfit = at.calc_objectivefun( + self.real_obs_data_conv, pred_data, self.cov_data) + self.data_misfit = np.mean(data_misfit) + self.data_misfit_std = np.std(data_misfit) + + # Logical variables for conv. criteria + why_stop = {'rel_data_misfit': 1 - (self.data_misfit / self.prev_data_misfit), + 'data_misfit': self.data_misfit, + 'prev_data_misfit': self.prev_data_misfit} + + if self.data_misfit < self.prev_data_misfit: + self.logger.info( + f'ML-MDA iteration number {self.iteration}! Objective function reduced from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') + else: + self.logger.info( + f'ML-MDA iteration number {self.iteration}! Objective function increased from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') + # Return conv = False, why_stop var. + self.current_state = deepcopy(self.state) + + return False, True, why_stop + +class esmda_h(multilevel,hybrid_update,esmdaMixIn): + ''' + A multilevel implementation of the ES-MDA algorithm with the hybrid gain + ''' + + def __init__(self,keys_da, keys_fwd, sim): + super().__init__(keys_da, keys_fwd, sim) + + self.proj = [(np.eye(self.ml_ne[l]) - (1 / self.ml_ne[l]) * + np.ones((self.ml_ne[l], self.ml_ne[l]))) / np.sqrt(self.ml_ne[l] - 1) for l in range(self.tot_level)] + + def calc_analysis(self): + self.aug_pred_data = [] + for l in range(self.tot_level): + self.aug_pred_data.append(at.aug_obs_pred_data(self.obs_data, [el[l] for el in self.pred_data], self.assim_index, + self.list_datatypes)[1]) + + init_en = Cholesky() # Initialize GeoStat class for generating realizations + if self.iteration == 1: # first iteration + # note, evaluate for high fidelity model + data_misfit = at.calc_objectivefun( + self.real_obs_data_conv, np.concatenate(self.aug_pred_data,axis=1), self.cov_data) + + # Store the (mean) data misfit (also for conv. check) + self.data_misfit = np.mean(data_misfit) + self.prior_data_misfit = np.mean(data_misfit) + self.prior_data_misfit_std = np.std(data_misfit) + self.data_misfit = np.mean(data_misfit) + self.data_misfit_std = np.std(data_misfit) + + self.logger.info( + f'Prior run complete with data misfit: {self.prior_data_misfit:0.1f}.') + self.data_random_state = deepcopy(np.random.get_state()) + self.real_obs_data = [] + self.scale_data = [] + for l in range(self.tot_level): + # populate the lists without unpacking the output form init_en.gen_real + (lambda x,y: (self.real_obs_data.append(x),self.scale_data.append(y)))(*init_en.gen_real(self.obs_data_vector, + self.alpha[self.iteration - 1] * + self.cov_data, self.ml_ne[l], + return_chol=True)) + self.E = [np.dot(self.real_obs_data[l], self.proj[l]) for l in range(self.tot_level)] + else: + self.data_random_state = deepcopy(np.random.get_state()) + # self.obs_data_vector, _ = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, + # self.list_datatypes) + for l in range(self.tot_level): + self.real_obs_data[l], self.scale_data[l] = init_en.gen_real(self.obs_data_vector, + self.alpha[self.iteration - + 1] * self.cov_data, + self.ml_ne[l], + return_chol=True) + self.E[l] = np.dot(self.real_obs_data[l], self.proj[l]) + + self.pert_preddata = [] + for l in range(self.tot_level): + if len(self.scale_data[l].shape) == 1: + self.pert_preddata.append(np.dot(np.expand_dims(self.scale_data[l] ** (-1), axis=1), + np.ones((1, self.ml_ne[l]))) * np.dot(self.aug_pred_data[l], self.proj[l])) + else: + self.pert_preddata.append(solve( + self.scale_data[l], np.dot(self.aug_pred_data[l], self.proj[l]))) + + aug_state= [] + for l in range(self.tot_level): + aug_state.append(at.aug_state(self.current_state[l], self.list_states)) + + self.update() + if hasattr(self, 'step'): + aug_state_upd = [aug_state[l] + self.step[l] for l in range(self.tot_level)] + # if hasattr(self, 'w_step'): + # self.W = self.current_W + self.w_step + # aug_prior_state = at.aug_state(self.prior_state, self.list_states) + # aug_state_upd = np.dot(aug_prior_state, (np.eye( + # self.ne) + self.W / np.sqrt(self.ne - 1))) + + # Extract updated state variables from aug_update + for l in range(self.tot_level): + self.state[l] = at.update_state(aug_state_upd[l], self.state[l], self.list_states) + self.state[l] = at.limits(self.state[l], self.prior_info) + + def check_convergence(self): + """ + Check ESMDA objective function for logging purposes. + """ + + self.prev_data_misfit = self.data_misfit + self.prev_data_misfit_std = self.data_misfit_std + + # Prelude to calc. conv. check (everything done below is from calc_analysis) + pred_data = [] + for l in range(self.tot_level): + pred_data.append(at.aug_obs_pred_data(self.obs_data, [el[l] for el in self.pred_data], self.assim_index, + self.list_datatypes)[1]) + + data_misfit = at.calc_objectivefun( + self.real_obs_data_conv, np.concatenate(pred_data,axis=1), self.cov_data) + self.data_misfit = np.mean(data_misfit) + self.data_misfit_std = np.std(data_misfit) + + # Logical variables for conv. criteria + why_stop = {'rel_data_misfit': 1 - (self.data_misfit / self.prev_data_misfit), + 'data_misfit': self.data_misfit, + 'prev_data_misfit': self.prev_data_misfit} + + if self.data_misfit < self.prev_data_misfit: + self.logger.info( + f'MDA iteration number {self.iteration}! Objective function reduced from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') + else: + self.logger.info( + f'MDA iteration number {self.iteration}! Objective function increased from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') + # Return conv = False, why_stop var. + self.current_state = deepcopy(self.state) + if hasattr(self, 'W'): + self.current_W = deepcopy(self.W) + + return False, True, why_stop + +class esmda_seq_h(multilevel,esmda_approx): + ''' + A multilevel implementation of the Sequeontial ES-MDA algorithm with the hybrid gain + ''' + + def __init__(self,keys_da, keys_fwd, sim): + super().__init__(keys_da, keys_fwd, sim) + + self.proj = (np.eye(self.ml_ne[0]) - (1 / self.ml_ne[0]) * + np.ones((self.ml_ne[0], self.ml_ne[0]))) / np.sqrt(self.ml_ne[0] - 1) + + self.multilevel['levels'] = [self.iteration] + + self.ne = self.ml_ne[0] + # adjust the real_obs_data to only containt the first ne samples + self.real_obs_data_conv = self.real_obs_data_conv[:,:self.ne] + + def calc_analysis(self): + + # collapse the level element of the predicted data + self.ml_pred = deepcopy(self.pred_data) + # concantenate the ml_pred data and state + self.pred_data = [] + curr_level = self.multilevel['levels'][0] + for level_pred_date in self.ml_pred: + keys = level_pred_date[curr_level].keys() + result ={} + for key in keys: + arrays = np.array([level_pred_date[curr_level][key]]) + result[key] = np.hstack(arrays) + self.pred_data.append(result) + + self.ml_state = deepcopy(self.state) + self.state = self.state[self.multilevel['levels'][0]] + self.current_state = self.current_state[self.multilevel['levels'][0]] + + super().calc_analysis() + + # Set the multilevel index and set the dimentions for all the states + self.multilevel['levels'][0] += 1 + self.ne = self.ml_ne[self.multilevel['levels'][0]] + self.proj =(np.eye(self.ne) - (1 / self.ne) * + np.ones((self.ne, self.ne))) / np.sqrt(self.ne - 1) + best_members = np.argsort(self.ensemble_misfit)[:self.ne] + self.ml_state[self.multilevel['levels'][0]] = {k: v[:, best_members] for k, v in self.state.items()} + self.state = deepcopy(self.ml_state) + + self.real_obs_data_conv = self.real_obs_data_conv[:,best_members] + + + + def check_convergence(self): + """ + Check ESMDA objective function for logging purposes. + """ + + self.prev_data_misfit = self.data_misfit + #self.prev_data_misfit_std = self.data_misfit_std + + # Prelude to calc. conv. check (everything done below is from calc_analysis) + pred_data = [] + for l in range(len(self.pred_data[0])): + level_pred = at.aug_obs_pred_data(self.obs_data, [el[l] for el in self.pred_data], self.assim_index, + self.list_datatypes)[1] + if level_pred is not None: # Can be None if level is not predicted + pred_data.append(level_pred) + + data_misfit = at.calc_objectivefun( + self.real_obs_data_conv, np.concatenate(pred_data,axis=1), self.cov_data) + self.ensemble_misfit = data_misfit + self.data_misfit = np.mean(data_misfit) + self.data_misfit_std = np.std(data_misfit) + + # Logical variables for conv. criteria + why_stop = {'rel_data_misfit': 1 - (self.data_misfit / self.prev_data_misfit), + 'data_misfit': self.data_misfit, + 'prev_data_misfit': self.prev_data_misfit} + + if self.data_misfit < self.prev_data_misfit: + self.logger.info( + f'MDA iteration number {self.iteration}! Objective function reduced from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') + else: + self.logger.info( + f'MDA iteration number {self.iteration}! Objective function increased from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') + # Return conv = False, why_stop var. + self.current_state = deepcopy(self.state) + if hasattr(self, 'W'): + self.current_W = deepcopy(self.W) + + return False, True, why_stop \ No newline at end of file diff --git a/pipt/update_schemes/update_methods_ns/hybrid_udpate.py b/pipt/update_schemes/update_methods_ns/hybrid_udpate.py new file mode 100644 index 0000000..5409914 --- /dev/null +++ b/pipt/update_schemes/update_methods_ns/hybrid_udpate.py @@ -0,0 +1,65 @@ +""" +ES, and Iterative ES updates with hybrid update matrix calculated from multi-fidelity runs. +""" + +import numpy as np +from scipy.linalg import solve +from pipt.misc_tools import analysis_tools as at + +class hybrid_update: + ''' + Class for hybrid update schemes as described in: Fossum, K., Mannseth, T., & Stordal, A. S. (2020). Assessment of + multilevel ensemble-based data assimilation for reservoir history matching. Computational Geosciences, 24(1), + 217–239. https://doi.org/10.1007/s10596-019-09911-x + + Note that the scheme is slightly modified to be inline with the standard (I)ES approximate update scheme. This + enables the scheme to efficiently be coupled with multiple updating strategies via class MixIn + ''' + + def update(self): + x_3 = [] + pert_state = [] + for l in range(self.tot_level): + aug_state = at.aug_state(self.current_state[l], self.list_states, self.cell_index) + mean_state = np.mean(aug_state, 1) + if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': + pert_state.append((self.state_scaling**(-1))[:, None] * (aug_state - np.dot(np.resize(mean_state, (len(mean_state), 1)), + np.ones((1, self.ml_ne[l]))))) + else: + pert_state.append((self.state_scaling**(-1) + )[:, None] * np.dot(aug_state, self.proj[l])) + + u_d, s_d, v_d = np.linalg.svd(self.pert_preddata[l], full_matrices=False) + if self.trunc_energy < 1: + ti = (np.cumsum(s_d) / sum(s_d)) <= self.trunc_energy + u_d, s_d, v_d = u_d[:, ti].copy(), s_d[ti].copy(), v_d[ti, :].copy() + + # x_1 = np.dot(u_d.T, solve(self.scale_data[l], + # (self.real_obs_data[l] - self.aug_pred_data[l]))) + + x_2 = solve(((self.lam + 1) * np.eye(len(s_d)) + np.diag(s_d ** 2)), u_d.T) + x_3.append(np.dot(np.dot(v_d.T, np.diag(s_d)), x_2)) + + # Calculate each row of self.step individually to avoid memory issues. + self.step = [np.empty(pert_state[l].shape) for l in range(self.tot_level)] + + # do maximum 1000 rows at a time. + step_size = min(1000, int(self.state_scaling.shape[0]/2)) + row_step = [np.arange(start, start+step_size) for start in + np.arange(0, self.state_scaling.shape[0]-step_size, step_size)] + #add the last rows + row_step.append(np.arange(row_step[-1][-1]+1, self.state_scaling.shape[0])) + + for row in row_step: + kg = sum([self.cov_wgt[indx_l]*np.dot(pert_state[indx_l][row, :], x_3[indx_l]) for indx_l in + range(self.tot_level)]) + for l in range(self.tot_level): + if len(self.scale_data[l].shape) == 1: + self.step[l][row, :] = np.dot(self.state_scaling[row, None] * kg, + np.dot(np.expand_dims(self.scale_data[l] ** (-1), axis=1), + np.ones((1, self.ml_ne[l]))) * + (self.real_obs_data[l] - self.aug_pred_data[l])) + else: + self.step[l][row, :] = np.dot(self.state_scaling[row, None] * kg, solve(self.scale_data[l], + (self.real_obs_data[l] - + self.aug_pred_data[l]))) \ No newline at end of file From a7e1af233a648e91d19e294010c7ffd202e51445 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 13 Jan 2026 13:59:06 +0100 Subject: [PATCH 80/94] Rewrite multilevel code --- ensemble/ensemble.py | 24 +- pipt/misc_tools/extract_tools.py | 38 +- pipt/update_schemes/multilevel.py | 1175 ++--------------- .../update_methods_ns/hybrid_udpate.py | 72 +- popt/loop/ensemble_gaussian.py | 3 +- 5 files changed, 209 insertions(+), 1103 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 74dc154..601f58c 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -161,9 +161,11 @@ def __init__(self, keys_en: dict, sim, redund_sim=None): self.list_states = list(self.keys_en['staticvar']) if 'multilevel' in self.keys_en: - ml_info = extract.extract_multilevel_info(self.keys_en) - self.multilevel, self.tot_level, self.ml_ne, self.ML_error_corr, self.error_comp_scheme, self.ML_corr_done = ml_info - + self.multilevel = extract.extract_multilevel_info(self.keys_en['multilevel']) + self.ml_ne = self.multilevel['ml_ne'] + self.tot_level = int(self.multilevel['levels']) + self.ml_corr_done = False + def get_list_assim_steps(self): """ @@ -508,16 +510,18 @@ def calc_ml_prediction(self, enX=None): return success def treat_modeling_error(self): - if not self.ML_error_corr=='none': - if self.error_comp_scheme=='sep': + if self.multilevel['ml_error_corr']: + scheme = self.multilevel['ml_error_corr'][1] + + if scheme =='sep': self.calc_modeling_error_sep() self.address_ML_error() - elif self.error_comp_scheme=='once': - if not self.ML_corr_done: + elif scheme =='once': + if not self.ml_corr_done: self.calc_modeling_error_ens() - self.ML_corr_done = True + self.ml_corr_done = True self.address_ML_error() - elif self.error_comp_scheme=='ens': + elif scheme =='ens': self.calc_modeling_error_ens() def calc_modeling_error_sep(self): @@ -525,7 +529,7 @@ def calc_modeling_error_sep(self): def calc_modeling_error_ens(self): - if self.ML_error_corr =='bias_corr': + if self.multilevel['ml_error_corr'][0] =='bias_corr': # modify self.pred_data without changing its structure. Hence, for each level (except the finest one) # we correct each data at each point in time. for assim_index in range(len(self.pred_data)): diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index 06865de..13c1fe8 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -289,30 +289,32 @@ def extract_multilevel_info(keys: Union[dict, list]) -> dict: such that we only have one level -- the high fidelity one ''' if isinstance(keys, list): - ml_info = list_to_dict(keys) - assert isinstance(ml_info, dict) + keys_ml = list_to_dict(keys) + assert isinstance(keys_ml, dict) # Set levels - levels = int(ml_info['levels']) - ml_info['levels'] = [elem for elem in range(levels)] + assert 'levels' in keys_ml, 'LEVELS keyword missing in MULTILEVEL!' + levels = int(keys_ml['levels']) + keys_ml['levels'] = [elem for elem in range(levels)] # Set multi-level ensemble size - en_size = ml_info.pop('en_size') - ml_info['ne'] = [range(int(elem)) for elem in en_size] - ml_ne = [int(elem) for elem in en_size] + assert 'en_size' in keys_ml, 'EN_SIZE keyword missing in MULTILEVEL!' + en_size = keys_ml.pop('en_size') + keys_ml['ne'] = [range(int(elem)) for elem in en_size] + keys_ml['ml_ne'] = [int(elem) for elem in en_size] + assert len(keys_ml['ml_ne']) == levels, 'The Ensemble Size must be specified for all levels!' + + # Set weights + assert 'ml_weights' in keys_ml or 'cov_wgt' in keys_ml, 'ML_WEIGHTS (or COV_WGT) keyword missing in MULTILEVEL!' + if 'cov_wgt' in keys_ml: + keys_ml['ml_weights'] = keys_ml.pop('cov_wgt') + if not np.sum(keys_ml['ml_weights']) == 1.0: + keys_ml['ml_weights'] = keys_ml['ml_weights']/np.sum(keys_ml['ml_weights']) # Set multi-level error - if not 'ml_error_corr' in ml_info: - ml_error_corr = 'none' - else: - ml_error_corr = ml_info['ml_error_corr'][0] - ml_corr_done = False - - if not ml_error_corr == 'none': - error_comp_scheme = ml_info['ml_error_corr'][1] - - # set attribute - return ml_info, levels, ml_ne, ml_error_corr, error_comp_scheme, ml_corr_done + keys_ml['ml_error_corr'] = keys_ml.get('ml_error_corr', None) + + return keys_ml def extract_local_analysis_info(keys: Union[dict, list], state: list) -> dict: diff --git a/pipt/update_schemes/multilevel.py b/pipt/update_schemes/multilevel.py index 4b4b610..1555ee3 100644 --- a/pipt/update_schemes/multilevel.py +++ b/pipt/update_schemes/multilevel.py @@ -8,6 +8,7 @@ from pipt.update_schemes.esmda import esmda_approx from pipt.update_schemes.esmda import esmdaMixIn from pipt.misc_tools import analysis_tools as at +import pipt.misc_tools.ensemble_tools as entools from geostat.decomp import Cholesky from misc import ecl @@ -31,20 +32,24 @@ import math +__all__ = ['multilevel', 'esmda_hybrid'] + class multilevel(Ensemble): """ Inititallize the multilevel class. Similar for all ML schemes, hence make one class for all. """ def __init__(self, keys_da,keys_fwd,sim): super().__init__(keys_da, keys_fwd, sim) - self._ext_ml_feat() - #self.ML_state = [{} for _ in range(self.tot_level)] - #self.ML_state[0] = deepcopy(self.state) - self.data_size = self.ext_data_size() - self.list_states = list(self.state.keys()) - self.init_ml_prior() - self.prior_state = deepcopy(self.state) + + self.list_states = list(self.idX.keys()) + + # Reorganize prior ensemble to multilevel structure if nested is true + self.enX = self.reorganize_ml_prior(self.enX) + self.prior_enX = deepcopy(self.enX) + + # Set ML specific options for simulator self._init_sim() + self.iteration = 0 self.lam = 0 # set LM lamda to zero as we are doing one full update. if 'energy' in self.keys_da: @@ -55,957 +60,72 @@ def __init__(self, keys_da,keys_fwd,sim): self.trunc_energy = 0.98 self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] - # define the list of states - # define the list of datatypes self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(self.obs_data, self.assim_index) - self.current_state = deepcopy(self.state) - self.cov_wgt = self.ext_cov_mat_wgt() self.cov_data = at.gen_covdata(self.datavar, self.assim_index, self.list_datatypes) - self.obs_data_vector, _ = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - self.list_datatypes) - - def _ext_ml_feat(self): - """ - Extract ML specific info from input. - """ - # Make sure ML is a list - if not isinstance(self.keys_da['multilevel'][0], list): - ml_opts = [self.keys_da['multilevel']] - else: - ml_opts = self.keys_da['multilevel'] - - # set default - self.ML_Nested = False - - # Check if 'levels' has been given; if not, give error (mandatory in MULTILEVEL) - assert 'levels' in list(zip(*ml_opts))[0], 'LEVELS has not been given in MULTILEVEL!' - # Check if ensemble size has been given; if not, give error (mandatory in MULTILEVEL) - assert 'en_size' in list(zip(*ml_opts))[0], 'En_Size has not been given in MULTILEVEL!' - # Check if the Hybrid Weights are provided. If not, give error - assert 'cov_wgt' in list(zip(*ml_opts))[0], 'COV_WGT has not been given in MLDA!' - - for i, opt in enumerate(list(zip(*self.keys_da['multilevel']))[0]): - if opt == 'levels': - self.tot_level = int(self.keys_da['multilevel'][i][1]) - if opt == 'nested_states': - if self.keys_da['multilevel'][i][1] == 'true': - self.ML_Nested = True - if opt == 'en_size': - self.ml_ne = [int(el) for el in self.keys_da['multilevel'][i][1]] - if opt == 'ml_error_corr': - #options for ML_error_corr are: bias_corr, deterministic, stochastic, telescopic - self.ML_error_corr = self.keys_da['multilevel'][i][1] - if not self.ML_error_corr=='none': - #options for error_comp_scheme are: once, ens, sep - self.error_comp_scheme = self.keys_da['multilevel'][i][2] - if opt == 'cov_wgt': - try: - cov_mat_wgt = [float(elem) for elem in [item for item in self.keys_da['multilevel'][i][1]]] - except: - cov_mat_wgt = [float(item) for item in self.keys_da['multilevel'][i][1]] - Sum = 0 - for i in range(len(cov_mat_wgt)): - Sum += cov_mat_wgt[i] - for i in range(len(cov_mat_wgt)): - cov_mat_wgt[i] /= Sum - self.cov_wgt = cov_mat_wgt - # Check that we have specified a size for all levels: - assert len(self.ml_ne) == self.tot_level, 'The Ensemble Size must be specified for all levels!' + self.vecObs, self.enObs = self.set_observations() def _init_sim(self): """ Ensure that the simulator is initiallized to handle ML forward simulation. """ self.sim.multilevel = [l for l in range(self.tot_level)] - # self.sim.mlne = - self.sim.rawmap = [None] * self.tot_level self.sim.ecl_coarse = [None] * self.tot_level self.sim.well_cells = [None] * self.tot_level - def ext_cov_mat_wgt(self): - # Make sure MULTILEVEL is a list - if not isinstance(self.keys_da['multilevel'][0], list): - mda_opts = [self.keys_da['multilevel']] - else: - mda_opts = self.keys_da['multilevel'] - - # Check if 'max_iter' has been given; if not, give error (mandatory in ITERATION) - assert 'cov_wgt' in list(zip(*mda_opts))[0], 'COV_WGT has not been given in MLDA!' - - # Extract max. iter - try: - cov_mat_wgt = [float(elem) for elem in [item[1] for item in mda_opts if item[0] == 'cov_wgt'][0]] - except: - cov_mat_wgt = [float(item[1]) for item in mda_opts if item[0]=='cov_wgt'] - Sum=0 - for i in range(len(cov_mat_wgt)): - Sum+=cov_mat_wgt[i] - for i in range(len(cov_mat_wgt)): - cov_mat_wgt[i]/=Sum - # Return max. iter - return cov_mat_wgt - - def init_ml_prior(self): + def reorganize_ml_prior(self, enX: np.ndarray) -> list: ''' - This function changes the structure of prior from independent - ensembles to a nested structure. + Reorganize prior ensemble to multilevel structure (list of matrices). ''' - if self.ML_Nested: - ''' - for i in range(self.tot_level-1,0,-1): - TMP = self.ML_state[i][self.keys_da['staticvar']].shape - for j in range(i): - self.ML_state[j][self.keys_da['staticvar']][0:TMP[0], 0:TMP[1]] = \ - self.ML_state[i][self.keys_da['staticvar']] - ''' - #for el in self.state.keys(): - # self.state[el] = np.repeat(self.state[el][np.newaxis,:,:], self.tot_level,axis=0) + ml_enX = [] + start = 0 + for l in self.multilevel['levels']: + stop = start + self.multilevel['ml_ne'][l] + ml_enX.append(enX[:, start:stop]) + start = stop + return ml_enX - self.state = [deepcopy(self.state) for _ in range(self.tot_level)] - for l in range(self.tot_level): - for el in self.state[0].keys(): - self.state[l][el] = self.state[l][el][:,:self.ml_ne[l]] - else: - # initiallize the state as an empty list of dictionaries with length equal self.tot_level - self.ml_state = [{} for _ in range(self.tot_level)] - # distribute the initial ensemble of states to the levels according to the given ensemble size. - start = 0 # intiallize - for l in range(self.tot_level): - stop = start + self.ml_ne[l] - for el in self.state.keys(): - self.ml_state[l][el] = self.state[el][:,start:stop] - start = stop - - del self.state - self.state = deepcopy(self.ml_state) - del self.ml_state - - def ext_data_size(self): - - # Make sure MULTILEVEL is a list - if not isinstance(self.keys_da['multilevel'][0], list): - mda_opts = [self.keys_da['multilevel']] - else: - mda_opts = self.keys_da['multilevel'] - - # Check if 'data_size' has been given - if not 'data_size' in list(zip(*mda_opts))[0]: # DATA_SIZE has not been given in MDA! - return None - # Extract data_size - try: - data_size = [int(elem) for elem in [item[1] for item in mda_opts if item[0] == 'data_size'][0]] - except: - data_size = [int(item[1]) for item in mda_opts if item[0]=='data_size'] - # Return data_size - return data_size - - -class mlhs_full(multilevel): - # sp_mfda_sim orig inherit this - ''' - Multilevel Hybrid Ensemble Smoother - ''' - - def __init__(self, keys_da,keys_fwd,sim): - """ - Standard initiallization - """ - super().__init__(keys_da,keys_fwd,sim) - - self.obs_data_BU = [] - for item in self.obs_data: - self.obs_data_BU.append(dict(item)) - - self.check_assimindex_simultaneous() - # define the assimilation index - - - self.check_fault() - - self.max_iter = 2 # No iterations - - def calc_analysis(self): - ''' - This class has been written based on the enkf class. It is designed for simultaneous assimilation of seismic - data with multilevel spatial data. - Using this calc_analysis tool we generate a seperate level-based covariance matrix for - every level and update them seperate from each other - This scheme updates each level based on the covariance and cross-covariance matrix - which are generated based on all the levels which are up/down-scaled to that specific level. - In short Modified Kristian's Idea - ''' - - # As the with statement in the ensemble code limits our access to the data and python does not seem to support - # pointers, I re-initialize some part of the code so that we'll access some essential data for the assimilaiton - #self.gen_ness_data() - self.obs_data = self.obs_data_BU - - self.B = [None] * self.tot_level - #self.B_gen(len(self.assim_index[1])) - - self.Dns_mat = [None] * self.tot_level - #self.Dns_mat_gen(len(self.assim_index[1])) - - #self.treat_modeling_error(0) - - # Generate the data auto-covariance matrix - #cov_data = self.gen_covdata(self.datavar, self.assim_index, self.list_datatypes) - - tot_pred = [] - self.Temp_State=[None]*self.tot_level - for level in range(self.tot_level): - obs_data_vector, pred = at.aug_obs_pred_data(self.obs_data, [time_dat[level] for time_dat in self.pred_data] - , self.assim_index, self.list_datatypes) # get some data - #pred = self.Dns_mat[level] * pred - tot_pred.append(pred) - - if not self.ML_error_corr == 'none': - if self.error_comp_scheme=='ens': - if self.ML_error_corr =='bias_corr': - L_mean = np.mean(tot_pred[-1], axis=1) - for l in range(self.tot_level-1): - tot_pred[l] += (L_mean - np.mean(tot_pred[l], axis=1))[:,np.newaxis] - - w_auto = self.cov_wgt - level_data_misfit = [None] * self.tot_level - if self.iteration == 1: # first iteration - misfit_data = 0 - for l in range(self.tot_level): - level_data_misfit[l] = at.calc_objectivefun(np.tile(obs_data_vector[:,np.newaxis],(1,self.ml_ne[l])), - tot_pred[l],self.cov_data) - misfit_data += w_auto[l] * np.mean(level_data_misfit[l]) - self.data_misfit = misfit_data - self.prior_data_misfit = misfit_data - self.prev_data_misfit = misfit_data - # Store the (mean) data misfit (also for conv. check) - if self.lam == 'auto': - self.lam = (0.5 * self.data_misfit)/len(self.obs_data_vector) - - self.logger.info(f'Prior run complete with data misfit: {self.prior_data_misfit:0.1f}. Lambda for initial analysis: {self.lam}') - - # Augment all the joint state variables (originally a dictionary) - aug_state = [at.aug_state(self.state[elem], self.list_states) for elem in range(self.tot_level)] - # concantenate all the elements - tot_aug_state = np.concatenate(aug_state, axis=1) - - # Mean state - mean_state = np.mean(tot_aug_state, 1) - - pert_state = [(aug_state[l] - np.dot(np.resize(mean_state, (len(mean_state), 1)), - np.ones((1, self.ml_ne[l])))) for l in - range(self.tot_level)] - - mean_preddata = [np.mean(tot_pred[elem], 1) for elem in range(self.tot_level)] - - tot_mean_preddata = sum([w_auto[elem] * np.mean(tot_pred[elem], 1) for elem in range(self.tot_level)]) - tot_mean_preddata /= sum([w_auto[elem] for elem in range(self.tot_level)]) - - # calculate the GMA covariances - pert_preddata = [(tot_pred[l] - np.dot(np.resize(mean_preddata[l], (len(mean_preddata[l]), 1)), - np.ones((1, self.ml_ne[l])))) for l in - range(self.tot_level)] - - self.update(pert_preddata,pert_state,mean_preddata,tot_mean_preddata,w_auto, tot_pred, aug_state) - - self.ML_state=self.Temp_State - - def gen_ness_data(self): - for i in range(self.tot_level): - self.ne=1 - self.sim.flow.level=i - self.level=i - assim_step=0 - assim_ind = [self.keys_da['obsname'], self.keys_da['assimindex'][assim_step]] - true_order = [self.keys_da['obsname'], self.keys_da['truedataindex']] - self.state=self.ML_state[i] - self.sim.setup_fwd_run(self.state, assim_ind, true_order) - os.mkdir(f'Test{i}') - folder=f'Test{i}'+os.sep - if self.Treat_Fault: - copyfile(f'IF/FL_{int(self.data_size[i])}.faults', 'IF/FL.faults') - self.sim.flow.run_fwd_sim(0, folder, wait_for_proc=True) - self.ecl_case = ecl.EclipseCase(f'Test{i}' + os.sep + self.sim.flow.file - + '.DATA') - tmp = self.ecl_case.cell_data('PORO') - self.sim.rawmap[self.level] = tmp - time.sleep(5) - for i in range(self.tot_level): - shutil.rmtree(f'Test{i}') - - def check_fault(self): - """ - Checks if there is a statement for generating a fault in the input file and if so - generates a synthetic fault based on the given input in there. - """ - # Make sure MULTILEVEL is a list - if not isinstance(self.keys_da['multilevel'][0], list): - fault_opts = [self.keys_da['multilevel']] - else: - fault_opts = self.keys_da['multilevel'] - - for i, opt in enumerate(list(zip(*self.keys_da['multilevel']))[0]): - if opt == 'generate_fault': - #options for ML_error_corr are: bias_corr, deterministic, stochastic, telescopic - fault_type=self.keys_da['multilevel'][i][1] - fault_dim=[float(item) for item in self.keys_da['multilevel'][i][2]] - if fault_type=='oblique': - self.generate_oblique_fault(fault_dim) - elif fault_type=='horizontal': - self.generate_horizontal_fault(fault_dim) - - self.Treat_Fault = False - for i, opt in enumerate(list(zip(*self.keys_da['multilevel']))[0]): - if opt=='treat_ml_fault': - self.Treat_Fault=True - - def generate_oblique_fault(self,fault_dim): - Dims=[int(self.prior_info[self.keys_da['staticvar']]['nx']), \ - int(self.prior_info[self.keys_da['staticvar']]['ny'])] - Margin=[int(np.floor(Dims[0]/10)),int(np.floor(Dims[1]/10))] - Temp_mat=np.zeros((Dims[0]-2*Margin[0],Dims[1]-2*Margin[1])) - for i in range(Temp_mat.shape[0]): - for j in range(Temp_mat.shape[1]): - if abs(i-j)<=np.floor(fault_dim[0]/2): - Temp_mat[i,Temp_mat.shape[1]-j-1]=fault_dim[1] - Temp_mat1=np.zeros((Dims[0],Dims[1])) - for i in range(Temp_mat.shape[0]): - for j in range(Temp_mat.shape[1]): - Temp_mat1[Margin[0]+i,Margin[1]+j]=Temp_mat[i,j] - Temp_mat = np.reshape(Temp_mat1, (np.product(Temp_mat1.shape), 1)) - for j in range(Temp_mat.shape[0]): - if Temp_mat[j, 0] != 0: - for l in range(self.tot_level): - for i in range(self.ml_ne[l]): - self.ML_state[l][self.keys_da['staticvar']][j,i]=Temp_mat[j] - - def generate_horizontal_fault(self,fault_dim): - Dims=[int(self.prior_info[self.keys_da['staticvar']]['nx']), \ - int(self.prior_info[self.keys_da['staticvar']]['ny'])] - Margin=[int(np.floor(Dims[0]/10)),int(np.floor(Dims[1]/10))] - Temp_mat=np.zeros((Dims[0]-2*Margin[0],Dims[1]-2*Margin[1])) - for i in range(int(fault_dim[0])): - for j in range(Temp_mat.shape[1]): - Temp_mat[int(Temp_mat.shape[0]/2)+i-int(fault_dim[0]/2),j]=fault_dim[1] - Temp_mat1=np.zeros((Dims[0],Dims[1])) - for i in range(Temp_mat.shape[0]): - for j in range(Temp_mat.shape[1]): - Temp_mat1[Margin[0]+i,Margin[1]+j]=Temp_mat[i,j] - Temp_mat = np.reshape(Temp_mat1, (np.product(Temp_mat1.shape), 1)) - for j in range(Temp_mat.shape[0]): - if Temp_mat[j, 0] != 0: - for l in range(self.tot_level): - for i in range(self.ml_ne[l]): - self.ML_state[l][self.keys_da['staticvar']][j,i]=Temp_mat[j] - - def B_gen(self,Multiplier): - for kk in range(self.tot_level): - self.level=kk - Ecl_coarse=self.sim.flow.ecl_coarse[self.level] - try: - Ecl_coarse = np.array(Ecl_coarse) - Ecl_coarse -= 1 - Rawmap_mask=self.sim.rawmap[self.level].mask - Rawmap_mask = Rawmap_mask[0, :, :] - ######### Notice !!!! - nx=self.prior_info[self.keys_da['staticvar']]['nx'] - ny=self.prior_info[self.keys_da['staticvar']]['ny'] - Shape=(nx,ny) - ######### - rows = np.zeros(Shape).flatten() - cols = np.zeros(Shape).flatten() - data = np.zeros(Shape).flatten() - mark = np.zeros(Shape).flatten() - Counter = 0 - for unit in range(Ecl_coarse.shape[0]): - I = None - J = None - Data = 1 / ((Ecl_coarse[unit, 3] - Ecl_coarse[unit, 2] + 1) * ( - Ecl_coarse[unit, 1] - Ecl_coarse[unit, 0] + 1)) - for i in range(Ecl_coarse[unit, 2], Ecl_coarse[unit, 3] + 1): - for j in range(Ecl_coarse[unit, 0], Ecl_coarse[unit, 1] + 1): - if Rawmap_mask[i, j] == False: - I = i - J = j - break - for i in range(Ecl_coarse[unit, 2], Ecl_coarse[unit, 3] + 1): - for j in range(Ecl_coarse[unit, 0], Ecl_coarse[unit, 1] + 1): - rows[Counter] = i * Shape[1] + j - cols[Counter] = I * Shape[1] + J - data[Counter] = Data - mark[i * Shape[1] + j] = 1 - Counter += 1 - - for i in range(Shape[0] * Shape[1]): - if mark[i] == 0: - rows[Counter] = i - cols[Counter] = i - data[Counter] = 1 - mark[i] = 1 - Counter += 1 - - rows = rows.reshape((Shape[0] * Shape[1], 1)) - cols = cols.reshape((Shape[0] * Shape[1], 1)) - data = data.reshape((Shape[0] * Shape[1], 1)) - COO = np.block([rows, cols, data]) - COO = COO[COO[:, 1].argsort(kind='mergesort')] - - Counter = 0 - for i in range(Shape[0] * Shape[1] - 1): - if COO[i, 1] != COO[i + 1, 1]: - Counter += 1 - - Counter = 0 - CXX = np.zeros(COO.shape) - CXX[0, 1] = Counter - CXX[0, 0] = COO[0, 0] - CXX[0, 2] = COO[0, 2] - for i in range(1, Shape[0] * Shape[1]): - if COO[i, 1] != COO[i - 1, 1]: - Counter += 1 - CXX[i, 1] = Counter - CXX[i, 0] = COO[i, 0] - CXX[i, 2] = COO[i, 2] - - S1 = self.data_size[self.level] - S2= CXX.shape[0] - Final_mat=np.zeros((CXX.shape[0]*Multiplier,CXX.shape[1])) - for i in range(Multiplier): - Final_mat[i*S2:(i+1)*S2,2]=CXX[:,2] - Final_mat[i*S2:(i+1)*S2,0]=CXX[:,0]+S2*i - Final_mat[i*S2:(i+1)*S2,1]=CXX[:,1]+S1*i - - rows = Final_mat[:, 1] - cols = Final_mat[:, 0] - data = Final_mat[:, 2] - - S2=np.product(Shape)*Multiplier - S1=self.data_size[self.level]*Multiplier - self.B[self.level]=coo_matrix((data,(rows,cols)),shape=(S1,S2)) - except: - COO=np.zeros((self.data_size[self.level]*Multiplier,3)) - S1=COO.shape[0] - for i in range(S1): - COO[i,0]=i - COO[i,1]=i - COO[i,2]=1 - self.B[self.level]=coo_matrix((COO[:,2],(COO[:,0],COO[:,1])),shape=(S1,S1)) - - def Dns_mat_gen(self, Multiplier): - for kk in range(self.tot_level): - self.level = kk - Ecl_coarse = self.sim.flow.ecl_coarse[self.level] - try: - Ecl_coarse = np.array(Ecl_coarse) - Ecl_coarse -= 1 - Rawmap_mask = self.sim.rawmap[self.level].mask - Rawmap_mask = Rawmap_mask[0, :, :] - ######### Notice !!!! - nx = self.prior_info[self.keys_da['staticvar']]['nx'] - ny = self.prior_info[self.keys_da['staticvar']]['ny'] - Shape = (nx, ny) - ######### - rows = np.zeros(Shape).flatten() - cols = np.zeros(Shape).flatten() - data = np.zeros(Shape).flatten() - mark = np.zeros(Shape).flatten() - Counter = 0 - for unit in range(Ecl_coarse.shape[0]): - I = None - J = None - for i in range(Ecl_coarse[unit, 2], Ecl_coarse[unit, 3] + 1): - for j in range(Ecl_coarse[unit, 0], Ecl_coarse[unit, 1] + 1): - if Rawmap_mask[i, j] == False: - I = i - J = j - break - for i in range(Ecl_coarse[unit, 2], Ecl_coarse[unit, 3] + 1): - for j in range(Ecl_coarse[unit, 0], Ecl_coarse[unit, 1] + 1): - rows[Counter] = i * Shape[1] + j - cols[Counter] = I * Shape[1] + J - data[Counter] = 1 - mark[i * Shape[1] + j] = 1 - Counter += 1 - - for i in range(Shape[0] * Shape[1]): - if mark[i] == 0: - rows[Counter] = i - cols[Counter] = i - data[Counter] = 1 - mark[i] = 1 - Counter += 1 - - rows = rows.reshape((Shape[0] * Shape[1], 1)) - cols = cols.reshape((Shape[0] * Shape[1], 1)) - data = data.reshape((Shape[0] * Shape[1], 1)) - COO = np.block([rows, cols, data]) - COO = COO[COO[:, 1].argsort(kind='mergesort')] - - Counter = 0 - for i in range(Shape[0] * Shape[1] - 1): - if COO[i, 1] != COO[i + 1, 1]: - Counter += 1 - - Counter = 0 - CXX = np.zeros(COO.shape) - CXX[0, 1] = Counter - CXX[0, 0] = COO[0, 0] - CXX[0, 2] = COO[0, 2] - for i in range(1, Shape[0] * Shape[1]): - if COO[i, 1] != COO[i - 1, 1]: - Counter += 1 - CXX[i, 1] = Counter - CXX[i, 0] = COO[i, 0] - CXX[i, 2] = COO[i, 2] - - S1 = self.data_size[self.level] - S2 = CXX.shape[0] - Final_mat = np.zeros((CXX.shape[0] * Multiplier, CXX.shape[1])) - for i in range(Multiplier): - Final_mat[i * S2:(i + 1) * S2, 2] = CXX[:, 2] - Final_mat[i * S2:(i + 1) * S2, 0] = CXX[:, 0] + S2 * i - Final_mat[i * S2:(i + 1) * S2, 1] = CXX[:, 1] + S1 * i - - rows = Final_mat[:, 0] - cols = Final_mat[:, 1] - data = Final_mat[:, 2] - - S2 = np.product(Shape) * Multiplier - S1 = self.data_size[self.level] * Multiplier - self.Dns_mat[self.level] = coo_matrix((data, (rows, cols)), shape=(S2, S1)) - except: - COO = np.zeros((self.data_size[self.level] * Multiplier, 3)) - S1 = COO.shape[0] - for i in range(S1): - COO[i, 0] = i - COO[i, 1] = i - COO[i, 2] = 1 - self.Dns_mat[self.level] = coo_matrix((COO[:, 2], (COO[:, 0], COO[:, 1])), shape=(S1, S1)) - - - def update(self,pert_preddata,pert_state,mean_preddata,tot_mean_preddata,w_auto, tot_pred, aug_state): - - #level_pert_preddata = [self.B[level] * pert_preddata[l] for l in range(self.tot_level)] - level_pert_preddata = [pert_preddata[l] for l in range(self.tot_level)] - #level_mean_preddata = [self.B[level] * mean_preddata[l] for l in range(self.tot_level)] - level_mean_preddata = [mean_preddata[l] for l in range(self.tot_level)] - #level_tot_mean_preddata = self.B[level] * tot_mean_preddata - level_tot_mean_preddata = tot_mean_preddata - - cov_auto = sum([w_auto[l] * at.calc_autocov(level_pert_preddata[l]) for l in range(self.tot_level)]) + \ - sum([w_auto[l] * np.outer((level_mean_preddata[l] - level_tot_mean_preddata), - (level_mean_preddata[l] - level_tot_mean_preddata)) for l in - range(self.tot_level)]) - cov_auto /= sum([w_auto[l] for l in range(self.tot_level)]) - - cov_cross = sum([w_auto[l] * at.calc_crosscov(pert_state[l], level_pert_preddata[l]) - for l in range(self.tot_level)]) - cov_cross /= sum([w_auto[l] for l in range(self.tot_level)]) - - #joint_data_cov = self.B[level] * self.cov_data * self.B[level].transpose() - joint_data_cov = self.cov_data - - kalman_gain_param = self.calc_kalmangain(cov_cross, cov_auto, - joint_data_cov) # global cov_cross and cov_auto - - for level in range(self.tot_level): - obs_data = self.efficient_real_gen(self.obs_data_vector, self.cov_data, self.ml_ne[level], \ - level) - - #level_tot_pred = self.B[level] * tot_pred[level] - level_tot_pred = tot_pred[level] - aug_state_upd = at.calc_kalman_filter_eq(aug_state[level], kalman_gain_param, obs_data, - level_tot_pred) # update levelwise - - self.Temp_State[level] = at.update_state(aug_state_upd, self.state[level], self.list_states) - - def efficient_real_gen(self, mean, var, number, level,original_size=False, limits=None, return_chol=False): - """ - This function is added to prevent additional computational cost if var is diagonal - MN 04/20 - """ - if not original_size: - var = np.array(var) #to enable var.shape - parsize = len(mean) - if parsize == 1 or len(var.shape) == 1: - l = np.sqrt(var) - # real = mean + L*np.random.randn(1, number) - else: - # Check if the covariance matrix is diagonal (only entries in the main diagonal). If so, we can use - # numpy.sqrt for efficiency - if 4==2: #np.count_nonzero(var - np.diagonal(var)) == 0: - l = np.sqrt(var) # only variance (diagonal) term - l=np.reshape(l,(l.size,1)) - else: - # Cholesky decomposition - l = linalg.cholesky(var) # cov. matrix has off-diag. terms - #Mean=deepcopy(mean) - Mean=np.reshape(mean,(mean.size,1)) - #Mean=self.B[level]*Mean - # Gen. realizations - # if len(var.shape) == 1: - # real = np.dot(Mean, np.ones((1, number))) + np.expand_dims((self.B[level]*l).flatten(), axis=1)*np.random.randn( - # np.size(Mean), number) - # else: - # real = np.tile(Mean, (1, number)) + np.dot(self.B[level]*l.T, np.random.randn(np.size(mean), - # number)) - if len(var.shape) == 1: - real = np.dot(Mean, np.ones((1, number))) + np.expand_dims((l).flatten(), axis=1)*np.random.randn( - np.size(Mean), number) - else: - real = np.tile(Mean, (1, number)) + np.dot(l.T, np.random.randn(np.size(mean), - number)) - - # Truncate values that are outside limits - # TODO: Make better truncation rules, or switch truncation on/off - if limits is not None: - # Truncate - real[real > limits['upper']] = limits['upper'] - real[real < limits['lower']] = limits['lower'] - - if return_chol: - return real, l - else: - return real - else: - var = np.array(var) # to enable var.shape - parsize = len(mean) - if parsize == 1 or len(var.shape) == 1: - l = np.sqrt(var) - # real = mean + L*np.random.randn(1, number) - else: - # Check if the covariance matrix is diagonal (only entries in the main diagonal). If so, we can use - # numpy.sqrt for efficiency - if 4 == 2: # np.count_nonzero(var - np.diagonal(var)) == 0: - l = np.sqrt(var) # only variance (diagonal) term - l = np.reshape(l, (l.size, 1)) - else: - # Cholesky decomposition - l = linalg.cholesky(var) # cov. matrix has off-diag. terms - # Mean=deepcopy(mean) - Mean = np.reshape(mean, (mean.size, 1)) - # Gen. realizations - if len(var.shape) == 1: - real = np.dot(Mean, np.ones((1, number))) + np.expand_dims((l).flatten(), - axis=1) * np.random.randn( - np.size(Mean), number) - else: - real = np.tile(Mean, (1, number)) + np.dot(l.T, np.random.randn(np.size(mean), - number)) - - # Truncate values that are outside limits - # TODO: Make better truncation rules, or switch truncation on/off - if limits is not None: - # Truncate - real[real > limits['upper']] = limits['upper'] - real[real < limits['lower']] = limits['lower'] - - if return_chol: - return real, l - else: - return real - def calc_kalmangain(self, cov_cross, cov_auto, cov_data, opt=None): - """ - Calculate the Kalman gain - Using mainly two options: linear soultion and pseudo inverse of the matrix - MN 04/2020 - """ - if opt is None: - calc_opt = 'lu' - - # Add data and predicted data auto-covariance matrices - if len(cov_data.shape)==1: - cov_data = np.diag(cov_data) - c_auto = cov_auto + cov_data - - if calc_opt == 'lu': - try: - kg = linalg.solve(c_auto.T, cov_cross.T) - kalman_gain = kg.T - except: - #Margin=10**5 - #kalman_gain = cov_cross * self.calc_pinv(c_auto, Margin=Margin) - #kalman_gain = cov_cross * self.calc_pinv(c_auto) - kalman_gain = cov_cross * np.linalg.pinv(c_auto) - #kalman_gain = cov_cross * np.linalg.pinv(c_auto, rcond=10**(-15)) - - elif calc_opt == 'chol': - # Cholesky decomp (upper triangular matrix) - u = linalg.cho_factor(c_auto.T, check_finite=False) - - # Solve linear system with cholesky square-root - kalman_gain = linalg.cho_solve(u, cov_cross.T, check_finite=False) - - # Return Kalman gain - return kalman_gain - - def check_convergence(self): - """ - Check if LM-EnRML have converged based on evaluation of change sizes of objective function, state and damping - parameter. - - Returns - ------- - conv: bool - Logic variable telling if algorithm has converged - why_stop: dict - Dict. with keys corresponding to conv. criteria, with logical variable telling which of them that has been - met - """ - success = False # init as false - - if hasattr(self, 'list_datatypes'): - assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] - list_datatypes = self.list_datatypes -# cov_data = self.gen_covdata(self.datavar, assim_index, list_datatypes) - pred_data = [None] * self.tot_level - level_mean_preddata = [None] * self.tot_level - for l in range(self.tot_level): - obs_data_vector, pred_data[l] = at.aug_obs_pred_data(self.obs_data, - [time_dat[l] for time_dat in self.pred_data], - assim_index, list_datatypes) - level_mean_preddata[l] = np.mean(pred_data[l], 1) - else: - assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] - list_datatypes, _ = at.get_list_data_types(self.obs_data, assim_index) - # cov_data = at.gen_covdata(self.datavar, assim_index, list_datatypes) - #cov_data = self.gen_covdata(self.datavar, assim_index, list_datatypes) - pred_data = [None] * self.tot_level - level_mean_preddata = [None] * self.tot_level - for l in range(self.tot_level): - obs_data_vector, pred_data[l] = at.aug_obs_pred_data(self.obs_data, - [time_dat[l] for time_dat in self.pred_data], - assim_index, list_datatypes) - level_mean_preddata[l] = np.mean(pred_data[l], 1) - - # self.prev_data_misfit_std = self.data_misfit_std - # if there was no reduction of the misfit, retain the old "valid" data misfit. - - # Calc. std dev of data misfit (used to update lamda) - # mat_obs = np.dot(obs_data_vector.reshape((len(obs_data_vector),1)), np.ones((1, self.ne))) # use the perturbed - # data instead. - # mat_obs = self.real_obs_data - level_data_misfit = [None] * self.tot_level - #list_states = list(self.state.keys()) - #cov_prior = at.block_diag_cov(self.cov_prior, list_states) - #ML_prior_state = [at.aug_state(self.ML_prior_state[elem], list_states) for elem in range(self.tot_level)] - #ML_state = [at.aug_state(self.state[elem], list_states) for elem in range(self.tot_level)] - # level_state_misfit = [None] * self.tot_level - # if len(self.cov_data.shape) == 1: - for l in range(self.tot_level): - - level_data_misfit[l] = at.calc_objectivefun(np.tile(obs_data_vector[:,np.newaxis],(1,self.ml_ne[l])), - pred_data[l],self.cov_data) - -# obs_data = self.Dns_mat[l] * self.obs_reals[l] - ##### This part is not done correctly as we do not need it now!!! ###### - # level_data_misfit[l] = np.diag(np.dot((pred_data[l] - obs_data).T * self.Dns_mat[l].transpose() * - # self.B[0].transpose(), - # np.dot(np.expand_dims(self.cov_data ** (-1), axis=1), - # np.ones((1, self.ne))) * self.B[0] * self.Dns_mat[l] * ( - # pred_data[l] - obs_data))) - #level_state_misfit[l] = np.diag(np.dot((ML_state[l] - ML_prior_state[l]).T, solve( - # cov_prior, (ML_state[l] - ML_prior_state[l])))) - # else: - # for l in range(self.tot_level): - # obs_data = self.Dns_mat[l]*self.obs_reals[l] - # obs_data = self.obs_reals[l] - # ''' - # level_data_misfit[l] = np.diag(np.dot((pred_data [l]- obs_data).T*self.Dns_mat[l].transpose()* - # self.B[0].transpose(),solve(self.B[0]*cov_data*self.B[0].transpose(), - # self.B[0]*self.Dns_mat[l]*(pred_data[l] - obs_data)))) - # level_state_misfit[l]=np.diag(np.dot((ML_state[l]-ML_prior_state[l]).T,solve( - # cov_prior,(ML_state[l]-ML_prior_state[l])))) - # ''' - # level_data_misfit[l] = np.diag(np.dot((pred_data[l] - obs_data).T * self.Dns_mat[l].transpose(), - # solve(self.cov_data, self.Dns_mat[l] * (pred_data[l] - obs_data)))) - - misfit_data = 0 -# misfit_state = 0 - w_auto = self.cov_wgt - for l in range(self.tot_level): - misfit_data += w_auto[l] * np.mean(level_data_misfit[l]) - # misfit_state+=w_auto[l]*np.mean(level_state_misfit[l]) - - self.data_misfit = misfit_data - # self.data_misfit_std = np.std(con_misfit) - - # # Calc. mean data misfit for convergence check, using the updated state variable - # self.data_misfit = np.dot((mean_preddata - obs_data_vector).T, - # solve(cov_data, (mean_preddata - obs_data_vector))) - - # Convergence check: Relative step size of data misfit or state change less than tolerance - why_stop = {} # todo: populate - - # update the last mismatch, only if this was a reduction of the misfit - if self.data_misfit < self.prev_data_misfit: - success = True - - - if success: - self.logger.info(f'ML Hybrid Smoother update complete! Objective function reduced from ' - f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') - # self.prev_data_misfit = self.data_misfit - # self.prev_data_misfit_std = self.data_misfit_std - else: - self.logger.info(f'ML Hybrid Smoother update complete! Objective function increased from ' - f'{self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') - - # Return conv = False, why_stop var. - return False, True, why_stop - -class smlses_s(multilevel,esmda_approx): - """ - The Sequential multilevel ensemble smoother with the "straightforward" flavour as descibed in Nezhadali, M., - Bhakta, T., Fossum, K., & Mannseth, T. (2023). Sequential multilevel assimilation of inverted seismic data. - Computational Geosciences, 27(2), 265–287. https://doi.org/10.1007/s10596-023-10191-9 - - Since the update schemes are basically a esmda update we inherit the esmda_approx method. Hence, we only have to - care about handling the multi-level features. - """ - - def __init__(self,keys_da, keys_fwd, sim): - super().__init__(keys_da, keys_fwd, sim) - - self.current_state = [self.current_state[0]] - self.state = [self.state[0]] - - # Overwrite the method for extracting ml_information. Here, we should only get the first level - def _ext_ml_info(self, grab_level=0): - ''' - Extract the info needed for ML simulations. Grab the first level info - ''' - - if 'multilevel' in self.keys_en: - # parse - self.multilevel = {} - for i, opt in enumerate(list(zip(*self.keys_en['multilevel']))[0]): - if opt == 'levels': - self.multilevel['levels'] = [elem for elem in range( - int(self.keys_en['multilevel'][i][1]))] - if opt == 'en_size': - self.multilevel['ne'] = [range(int(el)) - for el in self.keys_en['multilevel'][i][1]] - try: - self.multilevel['levels'] = [self.multilevel['levels'][grab_level]] - except IndexError: # When converged, we need to set the level to the final one - self.multilevel['levels'] = [self.multilevel['levels'][-1]] - #self.multilevel['ne'] = [self.multilevel['ne'][grab_level]] - def calc_analysis(self): - # Some preamble for multilevel - # Do this. - # flatten the level element of the predicted data - tmp = [] - for elem in self.pred_data: - tmp += elem - self.pred_data = tmp - - self.current_state = self.current_state[self.multilevel['levels'][0]] - self.state = self.state[self.multilevel['levels'][0]] - # call the inherited version via super() - super().calc_analysis() - - # Afterwork - self._ext_ml_info(grab_level=self.iteration) - - # Grab the prior for the next mda step. Draw the top scoring values. - self._update_ensemble() - - def _update_ensemble(self): - # Prelude to calc. conv. check (everything done below is from calc_analysis) - obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - self.list_datatypes) - - data_misfit = at.calc_objectivefun( - self.real_obs_data_conv, pred_data, self.cov_data) - - # sort the data_misfit after the percentile score - sort_ind = np.argsort(data_misfit)[self.multilevel['ne'][self.multilevel['levels'][0]]] - - # initialize self.current_state and self.state as empty lists with lenght equal to self.multilevel['levels'][0] - tmp_current_state = [[] for _ in range(self.multilevel['levels'][0]+1)] - tmp_state = [[] for _ in range(self.multilevel['levels'][0]+1)] - - tmp_current_state[self.multilevel['levels'][0]] = {el:self.current_state[el][:,sort_ind] for el in self.current_state.keys()} - tmp_state[self.multilevel['levels'][0]] = {el:self.state[el][:,sort_ind] for el in self.state.keys()} - - - #reduce the size of these ensembles as well - self.real_obs_data_conv = self.real_obs_data_conv[:,sort_ind] - self.real_obs_data = self.real_obs_data[:,sort_ind] - - # set the current state and state to the new values - self.current_state = tmp_current_state - self.state = tmp_state - - # update self.ne to be inline with new ensemble size - self.ne = len(self.multilevel['ne'][self.multilevel['levels'][0]]) - - # and update the projection to be inline with new ensemble size - self.proj = (np.eye(self.ne) - (1 / self.ne) * - np.ones((self.ne, self.ne))) / np.sqrt(self.ne - 1) - - def check_convergence(self): - """ - Check convergence for the smlses-s method - """ - - self.prev_data_misfit = self.data_misfit - self.prev_data_misfit_std = self.data_misfit_std - - # extract pred_data for the current level - level_pred_data = [el[0] for el in self.pred_data] - - # Prelude to calc. conv. check (everything done below is from calc_analysis) - obs_data_vector, pred_data = at.aug_obs_pred_data(self.obs_data, level_pred_data, self.assim_index, - self.list_datatypes) - - data_misfit = at.calc_objectivefun( - self.real_obs_data_conv, pred_data, self.cov_data) - self.data_misfit = np.mean(data_misfit) - self.data_misfit_std = np.std(data_misfit) - - # Logical variables for conv. criteria - why_stop = {'rel_data_misfit': 1 - (self.data_misfit / self.prev_data_misfit), - 'data_misfit': self.data_misfit, - 'prev_data_misfit': self.prev_data_misfit} - - if self.data_misfit < self.prev_data_misfit: - self.logger.info( - f'ML-MDA iteration number {self.iteration}! Objective function reduced from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') - else: - self.logger.info( - f'ML-MDA iteration number {self.iteration}! Objective function increased from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') - # Return conv = False, why_stop var. - self.current_state = deepcopy(self.state) - - return False, True, why_stop - -class esmda_h(multilevel,hybrid_update,esmdaMixIn): +class esmda_hybrid(multilevel,hybrid_update,esmdaMixIn): ''' A multilevel implementation of the ES-MDA algorithm with the hybrid gain ''' - def __init__(self,keys_da, keys_fwd, sim): super().__init__(keys_da, keys_fwd, sim) - self.proj = [(np.eye(self.ml_ne[l]) - (1 / self.ml_ne[l]) * - np.ones((self.ml_ne[l], self.ml_ne[l]))) / np.sqrt(self.ml_ne[l] - 1) for l in range(self.tot_level)] + self.proj = [] + for l in range(self.tot_level): + nl = self.ml_ne[l] + proj_l = (np.eye(nl) - np.ones((nl, nl))/nl) / np.sqrt(nl - 1) + self.proj.append(proj_l) + def calc_analysis(self): - self.aug_pred_data = [] + + # Get ensemble predictions at all levels + self.enPred = [] for l in range(self.tot_level): - self.aug_pred_data.append(at.aug_obs_pred_data(self.obs_data, [el[l] for el in self.pred_data], self.assim_index, - self.list_datatypes)[1]) + _, enPred_level = at.aug_obs_pred_data( + self.obs_data, + [el[l] for el in self.pred_data], + self.assim_index, + self.list_datatypes + ) + self.enPred.append(enPred_level) + + # Initialize GeoStat class for generating realizations + cholesky = Cholesky() - init_en = Cholesky() # Initialize GeoStat class for generating realizations if self.iteration == 1: # first iteration - # note, evaluate for high fidelity model + + # Note, evaluate for high fidelity model data_misfit = at.calc_objectivefun( - self.real_obs_data_conv, np.concatenate(self.aug_pred_data,axis=1), self.cov_data) + self.enObs, + np.concatenate(self.enPred,axis=1), # Is this correct, given the comment above?????? + self.cov_data + ) # Store the (mean) data misfit (also for conv. check) self.data_misfit = np.mean(data_misfit) @@ -1014,56 +134,50 @@ def calc_analysis(self): self.data_misfit = np.mean(data_misfit) self.data_misfit_std = np.std(data_misfit) - self.logger.info( - f'Prior run complete with data misfit: {self.prior_data_misfit:0.1f}.') + # Log initial data misfit + self.log_update(prior_run=True) self.data_random_state = deepcopy(np.random.get_state()) - self.real_obs_data = [] + + + self.ml_enObs = [] self.scale_data = [] + self.E = [] for l in range(self.tot_level): - # populate the lists without unpacking the output form init_en.gen_real - (lambda x,y: (self.real_obs_data.append(x),self.scale_data.append(y)))(*init_en.gen_real(self.obs_data_vector, - self.alpha[self.iteration - 1] * - self.cov_data, self.ml_ne[l], - return_chol=True)) - self.E = [np.dot(self.real_obs_data[l], self.proj[l]) for l in range(self.tot_level)] - else: - self.data_random_state = deepcopy(np.random.get_state()) - # self.obs_data_vector, _ = at.aug_obs_pred_data(self.obs_data, self.pred_data, self.assim_index, - # self.list_datatypes) - for l in range(self.tot_level): - self.real_obs_data[l], self.scale_data[l] = init_en.gen_real(self.obs_data_vector, - self.alpha[self.iteration - - 1] * self.cov_data, - self.ml_ne[l], - return_chol=True) - self.E[l] = np.dot(self.real_obs_data[l], self.proj[l]) - self.pert_preddata = [] - for l in range(self.tot_level): - if len(self.scale_data[l].shape) == 1: - self.pert_preddata.append(np.dot(np.expand_dims(self.scale_data[l] ** (-1), axis=1), - np.ones((1, self.ml_ne[l]))) * np.dot(self.aug_pred_data[l], self.proj[l])) - else: - self.pert_preddata.append(solve( - self.scale_data[l], np.dot(self.aug_pred_data[l], self.proj[l]))) + # Generate real data and scale data + enObs_level, scale_data_level = cholesky.gen_real( + self.vecObs, + self.alpha[self.iteration - 1] * self.cov_data, + self.ml_ne[l], + return_chol=True + ) + self.ml_enObs.append(enObs_level) + self.scale_data.append(scale_data_level) + self.E.append(np.dot(enObs_level, self.proj[l])) - aug_state= [] - for l in range(self.tot_level): - aug_state.append(at.aug_state(self.current_state[l], self.list_states)) + else: + self.data_random_state = deepcopy(np.random.get_state()) - self.update() + for l in range(self.tot_level): + self.ml_enObs[l], self.scale_data[l] = cholesky.gen_real( + self.ml_enObs[l], + self.alpha[self.iteration - 1] * self.cov_data, + self.ml_ne[l], + return_chol=True + ) + self.E[l] = np.dot(self.ml_enObs[l], self.proj[l]) + + # Calculate update step + self.update( + enX = self.enX, + enY = self.enPred, + enE = self.ml_enObs + ) if hasattr(self, 'step'): - aug_state_upd = [aug_state[l] + self.step[l] for l in range(self.tot_level)] - # if hasattr(self, 'w_step'): - # self.W = self.current_W + self.w_step - # aug_prior_state = at.aug_state(self.prior_state, self.list_states) - # aug_state_upd = np.dot(aug_prior_state, (np.eye( - # self.ne) + self.W / np.sqrt(self.ne - 1))) - - # Extract updated state variables from aug_update - for l in range(self.tot_level): - self.state[l] = at.update_state(aug_state_upd[l], self.state[l], self.list_states) - self.state[l] = at.limits(self.state[l], self.prior_info) + self.enX_temp = [self.enX[l] + self.step[l] for l in range(self.tot_level)] + # Enforce limits + limits = {key: self.prior_info[key].get('limits', (None, None)) for key in self.idX.keys()} + self.enX_temp = [entools.clip_matrix(self.enX_temp[l], limits, self.idX) for l in range(self.tot_level)] def check_convergence(self): """ @@ -1074,13 +188,21 @@ def check_convergence(self): self.prev_data_misfit_std = self.data_misfit_std # Prelude to calc. conv. check (everything done below is from calc_analysis) - pred_data = [] + enPred = [] for l in range(self.tot_level): - pred_data.append(at.aug_obs_pred_data(self.obs_data, [el[l] for el in self.pred_data], self.assim_index, - self.list_datatypes)[1]) + _, enPred_level = at.aug_obs_pred_data( + self.obs_data, + [el[l] for el in self.enX_temp], + self.assim_index, + self.list_datatypes + ) + enPred.append(enPred_level) data_misfit = at.calc_objectivefun( - self.real_obs_data_conv, np.concatenate(pred_data,axis=1), self.cov_data) + self.enObs, + np.concatenate(enPred,axis=1), + self.cov_data + ) self.data_misfit = np.mean(data_misfit) self.data_misfit_std = np.std(data_misfit) @@ -1089,105 +211,14 @@ def check_convergence(self): 'data_misfit': self.data_misfit, 'prev_data_misfit': self.prev_data_misfit} - if self.data_misfit < self.prev_data_misfit: - self.logger.info( - f'MDA iteration number {self.iteration}! Objective function reduced from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') - else: - self.logger.info( - f'MDA iteration number {self.iteration}! Objective function increased from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') - # Return conv = False, why_stop var. - self.current_state = deepcopy(self.state) - if hasattr(self, 'W'): - self.current_W = deepcopy(self.W) - - return False, True, why_stop - -class esmda_seq_h(multilevel,esmda_approx): - ''' - A multilevel implementation of the Sequeontial ES-MDA algorithm with the hybrid gain - ''' - - def __init__(self,keys_da, keys_fwd, sim): - super().__init__(keys_da, keys_fwd, sim) - - self.proj = (np.eye(self.ml_ne[0]) - (1 / self.ml_ne[0]) * - np.ones((self.ml_ne[0], self.ml_ne[0]))) / np.sqrt(self.ml_ne[0] - 1) - - self.multilevel['levels'] = [self.iteration] + # Log update results + success = self.data_misfit < self.prev_data_misfit + self.log_update(success=success) - self.ne = self.ml_ne[0] - # adjust the real_obs_data to only containt the first ne samples - self.real_obs_data_conv = self.real_obs_data_conv[:,:self.ne] - - def calc_analysis(self): - - # collapse the level element of the predicted data - self.ml_pred = deepcopy(self.pred_data) - # concantenate the ml_pred data and state - self.pred_data = [] - curr_level = self.multilevel['levels'][0] - for level_pred_date in self.ml_pred: - keys = level_pred_date[curr_level].keys() - result ={} - for key in keys: - arrays = np.array([level_pred_date[curr_level][key]]) - result[key] = np.hstack(arrays) - self.pred_data.append(result) - - self.ml_state = deepcopy(self.state) - self.state = self.state[self.multilevel['levels'][0]] - self.current_state = self.current_state[self.multilevel['levels'][0]] - - super().calc_analysis() - - # Set the multilevel index and set the dimentions for all the states - self.multilevel['levels'][0] += 1 - self.ne = self.ml_ne[self.multilevel['levels'][0]] - self.proj =(np.eye(self.ne) - (1 / self.ne) * - np.ones((self.ne, self.ne))) / np.sqrt(self.ne - 1) - best_members = np.argsort(self.ensemble_misfit)[:self.ne] - self.ml_state[self.multilevel['levels'][0]] = {k: v[:, best_members] for k, v in self.state.items()} - self.state = deepcopy(self.ml_state) - - self.real_obs_data_conv = self.real_obs_data_conv[:,best_members] - - - - def check_convergence(self): - """ - Check ESMDA objective function for logging purposes. - """ - - self.prev_data_misfit = self.data_misfit - #self.prev_data_misfit_std = self.data_misfit_std - - # Prelude to calc. conv. check (everything done below is from calc_analysis) - pred_data = [] - for l in range(len(self.pred_data[0])): - level_pred = at.aug_obs_pred_data(self.obs_data, [el[l] for el in self.pred_data], self.assim_index, - self.list_datatypes)[1] - if level_pred is not None: # Can be None if level is not predicted - pred_data.append(level_pred) - - data_misfit = at.calc_objectivefun( - self.real_obs_data_conv, np.concatenate(pred_data,axis=1), self.cov_data) - self.ensemble_misfit = data_misfit - self.data_misfit = np.mean(data_misfit) - self.data_misfit_std = np.std(data_misfit) - - # Logical variables for conv. criteria - why_stop = {'rel_data_misfit': 1 - (self.data_misfit / self.prev_data_misfit), - 'data_misfit': self.data_misfit, - 'prev_data_misfit': self.prev_data_misfit} - - if self.data_misfit < self.prev_data_misfit: - self.logger.info( - f'MDA iteration number {self.iteration}! Objective function reduced from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') - else: - self.logger.info( - f'MDA iteration number {self.iteration}! Objective function increased from {self.prev_data_misfit:0.1f} to {self.data_misfit:0.1f}.') # Return conv = False, why_stop var. - self.current_state = deepcopy(self.state) + self.enX = deepcopy(self.enX_temp) + self.enX_temp = None + if hasattr(self, 'W'): self.current_W = deepcopy(self.W) diff --git a/pipt/update_schemes/update_methods_ns/hybrid_udpate.py b/pipt/update_schemes/update_methods_ns/hybrid_udpate.py index 5409914..d53a6f0 100644 --- a/pipt/update_schemes/update_methods_ns/hybrid_udpate.py +++ b/pipt/update_schemes/update_methods_ns/hybrid_udpate.py @@ -16,7 +16,77 @@ class hybrid_update: enables the scheme to efficiently be coupled with multiple updating strategies via class MixIn ''' - def update(self): + def scale(self, data, scaling): + """ + Scale the data perturbations by the data error standard deviation. + + Args: + data (np.ndarray): data perturbations + scaling (np.ndarray): data error standard deviation + + Returns: + np.ndarray: scaled data perturbations + """ + + if len(scaling.shape) == 1: + return (scaling ** (-1))[:, None] * data + else: + return solve(scaling, data) + + def update(self, enX, enY, enE, **kwargs): + ''' + Perform the hybrid update. + + Parameters: + ---------- + enX : list of np.ndarray + List of state ensemble matrices for each level (nx, ne) + + enY : list of np.ndarray + List of predicted data ensemble matrices for each level (nd, ne) + + enE : list of np.ndarray + List of ensemble of perturbed observations for each level (nd, ne) + ''' + # Loop over levels to calculate the update step + X3 = [] + enXcentered = [] + for l in range(self.tot_level): + + # Get Perturbed state ensemble at level l + if ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes'): + enXcentered.append(self.scale(enX[l] - np.mean(enX[l], 1)[:,None], self.state_scaling)) + else: + enXcentered.append(self.scale(np.dot(enX[l], self.proj[l]), self.state_scaling)) + + # Calculet truncated SVD of predicted data ensemble at level l + enYcentered = self.scale(np.dot(enY[l], self.proj[l]), self.scale_data[l]) + Ud, Sd, VTd = at.truncSVD(enYcentered, energy=self.trunc_energy) + + X2 = solve(((self.lam + 1)*np.eye(len(Sd)) + np.diag(Sd**2)), Ud.T) + X3.append(np.dot(np.dot(VTd.T, np.diag(Sd)), X2)) + + # Calculate each row of self.step individually to avoid memory issues. + self.step = [np.empty(enXcentered[l].shape) for l in range(self.tot_level)] + step_size = min(1000, int(self.state_scaling.shape[0]/2)) # do maximum 1000 rows at a time. + + # Generate row batches + nrows = self.state_scaling.shape[0] + row_step = [np.arange(s, min(s + step_size, nrows)) for s in range(0, nrows, step_size)] + + # Loop over rows + for row in row_step: + ml_weights = self.multilevel['ml_weights'] + kg = sum([ml_weights[l]*np.dot(enXcentered[l][row, :], X3[l]) for l in range(self.tot_level)]) + + # Loop over levels + for l in range(self.tot_level): + enRes = self.scale(enE[l] - enY[l], self.scale_data[l]) + self.step[l][row, :] = np.dot(self.state_scaling[row, None] * kg, enRes) + + + + def _update(self): x_3 = [] pert_state = [] for l in range(self.tot_level): diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index c94c201..499fea4 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -119,8 +119,7 @@ def gradient(self, x, *args, **kwargs): index += ne if 'multilevel' in self.keys_en: - weight = ot.get_list_element(self.keys_en['multilevel'], 'cov_wgt') - weight = np.array(weight) + weight = np.array(self.keys_en['multilevel']['ml_weights']) if not np.sum(weight) == 1.0: weight = weight / np.sum(weight) grad = np.dot(grad_ml, weight) From 18fc38118340400f0ec2f02d34606904eb0b1b4b Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Tue, 13 Jan 2026 15:15:08 +0100 Subject: [PATCH 81/94] Fix multilevel bugs --- ensemble/ensemble.py | 24 ++++++++++--------- pipt/loop/ensemble.py | 2 +- pipt/misc_tools/analysis_tools.py | 6 ++--- pipt/misc_tools/extract_tools.py | 3 ++- .../{hybrid_udpate.py => hybrid_update.py} | 0 5 files changed, 19 insertions(+), 16 deletions(-) rename pipt/update_schemes/update_methods_ns/{hybrid_udpate.py => hybrid_update.py} (100%) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 601f58c..cc23516 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -163,7 +163,7 @@ def __init__(self, keys_en: dict, sim, redund_sim=None): if 'multilevel' in self.keys_en: self.multilevel = extract.extract_multilevel_info(self.keys_en['multilevel']) self.ml_ne = self.multilevel['ml_ne'] - self.tot_level = int(self.multilevel['levels']) + self.tot_level = len(self.multilevel['levels']) self.ml_corr_done = False @@ -211,19 +211,18 @@ def calc_prediction(self, enX=None, save_prediction=None): containing the responses at each time step given in PREDICTION. """ + # Use input state if given + if enX is None: + use_input_ensemble = False + enX = self.enX + self.enX = None # free memory + else: + use_input_ensemble = True - if isinstance(self.enX,list) and hasattr(self, 'multilevel'): # assume multilevel is used if state is a list + if isinstance(enX,list) and hasattr(self, 'multilevel'): # assume multilevel is used if state is a list success = self.calc_ml_prediction(enX) else: - # Use input state if given - if enX is None: - use_input_ensemble = False - enX = self.enX - self.enX = None # free memory - else: - use_input_ensemble = True - # Number of parallel runs nparallel = int(self.sim.input_dict.get('parallel', 1)) self.pred_data = [] @@ -321,6 +320,9 @@ def calc_prediction(self, enX=None, save_prediction=None): # some predicted data might need to be adjusted (e.g. scaled or compressed if it is 4D seis data). Do not # include this here. + if enX is not None: + self.enX = enX + enX = None # free memory # Store results if needed if save_prediction is not None: @@ -437,7 +439,7 @@ def calc_ml_prediction(self, enX=None): if ml_ne: level_enX = entools.matrix_to_list(enX[level], self.idX) - for n in range(ml_ne): + for n in ml_ne: if self.aux_input is not None: level_enX[n]['aux_input'] = self.aux_input[n] diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index d5cd22b..e8b74d6 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -561,7 +561,7 @@ def set_observations(self): def _ext_scaling(self): # get vector of scaling self.state_scaling = at.calc_scaling( - self.prior_enX, self.idX.keys(), self.prior_info) + self.prior_enX, self.idX, self.prior_info) self.Am = None diff --git a/pipt/misc_tools/analysis_tools.py b/pipt/misc_tools/analysis_tools.py index b48f290..5c9639c 100644 --- a/pipt/misc_tools/analysis_tools.py +++ b/pipt/misc_tools/analysis_tools.py @@ -1234,7 +1234,7 @@ def aug_state(state, list_state, cell_index=None): return aug -def calc_scaling(state, list_state, prior_info): +def calc_scaling(enX, idX, prior_info): """ Form the scaling to be used in svd related algoritms. Scaling consist of standard deviation for each `STATICVAR` It is important that this is formed in the same manner as the augmentet state vector is formed. Hence, with the same @@ -1256,7 +1256,7 @@ def calc_scaling(state, list_state, prior_info): """ scaling = [] - for elem in list_state: + for elem in idX.keys(): # more than single value. This is for multiple layers. Assume all values are active if len(prior_info[elem]['variance']) > 1: scaling.append(np.concatenate(tuple(np.sqrt(prior_info[elem]['variance'][z]) * @@ -1265,7 +1265,7 @@ def calc_scaling(state, list_state, prior_info): for z in range(prior_info[elem]['nz'])))) else: scaling.append(tuple(np.sqrt(prior_info[elem]['variance']) * - np.ones(state[elem].shape[0]))) + np.ones(enX[idX[elem][0]:idX[elem][1]].shape[0]))) return np.concatenate(scaling) diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index 13c1fe8..31624d2 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -281,13 +281,14 @@ def extract_initial_controls(keys: dict) -> dict: - + def extract_multilevel_info(keys: Union[dict, list]) -> dict: ''' Extract the info needed for ML simulations. Note if the ML keyword is not in keys_en we initialize such that we only have one level -- the high fidelity one ''' + keys_ml = keys if isinstance(keys, list): keys_ml = list_to_dict(keys) assert isinstance(keys_ml, dict) diff --git a/pipt/update_schemes/update_methods_ns/hybrid_udpate.py b/pipt/update_schemes/update_methods_ns/hybrid_update.py similarity index 100% rename from pipt/update_schemes/update_methods_ns/hybrid_udpate.py rename to pipt/update_schemes/update_methods_ns/hybrid_update.py From 5083684d01dadce88edc5f945a8112ea4d1eeb41 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 14 Jan 2026 09:31:43 +0100 Subject: [PATCH 82/94] Fix bugs --- pipt/update_schemes/multilevel.py | 4 +- popt/loop/ensemble_base.py | 23 ++-- popt/misc_tools/optim_tools.py | 16 +-- tests/test_toggle_ml_state.py | 168 ++++++++++++++++++++++++++++++ 4 files changed, 182 insertions(+), 29 deletions(-) create mode 100644 tests/test_toggle_ml_state.py diff --git a/pipt/update_schemes/multilevel.py b/pipt/update_schemes/multilevel.py index 1555ee3..2f9d970 100644 --- a/pipt/update_schemes/multilevel.py +++ b/pipt/update_schemes/multilevel.py @@ -160,7 +160,7 @@ def calc_analysis(self): for l in range(self.tot_level): self.ml_enObs[l], self.scale_data[l] = cholesky.gen_real( - self.ml_enObs[l], + self.vecObs, self.alpha[self.iteration - 1] * self.cov_data, self.ml_ne[l], return_chol=True @@ -192,7 +192,7 @@ def check_convergence(self): for l in range(self.tot_level): _, enPred_level = at.aug_obs_pred_data( self.obs_data, - [el[l] for el in self.enX_temp], + [el[l] for el in self.pred_data], self.assim_index, self.list_datatypes ) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index 32f98bb..842ae44 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -113,19 +113,11 @@ def function(self, x, *args, **kwargs): # Run simulation x = self.invert_scale_state(x) + x = self._reorganize_multilevel_ensemble(x) run_success = self.calc_prediction(enX=x, save_prediction=self.save_prediction) + x = self._reorganize_multilevel_ensemble(x) x = self.scale_state(x).squeeze() - # convert x (nparray) to state (dict) - #self.state = self.vec_to_state(x) - - # run the simulation - #self._invert_scale_state() # ensure that state is in [lb,ub] - #self._set_multilevel_state(self.state, x) # set multilevel state if applicable - #run_success = self.calc_prediction(save_prediction=self.save_prediction) # calculate flow data - #self._set_multilevel_state(self.state, x) # For some reason this has to be done again after calc_prediction - #self._scale_state() # scale back to [0, 1] - # Evaluate the objective function if run_success: func_values = self.obj_func( @@ -254,11 +246,12 @@ def save_stateX(self, path='./', filetype='npz'): np.savez_compressed(path + 'stateX.npz', **state_dict) elif filetype == 'npy': np.save(path + 'stateX.npy', stateX) - - def _set_multilevel_state(self, state, x): - if 'multilevel' in self.keys_en.keys() and len(x.shape) > 1: - en_size = ot.get_list_element(self.keys_en['multilevel'], 'en_size') - self.state = ot.toggle_ml_state(self.state, en_size) + + def _reorganize_multilevel_ensemble(self, x): + if ('multilevel' in self.keys_en) and (len(x.shape) > 1): + ml_ne = self.keys_en['multilevel']['ml_ne'] + x = ot.toggle_ml_state(x, ml_ne) + return x def _aux_input(self): diff --git a/popt/misc_tools/optim_tools.py b/popt/misc_tools/optim_tools.py index afe1908..735d9dc 100644 --- a/popt/misc_tools/optim_tools.py +++ b/popt/misc_tools/optim_tools.py @@ -120,28 +120,20 @@ def toggle_ml_state(state, ml_ne): """ if not isinstance(state,list): - L = len(ml_ne) # number of levels # initialize the state as an empty list of dictionaries with length equal self.tot_level - new_state = [{} for _ in range(L)] + new_state = [] # distribute the initial ensemble of states to the levels according to the given ensemble size. start = 0 # initialize - for l in range(L): + for l in range(len(ml_ne)): stop = start + ml_ne[l] - for el in state.keys(): - new_state[l][el] = state[el][:,start:stop] + new_state.append(state[:,start:stop]) start = stop del state else: # state is a list of levels - new_state = {} - for l in range(len(state)): - for el in state[l].keys(): - if el in new_state: - new_state[el] = np.hstack((new_state[el], state[l][el])) - else: - new_state[el] = state[l][el] + new_state = np.hstack(state) return new_state diff --git a/tests/test_toggle_ml_state.py b/tests/test_toggle_ml_state.py new file mode 100644 index 0000000..435bb6a --- /dev/null +++ b/tests/test_toggle_ml_state.py @@ -0,0 +1,168 @@ +""" +Test suite for toggle_ml_state function in popt.misc_tools.optim_tools +""" + +import numpy as np +import pytest +from popt.misc_tools.optim_tools import toggle_ml_state + + +def test_toggle_ml_state_matrix_to_list(): + """Test converting a matrix state to a list of levels""" + # Create a sample state matrix (rows=state_dim, cols=total_ensemble) + state_dim = 10 + total_ensemble = 15 + state = np.random.rand(state_dim, total_ensemble) + + # Define multilevel ensemble sizes + ml_ne = [5, 7, 3] # 3 levels with 5, 7, and 3 members respectively + + # Toggle to list format + result = toggle_ml_state(state, ml_ne) + + # Check that result is a list + assert isinstance(result, list) + + # Check that we have the correct number of levels + assert len(result) == len(ml_ne) + + # Check that each level has the correct ensemble size + for i, ne in enumerate(ml_ne): + assert result[i].shape == (state_dim, ne) + + # Check that the data is correctly distributed + start = 0 + for i, ne in enumerate(ml_ne): + stop = start + ne + np.testing.assert_array_equal(result[i], state[:, start:stop]) + start = stop + + +def test_toggle_ml_state_list_to_matrix(): + """Test converting a list of levels back to a matrix state""" + # Create sample state as list of levels + state_dim = 10 + ml_ne = [5, 7, 3] + + state_list = [ + np.random.rand(state_dim, ml_ne[0]), + np.random.rand(state_dim, ml_ne[1]), + np.random.rand(state_dim, ml_ne[2]) + ] + + # Toggle to matrix format + result = toggle_ml_state(state_list, ml_ne) + + # Check that result is a numpy array + assert isinstance(result, np.ndarray) + + # Check dimensions + total_ensemble = sum(ml_ne) + assert result.shape == (state_dim, total_ensemble) + + # Check that data is correctly concatenated + start = 0 + for i, ne in enumerate(ml_ne): + stop = start + ne + np.testing.assert_array_equal(result[:, start:stop], state_list[i]) + start = stop + + +def test_toggle_ml_state_roundtrip(): + """Test that toggling back and forth preserves the data""" + # Create initial state matrix + state_dim = 8 + total_ensemble = 12 + original_state = np.random.rand(state_dim, total_ensemble) + + ml_ne = [4, 5, 3] + + # Toggle to list then back to matrix + state_list = toggle_ml_state(original_state, ml_ne) + restored_state = toggle_ml_state(state_list, ml_ne) + + # Check that we get back the original state + np.testing.assert_array_equal(restored_state, original_state) + + +def test_toggle_ml_state_single_level(): + """Test with a single level (edge case)""" + state_dim = 5 + ensemble_size = 10 + state = np.random.rand(state_dim, ensemble_size) + + ml_ne = [ensemble_size] + + # Toggle to list + result = toggle_ml_state(state, ml_ne) + + assert isinstance(result, list) + assert len(result) == 1 + np.testing.assert_array_equal(result[0], state) + + # Toggle back + restored = toggle_ml_state(result, ml_ne) + np.testing.assert_array_equal(restored, state) + + +def test_toggle_ml_state_many_levels(): + """Test with many levels""" + state_dim = 6 + ml_ne = [2, 3, 1, 4, 2, 3] # 6 levels + total_ensemble = sum(ml_ne) + + state = np.random.rand(state_dim, total_ensemble) + + # Toggle to list + result = toggle_ml_state(state, ml_ne) + + assert len(result) == len(ml_ne) + for i, ne in enumerate(ml_ne): + assert result[i].shape[1] == ne + + # Toggle back and verify + restored = toggle_ml_state(result, ml_ne) + np.testing.assert_array_equal(restored, state) + + +def test_toggle_ml_state_preserves_values(): + """Test that specific values are preserved correctly""" + # Create a state with known values for verification + state = np.array([ + [1.0, 2.0, 3.0, 4.0, 5.0], + [10.0, 20.0, 30.0, 40.0, 50.0] + ]) + + ml_ne = [2, 3] + + # Toggle to list + result = toggle_ml_state(state, ml_ne) + + # Check first level + expected_level0 = np.array([[1.0, 2.0], [10.0, 20.0]]) + np.testing.assert_array_equal(result[0], expected_level0) + + # Check second level + expected_level1 = np.array([[3.0, 4.0, 5.0], [30.0, 40.0, 50.0]]) + np.testing.assert_array_equal(result[1], expected_level1) + + +def test_toggle_ml_state_empty_level(): + """Test behavior with empty levels (ensemble size = 0)""" + state_dim = 4 + ml_ne = [3, 0, 2] # Middle level has no members + total_ensemble = sum(ml_ne) + + state = np.random.rand(state_dim, total_ensemble) + + # Toggle to list + result = toggle_ml_state(state, ml_ne) + + assert len(result) == len(ml_ne) + assert result[0].shape == (state_dim, 3) + assert result[1].shape == (state_dim, 0) # Empty array + assert result[2].shape == (state_dim, 2) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From f88fe0e16ab2b3b573a0dbe0290b15684c214563 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 14 Jan 2026 09:58:32 +0100 Subject: [PATCH 83/94] Comment out log.info --- misc/ecl.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/misc/ecl.py b/misc/ecl.py index e645319..2051bdf 100644 --- a/misc/ecl.py +++ b/misc/ecl.py @@ -426,8 +426,8 @@ def __init__(self, root): self.ni = grid_head[1] # pylint: disable=invalid-name self.nj = grid_head[2] # pylint: disable=invalid-name self.nk = grid_head[3] # pylint: disable=invalid-name - log.info("Grid dimension is %d x %d x %d", - self.ni, self.nj, self.nk) + #log.info("Grid dimension is %d x %d x %d", + # self.ni, self.nj, self.nk) # also store a shape tuple which describes the grid cube self.shape = (self.nk, self.nj, self.ni) @@ -460,7 +460,7 @@ def __init__(self, root): # restart properties are only saved for the active elements, # so we can cache this number to compare self.num_active = numpy.sum(self.actnum) - log.info("Grid has %d active cells", self.num_active) + #log.info("Grid has %d active cells", self.num_active) def grid(self): """ From c6ed492a36add55ed83d69cf5d8327bd028076b9 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Wed, 14 Jan 2026 12:34:23 +0100 Subject: [PATCH 84/94] Fix small bug --- popt/loop/ensemble_base.py | 7 ------- popt/misc_tools/optim_tools.py | 2 +- 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index dd5203d..81d1256 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -141,13 +141,6 @@ def function(self, x, *args, **kwargs): self.enF = func_values return func_values - - # Add to covariance - self.cov = np.append(self.cov, var) - self.dim = self.cov.shape[0] - - # Make cov full covariance matrix - self.cov = np.diag(self.cov) def get_state(self): """ diff --git a/popt/misc_tools/optim_tools.py b/popt/misc_tools/optim_tools.py index bc5924b..5a8bdc3 100644 --- a/popt/misc_tools/optim_tools.py +++ b/popt/misc_tools/optim_tools.py @@ -341,7 +341,7 @@ def get_optimize_result(obj): if 'args' in savedata: for a, arg in enumerate(obj.args): - results[f'args[{a}]'] = arg + save_dict[f'args[{a}]'] = arg # Loop over variables to store in save list for save_typ in savedata: From fc4668ce4e8e4732629790afd89ccb4ac298c450 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 15 Jan 2026 10:11:01 +0100 Subject: [PATCH 85/94] Fix bug relating to data_mismatch for multilevel --- pipt/update_schemes/multilevel.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/pipt/update_schemes/multilevel.py b/pipt/update_schemes/multilevel.py index 8d90492..bf55148 100644 --- a/pipt/update_schemes/multilevel.py +++ b/pipt/update_schemes/multilevel.py @@ -62,8 +62,13 @@ def __init__(self, keys_da,keys_fwd,sim): self.assim_index = [self.keys_da['obsname'], self.keys_da['assimindex'][0]] self.list_datatypes, self.list_act_datatypes = at.get_list_data_types(self.obs_data, self.assim_index) - self.cov_data = at.gen_covdata(self.datavar, self.assim_index, self.list_datatypes) - self.vecObs, self.enObs = self.set_observations() + self.cov_data = at.gen_covdata(self.datavar, self.assim_index, self.list_datatypes) + self.vecObs, _ = at.aug_obs_pred_data( + self.obs_data, + self.pred_data, + self.assim_index, + self.list_datatypes + ) def _init_sim(self): """ @@ -122,7 +127,7 @@ def calc_analysis(self): # Note, evaluate for high fidelity model data_misfit = at.calc_objectivefun( - self.enObs, + self.enObs_conv, np.concatenate(self.enPred,axis=1), # Is this correct, given the comment above?????? self.cov_data ) @@ -199,7 +204,7 @@ def check_convergence(self): enPred.append(enPred_level) data_misfit = at.calc_objectivefun( - self.enObs, + self.enObs_conv, np.concatenate(enPred,axis=1), self.cov_data ) From 3cc09aa05886e7eab1e2d63b0e72c4d6d788763b Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 15 Jan 2026 10:13:58 +0100 Subject: [PATCH 86/94] Remove unused imports for multilevel --- pipt/update_schemes/multilevel.py | 22 +++------------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/pipt/update_schemes/multilevel.py b/pipt/update_schemes/multilevel.py index bf55148..cb8a493 100644 --- a/pipt/update_schemes/multilevel.py +++ b/pipt/update_schemes/multilevel.py @@ -3,33 +3,17 @@ inherit the ensemble class, hence the main loop is inherited. These classes will consider the analysis step. ''' -# local imports. Note, it is assumed that PET is installed and available in the path. +#────────────────────────────────────────────────────────────────────────────────────── from pipt.loop.ensemble import Ensemble -from pipt.update_schemes.esmda import esmda_approx from pipt.update_schemes.esmda import esmdaMixIn from pipt.misc_tools import analysis_tools as at import pipt.misc_tools.ensemble_tools as entools from geostat.decomp import Cholesky -from misc import ecl - from pipt.update_schemes.update_methods_ns.hybrid_update import hybrid_update -# system imports + import numpy as np -from scipy.sparse import coo_matrix -from scipy import linalg -import time -import shutil -import pickle -from scipy.linalg import solve # For linear system solvers -from scipy.stats import multivariate_normal -from scipy import sparse from copy import deepcopy -import random -import os -import sys -from scipy.stats import ortho_group -from shutil import copyfile -import math +#────────────────────────────────────────────────────────────────────────────────────── __all__ = ['multilevel', 'esmda_hybrid'] From 5f935aa0feaa8adc2dc1fa77ebd3e9254d8925c9 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 15 Jan 2026 10:15:29 +0100 Subject: [PATCH 87/94] Add progbar_settings to multilevel --- ensemble/ensemble.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 1699e3a..9d6f775 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -422,7 +422,7 @@ def calc_ml_prediction(self, enX=None): no_tot_run = int(self.sim.input_dict['parallel']) ml_pred_data = [] - for level in tqdm(self.multilevel['levels'], desc='Fidelity level', position=1): + for level in tqdm(self.multilevel['levels'], desc='Fidelity level', position=1, **progbar_settings): # Setup forward simulator and redundant simulator at the correct fidelity if self.sim.redund_sim is not None: From 295056812db554dd648a2754fb6561b7bd8ad1bb Mon Sep 17 00:00:00 2001 From: Kristian Fossum Date: Thu, 15 Jan 2026 10:29:29 +0100 Subject: [PATCH 88/94] Update popt/update_schemes/linesearch.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- popt/update_schemes/linesearch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index a8bb0f8..2c56887 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -199,7 +199,7 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c else: self.callback = None - # Remove 'datatype' form options if present (This is a temporary bugfix) + # Remove 'datatype' from options if present (This is a temporary bugfix) self.options.pop('datatype', None) # Custom convergence criteria (callable) From 542c51eb5d99f96a9365c02d8c627efb05741453 Mon Sep 17 00:00:00 2001 From: Kristian Fossum Date: Thu, 15 Jan 2026 10:30:49 +0100 Subject: [PATCH 89/94] Update pipt/update_schemes/update_methods_ns/hybrid_update.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- pipt/update_schemes/update_methods_ns/hybrid_update.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pipt/update_schemes/update_methods_ns/hybrid_update.py b/pipt/update_schemes/update_methods_ns/hybrid_update.py index b364209..e7d5bc5 100644 --- a/pipt/update_schemes/update_methods_ns/hybrid_update.py +++ b/pipt/update_schemes/update_methods_ns/hybrid_update.py @@ -59,7 +59,7 @@ def update(self, enX, enY, enE, **kwargs): else: enXcentered.append(self.scale(np.dot(enX[l], self.proj[l]), self.state_scaling)) - # Calculet truncated SVD of predicted data ensemble at level l + # Calculate truncated SVD of predicted data ensemble at level l enYcentered = self.scale(np.dot(enY[l], self.proj[l]), self.scale_data[l]) Ud, Sd, VTd = at.truncSVD(enYcentered, energy=self.trunc_energy) From 559c9bae3cc14e0a5cb233ae158d2f296e43cd6e Mon Sep 17 00:00:00 2001 From: Kristian Fossum Date: Thu, 15 Jan 2026 10:31:09 +0100 Subject: [PATCH 90/94] Update pipt/misc_tools/extract_tools.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- pipt/misc_tools/extract_tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pipt/misc_tools/extract_tools.py b/pipt/misc_tools/extract_tools.py index 58445be..4cb2326 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -1,4 +1,4 @@ -# This module inlcudes functions for extracting information from input dicts +# This module includes functions for extracting information from input dicts __all__ = [ 'extract_prior_info', From af200cafb0dfbe6204bdd9c5681ba6a7137d73a0 Mon Sep 17 00:00:00 2001 From: Kristian Fossum Date: Thu, 15 Jan 2026 10:32:45 +0100 Subject: [PATCH 91/94] Update pipt/loop/ensemble.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- pipt/loop/ensemble.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pipt/loop/ensemble.py b/pipt/loop/ensemble.py index f3b51d5..9187535 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -163,7 +163,7 @@ def check_assimindex_simultaneous(self): def _org_obs_data(self): """ Organize the input true observed data. The obs_data will be a list of length equal length of "TRUEDATAINDEX", - and each entery in the list will be a dictionary with keys equal to the "DATATYPE". + and each entry in the list will be a dictionary with keys equal to the "DATATYPE". Also, the pred_data variable (predicted data or forward simulation) will be initialized here with the same structure as the obs_data variable. From 16ed42ee357c93721facc54ca7d06473f5dec4f1 Mon Sep 17 00:00:00 2001 From: "Rolf J. Lorentzen" Date: Thu, 15 Jan 2026 13:57:33 +0100 Subject: [PATCH 92/94] Update optimize.py --- popt/loop/optimize.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/popt/loop/optimize.py b/popt/loop/optimize.py index 683c677..42a0cad 100644 --- a/popt/loop/optimize.py +++ b/popt/loop/optimize.py @@ -3,13 +3,14 @@ import numpy as np import time import pickle +from abc import ABC, abstractmethod # Internal imports import popt.misc_tools.optim_tools as ot from ensemble.logger import PetLogger -class Optimize: +class Optimize(ABC): """ Class for ensemble optimization algorithms. These are classified by calculating the sensitivity or gradient using ensemble instead of classical derivatives. The loop is else as a classic optimization loop: a state (or control From f8f1a74646c323a3e3452309895658af2d601bc3 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 15 Jan 2026 14:34:02 +0100 Subject: [PATCH 93/94] Fix bugs --- ensemble/ensemble.py | 2 +- popt/loop/ensemble_base.py | 2 ++ popt/loop/optimize.py | 4 ++-- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 9d6f775..3c892ab 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -233,7 +233,7 @@ def calc_prediction(self, enX=None, save_prediction=None): # Run setup function for simulator if hasattr(self.sim, 'setup_fwd_run'): self.sim.setup_fwd_run(redund_sim=self.sim.redund_sim) - + # Convert ensemble matrix to list of dictionaries enX = entools.matrix_to_list(enX, self.idX) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index 81d1256..f83f352 100644 --- a/popt/loop/ensemble_base.py +++ b/popt/loop/ensemble_base.py @@ -258,6 +258,8 @@ def _reorganize_multilevel_ensemble(self, x): ml_ne = self.keys_en['multilevel']['ml_ne'] x = ot.toggle_ml_state(x, ml_ne) return x + else: + return x def _aux_input(self): diff --git a/popt/loop/optimize.py b/popt/loop/optimize.py index 683c677..24b5e82 100644 --- a/popt/loop/optimize.py +++ b/popt/loop/optimize.py @@ -3,13 +3,14 @@ import numpy as np import time import pickle +from abc import ABC, abstractmethod # Internal imports import popt.misc_tools.optim_tools as ot from ensemble.logger import PetLogger -class Optimize: +class Optimize(ABC): """ Class for ensemble optimization algorithms. These are classified by calculating the sensitivity or gradient using ensemble instead of classical derivatives. The loop is else as a classic optimization loop: a state (or control @@ -94,7 +95,6 @@ def __init__(self, **options): self.options = None self.mean_state = None self.obj_func_values = None - self.fun = None # objective function self.obj_func_tol = None # objective tolerance limit # Initialize number of function and jacobi evaluations From b52f83f5309ce4ecefecaff0c456a8fa143fd83b Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Fri, 16 Jan 2026 12:19:15 +0100 Subject: [PATCH 94/94] Fix bugs --- ensemble/ensemble.py | 20 ++++++++++++++-- popt/update_schemes/linesearch.py | 24 +++++++++---------- .../update_schemes/subroutines/subroutines.py | 24 +++++++++++++------ 3 files changed, 47 insertions(+), 21 deletions(-) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 3c892ab..4f67361 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -209,6 +209,8 @@ def calc_prediction(self, enX=None, save_prediction=None): containing the responses at each time step given in PREDICTION. """ + one_state = False + # Use input state if given if enX is None: use_input_ensemble = False @@ -233,10 +235,20 @@ def calc_prediction(self, enX=None, save_prediction=None): # Run setup function for simulator if hasattr(self.sim, 'setup_fwd_run'): self.sim.setup_fwd_run(redund_sim=self.sim.redund_sim) - + + if enX.ndim == 1: + one_state = True + enX = enX[:, np.newaxis] + elif enX.shape[1] == 1: + one_state = True + + # If we have several models (num_models) but only one state input + if one_state and self.ne > 1: + enX = np.tile(enX, (1, self.ne)) + # Convert ensemble matrix to list of dictionaries enX = entools.matrix_to_list(enX, self.idX) - + if not (self.aux_input is None): for n in range(self.ne): enX[n]['aux_input'] = self.aux_input[n] @@ -268,6 +280,10 @@ def calc_prediction(self, enX=None, save_prediction=None): # Convert state enemble back to matrix form enX = entools.list_to_matrix(enX, self.idX) + # If only one state was inputted, keep only that state + if one_state and self.ne > 1: + enX = enX[:,0][:,np.newaxis] + # restore state ensemble if it was not inputted if not use_input_ensemble: self.enX = enX diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index 2c56887..abd546e 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -258,9 +258,9 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c # Check for initial inverse hessian for the BFGS method if self.method == 'BFGS': - self.Hk_inv = options.get('hess0_inv', np.eye(x.size)) + self._Hk_inv = options.get('hess0_inv', np.eye(x.size)) else: - self.Hk_inv = None + self._Hk_inv = None # Initialize some variables self.f_old = None @@ -276,8 +276,8 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c self.logger(f'\n \nUSER-SPECIFIED OPTIONS:\n{pprint.pformat(OptimizeResult(self.options))}\n') self.logger(**{ 'iter.': 0, - fun_xk_symbol: self.fk, - jac_inf_symbol: la.norm(self.jk, np.inf), + fun_xk_symbol: self._fk, + jac_inf_symbol: la.norm(self._jk, np.inf), 'step-size': self.step_size }) @@ -358,9 +358,9 @@ def calc_update(self, iter_resamp=0): if self.method == 'GD': pk = - self._jk if self.method == 'BFGS': - pk = - np.matmul(self.Hk_inv, self._jk) + pk = - np.matmul(self._Hk_inv, self._jk) if self.method == 'Newton-CG': - pk = newton_cg(self.jk, Hk=self.Hk, xk=self.xk, jac=self._jac, logger=self.logger) + pk = newton_cg(self._jk, Hk=self._Hk, xk=self._xk, jac=self._jk, logger=self.logger) # porject search direction onto the feasible set if self.bounds is not None: @@ -426,7 +426,7 @@ def calc_update(self, iter_resamp=0): if self.method == 'BFGS': yk = j_new - j_old if self.iteration == 1: self._Hk_inv = np.dot(yk,sk)/np.dot(yk,yk) * np.eye(sk.size) - self.Hk_inv = bfgs_update(self.Hk_inv, sk, yk) + self._Hk_inv = bfgs_update(self._Hk_inv, sk, yk) # Update status success = True @@ -440,8 +440,8 @@ def calc_update(self, iter_resamp=0): if self.logger is not None: self.logger(**{ 'iter.': self.iteration, - fun_xk_symbol: self.fk, - jac_inf_symbol: la.norm(self.jk, np.inf), + fun_xk_symbol: self._fk, + jac_inf_symbol: la.norm(self._jk, np.inf), 'step-size': step_size }) @@ -538,9 +538,9 @@ def _set_step_size(self, pk, amax): alpha = self.step_size else: - if (self.step_size_adapt == 1) and (np.dot(pk, self.jk) != 0): - alpha = 2*(self.fk - self.f_old)/np.dot(pk, self.jk) - elif (self.step_size_adapt == 2) and (np.dot(pk, self.jk) != 0): + if (self.step_size_adapt == 1) and (np.dot(pk, self._jk) != 0): + alpha = 2*(self._fk - self.f_old)/np.dot(pk, self._jk) + elif (self.step_size_adapt == 2) and (np.dot(pk, self._jk) != 0): slope_old = np.dot(self.p_old, self.j_old) slope_new = np.dot(pk, self._jk) alpha = self.step_size*slope_old/slope_new diff --git a/popt/update_schemes/subroutines/subroutines.py b/popt/update_schemes/subroutines/subroutines.py index 3700be2..17402f2 100644 --- a/popt/update_schemes/subroutines/subroutines.py +++ b/popt/update_schemes/subroutines/subroutines.py @@ -98,7 +98,7 @@ def phi(alpha): else: phi.fun_val = fk else: - logger(' Evaluating Armijo condition') + #logger(' Evaluating Armijo condition') phi.fun_val = fun(xk + alpha*pk) ls_nfev += 1 return phi.fun_val @@ -113,7 +113,7 @@ def dphi(alpha): else: dphi.jac_val = jk else: - logger(' Evaluating curvature condition') + #logger(' Evaluating curvature condition') dphi.jac_val = jac(xk + alpha*pk) ls_njev += 1 return np.dot(dphi.jac_val, pk) @@ -125,31 +125,37 @@ def dphi(alpha): # Start loop a = [0, step_size] for i in range(1, maxiter+1): - logger(f'Line search iteration: {i-1}') + logger(f'iteration: {i-1}') # Evaluate phi(ai) phi_i = phi(a[i]) # Check for sufficient decrease if (phi_i > phi_0 + c1*a[i]*dphi_0) or (phi_i >= phi(a[i-1]) and i>0): + logger(' Armijo condition: not satisfied') # Call zoom function - step_size = zoom(a[i-1], a[i], phi, dphi, phi_0, dphi_0, maxiter+1-i, c1, c2) + step_size = zoom(a[i-1], a[i], phi, dphi, phi_0, dphi_0, maxiter+1-i, c1, c2, iter_id=i) logger('──────────────────────────────────────────────────') return step_size, phi.fun_val, dphi.jac_val, ls_nfev, ls_njev + + logger(' Armijo condition: satisfied') # Evaluate dphi(ai) dphi_i = dphi(a[i]) # Check curvature condition if abs(dphi_i) <= -c2*dphi_0: + logger(' Curvature condition: satisfied') step_size = a[i] logger('──────────────────────────────────────────────────') return step_size, phi.fun_val, dphi.jac_val, ls_nfev, ls_njev + + logger(' Curvature condition: not satisfied') # Check for posetive derivative if dphi_i >= 0: # Call zoom function - step_size = zoom(a[i], a[i-1], phi, dphi, phi_0, dphi_0, maxiter+1-i, c1, c2) + step_size = zoom(a[i], a[i-1], phi, dphi, phi_0, dphi_0, maxiter+1-i, c1, c2, iter_id=i) logger('──────────────────────────────────────────────────') return step_size, phi.fun_val, dphi.jac_val, ls_nfev, ls_njev @@ -163,7 +169,7 @@ def dphi(alpha): return None, None, None, ls_nfev, ls_njev -def zoom(alo, ahi, f, df, f0, df0, maxiter, c1, c2): +def zoom(alo, ahi, f, df, f0, df0, maxiter, c1, c2, iter_id=0): '''Zoom function for line search algorithm. (This is the same as for scipy)''' phi_lo = f(alo) @@ -171,7 +177,7 @@ def zoom(alo, ahi, f, df, f0, df0, maxiter, c1, c2): dphi_lo = df(alo) for j in range(maxiter): - logger(f'Line search iteration: {j+1}') + logger(f'iteration: {iter_id+j+1}') tol_cubic = 0.2*(ahi-alo) tol_quad = 0.1*(ahi-alo) @@ -195,6 +201,7 @@ def zoom(alo, ahi, f, df, f0, df0, maxiter, c1, c2): # Check for sufficient decrease if (phi_j > f0 + c1*aj*df0) or (phi_j >= phi_lo): + logger(' Armijo condition: not satisfied') # store old values aold = ahi phi_old = phi_hi @@ -202,11 +209,14 @@ def zoom(alo, ahi, f, df, f0, df0, maxiter, c1, c2): ahi = aj phi_hi = phi_j else: + logger(' Armijo condition: satisfied') # check curvature condition dphi_j = df(aj) if abs(dphi_j) <= -c2*df0: + logger(' Curvature condition: satisfied') return aj + logger(' Curvature condition: not satisfied') if dphi_j*(ahi-alo) >= 0: # store old values aold = ahi