tope_v1 / generator.py
tope1129's picture
cascade generator submission: king
7d83938 verified
Raw
History Blame Contribute Delete
51.5 kB
"""cascade-fullctx-spectral-v13 β€” prefetched full-context mixture-of-priors generator.
This is the artifact a cascade miner actually competes with: a subclass of
``cascade.interface.DataGenerator`` that turns a single integer ``seed`` into a
corpus of univariate float series. The subnet holds the model, seeds, and
compute budget byte-identical between the king and every challenger, so the
*only* thing that moves the forecast score is the distribution this file emits.
The competitive lever is therefore **prior diversity + realism**: a corpus that
covers more of the shapes a real forecaster must handle (trend, multi-seasonal,
regime shifts, integrated/near-unit-root dynamics, smooth GP-like curves,
nonlinear/chaotic recurrences, mean-reverting stochastic volatility,
weekly/retail demand, intermittent demand, event recovery, and measurement
artifacts) trains a stronger zero-shot model than the reference generator's
trend+seasonal+AR(1) mix.
Design constraints this file respects (all from the contract in
``cascade.interface``):
* **Determinism is load-bearing.** Every value is drawn from one
``np.random.default_rng(seed)`` in a fixed draw order, so two runs at the same
seed produce byte-identical corpora β€” the property ``cascade verify`` audits
by building the corpus twice and comparing digests.
* **Code-only.** No shipped weights, no network, no clock, no un-seeded RNG.
Imports stay on the dependency allowlist (NumPy/SciPy only) and clear of the
static-guard blocklist.
* **Bounded + finite.** Each series is 1-D ``(L,)`` float64, length in
``[min_length, max_length]``, finite (no NaN/inf). ``_sanitize`` is the last
gate so a numerically unlucky draw can never poison a training run.
Everything is **vectorised per family** (a batched time-axis recurrence, never a
per-series Python loop over time). v13 keeps the v12 prefetch + dynamics-heavy
spine and adds: weekly demand, ARIMA-style integrated increments, RBF/RQ
spectral embedding, calendar-focused seasonality, heteroskedastic residuals,
prefix-causal measurement artifacts, and a fixed-length emit fast path.
"""
from __future__ import annotations
import json
from collections.abc import Iterator
from functools import lru_cache, partial
from pathlib import Path
from queue import Full, Queue
from threading import Event, Thread
import numpy as np
from scipy.signal import lfilter
from cascade.interface import DataGenerator
# Series generated per vectorised batch. Bounds peak memory to O(_CHUNK Β· max_len)
# so streaming feed modes (which request millions of series and stop early) never
# materialise the full corpus. Prefetching holds at most two completed chunks
# (current + queued) while the producer may build the next. The base block is
# 2048 Γ— 4096 Γ— 8 B = 64 MiB per base family block, plus temporary arrays.
# This remains comfortably below the 4 GiB sandbox cap. On the reference local
# A100 environment, 2048 rows generated ~6% more points/s than 1024 while 4096
# regressed slightly, so 2048 is the measured throughput sweet spot.
_CHUNK = 2048
# Multi-cadence seasonal bank. Full 4096-point contexts can identify several
# cycles even at 365/672/730-step periods, unlike short-crop generators.
_SEASONAL_PERIODS = np.array(
[4, 7, 12, 24, 30, 48, 52, 90, 96, 144, 168, 183, 288, 336, 365, 672, 730],
dtype=np.float64,
)
_SEASONAL_PROBS = np.array(
[0.04, 0.12, 0.04, 0.16, 0.05, 0.06, 0.04, 0.03, 0.07, 0.03,
0.13, 0.04, 0.04, 0.06, 0.07, 0.04, 0.05],
dtype=np.float64,
)
_SEASONAL_PROBS /= _SEASONAL_PROBS.sum()
# ── family mixture ──────────────────────────────────────────────────────────
# Names are the process families the corpus mixes over; the default weights are
# a deliberate spread (no single family dominates). Override with
# ``"family_weights": {"chaotic": 0.2, ...}`` in config.json to tune the prior
# without touching code β€” unspecified families keep their default weight.
_FAMILIES: tuple[str, ...] = (
"trend_seasonal_ar", # level + slope + multi-seasonal + AR(1) noise (rich reference)
"regime_shift", # piecewise level/variance regimes with structural breaks
"multiplicative", # positive level Γ— seasonal factor Γ— multiplicative noise
"ar2", # AR(2), stationarity-guaranteed, incl. near-unit-root
"integrated", # ARIMA-style I(1)/near-I(2)/I(2) with structured increments
"threshold_ar", # SETAR β€” regime-switching nonlinear recurrence
"chaotic", # bounded chaotic maps (logistic / sine / tent)
"spectral_gp", # RBF/RQ GP-like paths via 2L circulant embedding
"long_memory", # persistent/anti-persistent power-law spectra
"ou_stochastic_vol", # mean-reverting regimes + clustered/heavy-tailed volatility
"physical_sensors", # bounded/skewed/smooth physical measurement archetypes
"seasonal_counts", # seasonal Poisson/NB web and demand counts with bursts
"intermittent", # zero-inflated / intermittent demand
"pulse_outlier", # sharp/decaying events, outliers, and true flat runs
"weekly_demand", # period-7 retail/load demand with promotions and dips
)
# Dynamics-heavy spine from the v12 A/B (geomean 0.18431 vs 0.19097), with ~8%
# moved into weekly_demand (t_smo-validated sweet spot) by shaving core families.
_DEFAULT_WEIGHTS: dict[str, float] = {
"trend_seasonal_ar": 0.11,
"regime_shift": 0.11,
"multiplicative": 0.07,
"ar2": 0.14,
"integrated": 0.11,
"threshold_ar": 0.07,
"chaotic": 0.035,
"spectral_gp": 0.06,
"long_memory": 0.055,
"ou_stochastic_vol": 0.09,
"physical_sensors": 0.018,
"seasonal_counts": 0.018,
"intermittent": 0.009,
"pulse_outlier": 0.015,
"weekly_demand": 0.08,
}
class Generator(DataGenerator):
"""A mixture-of-priors generator. Submit as ``generator.Generator``."""
def __init__(self, config_dir: str, *, seed: int) -> None:
cfg_path = Path(config_dir) / "config.json"
cfg = json.loads(cfg_path.read_text(encoding="utf-8")) if cfg_path.is_file() else {}
self._cfg = cfg
self._seed = int(seed)
self._min_len = int(cfg.get("min_length", 64))
self._max_len = int(cfg.get("max_length", 4096)) # = [training] context_length (train on full context)
if self._min_len < 1 or self._max_len < self._min_len:
raise ValueError(f"invalid length band [{self._min_len}, {self._max_len}]")
weights = dict(_DEFAULT_WEIGHTS)
for k, v in dict(cfg.get("family_weights", {})).items():
if k in weights:
weights[k] = float(v)
w = np.asarray([weights[f] for f in _FAMILIES], dtype=np.float64)
if not np.all(np.isfinite(w)) or w.min() < 0 or w.sum() <= 0:
raise ValueError("family_weights must be finite, non-negative, and not all zero")
self._weights = w / w.sum()
# v3.9 length-NORMALIZED bimodal trend knobs (trend excursion is length-invariant;
# real trend-strength is ~0.02 and length-invariant, but v2's slope*t grows with L).
self._tr_hi_frac = float(cfg.get("tr_hi_frac", 0.25))
self._tr_exc_lo = float(cfg.get("tr_exc_lo", 0.4))
self._tr_exc_hi = float(cfg.get("tr_exc_hi", 3.0))
self._gr_exc_lo = float(cfg.get("gr_exc_lo", 0.3))
self._gr_exc_hi = float(cfg.get("gr_exc_hi", 2.0))
self._sa_clean_frac = float(cfg.get("sa_clean_frac", 0.4))
self._sa_clean_lo = float(cfg.get("sa_clean_lo", 0.02))
self._sa_clean_hi = float(cfg.get("sa_clean_hi", 0.12))
# Calendar focus + soft overlays (config-tunable without code edits).
self._seasonal_focus = float(np.clip(cfg.get("seasonal_focus", 0.75), 0.0, 1.0))
self._het_noise_frac = float(np.clip(cfg.get("het_noise_frac", 0.35), 0.0, 1.0))
self._ar2_seasonal_frac = float(np.clip(cfg.get("ar2_seasonal_frac", 0.25), 0.0, 1.0))
self._integrated_seasonal_frac = float(
np.clip(cfg.get("integrated_seasonal_frac", 0.30), 0.0, 1.0)
)
self._fixed_len = self._min_len == self._max_len
@property
def name(self) -> str:
return str(self._cfg.get("name", "cascade-fullctx-spectral-v13"))
def generate(self, n_series: int) -> Iterator[np.ndarray]:
# Lazy, chunked generation. This is REQUIRED for the streaming feed
# modes (chain.toml ``corpus_mode = "stream_cpu"``): the trainer calls
# ``generate(n_upper)`` with ``n_upper = token_budget // min_length + 2``
# β€” often millions β€” and stops pulling once the token budget is hit
# (see cascade/trainer/stream.py). Materialising all ``n_series`` up
# front would OOM before the first yield. Generating one CHUNK at a time
# keeps memory at O(CHUNK) and stops early when the consumer stops,
# while a fixed draw order keeps the whole sequence seed-deterministic.
if n_series <= 0:
return
rng = np.random.default_rng(self._seed)
max_len = self._max_len
fixed_len = self._fixed_len
# Blend uniform vs bank probs so calendar cadences get more mass without
# dropping long annual periods.
period_p = (1.0 - self._seasonal_focus) * (
np.ones_like(_SEASONAL_PROBS) / len(_SEASONAL_PROBS)
)
period_p = period_p + self._seasonal_focus * _SEASONAL_PROBS
period_p = period_p / period_p.sum()
# Bind the trend-excursion knobs as explicit builder arguments (no shared
# module state) so the corpus is a pure function of (seed, config).
builders = (
partial(
_trend_seasonal_ar,
hi_frac=self._tr_hi_frac,
exc_lo=self._tr_exc_lo,
exc_hi=self._tr_exc_hi,
clean_frac=self._sa_clean_frac,
clean_lo=self._sa_clean_lo,
clean_hi=self._sa_clean_hi,
het_frac=self._het_noise_frac,
period_p=period_p,
),
partial(_regime_shift, period_p=period_p),
partial(
_multiplicative,
hi_frac=self._tr_hi_frac,
exc_lo=self._gr_exc_lo,
exc_hi=self._gr_exc_hi,
period_p=period_p,
),
partial(_ar2, period_p=period_p, seasonal_frac=self._ar2_seasonal_frac),
partial(
_integrated,
period_p=period_p,
seasonal_frac=self._integrated_seasonal_frac,
),
_threshold_ar,
_chaotic,
_spectral_gp,
_long_memory,
partial(_ou_stochastic_vol, period_p=period_p),
partial(_physical_sensors, period_p=period_p),
partial(_seasonal_counts, period_p=period_p),
_intermittent,
partial(_pulse_outlier, period_p=period_p),
_weekly_demand,
)
# Generate one chunk ahead on a CPU thread while the consumer trains on
# the current chunk. The isolation benchmark measured 21.9% of training
# wall blocked in next(); a one-slot queue overlaps NumPy/SciPy work
# (which releases the GIL) without changing the RNG owner or draw order.
queue: Queue[object] = Queue(maxsize=1)
stop = Event()
done = object()
def put(item: object) -> bool:
while not stop.is_set():
try:
queue.put(item, timeout=0.1)
return True
except Full:
continue
return False
def produce() -> None:
try:
produced = 0
while produced < n_series and not stop.is_set():
# Always draw a FULL _CHUNK (yielding only what's still
# needed), so series i remains a pure function of (seed, i).
lengths = rng.integers(
self._min_len, max_len + 1, size=_CHUNK
)
fam_ids = rng.choice(
len(_FAMILIES), size=_CHUNK, p=self._weights
)
chunk: list[np.ndarray | None] = [None] * _CHUNK
for fam in range(len(_FAMILIES)):
idx = np.nonzero(fam_ids == fam)[0]
if idx.size == 0:
continue
block = builders[fam](rng, int(idx.size), max_len)
# Preserve positivity for count/magnitude families and
# exact integer structure for pure count processes.
preserve_nonnegative = fam in (2, 10, 11, 12, 14)
preserve_integers = fam == 11
block = _sanitize(
_measurement_artifacts(
rng,
block,
preserve_nonnegative=preserve_nonnegative,
preserve_integers=preserve_integers,
# Reverse only families whose laws remain valid
# under time reversal.
allow_reverse=fam in (0, 2, 7, 8, 14),
# Hard sensor bounds turn unbounded walks into
# absorbing flats β€” skip on integrated paths.
allow_range_artifacts=fam != 4,
)
)
for row, series_i in enumerate(idx):
if fixed_len:
chunk[series_i] = np.ascontiguousarray(
block[row], dtype=np.float64
)
else:
length = int(lengths[series_i])
chunk[series_i] = np.ascontiguousarray(
block[row, :length], dtype=np.float64
)
take = min(_CHUNK, n_series - produced)
if not put((chunk, take)):
return
produced += take
except BaseException as exc: # propagate producer failures
put(exc)
finally:
put(done)
producer = Thread(target=produce, name="cascade-generator", daemon=True)
producer.start()
try:
while True:
item = queue.get()
if item is done:
break
if isinstance(item, BaseException):
raise item
chunk, take = item
for arr in chunk[:take]:
# fam_ids partitions [0, _CHUNK); fail loud if that changes.
if arr is None: # pragma: no cover - defensive
raise RuntimeError("internal: unfilled series slot")
yield arr
finally:
stop.set()
producer.join(timeout=1.0)
# ── shared vectorised primitives ────────────────────────────────────────────
def _ar1_batch(innov: np.ndarray, phi: np.ndarray) -> np.ndarray:
"""AR(1) filter applied along the time axis of a (n, L) innovation block.
``x[:, t] = phi * x[:, t-1] + innov[:, t]``. The loop is over time (L
iterations, vectorised across the batch), never over the n series.
"""
n, L = innov.shape
x = np.empty((n, L), dtype=np.float64)
p = phi.reshape(n)
for i in range(n):
x[i] = lfilter([1.0], [1.0, -float(p[i])], innov[i])
return x
def _ar2_batch(innov: np.ndarray, a1: np.ndarray, a2: np.ndarray) -> np.ndarray:
"""AR(2) filter: ``x_t = a1 x_{t-1} + a2 x_{t-2} + e_t`` (batched over n)."""
n, L = innov.shape
x = np.empty((n, L), dtype=np.float64)
for i in range(n):
x[i] = lfilter(
[1.0], [1.0, -float(a1[i]), -float(a2[i])], innov[i]
)
return x
@lru_cache(maxsize=4)
def _seasonal_basis(L: int) -> tuple[np.ndarray, np.ndarray]:
"""Cached unit sine/cosine waves for the fixed cadence bank."""
angle = (
2.0
* np.pi
* np.arange(L, dtype=np.float64)[None, :]
/ _SEASONAL_PERIODS[:, None]
)
return np.sin(angle), np.cos(angle)
@lru_cache(maxsize=4)
def _rfftfreq_cached(L: int) -> np.ndarray:
return np.fft.rfftfreq(L)
def _prefix_standardize(
x: np.ndarray, *, calibration_points: int = 512, center: bool = True
) -> np.ndarray:
"""Location/scale from an initial prefix only (no future leakage)."""
prefix = x[:, : min(x.shape[1], calibration_points)]
if center:
mean = prefix.mean(axis=1, keepdims=True)
x = x - mean
std = prefix.std(axis=1, keepdims=True)
return x / np.where(std < 1e-12, 1.0, std)
def _seasonal(
rng: np.random.Generator,
n: int,
L: int,
k_max: int = 3,
period_p: np.ndarray | None = None,
) -> np.ndarray:
"""Sum of 1..k_max stationary or slowly modulated seasonal components."""
t = np.arange(L, dtype=np.float64)[None, :]
sin_basis, cos_basis = _seasonal_basis(L)
p = _SEASONAL_PROBS if period_p is None else period_p
k = rng.integers(1, k_max + 1, size=n)
out = np.zeros((n, L), dtype=np.float64)
for j in range(k_max):
active = np.nonzero(k > j)[0]
per = rng.choice(_SEASONAL_PERIODS, size=n, p=p)[:, None]
amp = rng.uniform(0.2, 2.0, size=n)[:, None]
phase = rng.uniform(0.0, 2.0 * np.pi, size=n)[:, None]
# Draw parameters for every row to preserve the fixed RNG sequence, but
# evaluate only active rows. Stationary components reuse the cadence
# bank via sin(a+b), avoiding a fresh transcendental pass over nΓ—L.
basis_idx = np.searchsorted(_SEASONAL_PERIODS, per[active, 0])
component = amp[active] * (
sin_basis[basis_idx] * np.cos(phase[active])
+ cos_basis[basis_idx] * np.sin(phase[active])
)
# Real seasonal strength and timing drift. TempoPFN's strongest
# non-SDE ablation was its complex-seasonality prior, so a minority of
# components receive slow amplitude and phase modulation while the
# stationary baseline remains well represented.
modulated = np.nonzero((k > j) & (rng.random(n) < 0.35))[0]
if modulated.size:
# Map global row indices into the active component block.
modulated_local = np.searchsorted(active, modulated)
# Slight period jitter approximates real cycle drift without leaving
# the cadence neighbourhood the bank already covers.
per_mod = per[modulated] * rng.uniform(
0.95, 1.05, size=(modulated.size, 1)
)
modulated_arg = (
2.0 * np.pi * t / per_mod + phase[modulated]
)
m_per = np.clip(
per_mod * rng.uniform(
4.0, 12.0, size=(modulated.size, 1)
),
32.0,
2.0 * L,
)
m_phase = rng.uniform(
0.0, 2.0 * np.pi, size=(modulated.size, 1)
)
slow = np.sin(2.0 * np.pi * t / m_per + m_phase)
amp_mod = 1.0 + rng.uniform(
0.05, 0.45, size=(modulated.size, 1)
) * slow
phase_mod = rng.uniform(
0.05, 0.75, size=(modulated.size, 1)
) * np.sin(2.0 * np.pi * t / (1.7 * m_per) - m_phase)
component[modulated_local] = (
amp[modulated]
* amp_mod
* np.sin(modulated_arg + phase_mod)
)
out[active] += component
return out
def _sparse_jumps(rng: np.random.Generator, n: int, L: int, rate: float, scale) -> np.ndarray:
"""A (n, L) block of mostly-zero values with occasional N(0, scale) jumps.
``cumsum`` over this yields a piecewise-constant level; ``exp(cumsum)`` of a
scaled version yields a piecewise-constant positive multiplier.
"""
mask = rng.random((n, L)) < rate
mask[:, 0] = False
rows, cols = np.nonzero(mask)
jumps = np.zeros((n, L), dtype=np.float64)
if rows.size == 0:
return jumps
# Rates are O(1/L), so draw magnitudes only for actual events rather than
# allocating and filling a second dense nΓ—L normal array.
s = np.asarray(scale, dtype=np.float64)
event_scale = s if s.ndim == 0 else s.reshape(n)[rows]
jumps[rows, cols] = rng.normal(0.0, 1.0, size=rows.size) * event_scale
return jumps
def _measurement_artifacts(
rng: np.random.Generator,
block: np.ndarray,
*,
preserve_nonnegative: bool,
preserve_integers: bool = False,
allow_reverse: bool = True,
allow_range_artifacts: bool = True,
) -> np.ndarray:
"""Apply sparse, cheap real-measurement effects to a generated block.
TempoPFN reports a 5.4% aggregate CRPS gain from its complete augmentation
pipeline, but does not isolate optimal probabilities for Toto2. These rates
are deliberately conservative: most rows remain untouched, and a selected
row receives only plausible reversal/sign, censoring, quantization, or
sample-and-hold behavior. Sensor thresholds are calibrated from a prefix so
history does not depend on unseen future values.
"""
original = np.asarray(block, dtype=np.float64)
out = original.copy()
n, L = out.shape
calibration_len = min(L, 512)
reverse = (rng.random(n) < 0.06) if allow_reverse else np.zeros(n, dtype=bool)
out[reverse] = out[reverse, ::-1]
if not preserve_nonnegative:
invert = rng.random(n) < 0.04
out[invert] *= -1.0
# Sensor saturation / floor effects from the observed prefix only.
for row in np.nonzero(rng.random(n) < 0.06)[0]:
q = float(rng.uniform(0.03, 0.18))
upper = rng.random() < 0.5
if not allow_range_artifacts:
continue
calibration = out[row, :calibration_len]
if upper:
out[row] = np.minimum(out[row], np.quantile(calibration, 1.0 - q))
else:
out[row] = np.maximum(out[row], np.quantile(calibration, q))
quantized = np.nonzero(rng.random(n) < 0.07)[0]
if quantized.size:
levels = rng.integers(16, 257, size=(quantized.size, 1))
if allow_range_artifacts:
x = out[quantized]
calibration = x[:, :calibration_len]
lo = calibration.min(axis=1, keepdims=True)
hi = calibration.max(axis=1, keepdims=True)
step = (hi - lo) / np.maximum(levels - 1, 1)
safe_step = np.where(step < 1e-12, 1.0, step)
clipped = np.clip(x, lo, hi)
out[quantized] = lo + np.rint((clipped - lo) / safe_step) * safe_step
# Zero-order-hold resampling approximates telemetry gathered at a lower
# cadence and forwarded at the nominal cadence.
held = np.nonzero(rng.random(n) < 0.04)[0]
if held.size:
factors = rng.choice([2, 4, 8], size=held.size, p=[0.55, 0.30, 0.15])
for factor in (2, 4, 8):
rows = held[factors == factor]
if rows.size:
out[rows] = np.repeat(
out[rows, ::factor], factor, axis=1
)[:, :L]
# Light TempoPFN-style damping / spike layer (skip integer count rows).
if not preserve_integers:
damped = np.nonzero(rng.random(n) < 0.06)[0]
if damped.size:
t = np.arange(L, dtype=np.float64)[None, :] / max(L - 1, 1)
strength = rng.uniform(0.15, 0.55, size=(damped.size, 1))
out[damped] *= 1.0 - strength * t
spiked = np.nonzero(rng.random(n) < 0.07)[0]
if spiked.size:
impulses = _sparse_jumps(
rng,
spiked.size,
L,
rate=2.5 / L,
scale=rng.uniform(2.0, 6.0, size=spiked.size),
)
out[spiked] += _ar1_batch(
impulses, rng.uniform(0.7, 0.95, size=spiked.size)
)
if preserve_integers:
out = np.maximum(np.rint(out), 0.0)
# Heavy zero inflation plus upper censoring can otherwise collapse a sparse
# row to its baseline. Such a row carries no forecasting signal.
degenerate = out[:, :calibration_len].std(axis=1) < 1e-9
out[degenerate] = original[degenerate]
return out
# ── family builders: each returns a (n, L) float64 block ────────────────────
def _trend_seasonal_ar(
rng: np.random.Generator,
n: int,
L: int,
*,
hi_frac: float = 0.25,
exc_lo: float = 0.4,
exc_hi: float = 3.0,
clean_frac: float = 0.4,
clean_lo: float = 0.02,
clean_hi: float = 0.12,
het_frac: float = 0.35,
period_p: np.ndarray | None = None,
) -> np.ndarray:
t = np.arange(L, dtype=np.float64)[None, :]
level = rng.normal(0.0, 1.0, size=(n, 1))
# v3: bimodal trend. The total trend EXCURSION over the series is drawn directly
# (0..exc across t/(L-1)), so the trend sits ~16x below v2's slope*t β€” v2's linear
# trend was a measured ~16x too strong vs real data at production lengths.
_hi = rng.random((n, 1)) < hi_frac
exc = np.where(_hi, rng.normal(0.0, exc_hi, size=(n, 1)),
rng.normal(0.0, exc_lo, size=(n, 1)))
tn = t / max(L - 1, 1)
series = level + exc * tn + _seasonal(rng, n, L, period_p=period_p)
phi = rng.uniform(0.0, 0.85, size=n)
clean = rng.random((n, 1)) < clean_frac
sigma = np.where(
clean,
rng.uniform(clean_lo, clean_hi, size=(n, 1)),
rng.uniform(0.1, 0.6, size=(n, 1)),
)
innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
# Heteroskedastic residual envelope on a minority of rows β€” real load/sensor
# noise is rarely homoskedastic.
log_vol = np.cumsum(_sparse_jumps(rng, n, L, rate=2.0 / L, scale=0.35), axis=1)
vol = np.exp(np.clip(log_vol, -2.0, 2.0))
het = (rng.random(n) < het_frac)[:, None]
innov = innov * np.where(het, vol, 1.0)
return series + _ar1_batch(innov, phi)
def _regime_shift(
rng: np.random.Generator,
n: int,
L: int,
*,
period_p: np.ndarray | None = None,
) -> np.ndarray:
# Piecewise-constant level via cumsum of sparse jumps, plus a piecewise
# variance regime (occasional volatility multiplier), plus mild seasonality.
level = np.cumsum(_sparse_jumps(rng, n, L, rate=3.0 / L, scale=2.0), axis=1)
log_vol = np.cumsum(_sparse_jumps(rng, n, L, rate=3.0 / L, scale=0.5), axis=1)
vol = np.exp(np.clip(log_vol, -3.0, 3.0)) * rng.uniform(0.1, 0.5, size=(n, 1))
noise = rng.normal(0.0, 1.0, size=(n, L)) * vol
seas = _seasonal(rng, n, L, k_max=2, period_p=period_p) * rng.uniform(
0.0, 1.0, size=(n, 1)
)
# Piecewise-affine drift complements abrupt level jumps. Sparse slope
# changes create ramps and recoveries without the explosive scale of an I(2)
# process, covering TempoPFN's high-impact Step/Sawtooth structures.
slope = rng.normal(0.0, 1.0 / L, size=(n, 1)) + np.cumsum(
_sparse_jumps(rng, n, L, rate=2.0 / L, scale=4.0 / L), axis=1
)
piecewise_trend = np.cumsum(slope, axis=1)
return level + piecewise_trend + seas + noise
def _multiplicative(
rng: np.random.Generator,
n: int,
L: int,
*,
hi_frac: float = 0.25,
exc_lo: float = 0.3,
exc_hi: float = 2.0,
period_p: np.ndarray | None = None,
) -> np.ndarray:
t = np.arange(L, dtype=np.float64)[None, :]
# v3: bimodal log-growth excursion (drawn directly), same rationale as the linear trend.
_hg = rng.random((n, 1)) < hi_frac
gexc = np.where(_hg, rng.normal(0.0, exc_hi, size=(n, 1)),
rng.normal(0.0, exc_lo, size=(n, 1)))
tn = t / max(L - 1, 1)
base_level = np.exp(gexc * tn + rng.normal(0.0, 0.3, size=(n, 1))) # positive, drifting
amp = rng.uniform(0.1, 0.6, size=(n, 1))
seasonal_shape = _seasonal(rng, n, L, k_max=1, period_p=period_p)
seasonal_sd = seasonal_shape.std(axis=1, keepdims=True)
seasonal_shape /= np.where(seasonal_sd < 1e-12, 1.0, seasonal_sd)
seas = 1.0 + amp * seasonal_shape
noise = 1.0 + rng.normal(0.0, 1.0, size=(n, L)) * rng.uniform(0.02, 0.15, size=(n, 1))
scale = rng.uniform(1.0, 50.0, size=(n, 1))
return scale * base_level * np.clip(seas, 0.05, None) * np.clip(noise, 0.05, None)
def _ar2(
rng: np.random.Generator,
n: int,
L: int,
*,
period_p: np.ndarray | None = None,
seasonal_frac: float = 0.25,
) -> np.ndarray:
# Draw partial autocorrelations in (-1, 1) and map to AR(2) coeffs via
# Levinson-Durbin, which guarantees stationarity. Bias p1 high for
# persistent (sometimes near-unit-root) series.
p1 = rng.uniform(0.3, 0.98, size=n)
p2 = rng.uniform(-0.6, 0.6, size=n)
a2 = p2
a1 = p1 * (1.0 - p2)
sigma = rng.uniform(0.2, 0.8, size=(n, 1))
innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
x = _ar2_batch(innov, a1, a2)
# Soft seasonal overlay on a minority β€” pure AR rarely matches weekly load.
seas = _seasonal(rng, n, L, k_max=2, period_p=period_p)
amp = rng.uniform(0.05, 0.5, size=(n, 1))
mask = (rng.random(n) < seasonal_frac)[:, None]
return x + np.where(mask, seas * amp, 0.0)
def _integrated(
rng: np.random.Generator,
n: int,
L: int,
*,
period_p: np.ndarray | None = None,
seasonal_frac: float = 0.30,
) -> np.ndarray:
"""Integrated paths with forecastable differenced dynamics.
Pure iid random walks mostly teach persistence. Real integrated series more
often have autocorrelated increments, recurring changes, or a persistent
local drift. Keep an iid minority, but give most rows ARIMA-like structure
in first differences that a forecaster can identify from context.
"""
branch = rng.random(n)
order2 = branch < 0.10
persistent_velocity = (branch >= 0.10) & (branch < 0.45)
drift = rng.normal(0.0, 0.02, size=(n, 1))
sigma = rng.uniform(0.2, 1.0, size=(n, 1))
raw = rng.normal(0.0, 1.0, size=(n, L)) * sigma
correlated = (rng.random((n, 1)) < 0.70) | persistent_velocity[:, None]
phi = rng.uniform(-0.35, 0.85, size=n)
phi[persistent_velocity] = rng.uniform(
0.97, 0.999, size=int(persistent_velocity.sum())
)
ar_steps = _ar1_batch(
raw * np.sqrt(np.maximum(1.0 - phi[:, None] ** 2, 1e-3)),
phi,
)
steps = np.where(correlated, ar_steps, raw)
seasonal_on = (rng.random((n, 1)) < seasonal_frac)
seasonal_steps = _prefix_standardize(
_seasonal(rng, n, L, k_max=1, period_p=period_p)
)
steps += seasonal_on * seasonal_steps * sigma * rng.uniform(
0.05, 0.35, size=(n, 1)
)
steps += drift
walk = np.cumsum(steps, axis=1)
walk2 = np.cumsum(walk, axis=1)
# I(2) variance grows cubically; divide by L to share scale with I(1).
return np.where(order2[:, None], walk2 / max(L, 1), walk)
def _threshold_ar(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
# SETAR(2): coefficient flips with the sign of the previous value β€” a simple
# nonlinear recurrence that produces asymmetric, regime-switching dynamics.
phi_hi = rng.uniform(0.3, 0.9, size=n)
phi_lo = rng.uniform(-0.9, 0.3, size=n)
const_hi = rng.normal(0.0, 0.3, size=n)
const_lo = rng.normal(0.0, 0.3, size=n)
sigma = rng.uniform(0.2, 0.7, size=(n, 1))
innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
x = np.empty((n, L), dtype=np.float64)
x[:, 0] = innov[:, 0]
for t in range(1, L):
prev = x[:, t - 1]
hi = prev >= 0.0
phi = np.where(hi, phi_hi, phi_lo)
const = np.where(hi, const_hi, const_lo)
x[:, t] = np.clip(const + phi * prev + innov[:, t], -1e6, 1e6)
return x
def _chaotic(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
# Bounded chaotic maps: logistic, sine, and tent. Burn off initial transients
# then emit; light observation noise on a majority of rows.
map_id = rng.integers(0, 3, size=n)
r_log = rng.uniform(3.6, 4.0, size=n)
r_sin = rng.uniform(0.85, 1.0, size=n)
r_tent = rng.uniform(1.2, 1.99, size=n)
r_a = np.where(map_id == 0, r_log, np.where(map_id == 1, r_sin, r_tent))
cur = rng.uniform(0.05, 0.95, size=n)
for _ in range(64):
nxt_log = r_a * cur * (1.0 - cur)
nxt_sin = r_a * np.sin(np.pi * cur)
nxt_tent = np.where(cur < 0.5, r_a * cur, r_a * (1.0 - cur))
cur = np.where(
map_id == 0, nxt_log, np.where(map_id == 1, nxt_sin, nxt_tent)
)
cur = np.clip(cur, 0.0, 1.0)
x = np.empty((n, L), dtype=np.float64)
x[:, 0] = cur
for t in range(1, L):
nxt_log = r_a * cur * (1.0 - cur)
nxt_sin = r_a * np.sin(np.pi * cur)
nxt_tent = np.where(cur < 0.5, r_a * cur, r_a * (1.0 - cur))
cur = np.where(
map_id == 0, nxt_log, np.where(map_id == 1, nxt_sin, nxt_tent)
)
cur = np.clip(cur, 0.0, 1.0)
x[:, t] = cur
x = _prefix_standardize(x)
noisy = rng.random(n) < 0.65
if noisy.any():
x[noisy] += rng.normal(0.0, 0.03, size=(int(noisy.sum()), L))
return x
def _spectral_gp(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
"""Composite RBF/Rational-Quadratic GP paths in O(n L log L).
Chronos KernelSynth uses both kernels. A 2L circulant embedding and retain
of the first L samples avoids making the two endpoints artificial neighbours
(which an L-periodic inverse FFT would do).
"""
embed_len = 2 * L
lag = np.minimum(np.arange(embed_len), embed_len - np.arange(embed_len))[None, :]
lengthscale = np.exp(rng.uniform(np.log(8.0), np.log(256.0), size=(n, 1)))
scaled_lag2 = (lag / lengthscale) ** 2
rbf_cov = np.exp(-0.5 * scaled_lag2)
alpha = np.exp(rng.uniform(np.log(0.1), np.log(10.0), size=(n, 1)))
rq_cov = (1.0 + scaled_lag2 / (2.0 * alpha)) ** (-alpha)
blend = rng.beta(0.7, 0.7, size=(n, 1))
covariance = blend * rbf_cov + (1.0 - blend) * rq_cov
spectrum = np.maximum(np.fft.rfft(covariance, axis=1).real, 0.0)
z = rng.standard_normal(spectrum.shape) + 1j * rng.standard_normal(spectrum.shape)
z[:, 0] = 0.0
x = np.fft.irfft(z * np.sqrt(spectrum), n=embed_len, axis=1)[:, :L]
return _prefix_standardize(x)
def _long_memory(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
"""Fractional power-law paths with both persistent and rough regimes.
Generate on a 2L embedding and keep the first L samples so evaluation
targets after a long context are not near an artificial wrap boundary.
"""
embed_len = 2 * L
f = _rfftfreq_cached(embed_len)
safe_f = np.maximum(f, 1.0 / embed_len)[None, :]
beta = rng.uniform(-0.6, 2.4, size=(n, 1))
amp = safe_f ** (-0.5 * beta)
# Some rows change roughness above a random frequency, giving smooth
# large-scale structure and rough local variation (or the reverse) without
# another FFT. Match amplitudes at the split to avoid a spectral jump.
multiscale = rng.random((n, 1)) < 0.4
split_idx = rng.integers(8, max(9, f.size // 3), size=(n, 1))
split_f = np.maximum(split_idx / embed_len, 1.0 / embed_len)
beta_hi = rng.uniform(-0.6, 2.8, size=(n, 1))
above = np.arange(f.size)[None, :] > split_idx
amp_hi = split_f ** (-0.5 * beta) \
* (safe_f / split_f) ** (-0.5 * beta_hi)
amp = np.where(multiscale & above, amp_hi, amp)
amp[:, 0] = 0.0
z = rng.standard_normal((n, f.size)) + 1j * rng.standard_normal((n, f.size))
x = np.fft.irfft(z * amp, n=embed_len, axis=1)[:, :L]
integrate = rng.random(n) < 0.25
if integrate.any():
x[integrate] = np.cumsum(x[integrate], axis=1)
return _prefix_standardize(x)
def _ou_stochastic_vol(
rng: np.random.Generator,
n: int,
L: int,
*,
period_p: np.ndarray | None = None,
) -> np.ndarray:
"""Regime-switching mean reversion with bounded stochastic volatility.
This is a CPU-cheap discrete Euler/AR analogue of TempoPFN's highest-impact
OU SDE prior. Regime paths, seasonal means, volatility envelopes, and
heavy-tail masks are sampled in whole blocks; only the state recurrence
scans time, vectorised across all rows.
"""
# Toggle between a fast/quiet and a slow/volatile regime. A cumulative XOR
# builds persistent Markov-like paths without a per-row Python loop.
switch_rate = np.exp(rng.uniform(np.log(0.001), np.log(0.15), size=(n, 1)))
switches = rng.random((n, L)) < switch_rate
switches[:, 0] = rng.random(n) < 0.5
regime = np.bitwise_and(np.cumsum(switches, axis=1), 1).astype(np.int8)
# One mean-reversion speed per row lets SciPy execute the recurrence in
# compiled code. Regime paths still switch equilibrium mean and volatility;
# rows span both fast/quiet and slow/persistent reversion rates.
slow = rng.random((n, 1)) < 0.5
phi = np.where(
slow,
rng.uniform(0.995, 0.9995, size=(n, 1)),
rng.uniform(0.90, 0.99, size=(n, 1)),
)
mu0 = rng.normal(-2.0, 1.0, size=(n, 1))
mu1 = rng.normal(2.0, 1.0, size=(n, 1))
mean = np.where(regime == 0, mu0, mu1)
seasonal_on = rng.random((n, 1)) < 0.6
mean += seasonal_on * _seasonal(rng, n, L, k_max=3, period_p=period_p) \
* rng.uniform(0.5, 3.0, size=(n, 1))
sigma0 = rng.lognormal(np.log(0.3), 0.3, size=(n, 1))
sigma1 = rng.lognormal(np.log(1.5), 0.5, size=(n, 1))
base_sigma = np.where(regime == 0, sigma0, sigma1)
# Mean-reverting log-vol (closer to TempoPFN OU vol than sparse-jump cumsum).
log_vol_innov = rng.normal(0.0, 0.15, size=(n, L))
log_vol_phi = rng.uniform(0.90, 0.995, size=n)
log_vol = _ar1_batch(log_vol_innov, log_vol_phi)
log_vol -= log_vol.mean(axis=1, keepdims=True)
vol = base_sigma * np.exp(np.clip(log_vol, -1.5, 1.5))
seasonal_vol = rng.random((n, 1)) < 0.30
if seasonal_vol.any():
seas_vol = _prefix_standardize(
_seasonal(rng, int(seasonal_vol.sum()), L, k_max=1, period_p=period_p),
center=False,
)
vol[seasonal_vol[:, 0]] *= np.exp(
0.25 * np.clip(seas_vol, -2.0, 2.0)
)
eps = rng.standard_normal((n, L))
heavy = np.nonzero(rng.random(n) < 0.35)[0]
if heavy.size:
# Replace only heavy-tailed rows; drawing Student-t noise for every row
# previously discarded 65% of that relatively expensive work.
eps[heavy] = (
rng.standard_t(4.0, size=(heavy.size, L)) / np.sqrt(2.0)
)
shocks = rng.random((n, L)) < (3.0 / L)
shock_rows, shock_cols = np.nonzero(shocks)
# As with sparse jumps, draw shock magnitudes only at the O(n) events.
eps[shock_rows, shock_cols] += rng.normal(
0.0, 5.0, size=shock_rows.size
)
innovation_scale = np.sqrt(np.maximum(1.0 - phi * phi, 1e-6))
drive = (1.0 - phi) * mean + innovation_scale * vol * eps
out = np.empty((n, L), dtype=np.float64)
out[:, 0] = mean[:, 0] + vol[:, 0] * eps[:, 0]
for i in range(n):
p = float(phi[i, 0])
out[i, 1:] = lfilter(
[1.0], [1.0, -p], drive[i, 1:], zi=[p * out[i, 0]]
)[0]
scale = np.exp(rng.uniform(np.log(0.1), np.log(50.0), size=(n, 1)))
shift = rng.uniform(-100.0, 100.0, size=(n, 1))
return out * scale + shift
def _physical_sensors(
rng: np.random.Generator,
n: int,
L: int,
*,
period_p: np.ndarray | None = None,
) -> np.ndarray:
"""Generic physical measurements without matching one private dataset.
Four row-level archetypes cover smooth signed measurements, bounded
percentages, pressure-like wandering levels, and non-negative skewed
magnitudes. All share multi-cadence seasonality, smooth synoptic variation,
and sparse fronts/gusts.
"""
seasonal = _seasonal(rng, n, L, k_max=2, period_p=period_p)
smooth = _spectral_gp(rng, n, L)
fronts = np.cumsum(
_sparse_jumps(rng, n, L, rate=5.0 / L, scale=1.0), axis=1
)
base = (
seasonal * rng.uniform(0.3, 2.0, size=(n, 1))
+ smooth * rng.uniform(0.2, 1.2, size=(n, 1))
+ fronts * rng.uniform(0.2, 1.0, size=(n, 1))
)
kind = rng.integers(0, 4, size=n)
out = base.copy()
bounded = kind == 1
if bounded.any():
gain = rng.uniform(0.8, 3.5, size=(int(bounded.sum()), 1))
midpoint = rng.uniform(-0.8, 0.8, size=(int(bounded.sum()), 1))
out[bounded] = 100.0 / (1.0 + np.exp(-gain * (base[bounded] - midpoint)))
pressure = kind == 2
if pressure.any():
count = int(pressure.sum())
walk = np.cumsum(rng.standard_normal((count, L)), axis=1) / np.sqrt(L)
level = rng.uniform(900.0, 1100.0, size=(count, 1))
out[pressure] = level + rng.uniform(2.0, 15.0, size=(count, 1)) * walk \
+ 2.0 * fronts[pressure] + 0.5 * seasonal[pressure]
magnitude = kind == 3
if magnitude.any():
count = int(magnitude.sum())
gusts = (rng.random((count, L)) < (8.0 / L)) \
* rng.lognormal(0.0, 0.8, size=(count, L))
power = rng.uniform(1.0, 1.6, size=(count, 1))
out[magnitude] = np.abs(base[magnitude]) ** power + gusts
return out
def _seasonal_counts(
rng: np.random.Generator,
n: int,
L: int,
*,
period_p: np.ndarray | None = None,
) -> np.ndarray:
"""Seasonal Poisson/negative-binomial counts with decaying bursts.
This keeps count positivity and discreteness intact while covering
overdispersion, cadence-linked rate variation, slow signed growth, and
release/news-like bursts. Computation remains batched across rows.
"""
t = np.arange(L, dtype=np.float64)[None, :]
p = _SEASONAL_PROBS if period_p is None else period_p
period = rng.choice(_SEASONAL_PERIODS, size=(n, 1), p=p)
phase = rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
amp = rng.uniform(0.15, 0.8, size=(n, 1))
log_rate = amp * np.sin(2.0 * np.pi * t / period + phase)
second = rng.random((n, 1)) < 0.55
log_rate += second * (0.5 * amp) * np.sin(
4.0 * np.pi * t / period + rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
)
# A minority carry explicit calendar interaction: intraday cadence plus
# seven day-specific factors, with a randomized weekend dip or lift.
calendar = rng.random((n, 1)) < 0.35
day_period = rng.choice([24, 48, 96, 144], size=(n, 1))
day_idx = (np.floor_divide(np.arange(L)[None, :], day_period) % 7).astype(np.int64)
day_factors = rng.normal(0.0, 0.12, size=(n, 7))
day_factors[:, 5:] += rng.uniform(-0.8, 0.3, size=(n, 1))
calendar_effect = np.take_along_axis(day_factors, day_idx, axis=1)
log_rate += calendar * calendar_effect
excursion = rng.uniform(-0.5, 0.5, size=(n, 1))
log_rate += excursion * t / max(L - 1, 1)
# Sparse positive impulses filtered by row-specific decay create bursts
# without a Python loop over timesteps.
impulses = (
(rng.random((n, L)) < (2.0 / L))
* rng.uniform(1.0, 10.0, size=(n, L))
)
burst = _ar1_batch(impulses, rng.uniform(0.85, 0.995, size=(n, 1)))
base = np.exp(rng.uniform(np.log(3.0), np.log(3000.0), size=(n, 1)))
lam = base * np.exp(np.clip(log_rate, -5.0, 5.0)) * (1.0 + burst)
np.clip(lam, 0.0, 1.0e7, out=lam)
# A gamma-mixed Poisson is negative-binomial marginally and provides
# realistic overdispersion. Half the rows remain ordinary Poisson.
overdispersed = rng.random((n, 1)) < 0.5
shape = rng.uniform(0.5, 4.0, size=(n, 1))
mixed = lam * rng.gamma(shape, 1.0 / shape, size=(n, L))
return rng.poisson(np.where(overdispersed, mixed, lam)).astype(np.float64)
def _intermittent(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
# Seasonal zero-inflated demand. Occurrence probabilities vary by cadence
# instead of being iid, teaching the model forecastable sparse structure.
t = np.arange(L, dtype=np.float64)[None, :]
base_p = rng.uniform(0.03, 0.35, size=(n, 1))
period = rng.choice([7.0, 12.0, 24.0, 48.0, 168.0], size=(n, 1))
season = rng.uniform(0.2, 1.2, size=(n, 1)) * np.sin(
2.0 * np.pi * t / period + rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
)
logit = np.log(base_p / (1.0 - base_p)) + season
p = 1.0 / (1.0 + np.exp(-logit))
occur = (rng.random((n, L)) < p).astype(np.float64)
magnitude = (
rng.gamma(shape=2.0, scale=1.0, size=(n, L))
* rng.uniform(1.0, 10.0, size=(n, 1))
* np.exp(0.25 * season)
)
baseline = rng.uniform(0.0, 0.5, size=(n, 1))
return baseline + occur * magnitude
def _pulse_outlier(
rng: np.random.Generator,
n: int,
L: int,
*,
period_p: np.ndarray | None = None,
) -> np.ndarray:
# Smooth base with forecastable and iid events, recovery, and held runs.
base = _spectral_gp(rng, n, L) * rng.uniform(0.5, 2.0, size=(n, 1))
base += _seasonal(rng, n, L, k_max=1, period_p=period_p) * rng.uniform(
0.0, 1.0, size=(n, 1)
)
# Mix iid sparse jumps with periodic/jittered event trains so intensity is
# partially predictable from context (TempoPFN spike prior gap).
mode = rng.random(n)
sharp = np.zeros((n, L), dtype=np.float64)
iid = mode < 0.45
if iid.any():
sharp[iid] = _sparse_jumps(
rng,
int(iid.sum()),
L,
rate=3.0 / L,
scale=rng.uniform(3.0, 8.0, size=int(iid.sum())),
)
periodic = ~iid
if periodic.any():
count = int(periodic.sum())
period = rng.choice(
np.array([24.0, 48.0, 168.0, 336.0], dtype=np.float64),
size=(count, 1),
)
phase = rng.uniform(0.0, 1.0, size=(count, 1)) * period
t = np.arange(L, dtype=np.float64)[None, :]
# Events near periodic anchors with small jitter.
dist = np.min(
np.stack(
[
np.abs(((t - phase) % period) - 0.0),
np.abs(((t - phase) % period) - period),
],
axis=0,
),
axis=0,
)
gate = dist <= rng.uniform(0.5, 2.5, size=(count, 1))
gate &= rng.random((count, L)) < 0.55
gate[:, 0] = False
rows, cols = np.nonzero(gate)
if rows.size:
scales = rng.uniform(3.0, 8.0, size=count)[rows]
sharp_periodic = np.zeros((count, L), dtype=np.float64)
sharp_periodic[rows, cols] = (
rng.normal(0.0, 1.0, size=rows.size) * scales
)
sharp[periodic] = sharp_periodic
impulses = _sparse_jumps(
rng, n, L, rate=2.0 / L, scale=rng.uniform(2.0, 7.0, size=n)
)
recovery = _ar1_batch(impulses, rng.uniform(0.75, 0.995, size=n))
series = base + sharp + recovery
# Sparse event loops, not a time-axis scan: typically two starts per row.
starts = rng.random((n, L)) < (2.0 / L)
starts[:, 0] = False
for row in range(n):
for start in np.nonzero(starts[row])[0]:
run = int(rng.integers(3, 65))
end = min(int(start) + run, L)
series[row, start:end] = series[row, start - 1]
return series
def _weekly_demand(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
"""Non-negative period-7 demand with promotions, dips, and count rows.
Adapted from the public demand prior validated around an 8% mixture share:
dedicated weekly structure beats relying only on seasonal_counts weekday
factors for retail/load-like series.
"""
time = np.arange(L, dtype=np.float64)[None, :]
normalized_time = time / max(L - 1, 1)
seasonal_amplitude = rng.uniform(0.03, 0.5, size=(n, 1))
profile = rng.normal(0.0, 1.0, size=(n, 7))
profile -= profile.mean(axis=1, keepdims=True)
has_weekend_dip = rng.random(n) < 0.5
dip_start = rng.integers(0, 7, size=n)
dip_depth = rng.uniform(0.4, 1.6, size=n)
weekend_profile = np.zeros((n, 7), dtype=np.float64)
rows = np.arange(n)
weekend_profile[rows, dip_start] -= dip_depth
weekend_profile[rows, (dip_start + 1) % 7] -= dip_depth
weekend_profile -= weekend_profile.mean(axis=1, keepdims=True)
profile += np.where(has_weekend_dip[:, None], weekend_profile, 0.0)
profile -= profile.mean(axis=1, keepdims=True)
phase = rng.integers(0, 7, size=(n, 1))
weekday_index = (np.arange(L)[None, :] + phase) % 7
weekly_log = seasonal_amplitude * np.take_along_axis(
profile, weekday_index, axis=1
)
excursion = (
rng.normal(0.0, 1.0, size=(n, 1))
* rng.uniform(0.3, 2.5, size=(n, 1))
)
trend = excursion * normalized_time
step_scale = rng.uniform(0.005, 0.05, size=(n, 1))
random_walk = np.clip(
np.cumsum(rng.normal(0.0, 1.0, size=(n, L)) * step_scale, axis=1),
-3.0,
3.0,
)
promotion_mask = rng.random((n, L)) < (
rng.uniform(1.0, 8.0, size=(n, 1)) / L
)
promotions = (
promotion_mask
* np.abs(rng.normal(0.0, 1.0, size=(n, L)))
* rng.uniform(0.5, 2.5, size=(n, 1))
)
echo = np.zeros_like(promotions)
echo[:, 1:] = promotions[:, :-1] * rng.uniform(0.2, 0.6, size=(n, 1))
promotions += echo
holiday_mask = rng.random((n, L)) < (
rng.uniform(0.0, 4.0, size=(n, 1)) / L
)
holiday_dips = (
holiday_mask
* np.abs(rng.normal(0.0, 1.0, size=(n, L)))
* rng.uniform(0.3, 1.5, size=(n, 1))
)
noise = rng.normal(0.0, 1.0, size=(n, L)) * rng.uniform(
0.02, 0.25, size=(n, 1)
)
base = rng.uniform(0.0, 8.0, size=(n, 1))
log_mean = np.clip(
base
+ trend
+ random_walk
+ weekly_log
+ promotions
- holiday_dips
+ noise,
-8.0,
13.0,
)
level = np.exp(log_mean)
is_count = rng.random(n) < 0.35
count_scale = rng.uniform(1.0, 60.0, size=(n, 1)) / np.clip(
level.mean(axis=1, keepdims=True), 1e-9, None
)
counts = rng.poisson(np.clip(level * count_scale, 0.0, 1e6)).astype(
np.float64
)
return np.where(is_count[:, None], counts, level)
# ── final safety gate ───────────────────────────────────────────────────────
def _sanitize(block: np.ndarray) -> np.ndarray:
"""Guarantee the contract: finite float64, no NaN/inf, bounded magnitude.
The trainer's ``check_series`` rejects any non-finite value, which would
fail the whole run β€” so this is the hard backstop after every family
builder. Replaces non-finite values and clips to a generous bound.
"""
x = np.asarray(block, dtype=np.float64)
np.nan_to_num(x, copy=False, nan=0.0, posinf=1e6, neginf=-1e6)
np.clip(x, -1e6, 1e6, out=x)
return x