Fixed class inheritance for popt ensembles (#100)

* update to TrustRegion

* some design changes to TrustRegion

* decoupled GenOpt from Ensemble

* decoupled GenOpt from Ensemble

* Cleaned up code duplication and renamed some stuff

* comments

---------

Co-authored-by: Mathias Methlie Nilsen <mani@cno-0006.ad.norceresearch.no>
Co-authored-by: Mathias Methlie Nilsen <mani@bgo-1714.ad.norceresearch.no>

Author: 3 people

popt/loop/__init__.py

@@ -1 +1 @@
-"""Main loop for running optimization."""
+"""Main loop for running optimization."""

popt/loop/ensemble_base.py

@@ -8,10 +8,12 @@
 # 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):
+__all__ = ['EnsembleOptimizationBaseClass']
+
+class EnsembleOptimizationBaseClass(SupEnsemble):
     '''
     Base class for the popt ensemble
     '''
@@ -33,61 +35,64 @@ def __init__(self, options, simulator, objective):
         else:
             sim = simulator
 
-        # Initialize PETEnsemble
+        # Initialize the PET Ensemble
         super().__init__(options, sim)
 
         # 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
+
+        # 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)]
             else:
                 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 +103,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 +128,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 +144,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):
         """