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

wums/boostHistHelpers.py

@@ -40,7 +40,7 @@ def broadcastSystHist(h1, h2, flow=True, by_ax_name=True):
             h2.ndim - 1 - i: h2.values(flow=flow).shape[h2.ndim - 1 - i]
             for i in range(h2.ndim - h1.ndim)
         }
-
+    
     broadcast_shape = list(moves.values()) + list(s1)
 
     try:
@@ -53,6 +53,10 @@ def broadcastSystHist(h1, h2, flow=True, by_ax_name=True):
             f"    h2.axes: {h2.axes}"
         )
 
+    if by_ax_name:
+        # We also have to move axes that are in common between h1 and h2 but in different order
+        moves.update({sum([k<i for k in moves.keys()]) + h1.axes.name.index(n2): None for i, n2 in enumerate(h2.axes.name) if n2 in h1.axes.name})
+
     # move back to original order
     new_vals = np.moveaxis(new_vals, np.arange(len(moves)), list(moves.keys()))
 

wums/fitutils.py

@@ -0,0 +1,989 @@
+import numpy as np
+import scipy
+import tensorflow as tf
+import math
+
+from numpy import (zeros, where, diff, floor, minimum, maximum, array, concatenate, logical_or, logical_xor,
+                   sqrt)
+
+def cubic_spline_interpolate(xi, yi, x, axis=-1, extrpl=[None, None]):
+
+    # natural cublic spline
+    # if extrpl is given, the spline is linearly extrapolated outside the given boundaries
+    
+    # https://www.math.ntnu.no/emner/TMA4215/2008h/cubicsplines.pdf
+    # https://random-walks.org/content/misc/ncs/ncs.html
+    
+    # move selected axis to the end
+    tensors = [xi, yi]
+    nelems = [tensor.shape.num_elements() for tensor in tensors]
+
+    max_nelems = max(nelems)
+    broadcast_shape = tensors[nelems.index(max_nelems)].shape
+
+    ndim = len(broadcast_shape)
+
+    if xi.shape.num_elements() < max_nelems:
+        xi = tf.broadcast_to(xi, broadcast_shape)
+    if yi.shape.num_elements() < max_nelems:
+        yi = tf.broadcast_to(yi, broadcast_shape)
+
+    # # permutation to move the selected axis to the end
+    selaxis = axis
+    if axis < 0:
+        selaxis = ndim + axis
+    axis = -1
+    permfwd = list(range(ndim))
+    permfwd.remove(selaxis)
+    permfwd.append(selaxis)
+
+    # reverse permutation to restore the original axis order
+    permrev = list(range(ndim))
+    permrev.remove(ndim-1)
+    permrev.insert(selaxis, ndim-1)
+
+
+    ## start spline construction
+    xi = tf.transpose(xi, permfwd)
+    yi = tf.transpose(yi, permfwd)
+    x = tf.transpose(x, permfwd)
+
+    h = tf.experimental.numpy.diff(xi, axis=axis)
+    b = tf.experimental.numpy.diff(yi, axis=axis) / h
+    v = 2 * (h[:, 1:] + h[:, :-1])
+    u = 6 * (b[:, 1:] - b[:, :-1])
+
+    shape = (xi.shape[0], xi.shape[-1]-2, xi.shape[-1]-2)
+    uu = u[:,:,None]
+
+    diag = v
+    superdiag = h[:, 1:-1]
+    subdiag = superdiag
+
+    z = tf.linalg.tridiagonal_solve(diagonals = [superdiag, diag, subdiag], rhs = uu, diagonals_format = "sequence")
+    z = tf.squeeze(z, axis=axis)
+    f = tf.zeros(xi.shape[0], dtype=tf.float64)[:,None]
+    z = tf.concat([f, z, f], axis=axis)
+
+    x_steps = tf.experimental.numpy.diff(x, axis=axis)
+    idx_zero_constant = tf.constant(0, dtype=tf.int64)
+    float64_zero_constant = tf.constant(0., dtype=tf.float64)
+
+    x_compare = x[...,None] < xi[..., None, :]
+    x_compare_all = tf.math.reduce_all(x_compare, axis=axis)
+    x_compare_none = tf.math.reduce_all(tf.logical_not(x_compare), axis=axis)
+    x_index = tf.argmax(x_compare, axis=axis) - 1
+    x_index = tf.where(x_compare_all, idx_zero_constant, x_index)
+    x_index = tf.where(x_compare_none, tf.constant(xi.shape[axis]-2, dtype=tf.int64), x_index)
+
+    nbatch = ndim - 1
+
+    z_xidx = tf.gather(z, x_index, batch_dims=nbatch, axis=-1)
+    z_1pxidx = tf.gather(z, x_index+1, batch_dims=nbatch, axis=-1)
+    h_xidx = tf.gather(h, x_index, batch_dims=nbatch, axis=-1)
+    xi_xidx = tf.gather(xi, x_index, batch_dims=nbatch, axis=-1)
+    xi_1pxidx = tf.gather(xi, x_index+1, batch_dims=nbatch, axis=-1)
+    dxxi = x - xi_xidx
+    dxix = xi_1pxidx - x
+
+    y = z_1pxidx / (6 * h_xidx) * dxxi ** 3 + \
+        z_xidx / (6 * h_xidx) * dxix ** 3 + \
+        (tf.gather(yi, x_index+1, batch_dims=nbatch, axis=-1) / h_xidx - \
+        z_1pxidx * h_xidx / 6) * dxxi + \
+        (tf.gather(yi, x_index, batch_dims=nbatch, axis=-1) / h_xidx - \
+        z_xidx * h_xidx / 6) * dxix
+
+
+    # right side linear extrapolation
+    if extrpl[1] != None:
+        x_max = tf.reshape(tf.constant(extrpl[1], dtype=tf.float64), (1,1))
+
+        # calculate derivative yp_max at boundary
+        x_compare_max = x_max[...,None] < xi[..., None, :]
+        x_compare_max_all = tf.math.reduce_all(x_compare_max, axis=axis)
+        x_compare_max_none = tf.math.reduce_all(tf.logical_not(x_compare_max), axis=axis)
+        x_index_max = tf.argmax(x_compare_max, axis=axis) - 1
+        x_index_max = tf.where(x_compare_max_all, idx_zero_constant, x_index_max)
+        x_index_max = tf.where(x_compare_max_none, tf.constant(xi.shape[axis]-2, dtype=tf.int64), x_index_max)
+
+        z_xidx = tf.gather(z, x_index_max, batch_dims=nbatch, axis=-1)
+        z_1pxidx =  tf.gather(z, x_index_max+1, batch_dims=nbatch, axis=-1)
+        hi_xidx = tf.gather(h, x_index_max, batch_dims=nbatch, axis=-1)
+        xi_xidx = tf.gather(xi, x_index_max, batch_dims=nbatch, axis=-1)
+        xi_1pxidx = tf.gather(xi, x_index_max+1, batch_dims=nbatch, axis=-1)
+        yi_xidx = tf.gather(yi, x_index_max, batch_dims=nbatch, axis=-1)
+        yi_1pxidx = tf.gather(yi, x_index_max+1, batch_dims=nbatch, axis=-1)
+
+        yp_max = z_1pxidx / (2 * hi_xidx) * (x_max - xi_xidx) ** 2 - \
+            z_xidx / (2 * hi_xidx) * (xi_1pxidx - x_max) ** 2 + \
+            1./hi_xidx*(yi_1pxidx - yi_xidx) - hi_xidx/6.*(z_1pxidx - z_xidx)
+
+        # evaluate spline at boundary
+        y_b = cubic_spline_interpolate(xi, yi, x_max, axis=axis) # check shape of x_max
+
+        # replace y by lin for x > x_max
+        extrpl_lin = yp_max*(x-x_max) + y_b
+        cond = x[0,:] >= x_max
+        cond = tf.broadcast_to(cond, (extrpl_lin.shape[0], extrpl_lin.shape[1]))
+        y = tf.where(cond, extrpl_lin, y)
+
+
+    # left side linear extrapolation
+    if extrpl[0] != None:
+        x_min = tf.reshape(tf.constant(extrpl[0], dtype=tf.float64), (1,1))
+
+        # calculate derivative yp_min at boundary
+        x_compare_min = x_min[...,None] >= xi[..., None, :]
+        x_compare_min_all = tf.math.reduce_all(x_compare_min, axis=axis)
+        x_compare_min_none = tf.math.reduce_all(tf.logical_not(x_compare_min), axis=axis)
+        x_index_min = tf.argmax(x_compare_min, axis=axis) 
+        x_index_min = tf.where(x_compare_min_all, idx_zero_constant, x_index_min)
+        x_index_min = tf.where(x_compare_min_none, tf.constant(0, dtype=tf.int64), x_index_min)
+
+        z_xidx = tf.gather(z, x_index_min, batch_dims=nbatch, axis=-1)
+        z_1pxidx =  tf.gather(z, x_index_min+1, batch_dims=nbatch, axis=-1)
+        hi_xidx = tf.gather(h, x_index_min, batch_dims=nbatch, axis=-1)
+        xi_xidx = tf.gather(xi, x_index_min, batch_dims=nbatch, axis=-1)
+        xi_1pxidx = tf.gather(xi, x_index_min+1, batch_dims=nbatch, axis=-1)
+        yi_xidx = tf.gather(yi, x_index_min, batch_dims=nbatch, axis=-1)
+        yi_1pxidx = tf.gather(yi, x_index_min+1, batch_dims=nbatch, axis=-1)
+
+        yp_min = z_1pxidx / (2 * hi_xidx) * (x_min - xi_xidx) ** 2 - \
+            z_xidx / (2 * hi_xidx) * (xi_1pxidx - x_min) ** 2 + \
+            1./hi_xidx*(yi_1pxidx - yi_xidx) - hi_xidx/6.*(z_1pxidx - z_xidx)
+
+        # evaluate spline at boundary
+        y_b = cubic_spline_interpolate(xi, yi, x_min, axis=axis) # check shape of x_max
+
+        # replace y by lin for x > x_min
+        extrpl_lin = yp_min*(x-x_min) + y_b
+        cond = x[0,:] <= x_min
+        cond = tf.broadcast_to(cond, (extrpl_lin.shape[0], extrpl_lin.shape[1]))
+        y = tf.where(cond, extrpl_lin, y)
+
+    ## convert back axes
+    y = tf.transpose(y, permrev)
+    return y
+
+
+
+def pchip_interpolate(xi, yi, x, axis=-1):
+    '''
+        Functionality:
+            1D PCHP interpolation
+        Authors:
+            Michael Taylor <mtaylor@atlanticsciences.com>
+            Mathieu Virbel <mat@meltingrocks.com>
+        Link:
+            https://gist.github.com/tito/553f1135959921ce6699652bf656150d
+            https://github.com/tensorflow/tensorflow/issues/46609#issuecomment-774573667
+    '''
+
+    tensors = [xi, yi]
+    nelems = [tensor.shape.num_elements() for tensor in tensors]
+
+    max_nelems = max(nelems)
+    broadcast_shape = tensors[nelems.index(max_nelems)].shape
+
+    ndim = len(broadcast_shape)
+
+    if xi.shape.num_elements() < max_nelems:
+        xi = tf.broadcast_to(xi, broadcast_shape)
+    if yi.shape.num_elements() < max_nelems:
+        yi = tf.broadcast_to(yi, broadcast_shape)
+
+    # # permutation to move the selected axis to the end
+    selaxis = axis
+    if axis < 0:
+        selaxis = ndim + axis
+
+    permfwd = list(range(ndim))
+    permfwd.remove(selaxis)
+    permfwd.append(selaxis)
+
+    # reverse permutation to restore the original axis order
+    permrev = list(range(ndim))
+    permrev.remove(ndim-1)
+    permrev.insert(selaxis, ndim-1)
+
+    xi = tf.transpose(xi, permfwd)
+    yi = tf.transpose(yi, permfwd)
+    x = tf.transpose(x, permfwd)
+    axis = -1
+
+    xi_steps = tf.experimental.numpy.diff(xi, axis=axis)
+
+
+    x_steps = tf.experimental.numpy.diff(x, axis=axis)
+
+    idx_zero_constant = tf.constant(0, dtype=tf.int64)
+    float64_zero_constant = tf.constant(0., dtype=tf.float64)
+
+    x_compare = x[...,None] < xi[..., None, :]
+    x_compare_all = tf.math.reduce_all(x_compare, axis=-1)
+    x_compare_none = tf.math.reduce_all(tf.logical_not(x_compare), axis=-1)
+    x_index = tf.argmax(x_compare, axis = -1) - 1
+
+    x_index = tf.where(x_compare_all, idx_zero_constant, x_index)
+    x_index = tf.where(x_compare_none, tf.constant(xi.shape[axis]-2, dtype=tf.int64), x_index)
+
+    # Calculate gradients d
+    h = tf.experimental.numpy.diff(xi, axis=axis)
+
+    d = tf.zeros_like(xi)
+
+    delta = tf.experimental.numpy.diff(yi, axis=axis) / h
+    # mode=='mono', Fritsch-Carlson algorithm from fortran numerical
+    # recipe
+
+    slice01 = [slice(None)]*ndim
+    slice01[axis] = slice(0,1)
+    slice01 = tuple(slice01)
+
+    slice0m1 = [slice(None)]*ndim
+    slice0m1[axis] = slice(0,-1)
+    slice0m1 = tuple(slice0m1)
+
+    slice1 = [slice(None)]*ndim
+    slice1[axis] = slice(1,None)
+    slice1 = tuple(slice1)
+
+    slicem1 = [slice(None)]*ndim
+    slicem1[axis] = slice(-1,None)
+    slicem1 = tuple(slicem1)
+
+    d = tf.concat(
+        (delta[slice01], 3 * (h[slice0m1] + h[slice1]) / ((h[slice0m1] + 2 * h[slice1]) / delta[slice0m1] +
+                                                (2 * h[slice0m1] + h[slice1]) / delta[slice1]), delta[slicem1]), axis=axis)
+
+    false_shape = [*xi.shape]
+    false_shape[axis] = 1
+    false_const = tf.fill(false_shape, False)
+
+    mask = tf.concat((false_const, tf.math.logical_xor(delta[slice0m1] > 0, delta[slice1] > 0), false_const), axis=axis)
+    d = tf.where(mask, float64_zero_constant, d)
+
+    mask = tf.math.logical_or(tf.concat((false_const, delta == 0), axis=axis), tf.concat((delta == 0, false_const), axis=axis))
+    d = tf.where(mask, float64_zero_constant, d)
+
+    xiperm = xi
+    yiperm = yi
+    dperm = d
+    hperm = h
+
+    nbatch = ndim - 1
+
+    xi_xidx = tf.gather(xiperm, x_index, batch_dims=nbatch, axis=-1)
+    xi_1pxidx = tf.gather(xiperm, 1 + x_index, batch_dims=nbatch, axis=-1)
+    yi_xidx = tf.gather(yiperm, x_index, batch_dims=nbatch, axis=-1)
+    yi_1pxidx = tf.gather(yiperm, 1 + x_index, batch_dims=nbatch, axis=-1)
+    d_xidx = tf.gather(dperm, x_index, batch_dims=nbatch, axis=-1)
+    d_1pxidx = tf.gather(dperm, 1 + x_index, batch_dims=nbatch, axis=-1)
+    h_xidx = tf.gather(hperm, x_index, batch_dims=nbatch, axis=-1)
+
+    dxxi = x - xi_xidx
+    dxxid = x - xi_1pxidx
+    dxxi2 = tf.math.pow(dxxi, 2)
+    dxxid2 = tf.math.pow(dxxid, 2)
+
+    y = (2 / tf.math.pow(h_xidx, 3) *
+            (yi_xidx * dxxid2 * (dxxi + h_xidx / 2) - yi_1pxidx * dxxi2 *
+            (dxxid - h_xidx / 2)) + 1 / tf.math.pow(h_xidx, 2) *
+            (d_xidx * dxxid2 * dxxi + d_1pxidx * dxxi2 * dxxid))
+
+    y = tf.transpose(y, permrev)
+
+    return y
+
+def pchip_interpolate_np(xi, yi, x, mode="mono", verbose=False):
+    '''
+        Functionality:
+            1D PCHP interpolation
+        Authors:
+            Michael Taylor <mtaylor@atlanticsciences.com>
+            Mathieu Virbel <mat@meltingrocks.com>
+        Link:
+            https://gist.github.com/tito/553f1135959921ce6699652bf656150d
+    '''
+
+    if mode not in ("mono", "quad"):
+        raise ValueError("Unrecognized mode string")
+
+    # Search for [xi,xi+1] interval for each x
+    xi = xi.astype("double")
+    yi = yi.astype("double")
+
+    x_index = zeros(len(x), dtype="int")
+    xi_steps = diff(xi)
+    if not all(xi_steps > 0):
+        raise ValueError("x-coordinates are not in increasing order.")
+
+    x_steps = diff(x)
+    if xi_steps.max() / xi_steps.min() < 1.000001:
+        # uniform input grid
+        if verbose:
+            print("pchip: uniform input grid")
+        xi_start = xi[0]
+        xi_step = (xi[-1] - xi[0]) / (len(xi) - 1)
+        x_index = minimum(maximum(floor((x - xi_start) / xi_step).astype(int), 0), len(xi) - 2)
+
+        # Calculate gradients d
+        h = (xi[-1] - xi[0]) / (len(xi) - 1)
+        d = zeros(len(xi), dtype="double")
+        if mode == "quad":
+            # quadratic polynomial fit
+            d[[0]] = (yi[1] - yi[0]) / h
+            d[[-1]] = (yi[-1] - yi[-2]) / h
+            d[1:-1] = (yi[2:] - yi[0:-2]) / 2 / h
+        else:
+            # mode=='mono', Fritsch-Carlson algorithm from fortran numerical
+            # recipe
+            delta = diff(yi) / h
+            d = concatenate((delta[0:1], 2 / (1 / delta[0:-1] + 1 / delta[1:]), delta[-1:]))
+            d[concatenate((array([False]), logical_xor(delta[0:-1] > 0, delta[1:] > 0), array([False])))] = 0
+            d[logical_or(concatenate((array([False]), delta == 0)), concatenate(
+                (delta == 0, array([False]))))] = 0
+        # Calculate output values y
+        dxxi = x - xi[x_index]
+        dxxid = x - xi[1 + x_index]
+        dxxi2 = pow(dxxi, 2)
+        dxxid2 = pow(dxxid, 2)
+        y = (2 / pow(h, 3) * (yi[x_index] * dxxid2 * (dxxi + h / 2) - yi[1 + x_index] * dxxi2 *
+                              (dxxid - h / 2)) + 1 / pow(h, 2) *
+             (d[x_index] * dxxid2 * dxxi + d[1 + x_index] * dxxi2 * dxxid))
+    else:
+        # not uniform input grid
+        if (x_steps.max() / x_steps.min() < 1.000001 and x_steps.max() / x_steps.min() > 0.999999):
+            # non-uniform input grid, uniform output grid
+            if verbose:
+                print("pchip: non-uniform input grid, uniform output grid")
+            x_decreasing = x[-1] < x[0]
+            if x_decreasing:
+                x = x[::-1]
+            x_start = x[0]
+            x_step = (x[-1] - x[0]) / (len(x) - 1)
+            x_indexprev = -1
+            for xi_loop in range(len(xi) - 2):
+                x_indexcur = max(int(floor((xi[1 + xi_loop] - x_start) / x_step)), -1)
+                x_index[1 + x_indexprev:1 + x_indexcur] = xi_loop
+                x_indexprev = x_indexcur
+            x_index[1 + x_indexprev:] = len(xi) - 2
+            if x_decreasing:
+                x = x[::-1]
+                x_index = x_index[::-1]
+        elif all(x_steps > 0) or all(x_steps < 0):
+            # non-uniform input/output grids, output grid monotonic
+            if verbose:
+                print("pchip: non-uniform in/out grid, output grid monotonic")
+            x_decreasing = x[-1] < x[0]
+            if x_decreasing:
+                x = x[::-1]
+            x_len = len(x)
+            x_loop = 0
+            for xi_loop in range(len(xi) - 1):
+                while x_loop < x_len and x[x_loop] < xi[1 + xi_loop]:
+                    x_index[x_loop] = xi_loop
+                    x_loop += 1
+            x_index[x_loop:] = len(xi) - 2
+            if x_decreasing:
+                x = x[::-1]
+                x_index = x_index[::-1]
+        else:
+            # non-uniform input/output grids, output grid not monotonic
+            if verbose:
+                print("pchip: non-uniform in/out grids, " "output grid not monotonic")
+            for index in range(len(x)):
+                loc = where(x[index] < xi)[0]
+                if loc.size == 0:
+                    x_index[index] = len(xi) - 2
+                elif loc[0] == 0:
+                    x_index[index] = 0
+                else:
+                    x_index[index] = loc[0] - 1
+        # Calculate gradients d
+        h = diff(xi)
+        d = zeros(len(xi), dtype="double")
+        delta = diff(yi) / h
+        if mode == "quad":
+            # quadratic polynomial fit
+            d[[0, -1]] = delta[[0, -1]]
+            d[1:-1] = (delta[1:] * h[0:-1] + delta[0:-1] * h[1:]) / (h[0:-1] + h[1:])
+        else:
+            # mode=='mono', Fritsch-Carlson algorithm from fortran numerical
+            # recipe
+            d = concatenate(
+                (delta[0:1], 3 * (h[0:-1] + h[1:]) / ((h[0:-1] + 2 * h[1:]) / delta[0:-1] +
+                                                      (2 * h[0:-1] + h[1:]) / delta[1:]), delta[-1:]))
+            d[concatenate((array([False]), logical_xor(delta[0:-1] > 0, delta[1:] > 0), array([False])))] = 0
+            d[logical_or(concatenate((array([False]), delta == 0)), concatenate(
+                (delta == 0, array([False]))))] = 0
+        dxxi = x - xi[x_index]
+        dxxid = x - xi[1 + x_index]
+        dxxi2 = pow(dxxi, 2)
+        dxxid2 = pow(dxxid, 2)
+        y = (2 / pow(h[x_index], 3) *
+             (yi[x_index] * dxxid2 * (dxxi + h[x_index] / 2) - yi[1 + x_index] * dxxi2 *
+              (dxxid - h[x_index] / 2)) + 1 / pow(h[x_index], 2) *
+             (d[x_index] * dxxid2 * dxxi + d[1 + x_index] * dxxi2 * dxxid))
+    return y
+
+
+
+def pchip_interpolate_np_forced(xi, yi, x, mode="mono", verbose=False):
+    '''
+        Functionality:
+            1D PCHP interpolation
+        Authors:
+            Michael Taylor <mtaylor@atlanticsciences.com>
+            Mathieu Virbel <mat@meltingrocks.com>
+        Link:
+            https://gist.github.com/tito/553f1135959921ce6699652bf656150d
+    '''
+
+    if mode not in ("mono", "quad"):
+        raise ValueError("Unrecognized mode string")
+
+    # Search for [xi,xi+1] interval for each x
+    xi = xi.astype("double")
+    yi = yi.astype("double")
+
+    x_index = zeros(len(x), dtype="int")
+    xi_steps = diff(xi)
+    if not all(xi_steps > 0):
+        raise ValueError("x-coordinates are not in increasing order.")
+
+    x_steps = diff(x)
+    # if xi_steps.max() / xi_steps.min() < 1.000001:
+    if False:
+        # uniform input grid
+        if verbose:
+            print("pchip: uniform input grid")
+        xi_start = xi[0]
+        xi_step = (xi[-1] - xi[0]) / (len(xi) - 1)
+        x_index = minimum(maximum(floor((x - xi_start) / xi_step).astype(int), 0), len(xi) - 2)
+
+        # Calculate gradients d
+        h = (xi[-1] - xi[0]) / (len(xi) - 1)
+        d = zeros(len(xi), dtype="double")
+        if mode == "quad":
+            # quadratic polynomial fit
+            d[[0]] = (yi[1] - yi[0]) / h
+            d[[-1]] = (yi[-1] - yi[-2]) / h
+            d[1:-1] = (yi[2:] - yi[0:-2]) / 2 / h
+        else:
+            # mode=='mono', Fritsch-Carlson algorithm from fortran numerical
+            # recipe
+            delta = diff(yi) / h
+            d = concatenate((delta[0:1], 2 / (1 / delta[0:-1] + 1 / delta[1:]), delta[-1:]))
+            d[concatenate((array([False]), logical_xor(delta[0:-1] > 0, delta[1:] > 0), array([False])))] = 0
+            d[logical_or(concatenate((array([False]), delta == 0)), concatenate(
+                (delta == 0, array([False]))))] = 0
+        # Calculate output values y
+        dxxi = x - xi[x_index]
+        dxxid = x - xi[1 + x_index]
+        dxxi2 = pow(dxxi, 2)
+        dxxid2 = pow(dxxid, 2)
+        y = (2 / pow(h, 3) * (yi[x_index] * dxxid2 * (dxxi + h / 2) - yi[1 + x_index] * dxxi2 *
+                              (dxxid - h / 2)) + 1 / pow(h, 2) *
+             (d[x_index] * dxxid2 * dxxi + d[1 + x_index] * dxxi2 * dxxid))
+    else:
+        # not uniform input grid
+        # if (x_steps.max() / x_steps.min() < 1.000001 and x_steps.max() / x_steps.min() > 0.999999):
+        if False:
+            # non-uniform input grid, uniform output grid
+            if verbose:
+                print("pchip: non-uniform input grid, uniform output grid")
+            x_decreasing = x[-1] < x[0]
+            if x_decreasing:
+                x = x[::-1]
+            x_start = x[0]
+            x_step = (x[-1] - x[0]) / (len(x) - 1)
+            x_indexprev = -1
+            for xi_loop in range(len(xi) - 2):
+                x_indexcur = max(int(floor((xi[1 + xi_loop] - x_start) / x_step)), -1)
+                x_index[1 + x_indexprev:1 + x_indexcur] = xi_loop
+                x_indexprev = x_indexcur
+            x_index[1 + x_indexprev:] = len(xi) - 2
+            if x_decreasing:
+                x = x[::-1]
+                x_index = x_index[::-1]
+        # elif all(x_steps > 0) or all(x_steps < 0):
+        elif True:
+            # non-uniform input/output grids, output grid monotonic
+            if verbose:
+                print("pchip: non-uniform in/out grid, output grid monotonic")
+            # x_decreasing = x[-1] < x[0]
+            x_decreasing = False
+            if x_decreasing:
+                x = x[::-1]
+            x_len = len(x)
+            x_loop = 0
+            for xi_loop in range(len(xi) - 1):
+                while x_loop < x_len and x[x_loop] < xi[1 + xi_loop]:
+                    x_index[x_loop] = xi_loop
+                    x_loop += 1
+            x_index[x_loop:] = len(xi) - 2
+
+            print("np_forced x_index", x_index)
+            if x_decreasing:
+                x = x[::-1]
+                x_index = x_index[::-1]
+        else:
+            # non-uniform input/output grids, output grid not monotonic
+            if verbose:
+                print("pchip: non-uniform in/out grids, " "output grid not monotonic")
+            for index in range(len(x)):
+                loc = where(x[index] < xi)[0]
+                if loc.size == 0:
+                    x_index[index] = len(xi) - 2
+                elif loc[0] == 0:
+                    x_index[index] = 0
+                else:
+                    x_index[index] = loc[0] - 1
+        # Calculate gradients d
+        h = diff(xi)
+        d = zeros(len(xi), dtype="double")
+        delta = diff(yi) / h
+        if mode == "quad":
+            # quadratic polynomial fit
+            d[[0, -1]] = delta[[0, -1]]
+            d[1:-1] = (delta[1:] * h[0:-1] + delta[0:-1] * h[1:]) / (h[0:-1] + h[1:])
+        else:
+            # mode=='mono', Fritsch-Carlson algorithm from fortran numerical
+            # recipe
+            d = concatenate(
+                (delta[0:1], 3 * (h[0:-1] + h[1:]) / ((h[0:-1] + 2 * h[1:]) / delta[0:-1] +
+                                                      (2 * h[0:-1] + h[1:]) / delta[1:]), delta[-1:]))
+            d[concatenate((array([False]), logical_xor(delta[0:-1] > 0, delta[1:] > 0), array([False])))] = 0
+            d[logical_or(concatenate((array([False]), delta == 0)), concatenate(
+                (delta == 0, array([False]))))] = 0
+        dxxi = x - xi[x_index]
+        dxxid = x - xi[1 + x_index]
+        dxxi2 = pow(dxxi, 2)
+        dxxid2 = pow(dxxid, 2)
+        y = (2 / pow(h[x_index], 3) *
+             (yi[x_index] * dxxid2 * (dxxi + h[x_index] / 2) - yi[1 + x_index] * dxxi2 *
+              (dxxid - h[x_index] / 2)) + 1 / pow(h[x_index], 2) *
+             (d[x_index] * dxxid2 * dxxi + d[1 + x_index] * dxxi2 * dxxid))
+    return y
+
+def qparms_to_quantiles(qparms, x_low = 0., x_high = 1., axis = -1):
+    deltax = tf.exp(qparms)
+    sumdeltax = tf.math.reduce_sum(deltax, axis=axis, keepdims=True)
+
+    deltaxnorm = deltax/sumdeltax
+
+    x0shape = list(deltaxnorm.shape)
+    x0shape[axis] = 1
+    x0 = tf.fill(x0shape, x_low)
+    x0 = tf.cast(x0, tf.float64)
+
+    deltaxfull = (x_high - x_low)*deltaxnorm
+    deltaxfull = tf.concat([x0, deltaxfull], axis = axis)
+
+    quants = tf.math.cumsum(deltaxfull, axis=axis)
+
+    return quants
+
+
+
+def quantiles_to_qparms(quants, quant_errs = None, x_low = 0., x_high = 1., axis = -1):
+
+    deltaxfull = tf.experimental.numpy.diff(quants, axis=axis)
+    deltaxnorm = deltaxfull/(x_high - x_low)
+    qparms = tf.math.log(deltaxnorm)
+
+    if quant_errs is not None:
+        quant_vars = tf.math.square(quant_errs)
+
+        ndim = len(quant_errs.shape)
+
+        slicem1 = [slice(None)]*ndim
+        slicem1[axis] = slice(None,-1)
+        slicem1 = tuple(slicem1)
+
+        slice1 = [slice(None)]*ndim
+        slice1[axis] = slice(1,None)
+        slice1 = tuple(slice1)
+
+        deltaxfull_vars = quant_vars[slice1] + quant_vars[slicem1]
+        deltaxfull_errs = tf.math.sqrt(deltaxfull_vars)
+
+        qparm_errs = deltaxfull_errs/deltaxfull
+
+        return qparms, qparm_errs
+    else:
+        return qparms
+
+
+def hist_to_quantiles(h, quant_cdfvals, axis = -1):
+    dtype = tf.float64
+
+    xvals = [tf.constant(center, dtype=dtype) for center in h.axes.centers]
+    xwidths = [tf.constant(width, dtype=dtype) for width in h.axes.widths]
+    xedges = [tf.constant(edge, dtype=dtype) for edge in h.axes.edges]
+    yvals = tf.constant(h.values(), dtype=dtype)
+
+    if not isinstance(quant_cdfvals, tf.Tensor):
+        quant_cdfvals = tf.constant(quant_cdfvals, tf.float64)
+
+    x_flat = tf.reshape(xedges[axis], (-1,))
+    x_low = x_flat[0]
+    x_high = x_flat[-1]
+
+    hist_cdfvals = tf.cumsum(yvals, axis=axis)/tf.reduce_sum(yvals, axis=axis, keepdims=True)
+
+    x0shape = list(hist_cdfvals.shape)
+    x0shape[axis] = 1
+    x0 = tf.zeros(x0shape, dtype = dtype)
+
+    hist_cdfvals = tf.concat([x0, hist_cdfvals], axis=axis)
+
+    quants = pchip_interpolate(hist_cdfvals, xedges[axis], quant_cdfvals, axis=axis)
+
+    quants = tf.where(quant_cdfvals == 0., x_low, quants)
+    quants = tf.where(quant_cdfvals == 1., x_high, quants)
+
+    ntot = tf.math.reduce_sum(yvals, axis=axis, keepdims=True)
+
+    quant_cdf_bar = ntot/(1.+ntot)*(quant_cdfvals + 0.5/ntot)
+    quant_cdfval_errs = ntot/(1.+ntot)*tf.math.sqrt(quant_cdfvals*(1.-quant_cdfvals)/ntot + 0.25/ntot/ntot)
+
+    quant_cdfvals_up = quant_cdf_bar + quant_cdfval_errs
+    quant_cdfvals_up = tf.clip_by_value(quant_cdfvals_up, 0., 1.)
+
+    quant_cdfvals_down = quant_cdf_bar - quant_cdfval_errs
+    quant_cdfvals_down = tf.clip_by_value(quant_cdfvals_down, 0., 1.)
+
+    quants_up = pchip_interpolate(hist_cdfvals, xedges[axis], quant_cdfvals_up, axis=axis)
+    quants_up = tf.where(quant_cdfvals_up == 0., x_low, quants_up)
+    quants_up = tf.where(quant_cdfvals_up == 1., x_high, quants_up)
+
+    quants_down = pchip_interpolate(hist_cdfvals, xedges[axis], quant_cdfvals_down, axis=axis)
+    quants_down = tf.where(quant_cdfvals_down == 0., x_low, quants_down)
+    quants_down = tf.where(quant_cdfvals_down == 1., x_high, quants_down)
+
+    quant_errs = 0.5*(quants_up - quants_down)
+
+    zero_const = tf.constant(0., dtype)
+
+    quant_errs = tf.where(quant_cdfvals == 0., zero_const, quant_errs)
+    quant_errs = tf.where(quant_cdfvals == 1., zero_const, quant_errs)
+
+    return quants.numpy(), quant_errs.numpy()
+
+def func_cdf_for_quantile_fit(xvals, xedges, qparms, quant_cdfvals, axis=-1, transform = None):
+    x_flat = tf.reshape(xedges[axis], (-1,))
+    x_low = x_flat[0]
+    x_high = x_flat[-1]
+
+    quants = qparms_to_quantiles(qparms, x_low = x_low, x_high = x_high, axis = axis)
+
+    spline_edges = xedges[axis]
+
+    ndim = len(xvals)
+
+    if transform is not None:
+        transform_cdf, transform_quantile = transform
+
+        slicelim = [slice(None)]*ndim
+        slicelim[axis] = slice(1, -1)
+        slicelim = tuple(slicelim)
+
+        quants = quants[slicelim]
+        quant_cdfvals = quant_cdfvals[slicelim]
+
+        quant_cdfvals = transform_quantile(quant_cdfvals)
+
+    cdfvals = pchip_interpolate(quants, quant_cdfvals, spline_edges, axis=axis)
+
+    if transform is not None:
+        cdfvals = transform_cdf(cdfvals)
+
+    slicefirst = [slice(None)]*ndim
+    slicefirst[axis] = slice(None, 1)
+    slicefirst = tuple(slicefirst)
+
+    slicelast = [slice(None)]*ndim
+    slicelast[axis] = slice(-1, None)
+    slicelast = tuple(slicelast)
+
+    cdfvals = (cdfvals - cdfvals[slicefirst])/(cdfvals[slicelast] - cdfvals[slicefirst])
+
+    return cdfvals
+
+def func_constraint_for_quantile_fit(xvals, xedges, qparms, axis=-1):
+    constraints = 0.5*tf.math.square(tf.math.reduce_sum(tf.exp(qparms), axis=axis) - 1.)
+    constraint = tf.math.reduce_sum(constraints)
+    return constraint
+
+@tf.function
+def val_grad(func, *args, **kwargs):
+    xdep = kwargs["parms"]
+    with tf.GradientTape() as t1:
+        t1.watch(xdep)
+        val = func(*args, **kwargs)
+    grad = t1.gradient(val, xdep)
+    return val, grad
+
+#TODO forward-over-reverse also here?
+@tf.function
+def val_grad_hess(func, *args, **kwargs):
+    xdep = kwargs["parms"]
+    with tf.GradientTape() as t2:
+        t2.watch(xdep)
+        with tf.GradientTape() as t1:
+            t1.watch(xdep)
+            val = func(*args, **kwargs)
+        grad = t1.gradient(val, xdep)
+    hess = t2.jacobian(grad, xdep)
+
+    return val, grad, hess    
+
+@tf.function
+def val_grad_hessp(func, p, *args, **kwargs):
+    xdep = kwargs["parms"]
+    with tf.autodiff.ForwardAccumulator(xdep, p) as acc:
+        with tf.GradientTape() as grad_tape:
+            grad_tape.watch(xdep)
+            val = func(*args, **kwargs)
+        grad = grad_tape.gradient(val, xdep)
+    hessp = acc.jvp(grad)
+  
+    return val, grad, hessp
+
+def loss_with_constraint(func_loss, parms, func_constraint = None, args_loss = (), extra_args_loss=(), args_constraint = (), extra_args_constraint = ()):
+    loss = func_loss(parms, *args_loss, *extra_args_loss)
+    if func_constraint is not None:
+        loss += func_constraint(*args_constraint, parms, *extra_args_constraint)
+
+    return loss
+
+def chisq_loss(parms, xvals, xwidths, xedges, yvals, yvariances, func, norm_axes = None, *args):
+    fvals = func(xvals, parms, *args)
+
+    # exclude zero-variance bins
+    variances_safe = tf.where(yvariances == 0., tf.ones_like(yvariances), yvariances)
+    chisqv = (fvals - yvals)**2/variances_safe
+    chisqv_safe = tf.where(yvariances == 0., tf.zeros_like(chisqv), chisqv)
+    return tf.reduce_sum(chisqv_safe)
+
+def chisq_normalized_loss(parms, xvals, xwidths, xedges, yvals, yvariances, func, norm_axes = None, *args):
+    fvals = func(xvals, parms, *args)
+    norm = tf.reduce_sum(fvals, keepdims=True, axis = norm_axes)
+    sumw = tf.reduce_sum(yvals, keepdims=True, axis = norm_axes)
+    if norm_axes is None:
+        for xwidth in xwidths:
+            norm *= xwidth
+    else:
+        for norm_axis in norm_axes:
+            norm *= xwidths[norm_axis]
+            
+    # exclude zero-variance bins
+    variances_safe = tf.where(yvariances == 0., tf.ones_like(yvariances), yvariances)
+    chisqv = (sumw*fvals/norm - yvals)**2/variances_safe
+    chisqv_safe = tf.where(yvariances == 0., tf.zeros_like(chisqv), chisqv)
+    return tf.reduce_sum(chisqv_safe)
+
+def nll_loss(parms, xvals, xwidths, xedges, yvals, yvariances, func, norm_axes = None, *args):
+    fvals = func(xvals, parms, *args)
+
+    # compute overall scaling needed to restore mean == variance condition
+    yval_total = tf.reduce_sum(yvals, keepdims = True, axis = norm_axes)
+    variance_total = tf.reduce_sum(yvariances, keepdims = True, axis = norm_axes)
+    isnull_total = variance_total == 0.
+    variance_total_safe = tf.where(isnull_total, tf.ones_like(variance_total), variance_total)
+    scale_total = yval_total/variance_total_safe
+    scale_total_safe = tf.where(isnull_total, tf.ones_like(scale_total), scale_total)
+
+    # skip likelihood calculation for empty bins to avoid inf or nan
+    # compute per-bin scaling needed to restore mean == variance condition, falling
+    # back to overall scaling for empty bins
+    isnull = tf.logical_or(yvals == 0., yvariances == 0.)
+    variances_safe = tf.where(isnull, tf.ones_like(yvariances), yvariances)
+    scale = yvals/variances_safe
+    scale_safe = tf.where(isnull, scale_total_safe*tf.ones_like(scale), scale)
+
+    norm = tf.reduce_sum(scale_safe*fvals, keepdims=True, axis = norm_axes)
+    if norm_axes is None:
+        for xwidth in xwidths:
+            norm *= xwidth
+    else:
+        for norm_axis in norm_axes:
+            norm *= xwidths[norm_axis]
+
+    fvalsnorm = fvals/norm
+
+    fvalsnorm_safe = tf.where(isnull, tf.ones_like(fvalsnorm), fvalsnorm)
+    nllv = -scale_safe*yvals*tf.math.log(scale_safe*fvalsnorm_safe)
+    nllv_safe = tf.where(isnull, tf.zeros_like(nllv), nllv)
+    nllsum = tf.reduce_sum(nllv_safe)
+    return nllsum
+
+def nll_loss_bin_integrated(parms, xvals, xwidths, xedges, yvals, yvariances, func, norm_axes = None, *args):
+    #TODO reduce code duplication with nll_loss_bin
+
+    norm_axis = 0
+    if norm_axes is not None:
+        if len(norm_axes) > 1:
+            raise ValueError("Only 1 nomralization access supported for bin-integrated nll")
+        norm_axis = norm_axes[0]
+
+    cdfvals = func(xvals, xedges, parms, *args)
+
+    slices_low = [slice(None)]*len(cdfvals.shape)
+    slices_low[norm_axis] = slice(None,-1)
+
+    slices_high = [slice(None)]*len(cdfvals.shape)
+    slices_high[norm_axis] = slice(1,None)
+
+    # bin_integrals = cdfvals[1:] - cdfvals[:-1]
+    bin_integrals = cdfvals[tuple(slices_high)] - cdfvals[tuple(slices_low)]
+    bin_integrals = tf.maximum(bin_integrals, tf.zeros_like(bin_integrals))
+
+    fvals = bin_integrals
+
+    # compute overall scaling needed to restore mean == variance condition
+    yval_total = tf.reduce_sum(yvals, keepdims = True, axis = norm_axes)
+    variance_total = tf.reduce_sum(yvariances, keepdims = True, axis = norm_axes)
+    isnull_total = variance_total == 0.
+    variance_total_safe = tf.where(isnull_total, tf.ones_like(variance_total), variance_total)
+    scale_total = yval_total/variance_total_safe
+    scale_total_safe = tf.where(isnull_total, tf.ones_like(scale_total), scale_total)
+
+    # skip likelihood calculation for empty bins to avoid inf or nan
+    # compute per-bin scaling needed to restore mean == variance condition, falling
+    # back to overall scaling for empty bins
+    isnull = tf.logical_or(yvals == 0., yvariances == 0.)
+    variances_safe = tf.where(isnull, tf.ones_like(yvariances), yvariances)
+    scale = yvals/variances_safe
+    scale_safe = tf.where(isnull, scale_total_safe*tf.ones_like(scale), scale)
+
+    norm = tf.reduce_sum(scale_safe*fvals, keepdims=True, axis = norm_axes)
+
+    fvalsnorm = fvals/norm
+
+    fvalsnorm_safe = tf.where(isnull, tf.ones_like(fvalsnorm), fvalsnorm)
+    nllv = -scale_safe*yvals*tf.math.log(scale_safe*fvalsnorm_safe)
+    nllv_safe = tf.where(isnull, tf.zeros_like(nllv), nllv)
+    nllsum = tf.reduce_sum(nllv_safe)
+    return nllsum
+
+def chisq_loss_bin_integrated(parms, xvals, xwidths, xedges, yvals, yvariances, func, norm_axes = None, *args):
+    #FIXME this is only defined in 1D for now
+    cdfvals = func(xedges, parms, *args)
+    bin_integrals = cdfvals[1:] - cdfvals[:-1]
+    bin_integrals = tf.maximum(bin_integrals, tf.zeros_like(bin_integrals))
+
+    fvals = bin_integrals
+
+    # exclude zero-variance bins
+    variances_safe = tf.where(yvariances == 0., tf.ones_like(yvariances), yvariances)
+    chisqv = (fvals - yvals)**2/variances_safe
+    chisqv_safe = tf.where(yvariances == 0., tf.zeros_like(chisqv), chisqv)
+    chisqsum = tf.reduce_sum(chisqv_safe)
+
+    return chisqsum
+
+
+def fit_hist(hist, func, initial_parmvals, max_iter = 5, edmtol = 1e-5, mode = "chisq", norm_axes = None, func_constraint = None,  args = (), args_constraint=()):
+
+    dtype = tf.float64
+
+    xvals = [tf.constant(center, dtype=dtype) for center in hist.axes.centers]
+    xwidths = [tf.constant(width, dtype=dtype) for width in hist.axes.widths]
+    xedges = [tf.constant(edge, dtype=dtype) for edge in hist.axes.edges]
+    yvals = tf.constant(hist.values(), dtype=dtype)
+    yvariances = tf.constant(hist.variances(), dtype=dtype)
+
+    covscale = 1.
+    if mode == "chisq":
+        floss = chisq_loss
+        covscale = 2.
+    elif mode == "nll":
+        floss = nll_loss
+    elif mode == "nll_bin_integrated":
+        floss = nll_loss_bin_integrated
+    elif mode == "chisq_normalized":
+        floss = chisq_normalized_loss
+        covscale = 2.
+    elif mode == "chisq_loss_bin_integrated":
+        floss = chisq_loss_bin_integrated
+        covscale = 2.
+    elif mode == "nll_extended":
+        raise Exception("Not Implemented")
+    else:
+        raise Exception("unsupported mode")
+
+    val_grad_args = { "func_loss" : floss,
+                    "func_constraint" : func_constraint,
+                    "args_loss" : (xvals, xwidths, xedges, yvals, yvariances, func, norm_axes),
+                    "extra_args_loss" : args,
+                    "args_constraint" : (xvals, xedges),
+                    "extra_args_constraint" : args_constraint}
+
+    def scipy_loss(parmvals, *args):
+        parms = tf.constant(parmvals, dtype=dtype)
+
+        # loss, grad = val_grad(floss, parms, xvals, xwidths, xedges, yvals, yvariances, func, norm_axes, *args)
+        loss, grad = val_grad(loss_with_constraint, parms=parms, **val_grad_args)
+        return loss.numpy(), grad.numpy()
+
+    def scipy_hessp(parmvals, p, *args):
+        parms = tf.constant(parmvals, dtype=dtype)
+
+        # loss, grad, hessp = val_grad_hessp(floss, p, parms, xvals, xwidths, xedges, yvals, yvariances, func, norm_axes, *args)
+        loss, grad, hessp = val_grad_hessp(loss_with_constraint, p, parms=parms, **val_grad_args)
+        return hessp.numpy()
+
+    current_parmvals = initial_parmvals
+    for iiter in range(max_iter):
+
+        res = scipy.optimize.minimize(scipy_loss, current_parmvals, method = "trust-krylov", jac = True, hessp = scipy_hessp, args = args)
+
+        current_parmvals = res.x
+
+        parms = tf.constant(current_parmvals, dtype=dtype)
+
+        # loss, grad, hess = val_grad_hess(floss, parms, xvals, xwidths, xedges, yvals, yvariances, func, norm_axes, *args)
+        loss, grad, hess = val_grad_hess(loss_with_constraint, parms=parms, **val_grad_args)
+        loss, grad, hess = loss.numpy(), grad.numpy(), hess.numpy()
+        
+        try:
+            eigvals = np.linalg.eigvalsh(hess)
+            gradv = grad[:, np.newaxis]
+            edmval = 0.5*gradv.transpose()@np.linalg.solve(hess, gradv)
+            edmval = edmval[0][0]
+        except np.linalg.LinAlgError:
+            eigvals = np.zeros_like(grad)
+            edmval = 99.
+
+        converged = edmval < edmtol and np.abs(edmval) >= 0. and eigvals[0] > 0.
+        if converged:
+            break
+
+    status = 1
+    covstatus = 1
+    
+    if edmval < edmtol and edmval >= -0.:
+        status = 0
+    if eigvals[0] > 0.:
+        covstatus = 0
+
+    try:
+        cov = covscale*np.linalg.inv(hess)
+    except np.linalg.LinAlgError:
+        cov = np.zeros_like(hess)
+        covstatus = 1
+
+    res = { "x" : current_parmvals,
+           "hess" : hess,
+           "cov" : cov,
+           "status" : status,
+           "covstatus" : covstatus,
+           "hess_eigvals" : eigvals,
+           "edmval" : edmval,
+           "loss_val" : loss }
+
+    return res
+

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/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.7.dist-info/METADATA

@@ -0,0 +1,54 @@
+Metadata-Version: 2.2
+Name: wums
+Version: 0.1.7
+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.7.dist-info/RECORD

@@ -0,0 +1,14 @@
+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=4iPx9Nr9y8c3p4ovy8XOS-xU_w11OyQEjISKkygxqcA,55918
+wums/tfutils.py,sha256=9efkkvxH7VtwJN2yBS6_-P9dLKs3CXdxMFdrEBNsna8,2892
+wums/Templates/index.php,sha256=9EYmfc0ltMqr5oOdA4_BVIHdSbef5aA0ORoRZBEADVw,4348
+wums-0.1.7.dist-info/METADATA,sha256=GrQyVuatMvHdallbstH7YdiACEMLIo5isHyugfFawW8,1784
+wums-0.1.7.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+wums-0.1.7.dist-info/top_level.txt,sha256=DCE1TVg7ySraosR3kYZkLIZ2w1Pwk2pVTdkqx6E-yRY,5
+wums-0.1.7.dist-info/RECORD,,

{wums-0.1.6.dist-info → wums-0.1.7.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.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,,