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

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
+