Skip to content

Piecewise function definitions in NumPy

18 Dec, 2009

In the spirit of documenting what I learn, to avoid having to search for it and re-learn it when I need it again, I want to make note of a nice function I found in NumPy that allows piecewise function definitions. The function is called… piecewise() (what a surprise!).

I use Python/NumPy/Matplotlib for post-processing and presenting my CFD results. Among the things I need to do is to reconstruct some of the quantities that go into the turbulence models (e.g. variable model coefficients, etc.), some of which have a piecewise definition depending on some parameter (e.g. Lumley flatness parameter, A), such as:
fA
While previously I would have used a selection statement like the following (where A is an array of floats corresponding to each node in a grid, and fA is some coefficient):

def fA(A):
    fA = np.zeros(np.shape(A))
    for i,a in enumerate(A):
    if a<0.05:
        fA[i]=np.sqrt(a/14.)
    elif a<0.7:
        fA[i]=a/0.8367
    else:
        fA[i]=sqrt(a)
    return(fA)

I now write something like this:

fA= np.piecewise(A,[A<0.05,(A<0.7) & (A>=0.05)],
                              [lambda x: np.sqrt(x/14.),
                               lambda x: x/0.8367,
                               lambda x: np.sqrt(x)])

The basic format of the piecewise statement, ignoring optional arguments, is numpy.piecewise(x, condlist, funclist) , where x is a NumPy ndarray, condlist is the list of boolean arrays upon which the piecwise function defenition depends, and funclist is a list of functions corresponding to each of the boolean conditons. The number of functions in funclist must be the same as the number of conditions, or one more, as in my example, to provide a default case to be used when all the conditions are false.

No comments yet

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s

%d bloggers like this: