####################
The CoreState object
####################

A `CoreState`_ is the object on which the computation steps (forecasts and analysis) are performed. It provides an interface to store and act on the data of the involved physical quantities (called *measures*) such as:

- the magnetic field (*MF*) convenience getter *B*, 
- the secular variation (*SV*) convenience getter *SV*, 
- the core flow (*U*) convenience getter *U*, 
- the subgrid errors (*ER*) convenience getter *ER*, 
- the augmented state (*Z*) ([U,ER] concatenated) convenience getter *Z*, 
- the shear at the CMB (*S*) convenience getter *S*, 
- the first derivative of *U* (*dUdt*) convenience getter *dUdt*, 
- the second derivative of *U* (*d2Udt2*) convenience getter *d2Udt2*,  
- the first derivative of *ER* (*dERdt*) convenience getter *dERdt*, 
- the second derivative of *ER* (*d2ERdt2*) convenience getter *d2ERdt2*.

In ``run_algo.py``, the `CoreState`_ is used to store measure data as 3D NumPy arrays under the format *(nb_realisations, nb_times, nb_coeffs)*. However, it can handle any shape of NumPy arrays.

CoreState operations
====================
The following preamble is assumed:

.. code-block:: python

    import numpy as np
    from pygeodyn.corestates import CoreState


Add a measure and implications on SH degrees
--------------------------------------------

The implementation of the CoreState is done with a dictionary of NumPy arrays dedicated to store spherical harmonic coefficients of the measures. Each entry of the dictionary contains the measure data, that can be any NumPy array, and a spherical harmonic degree associated with it.

* One dimension (1D): the number of coefficients :

.. code-block:: python

    >>> cs_1D = CoreState()
    # Add an array to store the Gauss coefficients up to degree 14 (224 = 14*(14+2))
    >>> data_mf_1D = np.zeros(224)
    >>> cs_1D.addMeasure('MF', data_mf_1D)
    >>> cs_1D.Lb # Returns the max_degree of 'MF'
    14


* Three dimensions (3D) : number of realisations, number of times, numbers of coefficients

.. code-block:: python

    >>> cs_3D = CoreState()
    >>> data_mf_3D = np.ones((10, 50, 224))
    >>> cs_3D.addMeasure('MF', data_mf_3D)
    # Adding the core_flow Lu=18 (720 = 2*18*(18+2))
    >>> data_u_3D = np.ones((10, 50, 720))
    >>> cs_3D.addMeasure('U', data_u_3D)
    >>> cs_3D.Lu # Returns the max_degree of 'U'
    18

The max degree :math:`L` is inferred from the size of the last dimension (:math:`N = 2L(L+2)` for *U* and :math:`N = L(L+2)` for other measures). If this behaviour is not desired, it is possible to force the value of the max degree when adding a measure:

.. code-block:: python

    >>> data_u_1D = np.zeros(800) # 800 does not correspond to 2*Lu(Lu+2)
    >>> cs_1D.addMeasure('U', data_u_1D) # This throws an error
    ValueError: Last dimension 800 of the data given for U does not lead to an integer max degree ! Got 19.024984394500787 instead.
    >>> cs_1D.addMeasure('U', data_u_1D, meas_max_degree=18) # Forcing the max_degree to 18 while keeping 800 coefficients in data_u_1D
    >>> cs_1D.Lu # Returns the max_degree of 'U' (here: 18 the forced value)
    18

Get/set measure data
--------------------

To get the measure data, it is possible to use either the `getMeasure`_ methods or access the data like an attribute with convenience getters (see list above).

* 1D :

.. code-block:: python

    >>> cs_1D.getMeasure('MF') # Returns data_mf_1D
    array([0., 0., ..., 0.])
    >>> cs_1D.B # Convenience getter: returns data_mf_1D as well
    array([0., 0., ..., 0.])
    >>> type(cs_1D.B)
    <class 'numpy.ndarray'>
    >>> cs_1D.B.shape
    (224,)
    >>> cs_1D.Lb # Returns the max_degree of 'MF'
    14

Once we have access to the data, it is possible to set the value like a Numpy array:

.. code-block:: python

    >>> cs_1D.B[0] = 2.
    >>> cs_1D.U
    array([., 0., ..., 0.])
    >>> cs_1D.B[:2] = 3 # With slicing
    array([3., 3., ..., 0.])

To change the value of the whole measure data array, use `setMeasure`_:

.. code-block:: python

    >>> cs_1D.setMeasure('U', np.ones(800))
    >>> cs_1D.U
    array([1., 1., ..., 1.])


* 3D :

It works exactly the same but CoreState also supports NumPy-like indexing. The indexing is delegated to the measures contained in the CoreState to return an indexed sub CoreState. In other words:

.. code-block:: python

    cs_2D = cs_3D[index]

is equivalent to

.. code-block:: python

    cs_2D = CoreState()
    for measure in cs_3D.measures:
        cs_2D.addMeasure(measure, cs_3D.getMeasure(measure)[index], meas_max_degree=cs_3D.getMaxDegree(measure))

An example on the ``cs_3D`` previously created:

.. code-block:: python

    >>> cs_2D = cs_3D[0] # Stores the indexed corestate in cs_2D
    >>> cs_2D.B # returns data_mf_3D[0]
    array([[1., 1., 1., ..., 1., 1., 1.],
           [1., 1., 1., ..., 1., 1., 1.],
           [1., 1., 1., ..., 1., 1., 1.],
           ...,
           [1., 1., 1., ..., 1., 1., 1.],
           [1., 1., 1., ..., 1., 1., 1.],
           [1., 1., 1., ..., 1., 1., 1.]])
    >>> cs_2D.B.shape
    (50, 224)
    >>> cs_2D = cs_3D[:, :10] # Works also with slices
    >>> cs_2D.B.shape # returns the shape of data_mf_3D[:, :10]
    (10, 10, 224)
    >>> cs_2D.U.shape # returns the shape of data_u_3D[:, :10]
    (10, 10, 720)
    >>> cs_2D.Lb # returns 14 as all max degrees are copied
    14

.. warning:: **Using indexing on the CoreState creates a new CoreState !**

    It means that to update a part of measure data, it is the measure data that must be indexed and not the **CoreState**. Say we want to change the value of the first coefficient of 'MF' for all times and realisations:

    .. code-block:: python

        >>> cs_3D.B[:,:,0] = -40 # OK: first it gets the B data and sets it.
        >>> cs_3D.B[0, 0, 0] # Returns the new value -40
        -40.0
        >>> cs_3D[:,:,0].B = 20 # WONT WORK: it first creates a new indexed CoreState then sets the B data of the indexed CoreState !
        >>> cs_3D.B[0, 0, 0] # Returns the old value: we only updated a transient copy of cs_3D
        -40.0

    The implementation of the indexing in CoreState is indeed done as presented above (create a new CoreState and loop through the measures to get the indexed data that are added to the new CoreState).


If the measures and their shapes match, it is possible to set the measures of a `CoreState`_ (or a part) using another `CoreState`_ :

.. code-block:: python

    >>> cs_1D = CoreState()
    >>> cs_1D.addMeasure('MF', np.ones(224))
    >>> cs_1D.addMeasure('SV', np.ones(224)*2)
    ['MF', 'SV']
    >>> cs_1D.measures
    >>> cs_3D = CoreState()
    >>> cs_3D.addMeasure('MF', np.zeros((10, 50, 224)))
    >>> cs_3D.addMeasure('SV', np.zeros((10, 50, 224)))
    >>> cs_3D.measures
    ['MF', 'SV']
    >>> cs_3D[0, 0] = cs_1D # Setting a subpart of cs_3D with cs_1D values
    >>> cs_3D.B[0, 0], cs_3D.shape # Returns np.ones(224)
    array([1., 1., ..., 1.])
    >>> cs_3D.SV[0, 0] # Returns np.ones(224)*2
    array([2., 2., ..., 2.])

.. _CoreState: https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.corestates.html
.. _getMeasure: https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.corestates.html#pygeodyn.corestates.CoreState.getMeasure
.. _setMeasure: https://geodynamo.gricad-pages.univ-grenoble-alpes.fr/pygeodyn/pygeodyn.corestates.html#pygeodyn.corestates.CoreState.setMeasure