Source code for pycif.plugins.controlvects.standard.utils.rescale_std

import numpy as np
from logging import debug
from .get_physical import get_physical
from .....utils.check.errclass import CifError




def rescale_std(controlvect, tracer, comp, trcr, glob_err, **kwargs):
    """Re-scale errors for a given tracer depending on the specied total budget


    :param controlvect: reference control vector
    :param tracer: tracer object
    :param glob_err: error on global budget
    :return:
    """

    debug("Computing global error for rescaling")

    # First try reading from pre-computed simulation
    file_scaling = f"{controlvect.workdir}/controlvect/globerror_scaling_{comp}_{trcr}.txt"
    try:
        with open(file_scaling, "r") as f:
            errscalar = float(f.read())

        controlvect.std[tracer.xpointer: tracer.xpointer +
                        tracer.dim] *= errscalar
        return

    except IOError:
        debug("No file to retrieve, computing from scratch")
        pass

    except Exception as e:
        debug(
            f"Could not retrieve information from {file_scaling}. "
            "Check the file and consider removing it to force re-computation of error scaling factor"
        )
        raise e

    # Fetch tracer std from global control vector
    # If scalar, re-build physical values
    if getattr(tracer, "type", "scalar") == "scalar":
        xb, std = get_physical(controlvect, tracer, comp, trcr)
    else:
        std = controlvect.std[tracer.xpointer: tracer.xpointer + tracer.dim]

    # Scale with areas
    if tracer.hresol not in ["hpixels", "global"]:
        areas = tracer.region_areas
    elif tracer.hresol == "hpixels":
        areas = tracer.domain.areas
    else:
        raise CifError("Check how to deal with global resolution")

    xerr = (np.reshape(std, (tracer.ndates, tracer.nlev, -1))
            * areas.flatten()[np.newaxis]).flatten()

    # Re-scaling by total areas
    if getattr(glob_err, "surface_unit", False):
        xerr /= areas.mean()

    # Put std in absolute value using "unit_scale"
    # Also include temporal period
    unit_scale = getattr(glob_err, "unit_scale", 1)
    xerr *= unit_scale

    # Re-create temporal correlation if any to aggregate temporally
    tcorr = np.eye(tracer.ndates)
    if hasattr(tracer, "tcorrelations") \
            and getattr(glob_err, "account_correlations", True):
        tcorr = tracer.tcorrelations
        tsqrt_evalues = tcorr.sqrt_evalues
        tevectors = tcorr.evectors
        tcorr = tevectors.dot(np.diag(tsqrt_evalues ** 2)).dot(tevectors.T)

    xerr = np.reshape(xerr, (tracer.ndates, -1))
    xerr = np.sqrt(np.sum(xerr * np.dot(tcorr, xerr), axis=0))

    # Rescale to period error
    if getattr(glob_err, "frequency_unit", False):
        xerr /= np.sqrt(tracer.ndates)

    # vertical correlation ?
    vcorr = np.eye(tracer.nlev)
    xerr = np.reshape(xerr, (tracer.nlev, -1))
    xerr = np.sqrt(np.sum(xerr * np.dot(vcorr, xerr), axis=0))

    # Re-create horizontal correlation if any
    hcorr = np.eye(xerr.size)
    if hasattr(tracer, "hcorrelations") \
            and getattr(glob_err, "account_correlations", True):
        hcorr = tracer.hcorrelations
        hsqrt_evalues = hcorr.sqrt_evalues
        hevectors = hcorr.evectors
        hcorr = hevectors.dot(np.diag(hsqrt_evalues ** 2)).dot(hevectors.T)

    # Computing covariances
    cov = np.outer(xerr, xerr) * hcorr

    # Scale to the chosen absolute unit, e.g., Tg/y
    toterr = np.sqrt(cov.sum())

    # Re-scale by number of pixels if surface unit
    if getattr(glob_err, "surface_unit", False):
        toterr /= np.sqrt(xerr.size)

    debug(f"Total errors after unit scaling: {toterr}")
    debug(f"Target total errors: {glob_err.total}")

    # Compute error
    errscalar = np.sqrt(glob_err.total / toterr)
    debug(f"Total error scaled by {errscalar} (sqrt)")

    controlvect.std[tracer.xpointer: tracer.xpointer + tracer.dim] *= errscalar

    # Dump scaling factor for later use
    with open(f"{controlvect.workdir}/controlvect/globerror_scaling_{comp}_{trcr}.txt", "w") as f:
        f.write(str(errscalar))