Skip to content

Demo utilities

Used to build the runnable examples in the Quickstart and Covariates pages.

tirex2.demo.Demo dataclass

A synthetic forecasting scenario for showcasing TiRex-2, built by :meth:create_nonstationary_demo or :meth:create_holidays_demo.

Bundles a target series (split into context and held-out future) with the covariates that explain it, and converts to a :class:~tirex2.model.types.TimeseriesType ready to pass to :meth:~tirex2.api_adapter.forecast.ForecastModel.forecast.

Examples:

>>> from tirex2 import load_model
>>> from tirex2.demo import Demo, plot_demo_forecast
>>> model = load_model("NX-AI/TiRex-2", device="cpu")
>>> demo = Demo.create_nonstationary_demo()
>>> ts_univariate = demo.to_timeseries_type(include_covariates=False)
>>> ts_multivariate = demo.to_timeseries_type(include_covariates=True)
>>> forecasts = model.forecast(
...     timeseries=[ts_univariate, ts_multivariate],
...     prediction_length=demo.horizon,
...     output_type="numpy",
... )
>>> fig = plot_demo_forecast(demo, *forecasts, engine="matplotlib")
Source code in src/tirex2/demo.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
@dataclass
class Demo:
    """A synthetic forecasting scenario for showcasing TiRex-2, built by
    :meth:`create_nonstationary_demo` or :meth:`create_holidays_demo`.

    Bundles a target series (split into context and held-out future) with the
    covariates that explain it, and converts to a :class:`~tirex2.model.types.TimeseriesType`
    ready to pass to :meth:`~tirex2.api_adapter.forecast.ForecastModel.forecast`.

    Examples
    --------
    >>> from tirex2 import load_model
    >>> from tirex2.demo import Demo, plot_demo_forecast
    >>> model = load_model("NX-AI/TiRex-2", device="cpu")
    >>> demo = Demo.create_nonstationary_demo()
    >>> ts_univariate = demo.to_timeseries_type(include_covariates=False)
    >>> ts_multivariate = demo.to_timeseries_type(include_covariates=True)
    >>> forecasts = model.forecast(
    ...     timeseries=[ts_univariate, ts_multivariate],
    ...     prediction_length=demo.horizon,
    ...     output_type="numpy",
    ... )
    >>> fig = plot_demo_forecast(demo, *forecasts, engine="matplotlib")
    """

    title: str
    description: str
    target_context: np.ndarray  # (context_length,)
    target_future: np.ndarray  # (horizon,)
    covariates: list[Covariate]

    @property
    def horizon(self) -> int:
        return len(self.target_future)

    def to_timeseries_type(self, include_covariates=True) -> TimeseriesType:
        # build target data
        target = torch.from_numpy(self.target_context).unsqueeze(0)

        past_covariates = None
        future_covariates = None
        if include_covariates:
            # build past covariates
            past_covariates_list = [c.context for c in self.covariates if c.future is None]
            if len(past_covariates_list) > 0:
                past_covariates = torch.from_numpy(np.stack(past_covariates_list).astype(np.float32))

            # build future covariates
            future_covariates_list = [
                np.concatenate([c.context, c.future]) for c in self.covariates if c.future is not None
            ]
            if len(future_covariates_list) > 0:
                future_covariates = torch.from_numpy(np.stack(future_covariates_list).astype(np.float32))

        return TimeseriesType(
            target=target,
            past_covariates=past_covariates,
            future_covariates=future_covariates,
        )

    def describe(self) -> str:
        print(f"{'-' * 70}")
        print("Demo Dataset")
        print(f"{'-' * 70}")
        print(f"Title:       {self.title}")
        print(f"Description: {self.description}")

    @classmethod
    def create_nonstationary_demo(cls, context_length: int = 540, horizon: int = 42, seed: int = 7) -> Demo:
        """Non-stationary demand: a CONTINUOUS future-known driver sets the wandering
        baseline level, plus a BINARY promotion flag that adds spikes.

        The baseline has no fixed mean (it follows a smoothed random walk + a slow
        swing that turns over inside the horizon), so the target's own history cannot
        say where the level is heading - only the continuous covariate can. The
        binary promotions add sharp spikes at irregular times. The model must use
        BOTH covariates: the continuous one to track the level, the flag for spikes.
        """
        rng = np.random.default_rng(seed)
        n_steps = context_length + horizon
        time_steps = np.arange(n_steps)

        # Continuous non-stationary driver (future-known): smoothed random walk plus a
        # slow swing phased to crest near t=0 and decline through the horizon.
        rw = _smooth(np.cumsum(rng.normal(0, 1, n_steps)), 31)
        rw = (rw - rw.mean()) / (rw.std() + 1e-8)
        swing = np.sin(2 * np.pi * (time_steps - (context_length - 70)) / 250.0)
        level = 120.0 + 40.0 * rw + 22.0 * swing  # wanders, no fixed mean

        weekly_profile = np.array([0.95, 0.92, 0.94, 0.98, 1.08, 1.32, 1.18])
        weekly = 34.0 * (weekly_profile[time_steps % 7] - 1.0)  # within-week deviation

        promo = _events(rng, n_steps, gap_lo=30, gap_hi=52, start=20, width_hi=3)
        promo = _force_events(promo, context_length, horizon, n_min=2, rng=rng)
        spike_abs = 55.0

        series = level + weekly + spike_abs * promo + rng.normal(0, 2.5, n_steps)

        return cls(
            title="TiRex-2 on a non-stationary series (continuous driver + promotion flag)",
            description="A continuous covariate sets the wandering level, while a binary flag adds spikes. The model model needs both covariates for solid forecasts.",
            target_context=series[:context_length].astype(np.float32),
            target_future=series[context_length:].astype(np.float32),
            covariates=[
                Covariate(
                    "demand driver (continuous, known ahead)",
                    level[:context_length],
                    level[context_length:],
                    kind="cont",
                ),
                Covariate(
                    "promotion (0/1, known ahead)",
                    promo[:context_length],
                    promo[context_length:],
                    kind="flag",
                ),
            ],
        )

    @classmethod
    def create_holidays_demo(cls, context_length: int = 540, horizon: int = 42, seed: int = 20) -> Demo:
        """ONE future-known holiday flag -> consistent multiplicative spike."""
        rng = np.random.default_rng(seed)
        base = _weekly_base(context_length, horizon, seed=seed)
        flag = _events(rng, context_length + horizon, gap_lo=26, gap_hi=46, start=8, width_hi=3)
        flag = _force_events(flag, context_length, horizon, n_min=3, rng=rng)
        spike = 0.80  # +80% demand on a holiday
        series = base * (1.0 + spike * flag)
        series = series * (1.0 + 0.012 * rng.normal(size=context_length + horizon))  # low noise

        return cls(
            title="TiRex-2 with a future-known covariate (holiday calendar)",
            description="Daily demand with known holidays. As they are irregular, they are unpredictable from history alone.",
            target_context=series[:context_length].astype(np.float32),
            target_future=series[context_length:].astype(np.float32),
            covariates=[
                Covariate(
                    "holiday flag (0/1, known ahead)",
                    flag[:context_length],
                    flag[context_length:],
                )
            ],
        )

create_nonstationary_demo classmethod 

create_nonstationary_demo(context_length: int = 540, horizon: int = 42, seed: int = 7) -> Demo

Non-stationary demand: a CONTINUOUS future-known driver sets the wandering baseline level, plus a BINARY promotion flag that adds spikes.

The baseline has no fixed mean (it follows a smoothed random walk + a slow swing that turns over inside the horizon), so the target's own history cannot say where the level is heading - only the continuous covariate can. The binary promotions add sharp spikes at irregular times. The model must use BOTH covariates: the continuous one to track the level, the flag for spikes.

Source code in src/tirex2/demo.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
@classmethod
def create_nonstationary_demo(cls, context_length: int = 540, horizon: int = 42, seed: int = 7) -> Demo:
    """Non-stationary demand: a CONTINUOUS future-known driver sets the wandering
    baseline level, plus a BINARY promotion flag that adds spikes.

    The baseline has no fixed mean (it follows a smoothed random walk + a slow
    swing that turns over inside the horizon), so the target's own history cannot
    say where the level is heading - only the continuous covariate can. The
    binary promotions add sharp spikes at irregular times. The model must use
    BOTH covariates: the continuous one to track the level, the flag for spikes.
    """
    rng = np.random.default_rng(seed)
    n_steps = context_length + horizon
    time_steps = np.arange(n_steps)

    # Continuous non-stationary driver (future-known): smoothed random walk plus a
    # slow swing phased to crest near t=0 and decline through the horizon.
    rw = _smooth(np.cumsum(rng.normal(0, 1, n_steps)), 31)
    rw = (rw - rw.mean()) / (rw.std() + 1e-8)
    swing = np.sin(2 * np.pi * (time_steps - (context_length - 70)) / 250.0)
    level = 120.0 + 40.0 * rw + 22.0 * swing  # wanders, no fixed mean

    weekly_profile = np.array([0.95, 0.92, 0.94, 0.98, 1.08, 1.32, 1.18])
    weekly = 34.0 * (weekly_profile[time_steps % 7] - 1.0)  # within-week deviation

    promo = _events(rng, n_steps, gap_lo=30, gap_hi=52, start=20, width_hi=3)
    promo = _force_events(promo, context_length, horizon, n_min=2, rng=rng)
    spike_abs = 55.0

    series = level + weekly + spike_abs * promo + rng.normal(0, 2.5, n_steps)

    return cls(
        title="TiRex-2 on a non-stationary series (continuous driver + promotion flag)",
        description="A continuous covariate sets the wandering level, while a binary flag adds spikes. The model model needs both covariates for solid forecasts.",
        target_context=series[:context_length].astype(np.float32),
        target_future=series[context_length:].astype(np.float32),
        covariates=[
            Covariate(
                "demand driver (continuous, known ahead)",
                level[:context_length],
                level[context_length:],
                kind="cont",
            ),
            Covariate(
                "promotion (0/1, known ahead)",
                promo[:context_length],
                promo[context_length:],
                kind="flag",
            ),
        ],
    )

create_holidays_demo classmethod 

create_holidays_demo(context_length: int = 540, horizon: int = 42, seed: int = 20) -> Demo

ONE future-known holiday flag -> consistent multiplicative spike.

Source code in src/tirex2/demo.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
@classmethod
def create_holidays_demo(cls, context_length: int = 540, horizon: int = 42, seed: int = 20) -> Demo:
    """ONE future-known holiday flag -> consistent multiplicative spike."""
    rng = np.random.default_rng(seed)
    base = _weekly_base(context_length, horizon, seed=seed)
    flag = _events(rng, context_length + horizon, gap_lo=26, gap_hi=46, start=8, width_hi=3)
    flag = _force_events(flag, context_length, horizon, n_min=3, rng=rng)
    spike = 0.80  # +80% demand on a holiday
    series = base * (1.0 + spike * flag)
    series = series * (1.0 + 0.012 * rng.normal(size=context_length + horizon))  # low noise

    return cls(
        title="TiRex-2 with a future-known covariate (holiday calendar)",
        description="Daily demand with known holidays. As they are irregular, they are unpredictable from history alone.",
        target_context=series[:context_length].astype(np.float32),
        target_future=series[context_length:].astype(np.float32),
        covariates=[
            Covariate(
                "holiday flag (0/1, known ahead)",
                flag[:context_length],
                flag[context_length:],
            )
        ],
    )

tirex2.demo.plot_demo_forecast 

plot_demo_forecast(demo: Demo, univariate_forecast: Tensor, multivariate_forecast: Tensor, max_context_to_show: int = 64, engine: str = 'plotly')
Source code in src/tirex2/demo.py 
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
def plot_demo_forecast(
    demo: Demo,
    univariate_forecast: torch.Tensor,
    multivariate_forecast: torch.Tensor,
    max_context_to_show: int = 64,
    engine: str = "plotly",
):
    plot_params = dict(
        demo=demo,
        univariate_forecast=univariate_forecast,
        multivariate_forecast=multivariate_forecast,
        max_context_to_show=max_context_to_show,
    )

    # call appropriate plotting function based on the specified engine
    match engine:
        case "matplotlib":
            return _plot_demo_forecast_matplotlib(**plot_params)
        case "plotly":
            return _plot_demo_forecast_plotly(**plot_params)
        case _:
            raise ValueError(f"Drawing {engine=} not supported.")