Function wrapper for uncertainties package for wrapping of functions accepting arrays

wrap_array_first_argument.py

def wrap_array_first_argument(f, derivatives_args=[], derivatives_kwargs={}):
    """
    Wraps a function f into a function that also accepts numbers with
    uncertainties (UFloat objects) using uncertainties.wrap().

    If the wrapped function is called with an array as first argument, the 
    wrap iterates over all values and therefore the result will also be an 
    array.
    """
    def f_with_array_first_argument(*args, **kwargs):
        args_values = list(args)
        # Get first argument.
        arr = args_values[0]
        # Check if the first argument is indeed an array.
        if not isinstance(arr, np.ndarray):
            # Otherwise try to call the function directly.
            # Further error handling is performed in there.
            return f(*args_values, **kwargs)
        # Loop over all entries in arr and pass on the rest of the arguments.
        return np.array([f(a, *args_values[1:], **kwargs) for a in arr])
    # Do the initial wrap of the function, which only accepts scalars.
    f_wrap = u.wrap(f, derivatives_args, derivatives_kwargs)
    # It is easier to work with f_with_affine_output, which represents
    # a wrapped version of 'f', when it bears the same name as 'f':
    # ! __name__ is read-only, in Python 2.3:
    f_with_array_first_argument.name = f_wrap.__name__
    return f_with_array_first_argument