Подбор гиперпараметров с Optuna

Туториал: оптимизация биннинга

Благодаря быстрому движку вывода на Rust NextStat удобно использовать как целевую функцию в black-box оптимизации. Один фит занимает порядка ~1–50 ms на CPU или ~0.5–5 ms на GPU, поэтому Optuna может оценивать сотни конфигураций за минуты, а не за часы.

В этом руководстве показано, как с помощью Optuna подобрать оптимальный биннинг гистограммы, максимизирующий значимость открытия Z₀. Тот же подход подходит и для других гиперпараметров: порогов по срезам, выборок сэмплов, стратегий учета систематик и т. п.

Установка

pip install nextstat optuna
# Опционально: pip install optuna-dashboard  (мониторинг в реальном времени)

Постановка задачи: оптимальный биннинг

Биннинг гистограммы напрямую влияет на статистическую чувствительность. Слишком мало бинов: теряется форма распределений. Слишком много: начинают доминировать статистические флуктуации. Оптимальный биннинг зависит от форм сигнала и фона, статистики и систематических неопределенностей.

Традиционно это делают вручную или простым перебором. С Optuna и NextStat можно автоматически искать параметры биннинга с помощью байесовской оптимизации (TPE).

Шаг 1: функция построения workspace

Сначала напишите функцию, которая строит HistFactory workspace по параметрам биннинга. Именно здесь находится логика, специфичная для вашего анализа.

import json
import numpy as np

def build_workspace(n_bins: int, lo: float, hi: float) -> dict:
    """Простой 1-канальный HistFactory workspace с заданным биннингом.

    В реальном анализе здесь обычно читаются ntuple и выполняется ребиннинг по заданным границам.
    Для иллюстрации используем аналитические формы.
    """
    edges = np.linspace(lo, hi, n_bins + 1)
    centers = 0.5 * (edges[:-1] + edges[1:])
    width = edges[1] - edges[0]

    # Аналитический сигнал: гауссов пик около 0.5
    signal = 50.0 * np.exp(-0.5 * ((centers - 0.5) / 0.08) ** 2) * width
    # Аналитический фон: убывающая экспонента
    background = 200.0 * np.exp(-2.0 * centers) * width

    # HistFactory workspace (формат pyhf)
    return {
        "channels": [{
            "name": "SR",
            "samples": [
                {
                    "name": "signal",
                    "data": signal.tolist(),
                    "modifiers": [
                        {"name": "mu", "type": "normfactor", "data": None}
                    ],
                },
                {
                    "name": "background",
                    "data": background.tolist(),
                    "modifiers": [
                        {"name": "bkg_norm", "type": "normsys",
                         "data": {"hi": 1.05, "lo": 0.95}},
                    ],
                },
            ],
        }],
        "observations": [{
            "name": "SR",
            "data": (signal + background).tolist(),
        }],
        "measurements": [{
            "name": "meas",
            "config": {
                "poi": "mu",
                "parameters": [],
            },
        }],
        "version": "1.0.0",
    }

Шаг 2: целевая функция Optuna

import nextstat

def objective(trial):
    """Целевая функция Optuna: максимизировать Z₀, подбирая биннинг."""

    # Пространство поиска
    n_bins = trial.suggest_int("n_bins", 3, 40)
    lo = trial.suggest_float("lo", 0.0, 0.3)
    hi = trial.suggest_float("hi", 0.7, 1.0)

    # Workspace с биннингом текущей пробы
    ws = build_workspace(n_bins=n_bins, lo=lo, hi=hi)

    # Загрузка в NextStat
    try:
        model = nextstat.from_pyhf(ws)
    except Exception:
        return 0.0  # некорректная конфигурация -> нулевая значимость

    # Проверка гипотезы при μ=0 -> значимость открытия Z₀
    hypo = nextstat.hypotest(model, mu=0.0)
    return float(hypo.significance)  # Z₀ в сигмах

Шаг 3: запуск исследования (study)

import optuna

# Создать study (максимизируем значимость)
study = optuna.create_study(
    direction="maximize",
    study_name="nextstat-binning",
    sampler=optuna.samplers.TPESampler(seed=42),
)

# Запустить 200 проб (порядка 10–30 секунд с NextStat)
study.optimize(objective, n_trials=200, show_progress_bar=True)

# Результаты
print(f"Best Z₀: {study.best_value:.3f}σ")
print(f"Best params: {study.best_params}")
# → Best Z₀: 3.142σ
# → Best params: {'n_bins': 15, 'lo': 0.12, 'hi': 0.92}

Шаг 4: анализ результатов

# Optuna built-in visualisations
from optuna.visualization import (
    plot_optimization_history,
    plot_param_importances,
    plot_contour,
)

# Сходимость оптимизации
fig1 = plot_optimization_history(study)
fig1.show()

# Какие параметры влияют сильнее всего?
fig2 = plot_param_importances(study)
fig2.show()

# 2D контур: n_bins vs lo
fig3 = plot_contour(study, params=["n_bins", "lo"])
fig3.show()

Шаг 5: логирование в W&B (опционально)

import wandb
from nextstat.mlops import metrics_dict

wandb.init(project="nextstat-optuna")

def objective_with_logging(trial):
    n_bins = trial.suggest_int("n_bins", 3, 40)
    lo = trial.suggest_float("lo", 0.0, 0.3)
    hi = trial.suggest_float("hi", 0.7, 1.0)

    ws = build_workspace(n_bins=n_bins, lo=lo, hi=hi)
    model = nextstat.from_pyhf(ws)
    result = nextstat.fit(model)

    # Логируем метрики фита в W&B
    wandb.log({
        "trial": trial.number,
        "n_bins": n_bins,
        "lo": lo,
        "hi": hi,
        **metrics_dict(result, prefix="fit/"),
    })

    return -float(result.nll) if result.converged else 0.0

Расширенное: цель с учетом GPU-ранкинга

Для моделей с большим числом каналов или бинов можно добавить в целевую функцию GPU-ускоренный ранкинг, чтобы оценивать чувствительность к систематикам одновременно со значимостью:

from nextstat.interpret import rank_impact

def objective_gpu(trial):
    n_bins = trial.suggest_int("n_bins", 5, 50)
    lo = trial.suggest_float("lo", 0.0, 0.2)
    hi = trial.suggest_float("hi", 0.8, 1.0)

    ws = build_workspace(n_bins=n_bins, lo=lo, hi=hi)
    model = nextstat.from_pyhf(ws)

    # Считаем Z₀ через проверку гипотезы
    hypo = nextstat.hypotest(model, mu=0.0)
    z0 = float(hypo.significance)

    # Вклад систематик как регуляризация
    ranking = rank_impact(model, top_n=5)
    total_syst = sum(r["total_impact"] for r in ranking)

    # Мультикритерий: максимизируем Z₀ и штрафуем чувствительность к систематикам
    return z0 - 0.1 * total_syst

Расширенное: параллельный поиск с Ray Tune

Для распределенного поиска на нескольких машинах можно обернуть Optuna в Ray Tune:

from ray import tune
from ray.tune.search.optuna import OptunaSearch

def trainable(config):
    """Trainable-функция Ray Tune, оборачивающая вызов NextStat."""
    ws = build_workspace(
        n_bins=config["n_bins"],
        lo=config["lo"],
        hi=config["hi"],
    )
    model = nextstat.from_pyhf(ws)
    hypo = nextstat.hypotest(model, mu=0.0)
    tune.report(z0=float(hypo.significance))

search = OptunaSearch(
    metric="z0",
    mode="max",
    sampler=optuna.samplers.TPESampler(seed=42),
)

tuner = tune.Tuner(
    trainable,
    param_space={
        "n_bins": tune.randint(3, 40),
        "lo": tune.uniform(0.0, 0.3),
        "hi": tune.uniform(0.7, 1.0),
    },
    tune_config=tune.TuneConfig(
        num_samples=200,
        search_alg=search,
    ),
)

results = tuner.fit()
print(results.get_best_result("z0", "max").config)

Производительность

Размер модели	Фит на CPU	Фит на GPU	200 проб
10 бинов, 2 NP	~1 ms	~0.5 ms	~1 s
100 бинов, 20 NP	~10 ms	~2 ms	~4 s
1000 бинов, 100 NP	~100 ms	~10 ms	~20 s

Советы

Прюнинг — MedianPruner в Optuna может досрочно останавливать бесперспективные пробы. Хорошо сочетается с NextStat, потому что каждая проба быстрая.
Несколько критериев — используйте optuna.create_study(directions=["maximize", "minimize"]), чтобы совместно оптимизировать Z₀ и время вычислений.
Сохранение прогресса — используйте optuna.create_study(storage="sqlite:///study.db"), чтобы сохранять результаты между сессиями.
Dashboard — запустите optuna-dashboard sqlite:///study.db для мониторинга в реальном времени в браузере.