&Faakhir30 [ENH] Update Datamodules according to new API design

[phoeenniixx]

Stacks on #2410 (and #2313 in coming commits FYI @Faakhir30 )
Also see #2407, #2398

This PR is aimed at v2-dev instead of main, so that we dont break current user code until this design is complete and stable. See #2407 for details. So I had to update the branch target in the workflows for the tests to run (which might fail due to the changes)

Stack

• TimeSeries - [ENH] Add TimeSeries datatype as per new API #2410
• Datamodule - &Faakhir30 [ENH] Update Datamodules according to new API design #2423 <- This PR
• Forecaster - [ENH] Add forecaster support and remove _pkg class for new API #2434

Note

The tutorials will fail for sure, will update the notebook in one of the branchs stacked over this PR.
Still need to add the Forecaster and tests - will add the notebook and docs changes on the upcoming branches as the vignette is incomplete without forecaster which is not in this PR.

Changes

• All metadata dataclasses (like EncDecDataModuleMetadata, TimeSeriesMetadata etc) are added to _metadata.py
• a new util for validation checks in utils/_validation
• Update the tests (needed to make some changes like updating the param name) in standalone model tests for v2 as well
  • Also had to add metadata in ignore of save_hyperparameters of Samformer. See this comment
• Updates to datamodules (currently duplicate in both classes)
  • Uses EncDecDataModuleMetadata (in case of EncoderDecoderTImeSeriesDatamodule) and tslib dm has its own metadata dataclass.
  • Data validation.
  • with_data method to create a new dm instance using the already present params (like already fitted scalers - only to create the transformed data).
    • Like creating a datamodule for prediction where we need the scalers from the datamodule we used for fitting.
    • Also useful to create the class "without data" which arrives at forecaster.fit() while we accept the dm class in __init__ of the forecaster. So, the class is actually just empty at this point without data.

[codecov]

Codecov Report

❌ Patch coverage is 98.01587% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.95%. Comparing base (524dcfd) to head (da46e88).

Files with missing lines	Patch %	Lines
...h_forecasting/data/data_module/base_data_module.py	97.76%	3 Missing ⚠️
...g/data/data_module/_encoder_decoder_data_module.py	98.36%	1 Missing ⚠️
...forecasting/data/data_module/_tslib_data_module.py	96.77%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##               ts    #2423      +/-   ##
==========================================
+ Coverage   87.79%   87.95%   +0.15%     
==========================================
  Files         202      203       +1     
  Lines       11292    11373      +81     
==========================================
+ Hits         9914    10003      +89     
+ Misses       1378     1370       -8     

Flag	Coverage Δ
cpu	87.95% <98.01%> (+0.15%)	⬆️
pytest	87.95% <98.01%> (+0.15%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[phoeenniixx]

The datamoudle_pkg classes and updates in test framework as done in #2313 will happen in another PR

[phoeenniixx]

need to add metadata to be ignored to prevent the torch UnpicklingError as metadata is now a dataclass

[CloseChoice]

This looks fine but there are a couple problems, I added them below

[CloseChoice]

here's a problem codex found and I verified:

from pathlib import Path
from tempfile import TemporaryDirectory

import lightning
import torch

from _common import make_module, make_timeseries
from pytorch_forecasting.metrics import MAE
from pytorch_forecasting.models.samformer._samformer_v2 import Samformer


module = make_module(make_timeseries())
model = Samformer(
    loss=MAE(), hidden_size=4, use_revin=False, metadata=module.metadata
)
print("Metadata saved in hyperparameters:", "metadata" in model.hparams)
# Minimal Lightning checkpoint: no training is needed to exercise reconstruction.
checkpoint = {
    "state_dict": model.state_dict(),
    "hyper_parameters": dict(model.hparams),
    "pytorch-lightning_version": lightning.__version__,
}
with TemporaryDirectory(prefix="dm-preview-") as directory:
    path = Path(directory) / "samformer.ckpt"
    torch.save(checkpoint, path)
    Samformer.load_from_checkpoint(path, loss=MAE())

returns an error.

[CloseChoice]

another problem
TorchNormalizer can be trained on one dim and applied to multiple.

"""Show that two-target prediction data can reuse a one-target normalizer."""

import warnings

import numpy as np
import pandas as pd
import torch

from pytorch_forecasting.data.data_module import EncoderDecoderTimeSeriesDataModule
from pytorch_forecasting.data.encoders import TorchNormalizer
from pytorch_forecasting.data.timeseries import TimeSeries


def make_timeseries(offset=0.0, num_groups=10, targets=None):
    groups = np.repeat(np.arange(num_groups), 40)
    times = np.tile(np.arange(40), num_groups)
    frame = pd.DataFrame(
        {
            "group": groups,
            "time": times,
            "target": offset + 10.0 + times + groups,
            "second_target": offset + 100.0 + times * 2 + groups,
            "feature": times.astype(float),
        }
    )
    return TimeSeries(
        frame,
        time="time",
        target=targets or ["target"],
        group=["group"],
        num=["feature"],
        known=["feature"],
    )


def make_module(data=None):
    torch.manual_seed(0)
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message=".*experimental.*")
        return EncoderDecoderTimeSeriesDataModule(
            time_series=data,
            max_encoder_length=10,
            max_prediction_length=5,
            target_normalizer=TorchNormalizer(),
        )


trained = make_module(make_timeseries())
trained.setup("fit")
prediction = trained.with_data(
    make_timeseries(targets=["target", "second_target"])
)
prediction.setup("predict")
x, _ = next(iter(prediction.predict_dataloader()))
print("Training target count:", trained.n_targets)
print("Prediction target count:", prediction.n_targets)
print("Prediction target_past shape:", tuple(x["target_past"].shape))
assert trained.n_targets == 1 and prediction.n_targets == 2
assert prediction._target_normalizer is trained._target_normalizer
print("BUG reproduced: incompatible targets are accepted and share the fitted normalizer.")

[CloseChoice]

also here is another bug, fitting a second config mutates state of the first. I think these should be handled separately.

"""Show that fitting a second configuration copy mutates the first copy."""

import torch

import warnings

import numpy as np
import pandas as pd

from pytorch_forecasting.data.data_module import EncoderDecoderTimeSeriesDataModule
from pytorch_forecasting.data.encoders import TorchNormalizer
from pytorch_forecasting.data.timeseries import TimeSeries


def make_timeseries(offset=0.0, num_groups=10, targets=None):
    groups = np.repeat(np.arange(num_groups), 40)
    times = np.tile(np.arange(40), num_groups)
    frame = pd.DataFrame(
        {
            "group": groups,
            "time": times,
            "target": offset + 10.0 + times + groups,
            "second_target": offset + 100.0 + times * 2 + groups,
            "feature": times.astype(float),
        }
    )
    return TimeSeries(
        frame,
        time="time",
        target=targets or ["target"],
        group=["group"],
        num=["feature"],
        known=["feature"],
    )


def make_module(data=None):
    torch.manual_seed(0)
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message=".*experimental.*")
        return EncoderDecoderTimeSeriesDataModule(
            time_series=data,
            max_encoder_length=10,
            max_prediction_length=5,
            target_normalizer=TorchNormalizer(),
        )


config = make_module()
first = config.with_data(make_timeseries())
first.setup("fit")
center_before = first._target_normalizer._scaler.center_.clone()

second = config.with_data(make_timeseries(offset=1000.0))
second.setup("fit")
center_after = first._target_normalizer._scaler.center_.clone()

print("First module's center before second fit:", center_before.item())
print("First module's center after second fit:", center_after.item())
assert first._target_normalizer is second._target_normalizer
assert not torch.equal(center_before, center_after)
print("BUG reproduced: fitting the second module changes the first module's statistics.")
``

[CloseChoice]

This could be one of the most serious bugs, there is an overlap between train and test overlap:

"""Show that single-series training and evaluation use identical windows."""

import warnings

import numpy as np
import pandas as pd
import torch

from pytorch_forecasting.data.data_module import EncoderDecoderTimeSeriesDataModule
from pytorch_forecasting.data.encoders import TorchNormalizer
from pytorch_forecasting.data.timeseries import TimeSeries


def make_timeseries(offset=0.0, num_groups=10, targets=None):
    groups = np.repeat(np.arange(num_groups), 40)
    times = np.tile(np.arange(40), num_groups)
    frame = pd.DataFrame(
        {
            "group": groups,
            "time": times,
            "target": offset + 10.0 + times + groups,
            "second_target": offset + 100.0 + times * 2 + groups,
            "feature": times.astype(float),
        }
    )
    return TimeSeries(
        frame,
        time="time",
        target=targets or ["target"],
        group=["group"],
        num=["feature"],
        known=["feature"],
    )


def make_module(data=None):
    torch.manual_seed(0)
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message=".*experimental.*")
        return EncoderDecoderTimeSeriesDataModule(
            time_series=data,
            max_encoder_length=10,
            max_prediction_length=5,
            target_normalizer=TorchNormalizer(),
        )


module = make_module(make_timeseries(num_groups=1))
module.setup("fit")
module.setup("test")
print("Train indices:", module._train_indices.tolist())
print("Validation indices:", module._val_indices.tolist())
print("Test indices:", module._test_indices.tolist())
assert module.train_windows
assert module.train_windows == module.val_windows == module.test_windows
print(f"BUG reproduced: all splits contain the same {len(module.train_windows)} windows.")