Altair: Extract and display regression coefficients

Viewed 266

This question addresses how to access and display the R2 value using mark_text()

I am interested in accessing and displaying the coefficients. Replacing rSquared with coef yields a flattened array of both the intercept and slope, as described in the documentation.

How can I index into this array to display only one of the values, e.g. the slope? I wondered if the mark_text() step should be preceded by a transform (possibly transform_filter(), or if altair.Text() could be used.

I am aware of other approaches which involve determining this information separately then adding it as an additional layer.

Apologies if this is a very straightforward question. Thanks in advance.

import altair as alt
import pandas as pd
import numpy as np

np.random.seed(42)
x = np.linspace(0, 10)
y = x - 5 + np.random.randn(len(x))

df = pd.DataFrame({'x': x, 'y': y})

chart = alt.Chart(df).mark_point().encode(
    x='x',
    y='y'
)
line = chart.transform_regression('x', 'y').mark_line()

params = alt.Chart(df).transform_regression(
    'x', 'y', params=True
).mark_text(align='left').encode(
    x=alt.value(20),  # pixels from left
    y=alt.value(20),  # pixels from top
    text='rSquared:N',

    # text='coef:N' # flattened array
    # text='coef[0]:N' # fails
)

chart + line + params
1 Answers

You can access this using a calculate transform:

params = alt.Chart(df).transform_regression(
    'x', 'y', params=True
).transform_calculate(
    intercept='datum.coef[0]',
    slope='datum.coef[1]',
).mark_text(align='left').encode(
    x=alt.value(20),  # pixels from left
    y=alt.value(20),  # pixels from top
    text='intercept:N'
)

chart + line + params

enter image description here

Related