Байесовское сэмплирование (NUTS & MAMS)

NextStat включает два градиентных MCMC-сэмплера на Rust: классический No-U-Turn Sampler (NUTS) и более новый Metropolis-Adjusted Microcanonical Sampler (MAMS, Robnik et al., 2025). Оба дают одинаковый формат Chainи совместимы с диагностикой ArviZ. Python-дополнение [bayes]устанавливает emcee (ансамблевый MCMC) и ArviZ (диагностика/визуализация) как опциональные Python-семплеры и инструменты анализа.

Установка

Это установит emcee (≥3.1.6), ArviZ (≥0.23.4) и numpy (≥2.0) вместе с основным пакетом.

Использование в Python

import json
from pathlib import Path
import nextstat

workspace = json.loads(Path("workspace.json").read_text())
model = nextstat.from_pyhf(json.dumps(workspace))

idata = nextstat.bayes.sample(
    model,
    n_chains=2,
    n_warmup=500,
    n_samples=1000,
    seed=42,
    target_accept=0.8,
)

print(idata)
# Возвращает arviz.InferenceData с постериорными семплами

Параметры

Диагностика

Возвращаемый объект InferenceData полностью совместим с ArviZ для диагностики:

import arviz as az

# Трейс-плот (trace plot)
az.plot_trace(idata)

# Таблица итогов (R-hat, ESS)
print(az.summary(idata))

# Попарные проекции (pair plot)
az.plot_pair(idata, var_names=["mu"])

MAMS: микроканонический сэмплер

MAMS (arXiv:2503.01707) — микроканонический MCMC-метод, использующий изокинетическую динамику на единичной сфере вместо древовидного выбора траектории. Статья сообщает об улучшении ESS на градиент в 2–7× по сравнению с NUTS на опубликованных бенчмарках. Ключевые отличия:

• Фиксированная длина траектории — без построения дерева, предсказуемая стоимость перехода
• Скорость живёт на S^d-1 с проективным обновлением градиента
• Частичное обновление скорости для декорреляции между переходами
• Overflow-устойчивый критерий MH через log-sum-exp формулировку
• 3-фазный авто-тюнинг warmup: размер шага → прекондиционер → длина декогеренции

Использование MAMS

import nextstat

model = nextstat.EightSchoolsModel(y, sigma, mu_prior=0.0, tau_prior=5.0)

# MAMS — замена NUTS в одну строку
result = nextstat.sample_mams(
    model,
    n_chains=4,
    n_warmup=1000,
    n_samples=2000,
    seed=42,
)

# Та же диагностика что и NUTS
print(result.summary())      # R-hat, ESS, дивергенции
idata = result.to_arviz()    # ArviZ InferenceData

NUTS vs MAMS

Оба сэмплера дают идентичный формат вывода и поддерживают одну и ту же диагностику (R-hat, bulk/tail ESS, дивергенции). Выбирайте NUTS для проверенной надёжности; пробуйте MAMS когда нужен потенциально более высокий ESS на градиент для сложных геометрий (воронки, мультимасштабные постериоры).

Единый интерфейс sample()

nextstat.sample() — единая точка входа для всех трёх MCMC-методов. Параметр method выбирает сэмплер; все специфичные kwargs передаются далее.

import nextstat as ns

model = ns.EightSchoolsModel([28,8,-3,7,-1,1,18,12], [15,10,16,11,9,11,10,18])

# NUTS (по умолчанию)
idata = ns.sample(model, method="nuts", n_samples=2000, return_idata=True)

# MAMS
idata = ns.sample(model, method="mams", n_samples=2000, return_idata=True)

# LAPS (GPU — требуется CUDA или Metal сборка)
result = ns.sample("eight_schools", method="laps",
                   model_data={"y": [28,8,-3,7,-1,1,18,12],
                               "sigma": [15,10,16,11,9,11,10,18]},
                   n_chains=4096, n_samples=2000)

# Сохранение трейса на диск
ns.sample(model, n_samples=2000, return_idata=True, out="trace.json")

Стратегия инициализации Pathfinder

NUTS и MAMS поддерживают init_strategy="pathfinder" — полный L-BFGS фит с гессианом для нахождения моды постериора. Диагональная обратная матрица Гессе используется как начальная масс-матрица, что позволяет сократить warmup (например, 200 вместо 1000 итераций). При неудаче фита — fallback на случайную инициализацию.

import nextstat as ns

# Pathfinder init для NUTS
idata = ns.sample(model, method="nuts", init_strategy="pathfinder", n_warmup=200)

# Pathfinder init для MAMS
idata = ns.sample(model, method="mams", init_strategy="pathfinder", n_warmup=200)

LAPS: GPU-ускоренное сэмплирование

LAPS (Late-Adjusted Parallel Sampler) — GPU-ускоренный вариант MAMS, запускающий 4096+ цепочек одновременно. Двухфазный warmup: Фаза 1 — нескорректированный MCLMC для быстрого исследования постериора, Фаза 2 — MH-корректировка для точного сэмплирования. Поддерживает CUDA (f64, NVIDIA GPU) и Metal (f32, Apple Silicon M1–M4). Автоматический fallback: CUDA > Metal > ошибка.

import nextstat as ns

# Встроенные модели: std_normal, eight_schools, neal_funnel, glm_logistic
result = ns.sample_laps(
    "eight_schools",
    model_data={"y": [28,8,-3,7,-1,1,18,12],
                "sigma": [15,10,16,11,9,11,10,18]},
    n_chains=4096,
    n_warmup=500,
    n_samples=2000,
    seed=42,
)
print(f"GPU цепочки: {result['n_gpu_chains']}")
print(f"Время: {result['wall_time_s']:.2f}s")

# Мульти-GPU (только CUDA)
result = ns.sample_laps("glm_logistic", model_data=data,
                        n_chains=65536, device_ids=[0,1,2,3])

Требуется NVIDIA GPU (CUDA-сборка) или Apple Silicon Mac (Metal-сборка). На CUDA: f64 точность. На Metal: f32 точность с fused multi-step ядром и SIMD-group cooperative ядром для тяжёлых GLM.

Оконная адаптация масс

Warmup фазы 2+3 LAPS теперь используют Stan-style удвоенные окна (по умолчанию 3) для оценки inv_mass. Каждое окно сбрасывает статистики Welford и dual averaging, улучшая сходимость на мультимасштабных моделях. Настраивается через n_mass_windows в LapsConfig (установите 1 для однопроходного поведения).

Римановский MAMS для Neal Funnel

Новая модель neal_funnel_riemannian с позиционно-зависимой метрикой Фишера G = diag(1/9, exp(−v), …). Использует эффективный потенциал, где члены лог-детерминанта сокращаются, давая масштабно-инвариантные натуральные градиенты. Доступно на Metal (f32) и CUDA (f64).

import nextstat as ns

# Римановский MAMS на GPU — естественная обработка воронкообразной геометрии
result = ns.sample_laps("neal_funnel_riemannian",
                        model_data={"dim": 10},
                        n_chains=4096, n_samples=2000)