ADMM-least-squares-method-answer.py

import numpy as np

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

b = np.array([[1], [1]])

N = A.shape[0]
M = A.shape[1]

class ADMM:
    max_loop = 100

    def __init__(self, x0, u0, rho):
        self.x = x0
        self.u = u0
        self.rho = rho

    def update_x(self):
        p = b - self.u
        Q = (2 * np.identity(M) + self.rho * np.dot(A.T, A))
        self.x = self.rho * np.dot(np.dot(np.linalg.inv(Q), A.T), p)

    def update_u(self):
        self.u = self.u + np.dot(A, self.x) - b

    def update(self):
        self.update_x()
        self.update_u()

    def fit(self):
        for loop in range(self.max_loop):
            self.update()

    def result_x(self, a):
        return np.dot(a, self.x)

print(A)
print(b)

admm = ADMM(x0=np.dot(A.T, b) / N, u0=0, rho=1)
admm.fit()
print("最小二乗法 by ADMM")

print("coef_")
print(np.round(admm.x, 10))

print('\n')

print("Result value : Actual value")
for i in range(N):
    print(admm.result_x(A[i]), " : ", b[i])

print(admm.x)