linear_regression

class braintools.optim.linear_regression(X, Y, *, fit_intercept=True)

Fit an ordinary least-squares linear regression.

Minimises \(\|Y - X W - b\|_F^2\) with jax.numpy.linalg.lstsq(). When X is rank deficient (for example more features than samples, or a constant feature with fit_intercept=True) the minimum-norm solution is returned, so constant features get a zero coefficient.

Parameters:

• X (array_like or brainunit.Quantity) – Design matrix, shape (n_samples, n_features).

• Y (array_like or brainunit.Quantity) – Targets, shape (n_samples,) or (n_samples, n_targets).

• fit_intercept (bool) – Fit an unpenalised intercept by centring X and Y first.

Returns:

coef of shape (n_features,) or (n_features, n_targets) (in the unit of Y over the unit of X) and intercept (in the unit of Y).

Return type:

RegressionResult

Raises:

ValueError – If X is not 2-D, Y is not 1-D or 2-D, or the sample counts differ.

See also

ridge_regression

L2-penalised least squares.

RLS

the same solution computed online, one batch at a time.

Notes

Matches sklearn.linear_model.LinearRegression. The function is pure and supports jax.jit, jax.vmap and jax.grad (gradients with respect to X require distinct non-zero singular values of the centred X; gradients with respect to Y always exist).

Examples

>>> import jax.numpy as jnp
>>> import braintools
>>> X = jnp.array([[0.], [1.], [2.], [3.]])
>>> Y = 2. * X[:, 0] + 1.
>>> res = braintools.optim.linear_regression(X, Y)
>>> print(jnp.round(res.coef, 4), round(float(res.intercept), 4))
[2.] 1.0