Beispiel #1
0
    def forecast(self):
        """
        Forecast step of the filter.
        Use the model object to propagate each ensemble member to the end of the given checkpoints to
        produce and ensemble of forecasts.
        Filter configurations dictionary is updated with the new results.
        If the prior is assumed to be a Gaussian, we are all set, otherwise we have to estimate it's
        parameters based on the provided forecast ensemble (e.g. in the case of 'prior_distribution' = GMM ).

        Args:
            None

        Returns:
            None

        """
        # generate the forecast states
        analysis_ensemble = self.filter_configs['analysis_ensemble']
        timespan = self.filter_configs['timespan']
        forecast_ensemble = utility.propagate_ensemble(
            ensemble=analysis_ensemble,
            model=self.filter_configs['model'],
            checkpoints=timespan,
            in_place=False)
        self.filter_configs['forecast_ensemble'] = forecast_ensemble
        forecast_state = utility.ensemble_mean(forecast_ensemble)
        self.filter_configs['forecast_state'] = forecast_state.copy()

        # Add more functionality after building the forecast ensemble if desired!
        # Obtain information about the prior distribution from the forecast ensemble
        # Covariance localization and hybridization are carried out if requested
        self.generate_prior_info()
Beispiel #2
0
def get_model_info(timespan,
                   ensemble_size=25,
                   model=None,
                   load_ensemble=False,
                   ignore_existing_repo=False,
                   repo_file=None):
    """
    """
    timespan = np.array(timespan).flatten()
    # print("Requesting timespan: ", timespan)
    # Check the experiment repository
    if repo_file is None:
        file_name = 'Coupled_Lorenz_Repository.pickle'
        file_name = os.path.abspath(
            os.path.join(os.path.dirname(__file__), file_name))
    else:
        file_name = os.path.abspath(repo_file)

    # Check repo if needed, and found:
    if ignore_existing_repo:
        ref_IC_np = ref_trajectory_np = init_ensemble_np = observations_np = None
        #
    elif os.path.isfile(file_name):
        print("Loading model repo from : %s" % file_name)
        cont = pickle.load(open(file_name, 'rb'))
        tspan = cont['tspan']

        # Get the initial condition anyways
        ref_IC_np = cont['ref_trajectory'][:, 0].copy()

        # Verify timespan, and get reference trajectory
        if tspan.size < timespan.size:
            ref_trajectory_np = None
            t_indexes = None
            save_info = True
            #
        elif tspan.size > timespan.size:
            # Look into the reference trajectory to find matches, if existing
            t_indexes = []
            for i, t in enumerate(timespan):
                loc = np.where(np.isclose(tspan, t, rtol=1e-12))[0]
                if loc.size > 0:
                    loc = loc[0]
                    t_indexes.append(loc)
                else:
                    t_indexes = None
                    break
            #
            if t_indexes is not None:
                t_indexes.sort()
                t_indexes = np.asarray(t_indexes)
                ref_trajectory_np = cont['ref_trajectory'][:, t_indexes]
            else:
                ref_trajectory_np = None
                #
        else:
            if np.isclose(tspan, timespan).all():
                ref_trajectory_np = cont['ref_trajectory']
                t_indexes = np.arange(timespan.size)
                #

        # Observations:
        if t_indexes is None or ref_trajectory_np is None:  # TODO: this is overkill!
            observations_np = None
        else:
            observations_np = cont['observations'][:, t_indexes]

        # Read the ensemble
        if load_ensemble:
            init_ensemble_np = cont['init_ensemble']
            if np.size(init_ensemble_np, 1) <= ensemble_size:
                init_ensemble_np = init_ensemble_np[:, :ensemble_size]
            else:
                init_ensemble_np = None
        else:
            init_ensemble_np = None

    else:
        print(
            "The coupled-lorenz model repo is not found. This should be the first time to run this function/script."
        )
        ref_IC_np = ref_trajectory_np = init_ensemble_np = observations_np = None

    # Check the integrity of the loaded data:
    if ref_trajectory_np is None or (load_ensemble and init_ensemble_np is None
                                     ) or observations_np is None:
        #
        # Generate stuff, and save them
        if model is None:
            print(
                "YOU MUST pass a MODEL object because the repo is not found in %s, or some information is missing!"
                % file_name)
            raise ValueError
        else:
            print("Got the model >> Creating experiemnt info...")

        # Reference IC
        if ref_IC_np is None:
            ref_IC = model._reference_initial_condition.copy()
            ref_IC_np = ref_IC.get_numpy_array()

        # Reference Trajectory (Truth)
        if ref_trajectory_np is None:
            if ref_IC_np is not None:
                ref_IC = model.state_vector(ref_IC_np.copy())
            else:
                ref_IC = model._reference_initial_condition.copy()
            # generate observations from coupled lorenz:
            ref_trajectory = model.integrate_state(ref_IC, timespan)
            ref_trajectory_np = utility.ensemble_to_np_array(ref_trajectory,
                                                             state_as_col=True)

        # Initial (forecast) Ensmble:
        if load_ensemble and init_ensemble_np is None:
            prior_noise_model = model.background_error_model
            prior_noise_sample = [
                prior_noise_model.generate_noise_vec()
                for i in xrange(ensemble_size)
            ]
            noise_mean = utility.ensemble_mean(prior_noise_sample)
            ic = prior_noise_model.generate_noise_vec().add(ref_IC)
            for i in xrange(ensemble_size):
                prior_noise_sample[i].axpy(-1.0, noise_mean)
                prior_noise_sample[i].add(ic, in_place=True)
            init_ensemble = prior_noise_sample
            init_ensemble_np = utility.ensemble_to_np_array(init_ensemble,
                                                            state_as_col=True)

        if observations_np is None:
            # Create observations' and assimilation checkpoints:
            if ref_trajectory is None:
                if ref_trajectory_np is None:
                    ref_trajectory = model.integrate_state(ref_IC, timespan)
                else:
                    ref_trajectory = [
                        model.state_vectory(ref_trajectory_np[i, :].copy())
                        for i in xrange(np.size(ref_trajectory, 1))
                    ]
            observations = [
                model.evaluate_theoretical_observation(x)
                for x in ref_trajectory
            ]
            # Perturb observations:
            obs_noise_model = model.observation_error_model
            for obs_ind in xrange(len(timespan)):
                observations[obs_ind].add(obs_noise_model.generate_noise_vec(),
                                          in_place=True)
            observations_np = utility.ensemble_to_np_array(observations,
                                                           state_as_col=True)

        print("Experiment repo created; saving for later use...")
        # save results for later use
        out_dict = dict(tspan=timespan,
                        ref_trajectory=ref_trajectory_np,
                        init_ensemble=init_ensemble_np,
                        observations=observations_np)
        pickle.dump(out_dict, open(file_name, 'wb'))
        print("...done...")
        #
    else:
        print("All information properly loaded from file")
        pass

    return timespan, ref_IC_np, ref_trajectory_np, init_ensemble_np, observations_np
Beispiel #3
0
    def filtering_cycle(self, update_reference=False):
        """
        Carry out a single assimilation cycle. This is just the analysis in this case
        All required variables are obtained from 'filter_configs' dictionary.
        
        Args:
            update_reference (default True): bool,
                A flag to decide whether to update the reference state in the filter or not.
                This should be removed!
                  
           
        """
        #
        model = self.model

        try:
            observation_time = self.filter_configs['observation_time']
        except (NameError, KeyError, AttributeError):
            observation_time = None

        try:
            forecast_time = self.filter_configs['forecast_time']
        except (NameError, KeyError, AttributeError):
            forecast_time = None

        try:
            analysis_time = self.filter_configs['analysis_time']
        except (NameError, KeyError, AttributeError):
            analysis_time = None

        try:
            reference_time = self.filter_configs['reference_time']
        except (NameError, KeyError, AttributeError):
            reference_time = None

        #
        try:
            # Retrieve the reference state and evaluate initial root-mean-squared error
            reference_state = self.filter_configs['reference_state'].copy()
        except (NameError, KeyError, AttributeError):
            reference_state = None
            #
            if self._verbose:
                print("Couldn't retrieve the reference state/time! ")
                print(
                    "True observations must be present in this case, and RMSE can't be evaluated!"
                )

        if reference_state is not None and reference_time is not None:
            # If the reference time is below window final limit, march it forward:
            if analysis_time is not None:
                new_time = analysis_time
            elif forecast_time is not None:
                new_time = forecast_time
            elif reference_time is not None:
                new_time = reference_time

            if (new_time - reference_time) > self._time_eps:
                # propagate forward the reference state, and update it
                local_trajectory = model.integrate_state(
                    reference_state, [reference_time, new_time])
                reference_time = new_time
                if isinstance(local_trajectory, list):
                    reference_state = local_trajectory[-1]
                else:
                    reference_state = local_trajectory

            elif (reference_time - new_time) > self._time_eps:
                print(
                    "Can't have the reference time after the forecast/analysis time!"
                )
                raise ValueError

            if update_reference:
                self.filter_configs.update(
                    {'reference_state': reference_state})
                self.filter_configs.update({'reference_time': reference_time})

        # Check time instances:
        times_array = np.array(
            [observation_time, analysis_time, forecast_time, reference_time])
        val_inds = np.where(times_array)[0]
        if len(val_inds) > 0:
            if not np.allclose(times_array[val_inds],
                               times_array[val_inds[0]],
                               atol=self._time_eps,
                               rtol=self._time_eps):
                print(
                    "Given time instances, e.g. forecast time, observation time, etc. MUST match! "
                )
                print(
                    "[observation_time, analysis_time, forecast_time, reference_time]",
                    [
                        observation_time, analysis_time, forecast_time,
                        reference_time
                    ])
                raise AssertionError
            else:
                # Good to go!
                pass
        else:
            # No times are given; just proceed
            pass

        #
        try:
            forecast_state = self.filter_configs['forecast_state'].copy()
        except (NameError, KeyError, AttributeError):
            print(
                "Couldn't retrieve the forecast state! This can't work; Aborting!"
            )
            raise AssertionError
        #
        try:
            observation = self.filter_configs['observation']
        except (NameError, KeyError, AttributeError):
            observation = None
        finally:
            if observation is None:
                print("The observation is not found!")
                raise AssertionError

        #
        # Forecast state/ensemble should be given:
        if forecast_state is None:
            try:
                forecast_state = utility.ensemble_mean(
                    self.filter_configs['forecast_ensemble'])
            except:
                print(
                    "Couldn't find either forecast state or forecast ensemble while analysis should be done first!"
                )
                raise AssertionError

        #
        # calculate the initial/forecast RMSE: this is the RMSE before assimilation at the first entry of the time-span

        state_size = model.state_size()

        #
        if forecast_state is not None and reference_state is not None:
            f_rmse = initial_rmse = utility.calculate_rmse(
                forecast_state, reference_state, state_size)
        else:
            f_rmse = initial_rmse = 0

        #
        # Start the filtering process: preprocessing -> filtering(forecast->+<-anslsysis) -> postprocessing
        if self.filter_configs['apply_preprocessing']:
            self.cycle_preprocessing()

        #
        # Analysis step (This calls the filter's analysis step.)
        # > --------------------------------------------||
        sep = "\n" + "*" * 80 + "\n"
        if self._verbose:
            print("%s ANALYSIS STEP %s" % (sep, sep))
        self.analysis()
        # > --------------------------------------------||
        #

        analysis_state = self.filter_configs['analysis_state'].copy()

        # Analysis RMSE:
        if reference_state is not None:
            a_rmse = utility.calculate_rmse(reference_state, analysis_state,
                                            state_size)
        else:
            a_rmse = 0

        #
        if True:
            print("Initial (f_rmse, a_rmse) :", (f_rmse, a_rmse))
            print("analysis_state", analysis_state)
            print("forecast_state", forecast_state)
            print("reference_state", reference_state)

        #
        # Apply post-processing if required
        if self.filter_configs['apply_postprocessing']:
            self.cycle_postprocessing()

        #
        # Update filter statistics (including RMSE)
        if 'filter_statistics' not in self.output_configs:
            self.output_configs.update(
                dict(filter_statistics=dict(initial_rmse=None,
                                            forecast_rmse=None,
                                            analysis_rmse=None)))
        else:
            if 'analysis_rmse' not in self.output_configs['filter_statistics']:
                self.output_configs['filter_statistics'].update(
                    dict(analysis_rmse=None))
            if 'forecast_rmse' not in self.output_configs['filter_statistics']:
                self.output_configs['filter_statistics'].update(
                    dict(forecast_rmse=None))
            if 'initial_rmse' not in self.output_configs['filter_statistics']:
                self.output_configs['filter_statistics'].update(
                    dict(initial_rmse=None))

        # now update the RMSE's
        self.output_configs['filter_statistics'].update(
            {'initial_rmse': initial_rmse})
        self.output_configs['filter_statistics'].update(
            {'forecast_rmse': f_rmse})
        #
        self.output_configs['filter_statistics'].update(
            {'analysis_rmse': a_rmse})

        # output and save results if requested
        if self.output_configs['scr_output']:
            self.print_cycle_results()
        if self.output_configs['file_output']:
            self.save_cycle_results()
Beispiel #4
0
    def analysis(self):
        """
        Analysis step:
        """

        #
        model = self.filter_configs['model']
        state_size = model.state_size()
        #

        # Check if the forecast ensemble should be inflated;
        f = self.filter_configs['inflation_factor']
        forecast_ensemble = utility.inflate_ensemble(
            self.filter_configs['forecast_ensemble'], f, in_place=False)
        forecast_state = utility.ensemble_mean(forecast_ensemble)
        ensemble_size = len(forecast_ensemble)

        # get the measurements vector
        observation = self.filter_configs['observation']
        observation_size = model.observation_vector_size()

        #
        normalized_weights = None
        self.filter_configs.update({'normalized_weights': normalized_weights})
        # Evaluate the likelihood for all particles
        particles_likelihood = np.asarray(
            self.evaluate_likelihood(model_states=forecast_ensemble,
                                     observation=observation))
        likelihood_sum = particles_likelihood.sum()
        if likelihood_sum == 0:  # This should not happen as long as all particles are passed to the likelihood
            # print particles_likelihood
            print(
                'All particles received negligible weights rounded down to zero!. Filter diverged; Ensemble is not '
                'updated!')
            return None
        else:
            normalized_weights = particles_likelihood / likelihood_sum

        # Monitors the effective particle number (effective sample size)
        ess = 1.0 / (np.asarray(normalized_weights)**2).sum()
        self.filter_configs['effective_sample_size'] = ess
        self.filter_configs.update({
            'particles_likelihood': particles_likelihood,
            'normalized_weights': normalized_weights
        })

        # Check the resampling flag and resampling scheme
        # Resampling step:
        if self._resample:
            resampling_scheme = self.filter_configs['resampling_scheme']
            normalize_weights = self.filter_configs['normalize_weights']
            resampling_indexes = self.resample_states(
                weights=normalized_weights,
                ensemble_size=ensemble_size,
                strategy=resampling_scheme,
                normalize_weights=normalize_weights)
            # Generate Analysis ensemble by weighted resampling
            print("Resampling...")
            if self.filter_configs['analysis_ensemble'] is None:
                analysis_ensemble = []
                for ens_ind in xrange(ensemble_size):
                    res_ind = resampling_indexes[ens_ind]
                    # print("Ens ind: %d ;  resampling from index %d" %(ens_ind, res_ind))
                    analysis_ensemble.append(forecast_ensemble[res_ind].copy())
            else:
                analysis_ensemble = self.filter_configs['analysis_ensemble']
                for ens_ind in xrange(ensemble_size):
                    res_ind = resampling_indexes[ens_ind]
                    # print("Ens ind: %d ;  resampling from index %d" %(ens_ind, res_ind))
                    analysis_ensemble[ens_ind] = forecast_ensemble[
                        res_ind].copy()
            #
        else:
            resampling_indexes = np.arange(ensemble_size)
            analysis_ensemble = [member.copy() for member in forecast_ensemble]

        self.filter_configs['resampling_indexes'] = resampling_indexes
        self.filter_configs['analysis_ensemble'] = analysis_ensemble
        self.filter_configs['analysis_state'] = utility.ensemble_mean(
            analysis_ensemble)
Beispiel #5
0
    def filtering_cycle(self, update_reference=True):
        """
        Carry out a single assimilation cycle. Forecast followed by analysis or the other way around.
        All required variables are obtained from 'filter_configs' dictionary.
        This base method is designed to work in both ensemble, and standard framework. 
        You can override it for sure in your filter.
        
        Args:
            update_reference (default True): bool,
                A flag to decide whether to update the reference state in the filter or not.
                  
        Returns:
            None
            
        """
        model = self.filter_configs['model']
        # print('reference_state', reference_state)
        timespan = self.filter_configs['timespan']
        observation_time = self.filter_configs['observation_time']
        forecast_time = self.filter_configs['forecast_time']
        analysis_time = self.filter_configs['analysis_time']
        forecast_first = self.filter_configs['forecast_first']
        apply_preprocessing = self.filter_configs['apply_preprocessing']
        apply_postprocessing = self.filter_configs['apply_postprocessing']

        # calculate the initial RMSE: this is the RMSE before assimilation at the first entry of the time-span
        try:
            reference_time = self.filter_configs['reference_time']
        except (KeyError, ValueError, AttributeError, NameError):
            raise ValueError(
                "Couldn't find the reference time in the configurations dictionary"
            )
        else:
            if reference_time is not None:
                if reference_time != timespan[0]:
                    raise ValueError(
                        "Reference time does not match the initial time of the reference time-span!"
                    )
        #
        #
        if forecast_first:
            # Analysis ensemble should be given first in this case
            try:
                initial_state = self.filter_configs['analysis_state']
            except (KeyError, ValueError, AttributeError, NameError):
                try:
                    initial_state = utility.ensemble_mean(
                        self.filter_configs['analysis_ensemble'])
                except:
                    raise ValueError(
                        "Couldn't find either analysis state or analysis ensemble while forecast should be done first!"
                    )
            finally:
                if initial_state is None:
                    try:
                        initial_state = utility.ensemble_mean(
                            self.filter_configs['analysis_ensemble'])
                    except:
                        raise ValueError(
                            "Couldn't find either analysis state or analysis ensemble while forecast should be done first!"
                        )
        else:
            # Forecast ensemble should be given first in this case
            try:
                initial_state = self.filter_configs['forecast_state']
            except:
                try:
                    initial_state = utility.ensemble_mean(
                        self.filter_configs['forecast_ensemble'])
                except:
                    raise ValueError(
                        "Couldn't find either forecast state or forecast ensemble while analysis should be done first!"
                    )
            if initial_state is None:
                try:
                    initial_state = utility.ensemble_mean(
                        self.filter_configs['forecast_ensemble'])
                except:
                    raise ValueError(
                        "Couldn't find either forecast state or forecast ensemble while analysis should be done first!"
                    )

        # Retrieve the reference state and evaluate initial root-mean-squared error
        reference_state = self.filter_configs[
            'reference_state']  # always given at the initial time of the timespan
        if reference_state is not None:
            reference_time = self.filter_configs['reference_time']

        try:
            initial_rmse = utility.calculate_rmse(initial_state,
                                                  reference_state)
        except:
            reference_state = None
            reference_time = None
            print("Failed to calculate Initial RMSE!")
            initial_rmse = 0.0
        #
        # Start the filtering process: preprocessing -> filtering(forecast->+<-anslsysis) -> postprocessing
        if apply_preprocessing:
            self.cycle_preprocessing()

        if not forecast_first and analysis_time != min(
                forecast_first, analysis_time, observation_time):
            # this is a double check!
            raise ValueError(
                "While ANALYSIS should be done first, confusion occurred with times given!\n"
                "\tCycle timespan:%s"
                "\tObservation time: %f\n"
                "\tForecast time: %f\n"
                "\tAnalysis time: %f" % (repr(timespan), observation_time,
                                         forecast_time, analysis_time))
        elif not forecast_first:
            #
            # print("\n\n\n\n\n\n ANALYSIS FIRTS \n\n\n\n\n")
            # analysis should be carried out first in this case
            state_size = self.filter_configs['model'].state_size()
            if reference_state is not None:
                try:
                    analysis_rmse = utility.calculate_rmse(
                        reference_state, analysis_state, state_size)
                except:
                    print("Failed to calculate Analysis RMSE!")
                    analysis_rmse = 0.0
            analysis_time = self.filter_configs['analysis_time']

            # Analysis step
            self.analysis()
            #

            # update the reference state
            if reference_time is not None:
                if reference_time != timespan[
                        0] or reference_time != analysis_time:
                    raise ValueError(
                        "Reference time does not match the initial time of the reference time-span!"
                    )
            local_checkpoints = [analysis_time, forecast_time]

            if reference_state is not None:
                tmp_trajectory = model.integrate_state(
                    initial_state=reference_state,
                    checkpoints=local_checkpoints)
                if isinstance(tmp_trajectory, list):
                    reference_state = tmp_trajectory[-1].copy()
                else:
                    reference_state = tmp_trajectory.copy()
                reference_time = local_checkpoints[-1]

            # forecast now:
            self.forecast()

            try:
                forecast_rmse = utility.calculate_rmse(reference_state,
                                                       forecast_state,
                                                       state_size)
            except:
                print("Failed to calculate Forecast RMSE!")
                forecast_rmse = 0.0
            # update the reference state Moved to the process
            if update_reference and reference_state is not None:
                self.filter_configs['reference_state'] = reference_state.copy()
                self.filter_configs['reference_time'] = reference_time

        else:
            # forecast should be done first
            # print("\n\n\n\n\n\n FORECAST FIRTS \n\n\n\n\n")
            state_size = self.filter_configs['model'].state_size()
            if reference_time is not None:
                if reference_time != timespan[0]:
                    raise ValueError(
                        "Reference time does not match the initial time of the reference time-span!"
                    )
                local_checkpoints = [reference_time, timespan[-1]]

            # Forecast step:
            self.forecast()
            #
            try:
                forecast_state = self.filter_configs['forecast_state']
            except (NameError, AttributeError):
                raise NameError("forecast_state must be updated by the filter "
                                "and added to 'self.filter_configs'!")

            if reference_state is not None:
                tmp_trajectory = model.integrate_state(
                    initial_state=reference_state, checkpoints=timespan)
                if isinstance(tmp_trajectory, list):
                    up_reference_state = tmp_trajectory[-1].copy()
                else:
                    up_reference_state = tmp_trajectory.copy()
                reference_time = local_checkpoints[-1]

                # Update the reference state
                if update_reference:
                    self.filter_configs[
                        'reference_state'] = up_reference_state.copy()
                    self.filter_configs['reference_time'] = reference_time

            # observation = self.filter_configs['observation']
            #
            try:
                forecast_rmse = utility.calculate_rmse(up_reference_state,
                                                       forecast_state,
                                                       state_size)
            except:
                print("Failed to calculate Forecast RMSE!")
                forecast_rmse = 0.0
            # Analysis step:
            self.analysis()
            #
            try:
                analysis_state = self.filter_configs['analysis_state']
            except (NameError, AttributeError):
                raise NameError("analysis_state must be updated by the filter "
                                "and added to 'self.filter_configs'!")

            try:
                analysis_rmse = utility.calculate_rmse(up_reference_state,
                                                       analysis_state,
                                                       state_size)
            except:
                print("Failed to calculate Analysis RMSE!")
                analysis_rmse = 0.0
            #

        # Apply post-processing if required
        if apply_postprocessing:
            self.cycle_postprocessing()

        # Update filter statistics (including RMSE)
        if 'filter_statistics' not in self.output_configs:
            self.output_configs.update(
                dict(filter_statistics=dict(initial_rmse=None,
                                            forecast_rmse=None,
                                            analysis_rmse=None)))
        else:
            if 'analysis_rmse' not in self.output_configs['filter_statistics']:
                self.output_configs['filter_statistics'].update(
                    dict(analysis_rmse=None))
            if 'forecast_rmse' not in self.output_configs['filter_statistics']:
                self.output_configs['filter_statistics'].update(
                    dict(forecast_rmse=None))
            if 'initial_rmse' not in self.output_configs['filter_statistics']:
                self.output_configs['filter_statistics'].update(
                    dict(initial_rmse=None))

        # now update the RMSE's
        self.output_configs['filter_statistics']['initial_rmse'] = initial_rmse
        self.output_configs['filter_statistics'][
            'forecast_rmse'] = forecast_rmse
        self.output_configs['filter_statistics'][
            'analysis_rmse'] = analysis_rmse

        # output and save results if requested
        if self.output_configs['scr_output']:
            self.print_cycle_results()
        if self.output_configs['file_output']:
            self.save_cycle_results()
Beispiel #6
0
    def smoothing_cycle(self, update_reference=True):
        """
        Carry out a single assimilation cycle. The analysis has to be carried out first given the forecast 
        at the assimilation time, and the observations (Real observations), or synthetic created using 
        the reference state of the model.
        All required variables are obtained from 'smoother_configs' dictionary.
        This base method is designed to work in both ensemble, and variational framework. 
        You can override it for sure in your smoother.
        
        Args:
            update_reference (default True): bool,
                A flag to decide whether to update the reference state in the smoother or not.
                  
        """
        #
        # > ===========================( Assimilation Cycle Starts )================================ <
        #
        model = self.model
        #
        try:
            # Retrieve the reference state and evaluate initial root-mean-squared error
            reference_state = self.smoother_configs['reference_state'].copy()
            reference_time = self.smoother_configs['reference_time']
        except (NameError, KeyError, AttributeError):
            reference_state = reference_time = None
        finally:
            if reference_state is None or reference_time is None:
                reference_state = reference_time = None
                if self._verbose:
                    print("Couldn't retrieve the reference state/time! ")
                    print(
                        "True observations must be present in this case, and RMSE can't be evaluated!"
                    )

        #
        try:
            forecast_state = self.smoother_configs['forecast_state'].copy()
        except (NameError, KeyError, AttributeError):
            print(
                "Couldn't retrieve the forecast state! This can't work; Aborting!"
            )
            raise AssertionError
        #
        try:
            observations_list = self.smoother_configs['observations_list']
        except (NameError, KeyError, AttributeError):
            observations_list = None
        finally:
            if observations_list is None:
                print(
                    "Neither the reference state is passed, nor true observations list is found!"
                )
                raise AssertionError
        #
        try:
            obs_checkpoints = self.smoother_configs['obs_checkpoints']
        except:
            print(
                "Either analysis_time or obs_checkpoints or both is/are missing! Aborting!"
            )

        try:
            analysis_timespan = self.smoother_configs['analysis_timespan']
        except (NameError, KeyError, AttributeError):
            analysis_timespan = None
        finally:
            if analysis_timespan is None:
                obs_checkpoints = self.smoother_configs['obs_checkpoints']
                bounds = self.smoother_configs['window_bounds']
                t0, t1 = bounds[0], bounds[-1]
                tf, ta = self.smoother_configs[
                    'forecast_time'], self.smoother_configs['analysis_time']
                analysis_timespan = np.asarray(
                    list(set.union(set(obs_checkpoints), {t0, t1, tf, ta})))
                analysis_timespan.sort()
                self.smoother_configs.update(
                    {'analysis_timespan': analysis_timespan})

        #
        # Forecast state/ensemble should be given:
        if forecast_state is None:
            try:
                forecast_state = utility.ensemble_mean(
                    self.smoother_configs['forecast_ensemble'])
            except:
                print(
                    "Couldn't find either forecast state or forecast ensemble while analysis should be done first!"
                )
                raise AssertionError

        # Check forecast, and analysis times:
        forecast_time = self.smoother_configs['forecast_time']
        analysis_time = self.smoother_configs['analysis_time']
        #
        if forecast_time is not None and analysis_time is not None:
            pass
            #
        elif forecast_time is None and analysis_time is None:
            print(
                "Both forecast and analysis time are not set. At least specify the forecast time!"
            )
            raise ValueError
            #
        elif analysis_time is None:
            if analysis_timespan is not None:
                analysis_time = analysis_timespan[0]
            else:
                # Analysis is carried out in the same place where forecast state is provided!
                analysis_time = forecast_time
        else:
            print(
                "You have to specify where the provided forecast state is calculated!"
            )
            raise ValueError

        #
        if analysis_time > forecast_time:
            print(
                "Analysis time > forecast time. Forward propagation of the forecast state is taking place now."
            )
            #
            local_checkpoints = [forecast_time, analysis_time]
            tmp_trajectory = model.integrate_state(
                initial_state=forecast_state, checkpoints=local_checkpoints)
            if isinstance(tmp_trajectory, list):
                forecast_state = tmp_trajectory[-1].copy()
            else:
                forecast_state = tmp_trajectory.copy()
            forecast_time = analysis_time
            self.smoother_configs['forecast_time'] = local_checkpoints[-1]
            self.smoother_configs['forecast_state'] = forecast_state
            #
        elif analysis_time < forecast_time:
            print(
                "Analysis time < Forecast time! Backward propagation of the state is Needed!"
            )
            raise AssertionError
            #
        else:
            # All is Good; analysis_time = forecast_time
            pass

        #
        if forecast_time is not None and reference_time is not None:
            #
            if reference_time < forecast_time:
                local_checkpoints = [reference_time, forecast_time]
                tmp_trajectory = model.integrate_state(
                    initial_state=reference_state,
                    checkpoints=local_checkpoints)
                if isinstance(tmp_trajectory, list):
                    reference_state = tmp_trajectory[-1].copy()
                else:
                    reference_state = tmp_trajectory.copy()
                    #
                self.smoother_configs['reference_state'] = reference_state
                self.smoother_configs['reference_time'] = local_checkpoints[-1]
                #
            elif reference_time == forecast_time:
                pass
            else:
                print("Time settings are not conformable!")
                print("Reference time:", reference_time)
                print("Forecast time:", forecast_time)
                print("Analysis time:", analysis_time)
                raise AssertionError
                #

        #
        # calculate the initial RMSE: this is the RMSE before assimilation at the first entry of the time-span
        if reference_time is not None:
            # print("analysis_timespan", analysis_timespan)
            # print("reference_time", reference_time)
            if (analysis_timespan[0] - reference_time) > self._time_eps:
                print(
                    "Reference time < initial time of the assimilation time-span. Forward propagation of the reference is taking place now."
                )
                print("reference_time", reference_time)
                print("analysis_timespan[0]", analysis_timespan[0])
                #
                local_checkpoints = [reference_time, analysis_timespan[0]]
                tmp_trajectory = model.integrate_state(
                    initial_state=reference_state,
                    checkpoints=local_checkpoints)
                if isinstance(tmp_trajectory, list):
                    reference_state = tmp_trajectory[-1].copy()
                else:
                    reference_state = tmp_trajectory.copy()
                reference_time = local_checkpoints[-1]
                reference_time = analysis_timespan[0]
                self.smoother_configs['reference_time'] = reference_time
                self.smoother_configs['reference_state'] = reference_state
                #
            elif (reference_time - analysis_timespan[0]) > self._time_eps:
                print(
                    "Reference time > the initial time of the assimilation time-span!"
                )
                raise ValueError
            else:
                # All good...
                pass
        else:
            self.smoother_configs['reference_time'] = None
            self.smoother_configs['reference_state'] = None

        #
        if abs(analysis_time - analysis_timespan[0]) > self._time_eps:
            print(
                "*" * 100 +
                "\nWARNING: The analysis time is not at the beginning of the assimilation timespan!\n"
                + "*" * 100)
            print("analysis_time", analysis_time)
            print("analysis_timespan[0]", analysis_timespan[0])
            print("\n" + "*" * 100 + "\n")

        #
        if (analysis_time - obs_checkpoints[0]) > self._time_eps:
            print(
                "First observation has to be at, or after, the analysis time")
            raise AssertionError
        if len(obs_checkpoints) != len(observations_list):
            print(
                "Observaions timespan must of be same length as the observations list!"
            )
            raise AssertionError

        #
        if forecast_state is not None and reference_state is not None:
            if abs(analysis_time - reference_time) > self._time_eps:
                print(
                    "While evaluating initial RMSE; analysis_time != reference_time"
                )
                print("reference_time: %f" % reference_time)
                print("Analysis time: %f" % analysis_time)
                raise ValueError
            else:
                initial_rmse = utility.calculate_rmse(forecast_state,
                                                      reference_state)
        else:
            initial_rmse = 0

        #
        # Start the smoothing process: preprocessing -> smoothing(forecast->+<-anslsysis) -> postprocessing
        if self.smoother_configs['apply_preprocessing']:
            self.cycle_preprocessing()

        #
        # Analysis step (This calls the smoother's analysis step.)
        # > --------------------------------------------||
        sep = "\n" + "*" * 80 + "\n"
        if self._verbose:
            print("%s ANALYSIS STEP %s" % (sep, sep))
        self.analysis()
        # > --------------------------------------------||
        #

        # Evaluate RMSEs: This is potential redundancy; just read analysis trajectory from the smoother)
        reference_state = self.smoother_configs['reference_state'].copy()
        reference_time = self.smoother_configs['reference_time']
        #
        forecast_state = self.smoother_configs['forecast_state'].copy()
        forecast_time = self.smoother_configs['forecast_time']
        #
        analysis_state = self.smoother_configs['analysis_state'].copy()
        analysis_time = self.smoother_configs['analysis_time']
        #
        analysis_trajectory = []
        analysis_trajectory.append(analysis_state)
        #

        state_size = model.state_size()

        # Forecast and analysis RMSE (at the beginning of the assimilation window)
        if abs(analysis_time - reference_time) > self._time_eps:
            print(
                "While evaluating initial RMSE; analysis_time != reference_time"
            )
            print("reference_time: %f" % reference_time)
            print("Analysis time: %f" % analysis_time)
            raise ValueError
        else:
            f_rmse = utility.calculate_rmse(reference_state, forecast_state,
                                            state_size)
            forecast_rmse_list = [f_rmse]
            forecast_times_list = [forecast_time]
            a_rmse = utility.calculate_rmse(reference_state, analysis_state,
                                            state_size)
            analysis_rmse_list = [a_rmse]
            analysis_times_list = [analysis_time]

        #
        if self._verbose:
            print("Initial (f_rmse, a_rmse) :", (f_rmse, a_rmse))
            print("analysis_state", analysis_state)
            print("forecast_state", forecast_state)
            print("reference_state", reference_state)

        #
        #
        for t0, t1 in zip(analysis_timespan[:-1], analysis_timespan[1:]):
            local_checkpoints = np.array([t0, t1])
            #
            # Propagate forecast
            local_trajectory = model.integrate_state(forecast_state,
                                                     local_checkpoints)
            if isinstance(local_trajectory, list):
                forecast_state = local_trajectory[-1].copy()
            else:
                forecast_state = local_trajectory.copy()
            # Propagate analysis
            local_trajectory = model.integrate_state(analysis_state,
                                                     local_checkpoints)
            if isinstance(local_trajectory, list):
                analysis_state = local_trajectory[-1].copy()
            else:
                analysis_state = local_trajectory.copy()
            analysis_trajectory.append(analysis_state)
            #

            #
            if reference_state is not None:
                if self._verbose:
                    print("\n Evaluating RMSE, inside smoothers_base...:")
                    print("Subinterval: ", [t0, t1])
                    print("reference_time: ", reference_time)
                    print("reference_state: ", reference_state)

                # Propagate reference
                local_trajectory = model.integrate_state(
                    reference_state, [reference_time, t1])
                reference_time = t1
                if isinstance(local_trajectory, list):
                    reference_state = local_trajectory[-1].copy()
                else:
                    reference_state = local_trajectory.copy()
                    #
                #
                f_rmse = utility.calculate_rmse(reference_state,
                                                forecast_state, state_size)
                a_rmse = utility.calculate_rmse(reference_state,
                                                analysis_state, state_size)
                forecast_rmse_list.append(f_rmse)
                analysis_rmse_list.append(a_rmse)
                forecast_times_list.append(t1)
                analysis_times_list.append(t1)
                #

                if self._verbose:
                    print("forecast_state: ", forecast_state)
                    print("analysis_state: ", analysis_state)

            else:
                # No reference state is provided. RMSE will be an empty list
                if self._verbose:
                    print(
                        "WARNING: No reference state is provided. RMSE will be an empty list!"
                    )

        #
        if reference_state is not None and update_reference:
            # If the reference time is below window final limit, march it forward:
            if (self.smoother_configs['window_bounds'][-1] -
                    reference_time) > self._time_eps:
                tf = self.smoother_configs['window_bounds'][-1]
                local_trajectory = model.integrate_state(
                    reference_state, [reference_time, tf])
                reference_time = tf
                if isinstance(local_trajectory, list):
                    reference_state = local_trajectory[-1]
                else:
                    reference_state = local_trajectory
                #
            self.smoother_configs.update({'reference_state': reference_state})
            self.smoother_configs.update({'reference_time': reference_time})
            #

        #
        # Apply post-processing if required
        if self.smoother_configs['apply_postprocessing']:
            self.cycle_postprocessing()

        #
        # forecast now  (This calls the smoother's forecast step):
        # > --------------------------------------------||
        if self._verbose:
            print("%s ANALYSIS STEP %s" % (sep, sep))
        self.forecast()
        # > --------------------------------------------||

        #
        # Update smoother statistics (including RMSE)
        if 'smoother_statistics' not in self.output_configs:
            self.output_configs.update(
                dict(smoother_statistics=dict(initial_rmse=None,
                                              forecast_rmse=None,
                                              analysis_rmse=None)))
        else:
            if 'analysis_rmse' not in self.output_configs[
                    'smoother_statistics']:
                self.output_configs['smoother_statistics'].update(
                    dict(analysis_rmse=None))
            if 'forecast_rmse' not in self.output_configs[
                    'smoother_statistics']:
                self.output_configs['smoother_statistics'].update(
                    dict(forecast_rmse=None))
            if 'initial_rmse' not in self.output_configs[
                    'smoother_statistics']:
                self.output_configs['smoother_statistics'].update(
                    dict(initial_rmse=None))

        # now update the RMSE's
        self.output_configs['smoother_statistics'].update(
            {'initial_rmse': initial_rmse})
        self.output_configs['smoother_statistics'].update(
            {'forecast_rmse': forecast_rmse_list})
        self.output_configs['smoother_statistics'].update(
            {'forecast_times': forecast_times_list})
        #
        self.output_configs['smoother_statistics'].update(
            {'analysis_rmse': analysis_rmse_list})
        self.output_configs['smoother_statistics'].update(
            {'analysis_times': analysis_times_list})

        # output and save results if requested
        if self.output_configs['scr_output']:
            self.print_cycle_results()
        if self.output_configs['file_output']:
            self.save_cycle_results()
Beispiel #7
0
    def EnKF_update(self, forecast_ensemble, observation, inflation_factor=1.0, in_place=False):
        """
        Calculate the kalman updated matrix, i.e. the update matrix of ensemble anomalies
            and calculate the (moving) analysis ensemble
        Args:
            state:

        Returns:
            analysis_ensemble: a list of model.state_vector objects
            X5: the state anomalies update matrix

        """
        model = self.smoother_configs['model']
        ensemble_size = len(forecast_ensemble)
        forecast_mean = utility.ensemble_mean(forecast_ensemble)
        forecast_anomalies = []

        if inflation_factor is None:
            inflation_factor = 1.0
        #
        for state in forecast_ensemble:
            if inflation_factor != 1:
                forecast_anomalies.append(state.axpy(-1, forecast_mean, in_place=False).scale(inflation_factor))
            else:
                forecast_anomalies.append(state.axpy(-1, forecast_mean, in_place=False))
        A_prime = utility.ensemble_to_np_array(forecast_anomalies)

        observation_perturbations = [model.evaluate_theoretical_observation(state) for state in forecast_anomalies]
        S = utility.ensemble_to_np_array(observation_perturbations)
        C = S.dot(S.T)

        try:
            C += (ensemble_size - 1) * model.observation_error_model.R[...]
        except:
            C += (ensemble_size - 1) * model.observation_error_model.R.get_numpy_array()

        C_eigs, Z = linalg.eigh(C)
        C_eigs = 1.0 / C_eigs

        obs_innov = observation.axpy(-1, model.evaluate_theoretical_observation(forecast_mean))
        y = Z.T.dot(obs_innov.get_numpy_array())
        y *= C_eigs
        y = Z.dot(y)
        y = S.T.dot(y)

        # Update ensemble mean
        analysis_mean = model.state_vector(A_prime.dot(y))
        analysis_mean = analysis_mean.add(forecast_mean)
        xa = analysis_mean.get_numpy_array()

        inv_sqrt_Ceig = sparse.spdiags(np.sqrt(C_eigs), 0, C_eigs.size, C_eigs.size)
        X2 = inv_sqrt_Ceig.dot(Z.T.dot(S))
        _, Sig2, V2 = linalg.svd(X2, full_matrices=False)

        # Calculate the square root of the matrix I-Sig2'*Sig2*theta
        Sqrtmat = sparse.spdiags(np.sqrt(1.0 - np.power(Sig2,2)), 0, Sig2.size, Sig2.size)
        Aa_prime = A_prime.dot(V2.T.dot(Sqrtmat.dot(V2)))
        
        if self._verbose:
            # For debugging
            print(A_prime, type(A_prime))
        
        # Update analysis ensemble:
        analysis_ensemble = [model.state_vector(xa+A_prime[:, j]) for j in xrange(np.size(A_prime, 1)) ]
        
        # Build the update kernel X5
        xx = V2.T.dot(Sqrtmat.dot(V2))

        for j in xrange(np.size(xx, 1)):
            xx[:, j] += y[:]
        IN = np.ones((ensemble_size, ensemble_size), dtype=np.float) / ensemble_size
        X5  = IN + inflation_factor * np.dot( (sparse.spdiags(np.ones(ensemble_size), 0, ensemble_size, ensemble_size) - IN), xx)

        return analysis_ensemble, X5
Beispiel #8
0
    def analysis(self):
        """
        Analysis step of the (Vanilla Ensemble Kalman Smoother (EnKS).
        In this case, the given forecast ensemble is propagated to the observation time instances to create model
        observations H(xk) that will be used in the assimilation process...

        """
        model = self.model

        # Timsespan info:
        forecast_time = self.smoother_configs['forecast_time']
        analysis_time = self.smoother_configs['analysis_time']
        if self.smoother_configs['obs_checkpoints'] is None:
            print("Couldn't find observation checkpoints in self.smoother_configs['obs_checkpoints']; None found!")
            raise ValueError
        else:
            obs_checkpoints = np.asarray(self.smoother_configs['obs_checkpoints'])

        if (analysis_time - obs_checkpoints[0] ) >= self._model_step_size:
            print("Observations MUST follow the assimilation times in this implementation!")
            raise ValueError

        if (analysis_time - obs_checkpoints[0] ) >= self._model_step_size:
            print("Observations MUST follow the assimilation times in this implementation!")
            raise ValueError

        # Get observations list:
        observations_list = self.smoother_configs['observations_list']
        #
        assim_flags = [True] * len(obs_checkpoints)
        if (forecast_time - obs_checkpoints[0]) >= self._model_step_size:
            print("forecast time can't be after the first observation time instance!")
            print("forecast_time", forecast_time)
            print("obs_checkpoints[0]", obs_checkpoints[0])
            raise AssertionError
            #
        elif abs(forecast_time - obs_checkpoints[0]) <= self._time_eps:
            # an observation exists at the analysis time
            pass
            #
        else:
            obs_checkpoints = np.insert(obs_checkpoints, 0, forecast_time)
            assim_flags.insert(0, False)

        # state and observation vector dimensions
        forecast_ensemble = self.smoother_configs['forecast_ensemble']
        if forecast_ensemble is None:
            print("Can't carry out the analysis step without a forecast ensemble; Found None!")
            raise ValueError

        if abs(forecast_time-analysis_time) > self._time_eps:
            print("At this point, the analysis and forecast times have to be the SAME!")
            print("Forecast Time=%f, Analysis time=%f" % (forecast_time, analysis_time))
            raise ValueError
        #
        elif (analysis_time-obs_checkpoints[0]) > self._time_eps:
            print("At this point, the analysis time should be equal to or less that the time of the first observation")
            print("Analysis time=%f Time of the first observation=%f" %(analysis_time, obs_checkpoints[0]))
            raise ValueError
        else:
            if len(observations_list) == 0:
                print("An empty list of observations is detected! No observations to assimilate;")
                print("Nothing to assimilate...")
                print("Setting the analysis state to the forecast state...")

                self.smoother_configs['analysis_state'] = self.smoother_configs['forecast_state'].copy()
                return

                #
            elif len(observations_list) == 1:
                if abs(obs_checkpoints[0]-forecast_time)<=self._time_eps and \
                    abs(analysis_time-forecast_time)<=self._time_eps:
                    print("A single observation is detected, and the assimilation time coincides with observation time;")
                    print("You are advised to use EnKF!")  # Todo, redirect automatically to 3D-Var
                    raise ValueError
                elif (obs_checkpoints[0]-forecast_time)>self._time_eps or (obs_checkpoints[0]-analysis_time)>self._time_eps:
                    pass
                else:
                    print("This settings instances are not acceptable:")
                    print("Observation time instance: %f" % obs_checkpoints[0])
                    print("Forecast time: %f" % forecast_time)
                    print("Observation time instance: %f" % analysis_time)
                    print("Terminating ")
                    raise ValueError
            else:
                # Good to go!
                pass

        #
        # > --------------------------------------------||
        # START the Smoothing  process:
        # > --------------------------------------------||
        # Forward propagation:
        # 1- Apply EnKF at each observation timepoint,
        # 2- Save Kalman gain matrix
        #
        cwd = os.getcwd()
        saved_kernels = []
        
        # Copy the initial forecast ensemble,
        moving_ensemble = [state.copy() for state in forecast_ensemble]
        ensemble_size = len(moving_ensemble)
        num_obs_points = len(observations_list)
        if self._verbose:
            print("Started the Analysis step of EnKS; with [num_obs_points] observation/assimilation time points;")
        inflation_factor = self.smoother_configs['inflation_factor']
        #
        # Forward to observation time instances and update observation term
        obs_ind = 0
        for iter_ind, t0, t1 in zip(xrange(num_obs_points), obs_checkpoints[: -1], obs_checkpoints[1: ]):
            local_ckeckpoints = np.array([t0, t1])

            # For initial subwindow only, check the flag on both and final assimilation flags
            if iter_ind == 0:
                assim_flag = assim_flags[iter_ind]
                if assim_flag:
                    observation = observations_list[obs_ind]
                    obs_ind += 1
                    print("TODO: Applying EnKF step at the initial time of the whole assimilation window")
                    moving_ensemble, X5 = self.EnKF_update(moving_ensemble, observation, inflation_factor, in_place=True)
                    saved_kernels = self._save_EnKF_update_kernel(X5, saved_kernels)
                # Copy initial ensemble; could be saved to file instead
                analysis_ensemble_np = utility.ensemble_to_np_array(moving_ensemble)

            # Now, let's do forecast/analysis step for each subwindow:
            for ens_ind in xrange(ensemble_size):
                state = moving_ensemble[ens_ind]
                tmp_trajectory = model.integrate_state(initial_state=state, checkpoints=local_ckeckpoints)
                if isinstance(tmp_trajectory, list):
                    moving_ensemble[ens_ind] = tmp_trajectory[-1].copy()
                else:
                    moving_ensemble[ens_ind] = tmp_trajectory.copy()
            #
            assim_flag = assim_flags[iter_ind+1]
            if assim_flag:
                observation = observations_list[obs_ind]
                obs_ind += 1
                moving_ensemble, X5 = self.EnKF_update(moving_ensemble, observation, inflation_factor, in_place=True)
            saved_kernels = self._save_EnKF_update_kernel(X5, saved_kernels)

        #
        # Backward propagation:
        for iter_ind in xrange(num_obs_points-2, -1, -1):
            assim_flag = assim_flags[iter_ind]
            try:
                next_assim_ind = iter_ind + 1 + assim_flags[iter_ind+1: ].index(True)
            except ValueError:
                # This is the latest assimilation cycle; proceed
                continue

            if assim_flag:
                kernel_file = saved_kernels[iter_ind]
                X5 = np.load(kernel_file)

                next_kernel = saved_kernels[next_assim_ind]
                X6 = np.load(next_kernel)

                X5 = X5.dot(X6)
                del X6

                # save/overwrite the updated kernel; no need to keep both versions!
                np.save(kernel_file, X5)

        # Update analysis ensemble at the initial time of the window:
        X5 = np.load(saved_kernels[0])
        analysis_ensemble_np = analysis_ensemble_np.dot(X5)
        del X5

        analysis_ensemble = moving_ensemble
        for ens_ind in xrange(ensemble_size):
            analysis_ensemble[ens_ind][:] = analysis_ensemble_np[:, ens_ind].copy()  # Need to copy?!

        self.smoother_configs.update({'analysis_ensemble': analysis_ensemble})
        self.smoother_configs.update({'analysis_state': utility.ensemble_mean(analysis_ensemble)})
Beispiel #9
0
                                  num_Ys=num_Ys,
                                  F=F,
                                  observation_size=observation_size,
                                  ensemble_size=ensemble_size,
                                  observation_opertor_type=observation_opertor_type,
                                  observation_noise_level=observation_noise_level,
                                  background_noise_level=background_noise_level,
                                  experiment_tspan=experiment_tspan
                                 )
        coupled_ref_trajectory, coupled_init_ensemble, coupled_observations = coupled_model_info
        ref_trajectory, init_ensemble, observations = model_info
        IC = ref_trajectory[0]


        # Forecast/Free-Run Trajectory:
        forecast_trajectory = model.integrate_state(utility.ensemble_mean(init_ensemble), checkpoints)
        forecast_trajectory_fine = model.integrate_state(utility.ensemble_mean(init_ensemble), np.arange(checkpoints[0], checkpoints[-1], 0.005))

        #
        print("Model Info Generated; Started Plotting...")
        enhance_plotter()
        if not os.path.isdir(plots_dir):
            os.makedirs(plots_dir)
        #

        # Coupled Lorenz:
        coupled_state_size = len(coupled_ref_trajectory[1])
        trajectory = np.asarray(coupled_ref_trajectory)
        fig = plt.figure(facecolor='white')
        ax = fig.add_subplot(111)
        extent = [checkpoints[0], checkpoints[-1], 1, coupled_state_size]