Conditional Piecewise Function

Question

This might be easier to explain by providing my attempt so far, then my intent+comments.

import numpy as np
from numpy import sqrt, arcsin, arcsinh

# Returns roots of quadratic (or lack of)
def roots(a, b, c):
    b24ac = b*b - 4*a*c
    if a == 0 or b24ac < 0:
        return np.nan, np.nan
    else:
        l = (-b - sqrt(b24ac))/(2*a)
        r = (-b + sqrt(b24ac))/(2*a)
        if l > r:
            l, r = r, l
        return l, r

# Some numeric functions
def pw1(z, a, b, c):
    return -arcsin((2*a*z+b)/sqrt(b*b - 4*a*c))/sqrt(-a)

def pw2(z, a, b, c):
    return arcsinh((2*a*z+b)/sqrt(4*a*c - b*b))/sqrt(a)

# Function incorporating above definitions w/ conditions/domains
def func(z, a, b, c):
    b24ac = b*b - 4*a*c
    l, r = roots(*abc)
    conditions = [(b24ac > 0 and a < 0) and (l < z and z < r),
                  (b24ac < 0 and a > 0)]
    choices = [pw1(z, a, b, c), 
               pw2(z, a, b, c)] 
    return np.select(conditions, choices)

This is my attempt at creating a python function that is a conditional piecewise function. For the mathematically curious, this is a portion of a full definition of the integral of $[ax^2+bx+c]^{-1/2}$. The necessary specifics are that I need a function that is conditional on the domain AND other parameters. I've looked into numpy's piecewise and select functions. Piecewise , for its condition list, only accepts logic on the domain (not the parameters). Unless I'm missing something, this seems like it won't work for me. Select has given me the most success. The only issues I've had are it not evaluating domain conditions, over its domain:

--->  conditions = [(b24ac > 0 and a < 0) and (l < z and z < r),
                  (b24ac < 0 and a > 0)]

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Lastly, it evaluates all choices (which are results as opposed to the functions that you give piecewise ), then selects the one per the conditions list to return (commenting out the and (l < z ... conditions):

c:\program files\python36\lib\site-packages\ipykernel_launcher.py:6: 
RuntimeWarning: invalid value encountered in sqrt

I started this weeks ago with a bunch of elif statements. This worked only for floats, not arrays. I'm using numpy, and having this function (or any really) evaluate over a domain is what I'm after. I then learned about piecewise and select and began playing with these, being from numpy.

I'd really like a pythonic way to do this. One that evaluates over numpy arrays, and only for the desired condition/domain. So, behaves like piecewise but has versatile conditioning like select . Any help and suggestions are greatly appreciated! Thanks.

Answer 1

The b24ac identifier is certainly descriptive, but most folks would probably name it discriminant .

use your own

You are complaining that you know how to compute the desired result, but the numpy tools at hand don't seem a good fit. So write your own function fn and use .apply(fn) , or nditer . Notice that you can arrange for all needed arguments to appear in extra columns, and have the function handle a row at a time.

I heard no concerns about speed, which would be the usual argument against such a custom function.

use np.piecewise

I agree with you that piecewise() appears to be appropriate to your needs. The aspect that seemed missing was tailoring your functions with four parameters before passing them to piecewise. A good match for that need would be functools.partial() .