test_data.py

from copy import deepcopy
import itertools
import pickle
from typing import Dict

import numpy as np
import pandas as pd
import pytest
from sklearn.preprocessing import StandardScaler
import torch

from pytorch_forecasting.data import (
    EncoderNormalizer,
    GroupNormalizer,
    NaNLabelEncoder,
    TimeSeriesDataSet,
    TimeSynchronizedBatchSampler,
)
from pytorch_forecasting.data.encoders import MultiNormalizer, TorchNormalizer
from pytorch_forecasting.data.examples import get_stallion_data
from pytorch_forecasting.data.timeseries import _find_end_indices
from pytorch_forecasting.utils import to_list

torch.manual_seed(23)


@pytest.mark.parametrize(
    "data,allow_nan",
    itertools.product(
        [
            (np.array([2, 3, 4]), np.array([1, 2, 3, 5, np.nan])),
            (np.array(["a", "b", "c"]), np.array(["q", "a", "nan"])),
        ],
        [True, False],
    ),
)
def test_NaNLabelEncoder(data, allow_nan):
    fit_data, transform_data = data
    encoder = NaNLabelEncoder(warn=False, add_nan=allow_nan)
    encoder.fit(fit_data)
    assert np.array_equal(
        encoder.inverse_transform(encoder.transform(fit_data)), fit_data
    ), "Inverse transform should reverse transform"
    if not allow_nan:
        with pytest.raises(KeyError):
            encoder.transform(transform_data)
    else:
        assert encoder.transform(transform_data)[0] == 0, "First value should be translated to 0 if nan"
        assert encoder.transform(transform_data)[-1] == 0, "Last value should be translated to 0 if nan"
        assert encoder.transform(fit_data)[0] > 0, "First value should not be 0 if not nan"


def test_NaNLabelEncoder_add():
    encoder = NaNLabelEncoder(add_nan=False)
    encoder.fit(np.array(["a", "b", "c"]))
    encoder2 = deepcopy(encoder)
    encoder2.fit(np.array(["d"]))
    assert encoder2.transform(np.array(["a"]))[0] == 0, "a must be encoded as 0"
    assert encoder2.transform(np.array(["d"]))[0] == 3, "d must be encoded as 3"


@pytest.mark.parametrize(
    "kwargs",
    [
        dict(method="robust"),
        dict(transformation="log"),
        dict(transformation="softplus"),
        dict(transformation="log1p"),
        dict(transformation="relu"),
        dict(center=False),
    ],
)
def test_EncoderNormalizer(kwargs):
    data = torch.rand(100)
    defaults = dict(method="standard", center=True)
    defaults.update(kwargs)
    kwargs = defaults
    normalizer = EncoderNormalizer(**kwargs)
    if kwargs.get("transformation") in ["relu", "softplus"]:
        data = data - 0.5

    if kwargs.get("transformation") in ["relu", "softplus", "log1p"]:
        assert (
            normalizer.inverse_transform(normalizer.fit_transform(data)) >= 0
        ).all(), "Inverse transform should yield only positive values"
    else:
        assert torch.isclose(
            normalizer.inverse_transform(normalizer.fit_transform(data)), data, atol=1e-5
        ).all(), "Inverse transform should reverse transform"


@pytest.mark.parametrize(
    "kwargs,groups",
    itertools.product(
        [
            dict(method="robust"),
            dict(transformation="log"),
            dict(transformation="relu"),
            dict(center=False),
            dict(transformation="log1p"),
            dict(transformation="softplus"),
            dict(scale_by_group=True),
        ],
        [[], ["a"]],
    ),
)
def test_GroupNormalizer(kwargs, groups):
    data = pd.DataFrame(dict(a=[1, 1, 2, 2, 3], b=[1.1, 1.1, 1.0, 5.0, 1.1]))
    defaults = dict(method="standard", transformation=None, center=True, scale_by_group=False)
    defaults.update(kwargs)
    kwargs = defaults
    kwargs["groups"] = groups
    kwargs["scale_by_group"] = kwargs["scale_by_group"] and len(kwargs["groups"]) > 0

    if kwargs.get("transformation") in ["relu", "softplus"]:
        data.b = data.b - 2.0
    normalizer = GroupNormalizer(**kwargs)
    encoded = normalizer.fit_transform(data["b"], data)

    test_data = dict(
        prediction=torch.tensor([encoded[0]]),
        target_scale=torch.tensor(normalizer.get_parameters([1])).unsqueeze(0),
    )

    if kwargs.get("transformation") in ["relu", "softplus", "log1p"]:
        assert (normalizer(test_data) >= 0).all(), "Inverse transform should yield only positive values"
    else:
        assert torch.isclose(
            normalizer(test_data), torch.tensor(data.b.iloc[0]), atol=1e-5
        ).all(), "Inverse transform should reverse transform"


def check_dataloader_output(dataset: TimeSeriesDataSet, out: Dict[str, torch.Tensor]):
    x, y = out

    assert isinstance(y, tuple), "y output should be tuple of wegith and target"

    # check for nans and finite
    for k, v in x.items():
        for vi in to_list(v):
            assert torch.isfinite(vi).all(), f"Values for {k} should be finite"
            assert not torch.isnan(vi).any(), f"Values for {k} should not be nan"

    # check weight
    assert y[1] is None or isinstance(y[1], torch.Tensor), "weights should be none or tensor"
    if isinstance(y[1], torch.Tensor):
        assert torch.isfinite(y[1]).all(), "Values for weight should be finite"
        assert not torch.isnan(y[1]).any(), "Values for weight should not be nan"

    # check target
    for targeti in to_list(y[0]):
        assert torch.isfinite(targeti).all(), "Values for target should be finite"
        assert not torch.isnan(targeti).any(), "Values for target should not be nan"

    # check shape
    assert x["encoder_cont"].size(2) == len(dataset.reals)
    assert x["encoder_cat"].size(2) == len(dataset.flat_categoricals)


@pytest.mark.parametrize(
    "kwargs",
    [
        dict(min_encoder_length=0, max_prediction_length=2),
        dict(static_categoricals=["agency", "sku"]),
        dict(static_reals=["avg_population_2017", "avg_yearly_household_income_2017"]),
        dict(time_varying_known_categoricals=["month"]),
        dict(
            time_varying_known_categoricals=["special_days", "month"],
            variable_groups=dict(
                special_days=[
                    "easter_day",
                    "good_friday",
                    "new_year",
                    "christmas",
                    "labor_day",
                    "independence_day",
                    "revolution_day_memorial",
                    "regional_games",
                    "fifa_u_17_world_cup",
                    "football_gold_cup",
                    "beer_capital",
                    "music_fest",
                ]
            ),
        ),
        dict(time_varying_known_reals=["time_idx", "price_regular", "discount_in_percent"]),
        dict(time_varying_unknown_reals=["volume", "log_volume", "industry_volume", "soda_volume", "avg_max_temp"]),
        dict(
            target_normalizer=GroupNormalizer(
                groups=["agency", "sku"],
                transformation="log1p",
                scale_by_group=True,
            )
        ),
        dict(target_normalizer=EncoderNormalizer(), min_encoder_length=2),
        dict(randomize_length=True, min_encoder_length=2, min_prediction_length=1),
        dict(predict_mode=True),
        dict(add_target_scales=True),
        dict(add_encoder_length=True),
        dict(add_encoder_length=True),
        dict(add_relative_time_idx=True),
        dict(weight="volume"),
        dict(
            scalers=dict(time_idx=GroupNormalizer(), price_regular=StandardScaler()),
            categorical_encoders=dict(month=NaNLabelEncoder()),
            time_varying_known_categoricals=["month"],
            time_varying_known_reals=["time_idx", "price_regular"],
        ),
        dict(categorical_encoders={"month": NaNLabelEncoder(add_nan=True)}, time_varying_known_categoricals=["month"]),
        dict(constant_fill_strategy=dict(volume=0.0), allow_missing_timesteps=True),
        dict(target_normalizer=None),
    ],
)
def test_TimeSeriesDataSet(test_data, kwargs):

    defaults = dict(
        time_idx="time_idx",
        target="volume",
        group_ids=["agency", "sku"],
        max_encoder_length=5,
        max_prediction_length=2,
    )
    defaults.update(kwargs)
    kwargs = defaults

    if kwargs.get("allow_missing_timesteps", False):
        np.random.seed(2)
        test_data = test_data.sample(frac=0.5)
        defaults["min_encoder_length"] = 0
        defaults["min_prediction_length"] = 1

    # create dataset and sample from it
    dataset = TimeSeriesDataSet(test_data, **kwargs)
    check_dataloader_output(dataset, next(iter(dataset.to_dataloader(num_workers=0))))


def test_from_dataset(test_dataset, test_data):
    dataset = TimeSeriesDataSet.from_dataset(test_dataset, test_data)
    check_dataloader_output(dataset, next(iter(dataset.to_dataloader(num_workers=0))))


def test_from_dataset_equivalence(test_data):
    training = TimeSeriesDataSet(
        test_data[lambda x: x.time_idx < x.time_idx.max() - 1],
        time_idx="time_idx",
        target="volume",
        time_varying_known_reals=["price_regular", "time_idx"],
        group_ids=["agency", "sku"],
        static_categoricals=["agency"],
        max_encoder_length=3,
        max_prediction_length=2,
        min_prediction_length=1,
        min_encoder_length=0,
        randomize_length=None,
        add_encoder_length=True,
        add_relative_time_idx=True,
        add_target_scales=True,
    )
    validation1 = TimeSeriesDataSet.from_dataset(training, test_data, predict=True)
    validation2 = TimeSeriesDataSet.from_dataset(
        training,
        test_data[lambda x: x.time_idx > x.time_idx.min() + 2],
        predict=True,
    )
    # ensure validation1 and validation2 datasets are exactly the same despite different data inputs
    for v1, v2 in zip(iter(validation1.to_dataloader(train=False)), iter(validation2.to_dataloader(train=False))):
        for k in v1[0].keys():
            if isinstance(v1[0][k], (tuple, list)):
                assert len(v1[0][k]) == len(v2[0][k])
                for idx in range(len(v1[0][k])):
                    assert torch.isclose(v1[0][k][idx], v2[0][k][idx]).all()
            else:
                assert torch.isclose(v1[0][k], v2[0][k]).all()
        assert torch.isclose(v1[1][0], v2[1][0]).all()


def test_dataset_index(test_dataset):
    index = []
    for x, _ in iter(test_dataset.to_dataloader()):
        index.append(test_dataset.x_to_index(x))
    index = pd.concat(index, axis=0, ignore_index=True)
    assert len(index) <= len(test_dataset), "Index can only be subset of dataset"


@pytest.mark.parametrize("min_prediction_idx", [0, 1, 3, 7])
def test_min_prediction_idx(test_dataset, test_data, min_prediction_idx):
    dataset = TimeSeriesDataSet.from_dataset(
        test_dataset, test_data, min_prediction_idx=min_prediction_idx, min_encoder_length=1, max_prediction_length=10
    )

    for x, _ in iter(dataset.to_dataloader(num_workers=0, batch_size=1000)):
        assert x["decoder_time_idx"].min() >= min_prediction_idx


@pytest.mark.parametrize(
    "value,variable,target",
    [
        (1.0, "price_regular", "encoder"),
        (1.0, "price_regular", "all"),
        (1.0, "price_regular", "decoder"),
        ("Agency_01", "agency", "all"),
        ("Agency_01", "agency", "decoder"),
    ],
)
def test_overwrite_values(test_dataset, value, variable, target):
    dataset = deepcopy(test_dataset)

    # create variables to check against
    control_outputs = next(iter(dataset.to_dataloader(num_workers=0, train=False)))
    dataset.set_overwrite_values(value, variable=variable, target=target)

    # test change
    outputs = next(iter(dataset.to_dataloader(num_workers=0, train=False)))
    check_dataloader_output(dataset, outputs)

    if variable in dataset.reals:
        output_name_suffix = "cont"
    else:
        output_name_suffix = "cat"

    if target == "all":
        output_names = [f"encoder_{output_name_suffix}", f"decoder_{output_name_suffix}"]
    else:
        output_names = [f"{target}_{output_name_suffix}"]

    for name in outputs[0].keys():
        changed = torch.isclose(outputs[0][name], control_outputs[0][name]).all()
        if name in output_names or (
            "cat" in name and variable == "agency"
        ):  # exception for static categorical which should always change
            assert not changed, f"Output {name} should change"
        else:
            assert changed, f"Output {name} should not change"

    # test resetting
    dataset.reset_overwrite_values()
    outputs = next(iter(dataset.to_dataloader(num_workers=0, train=False)))
    for name in outputs[0].keys():
        changed = torch.isclose(outputs[0][name], control_outputs[0][name]).all()
        assert changed, f"Output {name} should be reset"
    assert torch.isclose(outputs[1][0], control_outputs[1][0]).all(), "Target should be reset"


@pytest.mark.parametrize(
    "drop_last,shuffle,as_string,batch_size",
    [
        (True, True, True, 64),
        (False, False, False, 64),
        (True, False, False, 1000),
    ],
)
def test_TimeSynchronizedBatchSampler(test_dataset, shuffle, drop_last, as_string, batch_size):
    if as_string:
        dataloader = test_dataset.to_dataloader(
            batch_sampler="synchronized", shuffle=shuffle, drop_last=drop_last, batch_size=batch_size
        )
    else:
        sampler = TimeSynchronizedBatchSampler(
            data_source=test_dataset, shuffle=shuffle, drop_last=drop_last, batch_size=batch_size
        )
        dataloader = test_dataset.to_dataloader(batch_sampler=sampler)

    time_idx_pos = test_dataset.reals.index("time_idx")
    for x, _ in iter(dataloader):  # check all samples
        time_idx_of_first_prediction = x["decoder_cont"][:, 0, time_idx_pos]
        assert torch.isclose(
            time_idx_of_first_prediction, time_idx_of_first_prediction[0]
        ).all(), "Time index should be the same for the first prediction"


def test_find_end_indices():
    diffs = np.array([1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1])
    max_lengths = np.array([4, 4, 4, 4, 4, 4, 4, 4, 3, 2, 1, 4, 4, 4, 4, 4, 4, 4, 4, 3, 2, 1])

    ends, missings = _find_end_indices(diffs, max_lengths, min_length=3)
    ends_test = np.array([3, 4, 4, 5, 6, 8, 9, 10, 10, 10, 10, 14, 15, 15, 16, 17, 19, 20, 21, 21, 21, 21])
    missings_test = np.array([[0, 2], [5, 7], [11, 13], [16, 18]])
    np.testing.assert_array_equal(ends, ends_test)
    np.testing.assert_array_equal(missings, missings_test)


def test_raise_short_encoder_length(test_data):
    with pytest.warns(UserWarning):
        test_data = test_data[lambda x: ~((x.agency == "Agency_22") & (x.sku == "SKU_01") & (x.time_idx > 3))]
        TimeSeriesDataSet(
            test_data,
            time_idx="time_idx",
            target="volume",
            group_ids=["agency", "sku"],
            max_encoder_length=5,
            max_prediction_length=2,
            min_prediction_length=1,
            min_encoder_length=5,
        )


def test_categorical_target(test_data):
    dataset = TimeSeriesDataSet(
        test_data,
        time_idx="time_idx",
        target="agency",
        group_ids=["agency", "sku"],
        max_encoder_length=5,
        max_prediction_length=2,
        min_prediction_length=1,
        min_encoder_length=1,
    )

    _, y = next(iter(dataset.to_dataloader()))
    assert y[0].dtype is torch.long, "target must be of type long"


def test_pickle(test_dataset):
    pickle.dumps(test_dataset)
    pickle.dumps(test_dataset.to_dataloader())


@pytest.mark.parametrize(
    "kwargs",
    [
        {},
        dict(
            target_normalizer=GroupNormalizer(groups=["agency", "sku"], transformation="log1p", scale_by_group=True),
        ),
    ],
)
def test_new_group_ids(test_data, kwargs):
    """Test for new group ids in dataset"""
    train_agency = test_data["agency"].iloc[0]
    train_dataset = TimeSeriesDataSet(
        test_data[lambda x: x.agency == train_agency],
        time_idx="time_idx",
        target="volume",
        group_ids=["agency", "sku"],
        max_encoder_length=5,
        max_prediction_length=2,
        min_prediction_length=1,
        min_encoder_length=1,
        categorical_encoders=dict(agency=NaNLabelEncoder(add_nan=True), sku=NaNLabelEncoder(add_nan=True)),
        **kwargs,
    )

    # test sampling from training dataset
    next(iter(train_dataset.to_dataloader()))

    # create test dataset with group ids that have not been observed before
    test_dataset = TimeSeriesDataSet.from_dataset(train_dataset, test_data)

    # check that we can iterate through dataset without error
    for _ in iter(test_dataset.to_dataloader()):
        pass


def test_timeseries_columns_naming(test_data):
    with pytest.raises(ValueError):
        TimeSeriesDataSet(
            test_data.rename(columns=dict(agency="agency.2")),
            time_idx="time_idx",
            target="volume",
            group_ids=["agency.2", "sku"],
            max_encoder_length=5,
            max_prediction_length=2,
            min_prediction_length=1,
            min_encoder_length=1,
        )


def test_encoder_normalizer_for_covariates(test_data):
    dataset = TimeSeriesDataSet(
        test_data,
        time_idx="time_idx",
        target="volume",
        group_ids=["agency", "sku"],
        max_encoder_length=5,
        max_prediction_length=2,
        min_prediction_length=1,
        min_encoder_length=1,
        time_varying_known_reals=["price_regular"],
        scalers={"price_regular": EncoderNormalizer()},
    )
    next(iter(dataset.to_dataloader()))


@pytest.mark.parametrize(
    "kwargs",
    [
        {},
        dict(
            target_normalizer=MultiNormalizer(normalizers=[TorchNormalizer(), EncoderNormalizer()]),
        ),
        dict(add_target_scales=True),
        dict(weight="volume"),
    ],
)
def test_multitarget(test_data, kwargs):
    dataset = TimeSeriesDataSet(
        test_data.assign(volume1=lambda x: x.volume),
        time_idx="time_idx",
        target=["volume", "volume1"],
        group_ids=["agency", "sku"],
        max_encoder_length=5,
        max_prediction_length=2,
        min_prediction_length=1,
        min_encoder_length=1,
        time_varying_known_reals=["price_regular"],
        scalers={"price_regular": EncoderNormalizer()},
        **kwargs,
    )
    next(iter(dataset.to_dataloader()))


def test_check_nas(test_data):
    data = test_data.copy()
    data.loc[0, "volume"] = np.nan
    with pytest.raises(ValueError, match=r"1 \(.*infinite"):
        TimeSeriesDataSet(
            data,
            time_idx="time_idx",
            target=["volume"],
            group_ids=["agency", "sku"],
            max_encoder_length=5,
            max_prediction_length=2,
            min_prediction_length=1,
            min_encoder_length=1,
        )


@pytest.mark.parametrize(
    "kwargs",
    [
        dict(target="volume"),
        dict(target="agency", scalers={"volume": EncoderNormalizer()}),
        dict(target="volume", target_normalizer=EncoderNormalizer()),
        dict(target=["volume", "agency"]),
    ],
)
def test_lagged_variables(test_data, kwargs):
    dataset = TimeSeriesDataSet(
        test_data.copy(),
        time_idx="time_idx",
        group_ids=["agency", "sku"],
        max_encoder_length=5,
        max_prediction_length=2,
        min_prediction_length=1,
        min_encoder_length=3,  # one more than max lag for validation
        time_varying_unknown_reals=["volume"],
        time_varying_unknown_categoricals=["agency"],
        lags={"volume": [1, 2], "agency": [1, 2]},
        add_encoder_length=False,
        **kwargs,
    )

    x_all, _ = next(iter(dataset.to_dataloader()))

    for name in ["volume", "agency"]:
        if name in dataset.reals:
            vars = dataset.reals
            x = x_all["encoder_cont"]
        else:
            vars = dataset.flat_categoricals
            x = x_all["encoder_cat"]
        target_idx = vars.index(name)
        for lag in [1, 2]:
            lag_idx = vars.index(f"{name}_lagged_by_{lag}")
            target = x[..., target_idx][:, 0]
            lagged_target = torch.roll(x[..., lag_idx], -lag, dims=1)[:, 0]
            assert torch.isclose(target, lagged_target).all(), "lagged target must be the same as non-lagged target"


@pytest.mark.parametrize(
    "agency,first_prediction_idx,should_raise",
    [("Agency_01", 0, False), ("xxxxx", 0, True), ("Agency_01", 100, True), ("Agency_01", 4, False)],
)
def test_filter_data(test_dataset, agency, first_prediction_idx, should_raise):
    func = lambda x: (x.agency == agency) & (x.time_idx_first_prediction >= first_prediction_idx)
    if should_raise:
        with pytest.raises(ValueError):
            test_dataset.filter(func)
    else:
        filtered_dataset = test_dataset.filter(func)
        assert len(test_dataset.index) > len(
            filtered_dataset.index
        ), "filtered dataset should have less entries than original dataset"
        for x, _ in iter(filtered_dataset.to_dataloader()):
            index = test_dataset.x_to_index(x)
            assert (index["agency"] == agency).all(), "Agency filter has failed"
            assert index["time_idx"].min() == first_prediction_idx, "First prediction filter has failed"