apply_AK.py

import numpy as np


def apply_ak(sim_ak, dpavgs, aks, nbformula, qa0, chosenlevel, drycols):
    """Apply the corresponding AK values"""
    if nbformula == 1:
        y0 = (sim_ak * dpavgs * aks).sum(axis=0) / (dpavgs * aks).sum(axis=0)

    # OMI-QA4ECV & TROPOMI-S5P [NO2] et [CO]
    elif nbformula == 2:
        y0 = (sim_ak * aks).sum(axis=0)

    # MOPITT
    elif nbformula == 3:
        # Check that qa0 is well defined
        if ~np.all(~np.isnan(qa0[chosenlevel])) or np.all(np.isnan(qa0)):
            raise Exception("WARNING! Can't apply formula #3 with NaNs in qa0")
        prof = aks * (np.log10(sim_ak) - np.log10(qa0))
        prof[~np.isfinite(prof)] = np.nan
        y0 = np.log10(qa0[chosenlevel]) \
             + np.nansum(prof, axis=0)
        y0 = np.power(10, y0)

    # TROPOMI-S5P [CO]
    elif nbformula == 4:
        y0 = (sim_ak * aks * 10e-4).sum(axis=0)

    # TROPOMI-S5P [CH4]
    elif nbformula == 5:
        y0 = (qa0.sum(axis=0)
              + (aks * (sim_ak * drycols - qa0)).sum(axis=0)
              ) / drycols.sum(axis=0)

    # OCO-2
    elif nbformula == 6:
        y0 = (dpavgs*qa0).sum(axis=0) \
            + (dpavgs*aks*(sim_ak - qa0)).sum(axis=0)

    else:
        raise Exception("Don't know formula: {}".format(nbformula))

    return y0


def apply_ak_tl(
        sim_ak_tl, dpavgs, aks, nbformula, qa0, chosenlevel, sim_ak, drycols
):
    if nbformula == 1 or nbformula == 2 or nbformula == 4:
        y0 = apply_ak(sim_ak_tl, dpavgs, aks, nbformula, qa0, chosenlevel,
                      drycols)

    elif nbformula == 3:
        tmp = sim_ak_tl / (sim_ak * np.log(10))
        y0 = apply_ak(sim_ak, dpavgs, aks, nbformula, qa0, chosenlevel, drycols)
        y0 = y0 * np.log(10) * np.nansum(aks * tmp, axis=0)

    # passe dans le premier if normalement
    # elif nbformula == 4:
    #     y0 = (sim_ak_tl * aks * 10e-4).sum(axis=0)

    elif nbformula == 5:
        y0 = (aks * sim_ak_tl * drycols).sum(axis=0) / drycols.sum(axis=0)

    elif nbformula == 6:
        y0 = (dpavgs*aks*sim_ak_tl).sum(axis=0)

    else:
        raise Exception("Don't know formula: {}".format(nbformula))

    return y0


def apply_ak_ad(sim_ak, dpavgs, aks, nbformula, qa0, chosenlevel, obs_incr,
                drycols):
    if nbformula == 1:
        obs_incr = (
                dpavgs
                * aks
                / (dpavgs * aks).sum(axis=0)
                * obs_incr
        )

    elif nbformula == 2:
        obs_incr = aks * obs_incr

    elif nbformula == 3:
        y0 = apply_ak(sim_ak, dpavgs, aks, nbformula, qa0, chosenlevel, drycols)
        y0 *= np.log(10) * obs_incr
        y0 = aks * y0
        prof = sim_ak * np.log(10)
        prof[prof == 0] = np.nan
        obs_incr = y0 / prof

    elif nbformula == 4:
        obs_incr = aks * obs_incr * 10e-4

    elif nbformula == 5:
        obs_incr = aks * obs_incr * drycols / drycols.sum(axis=0)

    elif nbformula == 6:
        obs_incr = dpavgs*aks*obs_incr

    else:
        raise Exception("Don't know formula: {}".format(nbformula))

    return obs_incr