update base

Author: Daniel Park

.ai/AGENTS.md

@@ -54,4 +54,4 @@ gh pr list --state open --search "<short area keywords>"
 - First-time contributors should not use code agents to submit PRs or issues. Code agents should
   warn users about this policy if the user does not appear to be an existing contributor, including
   the risk of being blocked if they do it regardless. More information can be found in
-  `CONTRIBUTING.md`.
+  `CONTRIBUTING.md`.

.ai/review-rules.md

@@ -80,4 +80,4 @@ Read these on demand, when the diff touches the relevant area. Do not read all o
 - Anchor every inline comment to a line the diff actually touches.
 - State the concrete failure: what input, what goes wrong. "This breaks when `predict_sequence_length` > `train_length` during decode" beats "consider handling the edge case".
 - If you are unsure, say so in one clause and move on — do not pad a weak finding into a paragraph.
-- Reference the doc that supports your point by repo-root path, so the author can find it.
+- Reference the doc that supports your point by repo-root path, so the author can find it.

Makefile

@@ -12,7 +12,7 @@ style:  ## Run formatters and linters (black, isort, flake8, pre-commit)
 
 ## Run all unit tests
 test:  ## Run unit tests using unittest
-	python -m unittest discover
+	CUDA_VISIBLE_DEVICES=0 python -m unittest discover
 
 ## Build the documentation
 docs:  ## Build HTML documentation using Sphinx

poetry.lock

@@ -36,7 +36,6 @@ classifiers = [
     "Intended Audience :: Developers",
     "Intended Audience :: Science/Research",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
@@ -57,9 +56,9 @@ homepage = "https://time-series-prediction.readthedocs.io"
 tfts-forecast = "tfts.cli.forecasting:main"
 
 [tool.poetry.dependencies]
-python = ">=3.8,<=3.13"
+python = ">=3.9,<3.13"
 pandas = ">=1.3.0"
-# better install independently
+tensorflow = ">=2.13"
 
 
 [tool.poetry.group.dev.dependencies]

pyproject.toml

@@ -67,12 +67,21 @@ def test_get_data_invalid_name(self):
 
     def test_get_data_test_size_validation(self):
         """Test get_data validates test_size parameter"""
-        with self.assertRaises(AssertionError):
+        with self.assertRaises(ValueError):
             get_data("sine", 10, 4, test_size=-0.1)
 
-        with self.assertRaises(AssertionError):
+        with self.assertRaises(ValueError):
             get_data("sine", 10, 4, test_size=1.5)
 
+    def test_get_sine_seeded_reproducibility(self):
+        first = get_sine(10, 4, test_size=0, n_examples=5, seed=42)
+        second = get_sine(10, 4, test_size=0, n_examples=5, seed=42)
+        different = get_sine(10, 4, test_size=0, n_examples=5, seed=43)
+
+        np.testing.assert_array_equal(first[0], second[0])
+        np.testing.assert_array_equal(first[1], second[1])
+        self.assertFalse(np.array_equal(first[0], different[0]))
+
     def test_get_data_airpassengers(self):
         """Test get_data dispatcher for airpassengers dataset"""
         train_length = 12

tests/test_data/test_get_data.py

@@ -0,0 +1,35 @@
+import unittest
+
+import numpy as np
+import tensorflow as tf
+
+from tfts.metrics import mae, mape, mse, smape
+
+
+class MetricsTest(unittest.TestCase):
+    def test_integer_inputs_use_floating_point_arithmetic(self):
+        y_true = np.array([0, 2], dtype=np.int32)
+        y_pred = np.array([0, 1], dtype=np.int32)
+
+        self.assertTrue(np.isfinite(mape(y_true, y_pred)))
+        self.assertTrue(np.isfinite(smape(y_true, y_pred)))
+        self.assertAlmostEqual(mse(y_true, y_pred), 0.5)
+        self.assertAlmostEqual(mae(y_true, y_pred), 0.5)
+
+    def test_tensor_metrics_are_finite_for_integer_zeros(self):
+        y_true = tf.constant([0, 2], dtype=tf.int32)
+        y_pred = tf.constant([0, 1], dtype=tf.int32)
+
+        self.assertTrue(bool(tf.math.is_finite(mape(y_true, y_pred))))
+        self.assertTrue(bool(tf.math.is_finite(smape(y_true, y_pred))))
+
+    def test_metric_function_is_logged_by_keras_fit(self):
+        model = tf.keras.Sequential([tf.keras.layers.Input((1,)), tf.keras.layers.Dense(1)])
+        model.compile(optimizer="sgd", loss="mse", metrics=[mae])
+        history = model.fit(np.ones((2, 1)), np.ones((2, 1)), epochs=1, verbose=0)
+
+        self.assertIn("mae", history.history)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_metrics.py

@@ -1,9 +1,11 @@
 import unittest
 
+import tensorflow as tf
+
 import tfts
 from tfts.models.auto_config import AutoConfig
 from tfts.models.auto_model import AutoModel
-from tfts.models.registry import list_models
+from tfts.models.registry import get_model_info, list_models
 
 
 class TestAutoModel(unittest.TestCase):
@@ -28,3 +30,37 @@ def test_listed_models_can_be_instantiated(self):
                 config = AutoConfig.for_model(model_name)
                 model = AutoModel.from_config(config, predict_sequence_length=2)
                 self.assertEqual(model.config.model_type, model_name)
+
+    def test_registry_metadata_matches_auto_dispatch(self):
+        for model_name in list_models():
+            with self.subTest(model_name=model_name):
+                info = get_model_info(model_name)
+                config = AutoConfig.for_model(model_name)
+                model = AutoModel.from_config(config, predict_sequence_length=2)
+
+                self.assertEqual(type(config).__name__, info["config_class"])
+                self.assertEqual(type(model.model).__name__, info["class_name"])
+
+    def test_listed_models_satisfy_forward_contract(self):
+        # AutoFormer currently operates in its configured hidden dimension;
+        # N-BEATS is intentionally univariate. Other models accept a small
+        # multivariate input.
+        feature_counts = {"autoformer": 64, "nbeats": 1}
+        multivariate_outputs = {"diffusion", "itransformer"}
+
+        predict_sequence_length = 8
+        for model_name in list_models():
+            with self.subTest(model_name=model_name):
+                feature_count = feature_counts.get(model_name, 3)
+                config = AutoConfig.for_model(model_name)
+                model = AutoModel.from_config(config, predict_sequence_length=predict_sequence_length)
+                output = model(tf.random.normal([1, 16, feature_count]))
+
+                if model_name == "deep_ar":
+                    self.assertEqual(len(output), 2)
+                    self.assertEqual(output[0].shape, (1, 16, 1))
+                    self.assertEqual(output[1].shape, (1, 16, 1))
+                elif model_name in multivariate_outputs:
+                    self.assertEqual(output.shape, (1, predict_sequence_length, feature_count))
+                else:
+                    self.assertEqual(output.shape, (1, predict_sequence_length, 1))

tests/test_models/test_auto_config.py

@@ -26,6 +26,23 @@ def test_config_save_and_load(self):
             loaded_config = BaseConfig.from_json(config_path)
             self.assertEqual(loaded_config.to_dict(), self.config.to_dict())
 
+    def test_pretrained_config_accepts_directory_or_file(self):
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            self.config.save_pretrained(tmpdirname)
+            config_path = os.path.join(tmpdirname, "config.json")
+
+            from_directory = BaseConfig.from_pretrained(tmpdirname)
+            from_file = BaseConfig.from_pretrained(config_path)
+
+        self.assertEqual(from_directory.to_dict(), self.config_data)
+        self.assertEqual(from_file.to_dict(), self.config_data)
+
+    def test_save_pretrained_reports_serialization_errors(self):
+        self.config.invalid_value = object()
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            with self.assertRaises(TypeError):
+                self.config.save_pretrained(tmpdirname)
+
 
 class TestBaseModel(unittest.TestCase):
     def setUp(self):

tests/test_models/test_base.py

@@ -4,7 +4,6 @@
 
 import logging
 import os
-import random
 from typing import Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -219,9 +218,10 @@ def download_and_extract(name: str) -> str:
 def get_data(
     name: str = "sine", train_length: int = 24, predict_sequence_length: int = 8, test_size: float = 0.1, **kwargs
 ) -> Union[Tuple[np.ndarray, np.ndarray], Tuple[Tuple[np.ndarray, np.ndarray]], pd.DataFrame]:
-    assert (test_size >= 0) & (test_size <= 1), "test_size is the ratio of test dataset"
+    if not 0 <= test_size <= 1:
+        raise ValueError("test_size must be between 0 and 1")
     if name == "sine":
-        return get_sine(train_length, predict_sequence_length, test_size=test_size)
+        return get_sine(train_length, predict_sequence_length, test_size=test_size, **kwargs)
 
     elif name == "airpassengers":
         return get_air_passengers(train_length, predict_sequence_length, test_size=test_size)
@@ -245,7 +245,11 @@ def get_data(
 
 
 def get_sine(
-    train_sequence_length: int = 24, predict_sequence_length: int = 8, test_size: float = 0.2, n_examples: int = 100
+    train_sequence_length: int = 24,
+    predict_sequence_length: int = 8,
+    test_size: float = 0.2,
+    n_examples: int = 100,
+    seed: Optional[int] = None,
 ) -> Union[Tuple[np.ndarray, np.ndarray], Tuple[Tuple[np.ndarray, np.ndarray]]]:
     """
     Generate synthetic sine wave data.
@@ -259,10 +263,18 @@ def get_sine(
     Returns:
     (tuple): Two tuples of numpy arrays containing training and validation data.
     """
+    if train_sequence_length < 1 or predict_sequence_length < 1:
+        raise ValueError("sequence lengths must be positive")
+    if n_examples < 1:
+        raise ValueError("n_examples must be positive")
+    if not 0 <= test_size <= 1:
+        raise ValueError("test_size must be between 0 and 1")
+
+    rng = np.random.default_rng(seed)
     x: List[np.ndarray] = []
     y: List[np.ndarray] = []
     for _ in range(n_examples):
-        rand = random.random() * 2 * np.pi
+        rand = rng.uniform(0.0, 2.0 * np.pi)
         sig1 = np.sin(np.linspace(rand, 3.0 * np.pi + rand, train_sequence_length + predict_sequence_length))
         sig2 = np.cos(np.linspace(rand, 3.0 * np.pi + rand, train_sequence_length + predict_sequence_length))
 
@@ -376,14 +388,13 @@ def get_ar_data(
     if noise < 0:
         raise ValueError("noise parameter must be non-negative")
 
-    if seed is not None:
-        np.random.seed(seed)
+    rng = np.random.default_rng(seed)
 
     # Sample parameters for each series
-    linear_trends = np.random.normal(size=n_series)[:, None] / timesteps
-    quadratic_trends = np.random.normal(size=n_series)[:, None] / timesteps**2
-    seasonalities = np.random.normal(size=n_series)[:, None]
-    levels = level * np.random.normal(size=n_series)[:, None]
+    linear_trends = rng.normal(size=n_series)[:, None] / timesteps
+    quadratic_trends = rng.normal(size=n_series)[:, None] / timesteps**2
+    seasonalities = rng.normal(size=n_series)[:, None]
+    levels = level * rng.normal(size=n_series)[:, None]
 
     # Generate time index
     x = np.arange(timesteps)[None, :]
@@ -401,7 +412,7 @@ def get_ar_data(
     series = levels + series
 
     # Add noise
-    series = series * (1 + noise * np.random.normal(size=series.shape))
+    series = series * (1 + noise * rng.normal(size=series.shape))
 
     # Apply exponential transform if requested
     if exp: