From 8ab4ab9cbb0caa3c8ebdb54a94263ee960864a00 Mon Sep 17 00:00:00 2001 From: Lukas Heinrich Date: Sat, 30 Oct 2021 14:43:40 +0200 Subject: [PATCH 01/12] diffable model! --- src/pyhf/constraints.py | 13 +++++++++---- src/pyhf/modifiers/histosys.py | 8 +++++++- src/pyhf/modifiers/shapefactor.py | 4 ++-- src/pyhf/modifiers/shapesys.py | 4 ++-- src/pyhf/modifiers/staterror.py | 4 ++-- 5 files changed, 22 insertions(+), 11 deletions(-) diff --git a/src/pyhf/constraints.py b/src/pyhf/constraints.py index 30ef835e79..35c981f09c 100644 --- a/src/pyhf/constraints.py +++ b/src/pyhf/constraints.py @@ -10,7 +10,6 @@ def __dir__(): return __all__ - class gaussian_constraint_combined: def __init__(self, pdfconfig, batch_size=None): default_backend = pyhf.default_backend @@ -45,7 +44,9 @@ def __init__(self, pdfconfig, batch_size=None): if not parset.pdf_type == 'normal': continue - normal_constraint_data.append(thisauxdata) + normal_constraint_data.append( + default_backend.astensor(thisauxdata) + ) # many constraints are defined on a unit gaussian # but we reserved the possibility that a paramset @@ -53,9 +54,13 @@ def __init__(self, pdfconfig, batch_size=None): # by the paramset associated to staterror modifiers. # Such parsets define a 'sigmas' attribute try: - normal_constraint_sigmas.append(parset.sigmas) + normal_constraint_sigmas.append( + default_backend.astensor(parset.sigmas) + ) except AttributeError: - normal_constraint_sigmas.append([1.0] * len(thisauxdata)) + normal_constraint_sigmas.append( + default_backend.astensor([1.0] * len(thisauxdata)) + ) self._normal_data = None self._sigmas = None diff --git a/src/pyhf/modifiers/histosys.py b/src/pyhf/modifiers/histosys.py index d23e6c1ff3..4c4c0d2b81 100644 --- a/src/pyhf/modifiers/histosys.py +++ b/src/pyhf/modifiers/histosys.py @@ -31,11 +31,17 @@ def __init__(self, config): self.required_parsets = {} def collect(self, thismod, nom): + default_backend = pyhf.default_backend lo_data = thismod['data']['lo_data'] if thismod else nom hi_data = thismod['data']['hi_data'] if thismod else nom maskval = bool(thismod) mask = [maskval] * len(nom) - return {'lo_data': lo_data, 'hi_data': hi_data, 'mask': mask, 'nom_data': nom} + return { + 'lo_data': default_backend.astensor(lo_data), + 'hi_data': default_backend.astensor(hi_data), + 'mask': default_backend.astensor(mask, dtype = 'bool'), + 'nom_data': default_backend.astensor(nom) + } def append(self, key, channel, sample, thismod, defined_samp): self.builder_data.setdefault(key, {}).setdefault(sample, {}).setdefault( diff --git a/src/pyhf/modifiers/shapefactor.py b/src/pyhf/modifiers/shapefactor.py index 6c65c2e8c9..9359faa70e 100644 --- a/src/pyhf/modifiers/shapefactor.py +++ b/src/pyhf/modifiers/shapefactor.py @@ -142,9 +142,9 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] - global_concatenated_bin_indices = [ + global_concatenated_bin_indices = default_backend.astensor([ [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ] + ]) self._access_field = default_backend.tile( global_concatenated_bin_indices, diff --git a/src/pyhf/modifiers/shapesys.py b/src/pyhf/modifiers/shapesys.py index 5b5466857f..2341ff3a17 100644 --- a/src/pyhf/modifiers/shapesys.py +++ b/src/pyhf/modifiers/shapesys.py @@ -117,9 +117,9 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] ) - global_concatenated_bin_indices = [ + global_concatenated_bin_indices = default_backend.astensor([ [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ] + ]) self._access_field = default_backend.tile( global_concatenated_bin_indices, diff --git a/src/pyhf/modifiers/staterror.py b/src/pyhf/modifiers/staterror.py index e028b8f5db..2b3be5bdfc 100644 --- a/src/pyhf/modifiers/staterror.py +++ b/src/pyhf/modifiers/staterror.py @@ -144,9 +144,9 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] ) - global_concatenated_bin_indices = [ + global_concatenated_bin_indices = default_backend.astensor([ [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ] + ]) self._access_field = default_backend.tile( global_concatenated_bin_indices, From 471dcf0c0f97bc32131f006fa37b5437583fddac Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 30 Oct 2021 16:33:04 +0000 Subject: [PATCH 02/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/pyhf/constraints.py | 9 +++------ src/pyhf/modifiers/histosys.py | 4 ++-- src/pyhf/modifiers/shapefactor.py | 14 +++++++++++--- src/pyhf/modifiers/shapesys.py | 14 +++++++++++--- src/pyhf/modifiers/staterror.py | 14 +++++++++++--- 5 files changed, 38 insertions(+), 17 deletions(-) diff --git a/src/pyhf/constraints.py b/src/pyhf/constraints.py index 35c981f09c..73b45b05c5 100644 --- a/src/pyhf/constraints.py +++ b/src/pyhf/constraints.py @@ -10,6 +10,7 @@ def __dir__(): return __all__ + class gaussian_constraint_combined: def __init__(self, pdfconfig, batch_size=None): default_backend = pyhf.default_backend @@ -44,9 +45,7 @@ def __init__(self, pdfconfig, batch_size=None): if not parset.pdf_type == 'normal': continue - normal_constraint_data.append( - default_backend.astensor(thisauxdata) - ) + normal_constraint_data.append(default_backend.astensor(thisauxdata)) # many constraints are defined on a unit gaussian # but we reserved the possibility that a paramset @@ -54,9 +53,7 @@ def __init__(self, pdfconfig, batch_size=None): # by the paramset associated to staterror modifiers. # Such parsets define a 'sigmas' attribute try: - normal_constraint_sigmas.append( - default_backend.astensor(parset.sigmas) - ) + normal_constraint_sigmas.append(default_backend.astensor(parset.sigmas)) except AttributeError: normal_constraint_sigmas.append( default_backend.astensor([1.0] * len(thisauxdata)) diff --git a/src/pyhf/modifiers/histosys.py b/src/pyhf/modifiers/histosys.py index 4c4c0d2b81..31c059c68d 100644 --- a/src/pyhf/modifiers/histosys.py +++ b/src/pyhf/modifiers/histosys.py @@ -39,8 +39,8 @@ def collect(self, thismod, nom): return { 'lo_data': default_backend.astensor(lo_data), 'hi_data': default_backend.astensor(hi_data), - 'mask': default_backend.astensor(mask, dtype = 'bool'), - 'nom_data': default_backend.astensor(nom) + 'mask': default_backend.astensor(mask, dtype='bool'), + 'nom_data': default_backend.astensor(nom), } def append(self, key, channel, sample, thismod, defined_samp): diff --git a/src/pyhf/modifiers/shapefactor.py b/src/pyhf/modifiers/shapefactor.py index 9359faa70e..b0a0abb003 100644 --- a/src/pyhf/modifiers/shapefactor.py +++ b/src/pyhf/modifiers/shapefactor.py @@ -142,9 +142,17 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] - global_concatenated_bin_indices = default_backend.astensor([ - [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ]) + global_concatenated_bin_indices = default_backend.astensor( + [ + [ + [ + j + for c in pdfconfig.channels + for j in range(pdfconfig.channel_nbins[c]) + ] + ] + ] + ) self._access_field = default_backend.tile( global_concatenated_bin_indices, diff --git a/src/pyhf/modifiers/shapesys.py b/src/pyhf/modifiers/shapesys.py index 2341ff3a17..078db7253b 100644 --- a/src/pyhf/modifiers/shapesys.py +++ b/src/pyhf/modifiers/shapesys.py @@ -117,9 +117,17 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] ) - global_concatenated_bin_indices = default_backend.astensor([ - [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ]) + global_concatenated_bin_indices = default_backend.astensor( + [ + [ + [ + j + for c in pdfconfig.channels + for j in range(pdfconfig.channel_nbins[c]) + ] + ] + ] + ) self._access_field = default_backend.tile( global_concatenated_bin_indices, diff --git a/src/pyhf/modifiers/staterror.py b/src/pyhf/modifiers/staterror.py index 2b3be5bdfc..377610010b 100644 --- a/src/pyhf/modifiers/staterror.py +++ b/src/pyhf/modifiers/staterror.py @@ -144,9 +144,17 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] ) - global_concatenated_bin_indices = default_backend.astensor([ - [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ]) + global_concatenated_bin_indices = default_backend.astensor( + [ + [ + [ + j + for c in pdfconfig.channels + for j in range(pdfconfig.channel_nbins[c]) + ] + ] + ] + ) self._access_field = default_backend.tile( global_concatenated_bin_indices, From 84d40716f8f50972c0556b6c1d860adb1ffb7427 Mon Sep 17 00:00:00 2001 From: Lukas Heinrich Date: Sat, 30 Oct 2021 14:43:40 +0200 Subject: [PATCH 03/12] diffable model! --- src/pyhf/constraints.py | 13 +++++++++---- src/pyhf/modifiers/histosys.py | 8 +++++++- src/pyhf/modifiers/shapefactor.py | 4 ++-- src/pyhf/modifiers/shapesys.py | 4 ++-- src/pyhf/modifiers/staterror.py | 4 ++-- 5 files changed, 22 insertions(+), 11 deletions(-) diff --git a/src/pyhf/constraints.py b/src/pyhf/constraints.py index 30ef835e79..35c981f09c 100644 --- a/src/pyhf/constraints.py +++ b/src/pyhf/constraints.py @@ -10,7 +10,6 @@ def __dir__(): return __all__ - class gaussian_constraint_combined: def __init__(self, pdfconfig, batch_size=None): default_backend = pyhf.default_backend @@ -45,7 +44,9 @@ def __init__(self, pdfconfig, batch_size=None): if not parset.pdf_type == 'normal': continue - normal_constraint_data.append(thisauxdata) + normal_constraint_data.append( + default_backend.astensor(thisauxdata) + ) # many constraints are defined on a unit gaussian # but we reserved the possibility that a paramset @@ -53,9 +54,13 @@ def __init__(self, pdfconfig, batch_size=None): # by the paramset associated to staterror modifiers. # Such parsets define a 'sigmas' attribute try: - normal_constraint_sigmas.append(parset.sigmas) + normal_constraint_sigmas.append( + default_backend.astensor(parset.sigmas) + ) except AttributeError: - normal_constraint_sigmas.append([1.0] * len(thisauxdata)) + normal_constraint_sigmas.append( + default_backend.astensor([1.0] * len(thisauxdata)) + ) self._normal_data = None self._sigmas = None diff --git a/src/pyhf/modifiers/histosys.py b/src/pyhf/modifiers/histosys.py index 78b0f20074..c79fc25273 100644 --- a/src/pyhf/modifiers/histosys.py +++ b/src/pyhf/modifiers/histosys.py @@ -31,11 +31,17 @@ def __init__(self, config): self.required_parsets = {} def collect(self, thismod, nom): + default_backend = pyhf.default_backend lo_data = thismod['data']['lo_data'] if thismod else nom hi_data = thismod['data']['hi_data'] if thismod else nom maskval = bool(thismod) mask = [maskval] * len(nom) - return {'lo_data': lo_data, 'hi_data': hi_data, 'mask': mask, 'nom_data': nom} + return { + 'lo_data': default_backend.astensor(lo_data), + 'hi_data': default_backend.astensor(hi_data), + 'mask': default_backend.astensor(mask, dtype = 'bool'), + 'nom_data': default_backend.astensor(nom) + } def append(self, key, channel, sample, thismod, defined_samp): self.builder_data.setdefault(key, {}).setdefault(sample, {}).setdefault( diff --git a/src/pyhf/modifiers/shapefactor.py b/src/pyhf/modifiers/shapefactor.py index 6c65c2e8c9..9359faa70e 100644 --- a/src/pyhf/modifiers/shapefactor.py +++ b/src/pyhf/modifiers/shapefactor.py @@ -142,9 +142,9 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] - global_concatenated_bin_indices = [ + global_concatenated_bin_indices = default_backend.astensor([ [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ] + ]) self._access_field = default_backend.tile( global_concatenated_bin_indices, diff --git a/src/pyhf/modifiers/shapesys.py b/src/pyhf/modifiers/shapesys.py index 3d4d52e2b6..755d6cdebd 100644 --- a/src/pyhf/modifiers/shapesys.py +++ b/src/pyhf/modifiers/shapesys.py @@ -128,9 +128,9 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] ) - global_concatenated_bin_indices = [ + global_concatenated_bin_indices = default_backend.astensor([ [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ] + ]) self._access_field = default_backend.tile( global_concatenated_bin_indices, diff --git a/src/pyhf/modifiers/staterror.py b/src/pyhf/modifiers/staterror.py index 3cb1cd5c4e..3162ac505a 100644 --- a/src/pyhf/modifiers/staterror.py +++ b/src/pyhf/modifiers/staterror.py @@ -157,9 +157,9 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] ) - global_concatenated_bin_indices = [ + global_concatenated_bin_indices = default_backend.astensor([ [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ] + ]) self._access_field = default_backend.tile( global_concatenated_bin_indices, From ef34ab33485cc9225182d4306073367dd76756c6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 30 Oct 2021 16:33:04 +0000 Subject: [PATCH 04/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/pyhf/constraints.py | 9 +++------ src/pyhf/modifiers/histosys.py | 4 ++-- src/pyhf/modifiers/shapefactor.py | 14 +++++++++++--- src/pyhf/modifiers/shapesys.py | 14 +++++++++++--- src/pyhf/modifiers/staterror.py | 14 +++++++++++--- 5 files changed, 38 insertions(+), 17 deletions(-) diff --git a/src/pyhf/constraints.py b/src/pyhf/constraints.py index 35c981f09c..73b45b05c5 100644 --- a/src/pyhf/constraints.py +++ b/src/pyhf/constraints.py @@ -10,6 +10,7 @@ def __dir__(): return __all__ + class gaussian_constraint_combined: def __init__(self, pdfconfig, batch_size=None): default_backend = pyhf.default_backend @@ -44,9 +45,7 @@ def __init__(self, pdfconfig, batch_size=None): if not parset.pdf_type == 'normal': continue - normal_constraint_data.append( - default_backend.astensor(thisauxdata) - ) + normal_constraint_data.append(default_backend.astensor(thisauxdata)) # many constraints are defined on a unit gaussian # but we reserved the possibility that a paramset @@ -54,9 +53,7 @@ def __init__(self, pdfconfig, batch_size=None): # by the paramset associated to staterror modifiers. # Such parsets define a 'sigmas' attribute try: - normal_constraint_sigmas.append( - default_backend.astensor(parset.sigmas) - ) + normal_constraint_sigmas.append(default_backend.astensor(parset.sigmas)) except AttributeError: normal_constraint_sigmas.append( default_backend.astensor([1.0] * len(thisauxdata)) diff --git a/src/pyhf/modifiers/histosys.py b/src/pyhf/modifiers/histosys.py index c79fc25273..c406112c13 100644 --- a/src/pyhf/modifiers/histosys.py +++ b/src/pyhf/modifiers/histosys.py @@ -39,8 +39,8 @@ def collect(self, thismod, nom): return { 'lo_data': default_backend.astensor(lo_data), 'hi_data': default_backend.astensor(hi_data), - 'mask': default_backend.astensor(mask, dtype = 'bool'), - 'nom_data': default_backend.astensor(nom) + 'mask': default_backend.astensor(mask, dtype='bool'), + 'nom_data': default_backend.astensor(nom), } def append(self, key, channel, sample, thismod, defined_samp): diff --git a/src/pyhf/modifiers/shapefactor.py b/src/pyhf/modifiers/shapefactor.py index 9359faa70e..b0a0abb003 100644 --- a/src/pyhf/modifiers/shapefactor.py +++ b/src/pyhf/modifiers/shapefactor.py @@ -142,9 +142,17 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] - global_concatenated_bin_indices = default_backend.astensor([ - [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ]) + global_concatenated_bin_indices = default_backend.astensor( + [ + [ + [ + j + for c in pdfconfig.channels + for j in range(pdfconfig.channel_nbins[c]) + ] + ] + ] + ) self._access_field = default_backend.tile( global_concatenated_bin_indices, diff --git a/src/pyhf/modifiers/shapesys.py b/src/pyhf/modifiers/shapesys.py index 755d6cdebd..10c8de342e 100644 --- a/src/pyhf/modifiers/shapesys.py +++ b/src/pyhf/modifiers/shapesys.py @@ -128,9 +128,17 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] ) - global_concatenated_bin_indices = default_backend.astensor([ - [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ]) + global_concatenated_bin_indices = default_backend.astensor( + [ + [ + [ + j + for c in pdfconfig.channels + for j in range(pdfconfig.channel_nbins[c]) + ] + ] + ] + ) self._access_field = default_backend.tile( global_concatenated_bin_indices, diff --git a/src/pyhf/modifiers/staterror.py b/src/pyhf/modifiers/staterror.py index 3162ac505a..2730c02af8 100644 --- a/src/pyhf/modifiers/staterror.py +++ b/src/pyhf/modifiers/staterror.py @@ -157,9 +157,17 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] ) - global_concatenated_bin_indices = default_backend.astensor([ - [[j for c in pdfconfig.channels for j in range(pdfconfig.channel_nbins[c])]] - ]) + global_concatenated_bin_indices = default_backend.astensor( + [ + [ + [ + j + for c in pdfconfig.channels + for j in range(pdfconfig.channel_nbins[c]) + ] + ] + ] + ) self._access_field = default_backend.tile( global_concatenated_bin_indices, From 3831ebff78960ec418015f7eb2751eaa57a0d7d8 Mon Sep 17 00:00:00 2001 From: Nathan Simpson Date: Sun, 3 Jul 2022 10:04:20 +0200 Subject: [PATCH 05/12] jax backend fixes --- src/pyhf/tensor/jax_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pyhf/tensor/jax_backend.py b/src/pyhf/tensor/jax_backend.py index 5e4a65bc80..ee15a5e0d2 100644 --- a/src/pyhf/tensor/jax_backend.py +++ b/src/pyhf/tensor/jax_backend.py @@ -230,7 +230,7 @@ def astensor(self, tensor_in, dtype="float"): return jnp.asarray(tensor_in, dtype=dtype) def sum(self, tensor_in, axis=None): - return jnp.sum(tensor_in, axis=axis) + return jnp.sum(jnp.asarray(tensor_in), axis=axis) def product(self, tensor_in, axis=None): return jnp.prod(tensor_in, axis=axis) @@ -334,7 +334,7 @@ def concatenate(self, sequence, axis=0): output: the concatenated tensor """ - return jnp.concatenate(sequence, axis=axis) + return jnp.concatenate([jnp.array(x) for x in sequence], axis=axis) def simple_broadcast(self, *args): """ From 362521d8a1ca79afa0f9dd8e7437a6b8a85db5e3 Mon Sep 17 00:00:00 2001 From: Nathan Simpson Date: Sun, 3 Jul 2022 22:22:41 +0200 Subject: [PATCH 06/12] wow kinda works --- src/pyhf/modifiers/shapefactor.py | 4 ++-- src/pyhf/modifiers/shapesys.py | 10 ++++++---- src/pyhf/modifiers/staterror.py | 14 ++++++++------ src/pyhf/parameters/utils.py | 2 +- src/pyhf/pdf.py | 9 ++++++--- 5 files changed, 23 insertions(+), 16 deletions(-) diff --git a/src/pyhf/modifiers/shapefactor.py b/src/pyhf/modifiers/shapefactor.py index b0a0abb003..9de62c1249 100644 --- a/src/pyhf/modifiers/shapefactor.py +++ b/src/pyhf/modifiers/shapefactor.py @@ -182,8 +182,8 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for t, batch_access in enumerate(syst_access): selection = self.param_viewer.index_selection[s][t] for b, bin_access in enumerate(batch_access): - self._access_field[s, t, b] = ( - selection[bin_access] if bin_access < len(selection) else 0 + self._access_field = self._access_field.at[s, t, b].set( + selection[int(bin_access)] if bin_access < len(selection) else 0 ) self._precompute() diff --git a/src/pyhf/modifiers/shapesys.py b/src/pyhf/modifiers/shapesys.py index 10c8de342e..ea3e0067f1 100644 --- a/src/pyhf/modifiers/shapesys.py +++ b/src/pyhf/modifiers/shapesys.py @@ -171,10 +171,12 @@ def _reindex_access_field(self, pdfconfig): ) sample_mask = self._shapesys_mask[syst_index][singular_sample_index][0] - access_field_for_syst_and_batch[sample_mask] = selection - self._access_field[ - syst_index, batch_index - ] = access_field_for_syst_and_batch + access_field_for_syst_and_batch = access_field_for_syst_and_batch.at[ + sample_mask + ].set(selection) + self._access_field = self._access_field.at[syst_index, batch_index].set( + access_field_for_syst_and_batch + ) def _precompute(self): tensorlib, _ = get_backend() diff --git a/src/pyhf/modifiers/staterror.py b/src/pyhf/modifiers/staterror.py index 2730c02af8..edfadc1bbb 100644 --- a/src/pyhf/modifiers/staterror.py +++ b/src/pyhf/modifiers/staterror.py @@ -117,9 +117,9 @@ def finalize(self): modifier_data['data']['mask'] = masks[modname] sigmas = relerrs[masks[modname]] # list of bools, consistent with other modifiers (no numpy.bool_) - fixed = default_backend.tolist(sigmas == 0) + fixed = sigmas == 0 # ensures non-Nan constraint term, but in a future PR we need to remove constraints for these - sigmas[fixed] = 1.0 + sigmas = sigmas.at[fixed].set(1.0) self.required_parsets.setdefault(parname, [required_parset(sigmas, fixed)]) return self.builder_data @@ -197,10 +197,12 @@ def _reindex_access_field(self, pdfconfig): ) sample_mask = self._staterror_mask[syst_index][singular_sample_index][0] - access_field_for_syst_and_batch[sample_mask] = selection - self._access_field[ - syst_index, batch_index - ] = access_field_for_syst_and_batch + access_field_for_syst_and_batch = access_field_for_syst_and_batch.at[ + sample_mask + ].set(selection) + self._access_field = self._access_field.at[syst_index, batch_index].set( + access_field_for_syst_and_batch + ) def _precompute(self): if not self.param_viewer.index_selection: diff --git a/src/pyhf/parameters/utils.py b/src/pyhf/parameters/utils.py index 9f8e66e647..488ec4d017 100644 --- a/src/pyhf/parameters/utils.py +++ b/src/pyhf/parameters/utils.py @@ -37,7 +37,7 @@ def reduce_paramsets_requirements(paramsets_requirements, paramsets_user_configs for paramset_requirement in paramset_requirements: # undefined: the modifier does not support configuring that property v = paramset_requirement.get(k, 'undefined') - combined_paramset.setdefault(k, set()).add(v) + combined_paramset.setdefault(k, list()).append(v) if len(combined_paramset[k]) != 1: raise exceptions.InvalidNameReuse( diff --git a/src/pyhf/pdf.py b/src/pyhf/pdf.py index b145a89721..90054c7a66 100644 --- a/src/pyhf/pdf.py +++ b/src/pyhf/pdf.py @@ -49,8 +49,10 @@ def _create_parameters_from_spec(_reqs): paramset_type = getattr(pyhf.parameters, paramset_requirements['paramset_type']) paramset = paramset_type(**paramset_requirements) if paramset.constrained: # is constrained - auxdata += paramset.auxdata - auxdata_order.append(param_name) + print(paramset.auxdata) + if paramset.auxdata is not None: + auxdata += paramset.auxdata + auxdata_order.append(param_name) _sets[param_name] = paramset return _sets, auxdata, auxdata_order @@ -261,7 +263,8 @@ def suggested_init(self): """ init = [] for name in self.par_order: - init = init + self.par_map[name]['paramset'].suggested_init + if self.par_map[name]['paramset'].suggested_init is not None: + init = init + self.par_map[name]['paramset'].suggested_init return init def suggested_bounds(self): From 1704eb46474d8247a64d90f52f6d4f884c146505 Mon Sep 17 00:00:00 2001 From: Nathan Simpson Date: Mon, 4 Jul 2022 11:21:57 +0200 Subject: [PATCH 07/12] extra changes --- src/pyhf/parameters/utils.py | 4 +++- src/pyhf/pdf.py | 9 +++------ src/pyhf/tensor/jax_backend.py | 22 +++++++++++----------- 3 files changed, 17 insertions(+), 18 deletions(-) diff --git a/src/pyhf/parameters/utils.py b/src/pyhf/parameters/utils.py index 488ec4d017..1c6a5260b5 100644 --- a/src/pyhf/parameters/utils.py +++ b/src/pyhf/parameters/utils.py @@ -37,7 +37,9 @@ def reduce_paramsets_requirements(paramsets_requirements, paramsets_user_configs for paramset_requirement in paramset_requirements: # undefined: the modifier does not support configuring that property v = paramset_requirement.get(k, 'undefined') - combined_paramset.setdefault(k, list()).append(v) + # differentiable models mean that v could contain jax arrays/tracers + # neither of these are hashable, so the set-based logic here needs changing + combined_paramset.setdefault(k, set()).add(v) if len(combined_paramset[k]) != 1: raise exceptions.InvalidNameReuse( diff --git a/src/pyhf/pdf.py b/src/pyhf/pdf.py index 90054c7a66..b145a89721 100644 --- a/src/pyhf/pdf.py +++ b/src/pyhf/pdf.py @@ -49,10 +49,8 @@ def _create_parameters_from_spec(_reqs): paramset_type = getattr(pyhf.parameters, paramset_requirements['paramset_type']) paramset = paramset_type(**paramset_requirements) if paramset.constrained: # is constrained - print(paramset.auxdata) - if paramset.auxdata is not None: - auxdata += paramset.auxdata - auxdata_order.append(param_name) + auxdata += paramset.auxdata + auxdata_order.append(param_name) _sets[param_name] = paramset return _sets, auxdata, auxdata_order @@ -263,8 +261,7 @@ def suggested_init(self): """ init = [] for name in self.par_order: - if self.par_map[name]['paramset'].suggested_init is not None: - init = init + self.par_map[name]['paramset'].suggested_init + init = init + self.par_map[name]['paramset'].suggested_init return init def suggested_bounds(self): diff --git a/src/pyhf/tensor/jax_backend.py b/src/pyhf/tensor/jax_backend.py index ee15a5e0d2..b39b3e838a 100644 --- a/src/pyhf/tensor/jax_backend.py +++ b/src/pyhf/tensor/jax_backend.py @@ -453,18 +453,18 @@ def poisson(self, n, lam): lam = jnp.asarray(lam) return jnp.exp(xlogy(n, lam) - lam - gammaln(n + 1.0)) - def normal_logpdf(self, x, mu, sigma): - # this is much faster than - # norm.logpdf(x, loc=mu, scale=sigma) - # https://codereview.stackexchange.com/questions/69718/fastest-computation-of-n-likelihoods-on-normal-distributions - root2 = jnp.sqrt(2) - root2pi = jnp.sqrt(2 * jnp.pi) - prefactor = -jnp.log(sigma * root2pi) - summand = -jnp.square(jnp.divide((x - mu), (root2 * sigma))) - return prefactor + summand - # def normal_logpdf(self, x, mu, sigma): - # return norm.logpdf(x, loc=mu, scale=sigma) + # # this is much faster than + # # norm.logpdf(x, loc=mu, scale=sigma) + # # https://codereview.stackexchange.com/questions/69718/fastest-computation-of-n-likelihoods-on-normal-distributions + # root2 = jnp.sqrt(2) + # root2pi = jnp.sqrt(2 * jnp.pi) + # prefactor = -jnp.log(sigma * root2pi) + # summand = -jnp.square(jnp.divide((x - mu), (root2 * sigma))) + # return prefactor + summand + + def normal_logpdf(self, x, mu, sigma): + return norm.logpdf(x, loc=mu, scale=sigma) def normal(self, x, mu, sigma): r""" From 08c50ad64d19b616b94e3603cf54f96527f59bb5 Mon Sep 17 00:00:00 2001 From: Nathan Simpson Date: Mon, 4 Jul 2022 11:39:16 +0200 Subject: [PATCH 08/12] add test --- .../test_differentiable_model_construction.py | 115 ++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 tests/test_differentiable_model_construction.py diff --git a/tests/test_differentiable_model_construction.py b/tests/test_differentiable_model_construction.py new file mode 100644 index 0000000000..6922a32e76 --- /dev/null +++ b/tests/test_differentiable_model_construction.py @@ -0,0 +1,115 @@ +import pyhf +from jax import numpy as jnp +import jax + + +def test_model_building_grad(): + pyhf.set_backend('jax', default=True) + + def make_model( + nominal, + corrup_data, + corrdn_data, + stater_data, + normsys_up, + normsys_dn, + uncorr_data, + ): + return { + "channels": [ + { + "name": "achannel", + "samples": [ + { + "name": "background", + "data": nominal, + "modifiers": [ + {"name": "mu", "type": "normfactor", "data": None}, + {"name": "lumi", "type": "lumi", "data": None}, + { + "name": "mod_name", + "type": "shapefactor", + "data": None, + }, + { + "name": "corr_bkguncrt2", + "type": "histosys", + "data": { + 'hi_data': corrup_data, + 'lo_data': corrdn_data, + }, + }, + { + "name": "staterror2", + "type": "staterror", + "data": stater_data, + }, + { + "name": "norm", + "type": "normsys", + "data": {'hi': normsys_up, 'lo': normsys_dn}, + }, + ], + } + ], + }, + { + "name": "secondchannel", + "samples": [ + { + "name": "background", + "data": nominal, + "modifiers": [ + {"name": "mu", "type": "normfactor", "data": None}, + {"name": "lumi", "type": "lumi", "data": None}, + { + "name": "mod_name", + "type": "shapefactor", + "data": None, + }, + { + "name": "uncorr_bkguncrt2", + "type": "shapesys", + "data": uncorr_data, + }, + { + "name": "corr_bkguncrt2", + "type": "histosys", + "data": { + 'hi_data': corrup_data, + 'lo_data': corrdn_data, + }, + }, + { + "name": "staterror", + "type": "staterror", + "data": stater_data, + }, + { + "name": "norm", + "type": "normsys", + "data": {'hi': normsys_up, 'lo': normsys_dn}, + }, + ], + } + ], + }, + ], + } + + def pipe(x): + spec = make_model( + x * jnp.asarray([60.0, 62.0]), + x * jnp.asarray([60.0, 62.0]), + x * jnp.asarray([60.0, 62.0]), + x * jnp.asarray([5.0, 5.0]), + x * jnp.asarray(0.95), + x * jnp.asarray(1.05), + x * jnp.asarray([5.0, 5.0]), + ) + model = pyhf.Model(spec, validate=False) + nominal = jnp.array(model.config.suggested_init()) + data = model.expected_data(nominal) + return model.logpdf(nominal, data)[0] + + jax.grad(pipe)(3.4) # should work without error From c9b4981b2bd020c40baf504f9790bf538b997b1d Mon Sep 17 00:00:00 2001 From: Nathan Simpson Date: Mon, 4 Jul 2022 12:57:46 +0200 Subject: [PATCH 09/12] revert change --- src/pyhf/tensor/jax_backend.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/pyhf/tensor/jax_backend.py b/src/pyhf/tensor/jax_backend.py index b39b3e838a..ee15a5e0d2 100644 --- a/src/pyhf/tensor/jax_backend.py +++ b/src/pyhf/tensor/jax_backend.py @@ -453,18 +453,18 @@ def poisson(self, n, lam): lam = jnp.asarray(lam) return jnp.exp(xlogy(n, lam) - lam - gammaln(n + 1.0)) - # def normal_logpdf(self, x, mu, sigma): - # # this is much faster than - # # norm.logpdf(x, loc=mu, scale=sigma) - # # https://codereview.stackexchange.com/questions/69718/fastest-computation-of-n-likelihoods-on-normal-distributions - # root2 = jnp.sqrt(2) - # root2pi = jnp.sqrt(2 * jnp.pi) - # prefactor = -jnp.log(sigma * root2pi) - # summand = -jnp.square(jnp.divide((x - mu), (root2 * sigma))) - # return prefactor + summand - def normal_logpdf(self, x, mu, sigma): - return norm.logpdf(x, loc=mu, scale=sigma) + # this is much faster than + # norm.logpdf(x, loc=mu, scale=sigma) + # https://codereview.stackexchange.com/questions/69718/fastest-computation-of-n-likelihoods-on-normal-distributions + root2 = jnp.sqrt(2) + root2pi = jnp.sqrt(2 * jnp.pi) + prefactor = -jnp.log(sigma * root2pi) + summand = -jnp.square(jnp.divide((x - mu), (root2 * sigma))) + return prefactor + summand + + # def normal_logpdf(self, x, mu, sigma): + # return norm.logpdf(x, loc=mu, scale=sigma) def normal(self, x, mu, sigma): r""" From 8457f2c500d4b6134e7595d689361745ca2e743d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 1 Sep 2022 13:49:15 +0000 Subject: [PATCH 10/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/pyhf/modifiers/staterror.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/pyhf/modifiers/staterror.py b/src/pyhf/modifiers/staterror.py index 4167abf87a..22b4f1abe5 100644 --- a/src/pyhf/modifiers/staterror.py +++ b/src/pyhf/modifiers/staterror.py @@ -124,10 +124,9 @@ def finalize(self): fixed = sigmas == 0 # FIXME: sigmas that are zero will be fixed to 1.0 arbitrarily to ensure # non-Nan constraint term, but in a future PR need to remove constraints - # for these + # for these sigmas = sigmas.at[fixed].set(1.0) - self.required_parsets.setdefault(parname, [required_parset(sigmas, fixed)]) return self.builder_data @@ -178,7 +177,6 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): ] ) - self._access_field = default_backend.tile( global_concatenated_bin_indices, (len(self._staterr_mods), self.batch_size or 1, 1), From bb411d0a702c54a5168939b9b768b2dc49546845 Mon Sep 17 00:00:00 2001 From: Nathan Simpson Date: Thu, 1 Sep 2022 18:51:58 +0200 Subject: [PATCH 11/12] Remove self.__staterror_uncrt --- src/pyhf/modifiers/staterror.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/src/pyhf/modifiers/staterror.py b/src/pyhf/modifiers/staterror.py index 22b4f1abe5..86ca8072c5 100644 --- a/src/pyhf/modifiers/staterror.py +++ b/src/pyhf/modifiers/staterror.py @@ -153,18 +153,7 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] - self.__staterror_uncrt = default_backend.astensor( - [ - [ - [ - builder_data[m][s]['data']['uncrt'], - builder_data[m][s]['data']['nom_data'], - ] - for s in pdfconfig.samples - ] - for m in keys - ] - ) + global_concatenated_bin_indices = default_backend.astensor( [ [ From b30a8854c8f6a240964f2a0406fb688720e443a9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 1 Sep 2022 16:53:03 +0000 Subject: [PATCH 12/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/pyhf/modifiers/staterror.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/pyhf/modifiers/staterror.py b/src/pyhf/modifiers/staterror.py index 86ca8072c5..3c06ab6aa0 100644 --- a/src/pyhf/modifiers/staterror.py +++ b/src/pyhf/modifiers/staterror.py @@ -153,7 +153,6 @@ def __init__(self, modifiers, pdfconfig, builder_data, batch_size=None): for m in keys ] - global_concatenated_bin_indices = default_backend.astensor( [ [