diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index 48cd716..4f67361 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -15,17 +15,23 @@ from tqdm.auto import tqdm from p_tqdm import p_map import logging -from geostat.decomp import Cholesky # Making realizations # 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 +import pipt.misc_tools.ensemble_tools as entools from misc.system_tools.environ_var import OpenBlasSingleThread # Single threaded OpenBLAS runs +# Settings +####################################################################################################### +progbar_settings = { + '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 + 'unit': 'member', +} +####################################################################################################### class Ensemble: """ @@ -34,7 +40,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, @@ -51,22 +57,17 @@ 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. 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' - ) - self.logger = logging.getLogger('PET') - # Check if folder contains any En_ files, and remove them! for folder in glob('En_*'): try: @@ -86,13 +87,13 @@ def __init__(self, keys_en, 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 @@ -121,15 +122,22 @@ 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 = 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: 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 @@ -137,130 +145,26 @@ 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 loaded state variables - tmp_ne = [] - for tmp_state in self.state.keys(): - tmp_ne.extend([self.state[tmp_state].shape[1]]) - - if 'ne' not in self.keys_en: # NE not specified in input file - if max(tmp_ne) != min(tmp_ne): #Check loaded ensembles are the same size (if more than one state variable) - print('\033[1;33mInput states have different ensemble size\033[1;m') - sys.exit(1) - self.ne = min(tmp_ne) # Use the number of ensemble members in loaded ensemble - else: - # Use the number of ensemble members specified in input file (may be fewer than loaded) - self.ne = int(self.keys_en['ne']) - if self.ne <= min(tmp_ne): - # pick correct number of ensemble members - self.state = {key: val[:,:self.ne] for key, val in self.state.items()} + for key in self.keys_en['staticvar']: + if self.enX is None: + self.enX = tmp_load[key] + self.ne = self.enX.shape[1] else: - print('\033[1;33mInput states are smaller than NE\033[1;m') - 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() + 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])) - 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 - ''' + # 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: - # 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 + self.multilevel = extract.extract_multilevel_info(self.keys_en['multilevel']) + self.ml_ne = self.multilevel['ml_ne'] + self.tot_level = len(self.multilevel['levels']) + self.ml_corr_done = False + - # 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. @@ -286,7 +190,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. @@ -305,122 +209,91 @@ def calc_prediction(self, input_state=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 + enX = self.enX + self.enX = None # free memory + else: + use_input_ensemble = True - 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(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']) - else: - no_tot_run = 1 + 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) - - # 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)) - - # modified by xluo, for including the simulation of the mean reservoir model - # as used in the RLM-MAC algorithm - if 'daalg' in self.keys_en and self.keys_en['daalg'][1] == 'gies': - list_state.append({}) - list_member_index.append(self.ne) - - for key in self.state.keys(): - tmp_state = np.zeros(list_state[0][key].shape[0]) - - for i in range(self.ne): - tmp_state += list_state[i][key] - - list_state[self.ne][key] = tmp_state / 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 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) + + 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] + + ###################################################################################################################### + # 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 - 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) + 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 + 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, + **progbar_settings + ) + ###################################################################################################################### + + # 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 + 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() @@ -428,7 +301,7 @@ def calc_prediction(self, input_state=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 @@ -438,37 +311,90 @@ 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') - 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]) + 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.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((np.atleast_2d(el[ind][typ]).T) for el in en_pred), axis=1) + self.pred_data.extend([{typ: np.concatenate(tuple((el[ind][typ][:, np.newaxis]) for el in en_pred), axis=1) if any(elem is not None for elem in tuple((el[ind][typ]) for el in en_pred)) else None for typ in en_pred[0][0].keys()} for ind in range(len(en_pred[0]))]) # 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: 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(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): """ @@ -497,7 +423,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. @@ -505,95 +431,45 @@ def calc_ml_prediction(self, input_state=None): Parameters ---------- - input_state: + enX: If simulation is run stand-alone one can input any state. """ no_tot_run = int(self.sim.input_dict['parallel']) - ml_pred_data = ml_pred_data = [[] for _ in range(max(self.multilevel['levels'])+1)] # Ensure that the list is long enough + ml_pred_data = [] + + for level in tqdm(self.multilevel['levels'], desc='Fidelity level', position=1, **progbar_settings): - 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 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) - 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 - en_pred = [] - batch_en = [np.arange(start, start + batch_size) for start in - np.arange(0, len(ml_ne) - batch_size, batch_size)] - if len(batch_en): # if len(ml_ne) is less than batch_size - batch_en.append(np.arange(batch_en[-1][-1]+1, len(ml_ne))) - else: - batch_en.append(np.arange(0, len(ml_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 - 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) + # Run prediction in parallel using p_map + 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] @@ -607,7 +483,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 @@ -624,61 +500,44 @@ 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.info(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[level].extend([{typ: np.concatenate(tuple((np.atleast_2d(el[ind][typ]).T) for el in en_pred), axis=1) - if any(elem is not None for elem in tuple((el[ind][typ]) for el in en_pred)) - else None for typ in en_pred[0][0].keys()} for ind in range(len(en_pred[0]))]) - - empty_indices = [i for i, entry in enumerate(ml_pred_data) if not isinstance(entry, list) or not entry] - filled_idx = next( - (i for i, entry in enumerate(ml_pred_data) if isinstance(entry, list) and entry), - None - ) + ml_pred_data.append([{typ: np.concatenate(tuple((el[ind][typ][:, np.newaxis]) for el in en_pred), axis=1) + if any(elem is not None for elem in tuple((el[ind][typ]) for el in en_pred)) + else None for typ in en_pred[0][0].keys()} for ind in range(len(en_pred[0]))]) - if len(empty_indices): - # fill the emply list with correct structure - template = deepcopy(ml_pred_data[filled_idx]) - empty_struct = [#{ - #key: None - #for key, val in template_elem.items() - #} - None for template_elem in template] - for idx in empty_indices: - ml_pred_data[idx] = empty_struct # loop over time instance first, and the level instance. self.pred_data = np.array(ml_pred_data).T.tolist() - # Treat scaling - if 'scaling' in self.multilevel: - self.pred_data = self.treat_scaling(self.multilevel['scaling'],self.pred_data,self.multilevel['scaling_map']) - # this is for fidelity specific scaling of output - if hasattr(self,'treat_modeling_error'): self.treat_modeling_error() 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): @@ -686,7 +545,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/ensemble/logger.py b/ensemble/logger.py new file mode 100644 index 0000000..57010bc --- /dev/null +++ b/ensemble/logger.py @@ -0,0 +1,99 @@ +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 = 12 # Number of spaces for table formatting + + # Configurate logging + logging.basicConfig( + level=logging.INFO, + format='%(asctime)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. + + Example: + >>> logger = PetLogger() + >>> logger('This is a log message.') + 2024-06-01│12:00:00 : This is a log message. + >>> + >>> 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 │ + 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: + # Log message from args + msg = ' ' + ' '.join(str(arg) for arg in args) + self._logger.info(msg) + + if kwargs: + # Make strings for table logging + self._set_ns(**kwargs) + header = [] + values = [] + for key, value in kwargs.items(): + header.append(f'{key:^{self.ns}}') + try: + if isinstance(value, int) or isinstance(value, str): + values.append(f'{value:^{self.ns}}') + elif '%' in key: + values.append(f'{value:^{self.ns}.1f}') + else: + values.append(f'{value:^{self.ns}.3e}') + except: + 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('') + + 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/input_output/read_config.py b/input_output/read_config.py index 0d9d10d..c6af002 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) @@ -379,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/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): """ diff --git a/pipt/loop/assimilation.py b/pipt/loop/assimilation.py index 305f9a8..9266a45 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: @@ -46,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'): @@ -84,7 +92,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 @@ -96,12 +104,13 @@ 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 - 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 @@ -113,7 +122,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 @@ -123,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: @@ -159,20 +173,23 @@ 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 # 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 @@ -182,16 +199,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': @@ -200,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.state) - 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: - pickle.dump(self.ensemble.state, file) - with open('posterior_forecast.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(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. @@ -221,17 +238,17 @@ 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() + #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: 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): @@ -266,9 +283,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): @@ -306,35 +325,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. @@ -351,6 +341,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): @@ -361,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) @@ -440,7 +436,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 @@ -458,16 +457,9 @@ 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: + 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/loop/ensemble.py b/pipt/loop/ensemble.py index 839f259..9187535 100644 --- a/pipt/loop/ensemble.py +++ b/pipt/loop/ensemble.py @@ -1,7 +1,6 @@ """Descriptive description.""" # External import -import logging import os.path import numpy @@ -15,11 +14,15 @@ # 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 + +# 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): @@ -69,12 +72,9 @@ 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') - - # 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.') + # Setup logger + self.logger = PetLogger(filename='assim.log') + self.logger(f'=========== Running Data Assimilation - {keys_da["daalg"][0].upper()} ===========') # Internalize PIPT dictionary if not hasattr(self, 'keys_da'): @@ -99,19 +99,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': @@ -126,10 +115,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 = 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] @@ -500,150 +489,80 @@ def _org_data_var(self): self.datavar[i][datatype[j]] = est_noise # override the given value 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) - # Generate the data auto-covariance matrix - if 'emp_cov' in self.keys_da and self.keys_da['emp_cov'] == 'yes': + + 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, + 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 - tmp_E = np.dot(cholesky(self.cov_data).T, - np.random.randn(self.cov_data.shape[0], self.ne)) + # enObs: samples from N(0,Cd) + enObs = 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 + 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, + vecObs, + self.iteration + ) + + # Center the ensemble of perturbed observed data + enObs = vecObs[:, np.newaxis] - enObs + self.cov_data = np.var(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( - self.datavar, self.assim_index, self.list_datatypes) + 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': + 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) + data = self.cov_data, + aug_pred_data = self.enPred, + obs_data_vector = vecObs, + iteration = self.iteration + ) + + generator = Cholesky() # Initialize GeoStat class for generating realizations + enObs, self.scale_data = generator.gen_real( + mean = vecObs, + var = self.cov_data, + number = self.ne, + return_chol = True + ) + + return vecObs, enObs def _ext_scaling(self): # get vector of scaling self.state_scaling = at.calc_scaling( - self.prior_state, self.list_states, self.prior_info) - - delta_scaled_prior = self.state_scaling[:, None] * \ - np.dot(at.aug_state(self.prior_state, self.list_states), 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 - # zero. This part is a good place to include eventual additional truncation. - energy = 0 - trunc_index = len(s_d) - 1 # inititallize - for c, elem in enumerate(s_d): - energy += elem - if energy / sum(s_d) >= self.trunc_energy: - trunc_index = c # take the index where all energy is preserved - break - u_d, s_d, v_d = u_d[:, :trunc_index + - 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 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(self.state) - 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(self.state) - 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(self.state) - 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. + self.prior_enX, self.idX, self.prior_info) + + self.Am = None - !!! 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(self.state) - np.savez('temp_state_ml', self.temp_state) def compress_manager(self, data=None, vintage=0, aug_coeff=None): """ @@ -753,41 +672,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)) - - aug_state = at.aug_state(self.current_state, self.list_states) - self.update() + 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, + ) + + # 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 @@ -819,6 +761,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)] diff --git a/pipt/misc_tools/analysis_tools.py b/pipt/misc_tools/analysis_tools.py index de0140a..b30fd5d 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 @@ -655,10 +662,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) @@ -940,6 +948,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: @@ -1226,7 +1235,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 @@ -1248,7 +1257,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]) * @@ -1257,7 +1266,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) @@ -1474,91 +1483,63 @@ def subsample_state(index, aug_state, pert_state): return new_state -def init_local_analysis(init, state): - """Initialize local analysis. +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 + ] + +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. - Initialize the local analysis by reading the input variables, defining the parameter classes and search ranges. Build - the map of data/parameter positions. + r : int, optional + Rank to truncate the SVD to. If None, energy must be specified. - Args - ---- - init : dictionary containing the parsed information form the input file. - state : list of states that will be updated + 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 ------- - 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] + 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 < 1: + r = np.sum((np.cumsum(S) / sum(S)) <= energy) + else: + r = np.sum((np.cumsum(S) / sum(S)) <= energy/100) 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 + 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/misc_tools/ensemble_tools.py b/pipt/misc_tools/ensemble_tools.py new file mode 100644 index 0000000..25e409c --- /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(indecies, 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 5bfbde0..4cb2326 100644 --- a/pipt/misc_tools/extract_tools.py +++ b/pipt/misc_tools/extract_tools.py @@ -1,6 +1,8 @@ -# This module inlcudes fucntions for extracting information from input dicts +# This module includes functions for extracting information from input dicts + __all__ = [ 'extract_prior_info', + 'extract_initial_controls', 'extract_multilevel_info', 'extract_local_analysis_info', 'extract_maxiter', @@ -9,13 +11,17 @@ ] # Imports -import numpy as np +import numpy as np +import pandas as pd import pickle import os 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: ''' @@ -122,39 +128,194 @@ 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 such that we only have one level -- the high fidelity one ''' - ml_info = keys['multilevel'] - if isinstance(ml_info, list): - ml_info = list_to_dict(ml_info) - assert isinstance(ml_info, dict) + keys_ml = keys + if isinstance(keys, list): + 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' - error_comp_scheme = 'none' - ml_corr_done = True - 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/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!') diff --git a/pipt/update_schemes/enrml.py b/pipt/update_schemes/enrml.py index 8eb3dea..1c149de 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 @@ -58,13 +60,30 @@ 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) - - # Extract parameters like conv. tol. and damping param. from ITERATION keyword in DATAASSIM - self._ext_iter_param() - - # Within variables + self.prior_enX = cp.deepcopy(self.enX) # not sure if this is wise! + + # 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 + + # Load ACTNUM if given if 'actnum' in self.keys_da.keys(): try: self.actnum = np.load(self.keys_da['actnum'])['actnum'] @@ -72,36 +91,45 @@ 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 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) + 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() - self.current_state = cp.deepcopy(self.state) + + 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) @@ -109,35 +137,33 @@ 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}') + # Log initial data misfit + self.log_update(success=True, prior_run=True) 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() # run ordinary analysis + # Perform the 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))) + 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): """ @@ -153,8 +179,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 @@ -166,7 +198,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) @@ -187,13 +219,16 @@ def check_convergence(self): if self.data_misfit >= self.prev_data_misfit: success = False - 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( 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 @@ -214,13 +249,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) @@ -229,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 @@ -246,28 +287,27 @@ 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 + def log_update(self, success, prior_run=False): + ''' + Log the update results in a formatted table. + ''' + 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: + log_data["Reduction (%)"] = 'N/A' + + self.logger(**log_data) + - if self.trunc_energy > 1: # ensure that it is given as percentage - self.trunc_energy /= 100. class lmenrml_approx(lmenrmlMixIn, approx_update): @@ -298,7 +338,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 +360,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 +368,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/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/esmda.py b/pipt/update_schemes/esmda.py index 1e39929..52bdb29 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,19 +46,20 @@ 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() 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) @@ -66,12 +68,14 @@ 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() - self.current_state = deepcopy(self.state) + # Extract the inflation parameter from MDA keyword self.alpha = self._ext_inflation_param() @@ -102,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) @@ -119,49 +133,50 @@ def calc_analysis(self): self.data_misfit_std = np.std(data_misfit) self.ensemble_misfit = 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.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() 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)) - - aug_state = at.aug_state(self.current_state, self.list_states) - - self.update() + # Perform the 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))) + 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): """ @@ -180,12 +195,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) @@ -194,19 +212,41 @@ 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. - 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) return False, True, why_stop + def log_update(self, success=None, prior_run=False): + ''' + Log the update results in a formatted table. + ''' + 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: + log_data["Reduction (%)"] = 'N/A' + + self.logger(**log_data) + def _ext_inflation_param(self): r""" Extract the data covariance inflation parameter from the MDA keyword in DATAASSIM part. Also, we check that diff --git a/pipt/update_schemes/multilevel.py b/pipt/update_schemes/multilevel.py index 4b4b610..cb8a493 100644 --- a/pipt/update_schemes/multilevel.py +++ b/pipt/update_schemes/multilevel.py @@ -3,48 +3,37 @@ 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'] + 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 +44,77 @@ 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.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): """ 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) - - 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): + 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 - # 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_conv, + 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 +123,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.vecObs, + 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 +177,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.pred_data], + 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_conv, + np.concatenate(enPred,axis=1), + self.cov_data + ) self.data_misfit = np.mean(data_misfit) self.data_misfit_std = np.std(data_misfit) @@ -1089,106 +200,15 @@ 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) - return False, True, why_stop \ No newline at end of file + return False, True, why_stop diff --git a/pipt/update_schemes/update_methods_ns/approx_update.py b/pipt/update_schemes/update_methods_ns/approx_update.py index 1ec2194..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 @@ -16,128 +19,143 @@ 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) - - # 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() + 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 + Ud, Sd, VTd = at.truncSVD(enYcentered, energy=self.trunc_energy) + + # 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)) - 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)) + 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) + + # 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) + 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]] + reg_term = (self.lam + 1)*np.eye(Sd.size) + np.diag(Sd**2) + X = VTd.T @ np.diag(Sd) @ solve(reg_term, Ud.T) - #f = self.keys_da['localization'] + + # Check for adaptive localization + if 'autoadaloc' in loc_info: - if 'autoadaloc' in self.localization.loc_info: - - # 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)))) + # 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(enX - np.mean(enX, 1)[:,None], self.state_scaling) 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)) + enXcentered = self.scale(np.dot(enX, self.proj), self.state_scaling) + + # 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, enRes), + curr_param = self.list_states, + prior_info = self.prior_info + ) + + + # Check for local analysis + elif ('localanalysis' in loc_info) and (loc_info['localanalysis']): + + # Calculate weights + if 'distance' in loc_info: + weight = _calc_loc( + max_dist = loc_info['range'], + distance = loc_info['distance'], + prior_info = self.prior_info[self.list_states[0]], + 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(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 (residuals) + enRes = self.scale(enE - enY, self.scale_data) + + # Compute the update step with local analysis try: - self.step = weight.multiply( - np.dot(pert_state, X)).dot(scaled_delta_data) + self.step = weight.multiply(np.dot(enXcentered, X)).dot(enRes) except: - self.step = (weight*(np.dot(pert_state, X))).dot(scaled_delta_data) + self.step = (weight*(np.dot(enXcentered, X))).dot(enRes) + + # 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)) + # 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(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 (residuals) + enRes = self.scale(enE - enY, self.scale_data) - self.step = local_mask.multiply( - np.dot(pert_state, X)).dot(scaled_delta_data) + # Compute the update step with distance-based localization + self.step = mask.multiply(np.dot(enXcentered, X)).dot(enRes) + + + # 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 = [] @@ -145,60 +163,75 @@ def update(self): 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: - # 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 '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)) - 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)) - 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_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 + + if ('emp_cov' in self.keys_da) and (self.keys_da['emp_cov'] == 'yes'): + + # 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) + + Sinv = np.diag(1/Sd) + X0 = Sinv @ Ud.T @ enEcentered + eigval, eigvec = np.linalg.eig(X0 @ X0.T) + + # Calculate and scale difference between observations and predictions (residuals) + enRes = self.scale(enE - enY, self.scale_data) + + # 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) + 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)) - 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) + enXcentered = self.scale(np.dot(enX, self.proj), self.state_scaling) + 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): + """ + 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..406081c 100644 --- a/pipt/update_schemes/update_methods_ns/full_update.py +++ b/pipt/update_schemes/update_methods_ns/full_update.py @@ -18,14 +18,59 @@ 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 = 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)) + 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 +86,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 diff --git a/pipt/update_schemes/update_methods_ns/hybrid_update.py b/pipt/update_schemes/update_methods_ns/hybrid_update.py index 265a928..e7d5bc5 100644 --- a/pipt/update_schemes/update_methods_ns/hybrid_update.py +++ b/pipt/update_schemes/update_methods_ns/hybrid_update.py @@ -16,50 +16,72 @@ class hybrid_update: 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])) + 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 - 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() + Returns: + np.ndarray: scaled data perturbations + """ - # x_1 = np.dot(u_d.T, solve(self.scale_data[l], - # (self.real_obs_data[l] - self.aug_pred_data[l]))) + 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. - 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)) + 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)) + # 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) + + 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(pert_state[l].shape) for l in range(self.tot_level)] + 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. - # 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])) + # 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: - 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)]) + 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): - 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]))) + enRes = self.scale(enE[l] - enY[l], self.scale_data[l]) + self.step[l][row, :] = np.dot(self.state_scaling[row, None] * kg, enRes) + + diff --git a/pipt/update_schemes/update_methods_ns/subspace_update.py b/pipt/update_schemes/update_methods_ns/subspace_update.py index 0f3280b..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 @@ -18,58 +15,50 @@ 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 - # 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)) + # Compute scaled misfit (residual between predicted and observed data) + enRes = self.scale(enY - enE, self.scale_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() + # Truncate SVD of S + Us, Ss, VsT = at.truncSVD(S, energy=self.trunc_energy) + Sinv = np.diag(1/Ss) - 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)) + # 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 - X2 = np.dot(u, np.dot(ps_inv.T, z)) - X3 = np.dot(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) + - # 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))) + def scale(self, data, scaling): + """ + Scale the data perturbations by the data error standard deviation. - step_d = np.dot(X3, solve(np.eye(len(Lam)) + (1+self.lam) * - np.diag(Lam), np.dot(X2.T, scaled_misfit))) + Args: + data (np.ndarray): data perturbations + scaling (np.ndarray): data error standard deviation - # 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)) + Returns: + np.ndarray: scaled data perturbations + """ + + if len(scaling.shape) == 1: + return (scaling ** (-1))[:, None] * data + else: + return solve(scaling, data) 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']) diff --git a/popt/loop/ensemble_base.py b/popt/loop/ensemble_base.py index b3dd3fc..f83f352 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'] @@ -49,50 +51,96 @@ def __init__(self, options, simulator, 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 + # Initialize state-related attributes + self.stateX = np.array([]) # Current state vector, (nx,) + self.stateF = None # Function value(s) of current state + 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 (transformed) for state vector, (nx,) + self.ub = np.array([]) # Upper bounds (transformed) for state vector, (nx,) - def _initialize_state_info(self): - ''' - Initialize covariance and bounds based on prior information. - ''' - self.cov = np.array([]) - self.lb = [] - self.ub = [] - self.bounds = [] - + # 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 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.stateX = self.scale_state(self.stateX) - # Make cov full covariance matrix - self.cov = np.diag(self.cov) + 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() + + # check for ensmble + 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) + 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() + + # Evaluate the objective function + 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.stateF = func_values + else: + self.enF = func_values + + return func_values def get_state(self): """ @@ -101,7 +149,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): """ @@ -110,13 +158,7 @@ def get_cov(self): 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. - """ - return ot.update_optim_state(x, self.state, list(self.state.keys())) + return self.covX def get_bounds(self): """ @@ -127,58 +169,97 @@ def get_bounds(self): """ return self.bounds - - def function(self, x, *args, **kwargs): + + def scale_state(self, x): """ - This is the main function called during optimization. + Transform the internal state from [lb, ub] to [0, 1] Parameters ---------- - x : ndarray - Control vector, shape (number of controls, number of perturbations) + x : array_like + The input state Returns ------- - obj_func_values : numpy.ndarray - Objective function values, shape (number of perturbations, ) + x : array_like + The scaled state """ - self._aux_input() + 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 - # check for ensmble - if len(x.shape) == 1: self.ne = self.num_models - else: self.ne = x.shape[1] + def invert_scale_state(self, u): + """ + Transform the internal state from [0, 1] to [lb, ub] - # convert x (nparray) to state (dict) - self.state = self.vec_to_state(x) + Parameters + ---------- + u : array_like + The scaled state - # 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) # toggle back after calc_prediction + Returns + ------- + x : array_like + The unscaled state + """ + u = np.asarray(u) + x = np.zeros_like(u) - # Evaluate the objective function - if run_success: - func_values = self.obj_func( - self.pred_data, - input_dict=self.sim.input_dict, - true_order=self.sim.true_order, - state=self.state, # pass state for possible use in objective function - **kwargs - ) - else: - func_values = np.inf # the simulations have crashed + if self.transform is False: + return u - self._scale_state() # scale back to [0, 1] - if len(x.shape) == 1: self.state_func_values = func_values - else: self.ens_func_values = func_values + 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 + + def save_stateX(self, path='./', filetype='npz'): + ''' + Save the state vector. + + Parameters + ---------- + path : str + Path to save the state vector. Default is current directory. - return func_values + 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) - 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 + else: + return x def _aux_input(self): @@ -213,4 +294,4 @@ def _invert_scale_state(self): 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 + np.clip(self.state[key], self.lb[i], self.ub[i], out=self.state[key]) diff --git a/popt/loop/ensemble_gaussian.py b/popt/loop/ensemble_gaussian.py index 2d334ca..499fea4 100644 --- a/popt/loop/ensemble_gaussian.py +++ b/popt/loop/ensemble_gaussian.py @@ -1,51 +1,36 @@ # 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'] class GaussianEnsemble(EnsembleOptimizationBaseClass): """ - Class to store control states and evaluate objective functions. + Gaussian Ensemble class for ensemble-based optimization. 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 - """ 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 @@ -54,219 +39,172 @@ 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) """ # 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 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 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): - 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) - + Covarice matrix, shape (number of controls, number of controls) + Returns ------- - gradient : numpy.ndarray - The gradient evaluated at x, shape (number of controls, ) - """ - - # 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 covariance equal to the input - self.cov = args[0] + gradient : ndarray + Ensemble gradient, shape (number of controls, ) + ''' + # Update state vector + self.stateX = x - # 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 + # Set covariance equal to the input + self.covX = args[0] - # 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 = 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) else: - gradient = level_gradient[0] - - return gradient - - def hessian(self, x=None, *args): - r""" - Calculate the hessian matrix associated with ensemble, defined as: + grad = grad_ml[0] - $$ H = J(XX^T - \Sigma)/ (N_e-1) $$ + # 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) - where $X$ and $J$ are ensemble matrices of $x$ (or control variables) and objective function - perturbed by their respective means. + return grad - !!! 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) - - start_index += ml_ne - level_hessian.append(g_c / (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): - hessian += level_hessian[l]*cov_wgt[l] - hessian /= self.ne + # 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)) + + # 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 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): r""" Calculate weights used in sequential monte carlo optimization. + Updated version that accommodates new base class changes. Parameters ---------- @@ -281,54 +219,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]) + + # Evaluate objective function for ensemble + self.enF = self.function(self.enX, **kwargs) - state_ens = at.aug_state(self.state, list(self.state.keys())) - self.function(state_ens, **kwargs) + if not isinstance(self.enF, list): + self.enF = [self.enF] - if not isinstance(self.ens_func_values, list): - self.ens_func_values = [self.ens_func_values] - L = len(self.ens_func_values) + 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 @@ -338,29 +275,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 @@ -374,56 +311,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 - - 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 - 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) diff --git a/popt/loop/optimize.py b/popt/loop/optimize.py index 989c22d..24b5e82 100644 --- a/popt/loop/optimize.py +++ b/popt/loop/optimize.py @@ -1,24 +1,13 @@ # External imports import os import numpy as np -import logging import time import pickle from abc import ABC, abstractmethod # 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 : %(levelname)s : %(name)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) +from ensemble.logger import PetLogger class Optimize(ABC): @@ -73,8 +62,8 @@ def __init__(self, **options): options : dict Optimization options """ - # Set the logger - self.logger = logger + # Setup logger + self.logger = PetLogger('optim.log') # Save name for (potential) pickle dump/load self.pickle_restart_file = 'popt_restart_dump' @@ -104,7 +93,9 @@ def __init__(self, **options): # Initialize variables (set in subclasses) self.options = None + self.mean_state = None self.obj_func_values = None + self.obj_func_tol = None # objective tolerance limit # Initialize number of function and jacobi evaluations self.nfev = 0 @@ -156,8 +147,10 @@ def run_loop(self): epf_not_converged = True previous_state = None if self.epf: - previous_state = self.xk - logger.info(f' -----> EPF-EnOpt: {self.epf_iteration}, {self.epf["r"]} (outer iteration, penalty factor)') # print epf info + previous_state = self.mean_state + 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 @@ -175,53 +168,50 @@ 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', - 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('') + 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"]}') + self.logger(f'Total number of jacobi evaluations = {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(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 - 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.xk) / (np.abs(previous_state) + 1.0e-9) + 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.xk + previous_state = self.mean_state self.epf['r'] *= self.epf['r_factor'] # increase penalty factor - self.ftol *= self.epf['tol_factor'] # decrease tolerance - self.obj_func_values = self.fun(self.xk, epf = self.epf) + self.obj_func_tol *= self.epf['tol_factor'] # decrease tolerance + self.obj_func_values = self.fun(self.mean_state, **self.epf) self.iteration = 0 - info_str = ' {:<10} {:<10} {:<15} {:<15} {:<15} '.format('iter', 'alpha_iter', - 'obj_func', 'step-size', 'cov[0,0]') - self.logger.info(info_str) - self.logger.info(' {:<21} {:<15.4e}'.format(self.iteration, np.mean(self.obj_func_values))) 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 + self.logger(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(np.mean(self.fun(self.xk))),4)) - logger.info(f' -----> EPF-EnOpt: objective value without penalty = {final_obj_no_penalty}') # 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(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. @@ -241,7 +231,6 @@ def load(self): # Save in 'self' self.__dict__.update(tmp_load) - @abstractmethod def calc_update(self): """ This is an empty dummy function. Actual functionality must be defined by the subclasses. diff --git a/popt/misc_tools/optim_tools.py b/popt/misc_tools/optim_tools.py index a303430..5a8bdc3 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 @@ -349,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: diff --git a/popt/update_schemes/linesearch.py b/popt/update_schemes/linesearch.py index 43d71c1..9c9fe3f 100644 --- a/popt/update_schemes/linesearch.py +++ b/popt/update_schemes/linesearch.py @@ -199,6 +199,9 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c else: self.callback = None + # Remove 'datatype' from 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 @@ -255,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 @@ -265,22 +268,24 @@ def __init__(self, fun, x, jac, method='GD', hess=None, args=(), bounds=None, c self.p_old = None # Initial results - self.optimize_result = ot.get_optimize_result(self) + self.optimize_result = self.get_intermediate_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('\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 \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), + 'step-size': self.step_size + }) - self.run_loop() + self.run_loop() def fun(self, x, *args, **kwargs): self.nfev += 1 - x = ot.clip_state(x, self.bounds) # ensure bounds are respected + x = ot.clip_state(x, self.bounds) # ensure bounds are respected if self.args is None: f = np.mean(self.function(x, epf=self.epf)) else: @@ -353,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.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: @@ -368,7 +373,6 @@ 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.............') if self.lskwargs['method'] == 0: ls_res = line_search_backtracking( step_size=step_size, @@ -422,22 +426,24 @@ 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 # Save Results - self.optimize_result = ot.get_optimize_result(self) + self.optimize_result = self.get_intermediate_results() if self.saveit: ot.save_optimize_results(self.optimize_result) # 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): @@ -459,7 +465,7 @@ def calc_update(self, iter_resamp=0): # 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 @@ -472,7 +478,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 @@ -534,7 +540,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 diff --git a/popt/update_schemes/subroutines/subroutines.py b/popt/update_schemes/subroutines/subroutines.py index ddd1105..17402f2 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." @@ -95,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 @@ -110,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) @@ -122,40 +125,51 @@ 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 # 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 -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) @@ -163,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) @@ -179,11 +193,15 @@ 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) # 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 @@ -191,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 @@ -213,6 +234,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 +303,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 +330,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 +338,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 @@ -347,6 +383,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: 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"])