wums 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

wums/fitutilsjax.py

@@ -0,0 +1,86 @@
+import numpy as np
+import scipy
+import jax
+import jax.numpy as jnp
+
+def chisqloss(xvals, yvals, yvariances, func, parms):
+    return jnp.sum( (func(xvals, parms) - yvals)**2/yvariances )
+
+chisqloss_grad = jax.jit(jax.value_and_grad(chisqloss, argnums = 4), static_argnums = 3)
+
+def _chisqloss_grad_hess(xvals, yvals, yvariances, func, parms):
+    def lossf(parms):
+        return chisqloss(xvals, yvals, yvariances, func, parms)
+
+    gradf = jax.grad(lossf)
+    hessf = jax.jacfwd(gradf)
+
+    loss = lossf(parms)
+    grad = gradf(parms)
+    hess = hessf(parms)
+
+    return loss, grad, hess
+
+chisqloss_grad_hess = jax.jit(_chisqloss_grad_hess, static_argnums = 3)
+
+def _chisqloss_hessp(xvals, yvals, yvariances, func, parms, p):
+    def lossf(parms):
+        return chisqloss(xvals, yvals, yvariances, func, parms)
+
+    gradf = jax.grad(lossf)
+    hessp = jax.jvp(gradf, (parms,), (p,))[1]
+    return hessp
+
+chisqloss_hessp = jax.jit(_chisqloss_hessp, static_argnums = 3)
+
+def fit_hist_jax(hist, func, parmvals, max_iter = 5, edmtol = 1e-5):
+
+    xvals = [jnp.array(center) for center in hist.axes.centers]
+    yvals = jnp.array(hist.values())
+    yvariances = jnp.array(hist.variances())
+
+    def scipy_loss(parmvals):
+        parms = jnp.array(parmvals)
+        loss, grad = chisqloss_grad(xvals, yvals, yvariances, func, parms)
+        return np.asarray(loss).item(), np.asarray(grad)
+
+    def scipy_hessp(parmvals, p):
+        parms = jnp.array(parmvals)
+        tangent = jnp.array(p)
+        hessp = chisqloss_hessp(xvals, yvals, yvariances, func, parms, tangent)
+        return np.asarray(hessp)
+
+    for iiter in range(max_iter):
+        res = scipy.optimize.minimize(scipy_loss, parmvals, method = "trust-krylov", jac = True, hessp = scipy_hessp)
+
+        parms = jnp.array(res.x)
+        loss, grad, hess = chisqloss_grad_hess(xvals, yvals, yvariances, func, parms)
+        loss, grad, hess = np.asarray(loss).item(), np.asarray(grad), np.asarray(hess)
+
+        eigvals = np.linalg.eigvalsh(hess)
+        cov = 2.*np.linalg.inv(hess)
+
+        gradv = grad[:, np.newaxis]
+        edmval = 0.5*gradv.transpose()@cov@gradv
+        edmval = edmval[0][0]
+
+        converged = edmval < edmtol and np.abs(edmval) >= 0. and eigvals[0] > 0.
+        if converged:
+            break
+
+    status = 1
+    covstatus = 1
+    if edmval < edmtol and np.abs(edmval) >= 0.:
+        status = 0
+    if eigvals[0] > 0.:
+        covstatus = 0
+
+    res = { "x" : res.x,
+           "cov" : cov,
+           "status" : status,
+           "covstatus" : covstatus,
+           "hess_eigvals" : eigvals,
+           "edmval" : edmval,
+           "chisqval" : loss }
+
+    return res

wums/logging.py

@@ -42,19 +42,19 @@ def set_logging_level(log, verbosity):
     log.setLevel(logging_verboseLevel[max(0, min(4, verbosity))])
 
 
-def setup_logger(basefile, verbosity=3, no_colors=False, initName="wremnants"):
+def setup_logger(basefile, verbosity=3, no_colors=False, initName="wums"):
 
     setup_func = setup_base_logger if no_colors else setup_color_logger
     logger = setup_func(os.path.basename(basefile), verbosity, initName)
     # count messages of base logger
-    base_logger = logging.getLogger("wremnants")
+    base_logger = logging.getLogger("wums")
     add_logging_counter(base_logger)
     # stop total time
     add_time_info("Total time")
     return logger
 
 
-def setup_color_logger(name, verbosity, initName="wremnants"):
+def setup_color_logger(name, verbosity, initName="wums"):
     base_logger = logging.getLogger(initName)
     # set console handler
     ch = logging.StreamHandler()
@@ -65,14 +65,14 @@ def setup_color_logger(name, verbosity, initName="wremnants"):
     return base_logger.getChild(name)
 
 
-def setup_base_logger(name, verbosity, initName="wremnants"):
+def setup_base_logger(name, verbosity, initName="wums"):
     logging.basicConfig(format="%(levelname)s: %(message)s")
     base_logger = logging.getLogger(initName)
     set_logging_level(base_logger, verbosity)
     return base_logger.getChild(name)
 
 
-def child_logger(name, initName="wremnants"):
+def child_logger(name, initName="wums"):
     # count messages of child logger
     logger = logging.getLogger(initName).getChild(name)
     add_logging_counter(logger)
@@ -110,7 +110,7 @@ def print_logging_count(logger, verbosity=logging.WARNING):
                 )
 
 
-def add_time_info(tag, logger=logging.getLogger("wremnants")):
+def add_time_info(tag, logger=logging.getLogger("wums")):
     if not hasattr(logger, "times"):
         logger.times = {}
     logger.times[tag] = time.time()
@@ -125,7 +125,7 @@ def print_time_info(logger):
 
 
 def summary(verbosity=logging.WARNING, extended=True):
-    base_logger = logging.getLogger("wremnants")
+    base_logger = logging.getLogger("wums")
 
     base_logger.info(f"--------------------------------------")
     base_logger.info(f"----------- logger summary -----------")
@@ -141,5 +141,5 @@ def summary(verbosity=logging.WARNING, extended=True):
     # Iterate through all child loggers and print their names, levels, and counts
     all_loggers = logging.Logger.manager.loggerDict
     for logger_name, logger_obj in all_loggers.items():
-        if logger_name.startswith("wremnants."):
+        if logger_name.startswith("wums."):
             print_logging_count(logger_obj, verbosity=verbosity)

wums/plot_tools.py

@@ -683,7 +683,6 @@ def makeStackPlotWithRatio(
     stackedProcs,
     histName="nominal",
     unstacked=None,
-    fitresult=None,
     prefit=False,
     xlabel="",
     ylabel=None,
@@ -757,11 +756,6 @@ def makeStackPlotWithRatio(
         if xlim:
             h = h[complex(0, xlim[0]) : complex(0, xlim[1])]
 
-        # If plotting from combine, apply the action to the underlying hist.
-        # Don't do this for the generic case, as it screws up the ability to make multiple plots
-        if fitresult:
-            histInfo[k].hists[histName] = h
-
         if k != "Data":
             stack.append(h)
         else:
@@ -803,67 +797,6 @@ def makeStackPlotWithRatio(
         ratio_axes = None
         ax2 = None
 
-    if fitresult:
-        import uproot
-
-        combine_result = uproot.open(fitresult)
-
-        fittype = "prefit" if prefit else "postfit"
-
-        # set histograms to prefit/postfit values
-        for p in to_read:
-
-            hname = f"expproc_{p}_{fittype}" if p != "Data" else "obs"
-            vals = combine_result[hname].to_hist().values()
-            if len(histInfo[p].hists[histName].values()) != len(vals):
-                raise ValueError(
-                    f"The size of the combine histogram ({(vals.shape)}) is not consistent with the xlim or input hist ({histInfo[p].hists[histName].shape})"
-                )
-
-            histInfo[p].hists[histName].values()[...] = vals
-            if p == "Data":
-                histInfo[p].hists[histName].variances()[...] = vals
-
-        # for postfit uncertaity bands
-        axis = histInfo[to_read[0]].hists[histName].axes[0].edges
-
-        # need to divide by bin width
-        binwidth = axis[1:] - axis[:-1]
-        hexp = combine_result[f"expfull_{fittype}"].to_hist()
-        if hexp.storage_type != hist.storage.Weight:
-            raise ValueError(
-                f"Did not find uncertainties in {fittype} hist. Make sure you run combinetf with --computeHistErrors!"
-            )
-        nom = hexp.values() / binwidth
-        std = np.sqrt(hexp.variances()) / binwidth
-
-        hatchstyle = "///"
-        ax1.fill_between(
-            axis,
-            np.append(nom + std, (nom + std)[-1]),
-            np.append(nom - std, (nom - std)[-1]),
-            step="post",
-            facecolor="none",
-            zorder=2,
-            hatch=hatchstyle,
-            edgecolor="k",
-            linewidth=0.0,
-            label="Uncertainty",
-        )
-
-        if add_ratio:
-            ax2.fill_between(
-                axis,
-                np.append((nom + std) / nom, ((nom + std) / nom)[-1]),
-                np.append((nom - std) / nom, ((nom - std) / nom)[-1]),
-                step="post",
-                facecolor="none",
-                zorder=2,
-                hatch=hatchstyle,
-                edgecolor="k",
-                linewidth=0.0,
-            )
-
     opts = dict(stack=not no_stack, flow=flow)
     optsr = opts.copy()  # no binwnorm for ratio axis
     optsr["density"] = density
@@ -994,7 +927,7 @@ def makeStackPlotWithRatio(
 
         for i, (proc, style) in enumerate(zip(unstacked, linestyles)):
             unstack = histInfo[proc].hists[histName]
-            if not fitresult or proc not in to_read:
+            if proc not in to_read:
                 unstack = action(unstack)[select]
             if proc != "Data":
                 unstack = unstack * scale

wums/tfutils.py

@@ -0,0 +1,81 @@
+import tensorflow as tf
+
+def function_to_tflite(funcs, input_signatures, func_names=""):
+    """Convert function to tflite model using python dynamic execution trickery to ensure that inputs
+    and outputs are alphabetically ordered, since this is apparently the only way to prevent tflite from
+    scrambling them"""
+
+    if not isinstance(funcs, list):
+        funcs = [funcs]
+        input_signatures = [input_signatures]
+        func_names = [func_names]
+    func_names = [funcs[iif].__name__ if func_names[iif]=="" else func_names[iif] for iif in range(len(funcs))]
+
+    def wrapped_func(iif, *args):
+        outputs = funcs[iif](*args)
+
+        if not isinstance(outputs, tuple):
+            outputs = (outputs,)
+
+        output_dict = {}
+        for i,output in enumerate(outputs):
+            output_name = f"output_{iif:05d}_{i:05d}"
+            output_dict[output_name] = output
+
+        return output_dict
+
+    arg_string = []
+    for iif, input_signature in enumerate(input_signatures):
+        inputs = []
+        for i in range(len(input_signature)):
+            input_name = f"input_{iif:05d}_{i:05d}"
+            inputs.append(input_name)
+        arg_string.append(", ".join(inputs))
+
+    def_string = ""
+    def_string += "def make_module(wrapped_func, input_signatures):\n"
+    def_string += "    class Export_Module(tf.Module):\n"
+    for i, func in enumerate(funcs):
+        def_string += f"        @tf.function(input_signature = input_signatures[{i}])\n"
+        def_string += f"        def {func_names[i]}(self, {arg_string[i]}):\n"
+        def_string += f"            return wrapped_func({i}, {arg_string[i]})\n"
+    def_string += "    return Export_Module"
+
+    ldict = {}
+    exec(def_string, globals(), ldict)
+
+    make_module = ldict["make_module"]
+    Export_Module = make_module(wrapped_func, input_signatures)
+
+    module = Export_Module()
+    concrete_functions = [getattr(module, func_name).get_concrete_function() for func_name in func_names]
+    converter = tf.lite.TFLiteConverter.from_concrete_functions(concrete_functions, module)
+
+    # enable TenorFlow ops and DISABLE builtin TFLite ops since these apparently slow things down
+    converter.target_spec.supported_ops = [
+    tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.
+    ]
+
+    converter._experimental_allow_all_select_tf_ops = True
+
+    tflite_model = converter.convert()
+
+    test_interp = tf.lite.Interpreter(model_content = tflite_model)
+    print(test_interp.get_input_details())
+    print(test_interp.get_output_details())
+    print(test_interp.get_signature_list())
+
+    return tflite_model
+
+
+
+def function_to_saved_model(func, input_signature, output):
+
+    class Export_Module(tf.Module):
+        @tf.function(input_signature = input_signature)
+        def __call__(self, *args):
+            return func(*args)
+
+    model = Export_Module()
+
+    tf.saved_model.save(model, output)

wums-0.1.8.dist-info/METADATA

@@ -0,0 +1,54 @@
+Metadata-Version: 2.2
+Name: wums
+Version: 0.1.8
+Summary: .
+Author-email: David Walter <david.walter@cern.ch>, Josh Bendavid <josh.bendavid@cern.ch>, Kenneth Long <kenneth.long@cern.ch>, Jan Eysermans <jan.eysermans@cern.ch>
+License: MIT
+Project-URL: Homepage, https://github.com/WMass/wums
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: hist
+Requires-Dist: numpy
+Provides-Extra: plotting
+Requires-Dist: matplotlib; extra == "plotting"
+Requires-Dist: mplhep; extra == "plotting"
+Provides-Extra: fitting
+Requires-Dist: tensorflow; extra == "fitting"
+Requires-Dist: jax; extra == "fitting"
+Requires-Dist: scipy; extra == "fitting"
+Provides-Extra: pickling
+Requires-Dist: boost_histogram; extra == "pickling"
+Requires-Dist: h5py; extra == "pickling"
+Requires-Dist: hdf5plugin; extra == "pickling"
+Requires-Dist: lz4; extra == "pickling"
+Provides-Extra: all
+Requires-Dist: plotting; extra == "all"
+Requires-Dist: fitting; extra == "all"
+Requires-Dist: pickling; extra == "all"
+
+# WUMS: Wremnants Utilities, Modules, and other Stuff
+
+As the name suggests, this is a collection of different thins, all python based:
+- Fitting with tensorflow or jax
+- Custom pickling h5py objects 
+- Plotting functionality
+
+## Install
+
+The `wums` package can be pip installed with minimal dependencies:
+```bash
+pip install wums
+```
+Different dependencies can be added with `plotting`, `fitting`, `pickling` to use the corresponding scripts.
+For example, one can install with
+```bash
+pip install wums[plotting,fitting]
+```
+Or all dependencies with
+```bash
+pip install wums[all]
+```

wums-0.1.8.dist-info/RECORD

@@ -0,0 +1,16 @@
+scripts/test/testsplinepdf.py,sha256=sXnmDjEXiO0OIHAXLXU4UxTD4_nLwUpoojCecfjyT04,1964
+scripts/test/testsplinepdf2d.py,sha256=vGw9mq67f6aoymefLqv6CqF8teluva4Lx6tpbnC_NGU,8513
+wums/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+wums/boostHistHelpers.py,sha256=mgdPXAgmxriqoOhrhMctyZcfwEOPfV07V27CvGt2sk8,39260
+wums/fitutils.py,sha256=sPCMJqZGdXvDfc8OxjOB-Bpf45GWHKxmKkDV3SlMUQs,38297
+wums/fitutilsjax.py,sha256=HE1AcIZmI6N_xIHo8OHCPaYkHSnND_B-vI4Gl3vaUmA,2659
+wums/ioutils.py,sha256=ziyfQQ8CB3Ir2BJKJU3_a7YMF-Jd2nGXKoMQoJ2T8fo,12334
+wums/logging.py,sha256=L4514Xyq7L1z77Tkh8KE2HX88ZZ06o6SSRyQo96DbC0,4494
+wums/output_tools.py,sha256=SHcZqXAdqL9AkA57UF0b-R-U4u7rzDgL8Def4E-ulW0,6713
+wums/plot_tools.py,sha256=7GBQAO--wuP8aatkjy-ir1lQWpNrzMc1lSI6zSq3JXE,53502
+wums/tfutils.py,sha256=9efkkvxH7VtwJN2yBS6_-P9dLKs3CXdxMFdrEBNsna8,2892
+wums/Templates/index.php,sha256=9EYmfc0ltMqr5oOdA4_BVIHdSbef5aA0ORoRZBEADVw,4348
+wums-0.1.8.dist-info/METADATA,sha256=87fET64UzNDs6swv1-tcJWcuzVE5S3kEuLWOfy1JN6c,1784
+wums-0.1.8.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+wums-0.1.8.dist-info/top_level.txt,sha256=cGGeFZQ8IwVw-BhgxMCTu5zfkgQelfF1wEFFWGhycds,13
+wums-0.1.8.dist-info/RECORD,,

{wums-0.1.6.dist-info → wums-0.1.8.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.0)
+Generator: setuptools (75.8.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

wums-0.1.8.dist-info/top_level.txt

@@ -0,0 +1,2 @@
+scripts
+wums

wums-0.1.6.dist-info/METADATA

@@ -1,29 +0,0 @@
-Metadata-Version: 2.2
-Name: wums
-Version: 0.1.6
-Summary: .
-Author-email: David Walter <david.walter@cern.ch>, Josh Bendavid <josh.bendavid@cern.ch>, Kenneth Long <kenneth.long@cern.ch>
-License: MIT
-Project-URL: Homepage, https://github.com/WMass/wums
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Operating System :: OS Independent
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-Requires-Dist: boost_histogram
-Requires-Dist: h5py
-Requires-Dist: hdf5plugin
-Requires-Dist: hist
-Requires-Dist: lz4
-Requires-Dist: matplotlib
-Requires-Dist: mplhep
-Requires-Dist: numpy
-Requires-Dist: uproot
-
-# WUMS: Wremnants Utilities, Modules, and other Stuff
-
-The `wums` package can be pip installed:
-```bash
-pip install wums
-```

wums-0.1.6.dist-info/RECORD

@@ -1,11 +0,0 @@
-wums/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-wums/boostHistHelpers.py,sha256=F4SwQEVjNObFscfs0qrJEyOHYNKqUCmusW8HIF1o-0c,38993
-wums/ioutils.py,sha256=ziyfQQ8CB3Ir2BJKJU3_a7YMF-Jd2nGXKoMQoJ2T8fo,12334
-wums/logging.py,sha256=zNnLVJUwG3HMvr9NeXmiheX07VmsnSt8cQ6R4q4XBk4,4534
-wums/output_tools.py,sha256=SHcZqXAdqL9AkA57UF0b-R-U4u7rzDgL8Def4E-ulW0,6713
-wums/plot_tools.py,sha256=4iPx9Nr9y8c3p4ovy8XOS-xU_w11OyQEjISKkygxqcA,55918
-wums/Templates/index.php,sha256=9EYmfc0ltMqr5oOdA4_BVIHdSbef5aA0ORoRZBEADVw,4348
-wums-0.1.6.dist-info/METADATA,sha256=pTmIMc-rth2X53tju6Ef8WbJDda2zbr8isEqUpeqhDo,843
-wums-0.1.6.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-wums-0.1.6.dist-info/top_level.txt,sha256=DCE1TVg7ySraosR3kYZkLIZ2w1Pwk2pVTdkqx6E-yRY,5
-wums-0.1.6.dist-info/RECORD,,

wums-0.1.6.dist-info/top_level.txt

@@ -1 +0,0 @@
-wums