Quant Python

OLS From Scratch

Difficulty

Asked at Two Sigma, Jump Trading, Squarepoint Capital

Theory: Quant Python That Survives Review

Read the problem, hints and solution here. The editor needs a bigger screen: open this page on a laptop to write and run your code.

"Implement linear regression in numpy" is one of the most consistently reported coding tasks in quant researcher online assessments. Not because firms need another regression library: they want to see whether you know what the fit actually computes when the convenience wrapper is taken away.

Implement ols_fit(X, y).

  • X is a list of feature rows (each a list of floats), without an intercept column. It may also be a flat list of floats for the single-feature case.
  • y is a list of floats, one target per row.

Fit ordinary least squares with an intercept, and return a tuple (coeffs, r2):

  • coeffs: the fitted coefficients as a list of floats rounded to 4 decimals, intercept first, then one coefficient per feature in input order.
  • r2: the coefficient of determination \( R^2 = 1 - SS_{res}/SS_{tot} \), rounded to 4 decimals. If y is constant, define \( R^2 = 1.0 \).

Columns may be collinear. Solve by least squares (numpy.linalg.lstsq), which returns the minimum-norm solution in the degenerate case; the normal equations with a plain inverse will crash there.

Examples

ols_fit([[1], [2], [3]], [3, 5, 7])
# ([1.0, 2.0], 1.0)

ols_fit([[1, 2], [2, 1], [3, 4], [4, 3]], [10, 11, 20, 21])
# ([3.0, 3.0, 2.0], 1.0)

A hidden test runs 20,000 rows by 8 features, which any vectorised solution handles instantly.

Rate this problem
Language: Pythonols_fit
Rate this problem