diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 80dbb6b8..36a50b44 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -18,7 +18,7 @@ jobs: with: python-version: 3.8 - name: Install Black - run: pip install black[jupyter] + run: pip install "black[jupyter]==24.8.0" - name: Run Black run: black --check . diff --git a/.github/workflows/pypi_release.yml b/.github/workflows/pypi_release.yml index 9266f50d..a3e28652 100644 --- a/.github/workflows/pypi_release.yml +++ b/.github/workflows/pypi_release.yml @@ -23,11 +23,11 @@ jobs: - name: Install poetry shell: bash run: | - curl -sSL https://install.python-poetry.org | python3 - + curl -sSL https://install.python-poetry.org | python3 - --version 1.8.5 python -m pip install poetry-dynamic-versioning[plugin] - name: Set poetry path variable - run: echo "/Users/runner/.local/bin" >> $GITHUB_PATH + run: echo "$HOME/.local/bin" >> "$GITHUB_PATH" - name: Build run: | diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index d884f7a2..af19050d 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -21,6 +21,8 @@ jobs: exclude: - python-version: 3.9 tf-version: 2.13.1 + - os: macOS-latest + tf-version: 2.13.1 steps: - uses: actions/checkout@v4 @@ -37,9 +39,9 @@ jobs: - name: Install poetry shell: bash run: | - curl -sSL https://install.python-poetry.org | python3 - + curl -sSL https://install.python-poetry.org | python3 - --version 1.8.5 - name: Set poetry path variable - run: echo "/Users/runner/.local/bin" >> $GITHUB_PATH + run: echo "$HOME/.local/bin" >> "$GITHUB_PATH" - name: Configure poetry shell: bash diff --git a/.gitignore b/.gitignore index 11c5cd53..2ba0b50f 100755 --- a/.gitignore +++ b/.gitignore @@ -96,3 +96,4 @@ coverage.xml !/weights/.gitkeep CLAUDE.md temp/ +*.keras diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 45e54f6a..172be290 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -19,7 +19,7 @@ repos: hooks: - id: isort - repo: https://github.com/psf/black - rev: 22.3.0 + rev: 24.8.0 hooks: - id: black - repo: https://github.com/nbQA-dev/nbQA diff --git a/benchmark/README.md b/benchmark/README.md new file mode 100644 index 00000000..32758be2 --- /dev/null +++ b/benchmark/README.md @@ -0,0 +1,112 @@ +# TFTS Benchmark System + +A flexible benchmarking framework for evaluating TFTS models across multiple datasets with multiple metrics and multiple runs. Designed for reproducibility and paper-quality results. + +## Usage + +### Command-Line + +```bash +python -m benchmark.cli \ + --models rnn transformer dlinear \ + --datasets sine air_passengers \ + --metrics mae rmse mape \ + --runs 3 \ + --epochs 50 \ + --output-dir results/ +``` + +### Python API + +```python +from benchmark import BenchmarkConfig, BenchmarkRunner + +config = BenchmarkConfig( + models=["rnn", "transformer", "dlinear"], + datasets=["sine", "air_passengers"], + metrics=["mae", "rmse", "mape"], + runs=3, + epochs=50, + output_dir="benchmark_results", +) + +runner = BenchmarkRunner(config) +results = runner.run() + +results.print_table() # console +results.to_csv("results.csv") # CSV +results.to_latex("results.tex") # LaTeX for papers +``` + +## Adding a New Dataset + +```python +from benchmark import Dataset +import pandas as pd + +class MyDataset(Dataset): + name = "my_dataset" + description = "Description of my dataset" + train_length = 24 + predict_sequence_length = 8 + + def prepare_data(self, **kwargs): + # Load your data from any source (CSV, DB, API, etc.) + x, y = ... + return x, y + + def get_train_valid_split(self, **kwargs): + x, y = self.prepare_data(**kwargs) + # split into train/valid + return (x_train, y_train), (x_valid, y_valid) +``` + +Then register it: + +```python +from benchmark import BenchmarkRunner, DatasetRegistry + +registry = DatasetRegistry() +registry.register("my_dataset", MyDataset) + +config = BenchmarkConfig(datasets=["my_dataset"], ...) +runner = BenchmarkRunner(config, dataset_registry=registry) +results = runner.run() +``` + +## Architecture + +- **BenchmarkRunner**: Orchestrates running models on datasets, collecting results. +- **Dataset**: Abstract base; each dataset subclass implements `prepare_data()` and returns standardized format. +- **DatasetRegistry**: Maintains a registry of all available datasets. +- **ModelRegistry**: Wraps existing tfts model mapping. +- **BenchmarkMetrics**: Computes standard time-series metrics (MAE, MSE, RMSE, MAPE, etc.) +- **BenchmarkResults**: Formats and exports results (CSV, JSON, LaTeX, console table). + +## Metrics + +Available metrics: +- `mae`: Mean Absolute Error +- `mse`: Mean Squared Error +- `rmse`: Root Mean Squared Error +- `mape`: Mean Absolute Percentage Error +- `smape`: Symmetric MAPE +- `r2`: R-squared + +## Migrated Example Benchmarks + +The previous `examples/benchmarks` tasks are available as registered datasets: + +- `forecasting_sticker_sales` +- `CMI_detect_sleep_states` + +Both support `data_path` overrides through `per_dataset_config`. If the source +CSV is not available, they generate deterministic placeholder data so the +benchmark runner and CLI remain usable without Kaggle downloads. + +## Output Files + +After running, the following files are generated in `output_dir`: +- `results.json`: Raw results for each run. +- `results.csv`: Averaged results (mean/std per model-dataset). +- `results.tex`: LaTeX table for papers. diff --git a/benchmark/__init__.py b/benchmark/__init__.py new file mode 100644 index 00000000..7c9aad3c --- /dev/null +++ b/benchmark/__init__.py @@ -0,0 +1,38 @@ +"""TFTS Benchmark System. + +A flexible benchmarking framework for evaluating TFTS modelsacross multiple +datasets with multiple metrics and multiple runs. Designed for reproducibility +and paper-quality results. + +Example:: + + from benchmark import BenchmarkRunner, BenchmarkConfig + from benchmark.datasets import GrocerysalesDataset, RecruitRestaurantDataset + + config = BenchmarkConfig( + models=["rnn", "transformer", "dlinear"], + datasets=["grocery_sales", "recruit_restaurant"], + metrics=["mae", "rmse", "mape"], + runs=5, + output_dir="benchmark_results", + ) + + runner = BenchmarkRunner(config) + results = runner.run() + results.to_latex("benchmark_results.tex") + results.to_csv("benchmark_results.csv") +""" + +from benchmark.base import BenchmarkConfig, Dataset +from benchmark.formatter import BenchmarkResults +from benchmark.registry import DatasetRegistry, ModelRegistry +from benchmark.runner import BenchmarkRunner + +__all__ = [ + "BenchmarkRunner", + "BenchmarkConfig", + "BenchmarkResults", + "Dataset", + "DatasetRegistry", + "ModelRegistry", +] diff --git a/benchmark/base.py b/benchmark/base.py new file mode 100644 index 00000000..458cc288 --- /dev/null +++ b/benchmark/base.py @@ -0,0 +1,136 @@ +"""Base classes for the TFTS benchmark system.""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +import logging +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import numpy as np +import pandas as pd +import tensorflow as tf + +logger = logging.getLogger(__name__) + + +@dataclass +class BenchmarkConfig: + """Configuration for a benchmark run. + + Attributes: + models: List of model names to evaluate (e.g., ["rnn", "transformer"]). + Use ``["all"]`` to run all registered models. + datasets: List of dataset names to evaluate (e.g., ["sine", "grocery_sales"]). + Use ``["all"]`` to run all registered datasets. + metrics: List of metric names to compute. + Available: ``"mae"``, ``"mse"``, ``"rmse"``, ``"mape"``, ``"smape"``, ``"r2"``. + runs: Number of runs per model-dataset pair (for statistical significance). + epochs: Number of training epochs per run. + batch_size: Batch size for training. + learning_rate: Learning rate for the optimizer. + train_length: Lookback window length. If None, each dataset provides its own. + predict_sequence_length: Forecast horizon. If None, each dataset provides its own. + seed: Base seed. Each run uses ``seed + run_idx`` for reproducibility. + output_dir: Directory to save results. + save_models: Whether to save trained model weights. + verbose: Verbosity level (0=silent, 1=progress, 2=detailed). + device: Device to run on (e.g., ``"/gpu:0"`` or ``"/cpu:0"``). + per_dataset_config: Optional per-dataset configuration overrides. + Keys are dataset names, values are dicts with keys like ``train_length``, + ``predict_sequence_length``, ``epochs``, etc. + """ + + models: List[str] = field(default_factory=lambda: ["all"]) + datasets: List[str] = field(default_factory=lambda: ["all"]) + metrics: List[str] = field(default_factory=lambda: ["mae", "rmse", "mape"]) + runs: int = 1 + epochs: int = 50 + batch_size: int = 32 + learning_rate: float = 1e-3 + train_length: Optional[int] = None + predict_sequence_length: Optional[int] = None + seed: int = 42 + output_dir: str = "benchmark_results" + save_models: bool = False + verbose: int = 1 + device: str = "" + per_dataset_config: Dict[str, Dict[str, Any]] = field(default_factory=dict) + + def __post_init__(self): + if self.runs < 1: + raise ValueError("runs must be >= 1") + if self.epochs < 1: + raise ValueError("epochs must be >= 1") + + def get_dataset_config(self, dataset_name: str) -> Dict[str, Any]: + """Get configuration for a specific dataset, merging with defaults.""" + config = { + "epochs": self.epochs, + "batch_size": self.batch_size, + "learning_rate": self.learning_rate, + "train_length": self.train_length, + "predict_sequence_length": self.predict_sequence_length, + } + if dataset_name in self.per_dataset_config: + config.update(self.per_dataset_config[dataset_name]) + return config + + +class Dataset(ABC): + """Abstract base class for benchmark datasets. + + Each dataset subclass implements ``prepare_data()`` to load and format + the data. The class also provides metadata about the dataset. + + Attributes: + name: Unique identifier for the dataset. + description: Human-readable description. + train_length: Default lookback window length (can be overridden by config). + predict_sequence_length: Default forecast horizon (can be overridden by config). + """ + + name: str = "" + description: str = "" + train_length: int = 24 + predict_sequence_length: int = 8 + num_features: int = 1 + is_multivariate: bool = False + is_grouped: bool = False + target_column: str = "target" + time_column: str = "time" + group_column: Optional[str] = None + + @abstractmethod + def prepare_data(self, **kwargs) -> Union[ + Tuple[np.ndarray, np.ndarray], + Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]], + tf.data.Dataset, + ]: + """Prepare and return the dataset. + + Returns: + Either (x_train, y_train), ((x_train, y_train), (x_valid, y_valid)), + or a tf.data.Dataset. + """ + raise NotImplementedError + + @abstractmethod + def get_train_valid_split(self, **kwargs) -> Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: + """Return train and validation splits. + + Returns: + Tuple of (x_train, y_train), (x_valid, y_valid). + """ + raise NotImplementedError + + @classmethod + def list_params(cls) -> Dict[str, Any]: + """Return dataset parameters for display.""" + return { + "name": cls.name, + "description": cls.description, + "train_length": cls.train_length, + "predict_sequence_length": cls.predict_sequence_length, + "num_features": cls.num_features, + "is_multivariate": cls.is_multivariate, + "is_grouped": cls.is_grouped, + } diff --git a/benchmark/cli.py b/benchmark/cli.py new file mode 100644 index 00000000..43bb2434 --- /dev/null +++ b/benchmark/cli.py @@ -0,0 +1,196 @@ +"""Command-line interface for the TFTS benchmark system. + +Usage:: + + python -m benchmark.cli \ + --models rnn transformer dlinear \ + --datasets sine air_passengers \ + --metrics mae rmse mape \ + --runs 3 \ + --epochs 50 \ + --output-dir benchmark_results +""" + +import argparse +import logging +import sys +from typing import List + +from benchmark import BenchmarkConfig, BenchmarkRunner +from benchmark.registry import DatasetRegistry, ModelRegistry + +# Set up logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger("tfts.benchmark") + + +def get_parser() -> argparse.ArgumentParser: + """Build the argument parser.""" + parser = argparse.ArgumentParser( + prog="tfts-benchmark", + description="TFTS Benchmark: Run multiple models on multiple datasets.", + ) + parser.add_argument( + "--models", + nargs="+", + default=["all"], + help="Model names to benchmark (default: all). Use 'all' for every registered model.", + ) + parser.add_argument( + "--datasets", + nargs="+", + default=["all"], + help="Dataset names to benchmark (default: all). Use 'all' for every registered dataset.", + ) + parser.add_argument( + "--metrics", + nargs="+", + default=["mae", "rmse", "mape"], + choices=["mae", "mse", "rmse", "mape", "smape", "r2", "mape_pct"], + help="Metrics to compute.", + ) + parser.add_argument( + "--runs", + type=int, + default=1, + help="Number of runs per model-dataset pair (for statistical significance).", + ) + parser.add_argument( + "--epochs", + type=int, + default=50, + help="Number of training epochs.", + ) + parser.add_argument( + "--batch-size", + type=int, + default=32, + help="Batch size for training.", + ) + parser.add_argument( + "--learning-rate", + type=float, + default=1e-3, + help="Learning rate.", + ) + parser.add_argument( + "--train-length", + type=int, + default=None, + help="Lookback window. If not set, each dataset uses its default.", + ) + parser.add_argument( + "--predict-sequence-length", + type=int, + default=None, + help="Forecast horizon. If not set, each dataset uses its default.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Base random seed.", + ) + parser.add_argument( + "--output-dir", + type=str, + default="benchmark_results", + help="Directory to save results.", + ) + parser.add_argument( + "--verbose", + type=int, + choices=[0, 1, 2], + default=1, + help="Verbosity level (0=silent, 1=progress, 2=detailed).", + ) + parser.add_argument( + "--latex", + action="store_true", + help="Also generate a LaTeX table.", + ) + parser.add_argument( + "--list-models", + action="store_true", + help="List available models and exit.", + ) + parser.add_argument( + "--list-datasets", + action="store_true", + help="List available datasets and exit.", + ) + return parser + + +def main(argv: List[str] = None) -> int: + """Main entry point.""" + parser = get_parser() + args = parser.parse_args(argv) + + model_registry = ModelRegistry() + dataset_registry = DatasetRegistry() + # Load default datasets + from benchmark.datasets import ( + AirPassengersDataset, + CMIDetectSleepStatesDataset, + ForecastingStickerSalesDataset, + GrocerysalesDataset, + RecruitRestaurantDataset, + SineDataset, + ) + + dataset_registry.register("sine", SineDataset) + dataset_registry.register("air_passengers", AirPassengersDataset) + dataset_registry.register("grocery_sales", GrocerysalesDataset) + dataset_registry.register("recruit_restaurant", RecruitRestaurantDataset) + dataset_registry.register("forecasting_sticker_sales", ForecastingStickerSalesDataset) + dataset_registry.register("CMI_detect_sleep_states", CMIDetectSleepStatesDataset) + + if args.list_models: + print("Available models:") + for name in model_registry.available_models: + print(" " + f"- {name}") + return 0 + + if args.list_datasets: + print("Available datasets:") + for name in dataset_registry.list_datasets(): + print(" " + f"- {name}") + return 0 + + config = BenchmarkConfig( + models=args.models, + datasets=args.datasets, + metrics=args.metrics, + runs=args.runs, + epochs=args.epochs, + batch_size=args.batch_size, + learning_rate=args.learning_rate, + train_length=args.train_length, + predict_sequence_length=args.predict_sequence_length, + seed=args.seed, + output_dir=args.output_dir, + verbose=args.verbose, + ) + + runner = BenchmarkRunner(config, dataset_registry, model_registry) + results = runner.run() + + # Print to console + results.print_table() + + if args.latex: + import os + + latex_path = os.path.join(args.output_dir, "results.tex") + results.to_latex(latex_path) + print(f"\nLaTeX table saved to: {latex_path}") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmark/datasets/__init__.py b/benchmark/datasets/__init__.py new file mode 100644 index 00000000..ba9b0f0a --- /dev/null +++ b/benchmark/datasets/__init__.py @@ -0,0 +1,15 @@ +"""Built-in datasets for the TFTS benchmark system.""" + +from benchmark.datasets.grocery_sales import GrocerysalesDataset +from benchmark.datasets.legacy_examples import CMIDetectSleepStatesDataset, ForecastingStickerSalesDataset +from benchmark.datasets.recruit_restaurant import RecruitRestaurantDataset +from benchmark.datasets.synthetic import AirPassengersDataset, SineDataset + +__all__ = [ + "SineDataset", + "AirPassengersDataset", + "GrocerysalesDataset", + "RecruitRestaurantDataset", + "ForecastingStickerSalesDataset", + "CMIDetectSleepStatesDataset", +] diff --git a/benchmark/datasets/base.py b/benchmark/datasets/base.py new file mode 100644 index 00000000..3786b99b --- /dev/null +++ b/benchmark/datasets/base.py @@ -0,0 +1,28 @@ +"""Base dataset utilities for the benchmark system.""" + +import logging +from typing import Tuple + +import numpy as np + +logger = logging.getLogger(__name__) + + +def split_train_valid( + x: np.ndarray, + y: np.ndarray, + test_size: float = 0.2, +) -> Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: + """Split arrays into train/validation. + + Args: + x: Input array. + y: Target array. + test_size: Fraction of data to use for validation. + + Returns: + (x_train, y_train), (x_valid, y_valid). + """ + n = len(x) + split_idx = int(n * (1 - test_size)) + return (x[:split_idx], y[:split_idx]), (x[split_idx:], y[split_idx:]) diff --git a/benchmark/datasets/grocery_sales.py b/benchmark/datasets/grocery_sales.py new file mode 100644 index 00000000..f617b224 --- /dev/null +++ b/benchmark/datasets/grocery_sales.py @@ -0,0 +1,117 @@ +"""Kaggle Favorita grocery sales dataset for TFTS benchmark. + +Note: + This is a concrete example of how to add a real-world Kaggle dataset. + Users need to download the data from Kaggle and update ``data_path``. +""" + +import logging +import os +from typing import Any, Dict, Tuple + +import numpy as np +import pandas as pd + +from benchmark.base import Dataset +from benchmark.datasets.base import split_train_valid + +logger = logging.getLogger(__name__) + + +class GrocerysalesDataset(Dataset): + """Kaggle Favorita Grocery Sales dataset. + + Expected columns after preprocessing: + - ``date`` (datetime) + - ``store_nbr`` (int, group column) + - ``item_nbr`` (int, group column) + - ``unit_sales`` (float, target) + - additional feature columns (e.g., ``onpromotion``) + + Example: + >>> dataset = GrocerysalesDataset(data_path="/path/to/favorita") + >>> (x_train, y_train), (x_valid, y_valid) = dataset.get_train_valid_split() + """ + + name = "grocery_sales" + description = "Kaggle Favorita Grocery Sales Forecasting" + train_length = 28 + predict_sequence_length = 14 + num_features = 1 + is_grouped = True + target_column = "unit_sales" + time_column = "date" + group_column = ["store_nbr", "item_nbr"] + + def __init__(self, data_path: str = "", **kwargs): + self.data_path = data_path or os.environ.get("FAVORITA_PATH", "") + super().__init__() + + def _load_and_preprocess(self) -> pd.DataFrame: + """Load and preprocess the raw Favorita data. + + Returns: + Cleaned DataFrame ready for sequence generation. + """ + if not self.data_path: + logger.warning("No data_path provided for GrocerysalesDataset. " "Returning synthetic placeholder data.") + return self._synthetic_placeholder() + + # --- Stub for actual data loading --- + # Users should replace this with their actual preprocessing pipeline. + # Example: + # df = pd.read_csv(os.path.join(self.data_path, "train.csv")) + # ... preprocess ... + # return df + return self._synthetic_placeholder() + + def _synthetic_placeholder(self) -> pd.DataFrame: + """Generate a small synthetic placeholder when real data is missing.""" + logger.warning("Using synthetic placeholder for GrocerysalesDataset") + np.random.seed(42) + n_stores = 5 + n_items = 3 + days = 365 + records = [] + for store in range(n_stores): + for item in range(n_items): + trend = np.linspace(10, 20, days) + np.random.normal(0, 2, days) + seasonal = 5 * np.sin(2 * np.pi * np.arange(days) / 7) + for day, val in enumerate(trend + seasonal): + records.append( + { + "date": pd.Timestamp("2020-01-01") + pd.Timedelta(days=day), + "store_nbr": store, + "item_nbr": item, + "unit_sales": val, + "onpromotion": 0, + } + ) + return pd.DataFrame(records) + + def _generate_sequences(self, df: pd.DataFrame) -> Tuple[np.ndarray, np.ndarray]: + """Generate sliding-window sequences from a DataFrame. + + Returns: + x, y arrays where x.shape == (samples, train_length, 1), + y.shape == (samples, predict_sequence_length, 1). + """ + x_list, y_list = [], [] + for (store, item), group in df.groupby(["store_nbr", "item_nbr"]): + group = group.sort_values("date") + values = group["unit_sales"].values.astype(np.float32) + n = len(values) + train_len = self.train_length + pred_len = self.predict_sequence_length + for i in range(n - train_len - pred_len + 1): + x_list.append(values[i : i + train_len].reshape(-1, 1)) + y_list.append(values[i + train_len : i + train_len + pred_len].reshape(-1, 1)) + return np.array(x_list), np.array(y_list) + + def prepare_data(self, **kwargs) -> Tuple[np.ndarray, np.ndarray]: + df = self._load_and_preprocess() + return self._generate_sequences(df) + + def get_train_valid_split(self, **kwargs) -> Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: + x, y = self.prepare_data(**kwargs) + return split_train_valid(x, y, test_size=kwargs.get("test_size", 0.2)) diff --git a/benchmark/datasets/legacy_examples.py b/benchmark/datasets/legacy_examples.py new file mode 100644 index 00000000..dd21c760 --- /dev/null +++ b/benchmark/datasets/legacy_examples.py @@ -0,0 +1,227 @@ +"""Datasets migrated from the legacy ``examples/benchmarks`` folder.""" + +import logging +import os +from typing import List, Sequence, Tuple + +import numpy as np +import pandas as pd + +from benchmark.base import Dataset +from benchmark.datasets.base import split_train_valid + +logger = logging.getLogger(__name__) + + +def _resolve_data_path(value: str, env_name: str) -> str: + path = value or os.environ.get(env_name, "") + if path and os.path.isdir(path): + return os.path.join(path, "train.csv") + return path + + +def _window_arrays( + df: pd.DataFrame, + group_columns: Sequence[str], + feature_columns: Sequence[str], + target_column: str, + time_column: str, + train_length: int, + predict_sequence_length: int, +) -> Tuple[np.ndarray, np.ndarray]: + x_list: List[np.ndarray] = [] + y_list: List[np.ndarray] = [] + group_key = list(group_columns) if len(group_columns) > 1 else group_columns[0] + + for _, group in df.groupby(group_key, sort=False): + group = group.sort_values(time_column) + features = group[list(feature_columns)].to_numpy(dtype=np.float32) + target = group[target_column].to_numpy(dtype=np.float32).reshape(-1, 1) + limit = len(group) - train_length - predict_sequence_length + 1 + for start in range(max(0, limit)): + x_list.append(features[start : start + train_length]) + y_list.append(target[start + train_length : start + train_length + predict_sequence_length]) + + if not x_list: + raise ValueError( + "No benchmark windows were generated. " + "Reduce train_length/predict_sequence_length or provide a longer dataset." + ) + return np.asarray(x_list, dtype=np.float32), np.asarray(y_list, dtype=np.float32) + + +class ForecastingStickerSalesDataset(Dataset): + """Kaggle Playground sticker sales forecasting benchmark. + + This replaces ``examples/benchmarks/forecasting_sticker_sales`` with a + standard benchmark dataset. Provide ``data_path`` as a CSV file or directory + containing ``train.csv``. Expected Kaggle columns are ``date``, ``country``, + ``store``, ``product`` and ``num_sold``. + """ + + name = "forecasting_sticker_sales" + description = "Kaggle Playground sticker sales forecasting" + train_length = 144 + predict_sequence_length = 32 + num_features = 4 + is_grouped = True + target_column = "num_sold" + time_column = "date" + group_column = ["country", "store", "product"] + + def _load_dataframe(self, data_path: str = "", **kwargs) -> pd.DataFrame: + path = _resolve_data_path(data_path, "STICKER_SALES_PATH") + if path and os.path.exists(path): + return pd.read_csv(path) + + logger.warning("Sticker sales data not found. Using deterministic synthetic placeholder data.") + rng = np.random.default_rng(kwargs.get("seed", 42)) + dates = pd.date_range("2017-01-01", periods=240, freq="D") + rows = [] + for country in ["Canada", "Finland", "Kenya"]: + for store in ["KaggleMart", "KaggleRama"]: + for product in ["Sticker A", "Sticker B"]: + base = 120 + 20 * (country == "Canada") + 12 * (store == "KaggleRama") + product_shift = 15 * (product == "Sticker B") + seasonal = 18 * np.sin(2 * np.pi * np.arange(len(dates)) / 7) + trend = np.linspace(0, 30, len(dates)) + noise = rng.normal(0, 4, len(dates)) + values = np.maximum(1, base + product_shift + seasonal + trend + noise) + for date, value in zip(dates, values): + rows.append( + { + "date": date, + "country": country, + "store": store, + "product": product, + "num_sold": value, + } + ) + return pd.DataFrame(rows) + + def _preprocess(self, df: pd.DataFrame) -> pd.DataFrame: + df = df.copy() + df["date"] = pd.to_datetime(df["date"]) + df["num_sold"] = pd.to_numeric(df["num_sold"], errors="coerce") + df = df.dropna(subset=["date", "num_sold"]) + + group_cols = ["country", "store", "product"] + for column in group_cols: + if column not in df: + df[column] = "series" + + dayofweek = df["date"].dt.dayofweek.astype(np.float32) + dayofyear = df["date"].dt.dayofyear.astype(np.float32) + df["target_scaled"] = df.groupby(group_cols)["num_sold"].transform( + lambda values: (values - values.mean()) / (values.std(ddof=0) + 1e-6) + ) + df["dow_sin"] = np.sin(2 * np.pi * dayofweek / 7.0) + df["dow_cos"] = np.cos(2 * np.pi * dayofweek / 7.0) + df["doy_sin"] = np.sin(2 * np.pi * dayofyear / 365.25) + return df + + def prepare_data(self, **kwargs) -> Tuple[np.ndarray, np.ndarray]: + train_length = kwargs.get("train_length") or self.train_length + predict_length = kwargs.get("predict_sequence_length") or self.predict_sequence_length + df = self._preprocess(self._load_dataframe(**kwargs)) + return _window_arrays( + df=df, + group_columns=["country", "store", "product"], + feature_columns=["target_scaled", "dow_sin", "dow_cos", "doy_sin"], + target_column="target_scaled", + time_column="date", + train_length=train_length, + predict_sequence_length=predict_length, + ) + + def get_train_valid_split(self, **kwargs) -> Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: + x, y = self.prepare_data(**kwargs) + return split_train_valid(x, y, test_size=kwargs.get("test_size", 0.2)) + + +class CMIDetectSleepStatesDataset(Dataset): + """CMI detect sleep states sequence benchmark. + + The legacy folder only provided a benchmark slot. This adapter supports + common CMI-style CSV columns (``series_id``, ``step``/``timestamp``, + ``anglez``, ``enmo`` and ``awake``/``target``) and falls back to synthetic + sleep-state sequences when no local data is supplied. + """ + + name = "CMI_detect_sleep_states" + description = "Child Mind Institute detect sleep states sequence benchmark" + train_length = 144 + predict_sequence_length = 32 + num_features = 3 + is_grouped = True + target_column = "awake" + time_column = "step" + group_column = "series_id" + + def _load_dataframe(self, data_path: str = "", **kwargs) -> pd.DataFrame: + path = _resolve_data_path(data_path, "CMI_SLEEP_STATES_PATH") + if path and os.path.exists(path): + return pd.read_csv(path) + + logger.warning("CMI sleep states data not found. Using deterministic synthetic placeholder data.") + rng = np.random.default_rng(kwargs.get("seed", 42)) + rows = [] + steps = np.arange(360) + for series_idx in range(8): + phase = series_idx * 0.35 + circadian = np.sin(2 * np.pi * steps / 96 + phase) + awake = (circadian > -0.15).astype(np.float32) + enmo = np.maximum(0, 0.04 + 0.18 * awake + rng.normal(0, 0.015, len(steps))) + anglez = 20 * np.sin(2 * np.pi * steps / 48 + phase) + rng.normal(0, 3, len(steps)) + for step, angle, motion, state in zip(steps, anglez, enmo, awake): + rows.append( + { + "series_id": f"series_{series_idx}", + "step": step, + "anglez": angle, + "enmo": motion, + "awake": state, + } + ) + return pd.DataFrame(rows) + + def _preprocess(self, df: pd.DataFrame) -> pd.DataFrame: + df = df.copy() + if "series_id" not in df: + df["series_id"] = "series" + if "step" not in df: + if "timestamp" in df: + df["step"] = pd.to_datetime(df["timestamp"]).astype("int64") // 10**9 + else: + df["step"] = df.groupby("series_id").cumcount() + + target_column = "awake" if "awake" in df else "target" + if target_column not in df: + raise ValueError("CMI sleep states data must include an 'awake' or 'target' column.") + + for column in ["anglez", "enmo"]: + if column not in df: + df[column] = 0.0 + + df["awake"] = pd.to_numeric(df[target_column], errors="coerce").fillna(0).astype(np.float32) + df["anglez"] = pd.to_numeric(df["anglez"], errors="coerce").fillna(0).astype(np.float32) / 90.0 + df["enmo"] = pd.to_numeric(df["enmo"], errors="coerce").fillna(0).astype(np.float32) + return df + + def prepare_data(self, **kwargs) -> Tuple[np.ndarray, np.ndarray]: + train_length = kwargs.get("train_length") or self.train_length + predict_length = kwargs.get("predict_sequence_length") or self.predict_sequence_length + df = self._preprocess(self._load_dataframe(**kwargs)) + return _window_arrays( + df=df, + group_columns=["series_id"], + feature_columns=["anglez", "enmo", "awake"], + target_column="awake", + time_column="step", + train_length=train_length, + predict_sequence_length=predict_length, + ) + + def get_train_valid_split(self, **kwargs) -> Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: + x, y = self.prepare_data(**kwargs) + return split_train_valid(x, y, test_size=kwargs.get("test_size", 0.2)) diff --git a/benchmark/datasets/recruit_restaurant.py b/benchmark/datasets/recruit_restaurant.py new file mode 100644 index 00000000..2092935c --- /dev/null +++ b/benchmark/datasets/recruit_restaurant.py @@ -0,0 +1,112 @@ +"""Kaggle Recruit Restaurant Forecast dataset for TFTS benchmark. + +Note: + This is a concrete example of how to add a real-world Kaggle dataset. + Users need to download the data from Kaggle and update ``data_path``. +""" + +import logging +import os +from typing import Any, Tuple + +import numpy as np +import pandas as pd + +from benchmark.base import Dataset +from benchmark.datasets.base import split_train_valid + +logger = logging.getLogger(__name__) + + +class RecruitRestaurantDataset(Dataset): + """Kaggle Recruit Restaurant Forecast dataset. + + Expected columns after preprocessing: + - ``visit_date`` (datetime) + - ``id`` (str, group column) + - ``visitors`` (float, target) + - additional features: ``genre_name``, ``area_name``, etc. + + Example: + >>> dataset = RecruitRestaurantDataset(data_path="/path/to/recruit") + >>> (x_train, y_train), (x_valid, y_valid) = dataset.get_train_valid_split() + """ + + name = "recruit_restaurant" + description = "Kaggle Recruit Restaurant Forecasting" + train_length = 28 + predict_sequence_length = 14 + num_features = 1 + is_grouped = True + target_column = "visitors" + time_column = "visit_date" + group_column = "id" + + def __init__(self, data_path: str = "", **kwargs): + self.data_path = data_path or os.environ.get("RECRUIT_PATH", "") + super().__init__() + + def _load_and_preprocess(self) -> pd.DataFrame: + """Load and preprocess the raw recruit data. + + Returns: + Cleaned DataFrame ready for sequence generation. + """ + if not self.data_path: + logger.warning( + "No data_path provided for RecruitRestaurantDataset. " "Returning synthetic placeholder data." + ) + return self._synthetic_placeholder() + + # --- Stub for actual data loading --- + # Users should replace this with their actual preprocessing pipeline. + return self._synthetic_placeholder() + + def _synthetic_placeholder(self) -> pd.DataFrame: + """Generate a small synthetic placeholder when real data is missing.""" + logger.warning("Using synthetic placeholder for RecruitRestaurantDataset") + np.random.seed(42) + n_restaurants = 5 + days = 365 + records = [] + for rid in range(n_restaurants): + trend = np.linspace(20, 40, days) + np.random.normal(0, 5, days) + seasonal = 10 * np.sin(2 * np.pi * np.arange(days) / 7) + for day, val in enumerate(trend + seasonal): + records.append( + { + "visit_date": pd.Timestamp("2020-01-01") + pd.Timedelta(days=day), + "id": f"restaurant_{rid}", + "visitors": max(0, val), + "genre_name": np.random.choice(["Italian", "Japanese", "French"]), + "area_name": np.random.choice(["Shibuya", "Shinjuku", "Ginza"]), + } + ) + return pd.DataFrame(records) + + def _generate_sequences(self, df: pd.DataFrame) -> Tuple[np.ndarray, np.ndarray]: + """Generate sliding-window sequences from a DataFrame. + + Returns: + x, y arrays where x.shape == (samples, train_length, 1), + y.shape == (samples, predict_sequence_length, 1). + """ + x_list, y_list = [], [] + for rid, group in df.groupby("id"): + group = group.sort_values("visit_date") + values = group["visitors"].values.astype(np.float32) + n = len(values) + train_len = self.train_length + pred_len = self.predict_sequence_length + for i in range(n - train_len - pred_len + 1): + x_list.append(values[i : i + train_len].reshape(-1, 1)) + y_list.append(values[i + train_len : i + train_len + pred_len].reshape(-1, 1)) + return np.array(x_list), np.array(y_list) + + def prepare_data(self, **kwargs) -> Tuple[np.ndarray, np.ndarray]: + df = self._load_and_preprocess() + return self._generate_sequences(df) + + def get_train_valid_split(self, **kwargs) -> Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: + x, y = self.prepare_data(**kwargs) + return split_train_valid(x, y, test_size=kwargs.get("test_size", 0.2)) diff --git a/benchmark/datasets/synthetic.py b/benchmark/datasets/synthetic.py new file mode 100644 index 00000000..20bc2327 --- /dev/null +++ b/benchmark/datasets/synthetic.py @@ -0,0 +1,68 @@ +"""Synthetic benchmark datasets.""" + +import random +from typing import Any, Dict, Tuple + +import numpy as np + +from benchmark.base import Dataset +from benchmark.datasets.base import split_train_valid +from tfts.data.get_data import get_air_passengers, get_sine + + +class SineDataset(Dataset): + """Synthetic sine wave benchmark dataset.""" + + name = "sine" + description = "Synthetic sine wave data." + train_length = 24 + predict_sequence_length = 8 + num_features = 1 + + def prepare_data(self, **kwargs) -> Tuple[np.ndarray, np.ndarray]: + x, y = get_sine( + train_sequence_length=self.train_length, + predict_sequence_length=self.predict_sequence_length, + test_size=0.0, + n_examples=kwargs.get("n_examples", 100), + ) + return x, y + + def get_train_valid_split(self, **kwargs) -> Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: + train_len = kwargs.get("train_length") or self.train_length + pred_len = kwargs.get("predict_sequence_length") or self.predict_sequence_length + (x_train, y_train), (x_valid, y_valid) = get_sine( + train_sequence_length=train_len, + predict_sequence_length=pred_len, + test_size=kwargs.get("test_size", 0.2), + n_examples=kwargs.get("n_examples", 100), + ) + return (x_train, y_train), (x_valid, y_valid) + + +class AirPassengersDataset(Dataset): + """Air passengers benchmark dataset.""" + + name = "air_passengers" + description = "Airline passenger counts (1949-1960)." + train_length = 24 + predict_sequence_length = 8 + num_features = 1 + + def prepare_data(self, **kwargs) -> Tuple[np.ndarray, np.ndarray]: + x, y = get_air_passengers( + train_sequence_length=self.train_length, + predict_sequence_length=self.predict_sequence_length, + test_size=0.0, + ) + return x, y + + def get_train_valid_split(self, **kwargs) -> Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: + train_len = kwargs.get("train_length") or self.train_length + pred_len = kwargs.get("predict_sequence_length") or self.predict_sequence_length + (x_train, y_train), (x_valid, y_valid) = get_air_passengers( + train_sequence_length=train_len, + predict_sequence_length=pred_len, + test_size=kwargs.get("test_size", 0.2), + ) + return (x_train, y_train), (x_valid, y_valid) diff --git a/benchmark/formatter.py b/benchmark/formatter.py new file mode 100644 index 00000000..479b6aa9 --- /dev/null +++ b/benchmark/formatter.py @@ -0,0 +1,190 @@ +"""Result formatting for the TFTS benchmark system. + +Supports console tables, CSV, JSON, and LaTeX output for papers.""" + +import csv +import json +import logging +import os +from typing import Any, Dict, List, Optional + +import numpy as np + +logger = logging.getLogger(__name__) + + +def _format_value(val: Any) -> str: + """Pretty formatter for values.""" + if isinstance(val, float): + return f"{val:.4f}" + return str(val) + + +def _avg(values: List[float]) -> float: + """Return mean, ignoring NaNs.""" + arr = np.array(values) + arr = arr[~np.isnan(arr)] + return float(np.mean(arr)) if len(arr) > 0 else float("nan") + + +def _std(values: List[float]) -> float: + """Return std dev, ignoring NaNs.""" + arr = np.array(values) + arr = arr[~np.isnan(arr)] + return float(np.std(arr)) if len(arr) > 0 else float("nan") + + +class BenchmarkResults: + """Container for benchmark results with export helpers. + + Attributes: + results: Raw list of per-run result dictionaries. + """ + + def __init__(self, results: List[Dict[str, Any]]): + self.results = results + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _pivot(self) -> Dict[str, Dict[str, Dict[str, List[float]]]]: + """Pivot raw results into {dataset: {model: {metric: [values]}}}.""" + out: Dict[str, Dict[str, Dict[str, List[float]]]] = {} + for row in self.results: + ds_name = row.get("dataset", "unknown") + model_name = row.get("model", "unknown") + metrics = row.get("metrics", {}) + if ds_name not in out: + out[ds_name] = {} + if model_name not in out[ds_name]: + out[ds_name][model_name] = {} + for metric, value in metrics.items(): + if metric not in out[ds_name][model_name]: + out[ds_name][model_name][metric] = [] + try: + out[ds_name][model_name][metric].append(float(value)) + except (TypeError, ValueError): + pass + return out + + # ------------------------------------------------------------------ + # Public export API + # ------------------------------------------------------------------ + + def to_dataframe(self): + """Return a pandas DataFrame with averaged results (mean std columns).""" + try: + import pandas as pd + except ImportError: + raise ImportError("pandas is required for to_dataframe()") + + rows = [] + for ds_name, models in self._pivot().items(): + for model_name, metrics in models.items(): + flat: Dict[str, Any] = {"dataset": ds_name, "model": model_name} + for metric_name, values in metrics.items(): + flat[f"{metric_name}_mean"] = _avg(values) + flat[f"{metric_name}_std"] = _std(values) + rows.append(flat) + return pd.DataFrame(rows) + + def to_csv(self, path: str) -> None: + """Export results to a CSV file.""" + try: + import pandas as pd + + self.to_dataframe().to_csv(path, index=False) + logger.info("Results saved to %s", path) + except ImportError: + # Fallback with csv module + _dicts = [dict(r) for r in self.results] + if not _dicts: + return + keys = _dicts[0].keys() + with open(path, "w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=keys) + writer.writeheader() + writer.writerows(_dicts) + logger.info("Results saved to %s", path) + + def to_json(self, path: str) -> None: + """Export raw results to a JSON file.""" + with open(path, "w", encoding="utf-8") as fh: + json.dump(self.results, fh, indent=2, default=str) + logger.info("Results saved to %s", path) + + def to_latex(self, path: str, metric: Optional[str] = None) -> None: + """Export averaged results to a LaTeX table suitable for papers. + + Args: + path: Output .tex file path. + metric: If given, only include this metric. Otherwise include all. + """ + lines: List[str] = [ + r"\begin{table}[ht]", + r"\centering", + r"\caption{Benchmark Results}", + r"\label{tab:benchmark}", + ] + + pivot = self._pivot() + first_dataset = next(iter(pivot)) if pivot else None + if first_dataset is None: + lines.extend([r"\begin{tabular}{c}", "No results", r"\end{tabular}", r"\end{table}"]) + with open(path, "w", encoding="utf-8") as fh: + fh.write("\n".join(lines) + "\n") + return + + models = sorted(next(iter(pivot.values())).keys()) + metrics = sorted(next(iter(next(iter(pivot.values())).values())).keys()) + target_metrics = [metric] if metric else metrics + + # Build header + header = "Dataset & Model & " + " & ".join(target_metrics) + r" \\\\" + lines.append(r"\begin{tabular}{ll" + "r" * len(target_metrics) + "}") + lines.append(r"\toprule") + lines.append(header) + lines.append(r"\midrule") + + for ds_name, model_data in pivot.items(): + for model_name in models: + vals = model_data.get(model_name, {}) + row_vals = [] + for m in target_metrics: + values = vals.get(m, []) + if values: + mean = _avg(values) + std = _std(values) + row_vals.append(f"{mean:.4f} $\\pm$ {std:.4f}") + else: + row_vals.append("-") + lines.append(f"{ds_name} & {model_name} & " + " & ".join(row_vals) + r" \\\\") + + lines.extend([r"\bottomrule", r"\end{tabular}", r"\end{table}"]) + + with open(path, "w", encoding="utf-8") as fh: + fh.write("\n".join(lines) + "\n") + + logger.info("LaTeX table saved to %s", path) + + def print_table(self) -> None: + """Print results to the console as a formatted table.""" + pivot = self._pivot() + if not pivot: + print("No results to display.") + return + + for ds_name, models in pivot.items(): + print(f"\n{'=' * 60}") + print(f"Dataset: {ds_name}") + print("=" * 60) + for model_name, metrics in models.items(): + print(f"\n Model: {model_name}") + for metric_name, values in metrics.items(): + if values: + mean = _avg(values) + std = _std(values) + print(f" {metric_name:10s}: {mean:10.4f} ± {std:8.4f} (n={len(values)})") + else: + print(f" {metric_name:10s}: N/A") diff --git a/benchmark/metrics.py b/benchmark/metrics.py new file mode 100644 index 00000000..88006c7f --- /dev/null +++ b/benchmark/metrics.py @@ -0,0 +1,118 @@ +"""Metrics computation for the TFTS benchmark system. + +Supports standard time-series forecasting metrics, with both NumPy and +tf.keras metric implementations.""" + +import logging +from typing import Dict, List, Optional, Union + +import numpy as np +import tensorflow as tf + +logger = logging.getLogger(__name__) + + +class BenchmarkMetrics: + """Compute and manage time-series forecasting metrics. + + Attributes: + metrics: List of metric names to compute. + """ + + AVAILABLE_METRICS = { + "mae", + "mse", + "rmse", + "mape", + "smape", + "r2", + "mape_pct", + } + + def __init__(self, metrics: Optional[List[str]] = None): + self.metrics = metrics or ["mae", "rmse", "mape"] + invalid = set(self.metrics) - self.AVAILABLE_METRICS + if invalid: + raise ValueError(f"Invalid metrics: {invalid}. " f"Available: {self.AVAILABLE_METRICS}") + + def compute( + self, + y_true: Union[np.ndarray, tf.Tensor], + y_pred: Union[np.ndarray, tf.Tensor], + metrics: Optional[List[str]] = None, + ) -> Dict[str, float]: + """Compute all requested metrics between y_true and y_pred. + + Args: + y_true: Ground truth values. + y_pred: Predicted values. + metrics: Optional subset of metrics to compute (uses self.metrics if None). + + Returns: + Dictionary mapping metric name to float value. + """ + y_true = np.asarray(y_true) + y_pred = np.asarray(y_pred) + + if y_true.shape != y_pred.shape: + raise ValueError(f"Shape mismatch: y_true {y_true.shape} vs y_pred {y_pred.shape}") + + to_compute = metrics if metrics is not None else self.metrics + results: Dict[str, float] = {} + for metric in to_compute: + fn = getattr(self, metric, None) + if fn is None: + logger.warning("Unknown metric: %s", metric) + continue + try: + results[metric] = float(fn(y_true, y_pred)) + except Exception as exc: + logger.warning("Metric %s failed: %s", metric, exc) + results[metric] = float("nan") + return results + + @staticmethod + def mae(y_true: np.ndarray, y_pred: np.ndarray) -> float: + """Mean Absolute Error.""" + return float(np.mean(np.abs(y_true - y_pred))) + + @staticmethod + def mse(y_true: np.ndarray, y_pred: np.ndarray) -> float: + """Mean Squared Error.""" + return float(np.mean((y_true - y_pred) ** 2)) + + @staticmethod + def rmse(y_true: np.ndarray, y_pred: np.ndarray) -> float: + """Root Mean Squared Error.""" + return float(np.sqrt(np.mean((y_true - y_pred) ** 2))) + + @staticmethod + def mape(y_true: np.ndarray, y_pred: np.ndarray) -> float: + """Mean Absolute Percentage Error (avoids division by zero).""" + mask = y_true != 0 + if not np.any(mask): + return float("nan") + return float(np.mean(np.abs((y_true[mask] - y_pred[mask]) / y_true[mask])) * 100.0) + + @staticmethod + def smape(y_true: np.ndarray, y_pred: np.ndarray) -> float: + """Symmetric Mean Absolute Percentage Error.""" + denom = (np.abs(y_true) + np.abs(y_pred)) / 2.0 + mask = denom != 0 + if not np.any(mask): + return float("nan") + return float(np.mean(np.abs(y_true[mask] - y_pred[mask]) / denom[mask]) * 100.0) + + @staticmethod + def r2(y_true: np.ndarray, y_pred: np.ndarray) -> float: + """R-squared.""" + ss_res = np.sum((y_true - y_pred) ** 2) + ss_tot = np.sum((y_true - np.mean(y_true)) ** 2) + if ss_tot == 0: + return float("nan") + return float(1.0 - ss_res / ss_tot) + + @staticmethod + def mape_pct(y_true: np.ndarray, y_pred: np.ndarray) -> float: + """Mean Absolute Percentage Error as a percentage (0-100 scale).""" + return BenchmarkMetrics.mape(y_true, y_pred) diff --git a/benchmark/registry.py b/benchmark/registry.py new file mode 100644 index 00000000..e0e9b780 --- /dev/null +++ b/benchmark/registry.py @@ -0,0 +1,108 @@ +"""Registry for datasets and models in the TFTS benchmark system.""" + +import logging +from typing import Callable, Dict, List, Optional, Type, Union + +from benchmark.base import Dataset +from tfts.models.auto_config import CONFIG_MAPPING_NAMES +from tfts.models.auto_model import MODEL_MAPPING_NAMES + +logger = logging.getLogger(__name__) + + +class _Registry: + """Internal registry base class.""" + + def __init__(self): + self._items: Dict[str, type] = {} + + def register(self, name: str, item: type) -> None: + if name in self._items: + logger.warning("Overwriting existing registration: %s", name) + self._items[name] = item + logger.debug("Registered: %s", name) + + def get(self, name: str) -> type: + if name not in self._items: + raise KeyError(f"'{name}' not found. Available: {list(self._items.keys())}") + return self._items[name] + + def list_items(self) -> Dict[str, type]: + return dict(self._items) + + def __contains__(self, item: str) -> bool: + return item in self._items + + +class DatasetRegistry(_Registry): + """Registry for benchmark datasets. + + Usage:: + + from tfts.benchmark.registry import DatasetRegistry + from tfts.benchmark.datasets import SineDataset + + registry = DatasetRegistry() + registry.register("sine", SineDataset) + + ds_cls = registry.get("sine") + ds = ds_cls() + print(ds.name) + """ + + def __init__(self): + super().__init__() + self._lazy: Dict[str, Callable[[], Dataset]] = {} + + def register_lazy(self, name: str, factory: Callable[[], Dataset]) -> None: + """Register a lazy factory so datasets are only instantiated on demand.""" + self._lazy[name] = factory + + def get(self, name: str) -> Type[Dataset]: + if name in self._lazy: + # Return a tiny wrapper that calls the factory + factory = self._lazy[name] + + class _LazyDataset(Dataset): + """Lazy-loaded dataset wrapper.""" + + def prepare_data(self, **kwargs): + return factory().prepare_data(**kwargs) + + def get_train_valid_split(self, **kwargs): + return factory().get_train_valid_split(**kwargs) + + _LazyDataset.name = name + return _LazyDataset + return super().get(name) + + def list_datasets(self) -> List[str]: + return sorted(set(list(self._items.keys())) | set(self._lazy.keys())) + + +class ModelRegistry: + """Registry for models available in the benchmark. + + Wraps the existing :mod:`tfts.models` mapping. + """ + + def __init__(self): + # Synchronized with tfts.models.auto_model.MODEL_MAPPING_NAMES + self._models: Dict[str, str] = dict(MODEL_MAPPING_NAMES) + + @property + def available_models(self) -> List[str]: + return list(self._models.keys()) + + def get(self, name: str) -> str: + if name not in self._models: + raise KeyError(f"Model '{name}' not found. Available: {self.available_models}") + return self._models[name] + + def register(self, name: str, class_name: str) -> None: + """Register a custom model.""" + self._models[name] = class_name + logger.debug("Registered model: %s -> %s", name, class_name) + + def __contains__(self, item: str) -> bool: + return item in self._models diff --git a/benchmark/runner.py b/benchmark/runner.py new file mode 100644 index 00000000..cabd2025 --- /dev/null +++ b/benchmark/runner.py @@ -0,0 +1,216 @@ +"""Benchmark runner for the TFTS benchmark system. + +Orchestrates running multiple models on multiple datasets with multiple runs +and collects results.""" + +import json +import logging +import os +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import tensorflow as tf + +from benchmark.base import BenchmarkConfig, Dataset +from benchmark.formatter import BenchmarkResults +from benchmark.metrics import BenchmarkMetrics +from benchmark.registry import DatasetRegistry, ModelRegistry +from tfts import AutoConfig, AutoModel, Trainer, set_seed + +logger = logging.getLogger(__name__) + + +class BenchmarkRunner: + """Run a benchmark experiment. + + Example:: + + from tfts.benchmark import BenchmarkRunner, BenchmarkConfig + + config = BenchmarkConfig( + models=["rnn", "transformer"], + datasets=["sine", "air_passengers"], + metrics=["mae", "rmse"], + runs=3, + epochs=50, + ) + runner = BenchmarkRunner(config) + results = runner.run() + results.to_latex("results.tex") + """ + + def __init__( + self, + config: BenchmarkConfig, + dataset_registry: Optional[DatasetRegistry] = None, + model_registry: Optional[ModelRegistry] = None, + ): + self.config = config + self.dataset_registry = dataset_registry or _default_dataset_registry() + self.model_registry = model_registry or ModelRegistry() + self.metrics = BenchmarkMetrics(config.metrics) + self.results: List[Dict[str, Any]] = [] + + def run(self) -> BenchmarkResults: + """Execute the benchmark and return results. + + Returns: + BenchmarkResults: Container with raw and formatted results. + """ + datasets = self._resolve_datasets() + models = self._resolve_models() + + logger.info("=" * 60) + logger.info("Starting TFTS Benchmark") + logger.info("Models: %s", models) + logger.info("Datasets: %s", datasets) + logger.info("Runs per experiment: %d", self.config.runs) + logger.info("=" * 60) + + for dataset_name in datasets: + for model_name in models: + self._run_experiment(dataset_name, model_name) + + results = BenchmarkResults(self.results) + self._save_results(results) + return results + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _resolve_datasets(self) -> List[str]: + """Return the actual list of dataset names to run.""" + registered = self.dataset_registry.list_datasets() + if self.config.datasets == ["all"]: + return registered + missing = set(self.config.datasets) - set(registered) + if missing: + raise ValueError(f"Unknown datasets: {missing}. Available: {registered}") + return self.config.datasets + + def _resolve_models(self) -> List[str]: + """Return the actual list of model names to run.""" + available = self.model_registry.available_models + if self.config.models == ["all"]: + return available + missing = set(self.config.models) - set(available) + if missing: + raise ValueError(f"Unknown models: {missing}. Available: {available}") + return self.config.models + + def _run_experiment(self, dataset_name: str, model_name: str) -> None: + """Run all trials for a single dataset-model pair.""" + logger.info("-" * 60) + logger.info("Experiment: %s / %s", dataset_name, model_name) + + ds_config = self.config.get_dataset_config(dataset_name) + cls_ = self.dataset_registry.get(dataset_name) + dataset = cls_() + + for run_idx in range(self.config.runs): + seed = self.config.seed + run_idx + set_seed(seed) + + result = self._run_single_trial( + dataset=dataset, + dataset_name=dataset_name, + model_name=model_name, + run_idx=run_idx, + seed=seed, + ds_config=ds_config, + ) + self.results.append(result) + + def _run_single_trial( + self, + dataset: Dataset, + dataset_name: str, + model_name: str, + run_idx: int, + seed: int, + ds_config: Dict[str, Any], + ) -> Dict[str, Any]: + """Run a single trial and return the result dict.""" + logger.info(" Run %d/%d (seed=%d)", run_idx + 1, self.config.runs, seed) + + train_data, valid_data = dataset.get_train_valid_split(**ds_config) + + train_length = ds_config.get("train_length") or dataset.train_length + predict_length = ds_config.get("predict_sequence_length") or dataset.predict_sequence_length + epochs = ds_config.get("epochs", self.config.epochs) + batch_size = ds_config.get("batch_size", self.config.batch_size) + learning_rate = ds_config.get("learning_rate", self.config.learning_rate) + + # Build model + model_config = AutoConfig.for_model(model_name) + # Adjust input shape if known + if hasattr(model_config, "input_shape") and train_data[0].ndim == 3: + model_config.input_shape = train_data[0].shape[1:] + + model = AutoModel.from_config(model_config, predict_sequence_length=predict_length) + trainer = Trainer(model) + + # Train + history = trainer.train( + train_dataset=train_data, + valid_dataset=valid_data, + epochs=epochs, + batch_size=batch_size, + verbose=0 if self.config.verbose < 2 else 1, + ) + + # Evaluate + x_valid, y_valid = valid_data + y_pred = trainer.predict(x_valid) + metrics = self.metrics.compute(y_valid, y_pred) + + result = { + "dataset": dataset_name, + "model": model_name, + "run": run_idx, + "seed": seed, + "train_length": train_length, + "predict_sequence_length": predict_length, + "epochs": epochs, + "batch_size": batch_size, + "learning_rate": learning_rate, + "metrics": metrics, + "history": {k: [float(v) for v in vals] for k, vals in (history.history if history else {}).items()}, + } + return result + + def _save_results(self, results: BenchmarkResults) -> None: + """Save results to the output directory.""" + os.makedirs(self.config.output_dir, exist_ok=True) + results.to_json(os.path.join(self.config.output_dir, "results.json")) + results.to_csv(os.path.join(self.config.output_dir, "results.csv")) + results.to_latex(os.path.join(self.config.output_dir, "results.tex")) + + +# -------------------------------------------------------------------------- +# Default registry population +# -------------------------------------------------------------------------- + + +def _default_dataset_registry() -> DatasetRegistry: + """Build a :class:`DatasetRegistry` with built-in datasets.""" + registry = DatasetRegistry() + + # Lazy import to avoid circular dependency at package top-level + from benchmark.datasets import ( + AirPassengersDataset, + CMIDetectSleepStatesDataset, + ForecastingStickerSalesDataset, + GrocerysalesDataset, + RecruitRestaurantDataset, + SineDataset, + ) + + registry.register("sine", SineDataset) + registry.register("air_passengers", AirPassengersDataset) + registry.register("grocery_sales", GrocerysalesDataset) + registry.register("recruit_restaurant", RecruitRestaurantDataset) + registry.register("forecasting_sticker_sales", ForecastingStickerSalesDataset) + registry.register("CMI_detect_sleep_states", CMIDetectSleepStatesDataset) + return registry diff --git a/examples/benchmarks/CMI_detect_sleep_states/README.md b/examples/benchmarks/CMI_detect_sleep_states/README.md deleted file mode 100644 index e69de29b..00000000 diff --git a/examples/benchmarks/CMI_detect_sleep_states/conf.yaml b/examples/benchmarks/CMI_detect_sleep_states/conf.yaml deleted file mode 100644 index e69de29b..00000000 diff --git a/examples/benchmarks/CMI_detect_sleep_states/dataset.py b/examples/benchmarks/CMI_detect_sleep_states/dataset.py deleted file mode 100644 index e69de29b..00000000 diff --git a/examples/benchmarks/forecasting_sticker_sales/README.md b/examples/benchmarks/forecasting_sticker_sales/README.md deleted file mode 100644 index e69de29b..00000000 diff --git a/examples/benchmarks/forecasting_sticker_sales/conf.yaml b/examples/benchmarks/forecasting_sticker_sales/conf.yaml deleted file mode 100644 index 883f3cee..00000000 --- a/examples/benchmarks/forecasting_sticker_sales/conf.yaml +++ /dev/null @@ -1,21 +0,0 @@ -seed: 315 - -data: - data_path: "data.csv" - target_column: "target" - freq: 'h' - -model: - name: bert - predict_sequence_length: 32 - n_layers: 2 - hidden_size: 128 - n_features: 10 - n_output: 1 - - -training: - batch_size: 128 - epochs: 30 - learning_rate: 0.001 - loss: "MSE" diff --git a/examples/benchmarks/forecasting_sticker_sales/dataset.py b/examples/benchmarks/forecasting_sticker_sales/dataset.py deleted file mode 100644 index daf41971..00000000 --- a/examples/benchmarks/forecasting_sticker_sales/dataset.py +++ /dev/null @@ -1,215 +0,0 @@ -import warnings - -from joblib import Parallel, delayed -import numpy as np -from omegaconf import OmegaConf -import pandas as pd -import requests -from sklearn.preprocessing import StandardScaler -from tensorflow.keras.utils import Sequence - -warnings.filterwarnings("ignore") - - -# https://www.kaggle.com/code/cdeotte/transformer-starter-lb-0-052 -class TimeSeriesProcessor: - def __init__(self, use_internet=True, path="./"): - self.use_internet = use_internet - self.path = path - self.scales = {} - self.gdp_data = None - - def fetch_gdp(self, df): - """Unified GDP fetching logic.""" - alpha3_map = { - "Canada": "CAN", - "Finland": "FIN", - "Italy": "ITA", - "Kenya": "KEN", - "Norway": "NOR", - "Singapore": "SGP", - } - df["alpha3"] = df["country"].map(alpha3_map) - df["year"] = df["date"].dt.year - years = df["year"].unique() - - if self.use_internet: - gdp_dict = {} - for country, a3 in alpha3_map.items(): - try: - url = f"https://api.worldbank.org/v2/country/{a3}/indicator/NY.GDP.PCAP.CD?date={min(years)}:{max(years)}&format=json" # noqa: E501,E231 - res = requests.get(url).json()[1] - for entry in res: - gdp_dict[(a3, int(entry["date"]))] = entry["value"] - except Exception as e: - print(f"Error fetching GDP for {a3}: {e}") - self.gdp_data = gdp_dict - else: - # Assume local file exists - gdp_df = pd.read_csv(f"{self.path}gdp.csv").set_index("alpha3") - self.gdp_data = gdp_df.to_dict() - - return df - - def process_features(self, df, is_train=True): - """Calculates GDP ratios and store-based normalization.""" - df = df.copy() - df["date"] = pd.to_datetime(df["date"]) - - if self.gdp_data is None: - df = self.fetch_gdp(df) - - df["GDP"] = df.apply(lambda x: self.gdp_data.get((x["alpha3"], x["year"]), 1.0), axis=1) - - # 1. GDP Normalization - df["scaled_target"] = df["num_sold"] / df["GDP"] - - # 2. Store Ratio (calculate during train, apply during test) - if is_train: - self.store_ratios = df.groupby("store")["scaled_target"].mean().to_dict() - - df["scaled_target"] /= df["store"].map(self.store_ratios) - - # 3. Kenya Fudge Factor - df.loc[df["country"] == "Kenya", "scaled_target"] *= 1.15 - - return df - - def dataframe_to_tensor(self, df): - """ - Pivots the dataframe into a 3D tensor: (Products, Time, Series) - Series = Country + Store combinations. - """ - # Create a unique key for each Country/Store combination - df["series_key"] = df["country"] + "_" + df["store"] - - products = sorted(df["product"].unique()) - series_keys = sorted(df["series_key"].unique()) - - tensor_list = [] - for prod in products: - # Efficient pivoting instead of nested loops - subset = df[df["product"] == prod].pivot(index="date", columns="series_key", values="scaled_target") - - # Save scaling params per product - if prod not in self.scales: - self.scales[prod] = {"mean": subset.values.mean(), "std": subset.values.std()} - - # Standard Scale - scaled_val = (subset.values - self.scales[prod]["mean"]) / self.scales[prod]["std"] - tensor_list.append(scaled_val) - - return np.stack(tensor_list), products, series_keys - - def inverse_transform(self, pred, product_name, country, store, date): - """Reverses all transformations to get the original num_sold scale.""" - # 1. Reverse Standard Scale - val = (pred * self.scales[product_name]["std"]) + self.scales[product_name]["mean"] - - # 2. Reverse Kenya Factor - if country == "Kenya": - val /= 1.15 - - # 3. Reverse Store Ratio - val *= self.store_ratios[store] - - # 4. Reverse GDP - year = pd.to_datetime(date).year - # Note: You'd need a helper to get alpha3 from country - alpha3 = { - "Canada": "CAN", - "Finland": "FIN", - "Italy": "ITA", - "Kenya": "KEN", - "Norway": "NOR", - "Singapore": "SGP", - }[country] - val *= self.gdp_data.get((alpha3, year), 1.0) - - return val - - -class TimeSeriesDataset(Sequence): - def __init__( - self, - data, - mode="train", # "train" or "test" - product_idx=0, - train_sequence_length=1440, - predict_sequence_length=32, - batch_size=32, - ): - self.data = data[product_idx] # Shape: (Time, Series) - self.mode = mode - self.product_idx = product_idx - self.train_sequence_length = train_sequence_length - self.predict_sequence_length = predict_sequence_length - self.batch_size = batch_size - - nans = np.isnan(self.data).astype("float32") - self.combined_data = np.stack([np.nan_to_num(self.data), nans], axis=-1) - - def __len__(self): - return int(np.ceil(self.data.shape[1] / self.batch_size)) - - def __getitem__(self, idx): - if self.mode == "train": - return self._get_train_batch() - else: - return self._get_test_batch(idx) - - def _get_train_batch(self): - X = np.zeros((self.batch_size, self.train_sequence_length, 2), dtype="float32") - y = np.zeros((self.batch_size, self.predict_sequence_length), dtype="float32") - - for i in range(self.batch_size): - series_idx = np.random.randint(0, self.data.shape[1]) - start = np.random.randint(0, self.data.shape[0] - self.train_sequence_length - self.predict_sequence_length) - - X[i] = self.combined_data[start : start + self.train_sequence_length, series_idx, :] - y[i] = self.combined_data[ - start + self.train_sequence_length : start + self.train_sequence_length + self.predict_sequence_length, - series_idx, - 0, - ] - return X, y - - def _get_test_batch(self, idx): - """Returns the LAST train_len for each category for prediction.""" - start_series = idx * self.batch_size - end_series = min((idx + 1) * self.batch_size, self.data.shape[1]) - actual_bs = end_series - start_series - - X = np.zeros((actual_bs, self.train_sequence_length, 2), dtype="float32") - - for i, s_idx in enumerate(range(start_series, end_series)): - # Always take the very tail of the data - X[i] = self.combined_data[-self.train_sequence_length :, s_idx, :] - - return X - - -if __name__ == "__main__": - # 1. Process Data - processor = TimeSeriesProcessor(use_internet=True) - df_train = pd.read_csv("/kaggle/input/playground-series-s5e1/train.csv") - df_processed = processor.process_features(df_train, is_train=True) - tensor, product_names, series_names = processor.dataframe_to_tensor(df_processed) - - # 2. Create Train Dataset for Product 0 - train_gen = TimeSeriesDataset(tensor, mode="train", product_idx=0) - - # 3. Create Test Dataset (the last window for all series in Product 0) - test_gen = TimeSeriesDataset(tensor, mode="test", product_idx=0) - - # # 4. Predict - # predictions = model.predict(test_gen) # (Total Series, pred_len) - - # # 5. Reverse Scaling for a specific prediction - # raw_pred = processor.inverse_transform( - # pred=predictions[0, 0], - # product_name=product_names[0], - # country="Canada", - # store="KaggleMart", - # date="2026-01-01" - # ) diff --git a/examples/benchmarks/forecasting_sticker_sales/run.py b/examples/benchmarks/forecasting_sticker_sales/run.py deleted file mode 100644 index 2c2f3704..00000000 --- a/examples/benchmarks/forecasting_sticker_sales/run.py +++ /dev/null @@ -1,96 +0,0 @@ -import argparse -import math -import random - -from dataset import DataReader, TrainDataset -import numpy as np -from omegaconf import OmegaConf -import pandas as pd -import tensorflow as tf - -from tfts import AutoConfig, AutoModel, Pipeline, set_seed - - -def parse_args(): - parser = argparse.ArgumentParser(description="tfts forecasting") - parser.add_argument("--config_path", type=str, default="conf.yaml", help="Path to base config file") - parser.add_argument("--debug", type=bool, default=False, help="Enable debug mode") - parser.add_argument("--is_training", type=bool, default=True, help="Whether to train or predict") - parser.add_argument("--model_name", type=str, default=None, help="Model name, e.g., BERT, LSTM") - parser.add_argument("--batch_size", type=int, default=None, help="Batch size") - parser.add_argument("--epochs", type=int, default=None, help="Number of epochs") - - args = parser.parse_args() - return args - - -# def run_inference(product_idx): -# """Runs the recursive prediction for a specific product.""" -# # Ensure history has the 2nd channel (NaN indicator) -# # history_tensor shape: (5, 2557, 18) -> Needs expansion to (1, LEN, 18, 2) -# data = np.expand_dims(self.history_tensor, axis=-1) -# nans = np.isnan(data).astype('float32') -# data = np.concatenate([data, nans], axis=-1) - -# product_preds = np.zeros((18, self.PRED_LEN * self.STEPS)) -# bad_rows = [] - -# for jj in range(18): -# # Get last window of training data for this series -# # Shape: (1, LEN, 2) -# current_window = data[product_idx:product_idx+1, -self.LEN:, jj, :].copy() - -# if np.isnan(current_window[:, :, 0]).sum() == self.LEN: -# bad_rows.append(jj) -# continue - -# series_predictions = [] - -# for step in range(self.STEPS): -# # Predict next 32 days -# # Input shape: (1, LEN, 2) -# p2 = self.model(np.nan_to_num(current_window)) -# p2 = p2.numpy().reshape((1, self.PRED_LEN, 1)) - -# # Add dummy NaN indicator (0.0) to predictions for the next step -# p2_with_nan = np.concatenate([p2, np.zeros_like(p2)], axis=-1) -# series_predictions.append(p2_with_nan) - -# # Update window: Slide window forward -# # Remove oldest 32, append newest 32 -# current_window = np.concatenate([current_window[:, self.PRED_LEN:, :], p2_with_nan], axis=1) - -# # Combine all steps and remove the NaN indicator channel -# product_preds[jj, :] = np.concatenate([z[:, :, 0] for z in series_predictions], axis=1).flatten() - -# # Handle bad rows (series with no training data) -# if bad_rows: -# fill_val = np.nanmean(product_preds, axis=0) -# for r in bad_rows: -# product_preds[r, :] = fill_val - -# return product_preds - - -def main(): - args = parse_args() - cfg = OmegaConf.load(args.config_path) - - set_seed(cfg.seed) - - data_reader = DataReader() - train_df = data_reader.load_data("/kaggle/input/playground-series-s5e1/train.csv") - train_df = data_reader.add_features(train_df) - data_tensor = data_reader.reshape_to_tensor(train_df) - - train_dataset = TrainDataset( - data=data_tensor, product_idx=0, batch_size=64, train_sequence_length=1440, predict_sequence_length=32 - ) - - forecaster = Pipeline(cfg) - - forecaster.train(train_dataset=train_dataset) - - -if __name__ == "__main__": - main() diff --git a/examples/run_anomaly.py b/examples/run_anomaly.py index 8f1b2054..4a109f72 100644 --- a/examples/run_anomaly.py +++ b/examples/run_anomaly.py @@ -8,8 +8,9 @@ import numpy as np import pandas as pd from sklearn.preprocessing import StandardScaler +import tensorflow as tf -from tfts import AutoConfig, AutoModel, AutoModelForAnomaly, KerasTrainer +from tfts import AutoConfig, AutoModelForAnomaly, KerasTrainer, set_seed def parse_args(): @@ -28,7 +29,7 @@ def parse_args(): def create_subsequences(time_series, train_length, pred_length): """Create subsequences for training and prediction.""" subsequences, next_values = [], [] - for i in range(len(time_series) - train_length - pred_length): + for i in range(len(time_series) - train_length - pred_length + 1): subsequences.append(time_series[i : i + train_length]) next_values.append(time_series[i + train_length : i + train_length + pred_length].T[0]) return subsequences, next_values @@ -53,6 +54,7 @@ def load_and_preprocess_data(args): def train_model(args): """Train the model using the specified arguments.""" + set_seed(args.seed) x_train, y_train, _ = load_and_preprocess_data(args) config = AutoConfig.for_model(args.use_model) @@ -60,7 +62,13 @@ def train_model(args): model = AutoModelForAnomaly.from_config(config) trainer = KerasTrainer(model) - trainer.train((x_train, y_train), (x_train, y_train), epochs=args.epochs) + trainer.train( + (x_train, y_train), + (x_train, y_train), + optimizer=tf.keras.optimizers.Adam(args.learning_rate), + epochs=args.epochs, + batch_size=args.batch_size, + ) trainer.save_model(args.output_dir) print(f"Model trained and saved to {args.output_dir}") diff --git a/examples/run_classification.py b/examples/run_classification.py index 4c56aaa2..b23379b7 100644 --- a/examples/run_classification.py +++ b/examples/run_classification.py @@ -8,7 +8,7 @@ from sklearn.model_selection import train_test_split import tensorflow as tf -from tfts import AutoConfig, AutoModelForClassification, KerasTrainer +from tfts import AutoConfig, AutoModelForClassification, KerasTrainer, set_seed logging.getLogger("tensorflow").setLevel(logging.ERROR) @@ -49,6 +49,7 @@ def readucr(filename): def run_train(args): + set_seed(args.seed) x_train, y_train, x_test, y_test = prepare_data() x_train, x_val, y_train, y_val = train_test_split(x_train, y_train, test_size=0.2, random_state=42) diff --git a/examples/run_prediction_simple.py b/examples/run_prediction_simple.py index ffba3251..c2d9bb89 100644 --- a/examples/run_prediction_simple.py +++ b/examples/run_prediction_simple.py @@ -1,5 +1,10 @@ """Demo of time series prediction by tfts -python run_prediction_simple.py --use_model rnn + +Two equivalent approaches: + 1. Simple pipeline API (recommended) + python run_prediction_simple.py --use_model dlinear + 2. Manual API (for full control, original style) + python run_prediction_simple.py --use_model rnn --manual """ import argparse @@ -9,50 +14,88 @@ import matplotlib.pyplot as plt import numpy as np import tensorflow as tf -from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint -from tensorflow.keras.optimizers.schedules import LearningRateSchedule +from tensorflow.keras.callbacks import EarlyStopping import tfts -from tfts import AutoConfig, AutoModel, KerasTrainer def parse_args(): parser = argparse.ArgumentParser() - parser.add_argument("--seed", type=int, default=315, required=False, help="seed") - parser.add_argument("--use_model", type=str, default="bert", help="model for train") + parser.add_argument("--seed", type=int, default=315, help="seed") + parser.add_argument("--use_model", type=str, default="dlinear", help="model for train") parser.add_argument("--use_data", type=str, default="sine", help="dataset: sine or air passengers") parser.add_argument("--train_length", type=int, default=24, help="sequence length for train") parser.add_argument("--predict_sequence_length", type=int, default=12, help="sequence length for predict") - parser.add_argument("--epochs", type=int, default=100, help="Number of training epochs") + parser.add_argument("--epochs", type=int, default=50, help="Number of training epochs") parser.add_argument("--batch_size", type=int, default=16, help="Batch size for training") - parser.add_argument("--learning_rate", type=float, default=5e-4, help="learning rate for training") - + parser.add_argument("--learning_rate", type=float, default=1e-3, help="learning rate") + parser.add_argument("--manual", action="store_true", help="Use the manual API instead of pipeline") return parser.parse_args() -def set_seed(seed): - random.seed(seed) - np.random.seed(seed) - os.environ["PYTHONHASHSEED"] = str(seed) - tf.random.set_seed(seed) +# ============================================================== +# Approach 1: Simple Pipeline API (recommended for most users) +# ============================================================== + +def run_pipeline(args): + """3-line forecasting with the pipeline API.""" + tfts.set_seed(args.seed) -def run_train(args): - set_seed(args.seed) + # Get data train, valid = tfts.get_data(args.use_data, args.train_length, args.predict_sequence_length, test_size=0.2) - optimizer = tf.keras.optimizers.Adam(args.learning_rate) + + # Create pipeline + pipe = tfts.pipeline( + "forecasting", + model=args.use_model, + lookback=args.train_length, + horizon=args.predict_sequence_length, + learning_rate=args.learning_rate, + epochs=args.epochs, + batch_size=args.batch_size, + early_stopping_patience=5, + seed=args.seed, + ) + pipe.summary() + + # Train (using raw arrays — pipeline handles the rest) + _ = pipe.trainer.train( + train, + valid, + epochs=args.epochs, + batch_size=args.batch_size, + verbose=1, + early_stopping_patience=5, + ) + + pred = pipe.trainer.predict(valid[0]) + pipe.trainer.plot(history=valid[0], true=valid[1], pred=pred) + return pred + + +# ============================================================== +# Approach 2: Manual API (full control — original style) +# ============================================================== + + +def run_manual(args): + """Step-by-step control with AutoConfig / AutoModel / Trainer.""" + tfts.set_seed(args.seed) + train, valid = tfts.get_data(args.use_data, args.train_length, args.predict_sequence_length, test_size=0.2) + loss_fn = tf.keras.losses.MeanSquaredError() + optimizer = tf.keras.optimizers.Adam(args.learning_rate) - # for strong seasonality data like sine or air passengers, set up skip_connect_circle True - config = AutoConfig.for_model(args.use_model) - model = AutoModel.from_config(config, predict_sequence_length=args.predict_sequence_length) + config = tfts.AutoConfig.for_model(args.use_model) + model = tfts.AutoModel.from_config(config, predict_sequence_length=args.predict_sequence_length) - trainer = KerasTrainer(model) + trainer = tfts.Trainer(model) trainer.train( train, valid, - optimizer=optimizer, loss_fn=loss_fn, + optimizer=optimizer, epochs=args.epochs, callbacks=[EarlyStopping("val_loss", patience=5)], ) @@ -60,8 +103,16 @@ def run_train(args): pred = trainer.predict(valid[0]) trainer.plot(history=valid[0], true=valid[1], pred=pred) + # Evaluate + metrics = trainer.evaluate(valid) + print(f"\nEvaluation: {metrics}") + return pred + if __name__ == "__main__": args = parse_args() - run_train(args) + if args.manual: + run_manual(args) + else: + run_pipeline(args) plt.show() diff --git a/examples/run_tuner.py b/examples/run_tuner.py index a61096e1..cf79b07e 100644 --- a/examples/run_tuner.py +++ b/examples/run_tuner.py @@ -1,6 +1,7 @@ """Demo to tune the model parameters by Autotune""" import numpy as np +import tensorflow as tf from tfts import AutoConfig, AutoModel, KerasTrainer, get_data @@ -30,9 +31,15 @@ def objective(self, trial): config.num_stacked_layers = num_layers model = AutoModel.from_config(config, predict_sequence_length=self.predict_sequence_length) - trainer = KerasTrainer(model, optimizer_config={"learning_rate": learning_rate}) - - trainer.train(self.train_data, self.valid_data, epochs=epochs, verbose=0) + trainer = KerasTrainer(model) + + trainer.train( + self.train_data, + self.valid_data, + optimizer=tf.keras.optimizers.Adam(learning_rate), + epochs=epochs, + verbose=0, + ) x_valid, y_valid = self.valid_data predictions = trainer.predict(x_valid) diff --git a/pyproject.toml b/pyproject.toml index e1b44f61..24ad58d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,9 @@ repository = "https://github.com/LongxingTan/Time-series-prediction" documentation = "https://time-series-prediction.readthedocs.io" homepage = "https://time-series-prediction.readthedocs.io" +[tool.poetry.scripts] +tfts-forecast = "tfts.cli.forecasting:main" + [tool.poetry.dependencies] python = ">=3.8,<=3.13" pandas = ">=1.3.0" diff --git a/setup.cfg b/setup.cfg index 2b2003bf..6f21578a 100644 --- a/setup.cfg +++ b/setup.cfg @@ -4,6 +4,8 @@ show-source = true ignore = # space before : (needed for how black formats slicing) E203, + # missing whitespace after : (false positive in f-string format specs like {x:.4f}) + E231, # line break before binary operator W503, # line break after binary operator diff --git a/tests/test_benchmark_formatter.py b/tests/test_benchmark_formatter.py new file mode 100644 index 00000000..292e8dee --- /dev/null +++ b/tests/test_benchmark_formatter.py @@ -0,0 +1,340 @@ +import builtins +from contextlib import redirect_stdout +import csv +import io +import json +import math +import os +import tempfile +import unittest +from unittest.mock import patch + +import numpy as np + +from benchmark.base import BenchmarkConfig +from benchmark.formatter import BenchmarkResults +from benchmark.metrics import BenchmarkMetrics +from benchmark.registry import DatasetRegistry, ModelRegistry +from benchmark.runner import BenchmarkRunner + + +class BenchmarkResultsTest(unittest.TestCase): + def test_to_csv_fallback_handles_dict_results_without_pandas(self): + results = BenchmarkResults( + [ + {"dataset": "synthetic", "model": "rnn", "metrics": {"mae": 0.1}}, + {"dataset": "synthetic", "model": "tcn", "metrics": {"mae": 0.2}}, + ] + ) + + original_import = builtins.__import__ + + def import_without_pandas(name, *args, **kwargs): + if name == "pandas": + raise ImportError("pandas disabled for fallback test") + return original_import(name, *args, **kwargs) + + with tempfile.NamedTemporaryFile(suffix=".csv") as tmp: + try: + builtins.__import__ = import_without_pandas + results.to_csv(tmp.name) + finally: + builtins.__import__ = original_import + + tmp.seek(0) + rows = list(csv.DictReader(line.decode("utf-8") for line in tmp.readlines())) + + self.assertEqual(rows[0]["dataset"], "synthetic") + self.assertEqual(rows[0]["model"], "rnn") + self.assertEqual(rows[0]["metrics"], "{'mae': 0.1}") + + def test_dataframe_and_export_formats(self): + results = BenchmarkResults( + [ + {"dataset": "a", "model": "rnn", "metrics": {"mae": 1.0, "rmse": 2.0}}, + {"dataset": "a", "model": "rnn", "metrics": {"mae": 3.0, "rmse": "bad"}}, + {"dataset": "b", "model": "tcn", "metrics": {"mae": 4.0}}, + ] + ) + + frame = results.to_dataframe() + self.assertEqual(set(frame["dataset"]), {"a", "b"}) + row = frame[(frame["dataset"] == "a") & (frame["model"] == "rnn")].iloc[0] + self.assertEqual(row["mae_mean"], 2.0) + self.assertEqual(row["mae_std"], 1.0) + + with tempfile.TemporaryDirectory() as tmpdir: + csv_path = f"{tmpdir}/results.csv" + json_path = f"{tmpdir}/results.json" + results.to_csv(csv_path) + results.to_json(json_path) + with open(csv_path, encoding="utf-8") as fh: + self.assertIn("mae_mean", fh.readline()) + with open(json_path, encoding="utf-8") as fh: + self.assertEqual(json.load(fh)[0]["dataset"], "a") + + def test_dataframe_reports_missing_pandas_dependency(self): + results = BenchmarkResults([]) + original_import = builtins.__import__ + + def import_without_pandas(name, *args, **kwargs): + if name == "pandas": + raise ImportError("pandas disabled for dataframe test") + return original_import(name, *args, **kwargs) + + try: + builtins.__import__ = import_without_pandas + with self.assertRaisesRegex(ImportError, "pandas is required"): + results.to_dataframe() + finally: + builtins.__import__ = original_import + + def test_pivot_and_latex_cover_missing_and_invalid_values(self): + results = BenchmarkResults( + [ + {"metrics": {"mae": 1.0, "invalid": "not-a-number"}}, + {"dataset": "second", "model": "only", "metrics": {"mae": 2.0}}, + ] + ) + pivot = results._pivot() + self.assertEqual(pivot["unknown"]["unknown"]["mae"], [1.0]) + self.assertEqual(pivot["unknown"]["unknown"]["invalid"], []) + + with tempfile.NamedTemporaryFile(suffix=".tex") as tmp: + results.to_latex(tmp.name, metric="rmse") + tmp.seek(0) + latex = tmp.read().decode("utf-8") + self.assertIn("rmse", latex) + self.assertIn("-", latex) + + with tempfile.NamedTemporaryFile(suffix=".tex") as tmp: + BenchmarkResults([{"dataset": "first", "model": "only", "metrics": {"mae": 1.0}}]).to_latex( + tmp.name, metric="mae" + ) + tmp.seek(0) + self.assertIn("1.0000", tmp.read().decode("utf-8")) + + def test_empty_exports_and_console_output(self): + empty = BenchmarkResults([]) + with tempfile.TemporaryDirectory() as tmpdir: + csv_path = f"{tmpdir}/empty.csv" + original_import = builtins.__import__ + + def import_without_pandas(name, *args, **kwargs): + if name == "pandas": + raise ImportError("pandas disabled for empty fallback test") + return original_import(name, *args, **kwargs) + + try: + builtins.__import__ = import_without_pandas + empty.to_csv(csv_path) + finally: + builtins.__import__ = original_import + self.assertFalse(os.path.exists(csv_path)) + latex_path = f"{tmpdir}/empty.tex" + empty.to_latex(latex_path) + with open(latex_path, encoding="utf-8") as fh: + self.assertIn("No results", fh.read()) + + output = io.StringIO() + with redirect_stdout(output): + empty.print_table() + self.assertIn("No results to display", output.getvalue()) + + def test_print_table_formats_nonempty_results(self): + results = BenchmarkResults([{"dataset": "synthetic", "model": "rnn", "metrics": {"mae": 0.25}}]) + output = io.StringIO() + with redirect_stdout(output): + results.print_table() + self.assertIn("Dataset: synthetic", output.getvalue()) + self.assertIn("mae", output.getvalue()) + + with patch.object(results, "_pivot", return_value={"synthetic": {"rnn": {"mae": []}}}): + output = io.StringIO() + with redirect_stdout(output): + results.print_table() + self.assertIn("N/A", output.getvalue()) + + +class BenchmarkHelpersTest(unittest.TestCase): + def test_formatter_helpers_ignore_nan_values(self): + from benchmark.formatter import _avg, _format_value, _std + + self.assertEqual(_format_value(1.23456), "1.2346") + self.assertEqual(_format_value("value"), "value") + self.assertEqual(_avg([1.0, float("nan"), 3.0]), 2.0) + self.assertEqual(_std([1.0, float("nan"), 3.0]), 1.0) + self.assertTrue(math.isnan(_avg([]))) + self.assertTrue(math.isnan(_std([float("nan")]))) + + def test_benchmark_metrics_cover_standard_and_edge_cases(self): + y_true = np.array([0.0, 2.0, 4.0]) + y_pred = np.array([0.0, 1.0, 2.0]) + metrics = BenchmarkMetrics(["mae", "mse", "rmse", "mape", "smape", "r2", "mape_pct"]) + values = metrics.compute(y_true, y_pred) + self.assertEqual(values["mae"], 1.0) + self.assertEqual(values["mse"], 5.0 / 3.0) + self.assertAlmostEqual(values["rmse"], np.sqrt(5.0 / 3.0)) + self.assertIn("mape_pct", values) + + self.assertTrue(math.isnan(BenchmarkMetrics.mape(np.zeros(2), np.ones(2)))) + self.assertTrue(math.isnan(BenchmarkMetrics.smape(np.zeros(2), np.zeros(2)))) + self.assertTrue(math.isnan(BenchmarkMetrics.r2(np.ones(2), np.zeros(2)))) + self.assertEqual(metrics.compute(y_true, y_pred, metrics=["does_not_exist"]), {}) + metrics.mae = lambda *_: (_ for _ in ()).throw(RuntimeError("broken metric")) + self.assertTrue(math.isnan(metrics.compute(y_true, y_pred, metrics=["mae"])["mae"])) + with self.assertRaises(ValueError): + metrics.compute(y_true, np.zeros(2)) + with self.assertRaises(ValueError): + BenchmarkMetrics(["does_not_exist"]) + + +class BenchmarkRunnerTest(unittest.TestCase): + def test_registry_resolution_and_runner_validation(self): + config = BenchmarkConfig(models=["rnn"], datasets=["toy"], output_dir="unused") + dataset_registry = DatasetRegistry() + dataset_registry.register("toy", object) + model_registry = ModelRegistry() + runner = BenchmarkRunner(config, dataset_registry, model_registry) + + self.assertEqual(runner._resolve_datasets(), ["toy"]) + self.assertEqual(runner._resolve_models(), ["rnn"]) + with self.assertRaises(ValueError): + BenchmarkRunner( + BenchmarkConfig(models=["missing"], datasets=["toy"]), dataset_registry, model_registry + )._resolve_models() + with self.assertRaises(ValueError): + BenchmarkRunner( + BenchmarkConfig(models=["rnn"], datasets=["missing"]), dataset_registry, model_registry + )._resolve_datasets() + + all_config = BenchmarkConfig(models=["all"], datasets=["all"]) + all_runner = BenchmarkRunner(all_config, dataset_registry, model_registry) + self.assertEqual(all_runner._resolve_datasets(), ["toy"]) + self.assertIn("rnn", all_runner._resolve_models()) + + def test_single_trial_uses_dataset_overrides(self): + config = BenchmarkConfig(models=["rnn"], datasets=["toy"], epochs=3, batch_size=4, learning_rate=0.1) + dataset_registry = DatasetRegistry() + runner = BenchmarkRunner(config, dataset_registry, ModelRegistry()) + dataset = type("ToyDataset", (), {"train_length": 8, "predict_sequence_length": 2})() + train = (np.zeros((2, 4, 1)), np.zeros((2, 2, 1))) + valid = (np.zeros((1, 4, 1)), np.zeros((1, 2, 1))) + dataset.get_train_valid_split = lambda **kwargs: (train, valid) + history = type("History", (), {"history": {"loss": [np.float32(0.5)]}})() + + with patch("benchmark.runner.AutoConfig") as auto_config, patch( + "benchmark.runner.AutoModel" + ) as auto_model, patch("benchmark.runner.Trainer") as trainer_cls: + auto_config.for_model.return_value = type("Config", (), {"input_shape": None})() + auto_model.from_config.return_value = object() + trainer_cls.return_value.train.return_value = history + trainer_cls.return_value.predict.return_value = valid[1] + result = runner._run_single_trial( + dataset, + "toy", + "rnn", + run_idx=0, + seed=42, + ds_config={"train_length": 5, "predict_sequence_length": 1, "epochs": 1, "batch_size": 2}, + ) + + self.assertEqual(result["train_length"], 5) + self.assertEqual(result["predict_sequence_length"], 1) + self.assertEqual(result["history"]["loss"], [0.5]) + + def test_run_and_save_results(self): + with tempfile.TemporaryDirectory() as tmpdir: + config = BenchmarkConfig(models=["rnn"], datasets=["toy"], output_dir=tmpdir) + dataset_registry = DatasetRegistry() + dataset_registry.register("toy", object) + runner = BenchmarkRunner(config, dataset_registry, ModelRegistry()) + with patch.object(runner, "_run_experiment") as run_experiment: + result = runner.run() + run_experiment.assert_called_once_with("toy", "rnn") + self.assertEqual(result.results, []) + + def test_experiment_runs_each_seed_and_default_registry_is_available(self): + config = BenchmarkConfig(models=["rnn"], datasets=["toy"], runs=2, seed=10) + dataset_registry = DatasetRegistry() + + class ToyDataset: + train_length = 4 + predict_sequence_length = 1 + + dataset_registry.register("toy", ToyDataset) + runner = BenchmarkRunner(config, dataset_registry, ModelRegistry()) + trial_results = [{"run": 0}, {"run": 1}] + with patch("benchmark.runner.set_seed") as set_seed, patch.object( + runner, "_run_single_trial", side_effect=trial_results + ) as run_trial: + runner._run_experiment("toy", "rnn") + self.assertEqual(runner.results, trial_results) + self.assertEqual(set_seed.call_args_list[0].args, (10,)) + self.assertEqual(set_seed.call_args_list[1].args, (11,)) + self.assertEqual(run_trial.call_count, 2) + + from benchmark.runner import _default_dataset_registry + + default_names = _default_dataset_registry().list_datasets() + self.assertIn("sine", default_names) + self.assertIn("grocery_sales", default_names) + + +class BenchmarkConfigTest(unittest.TestCase): + def test_validation_and_dataset_overrides(self): + with self.assertRaises(ValueError): + BenchmarkConfig(runs=0) + with self.assertRaises(ValueError): + BenchmarkConfig(epochs=0) + config = BenchmarkConfig( + epochs=3, + batch_size=4, + per_dataset_config={"toy": {"epochs": 1, "train_length": 5}}, + ) + self.assertEqual(config.get_dataset_config("toy")["epochs"], 1) + self.assertEqual(config.get_dataset_config("other")["epochs"], 3) + + +class RegistryTest(unittest.TestCase): + def test_base_registry_and_model_registry_operations(self): + registry = DatasetRegistry() + registry.register("toy", int) + self.assertIn("toy", registry) + self.assertEqual(registry.get("toy"), int) + self.assertEqual(registry.list_items(), {"toy": int}) + with self.assertLogs("benchmark.registry", level="WARNING"): + registry.register("toy", str) + with self.assertRaises(KeyError): + registry.get("missing") + + models = ModelRegistry() + models.register("custom", "CustomModel") + self.assertIn("custom", models) + self.assertEqual(models.get("custom"), "CustomModel") + with self.assertRaises(KeyError): + models.get("missing") + + +class DatasetRegistryTest(unittest.TestCase): + def test_lazy_dataset_registration_returns_instantiable_wrapper(self): + class ToyDataset: + def prepare_data(self, **kwargs): + return "prepared" + + def get_train_valid_split(self, **kwargs): + return "split" + + registry = DatasetRegistry() + registry.register_lazy("toy", lambda: ToyDataset()) + + dataset_cls = registry.get("toy") + dataset = dataset_cls() + + self.assertEqual(dataset_cls.name, "toy") + self.assertEqual(dataset.prepare_data(), "prepared") + self.assertEqual(dataset.get_train_valid_split(), "split") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_data/test_processor.py b/tests/test_data/test_processor.py new file mode 100644 index 00000000..34e23322 --- /dev/null +++ b/tests/test_data/test_processor.py @@ -0,0 +1,169 @@ +import unittest + +import numpy as np +import pandas as pd +import tensorflow as tf + +from tfts.data.auto_preprocessor import AutoPreprocessor +from tfts.data.processor import DataProcessor, _looks_like_time + + +class DataProcessorTest(unittest.TestCase): + def setUp(self): + self.df = pd.DataFrame( + { + "date": pd.date_range("2024-01-01", periods=20, freq="D"), + "value": np.arange(20, dtype=float), + } + ) + + def test_standard_inverse_transform(self): + processor = DataProcessor(lookback=4, horizon=2, normalize="standard", validation_split=0) + processor.prepare(self.df, target_col="value", time_col="date") + normalized = (self.df["value"].to_numpy() - self.df["value"].mean()) / self.df["value"].std() + np.testing.assert_allclose(processor.inverse_transform(normalized), self.df["value"], rtol=1e-6) + + def test_inference_reuses_training_scaler_and_latest_window(self): + processor = DataProcessor(lookback=4, horizon=2, normalize="minmax", validation_split=0) + processor.prepare(self.df, target_col="value", time_col="date") + scaler = dict(processor._scaler_params) + + inference_df = self.df.tail(4).copy() + inference_ds = processor.prepare_for_inference(inference_df, target_col="value", time_col="date") + batches = list(inference_ds.as_numpy_iterator()) + + self.assertEqual(processor._scaler_params, scaler) + self.assertEqual(len(batches), 1) + self.assertEqual(batches[0][0].shape, (1, 4, 1)) + + def test_validation_split_is_chronological(self): + processor = DataProcessor( + lookback=3, + horizon=1, + normalize=None, + validation_split=0.25, + shuffle=False, + batch_size=64, + ) + train_ds, valid_ds = processor.prepare(self.df, target_col="value", time_col="date") + train_x = next(iter(train_ds))[0].numpy() + valid_x = next(iter(valid_ds))[0].numpy() + self.assertLess(train_x[-1, -1, 0], valid_x[0, -1, 0]) + + def test_scaler_is_fitted_without_validation_leakage(self): + processor = DataProcessor(lookback=3, horizon=1, normalize="minmax", validation_split=0.25) + processor.prepare(self.df, target_col="value", time_col="date") + self.assertEqual(processor._scaler_params["max"], 14.0) + + def test_prepare_without_validation_and_standard_inverse(self): + processor = DataProcessor(lookback=4, horizon=2, normalize="standard", validation_split=0, shuffle=False) + dataset = processor.prepare(self.df, target_col="value", time_col="date") + self.assertIsInstance(dataset, tf.data.Dataset) + np.testing.assert_allclose(processor.inverse_transform(np.array([0.0])), np.array([self.df["value"].mean()])) + + minmax = DataProcessor(lookback=4, horizon=2, normalize="minmax", validation_split=0) + minmax.prepare(self.df, target_col="value", time_col="date") + np.testing.assert_allclose(minmax.inverse_transform(np.array([0.0, 1.0])), np.array([0.0, 19.0])) + unfitted = DataProcessor(normalize="minmax") + np.testing.assert_array_equal(unfitted.inverse_transform(np.array([1.0])), np.array([1.0])) + + def test_inference_requires_fitted_normalizer(self): + processor = DataProcessor(lookback=4, horizon=2, normalize="minmax") + with self.assertRaisesRegex(RuntimeError, "fitted"): + processor.prepare_for_inference(self.df.tail(4), target_col="value", time_col="date") + + def test_validation_and_target_time_inference_helpers(self): + for kwargs in [ + {"lookback": 0}, + {"horizon": 0}, + {"batch_size": 0}, + {"stride": 0}, + {"validation_split": -0.1}, + {"validation_split": 1.0}, + {"normalize": "bad"}, + ]: + with self.assertRaises(ValueError): + DataProcessor(**kwargs) + + numeric = pd.DataFrame({"index": [1, 2, 3], "value": [4.0, 5.0, 6.0]}) + self.assertEqual(DataProcessor._infer_target(numeric), "value") + self.assertEqual(DataProcessor._infer_target(pd.DataFrame({"first": [1, 2, 3]})), "first") + self.assertEqual(DataProcessor._infer_time(self.df), "date") + self.assertEqual(DataProcessor._infer_time(pd.DataFrame({"date": [1, 2]})), "date") + self.assertEqual(DataProcessor._infer_time(pd.DataFrame({"first": [1, 2]})), "first") + self.assertTrue(_looks_like_time(pd.Series(pd.date_range("2024-01-01", periods=2)))) + with self.assertRaises(ValueError): + DataProcessor._infer_target(pd.DataFrame({"text": ["a", "b"]})) + with self.assertRaises(ValueError): + DataProcessor(lookback=2, horizon=2, validation_split=0.5)._split_dataset( + tf.data.Dataset.from_tensor_slices((np.zeros((0, 2, 1)), np.zeros((0, 2, 1)))) + ) + + def test_group_normalization_fit_frame_and_time_like_heuristic(self): + grouped = self.df.assign(group=["a"] * 10 + ["b"] * 10) + processor = DataProcessor(group_col="group", validation_split=0.25) + fit_frame = processor._normalization_fit_frame(grouped) + self.assertEqual(len(fit_frame), 14) + self.assertTrue( + DataProcessor._infer_time(pd.DataFrame(index=pd.date_range("2024-01-01", periods=2))).endswith("idx") + ) + self.assertTrue(_looks_like_time(pd.Series([1, 2, 3]))) + self.assertFalse(_looks_like_time(pd.Series([1, 3, 2]))) + + +class AutoPreprocessorTest(unittest.TestCase): + def test_forward_fill_preserves_unprocessed_columns(self): + df = pd.DataFrame({"time": [1, 2, 3], "value": [1.0, np.nan, 3.0]}) + result = AutoPreprocessor(handle_missing="ffill", columns=["value"]).fit_transform(df) + self.assertEqual(list(result.columns), ["time", "value"]) + self.assertEqual(result.loc[1, "value"], 1.0) + + def test_interpolate_drop_and_no_missing_strategies(self): + df = pd.DataFrame({"value": [np.nan, 2.0, np.nan, 6.0, np.nan], "label": [1, 2, 3, 4, 5]}) + interpolated = AutoPreprocessor(handle_missing="interpolate", columns=["value"]).fit_transform(df) + np.testing.assert_allclose(interpolated["value"], [2.0, 2.0, 4.0, 6.0, 6.0]) + + dropped = AutoPreprocessor(handle_missing="drop", columns=["value"]).fit_transform(df) + self.assertEqual(len(dropped), 2) + untouched = AutoPreprocessor(handle_missing=None, columns=["value"]).fit_transform(df) + self.assertTrue(untouched["value"].isna().any()) + + def test_clip_normalize_and_inverse_transform(self): + df = pd.DataFrame({"value": [1.0, 2.0, 3.0, 100.0], "other": [10, 20, 30, 40]}) + preprocessor = AutoPreprocessor(handle_outliers="clip", normalize="minmax", columns=["value"]) + transformed = preprocessor.fit_transform(df) + self.assertLess(transformed.loc[3, "value"], 1.0) + restored = preprocessor.inverse_transform(transformed) + self.assertAlmostEqual(restored.loc[0, "value"], 1.0) + self.assertEqual(preprocessor.get_fitted_columns(), ["value"]) + self.assertIn("fitted", repr(preprocessor)) + + standard = AutoPreprocessor(normalize="standard", columns=["value"]).fit(df) + standard_values = standard.transform(df) + np.testing.assert_allclose(standard.inverse_transform(standard_values)["value"], df["value"]) + + def test_validation_and_unfitted_errors(self): + with self.assertRaises(ValueError): + AutoPreprocessor(handle_missing="invalid") + with self.assertRaises(ValueError): + AutoPreprocessor(handle_outliers="invalid") + with self.assertRaises(ValueError): + AutoPreprocessor(normalize="invalid") + + preprocessor = AutoPreprocessor() + with self.assertRaisesRegex(RuntimeError, "not fitted"): + preprocessor.transform(pd.DataFrame({"value": [1.0]})) + with self.assertRaisesRegex(RuntimeError, "not fitted"): + preprocessor.inverse_transform(pd.DataFrame({"value": [1.0]})) + + def test_missing_requested_columns_are_ignored(self): + preprocessor = AutoPreprocessor( + handle_missing=None, handle_outliers="clip", normalize="standard", columns=["missing", "value"] + ) + result = preprocessor.fit_transform(pd.DataFrame({"value": [1.0, 2.0, 3.0]})) + self.assertEqual(list(result.columns), ["value"]) + self.assertIn("value", preprocessor.inverse_transform(result)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_data/test_timeseries.py b/tests/test_data/test_timeseries.py index 2e3fe042..0975b824 100644 --- a/tests/test_data/test_timeseries.py +++ b/tests/test_data/test_timeseries.py @@ -551,6 +551,66 @@ def test_multiple_targets(self): self.assertEqual(len(seq.target), 2) self.assertIn("value", seq.target) self.assertIn("value2", seq.target) + x, y = seq[0] + self.assertEqual(x.shape[-1], 2) + self.assertEqual(y.shape[-1], 2) + + def test_feature_columns_are_included_in_encoder_inputs(self): + seq = TimeSeriesSequence( + data=self.data, + time_idx="date", + target_column="value", + feature_columns=["feature1", "feature2"], + train_sequence_length=10, + predict_sequence_length=1, + ) + x, y = seq[0] + self.assertEqual(x.shape[-1], 3) + self.assertEqual(y.shape[-1], 1) + np.testing.assert_allclose(x[0, :, 1], self.data["feature1"].iloc[:10]) + + def test_generated_feature_columns_are_included_in_encoder_inputs(self): + config = {"date_features": {"type": "datetime", "features": ["dayofweek"], "time_col": "date"}} + seq = TimeSeriesSequence( + data=self.data, + time_idx="date", + target_column="value", + feature_columns=["date_dayofweek"], + train_sequence_length=10, + predict_sequence_length=1, + feature_config=config, + ) + x, y = seq[0] + self.assertEqual(x.shape[-1], 2) + self.assertEqual(y.shape[-1], 1) + np.testing.assert_allclose(x[0, :, 1], self.data["date"].dt.dayofweek.iloc[:10]) + + def test_one_step_horizon_and_boundary_continuity(self): + data = pd.DataFrame({"time": [0, 1, 2, 4, 5], "value": np.arange(5)}) + seq = TimeSeriesSequence( + data=data, + time_idx="time", + target_column="value", + train_sequence_length=2, + predict_sequence_length=1, + ) + self.assertEqual(len(seq.sequences), 1) + np.testing.assert_array_equal(seq.sequences[0][0][:, 0], [0, 1]) + np.testing.assert_array_equal(seq.sequences[0][1][:, 0], [2]) + + def test_inference_mode_uses_latest_complete_window(self): + data = pd.DataFrame({"time": range(5), "value": np.arange(5)}) + seq = TimeSeriesSequence( + data=data, + time_idx="time", + target_column="value", + train_sequence_length=3, + predict_sequence_length=2, + mode="inference", + ) + self.assertEqual(len(seq.sequences), 3) + np.testing.assert_array_equal(seq.sequences[-1][0][:, 0], [2, 3, 4]) + self.assertEqual(seq.sequences[-1][1].shape, (2, 1)) def test_multiple_targets_as_list(self): """Test target column provided as list.""" diff --git a/tests/test_examples/test_prediction.py b/tests/test_examples/test_prediction.py index 31a78a63..9f28d8ba 100644 --- a/tests/test_examples/test_prediction.py +++ b/tests/test_examples/test_prediction.py @@ -3,7 +3,8 @@ import tensorflow as tf -from examples.run_prediction_simple import parse_args, run_train, set_seed +from examples.run_prediction_simple import parse_args, run_manual, run_pipeline +from tfts import set_seed class PredictionTest(unittest.TestCase): @@ -24,4 +25,4 @@ class args(object): learning_rate = 0.003 set_seed(args.seed) - run_train(args) + run_manual(args) diff --git a/tests/test_features/test_auto_feature.py b/tests/test_features/test_auto_feature.py new file mode 100644 index 00000000..b4d14f6c --- /dev/null +++ b/tests/test_features/test_auto_feature.py @@ -0,0 +1,68 @@ +import unittest + +import numpy as np +import pandas as pd + +from tfts.features.auto_feature import AutoFeatureEngineer, _default_datetime_features, _default_fourier_features + + +class AutoFeatureEngineerTest(unittest.TestCase): + def setUp(self): + self.df = pd.DataFrame( + { + "date": pd.date_range("2024-01-01", periods=12, freq="D"), + "value": np.arange(12, dtype=float), + } + ) + + def test_requires_fit_before_transform(self): + engineer = AutoFeatureEngineer(lags=[1], windows=[2]) + with self.assertRaisesRegex(RuntimeError, "not fitted"): + engineer.transform(self.df) + + def test_fit_transform_adds_datetime_and_fourier_features(self): + engineer = AutoFeatureEngineer( + lags=[1], + windows=[2], + rolling_functions="all", + add_datetime=True, + add_fourier=True, + ) + result = engineer.fit_transform(self.df, time_col="date", target_col="value") + + self.assertEqual(len(result), len(self.df) - 1) + self.assertIn("value_lag_1", result.columns) + self.assertIn("value_roll_2_max", result.columns) + self.assertIn("date_month", result.columns) + self.assertIn("date_month_sin", result.columns) + self.assertEqual(len(engineer.get_feature_names()), len(result.columns) - 2) + self.assertIn("fitted", repr(engineer)) + + def test_rolling_function_variants_and_datetime_defaults(self): + self.assertEqual(AutoFeatureEngineer(rolling_functions="median")._resolve_rolling_functions(), ["median"]) + self.assertEqual( + AutoFeatureEngineer(rolling_functions=["min", "max"])._resolve_rolling_functions(), ["min", "max"] + ) + self.assertEqual(_default_fourier_features(), ["month_sin", "month_cos", "dayofweek_sin", "dayofweek_cos"]) + self.assertEqual(_default_datetime_features(pd.Series([1, 2, 3])), ["month", "dayofweek"]) + self.assertNotIn("hour", _default_datetime_features(self.df["date"])) + + subdaily = pd.Series(pd.date_range("2024-01-01", periods=3, freq="h")) + self.assertIn("hour", _default_datetime_features(subdaily)) + + def test_grouped_features_preserve_group_boundaries(self): + df = pd.DataFrame( + { + "date": list(pd.date_range("2024-01-01", periods=4, freq="D")) * 2, + "group": ["a"] * 4 + ["b"] * 4, + "value": [1.0, 2.0, 3.0, 4.0, 10.0, 20.0, 30.0, 40.0], + } + ) + engineer = AutoFeatureEngineer(lags=[1], windows=[2], group_cols=["group"]) + result = engineer.fit_transform(df, time_col="date", target_col="value") + self.assertEqual(set(result["group"]), {"a", "b"}) + self.assertEqual(result.loc[result["group"] == "b", "value_lag_1"].iloc[0], 10.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_init.py b/tests/test_init.py new file mode 100644 index 00000000..97649c31 --- /dev/null +++ b/tests/test_init.py @@ -0,0 +1,21 @@ +import importlib +import unittest + +import tfts + + +class PackageInitTest(unittest.TestCase): + def test_public_exports_and_lazy_forecasting_pipeline(self): + tfts = importlib.reload(__import__("tfts")) + self.assertEqual(tfts.__version__, "0.0.5") + self.assertIn("Trainer", tfts.__all__) + self.assertIn("BenchmarkRunner", tfts.__all__) + self.assertIs(tfts.ForecastingPipeline, tfts.ForecastingPipeline) + + def test_unknown_attribute_raises_attribute_error(self): + with self.assertRaisesRegex(AttributeError, "missing_attribute"): + _ = tfts.missing_attribute + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_models/test_auto_config.py b/tests/test_models/test_auto_config.py index fb5a0569..ed48fa68 100644 --- a/tests/test_models/test_auto_config.py +++ b/tests/test_models/test_auto_config.py @@ -1,9 +1,30 @@ import unittest +import tfts from tfts.models.auto_config import AutoConfig +from tfts.models.auto_model import AutoModel +from tfts.models.registry import list_models class TestAutoModel(unittest.TestCase): def test_auto_config(self): config = AutoConfig.for_model("bert") print(config.hidden_size) + + def test_top_level_import_exports_public_api(self): + self.assertTrue(hasattr(tfts, "AutoConfig")) + self.assertTrue(hasattr(tfts, "AutoModel")) + self.assertIn("bert", tfts.list_models()) + + def test_listed_models_have_auto_configs(self): + for model_name in list_models(): + with self.subTest(model_name=model_name): + config = AutoConfig.for_model(model_name) + self.assertEqual(config.model_type, model_name) + + def test_listed_models_can_be_instantiated(self): + for model_name in list_models(): + with self.subTest(model_name=model_name): + config = AutoConfig.for_model(model_name) + model = AutoModel.from_config(config, predict_sequence_length=2) + self.assertEqual(model.config.model_type, model_name) diff --git a/tests/test_models/test_auto_model.py b/tests/test_models/test_auto_model.py index 6da0d32b..dcbed443 100644 --- a/tests/test_models/test_auto_model.py +++ b/tests/test_models/test_auto_model.py @@ -1,3 +1,4 @@ +import tempfile import unittest import numpy as np @@ -70,3 +71,16 @@ def test_auto_model_for_uncertainty(self): x = tf.random.normal([2, 14, 4]) output = model(x) print(output.shape) + + def test_save_and_load_preserves_prediction_length(self): + config = AutoConfig.for_model("rnn") + model = AutoModel.from_config(config, predict_sequence_length=3) + model.build_model(tf.keras.Input(shape=(8, 2))) + + with tempfile.TemporaryDirectory() as tmpdir: + model.save_pretrained(tmpdir) + loaded = AutoModel.from_pretrained(tmpdir) + output = loaded(tf.random.normal([2, 8, 2])) + + self.assertEqual(loaded.predict_sequence_length, 3) + self.assertEqual(output.shape, (2, 3, 1)) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index a798582d..9de78d6a 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -8,7 +8,9 @@ import tensorflow as tf from tfts import AutoConfig, AutoModel -from tfts.trainer import BaseTrainer, KerasTrainer, Seq2seqKerasTrainer, Trainer, set_seed +from tfts.trainer import BaseTrainer, EagerTrainer, KerasTrainer, Seq2seqKerasTrainer, Trainer, set_seed +from tfts.training.runtime import configure_precision, create_distribution_strategy +from tfts.training_args import TrainingArguments class SetSeedTest(unittest.TestCase): @@ -54,8 +56,6 @@ def test_initialization_with_defaults(self): def test_initialization_with_custom_args(self): """Test BaseTrainer initialization with custom training arguments.""" - from tfts.training_args import TrainingArguments - custom_args = TrainingArguments( output_dir="./custom_output", learning_rate=0.001, per_device_train_batch_size=16 ) @@ -78,8 +78,6 @@ def test_create_optimizer(self): def test_create_lr_scheduler_linear(self): """Test linear learning rate scheduler creation.""" - from tfts.training_args import TrainingArguments - args = TrainingArguments(output_dir="./test", lr_scheduler_type="linear", max_steps=100) trainer = BaseTrainer(self.model, args=args) scheduler = trainer._create_lr_scheduler() @@ -87,8 +85,6 @@ def test_create_lr_scheduler_linear(self): def test_create_lr_scheduler_none(self): """Test that no scheduler is created when type is not specified.""" - from tfts.training_args import TrainingArguments - args = TrainingArguments(output_dir="./test", lr_scheduler_type="none", max_steps=100) trainer = BaseTrainer(self.model, args=args) scheduler = trainer._create_lr_scheduler() @@ -165,24 +161,89 @@ def test_get_inputs_invalid_type(self): def test_global_batch_size(self): """Test global batch size calculation.""" - from tfts.training_args import TrainingArguments - args = TrainingArguments(output_dir="./test", per_device_train_batch_size=8) trainer = BaseTrainer(self.model, args=args) batch_size = trainer.global_batch_size self.assertGreater(batch_size, 0) + def test_create_optimizer_uses_training_arguments(self): + """Test optimizer creation respects TrainingArguments.""" + args = TrainingArguments(output_dir="./test", learning_rate=0.002, weight_decay=0.01, lr_scheduler_type="none") + trainer = BaseTrainer(self.model, args=args) + optimizer = trainer._create_optimizer() + self.assertAlmostEqual(float(tf.keras.backend.get_value(optimizer.learning_rate)), 0.002) + def test_save_model(self): """Test model saving functionality.""" with tempfile.TemporaryDirectory() as tmpdir: trainer = BaseTrainer(self.model) + inputs = tf.keras.Input(shape=(10, 1)) + trainer.model = trainer.model.build_model(inputs) trainer._save(tmpdir) # Check that config file exists config_path = os.path.join(tmpdir, "config.json") self.assertTrue(os.path.exists(config_path)) + weights_path = os.path.join(tmpdir, "tf_model.weights.h5") + self.assertTrue(os.path.exists(weights_path)) + + def test_save_unbuilt_model_raises_clear_error(self): + """Test saving an unbuilt model fails before writing partial files.""" + with tempfile.TemporaryDirectory() as tmpdir: + trainer = BaseTrainer(self.model) + with self.assertRaisesRegex(ValueError, "cannot be saved before the model is built"): + trainer._save(tmpdir) + self.assertFalse(os.path.exists(os.path.join(tmpdir, "config.json"))) + self.assertFalse(os.path.exists(os.path.join(tmpdir, "tf_model.weights.h5"))) -class TrainerTest(unittest.TestCase): +class TrainingRuntimeTest(unittest.TestCase): + """Test training runtime configuration helpers.""" + + def tearDown(self): + tf.keras.mixed_precision.set_global_policy("float32") + + def test_training_arguments_fp16_sets_precision(self): + """Test fp16 compatibility flag maps to mixed_float16 policy.""" + args = TrainingArguments(output_dir="./test", fp16=True) + self.assertEqual(args.precision, "mixed_float16") + + def test_training_arguments_bf16_sets_precision(self): + """Test bf16 compatibility flag maps to mixed_bfloat16 policy.""" + args = TrainingArguments(output_dir="./test", bf16=True) + self.assertEqual(args.precision, "mixed_bfloat16") + + def test_training_arguments_rejects_conflicting_precision_flags(self): + """Test fp16 and bf16 cannot both be enabled.""" + with self.assertRaises(ValueError): + TrainingArguments(output_dir="./test", fp16=True, bf16=True) + + def test_configure_precision(self): + """Test precision policy is applied globally.""" + args = TrainingArguments(output_dir="./test", precision="mixed_float16") + policy = configure_precision(args) + self.assertEqual(policy.name, "mixed_float16") + self.assertEqual(tf.keras.mixed_precision.global_policy().name, "mixed_float16") + + @patch("tfts.training.runtime.tf.distribute.MirroredStrategy", create=True) + @patch("tensorflow.config.list_physical_devices") + def test_create_distribution_strategy_auto_multi_gpu(self, mock_list_devices, mock_mirrored_strategy): + """Test automatic strategy selection for multiple GPUs.""" + mock_list_devices.return_value = ["GPU:0", "GPU:1"] + strategy = create_distribution_strategy(TrainingArguments(output_dir="./test")) + mock_mirrored_strategy.assert_called_once_with() + self.assertEqual(strategy, mock_mirrored_strategy.return_value) + + @patch("tensorflow.config.list_physical_devices") + def test_create_distribution_strategy_auto_cpu(self, mock_list_devices): + """Test automatic strategy selection on CPU.""" + mock_list_devices.return_value = [] + strategy = create_distribution_strategy(TrainingArguments(output_dir="./test")) + self.assertIsInstance(strategy, tf.distribute.Strategy) + + +class EagerTrainerTest(unittest.TestCase): + """Tests for EagerTrainer (legacy custom training loop).""" + def setUp(self): self.fit_config = { "epochs": 2, @@ -210,7 +271,7 @@ def test_trainer_basic(self): # 1gpu, no dist config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer( + trainer = EagerTrainer( model, ) trainer.train( @@ -226,7 +287,7 @@ def test_trainer_fit_alias(self): """Test that fit() is an alias for train().""" config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer(model) + trainer = EagerTrainer(model) # fit should work the same as train trainer.fit( @@ -240,7 +301,7 @@ def test_trainer_without_validation(self): """Test training without validation data.""" config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer(model) + trainer = EagerTrainer(model) trainer.train( train_loader=self.train_loader, valid_loader=None, optimizer=tf.keras.optimizers.Adam(0.003), epochs=1 @@ -250,7 +311,7 @@ def test_trainer_with_lr_scheduler(self): """Test trainer with learning rate scheduler.""" config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer(model) + trainer = EagerTrainer(model) lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay( initial_learning_rate=0.003, decay_steps=10, decay_rate=0.9 @@ -264,7 +325,7 @@ def test_trainer_with_ema(self): """Test trainer with exponential moving average.""" config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer(model) + trainer = EagerTrainer(model) trainer.train(train_loader=self.train_loader, valid_loader=self.valid_loader, use_ema=True, epochs=1) @@ -272,7 +333,7 @@ def test_trainer_with_multiple_metrics(self): """Test trainer with multiple evaluation metrics.""" config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer(model) + trainer = EagerTrainer(model) metrics = [ lambda x, y: np.mean(np.abs(x.numpy() - y.numpy())), @@ -285,7 +346,7 @@ def test_trainer_early_stopping(self): """Test early stopping functionality.""" config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer(model) + trainer = EagerTrainer(model) trainer.train( train_loader=self.train_loader, @@ -299,7 +360,7 @@ def test_trainer_gradient_clipping(self): """Test gradient clipping with custom max_grad_norm.""" config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer(model) + trainer = EagerTrainer(model) trainer.train(train_loader=self.train_loader, valid_loader=self.valid_loader, max_grad_norm=1.0, epochs=1) @@ -307,7 +368,7 @@ def test_trainer_custom_loss(self): """Test trainer with custom loss function.""" config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer(model) + trainer = EagerTrainer(model) custom_loss = tf.keras.losses.MeanAbsoluteError() @@ -317,14 +378,14 @@ def test_trainer_2gpu(self): strategy = tf.distribute.MirroredStrategy() config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer(model, strategy=strategy) + trainer = EagerTrainer(model, strategy=strategy) trainer.train(self.train_loader, self.valid_loader, **self.fit_config) def test_trainer_kwargs(self): """Test that custom kwargs are set as attributes.""" config = AutoConfig.for_model("rnn") model = AutoModel.from_config(config, predict_sequence_length=2) - trainer = Trainer(model, custom_param="test_value", another_param=42) + trainer = EagerTrainer(model, custom_param="test_value", another_param=42) self.assertEqual(trainer.custom_param, "test_value") self.assertEqual(trainer.another_param, 42) @@ -491,6 +552,40 @@ def test_get_model(self): retrieved_model = trainer.get_model() self.assertIsInstance(retrieved_model, tf.keras.Model) + def test_evaluate_predict_and_default_task_helpers(self): + model = tf.keras.Sequential([tf.keras.Input(shape=(4, 1)), tf.keras.layers.Dense(1)]) + trainer = KerasTrainer(model) + x = np.random.random((2, 4, 1)).astype(np.float32) + y = np.random.random((2, 4, 1)).astype(np.float32) + + list_metrics = trainer.evaluate((x, y), metrics=["mae"]) + dataset = tf.data.Dataset.from_tensor_slices((x, y)).batch(1) + dataset_metrics = trainer.evaluate(dataset, metrics=["mse"]) + np.testing.assert_equal(set(list_metrics), {"mae"}) + np.testing.assert_equal(set(dataset_metrics), {"mse"}) + self.assertEqual(trainer.predict(dataset).shape, x.shape) + with self.assertRaises(TypeError): + trainer.evaluate("invalid") + + trainer._task = "classification" + self.assertIsInstance(trainer._default_loss(), tf.keras.losses.SparseCategoricalCrossentropy) + self.assertEqual(trainer._default_metrics(), ["accuracy"]) + trainer._task = "forecasting" + self.assertIsInstance(trainer._default_loss(), tf.keras.losses.MeanSquaredError) + self.assertEqual(trainer._default_metrics(), ["mae"]) + + def test_build_callbacks_covers_optional_callbacks(self): + with tempfile.TemporaryDirectory() as tmpdir: + callbacks = KerasTrainer._build_callbacks( + early_stopping_patience=1, + checkpoint_dir=tmpdir, + reduce_lr_patience=2, + ) + self.assertEqual(len(callbacks), 3) + self.assertIsInstance(callbacks[0], tf.keras.callbacks.EarlyStopping) + self.assertIsInstance(callbacks[1], tf.keras.callbacks.ModelCheckpoint) + self.assertIsInstance(callbacks[2], tf.keras.callbacks.ReduceLROnPlateau) + def test_plot(self): """Test plot functionality.""" config = AutoConfig.for_model("rnn") diff --git a/tests/test_tuner.py b/tests/test_tuner.py new file mode 100644 index 00000000..3cf0f55e --- /dev/null +++ b/tests/test_tuner.py @@ -0,0 +1,100 @@ +import unittest +from unittest.mock import Mock, patch + +import numpy as np +import tensorflow as tf + +from tfts.tuner import optuna_tuner +from tfts.tuner.optuna_tuner import OptunaTuner + + +class FakeTrial: + def suggest_categorical(self, name, choices): + return choices[0] + + def suggest_float(self, name, low, high, log=False): + return (low, high, log) + + def suggest_int(self, name, low, high): + return low + + +class OptunaTunerTest(unittest.TestCase): + def setUp(self): + data = (np.zeros((2, 4, 1)), np.zeros((2, 1, 1))) + self.tuner = OptunaTuner(data, data, predict_sequence_length=1) + + def test_parameter_suggestions_cover_supported_specs(self): + params = self.tuner._suggest_params( + FakeTrial(), + { + "model_type": ["rnn", "dlinear"], + "learning_rate": [0.0001, 0.01], + "dropout": [1.0, 2.0], + "num_layers": [1, 3], + }, + ) + self.assertEqual(params["model_type"], "rnn") + self.assertEqual(params["learning_rate"], (0.0001, 0.01, True)) + self.assertEqual(params["dropout"], (1.0, 2.0, False)) + self.assertEqual(params["num_layers"], 1) + + with self.assertRaises(ValueError): + self.tuner._suggest_params(FakeTrial(), {"bad": [1, 2, 3]}) + + def test_best_accessors_and_score_extraction(self): + self.assertIsNone(self.tuner.get_study()) + with self.assertRaisesRegex(RuntimeError, "No search"): + self.tuner.get_best_params() + with self.assertRaisesRegex(RuntimeError, "No search"): + self.tuner.get_best_score() + + self.assertEqual(self.tuner._extract_score({"mse": 0.25}), 0.25) + callable_tuner = OptunaTuner(self.tuner.train_data, self.tuner.valid_data, metric=lambda y, p: 1.0) + self.assertEqual(callable_tuner._extract_score({"mae": 0.5}), 0.5) + with self.assertRaises(ValueError): + self.tuner._extract_score({}) + self.assertEqual(repr(self.tuner), "OptunaTuner(metric='mse', direction='minimize')") + + def test_objective_builds_and_evaluates_a_trial(self): + config = type("Config", (), {"hidden_size": 8})() + fake_trainer = Mock() + fake_trainer.evaluate.return_value = {"mse": 0.125} + with patch.object(optuna_tuner.AutoConfig, "for_model", return_value=config), patch.object( + optuna_tuner.AutoModel, "from_config", return_value=object() + ), patch.object(optuna_tuner, "Trainer", return_value=fake_trainer), patch.object( + optuna_tuner, "_default_optimizer", return_value=Mock() + ): + score = self.tuner._objective( + FakeTrial(), {"model_type": ["rnn"], "hidden_size": [4, 8], "unknown": [1, 2]}, 1, 0 + ) + + self.assertEqual(score, 0.125) + self.assertEqual(config.hidden_size, 4) + fake_trainer.train.assert_called_once() + + def test_search_stores_study_and_supports_optional_dependency_error(self): + study = Mock() + study.best_params = {"model_type": "rnn"} + study.best_value = 0.2 + fake_optuna = Mock() + fake_optuna.logging.WARNING = "warning" + fake_optuna.create_study.return_value = study + with patch.object(optuna_tuner, "_require_optuna", return_value=fake_optuna), patch.object( + self.tuner, "_objective", return_value=0.2 + ): + self.assertEqual(self.tuner.search({"model_type": ["rnn"]}, n_trials=1), ({"model_type": "rnn"}, 0.2)) + self.assertIs(self.tuner.get_study(), study) + self.assertEqual(self.tuner.get_best_params(), {"model_type": "rnn"}) + self.assertEqual(self.tuner.get_best_score(), 0.2) + + with self.assertRaisesRegex(ImportError, "optuna is required"): + optuna_tuner._require_optuna() + + def test_default_optimizer_is_a_tensorflow_optimizer(self): + optimizer = optuna_tuner._default_optimizer(0.001) + self.assertIsInstance(optimizer, tf.keras.optimizers.Optimizer) + + +if __name__ == "__main__": + unittest.main() diff --git a/tfts/__init__.py b/tfts/__init__.py index 652fad8f..009164ee 100644 --- a/tfts/__init__.py +++ b/tfts/__init__.py @@ -1,6 +1,17 @@ -"""tfts package for time series prediction with TensorFlow""" +"""TFTS — Deep Learning for Time Series. -from tfts.data import TimeSeriesSequence, get_data +Usage: + >>> import tfts + >>> pipe = tfts.pipeline("forecasting", model="dlinear", lookback=96, horizon=24) + >>> pipe.fit(df, target_col="value", epochs=50) + >>> preds = pipe.predict(steps=24) + >>> tfts.list_models() +""" + +from tfts.cli import pipeline +from tfts.data import AutoPreprocessor, DataProcessor, TimeSeriesSequence, get_data +from tfts.features import AutoFeatureEngineer, FeatureRegistry +from tfts.metrics import evaluate as evaluate_metrics from tfts.models.auto_config import AutoConfig from tfts.models.auto_model import ( AutoModel, @@ -10,11 +21,56 @@ AutoModelForSegmentation, AutoModelForUncertainty, ) +from tfts.models.registry import list_models + +# Legacy compatibility from tfts.tasks.pipeline import Pipeline -from tfts.trainer import KerasTrainer, Trainer, set_seed +from tfts.trainer import EagerTrainer, KerasTrainer, Trainer, set_seed from tfts.training_args import TrainingArguments +from tfts.tuner import OptunaTuner + +try: + import sys + + import benchmark as _benchmark + from benchmark import BenchmarkConfig, BenchmarkResults, BenchmarkRunner, Dataset, DatasetRegistry, ModelRegistry + import benchmark.base as _benchmark_base + import benchmark.datasets as _benchmark_datasets + import benchmark.formatter as _benchmark_formatter + import benchmark.metrics as _benchmark_metrics + import benchmark.registry as _benchmark_registry + import benchmark.runner as _benchmark_runner + + sys.modules.setdefault("tfts.benchmark", _benchmark) + sys.modules.setdefault("tfts.benchmark.base", _benchmark_base) + sys.modules.setdefault("tfts.benchmark.datasets", _benchmark_datasets) + sys.modules.setdefault("tfts.benchmark.formatter", _benchmark_formatter) + sys.modules.setdefault("tfts.benchmark.metrics", _benchmark_metrics) + sys.modules.setdefault("tfts.benchmark.registry", _benchmark_registry) + sys.modules.setdefault("tfts.benchmark.runner", _benchmark_runner) + + _BENCHMARK_EXPORTS = [ + "BenchmarkConfig", + "BenchmarkResults", + "BenchmarkRunner", + "Dataset", + "DatasetRegistry", + "ModelRegistry", + ] +except (ImportError, ModuleNotFoundError): + _BENCHMARK_EXPORTS = [] __all__ = [ + # -- Primary API -- + "pipeline", + "ForecastingPipeline", + "DataProcessor", + # -- Preprocessing -- + "AutoPreprocessor", + # -- Features -- + "AutoFeatureEngineer", + "FeatureRegistry", + # -- Models -- "AutoModel", "AutoModelForPrediction", "AutoModelForClassification", @@ -22,12 +78,30 @@ "AutoModelForAnomaly", "AutoModelForUncertainty", "AutoConfig", + "list_models", + # -- Training -- "Trainer", "KerasTrainer", + "EagerTrainer", "TrainingArguments", - "set_seed" "Pipeline", + "set_seed", + # -- Tuning -- + "OptunaTuner", + # -- Data -- "get_data", "TimeSeriesSequence", -] + # -- Evaluation -- + "evaluate_metrics", + # -- Legacy -- + "Pipeline", +] + _BENCHMARK_EXPORTS + +__version__ = "0.0.5" + + +def __getattr__(name: str): + if name == "ForecastingPipeline": + from tfts.cli.forecasting import ForecastingPipeline -__version__ = "0.0.0" + return ForecastingPipeline + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tfts/cli/__init__.py b/tfts/cli/__init__.py new file mode 100644 index 00000000..ddcc8509 --- /dev/null +++ b/tfts/cli/__init__.py @@ -0,0 +1,101 @@ +"""TFTS Pipeline API — the primary entry point for users. + +Provides a ``pipeline()`` factory function that returns the right +Pipeline subclass for the requested task. + +Examples: + >>> import tfts + >>> forecaster = tfts.pipeline("forecasting", model="patch_tst", + ... lookback=96, horizon=24) + >>> forecaster.fit(df, target_col="sales") + >>> preds = forecaster.predict(steps=24) +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Union + +from ..models.auto_config import AutoConfig +from ..models.auto_model import AutoModel + +if TYPE_CHECKING: + from .forecasting import ForecastingPipeline + +__all__ = ["pipeline", "ForecastingPipeline"] + + +def __getattr__(name: str): + if name == "ForecastingPipeline": + from .forecasting import ForecastingPipeline + + return ForecastingPipeline + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def pipeline( + task: str = "forecasting", + model: Union[str, AutoModel] = "dlinear", + lookback: int = 96, + horizon: int = 24, + config: Optional[AutoConfig] = None, + batch_size: int = 32, + normalize: Optional[str] = "minmax", + learning_rate: float = 1e-3, + epochs: int = 50, + early_stopping_patience: Optional[int] = 5, + seed: int = 42, + **kwargs, +) -> "ForecastingPipeline": + """Create a pipeline for a time series task. + + This is the recommended entry point. It returns a Pipeline object that + handles data preparation, model building, training, and prediction. + + Args: + task: Task type — ``'forecasting'`` (more tasks coming soon). + model: Model name (e.g. ``'patch_tst'``, ``'transformer'``, + ``'dlinear'``, ``'nbeats'``) or an ``AutoModel`` instance. + lookback: Number of past steps to use as input. + horizon: Number of future steps to predict. + config: Optional ``AutoConfig`` for fine-grained control. + batch_size: Batch size for training. + normalize: Normalization — ``'minmax'``, ``'standard'``, or ``None``. + learning_rate: Optimizer learning rate. + epochs: Number of training epochs. + early_stopping_patience: Early stopping patience. + seed: Random seed. + **kwargs: Passed to the model config. + + Returns: + A Pipeline object ready for ``.fit()`` and ``.predict()``. + + Raises: + ValueError: If the task is unknown. + + Examples: + >>> pipe = tfts.pipeline("forecasting", model="dlinear", + ... lookback=96, horizon=24) + >>> pipe.fit(df, target_col="sales", epochs=50) + >>> preds = pipe.predict(steps=24) + """ + from .forecasting import ForecastingPipeline + + if task == "forecasting": + return ForecastingPipeline( + model=model, + lookback=lookback, + horizon=horizon, + config=config, + batch_size=batch_size, + normalize=normalize, + learning_rate=learning_rate, + epochs=epochs, + early_stopping_patience=early_stopping_patience, + seed=seed, + **kwargs, + ) + + raise ValueError( + f"Unknown task '{task}'. Currently supported: 'forecasting'. " + f"More tasks (classification, anomaly) are coming soon." + ) diff --git a/tfts/cli/forecasting.py b/tfts/cli/forecasting.py new file mode 100644 index 00000000..c57462b1 --- /dev/null +++ b/tfts/cli/forecasting.py @@ -0,0 +1,304 @@ +"""Forecasting CLI and end-to-end forecasting workflow. + +Provides the highest-level API for time series forecasting, wrapping +DataProcessor, AutoModel, and Trainer behind a single interface. +""" + +import argparse +import logging +from typing import Dict, Optional, Sequence, Union + +import numpy as np +import pandas as pd +import tensorflow as tf + +from ..data import get_data +from ..data.processor import DataProcessor +from ..models.auto_config import AutoConfig +from ..models.auto_model import AutoModel +from ..trainer import Trainer, set_seed + +logger = logging.getLogger(__name__) + +__all__ = ["ForecastingPipeline", "main"] + + +class ForecastingPipeline: + """End-to-end forecasting pipeline with a transformers-like API. + + Handles data preparation, model building, training, prediction, and + evaluation in a single object. + + Args: + model: Model name (e.g. ``'patch_tst'``, ``'transformer'``) or an + ``AutoModel`` instance. + lookback: Number of past time steps used as input. + horizon: Number of future time steps to predict. + config: Optional AutoConfig for fine-grained control. + batch_size: Batch size for training. + normalize: Normalization — ``'minmax'``, ``'standard'``, or ``None``. + learning_rate: Learning rate for the optimizer. + epochs: Default number of training epochs. + early_stopping_patience: Patience for early stopping. + seed: Random seed for reproducibility. + **kwargs: Additional arguments passed to the model config. + + Examples: + >>> import tfts + >>> pipe = tfts.pipeline("forecasting", model="dlinear", + ... lookback=96, horizon=24) + >>> pipe.fit(df, target_col="value", epochs=50) + >>> preds = pipe.predict(steps=24) + + >>> # Evaluate on holdout + >>> metrics = pipe.evaluate(test_df) + """ + + def __init__( + self, + model: Union[str, AutoModel] = "dlinear", + lookback: int = 96, + horizon: int = 24, + config: Optional[AutoConfig] = None, + batch_size: int = 32, + normalize: Optional[str] = "minmax", + learning_rate: float = 1e-3, + epochs: int = 50, + early_stopping_patience: Optional[int] = 5, + seed: int = 42, + **kwargs, + ): + self.lookback = lookback + self.horizon = horizon + self.batch_size = batch_size + self.normalize = normalize + self.learning_rate = learning_rate + self.epochs = epochs + self.early_stopping_patience = early_stopping_patience + self.seed = seed + + set_seed(seed) + + # Model setup + if isinstance(model, str): + self.config = config or AutoConfig.for_model(model) + self.config.update(kwargs) + self._model = AutoModel.from_config(self.config, predict_sequence_length=horizon) + else: + self._model = model + self.config = getattr(model, "config", config) + + self.model_name = getattr(self.config, "model_type", model if isinstance(model, str) else "custom") + + # Data processor + self.processor = DataProcessor( + lookback=lookback, + horizon=horizon, + batch_size=batch_size, + normalize=normalize, + ) + + # Will be set during fit() + self.trainer: Trainer = Trainer(self._model) + self._target_col: Optional[str] = None + self._time_col: Optional[str] = None + self._fitted: bool = False + + logger.info(f"Pipeline ready: model={self.model_name}, lookback={lookback}, horizon={horizon}") + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def fit( + self, + df: pd.DataFrame, + target_col: Optional[str] = None, + time_col: Optional[str] = None, + validation_split: float = 0.2, + epochs: Optional[int] = None, + verbose: int = 1, + **trainer_kwargs, + ) -> tf.keras.callbacks.History: + """Train the forecasting model. + + Args: + df: Time series DataFrame. + target_col: Column to forecast. Auto-detected if None. + time_col: Time column. Auto-detected if None. + validation_split: Fraction of data for validation. + epochs: Training epochs (overrides constructor default). + verbose: 0=silent, 1=progress bar, 2=one line. + **trainer_kwargs: Passed to ``Trainer.train()``. + + Returns: + Keras History object. + """ + self._target_col = target_col + self._time_col = time_col + + # Prepare data + self.processor.validation_split = validation_split + result = self.processor.prepare(df, target_col=target_col, time_col=time_col) + + if isinstance(result, tuple): + train_ds, valid_ds = result + else: + train_ds, valid_ds = result, None + + # Train + epochs = epochs or self.epochs + history = self.trainer.train( + train_ds, + valid_dataset=valid_ds, + epochs=epochs, + batch_size=self.batch_size, + verbose=verbose, + early_stopping_patience=self.early_stopping_patience, + **{k: v for k, v in trainer_kwargs.items() if k != "early_stopping_patience"}, + ) + + self._fitted = True + return history + + def predict(self, steps: Optional[int] = None, df: Optional[pd.DataFrame] = None) -> np.ndarray: + """Generate forecasts. + + Args: + steps: Number of steps to predict. Defaults to ``self.horizon``. + df: DataFrame with new data to predict from (uses training data if None). + + Returns: + Numpy array of shape ``(n_series, steps, n_targets)``. + """ + if not self._fitted: + raise RuntimeError("Pipeline must be fitted before prediction. Call .fit() first.") + + steps = steps or self.horizon + + if df is not None: + ds = self.processor.prepare_for_inference(df, target_col=self._target_col, time_col=self._time_col) + preds = self.trainer.predict(ds) + else: + # Use the model directly — user is responsible for input shape + raise ValueError( + "Please pass `df` with the latest lookback-length data for prediction, e.g. " + "pipeline.predict(steps=24, df=recent_data)" + ) + + # Inverse transform + preds = self.processor.inverse_transform(preds) + + # Trim to requested steps + if preds.shape[1] > steps: + preds = preds[:, :steps, :] + + return preds + + def evaluate( + self, df: pd.DataFrame, target_col: Optional[str] = None, time_col: Optional[str] = None + ) -> Dict[str, float]: + """Evaluate the model on a test DataFrame. + + Args: + df: Test DataFrame. + target_col: Target column name. + time_col: Time column name. + + Returns: + Dict of metric_name -> value. + """ + if self.trainer is None: + raise RuntimeError("Pipeline must be fitted before evaluation. Call .fit() first.") + + # Prepare the test data using the same processor params + processor = DataProcessor( + lookback=self.lookback, horizon=self.horizon, batch_size=self.batch_size, normalize=None + ) + ds = processor.prepare(df, target_col=target_col or self._target_col, time_col=time_col or self._time_col) + if isinstance(ds, tuple): + ds = ds[0] # use first split + + return self.trainer.evaluate(ds) + + def save(self, path: str) -> None: + """Save the trained model.""" + if self.trainer is None: + raise RuntimeError("Nothing to save — train the pipeline first.") + self.trainer.save_model(path) + logger.info(f"Pipeline saved to {path}") + + def summary(self) -> None: + """Print a summary of the pipeline.""" + line = "=" * 60 + print(line) + print("TFTS Forecasting Pipeline") + print(line) + print(f" Model : {self.model_name}") + print(f" Lookback : {self.lookback}") + print(f" Horizon : {self.horizon}") + print(f" Batch size : {self.batch_size}") + print(f" Normalize : {self.normalize}") + print(f" Fitted : {self._fitted}") + print(line) + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + """Parse forecasting CLI arguments.""" + parser = argparse.ArgumentParser(description="Train a TFTS forecasting model on a built-in dataset.") + parser.add_argument("--model", type=str, default="dlinear", help="Model name, e.g. dlinear, rnn, transformer.") + parser.add_argument("--data", type=str, default="sine", help="Built-in dataset name, e.g. sine or airpassengers.") + parser.add_argument("--lookback", type=int, default=24, help="Input sequence length.") + parser.add_argument("--horizon", type=int, default=12, help="Prediction sequence length.") + parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs.") + parser.add_argument("--batch-size", type=int, default=16, help="Training batch size.") + parser.add_argument("--learning-rate", type=float, default=1e-3, help="Adam learning rate.") + parser.add_argument("--test-size", type=float, default=0.2, help="Validation split ratio.") + parser.add_argument("--early-stopping-patience", type=int, default=5, help="Early stopping patience.") + parser.add_argument("--seed", type=int, default=315, help="Random seed.") + parser.add_argument("--output-dir", type=str, default=None, help="Optional directory for saving model weights.") + parser.add_argument("--verbose", type=int, default=1, choices=[0, 1, 2], help="Keras training verbosity.") + return parser.parse_args(argv) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + """Train and evaluate a forecasting model from the command line.""" + args = parse_args(argv) + set_seed(args.seed) + + train_data, valid_data = get_data( + args.data, + train_length=args.lookback, + predict_sequence_length=args.horizon, + test_size=args.test_size, + ) + + config = AutoConfig.for_model(args.model) + model = AutoModel.from_config(config, predict_sequence_length=args.horizon) + trainer = Trainer(model) + optimizer = tf.keras.optimizers.Adam(args.learning_rate) + + trainer.train( + train_data, + valid_data, + optimizer=optimizer, + epochs=args.epochs, + batch_size=args.batch_size, + verbose=args.verbose, + early_stopping_patience=args.early_stopping_patience, + ) + + metrics = trainer.evaluate(valid_data) + preds = trainer.predict(valid_data[0]) + print(f"metrics: {metrics}") + print(f"predictions shape: {preds.shape}") + + if args.output_dir: + trainer.save_model(args.output_dir) + print(f"saved model to: {args.output_dir}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tfts/constants.py b/tfts/constants.py index 0321f8f1..fb436080 100644 --- a/tfts/constants.py +++ b/tfts/constants.py @@ -1,3 +1,10 @@ +"""TFTS constants — cache paths and weight file names. + +All paths respect the ``TFTS_HOME`` environment variable (defaults to +``~/.cache/tfts``). Model weights, datasets, and assets are stored in +subdirectories under this root. +""" + import os # default cache @@ -15,7 +22,7 @@ default_assets_cache_path = os.path.join(TFTS_HOME, "assets") TFTS_HUB_CACHE = os.getenv("TFTS_HUB_CACHE", default_cache_path) -TFTS_DATASETS_CACHE = os.getenv("TFTS_DATASETS_CACHE", default_assets_cache_path) +TFTS_DATASETS_CACHE = os.getenv("TFTS_DATASETS_CACHE", default_datasets_path) TFTS_ASSETS_CACHE = os.getenv("TFTS_ASSETS_CACHE", default_assets_cache_path) TF2_WEIGHTS_NAME = "tf_model.weights.h5" diff --git a/tfts/data/__init__.py b/tfts/data/__init__.py index 91147013..756c70bf 100644 --- a/tfts/data/__init__.py +++ b/tfts/data/__init__.py @@ -1,4 +1,15 @@ """tfts data""" +from .auto_preprocessor import AutoPreprocessor from .get_data import get_air_passengers, get_data, get_sine +from .processor import DataProcessor from .timeseries import TimeSeriesSequence + +__all__ = [ + "AutoPreprocessor", + "DataProcessor", + "TimeSeriesSequence", + "get_air_passengers", + "get_data", + "get_sine", +] diff --git a/tfts/data/auto_preprocessor.py b/tfts/data/auto_preprocessor.py new file mode 100644 index 00000000..2ae05473 --- /dev/null +++ b/tfts/data/auto_preprocessor.py @@ -0,0 +1,241 @@ +"""AutoPreprocessor — sklearn-style data preprocessing for time series. + +Provides a single class that handles missing values, outlier clipping, +and normalization with fit / transform / inverse_transform semantics. +Fitted parameters are stored so that transform can be applied to new +data (e.g. inference) and inverse_transform can return predictions +to the original scale. +""" + +import logging +from typing import Dict, List, Optional, Union + +import numpy as np +import pandas as pd + +logger = logging.getLogger(__name__) + + +class AutoPreprocessor: + """Automatic data preprocessor for time series DataFrames. + + Handles three common preprocessing steps in order: + 1. Missing values (``handle_missing``) + 2. Outliers (``handle_outliers``) + 3. Normalization (``normalize``) + + Follow the sklearn convention: call :meth:`fit_transform` on training + data, then :meth:`transform` on test / inference data, and + :meth:`inverse_transform` on predictions to get the original scale. + + Args: + handle_missing: How to fill NaN values. + - ``'ffill'``: forward fill, then back-fill leading NaNs. + - ``'interpolate'``: linear interpolation, then edge-fill. + - ``'drop'``: drop rows containing NaN. + - ``None``: skip. + handle_outliers: How to cap extreme values. + - ``'clip'``: clip to ``[Q1 - 1.5*IQR, Q3 + 1.5*IQR]`` per column. + - ``None``: skip. + normalize: Normalization method. + - ``'standard'``: ``(x - mean) / std``. + - ``'minmax'``: ``(x - min) / (max - min)``. + - ``None``: skip. + columns: Column names to process. If ``None``, all numeric columns + are selected during :meth:`fit`. + + Examples: + >>> from tfts.data import AutoPreprocessor + >>> pre = AutoPreprocessor(handle_missing="interpolate", + ... handle_outliers="clip", + ... normalize="standard") + >>> df_clean = pre.fit_transform(df) + >>> df_orig = pre.inverse_transform(df_clean) + """ + + _VALID_MISSING = ("ffill", "interpolate", "drop", None) + _VALID_OUTLIERS = ("clip", None) + _VALID_NORMALIZE = ("standard", "minmax", None) + + def __init__( + self, + handle_missing: Optional[str] = "ffill", + handle_outliers: Optional[str] = None, + normalize: Optional[str] = None, + columns: Optional[List[str]] = None, + ) -> None: + if handle_missing not in self._VALID_MISSING: + raise ValueError(f"handle_missing must be one of {self._VALID_MISSING}, got {handle_missing!r}") + if handle_outliers not in self._VALID_OUTLIERS: + raise ValueError(f"handle_outliers must be one of {self._VALID_OUTLIERS}, got {handle_outliers!r}") + if normalize not in self._VALID_NORMALIZE: + raise ValueError(f"normalize must be one of {self._VALID_NORMALIZE}, got {normalize!r}") + + self.handle_missing = handle_missing + self.handle_outliers = handle_outliers + self.normalize = normalize + self.columns = columns + + # Fitted parameters (populated by fit) + self._fitted_columns: List[str] = [] + self._clip_bounds: Dict[str, tuple] = {} + self._norm_params: Dict[str, dict] = {} + self._fitted: bool = False + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def fit(self, df: pd.DataFrame) -> "AutoPreprocessor": + """Fit the preprocessor on *df* (learn clip bounds & norm params). + + Args: + df: Training DataFrame. + + Returns: + self (for chaining). + """ + cols = self.columns if self.columns is not None else df.select_dtypes(include=[np.number]).columns.tolist() + self._fitted_columns = list(cols) + + # Learn outlier bounds + if self.handle_outliers == "clip": + for col in self._fitted_columns: + if col not in df.columns: + continue + q1 = df[col].quantile(0.25) + q3 = df[col].quantile(0.75) + iqr = q3 - q1 + self._clip_bounds[col] = (q1 - 1.5 * iqr, q3 + 1.5 * iqr) + + # Learn normalization parameters on the *already cleaned* data + # so that outliers don't skew mean/std. We apply the same + # cleaning steps first, then compute stats. + df_clean = self._fill_missing(df) + df_clean = self._clip_outliers(df_clean) + + if self.normalize == "standard": + for col in self._fitted_columns: + if col not in df_clean.columns: + continue + self._norm_params[col] = {"mean": df_clean[col].mean(), "std": df_clean[col].std()} + elif self.normalize == "minmax": + for col in self._fitted_columns: + if col not in df_clean.columns: + continue + self._norm_params[col] = {"min": df_clean[col].min(), "max": df_clean[col].max()} + + self._fitted = True + return self + + def transform(self, df: pd.DataFrame) -> pd.DataFrame: + """Apply fitted preprocessing to *df*. + + Args: + df: DataFrame to transform. + + Returns: + Transformed copy of *df*. + + Raises: + RuntimeError: If :meth:`fit` has not been called. + """ + self._check_fitted() + result = df.copy() + result = self._fill_missing(result) + result = self._clip_outliers(result) + result = self._apply_normalization(result) + return result + + def fit_transform(self, df: pd.DataFrame) -> pd.DataFrame: + """Convenience: fit then transform in one call.""" + return self.fit(df).transform(df) + + def inverse_transform(self, df: pd.DataFrame) -> pd.DataFrame: + """Reverse the normalization step only. + + Missing-value filling and outlier clipping are not reversible, + so only the normalization is undone. + + Args: + df: Normalized DataFrame (or a subset of columns). + + Returns: + DataFrame with normalization reversed. + """ + self._check_fitted() + result = df.copy() + for col in self._fitted_columns: + if col not in result.columns or col not in self._norm_params: + continue + params = self._norm_params[col] + if self.normalize == "standard": + result[col] = result[col] * params["std"] + params["mean"] + elif self.normalize == "minmax": + result[col] = result[col] * (params["max"] - params["min"]) + params["min"] + return result + + def get_fitted_columns(self) -> List[str]: + """Return the list of columns this preprocessor was fitted on.""" + return list(self._fitted_columns) + + # ------------------------------------------------------------------ + # Internal steps + # ------------------------------------------------------------------ + + def _fill_missing(self, df: pd.DataFrame) -> pd.DataFrame: + """Handle missing values in the fitted columns only.""" + if self.handle_missing is None: + return df + + if self.handle_missing == "ffill": + result = df.copy() + columns = self._fitted_columns or list(result.columns) + result[columns] = result[columns].ffill().bfill() + return result + elif self.handle_missing == "interpolate": + result = df.copy() + result[self._fitted_columns] = ( + result[self._fitted_columns].interpolate(limit_direction="both").bfill().ffill() + ) + return result + elif self.handle_missing == "drop": + return df.dropna(subset=self._fitted_columns if self._fitted_columns else None) + return df + + def _clip_outliers(self, df: pd.DataFrame) -> pd.DataFrame: + """Clip outliers using fitted bounds.""" + if self.handle_outliers != "clip": + return df + + for col, (lower, upper) in self._clip_bounds.items(): + if col in df.columns: + df[col] = df[col].clip(lower, upper) + return df + + def _apply_normalization(self, df: pd.DataFrame) -> pd.DataFrame: + """Normalize using fitted parameters.""" + if self.normalize is None: + return df + + for col in self._fitted_columns: + if col not in df.columns or col not in self._norm_params: + continue + params = self._norm_params[col] + if self.normalize == "standard": + df[col] = (df[col] - params["mean"]) / (params["std"] + 1e-8) + elif self.normalize == "minmax": + df[col] = (df[col] - params["min"]) / (params["max"] - params["min"] + 1e-8) + return df + + def _check_fitted(self) -> None: + if not self._fitted: + raise RuntimeError("AutoPreprocessor is not fitted yet. Call .fit() or .fit_transform() first.") + + def __repr__(self) -> str: + status = "fitted" if self._fitted else "not fitted" + return ( + f"AutoPreprocessor(handle_missing={self.handle_missing!r}, " + f"handle_outliers={self.handle_outliers!r}, " + f"normalize={self.normalize!r}, {status})" + ) diff --git a/tfts/data/get_data.py b/tfts/data/get_data.py index 57d1a2c6..308611c3 100644 --- a/tfts/data/get_data.py +++ b/tfts/data/get_data.py @@ -16,6 +16,156 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) +AIR_PASSENGERS_VALUES = np.array( + [ + 112, + 118, + 132, + 129, + 121, + 135, + 148, + 148, + 136, + 119, + 104, + 118, + 115, + 126, + 141, + 135, + 125, + 149, + 170, + 170, + 158, + 133, + 114, + 140, + 145, + 150, + 178, + 163, + 172, + 178, + 199, + 199, + 184, + 162, + 146, + 166, + 171, + 180, + 193, + 181, + 183, + 218, + 230, + 242, + 209, + 191, + 172, + 194, + 196, + 196, + 236, + 235, + 229, + 243, + 264, + 272, + 237, + 211, + 180, + 201, + 204, + 188, + 235, + 227, + 234, + 264, + 302, + 293, + 259, + 229, + 203, + 229, + 242, + 233, + 267, + 269, + 270, + 315, + 364, + 347, + 312, + 274, + 237, + 278, + 284, + 277, + 317, + 313, + 318, + 374, + 413, + 405, + 355, + 306, + 271, + 306, + 315, + 301, + 356, + 348, + 355, + 422, + 465, + 467, + 404, + 347, + 305, + 336, + 340, + 318, + 362, + 348, + 363, + 435, + 491, + 505, + 404, + 359, + 310, + 337, + 360, + 342, + 406, + 396, + 420, + 472, + 548, + 559, + 463, + 407, + 362, + 405, + 417, + 391, + 419, + 461, + 472, + 535, + 622, + 606, + 508, + 461, + 390, + 432, + ], + dtype=np.float32, +) + TS_DATASETS_URL = { "air_passengers": { @@ -129,14 +279,14 @@ def get_sine( x_array = np.array(x)[:, :, 0:1] y_array = np.array(y)[:, :, 0:1] - logging.info("Load sine data", x_array.shape, y_array.shape) + logging.info(f"Load sine data {x_array.shape} {y_array.shape}") if test_size > 0: - slice = int(n_examples * (1 - test_size)) - x_train = x_array[:slice] - y_train = y_array[:slice] - x_valid = x_array[slice:] - y_valid = y_array[slice:] + split_idx = int(n_examples * (1 - test_size)) + x_train = x_array[:split_idx] + y_train = y_array[:split_idx] + x_valid = x_array[split_idx:] + y_valid = y_array[split_idx:] return (x_train, y_train), (x_valid, y_valid) return x_array, y_array @@ -154,31 +304,31 @@ def get_air_passengers(train_sequence_length: int = 24, predict_sequence_length: Tuple of training and validation data, each containing inputs and outputs. """ - df = pd.read_csv(TS_DATASETS_URL["air_passengers"]["url"], parse_dates=None, date_parser=None, nrows=144) - v = df.iloc[:, 1:2].values - v = (v - np.max(v)) / (np.max(v) - np.min(v)) # MinMaxScaler - - x: List[np.ndarray] = [] - y: List[np.ndarray] = [] - for seq in range(1, train_sequence_length + 1): - x_roll = np.roll(v, seq, axis=0) - x.append(x_roll) - x_array = np.stack(x, axis=1) - x_array = x_array[train_sequence_length:-predict_sequence_length, ::-1, :] - - for seq in range(predict_sequence_length): - y_roll = np.roll(v, -seq) - y.append(y_roll) - y_array = np.stack(y, axis=1) - y_array = y_array[train_sequence_length:-predict_sequence_length] - logging.info("Load air passenger data", x_array.shape, y_array.shape) + if train_sequence_length < 1 or predict_sequence_length < 1: + raise ValueError("train_sequence_length and predict_sequence_length must be positive") + if train_sequence_length + predict_sequence_length > len(AIR_PASSENGERS_VALUES): + raise ValueError("Requested sequence lengths exceed the AirPassengers dataset length") + + # Keep this canonical small dataset in-package so examples and tests work offline. + v = AIR_PASSENGERS_VALUES.reshape(-1, 1).copy() + v = (v - np.min(v)) / (np.max(v) - np.min(v)) # MinMaxScaler + + window_count = len(v) - train_sequence_length - predict_sequence_length + 1 + x_array = np.stack([v[i : i + train_sequence_length] for i in range(window_count)]) + y_array = np.stack( + [ + v[i + train_sequence_length : i + train_sequence_length + predict_sequence_length] + for i in range(window_count) + ] + ) + logging.info(f"Load air passenger data {x_array.shape} {y_array.shape}") if test_size > 0: - slice = int(len(x_array) * (1 - test_size)) - x_train = x_array[:slice] - y_train = y_array[:slice] - x_valid = x_array[slice:] - y_valid = y_array[slice:] + split_idx = int(len(x_array) * (1 - test_size)) + x_train = x_array[:split_idx] + y_train = y_array[:split_idx] + x_valid = x_array[split_idx:] + y_valid = y_array[split_idx:] return (x_train, y_train), (x_valid, y_valid) return x_array, y_array diff --git a/tfts/data/processor.py b/tfts/data/processor.py new file mode 100644 index 00000000..67ccbd3d --- /dev/null +++ b/tfts/data/processor.py @@ -0,0 +1,294 @@ +"""DataProcessor — the "Tokenizer" for time series. + +Provides a clean, high-level interface for preparing time series data for +training, validation, and prediction. +""" + +import logging +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import pandas as pd +import tensorflow as tf + +from .timeseries import TimeSeriesSequence + +logger = logging.getLogger(__name__) + + +class DataProcessor: + """Unified data preprocessor for time series tasks. + + Wraps TimeSeriesSequence and provides a transformers-like interface. + Handles sliding windows, normalization, train/valid/test splits + automatically. + + Args: + lookback: Number of past time steps used as input. + horizon: Number of future time steps to predict. + batch_size: Batch size for tf.data datasets. + stride: Step size for sliding window (1 = every step, >1 = downsampling). + normalize: Normalization method — ``'minmax'``, ``'standard'``, or ``None``. + group_col: Column(s) to group multiple time series. + feature_cols: Additional feature columns to include. + fill_missing_dates: Whether to fill gaps in the time index. + freq: Frequency string (e.g. ``'D'``, ``'H'``) when filling dates. + validation_split: Fraction of training data to use for validation. + shuffle: Whether to shuffle the training dataset. + seed: Random seed for reproducibility. + + Examples: + >>> df = pd.DataFrame({ + ... 'date': pd.date_range('2023-01-01', periods=365), + ... 'sales': np.random.randn(365).cumsum(), + ... }) + >>> processor = DataProcessor(lookback=30, horizon=7) + >>> train_ds, valid_ds = processor.prepare(df, target_col='sales') + >>> for x, y in train_ds.take(1): + ... print(x.shape, y.shape) + """ + + def __init__( + self, + lookback: int = 96, + horizon: int = 24, + batch_size: int = 32, + stride: int = 1, + normalize: Optional[str] = "minmax", + group_col: Optional[Union[str, List[str]]] = None, + feature_cols: Optional[List[str]] = None, + fill_missing_dates: bool = False, + freq: Optional[str] = None, + validation_split: float = 0.2, + shuffle: bool = True, + seed: int = 42, + ): + if lookback < 1: + raise ValueError(f"lookback must be >= 1, got {lookback}") + if horizon < 1: + raise ValueError(f"horizon must be >= 1, got {horizon}") + if batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {batch_size}") + if stride < 1: + raise ValueError(f"stride must be >= 1, got {stride}") + if not 0 <= validation_split < 1: + raise ValueError(f"validation_split must be in [0, 1), got {validation_split}") + if normalize not in (None, "minmax", "standard"): + raise ValueError(f"normalize must be 'minmax', 'standard', or None, got {normalize}") + + self.lookback = lookback + self.horizon = horizon + self.batch_size = batch_size + self.stride = stride + self.normalize = normalize + self.group_col = group_col + self.feature_cols = feature_cols or [] + self.fill_missing_dates = fill_missing_dates + self.freq = freq + self.validation_split = validation_split + self.shuffle = shuffle + self.seed = seed + + # Set during prepare() + self._scaler_params: Optional[Dict] = None + self._feature_names: List[str] = [] + self._target_col: Optional[str] = None + self._time_col: Optional[str] = None + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def prepare( + self, + df: pd.DataFrame, + target_col: Optional[str] = None, + time_col: Optional[str] = None, + ) -> Union[ + Tuple[tf.data.Dataset, tf.data.Dataset], + Tuple[tf.data.Dataset, tf.data.Dataset, tf.data.Dataset], + tf.data.Dataset, + ]: + """Prepare data and return tf.data.Dataset(s). + + Args: + df: Input DataFrame. + target_col: Name of the column to forecast. Auto-detected if None. + time_col: Name of the time column. Auto-detected if None. + + Returns: + If ``validation_split > 0``: + ``(train_ds, valid_ds)`` + If ``validation_split == 0``: + ``train_ds`` + """ + df = df.copy() + + self._target_col = target_col or self._infer_target(df) + self._time_col = time_col or self._infer_time(df) + + # Normalize + if self.normalize is not None: + fit_df = self._normalization_fit_frame(df) + self._apply_normalization(fit_df, fit=True) + df = self._apply_normalization(df, fit=False) + + # Build sequence + seq = self._build_sequence(df) + ds = seq.get_tf_dataset() + + if self.validation_split > 0: + return self._split_dataset(ds) + return ds + + def prepare_for_inference( + self, + df: pd.DataFrame, + target_col: Optional[str] = None, + time_col: Optional[str] = None, + ) -> tf.data.Dataset: + """Prepare data for inference (no shuffle, batch_size=1 by default).""" + df = df.copy() + target = target_col or self._target_col or self._infer_target(df) + time = time_col or self._time_col or self._infer_time(df) + if self.normalize is not None: + if self._scaler_params is None: + raise RuntimeError("DataProcessor must be fitted with prepare() before normalized inference") + df = self._apply_normalization(df, fit=False) + return self._build_sequence(df, target_col=target, time_col=time, mode="inference").get_tf_dataset() + + def inverse_transform(self, values: Union[np.ndarray, tf.Tensor]) -> Union[np.ndarray, tf.Tensor]: + """Reverse the normalization applied during prepare(). + + Args: + values: Normalized predictions or targets. + + Returns: + Values in the original scale. + """ + if self._scaler_params is None: + return values + if self.normalize == "standard": + return values * self._scaler_params["std"] + self._scaler_params["mean"] + min_val = self._scaler_params["min"] + max_val = self._scaler_params["max"] + return values * (max_val - min_val) + min_val + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _infer_target(df: pd.DataFrame) -> str: + """Auto-detect target column (last numeric column).""" + numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist() + if not numeric_cols: + raise ValueError("No numeric columns found in DataFrame. Please specify target_col.") + # Exclude obvious time/index columns + candidates = [c for c in numeric_cols if not _looks_like_time(df[c])] + if candidates: + return candidates[-1] # usually the value column is last + return numeric_cols[-1] + + @staticmethod + def _infer_time(df: pd.DataFrame) -> str: + """Auto-detect time column.""" + # Check index first + if isinstance(df.index, pd.DatetimeIndex): + col_name = df.index.name or "time_idx" + df[col_name] = df.index + return col_name + # Look for datetime columns + for col in df.columns: + if pd.api.types.is_datetime64_any_dtype(df[col]): + return col + # Look for columns named like time + for name in ("date", "time", "datetime", "timestamp", "ds"): + if name in df.columns: + return name + # Fallback to first column + return df.columns[0] + + def _apply_normalization(self, df: pd.DataFrame, fit: bool) -> pd.DataFrame: + """Apply min-max or standard normalization to target column.""" + target = self._target_col + if self.normalize == "minmax": + if fit: + self._scaler_params = {"min": df[target].min(), "max": df[target].max()} + min_val = self._scaler_params["min"] + max_val = self._scaler_params["max"] + df[target] = (df[target] - min_val) / (max_val - min_val + 1e-8) + elif self.normalize == "standard": + if fit: + self._scaler_params = {"mean": df[target].mean(), "std": df[target].std()} + mean = self._scaler_params["mean"] + std = self._scaler_params["std"] + df[target] = (df[target] - mean) / (std + 1e-8) + return df + + def _normalization_fit_frame(self, df: pd.DataFrame) -> pd.DataFrame: + """Return only the chronological training portion used to fit scaling.""" + if self.validation_split <= 0: + return df.copy() + if self.group_col: + group_cols = [self.group_col] if isinstance(self.group_col, str) else self.group_col + parts = [] + for _, group in df.groupby(group_cols, observed=True, sort=False): + split_idx = max(1, int(len(group) * (1 - self.validation_split))) + parts.append(group.iloc[:split_idx]) + return pd.concat(parts, axis=0).copy() + split_idx = max(1, int(len(df) * (1 - self.validation_split))) + return df.iloc[:split_idx].copy() + + def _build_sequence( + self, + df: pd.DataFrame, + target_col: Optional[str] = None, + time_col: Optional[str] = None, + mode: str = "train", + ) -> TimeSeriesSequence: + """Build a TimeSeriesSequence from the DataFrame.""" + return TimeSeriesSequence( + data=df, + time_idx=time_col or self._time_col, + target_column=target_col or self._target_col, + train_sequence_length=self.lookback, + predict_sequence_length=self.horizon, + batch_size=self.batch_size, + group_column=self.group_col, + feature_columns=self.feature_cols if self.feature_cols else None, + stride=self.stride, + mode=mode, + ) + + def _split_dataset(self, ds: tf.data.Dataset) -> Tuple[tf.data.Dataset, tf.data.Dataset]: + """Split a tf.data.Dataset into train / validation.""" + samples = list(ds.unbatch().as_numpy_iterator()) + if len(samples) < 2: + raise ValueError("At least two windows are required when validation_split is greater than zero") + + split_idx = min(len(samples) - 1, max(1, int(len(samples) * (1 - self.validation_split)))) + train_x = np.stack([sample[0] for sample in samples[:split_idx]]) + train_y = np.stack([sample[1] for sample in samples[:split_idx]]) + valid_x = np.stack([sample[0] for sample in samples[split_idx:]]) + valid_y = np.stack([sample[1] for sample in samples[split_idx:]]) + + train_ds = tf.data.Dataset.from_tensor_slices((train_x, train_y)) + if self.shuffle: + train_ds = train_ds.shuffle(buffer_size=len(train_x), seed=self.seed) + train_ds = train_ds.batch(self.batch_size).prefetch(tf.data.AUTOTUNE) + valid_ds = tf.data.Dataset.from_tensor_slices((valid_x, valid_y)).batch(self.batch_size) + return train_ds, valid_ds + + +def _looks_like_time(series: pd.Series) -> bool: + """Heuristic to detect time-like columns.""" + if pd.api.types.is_datetime64_any_dtype(series): + return True + # Check if values look like timestamps or sequential integers + if series.dtype.kind in "iu": # integer + if len(series) > 1 and series.is_monotonic_increasing: + diffs = series.diff().dropna() + if diffs.nunique() <= 3: # regular step + return True + return False diff --git a/tfts/data/timeseries.py b/tfts/data/timeseries.py index c2557f10..5c5f4c8a 100644 --- a/tfts/data/timeseries.py +++ b/tfts/data/timeseries.py @@ -91,11 +91,12 @@ def __init__( self.predict_sequence_length = predict_sequence_length self.stride = stride self.batch_size = batch_size - self.group_column = group_column + self.group_column = [group_column] if isinstance(group_column, str) else group_column + self.feature_columns = [feature_columns] if isinstance(feature_columns, str) else list(feature_columns or []) self.drop_last = drop_last self.feature_config = feature_config or {} self.mode = mode - self.group_ids = group_column or [] + self.group_ids = self.group_column or [] # Initialize feature registry self.feature_registry = FeatureRegistry() @@ -103,53 +104,21 @@ def __init__( # Validate inputs and apply feature transformations self._validate_inputs() self._apply_feature_transforms() + self._validate_feature_columns() # Generate sequences self.sequences = [] - if group_column is not None: - for _, group in data.groupby(group_column, observed=True): + if self.group_column: + for _, group in self.data.groupby(self.group_column, observed=True): self.sequences.extend(self._generate_sequences(group, time_idx=time_idx, target_column=target_column)) else: - self.sequences.extend(self._generate_sequences(data, time_idx=time_idx, target_column=target_column)) + self.sequences.extend(self._generate_sequences(self.data, time_idx=time_idx, target_column=target_column)) logger.info( f"Initialized TimeSeriesSequence with {len(self.sequences)} sequences, " f"batch_size={batch_size}, mode={mode}" ) - def _build_sequences(self): - """Builds a lookup table for sequences to avoid heavy DataFrame slicing during training.""" - sequence_indices = [] - - if self.group_column: - grouped = self.data.groupby(self.group_column) - else: - grouped = [("all", self.data)] - - for _, group in grouped: - group = group.sort_values(self.time_idx) - n_rows = len(group) - max_idx = n_rows - self.train_sequence_length - self.predict_sequence_length + 1 - - # Pre-extract numpy arrays for speed - feature_data = group[self.features].values.astype(np.float32) - target_data = group[self.target].values.astype(np.float32) - - for i in range(0, max_idx, self.stride): - sequence_indices.append( - { - "x": feature_data[i : i + self.train_sequence_length], - "y": target_data[ - i - + self.train_sequence_length : i - + self.train_sequence_length - + self.predict_sequence_length - ], - } - ) - - return sequence_indices - def __len__(self) -> int: """Get the number of batches in the sequence. @@ -175,10 +144,12 @@ def __getitem__(self, idx: int) -> Tuple[np.ndarray, np.ndarray]: end_idx = min(start_idx + self.batch_size, len(self.sequences)) batch_sequences = self.sequences[start_idx:end_idx] + if not batch_sequences: + raise IndexError(f"Batch index {idx} is out of range") # Stack encoder inputs and decoder targets using np.stack - encoder_inputs = np.stack([seq[0] for seq in batch_sequences]) - decoder_targets = np.stack([seq[1] for seq in batch_sequences]) + encoder_inputs = np.stack([seq[0] for seq in batch_sequences]).astype(np.float32) + decoder_targets = np.stack([seq[1] for seq in batch_sequences]).astype(np.float32) return encoder_inputs, decoder_targets @@ -200,7 +171,7 @@ def get_tf_dataset(self) -> tf.data.Dataset: def _generate_sequences( self, group: pd.DataFrame, time_idx: str, target_column: str - ) -> List[Tuple[np.ndarray, np.ndarray, int]]: + ) -> List[Tuple[np.ndarray, np.ndarray]]: """Generate sequences from a group of data. Args: @@ -213,7 +184,9 @@ def _generate_sequences( Each sequence is a 2D array with shape (length, num_features) """ group = group.sort_values(by=time_idx) - target_values = group[target_column].values + input_columns = list(dict.fromkeys(self.target + self.feature_columns)) + input_values = group[input_columns].to_numpy() + target_values = group[self.target].to_numpy() time_values = group[time_idx].values # Convert time values to numeric if they are datetime @@ -221,42 +194,41 @@ def _generate_sequences( time_values = time_values.astype(np.int64) // 10**9 # Convert to seconds sequences = [] - max_start_idx = len(group) - self.train_sequence_length - self.predict_sequence_length + 1 + if self.mode == "inference": + max_start_idx = len(group) - self.train_sequence_length + 1 + else: + max_start_idx = len(group) - self.train_sequence_length - self.predict_sequence_length + 1 - for i in range(0, max_start_idx, self.stride): + for i in range(0, max(0, max_start_idx), self.stride): # Get indices for encoder sequence encoder_start = i encoder_end = i + self.train_sequence_length encoder_indices = np.arange(encoder_start, encoder_end) # Get indices for decoder sequence - decoder_start = encoder_end - decoder_end = decoder_start + self.predict_sequence_length - decoder_indices = np.arange(decoder_start, decoder_end) - - # Check if sequences are continuous - encoder_time_diffs = np.diff(time_values[encoder_indices]) - decoder_time_diffs = np.diff(time_values[decoder_indices]) - - # For datetime values, check if differences are consistent - if pd.api.types.is_datetime64_any_dtype(group[time_idx]): - expected_diff = (group[time_idx].iloc[1] - group[time_idx].iloc[0]).total_seconds() - is_continuous = np.all(np.abs(encoder_time_diffs - expected_diff) < 1e-6) and np.all( - np.abs(decoder_time_diffs - expected_diff) < 1e-6 - ) + if self.mode == "inference": + decoder_indices = np.array([], dtype=int) + window_indices = encoder_indices else: - is_continuous = np.all(encoder_time_diffs == encoder_time_diffs[0]) and np.all( - decoder_time_diffs == decoder_time_diffs[0] - ) + decoder_start = encoder_end + decoder_end = decoder_start + self.predict_sequence_length + decoder_indices = np.arange(decoder_start, decoder_end) + window_indices = np.concatenate([encoder_indices, decoder_indices]) + + # Check the entire encoder/decoder window, including their boundary. + time_diffs = np.diff(time_values[window_indices]) + is_continuous = len(time_diffs) == 0 or np.all(time_diffs == time_diffs[0]) if ( len(encoder_indices) == self.train_sequence_length - and len(decoder_indices) == self.predict_sequence_length + and (self.mode == "inference" or len(decoder_indices) == self.predict_sequence_length) and is_continuous ): - # Ensure 2D arrays with shape (length, 1) for single feature - encoder_sequence = target_values[encoder_indices].reshape(-1, 1) - decoder_sequence = target_values[decoder_indices].reshape(-1, 1) + encoder_sequence = input_values[encoder_indices] + if self.mode == "inference": + decoder_sequence = np.zeros((self.predict_sequence_length, len(self.target)), dtype=np.float32) + else: + decoder_sequence = target_values[decoder_indices] sequences.append((encoder_sequence, decoder_sequence)) return sequences @@ -276,6 +248,8 @@ def _validate_inputs(self) -> None: raise ValueError("predict_sequence_length must be at least 1") if self.stride < 1: raise ValueError("stride must be at least 1") + if self.batch_size < 1: + raise ValueError("batch_size must be at least 1") # Validate mode valid_modes = ["train", "validation", "test", "inference"] @@ -354,6 +328,12 @@ def _apply_feature_transforms(self) -> None: logger.error(f"Error applying feature transform {transform_type} for {feature_name}: {str(e)}") raise + def _validate_feature_columns(self) -> None: + """Validate encoder feature columns after configured transforms run.""" + missing_features = [col for col in self.feature_columns if col not in self.data.columns] + if missing_features: + raise ValueError(f"Data is missing feature columns: {missing_features}") + @classmethod def from_df( cls, diff --git a/tfts/features/__init__.py b/tfts/features/__init__.py index cb711096..f606dafc 100644 --- a/tfts/features/__init__.py +++ b/tfts/features/__init__.py @@ -1,6 +1,18 @@ """tfts features""" +from .auto_feature import AutoFeatureEngineer from .datetime_feature import add_datetime_feature from .one_order_feature import add_lag_feature, add_moving_average_feature, add_roll_feature, add_transform_feature from .registry import FeatureRegistry from .two_order_feature import add_2order_feature + +__all__ = [ + "AutoFeatureEngineer", + "FeatureRegistry", + "add_2order_feature", + "add_datetime_feature", + "add_lag_feature", + "add_moving_average_feature", + "add_roll_feature", + "add_transform_feature", +] diff --git a/tfts/features/auto_feature.py b/tfts/features/auto_feature.py new file mode 100644 index 00000000..d98743d4 --- /dev/null +++ b/tfts/features/auto_feature.py @@ -0,0 +1,208 @@ +"""AutoFeatureEngineer — sklearn-style automatic feature engineering for time series. + +Wraps the existing feature functions (lag, rolling, datetime, …) behind a +single class with fit / transform semantics so users don't need to call +individual functions manually. +""" + +import logging +from typing import List, Optional, Union + +import numpy as np +import pandas as pd + +from .datetime_feature import add_datetime_feature +from .one_order_feature import add_lag_feature, add_moving_average_feature, add_roll_feature, add_transform_feature +from .registry import FeatureRegistry + +logger = logging.getLogger(__name__) + + +class AutoFeatureEngineer: + """Automatic feature engineering for time series DataFrames. + + Generates lag features, rolling-window features, datetime features, + and optional Fourier terms behind a simple ``fit_transform`` API. + Tracks generated feature names via an internal + :class:`~tfts.features.registry.FeatureRegistry`. + + Args: + lags: Lag offsets to generate (e.g. ``[1, 7, 14]``). + windows: Rolling window sizes (e.g. ``[7, 30]``). + rolling_functions: Aggregations for rolling windows. + Defaults to ``['mean', 'std']`` for speed. Use ``'all'`` + to get ``['mean', 'std', 'min', 'max']``. + add_datetime: If ``True``, add calendar features (month, day-of- + week, …). + datetime_features: Specific datetime features to generate. + ``None`` means ``['month', 'dayofweek', 'hour']`` (where + applicable). + add_fourier: If ``True``, add cyclical (sin/cos) features for + month and day-of-week. + group_cols: Columns to group by (for multi-series data). + + Examples: + >>> from tfts.features import AutoFeatureEngineer + >>> eng = AutoFeatureEngineer(lags=[1, 7], windows=[7, 30], + ... add_datetime=True) + >>> df_feat = eng.fit_transform(df, time_col="date", target_col="value") + >>> eng.get_feature_names()[:3] + ['value_lag_1', 'value_lag_7', 'value_roll_7_mean'] + """ + + def __init__( + self, + lags: Optional[List[int]] = None, + windows: Optional[List[int]] = None, + rolling_functions: Union[str, List[str]] = "default", + add_datetime: bool = False, + datetime_features: Optional[List[str]] = None, + add_fourier: bool = False, + group_cols: Optional[List[str]] = None, + ) -> None: + self.lags = lags or [1, 7] + self.windows = windows or [7] + self.rolling_functions = rolling_functions + self.add_datetime = add_datetime + self.datetime_features = datetime_features + self.add_fourier = add_fourier + self.group_cols = group_cols + + self._registry = FeatureRegistry() + self._original_columns: List[str] = [] + self._fitted: bool = False + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def fit(self, df: pd.DataFrame, time_col: str, target_col: str) -> "AutoFeatureEngineer": + """Fit the engineer — records original columns so ``transform`` + can later identify which columns are new features. + + Args: + df: Training DataFrame. + time_col: Name of the time column. + target_col: Name of the target column. + + Returns: + self (for chaining). + """ + self._time_col = time_col + self._target_col = target_col + self._original_columns = df.columns.tolist() + self._fitted = True + return self + + def transform(self, df: pd.DataFrame) -> pd.DataFrame: + """Apply feature engineering to *df*. + + Must be called after :meth:`fit`. + + Args: + df: DataFrame with the same time/target columns used in + :meth:`fit`. + + Returns: + DataFrame with new feature columns appended. + """ + self._check_fitted() + result = df.copy() + + # 1. Lag features + result = add_lag_feature( + data=result, + columns=self._target_col, + lags=self.lags, + time_col=self._time_col, + group_cols=self.group_cols, + ) + + # 2. Rolling features + roll_funcs = self._resolve_rolling_functions() + result = add_roll_feature( + data=result, + columns=self._target_col, + windows=self.windows, + functions=roll_funcs, + time_col=self._time_col, + group_cols=self.group_cols, + ) + + # 3. Datetime features + if self.add_datetime: + features = self.datetime_features + if features is None: + # Pick sensible defaults depending on the data's resolution + features = _default_datetime_features(result[self._time_col]) + result = add_datetime_feature(data=result, time_col=self._time_col, features=features) + + # 4. Fourier (cyclical) features + if self.add_fourier: + fourier_features = _default_fourier_features() + result = add_datetime_feature(data=result, time_col=self._time_col, features=fourier_features) + + # Register the new feature columns + new_cols = [c for c in result.columns if c not in self._original_columns] + self._registry.register(new_cols) + + # Drop rows with NaN introduced by lag/rolling + result = result.dropna().reset_index(drop=True) + + logger.info(f"AutoFeatureEngineer added {len(new_cols)} features") + return result + + def fit_transform(self, df: pd.DataFrame, time_col: str, target_col: str) -> pd.DataFrame: + """Convenience: fit then transform in one call.""" + return self.fit(df, time_col, target_col).transform(df) + + def get_feature_names(self) -> List[str]: + """Return names of the generated features (excluding originals).""" + return self._registry.get_features() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _resolve_rolling_functions(self) -> List[str]: + if self.rolling_functions == "all": + return ["mean", "std", "min", "max"] + if self.rolling_functions == "default": + return ["mean", "std"] + if isinstance(self.rolling_functions, str): + return [self.rolling_functions] + return list(self.rolling_functions) + + def _check_fitted(self) -> None: + if not self._fitted: + raise RuntimeError("AutoFeatureEngineer is not fitted yet. Call .fit() or .fit_transform() first.") + + def __repr__(self) -> str: + status = "fitted" if self._fitted else "not fitted" + return ( + f"AutoFeatureEngineer(lags={self.lags}, windows={self.windows}, " + f"add_datetime={self.add_datetime}, add_fourier={self.add_fourier}, {status})" + ) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _default_datetime_features(time_series: pd.Series) -> List[str]: + """Pick sensible datetime features based on the time resolution.""" + if not pd.api.types.is_datetime64_any_dtype(time_series): + return ["month", "dayofweek"] + + sample = time_series.dropna().iloc[:3] + has_sub_daily = any(ts.hour != 0 or ts.minute != 0 for ts in sample if pd.notna(ts)) + features = ["month", "dayofweek", "day"] + if has_sub_daily: + features.append("hour") + return features + + +def _default_fourier_features() -> List[str]: + """Default cyclical (sin/cos) features.""" + return ["month_sin", "month_cos", "dayofweek_sin", "dayofweek_cos"] diff --git a/tfts/features/registry.py b/tfts/features/registry.py index 315a636f..fb1292b6 100644 --- a/tfts/features/registry.py +++ b/tfts/features/registry.py @@ -57,6 +57,9 @@ def register(self, cols: Union[str, List[str]]) -> None: raise ValueError(f"Column name contains invalid characters: {col}") self.columns.extend(cols) + # Deduplicate while preserving order + seen = set() + self.columns = [c for c in self.columns if not (c in seen or seen.add(c))] logger.debug(f"Registered {len(cols)} features: {cols}") def get_features(self) -> List[str]: @@ -147,7 +150,7 @@ def __repr__(self) -> str: feature_registry = FeatureRegistry() -def registry(func: Callable) -> Callable: +def registry(func: Callable[..., Any]) -> Callable[..., Any]: """Decorator to register features returned by a function. This decorator automatically registers any features returned by the decorated diff --git a/tfts/generator.py b/tfts/generator.py index f7bc5b61..79f5ccf7 100644 --- a/tfts/generator.py +++ b/tfts/generator.py @@ -1,6 +1,16 @@ -"""tfts Generator""" +"""tfts Generator -from typing import Any, Dict, Optional, Union +This module provides auto-regressive generation utilities for time series models. + +When using GenerationMixin as a mixin, the host class should provide these attributes: +- time_idx: Name of the time index column (str or None, defaults to "time_idx") +- group_column: List of group identifier column names (list or None) +- target: List of target column names +- train_sequence_length: Length of input sequences +- get_feature_names(): Method returning list of all feature column names +""" + +from typing import Any, Dict, List, Optional, Union import numpy as np import pandas as pd @@ -15,8 +25,38 @@ def __init__(self, **kwargs) -> None: class GenerationMixin: """ A class containing auto-regressive generation, to be used as a mixin. + + Required host class attributes: + - time_idx (str): Name of the time index column + - group_column (list[str]): List of group identifier columns + - target (list[str]): List of target column names + - train_sequence_length (int): Length of input sequences for the model + - get_feature_names() -> list[str]: Returns list of all feature column names """ + _REQUIRED_ATTRS = ["time_idx", "group_column", "target", "train_sequence_length"] + _REQUIRED_METHODS = ["get_feature_names"] + + def _validate_generation_attrs(self) -> None: + """Validate that required attributes exist on the host class.""" + missing = [] + for attr in self._REQUIRED_ATTRS: + if not hasattr(self, attr): + missing.append(attr) + # Allow time_idx and group_column to be None, but enforce others + elif getattr(self, attr) is None and attr not in ["time_idx", "group_column"]: + missing.append(attr) + + for method in self._REQUIRED_METHODS: + if not hasattr(self, method): + missing.append(method + "()") + + if missing: + raise AttributeError( + f"GenerationMixin requires the following attributes/methods on the host class: " + f"{', '.join(missing)}. See class docstring for details." + ) + def prepare_inputs_for_generation(self, *args, **kwargs): return @@ -44,11 +84,13 @@ def generate( Returns: DataFrame with original inputs and generated predictions """ + self._validate_generation_attrs() + generation_config = generation_config or {} steps = generation_config.get("steps", 1) - time_idx = generation_config.get("time_idx", self.time_idx) + time_idx = generation_config.get("time_idx", getattr(self, "time_idx", "time_idx")) time_step = generation_config.get("time_step", 1) - group_columns = generation_config.get("group_columns", self.group_column) + group_columns = generation_config.get("group_columns", getattr(self, "group_column", None)) add_features_func = generation_config.get("add_features_func", None) # Convert inputs to DataFrame if needed @@ -57,24 +99,19 @@ def generate( if len(features) != inputs.shape[1]: raise ValueError(f"Input array shape {inputs.shape} doesn't match feature count {len(features)}") inputs_df = pd.DataFrame(inputs, columns=features) - - # Add time index if not present - if time_idx not in inputs_df.columns: - # Create a time index - max_time = inputs_df[time_idx].max() if time_idx in inputs_df.columns else 0 - inputs_df[time_idx] = np.arange(max_time + 1, max_time + len(inputs_df) + 1) else: inputs_df = inputs.copy() # Ensure inputs are sorted by time index - inputs_df = inputs_df.sort_values(by=time_idx) + if time_idx in inputs_df.columns: + inputs_df = inputs_df.sort_values(by=time_idx) # Create a copy to store results results_df = inputs_df.copy() - last_time_idx = results_df[time_idx].max() + last_time_idx = results_df[time_idx].max() if time_idx in results_df.columns else 0 # Get the sequence length for the model - seq_length = self.train_sequence_length + seq_length = getattr(self, "train_sequence_length", max_steps) # Predict one step at a time and add to results for step in range(steps): @@ -108,8 +145,9 @@ def generate( new_row[col] = results_df[col].iloc[-1] # Add prediction for target columns - for i, target in enumerate(self.target): - new_row[target] = prediction[0, 0, i] # Assuming [batch, time, feature] format + target = getattr(self, "target", []) + for i, tgt in enumerate(target): + new_row[tgt] = prediction[0, 0, i] # Assuming [batch, time, feature] format # Add new row to results new_df = pd.DataFrame([new_row]) diff --git a/tfts/layers/__init__.py b/tfts/layers/__init__.py index e886dc42..04719187 100644 --- a/tfts/layers/__init__.py +++ b/tfts/layers/__init__.py @@ -1 +1,84 @@ -"""tfts layers""" +"""Neural network layers for time series prediction models.""" + +# Attention layers +from .attention_layer import Attention, ProbAttention, SelfAttention, SparseAttention + +# Decomposition layers +from .autoformer_layer import MovingAvg, SeriesDecomp + +# Convolution layers +from .cnn_layer import ConvTemp + +# Dense/Feedforward layers +from .dense_layer import DenseTemp, FeedForwardNetwork, MoeMLP + +# Embedding layers +from .embed_layer import DataEmbedding, TokenEmbedding + +# Graph layers +from .graph_layer import GraphAttention, GraphConv +from .mask_layer import CausalMask, ProbMask + +# MoE layers +from .moe_layer import SparseMoe + +# NBeats layers +from .nbeats_layer import GenericBlock, SeasonalityBlock, TrendBlock + +# Position encoding +from .position_layer import PositionalEmbedding, PositionalEncoding + +# RWKV layers +from .rwkv_layer import ChannelMixing, TimeMixing + +# UNet layers +from .unet_layer import ConvbrLayer, ReBlock, SeBlock + +# Utility layers +from .util_layer import CreateDecoderFeature, ShapeLayer, ZerosLayer + +__all__ = [ + # Attention + "Attention", + "ProbAttention", + "SelfAttention", + "SparseAttention", + # Masks + "CausalMask", + "ProbMask", + # Dense/Feedforward + "DenseTemp", + "FeedForwardNetwork", + "MoeMLP", + # Embedding + "DataEmbedding", + "TokenEmbedding", + # Convolution + "ConvTemp", + # Decomposition + "MovingAvg", + "SeriesDecomp", + # Position encoding + "PositionalEmbedding", + "PositionalEncoding", + # RWKV + "ChannelMixing", + "TimeMixing", + # Utility + "CreateDecoderFeature", + "ShapeLayer", + "ZerosLayer", + # Graph + "GraphAttention", + "GraphConv", + # NBeats + "GenericBlock", + "SeasonalityBlock", + "TrendBlock", + # UNet + "ConvbrLayer", + "ReBlock", + "SeBlock", + # MoE + "SparseMoe", +] diff --git a/tfts/layers/attention_layer.py b/tfts/layers/attention_layer.py index b087a3ef..cf2467f1 100644 --- a/tfts/layers/attention_layer.py +++ b/tfts/layers/attention_layer.py @@ -234,7 +234,7 @@ def _get_initial_context(self, v, L_Q): B = tf.shape(v)[0] if not self.mask_flag: v_sum = tf.math.reduce_sum(v, axis=-2) - context = tf.identity(tf.boradcast_to(tf.expand_dims(v_sum, -2), [B, H, L_Q, v_sum.shape[-1]])) + context = tf.identity(tf.broadcast_to(tf.expand_dims(v_sum, -2), [B, H, L_Q, v_sum.shape[-1]])) else: assert L_Q == L_V context = tf.math.cumsum(v, axis=-2) diff --git a/tfts/losses/__init__.py b/tfts/losses/__init__.py index e69de29b..186f17e9 100644 --- a/tfts/losses/__init__.py +++ b/tfts/losses/__init__.py @@ -0,0 +1,5 @@ +"""Loss functions for time series prediction.""" + +from .loss import MultiQuantileLoss + +__all__ = ["MultiQuantileLoss"] diff --git a/tfts/losses/loss.py b/tfts/losses/loss.py index 33e28841..532fc249 100644 --- a/tfts/losses/loss.py +++ b/tfts/losses/loss.py @@ -1,26 +1,31 @@ -from typing import List +from typing import List, Union import tensorflow as tf class MultiQuantileLoss(tf.keras.losses.Loss): - def __init__(self, quantiles: List[float], name="multi_quantile_loss"): + """Multi-quantile loss using the pinball loss function. + + Computes the pinball (quantile) loss for each specified quantile and + sums across quantiles. Handles multi-horizon, multi-target outputs. + + Args: + quantiles: List of quantile fractions, e.g. ``[0.1, 0.5, 0.9]``. + name: Loss name. + + Shape: + - ``y_true``: ``(batch, pred_len, num_labels)`` + - ``y_pred``: ``(batch, pred_len, num_labels * len(quantiles))`` + """ + + def __init__(self, quantiles: List[float], name: str = "multi_quantile_loss"): super().__init__(name=name) self.quantiles = quantiles - def call(self, y_true, y_pred): - """ - y_true: [batch, pred_len, num_labels] - y_pred: [batch, pred_len, num_labels * num_quantiles] - """ - # Reshape y_pred to [batch, pred_len, num_labels, num_quantiles] - # and y_true to [batch, pred_len, num_labels, 1] + def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: y_true = tf.expand_dims(y_true, axis=-1) - - # Split y_pred into the different quantiles - # Assuming the head outputs quantiles stacked in the last dimension num_labels = y_true.shape[-2] - y_pred = tf.reshape(y_pred, [-1, y_pred.shape[1], num_labels, len(self.quantiles)]) + y_pred = tf.reshape(y_pred, [-1, tf.shape(y_pred)[1], num_labels, len(self.quantiles)]) losses = [] for i, q in enumerate(self.quantiles): diff --git a/tfts/metrics.py b/tfts/metrics.py new file mode 100644 index 00000000..b39146ef --- /dev/null +++ b/tfts/metrics.py @@ -0,0 +1,166 @@ +"""Time series evaluation metrics. + +Standard point-forecast metrics usable with numpy arrays or TensorFlow tensors. +""" + +from typing import Union + +import numpy as np +import tensorflow as tf + + +def mse(y_true: Union[np.ndarray, tf.Tensor], y_pred: Union[np.ndarray, tf.Tensor]) -> Union[np.ndarray, tf.Tensor]: + """Mean Squared Error. + + Args: + y_true: Ground truth values. + y_pred: Predicted values. + + Returns: + Scalar or array of MSE values. + """ + return _reduce(np.square(_sub(y_true, y_pred))) + + +def mae(y_true: Union[np.ndarray, tf.Tensor], y_pred: Union[np.ndarray, tf.Tensor]) -> Union[np.ndarray, tf.Tensor]: + """Mean Absolute Error.""" + return _reduce(np.abs(_sub(y_true, y_pred))) + + +def rmse(y_true: Union[np.ndarray, tf.Tensor], y_pred: Union[np.ndarray, tf.Tensor]) -> Union[np.ndarray, tf.Tensor]: + """Root Mean Squared Error.""" + return _sqrt(mse(y_true, y_pred)) + + +def mape( + y_true: Union[np.ndarray, tf.Tensor], y_pred: Union[np.ndarray, tf.Tensor], eps: float = 1e-8 +) -> Union[np.ndarray, tf.Tensor]: + """Mean Absolute Percentage Error. + + Args: + y_true: Ground truth values. + y_pred: Predicted values. + eps: Small constant to avoid division by zero. + + Returns: + MAPE as a percentage (0-100 scale). + """ + backend = _backend(y_true) + denominator = backend.maximum(backend.abs(y_true), backend.array(eps, dtype=y_true.dtype)) + return 100.0 * _reduce(backend.abs(_sub(y_true, y_pred)) / denominator) + + +def smape( + y_true: Union[np.ndarray, tf.Tensor], y_pred: Union[np.ndarray, tf.Tensor], eps: float = 1e-8 +) -> Union[np.ndarray, tf.Tensor]: + """Symmetric Mean Absolute Percentage Error. + + Args: + y_true: Ground truth values. + y_pred: Predicted values. + eps: Small constant to avoid division by zero. + + Returns: + SMAPE as a percentage (0-200 scale). + """ + backend = _backend(y_true) + numerator = backend.abs(_sub(y_true, y_pred)) + denominator = (backend.abs(y_true) + backend.abs(y_pred)) / 2.0 + backend.array(eps, dtype=y_true.dtype) + return 100.0 * _reduce(numerator / denominator) + + +def r2_score( + y_true: Union[np.ndarray, tf.Tensor], y_pred: Union[np.ndarray, tf.Tensor] +) -> Union[np.ndarray, tf.Tensor]: + """R² coefficient of determination.""" + backend = _backend(y_true) + ss_res = backend.sum(backend.square(_sub(y_true, y_pred))) + ss_tot = backend.sum(backend.square(_sub(y_true, backend.mean(y_true)))) + return 1.0 - ss_res / ss_tot + + +def evaluate( + y_true: Union[np.ndarray, tf.Tensor], + y_pred: Union[np.ndarray, tf.Tensor], + metrics: Union[str, list] = "all", +) -> dict: + """Evaluate predictions with one or more metrics. + + Args: + y_true: Ground truth values. + y_pred: Predicted values. + metrics: Metric name, list of names, or "all" for all metrics. + + Returns: + Dictionary mapping metric names to their values. + """ + _METRICS = { + "mse": mse, + "mae": mae, + "rmse": rmse, + "mape": mape, + "smape": smape, + "r2": r2_score, + } + if metrics == "all": + names = list(_METRICS.keys()) + elif isinstance(metrics, str): + names = [metrics] + else: + names = metrics + + results = {} + for name in names: + if name not in _METRICS: + raise ValueError(f"Unknown metric '{name}'. Available: {list(_METRICS.keys())}") + results[name] = float(_METRICS[name](y_true, y_pred)) + return results + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +class _NumpyBackend: + @staticmethod + def array(x, dtype=None): + return np.array(x, dtype=dtype) + + abs = staticmethod(np.abs) + square = staticmethod(np.square) + sqrt = staticmethod(np.sqrt) + maximum = staticmethod(np.maximum) + sum = staticmethod(np.sum) + mean = staticmethod(np.mean) + + +class _TfBackend: + @staticmethod + def array(x, dtype=None): + return tf.constant(x, dtype=dtype) + + abs = staticmethod(tf.abs) + square = staticmethod(tf.square) + sqrt = staticmethod(tf.sqrt) + maximum = staticmethod(tf.maximum) + sum = staticmethod(tf.reduce_sum) + mean = staticmethod(tf.reduce_mean) + + +def _backend(x): + return _TfBackend if isinstance(x, tf.Tensor) else _NumpyBackend + + +def _sub(a, b): + return a - b + + +def _sqrt(x): + backend = _backend(x) + return backend.sqrt(x) + + +def _reduce(x): + backend = _backend(x) + return backend.mean(x) diff --git a/tfts/models/__init__.py b/tfts/models/__init__.py index e09bfea5..d7062822 100644 --- a/tfts/models/__init__.py +++ b/tfts/models/__init__.py @@ -10,3 +10,4 @@ AutoModelForUncertainty, ) from .base import BaseConfig, BaseModel +from .registry import list_models diff --git a/tfts/models/auto_config.py b/tfts/models/auto_config.py index d5033d61..8f8e361a 100644 --- a/tfts/models/auto_config.py +++ b/tfts/models/auto_config.py @@ -23,6 +23,11 @@ ("rwkv", "RWKVConfig"), ("patch_tst", "PatchTSTConfig"), ("deep_ar", "DeepARConfig"), + ("itransformer", "ITransformerConfig"), + ("timesfm", "TimesFmConfig"), + ("gpt", "GPTConfig"), + ("diffusion", "DiffusionConfig"), + ("tide", "TideConfig"), ] ) diff --git a/tfts/models/auto_model.py b/tfts/models/auto_model.py index c1fe09b5..1fc454b5 100644 --- a/tfts/models/auto_model.py +++ b/tfts/models/auto_model.py @@ -46,6 +46,11 @@ ("rwkv", "RWKV"), ("patch_tst", "PatchTST"), ("deep_ar", "DeepAR"), + ("itransformer", "ITransformer"), + ("timesfm", "TimesFm"), + ("gpt", "GPT"), + ("diffusion", "Diffusion"), + ("tide", "Tide"), ] ) @@ -56,8 +61,9 @@ class AutoModel(BaseModel): output tensor: [batch_size, predict_sequence_length, num_labels] """ - def __init__(self, model, config): - super().__init__(config=config) + def __init__(self, model, config, predict_sequence_length: Optional[int] = None): + predict_sequence_length = predict_sequence_length or getattr(model, "predict_sequence_length", 1) + super().__init__(predict_sequence_length=predict_sequence_length, config=config) self.model = model self.config = config @@ -95,13 +101,17 @@ def __call__( @classmethod def from_config(cls, config, predict_sequence_length: int = 1): model_name = config.model_type + if model_name not in MODEL_MAPPING_NAMES: + raise ValueError( + f"Unrecognized model: {model_name}. Should contain one of {', '.join(MODEL_MAPPING_NAMES.keys())}" + ) class_name = MODEL_MAPPING_NAMES[model_name] module = importlib.import_module(f".{model_name}", "tfts.models") model = getattr(module, class_name)(config=config, predict_sequence_length=predict_sequence_length) - return cls(model, config) + return cls(model, config, predict_sequence_length=predict_sequence_length) @classmethod - def from_pretrained(cls, weights_dir: Union[str, os.PathLike], predict_sequence_length: int = 1): + def from_pretrained(cls, weights_dir: Union[str, os.PathLike], predict_sequence_length: Optional[int] = None): config_path = os.path.join(weights_dir, "config.json") if not os.path.exists(config_path): raise FileNotFoundError(f"Config file not found at {config_path}") @@ -119,6 +129,7 @@ def from_pretrained(cls, weights_dir: Union[str, os.PathLike], predict_sequence_ # Dynamically get the correct Config subclass config = AutoConfig.for_model(model_type) config.update(config_dict) # update with the saved values + predict_sequence_length = predict_sequence_length or getattr(config, "predict_sequence_length", 1) # Build model and load weights model = cls.from_config(config, predict_sequence_length=predict_sequence_length) @@ -133,7 +144,7 @@ def from_pretrained(cls, weights_dir: Union[str, os.PathLike], predict_sequence_ model.build_model(inputs) model.model.load_weights(os.path.join(weights_dir, TF2_WEIGHTS_NAME)) - return cls(model, config) + return model except Exception as e: raise OSError( f"Error loading model weights from {weights_dir}. " @@ -156,10 +167,10 @@ def __call__( model_output = self.model(x, output_hidden_states=output_hidden_states, return_dict=return_dict) - if self.config.skip_connect_circle: + if getattr(self.config, "skip_connect_circle", False): x_mean = x[:, -self.predict_sequence_length :, 0:1] model_output = model_output + x_mean - elif self.config.skip_connect_mean: + elif getattr(self.config, "skip_connect_mean", False): x_mean = tf.tile(tf.reduce_mean(x[..., 0:1], axis=1, keepdims=True), [1, self.predict_sequence_length, 1]) model_output = model_output + x_mean return model_output @@ -225,22 +236,16 @@ def detect( @classmethod def from_pretrained(cls, weights_dir: Union[str, os.PathLike]): - model_path = os.path.join(weights_dir, "model.h5") - model = tf.keras.models.load_model(model_path) - logger.info(f"Load model from {weights_dir}") - config_path = os.path.join(weights_dir, "config.json") - if not os.path.exists(config_path): - raise FileNotFoundError(f"Config file not found at {config_path}") - - config = BaseConfig.from_json(config_path) # Load config from JSON - return cls(model, config) + model = AutoModel.from_pretrained(weights_dir) + logger.info(f"Loaded anomaly model from {weights_dir}") + return cls(model, model.config) @classmethod - def from_config(cls, config): + def from_config(cls, config, predict_sequence_length: int = 1): model_name = config.model_type class_name = MODEL_MAPPING_NAMES[model_name] module = importlib.import_module(f".{model_name}", "tfts.models") - model = getattr(module, class_name)(config=config) + model = getattr(module, class_name)(config=config, predict_sequence_length=predict_sequence_length) return cls(model, config) @@ -262,11 +267,11 @@ def __call__( return model_output @classmethod - def from_config(cls, config): + def from_config(cls, config, predict_sequence_length: int = 1): model_name = config.model_type class_name = MODEL_MAPPING_NAMES[model_name] module = importlib.import_module(f".{model_name}", "tfts.models") - model = getattr(module, class_name)(config=config) + model = getattr(module, class_name)(config=config, predict_sequence_length=predict_sequence_length) return cls(model, config) @@ -288,11 +293,11 @@ def __call__( return model_output @classmethod - def from_config(cls, config): + def from_config(cls, config, predict_sequence_length: int = 1): model_name = config.model_type class_name = MODEL_MAPPING_NAMES[model_name] module = importlib.import_module(f".{model_name}", "tfts.models") - model = getattr(module, class_name)(config=config) + model = getattr(module, class_name)(config=config, predict_sequence_length=predict_sequence_length) return cls(model, config) diff --git a/tfts/models/autoformer.py b/tfts/models/autoformer.py index 7494baf4..62a9871d 100644 --- a/tfts/models/autoformer.py +++ b/tfts/models/autoformer.py @@ -124,7 +124,7 @@ def __call__( # batch_size, _, n_feature = self.shape_layer(encoder_feature) # Encoder - encoder_output = self.encoder(x) + encoder_output = self.encoder(encoder_feature) encoder_output = self.dense1(encoder_output) encoder_output = self.dense2(encoder_output) diff --git a/tfts/models/base.py b/tfts/models/base.py index a8ddb106..63f70372 100644 --- a/tfts/models/base.py +++ b/tfts/models/base.py @@ -1,7 +1,7 @@ """Base class for config and model""" from abc import ABC, abstractmethod -import collections +from collections.abc import Mapping import json import logging import os @@ -17,23 +17,26 @@ class BaseModel(ABC): - """Bert model for time series forecasting. + """Base model for time series forecasting. - This model implements a transformer-based architecture (BERT) adapted for time series data. - It processes time series inputs through a transformer encoder and produces predictions - for future time steps. + Abstract base class that all tfts models inherit from. + Subclasses must implement __call__ and can optionally override build_model. Parameters ---------- predict_sequence_length : int, optional Number of future time steps to predict, by default 1 - config : BertConfig, optional + config : BaseConfig, optional Configuration parameters for the model, by default None """ def __init__(self, predict_sequence_length: int = 1, config: Optional["BaseConfig"] = None): self.config = config self.predict_sequence_length = predict_sequence_length + if isinstance(self.config, dict): + self.config["predict_sequence_length"] = predict_sequence_length + elif self.config is not None: + self.config.predict_sequence_length = predict_sequence_length self.model = None # Model should be defined later (may not be directly used in all subclasses) def build_model(self, inputs: tf.keras.layers.Input) -> tf.keras.Model: @@ -54,7 +57,18 @@ def build_model(self, inputs: tf.keras.layers.Input) -> tf.keras.Model: return self.model else: outputs = self(inputs) - return tf.keras.Model(inputs, outputs) + self.model = tf.keras.Model(inputs, outputs) + return self.model + + def _keras_model_for_saving(self) -> tf.keras.Model: + if isinstance(self.model, tf.keras.Model): + return self.model + if isinstance(self.model, BaseModel) and isinstance(self.model.model, tf.keras.Model): + return self.model.model + raise ValueError( + "Model weights cannot be saved before the model is built. " + "Call `build_model(...)` or train the model before saving weights." + ) def to_model(self): inputs = tf.keras.Input(shape=(self.config.input_shape)) @@ -109,37 +123,40 @@ def save_pretrained( logger.error(f"Provided path ({save_directory}) should be a directory, not a file") return + keras_model = self._keras_model_for_saving() + os.makedirs(save_directory, exist_ok=True) - self.config.architectures = [self.__class__.__name__[2:]] + # Use model_type from config if available, otherwise derive from class name + name = self.__class__.__name__ + architecture = getattr(self.config, "model_type", name) + self.config.architectures = [architecture] self.config.save_pretrained(save_directory) weights_file = os.path.join(save_directory, TF2_WEIGHTS_NAME) # Or the appropriate extension - try: - self.model.save_weights(weights_file) - logging.info(f"Model weights successfully saved in {weights_file}") - except Exception as e: - logging.error(f"Failed to save model weights to {weights_file}: {e}") - return + keras_model.save_weights(weights_file) + logging.info(f"Model weights successfully saved in {weights_file}") def save_weights(self, weights_path: str): if weights_path.endswith(".h5"): # User passed a full filepath weights_file = weights_path config_file = weights_path.replace(".h5", ".config.json") - os.makedirs(os.path.dirname(weights_file), exist_ok=True) + weights_dir = os.path.dirname(weights_file) + if weights_dir: + os.makedirs(weights_dir, exist_ok=True) else: # User passed a directory os.makedirs(weights_path, exist_ok=True) weights_file = os.path.join(weights_path, TF2_WEIGHTS_NAME) config_file = os.path.join(weights_path, CONFIG_NAME) - self.model.save_weights(weights_file) + self._keras_model_for_saving().save_weights(weights_file) self.config.to_json(config_file) logger.info(f"Model weights successfully saved in {weights_file}") def save_model(self, weights_dir: str): - self.model.save(weights_dir) + self._keras_model_for_saving().save(weights_dir) logger.info(f"Protobuf model successfully saved in {weights_dir}") def summary(self): @@ -240,7 +257,7 @@ def rec(nest, prefix, into): for k, v in nest.items(): if sep in k: raise ValueError(f"separator '{sep}' not allowed to be in key '{k}'") - if isinstance(v, collections.Mapping): + if isinstance(v, Mapping): rec(v, prefix + k + sep, into) else: into[prefix + k] = v diff --git a/tfts/models/bert.py b/tfts/models/bert.py index cd546aae..c1a308f1 100644 --- a/tfts/models/bert.py +++ b/tfts/models/bert.py @@ -181,7 +181,7 @@ def __call__( elif self.config.pooling_method == "last": encoder_output = memory[:, -1] else: - raise ValueError(f"Pooling method should be mean or last, while received {self.config.poolint_method}") + raise ValueError(f"Pooling method should be mean or last, while received {self.config.pooling_method}") for layer in self.dense_layers: encoder_output = layer(encoder_output) diff --git a/tfts/models/dlinear.py b/tfts/models/dlinear.py index ede55c32..3462b124 100644 --- a/tfts/models/dlinear.py +++ b/tfts/models/dlinear.py @@ -21,11 +21,13 @@ def __init__( kernel_size: int = 25, channels: int = 3, individual: bool = False, + dropout_rate: float = 0.0, ): super().__init__() self.kernel_size = kernel_size self.channels = channels # number of input features self.individual = individual + self.dropout_rate = dropout_rate self.activation: Optional[str] = None self.initializer: str = "glorot_uniform" diff --git a/tfts/models/informer.py b/tfts/models/informer.py index 60a4a49a..8c77f1f7 100644 --- a/tfts/models/informer.py +++ b/tfts/models/informer.py @@ -28,15 +28,15 @@ class InformerConfig(BaseConfig): def __init__( self, - hidden_size=64, - num_layers=1, - num_decoder_layers=None, - num_attention_heads=1, - attention_probs_dropout_prob=0.0, - ffn_intermediate_size=128, - hidden_dropout_prob=0.0, - prob_attention=False, - distil_conv=False, + hidden_size: int = 64, + num_layers: int = 1, + num_decoder_layers: Optional[int] = None, + num_attention_heads: int = 1, + attention_probs_dropout_prob: float = 0.0, + ffn_intermediate_size: int = 128, + hidden_dropout_prob: float = 0.0, + prob_attention: bool = False, + distil_conv: bool = False, ): super().__init__() self.hidden_size = hidden_size @@ -93,7 +93,7 @@ def __call__( teacher: Optional[tf.Tensor] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, - ): + ) -> tf.Tensor: """Informer call function""" x, encoder_feature, decoder_feature = self._prepare_3d_inputs(inputs, ignore_decoder_inputs=False) encoder_feature = self.encoder_embedding(encoder_feature) # batch * seq * embedding_size @@ -131,7 +131,7 @@ def __init__( self.prob_attention = prob_attention self.distil_conv = distil_conv - def build(self, input_shape): + def build(self, input_shape: tf.TensorShape) -> None: if not self.prob_attention: attn_layer = Attention(self.hidden_size, self.num_attention_heads, self.attention_probs_dropout_prob) else: @@ -155,7 +155,7 @@ def build(self, input_shape): self.norm_layer = LayerNormalization() super(Encoder, self).build(input_shape) - def call(self, x, mask=None): + def call(self, x: tf.Tensor, mask: Optional[tf.Tensor] = None) -> tf.Tensor: """Informer encoder call function""" if self.conv_layers is not None: for attn_layer, conv_layer in zip(self.layers, self.conv_layers): @@ -170,7 +170,7 @@ def call(self, x, mask=None): x = self.norm_layer(x) return x - def get_config(self): + def get_config(self) -> Dict[str, Any]: config = { "hidden_size": self.hidden_size, "num_layers": self.num_layers, diff --git a/tfts/models/pfn.py b/tfts/models/pfn.py index 8c5418ef..7438ad45 100644 --- a/tfts/models/pfn.py +++ b/tfts/models/pfn.py @@ -1,4 +1,33 @@ """ `ForecastPFN: Synthetically-Trained Zero-Shot Forecasting `_ + +TODO: Implement ForecastPFN model. This is currently a stub. """ + +import logging +from typing import Optional + +from .base import BaseConfig, BaseModel + +logger = logging.getLogger(__name__) + + +class PFNConfig(BaseConfig): + model_type: str = "pfn" + + def __init__(self, **kwargs): + super().__init__() + self.update(kwargs) + logger.warning("PFNConfig is a stub — the PFN model is not yet implemented.") + + +class PFN(BaseModel): + """ForecastPFN model — Not yet implemented.""" + + def __init__(self, predict_sequence_length: int = 1, config: Optional[PFNConfig] = None): + super().__init__() + raise NotImplementedError( + "PFN (ForecastPFN) is not implemented yet. " + "See https://github.com/LongxingTan/Time-series-prediction/issues for progress." + ) diff --git a/tfts/models/registry.py b/tfts/models/registry.py new file mode 100644 index 00000000..848128fc --- /dev/null +++ b/tfts/models/registry.py @@ -0,0 +1,266 @@ +"""Model Registry — centralized catalog of all available models. + +Provides metadata and discovery for every model in TFTS, similar to how +transformers maintains its model hub. +""" + +from collections import OrderedDict +from typing import Any, Dict, List, Optional + +# --------------------------------------------------------------------------- +# Model metadata registry +# --------------------------------------------------------------------------- + +MODEL_REGISTRY: Dict[str, Dict[str, Any]] = OrderedDict( + [ + ( + "seq2seq", + { + "class_name": "Seq2seq", + "config_class": "Seq2seqConfig", + "description": "Basic encoder-decoder sequence-to-sequence model with attention.", + "paper": "", + "tags": ["baseline", "encoder-decoder"], + }, + ), + ( + "rnn", + { + "class_name": "RNN", + "config_class": "RNNConfig", + "description": "Stacked LSTM/GRU with optional attention for time series.", + "paper": "", + "tags": ["baseline", "recurrent"], + }, + ), + ( + "wavenet", + { + "class_name": "WaveNet", + "config_class": "WaveNetConfig", + "description": "Dilated causal convolutions for long-range temporal dependencies.", + "paper": "https://arxiv.org/abs/1609.03499", + "tags": ["convolutional", "long-range"], + }, + ), + ( + "tcn", + { + "class_name": "TCN", + "config_class": "TCNConfig", + "description": "Temporal Convolutional Network with dilated causal convolutions.", + "paper": "https://arxiv.org/abs/1803.01271", + "tags": ["convolutional", "efficient"], + }, + ), + ( + "transformer", + { + "class_name": "Transformer", + "config_class": "TransformerConfig", + "description": "Classic encoder-decoder Transformer for time series forecasting.", + "paper": "https://arxiv.org/abs/1706.03762", + "tags": ["attention", "encoder-decoder"], + }, + ), + ( + "bert", + { + "class_name": "Bert", + "config_class": "BertConfig", + "description": "BERT-style masked pre-training adapted for time series.", + "paper": "https://arxiv.org/abs/1810.04805", + "tags": ["pretraining", "attention", "encoder-only"], + }, + ), + ( + "informer", + { + "class_name": "Informer", + "config_class": "InformerConfig", + "description": "Efficient Transformer with ProbSparse self-attention for long sequences.", + "paper": "https://arxiv.org/abs/2012.07436", + "tags": ["attention", "long-sequence", "efficient", "SOTA"], + }, + ), + ( + "autoformer", + { + "class_name": "AutoFormer", + "config_class": "AutoFormerConfig", + "description": "Auto-correlation mechanism with series decomposition for seasonal-trend modeling.", + "paper": "https://arxiv.org/abs/2106.13008", + "tags": ["decomposition", "seasonal", "SOTA"], + }, + ), + ( + "tft", + { + "class_name": "TFTransformer", + "config_class": "TFTransformerConfig", + "description": "Temporal Fusion Transformer — interpretable multi-horizon forecasting.", + "paper": "https://arxiv.org/abs/1912.09363", + "tags": ["interpretable", "multi-horizon", "attention", "SOTA"], + }, + ), + ( + "unet", + { + "class_name": "Unet", + "config_class": "UnetConfig", + "description": "U-Net style architecture with skip connections for time series.", + "paper": "https://arxiv.org/abs/1505.04597", + "tags": ["convolutional", "skip-connection"], + }, + ), + ( + "nbeats", + { + "class_name": "NBeats", + "config_class": "NBeatsConfig", + "description": "Neural basis expansion — pure MLP stack with interpretable basis functions.", + "paper": "https://arxiv.org/abs/1905.10437", + "tags": ["mlp", "interpretable", "SOTA", "basis-expansion"], + }, + ), + ( + "dlinear", + { + "class_name": "DLinear", + "config_class": "DLinearConfig", + "description": "Surprisingly strong linear baseline — questions Transformer necessity.", + "paper": "https://arxiv.org/abs/2205.13504", + "tags": ["linear", "simple", "baseline", "SOTA"], + }, + ), + ( + "rwkv", + { + "class_name": "RWKV", + "config_class": "RWKVConfig", + "description": "RNN-style efficient attention — linear complexity with Transformer quality.", + "paper": "https://arxiv.org/abs/2305.13048", + "tags": ["efficient", "attention", "recurrent"], + }, + ), + ( + "patch_tst", + { + "class_name": "PatchTST", + "config_class": "PatchTSTConfig", + "description": "Patch-based time series Transformer — segments time series into patches.", + "paper": "https://arxiv.org/abs/2211.14730", + "tags": ["patching", "attention", "SOTA"], + }, + ), + ( + "deep_ar", + { + "class_name": "DeepAR", + "config_class": "DeepARConfig", + "description": "Probabilistic autoregressive RNN for uncertainty-aware forecasting.", + "paper": "https://arxiv.org/abs/1704.04110", + "tags": ["probabilistic", "recurrent", "uncertainty"], + }, + ), + ( + "itransformer", + { + "class_name": "iTransformer", + "config_class": "iTransformerConfig", + "description": "Inverted Transformer — applies attention across variates instead of time.", + "paper": "https://arxiv.org/abs/2310.06625", + "tags": ["attention", "multivariate", "SOTA"], + }, + ), + ( + "timesfm", + { + "class_name": "TimesFM", + "config_class": "TimesFMConfig", + "description": "Google's foundation model for time series — decoder-only with patching.", + "paper": "https://arxiv.org/abs/2310.10688", + "tags": ["foundation-model", "patching", "decoder-only", "SOTA"], + }, + ), + ( + "gpt", + { + "class_name": "Gpt", + "config_class": "GptConfig", + "description": "GPT-style decoder-only Transformer adapted for time series.", + "paper": "", + "tags": ["decoder-only", "attention", "generative"], + }, + ), + ( + "diffusion", + { + "class_name": "Diffusion", + "config_class": "DiffusionConfig", + "description": "Denoising diffusion probabilistic model for time series generation.", + "paper": "https://arxiv.org/abs/2006.11239", + "tags": ["generative", "diffusion", "probabilistic"], + }, + ), + ( + "tide", + { + "class_name": "TiDE", + "config_class": "TiDEConfig", + "description": "Time-series Dense Encoder — simple MLP with covariate projection.", + "paper": "https://arxiv.org/abs/2304.08424", + "tags": ["mlp", "efficient", "covariates"], + }, + ), + ] +) + + +def list_models(tag: Optional[str] = None) -> List[str]: + """List all available model names, optionally filtered by tag. + + Args: + tag: If provided, only return models matching this tag + (e.g. ``'SOTA'``, ``'attention'``, ``'convolutional'``). + + Returns: + Sorted list of model names. + + Examples: + >>> tfts.list_models() + ['autoformer', 'bert', 'deep_ar', 'diffusion', 'dlinear', ...] + + >>> tfts.list_models(tag='SOTA') + ['autoformer', 'dlinear', 'informer', 'itransformer', 'nbeats', 'patch_tst', 'tft', 'timemixer', 'timesfm'] + """ + if tag is None: + return sorted(MODEL_REGISTRY.keys()) + return sorted(k for k, v in MODEL_REGISTRY.items() if tag in v.get("tags", [])) + + +def get_model_info(model_name: str) -> Dict[str, Any]: + """Get metadata for a specific model. + + Args: + model_name: Name of the model as used in the registry. + + Returns: + Dictionary with keys: class_name, config_class, description, paper, tags. + + Raises: + ValueError: If the model name is not recognized. + """ + if model_name not in MODEL_REGISTRY: + raise ValueError(f"Unknown model '{model_name}'. Available: {list_models()}") + return dict(MODEL_REGISTRY[model_name]) + + +def get_model_class_name(model_name: str) -> str: + """Resolve a model name to its Python class name.""" + return MODEL_REGISTRY[model_name]["class_name"] + + +def get_config_class_name(model_name: str) -> str: + """Resolve a model name to its Config class name.""" + return MODEL_REGISTRY[model_name]["config_class"] diff --git a/tfts/models/tcn.py b/tfts/models/tcn.py index 9a50889f..dccb33ad 100644 --- a/tfts/models/tcn.py +++ b/tfts/models/tcn.py @@ -1,12 +1,12 @@ """ -`WaveNet: A Generative Model for Raw Audio -`_ +`Temporal Convolutional Networks +`_ """ from typing import List, Optional, Tuple import tensorflow as tf -from tensorflow.keras.layers import Concatenate, Conv1D, Dense, Dropout, Lambda, ReLU, Reshape +from tensorflow.keras.layers import Concatenate, Conv1D, Dense, Dropout, Lambda, ReLU from tfts.layers.cnn_layer import ConvTemp from tfts.layers.dense_layer import DenseTemp @@ -19,20 +19,15 @@ class TCNConfig(BaseConfig): def __init__( self, - dilation_rates: List[int] = [2**i for i in range(4)], - kernel_sizes: List[int] = [2 for _ in range(4)], + dilation_rates: Optional[List[int]] = None, + kernel_sizes: Optional[List[int]] = None, filters: int = 128, dense_hidden_size: int = 64, ): - """ - Initializes the configuration for the Temporal Convolutional Network (TCN) model with the specified parameters. - - Args: - dilation_rates: List of dilation rates for each layer. - kernel_sizes: List of kernel sizes for each convolutional layer. - filters: The number of filters (channels) in each convolutional layer. - dense_hidden_size: The size of the dense hidden layer. - """ + if dilation_rates is None: + dilation_rates = [2**i for i in range(4)] + if kernel_sizes is None: + kernel_sizes = [2 for _ in range(4)] super().__init__() self.dilation_rates: List[int] = dilation_rates self.kernel_sizes: List[int] = kernel_sizes @@ -56,9 +51,9 @@ def __init__(self, predict_sequence_length: int = 1, config: Optional[TCNConfig] self.project1 = Dense(predict_sequence_length, activation=None) self.drop1 = Dropout(0.0) - self.dense1 = Dense(512, activation="relu") + self.dense1 = Dense(self.config.dense_hidden_size * 8, activation="relu") self.drop2 = Dropout(0.0) - self.dense2 = Dense(1024, activation="relu") + self.dense2 = Dense(self.config.dense_hidden_size * 16, activation="relu") def __call__( self, @@ -96,10 +91,10 @@ def __call__( encoder_output = self.dense1(encoder_output) encoder_output = self.drop2(encoder_output) encoder_output = self.dense2(encoder_output) - encoder_output = self.drop2(encoder_output) + encoder_output = self.drop1(encoder_output) outputs = self.project1(encoder_output) - outputs = Reshape((outputs.shape[1], 1))(outputs) + outputs = tf.keras.layers.Reshape((self.predict_sequence_length, 1))(outputs) # outputs = tf.tile(outputs, (1, self.predict_sequence_length, 1)) # outputs = self.dense3(encoder_outputs) diff --git a/tfts/models/timemixer.py b/tfts/models/timemixer.py index 2c14df3c..a782fe3f 100644 --- a/tfts/models/timemixer.py +++ b/tfts/models/timemixer.py @@ -1,4 +1,33 @@ """ `TimeMixer: Decomposable Multiscale Mixing for Time Series Forecasting `_ + +TODO: Implement TimeMixer model. This is currently a stub. """ + +import logging +from typing import Optional + +from .base import BaseConfig, BaseModel + +logger = logging.getLogger(__name__) + + +class TimeMixerConfig(BaseConfig): + model_type: str = "timemixer" + + def __init__(self, **kwargs): + super().__init__() + self.update(kwargs) + logger.warning("TimeMixerConfig is a stub — the TimeMixer model is not yet implemented.") + + +class TimeMixer(BaseModel): + """TimeMixer model — Not yet implemented.""" + + def __init__(self, predict_sequence_length: int = 1, config: Optional[TimeMixerConfig] = None): + super().__init__() + raise NotImplementedError( + "TimeMixer is not implemented yet. " + "See https://github.com/LongxingTan/Time-series-prediction/issues for progress." + ) diff --git a/tfts/models/timesfm.py b/tfts/models/timesfm.py index 2def7234..9f10c3d6 100644 --- a/tfts/models/timesfm.py +++ b/tfts/models/timesfm.py @@ -9,7 +9,7 @@ import tensorflow as tf from tensorflow.keras.layers import Dense, Dropout, LayerNormalization -from ..layers.attention_layer import Attention, SelfAttention +from ..layers.attention_layer import SelfAttention from ..layers.dense_layer import FeedForwardNetwork from ..layers.embed_layer import DataEmbedding from .base import BaseConfig, BaseModel diff --git a/tfts/models/transformer.py b/tfts/models/transformer.py index af21124e..70a57e63 100644 --- a/tfts/models/transformer.py +++ b/tfts/models/transformer.py @@ -4,7 +4,7 @@ """ import logging -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import numpy as np import tensorflow as tf @@ -13,7 +13,6 @@ from tfts.layers.attention_layer import Attention, SelfAttention from tfts.layers.dense_layer import FeedForwardNetwork from tfts.layers.embed_layer import DataEmbedding -from tfts.layers.mask_layer import CausalMask from .base import BaseConfig, BaseModel @@ -42,7 +41,7 @@ def __init__( classifier_dropout: Optional[float] = None, layer_norm_eps: float = 1e-12, pad_token_id: int = 0, - **kwargs: Dict[str, object] + **kwargs: Any ) -> None: """ Initializes the configuration for the Transformer model with the specified parameters. @@ -151,11 +150,6 @@ def __call__( decoder_feature, init_input=x[:, -1:, 0:1], encoder_memory=memory, teacher=teacher ) - # Example for new CausalMask usage: - # dummy = tf.zeros((B, L, 1)) - # mask_layer = CausalMask(num_attention_heads=1) - # casual_mask = mask_layer(dummy) - return decoder_outputs @@ -392,7 +386,7 @@ def build(self, input_shape): cross_attention_layer = Attention( self.hidden_size, self.num_attention_heads, self.attention_probs_dropout_prob ) - ffn_layer = FeedForwardNetwork(self.ffn_intermediate_size, self.hidden_size, self.hidden_dropout_prob) + ffn_layer = FeedForwardNetwork(self.hidden_size, self.ffn_intermediate_size, self.hidden_dropout_prob) ln_layer1 = LayerNormalization(epsilon=self.layer_norm_eps, dtype="float32") ln_layer2 = LayerNormalization(epsilon=self.layer_norm_eps, dtype="float32") ln_layer3 = LayerNormalization(epsilon=self.layer_norm_eps, dtype="float32") diff --git a/tfts/tasks/__init__.py b/tfts/tasks/__init__.py index e69de29b..19430619 100644 --- a/tfts/tasks/__init__.py +++ b/tfts/tasks/__init__.py @@ -0,0 +1,16 @@ +"""Task-specific handlers and pipelines for time series prediction.""" + +from .auto_task import AnomalyHead, ClassificationHead, GaussianHead, PredictionHead, SegmentationHead +from .base import BaseTask, ModelOutput +from .pipeline import Pipeline + +__all__ = [ + "AnomalyHead", + "BaseTask", + "ClassificationHead", + "GaussianHead", + "ModelOutput", + "Pipeline", + "PredictionHead", + "SegmentationHead", +] diff --git a/tfts/tasks/pipeline.py b/tfts/tasks/pipeline.py index 3175301a..07a8eb70 100644 --- a/tfts/tasks/pipeline.py +++ b/tfts/tasks/pipeline.py @@ -5,7 +5,7 @@ import numpy as np import tensorflow as tf -from ..models import AutoConfig, AutoModel +from ..training.runtime import create_distribution_strategy logger = logging.getLogger(__name__) @@ -23,16 +23,7 @@ def __init__(self, cfg, processor: Optional[Callable] = None, strategy: Optional def _setup_strategy(self): """Detects GPUs and returns the appropriate distribution strategy.""" - gpus = tf.config.list_physical_devices("GPU") - if len(gpus) > 1: - logger.info(f"Using MirroredStrategy with {len(gpus)} GPUs") - return tf.distribute.MirroredStrategy() - elif len(gpus) == 1: - logger.info("Using OneDeviceStrategy (1 GPU)") - return tf.distribute.OneDeviceStrategy(device="/gpu:0") - else: - logger.info("Using default strategy (CPU)") - return tf.distribute.get_strategy() + return create_distribution_strategy() def build_model(self, n_features, n_outputs): # Update model config with actual data dimensions @@ -40,6 +31,8 @@ def build_model(self, n_features, n_outputs): self.cfg.model.n_features = n_features self.cfg.model.n_outputs = n_outputs + from ..models import AutoConfig, AutoModel + config = AutoConfig()(self.cfg.model.name) config.output_size = n_outputs diff --git a/tfts/trainer.py b/tfts/trainer.py index 3a98ce20..6dda0b53 100644 --- a/tfts/trainer.py +++ b/tfts/trainer.py @@ -12,15 +12,16 @@ from .constants import CONFIG_NAME, TF2_WEIGHTS_INDEX_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME, TFTS_HOME, TFTS_HUB_CACHE from .models.base import BaseModel +from .training.runtime import configure_precision, create_distribution_strategy from .training_args import TrainingArguments -__all__ = ["Trainer", "KerasTrainer", "Seq2seqKerasTrainer", "set_seed"] +__all__ = ["Trainer", "KerasTrainer", "EagerTrainer", "Seq2seqKerasTrainer", "set_seed"] logger = logging.getLogger(__name__) -def set_seed(seed): +def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) os.environ["PYTHONHASHSEED"] = str(seed) @@ -40,66 +41,65 @@ def __init__( self.model = model self.config = model.config if hasattr(model, "config") else None self.args = args or TrainingArguments(output_dir=TFTS_HUB_CACHE) - self.strategy = strategy or tf.distribute.get_strategy() + self.strategy = strategy or create_distribution_strategy(self.args) - # with self.get_strategy_scope(strategy): - # self.model = self._setup_model(model) - # self.loss_fn = loss_fn - # self.metrics = metrics or [] - # self.optimizer = optimizer or self._create_optimizer() - # self.lr_scheduler = lr_scheduler or self._create_lr_scheduler() - # - # # Training state - # self.global_step = tf.Variable(0, trainable=False, dtype=tf.int32) - # if self.args.fp16: - # self._setup_mixed_precision() - - def evaluate(self): + def evaluate(self) -> None: pass - def get_train_dataloader(self): + def get_train_dataloader(self) -> Any: return - def get_eval_dataloader(self): + def get_eval_dataloader(self) -> Any: return - def get_test_dataloader(self): + def get_test_dataloader(self) -> Any: return - def get_learning_rates(self): + def get_learning_rates(self) -> Any: return - def create_accelerator_and_postprocess(self): + def create_accelerator_and_postprocess(self) -> Any: return - def get_distribution_strategy(): - gpus = tf.config.list_physical_devices("GPU") - if len(gpus) > 1: - return tf.distribute.MirroredStrategy() - elif len(gpus) == 1: - return tf.distribute.OneDeviceStrategy(device="/gpu:0") - else: - return tf.distribute.OneDeviceStrategy(device="/cpu:0") + @staticmethod + def get_distribution_strategy() -> tf.distribute.Strategy: + return create_distribution_strategy() - def get_strategy_scope(self): + def get_strategy_scope(self) -> Union[tf.distribute.Strategy.scope, nullcontext]: return self.strategy.scope() if self.strategy else nullcontext() - def _create_optimizer(self) -> tf.keras.optimizers.Optimizer: + def _create_optimizer( + self, + learning_rate: Optional[Union[float, tf.keras.optimizers.schedules.LearningRateSchedule]] = None, + ) -> tf.keras.optimizers.Optimizer: """Create optimizer with specified parameters.""" - return tf.keras.optimizers.Adam( - learning_rate=self.args.learning_rate, - beta_1=self.args.adam_beta1, - beta_2=self.args.adam_beta2, - epsilon=self.args.adam_epsilon, - weight_decay=self.args.weight_decay, - ) + learning_rate = learning_rate if learning_rate is not None else self.args.learning_rate + # tf.keras.optimizers.Adam does not support weight_decay directly. + # Use AdamW if available, otherwise fall back to standard Adam. + try: + return tf.keras.optimizers.AdamW( + learning_rate=learning_rate, + beta_1=self.args.adam_beta1, + beta_2=self.args.adam_beta2, + epsilon=self.args.adam_epsilon, + weight_decay=self.args.weight_decay, + ) + except AttributeError: + return tf.keras.optimizers.Adam( + learning_rate=learning_rate, + beta_1=self.args.adam_beta1, + beta_2=self.args.adam_beta2, + epsilon=self.args.adam_epsilon, + ) def _create_lr_scheduler(self) -> Optional[tf.keras.optimizers.schedules.LearningRateSchedule]: """Create learning rate scheduler based on arguments.""" + decay_steps = self.args.max_steps if self.args.max_steps > 0 else self.args.num_train_epochs + decay_steps = max(1, int(decay_steps)) if self.args.lr_scheduler_type == "linear": return tf.keras.optimizers.schedules.PolynomialDecay( initial_learning_rate=self.args.learning_rate, - decay_steps=self.args.max_steps if self.args.max_steps > 0 else self.args.num_train_epochs, + decay_steps=decay_steps, end_learning_rate=0, power=1.0, ) @@ -107,15 +107,7 @@ def _create_lr_scheduler(self) -> Optional[tf.keras.optimizers.schedules.Learnin def _setup_mixed_precision(self) -> None: """Configure mixed precision training.""" - policy = tf.keras.mixed_precision.Policy("mixed_float16") - tf.keras.mixed_precision.set_global_policy(policy) - logger.info("Mixed precision enabled.") - - # def _setup_ema(self) -> None: - # """Configure Exponential Moving Average if enabled.""" - # self.ema = None - # if self.config.use_ema: - # self.ema = tf.train.ExponentialMovingAverage(self.config.ema_decay) + configure_precision(self.args) def get_inputs(self, train_dataset): if isinstance(train_dataset, tf.data.Dataset): @@ -134,6 +126,18 @@ def get_inputs(self, train_dataset): raise ValueError("Unsupported dataset type. Expected tf.data.Dataset, keras.utils.Sequence, or list/tuple.") return inputs + def _keras_model_for_saving(self) -> tf.keras.Model: + if isinstance(self.model, tf.keras.Model): + return self.model + if isinstance(self.model, BaseModel): + model = self.model.model + if isinstance(model, tf.keras.Model): + return model + raise ValueError( + "Model weights cannot be saved before the model is built. " + "Call `train(...)`, or build the model with input shapes before `save_model(...)`." + ) + def _prepare_inputs_for_model( self, x: Union[np.ndarray, pd.DataFrame] ) -> Union[Dict[str, tf.keras.layers.Input], List[tf.keras.layers.Input], tf.keras.layers.Input]: @@ -158,7 +162,6 @@ def _prepare_inputs_for_model( def _save(self, output_dir: Optional[str] = None): output_dir = output_dir if output_dir is not None else TFTS_HOME - os.makedirs(output_dir, exist_ok=True) logger.info(f"Saving model checkpoint to {output_dir}") # save_model = self.model.model if hasattr(self.model, "model") else self.model # self.model.save_pretrained(output_dir) @@ -169,26 +172,39 @@ def _save(self, output_dir: Optional[str] = None): logger.error(f"Provided path ({save_directory}) should be a directory, not a file") return + keras_model = self._keras_model_for_saving() + os.makedirs(save_directory, exist_ok=True) - self.config.architectures = [self.model.__class__.__name__[2:]] - self.config.save_pretrained(save_directory) + # Use model_type from config if available, otherwise derive from class name + name = self.model.__class__.__name__ + architecture = getattr(self.config, "model_type", name) + if self.config is not None: + self.config.architectures = [architecture] + self.config.save_pretrained(save_directory) weights_file = os.path.join(save_directory, TF2_WEIGHTS_NAME) # Or the appropriate extension - try: - self.model.save_weights(weights_file) - logging.info(f"Model weights successfully saved in {weights_file}") - except Exception as e: - logging.error(f"Failed to save model weights to {weights_file}: {e}") - return + keras_model.save_weights(weights_file) + logging.info(f"Model weights successfully saved in {weights_file}") @property def global_batch_size(self): return self.args.per_device_train_batch_size * self.strategy.num_replicas_in_sync -class KerasTrainer(BaseTrainer): - """Keras trainer from tf.keras""" +class Trainer(BaseTrainer): + """Unified Trainer for time series tasks. + + Automatically selects loss, optimizer, and metrics based on task type. + Supports both eager and compiled (model.fit) training. + + Examples: + >>> from tfts import AutoModel, AutoConfig, Trainer + >>> config = AutoConfig.for_model("transformer") + >>> model = AutoModel.from_config(config, predict_sequence_length=12) + >>> trainer = Trainer(model) + >>> trainer.train(train_dataset, valid_dataset, epochs=50) + """ def __init__( self, @@ -197,81 +213,103 @@ def __init__( args: Optional[TrainingArguments] = None, **kwargs: Dict[str, object], ) -> None: - """ - Initializes the trainer with the model, loss function, optimizer, and other optional parameters. - - Args: - model: A Keras Model or Sequential instance to train. - strategy: Optional distribution strategy for multi-GPU or multi-node training. - **kwargs: Additional arguments that are passed to the instance as attributes. - """ super().__init__(model, args, strategy, **kwargs) self.model = model self.config = model.config if hasattr(model, "config") else None + self._task: str = "forecasting" # 'forecasting', 'classification', 'anomaly' for key, value in kwargs.items(): setattr(self, key, value) + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + def train( self, train_dataset: Union[tf.data.Dataset, List[tf.Tensor], Tuple[tf.Tensor, tf.Tensor]], valid_dataset: Optional[Union[tf.data.Dataset, List[tf.Tensor], Tuple[tf.Tensor, tf.Tensor]]] = None, - loss_fn: Union[Callable, tf.keras.losses.Loss, str] = "mse", - optimizer: Union[tf.keras.optimizers.Optimizer, str, Dict] = "adam", - epochs: int = 10, - batch_size: int = 64, + loss_fn: Union[Callable, tf.keras.losses.Loss, str, None] = None, + optimizer: Union[tf.keras.optimizers.Optimizer, str, Dict, None] = None, + epochs: Optional[int] = None, + batch_size: Optional[int] = None, steps_per_epoch: Optional[int] = None, metrics: Optional[Union[List[tf.keras.metrics.Metric], List[str]]] = None, callbacks: Optional[List[tf.keras.callbacks.Callback]] = None, + early_stopping_patience: Optional[int] = None, + checkpoint_dir: Optional[str] = None, + reduce_lr_patience: Optional[int] = None, run_eagerly: bool = True, verbose: int = 1, **kwargs: Dict[str, object], ) -> tf.keras.callbacks.History: - """ - Trains the model on the provided dataset. + """Train the model. Args: - train_dataset: A tf.data.Dataset or list/tuple of tensors (x_train, y_train). - valid_dataset: A tf.data.Dataset or list/tuple of tensors (x_valid, y_valid), optional. - loss_fn: A callable or Keras loss function. Default is MeanSquaredError. - optimizer: A Keras optimizer instance. Default is Adam with learning rate 0.003. - epochs: Number of epochs to train for. Default is 10. - batch_size: Number of samples per batch. Default is 64. - steps_per_epoch: Number of steps per epoch. Optional. - metrics: List of metrics for monitoring during training. Optional. - callbacks: List of keras callbacks during training. Optional. - run_eagerly: Whether to run eagerly. Default is True. - verbose: Verbosity level. Default is 1. - **kwargs: Additional keyword arguments for callbacks. + train_dataset: Training data — tf.data.Dataset or (x, y) tuple. + valid_dataset: Validation data (optional). + loss_fn: Loss function. Auto-selected if None. + optimizer: Optimizer. Defaults to Adam with config learning rate. + epochs: Number of training epochs. + batch_size: Samples per batch. + steps_per_epoch: Steps per epoch (for infinite datasets). + metrics: Keras metrics to track. Auto-selected if None. + callbacks: Keras callbacks. Built-in callbacks are added automatically. + early_stopping_patience: If set, adds EarlyStopping callback. + checkpoint_dir: If set, saves best model weights. + reduce_lr_patience: If set, adds ReduceLROnPlateau callback. + run_eagerly: Run eagerly (True) or with tf.function (False). + verbose: 0=silent, 1=progress bar, 2=one line per epoch. Returns: - A History object containing training logs. + A History object. """ - - if not callbacks: - callbacks: List[tf.keras.callbacks.Callback] = [] + callbacks = list(callbacks) if callbacks else [] + configure_precision(self.args) + epochs = int(epochs if epochs is not None else self.args.num_train_epochs) + batch_size = int(batch_size if batch_size is not None else self.global_batch_size) + + # Auto-build callbacks + callbacks += self._build_callbacks( + early_stopping_patience=early_stopping_patience, + checkpoint_dir=checkpoint_dir, + reduce_lr_patience=reduce_lr_patience, + ) with self.get_strategy_scope(): - # if lr_scheduler: # just set Optimizer(learning_rate=lr_scheduler) - # callbacks.append(tf.keras.callbacks.LearningRateScheduler(lr_scheduler, verbose=True)) + # Auto-select loss & optimizer + if loss_fn is None: + loss_fn = self._default_loss() + if optimizer is None: + optimizer = self._create_optimizer(learning_rate=self._create_lr_scheduler()) + if metrics is None: + metrics = self._default_metrics() if isinstance(optimizer, (str, dict)): optimizer = tf.keras.optimizers.get(optimizer) + # Build model if needed if not isinstance(self.model, tf.keras.Model): inputs = self.get_inputs(train_dataset) if "build_model" not in dir(self.model): - raise TypeError("Trainer model should either be `tf.keras.Model` or has `build_model()` method") + raise TypeError("Trainer model must be `tf.keras.Model` or have `build_model()`") self.model = self.model.build_model(inputs=inputs) - self.model.compile(loss=loss_fn, optimizer=optimizer, metrics=metrics, run_eagerly=run_eagerly) + compile_kwargs = { + "loss": loss_fn, + "optimizer": optimizer, + "metrics": metrics, + "run_eagerly": run_eagerly, + } + if self.args.jit_compile: + compile_kwargs["jit_compile"] = True + self.model.compile(**compile_kwargs) - trainable_params = np.sum([tf.keras.backend.count_params(w) for w in self.model.trainable_weights]) - tf.print(f"Trainable parameters: {trainable_params}") + trainable_params = int(np.sum([tf.keras.backend.count_params(w) for w in self.model.trainable_weights])) + logger.info(f"Trainable parameters: {trainable_params:,}") if isinstance(train_dataset, (list, tuple)): x_train, y_train = train_dataset - history = self.model.fit( x_train, y_train, @@ -295,23 +333,72 @@ def train( return history def fit(self, **params): + """Alias for train().""" return self.train(**params) - def predict(self, x_test: tf.Tensor) -> tf.Tensor: - return self.model(x_test) + def evaluate( + self, + dataset: Union[tf.data.Dataset, List[tf.Tensor], Tuple[tf.Tensor, tf.Tensor]], + metrics: Optional[List[str]] = None, + ) -> Dict[str, float]: + """Evaluate the model on a dataset. + + Args: + dataset: Evaluation data. + metrics: Metric names (e.g. ['mse', 'mae']). Uses built-in metrics if None. + + Returns: + Dictionary of metric_name -> value. + """ + from .metrics import evaluate as compute_metrics + + if isinstance(dataset, (list, tuple)): + x, y_true = dataset + y_pred = self.model(x, training=False) + elif isinstance(dataset, tf.data.Dataset): + y_true_list, y_pred_list = [], [] + for x_batch, y_batch in dataset: + y_pred_list.append(self.model(x_batch, training=False)) + y_true_list.append(y_batch) + y_true = tf.concat(y_true_list, axis=0) + y_pred = tf.concat(y_pred_list, axis=0) + else: + raise TypeError(f"Unsupported dataset type: {type(dataset)}") + + return compute_metrics(y_true, y_pred, metrics=metrics or "all") + + def predict(self, x: Union[tf.Tensor, np.ndarray, tf.data.Dataset]) -> np.ndarray: + """Make predictions. + + Args: + x: Single input tensor/array or a tf.data.Dataset. + + Returns: + Numpy array of predictions. + """ + if isinstance(x, tf.data.Dataset): + preds = [] + for batch in x: + inp = batch[0] if isinstance(batch, (tuple, list)) else batch + preds.append(self.model(inp, training=False)) + return tf.concat(preds, axis=0).numpy() + return self.model(x, training=False).numpy() def get_model(self) -> tf.keras.Model: + """Return the underlying Keras model.""" return self.model def save_model(self, output_dir: Optional[str] = None): - # save the model, checkpoint_dir if you use Checkpoint callback to save your best weights - if self.strategy.cluster_resolver and not self.strategy.cluster_resolver.task_type == "chief": - return - + """Save model weights and config.""" + if hasattr(self.strategy, "cluster_resolver") and self.strategy.cluster_resolver: + if self.strategy.cluster_resolver.task_type != "chief": + return output_dir = TFTS_HOME if output_dir is None else output_dir self._save(output_dir) - def plot(self, history, true: np.ndarray, pred: np.ndarray): + @staticmethod + def plot(history: np.ndarray, true: np.ndarray, pred: np.ndarray): + """Quick plot of history, ground truth, and predictions.""" import matplotlib.pyplot as plt train_length = history.shape[1] @@ -323,17 +410,88 @@ def plot(self, history, true: np.ndarray, pred: np.ndarray): plt.plot(range(train_length, train_length + pred_length), pred[example, :, 0], label="Predicted") plt.legend() + # ------------------------------------------------------------------ + # Auto-config helpers + # ------------------------------------------------------------------ + + def _default_loss(self) -> tf.keras.losses.Loss: + """Return a sensible default loss for the current task.""" + if self._task == "classification": + return tf.keras.losses.SparseCategoricalCrossentropy() + return tf.keras.losses.MeanSquaredError() + + def _default_optimizer(self) -> tf.keras.optimizers.Optimizer: + """Return a sensible default optimizer.""" + return self._create_optimizer(learning_rate=self._create_lr_scheduler()) + + def _default_metrics(self) -> List[str]: + """Return default metrics for monitoring.""" + if self._task == "classification": + return ["accuracy"] + return ["mae"] + + @staticmethod + def _build_callbacks( + early_stopping_patience: Optional[int] = None, + checkpoint_dir: Optional[str] = None, + reduce_lr_patience: Optional[int] = None, + ) -> List[tf.keras.callbacks.Callback]: + """Build standard Keras callbacks from simple parameters.""" + callbacks: List[tf.keras.callbacks.Callback] = [] + + if early_stopping_patience is not None: + callbacks.append( + tf.keras.callbacks.EarlyStopping( + monitor="val_loss", + patience=early_stopping_patience, + restore_best_weights=True, + ) + ) + if checkpoint_dir is not None: + os.makedirs(checkpoint_dir, exist_ok=True) + callbacks.append( + tf.keras.callbacks.ModelCheckpoint( + filepath=os.path.join(checkpoint_dir, "best_weights.weights.h5"), + monitor="val_loss", + save_best_only=True, + save_weights_only=True, + ) + ) + if reduce_lr_patience is not None: + callbacks.append( + tf.keras.callbacks.ReduceLROnPlateau( + monitor="val_loss", + patience=reduce_lr_patience, + factor=0.5, + min_lr=1e-7, + ) + ) + return callbacks + + +# --------------------------------------------------------------------------- +# Backward-compatible aliases +# --------------------------------------------------------------------------- + +KerasTrainer = Trainer # legacy alias + -class Seq2seqKerasTrainer(KerasTrainer): - """As the transformers forum mentioned: https://discuss.huggingface.co/t/trainer-vs-seq2seqtrainer/3145/2 - Seq2SeqTrainer is mostly about predict_with_generate.""" +class Seq2seqKerasTrainer(Trainer): + """Seq2SeqTrainer — supports predict_with_generate. + + See: https://discuss.huggingface.co/t/trainer-vs-seq2seqtrainer/3145/2 + """ def __init__(self, *args, **kwargs): - super(Seq2seqKerasTrainer, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) + +class EagerTrainer(object): + """Low-level custom training loop trainer (legacy). -class Trainer(object): - """Custom trainer for tensorflow with support for CPU, GPU, and multi-GPU.""" + Use ``Trainer`` for the standard high-level API. This class provides + manual gradient-tape-based training for users who need full control. + """ def __init__( self, @@ -364,35 +522,7 @@ def train( max_grad_norm: float = 5.0, transform: Optional[Callable] = None, ) -> None: - """ - Trains the model using the provided data loaders. - - Parameters - ---------- - train_loader : Union[tf.data.Dataset, Generator] - The training data loader, which can be a `tf.data.Dataset` or a Python generator. - valid_loader : Union[tf.data.Dataset, Generator, None], optional - The validation data loader, by default None. - epochs : int, optional - The number of epochs to train the model, by default 10. - learning_rate : float, optional - The initial learning rate for the optimizer, by default 3e-4. - verbose : int, optional - The verbosity level (0 = silent, 1 = progress bar, 2 = one line per epoch), by default 1. - eval_metric : Union[Callable, List[Callable], None], optional - The evaluation metric(s) to use for validation, by default None. - model_dir : Optional[str], optional - The directory to save the model weights, by default "../weights". - use_ema : bool, optional - Whether to use exponential moving average (EMA) for the model weights, by default False. - stop_no_improve_epochs : Optional[int], optional - If provided, training will stop if the validation metric does not improve for the specified - number of epochs, by default None. - max_grad_norm : float, optional - the max gradient while backprop. - transform : Optional[Callable], optional - A function to transform the data before feeding it to the model, by default None. - """ + """Train with manual gradient tape loop.""" self.loss_fn = loss_fn if optimizer is None: optimizer = tf.keras.optimizers.Adam(0.003) @@ -411,9 +541,8 @@ def train( if model_dir is None: model_dir = TFTS_HUB_CACHE - if stop_no_improve_epochs is not None: - no_improve_epochs: int = 0 - best_metric: float = float("inf") + no_improve_epochs: int = 0 + best_metric: float = float("inf") if not isinstance(self.model, tf.keras.Model): if "build_model" not in dir(self.model): @@ -433,13 +562,13 @@ def train( self.ema = None for epoch in range(epochs): - train_loss, train_scores = self.train_loop(train_loader) - log_str = f"Epoch: {epoch + 1}, Train Loss: {train_loss:.4f}" # noqa + train_loss, train_scores = self._train_loop(train_loader) + log_str = f"Epoch: {epoch + 1}, Train Loss: {train_loss:.4f}" if valid_loader is not None: - valid_loss, valid_scores = self.valid_loop(valid_loader) - log_str += f", Valid Loss: {valid_loss:.4f}" # noqa - log_str + ",".join([" Valid Metrics{}: {:.4f}".format(i, me) for i, me in enumerate(valid_scores)]) + valid_loss, valid_scores = self._valid_loop(valid_loader) + log_str += f", Valid Loss: {valid_loss:.4f}" + log_str += ",".join([f" Valid Metrics{i}: {me:.4f}" for i, me in enumerate(valid_scores)]) if (stop_no_improve_epochs is not None) and (eval_metric is not None): if valid_scores[0] >= best_metric: @@ -453,39 +582,32 @@ def train( logger.info(log_str) - # self.export_model(model_dir, only_pb=True) # save the model - def fit(self, **params): return self.train(**params) - def train_loop(self, train_loader): + def _train_loop(self, train_loader: Any) -> tuple[float, list[Any]]: train_loss: float = 0.0 y_trues, y_preds = [], [] - for step, (x_train, y_train) in enumerate(train_loader): - y_pred, step_loss = self.train_step(x_train, y_train) + y_pred, step_loss = self._train_step(x_train, y_train) train_loss += step_loss y_preds.append(y_pred) y_trues.append(y_train) - scores = [] if self.eval_metric: y_preds = tf.concat(y_preds, axis=0) y_trues = tf.concat(y_trues, axis=0) - for metric in self.eval_metric: scores.append(metric(y_trues, y_preds)) return train_loss / (step + 1), scores - def train_step(self, x_train, y_train): + def _train_step(self, x_train: tf.Tensor, y_train: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor]: with tf.GradientTape() as tape: y_pred = self.model(x_train, training=True) loss = self.loss_fn(y_train, y_pred) - gradients = tape.gradient(loss, self.model.trainable_variables) gradients = [(tf.clip_by_value(grad, -self.max_grad_norm, self.max_grad_norm)) for grad in gradients] _ = self.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables)) - if self.lr_scheduler is not None: lr = self.lr_scheduler(self.global_step) self.optimizer.learning_rate.assign(lr) @@ -493,54 +615,45 @@ def train_step(self, x_train, y_train): lr = self.learning_rate self.optimizer.learning_rate.assign(lr) self.global_step.assign_add(1) - # logger.info(f'Step: {self.global_step.numpy()}, Loss: {loss}' return y_pred, loss - def valid_loop(self, valid_loader): + def _valid_loop(self, valid_loader: Any) -> tuple[float, list[Any]]: valid_loss: float = 0.0 y_valid_trues, y_valid_preds = [], [] - for valid_step, (x_valid, y_valid) in enumerate(valid_loader): - y_valid_pred, valid_step_loss = self.valid_step(x_valid, y_valid) + y_valid_pred, valid_step_loss = self._valid_step(x_valid, y_valid) valid_loss += valid_step_loss y_valid_trues.append(y_valid) y_valid_preds.append(y_valid_pred) - valid_scores = [] if self.eval_metric: y_valid_preds = tf.concat(y_valid_preds, axis=0) y_valid_trues = tf.concat(y_valid_trues, axis=0) - for metric in self.eval_metric: valid_scores.append(metric(y_valid_trues, y_valid_preds)) return valid_loss / (valid_step + 1), valid_scores - def valid_step(self, x_valid, y_valid): - + def _valid_step(self, x_valid: tf.Tensor, y_valid: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor]: y_valid_pred = self.model(x_valid, training=False) valid_loss = self.loss_fn(y_valid, y_valid_pred) return y_valid_pred, valid_loss - def predict(self, test_loader): + def predict(self, test_loader: Any) -> tuple[tf.Tensor, tf.Tensor]: y_test_trues, y_test_preds = [], [] for x_test, y_test in test_loader: y_test_pred = self.model(x_test, training=False) y_test_preds.append(y_test_pred) y_test_trues.append(y_test) - y_test_trues = tf.concat(y_test_trues, axis=0) y_test_preds = tf.concat(y_test_preds, axis=0) return tf.squeeze(y_test_trues, axis=-1), y_test_preds def save_model(self, model_dir, only_pb=True): - # save the model if not model_dir.endswith(".keras"): model_dir = f"{model_dir}.keras" - os.makedirs(os.path.dirname(model_dir), exist_ok=True) self.model.save(model_dir) logger.info(f"Model successfully saved in {model_dir}") - if not only_pb: self.model.save_weights(f"{model_dir}.ckpt") logger.info(f"Model weights successfully saved in {model_dir}.ckpt") diff --git a/tfts/training/__init__.py b/tfts/training/__init__.py new file mode 100644 index 00000000..f96587ab --- /dev/null +++ b/tfts/training/__init__.py @@ -0,0 +1,5 @@ +"""Training runtime helpers.""" + +from .runtime import configure_precision, create_distribution_strategy + +__all__ = ["configure_precision", "create_distribution_strategy"] diff --git a/tfts/training/runtime.py b/tfts/training/runtime.py new file mode 100644 index 00000000..f3b251e0 --- /dev/null +++ b/tfts/training/runtime.py @@ -0,0 +1,78 @@ +"""Runtime setup for training. + +This module keeps accelerator and precision choices in one place so trainers, +pipelines, and future custom loops can share the same behavior. +""" + +import logging +from typing import Optional + +import tensorflow as tf + +from ..training_args import TrainingArguments + +logger = logging.getLogger(__name__) + + +def _create_mirrored_strategy() -> tf.distribute.Strategy: + """Create MirroredStrategy when this TensorFlow build supports it.""" + mirrored_strategy = getattr(tf.distribute, "MirroredStrategy", None) + if mirrored_strategy is None: + logger.warning("MirroredStrategy is not available in this TensorFlow build; using default strategy.") + return tf.distribute.get_strategy() + return mirrored_strategy() + + +def create_distribution_strategy(args: Optional[TrainingArguments] = None) -> tf.distribute.Strategy: + """Create a TensorFlow distribution strategy from training arguments. + + Args: + args: Training arguments. If omitted, uses automatic local device detection. + + Returns: + A TensorFlow distribution strategy. + """ + strategy_name = args.strategy if args is not None else "auto" + + if strategy_name == "default": + return tf.distribute.get_strategy() + + if strategy_name == "multi_worker": + logger.info("Using MultiWorkerMirroredStrategy") + return tf.distribute.MultiWorkerMirroredStrategy() + + gpus = tf.config.list_physical_devices("GPU") + + if strategy_name == "mirrored": + logger.info("Using MirroredStrategy") + return _create_mirrored_strategy() + + if strategy_name == "one_device": + device = "/gpu:0" if gpus else "/cpu:0" + logger.info("Using OneDeviceStrategy on %s", device) + return tf.distribute.OneDeviceStrategy(device=device) + + if len(gpus) > 1: + logger.info("Using MirroredStrategy with %s GPUs", len(gpus)) + return _create_mirrored_strategy() + if len(gpus) == 1: + logger.info("Using OneDeviceStrategy on /gpu:0") + return tf.distribute.OneDeviceStrategy(device="/gpu:0") + + logger.info("Using default TensorFlow strategy") + return tf.distribute.get_strategy() + + +def configure_precision(args: TrainingArguments) -> tf.keras.mixed_precision.Policy: + """Apply the configured Keras mixed precision policy. + + Args: + args: Training arguments containing the precision setting. + + Returns: + The active mixed precision policy. + """ + policy = tf.keras.mixed_precision.Policy(args.precision) + tf.keras.mixed_precision.set_global_policy(policy) + logger.info("Using precision policy: %s", policy.name) + return policy diff --git a/tfts/training_args.py b/tfts/training_args.py index 6b9042d1..b17086d8 100644 --- a/tfts/training_args.py +++ b/tfts/training_args.py @@ -34,7 +34,7 @@ class TrainingArguments: default=-1, metadata={"help": "If > 0: set total number of training steps to perform. Override num_train_epochs."}, ) - lr_scheduler_type: Union[str] = field( + lr_scheduler_type: str = field( default="linear", metadata={"help": "The scheduler type to use."}, ) @@ -50,6 +50,17 @@ class TrainingArguments: default=0.0, metadata={"help": "Linear warmup over warmup_ratio fraction of total steps."} ) warmup_steps: int = field(default=0, metadata={"help": "Linear warmup over warmup_steps."}) + strategy: str = field( + default="auto", + metadata={ + "help": "Distribution strategy to use: 'auto', 'default', 'one_device', 'mirrored', or 'multi_worker'." + }, + ) + precision: str = field( + default="float32", + metadata={"help": "Numerical precision policy: 'float32', 'mixed_float16', or 'mixed_bfloat16'."}, + ) + jit_compile: bool = field(default=False, metadata={"help": "Whether to enable XLA compilation in Keras compile."}) bf16: bool = field( default=False, @@ -66,4 +77,18 @@ class TrainingArguments: ) def __post_init__(self): - pass + valid_strategies = {"auto", "default", "one_device", "mirrored", "multi_worker"} + if self.strategy not in valid_strategies: + raise ValueError(f"strategy must be one of {sorted(valid_strategies)}, got {self.strategy}") + + if self.fp16 and self.bf16: + raise ValueError("fp16 and bf16 cannot both be enabled") + + if self.fp16: + self.precision = "mixed_float16" + elif self.bf16: + self.precision = "mixed_bfloat16" + + valid_precision = {"float32", "mixed_float16", "mixed_bfloat16"} + if self.precision not in valid_precision: + raise ValueError(f"precision must be one of {sorted(valid_precision)}, got {self.precision}") diff --git a/tfts/tuner/__init__.py b/tfts/tuner/__init__.py new file mode 100644 index 00000000..0c6b5236 --- /dev/null +++ b/tfts/tuner/__init__.py @@ -0,0 +1,5 @@ +"""TFTS Tuner — hyperparameter search utilities.""" + +from .optuna_tuner import OptunaTuner + +__all__ = ["OptunaTuner"] diff --git a/tfts/tuner/optuna_tuner.py b/tfts/tuner/optuna_tuner.py new file mode 100644 index 00000000..086394a4 --- /dev/null +++ b/tfts/tuner/optuna_tuner.py @@ -0,0 +1,244 @@ +"""OptunaTuner — hyperparameter search for TFTS models. + +Wraps Optuna to tune model configs and training hyperparameters +with a simple ``search()`` API. Optuna is an optional dependency; +importing this module raises a helpful error if ``optuna`` is not +installed. +""" + +import logging +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import numpy as np + +from ..models.auto_config import AutoConfig +from ..models.auto_model import AutoModel +from ..trainer import Trainer + +logger = logging.getLogger(__name__) + + +def _require_optuna() -> "optuna": # type: ignore[name-defined] # noqa: F821 + """Lazy-import optuna with a clear error message.""" + try: + import optuna + + return optuna + except ImportError: + raise ImportError("optuna is required for OptunaTuner. " "Install it with: pip install optuna") + + +class OptunaTuner: + """Hyperparameter tuner powered by Optuna. + + Given a list of model names and a parameter search space, builds and + trains TFTS models inside an Optuna study and returns the best + configuration. + + Args: + train_data: Training dataset — ``(x_train, y_train)`` tuple or + ``tf.data.Dataset``. + valid_data: Validation dataset — same format. + predict_sequence_length: Forecast horizon. + metric: Metric to optimize. ``'mse'``, ``'mae'``, or a callable + ``f(y_true, y_pred) -> float``. + direction: ``'minimize'`` (default) or ``'maximize'``. + + Examples: + >>> from tfts.tuner import OptunaTuner + >>> tuner = OptunaTuner(train_data, valid_data, predict_sequence_length=7) + >>> best_params, best_score = tuner.search( + ... param_space={ + ... "model_type": ["rnn", "dlinear"], + ... "learning_rate": [1e-4, 1e-2], + ... }, + ... n_trials=20, + ... ) + """ + + def __init__( + self, + train_data: Any, + valid_data: Any, + predict_sequence_length: int = 1, + metric: Union[str, Callable] = "mse", + direction: str = "minimize", + ) -> None: + self.train_data = train_data + self.valid_data = valid_data + self.predict_sequence_length = predict_sequence_length + self.metric = metric + self.direction = direction + + self._study: Optional[Any] = None # optuna.Study + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def search( + self, + param_space: Dict[str, Any], + n_trials: int = 20, + epochs: int = 10, + verbose: int = 0, + ) -> Tuple[Dict[str, Any], float]: + """Run the hyperparameter search. + + Args: + param_space: Dictionary mapping parameter names to search + ranges:: + + { + "model_type": ["rnn", "dlinear"], # categorical + "learning_rate": [1e-4, 1e-2], # log-uniform + "hidden_size": [32, 256], # int uniform + "num_layers": [1, 4], # int uniform + } + + - **list of strings** → categorical choice + - **list of two floats** `[lo, hi]` → + - float log-uniform if both > 0 and lo < 1 + - int uniform otherwise + n_trials: Number of Optuna trials. + epochs: Training epochs per trial. + verbose: Keras verbosity (0 = silent). + + Returns: + ``(best_params, best_score)`` tuple. + """ + optuna = _require_optuna() + + # Silence optuna logs unless user wants them + optuna.logging.set_verbosity(optuna.logging.WARNING) + + study = optuna.create_study(direction=self.direction) + study.optimize( + func=lambda trial: self._objective(trial, param_space, epochs, verbose), + n_trials=n_trials, + ) + + self._study = study + return study.best_params, study.best_value + + def get_best_params(self) -> Dict[str, Any]: + """Return the best parameters found so far. + + Raises: + RuntimeError: If :meth:`search` has not been called. + """ + if self._study is None: + raise RuntimeError("No search has been run yet. Call .search() first.") + return dict(self._study.best_params) + + def get_best_score(self) -> float: + """Return the best score found so far.""" + if self._study is None: + raise RuntimeError("No search has been run yet. Call .search() first.") + return float(self._study.best_value) + + def get_study(self) -> Any: + """Return the underlying Optuna study (for advanced plotting). + + Returns ``None`` before :meth:`search` is called. + """ + return self._study + + # ------------------------------------------------------------------ + # Internal + # ------------------------------------------------------------------ + + def _objective( + self, + trial: Any, + param_space: Dict[str, Any], + epochs: int, + verbose: int, + ) -> float: + """Optuna objective: build model, train, evaluate.""" + # Suggest parameters + params = self._suggest_params(trial, param_space) + + # Extract model_type (required) + model_type = params.pop("model_type") + + # Build config with suggested overrides + config = AutoConfig.for_model(model_type) + for key, value in params.items(): + if hasattr(config, key): + setattr(config, key, value) + else: + logger.debug(f"Config key {key!r} not found on {model_type} config, skipping") + + # Build & train model + model = AutoModel.from_config(config, predict_sequence_length=self.predict_sequence_length) + trainer = Trainer(model) + + # Unpack learning_rate if present — pass via optimizer + lr = params.get("learning_rate", 1e-3) + optimizer = _default_optimizer(lr) + + trainer.train( + self.train_data, + valid_dataset=self.valid_data, + epochs=epochs, + optimizer=optimizer, + verbose=verbose, + ) + + # Evaluate on validation data + metrics = trainer.evaluate(self.valid_data, metrics=[self.metric] if isinstance(self.metric, str) else None) + score = self._extract_score(metrics) + return score + + def _suggest_params(self, trial: Any, param_space: Dict[str, Any]) -> Dict[str, Any]: + """Convert param_space into optuna suggestions.""" + params: Dict[str, Any] = {} + + for name, spec in param_space.items(): + # Categorical: list of strings + if isinstance(spec, list) and len(spec) > 0 and isinstance(spec[0], str): + params[name] = trial.suggest_categorical(name, spec) + continue + + # Numeric range: [lo, hi] + if isinstance(spec, (list, tuple)) and len(spec) == 2: + lo, hi = spec + if isinstance(lo, float) or isinstance(hi, float): + # log-uniform for learning_rate style params + if lo > 0 and lo < 1: + params[name] = trial.suggest_float(name, lo, hi, log=True) + else: + params[name] = trial.suggest_float(name, lo, hi) + else: + params[name] = trial.suggest_int(name, int(lo), int(hi)) + continue + + raise ValueError( + f"Cannot infer suggestion type for param {name!r} with spec {spec!r}. " + "Use a list of strings for categorical or [lo, hi] for numeric." + ) + + return params + + def _extract_score(self, metrics: Dict[str, float]) -> float: + """Extract a single scalar score from the metrics dict.""" + if isinstance(self.metric, str) and self.metric in metrics: + return float(metrics[self.metric]) + # Fallback: return first value + if metrics: + return float(next(iter(metrics.values()))) + raise ValueError("No metric value could be extracted from the evaluation results.") + + def __repr__(self) -> str: + return f"OptunaTuner(metric={self.metric!r}, direction={self.direction!r})" + + +def _default_optimizer(lr: float): + """Create a default optimizer with the given learning rate.""" + import tensorflow as tf + + try: + return tf.keras.optimizers.AdamW(learning_rate=lr, weight_decay=1e-4) + except AttributeError: + return tf.keras.optimizers.Adam(learning_rate=lr)