[SciPy-User] Vectorizing functions where not known if each arg is (broadcast compatible) scalar or ndarray

William Furnass will@thearete.co...
Tue Sep 18 09:04:28 CDT 2012


Hi,

I'm looking for a simple way to create a vectorized version of the
following function that flexibly allows one or more of the inputs to
be either an ndarray of constant length l or a scalar.

With simpler functions (such as the 'reynolds' function mentioned
below) the np.vectorize function does the job but I'm not sure how to
vectorize the following given the conditionals involving Re (which
could be vector or scalar) e.g. how should the case for Re > 4000 be
calculated and D or k_s be indexed if it is not without introducing
quite a number of checks whether Re, D or k_s are ndarrays or scalar
floats?

Also, is there a numpy function that allows one to check whether a
number of scalar and ndarray are broadcast compatible?

Cheers,

Will

def friction_factor(D, Q, k_s, T = 10.0, den = 1000.0):
    Re = reynolds(D, Q, T, den)
    if Re == 0:
        f = 0
    elif Re < 2000:
        f = 64 / Re
    elif 2000 <= Re < 4000:
        y3 = -0.86859 * np.log((k_s / (3.7 * D)) + (5.74 / (4000**0.9)))
        y2 = (k_s / (3.7 * D)) + (5.74 / (Re**0.9))
        fa = y3**-2
        fb = fa * (2 - (0.00514215 / (y2*y3)))
        r = Re / 2000.
        x4 = r * (0.032 - (3. * fa) + (0.5 * fb))
        x3 = -0.128 + (13. * fa) - (2.0 * fb)
        x2 =  0.128 - (17. * fa) + (2.5 * fb)
        x1 = (7 * fa) - fb
        f = x1 + r * (x2 + r * (x3 + x4))
    elif Re >= 4000:
        f = 0.25 / (np.log10((k_s / (3.7 * D)) + (5.74 / (Re**0.9))))**2
    return f
friction_factor = np.vectorized(friction_factor)


More information about the SciPy-User mailing list