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

wums/boostHistHelpers.py

@@ -4,6 +4,7 @@ from functools import reduce
 
 import hist
 import numpy as np
+from scipy.interpolate import make_smoothing_spline
 
 from wums import logging
 
@@ -60,7 +61,7 @@ def broadcastSystHist(h1, h2, flow=True, by_ax_name=True):
     # move back to original order
     new_vals = np.moveaxis(new_vals, np.arange(len(moves)), list(moves.keys()))
 
-    if new_vals.shape != h2.values(flow=flow).shape:
+    if new_vals.shape != s2:
         raise ValueError(
             f"Broadcast shape {new_vals.shape} (from h1.shape={h1.values(flow=flow).shape}, axes={h1.axes.name}) "
             f"does not match desired shape {h2.view(flow=flow).shape} (axes={h2.axes.name})"
@@ -193,10 +194,10 @@ def multiplyWithVariance(vals1, vals2, vars1=None, vars2=None):
     return outvals, outvars
 
 
-def multiplyHists(h1, h2, allowBroadcast=True, createNew=True, flow=True):
+def multiplyHists(h1, h2, allowBroadcast=True, createNew=True, flow=True, broadcast_by_ax_name=True):
     if allowBroadcast:
-        h1 = broadcastSystHist(h1, h2, flow=flow)
-        h2 = broadcastSystHist(h2, h1, flow=flow)
+        h1 = broadcastSystHist(h1, h2, flow=flow, by_ax_name=broadcast_by_ax_name)
+        h2 = broadcastSystHist(h2, h1, flow=flow, by_ax_name=broadcast_by_ax_name)
 
     if (
         h1.storage_type == hist.storage.Double
@@ -233,6 +234,7 @@ def concatenateHists(h1, h2, allowBroadcast=True, by_ax_name=True, flow=False):
         h2 = broadcastSystHist(h2, h1, flow=flow, by_ax_name=by_ax_name)
 
     axes = []
+
     for ax1, ax2 in zip(h1.axes, h2.axes):
         if ax1 == ax2:
             axes.append(ax1)
@@ -263,7 +265,7 @@ def concatenateHists(h1, h2, allowBroadcast=True, by_ax_name=True, flow=False):
                 )
             else:
                 raise ValueError(
-                    f"Cannot concatenate hists with inconsistent axes: {ax1.name} and {ax2.name}"
+                    f"Cannot concatenate hists with inconsistent axes: {ax1.name}: ({ax1.edges}) and {ax2.name}: ({ax2.edges})"
                 )
 
     newh = hist.Hist(*axes, storage=h1.storage_type())
@@ -407,6 +409,10 @@ def normalize(h, scale=1e6, createNew=True, flow=True):
     return scaleHist(h, scale, createNew, flow)
 
 
+def renameAxis(h, axis_name, new_name):
+    h.axes[axis_name].__dict__['name'] = new_name
+
+
 def makeAbsHist(h, axis_name, rename=True):
     ax = h.axes[axis_name]
     axidx = list(h.axes).index(ax)
@@ -455,8 +461,8 @@ def compatibleBins(edges1, edges2):
 
 def rebinHistMultiAx(h, axes, edges=[], lows=[], highs=[]):
     # edges: lists of new edges or integers to merge bins, in case new edges are given the lows and highs will be ignored
-    # lows: list of new lower boundaries
-    # highs: list of new upper boundaries
+    # lows: list of new lower boundaries or bins
+    # highs: list of new upper boundaries or bins
 
     sel = {}
     for ax, low, high, rebin in itertools.zip_longest(axes, lows, highs, edges):
@@ -467,14 +473,21 @@ def rebinHistMultiAx(h, axes, edges=[], lows=[], highs=[]):
             h = rebinHist(h, ax, rebin)
         elif low is not None and high is not None:
             # in case high edge is upper edge of last bin we need to manually set the upper limit
-            upper = hist.overflow if high == h.axes[ax].edges[-1] else complex(0, high)
-            logger.info(f"Restricting the axis '{ax}' to range [{low}, {high}]")
+            # distinguish case with pure real or imaginary number (to select index or value, respectively)
+            # in the former case, force casting into integer
+            if isinstance(high, int):
+                upper = hist.overflow if high == h.axes[ax].size else high
+            elif isinstance(high, complex):
+                high_imag = high.imag
+                upper = hist.overflow if high_imag == h.axes[ax].edges[-1] else high
+            logger.info(f"Slicing the axis '{ax}' to [{low}, {upper}]")
             sel[ax] = slice(
-                complex(0, low), upper, hist.rebin(rebin) if rebin else None
+                low, upper, hist.rebin(rebin) if rebin else None
             )
         elif type(rebin) == int and rebin > 1:
             logger.info(f"Rebinning the axis '{ax}' by [{rebin}]")
             sel[ax] = slice(None, None, hist.rebin(rebin))
+
     return h[sel] if len(sel) > 0 else h
 
 
@@ -499,9 +512,9 @@ def mirrorAxes(h, axes, flow=True):
     return h
 
 
-def disableAxisFlow(ax):
+def disableAxisFlow(ax, under=False, over=False):
     if isinstance(ax, hist.axis.Integer):
-        args = [ax.edges[0], ax.edges[-1]]
+        args = [int(ax.edges[0]), int(ax.edges[-1])]
     elif isinstance(ax, hist.axis.Regular):
         args = [ax.size, ax.edges[0], ax.edges[-1]]
     else:
@@ -510,17 +523,24 @@ def disableAxisFlow(ax):
     return type(ax)(
         *args,
         name=ax.name,
-        overflow=False,
-        underflow=False,
+        overflow=over,
+        underflow=under,
         circular=ax.traits.circular,
     )
 
 
-def disableFlow(h, axis_name):
-    # disable the overflow and underflow bins of a single axes, while keeping the flow bins of other axes
+def disableFlow(h, axis_name, under=False, over=False):
+    # axes_name can be either string or a list of strings with the axis name(s) to disable the flow
+    if not isinstance(axis_name, str):
+        for var in axis_name:
+            if var in h.axes.name:
+                h = disableFlow(h, var)
+        return h
+
+    # disable the overflow and underflow bins of a single axis, while keeping the flow bins of other axes
     ax = h.axes[axis_name]
     ax_idx = [a.name for a in h.axes].index(axis_name)
-    new_ax = disableAxisFlow(ax)
+    new_ax = disableAxisFlow(ax, under=under, over=over)
     axes = list(h.axes)
     axes[ax_idx] = new_ax
     hnew = hist.Hist(*axes, name=h.name, storage=h.storage_type())
@@ -528,7 +548,7 @@ def disableFlow(h, axis_name):
         (
             slice(None)
             if i != ax_idx
-            else slice(ax.traits.underflow, new_ax.size + ax.traits.underflow)
+            else slice(ax.traits.underflow * (not under), ax.size + ax.traits.underflow + ax.traits.overflow * over)
         )
         for i in range(len(axes))
     ]
@@ -773,9 +793,9 @@ def unrolledHist(h, obs=None, binwnorm=None, add_flow_bins=False):
 
     if binwnorm:
         edges = (
-            plot_tools.extendEdgesByFlow(hproj) if add_flow_bins else hproj.axes.edges
+            extendEdgesByFlow(hproj) if add_flow_bins else hproj.axes.edges
         )
-        binwidths = np.outer(*[np.diff(e.squeeze()) for e in edges]).flatten()
+        binwidths = np.array(list(itertools.product(*[np.diff(e.squeeze()) for e in edges]))).prod(axis=1)
         scale = binwnorm / binwidths
     else:
         scale = 1
@@ -1133,3 +1153,42 @@ def rssHistsMid(h, syst_axis, scale=1.0):
     hDown = addHists(hnom, hrss[{"downUpVar": -1j}], scale2=-1.0)
 
     return hUp, hDown
+
+def smooth_hist(h, smooth_ax_name, exclude_axes=[], start_bin=0, end_bin=None):
+
+    hnew = h.copy()
+
+    smooth_ax = h.axes[smooth_ax_name]
+    hproj = h.project(smooth_ax_name, *[ax for ax in h.axes.name if ax not in [smooth_ax_name, *exclude_axes]])
+    smoothh = hproj.copy()
+
+    # Reshape before looping over all other bins and smoothing along the relevant axis
+    vals = smoothh.values().reshape(smooth_ax.size, -1)
+    if not end_bin:
+        end_bin = vals.shape[0]
+
+    # Correct for bin width
+    binw = np.diff(smoothh.axes[smooth_ax_name].edges)
+    vals = (vals.T/binw).T
+
+    for b in range(vals.shape[-1]):
+        spl = make_smoothing_spline(smooth_ax.centers[start_bin:end_bin], vals[start_bin:end_bin,b])
+        vals[start_bin:end_bin,b] = spl(smooth_ax.centers[start_bin:end_bin])
+
+    #Recorrect for bin width
+    vals = (vals.T*binw).T
+
+    smoothh.values()[...] = vals.reshape(smoothh.shape)
+
+    if not exclude_axes:
+        return smoothh.project(*hnew.axes.name)
+
+    # If some axis has been excluded, broadcast it back 
+    smoothfac = divideHists(smoothh, hproj).project(*[ax for ax in h.axes.name if ax not in exclude_axes])
+
+    # Broadcast over excluded axes
+    indices = tuple(
+        None if ax in exclude_axes else slice(None) for ax in h.axes.name
+    )
+    hnew.values()[...] = h.values() * smoothfac.values()[indices]
+    return hnew

wums/ioutils.py

@@ -325,8 +325,12 @@ class H5Unpickler(pickle.Unpickler):
             raise pickle.UnpicklingError("unsupported persistent object")
 
 
-def pickle_dump_h5py(name, obj, h5out):
+def pickle_dump_h5py(name, obj, h5out, override=False):
     """Write an object to a new h5py group which will be created within the provided group."""
+    if name in h5out and override:
+        # If the group already exists, delete it first
+        del h5out[name]
+
     obj_group = h5out.create_group(name)
     try:
         obj_group.attrs["narf_h5py_pickle_protocol_version"] = CURRENT_PROTOCOL_VERSION

wums/output_tools.py

@@ -165,6 +165,12 @@ def make_meta_info_dict(
     return meta_data
 
 
+def encode_complex(obj):
+    if isinstance(obj, complex):
+        return {"real": obj.real, "imag": obj.imag}
+    raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable")
+
+
 def write_logfile(
     outpath,
     logname,
@@ -181,7 +187,8 @@ def write_logfile(
             logf.write("\n" + "-" * 80 + "\n")
             if isinstance(v, dict):
                 logf.write(k)
-                logf.write(json.dumps(v, indent=5).replace("\\n", "\n"))
+                # String conversion needed for non-primitive types (e.g., complex number)
+                logf.write(json.dumps(v, default=encode_complex, indent=5).replace("\\n", "\n"))
             else:
                 logf.write(f"{k}: {v}\n")
 
@@ -202,3 +209,19 @@ def write_index_and_log(
 ):
     write_indexfile(outpath, template_dir)
     write_logfile(outpath, logname, args, analysis_meta_info)
+
+def write_lz4_pkl_output(
+    outfile, outfolder, output_dict, basedir, args=None, file_meta_data=None
+):
+    if not outfile.endswith(".pkl.lz4"):
+        outfile += ".pkl.lz4"
+
+    logger.info(f"Write file {outfile}")
+    result_dict = {
+        outfolder : output_dict,
+        "meta_data": make_meta_info_dict(args, wd=basedir),
+    }
+    if file_meta_data is not None:
+        result_dict["file_meta_data"] = file_meta_data
+    with lz4.frame.open(outfile, "wb") as f:
+        pickle.dump(result_dict, f, protocol=pickle.HIGHEST_PROTOCOL)

wums/plot_tools.py

@@ -1,12 +1,7 @@
-import datetime
-import json
+import importlib
+import inspect
 import math
-import pathlib
-import shutil
-import socket
-import sys
 import textwrap
-import importlib
 
 import hist
 import matplotlib as mpl
@@ -22,7 +17,7 @@ from matplotlib.patches import Polygon
 from matplotlib.ticker import StrMethodFormatter
 
 from wums import boostHistHelpers as hh
-from wums import ioutils, logging
+from wums import logging
 
 hep.style.use(hep.style.ROOT)
 
@@ -466,6 +461,7 @@ def addLegend(
     reverse=True,
     labelcolor=None,
     padding_loc="auto",
+    title=None,
 ):
     handles, labels = ax.get_legend_handles_labels()
     if extra_entries_first:
@@ -496,6 +492,7 @@ def addLegend(
     text_size = get_textsize(ax, text_size)
     handler_map = get_custom_handler_map(custom_handlers)
     leg = ax.legend(
+        title=title,
         handles=handles,
         labels=labels,
         prop={"size": text_size},
@@ -507,6 +504,8 @@ def addLegend(
         markerfirst=markerfirst,
         labelcolor=labelcolor,
     )
+    if title is not None:
+        leg.set_title(title, prop={"size": text_size})
 
     if extra_text is not None:
         if extra_text_loc is None:
@@ -609,18 +608,18 @@ def add_decor(
     ax, title, label=None, lumi=None, loc=2, data=True, text_size=None, no_energy=False
 ):
     text_size = get_textsize(ax, text_size)
-
     if title in ["CMS", "ATLAS", "LHCb", "ALICE"]:
         module = getattr(hep, title.lower())
         make_text = module.text
         make_label = module.label
     else:
+
         def make_text(text=None, **kwargs):
             for key, value in dict(hep.rcParams.text._get_kwargs()).items():
                 if (
                     value is not None
                     and key not in kwargs
-                    and key in inspect.getfullargspec(label_base.exp_text).kwonlyargs
+                    and key in inspect.getfullargspec(hep.label.exp_text).kwonlyargs
                 ):
                     kwargs.setdefault(key, value)
             kwargs.setdefault("italic", (False, True, False))
@@ -632,7 +631,7 @@ def add_decor(
                 if (
                     value is not None
                     and key not in kwargs
-                    and key in inspect.getfullargspec(label_base.exp_text).kwonlyargs
+                    and key in inspect.getfullargspec(hep.label.exp_text).kwonlyargs
                 ):
                     kwargs.setdefault(key, value)
             kwargs.setdefault("italic", (False, True, False))
@@ -651,7 +650,7 @@ def add_decor(
             data=data,
             loc=loc,
         )
-    
+
     # else:
     #     if loc==0:
     #         # above frame
@@ -666,7 +665,7 @@ def add_decor(
     #         x = 0.05
     #         y = 0.88
     #     elif loc==2:
-    #         #     
+    #         #
     #     ax.text(
     #         x,
     #         y,
@@ -678,12 +677,12 @@ def add_decor(
     #     if label is not None:
     #         ax.text(0.05, 0.80, label, transform=ax.transAxes, fontstyle="italic")
 
+
 def makeStackPlotWithRatio(
     histInfo,
     stackedProcs,
     histName="nominal",
     unstacked=None,
-    prefit=False,
     xlabel="",
     ylabel=None,
     rlabel="Data/Pred.",
@@ -726,12 +725,14 @@ def makeStackPlotWithRatio(
     alpha=0.7,
     legPos="upper right",
     leg_padding="auto",
+    lowerLeg=True,
     lowerLegCols=2,
     lowerLegPos="upper right",
     lower_panel_variations=0,
     lower_leg_padding="auto",
     scaleRatioUnstacked=[],
     subplotsizes=[4, 2],
+    x_vertLines_edges=[],
 ):
     add_ratio = not (no_stack or no_ratio)
     if ylabel is None:
@@ -799,9 +800,8 @@ def makeStackPlotWithRatio(
 
     opts = dict(stack=not no_stack, flow=flow)
     optsr = opts.copy()  # no binwnorm for ratio axis
-    optsr["density"] = density
     if density:
-        opts["density"] = True
+        opts["density"] = True            
     else:
         opts["binwnorm"] = binwnorm
 
@@ -822,13 +822,14 @@ def makeStackPlotWithRatio(
             for x in (data_hist.sum(), hh.sumHists(stack).sum())
         ]
         scale = vals[0] / vals[1]
-        unc = scale * (varis[0] / vals[0] ** 2 + varis[1] / vals[1] ** 2)**0.5
+        unc = scale * (varis[0] / vals[0] ** 2 + varis[1] / vals[1] ** 2) ** 0.5
         ndigits = -math.floor(math.log10(abs(unc))) + 1
         logger.info(
             f"Rescaling all processes by {round(scale,ndigits)} +/- {round(unc,ndigits)} to match data norm"
         )
         stack = [s * scale for s in stack]
 
+    
     hep.histplot(
         stack,
         histtype="fill" if not no_fill else "step",
@@ -927,8 +928,7 @@ def makeStackPlotWithRatio(
 
         for i, (proc, style) in enumerate(zip(unstacked, linestyles)):
             unstack = histInfo[proc].hists[histName]
-            if proc not in to_read:
-                unstack = action(unstack)[select]
+            unstack = action(unstack)[select]
             if proc != "Data":
                 unstack = unstack * scale
             if len(scaleRatioUnstacked) > i:
@@ -975,6 +975,11 @@ def makeStackPlotWithRatio(
                 extra_labels.append(histInfo[proc].label)
             if ratio_to_data and proc == "Data" or not add_ratio:
                 continue
+            if xlim:
+                unstack = unstack[complex(0, xlim[0]) : complex(0, xlim[1])]
+            if density:
+                unstack = hh.scaleHist(unstack, np.sum(ratio_ref.values())/np.sum(unstack.values()))
+
             stack_ratio = hh.divideHists(
                 unstack,
                 ratio_ref,
@@ -1001,6 +1006,14 @@ def makeStackPlotWithRatio(
                 **optsr,
             )
 
+    if len(x_vertLines_edges):
+        h_inclusive = hh.sumHists(stack)
+        max_y = 1.05 * np.max(h_inclusive.values() + h_inclusive.variances() ** 0.5)
+        min_y = np.min(h_inclusive.values() - h_inclusive.variances() ** 0.5)
+        for x in x_vertLines_edges:
+            ax1.plot([x, x], [min_y, max_y], linestyle="--", color="black")
+            ax2.plot([x, x], [rrange[0], rrange[1]], linestyle="--", color="black")
+
     addLegend(
         ax1,
         nlegcols,
@@ -1010,7 +1023,7 @@ def makeStackPlotWithRatio(
         text_size=legtext_size,
         padding_loc=leg_padding,
     )
-    if add_ratio:
+    if add_ratio and lowerLeg:
         addLegend(
             ax2,
             lowerLegCols,
@@ -1038,7 +1051,7 @@ def makeStackPlotWithRatio(
 def makePlotWithRatioToRef(
     hists,
     labels,
-    colors,
+    colors=None,
     hists_ratio=None,
     midratio_idxs=None,
     linestyles=[],
@@ -1074,6 +1087,7 @@ def makePlotWithRatioToRef(
     cms_label=None,
     cutoff=1e-6,
     only_ratio=False,
+    ratio_legend=True,
     width_scale=1,
     automatic_scale=True,
     base_size=8,
@@ -1095,9 +1109,19 @@ def makePlotWithRatioToRef(
     elif select is not None:
         hists_ratio = [h[select] for h in hists_ratio]
 
-    if len(hists_ratio) != len(labels) or len(hists_ratio) != len(colors):
+    if colors is None:
+        colors = plt.rcParams["axes.prop_cycle"].by_key()["color"][: len(hists)]
+    if len(colors) < len(hists):
+        colors = (
+            colors
+            + plt.rcParams["axes.prop_cycle"].by_key()["color"][
+                : (len(hists) - len(colors))
+            ]
+        )
+
+    if len(hists_ratio) != len(labels):
         raise ValueError(
-            f"Number of hists ({len(hists_ratio)}), colors ({len(colors)}), and labels ({len(labels)}) must agree!"
+            f"Number of hists ({len(hists_ratio)}) and labels ({len(labels)}) must agree!"
         )
     ratio_hists = [
         hh.divideHists(
@@ -1232,7 +1256,7 @@ def makePlotWithRatioToRef(
             fill_between=fill_between,
             dataIdx=dataIdx,
             baseline=baseline,
-            add_legend=not only_ratio,
+            add_legend=ratio_legend and not only_ratio,
         )
         if midratio_hists:
             plotRatio(
@@ -1513,7 +1537,7 @@ def fix_axes(
     if noSci and not logy:
         redo_axis_ticks(ax1, "y")
     elif not logy:
-        ax1.ticklabel_format(style="sci", useMathText=True, axis="y", scilimits=(0, 0))
+        ax1.ticklabel_format(style="sci", useMathText=True, axis="y", scilimits=(-2, 2))
 
     if ratio_axes is not None:
         if not isinstance(ratio_axes, (list, tuple, np.ndarray)):
@@ -1556,6 +1580,8 @@ def redo_axis_ticks(ax, axlabel, no_labels=False):
     fixedloc = ticker.FixedLocator(
         autoloc.tick_values(*getattr(ax, f"get_{axlabel}lim")())
     )
+    if ax.get_xscale() == 'log':
+        fixedloc = ticker.LogLocator(base=10, numticks=5)
     getattr(ax, f"{axlabel}axis").set_major_locator(fixedloc)
     ticks = getattr(ax, f"get_{axlabel}ticks")()
     labels = [format_axis_num(x, ticks[-1]) for x in ticks] if not no_labels else []
@@ -1798,19 +1824,21 @@ def read_axis_label(x, labels, with_unit=True):
         return x
 
 
-def get_axis_label(config, default_keys=None, label=None, is_bin=False):
+def get_axis_label(config, default_keys=None, label=None, is_bin=False, with_unit=True):
     if label is not None:
         return label
 
     if default_keys is None:
         return "Bin index"
+    elif isinstance(default_keys, str):
+        default_keys = [default_keys]
 
     labels = getattr(config, "axis_labels", {})
 
-    if len(default_keys) == 1:        
+    if len(default_keys) == 1:
         if is_bin:
             return f"{read_axis_label(default_keys[0], labels, False)} bin"
         else:
-            return read_axis_label(default_keys[0], labels)
+            return read_axis_label(default_keys[0], labels, with_unit)
     else:
         return f"({', '.join([read_axis_label(a, labels, False) for a in default_keys])}) bin"

{wums-0.1.8.dist-info → wums-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: wums
-Version: 0.1.8
+Version: 0.1.9
 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

wums-0.1.9.dist-info/RECORD

@@ -0,0 +1,14 @@
+wums/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+wums/boostHistHelpers.py,sha256=KMZXVPLlx7_v0XgGEQg6mZmuzEOQ6AWO5Wd_mPaAPjg,41636
+wums/fitutils.py,sha256=sPCMJqZGdXvDfc8OxjOB-Bpf45GWHKxmKkDV3SlMUQs,38297
+wums/fitutilsjax.py,sha256=HE1AcIZmI6N_xIHo8OHCPaYkHSnND_B-vI4Gl3vaUmA,2659
+wums/ioutils.py,sha256=EyCOBin7ifStLtdgKl7J1_0VWB6RnXWtIWFge9x73Ow,12465
+wums/logging.py,sha256=L4514Xyq7L1z77Tkh8KE2HX88ZZ06o6SSRyQo96DbC0,4494
+wums/output_tools.py,sha256=89rQPOWpwGuzJK5ZQBDv38rmO9th1D2206QOe9PE-gY,7572
+wums/plot_tools.py,sha256=0olJuXAnuDDXJrbF9GF_ZaWBJ14ljOzF5ww19VGCG_g,54662
+wums/tfutils.py,sha256=9efkkvxH7VtwJN2yBS6_-P9dLKs3CXdxMFdrEBNsna8,2892
+wums/Templates/index.php,sha256=9EYmfc0ltMqr5oOdA4_BVIHdSbef5aA0ORoRZBEADVw,4348
+wums-0.1.9.dist-info/METADATA,sha256=TgZpEIJTjE_xZ3J_cbrXaD2EClyo3iEKS6MHl5lanLU,1784
+wums-0.1.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+wums-0.1.9.dist-info/top_level.txt,sha256=DCE1TVg7ySraosR3kYZkLIZ2w1Pwk2pVTdkqx6E-yRY,5
+wums-0.1.9.dist-info/RECORD,,

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

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

wums-0.1.9.dist-info/top_level.txt

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

scripts/test/testsplinepdf.py

@@ -1,90 +0,0 @@
-import wums.fitutils
-
-import tensorflow as tf
-
-import matplotlib.pyplot as plt
-
-import numpy as np
-import hist
-import math
-
-np.random.seed(1234)
-
-nevt = 100000
-
-rgaus = np.random.normal(size=(nevt,))
-
-print(rgaus.dtype)
-print(rgaus)
-
-axis0 = hist.axis.Regular(100, -5., 5.)
-
-htest = hist.Hist(axis0)
-htest.fill(rgaus)
-
-print(htest)
-
-
-quant_cdfvals = tf.constant([0.0, 1e-3, 0.02, 0.05, 0.10, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 0.90, 0.95, 0.98, 1.0-1e-3, 1.0], tf.float64)
-
-nquants = quant_cdfvals.shape.num_elements()
-
-def func_transform_cdf(quantile):
-    const_sqrt2 = tf.constant(math.sqrt(2.), quantile.dtype)
-    return 0.5*(1. + tf.math.erf(quantile/const_sqrt2))
-
-def func_transform_quantile(cdf):
-    const_sqrt2 = tf.constant(math.sqrt(2.), cdf.dtype)
-    return const_sqrt2*tf.math.erfinv(2*cdf - 1.)
-
-
-
-def func_cdf(xvals, xedges, parms, quant_cdfvals):
-    qparms = parms
-
-    cdf = narf.fitutils.func_cdf_for_quantile_fit(xvals, xedges, qparms, quant_cdfvals, transform = (func_transform_cdf, func_transform_quantile))
-
-    return cdf
-
-
-#this is just for plotting
-def func_pdf(h, parms):
-    dtype = tf.float64
-    xvals = [tf.constant(center, dtype=dtype) for center in h.axes.centers]
-    xedges = [tf.constant(edge, dtype=dtype) for edge in h.axes.edges]
-
-    tfparms = tf.constant(parms)
-
-    cdf = func_cdf(xvals, xedges, tfparms, quant_cdfvals)
-
-    pdf = cdf[1:] - cdf[:-1]
-    pdf = tf.maximum(pdf, tf.zeros_like(pdf))
-
-    return pdf
-
-nparms = nquants-1
-
-
-initial_parms = np.array([np.log(1./nparms)]*nparms)
-
-res = narf.fitutils.fit_hist(htest, func_cdf, initial_parms, mode="nll_bin_integrated", func_constraint=narf.fitutils.func_constraint_for_quantile_fit, args = (quant_cdfvals,))
-
-print(res)
-
-
-parmvals = res["x"]
-
-
-pdfvals = func_pdf(htest, parmvals)
-pdfvals *= htest.sum()/np.sum(pdfvals)
-
-#
-plot = plt.figure()
-plt.yscale("log")
-htest.plot()
-plt.plot(htest.axes[0].centers, pdfvals)
-# plt.show()
-plot.savefig("test.png")
-
-
-

scripts/test/testsplinepdf2d.py

@@ -1,323 +0,0 @@
-import wums.fitutils
-
-import tensorflow as tf
-
-import matplotlib.pyplot as plt
-
-import numpy as np
-import hist
-import math
-
-import onnx
-import tf2onnx
-
-np.random.seed(1234)
-
-nevt = 20000
-
-runiform = np.random.random((nevt,))
-rgaus = np.random.normal(size=(nevt,))
-
-data = np.stack([runiform, rgaus], axis=-1)
-
-# "pt"-dependent mean and sigma
-data[:,1] = -0.1 + 0.1*data[:,0] + (1. + 0.2*data[:,0])*data[:,1]
-
-
-# print(rgaus.dtype)
-# print(rgaus)
-
-axis0 = hist.axis.Regular(50, 0., 1., name="pt")
-axis1 = hist.axis.Regular(100, -5., 5., name="recoil")
-
-htest_data = hist.Hist(axis0, axis1)
-htest_mc = hist.Hist(axis0, axis1)
-
-# print("data.shape", data.shape)
-htest_data.fill(data[:nevt//2,0], data[:nevt//2, 1])
-htest_mc.fill(data[nevt//2:,0], data[nevt//2:, 1])
-
-
-
-
-quant_cdfvals = tf.constant([0.0, 1e-3, 0.02, 0.05, 0.10, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 0.90, 0.95, 0.98, 1.0-1e-3, 1.0], dtype = tf.float64)
-nquants = quant_cdfvals.shape.num_elements()
-
-print("nquants", nquants)
-
-#cdf is in terms of axis1, so shapes need to be compatible
-quant_cdfvals = quant_cdfvals[None, :]
-
-
-# get quantiles from histogram, e.g. to help initialize the parameters for the fit (not actually used here)
-
-# hist_quantiles, hist_quantile_errs = wums.fitutils.hist_to_quantiles(htest, quant_cdfvals, axis=1)
-#
-# print(hist_quantiles)
-# print(hist_quantile_errs)
-#
-# hist_qparms, hist_qparm_errs = wums.fitutils.quantiles_to_qparms(hist_quantiles, hist_quantile_errs)
-#
-# print(hist_qparms)
-# print(hist_qparm_errs)
-
-def parms_to_qparms(xvals, parms):
-
-    parms_2d = tf.reshape(parms, (-1, 2))
-    parms_const = parms_2d[:,0]
-    parms_slope = parms_2d[:,1]
-
-    #cdf is in terms of axis1, so shapes need to be compatible
-    parms_const = parms_const[None, :]
-    parms_slope = parms_slope[None, :]
-
-    qparms = parms_const + parms_slope*xvals[0]
-
-    return qparms
-
-
-def func_transform_cdf(quantile):
-    const_sqrt2 = tf.constant(math.sqrt(2.), quantile.dtype)
-    return 0.5*(1. + tf.math.erf(quantile/const_sqrt2))
-
-def func_transform_quantile(cdf):
-    const_sqrt2 = tf.constant(math.sqrt(2.), cdf.dtype)
-    return const_sqrt2*tf.math.erfinv(2.*cdf - 1.)
-
-# def func_transform_cdf(quantile):
-#     return tf.math.log(quantile/(1.-quantile))
-#
-# def func_transform_quantile(cdf):
-#     return tf.math.sigmoid(cdf)
-
-def func_cdf(xvals, xedges, parms):
-    qparms = parms_to_qparms(xvals, parms)
-    # return wums.fitutils.func_cdf_for_quantile_fit(xvals, xedges, qparms, quant_cdfvals, axis=1)
-
-    return wums.fitutils.func_cdf_for_quantile_fit(xvals, xedges, qparms, quant_cdfvals, axis=1, transform = (func_transform_cdf, func_transform_quantile))
-
-def func_constraint(xvals, xedges, parms):
-    qparms = parms_to_qparms(xvals, parms)
-    return wums.fitutils.func_constraint_for_quantile_fit(xvals, xedges, qparms)
-
-#this is just for plotting
-def func_pdf(h, parms):
-    dtype = tf.float64
-    xvals = [tf.constant(center, dtype=dtype) for center in h.axes.centers]
-    xedges = [tf.constant(edge, dtype=dtype) for edge in h.axes.edges]
-
-    tfparms = tf.constant(parms)
-
-    cdf = func_cdf(xvals, xedges, tfparms)
-
-    pdf = cdf[:,1:] - cdf[:,:-1]
-    pdf = tf.maximum(pdf, tf.zeros_like(pdf))
-
-    return pdf
-
-nparms = nquants-1
-
-
-# print("edges", htest.edges)
-
-# assert(0)
-
-initial_parms_const = np.array([np.log(1./nparms)]*nparms)
-initial_parms_slope = np.zeros_like(initial_parms_const)
-
-initial_parms = np.stack([initial_parms_const, initial_parms_slope], axis=-1)
-initial_parms = np.reshape(initial_parms, (-1,))
-
-res_data = wums.fitutils.fit_hist(htest_data, func_cdf, initial_parms, mode="nll_bin_integrated", norm_axes=[1])
-
-res_mc = wums.fitutils.fit_hist(htest_mc, func_cdf, initial_parms, mode="nll_bin_integrated", norm_axes=[1])
-
-print(res_data)
-
-
-parmvals_data = tf.constant(res_data["x"], tf.float64)
-parmvals_mc = tf.constant(res_mc["x"], tf.float64)
-
-hess_data = res_data["hess"]
-hess_mc = res_mc["hess"]
-
-def get_scaled_eigenvectors(hess, num_null = 2):
-    e,v = np.linalg.eigh(hess)
-
-    # remove the null eigenvectors
-    e = e[None, num_null:]
-    v = v[:, num_null:]
-
-    # scale the eigenvectors
-    vscaled = v/np.sqrt(e)
-
-    return vscaled
-
-vscaled_data = tf.constant(get_scaled_eigenvectors(hess_data), tf.float64)
-vscaled_mc = tf.constant(get_scaled_eigenvectors(hess_data), tf.float64)
-
-print("vscaled_data.shape", vscaled_data.shape)
-
-ut_flat = np.reshape(htest_data.axes.edges[1], (-1,))
-ut_low = tf.constant(ut_flat[0], tf.float64)
-ut_high = tf.constant(ut_flat[-1], tf.float64)
-
-def func_cdf_mc(pt, ut):
-    pts = tf.reshape(pt, (1,1))
-    uts = tf.reshape(ut, (1,1))
-
-    xvals = [pts, None]
-    xedges = [None, uts]
-
-    parms = parmvals_mc
-
-    qparms = parms_to_qparms(xvals, parms)
-
-    ut_axis = 1
-
-    quants = wums.fitutils.qparms_to_quantiles(qparms, x_low = ut_low, x_high = ut_high, axis = ut_axis)
-    spline_edges = xedges[ut_axis]
-
-    cdfvals = wums.fitutils.pchip_interpolate(quants, quant_cdfvals, spline_edges, axis=ut_axis)
-
-    return cdfvals
-
-def func_cdfinv_data(pt, quant):
-    pts = tf.reshape(pt, (1,1))
-    quant_outs = tf.reshape(quant, (1,1))
-
-    xvals = [pts, None]
-    xedges = [None, quant_outs]
-
-    parms = parmvals_data
-
-    qparms = parms_to_qparms(xvals, parms)
-
-    ut_axis = 1
-
-    quants = wums.fitutils.qparms_to_quantiles(qparms, x_low = ut_low, x_high = ut_high, axis = ut_axis)
-    spline_edges = xedges[ut_axis]
-
-    cdfinvvals = wums.fitutils.pchip_interpolate(quant_cdfvals, quants, spline_edges, axis=ut_axis)
-
-    return cdfinvvals
-
-def func_cdfinv_pdf_data(pt, quant):
-    with tf.GradientTape() as t:
-        t.watch(quant)
-        cdfinv = func_cdfinv_data(pt, quant)
-    pdfreciprocal = t.gradient(cdfinv, quant)
-    pdf = 1./pdfreciprocal
-    return cdfinv, pdf
-
-scalar_spec = tf.TensorSpec([], tf.float64)
-
-
-def transform_mc(pt, ut):
-    with tf.GradientTape(persistent=True) as t:
-        t.watch(parmvals_mc)
-        t.watch(parmvals_data)
-
-        cdf_mc = func_cdf_mc(pt, ut)
-        ut_transformed, pdf = func_cdfinv_pdf_data(pt, cdf_mc)
-
-        ut_transformed = tf.reshape(ut_transformed, [])
-        pdf = tf.reshape(pdf, [])
-
-    pdf_grad_mc = t.gradient(pdf, parmvals_mc)
-    pdf_grad_data = t.gradient(pdf, parmvals_data)
-
-    del t
-
-    weight_grad_mc = pdf_grad_mc/pdf
-    weight_grad_data = pdf_grad_data/pdf
-
-    weight_grad_mc = weight_grad_mc[None, :]
-    weight_grad_data = weight_grad_data[None, :]
-
-    weight_grad_mc_eig = weight_grad_mc @ vscaled_mc
-    weight_grad_data_eig = weight_grad_data @ vscaled_data
-
-    weight_grad_mc_eig = tf.reshape(weight_grad_mc_eig, [-1])
-    weight_grad_data_eig = tf.reshape(weight_grad_data_eig, [-1])
-
-    weight_grad_eig = tf.concat([weight_grad_mc_eig, weight_grad_data_eig], axis=0)
-
-    return ut_transformed, weight_grad_eig
-    # return ut_transformed
-
-@tf.function
-def transform_mc_simple(pt, ut):
-    cdf_mc = func_cdf_mc(pt, ut)
-    ut_transformed, pdf = func_cdfinv_pdf_data(pt, cdf_mc)
-
-    ut_transformed = tf.reshape(ut_transformed, [])
-
-    return ut_transformed
-
-
-
-pt_test = tf.constant(0.2, tf.float64)
-ut_test = tf.constant(1.0, tf.float64)
-
-ut, grad = transform_mc(pt_test, ut_test)
-# ut = transform_mc(pt_test, ut_test)
-
-print("shapes", ut.shape, grad.shape)
-
-print("ut", ut)
-print("grad", grad)
-
-input_signature = [scalar_spec, scalar_spec]
-
-class TestMod(tf.Module):
-
-    @tf.function(input_signature =  [scalar_spec, scalar_spec])
-    def __call__(self, pt, ut):
-        return transform_mc(pt, ut)
-
-module = TestMod()
-# tf.saved_model.save(module, "test")
-
-concrete_function = module.__call__.get_concrete_function()
-
-# Convert the model
-converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_function], module)
-
-# converter = tf.lite.TFLiteConverter.from_saved_model("test") # path to the SavedModel directory
-converter.target_spec.supported_ops = [
-  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.
-  tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.
-]
-
-tflite_model = converter.convert()
-
-# print(tflite_model)
-
-# Save the model.
-with open('model.tflite', 'wb') as f:
-  f.write(tflite_model)
-
-
-# onnx_model, _ = tf2onnx.convert.from_function(transform_mc, input_signature)
-# onnx.save(onnx_model, "test.onnx")
-
-
-parmvals = res_data["x"]
-
-
-pdfvals = func_pdf(htest_data, parmvals)
-pdfvals *= htest_data.sum()/np.sum(pdfvals)
-
-
-# hplot = htest[5]
-
-plot = plt.figure()
-plt.yscale("log")
-htest_data[5,:].plot()
-plt.plot(htest_data.axes[1].centers, pdfvals[5])
-# plt.show()
-plot.savefig("test.png")
-
-
-

wums-0.1.8.dist-info/RECORD

@@ -1,16 +0,0 @@
-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.8.dist-info/top_level.txt

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