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

import numpy as np

from ....utils.iterators import iter_tracers




def build_b_block(controlvect, tracer):
    """Build the B matrix block corresponding to 'tracer'

    Args:
        controlvect (ControlVect): the controlvector plugin
        tracer (tracer plugin): the tracer to compute the corresping B block

    Returns:
        2D array: B matrix block
    """
    b_horiz = np.eye(tracer.dim)

    # Deals with non-diagonal horizontal correlations
    if hasattr(tracer, "hcorrelations") and tracer.hresol == "hpixels":
        corr = tracer.hcorrelations
        sqrt_evalues = corr.sqrt_evalues
        evectors = corr.evectors

        btmp = evectors.dot(np.diag(sqrt_evalues ** 2)).dot(evectors.T)

        for i in range(tracer.ndates * tracer.vresoldim):
            slc = slice(i * tracer.hresoldim, (i+1) * tracer.hresoldim)
            b_horiz[slc, slc] = btmp

    b_temp = np.eye(tracer.dim)

    # Deals with non-diagonal temporal correlations
    if hasattr(tracer, "tcorrelations"):
        corr = tracer.tcorrelations
        sqrt_evalues = corr.sqrt_evalues
        evectors = corr.evectors

        btmp = evectors.dot(np.diag(sqrt_evalues ** 2)).dot(evectors.T)

        for i in range(tracer.hresoldim * tracer.vresoldim):
            rdim = range(i, tracer.dim,
                         tracer.hresoldim * tracer.vresoldim)
            slc = np.ix_(rdim, rdim)
            b_temp[slc] = btmp

    start = tracer.xpointer
    stop = tracer.xpointer + tracer.dim

    b_block = (
        controlvect.std[start:stop, np.newaxis] *
        controlvect.std[start:stop] *
        b_temp.dot(b_horiz)
    )

    return b_block





def build_b(controlvect, component=None, parameter=None):
    """Compute the full B matrix

    Args:
        controlvect (ControlVect): the controlvector plugin
        component (str, optional): only compute B block for a given tracer ('component', 'tracer').
            Must be used with the 'tracer' argument. Defaults to None.
        parameter (str, optional): only compute B block for a given tracer ('component', 'tracer').
            Must be used with the 'component' argument. Defaults to None.Defaults to None.

    Returns:
        2D array: B matrix
    """
    # Checking if 'component' and 'parameter' arguments are all None or if none is 'None'
    if ((component is None and parameter is not None) or
            (component is not None and parameter is None)):
        raise ValueError(
            "'component' and 'tracer' arguments must used together, "
            "one can not be 'None' while the other is not"
        )

    datavect = controlvect.datavect

    # Building only one block
    if component is not None:
        # At this point 'parameter' can not be 'None'
        try:
            # Directly getting tracer plugin
            comp = getattr(datavect.components, component)
            tracer = getattr(comp, parameter)
        except AttributeError as e:
            raise ValueError(
                f"tracer '{component}/{parameter}' not found") from e

        if not tracer.iscontrol:
            raise ValueError(
                f"tracer '{component}/{parameter}' is not in the control vector")

        return build_b_block(controlvect, tracer)

    # Building the full B matrix
    b_full = np.eye(controlvect.dim)

    for _, tracer in iter_tracers(datavect):
        if not tracer.iscontrol:
            continue

        start = tracer.xpointer
        stop = tracer.xpointer + tracer.dim

        b_full[start:stop, start:stop] = build_b_block(controlvect, tracer)

    return b_full