There is a package named segmented in R. Is there a similar package in python?

  • $\begingroup$ Is there anything specific that you need? For known break points, this can just be modeled by an interaction with a indicator function (0 before, 1 after the break) or a linear spline. The first approach has a jump, the second approach results in a connected piecewise regression line. $\endgroup$
    – Josef
    Oct 2 '15 at 4:28
  • $\begingroup$ LINEAR-TREE (sklearn wrapper): github.com/cerlymarco/linear-tree $\endgroup$ May 19 at 15:13

No, currently there isn't a package in Python that does segmented linear regression as thoroughly as those in R (e.g. R packages listed in this blog post). Alternatively, you can use a Bayesian Markov Chain Monte Carlo algorithm in Python to create your segmented model.

Segmented linear regression, as implemented by all the R packages in the above link, doesn't permit extra parameter constraints (i.e. priors), and because these packages take a frequentist approach, the resulting model doesn't give you probability distributions for the model parameters (i.e. breakpoints, slopes, etc). Defining a segmented model in statsmodels, which is frequentist, is even more restrictive because the model requires a fixed x-coordinate breakpoint.

You can design a segmented model in Python using the Bayesian Markov Chain Monte Carlo algorithm emcee. Jake Vanderplas wrote a useful blog post and paper for how to implement emcee with comparisons to PyMC and PyStan.


  • Segmented model with data:

Segmented regression

  • Probability distributions of fit parameters:

enter image description here


enter image description here

This is an implementation of my own.

import numpy as np
import matplotlib.pylab as plt
from sklearn.tree import DecisionTreeRegressor
from sklearn.linear_model import LinearRegression

# parameters for setup
n_data = 20

# segmented linear regression parameters
n_seg = 3

fig, (ax0, ax1) = plt.subplots(1, 2)

# example 1
#xs = np.sort(np.random.rand(n_data))
#ys = np.random.rand(n_data) * .3 + np.tanh(5* (xs -.5))

# example 2
xs = np.linspace(-1, 1, 20)
ys = np.random.rand(n_data) * .3 + np.tanh(3*xs)

dys = np.gradient(ys, xs)

rgr = DecisionTreeRegressor(max_leaf_nodes=n_seg)
rgr.fit(xs.reshape(-1, 1), dys.reshape(-1, 1))
dys_dt = rgr.predict(xs.reshape(-1, 1)).flatten()

ys_sl = np.ones(len(xs)) * np.nan
for y in np.unique(dys_dt):
    msk = dys_dt == y
    lin_reg = LinearRegression()
    lin_reg.fit(xs[msk].reshape(-1, 1), ys[msk].reshape(-1, 1))
    ys_sl[msk] = lin_reg.predict(xs[msk].reshape(-1, 1)).flatten()
    ax0.plot([xs[msk][0], xs[msk][-1]],
             [ys_sl[msk][0], ys_sl[msk][-1]],
             color='r', zorder=1)

ax0.scatter(xs, ys, label='data')
ax0.scatter(xs, ys_sl, s=3**2, label='seg lin reg', color='g', zorder=5)

ax1.scatter(xs, dys, label='data')
ax1.scatter(xs, dys_dt, label='DecisionTree', s=2**2)


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.