Python: Finding multiple linear trend lines in a scatter plot

My following solution proposal is based on the RANSAC algorithm. It is method to fit a mathematical model (e.g. a line) to data with heavy of outliers.

RANSAC is one specific method from the field of robust regression.

My solution below first fits a line with RANSAC. Then you remove the data points close to this line from your data set (which is the same as keeping the outliers), fit RANSAC again, remove data, etc until only very few points are left.

Such approaches always have parameters which are data dependent (e.g. noise level or proximity of the lines). In the following solution and MIN_SAMPLES and residual_threshold are parameters which might require some adaption to the structure of your data:

import matplotlib.pyplot as plt
import numpy as np
from sklearn import linear_model

MIN_SAMPLES = 3

x = np.linspace(0, 2, 100)

xs, ys = [], []

# generate points for thee lines described by a and b,
# we also add some noise:
for a, b in [(1.0, 2), (0.5, 1), (1.2, -1)]:
    xs.extend(x)
    ys.extend(a * x + b + .1 * np.random.randn(len(x)))

xs = np.array(xs)
ys = np.array(ys)
plt.plot(xs, ys, "r.")

colors = "rgbky"
idx = 0

while len(xs) > MIN_SAMPLES:

    # build design matrix for linear regressor
    X = np.ones((len(xs), 2))
    X[:, 1] = xs

    ransac = linear_model.RANSACRegressor(
        residual_threshold=.3, min_samples=MIN_SAMPLES
    )

    res = ransac.fit(X, ys)

    # vector of boolean values, describes which points belong
    # to the fitted line:
    inlier_mask = ransac.inlier_mask_

    # plot point cloud:
    xinlier = xs[inlier_mask]
    yinlier = ys[inlier_mask]

    # circle through colors:
    color = colors[idx % len(colors)]
    idx += 1
    plt.plot(xinlier, yinlier, color + "*")

    # only keep the outliers:
    xs = xs[~inlier_mask]
    ys = ys[~inlier_mask]

plt.show()

In the following plot points shown as stars belong to the clusters detected by my code. You also see a few points depicted as circles which are the points remaining after the iterations. The few black stars form a cluster which you could get rid of by increasing MIN_SAMPLES and / or residual_threshold.