test_modelling.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-


# see test/functs/fixtures.py for global fixtures "course_..."

import dios

from saqc import BAD, UNFLAGGED
from saqc.core import initFlagsLike
from saqc.funcs.tools import mask
from saqc.funcs.residues import calculatePolynomialResidues, calculateRollingResidues

from tests.fixtures import *


@pytest.mark.parametrize("dat", [pytest.lazy_fixture("course_2")])
def test_modelling_polyFit_forRegular(dat):
    data, _ = dat(freq="10min", periods=30, initial_level=0, final_level=100, out_val=-100)
    # add some nice sine distortion
    data = data + 10 * np.sin(np.arange(0, len(data.indexes[0])))
    data = dios.DictOfSeries(data)
    flagger = initFlagsLike(data)
    result1, _ = calculatePolynomialResidues(data, "data", flagger, 11, 2, numba=False)
    result2, _ = calculatePolynomialResidues(data, "data", flagger, 11, 2, numba=True)
    assert (result1["data"] - result2["data"]).abs().max() < 10 ** -10
    result3, _ = calculatePolynomialResidues(data, "data", flagger, "110min", 2, numba=False)
    assert result3["data"].equals(result1["data"])
    result4, _ = calculatePolynomialResidues(data, "data", flagger, 11, 2, numba=True, min_periods=11)
    assert (result4["data"] - result2["data"]).abs().max() < 10 ** -10
    data.iloc[13:16] = np.nan
    result5, _ = calculatePolynomialResidues(data, "data", flagger, 11, 2, numba=True, min_periods=9)
    assert result5["data"].iloc[10:19].isna().all()


@pytest.mark.parametrize("dat", [pytest.lazy_fixture("course_2")])
def test_modelling_rollingMean_forRegular(dat):
    data, _ = dat(freq="10min", periods=30, initial_level=0, final_level=100, out_val=-100)
    data = dios.DictOfSeries(data)
    flagger = initFlagsLike(data)
    calculateRollingResidues(data, "data", flagger, 5, func=np.mean, eval_flags=True, min_periods=0, center=True)
    calculateRollingResidues(data, "data", flagger, 5, func=np.mean, eval_flags=True, min_periods=0, center=False)


@pytest.mark.parametrize("dat", [pytest.lazy_fixture("course_1")])
def test_modelling_mask(dat):
    data, _ = dat()
    data = dios.DictOfSeries(data)
    flagger = initFlagsLike(data)
    field = "data"

    # set flags everywhere to test unflagging
    flagger[:, field] = BAD

    common = dict(data=data, field=field, flagger=flagger, mode='periodic')
    data_seasonal, flagger_seasonal = mask(**common, period_start="20:00", period_end="40:00", include_bounds=False)
    flags = flagger_seasonal[field]
    m = (20 <= flags.index.minute) & (flags.index.minute <= 40)
    assert all(flagger_seasonal[field][m] == UNFLAGGED)
    assert all(data_seasonal[field][m].isna())

    data_seasonal, flagger_seasonal = mask(**common, period_start="15:00:00", period_end="02:00:00")
    flags = flagger_seasonal[field]
    m = (15 <= flags.index.hour) & (flags.index.hour <= 2)
    assert all(flagger_seasonal[field][m] == UNFLAGGED)
    assert all(data_seasonal[field][m].isna())

    data_seasonal, flagger_seasonal = mask(**common, period_start="03T00:00:00", period_end="10T00:00:00")
    flags = flagger_seasonal[field]
    m = (3 <= flags.index.hour) & (flags.index.hour <= 10)
    assert all(flagger_seasonal[field][m] == UNFLAGGED)
    assert all(data_seasonal[field][m].isna())

    mask_ser = pd.Series(False, index=data["data"].index)
    mask_ser[::5] = True
    data["mask_ser"] = mask_ser
    flagger = initFlagsLike(data)
    data_masked, flagger_masked = mask(data, "data", flagger, mode='mask_var', mask_var="mask_ser")
    m = mask_ser
    assert all(flagger_masked[field][m] == UNFLAGGED)
    assert all(data_masked[field][m].isna())