Search code examples
pythonnumpypiecewise

piecewise linear function with numpy.piecewise


I am trying to use the data from two x0 and y0 coordinate arrays to create a function that uses the provided x0 and y0 to compute a piecewise series of segments.

For doing that, I create a function

import numpy as np
import matplotlib.pylab as pl

def broken_line(x, x0, y0):
    cl = []
    fl = []
    for i in range(len(x0)-1):
        ab = np.polyfit(x0[i:i+2], y0[i:i+2], 1)
        # Compute and append a "condition" interval
        cl.append(np.logical_and(x >= x0[i], x <= x0[i+1]))
        # Create a line function for the interval
        fl.append(lambda x: x*ab[0] + ab[1])
    return(np.piecewise(x, condlist=cl, funclist=fl))

and then to test it I plot the results of

x0 = np.array([1, 3, 5, 10])
y0 = np.array([2, 1, 5, 7])

x = np.linspace(1, 10, 30)

pl.plot(x, broken_line(x, x0, y0))
pl.plot(x0, y0)
pl.show()

enter image description here

However, the results is not as I would expect. I had a look at other posts on the topic, including this and this other, together with the numpy.piecewise documentation. However, I was not able to figure out why the code is not working as expected. It looks like only the last definition of lambda is considered. Suggestions are all welcome.


Solution

  • The ab in the lambda definition is defined in the surrounding scope, thus changes with every iteration. Only the ab values of the last iteration are reflected into all the lambda funtions.

    One possible solution is to use a factory method for the creation of the lambda function:

    import numpy as np
    import matplotlib.pylab as pl
    
    def lambda_factory(ab):
        return lambda x:x*ab[0]+ab[1]
    
    def broken_line(x, x0, y0):
        cl = []
        fl = []
        for i in range(len(x0)-1):
            ab = np.polyfit(x0[i:i+2], y0[i:i+2], 1)
            # Compute and append a "condition" interval
            cl.append(np.logical_and(x >= x0[i], x <= x0[i+1]))
            # Create a line function for the interval
            fl.append(lambda_factory(ab))
        return(np.piecewise(x, condlist=cl, funclist=fl))
    
    x0 = np.array([1, 3, 5, 10])
    y0 = np.array([2, 1, 5, 7])
    
    x = np.linspace(1, 10, 30)
    
    pl.plot(x, broken_line(x, x0, y0))
    pl.plot(x0, y0)
    pl.show()
    

    Another solution is to save ab in a variable local to the lambda, thus using

    fl.append(lambda x, ab=ab:x*ab[0]+ab[1])
    

    within the loop. Here you create a local variable ab of the outer scope variable ab.

    In both cases the result looks like this:

    Resulting piecewise fit

    For further reference see the python faq