From 9e3a8216a6639f41e3430773aa06f82ff390f797 Mon Sep 17 00:00:00 2001
From: Mathias Methlie Nilsen <mani@cno-0006.ad.norceresearch.no>
Date: Wed, 9 Jul 2025 13:05:24 +0200
Subject: [PATCH 1/6] 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 43f99c61..6fa80a44 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 b8fe92f3..956cd9c6 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 5272d750..9ed6e5e3 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 <mani@bgo-1714.ad.norceresearch.no>
Date: Wed, 16 Jul 2025 14:13:39 +0200
Subject: [PATCH 2/6] 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 9ed6e5e3..cd616f93 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 <mani@bgo-1714.ad.norceresearch.no>
Date: Tue, 5 Aug 2025 13:57:56 +0200
Subject: [PATCH 3/6] 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 956cd9c6..7bf1081f 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 <mani@bgo-1714.ad.norceresearch.no>
Date: Tue, 5 Aug 2025 14:00:39 +0200
Subject: [PATCH 4/6] 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 77ff998d..507ad8ca 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 33d4497c..2500fc2b 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 cacbe992..81809c26 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 <mani@bgo-1714.ad.norceresearch.no>
Date: Thu, 7 Aug 2025 10:19:10 +0200
Subject: [PATCH 5/6] 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 2500fc2b..951a2be3 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 62f73892..f4bf8326 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 5bcba7d4..2dc7536e 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 <mani@bgo-1714.ad.norceresearch.no>
Date: Thu, 7 Aug 2025 13:33:14 +0200
Subject: [PATCH 6/6] 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 ead116b3..5ded1783 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 951a2be3..e5363a2a 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 f4bf8326..2d334ca6 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 81809c26..c786627c 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')