pythonnumpynumpy-ndarray

How to write a function for numpy array input whose action depends on the input's numerical value


The function has this mathematical form

f(x) = 1, if x<1
     = g(x), for 1<=x<10
     = 0, for x >=10

where g(x) is a simple function.

It is straightforward to write such a function if the input is a float using if/else. If the input is a numpy array, is there a neat/efficient implementation without explicit loop?


Solution

  • You can use numpy.piecewise:

    x = np.linspace(0, 20, 30)
    g = lambda x: -x
    
    y = np.piecewise(x, [x<1, (1<=x)&(x<10), x>=10],
                        [1, g, 0])
    

    NB. In contrast to numpy.where, it will only evaluate the function for each valid chunk of x.

    Output:

    array([ 1.        ,  1.        , -1.37931034, -2.06896552, -2.75862069,
           -3.44827586, -4.13793103, -4.82758621, -5.51724138, -6.20689655,
           -6.89655172, -7.5862069 , -8.27586207, -8.96551724, -9.65517241,
            0.        ,  0.        ,  0.        ,  0.        ,  0.        ,
            0.        ,  0.        ,  0.        ,  0.        ,  0.        ,
            0.        ,  0.        ,  0.        ,  0.        ,  0.        ])
    

    As a function:

    def f(x):
        return np.piecewise(x,
                            [x<1, (1<=x)&(x<10), x>=10],
                            [1, np.log, 0])
    
    x = np.linspace(-10, 20, num=200)
    
    plt.plot(x, f(x))
    

    Output:

    enter image description here