From dc71ebcac0a2d335e21d646f017236da25d2a1b3 Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Mon, 22 Apr 2024 21:45:20 -0400 Subject: [PATCH] wip --- .idea/.gitignore | 2 + .pre-commit-config.yaml | 23 +++++- pyproject.toml | 135 +++++++++++++++--------------- src/zfit_pwa/__init__.py | 1 - src/zfit_pwa/compwa/__init__.py | 0 src/zfit_pwa/compwa/data.py | 0 src/zfit_pwa/compwa/variables.py | 42 ++++++++++ src/zfit_pwa/compwa/wrapper.py | 60 ++++++++++++++ tests/compwa/test_wrapper.py | 136 +++++++++++++++++++++++++++++++ 9 files changed, 330 insertions(+), 69 deletions(-) create mode 100644 src/zfit_pwa/compwa/__init__.py create mode 100644 src/zfit_pwa/compwa/data.py create mode 100644 src/zfit_pwa/compwa/variables.py create mode 100644 src/zfit_pwa/compwa/wrapper.py create mode 100644 tests/compwa/test_wrapper.py diff --git a/.idea/.gitignore b/.idea/.gitignore index 13566b8..a9d7db9 100644 --- a/.idea/.gitignore +++ b/.idea/.gitignore @@ -6,3 +6,5 @@ # Datasource local storage ignored files /dataSources/ /dataSources.local.xml +# GitHub Copilot persisted chat sessions +/copilot/chatSessions diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index fbe6613..a0f276e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,3 +1,8 @@ +ci: + autoupdate_commit_msg: "chore: update pre-commit hooks" + autofix_commit_msg: "style: pre-commit fixes" + autoupdate_schedule: quarterly + repos: - repo: https://github.com/adamchainz/blacken-docs rev: "1.16.0" @@ -137,12 +142,24 @@ repos: rev: 0.7.1 hooks: - id: nbstripout + - repo: https://github.com/MarcoGorelli/auto-walrus rev: 0.3.3 hooks: - id: auto-walrus - - repo: https://github.com/shssoichiro/oxipng - rev: v9.1.0 + # uncomment locally if needed, currently needs rust version not available on pre-commit.ci + # - repo: https://github.com/shssoichiro/oxipng + # rev: v9.1.0 + # hooks: + # - id: oxipng + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: "v0.4.1" hooks: - - id: oxipng + - id: ruff + types_or: [python, pyi, jupyter] + args: [--fix, --unsafe-fixes, --show-fixes] + # Run the formatter. + - id: ruff-format + types_or: [python, pyi, jupyter] diff --git a/pyproject.toml b/pyproject.toml index 52f29b2..a3cf1f6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,47 +6,52 @@ build-backend = "hatchling.build" [project] name = "zfit-pwa" authors = [ - { name = "Jonas Eschle", email = "Jonas.Eschle@cern.ch" }, + { name = "Jonas Eschle", email = "Jonas.Eschle@cern.ch" }, ] description = "Tools to adapt to Partial Wave Analysis packages" readme = "README.md" license.file = "LICENSE" requires-python = ">=3.8" classifiers = [ - "Development Status :: 1 - Planning", - "Intended Audience :: Science/Research", - "Intended Audience :: Developers", - "License :: OSI Approved :: BSD License", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Topic :: Scientific/Engineering", - "Typing :: Typed", + "Development Status :: 1 - Planning", + "Intended Audience :: Science/Research", + "Intended Audience :: Developers", + "License :: OSI Approved :: BSD License", + "Operating System :: OS Independent", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering", + "Typing :: Typed", ] dynamic = ["version"] -dependencies = [] +dependencies = ["zfit"] [project.optional-dependencies] test = [ - "pytest >=6", - "pytest-cov >=3", -] -dev = [ - "pytest >=6", - "pytest-cov >=3", + "pytest >=6", + "pytest-cov >=3", ] + docs = [ - "sphinx>=7.0", - "myst_parser>=0.13", - "sphinx_copybutton", - "sphinx_autodoc_typehints", - "furo>=2023.08.17", + "sphinx>=7.0", + "myst_parser>=0.13", + "sphinx_copybutton", + "sphinx_autodoc_typehints", + "furo>=2023.08.17", +] +compwa = [ + "qrules", + "ampform", + "tensorwaves", +] +dev = [ + "zfit-pwa[test,docs,compwa]", + "tensorwaves[phsp]", ] [project.urls] @@ -67,22 +72,22 @@ scripts.test = "pytest {args}" [tool.pytest.ini_options] minversion = "6.0" -addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"] -xfail_strict = true -filterwarnings = [ - "error", -] +#addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"] +#xfail_strict = true +#filterwarnings = [ +# "error", +#] log_cli_level = "INFO" testpaths = [ - "tests", + "tests", ] [tool.coverage] run.source = ["zfit_pwa"] report.exclude_also = [ - '\.\.\.', - 'if typing.TYPE_CHECKING:', + '\.\.\.', + 'if typing.TYPE_CHECKING:', ] [tool.mypy] @@ -106,32 +111,32 @@ src = ["src"] [tool.ruff.lint] extend-select = [ - "B", # flake8-bugbear - "I", # isort - "ARG", # flake8-unused-arguments - "C4", # flake8-comprehensions - "EM", # flake8-errmsg - "ICN", # flake8-import-conventions - "G", # flake8-logging-format - "PGH", # pygrep-hooks - "PIE", # flake8-pie - "PL", # pylint - "PT", # flake8-pytest-style - "PTH", # flake8-use-pathlib - "RET", # flake8-return - "RUF", # Ruff-specific - "SIM", # flake8-simplify - "T20", # flake8-print - "UP", # pyupgrade - "YTT", # flake8-2020 - "EXE", # flake8-executable - "NPY", # NumPy specific rules - "PD", # pandas-vet + "B", # flake8-bugbear + "I", # isort + "ARG", # flake8-unused-arguments + "C4", # flake8-comprehensions + "EM", # flake8-errmsg + "ICN", # flake8-import-conventions + "G", # flake8-logging-format + "PGH", # pygrep-hooks + "PIE", # flake8-pie + "PL", # pylint + "PT", # flake8-pytest-style + "PTH", # flake8-use-pathlib + "RET", # flake8-return + "RUF", # Ruff-specific + "SIM", # flake8-simplify + "T20", # flake8-print + "UP", # pyupgrade + "YTT", # flake8-2020 + "EXE", # flake8-executable + "NPY", # NumPy specific rules + "PD", # pandas-vet ] ignore = [ - "PLR09", # Too many <...> - "PLR2004", # Magic value used in comparison - "ISC001", # Conflicts with formatter + "PLR09", # Too many <...> + "PLR2004", # Magic value used in comparison + "ISC001", # Conflicts with formatter ] isort.required-imports = ["from __future__ import annotations"] # Uncomment if using a _compat.typing backport @@ -148,9 +153,9 @@ ignore-paths = [".*/_version.py"] reports.output-format = "colorized" similarities.ignore-imports = "yes" messages_control.disable = [ - "design", - "fixme", - "line-too-long", - "missing-module-docstring", - "wrong-import-position", + "design", + "fixme", + "line-too-long", + "missing-module-docstring", + "wrong-import-position", ] diff --git a/src/zfit_pwa/__init__.py b/src/zfit_pwa/__init__.py index 252815e..6169293 100644 --- a/src/zfit_pwa/__init__.py +++ b/src/zfit_pwa/__init__.py @@ -4,7 +4,6 @@ zfit-pwa: Tools to adapt to Partial Wave Analysis packages """ - from __future__ import annotations from ._version import version as __version__ diff --git a/src/zfit_pwa/compwa/__init__.py b/src/zfit_pwa/compwa/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/zfit_pwa/compwa/data.py b/src/zfit_pwa/compwa/data.py new file mode 100644 index 0000000..e69de29 diff --git a/src/zfit_pwa/compwa/variables.py b/src/zfit_pwa/compwa/variables.py new file mode 100644 index 0000000..e2dcb48 --- /dev/null +++ b/src/zfit_pwa/compwa/variables.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from typing import Mapping + +import numpy as np +import pandas as pd +import zfit +from zfit.core.interfaces import ZfitUnbinnedData + + +def obs_from_frame(frame1, frame2=None, bufferfactor=0.01): + obs = [] + if frame2 is None: + frame2 = frame1 + + if isinstance(frame1, ZfitUnbinnedData) or isinstance(frame2, ZfitUnbinnedData): + return frame1.space + + if not isinstance(frame1, (Mapping, pd.DataFrame)) or not isinstance( + frame2, (Mapping, pd.DataFrame) + ): + raise ValueError( + "frame1 and frame2 have to be either a mapping or a pandas DataFrame, or a zfit Data object. They are currently of type: ", + type(frame1), + type(frame2), + ) + for ob in frame2: + minimum = np.min([np.min(frame1[ob]), np.min(frame2[ob])]) + maximum = np.max([np.max(frame1[ob]), np.max(frame2[ob])]) + dist = maximum - minimum + buffer = bufferfactor * dist + obs.append( + zfit.Space( + ob, + limits=( + minimum - buffer, + maximum + buffer, + ), + ) + ) + obsall = zfit.dimension.combine_spaces(*obs) + return obsall diff --git a/src/zfit_pwa/compwa/wrapper.py b/src/zfit_pwa/compwa/wrapper.py new file mode 100644 index 0000000..1b247ef --- /dev/null +++ b/src/zfit_pwa/compwa/wrapper.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +import types + +import numpy as np +import zfit # suppress tf warnings +import zfit.z.numpy as znp +from zfit import supports, z + +from zfit_pwa.compwa.variables import obs_from_frame + + +def patched_call(self, data) -> np.ndarray: + extended_data = {**self.__parameters, **data} # type: ignore[arg-type] + return self.__function(extended_data) # type: ignore[arg-type] + + +class ComPWAPDF(zfit.pdf.BasePDF): + def __init__( + self, intensity, norm, obs=None, params=None, extended=None, name="ComPWA" + ): + """ComPWA intensity normalized over the *norm* dataset.""" + if params is None: + params = { + name: zfit.param.convert_to_parameter( + val, name=name, prefer_constant=False + ) + for name, val in intensity.parameters.items() + } + if obs is None: + obs = obs_from_frame(norm) + intensity.__call__ = types.MethodType(patched_call, intensity) + super().__init__(obs, params=params, name=name, extended=extended) + self.intensity = intensity + norm = {ob: znp.array(ar) for ob, ar in zip(self.obs, z.unstack_x(norm))} + self.norm_sample = norm + + @supports(norm=True) + def _pdf(self, x, norm): + data = {ob: znp.array(ar) for ob, ar in zip(self.obs, z.unstack_x(x))} + params = {p.name: znp.array(p) for p in self.params.values()} + data |= params + + unnormalized_pdf = self._jitted_unnormalized_pdf(data) + + if norm is False: + return unnormalized_pdf + else: + norm_sample = self.norm_sample | params + return unnormalized_pdf / self._jitted_normalization(norm_sample) + + @z.function(wraps="tensorwaves") + def _jitted_unnormalized_pdf(self, data): + unnormalized_pdf = self.intensity(data) + + return unnormalized_pdf + + @z.function(wraps="tensorwaves") + def _jitted_normalization(self, norm_sample): + return znp.mean(self._jitted_unnormalized_pdf(norm_sample)) diff --git a/tests/compwa/test_wrapper.py b/tests/compwa/test_wrapper.py new file mode 100644 index 0000000..91ccb97 --- /dev/null +++ b/tests/compwa/test_wrapper.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import numpy as np +import pandas as pd + +from zfit_pwa.compwa.wrapper import ComPWAPDF + + +# @pytest.fixture() +def create_amplitude(): + import qrules + + reaction = qrules.generate_transitions( + initial_state=("J/psi(1S)", [-1, +1]), + final_state=["gamma", "pi0", "pi0"], + allowed_intermediate_particles=["f(0)"], + allowed_interaction_types=["strong", "EM"], + formalism="helicity", + ) + + import ampform + from ampform.dynamics.builder import ( + create_non_dynamic_with_ff, + create_relativistic_breit_wigner_with_ff, + ) + + model_builder = ampform.get_builder(reaction) + model_builder.scalar_initial_state_mass = True + model_builder.stable_final_state_ids = [0, 1, 2] + model_builder.set_dynamics("J/psi(1S)", create_non_dynamic_with_ff) + for name in reaction.get_intermediate_particles().names: + model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff) + model = model_builder.formulate() + + from tensorwaves.data import TFPhaseSpaceGenerator, TFUniformRealNumberGenerator + + rng = TFUniformRealNumberGenerator(seed=0) + phsp_generator = TFPhaseSpaceGenerator( + initial_state_mass=reaction.initial_state[-1].mass, + final_state_masses={i: p.mass for i, p in reaction.final_state.items()}, + ) + phsp_momenta = phsp_generator.generate(100_000, rng) + + unfolded_expression = model.expression.doit() + + return model, reaction + + +def test_wrapper_simple(): + import zfit + + model, reaction = create_amplitude() + + from tensorwaves.function.sympy import create_parametrized_function + + unfolded_expression = model.expression.doit() + intensity_func = create_parametrized_function( + expression=unfolded_expression, + parameters=model.parameter_defaults, + backend="tensorflow", + ) + from tensorwaves.data import SympyDataTransformer + + helicity_transformer = SympyDataTransformer.from_sympy( + model.kinematic_variables, backend="numpy" + ) + from tensorwaves.data import ( + IntensityDistributionGenerator, + TFPhaseSpaceGenerator, + TFUniformRealNumberGenerator, + TFWeightedPhaseSpaceGenerator, + ) + + rng = TFUniformRealNumberGenerator(seed=0) + phsp_generator = TFPhaseSpaceGenerator( + initial_state_mass=reaction.initial_state[-1].mass, + final_state_masses={i: p.mass for i, p in reaction.final_state.items()}, + ) + phsp_momenta = phsp_generator.generate(100_000, rng) + + weighted_phsp_generator = TFWeightedPhaseSpaceGenerator( + initial_state_mass=reaction.initial_state[-1].mass, + final_state_masses={i: p.mass for i, p in reaction.final_state.items()}, + ) + data_generator = IntensityDistributionGenerator( + domain_generator=weighted_phsp_generator, + function=intensity_func, + domain_transformer=helicity_transformer, + ) + data_momenta = data_generator.generate(10_000, rng) + + phsp = helicity_transformer(phsp_momenta) + data = helicity_transformer(data_momenta) + data_frame = pd.DataFrame(data) + phsp_frame = pd.DataFrame(phsp) + initial_parameters = { + R"C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}}": ( + 1.0 + ), + "m_{f_{0}(500)}": 0.4, + "m_{f_{0}(980)}": 0.88, + "m_{f_{0}(1370)}": 1.22, + "m_{f_{0}(1500)}": 1.45, + "m_{f_{0}(1710)}": 1.83, + R"\Gamma_{f_{0}(500)}": 0.3, + R"\Gamma_{f_{0}(980)}": 0.1, + R"\Gamma_{f_{0}(1710)}": 0.3, + } + + # data conversion + # phsp_zfit = zfit.Data.from_pandas(phsp_frame) + # data_zfit = zfit.Data.from_pandas(data_frame) + data_frame = data_frame.astype(np.float64) + phsp_frame = phsp_frame.astype(np.float64) + intensity = intensity_func + + pdf = ComPWAPDF( + intensity=intensity, + norm=phsp_frame, + ) + # for p in pdf.get_params(): # returns the free params automatically + # p.set_value(p + np.random.normal(0, 0.01)) + zfit.run.set_graph_mode(False) + loss = zfit.loss.UnbinnedNLL(pdf, data_frame) + + # ok, here I was caught up playing around :) Minuit seems to perform the best though + minimizer = zfit.minimize.Minuit(verbosity=7, gradient=True) + # minimizer = zfit.minimize.ScipyLBFGSBV1(verbosity=8) + # minimizer = zfit.minimize.ScipyTrustKrylovV1(verbosity=8) + # minimizer = zfit.minimize.NLoptMMAV1(verbosity=9) + # minimizer = zfit.minimize.IpyoptV1(verbosity=8) + result = minimizer.minimize(loss) + print(result) + result.hesse() + print(result) + assert result.valid