Source code for pycif.plugins.controlvects.standard.init_bprod

import numpy as np
from logging import debug

try:
    import cPickle as pickle
except ImportError:
    import pickle

from .utils.rescale_std import rescale_std
from .utils.get_correlations import get_vcorr, get_tcorr, get_hcorr




def init_bprod(cntrlv, options={}, **kwargs):
    """Initilializes the product of chi by sqrt-B. It allows translating
    information from the minimization space to the control space. This first
    needs to initialize correlation matrices

    Args:
        cntrlv (dict): definition of the control vector

    Returns:
        updated control vector:

    """

    datavect = cntrlv.datavect

    # Initializing sqrt-B
    cntrlv.hcorrelations = {}
    cntrlv.tcorrelations = {}
    cntrlv.vcorrelations = {}
    cntrlv.chi_dim = 0
    components = datavect.components
    for comp in components.attributes:
        component = getattr(components, comp)

        # Skip if component does not have parameters
        if not hasattr(component, "parameters"):
            continue

        for trcr in component.parameters.attributes:
            tracer = getattr(component.parameters, trcr)

            # Skip tracers that are not control variables
            if not hasattr(tracer, "hresol"):
                continue

            # Deals with horizontal correlations if any
            hresol = tracer.hresol
            if hasattr(tracer, "hcorrelations") \
                    and hresol != "global":
                debug(f"Fetching horizontal correlation for {comp}/{trcr}")
                sqrt_evalues, evectors = get_hcorr(cntrlv, tracer)
                tracer.chi_hresoldim = sqrt_evalues.size

            else:
                tracer.chi_hresoldim = tracer.hresoldim

            # Initializes temporal correlations if any
            if hasattr(tracer, "tcorrelations"):
                corr = tracer.tcorrelations
                if getattr(corr, "multi_sigmas", False):
                    sigmas = corr.sigmas.attributes
                    tresols = []
                    ndates = getattr(tracer, "dates").size
                    B = np.ones((ndates, ndates))
                    for sigma in sigmas:
                        tmp_corr = getattr(corr.sigmas, sigma)
                        sqrt_evalues, evectors = get_tcorr(
                            cntrlv, tracer, tmp_corr)
                        tresols.append(sqrt_evalues.size)

                        # Aggregate to full B
                        B *= evectors.dot(
                            np.diag(sqrt_evalues ** 2)).dot(evectors.T)

                    # Check that all temporal resolution from sub-sigmas are
                    # consistent
                    tresols = list(set(tresols))
                    if len(tresols) != 1:
                        raise Exception("Error with multiple temporal "
                                        "correlations, the length of all "
                                        "sub-correlation matrices is not "
                                        "consistent.")

                    chi_tresoldim = tresols[0]

                    # Aggregate multiple B
                    evalues, evectors = np.linalg.eigh(B)
                    corr.evectors = evectors
                    corr.sqrt_evalues = np.sqrt(evalues)

                else:
                    corr.type = getattr(corr, "type", "isotrope")
                    sqrt_evalues, evectors = get_tcorr(cntrlv, tracer, corr)
                    chi_tresoldim = sqrt_evalues.size

                tracer.chi_tresoldim = chi_tresoldim

            else:
                tracer.chi_tresoldim = tracer.ndates

            # Deal with vertical correlations
            if hasattr(tracer, "vcorrelations") \
                    and tracer.vresol != "column":
                sqrt_evalues, evectors = get_vcorr(cntrlv, tracer)
                tracer.chi_vresoldim = sqrt_evalues.size

            else:
                tracer.chi_vresoldim = tracer.vresoldim

            # Update the dimension of the full control vector
            tracer.chi_pointer = cntrlv.chi_dim
            tracer.chi_dim = \
                tracer.chi_hresoldim \
                * tracer.chi_tresoldim \
                * tracer.chi_vresoldim
            cntrlv.chi_dim += tracer.chi_dim

            # Rescale std depending on the global error if specified
            glob_err = getattr(tracer, "glob_err", None)
            if glob_err is None:
                continue

            if glob_err.total < 0:
                raise Exception("glob_err must be > 0")

            rescale_std(cntrlv, tracer, comp, trcr, glob_err)

    # Defining chi from the total dimension
    cntrlv.chi = np.zeros((cntrlv.chi_dim,))

    # Save the full B matrix if required
    # Be careful with the size of your set-up
    if cntrlv.save_full_B:
        bfull = cntrlv.build_b(cntrlv)
        bfile = "{}/controlvect/bfull.pickle".format(cntrlv.workdir)
        with open(bfile, "wb") as f:
            pickle.dump(bfull, f)

    # Random perturbation of the control vector
    if cntrlv.perturb_xb:
        np.random.seed()
        x = np.random.normal(size=(1, cntrlv.dim))
        x_sample = cntrlv.sqrtbprod(x)
        cntrlv.xb = x_sample

    return cntrlv