osqp admm algorithm written in python with numpy

osqp.py

# %%
# Copyright [2024] Taylor Howell

import numpy as np
# %%
# OSQP: https://web.stanford.edu/~boyd/papers/pdf/osqp.pdf

# problem setup:
# minimize   0.5 x' P x + q' x
# subject to     l <= A x <= u

# dimensions
n = 10
m = 5

# data
_P = np.random.normal(size=n**2).reshape((n, n))
P = _P.T @ _P

q = np.random.normal(size=n)

A = np.hstack([np.zeros((m, n - m)), np.eye(m)])
A = np.random.normal(size=m * n).reshape((m, n))

l = -1.0 * np.ones(m)
u = 1.0 * np.ones(m)


# %%
def projection(z, l, u):
  return np.minimum(np.maximum(l, z), u)


# %%
def admm_iteration(x, y, z, P, q, A, l, u, rho=0.1, sigma=1.0e-6, alpha=1.6):
  # dimension
  n = len(x)
  m = len(l)

  # primal update (3)
  w = np.hstack([sigma * x - q, z - y / rho])
  W = np.vstack([
      np.hstack([P + sigma * np.eye(n), A.T]),
      np.hstack([A, -np.eye(m) / rho]),
  ])
  xv = np.linalg.solve(W, w)
  xnt = xv[:n]
  vn = xv[n:]

  # (4)
  znt = z + (vn - y) / rho

  # (5)
  xn = alpha * xnt + (1 - alpha) * x

  # (6)
  zn = projection(alpha * znt + (1 - alpha) * z + y / rho, l, u)

  # (7)
  yn = y + rho * (alpha * znt + (1 - alpha) * z - zn)

  return xn, yn, zn


# %%
def residual_primal(x, z, A):
  return A @ x - z


def residual_dual(x, y, P, q, A):
  return P @ x + q + A.T @ y


# %%
def primal_infeasible(y, A, l, u, eps=1.0e-3):
  if (
      np.linalg.norm(A.T @ y) / len(y) < eps
      and np.dot(u, np.maximum(y, 0.0)) + np.dot(l, np.minimum(y, 0.0)) < 0.0
  ):
    return True
  return False


def dual_infeasible(x, P, q, A, l, u, eps=1.0e-3):
  if (
      np.linalg.norm(P @ x) / len(x) < eps
      and np.dot(q, x) < 0.0
      and np.linalg.norm(A @ x) / len(x) < eps
  ):
    return True
  return False


# %%
# settings
rho = 0.1
sigma = 1.0e-6
alpha = 1.6
eps_rp = 1.0e-5
eps_rd = 1.0e-5
max_iter = 1000
# %%
# solve
x = np.zeros(n)
y = np.zeros(m)
z = np.zeros(m)

for i in range(max_iter):
  # iteration
  x, y, z = admm_iteration(x, y, z, P, q, A, l, u)

  # residuals
  rp = np.linalg.norm(residual_primal(x, z, A), np.inf)
  rd = np.linalg.norm(residual_dual(x, y, P, q, A), np.inf)

  # infeasible
  if primal_infeasible(y, A, l, u):
    print("primal infeasible")
    break
  if dual_infeasible(x, P, q, A, l, u):
    print("dual infeasible")
    break

  # status
  solved = rp < eps_rp and rd < eps_rd
  if i % max_iter / 10 == 0 or solved:
    print(f"residual ({i})")
    print(" primal: ", rp)
    print(" dual: ", rd)

  # convergence
  if solved:
    print("tolerance achieved")
    break
# %%
print("solution: \n", x)