How do I add a linear regression line to each scatterplot in my scatterplot matrix?

user_6205041 0
Viewed 231

I have code that creates a scatterplot matrix and I would like to add a linear regression line to each facet. Code and the current graph are shown below. I currently have a scatterplot for each of the variable combinations for the first five variables in the dataset. I would like to add the regression lines so that when the individual hovers over the line they can also see the correlation.

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import plotly.express as px
from sklearn import datasets
from typing import Tuple, List
import plotly.graph_objects as go
from plotly.subplots import make_subplots

def load_data() -> Tuple[np.ndarray, np.ndarray, List[str]]:
    """Load the wine dataset

    Returns:
        features: the dataset features
        target: the labels of the dataset
        feature_names: names of each feature
    """
    wine = datasets.load_wine()
    features = wine['data']
    target = wine['target']
    feature_names = wine['feature_names']
    return features, target, feature_names

features, target, feature_names = load_data()
Data = {
    feature_names[0]:features[:,0].tolist(),
    feature_names[1]:features[:,1].tolist(),
    feature_names[2]:features[:,2].tolist(),
    'Target': target.tolist()
}
Data = pd.DataFrame(data = Data)

index_vals = Data['Target'].astype('category').cat.codes

fig = go.Figure(data = go.Splom(dimensions = [
    dict(label = feature_names[0],values = Data[feature_names[0]]),
    dict(label = feature_names[1],values = Data[feature_names[1]]),
    dict(label = feature_names[2],values = Data[feature_names[2]])],
   text = Data['Target'],
    marker = dict(color = index_vals,showscale = False,size = 8)
))

fig.update_layout(
    title='Wine Dataset',
    dragmode='select',
    width=900,
    height=600,
    hovermode='closest',
)

fig.show()
plotly python
1 Answers

With newer versions of plotly, all you need to do is include trendline="ols" in px.scatter. Here's an example that builds on the dataset px.data.tips():

enter image description here

Complete code:

import plotly.express as px
df = px.data.tips()
fig = px.scatter(df, x="total_bill", y="tip",
                 trendline = 'ols',
                 trendline_color_override = 'red',
                 facet_col="day",
                 facet_col_wrap = 2)
fig.update_xaxes(matches=None)
fig.show()
user_34377870
Related