tirex-mirror 2025.10.28__tar.gz → 2025.11.4__tar.gz

{tirex_mirror-2025.10.28 → tirex_mirror-2025.11.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.10.28
+Version: 2025.11.4
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT

{tirex_mirror-2025.10.28 → tirex_mirror-2025.11.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tirex-mirror"
-version = "2025.10.28"
+version = "2025.11.04"
 description = "Unofficial mirror of NX-AI/tirex for packaging"
 readme = "README.md"
 requires-python = ">=3.11"

{tirex_mirror-2025.10.28 → tirex_mirror-2025.11.4}/src/tirex/api_adapter/forecast.py

@@ -23,7 +23,6 @@ def _format_output(
     quantiles: torch.Tensor,
     means: torch.Tensor,
     sample_meta: list[dict],
-    quantile_levels: list[float],
     output_type: Literal["torch", "numpy", "gluonts"],
 ):
     if output_type == "torch":
@@ -35,7 +34,7 @@ def _format_output(
             from .gluon import format_gluonts_output
         except ImportError:
             raise ValueError("output_type glutonts needs GluonTs but GluonTS is not available (not installed)!")
-        return format_gluonts_output(quantiles, means, sample_meta, quantile_levels)
+        return format_gluonts_output(quantiles, means, sample_meta)
     else:
         raise ValueError(f"Invalid output type: {output_type}")
 
@@ -61,14 +60,13 @@ def _pad_time_series_batch(
     return padded
 
 
-def _as_generator(batches, fc_func, quantile_levels, output_type, **predict_kwargs):
+def _as_generator(batches, fc_func, output_type, **predict_kwargs):
     for batch_ctx, batch_meta in batches:
         quantiles, mean = fc_func(batch_ctx, **predict_kwargs)
         yield _format_output(
             quantiles=quantiles,
             means=mean,
             sample_meta=batch_meta,
-            quantile_levels=quantile_levels,
             output_type=output_type,
         )
 
@@ -152,7 +150,6 @@ def _gen_forecast(
     fc_func,
     batches,
     output_type,
-    quantile_levels,
     yield_per_batch,
     resample_strategy: str | None = None,
     max_context: int = 2016,
@@ -173,7 +170,7 @@ def _gen_forecast(
         fc_func = partial(_call_fc_with_padding, base_fc_func)
 
     if yield_per_batch:
-        return _as_generator(batches, fc_func, quantile_levels, output_type, **predict_kwargs)
+        return _as_generator(batches, fc_func, output_type, **predict_kwargs)
 
     prediction_q = []
     prediction_m = []
@@ -191,34 +188,48 @@ def _gen_forecast(
         quantiles=prediction_q,
         means=prediction_m,
         sample_meta=sample_meta,
-        quantile_levels=quantile_levels,
         output_type=output_type,
     )
 
 
-def _common_forecast_doc():
-    common_doc = f"""
+class ForecastModel(ABC):
+    @abstractmethod
+    def _forecast_quantiles(self, batch, **predict_kwargs):
+        pass
+
+    @property
+    def max_context_length(self) -> int:
+        # retrieve the max_context attribute of the model configuration if present
+        return getattr(getattr(self, "config", None), "max_context", 2016)
+
+    def forecast(
+        self,
+        context: ContextType,
+        output_type: Literal["torch", "numpy", "gluonts"] = "torch",
+        batch_size: int = 512,
+        yield_per_batch: bool = False,
+        resample_strategy: Literal["frequency"] | None = None,
+        **predict_kwargs,
+    ):
+        """
         This method takes historical context data as input and outputs probabilistic forecasts.
 
         Args:
             output_type (Literal["torch", "numpy", "gluonts"], optional):
                 Specifies the desired format of the returned forecasts:
-                - "torch": Returns forecasts as `torch.Tensor` objects [batch_dim, forecast_len, |quantile_levels|]
-                - "numpy": Returns forecasts as `numpy.ndarray` objects [batch_dim, forecast_len, |quantile_levels|]
+                - "torch": Returns forecasts as `torch.Tensor` objects [batch_dim, forecast_len, quantile_count]
+                - "numpy": Returns forecasts as `numpy.ndarray` objects [batch_dim, forecast_len, quantile_count]
                 - "gluonts": Returns forecasts as a list of GluonTS `Forecast` objects.
                 Defaults to "torch".
 
             batch_size (int, optional): The number of time series instances to process concurrently by the model.
                                         Defaults to 512. Must be $>= 1$.
 
-            quantile_levels (List[float], optional): Quantile levels for which predictions should be generated.
-                                                     Defaults to (0.1, 0.2, ..., 0.9).
-
             yield_per_batch (bool, optional): If `True`, the method will act as a generator, yielding
                                               forecasts batch by batch as they are computed.
                                               Defaults to `False`.
 
-            resample_strategy (Optional[str], optional): Choose a resampling strategy. Allowed values: {RESAMPLE_STRATEGIES}.
+            resample_strategy (Optional[str], optional): Choose a resampling strategy. Allowed values: "frequency".
                                                 If `None`, no resampling is applied. Currently only "frequency" is supported.
 
             **predict_kwargs: Additional keyword arguments that are passed directly to the underlying
@@ -233,32 +244,7 @@ def _common_forecast_doc():
                   - If `output_type="torch"`: `Tuple[torch.Tensor, torch.Tensor]` (quantiles, mean).
                   - If `output_type="numpy"`: `Tuple[numpy.ndarray, numpy.ndarray]` (quantiles, mean).
                   - If `output_type="gluonts"`: A `List[gluonts.model.forecast.Forecast]` of all forecasts.
-        """
-    return common_doc
-
 
-class ForecastModel(ABC):
-    @abstractmethod
-    def _forecast_quantiles(self, batch, **predict_kwargs):
-        pass
-
-    @property
-    def max_context_length(self) -> int:
-        # retrieve the max_context attribute of the model configuration if present
-        return getattr(getattr(self, "config", None), "max_context", 2016)
-
-    def forecast(
-        self,
-        context: ContextType,
-        output_type: Literal["torch", "numpy", "gluonts"] = "torch",
-        batch_size: int = 512,
-        quantile_levels: list[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
-        yield_per_batch: bool = False,
-        resample_strategy: Literal["frequency"] | None = None,
-        **predict_kwargs,
-    ):
-        f"""
-        {_common_forecast_doc}
         Args:
             context (ContextType): The historical "context" data of the time series:
                 - `torch.Tensor`: 1D `[context_length]` or 2D `[batch_dim, context_length]` tensor
@@ -272,7 +258,6 @@ class ForecastModel(ABC):
             self._forecast_quantiles,
             batches,
             output_type,
-            quantile_levels,
             yield_per_batch,
             resample_strategy=resample_strategy,
             max_context=self.max_context_length,
@@ -284,14 +269,44 @@ class ForecastModel(ABC):
         gluonDataset,
         output_type: Literal["torch", "numpy", "gluonts"] = "torch",
         batch_size: int = 512,
-        quantile_levels: list[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
         yield_per_batch: bool = False,
         resample_strategy: Literal["frequency"] | None = None,
         data_kwargs: dict = {},
         **predict_kwargs,
     ):
-        f"""
-        {_common_forecast_doc()}
+        """
+        This method takes historical context data as input and outputs probabilistic forecasts.
+
+        Args:
+            output_type (Literal["torch", "numpy", "gluonts"], optional):
+                Specifies the desired format of the returned forecasts:
+                - "torch": Returns forecasts as `torch.Tensor` objects [batch_dim, forecast_len, quantile_count]
+                - "numpy": Returns forecasts as `numpy.ndarray` objects [batch_dim, forecast_len, quantile_count]
+                - "gluonts": Returns forecasts as a list of GluonTS `Forecast` objects.
+                Defaults to "torch".
+
+            batch_size (int, optional): The number of time series instances to process concurrently by the model.
+                                        Defaults to 512. Must be $>= 1$.
+
+            yield_per_batch (bool, optional): If `True`, the method will act as a generator, yielding
+                                              forecasts batch by batch as they are computed.
+                                              Defaults to `False`.
+
+            resample_strategy (Optional[str], optional): Choose a resampling strategy. Allowed values: "frequency".
+                                                If `None`, no resampling is applied. Currently only "frequency" is supported.
+
+            **predict_kwargs: Additional keyword arguments that are passed directly to the underlying
+                              prediction mechanism of the pre-trained model. Refer to the model's
+                              internal prediction method documentation for available options.
+
+        Returns:
+            The return type depends on `output_type` and `yield_per_batch`:
+                - If `yield_per_batch` is `True`: An iterator that yields forecasts. Each yielded item
+                  will correspond to a batch of forecasts in the format specified by `output_type`.
+                - If `yield_per_batch` is `False`: A single object containing all forecasts.
+                  - If `output_type="torch"`: `Tuple[torch.Tensor, torch.Tensor]` (quantiles, mean).
+                  - If `output_type="numpy"`: `Tuple[numpy.ndarray, numpy.ndarray]` (quantiles, mean).
+                  - If `output_type="gluonts"`: A `List[gluonts.model.forecast.Forecast]` of all forecasts.
 
         Args:
             gluonDataset (gluon_ts.dataset.common.Dataset): A GluonTS dataset object containing the
@@ -311,7 +326,6 @@ class ForecastModel(ABC):
             self._forecast_quantiles,
             batches,
             output_type,
-            quantile_levels,
             yield_per_batch,
             resample_strategy=resample_strategy,
             max_context=self.max_context_length,
@@ -323,14 +337,44 @@ class ForecastModel(ABC):
         hf_dataset,
         output_type: Literal["torch", "numpy", "gluonts"] = "torch",
         batch_size: int = 512,
-        quantile_levels: list[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
         yield_per_batch: bool = False,
         resample_strategy: Literal["frequency"] | None = None,
         data_kwargs: dict = {},
         **predict_kwargs,
     ):
-        f"""
-        {_common_forecast_doc()}
+        """
+        This method takes historical context data as input and outputs probabilistic forecasts.
+
+        Args:
+            output_type (Literal["torch", "numpy", "gluonts"], optional):
+                Specifies the desired format of the returned forecasts:
+                - "torch": Returns forecasts as `torch.Tensor` objects [batch_dim, forecast_len, quantile_count]
+                - "numpy": Returns forecasts as `numpy.ndarray` objects [batch_dim, forecast_len, quantile_count]
+                - "gluonts": Returns forecasts as a list of GluonTS `Forecast` objects.
+                Defaults to "torch".
+
+            batch_size (int, optional): The number of time series instances to process concurrently by the model.
+                                        Defaults to 512. Must be $>= 1$.
+
+            yield_per_batch (bool, optional): If `True`, the method will act as a generator, yielding
+                                              forecasts batch by batch as they are computed.
+                                              Defaults to `False`.
+
+            resample_strategy (Optional[str], optional): Choose a resampling strategy. Allowed values: "frequency".
+                                                If `None`, no resampling is applied. Currently only "frequency" is supported.
+
+            **predict_kwargs: Additional keyword arguments that are passed directly to the underlying
+                              prediction mechanism of the pre-trained model. Refer to the model's
+                              internal prediction method documentation for available options.
+
+        Returns:
+            The return type depends on `output_type` and `yield_per_batch`:
+                - If `yield_per_batch` is `True`: An iterator that yields forecasts. Each yielded item
+                  will correspond to a batch of forecasts in the format specified by `output_type`.
+                - If `yield_per_batch` is `False`: A single object containing all forecasts.
+                  - If `output_type="torch"`: `Tuple[torch.Tensor, torch.Tensor]` (quantiles, mean).
+                  - If `output_type="numpy"`: `Tuple[numpy.ndarray, numpy.ndarray]` (quantiles, mean).
+                  - If `output_type="gluonts"`: A `List[gluonts.model.forecast.Forecast]` of all forecasts.
 
         Args:
             hf_dataset (datasets.Dataset): A Hugging Face `Dataset` object containing the
@@ -352,7 +396,6 @@ class ForecastModel(ABC):
             self._forecast_quantiles,
             batches,
             output_type,
-            quantile_levels,
             yield_per_batch,
             resample_strategy=resample_strategy,
             max_context=self.max_context_length,

{tirex_mirror-2025.10.28 → tirex_mirror-2025.11.4}/src/tirex/api_adapter/gluon.py

@@ -30,7 +30,8 @@ def get_gluon_batches(gluonDataset: Dataset, batch_size: int, **gluon_kwargs):
     return _batch_iterable(map(_get_gluon_ts_map(**gluon_kwargs), gluonDataset), batch_size)
 
 
-def format_gluonts_output(quantile_forecasts: torch.Tensor, mean_forecasts, meta: list[dict], quantile_levels):
+def format_gluonts_output(quantile_forecasts: torch.Tensor, mean_forecasts, meta: list[dict]):
+    quantile_levels: list[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
     forecasts = []
     for i in range(quantile_forecasts.shape[0]):
         start_date = meta[i].get(FieldName.START, pd.Period("01-01-2000", freq=meta[i].get("freq", "h")))

{tirex_mirror-2025.10.28 → tirex_mirror-2025.11.4}/src/tirex/util.py

@@ -419,7 +419,7 @@ def run_fft_analysis(
     scaling : {'amplitude', 'power', 'raw'}
         - 'amplitude': one-sided amplitude spectrum with window-power compensation.
         - 'power'    : one-sided power (not density) with window-power compensation.
-        - 'raw'      : |rfft(yw)| (no normalization, mostly for debugging).
+        - 'raw'      : rfft(yw) (no normalization, mostly for debugging).
     peak_prominence : float
         Absolute threshold on the normalized spectrum for peak detection.
 
@@ -527,7 +527,6 @@ def resampling_factor(inverted_freq, path_size):
 def custom_find_peaks(
     f,
     spec,
-    *,
     max_peaks=5,
     prominence_threshold=0.1,
     min_period=64,
@@ -615,3 +614,16 @@ def round_up_to_next_multiple_of(x: int, multiple_of: int) -> int:
 def dataclass_from_dict(cls, dict: dict):
     class_fields = {f.name for f in fields(cls)}
     return cls(**{k: v for k, v in dict.items() if k in class_fields})
+
+
+def select_quantile_subset(quantiles: torch.Tensor, quantile_levels: list[float]):
+    """
+    Select specified quantile levels from the quantiles.
+    """
+    trained_quantiles = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+    assert set(quantile_levels).issubset(trained_quantiles), (
+        f"Only the following quantile_levels are supported: {quantile_levels}"
+    )
+    quantile_levels_idx = [trained_quantiles.index(q) for q in quantile_levels]
+    quantiles_idx = torch.tensor(quantile_levels_idx, dtype=torch.long, device=quantiles.device)
+    return torch.index_select(quantiles, dim=-1, index=quantiles_idx).squeeze(-1)

{tirex_mirror-2025.10.28 → tirex_mirror-2025.11.4}/src/tirex_mirror.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.10.28
+Version: 2025.11.4
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT

{tirex_mirror-2025.10.28 → tirex_mirror-2025.11.4}/tests/test_chronos_zs.py

@@ -10,6 +10,7 @@ import fev
 import pytest
 
 from tirex import ForecastModel, load_model
+from tirex.util import select_quantile_subset
 
 
 def geometric_mean(s):
@@ -25,14 +26,13 @@ def eval_task(model, task):
         loaded_targets = [t for t in past_data["target"]]
 
         start_time = time.monotonic()
-        quantiles, means = model.forecast(
-            loaded_targets, quantile_levels=task.quantile_levels, prediction_length=task.horizon
-        )
+        quantiles, means = model.forecast(loaded_targets, prediction_length=task.horizon)
         inference_time += time.monotonic() - start_time
 
         predictions_dict = {"predictions": means}
+        quantiles_subset = select_quantile_subset(quantiles, task.quantile_levels)
         for idx, level in enumerate(task.quantile_levels):
-            predictions_dict[str(level)] = quantiles[:, :, idx]
+            predictions_dict[str(level)] = quantiles_subset[:, :, idx]
 
         predictions_per_window.append(
             fev.combine_univariate_predictions_to_multivariate(

{tirex_mirror-2025.10.28 → tirex_mirror-2025.11.4}/tests/test_forecast_adapter.py

@@ -25,18 +25,15 @@ class DummyForecaster(ForecastModel):
         return fc_random_from_tensor(batch, **kwargs)
 
 
-quantile_levels = list(np.linspace(0.1, 0.9, 9))
-
-
 # ----- Tests: Output formatting -----
 def test_format_output_shapes(dummy_fc_func):
     B, L = 2, 5
     PL = 10
     q, m = dummy_fc_func(torch.rand(B, L), prediction_length=PL)
-    out_q, out_m = _format_output(q, m, [{}] * B, quantile_levels, "torch")
+    out_q, out_m = _format_output(q, m, [{}] * B, "torch")
     assert isinstance(out_q, torch.Tensor)
     assert isinstance(out_m, torch.Tensor)
-    assert out_q.shape == (B, PL, len(quantile_levels))
+    assert out_q.shape == (B, PL, 9)
     assert out_m.shape == (B, PL)
 
 
@@ -44,10 +41,10 @@ def test_format_output_shapes(dummy_fc_func):
     B, L = 2, 5
     PL = 10
     q, m = dummy_fc_func(torch.rand(B, L), prediction_length=PL)
-    out_q, out_m = _format_output(q, m, [{}] * B, quantile_levels, "numpy")
+    out_q, out_m = _format_output(q, m, [{}] * B, "numpy")
     assert isinstance(out_q, np.ndarray)
     assert isinstance(out_m, np.ndarray)
-    assert out_q.shape == (B, PL, len(quantile_levels))
+    assert out_q.shape == (B, PL, 9)
     assert out_m.shape == (B, PL)
 
 
@@ -55,7 +52,7 @@ def test_format_output_shapes(dummy_fc_func):
 def test_gen_forecast_single_batch(dummy_fc_func):
     context = torch.rand((5, 20))
     batches = get_batches(context, batch_size=2)
-    q, m = _gen_forecast(dummy_fc_func, batches, "torch", quantile_levels, prediction_length=10, yield_per_batch=False)
+    q, m = _gen_forecast(dummy_fc_func, batches, "torch", prediction_length=10, yield_per_batch=False)
     assert q.shape == (5, 10, 9)
     assert m.shape == (5, 10)
 
@@ -63,9 +60,7 @@ def test_gen_forecast_single_batch(dummy_fc_func):
 def test_gen_forecast_iterator(dummy_fc_func):
     context = torch.rand((5, 20))
     batches = get_batches(context, batch_size=2)
-    iterator = _gen_forecast(
-        dummy_fc_func, batches, "torch", quantile_levels, prediction_length=10, yield_per_batch=True
-    )
+    iterator = _gen_forecast(dummy_fc_func, batches, "torch", prediction_length=10, yield_per_batch=True)
     outputs = list(iterator)
     assert len(outputs) == 3
     for i, (q, m) in enumerate(outputs):