Skip to content

Commit f4bf46f

Browse files
lorenzenniolorenzennio
committed
personal use case of custom modifier
1 parent bb302c2 commit f4bf46f

File tree

1 file changed

+70
-149
lines changed

1 file changed

+70
-149
lines changed

‎src/pyhf/experimental/modifiers.py‎

Lines changed: 70 additions & 149 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
from pyhf import events, get_backend
88
from pyhf.parameters import ParamViewer
99

10+
import numpy as np
11+
1012
log = logging.getLogger(__name__)
1113

1214
__all__ = ["add_custom_modifier"]
@@ -35,81 +37,83 @@ class BaseBuilder:
3537
...
3638

3739

38-
def _allocate_new_param(
39-
p: dict[str, Sequence[float]]
40-
) -> dict[str, str | bool | int | Sequence[float]]:
41-
return {
42-
"paramset_type": "unconstrained",
43-
"n_parameters": 1,
44-
"is_shared": True,
45-
"inits": p["inits"],
46-
"bounds": p["bounds"],
47-
"is_scalar": True,
48-
"fixed": False,
49-
}
50-
51-
52-
def make_func(expression: str, deps: list[str]) -> Callable[[Sequence[float]], Any]:
53-
def func(d: Sequence[float]) -> Any:
54-
return ne.evaluate(expression, local_dict=dict(zip(deps, d)))
55-
56-
return func
40+
def add(funcname, par_names, newparams, input_set=None, namespace=None):
41+
namespace = namespace or {}
5742

43+
def make_func(
44+
expression: str, namespace=namespace
45+
) -> Callable[[Sequence[float]], Any]:
46+
def func(deps: Sequence[float]) -> Any:
47+
if expression in namespace:
48+
parvals = dict(zip(par_names, deps))
49+
return namespace[expression](parvals)()
50+
return ne.evaluate(
51+
expression, local_dict=dict(zip(par_names, deps), **namespace)
52+
)
5853

59-
def make_builder(
60-
func_name: str, deps: list[str], new_params: dict[str, dict[str, Sequence[float]]]
61-
) -> BaseBuilder:
62-
class _builder(BaseBuilder):
63-
is_shared = False
54+
return func
55+
56+
def _allocate_new_param(p):
57+
param_dict = {
58+
'paramset_type': p['paramset_type']
59+
if 'paramset_type' in p.keys()
60+
else 'unconstrained',
61+
'n_parameters': 1,
62+
'is_shared': True,
63+
'inits': p['inits'],
64+
'bounds': p['bounds'],
65+
'is_scalar': True,
66+
'fixed': False,
67+
'auxdata': p['auxdata'] if 'auxdata' in p.keys() else (0.0,),
68+
}
69+
return param_dict
70+
71+
class _builder:
72+
is_shared = True
6473

6574
def __init__(self, config):
66-
self.builder_data = {"funcs": {}}
75+
self.builder_data = {'funcs': {}}
6776
self.config = config
77+
self.required_parsets = {}
6878

6979
def collect(self, thismod, nom):
70-
maskval = bool(thismod)
80+
maskval = True if thismod else False
7181
mask = [maskval] * len(nom)
72-
return {"mask": mask}
82+
return {'mask': mask}
7383

74-
def append(self, key, channel, sample, thismod, defined_sample):
84+
def append(self, key, channel, sample, thismod, defined_samp):
7585
self.builder_data.setdefault(key, {}).setdefault(sample, {}).setdefault(
76-
"data", {"mask": []}
86+
'data', {'mask': []}
7787
)
7888
nom = (
79-
defined_sample["data"]
80-
if defined_sample
89+
defined_samp['data']
90+
if defined_samp
8191
else [0.0] * self.config.channel_nbins[channel]
8292
)
83-
mod_data = self.collect(thismod, nom)
84-
self.builder_data[key][sample]["data"]["mask"] += mod_data["mask"]
93+
moddata = self.collect(thismod, nom)
94+
self.builder_data[key][sample]['data']['mask'] += moddata['mask']
8595
if thismod:
86-
if thismod["name"] != func_name:
87-
self.builder_data["funcs"].setdefault(
88-
thismod["name"], thismod["data"]["expr"]
96+
if thismod['name'] != funcname:
97+
self.builder_data['funcs'].setdefault(
98+
thismod['name'], thismod['data']['expr']
8999
)
90100
self.required_parsets = {
91-
k: [_allocate_new_param(v)] for k, v in new_params.items()
101+
k: [_allocate_new_param(v)] for k, v in newparams.items()
92102
}
93103

94104
def finalize(self):
95105
return self.builder_data
96106

97-
return _builder
98-
99-
100-
def make_applier(
101-
func_name: str, deps: list[str], new_params: dict[str, dict[str, Sequence[float]]]
102-
) -> BaseApplier:
103-
class _applier(BaseApplier):
104-
name = func_name
105-
op_code = "multiplication"
107+
class _applier:
108+
name = funcname
109+
op_code = 'addition'
106110

107111
def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None):
108-
self.funcs = [make_func(v, deps) for v in builder_data["funcs"].values()]
112+
self.funcs = [make_func(f) for f in builder_data['funcs'].values()]
109113

110114
self.batch_size = batch_size
111-
pars_for_applier = deps
112-
_mod_names = [f"{mtype}/{m}" for m, mtype in modifiers]
115+
pars_for_applier = par_names
116+
_modnames = [f'{mtype}/{m}' for m, mtype in modifiers]
113117

114118
parfield_shape = (
115119
(self.batch_size, pdfconfig.npars)
@@ -119,25 +123,25 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None):
119123
self.param_viewer = ParamViewer(
120124
parfield_shape, pdfconfig.par_map, pars_for_applier
121125
)
122-
self._custom_mod_mask = [
123-
[[builder_data[mod_name][s]["data"]["mask"]] for s in pdfconfig.samples]
124-
for mod_name in _mod_names
126+
self._custommod_mask = [
127+
[[builder_data[modname][s]['data']['mask']] for s in pdfconfig.samples]
128+
for modname in _modnames
125129
]
126130
self._precompute()
127-
events.subscribe("tensorlib_changed")(self._precompute)
131+
events.subscribe('tensorlib_changed')(self._precompute)
128132

129133
def _precompute(self):
130134
tensorlib, _ = get_backend()
131135
if not self.param_viewer.index_selection:
132136
return
133-
self.custom_mod_mask = tensorlib.tile(
134-
tensorlib.astensor(self._custom_mod_mask),
137+
self.custommod_mask = tensorlib.tile(
138+
tensorlib.astensor(self._custommod_mask),
135139
(1, 1, self.batch_size or 1, 1),
136140
)
137-
self.custom_mod_mask_bool = tensorlib.astensor(
138-
self.custom_mod_mask, dtype="bool"
141+
self.custommod_mask_bool = tensorlib.astensor(
142+
self.custommod_mask, dtype="bool"
139143
)
140-
self.custom_mod_default = tensorlib.ones(self.custom_mod_mask.shape)
144+
self.custommod_default = tensorlib.zeros(self.custommod_mask.shape)
141145

142146
def apply(self, pars):
143147
"""
@@ -147,100 +151,17 @@ def apply(self, pars):
147151
if not self.param_viewer.index_selection:
148152
return
149153
tensorlib, _ = get_backend()
150-
if self.batch_size is None:
151-
deps = self.param_viewer.get(pars)
152-
results = tensorlib.astensor([f(deps) for f in self.funcs])
153-
results = tensorlib.einsum(
154-
"msab,m->msab", self.custom_mod_mask, results
155-
)
156-
else:
157-
deps = self.param_viewer.get(pars)
158-
results = tensorlib.astensor([f(deps) for f in self.funcs])
159-
results = tensorlib.einsum(
160-
"msab,ma->msab", self.custom_mod_mask, results
161-
)
154+
deps = self.param_viewer.get(pars)
155+
out = tensorlib.astensor([f(deps) for f in self.funcs])
156+
results = np.zeros_like(self.custommod_mask)
157+
np.place(results, self.custommod_mask, out)
162158
results = tensorlib.where(
163-
self.custom_mod_mask_bool, results, self.custom_mod_default
159+
self.custommod_mask_bool, results, self.custommod_default
164160
)
165161
return results
166162

167-
return _applier
168-
169-
170-
def add_custom_modifier(
171-
func_name: str, deps: list[str], new_params: dict[str, dict[str, Sequence[float]]]
172-
) -> dict[str, tuple[BaseBuilder, BaseApplier]]:
173-
r"""
174-
Add a custom modifier type with the modifier data defined through a custom
175-
numexpr string expression.
176-
177-
Example:
178-
179-
>>> import pyhf
180-
>>> import pyhf.experimental.modifiers
181-
>>> pyhf.set_backend("numpy")
182-
>>> new_params = {
183-
... "m1": {"inits": (1.0,), "bounds": ((-5.0, 5.0),)},
184-
... "m2": {"inits": (1.0,), "bounds": ((-5.0, 5.0),)},
185-
... }
186-
>>> expanded_pyhf = pyhf.experimental.modifiers.add_custom_modifier(
187-
... "custom", ["m1", "m2"], new_params
188-
... )
189-
>>> model = pyhf.Model(
190-
... {
191-
... "channels": [
192-
... {
193-
... "name": "singlechannel",
194-
... "samples": [
195-
... {
196-
... "name": "signal",
197-
... "data": [10, 20],
198-
... "modifiers": [
199-
... {
200-
... "name": "f2",
201-
... "type": "custom",
202-
... "data": {"expr": "m1"},
203-
... },
204-
... ],
205-
... },
206-
... {
207-
... "name": "background",
208-
... "data": [100, 150],
209-
... "modifiers": [
210-
... {
211-
... "name": "f1",
212-
... "type": "custom",
213-
... "data": {"expr": "m1+(m2**2)"},
214-
... },
215-
... ],
216-
... },
217-
... ],
218-
... }
219-
... ]
220-
... },
221-
... modifier_set=expanded_pyhf,
222-
... poi_name="m1",
223-
... validate=False,
224-
... batch_size=1,
225-
... )
226-
>>> model.config.modifiers
227-
[('f1', 'custom'), ('f2', 'custom')]
228-
229-
Args:
230-
func_name (:obj:`str`): The name of the custom modifier type.
231-
deps (:obj:`list`): The names of the new parameters of the modifier
232-
function.
233-
new_params (:obj:`dict`): The new parameters.
234-
235-
Returns:
236-
:obj:`dict`: The updated ``pyhf.modifiers.histfactory_set`` with the added
237-
custom modifier type.
238-
239-
.. versionadded:: 0.8.0
240-
"""
241-
_builder = make_builder(func_name, deps, new_params)
242-
_applier = make_applier(func_name, deps, new_params)
243-
244163
modifier_set = {_applier.name: (_builder, _applier)}
245-
modifier_set.update(**pyhf.modifiers.histfactory_set)
164+
modifier_set.update(
165+
**(input_set if input_set is not None else pyhf.modifiers.histfactory_set)
166+
)
246167
return modifier_set

0 commit comments

Comments
 (0)