扩展分箱模型使用compute_sweights计算s权重时遇异常

问题 #

使用compute_sweights()计算扩展分箱模型的s权重时，无分箱PDF场景下运行正常，但分箱PDF即便拟合效果良好，仍会抛出ModelNotFittedToData错误。示例代码如下：

import numpy as np
import os
os.environ["ZFIT_DISABLE_TF_WARNINGS"] = "1"
import zfit
from hepstats.splot import compute_sweights

class FirstOrderPoly(zfit.pdf.ZPDF):
    """First order polynomial `a * x + b`"""
    _PARAMS = ['a']
    @zfit.supports(norm=False)
    def _pdf(self, x, norm, params):
        del norm
        data = x[0]
        a = params['a']
        return 1 + a * data
def cdf_poly(limit, a):
    return limit + 0.5 * a * limit ** 2
def integral_func(limits, norm, params, model):
    del norm, model
    a = params['a']
    lower, upper = limits.v1.limits
    integral = cdf_poly(upper, a) - cdf_poly(lower, a)
    return integral
integral_limits = zfit.Space(axes=(0,), limits=(zfit.Space.ANY, zfit.Space.ANY))
FirstOrderPoly.register_analytic_integral(func=integral_func, limits=integral_limits)

binned = True

nbkg = 20000
nsig = 45000
bounds = (1920, 2050)
np.random.seed(0)
bkg = np.random.uniform(1920, 2050, nbkg)
peak = np.random.normal(1969, 6.75, nsig)
mass = np.concatenate((bkg,peak))
sel = (mass > bounds[0]) & (mass < bounds[1])
mass = mass[sel]
sorter = np.argsort(mass)
mass = mass[sorter]

obs = zfit.Space("x", limits=bounds)
mean = zfit.Parameter("mean", 1969, 1920, 2050)
sigma = zfit.Parameter("sigma", 6.75, 0.02, 10)
bgA = zfit.Parameter("bgA", 0, -1.0, 1.0)
Nsig = zfit.Parameter("Nsig", nsig, 0.0, len(mass))
Nbkg = zfit.Parameter("Nbkg", nbkg, 0.0, len(mass))

if binned:
    data = zfit.Data(data=mass, obs=obs).to_binned(100)
    signal = zfit.pdf.Gauss(obs=obs, mu=mean, sigma=sigma).create_extended(Nsig).to_binned(data.space)
    bg = FirstOrderPoly(obs=obs, a=bgA).create_extended(Nbkg).to_binned(data.space)
    tot_model = zfit.pdf.BinnedSumPDF(pdfs=[signal, bg])
    nll = zfit.loss.ExtendedBinnedNLL(model=tot_model, data=data)
else:
    data = zfit.Data(data=mass, obs=obs)
    signal = zfit.pdf.Gauss(obs=obs, mu=mean, sigma=sigma).create_extended(Nsig)
    bg = FirstOrderPoly(obs=obs, a=bgA).create_extended(Nbkg)
    tot_model = zfit.pdf.SumPDF(pdfs=[signal, bg])
    nll = zfit.loss.ExtendedUnbinnedNLL(model=tot_model, data=data)
minimizer = zfit.minimize.Minuit()
result = minimizer.minimize(loss=nll)
result.hesse()
sweights = compute_sweights(tot_model, data)[Nsig]

错误由compute_sweights()内部依赖的eval_pdf()接口触发，该接口无法正确处理分箱PDF的输入逻辑。请问这是bug，还是需要换一种方式定义扩展无分箱复合模型？

解答 #

这不是bug，而是compute_sweights()当前仅支持无分箱的扩展模型，它依赖的底层接口无法直接兼容分箱PDF（如BinnedSumPDF）和分箱数据。

推荐两种解决方法：

方法一：分箱拟合+无分箱计算s权重 #

保持分箱拟合的代码不变，在计算s权重时，基于拟合得到的参数值重新定义无分箱扩展模型，传入原始无分箱数据计算权重：

# 拟合完成后替换原sweights计算代码
if binned:
    # 重新构建无分箱扩展模型
    unbinned_signal = zfit.pdf.Gauss(obs=obs, mu=mean, sigma=sigma).create_extended(Nsig)
    unbinned_bg = FirstOrderPoly(obs=obs, a=bgA).create_extended(Nbkg)
    unbinned_model = zfit.pdf.SumPDF(pdfs=[unbinned_signal, unbinned_bg])
    # 使用原始无分箱数据
    unbinned_data = zfit.Data(data=mass, obs=obs)
    sweights = compute_sweights(unbinned_model, unbinned_data)[Nsig]
else:
    sweights = compute_sweights(tot_model, data)[Nsig]

这种方法既保留了分箱拟合的精度优势，又满足compute_sweights()对无分箱模型的要求。

方法二：手动实现分箱模型的s权重计算 #

若必须基于分箱数据计算，可参考s权重的数学公式手动实现：

• 计算每个分箱内信号、背景的PDF值
• 构建并求逆协方差矩阵
• 结合分箱计数计算对应s权重，再映射到原始数据点

此方法复杂度较高，仅在特殊场景下使用。

内容的提问来源于stack exchange，提问作者Wren