diff --git a/.gitignore b/.gitignore index 4fd5169..a07556d 100644 --- a/.gitignore +++ b/.gitignore @@ -103,5 +103,6 @@ venv.bak/ # mypy .mypy_cache/ /.idea/sonarlint/* +/src/zfit_physics/_version.py /tests/tfpwa/data/ /src/zfit_physics/_version.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b9159a2..fb0985b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,5 @@ ci: - autoupdate_schedule: quarterly - + autoupdate_schedule: quarterly repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 @@ -39,62 +38,70 @@ repos: - repo: https://github.com/pre-commit/pygrep-hooks rev: v1.10.0 hooks: - - id: python-use-type-annotations - - id: python-check-mock-methods - - id: python-no-eval - - id: rst-directive-colons - - - repo: https://github.com/PyCQA/isort - rev: 5.13.2 - hooks: - - id: isort - + - id: python-use-type-annotations + - id: python-check-mock-methods + - id: python-no-eval + - id: rst-directive-colons + - repo: https://github.com/PyCQA/isort + rev: 5.13.2 + hooks: + - id: isort - repo: https://github.com/asottile/pyupgrade rev: v3.17.0 hooks: - id: pyupgrade - args: [ --py38-plus ] - + args: + - --py39-plus - repo: https://github.com/asottile/setup-cfg-fmt rev: v2.5.0 hooks: - id: setup-cfg-fmt - args: [ --max-py-version=3.12, --include-version-classifiers ] - - - # Notebook formatting + args: + - --max-py-version=3.12 + - --include-version-classifiers - repo: https://github.com/nbQA-dev/nbQA rev: 1.8.7 hooks: - id: nbqa-isort - additional_dependencies: [ isort ] - + additional_dependencies: + - isort - id: nbqa-pyupgrade - additional_dependencies: [ pyupgrade ] - args: [ --py38-plus ] - + additional_dependencies: + - pyupgrade + args: + - --py39-plus - repo: https://github.com/mgedmin/check-manifest rev: '0.49' hooks: - id: check-manifest - stages: [ manual ] + stages: + - manual - repo: https://github.com/sondrelg/pep585-upgrade - rev: 'v1.0' + rev: v1.0 hooks: - - id: upgrade-type-hints - args: [ '--futures=true' ] - + - id: upgrade-type-hints + args: + - --futures=true - repo: https://github.com/MarcoGorelli/auto-walrus rev: 0.3.4 hooks: - - id: auto-walrus + - id: auto-walrus - repo: https://github.com/astral-sh/ruff-pre-commit rev: "v0.6.9" hooks: - id: ruff - types_or: [ python, pyi, jupyter ] - args: [ --fix, --unsafe-fixes, --show-fixes , --line-length=120] - # Run the formatter. + types_or: + - python + - pyi + - jupyter + args: + - --fix + - --unsafe-fixes + - --show-fixes + - --line-length=120 - id: ruff-format - types_or: [ python, pyi, jupyter ] + types_or: + - python + - pyi + - jupyter diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5be70da..20d82d7 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -7,6 +7,7 @@ Develop Major Features and Improvements ------------------------------- +- add a RooFit compatibility layer and automatically convert losses, also inside minimizers (through ``SimpleLoss.from_any``) - `TF-PWA `_ support for loss functions. Minimizer can directly minimize the loss function of a model. Breaking changes diff --git a/docs/api/static/zfit_physics.roofit.rst b/docs/api/static/zfit_physics.roofit.rst new file mode 100644 index 0000000..e6a0a77 --- /dev/null +++ b/docs/api/static/zfit_physics.roofit.rst @@ -0,0 +1,52 @@ +RooFit +======================= + +ROOT provides with the `RooFit library `_ a toolkit for modeling the expected distribution of events in a physics analysis. +It can be connected with zfit, currently by providing a loss function that can be minimized by a zfit minimizer. + +This requires the `ROOT framework `_ to be installed and available in the python environment. +For example via conda: + +.. code-block:: console + + $ mamba install -c conda-forge root + +Import the module with: + +.. code-block:: python + + import zfit_physics.roofit as ztfroofit + +this will enable the RooFit functionality in zfit. + +We can create a RooFit NLL as ``RooFit_nll`` and use it as a loss function in zfit. For example, with a Gaussian model ``RooFit_gauss`` and a dataset ``RooFit_data``, both created with RooFit: + +.. code-block:: python + + RooFit_nll = RooFit_gauss.createNLL(RooFit_data) + minimizer.minimize(loss=RooFit_nll) + +More explicitly, the loss function can be created with + +.. code-block:: python + + nll = zroofit.loss.nll_from_roofit(fcn) + + +Variables +++++++++++++ + + +.. automodule:: zfit_physics.roofit.variables + :members: + :undoc-members: + :show-inheritance: + + +Loss +++++++++++++ + +.. automodule:: zfit_physics.roofit.loss + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/index.rst b/docs/index.rst index f68d2bf..1f79c35 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -31,3 +31,4 @@ Extensions :maxdepth: 1 api/static/zfit_physics.tfpwa.rst + api/static/zfit_physics.roofit.rst diff --git a/src/zfit_physics/roofit/__init__.py b/src/zfit_physics/roofit/__init__.py new file mode 100644 index 0000000..e66f8bb --- /dev/null +++ b/src/zfit_physics/roofit/__init__.py @@ -0,0 +1 @@ +from . import loss, variables diff --git a/src/zfit_physics/roofit/loss.py b/src/zfit_physics/roofit/loss.py new file mode 100644 index 0000000..12fb73a --- /dev/null +++ b/src/zfit_physics/roofit/loss.py @@ -0,0 +1,79 @@ +# Copyright (c) 2024 zfit +from __future__ import annotations + +from collections.abc import Iterable +from typing import TYPE_CHECKING + +from zfit.core.interfaces import ZfitParameter + +if TYPE_CHECKING: + try: + import ROOT + except ImportError: + ROOT = None +import warnings + +import zfit +from zfit.util.container import convert_to_container + +from .variables import roo2z_param + + +def nll_from_roofit(nll: ROOT.RooAbsReal, params: ZfitParameter | Iterable[ZfitParameter] = None): + """ + Converts a RooFit NLL (negative log-likelihood) to a Zfit loss object. + + Args: + nll: The RooFit NLL object to be converted. + params: The ``zfit.Parameter`` to be used in the loss. If None, all parameters in the NLL will be used + + Returns: + zfit.loss.SimpleLoss: The converted Zfit loss object. + + Raises: + TypeError: If the provided RooFit loss does not have an error level. + """ + params = {} if params is None else {p.name: p for p in convert_to_container(params)} + + import zfit + + def roofit_eval(x): + for par, arg in zip(nll.getVariables(), x): + par.setVal(arg) + # following RooMinimizerFcn.cxx + nll.setHideOffset(False) + r = nll.getVal() + nll.setHideOffset(True) + return r + + paramsall = [] + for v in nll.getVariables(): + param = params[name] if (name := v.GetName()) in params else roo2z_param(v) + paramsall.append(param) + + if (errordef := getattr(nll, "defaultErrorLevel", lambda: None)()) is None and ( + errordef := getattr(nll, "errordef", lambda: None)() + ) is None: + msg = ( + "Provided loss is RooFit loss but has not error level. " + "Either set it or create an attribute on the fly (like `nllroofit.errordef = 0.5`) " + ) + raise TypeError(msg) + return zfit.loss.SimpleLoss(roofit_eval, paramsall, errordef=errordef, jit=False, gradient="num", hessian="num") + + +def _nll_from_roofit_or_false(nll, params=None): + ROOT = None + if "RooAbsReal" in str(type(nll)): + try: + import ROOT + except ImportError: + warnings.warn( + f"nll ({nll}) seems to be of type RooAbsReal but ROOT is not available, skipping.", stacklevel=2 + ) + if ROOT is None or not isinstance(nll, ROOT.RooAbsReal): + return False # not a RooFit loss + return nll_from_roofit(nll, params=params) + + +zfit.loss.SimpleLoss.register_convertable_loss(nll_from_roofit, priority=50) diff --git a/src/zfit_physics/roofit/variables.py b/src/zfit_physics/roofit/variables.py new file mode 100644 index 0000000..ba002cf --- /dev/null +++ b/src/zfit_physics/roofit/variables.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import zfit + +if TYPE_CHECKING: + try: + import ROOT + except ImportError: + ROOT = None + + +def roo2z_param(v: ROOT.RooRealVar) -> zfit.Parameter: + """ + Converts a RooFit RooRealVar to a zfit parameter. + + Args: + v: RooFit RooRealVar to convert. + + Returns: + A zfit.Parameter object with properties copied from the RooFit variable. + """ + + name = v.GetName() + value = v.getVal() + label = v.GetTitle() + lower = v.getMin() + upper = v.getMax() + floating = not v.isConstant() + stepsize = None + if v.hasError(): + stepsize = v.getError() + elif v.hasAsymError(): # just take average + stepsize = (v.getErrorHi() - v.getErrorLo()) / 2 + return zfit.Parameter(name, value, lower=lower, upper=upper, floating=floating, step_size=stepsize, label=label) diff --git a/tests/roofit/test_loss_compat_roofit.py b/tests/roofit/test_loss_compat_roofit.py new file mode 100644 index 0000000..eb1a3ab --- /dev/null +++ b/tests/roofit/test_loss_compat_roofit.py @@ -0,0 +1,70 @@ +import numpy as np +import pytest + +_ = pytest.importorskip("ROOT") + + +def test_loss_registry(): + import zfit + + import zfit_physics.roofit as zroofit + + # create space + obs = zfit.Space("x", -2, 3) + + # parameters + mu = zfit.Parameter("mu", 1.2, -4, 6) + sigma = zfit.Parameter("sigma", 1.3, 0.5, 10) + + # model building, pdf creation + gauss = zfit.pdf.Gauss(mu=mu, sigma=sigma, obs=obs) + + # data + ndraw = 10_000 + data = np.random.normal(loc=2.0, scale=3.0, size=ndraw) + data = obs.filter(data) # works also for pandas DataFrame + + from ROOT import RooArgSet, RooDataSet, RooGaussian, RooRealVar + + mur = RooRealVar("mu", "mu", 1.2, -4, 6) + sigmar = RooRealVar("sigma", "sigma", 1.3, 0.5, 10) + obsr = RooRealVar("x", "x", -2, 3) + gaussr = RooGaussian("gauss", "gauss", obsr, mur, sigmar) + + datar = RooDataSet("data", "data", {obsr}) + for d in data: + obsr.setVal(d) + datar.add(RooArgSet(obsr)) + + # create a loss function + nll = gaussr.createNLL(datar) + nll_fromroofit = zroofit.loss.nll_from_roofit(nll) + + nllz = zfit.loss.UnbinnedNLL(model=gauss, data=data) + + # create a minimizer + tol = 1e-3 + verbosity = 0 + minimizer = zfit.minimize.Minuit(gradient=True, verbosity=verbosity, tol=tol, mode=1) + minimizerzgrad = zfit.minimize.Minuit(gradient=False, verbosity=verbosity, tol=tol, mode=1) + + params = nllz.get_params() + initvals = np.array(params) + + with zfit.param.set_values(params, initvals): + result = minimizer.minimize(nllz) + + with zfit.param.set_values(params, initvals): + result2 = minimizer.minimize(nll) + + assert result.params['mu']['value'] == pytest.approx(result2.params['mu']['value'], rel=1e-3) + assert result.params['sigma']['value'] == pytest.approx(result2.params['sigma']['value'], rel=1e-3) + + with zfit.param.set_values(params, params): + result4 = minimizerzgrad.minimize(nll) + + assert result.params['mu']['value'] == pytest.approx(result4.params['mu']['value'], rel=1e-3) + assert result.params['sigma']['value'] == pytest.approx(result4.params['sigma']['value'], rel=1e-3) + + with zfit.param.set_values(params, params): + result5 = minimizerzgrad.minimize(nll_fromroofit)