Least Squares in NumPy

This article will introduce how to calculate AX = B with the least-squares method in Python.

Least Squares NumPy With numpy.linalg.lstsq() Function in Python

The equation AX = B is known as the linear matrix equation. The numpy.linalg.lstsq() function can be used to solve the linear matrix equation AX = B with the least-squares method in Python. Actually, it is pretty straightforward. This function takes the matrices and returns the least square solution to the linear matrix equation in the form of another matrix. See the following code example.

import numpy as np

A = [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 0, 0]]
B = [1, 1, 1, 1, 1]

X = np.linalg.lstsq(A, B, rcond=-1)
print(X[0])

Output:

[5.00000000e-01 5.00000000e-01 1.09109979e-16 1.64621130e-16]

In the above code, we calculated the solution to the linear matrix equation AX = B with the np.linalg.lstsq() function in Python. This method gets a little tricky when we start to add weights to our matrices. There are two main methods that we can use to find the solution to this kind of problem.

The first solution involves using array indexing with the np.newaxis specifier to add a new dimension to the weights. It is illustrated in the coding example below.

import numpy as np

A = np.array([[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 0, 0]])
B = np.array([1, 1, 1, 1, 1])

W = np.array([1, 2, 3, 4, 5])

Aw = A * np.sqrt(W[:, np.newaxis])
Bw = B * np.sqrt(W)

X = np.linalg.lstsq(Aw, Bw, rcond=-1)
print(X[0])

Output:

[ 5.00000000e-01  5.00000000e-01 -4.40221936e-17  1.14889576e-17]

In the above code, we calculated the solution to the linear matrix equation AX = B along with the weights W using the np.newaxis and np.linalg.lstsq() function in Python. This method works fine but is not very easy to understand and readable.

The second solution is a bit more readable and easy to understand. It involves transforming the weights into a diagonal matrix and then using it. It is demonstrated in the coding example below.

import numpy as np

A = [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 0, 0]]
B = [1, 1, 1, 1, 1]
W = [1, 2, 3, 4, 5]

W = np.sqrt(np.diag(W))

Aw = np.dot(W, A)
Bw = np.dot(B, W)

X = np.linalg.lstsq(Aw, Bw, rcond=-1)
print(X[0])

Output:

[ 5.00000000e-01  5.00000000e-01 -4.40221936e-17  1.14889576e-17]

In the above code, we calculated the solution to the linear matrix equation AX = B along with the weights W by converting the weights into a diagonal matrix and then using the np.linalg.lstsq() function.