fixed some issues

ensemble/__init__.py

@@ -1 +1 @@
-"""Multiple realisations management."""
+"""Multiple realisations management."""

ensemble/ensemble.py

@@ -18,13 +18,15 @@
 
 # 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
 from misc import read_input_csv as rcsv
 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 +58,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 +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_<STATICVAR-name> keyword.
@@ -143,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()
+
+        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 +180,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_<STATE entries>.
-        '''
-
-        # Get state names as list
-        state_names = self.keys_en['state']
-        if not isinstance(state_names, list): state_names = [state_names]
-
-        # Check if PRIOR_<state names> 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 +251,7 @@ def gen_init_ensemble(self):
         # 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.

input_output/read_config.py

@@ -55,13 +55,13 @@ 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']
         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
@@ -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)

pipt/loop/assimilation.py

@@ -16,10 +16,11 @@
 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
+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
@@ -83,15 +84,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 +113,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:
@@ -279,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