Linear Regression: Predicting Numbers — Machine Learning Fundamentals

You're predicting a number. Sale price. Number of clicks. Tomorrow's temperature. The target is continuous: it could be 47, 47.3, or 47.291.

Try linear regression first. It's one of the oldest methods, it trains in milliseconds, you can interpret its coefficients, and on many real problems it does better than more complex models. If linear regression fails, you have still learned something: the relationship probably isn't linear, so it is reasonable to try a more complex model.

This lesson is about that first step.

Run these right here: scikit-learn, pandas and matplotlib all work in the browser. The first scikit-learn import takes a few seconds while it downloads. After that it's instant. Expected output is also shown in comments below each block.

1. When Linear Regression Is the Right Tool

Check three things:

The target is continuous: a real number, not a category.
You expect a roughly linear relationship: a bigger house sells for a higher price, and more ad spend brings more sales. Not "U-shaped" or "explodes after a threshold".
You want to interpret the model: "for every extra bedroom, price goes up by £30k". Linear coefficients tell you exactly that.

If the target is a category (spam/not-spam), you want classification, not regression. If the relationship is strongly non-linear, you'll want trees.

2. The Mental Model

You have a cloud of points. Linear regression draws the one straight line that is closest to all of them. "Closest" means the smallest sum of squared vertical distances from the points to the line. In simple words: the line that passes through the middle of the cloud.

python

price
  ▲                          ●
  │                       ●     ●
  │                   ●●      ●
  │              ●  ●  ●
  │        ●  ●●●
  │   ●●●●        ╲       ← the line
  │●●●            (best fit)
  └────────────────────────► square footage

That's the whole idea. One feature → one line. Two features → a tilted plane in 3D. Many features → a "hyperplane" that you can't draw, but the maths works the same way.

3. The Equation Behind It

One feature:

python

y = m·x + b

+ setup added so this can run · defines m·x, b

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

m·x = _AutoMock('m·x')
b = _AutoMock('b')

Same y = mx + b you saw in school. m is the slope (how much y changes when x goes up by 1), b is the intercept (y when x is 0). Linear regression only has to learn the best m and b from the data.

Multiple features:

python

y = w₁·x₁ + w₂·x₂ + w₃·x₃ + ... + b

Each feature gets its own weight (w_i, also called a coefficient). The model learns one weight per feature, plus one intercept. It is still a straight line, but in more dimensions.

You won't write this maths yourself. scikit-learn fits the weights for you in a single .fit() call.

4. Code: House Prices in 12 Lines

python

import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split

df = pd.DataFrame({
    "sqft":     [800, 1200, 1500, 1800, 2200, 2600, 3000, 3500],
    "bedrooms": [1,   2,    2,    3,    3,    4,    4,    5],
    "price":    [180, 250,  300,  370,  430,  520,  590,  690],   # in thousands of pounds
})

X = df[["sqft", "bedrooms"]]
y = df["price"]

X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)

model = LinearRegression().fit(X_train, y_train)

print("Coefficients:", model.coef_)
print("Intercept:   ", model.intercept_)
print("R²:          ", model.score(X_test, y_test))
# → Coefficients: [ 0.16914661 14.28884026]
# → Intercept:    22.95404814004371
# → R²:           0.999442223406977

The model learned: price ≈ 0.169·sqft + 14.29·bedrooms + 22.95 (in £k).

Translate that:

Every extra square foot adds about £169 to the price.

Every extra bedroom adds about £14,290, holding sqft constant.

A 0-sqft, 0-bedroom house would cost £22,950. (That house can't exist, of course. The intercept is just the value the line gives when every feature is 0.)

R² of 0.999 means the model explains 99.9% of the variance in price on the test set. With 8 rows of clean synthetic data, that's expected. Real data is messier.

5. Interpreting Coefficients: an Important Warning

Reading one coefficient on its own ("every extra bedroom adds £14,290") is fine. But comparing coefficients to decide which feature matters more only works when the features are on comparable scales. Sqft ranges from 800 to 3500. Bedrooms range from 1 to 5. You can't compare the coefficients directly, because the units are so different.

To compare relative importance, scale the features first:

python

from sklearn.preprocessing import StandardScaler

scaler = StandardScaler()
X_scaled = scaler.fit_transform(X_train)

model_scaled = LinearRegression().fit(X_scaled, y_train)     # a new name, so `model` stays the unscaled model
print(dict(zip(X.columns, model_scaled.coef_)))
# → {'sqft': np.float64(153.06...), 'bedrooms': np.float64(18.45...)}

+ setup added so this can run · re-runs 1 earlier block on this page first

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

X_train = _AutoMock('X_train')
y_train = _AutoMock('y_train')
def LinearRegression(*_a, **_kw):
    print('-> LinearRegression() called')
    return _AutoMock('LinearRegression()')
X = _AutoMock('X')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import pandas as pd
    from sklearn.linear_model import LinearRegression
    from sklearn.model_selection import train_test_split
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport pandas as pd\nfrom sklearn.linear_model import LinearRegression\nfrom sklearn.model_selection import train_test_split\n\ndf = pd.DataFrame({\n    "sqft":     [800, 1200, 1500, 1800, 2200, 2600, 3000, 3500],\n    "bedrooms": [1,   2,    2,    3,    3,    4,    4,    5],\n    "price":    [180, 250,  300,  370,  430,  520,  590,  690],   # in thousands of pounds\n})\n\nX = df[["sqft", "bedrooms"]]\ny = df["price"]\n\nX_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)\n\nmodel = LinearRegression().fit(X_train, y_train)\n\nprint("Coefficients:", model.coef_)\nprint("Intercept:   ", model.intercept_)\nprint("R²:          ", model.score(X_test, y_test))\n# → Coefficients: [ 0.16914661 14.28884026]\n# → Intercept:    22.95404814004371\n# → R²:           0.999442223406977',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

Now the coefficients use the same unit ("1 standard deviation"), so you can say that sqft has roughly 8× the effect of the number of bedrooms on price.

The correct way is to do the scaling inside a Pipeline, so that information can't leak from the test data into training by accident. That's lesson 6.

6. Metrics for Regression

.score() gives you R². That's a fine summary number, but you usually want at least one error metric in the units of your target.

python

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import numpy as np

y_pred = model.predict(X_test)

mae  = mean_absolute_error(y_test, y_pred)
rmse = np.sqrt(mean_squared_error(y_test, y_pred))
r2   = r2_score(y_test, y_pred)

print(f"MAE  = £{mae:,.0f}k")
print(f"RMSE = £{rmse:,.0f}k")
print(f"R²   = {r2:.3f}")
# → MAE  = £2k
# → RMSE = £3k
# → R²   = 0.999

+ setup added so this can run · re-runs 1 earlier block on this page first

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

X_test = _AutoMock('X_test')
y_test = _AutoMock('y_test')
model = _AutoMock('model')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import pandas as pd
    from sklearn.linear_model import LinearRegression
    from sklearn.model_selection import train_test_split
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport pandas as pd\nfrom sklearn.linear_model import LinearRegression\nfrom sklearn.model_selection import train_test_split\n\ndf = pd.DataFrame({\n    "sqft":     [800, 1200, 1500, 1800, 2200, 2600, 3000, 3500],\n    "bedrooms": [1,   2,    2,    3,    3,    4,    4,    5],\n    "price":    [180, 250,  300,  370,  430,  520,  590,  690],   # in thousands of pounds\n})\n\nX = df[["sqft", "bedrooms"]]\ny = df["price"]\n\nX_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)\n\nmodel = LinearRegression().fit(X_train, y_train)\n\nprint("Coefficients:", model.coef_)\nprint("Intercept:   ", model.intercept_)\nprint("R²:          ", model.score(X_test, y_test))\n# → Coefficients: [ 0.16914661 14.28884026]\n# → Intercept:    22.95404814004371\n# → R²:           0.999442223406977',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

Metric	Reads as	When to prefer it
MAE	"Average error of £X in the target unit."	Easy to explain to non-technical people.
RMSE	"Like MAE, but big errors count more."	When occasional huge errors are unacceptable.
R²	"% of variance the model explains."	Quick check that the model is useful at all.

A common interview question: "why is RMSE always ≥ MAE?" Because squaring the errors before averaging gives big errors more weight. Taking the square root brings you back to the original units, but the extra weight on big errors stays in the result.

7. When Linear Regression Fails

The model can only change one slope per feature. If the real relationship is curved, a straight line can't follow it.

python

y
▲       ●●●
│     ●●   ●●
│   ●●       ●●          ← real pattern (curved)
│  ●           ●
│ ●             ●        ── the best straight line LR can do
│●               ●
└─────────────────► x

There are two ways around this:

1. Polynomial features. Build x², x³ columns yourself. Linear regression on those columns fits a curve.

python

from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import make_pipeline

poly_model = make_pipeline(PolynomialFeatures(degree=3), LinearRegression())

+ setup added so this can run · re-runs 1 earlier block on this page first

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

def LinearRegression(*_a, **_kw):
    print('-> LinearRegression() called')
    return _AutoMock('LinearRegression()')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import pandas as pd
    from sklearn.linear_model import LinearRegression
    from sklearn.model_selection import train_test_split
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport pandas as pd\nfrom sklearn.linear_model import LinearRegression\nfrom sklearn.model_selection import train_test_split\n\ndf = pd.DataFrame({\n    "sqft":     [800, 1200, 1500, 1800, 2200, 2600, 3000, 3500],\n    "bedrooms": [1,   2,    2,    3,    3,    4,    4,    5],\n    "price":    [180, 250,  300,  370,  430,  520,  590,  690],   # in thousands of pounds\n})\n\nX = df[["sqft", "bedrooms"]]\ny = df["price"]\n\nX_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)\n\nmodel = LinearRegression().fit(X_train, y_train)\n\nprint("Coefficients:", model.coef_)\nprint("Intercept:   ", model.intercept_)\nprint("R²:          ", model.score(X_test, y_test))\n# → Coefficients: [ 0.16914661 14.28884026]\n# → Intercept:    22.95404814004371\n# → R²:           0.999442223406977',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

2. Switch to trees. Decision trees and random forests handle non-linear relationships without extra feature engineering. But they are harder to interpret.

How do you know you have this problem? Plot the residuals.

8. The Residual Plot: Checking the Fit

A residual is actual - predicted for each point. If linear regression fits well, the residuals should form a random cloud around zero. If there's a pattern in the residuals, the model has missed some of the structure in the data.

python

import matplotlib.pyplot as plt

y_pred = model.predict(X_test)
residuals = y_test - y_pred

plt.scatter(y_pred, residuals)
plt.axhline(0, color="red", linestyle="--")
plt.xlabel("Predicted price"); plt.ylabel("Residual (actual - predicted)")
plt.title("Residual plot: look for randomness")
plt.show()

+ setup added so this can run · re-runs 2 earlier blocks on this page first

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

X_test = _AutoMock('X_test')
y_test = _AutoMock('y_test')
model = _AutoMock('model')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import pandas as pd
    from sklearn.linear_model import LinearRegression
    from sklearn.model_selection import train_test_split
    from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport pandas as pd\nfrom sklearn.linear_model import LinearRegression\nfrom sklearn.model_selection import train_test_split\n\ndf = pd.DataFrame({\n    "sqft":     [800, 1200, 1500, 1800, 2200, 2600, 3000, 3500],\n    "bedrooms": [1,   2,    2,    3,    3,    4,    4,    5],\n    "price":    [180, 250,  300,  370,  430,  520,  590,  690],   # in thousands of pounds\n})\n\nX = df[["sqft", "bedrooms"]]\ny = df["price"]\n\nX_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)\n\nmodel = LinearRegression().fit(X_train, y_train)\n\nprint("Coefficients:", model.coef_)\nprint("Intercept:   ", model.intercept_)\nprint("R²:          ", model.score(X_test, y_test))\n# → Coefficients: [ 0.16914661 14.28884026]\n# → Intercept:    22.95404814004371\n# → R²:           0.999442223406977',
    'from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score\nimport numpy as np\n\ny_pred = model.predict(X_test)\n\nmae  = mean_absolute_error(y_test, y_pred)\nrmse = np.sqrt(mean_squared_error(y_test, y_pred))\nr2   = r2_score(y_test, y_pred)\n\nprint(f"MAE  = £{mae:,.0f}k")\nprint(f"RMSE = £{rmse:,.0f}k")\nprint(f"R²   = {r2:.3f}")\n# → MAE  = £2k\n# → RMSE = £3k\n# → R²   = 0.999',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

What you want to see: random scatter around 0. What shows a non-linear pattern: a clear curve or fan shape in the residuals. That tells you to try polynomial features or a tree-based model.

Common Mistakes

Using accuracy on a regression problem. Accuracy is for classification. You'll get either an error or a meaningless score. Use MAE / RMSE / R².
Interpreting unscaled coefficients as importance. A coefficient of 0.169 looks tiny next to 14.29. But one is per square foot and the other is per bedroom. Scale first.
Forgetting the linearity assumption. Linear regression on a clearly curved relationship will run and return a number, and it will be wrong without any warning. Always plot residuals.
Including the target as a feature. This happens. You make a mistake when dropping columns, and price ends up in X. The model gets R² = 1.0, and you release a model that is useless.

🎯 Your Turn: Salary vs Years of Experience

Inline data:

python

years   = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
salary  = [35, 42, 47, 53, 60, 64, 71, 77, 82, 89]   # in thousands of pounds

Fit a linear regression. Then:

1. Predict the salary for 5 years of experience.
2. Report the RMSE on the full dataset (we're not splitting, only 10 rows).
3. Return (predicted_salary_for_5_years, rmse) rounded to 2 decimals.

python

import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error

def fit_and_predict():
    years  = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).reshape(-1, 1)
    salary = np.array([35, 42, 47, 53, 60, 64, 71, 77, 82, 89])

    # TODO 1: instantiate and fit LinearRegression
    # TODO 2: predict for years = 5
    # TODO 3: compute RMSE on the full dataset
    # TODO 4: return (round(prediction, 2), round(rmse, 2))

    pass

print(fit_and_predict())
# → expected something near (59.05, 0.57)

Hint 1: predict expects 2D input

model.predict([[5]]): note the double brackets. predict needs a 2D array even for one prediction. The result is a 1D array, so use [0] to get the single number.

Hint 2: RMSE in one line

rmse = np.sqrt(mean_squared_error(salary, model.predict(years))). Compare actuals to predictions on the same data the model was trained on (fine here because we're only checking fit quality, not generalisation).

Show full solution

python

import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error

def fit_and_predict():
    years  = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).reshape(-1, 1)
    salary = np.array([35, 42, 47, 53, 60, 64, 71, 77, 82, 89])

    model = LinearRegression().fit(years, salary)

    prediction = model.predict([[5]])[0]
    rmse = np.sqrt(mean_squared_error(salary, model.predict(years)))

    return (round(prediction, 2), round(rmse, 2))

print(fit_and_predict())
# → (np.float64(59.05), np.float64(0.57))   (NumPy 2 shows the number type. The values are 59.05 and 0.57.)

The model learned a slope of about £5.90k per year and an intercept of about £29.5k. The fit is very close: the RMSE is under £1k, for salaries in the £35-£89k range. Linear regression works well here.

What You Learned

Linear regression predicts a continuous target by fitting one straight line (or hyperplane).
The equation y = w₁x₁ + ... + b: the model learns the weights.
LinearRegression().fit(X, y), then .coef_, .intercept_, .predict(), .score().
Compare coefficients only after scaling the features to comparable units.
Three metrics: MAE (avg error), RMSE (penalises big misses), R² (variance explained).
When residuals show a pattern, the relationship is probably non-linear: try polynomial features or trees.

Next: Classification: when the answer is a category, not a number.