tirex-mirror 2025.10.2__py3-none-any.whl → 2025.10.7__py3-none-any.whl

tirex/api_adapter/forecast.py

@@ -2,15 +2,22 @@
 # This software may be used and distributed according to the terms of the NXAI Community License Agreement.
 
 from abc import ABC, abstractmethod
-from typing import Literal
+from functools import partial
+from math import ceil
+from typing import Literal, Optional
 
 import torch
 
+from tirex.util import frequency_resample
+
 from .standard_adapter import ContextType, get_batches
 
 DEF_TARGET_COLUMN = "target"
 DEF_META_COLUMNS = ("start", "item_id")
 
+# Allowed resampling strategies (extend as new strategies are implemented)
+RESAMPLE_STRATEGIES: list[str] = ["frequency"]
+
 
 def _format_output(
     quantiles: torch.Tensor,
@@ -33,6 +40,27 @@ def _format_output(
         raise ValueError(f"Invalid output type: {output_type}")
 
 
+def _pad_time_series_batch(
+    batch_series: list[torch.Tensor],
+    max_length: int,
+) -> torch.Tensor:
+    if not batch_series:
+        return torch.empty((0, max_length))
+
+    first = batch_series[0]
+    dtype = first.dtype if first.is_floating_point() else torch.float32
+    device = first.device
+
+    padded = torch.full((len(batch_series), max_length), float("nan"), dtype=dtype, device=device)
+
+    for idx, series in enumerate(batch_series):
+        series = series.to(padded.dtype)
+        series_len = series.shape[0]
+        padded[idx, max_length - series_len :] = series
+
+    return padded
+
+
 def _as_generator(batches, fc_func, quantile_levels, output_type, **predict_kwargs):
     for batch_ctx, batch_meta in batches:
         quantiles, mean = fc_func(batch_ctx, **predict_kwargs)
@@ -45,7 +73,105 @@ def _as_generator(batches, fc_func, quantile_levels, output_type, **predict_kwar
         )
 
 
-def _gen_forecast(fc_func, batches, output_type, quantile_levels, yield_per_batch, **predict_kwargs):
+def _call_fc_with_padding(fc_func, batch_series: list[torch.Tensor], **predict_kwargs):
+    if not batch_series:
+        raise ValueError("Received empty batch for forecasting")
+
+    max_len = max(series.shape[0] for series in batch_series)
+    padded_ts = _pad_time_series_batch(batch_series, max_len)
+
+    return fc_func(padded_ts, **predict_kwargs)
+
+
+def _resample_fc_func_wrapper(
+    fc_func,
+    batch,
+    resample_strategy: str,
+    max_context: int = 2016,
+    **predict_kwargs,
+):
+    # downsample the time series based on the dominant frequencies, if enabled
+    max_period = (max_context // 1000) * 500
+    prediction_length = predict_kwargs.get("prediction_length", 100)
+    batch_resampled_ts: list[torch.Tensor] = []
+    fc_resample_fns = []
+    scaling_factors = []
+
+    # select the function doing the resampling
+    ctx_resample_fn = lambda x: (x, 1.0, (lambda y: y))
+    match resample_strategy:
+        case "frequency":
+            ctx_resample_fn = frequency_resample
+        case _:
+            raise RuntimeError("This shouldn't happen.")
+
+    for series in batch:
+        resampled_ts, _sample_factor, fc_resample_fn = ctx_resample_fn(
+            series,
+            prediction_length=prediction_length,
+            max_period=max_period,
+        )
+
+        batch_resampled_ts.append(resampled_ts)
+        fc_resample_fns.append(fc_resample_fn)
+        scaling_factors.append(_sample_factor)
+
+    # Compute per-item required horizons (in downsampled domain)
+    per_item_pred_lens = [int(ceil(prediction_length * sf)) for sf in scaling_factors]
+    max_pred_len = max(per_item_pred_lens) if per_item_pred_lens else int(prediction_length)
+    predict_kwargs.update(prediction_length=max_pred_len)
+
+    max_ts_length = max(ts.shape[0] for ts in batch_resampled_ts)
+    padded_ts = _pad_time_series_batch(batch_resampled_ts, max_ts_length)
+    print(f"Average sample batch factor: {sum(scaling_factors) / len(scaling_factors)}")
+
+    # generate prediction
+    fc_quantiles, fc_mean = fc_func(padded_ts, **predict_kwargs)
+
+    batch_prediction_q = []
+    batch_prediction_m = []
+    for el_q, el_m, fc_resample_fn, item_pred_len in zip(fc_quantiles, fc_mean, fc_resample_fns, per_item_pred_lens):
+        # truncate the forecasts to their individual sample factor adjusted prediction lengths
+        el_q = el_q[:item_pred_len, ...]  # [T, Q]
+        el_m = el_m[:item_pred_len]  # [T]
+
+        # upsample prediction
+        quantiles = fc_resample_fn(el_q.squeeze(0).transpose(0, 1)).transpose(0, 1)  # [T, Q]
+        mean = fc_resample_fn(el_m.squeeze(0))
+
+        quantiles = quantiles[:prediction_length, ...]
+        mean = mean[:prediction_length]
+
+        batch_prediction_q.append(quantiles)
+        batch_prediction_m.append(mean)
+
+    return torch.stack(batch_prediction_q, dim=0), torch.stack(batch_prediction_m, dim=0)
+
+
+def _gen_forecast(
+    fc_func,
+    batches,
+    output_type,
+    quantile_levels,
+    yield_per_batch,
+    resample_strategy: str | None = None,
+    max_context: int = 2016,
+    **predict_kwargs,
+):
+    base_fc_func = fc_func
+
+    if resample_strategy is not None:
+        if resample_strategy not in RESAMPLE_STRATEGIES:
+            raise ValueError(f"Invalid resample strategy: {resample_strategy}. Allowed: {RESAMPLE_STRATEGIES}")
+        fc_func = partial(
+            _resample_fc_func_wrapper,
+            base_fc_func,
+            resample_strategy=resample_strategy,
+            max_context=max_context,
+        )
+    else:
+        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)
 
@@ -92,6 +218,9 @@ def _common_forecast_doc():
                                               forecasts batch by batch as they are computed.
                                               Defaults to `False`.
 
+            resample_strategy (Optional[str], optional): Choose a resampling strategy. Allowed values: {RESAMPLE_STRATEGIES}.
+                                                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.
@@ -113,6 +242,11 @@ class ForecastModel(ABC):
     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,
@@ -120,6 +254,7 @@ class ForecastModel(ABC):
         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"""
@@ -134,7 +269,14 @@ class ForecastModel(ABC):
         assert batch_size >= 1, "Batch size must be >= 1"
         batches = get_batches(context, batch_size)
         return _gen_forecast(
-            self._forecast_quantiles, batches, output_type, quantile_levels, yield_per_batch, **predict_kwargs
+            self._forecast_quantiles,
+            batches,
+            output_type,
+            quantile_levels,
+            yield_per_batch,
+            resample_strategy=resample_strategy,
+            max_context=self.max_context_length,
+            **predict_kwargs,
         )
 
     def forecast_gluon(
@@ -144,6 +286,7 @@ class ForecastModel(ABC):
         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,
     ):
@@ -165,7 +308,14 @@ class ForecastModel(ABC):
 
         batches = get_gluon_batches(gluonDataset, batch_size, **data_kwargs)
         return _gen_forecast(
-            self._forecast_quantiles, batches, output_type, quantile_levels, yield_per_batch, **predict_kwargs
+            self._forecast_quantiles,
+            batches,
+            output_type,
+            quantile_levels,
+            yield_per_batch,
+            resample_strategy=resample_strategy,
+            max_context=self.max_context_length,
+            **predict_kwargs,
         )
 
     def forecast_hfdata(
@@ -175,6 +325,7 @@ class ForecastModel(ABC):
         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,
     ):
@@ -198,5 +349,12 @@ class ForecastModel(ABC):
 
         batches = get_hfdata_batches(hf_dataset, batch_size, **data_kwargs)
         return _gen_forecast(
-            self._forecast_quantiles, batches, output_type, quantile_levels, yield_per_batch, **predict_kwargs
+            self._forecast_quantiles,
+            batches,
+            output_type,
+            quantile_levels,
+            yield_per_batch,
+            resample_strategy=resample_strategy,
+            max_context=self.max_context_length,
+            **predict_kwargs,
         )

tirex/api_adapter/gluon.py

@@ -7,7 +7,7 @@ from gluonts.dataset.common import Dataset
 from gluonts.dataset.field_names import FieldName
 from gluonts.model.forecast import QuantileForecast
 
-from .standard_adapter import _batch_pad_iterable
+from .standard_adapter import _batch_iterable
 
 DEF_TARGET_COLUMN = FieldName.TARGET  # target
 DEF_META_COLUMNS = (FieldName.START, FieldName.ITEM_ID)
@@ -27,7 +27,7 @@ def _get_gluon_ts_map(**gluon_kwargs):
 
 
 def get_gluon_batches(gluonDataset: Dataset, batch_size: int, **gluon_kwargs):
-    return _batch_pad_iterable(map(_get_gluon_ts_map(**gluon_kwargs), gluonDataset), batch_size)
+    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):

tirex/api_adapter/hf_data.py

@@ -4,7 +4,7 @@
 import datasets
 import torch
 
-from .standard_adapter import _batch_pad_iterable
+from .standard_adapter import _batch_iterable
 
 DEF_TARGET_COLUMN = "target"
 
@@ -35,4 +35,4 @@ def _get_hf_map(dataset: datasets.Dataset, **hf_kwargs):
 
 def get_hfdata_batches(hf_dataset: datasets.Dataset, batch_size: int, **hf_kwargs):
     dataset, map_func = _get_hf_map(hf_dataset, **hf_kwargs)
-    return _batch_pad_iterable(map(map_func, dataset), batch_size)
+    return _batch_iterable(map(map_func, dataset), batch_size)

tirex/api_adapter/standard_adapter.py

@@ -16,13 +16,36 @@ ContextType = Union[
 ]
 
 
-def _batched_slice(full_batch, full_meta: list[dict] | None, batch_size: int) -> Iterator[tuple[Sequence, list[dict]]]:
-    if len(full_batch) <= batch_size:
-        yield full_batch, full_meta if full_meta is not None else [{} for _ in range(len(full_batch))]
+def _ensure_1d_tensor(sample) -> torch.Tensor:
+    if isinstance(sample, torch.Tensor):
+        tensor = sample
     else:
-        for i in range(0, len(full_batch), batch_size):
-            batch = full_batch[i : i + batch_size]
-            yield batch, (full_meta[i : i + batch_size] if full_meta is not None else [{} for _ in range(len(batch))])
+        tensor = torch.as_tensor(sample)
+
+    if tensor.ndim > 1:
+        tensor = tensor.squeeze()
+
+    assert tensor.ndim == 1, "Each sample must be one-dimensional"
+    return tensor
+
+
+def _batched_slice(
+    full_batch,
+    full_meta: list[dict] | None,
+    batch_size: int,
+) -> Iterator[tuple[list[torch.Tensor], list[dict]]]:
+    total = len(full_batch)
+    for start in range(0, total, batch_size):
+        batch = full_batch[start : start + batch_size]
+        meta = full_meta[start : start + batch_size] if full_meta is not None else [{} for _ in range(len(batch))]
+
+        batch_series = []
+        for idx in range(len(batch)):
+            sample = batch[idx]
+            tensor = _ensure_1d_tensor(sample)
+            batch_series.append(tensor)
+
+        yield batch_series, meta
 
 
 def _batched(iterable: Iterable, n: int):
@@ -31,21 +54,21 @@ def _batched(iterable: Iterable, n: int):
         yield batch
 
 
-def _batch_pad_iterable(iterable: Iterable[tuple[torch.Tensor, dict]], batch_size: int):
+def _batch_iterable(
+    iterable: Iterable[tuple[torch.Tensor, dict | None]],
+    batch_size: int,
+) -> Iterator[tuple[list[torch.Tensor], list[dict]]]:
     for batch in _batched(iterable, batch_size):
-        # ctx_it_len, ctx_it_data, it_meta = itertools.tee(batch, 3)
-        max_len = max(len(el[0]) for el in batch)
-        padded_batch = []
-        meta = []
-        for el in batch:
-            sample = el[0]
-            assert isinstance(sample, torch.Tensor)
-            assert sample.ndim == 1
-            assert len(sample) > 0, "Each sample needs to have a length > 0"
-            padding = torch.full(size=(max_len - len(sample),), fill_value=torch.nan, device=sample.device)
-            padded_batch.append(torch.cat((padding, sample)))
-            meta.append(el[1])
-        yield torch.stack(padded_batch), meta
+        series_list: list[torch.Tensor] = []
+        meta_list: list[dict] = []
+
+        for sample, meta in batch:
+            tensor = _ensure_1d_tensor(sample)
+            assert len(tensor) > 0, "Each sample needs to have a length > 0"
+            series_list.append(tensor)
+            meta_list.append(meta if meta is not None else {})
+
+        yield series_list, meta_list
 
 
 def get_batches(context: ContextType, batch_size: int):
@@ -59,9 +82,9 @@ def get_batches(context: ContextType, batch_size: int):
         if context.ndim == 1:
             context = np.expand_dims(context, axis=0)
         assert context.ndim == 2
-        batches = map(lambda x: (torch.Tensor(x[0]), x[1]), _batched_slice(context, None, batch_size))
+        batches = _batched_slice(context, None, batch_size)
     elif isinstance(context, (list, Iterable)):
-        batches = _batch_pad_iterable(map(lambda x: (torch.Tensor(x), None), context), batch_size)
+        batches = _batch_iterable(map(lambda x: (torch.Tensor(x), None), context), batch_size)
     if batches is None:
         raise ValueError(f"Context type {type(context)} not supported! Supported Types: {ContextType}")
     return batches

tirex/base.py

@@ -1,6 +1,7 @@
 # Copyright (c) NXAI GmbH.
 # This software may be used and distributed according to the terms of the NXAI Community License Agreement.
 
+import logging
 import os
 from abc import ABC, abstractmethod
 from typing import Literal, TypeVar
@@ -8,7 +9,10 @@ from typing import Literal, TypeVar
 import torch
 from huggingface_hub import hf_hub_download
 
+from tirex.models.slstm.cell import sLSTMCellTorch
+
 T = TypeVar("T", bound="PretrainedModel")
+VERSION_DELIMITER = "-"
 
 
 def skip_cuda():
@@ -29,6 +33,17 @@ def parse_hf_repo_id(path):
     return "/".join(parts[0:2])
 
 
+def parse_model_string(model_string):
+    if VERSION_DELIMITER in model_string:
+        parts = model_string.split(VERSION_DELIMITER)
+        model_id, version = parts[0], parts[0]
+    else:
+        model_id = model_string
+        version = None
+
+    return model_id, version
+
+
 class PretrainedModel(ABC):
     REGISTRY: dict[str, "PretrainedModel"] = {}
 
@@ -38,7 +53,7 @@ class PretrainedModel(ABC):
 
     @classmethod
     def from_pretrained(
-        cls: type[T], path: str, backend: str, device: str | None = None, hf_kwargs=None, ckp_kwargs=None
+        cls: type[T], path: str, backend: str, device: str | None = None, compile=False, hf_kwargs=None, ckp_kwargs=None
     ) -> T:
         if hf_kwargs is None:
             hf_kwargs = {}
@@ -58,9 +73,10 @@ class PretrainedModel(ABC):
         model: T = cls(backend=backend, **checkpoint["hyper_parameters"])
         model.on_load_checkpoint(checkpoint)
         model.load_state_dict(checkpoint["state_dict"])
+        model = model.to(device)
 
-        if backend == "cuda":
-            model = model.to(device)
+        if compile and backend == "torch":
+            sLSTMCellTorch.slstm_forward = torch.compile(sLSTMCellTorch.slstm_forward, mode="max-autotune")
         return model
 
     @classmethod
@@ -76,6 +92,7 @@ def load_model(
     path: str,
     device: str | None = None,
     backend: Literal["torch", "cuda"] | None = None,
+    compile: bool = False,
     hf_kwargs=None,
     ckp_kwargs=None,
 ) -> PretrainedModel:
@@ -85,6 +102,7 @@ def load_model(
         path (str): Hugging Face path to the model (e.g. NX-AI/TiRex)
         device (str, optional): The device on which to load the model (e.g., "cuda:0", "cpu").
         backend (torch | cuda): What backend to use, torch or the custom CUDA kernels. Defaults to cuda when xlstm is installed, else torch.
+        compile (bool, optional): toch.compile the sLSTM cells, only works with the torch backend
         hf_kwargs (dict, optional): Keyword arguments to pass to the Hugging Face Hub download method.
         ckp_kwargs (dict, optional): Keyword arguments to pass when loading the checkpoint.
 
@@ -99,11 +117,14 @@ def load_model(
         backend = "torch" if skip_cuda() or not xlstm_available() else "cuda"
 
     try:
-        _, model_id = parse_hf_repo_id(path).split("/")
+        _, model_string = parse_hf_repo_id(path).split("/")
+        model_id, version = parse_model_string(model_string)
     except:
         raise ValueError(f"Invalid model path {path}")
     model_cls = PretrainedModel.REGISTRY.get(model_id, None)
     if model_cls is None:
         raise ValueError(f"Invalid model id {model_id}")
 
-    return model_cls.from_pretrained(path, device=device, backend=backend, hf_kwargs=hf_kwargs, ckp_kwargs=ckp_kwargs)
+    return model_cls.from_pretrained(
+        path, device=device, backend=backend, compile=compile, hf_kwargs=hf_kwargs, ckp_kwargs=ckp_kwargs
+    )

tirex/models/slstm/cell.py

@@ -43,13 +43,11 @@ class sLSTMCell(nn.Module):
         state = self._get_state(input, state)
 
         if self.backend == "torch":
-            all_states = self._impl_torch(input, state)
+            output, state = self._impl_torch(input, state)
         elif self.backend == "cuda":
-            all_states = self._impl_cuda(input, state)
+            output, state = self._impl_cuda(input, state)
 
-        state = all_states[:, -1]
-        output = self._permute_output(all_states[0][1:])
-        return output.to(input.dtype), state.to(input.dtype)
+        return self._permute_output(output).to(input.dtype), state.to(input.dtype)
 
     def _impl_torch(self, input: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
         input = input.to(dtype=torch.bfloat16)
@@ -64,7 +62,7 @@ class sLSTMCell(nn.Module):
             .reshape(-1)
         )
 
-        return slstm_forward(input, state, recurrent_kernel, bias)[0]
+        return sLSTMCellTorch.slstm_forward(input, state, recurrent_kernel, bias)
 
     def _impl_cuda(self, input: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
         if input.device.type != "cuda":
@@ -88,7 +86,7 @@ class sLSTMCell(nn.Module):
 
         input = input.permute(0, 1, 3, 2, 4).reshape(input.shape[0], input.shape[1], -1)
 
-        return self.func.apply(
+        all_states = self.func.apply(
             False,
             input.contiguous(),
             state.contiguous(),
@@ -96,6 +94,10 @@ class sLSTMCell(nn.Module):
             self._bias_.contiguous(),
         )
 
+        state = all_states[:, -1]
+        output = all_states[0][1:]
+        return output, state
+
     def _get_input(self, x: torch.Tensor) -> torch.Tensor:
         assert x.shape[-1] == self.config.embedding_dim * self.config.num_gates, (
             f"Input size mismatch: Expected input size {self.config.embedding_dim * self.config.num_gates}, but got {input.size(-1)}."
@@ -119,73 +121,60 @@ class sLSTMCell(nn.Module):
         return output.permute(1, 2, 0, 3)
 
 
-def slstm_forward(
-    x: torch.Tensor,  # [S, B, G*I]
-    states: torch.Tensor,  # [4, B, H] only the first is used for recurrence!
-    R: torch.Tensor,  # [K, R*H, H] - K num_heads
-    b: torch.Tensor,  # [T*H]
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    num_states = states.shape[0]
-    sequence_dim = x.shape[0]
-    # this only works for a fully-connected RNN, for a hin change this
-    num_gates_r = R.shape[2] // R.shape[1]
-    hidden_dim = R.shape[1] * R.shape[0]
-    batch_dim = x.shape[1]
-    num_heads = R.shape[0]
-
-    assert batch_dim == states.shape[1]
-    assert hidden_dim == states.shape[2]
-
-    states_all = torch.zeros(
-        [num_states, sequence_dim + 1, batch_dim, hidden_dim],
-        device=x.device,
-        dtype=x.dtype,
-    )
-    states_all[:, 0] = states
-    for i, Wx_t in enumerate(x.unbind(dim=0)):
-        Ry = (
-            states[0]
-            .reshape(batch_dim, num_heads, 1, -1)
-            .matmul(R.unsqueeze(0))
-            .reshape(batch_dim, num_heads, num_gates_r, -1)
-            .transpose(1, 2)
-            .reshape(batch_dim, -1)
-        )
-        sdtype = states.dtype
-        Wx_t, Ry, b, states = Wx_t.float(), Ry.float(), b.float(), states.float()
-        states, gates = slstm_forward_pointwise(Wx_t, Ry, b, states)
-        states = states.to(dtype=sdtype)
-        states_all[:, i + 1] = states
-
-    # shapes ([S, B, H], ([B,H], [B,H], [B,H])
-    return states_all, states
-
-
-def slstm_forward_pointwise(
-    Wx: torch.Tensor,  # dim [B, 4*H]
-    Ry: torch.Tensor,  # dim [B, 4*H]
-    b: torch.Tensor,  # dim [1, 4*H]
-    states: torch.Tensor,  # dim [4, B, H]
-) -> tuple[torch.Tensor, torch.Tensor]:
-    raw = Wx + Ry + b
-
-    iraw, fraw, zraw, oraw = torch.unbind(raw.view(raw.shape[0], 4, -1), dim=1)
-    y, c, n, m = torch.unbind(states.view(4, states.shape[1], -1), dim=0)
-
-    # with torch.no_grad():  # THE difference to maxg aka max_gradient (here max / max_static)
-    # Equations reference the xlstm paper on page 4: https://arxiv.org/pdf/2405.04517
-    logfplusm = m + F.logsigmoid(fraw)  # eq 15
-    if torch.all(n == 0.0):
-        mnew = iraw
-    else:
-        mnew = torch.max(iraw, logfplusm)  # eq 15
-    ogate = torch.sigmoid(oraw)  # eq 14
-    igate = torch.minimum(torch.exp(iraw - mnew), torch.ones_like(iraw))  # eq 16
-    fgate = torch.minimum(torch.exp(logfplusm - mnew), torch.ones_like(iraw))  # eq 17
-    zgate = torch.tanh(zraw)  # eq 11
-    cnew = fgate * c + igate * zgate  # eq 8
-    nnew = fgate * n + igate  # eq 9
-    hnew = ogate * cnew / nnew  # eq 10
-
-    # y (4, B, H), state (4, B, H)
-    return torch.stack((hnew, cnew, nnew, mnew), dim=0), torch.stack((igate, fgate, zraw, ogate), dim=0)
+class sLSTMCellTorch:
+    @staticmethod
+    def slstm_forward(
+        x: torch.Tensor,  # [S, B, G*I]
+        states: torch.Tensor,  # [4, B, H] only the first is used for recurrence!
+        R: torch.Tensor,  # [K, R*H, H] - K num_heads
+        b: torch.Tensor,  # [T*H]
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        num_gates = 4
+        num_heads = R.shape[0]
+        S, B, _ = x.shape
+        H = R.shape[1] * num_heads
+        assert states.shape == (num_gates, B, H)
+
+        states = states.to(R.dtype).unbind(dim=0)
+        output = []
+        for i in range(S):
+            Ry = (
+                states[0]
+                .reshape(B, num_heads, 1, -1)
+                .matmul(R.unsqueeze(0))
+                .reshape(B, num_heads, num_gates, -1)
+                .transpose(1, 2)
+                .reshape(B, -1)
+            )
+            states = sLSTMCellTorch.slstm_forward_pointwise(
+                x[i].float(), Ry.float(), b.float(), [s.float() for s in states]
+            )
+            states = [s.to(dtype=R.dtype) for s in states]
+            output.append(states[0])
+
+        return torch.stack(output), torch.stack(states)  # (S, B, H), 4 x (B, H)
+
+    @staticmethod
+    def slstm_forward_pointwise(
+        Wx: torch.Tensor,  # dim [B, 4*H]
+        Ry: torch.Tensor,  # dim [B, 4*H]
+        b: torch.Tensor,  # dim [1, 4*H]
+        states: torch.Tensor,  # dim 4 x [B, H]
+    ) -> list[torch.Tensor]:
+        y, c, n, m = states
+
+        raw = Wx + Ry + b
+        iraw, fraw, zraw, oraw = torch.unbind(raw.view(raw.shape[0], 4, -1), dim=1)
+
+        # Equations reference the xlstm paper on page 4: https://arxiv.org/pdf/2405.04517
+        logfplusm = m + F.logsigmoid(fraw)  # eq 15
+        mnew = torch.where(torch.all(n == 0.0), iraw, torch.max(iraw, logfplusm))  # eq 15
+        ogate = torch.sigmoid(oraw)  # eq 14
+        igate = torch.minimum(torch.exp(iraw - mnew), torch.ones_like(iraw))  # eq 16
+        fgate = torch.minimum(torch.exp(logfplusm - mnew), torch.ones_like(iraw))  # eq 17
+        zgate = torch.tanh(zraw)  # eq 11
+        cnew = fgate * c + igate * zgate  # eq 8
+        nnew = fgate * n + igate  # eq 9
+        hnew = ogate * cnew / nnew  # eq 10
+
+        return [hnew, cnew, nnew, mnew]  # 4 x (B, H)

tirex/models/tirex.py

@@ -179,8 +179,25 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
         quantile_preds = torch.transpose(quantile_preds, 1, 2)  # switch quantile and num_token_dimension
         # quantile_preds: [batch_size, num_quantiles, num_token, output_patch_size]
 
+        quantile_preds = self._forward_model(torch.cat((input_token, input_mask), dim=2))
+
+        quantile_preds = torch.unflatten(
+            quantile_preds, -1, (len(self.config.quantiles), self.config.output_patch_size)
+        )
+        quantile_preds = torch.transpose(quantile_preds, 1, 2)  # switch quantile and num_token_dimension
+        # quantile_preds: [batch_size, num_quantiles, num_token, output_patch_size]
         return quantile_preds, hidden_states
 
+    def _forward_model(self, input: torch.Tensor):
+        hidden_states = self.input_patch_embedding(input)
+
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+
+        hidden_states = self.out_norm(hidden_states)
+
+        return self.output_patch_embedding(hidden_states)
+
     def _interpolate_quantiles(self, predictions: torch.Tensor, quantile_levels: list[float]):
         training_quantile_levels = self.config.quantiles
         if min(quantile_levels) < min(training_quantile_levels) or max(quantile_levels) > max(training_quantile_levels):

tirex/util.py

@@ -1,7 +1,611 @@
 # Copyright (c) NXAI GmbH.
 # This software may be used and distributed according to the terms of the NXAI Community License Agreement.
 
+from collections.abc import Callable
 from dataclasses import fields
+from functools import partial
+from math import ceil
+from typing import Literal, Optional
+
+import numpy as np
+import torch
+
+
+def frequency_resample(
+    ts: torch.Tensor,
+    prediction_length: int,
+    patch_size: int = 64,
+    peak_prominence: float = 0.1,
+    selection_method: Literal["low_harmonic", "high_harmonic", "highest_amplitude"] = "low_harmonic",
+    min_period: int | None = None,
+    max_period: int = 1000,
+    bandpass_filter: bool = True,
+    nifr_enabled: bool = True,
+    nifr_start_integer: int = 2,
+    nifr_end_integer: int = 12,
+    nifr_clamp_large_factors: bool = False,
+) -> tuple[torch.Tensor, float, Callable[[torch.Tensor], torch.Tensor]]:
+    """
+    Downsample a time series according to a frequency-based strategy and return a helper to upsample
+    forecasts back to the original resolution.
+
+    Parameters
+    ----------
+    ts : torch.Tensor
+        1D time series of shape [T].
+    prediction_length : int
+        Requested forecast horizon; short horizons (<100) skip resampling.
+    patch_size : int, default 64
+        Nominal patch size used to align one dominant period to one patch.
+    peak_prominence : float, default 0.1
+        Threshold for FFT peak detection on the normalized spectrum.
+    selection_method : {"low_harmonic", "high_harmonic", "highest_amplitude"}, default "low_harmonic"
+        How to resolve two dominant peaks in ~2x harmonic relation.
+    min_period : int or None, optional
+        Minimum period to consider; if None, defaults to `patch_size`.
+    max_period : int, default 1000
+        Maximum period to consider for FFT peak search.
+    bandpass_filter : bool, default True
+        If True, suppresses very low frequencies before peak search.
+    nifr_enabled : bool, default True
+        Enable nearest-integer-fraction rounding of the factor.
+    nifr_start_integer : int, default 2
+        Smallest integer k used for 1/k grid when NIFR is enabled.
+    nifr_end_integer : int, default 12
+        Largest integer k used for 1/k grid when NIFR is enabled.
+    nifr_clamp_large_factors : bool, default False
+        If True, clamps large factors in [1, 1/nifr_start_integer] to 1.0.
+
+    Returns
+    -------
+    resampled_ts : torch.Tensor
+        The resampled input series.
+    sample_factor : float
+        Applied sampling factor (<= 1 means downsampling; 1.0 = identity).
+    fc_resample_fn : Callable[[torch.Tensor], torch.Tensor]
+        Function that upsamples a forecast back to the original resolution using the inverse factor.
+
+    Notes
+    -----
+    - For short horizons (prediction_length < 100), resampling is disabled and the factor is set to 1.0.
+    - The factor is clamped to at most 1.0 to avoid upsampling the context.
+    """
+    sample_factor = frequency_factor(
+        ts,
+        max_period=max_period,
+        min_period=min_period,
+        bandpass_filter=bandpass_filter,
+        selection_method=selection_method,
+        peak_prominence=peak_prominence,
+        patch_size=patch_size,
+        nifr_enabled=nifr_enabled,
+        nifr_start_integer=nifr_start_integer,
+        nifr_end_integer=nifr_end_integer,
+        nifr_clamp_large_factors=nifr_clamp_large_factors,
+    )
+
+    sample_factor = min(1, sample_factor)
+
+    if prediction_length < 100:
+        # do not resample for short forecasts
+        sample_factor = 1.0
+
+    fc_resample_factor = 1 / sample_factor
+    fc_resample_fn = partial(resample, sample_rate=fc_resample_factor)
+    resampled_ts = resample(ts, sample_rate=sample_factor)
+
+    return resampled_ts, sample_factor, fc_resample_fn
+
+
+def frequency_factor(
+    ts: torch.Tensor,
+    patch_size: int = 64,  # This doesn't have to match model patch size, but rather the 'target frequency'
+    peak_prominence: float = 0.1,
+    selection_method: Literal["low_harmonic", "high_harmonic", "highest_amplitude"] = "low_harmonic",
+    min_period: int | None = None,
+    max_period: int = 1000,
+    bandpass_filter: bool = True,
+    nifr_enabled: bool = False,
+    nifr_start_integer: int = 2,
+    nifr_end_integer: int = 12,
+    nifr_clamp_large_factors: bool = False,
+) -> float:
+    """
+    Estimate a sampling factor from the dominant frequency of a 1D series so that one period
+    approximately fits into one patch of length `patch_size`.
+
+    The factor is computed as `patch_size / period`, where `period = 1 / f*` and `f*`
+    is the selected dominant frequency from the one-sided FFT of the series (NaNs are
+    linearly interpolated for analysis). If two prominent peaks are detected whose
+    frequencies are in roughly a 2x harmonic relation (ratio in [1.5, 2.5]),
+    `selection_method` determines whether to select the lower or higher harmonic. A set of
+    guards returns 1.0 (identity) for short series, invalid/non-finite results, or when no
+    prominent peak is found. Optional nearest-integer-fraction rounding (NIFR) can snap the
+    factor to the closest value in {1} ∪ {1/k | k ∈ [nifr_start_integer, nifr_end_integer]}.
+
+    Parameters
+    ----------
+    ts : torch.Tensor
+        Input 1D series (last dim is time). NaNs are linearly interpolated for FFT analysis only;
+        the original series is not modified.
+    patch_size : int, default 64
+        Target number of samples per period.
+    peak_prominence : float, default 0.1
+        Minimum normalized spectrum height to treat a bin as a peak.
+    selection_method : {"low_harmonic", "high_harmonic", "highest_amplitude"}, default "low_harmonic"
+        Rule for picking between two ~2x related peaks.
+    min_period : int or None, optional
+        Minimum period to consider. If None, defaults to `patch_size`.
+    max_period : int, default 1000
+        Series shorter than `2 * max_period` return 1.0.
+    bandpass_filter : bool, default True
+        If True, very low frequencies below 1 / max_period are suppressed.
+    nifr_enabled : bool, default False
+        Enable nearest-integer-fraction rounding of the factor.
+    nifr_start_integer : int, default 2
+        Smallest integer k used for the 1/k grid when NIFR is enabled.
+    nifr_end_integer : int, default 12
+        Largest integer k used for the 1/k grid when NIFR is enabled.
+    nifr_clamp_large_factors : bool, default False
+        If True, clamps factors in [1, 1/nifr_start_integer] to 1.0.
+
+    Returns
+    -------
+    float
+        The sampling factor. Values <= 0 or non-finite are mapped to 1.0. If no valid
+        dominant frequency is found or the series is too short, returns 1.0.
+
+    Notes
+    -----
+    - The factor is computed as `patch_size / period`, where `period = 1 / f*` and `f*` is the selected
+      dominant FFT frequency.
+    - If two prominent peaks are detected ~2x apart, `selection_method` determines whether to select the lower or higher harmonic.
+    - Optional nearest-integer-fraction rounding (NIFR) can snap the factor to the closest value in {1} ∪ {1/k | k ∈ [nifr_start_integer, nifr_end_integer]}.
+    """
+    if min_period is None:
+        # NOTE: be careful when min_period is not matching patch_size, it can create unexpected scaling factors!
+        min_period = patch_size
+
+    # Ensure CPU numpy array for FFT analysis
+    ts_np = ts.detach().cpu().numpy() if isinstance(ts, torch.Tensor) else np.asarray(ts)
+
+    # NOTE: If the series is shorter than max_period *2, FFT may not be accurate, to avoid detecting these peaks, we don't scale
+    if ts_np.size < max_period * 2:
+        return 1.0
+
+    freqs, specs, peak_idc = run_fft_analysis(
+        ts_np,
+        scaling="amplitude",
+        peak_prominence=peak_prominence,
+        min_period=min_period,
+        max_period=max_period,
+        bandpass_filter=bandpass_filter,
+    )
+
+    # No detectable peaks -> keep original sampling
+    if peak_idc.size == 0:
+        return 1.0
+
+    # Choose initial candidate as the highest-amplitude peak
+    chosen_idx = int(peak_idc[0])
+
+    # If two peaks exist, check for ~2x harmonic relation and prefer the higher/lower one
+    if peak_idc.size >= 2:
+        idx_a = int(peak_idc[0])  # highest amplitude
+        idx_b = int(peak_idc[1])  # second highest amplitude
+        f_a = float(freqs[idx_a])
+        f_b = float(freqs[idx_b])
+
+        # Determine lower/higher frequency
+        low_f = min(f_a, f_b)
+        high_f = max(f_a, f_b)
+
+        if low_f > 0:
+            ratio = high_f / low_f
+            # Roughly half relation
+            if 1.5 <= ratio <= 2.5:
+                if selection_method == "low_harmonic":
+                    chosen_idx = idx_a if f_a < f_b else idx_b
+                elif selection_method == "high_harmonic":
+                    chosen_idx = idx_a if f_a > f_b else idx_b
+
+    chosen_freq = float(freqs[chosen_idx])
+
+    # Guard against zero or non-finite frequency
+    if not np.isfinite(chosen_freq) or chosen_freq <= 0:
+        return 1.0
+
+    # Convert to period and compute scaling factor so one period fits one patch
+    period = 1.0 / chosen_freq
+    factor = resampling_factor(period, patch_size)
+    factor = round(factor, 4)
+
+    # Guard against factor being negative
+    if not np.isfinite(factor) or factor <= 0:
+        return 1.0
+
+    # nearest interger fraction rounding (nifr)
+    if nifr_enabled:
+        int_fractions = np.concatenate([[1], 1 / np.arange(nifr_start_integer, nifr_end_integer + 1)])
+        diff = np.abs(factor - int_fractions)
+        min_diff_idc = np.argmin(diff)
+        factor = int_fractions[min_diff_idc]
+
+        if nifr_clamp_large_factors:
+            # Clamp everything between 1 and 1/nifr_start_integer to 1, that is no scaling
+            factor = factor if factor < int_fractions[1] else 1
+
+    return float(factor)
+
+
+def resample(ts: torch.Tensor, sample_rate: float, window_position: str = "center") -> torch.Tensor:
+    """
+    Resample the time series using NaN-tolerant window averaging with size 1/sample_rate.
+
+    - If sample_rate > 1 the series is upsampled; windows may collapse to a single index.
+    - If sample_rate < 1 the series is downsampled; windows span multiple indices.
+    - If sample_rate == 1 the series is returned unchanged (cast to float for NaN support).
+
+    Window alignment controlled by `window_position`:
+    - "center": average over [c - L/2, c + L/2]
+    - "left"  : average over [c - L,   c]
+    - "right" : average over [c,       c + L]
+
+    The window is truncated at the boundaries. NaNs are ignored via nan-mean semantics;
+    if a window contains only NaNs, the output is NaN.
+
+    Arguments:
+    ----------
+    ts: torch.Tensor of shape [..., T]
+        The time series to be rescaled (last dim is time).
+    sample_rate: float
+        The factor determining the final number of timesteps in the series, i.e., T' = ceil(T * sample_rate).
+    window_position: {"center", "left", "right"}
+        Placement of the averaging window relative to each target coordinate.
+
+    Returns:
+    --------
+    torch.Tensor of shape [..., ceil(T * sample_rate)] with dtype float.
+    """
+    # Validate inputs
+    if sample_rate <= 0 or sample_rate == 1:
+        # Invalid or no scaling; return original as float
+        return ts.to(torch.float)
+
+    src_num_timesteps = ts.shape[-1]
+    tgt_num_timesteps = ceil(src_num_timesteps * sample_rate)
+
+    # Do not change coordinate creation logic
+    src_coords = torch.arange(src_num_timesteps, device=ts.device)
+    tgt_coords = torch.linspace(0, src_num_timesteps - 1, tgt_num_timesteps, device=ts.device)
+
+    if sample_rate == 1:
+        return ts.to(torch.float)
+
+    # Branch: upsampling -> linear interpolation between nearest neighbors (NaN-aware)
+    if sample_rate > 1:
+        # Neighbour indices for each target coordinate along the last dimension
+        tgt_in_src_idx_lo = tgt_coords.floor().to(torch.long)
+        tgt_in_src_idx_hi = tgt_coords.ceil().to(torch.long)
+
+        # Distances in index space and offsets from lower index
+        dist = src_coords[tgt_in_src_idx_hi] - src_coords[tgt_in_src_idx_lo]
+
+        # Work in float for NaN support; gather neighbour values
+        src_lo_vals = ts[..., tgt_in_src_idx_lo].to(torch.float)
+        src_hi_vals = ts[..., tgt_in_src_idx_hi].to(torch.float)
+        diff = src_hi_vals - src_lo_vals
+        offset = tgt_coords - src_coords[tgt_in_src_idx_lo]
+
+        # Allocate output
+        tgt_values = torch.empty(*ts.shape[:-1], tgt_num_timesteps, dtype=torch.float, device=ts.device)
+
+        # Masks
+        exact_mask = dist == 0
+        interp_mask = ~exact_mask
+
+        # Exact source index -> take the source value
+        if exact_mask.any():
+            tgt_values[..., exact_mask] = src_lo_vals[..., exact_mask]
+
+        # Linear interpolate where indices differ
+        if interp_mask.any():
+            tgt_values[..., interp_mask] = (
+                diff[..., interp_mask] / dist[interp_mask].to(torch.float) * offset[interp_mask]
+                + src_lo_vals[..., interp_mask]
+            )
+
+        # Propagate NaNs from either neighbour
+        nan_mask = torch.isnan(src_lo_vals) | torch.isnan(src_hi_vals)
+        if nan_mask.any():
+            tgt_values[..., nan_mask] = torch.nan
+
+        return tgt_values
+
+    # Window length in source-index units
+    L = 1.0 / sample_rate
+    half_L = 0.5 * L
+
+    if window_position == "center":
+        left_f = tgt_coords - half_L
+        right_f = tgt_coords + half_L
+    elif window_position == "left":
+        left_f = tgt_coords - L
+        right_f = tgt_coords
+    elif window_position == "right":
+        left_f = tgt_coords
+        right_f = tgt_coords + L
+    else:
+        raise ValueError("window_position must be one of {'center','left','right'}")
+
+    # Convert to integer indices, inclusive bounds
+    left_idx = torch.ceil(left_f).to(torch.long)
+    right_idx = torch.floor(right_f).to(torch.long)
+
+    # Clip to valid range and ensure non-empty windows (at least one index)
+    left_idx = torch.clamp(left_idx, 0, src_num_timesteps - 1)
+    right_idx = torch.clamp(right_idx, 0, src_num_timesteps - 1)
+    right_idx = torch.maximum(right_idx, left_idx)
+
+    # Prepare cumulative sums for fast [l, r] segment nan-mean along the last dim
+    ts_float = ts.to(torch.float)
+    valid_mask = ~torch.isnan(ts_float)
+
+    values_filled = torch.where(valid_mask, ts_float, torch.zeros_like(ts_float))
+    counts = valid_mask.to(torch.float)
+
+    cumsum_vals = values_filled.cumsum(dim=-1)
+    cumsum_cnts = counts.cumsum(dim=-1)
+
+    # Pad a leading zero to make inclusive range sums easy: sum[l:r] = cs[r] - cs[l-1]
+    pad_shape = (*ts.shape[:-1], 1)
+    zeros_vals = torch.zeros(pad_shape, dtype=cumsum_vals.dtype, device=ts.device)
+    zeros_cnts = torch.zeros(pad_shape, dtype=cumsum_cnts.dtype, device=ts.device)
+    cumsum_vals = torch.cat([zeros_vals, cumsum_vals], dim=-1)
+    cumsum_cnts = torch.cat([zeros_cnts, cumsum_cnts], dim=-1)
+
+    # Build broadcastable indices for gather along the last dim
+    prefix_shape = ts.shape[:-1]
+    target_len = tgt_num_timesteps
+
+    def _expand_index(idx: torch.Tensor) -> torch.Tensor:
+        # idx shape [target_len] -> [..., target_len]
+        view_shape = (1,) * len(prefix_shape) + (target_len,)
+        return idx.view(view_shape).expand(*prefix_shape, target_len)
+
+    # For inclusive [l, r], use cumsum at (r+1) and (l)
+    r_plus1 = torch.clamp(right_idx + 1, 0, src_num_timesteps)
+    l_idx = left_idx
+
+    r_plus1_exp = _expand_index(r_plus1)
+    l_exp = _expand_index(l_idx)
+
+    seg_sums = cumsum_vals.gather(dim=-1, index=r_plus1_exp) - cumsum_vals.gather(dim=-1, index=l_exp)
+    seg_cnts = cumsum_cnts.gather(dim=-1, index=r_plus1_exp) - cumsum_cnts.gather(dim=-1, index=l_exp)
+
+    # Compute nan-mean: where count==0 -> NaN
+    with torch.no_grad():
+        safe_cnts = torch.where(seg_cnts > 0, seg_cnts, torch.ones_like(seg_cnts))
+    averages = seg_sums / safe_cnts
+    averages = torch.where(seg_cnts > 0, averages, torch.full_like(averages, float("nan")))
+
+    return averages
+
+
+def run_fft_analysis(
+    y,
+    dt: float = 1.0,
+    window: str = "hann",
+    detrend: bool = True,
+    scaling: str = "amplitude",
+    peak_prominence: float = 0.1,
+    min_period: int = 64,
+    max_period: int = 1000,
+    bandpass_filter: bool = True,
+):
+    """
+    Compute one-sided FFT frequencies and spectrum magnitude for a real 1D signal.
+
+    Parameters
+    ----------
+    y : array_like
+        1D time series (regularly sampled). NaNs will be linearly interpolated.
+    dt : float
+        Sampling period (time between samples). Frequencies are cycles per unit of dt.
+    window : {'hann', None}
+        Optional taper to reduce leakage.
+    detrend : bool
+        If True, remove the mean before FFT.
+    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).
+    peak_prominence : float
+        Absolute threshold on the normalized spectrum for peak detection.
+
+    Returns
+    -------
+    f : ndarray
+        Frequencies (non-negative), length N//2 + 1, in cycles per unit (1/dt).
+    spec : ndarray
+        Spectrum corresponding to `f` under the chosen `scaling`.
+    peaks_idx : ndarray
+        Indices into f of detected peaks.
+    """
+    y = np.asarray(y, dtype=float)
+    if y.ndim != 1:
+        y = y.reshape(-1)
+    n = y.size
+    if n < 2:
+        return np.array([]), np.array([]), np.array([])
+
+    # Fill NaNs linearly (handles edge NaNs as well)
+    y = _nan_linear_interpolate(y)
+
+    if detrend:
+        y = y - np.mean(y)
+
+    # Windowing
+    if window == "hann":
+        w = np.hanning(n)
+        yw = y * w
+        # average window power (for proper amplitude/power normalization)
+        w_power = np.sum(w**2) / n
+    elif window is None:
+        yw = y
+        w_power = 1.0
+    else:
+        raise ValueError("window must be either 'hann' or None")
+
+    # FFT (one-sided)
+    Y = np.fft.rfft(yw)
+    f = np.fft.rfftfreq(n, d=dt)  # cycles per unit time
+
+    if scaling == "raw":
+        spec = np.abs(Y)
+    elif scaling == "amplitude":
+        # One-sided amplitude with window power compensation
+        spec = np.abs(Y) / (n * np.sqrt(w_power))
+        if n % 2 == 0:
+            spec[1:-1] *= 2.0
+        else:
+            spec[1:] *= 2.0
+    elif scaling == "power":
+        # One-sided power (not PSD)
+        spec = (np.abs(Y) ** 2) / (n**2 * w_power)
+        if n % 2 == 0:
+            spec[1:-1] *= 2.0
+        else:
+            spec[1:] *= 2.0
+    else:
+        raise ValueError("scaling must be 'amplitude', 'power', or 'raw'")
+
+    # Normalize the spectrum by its maximum value
+    if spec.max() > 0:
+        spec = spec / spec.max()
+
+    # Find peaks in the spectrum
+    peaks_idx = custom_find_peaks(
+        f,
+        spec,
+        max_peaks=2,
+        prominence_threshold=peak_prominence,
+        min_period=min_period,
+        max_period=max_period,
+        bandpass_filter=bandpass_filter,
+    )
+
+    return f, spec, peaks_idx
+
+
+def _nan_linear_interpolate(y: np.ndarray) -> np.ndarray:
+    y = y.astype(np.float32)
+    if y.ndim != 1:
+        y = y.reshape(-1)
+    n = y.size
+    mask = np.isfinite(y)
+    if mask.all():
+        return y
+    if (~mask).all():
+        return np.zeros(n, dtype=np.float32)
+    idx = np.arange(n)
+    y_interp = y.copy()
+    y_interp[~mask] = np.interp(idx[~mask], idx[mask], y[mask])
+    return y_interp
+
+
+def resampling_factor(inverted_freq, path_size):
+    """
+    Compute the resampling factor based on the inverted frequency and path size.
+    """
+    if inverted_freq <= 0:
+        return 1.0
+    factor = path_size / inverted_freq
+    return factor
+
+
+def custom_find_peaks(
+    f,
+    spec,
+    *,
+    max_peaks=5,
+    prominence_threshold=0.1,
+    min_period=64,
+    max_period=1000,
+    bandpass_filter=True,
+):
+    """
+    Finds prominent peaks in a spectrum using a simple custom logic.
+
+    A peak is a local maximum. A peak is considered prominent if its height
+    (on a normalized spectrum) is greater than a given threshold.
+
+    Parameters
+    ----------
+    f : np.ndarray
+        Frequency array (currently unused but kept for API consistency).
+    spec : np.ndarray
+        The normalized spectrum.
+    max_peaks : int
+        The maximum number of peaks to return.
+    prominence_threshold : float
+        The minimum height for a peak to be considered prominent.
+
+    Returns
+    -------
+    np.ndarray
+        An array of indices of the detected peaks in the spectrum. Returns an
+        empty array if no prominent peaks are found.
+    """
+    if len(spec) < 5:  # Need at least 5 points to exclude last two bins
+        return np.array([], dtype=int)
+
+    if bandpass_filter:  # only truly filter low frequencies, high frequencies are dealt with later
+        min_freq = 1 / max_period
+        freq_mask = f >= min_freq
+        spec = spec * freq_mask
+
+    # Find all local maxima, excluding the last two bins
+    local_maxima_indices = []
+    for i in range(1, len(spec) - 2):
+        if spec[i] > spec[i - 1] and spec[i] > spec[i + 1]:
+            local_maxima_indices.append(i)
+
+    if not local_maxima_indices:
+        return np.array([], dtype=int)
+
+    # Filter by prominence (height)
+    prominent_peaks = []
+    for idx in local_maxima_indices:
+        if spec[idx] > prominence_threshold:
+            prominent_peaks.append((idx, spec[idx]))
+
+    # If no peaks are above the threshold, return an empty list
+    if not prominent_peaks:
+        return np.array([], dtype=int)
+
+    # Check for clear peaks below min_period (do lowpass filter)
+    for idx, _ in prominent_peaks:
+        period = 1 / f[idx]
+        if period < min_period:
+            return np.array([], dtype=int)
+
+    # Filter by period
+    period_filtered_peaks = []
+    for idx, prominence in prominent_peaks:
+        period = 1 / f[idx]
+
+        if min_period <= period <= max_period:
+            period_filtered_peaks.append((idx, prominence))
+
+    if not period_filtered_peaks:
+        return np.array([], dtype=int)
+
+    # Sort by height and return the top `max_peaks`
+    period_filtered_peaks.sort(key=lambda x: x[1], reverse=True)
+    peak_indices = np.array([p[0] for p in period_filtered_peaks[:max_peaks]], dtype=int)
+
+    return peak_indices
 
 
 def round_up_to_next_multiple_of(x: int, multiple_of: int) -> int:

{tirex_mirror-2025.10.2.dist-info → tirex_mirror-2025.10.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.10.2
+Version: 2025.10.7
 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.7.dist-info/RECORD

@@ -0,0 +1,21 @@
+tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
+tirex/base.py,sha256=JZuWzezuDenrlDekMaaPLBWzB9mw6vJ_qF2GWjou8ws,4229
+tirex/util.py,sha256=AQ8eM-7IpGeAqegVqflrjPQITzIEKltRcJOBbJmikcA,22758
+tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
+tirex/api_adapter/forecast.py,sha256=CpX9YfiQ1X6nvzODklCPzuKPlYidJDha8FnEU-6zr1Q,13919
+tirex/api_adapter/gluon.py,sha256=CbqL8ZgSRvxWAHK6TWqKdxUrMVWmGskRWOHNF84Lh1U,1819
+tirex/api_adapter/hf_data.py,sha256=TRyys2xKIGZS0Yhq2Eb61lWCMg5CWWn1yRlLIN1mU7o,1369
+tirex/api_adapter/standard_adapter.py,sha256=vdlxNs8mTUtPgK_5WMqYqNdMj8W44igqWsAgtggt_xk,2809
+tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
+tirex/models/patcher.py,sha256=EOXFkHsPkq0nuxRNLAbnrgJtcYq0IMC3YIg_16WArg4,3213
+tirex/models/tirex.py,sha256=JKNuCzTI6B9_yCbcmTf2UFjAQXulLNEmloqtAhKJKjQ,9830
+tirex/models/slstm/block.py,sha256=V91Amgz8WAOOHo4fK1UZxd4Dgbx4-X6kUBS6X4m0tKQ,2006
+tirex/models/slstm/cell.py,sha256=JfCs1aUy9IHuz9RwExhUwiUtbg8WmbEg4upcO7hA5Rg,7229
+tirex/models/slstm/layer.py,sha256=93CAYuG-HmUpF7mBAQ-z1S1u2__W10EW5jPToR57qqM,2747
+tirex_mirror-2025.10.7.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
+tirex_mirror-2025.10.7.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
+tirex_mirror-2025.10.7.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
+tirex_mirror-2025.10.7.dist-info/METADATA,sha256=KdlklRtYjCWVG3GB12RFNi1fPxYO0qmVqlrq6r7mURs,11443
+tirex_mirror-2025.10.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tirex_mirror-2025.10.7.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
+tirex_mirror-2025.10.7.dist-info/RECORD,,

tirex_mirror-2025.10.2.dist-info/RECORD

@@ -1,21 +0,0 @@
-tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
-tirex/base.py,sha256=fwyUTGL103kK5jgK5MoLSIHQcZb4lrox_D9fNbY1W1k,3507
-tirex/util.py,sha256=7DFVBXwGQA4niT9VhYbt8iKMBINJVW4LfwwpggFS0Us,469
-tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
-tirex/api_adapter/forecast.py,sha256=snv0sT1_1WzjkhP1YV-I7CMQmSChl93qFc3b6fwUAS0,8502
-tirex/api_adapter/gluon.py,sha256=faiYyn0kBBVQKbpWqrVoyylxZUrmr-qce66twpguVds,1827
-tirex/api_adapter/hf_data.py,sha256=T1eaxqC3OO9yOzIvw4sr55x6iA2AHKJTZd36rROM4fQ,1377
-tirex/api_adapter/standard_adapter.py,sha256=bI3XGYlWQu5EDyhDZyYqOJMbwi5h1aovPQvfHuWETJk,2618
-tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
-tirex/models/patcher.py,sha256=EOXFkHsPkq0nuxRNLAbnrgJtcYq0IMC3YIg_16WArg4,3213
-tirex/models/tirex.py,sha256=Kglea86t_f3nXXHSjFgssxxrd1Qbwfr1eB_5gKfWYxM,9098
-tirex/models/slstm/block.py,sha256=V91Amgz8WAOOHo4fK1UZxd4Dgbx4-X6kUBS6X4m0tKQ,2006
-tirex/models/slstm/cell.py,sha256=ippaAPKI83j3_1l3pu9ks-iBGO641Elm1W4HsHgVu-c,7601
-tirex/models/slstm/layer.py,sha256=93CAYuG-HmUpF7mBAQ-z1S1u2__W10EW5jPToR57qqM,2747
-tirex_mirror-2025.10.2.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
-tirex_mirror-2025.10.2.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
-tirex_mirror-2025.10.2.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
-tirex_mirror-2025.10.2.dist-info/METADATA,sha256=Aq9VAU0pojVsrxwbfFLCmwmk1Gfl_Z_49G4yB-Z9eLY,11443
-tirex_mirror-2025.10.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tirex_mirror-2025.10.2.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
-tirex_mirror-2025.10.2.dist-info/RECORD,,