| 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 |
|
|
| try: |
| from numba import njit |
| _HAS_NUMBA = True |
| except ImportError: |
| _HAS_NUMBA = False |
|
|
| def njit(*args, **kwargs): |
| def wrap(fn): |
| return fn |
| if args and callable(args[0]) and not kwargs: |
| return args[0] |
| return wrap |
|
|
| _CHUNK = 2048 |
| _STARTUP_CHUNK = 256 |
| _RAMP_CHUNK = 1024 |
| _SEASONAL_PERIODS = np.array([4, 7, 12, 15, 24, 30, 48, 52, 60, 90, 96, 144, 168, 183, 240, 288, 336, 365, 672, 730], dtype=np.float64) |
| _SEASONAL_PROBS = np.array([0.01, 0.23, 0.02, 0.02, 0.07, 0.01, 0.06, 0.01, 0.08, 0.01, 0.14, 0.07, 0.04, 0.01, 0.08, 0.08, 0.02, 0.02, 0.01, 0.01], dtype=np.float64) |
| _SEASONAL_PROBS /= _SEASONAL_PROBS.sum() |
| _SEASONAL_PAIRS = np.array([[15, 60], [60, 240], [24, 168], [48, 336], [96, 672], [7, 365], [12, 52]], dtype=np.float64) |
| _FAMILIES: tuple[str, ...] = ('k00', 'k01', 'k02', 'k03', 'k04', 'k05', 'k06', 'k07', 'k08', 'k09', 'k10', 'k11', 'k12', 'k13', 'k14', 'k15', 'k16', 'k17', 'k18', 'k19', 'k20') |
| _DEFAULT_WEIGHTS: dict[str, float] = {'k00': 0.095, 'k01': 0.095, 'k02': 0.06, 'k03': 0.105, 'k04': 0.095, 'k05': 0.06, 'k06': 0.02, 'k07': 0.07, 'k08': 0.07, 'k09': 0.08, 'k10': 0.08, 'k11': 0.06, 'k12': 0.02, 'k13': 0.02, 'k14': 0.07, 'k15': 0.0, 'k16': 0.0, 'k17': 0.0, 'k18': 0.0, 'k19': 0.0, 'k20': 0.0} |
| _CLEAN: frozenset[str] = frozenset({'k15', 'k18', 'k17', 'k19'}) |
| _NONNEG: frozenset[str] = frozenset({'k02', 'k10', 'k11', 'k12'}) |
| _INT_FAM: frozenset[str] = frozenset({'k11', 'k12'}) |
| _REVERSE_FAM: frozenset[str] = frozenset({'k00', 'k02', 'k07', 'k08'}) |
| _KERNELS_WARMED = False |
|
|
| def _k15(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| jumps = rng.uniform(1.0, 8.0, size=(n, 1)) |
| at = rng.random((n, L)) < jumps / max(L, 1) |
| at[:, 0] = True |
| size = rng.normal(0.0, 1.0, size=(n, L)) * rng.uniform(0.4, 3.0, size=(n, 1)) |
| level = np.cumsum(at * size, axis=1) |
| scale = np.exp(rng.uniform(np.log(1.0), np.log(2000.0), size=(n, 1))) |
| exact = rng.random((n, 1)) < 0.5 |
| sd = np.where(exact, 0.0, rng.uniform(0.002, 0.03, size=(n, 1))) |
| out = (rng.uniform(-2.0, 2.0, size=(n, 1)) + level) * scale |
| out = out + rng.normal(0.0, 1.0, size=(n, L)) * sd * scale |
| integral = rng.random((n, 1)) < 0.4 |
| return np.where(integral, np.rint(out), out) |
|
|
| def _k16(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| drive = _ar1_batch(rng.normal(0.0, 1.0, size=(n, L)), rng.uniform(0.99, 0.9995, size=(n, 1))) |
| drive = (drive - drive.mean(axis=1, keepdims=True)) / np.maximum(drive.std(axis=1, keepdims=True), 1e-09) |
| hot = drive > rng.uniform(0.2, 1.2, size=(n, 1)) |
| lo = rng.uniform(0.02, 0.2, size=(n, 1)) |
| ratio = rng.uniform(4.0, 25.0, size=(n, 1)) |
| sd = np.where(hot, lo * ratio, lo) |
| x = _ar1_batch(rng.normal(0.0, 1.0, size=(n, L)) * sd, rng.uniform(0.9, 0.999, size=(n, 1))) |
| scale = np.exp(rng.uniform(np.log(1.0), np.log(500.0), size=(n, 1))) |
| return x * scale + rng.uniform(-1.0, 1.0, size=(n, 1)) * scale |
|
|
| def _k17(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| t = np.arange(L, dtype=np.float64)[None, :] |
| day = rng.choice(np.array([24.0, 48.0, 96.0, 144.0]), size=(n, 1)) |
| week = day * 7.0 |
| amp_w = rng.uniform(0.4, 1.6, size=(n, 1)) |
| amp_d = rng.uniform(0.3, 1.4, size=(n, 1)) |
| y = amp_w * np.sin(2.0 * np.pi * t / week + rng.uniform(0, 2 * np.pi, (n, 1))) |
| y = y + amp_d * np.sin(2.0 * np.pi * t / day + rng.uniform(0, 2 * np.pi, (n, 1))) |
| y = y + 0.35 * amp_d * np.sin(4.0 * np.pi * t / day + rng.uniform(0, 2 * np.pi, (n, 1))) |
| drift = rng.uniform(-0.3, 0.3, size=(n, 1)) * t / max(L - 1, 1) |
| noise = rng.normal(0.0, 1.0, size=(n, L)) * rng.uniform(0.02, 0.15, size=(n, 1)) |
| base = np.exp(rng.uniform(np.log(5.0), np.log(5000.0), size=(n, 1))) |
| out = base * np.exp(np.clip(y * 0.4 + drift + noise, -6.0, 6.0)) |
| counts = rng.random((n, 1)) < 0.45 |
| return np.where(counts, np.rint(out), out) |
|
|
| def _k18(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| t = np.arange(L, dtype=np.float64)[None, :] |
| k = int(rng.integers(3, 6)) |
| out = np.zeros((n, L), dtype=np.float64) |
| base = rng.uniform(10.0, 400.0, size=(n, 1)) |
| for _ in range(k): |
| period = base * rng.uniform(0.31, 2.7, size=(n, 1)) |
| out += rng.uniform(0.2, 1.0, size=(n, 1)) * np.sin(2.0 * np.pi * t / period + rng.uniform(0, 2 * np.pi, size=(n, 1))) |
| sd = rng.uniform(0.005, 0.05, size=(n, 1)) |
| scale = np.exp(rng.uniform(np.log(1.0), np.log(1000.0), size=(n, 1))) |
| out = out + rng.normal(0.0, 1.0, size=(n, L)) * sd |
| return out * scale + rng.uniform(-1.0, 1.0, size=(n, 1)) * scale |
|
|
| def _k19(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| rate = rng.uniform(1.0, 25.0, size=(n, 1)) / max(L, 1) |
| hits = (rng.random((n, L)) < rate).astype(np.float64) |
| mag = rng.gamma(2.0, 1.0, size=(n, L)) * rng.uniform(1.0, 12.0, size=(n, 1)) |
| decay = rng.uniform(0.9, 0.998, size=(n, 1)) |
| flow = _ar1_batch(hits * mag, decay) |
| baseflow = rng.uniform(0.03, 0.6, size=(n, 1)) |
| scale = np.exp(rng.uniform(np.log(1.0), np.log(800.0), size=(n, 1))) |
| sd = rng.uniform(0.0, 0.02, size=(n, 1)) |
| out = (flow + baseflow) * scale |
| return np.maximum(out * (1.0 + rng.normal(0.0, 1.0, (n, L)) * sd), 0.0) |
|
|
| def _k20(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| cap = rng.integers(4, 80, size=(n, 1)).astype(np.float64) |
| step = rng.normal(0.0, 1.0, size=(n, L)) * rng.uniform(0.01, 0.12, size=(n, 1)) * cap |
| walk = np.cumsum(step, axis=1) + rng.uniform(0.0, 1.0, size=(n, 1)) * cap |
| span = 2.0 * cap |
| folded = cap - np.abs(np.mod(walk, span) - cap) |
| quiet = rng.random((n, 1)) < 0.35 |
| folded = np.where(quiet, folded, folded + rng.normal(0.0, 0.35, size=(n, L))) |
| return np.clip(np.rint(folded), 0.0, cap) |
|
|
| def _validate_parameters(parameters: dict[str, float], label: str) -> None: |
| if not all((np.isfinite(value) for value in parameters.values())): |
| raise ValueError(f'{label} must contain only finite values') |
| probability_names = {'sa_clean_frac', 'integrated_heavy_frac', 'integrated_sv_frac', *(key for key in parameters if key.startswith(('observation.', 'augment.')))} |
| for name in probability_names: |
| if not 0.0 <= parameters[name] <= 1.0: |
| raise ValueError(f'{label}.{name} must be in [0, 1]') |
| for name in ('tr_exc_lo', 'tr_exc_hi', 'gr_exc_lo', 'gr_exc_hi', 'sa_clean_lo', 'sa_clean_hi'): |
| if parameters[name] < 0.0: |
| raise ValueError(f'{label}.{name} must be non-negative') |
| for lo_name, hi_name in (('tr_exc_lo', 'tr_exc_hi'), ('gr_exc_lo', 'gr_exc_hi'), ('sa_clean_lo', 'sa_clean_hi'), ('observation.irregular_hold_prob_lo', 'observation.irregular_hold_prob_hi'), ('observation.shock_prob_lo', 'observation.shock_prob_hi')): |
| if parameters[lo_name] > parameters[hi_name]: |
| raise ValueError(f'{label}.{lo_name} must be <= {hi_name}') |
|
|
| class Generator(DataGenerator): |
|
|
| 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)) |
| 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() |
| curriculum = dict(cfg.get('curriculum', {})) |
| self._curriculum_enabled = bool(curriculum.get('enabled', False)) |
| self._expected_budget_fraction = float(curriculum.get('expected_budget_fraction', 1.0)) |
| if not np.isfinite(self._expected_budget_fraction) or not 0.0 < self._expected_budget_fraction <= 1.0: |
| raise ValueError('curriculum.expected_budget_fraction must be in (0, 1]') |
| self._curriculum_start = float(curriculum.get('start_fraction', 0.1)) |
| self._curriculum_end = float(curriculum.get('end_fraction', 0.7)) |
| if not 0.0 <= self._curriculum_start < self._curriculum_end <= 1.0: |
| raise ValueError('curriculum fractions must satisfy 0 <= start_fraction < end_fraction <= 1') |
| start_weights = dict(weights) |
| for k, v in dict(curriculum.get('start_family_weights', {})).items(): |
| if k in start_weights: |
| start_weights[k] = float(v) |
| start_w = np.asarray([start_weights[f] for f in _FAMILIES], dtype=np.float64) |
| if not np.all(np.isfinite(start_w)) or start_w.min() < 0 or start_w.sum() <= 0: |
| raise ValueError('curriculum.start_family_weights must be finite, non-negative, and not all zero') |
| self._start_weights = start_w / start_w.sum() |
| self._tr_hi_frac = float(cfg.get('tr_hi_frac', 0.25)) |
| self._prefetch_depth = int(cfg.get('prefetch_depth', 2)) |
| if not 1 <= self._prefetch_depth <= 4: |
| raise ValueError('prefetch_depth must be in [1, 4]') |
| augment = dict(cfg.get('augment', {})) |
| observation = dict(cfg.get('observation', {})) |
| self._parameters = {'tr_exc_lo': float(cfg.get('tr_exc_lo', 0.4)), 'tr_exc_hi': float(cfg.get('tr_exc_hi', 3.0)), 'gr_exc_lo': float(cfg.get('gr_exc_lo', 0.3)), 'gr_exc_hi': float(cfg.get('gr_exc_hi', 2.0)), 'sa_clean_frac': float(cfg.get('sa_clean_frac', 0.4)), 'sa_clean_lo': float(cfg.get('sa_clean_lo', 0.02)), 'sa_clean_hi': float(cfg.get('sa_clean_hi', 0.12)), 'integrated_heavy_frac': float(cfg.get('integrated_heavy_frac', 0.25)), 'integrated_sv_frac': float(cfg.get('integrated_sv_frac', 0.3)), 'observation.censor_rate': float(observation.get('censor_rate', 0.06)), 'observation.quantize_rate': float(observation.get('quantize_rate', 0.07)), 'observation.regular_hold_rate': float(observation.get('regular_hold_rate', 0.04)), 'observation.irregular_hold_rate': float(observation.get('irregular_hold_rate', 0.0)), 'observation.irregular_hold_prob_lo': float(observation.get('irregular_hold_prob_lo', 0.01)), 'observation.irregular_hold_prob_hi': float(observation.get('irregular_hold_prob_hi', 0.1)), 'observation.shock_row_rate': float(observation.get('shock_row_rate', 0.0)), 'observation.shock_prob_lo': float(observation.get('shock_prob_lo', 0.001)), 'observation.shock_prob_hi': float(observation.get('shock_prob_hi', 0.015)), 'augment.tsmixup': float(augment.get('tsmixup', 0.0)), 'augment.pad_prefix': float(augment.get('pad_prefix', 0.0))} |
| start_parameters = dict(curriculum.get('start_parameters', {})) |
| start_observation = dict(start_parameters.pop('observation', {})) |
| start_augment = dict(start_parameters.pop('augment', {})) |
| known_top_level = {key for key in self._parameters if '.' not in key} |
| unknown = set(start_parameters) - known_top_level |
| unknown.update((f'observation.{key}' for key in start_observation if f'observation.{key}' not in self._parameters)) |
| unknown.update((f'augment.{key}' for key in start_augment if f'augment.{key}' not in self._parameters)) |
| if unknown: |
| names = ', '.join(sorted(unknown)) |
| raise ValueError(f'unknown curriculum.start_parameters: {names}') |
| self._start_parameters = dict(self._parameters) |
| for key, value in start_parameters.items(): |
| self._start_parameters[key] = float(value) |
| for key, value in start_observation.items(): |
| self._start_parameters[f'observation.{key}'] = float(value) |
| for key, value in start_augment.items(): |
| self._start_parameters[f'augment.{key}'] = float(value) |
| _validate_parameters(self._parameters, 'final parameters') |
| _validate_parameters(self._start_parameters, 'curriculum.start_parameters') |
|
|
| @property |
| def name(self) -> str: |
| return str(self._cfg.get('name', '')) |
|
|
| def _blend_at(self, token_progress: float) -> float: |
| if not self._curriculum_enabled: |
| return 1.0 |
| position = (token_progress - self._curriculum_start) / (self._curriculum_end - self._curriculum_start) |
| position = float(np.clip(position, 0.0, 1.0)) |
| return position * position * (3.0 - 2.0 * position) |
|
|
| def _weights_at(self, token_progress: float) -> np.ndarray: |
| blend = self._blend_at(token_progress) |
| if blend >= 1.0: |
| return self._weights |
| if blend <= 0.0: |
| return self._start_weights |
| return (1.0 - blend) * self._start_weights + blend * self._weights |
|
|
| def _parameters_at(self, token_progress: float) -> dict[str, float]: |
| blend = self._blend_at(token_progress) |
| if blend >= 1.0: |
| return dict(self._parameters) |
| if blend <= 0.0: |
| return dict(self._start_parameters) |
| return {key: (1.0 - blend) * self._start_parameters[key] + blend * final_value for key, final_value in self._parameters.items()} |
|
|
| def _progress_at(self, emitted_points: float, target_points: int) -> float: |
| return emitted_points / (target_points * self._expected_budget_fraction) |
|
|
| def generate(self, n_series: int) -> Iterator[np.ndarray]: |
| if n_series <= 0: |
| return |
| _warmup_kernels() |
| rng = np.random.default_rng(self._seed) |
| max_len = self._max_len |
| fixed_len = self._min_len == max_len |
| target_points = max(1, max(n_series - 2, 1) * self._min_len) |
| queue: Queue[object] = Queue(maxsize=self._prefetch_depth) |
| 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 |
| emitted_points = 0 |
| while produced < n_series and (not stop.is_set()): |
| if produced == 0: |
| batch_size = _STARTUP_CHUNK |
| elif produced == _STARTUP_CHUNK: |
| batch_size = _RAMP_CHUNK |
| else: |
| batch_size = _CHUNK |
| lengths = rng.integers(self._min_len, max_len + 1, size=batch_size) |
| take = min(batch_size, n_series - produced) |
| chunk_points = int(lengths[:take].sum()) |
| midpoint_progress = self._progress_at(emitted_points + 0.5 * chunk_points, target_points) |
| family_weights = self._weights_at(midpoint_progress) |
| parameters = self._parameters_at(midpoint_progress) |
| builders = (partial(_k00, hi_frac=self._tr_hi_frac, exc_lo=parameters['tr_exc_lo'], exc_hi=parameters['tr_exc_hi'], clean_frac=parameters['sa_clean_frac'], clean_lo=parameters['sa_clean_lo'], clean_hi=parameters['sa_clean_hi']), _k01, partial(_k02, hi_frac=self._tr_hi_frac, exc_lo=parameters['gr_exc_lo'], exc_hi=parameters['gr_exc_hi']), _k03, partial(_k04, heavy_frac=parameters['integrated_heavy_frac'], sv_frac=parameters['integrated_sv_frac']), _k05, _k06, _k07, _k08, _k09, _k10, _k11, _k12, _k13, _k14, _k15, _k16, _k17, _k18, _k19, _k20) |
| current_observation = {key.removeprefix('observation.'): value for key, value in parameters.items() if key.startswith('observation.')} |
| fam_ids = rng.choice(len(_FAMILIES), size=batch_size, p=family_weights) |
| packed = np.empty((batch_size, max_len), dtype=np.float64) if fixed_len else None |
| chunk: list[np.ndarray | None] | None = None if fixed_len else [None] * batch_size |
| 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) |
| family = _FAMILIES[fam] |
| if family in _CLEAN: |
| block = _sanitize(block) |
| else: |
| block = _sanitize(_measurement_artifacts(rng, block, preserve_nonnegative=family in _NONNEG, preserve_integers=family in _INT_FAM, allow_reverse=family in _REVERSE_FAM, allow_range_artifacts=family != 'k04', **current_observation)) |
| if fixed_len: |
| packed[idx] = block |
| else: |
| for row, series_i in enumerate(idx): |
| length = int(lengths[series_i]) |
| chunk[series_i] = np.ascontiguousarray(block[row, :length], dtype=np.float64) |
| if fixed_len: |
| mix_rate = parameters['augment.tsmixup'] |
| if mix_rate > 0.0: |
| mixed = np.nonzero(rng.random(batch_size) < mix_rate)[0] |
| for series_i in mixed: |
| n_other = int(rng.integers(1, 3)) |
| others = rng.integers(0, batch_size, size=n_other) |
| weights = rng.dirichlet(np.ones(n_other + 1)) |
| combined = weights[0] * packed[series_i] |
| for j, other_i in enumerate(others): |
| combined = combined + weights[j + 1] * packed[int(other_i)] |
| packed[series_i] = _sanitize(combined) |
| pad_rate = parameters['augment.pad_prefix'] |
| if pad_rate > 0.0: |
| padded = np.nonzero(rng.random(batch_size) < pad_rate)[0] |
| for series_i in padded: |
| series = packed[series_i] |
| if series.size < 8: |
| continue |
| cut = int(rng.integers(series.size // 8, 3 * series.size // 4)) |
| series[:cut] = series[cut] |
| if not put(('packed', packed, take)): |
| return |
| else: |
| pad_rate = parameters['augment.pad_prefix'] |
| if pad_rate > 0.0: |
| padded = np.nonzero(rng.random(batch_size) < pad_rate)[0] |
| for series_i in padded: |
| series = chunk[series_i] |
| if series is None or series.size < 8: |
| continue |
| cut = int(rng.integers(series.size // 8, 3 * series.size // 4)) |
| series[:cut] = series[cut] |
| if not put(('list', chunk, take)): |
| return |
| produced += take |
| emitted_points += chunk_points |
| except BaseException as exc: |
| put(exc) |
| finally: |
| put(done) |
| producer = Thread(target=produce, name='', daemon=True) |
| producer.start() |
| try: |
| while True: |
| item = queue.get() |
| if item is done: |
| break |
| if isinstance(item, BaseException): |
| raise item |
| kind, payload, take = item |
| if kind == 'packed': |
| for i in range(take): |
| yield np.ascontiguousarray(payload[i], dtype=np.float64) |
| else: |
| for arr in payload[:take]: |
| if arr is None: |
| raise RuntimeError('internal: unfilled series slot') |
| yield arr |
| finally: |
| stop.set() |
| producer.join(timeout=1.0) |
|
|
| @njit(cache=False) |
| def _ar1_numba(innov: np.ndarray, phi: np.ndarray) -> np.ndarray: |
| n, L = innov.shape |
| x = np.empty((n, L), dtype=np.float64) |
| for i in range(n): |
| p = phi[i] |
| prev = innov[i, 0] |
| x[i, 0] = prev |
| for t in range(1, L): |
| prev = p * prev + innov[i, t] |
| x[i, t] = prev |
| return x |
|
|
| @njit(cache=False) |
| def _ar2_numba(innov: np.ndarray, a1: np.ndarray, a2: np.ndarray) -> np.ndarray: |
| n, L = innov.shape |
| x = np.empty((n, L), dtype=np.float64) |
| for i in range(n): |
| aa1 = a1[i] |
| aa2 = a2[i] |
| x0 = innov[i, 0] |
| x[i, 0] = x0 |
| if L > 1: |
| x1 = aa1 * x0 + innov[i, 1] |
| x[i, 1] = x1 |
| prev2 = x0 |
| prev1 = x1 |
| for t in range(2, L): |
| cur = aa1 * prev1 + aa2 * prev2 + innov[i, t] |
| x[i, t] = cur |
| prev2 = prev1 |
| prev1 = cur |
| return x |
|
|
| @njit(cache=False) |
| def _ar1_zi_numba(drive: np.ndarray, phi: np.ndarray, x0: np.ndarray) -> np.ndarray: |
| n, L = drive.shape |
| x = np.empty((n, L), dtype=np.float64) |
| for i in range(n): |
| p = phi[i] |
| prev = x0[i] |
| x[i, 0] = prev |
| for t in range(1, L): |
| prev = p * prev + drive[i, t] |
| x[i, t] = prev |
| return x |
|
|
| def _warmup_kernels() -> None: |
| global _KERNELS_WARMED |
| if _KERNELS_WARMED or not _HAS_NUMBA: |
| return |
| z = np.zeros((2, 8), dtype=np.float64) |
| p = np.array([0.5, 0.4], dtype=np.float64) |
| _ar1_numba(z, p) |
| _ar2_numba(z, p, p) |
| _ar1_zi_numba(z, p, p) |
| _KERNELS_WARMED = True |
|
|
| def _ar1_batch(innov: np.ndarray, phi: np.ndarray) -> np.ndarray: |
| n, L = innov.shape |
| if n == 0: |
| return np.empty((0, L), dtype=np.float64) |
| p = np.asarray(phi, dtype=np.float64).reshape(n) |
| if _HAS_NUMBA: |
| return _ar1_numba(np.ascontiguousarray(innov, dtype=np.float64), np.ascontiguousarray(p)) |
| x = np.empty((n, L), dtype=np.float64) |
| 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: |
| n, L = innov.shape |
| if n == 0: |
| return np.empty((0, L), dtype=np.float64) |
| aa1 = np.asarray(a1, dtype=np.float64).reshape(n) |
| aa2 = np.asarray(a2, dtype=np.float64).reshape(n) |
| if _HAS_NUMBA: |
| return _ar2_numba(np.ascontiguousarray(innov, dtype=np.float64), np.ascontiguousarray(aa1), np.ascontiguousarray(aa2)) |
| x = np.empty((n, L), dtype=np.float64) |
| for i in range(n): |
| x[i] = lfilter([1.0], [1.0, -float(aa1[i]), -float(aa2[i])], innov[i]) |
| return x |
|
|
| def _ar1_zi_batch(drive: np.ndarray, phi: np.ndarray, x0: np.ndarray) -> np.ndarray: |
| n, L = drive.shape |
| if n == 0: |
| return np.empty((0, L), dtype=np.float64) |
| p = np.asarray(phi, dtype=np.float64) |
| x0v = np.asarray(x0, dtype=np.float64) |
| if p.ndim > 1: |
| p = p[:, 0] |
| if x0v.ndim > 1: |
| x0v = x0v[:, 0] |
| p = np.ascontiguousarray(p.reshape(n)) |
| x0v = np.ascontiguousarray(x0v.reshape(n)) |
| if _HAS_NUMBA: |
| return _ar1_zi_numba(np.ascontiguousarray(drive, dtype=np.float64), p, x0v) |
| x = np.empty((n, L), dtype=np.float64) |
| x[:, 0] = x0v |
| for i in range(n): |
| x[i, 1:] = lfilter([1.0], [1.0, -float(p[i])], drive[i, 1:], zi=[p[i] * x0v[i]])[0] |
| return x |
|
|
| def _prefix_mean_std(x: np.ndarray, *, calibration_points: int=512) -> tuple[np.ndarray, np.ndarray]: |
| prefix = x[:, :min(x.shape[1], calibration_points)] |
| mean = prefix.mean(axis=1, keepdims=True) |
| std = prefix.std(axis=1, keepdims=True) |
| return (mean, np.where(std < 1e-12, 1.0, std)) |
|
|
| def _prefix_standardize(x: np.ndarray, *, center: bool=True, calibration_points: int=512) -> np.ndarray: |
| mean, std = _prefix_mean_std(x, calibration_points=calibration_points) |
| return (x - mean) / std if center else x / std |
|
|
| @lru_cache(maxsize=4) |
| def _seasonal_basis(L: int) -> tuple[np.ndarray, np.ndarray]: |
| angle = 2.0 * np.pi * np.arange(L, dtype=np.float64)[None, :] / _SEASONAL_PERIODS[:, None] |
| return (np.sin(angle), np.cos(angle)) |
|
|
| def _seasonal(rng: np.random.Generator, n: int, L: int, k_max: int=3) -> np.ndarray: |
| t = np.arange(L, dtype=np.float64)[None, :] |
| sin_basis, cos_basis = _seasonal_basis(L) |
| k = rng.integers(1, k_max + 1, size=n) |
| pair = _SEASONAL_PAIRS[rng.integers(0, len(_SEASONAL_PAIRS), size=n)] |
| use_pair = rng.random(n) < 0.35 |
| 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=_SEASONAL_PROBS) |
| if j < 2: |
| per = np.where(use_pair, pair[:, j], per) |
| per = per[:, None] |
| amp = rng.uniform(0.2, 2.0, size=n)[:, None] |
| phase = rng.uniform(0.0, 2.0 * np.pi, size=n)[:, None] |
| 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])) |
| modulated = np.nonzero((k > j) & (rng.random(n) < 0.35))[0] |
| if modulated.size: |
| modulated_local = np.searchsorted(active, modulated) |
| modulated_arg = 2.0 * np.pi * t / per[modulated] + phase[modulated] |
| m_per = np.clip(per[modulated] * 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: |
| 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 |
| 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, censor_rate: float=0.06, quantize_rate: float=0.07, regular_hold_rate: float=0.04, irregular_hold_rate: float=0.0, irregular_hold_prob_lo: float=0.01, irregular_hold_prob_hi: float=0.1, shock_row_rate: float=0.0, shock_prob_lo: float=0.001, shock_prob_hi: float=0.015) -> np.ndarray: |
| original = np.asarray(block, dtype=np.float64) |
| n, L = original.shape |
| if n == 0: |
| return original |
| out = original.copy() |
| reverse = rng.random(n) < 0.06 if allow_reverse else np.zeros(n, dtype=bool) |
| if reverse.any(): |
| out[reverse] = out[reverse, ::-1] |
| if not preserve_nonnegative: |
| invert = rng.random(n) < 0.04 |
| if invert.any(): |
| out[invert] *= -1.0 |
| calibration_len = min(L, 512) |
| if shock_row_rate > 0.0: |
| shocked = np.nonzero(rng.random(n) < shock_row_rate)[0] |
| if shocked.size and L > 1: |
| diff = np.diff(out[shocked, :calibration_len], axis=1) |
| center = np.median(diff, axis=1, keepdims=True) |
| robust_scale = 1.4826 * np.median(np.abs(diff - center), axis=1, keepdims=True) |
| fallback = np.maximum(np.std(diff, axis=1, keepdims=True), 1e-09) |
| robust_scale = np.where(robust_scale > 1e-09, robust_scale, fallback) |
| event_prob = rng.uniform(shock_prob_lo, shock_prob_hi, size=(shocked.size, 1)) |
| event_rows, event_cols = np.nonzero(rng.random((shocked.size, L)) < event_prob) |
| if event_rows.size: |
| favored_sign = rng.choice([-1.0, 1.0], size=(shocked.size, 1)) |
| sign = np.where(rng.random(event_rows.size) < 0.75, favored_sign[event_rows, 0], -favored_sign[event_rows, 0]) |
| magnitude = rng.lognormal(mean=np.log(4.0), sigma=0.6, size=event_rows.size) |
| out[shocked[event_rows], event_cols] += sign * magnitude * robust_scale[event_rows, 0] |
| if preserve_nonnegative: |
| np.maximum(out, 0.0, out=out) |
| if censor_rate > 0.0: |
| for row in np.nonzero(rng.random(n) < censor_rate)[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: |
| threshold = np.quantile(calibration, 1.0 - q) |
| out[row] = np.minimum(out[row], threshold) |
| else: |
| threshold = np.quantile(calibration, q) |
| out[row] = np.maximum(out[row], threshold) |
| if quantize_rate > 0.0: |
| quantized = np.nonzero(rng.random(n) < quantize_rate)[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 |
| if regular_hold_rate > 0.0: |
| held = np.nonzero(rng.random(n) < regular_hold_rate)[0] |
| if held.size: |
| factors = rng.choice([2, 4, 8], size=held.size, p=[0.55, 0.3, 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] |
| if irregular_hold_rate > 0.0 and L > 1: |
| irregular = np.nonzero(rng.random(n) < irregular_hold_rate)[0] |
| if irregular.size: |
| hold_prob = rng.uniform(irregular_hold_prob_lo, irregular_hold_prob_hi, size=(irregular.size, 1)) |
| hold = rng.random((irregular.size, L)) < hold_prob |
| hold[:, 0] = False |
| source_index = np.where(~hold, np.arange(L, dtype=np.int64)[None, :], 0) |
| np.maximum.accumulate(source_index, axis=1, out=source_index) |
| out[irregular] = np.take_along_axis(out[irregular], source_index, axis=1) |
| if preserve_integers: |
| out = np.maximum(np.rint(out), 0.0) |
| degenerate = out[:, :calibration_len].std(axis=1) < 1e-09 |
| if degenerate.any(): |
| out[degenerate] = original[degenerate] |
| return out |
|
|
| def _k00(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) -> np.ndarray: |
| t = np.arange(L, dtype=np.float64)[None, :] |
| level = rng.normal(0.0, 1.0, size=(n, 1)) |
| _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) |
| 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 |
| return series + _ar1_batch(innov, phi) |
|
|
| def _k01(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| 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) * rng.uniform(0.0, 1.0, size=(n, 1)) |
| 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 _k02(rng: np.random.Generator, n: int, L: int, *, hi_frac: float=0.25, exc_lo: float=0.3, exc_hi: float=2.0) -> np.ndarray: |
| t = np.arange(L, dtype=np.float64)[None, :] |
| _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))) |
| amp = rng.uniform(0.1, 0.6, size=(n, 1)) |
| seasonal_shape = _seasonal(rng, n, L, k_max=1) |
| seasonal_shape = _prefix_standardize(seasonal_shape, center=False) |
| 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 _k03(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| 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)) |
| burn = 512 |
| innov = rng.normal(0.0, 1.0, size=(n, L + burn)) * sigma |
| return _ar2_batch(innov, a1, a2)[:, burn:] |
|
|
| def _k04(rng: np.random.Generator, n: int, L: int, *, heavy_frac: float=0.25, sv_frac: float=0.3) -> np.ndarray: |
| order2 = rng.random(n) < 0.35 |
| drift = rng.normal(0.0, 0.02, size=(n, 1)) |
| sigma = rng.uniform(0.2, 1.0, size=(n, 1)) |
| eps = rng.normal(0.0, 1.0, size=(n, L)) |
| heavy = np.nonzero(rng.random(n) < heavy_frac)[0] |
| if heavy.size: |
| df = rng.uniform(3.0, 12.0, size=(heavy.size, 1)) |
| eps[heavy] = rng.standard_t(df, size=(heavy.size, L)) / np.sqrt(df / (df - 2.0)) |
| stochastic = np.nonzero(rng.random(n) < sv_frac)[0] |
| if stochastic.size: |
| phi = 0.995 |
| burn = 256 |
| vol_innov = rng.standard_normal((stochastic.size, L + burn)) * np.sqrt(1.0 - phi * phi) |
| log_vol = lfilter([1.0], [1.0, -phi], vol_innov, axis=1)[:, burn:] |
| log_vol *= rng.uniform(0.1, 0.55, size=(stochastic.size, 1)) |
| eps[stochastic] *= np.exp(np.clip(log_vol, -2.0, 2.0)) |
| steps = eps * sigma + drift |
| walk = np.cumsum(steps, axis=1) |
| walk2 = np.cumsum(walk, axis=1) |
| o2 = order2[:, None] |
| return np.where(o2, walk2 / max(L, 1), walk) |
|
|
| def _k05(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| 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)) |
| burn = 256 |
| total = L + burn |
| innov = rng.normal(0.0, 1.0, size=(n, total)) * sigma |
| x = np.empty((n, total), dtype=np.float64) |
| x[:, 0] = innov[:, 0] |
| for t in range(1, total): |
| 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], -1000000.0, 1000000.0) |
| return x[:, burn:] |
|
|
| def _k06(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| use_sine = rng.random(n) < 0.5 |
| r_log = rng.uniform(3.6, 4.0, size=n) |
| r_sin = rng.uniform(0.85, 1.0, size=n) |
| x0 = rng.uniform(0.05, 0.95, size=n) |
| cur = x0.copy() |
| for _ in range(64): |
| nxt_log = r_log * cur * (1.0 - cur) |
| nxt_sin = r_sin * np.sin(np.pi * cur) |
| cur = np.clip(np.where(use_sine, nxt_sin, nxt_log), 0.0, 1.0) |
| x = np.empty((n, L), dtype=np.float64) |
| x[:, 0] = cur |
| for t in range(1, L): |
| nxt_log = r_log * cur * (1.0 - cur) |
| nxt_sin = r_sin * np.sin(np.pi * cur) |
| cur = np.where(use_sine, nxt_sin, nxt_log) |
| cur = np.clip(cur, 0.0, 1.0) |
| x[:, t] = cur |
| return _prefix_standardize(x) |
|
|
| def _k07(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| embed = 2 * L |
| f = np.fft.rfftfreq(embed)[None, :] |
| lengthscale = np.exp(rng.uniform(np.log(8.0), np.log(256.0), size=(n, 1))) |
| envelope = np.exp(-0.5 * (2.0 * np.pi * lengthscale * f) ** 2) |
| z = rng.standard_normal((n, f.shape[1])) + 1j * rng.standard_normal((n, f.shape[1])) |
| z[:, 0] = 0.0 |
| x = np.fft.irfft(z * np.sqrt(envelope), n=embed, axis=1)[:, :L] |
| return _prefix_standardize(x) |
|
|
| def _k08(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| embed = 2 * L |
| f = np.fft.rfftfreq(embed) |
| safe_f = np.maximum(f, 1.0 / embed)[None, :] |
| beta = rng.uniform(-0.6, 2.4, size=(n, 1)) |
| amp = safe_f ** (-0.5 * beta) |
| 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, 1.0 / embed) |
| 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, 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 _k09(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| 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) |
| slow = rng.random((n, 1)) < 0.5 |
| phi = np.where(slow, rng.uniform(0.995, 0.9995, size=(n, 1)), rng.uniform(0.9, 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) * rng.uniform(0.5, 3.0, size=(n, 1)) |
| log_sigma0 = rng.normal(np.log(0.3), 0.3, size=(n, 1)) |
| log_sigma1 = rng.normal(np.log(1.5), 0.5, size=(n, 1)) |
| log_sigma_mean = np.where(regime == 0, log_sigma0, log_sigma1) |
| vol_rho = rng.uniform(0.951, 0.995, size=(n, 1)) |
| vol_eta = rng.uniform(0.03, 0.2, size=(n, 1)) |
| vol_eps = rng.standard_normal((n, L)) |
| vol_drive = (1.0 - vol_rho) * log_sigma_mean + np.sqrt(1.0 - vol_rho * vol_rho) * vol_eta * vol_eps |
| log_vol = _ar1_zi_batch(vol_drive, vol_rho[:, 0], log_sigma_mean[:, 0]) |
| vol = np.exp(np.clip(log_vol, -5.0, 5.0)) |
| eps = rng.standard_normal((n, L)) |
| heavy = np.nonzero(rng.random(n) < 0.35)[0] |
| if heavy.size: |
| 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) |
| 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-06)) |
| drive = (1.0 - phi) * mean + innovation_scale * vol * eps |
| x0 = mean[:, 0] + vol[:, 0] * eps[:, 0] |
| out = _ar1_zi_batch(drive, phi[:, 0], x0) |
| 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 _k10(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| kind = rng.integers(0, 4, size=n) |
| seasonal = _seasonal(rng, n, L, k_max=2) |
| fronts = np.cumsum(_sparse_jumps(rng, n, L, rate=5.0 / L, scale=1.0), axis=1) |
| need_smooth = kind != 2 |
| smooth = np.zeros((n, L), dtype=np.float64) |
| if need_smooth.any(): |
| idx = np.nonzero(need_smooth)[0] |
| smooth[idx] = _k07(rng, int(idx.size), L) |
| 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)) |
| 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()) |
| diffusion = np.exp(rng.uniform(np.log(0.03), np.log(0.2), size=(count, 1))) |
| walk = np.cumsum(rng.standard_normal((count, L)) * diffusion, axis=1) |
| level = rng.uniform(900.0, 1100.0, size=(count, 1)) |
| out[pressure] = level + 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 _k11(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| t = np.arange(L, dtype=np.float64)[None, :] |
| period = rng.choice(_SEASONAL_PERIODS, size=(n, 1), p=_SEASONAL_PROBS) |
| 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))) |
| 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) |
| 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, 10000000.0, out=lam) |
| 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 _k12(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| 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 = np.maximum(1.0, np.rint(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))) |
| return occur * magnitude |
|
|
| def _k13(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| base = _k07(rng, n, L) * rng.uniform(0.5, 2.0, size=(n, 1)) |
| base += _seasonal(rng, n, L, k_max=1) * rng.uniform(0.0, 1.0, size=(n, 1)) |
| sharp = _sparse_jumps(rng, n, L, rate=3.0 / L, scale=rng.uniform(3.0, 8.0, size=n)) |
| 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 |
| 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 |
| _CS_CALM_LO = 192 |
| _CS_CALM_HI = 1024 |
| _CS_DYNAMIC_LO = 96 |
| _CS_DYNAMIC_HI = 512 |
|
|
| def _k14(rng: np.random.Generator, n: int, L: int) -> np.ndarray: |
| if n <= 0: |
| return np.empty((0, L), dtype=np.float64) |
| if L <= 0: |
| return np.empty((n, 0), dtype=np.float64) |
| kind = rng.integers(0, 4, size=n) |
| dynamic = np.empty((n, L), dtype=np.float64) |
| mask = kind == 0 |
| if mask.any(): |
| dynamic[mask] = _seasonal(rng, int(mask.sum()), L, k_max=2) |
| mask = kind == 1 |
| if mask.any(): |
| m = int(mask.sum()) |
| dynamic[mask] = _ar1_batch(rng.normal(size=(m, L)) * rng.uniform(0.12, 0.55, size=(m, 1)), rng.uniform(0.35, 0.92, size=m)) |
| mask = kind == 2 |
| if mask.any(): |
| dynamic[mask] = _k07(rng, int(mask.sum()), L) |
| mask = kind == 3 |
| if mask.any(): |
| m = int(mask.sum()) |
| dynamic[mask] = np.cumsum(rng.normal(size=(m, L)) * rng.uniform(0.025, 0.16, size=(m, 1)), axis=1) |
| calm_kind = rng.integers(0, 4, size=n) |
| level = rng.normal(0.0, 2.0, size=n) |
| scale = np.exp(rng.uniform(np.log(0.4), np.log(12.0), size=n)) |
| dynamic_amp = rng.uniform(0.6, 2.2, size=n) |
| start_calm = rng.random(n) < 0.65 |
| cal_n = min(L, 512) |
| means = dynamic[:, :cal_n].mean(axis=1, keepdims=True) |
| stds = dynamic[:, :cal_n].std(axis=1, keepdims=True) |
| stds = np.where(stds > 1e-12, stds, 1.0) |
| dynamic -= means |
| dynamic /= stds |
| out = np.empty((n, L), dtype=np.float64) |
| for row in range(n): |
| current = float(level[row]) |
| calm = bool(start_calm[row]) |
| pos = 0 |
| segment_index = 0 |
| amp = float(dynamic_amp[row]) |
| mode = int(calm_kind[row]) |
| while pos < L: |
| if calm: |
| seg_len = int(rng.integers(_CS_CALM_LO, _CS_CALM_HI + 1)) |
| else: |
| seg_len = int(rng.integers(_CS_DYNAMIC_LO, _CS_DYNAMIC_HI + 1)) |
| if segment_index == 0 and L >= 2 * _CS_DYNAMIC_LO: |
| seg_len = min(seg_len, L - _CS_DYNAMIC_LO) |
| end = min(pos + max(seg_len, 1), L) |
| span = end - pos |
| if calm: |
| if mode == 0: |
| out[row, pos:end] = current |
| values_last = current |
| elif mode == 1: |
| drift = rng.normal(0.0, 0.0025, size=span).cumsum() |
| drift += np.linspace(0.0, float(rng.normal(0.0, 0.025)), span) |
| values = current + drift |
| out[row, pos:end] = values |
| values_last = float(values[-1]) |
| elif mode == 2: |
| count_level = max(0.0, float(np.rint(abs(current) * 8.0))) |
| updates = rng.random(span) < 0.025 |
| changes = updates * rng.choice([-1.0, 1.0], size=span) |
| values = np.maximum(count_level + np.cumsum(changes), 0.0) |
| out[row, pos:end] = values |
| values_last = float(values[-1]) |
| else: |
| events = rng.random(span) < 0.012 |
| values = events * rng.gamma(1.5, 0.35, size=span) |
| out[row, pos:end] = values |
| values_last = float(values[-1]) |
| else: |
| piece = dynamic[row, pos:end] * amp |
| values = piece - piece[0] + current |
| if span > 1 and np.ptp(values) < 1e-10: |
| values = current + np.linspace(0.0, 1.0, span) |
| out[row, pos:end] = values |
| values_last = float(values[-1]) |
| current = values_last |
| pos = end |
| calm = not calm |
| segment_index += 1 |
| row_scale = 1.0 if mode == 2 else float(scale[row]) |
| out[row] *= row_scale |
| if L > 1 and np.ptp(out[row]) < 1e-10: |
| out[row, -1] += max(0.001, 0.01 * row_scale) |
| return out |
|
|
| def _sanitize(block: np.ndarray) -> np.ndarray: |
| x = np.asarray(block, dtype=np.float64) |
| np.nan_to_num(x, copy=False, nan=0.0, posinf=1000000.0, neginf=-1000000.0) |
| if x.ndim == 1: |
| peak = float(np.max(np.abs(x))) |
| if peak > 1000000.0: |
| x *= 1000000.0 / peak |
| else: |
| peak = np.max(np.abs(x), axis=1, keepdims=True) |
| scale = np.where(peak > 1000000.0, 1000000.0 / np.maximum(peak, 1e-12), 1.0) |
| x *= scale |
| return x |
