feat: Experimental implementation of custom modifiers

pyproject.toml

@@ -82,8 +82,9 @@ contrib = [
     "matplotlib>=3.0.0",
     "requests>=2.22.0",
 ]
+experimental = ["numexpr>=2.8.0"]
 backends = ["pyhf[tensorflow,torch,jax,minuit]"]
-all = ["pyhf[backends,xmlio,contrib,shellcomplete]"]
+all = ["pyhf[backends,xmlio,contrib,experimental,shellcomplete]"]
 
 # Developer extras
 test = [
@@ -105,7 +106,7 @@ test = [
     "pytest-socket>=0.2.0",  # c.f. PR #1917
 ]
 docs = [
-    "pyhf[xmlio,contrib]",
+    "pyhf[xmlio,contrib,experimental]",
     "sphinx>=7.0.0",  # c.f. https://github.com/scikit-hep/pyhf/pull/2271
     "sphinxcontrib-bibtex~=2.1",
     "sphinx-click",
@@ -244,6 +245,7 @@ warn_unreachable = true
 module = [
   'jax.*',
   'matplotlib.*',
+  'numexpr.*',
   'scipy.*',
   'tensorflow.*',
   'tensorflow_probability.*',
@@ -262,6 +264,7 @@ module = [
   'pyhf.cli.*',
   'pyhf.modifiers.*',
   'pyhf.exceptions.*',
+  'pyhf.experimental.*',
   'pyhf.parameters.*',
   'pyhf.schema.*',
   'pyhf.writexml',

src/pyhf/experimental/__init__.py

@@ -0,0 +1,5 @@
+"""
+Experimental features for pyhf.
+
+Modules in experimental may rapidly change with API breaking changes.
+"""

src/pyhf/experimental/modifiers.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Sequence
+
+import pyhf
+from pyhf import events, get_backend
+from pyhf.parameters import ParamViewer
+
+log = logging.getLogger(__name__)
+
+__all__ = ["add_custom_modifier"]
+
+
+def __dir__():
+    return __all__
+
+
+try:
+    import numexpr as ne
+except ModuleNotFoundError:
+    log.error(
+        "\nInstallation of the experimental extra is required to use pyhf.experimental.modifiers"
+        + "\nPlease install with: python -m pip install 'pyhf[experimental]'\n",
+        exc_info=True,
+    )
+    raise
+
+
+class BaseApplier:
+    ...
+
+
+class BaseBuilder:
+    ...
+
+
+def _allocate_new_param(
+    p: dict[str, Sequence[float]]
+) -> dict[str, str | bool | int | Sequence[float]]:
+    return {
+        "paramset_type": "unconstrained",
+        "n_parameters": 1,
+        "is_shared": True,
+        "inits": p["inits"],
+        "bounds": p["bounds"],
+        "is_scalar": True,
+        "fixed": False,
+    }
+
+
+def make_func(expression: str, deps: list[str]) -> Callable[[Sequence[float]], Any]:
+    def func(d: Sequence[float]) -> Any:
+        return ne.evaluate(expression, local_dict=dict(zip(deps, d)))
+
+    return func
+
+
+def make_builder(
+    func_name: str, deps: list[str], new_params: dict[str, dict[str, Sequence[float]]]
+) -> BaseBuilder:
+    class _builder(BaseBuilder):
+        is_shared = False
+
+        def __init__(self, config):
+            self.builder_data = {"funcs": {}}
+            self.config = config
+
+        def collect(self, thismod, nom):
+            maskval = bool(thismod)
+            mask = [maskval] * len(nom)
+            return {"mask": mask}
+
+        def append(self, key, channel, sample, thismod, defined_samp):
+            self.builder_data.setdefault(key, {}).setdefault(sample, {}).setdefault(
+                "data", {"mask": []}
+            )
+            nom = (
+                defined_samp["data"]
+                if defined_samp
+                else [0.0] * self.config.channel_nbins[channel]
+            )
+            moddata = self.collect(thismod, nom)
+            self.builder_data[key][sample]["data"]["mask"] += moddata["mask"]
+            if thismod:
+                if thismod["name"] != func_name:
+                    print(thismod)
+                    self.builder_data["funcs"].setdefault(
+                        thismod["name"], thismod["data"]["expr"]
+                    )
+                self.required_parsets = {
+                    k: [_allocate_new_param(v)] for k, v in new_params.items()
+                }
+
+        def finalize(self):
+            return self.builder_data
+
+    return _builder
+
+
+def make_applier(
+    func_name: str, deps: list[str], new_params: dict[str, dict[str, Sequence[float]]]
+) -> BaseApplier:
+    class _applier(BaseApplier):
+        name = func_name
+        op_code = "multiplication"
+
+        def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None):
+            self.funcs = [make_func(v, deps) for v in builder_data["funcs"].values()]
+
+            self.batch_size = batch_size
+            pars_for_applier = deps
+            _modnames = [f"{mtype}/{m}" for m, mtype in modifiers]
+
+            parfield_shape = (
+                (self.batch_size, pdfconfig.npars)
+                if self.batch_size
+                else (pdfconfig.npars,)
+            )
+            self.param_viewer = ParamViewer(
+                parfield_shape, pdfconfig.par_map, pars_for_applier
+            )
+            self._custom_mod_mask = [
+                [[builder_data[modname][s]["data"]["mask"]] for s in pdfconfig.samples]
+                for modname in _modnames
+            ]
+            self._precompute()
+            events.subscribe("tensorlib_changed")(self._precompute)
+
+        def _precompute(self):
+            tensorlib, _ = get_backend()
+            if not self.param_viewer.index_selection:
+                return
+            self.custom_mod_mask = tensorlib.tile(
+                tensorlib.astensor(self._custom_mod_mask),
+                (1, 1, self.batch_size or 1, 1),
+            )
+            self.custom_mod_mask_bool = tensorlib.astensor(
+                self.custom_mod_mask, dtype="bool"
+            )
+            self.custom_mod_default = tensorlib.ones(self.custom_mod_mask.shape)
+
+        def apply(self, pars):
+            """
+            Returns:
+                modification tensor: Shape (n_modifiers, n_global_samples, n_alphas, n_global_bin)
+            """
+            if not self.param_viewer.index_selection:
+                return
+            tensorlib, _ = get_backend()
+            if self.batch_size is None:
+                deps = self.param_viewer.get(pars)
+                print("deps", deps.shape)
+                results = tensorlib.astensor([f(deps) for f in self.funcs])
+                results = tensorlib.einsum(
+                    "msab,m->msab", self.custom_mod_mask, results
+                )
+            else:
+                deps = self.param_viewer.get(pars)
+                print("deps", deps.shape)
+                results = tensorlib.astensor([f(deps) for f in self.funcs])
+                results = tensorlib.einsum(
+                    "msab,ma->msab", self.custom_mod_mask, results
+                )
+            results = tensorlib.where(
+                self.custom_mod_mask_bool, results, self.custom_mod_default
+            )
+            return results
+
+    return _applier
+
+
+def add_custom_modifier(
+    func_name: str, deps: list[str], new_params: dict[str, dict[str, Sequence[float]]]
+) -> dict[str, tuple[BaseBuilder, BaseApplier]]:
+    _builder = make_builder(func_name, deps, new_params)
+    _applier = make_applier(func_name, deps, new_params)
+
+    modifier_set = {_applier.name: (_builder, _applier)}
+    modifier_set.update(**pyhf.modifiers.histfactory_set)
+    return modifier_set

tests/test_experimental.py

@@ -0,0 +1,94 @@
+import sys
+from importlib import reload
+from unittest import mock
+
+import pytest
+
+import pyhf
+import pyhf.experimental.modifiers
+
+
+def test_missing_experimental_extra():
+    """
+    Verify ModuleNotFoundError if dependencies required of the experimental
+    extra are not installed.
+    """
+    with mock.patch.dict(sys.modules):
+        sys.modules["numexpr"] = None
+        with pytest.raises(ModuleNotFoundError):
+            reload(sys.modules["pyhf.experimental.modifiers"])
+
+
+def test_add_custom_modifier(backend):
+    tensorlib, _ = backend
+
+    new_params = {
+        "m1": {"inits": (1.0,), "bounds": ((-5.0, 5.0),)},
+        "m2": {"inits": (1.0,), "bounds": ((-5.0, 5.0),)},
+    }
+
+    expanded_pyhf = pyhf.experimental.modifiers.add_custom_modifier(
+        "customfunc", ["m1", "m2"], new_params
+    )
+    model = pyhf.Model(
+        {
+            "channels": [
+                {
+                    "name": "singlechannel",
+                    "samples": [
+                        {
+                            "name": "signal",
+                            "data": [10] * 20,
+                            "modifiers": [
+                                {
+                                    "name": "f2",
+                                    "type": "customfunc",
+                                    "data": {"expr": "m1"},
+                                },
+                            ],
+                        },
+                        {
+                            "name": "background",
+                            "data": [100] * 20,
+                            "modifiers": [
+                                {
+                                    "name": "f1",
+                                    "type": "customfunc",
+                                    "data": {"expr": "m1+(m2**2)"},
+                                },
+                            ],
+                        },
+                    ],
+                }
+            ]
+        },
+        modifier_set=expanded_pyhf,
+        poi_name="m1",
+        validate=False,
+        batch_size=1,
+    )
+
+    assert tensorlib.tolist(model.expected_actualdata([[1.0, 2.0]])) == [
+        [
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+            510.0,
+        ]
+    ]