Fitting piecewise function in Python

Question

I'm trying to fit a piecewise defined function to a data set in Python. I've searched for quite a while now, but I haven't found an answer whether it is possible or not.

To get an impression of what I am trying to do, look at the following example (which is not working for me). Here I'm trying to fit a shifted absolute value function (f(x) = |x-p|) to a dataset with p as the fit parameter.

import scipy.optimize as so
import numpy as np

def fitfunc(x,p):
   if x>p:
      return x-p
   else:
      return -(x-p)

fitfunc = np.vectorize(fitfunc) #vectorize so you can use func with array

x=np.arange(1,10)
y=fitfunc(x,6)+0.1*np.random.randn(len(x))

popt, pcov = so.curve_fit(fitfunc, x, y) #fitting routine that gives error

Is there any way of accomplishing this in Python?

A way of doing this in R is :

# Fit of a absolute value function f(x)=|x-p|

f.lr <- function(x,p) {
    ifelse(x>p, x-p,-(x-p))
}
x <- seq(0,10)  #
y <- f.lr(x,6) + rnorm (length(x),0,2)
plot(y ~ x)
fit.lr <- nls(y ~ f.lr(x,p), start = list(p = 0), trace = T, control = list(warnOnly = T,minFactor = 1/2048))
summary(fit.lr)
coefficients(fit.lr)
p.fit <- coefficients(fit.lr)["p"]
x_fine <- seq(0,10,length.out=1000)
lines(x_fine,f.lr(x_fine,p.fit),type='l',col='red')
lines(x,f.lr(x,6),type='l',col='blue')

After even more research I found a way of doing it. In this solution, I don't like the fact that I have to define the error function myself. Further I'm not really sure why it has to be in this lambda-style. Therefore any kind of suggestions or more sophisticated solutions are very welcome.

import scipy.optimize as so
import numpy as np
import matplotlib.pyplot as plt

def fitfunc(p,x): return x - p if x > p else p - x 

def array_fitfunc(p,x):
    y = np.zeros(x.shape)
    for i in range(len(y)):
        y[i]=fitfunc(x[i],p)
    return y

errfunc = lambda p, x, y: array_fitfunc(p, x) - y # Distance to the target function

x=np.arange(1,10)
x_fine=np.arange(1,10,0.1)
y=array_fitfunc(6,x)+1*np.random.randn(len(x)) #data with noise

p1, success = so.leastsq(errfunc, -100, args=(x, y), epsfcn=1.) # -100 is the initial value for p; epsfcn sets the step width

plt.plot(x,y,'o') # fit data
plt.plot(x_fine,array_fitfunc(6,x_fine),'r-') #original function
plt.plot(x_fine,array_fitfunc(p1[0],x_fine),'b-') #fitted version
plt.show()

Answer 1

To finish this up here, I'll share my own final solution to the problem. In order to stay close to my original question, you just have to define the vectorized function yourself and not use np.vectorize.

import scipy.optimize as so
import numpy as np

def fitfunc(x,p):
   if x>p:
      return x-p
   else:
      return -(x-p)

fitfunc_vec = np.vectorize(fitfunc) #vectorize so you can use func with array

def fitfunc_vec_self(x,p):
  y = np.zeros(x.shape)
  for i in range(len(y)):
    y[i]=fitfunc(x[i],p)
  return y


x=np.arange(1,10)
y=fitfunc_vec_self(x,6)+0.1*np.random.randn(len(x))

popt, pcov = so.curve_fit(fitfunc_vec_self, x, y) #fitting routine that gives error
print popt
print pcov

Output:

[ 6.03608994]
[[ 0.00124934]]

Answer 2

Couldn't you simply replace fitfunc with

def fitfunc2(x, p):
    return np.abs(x-p)

which then produces something like

>>> x = np.arange(1,10)
>>> y = fitfunc2(x,6) + 0.1*np.random.randn(len(x))
>>> 
>>> so.curve_fit(fitfunc2, x, y) 
(array([ 5.98273313]), array([[ 0.00101859]]))

Using a switch function and/or building blocks like where to replace branches, this should scale up to more complicated expressions without needing to call vectorize.

[PS: the errfunc in your least squares example doesn't need to be a lambda. You could write

def errfunc(p, x, y):
    return array_fitfunc(p, x) - y

instead, if you liked.]