#######################################
Using low-level forecast/analysis steps
#######################################

.. admonition:: Work in progress

    This section should be expanded. Give a complete snippet ?

AugKF
=====
We present here how to use the forecast and analysis steps of the AugKF algorithm.

Forecast/Analyse a CoreState
----------------------------
Principle
^^^^^^^^^
An algorithm is implemented through a `AugkfAlgo`_ object. It contains all the informations needed for the algorithm (configuration, priors, correlation matrices, drift_matrices). It uses a method *parallel_forecast_step* to forecast CoreStates that is, under the hood, delegated to a `AugkfForecaster`_ that is an object created at the initialisation of `AugkfAlgo`_.

Similarly, the `AugkfAlgo`_ also implements *analysis_step* that is delegated to an `AugkfAnalyser`_ object, that also stores the observations needed for the *analysis_step*.

The easiest way to set up these objects is to use the an `create_augkf_algo`_ function from a Config that will take care of initialising all objects and data needed:

.. code-block:: python

    from pygeodyn.inout.config import ComputationConfig
    from pygeodyn.augkf.algo import create_augkf_algo

    config = ComputationConfig(0,'augkf.conf') # Or whatever config path
    nb_realisations = 10
    nb_times = 3
    seed = 1
    attributed_models = np.array(list(range(nb_realisations)))
    algo = create_augkf_algo(config, nb_realisations, seed, attributed_models)

Then, we can create CoreState objects to perform the forecast/analysis on them.

AR1 Example
^^^^^^^^^^^

Both *parallel_forecast_step* and *analysis_step* take a 2D CoreState as input (dim: nb_reals x nb_coefs). 

The following code shows an example of a forecast step followed by an analysis step using an AR1 process:

.. code-block:: python
    
    # Build a complete CoreState with 10 realisations, 3 times and the number of coefficients given in the config.

    # You can use your own corestate initialization or use the algo.init_corestate() methode provided
    
    ###################################################
    #PROVIDED#

    cs, _, _, _, Z_AR3 = algo.init_corestates()

    ###################################################
    #OWN#

    cs = CoreState()
    corestate_measures={}
    corestate_measures.update({'MF': self.config.Nb,
                                'U': self.config.Nu2,
                                'SV': self.config.Nsv,
                                'ER': self.config.Nsv,
                                'Z': self.config.Nz})
    
    # Do whatever operation to initialise the CoreState at t=0 (ex: set all realisations of B according to a magnetic model)
    cs.B[:, 0] = my_awesome_magnetic_model
    # Z_AR3 is None if AR1
    Z_AR3 = None

    ###################################################
    #check if observation is found

    algo.analyser.check_if_analysis_data(0)

    ###################################################
    #FORECAST

    # forecast to t=1
    cs[:, 1] = algo.forecaster.parallel_forecast_step(cs[:,0], seed, 1)

    ##################################################
    #ANALYSIS

    # Now do an analysis step on the forecast CoreState at t=1
    # if AR1: requires 1 observation date at t=1 for analysis 
    analysis_result = algo.analyser.analysis_step(cs[:, 1])
        
    # The result is a CoreState containing all analysed realisations that we can again use to update our CoreState (stored here at t=2)
    cs[:, 2] = analysis_result

    # Then we can access the data as previously:
    cs.B[:, 0, 0] # Initial value of g10 for all realisations
    cs.B[:, 1, 0] # Forecast value of g10 for all realisations
    cs.U[0, 2] # Analysed values of all core flow coefs for the realisation 0
    # and so on...


AR3 Example
^^^^^^^^^^^


*parallel_forecast_step* takes a 2D CoreState as input (dim: nb_reals x nb_coefs). 

*analysis_step* takes a 3D CoreState as input (dim: nb_reals x 2*dt_a_f ratio + 1 x nb_coefs).

The following code shows an example of a forecast step followed by an analysis step using an AR3 process:

.. code-block:: python
    
    # Build a complete CoreState with 10 realisations, 3 times and the number of coefficients given in the config.

    # You can use your own corestate initialization or use the algo.init_corestate() methode provided
    
    ###################################################
    #PROVIDED#

    cs, _, _, _, Z_AR3 = algo.init_corestates()

    ###################################################
    #OWN#

    cs = CoreState()
    corestate_measures={}
    corestate_measures.update({'MF': self.config.Nb,
                                'U': self.config.Nu2,
                                'SV': self.config.Nsv,
                                'ER': self.config.Nsv,
                                'Z': self.config.Nz,
                                'dUdt': self.config.Nu2,
                                'd2Udt2': self.config.Nu2,
                                'dERdt': self.config.Nsv,
                                'd2ERdt2': self.config.Nsv})
    
    # Do whatever operation to initialise the CoreState at t=0 (ex: set all realisations of B according to a magnetic model)
    cs.B[:, 0] = my_awesome_magnetic_model
    # init Z_AR3 at t=0 (Z_AR3 stores states at 3 last times. dimension: nb_reals x 3 x nb_coefs)
    Z_AR3 = my_init_Z_AR3
    
    ###################################################
    #check if observation is found

    algo.analyser.check_if_analysis_data(0)

    ###################################################
    #FORECAST

    # forecast to t=2
    cs[:, 1], Z_AR3 = algo.forecaster.parallel_forecast_step(cs[:, 0], Z_AR3, seed, 1)
    cs[:, 2], Z_AR3 = algo.forecaster.parallel_forecast_step(cs[:, 1], Z_AR3, seed, 1)
    
    ##################################################
    #ANALYSIS

    # Now do an analysis step on the forecast CoreState at t=1
    # if AR3: requires 3 observation dates at t=0,1,2 for analysis
    analysis_result, Z_AR3 = algo.analyser.analysis_step(cs, Z_AR3)
   
    # The result is a CoreState containing all analysed realisations that we can again use to update our CoreState (stored here at t=2)
    cs[:, 2] = analysis_result

    # Then we can access the data as previously:
    cs.B[:, 0, 0] # Initial value of g10 for all realisations
    cs.B[:, 1, 0] # Forecast value of g10 for all realisations
    cs.U[0, 2] # Analysed values of all core flow coefs for the realisation 0
    # and so on...


More information
----------------
See the API reference of the Augkf objects:

* `AugkfAlgo`_
* `AugkfForecaster`_
* `AugkfAnalyser`_

Common operations are implemented in the motherclass `Computer <https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.generic.computer.html>`_ from which AugkfForecaster and AugkfAnalyser inherit. The version of the algorithm with a PCA is implemented in PCA-specific objects (`AugkfAlgoPCA`_, `AugkfForecasterPCA`_, `AugkfAnalyserPCA`_).

.. _create_augkf_algo: https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.augkf.algo.html#pygeodyn.augkf.algo.create_augkf_algo
.. _AugkfAlgo: https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.augkf.algo.html#pygeodyn.augkf.algo.AugkfAlgo
.. _AugkfForecaster: https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.augkf.forecaster.html#pygeodyn.augkf.forecaster.AugkfForecaster
.. _AugkfAnalyser: https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.augkf.analyser.html#pygeodyn.augkf.analyser.AugkfAnalyser
.. _AugkfAlgoPCA: https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.augkf.algo.html#pygeodyn.augkf.algo.AugkfAlgoPCA
.. _AugkfForecasterPCA: https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.augkf.forecaster.html#pygeodyn.augkf.forecaster.AugkfForecasterPCA
.. _AugkfAnalyserPCA: https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.augkf.analyser.html#pygeodyn.augkf.analyser.AugkfAnalyserPCA