admm.py

"""ADMM solver."""
import enum
from typing import Tuple
import numpy as np
from copy import deepcopy


class ADMMProjection(enum.IntEnum):
    """Choice of ADMM projection scheme."""
    STANDARD = 0
    STANDARDSLACK = 1
    ADMMSLACK = 2

class ADMMRhoUpdate(enum.IntEnum):
    """Choice of ADMM rho update rule."""
    NONE = 0
    BASIC = 1
    SQRT = 2
    OSQP = 3


class ADMM:
    def __init__(self, Q = None, q = None, A = None, l = None, u = None, options = {}):
        if Q is None or q is None or A is None or l is None or u is None:
            print("Problem data not set. Call update_problem_info() before use.")
            self.ready = False
        else:
            self.update_problem_info(Q, q, A, l, u, options)

    def update_problem_info(self, Q = None, q = None, A = None, l = None, u = None, options = {}):
        self.ready = True
        """
        Update the problem data
        Args:
            Q: (n, n) positive semidefinite matrix
            q: (n,) vector
            A: (m, n) constraint matrix
            l: (m,) lower bounds
            u: (m,) upper bounds
        """
        if Q is not None:
            self.Q = deepcopy(Q)
        if q is not None:
            self.q = deepcopy(q)
        if A is not None:
            self.A = deepcopy(A)
        if l is not None:
            self.l = deepcopy(l)
        if u is not None:
            self.u = deepcopy(u)
        self.n = Q.shape[0]
        self.m = A.shape[0]

        self.rho = options.get('rho', 0.1)   
        self.max_rho = options.get('max_rho', 1e4)
        self.min_rho = options.get('min_rho', 1e-4)
        self.rho = max(min(self.rho, self.max_rho), self.min_rho)
        self.rho_update_ratio = options.get('rho_update_ratio', 5)
        self.rho_update_steps = options.get('rho_update_steps', 25)

        self.alpha = options.get('alpha', 1000)
        self.max_iter = options.get('max_iter', 1000)
        self.primal_tol = options.get('primal_tol', 1e-10)
        self.dual_tol = options.get('dual_tol', 1e-10)
        self.check_for_exit = options.get('check_for_exit', True)

        # also set up the Q_bar for standard slacks
        self.Qbar = np.block([[self.Q,                     np.zeros((self.n, self.m))],
                              [np.zeros((self.m, self.n)), self.alpha * np.eye(self.m)]])
        self.qbar = np.concatenate([self.q, np.zeros(self.m)])
        self.Abar = np.hstack([self.A, np.eye(self.m)])
        '''
        uncomment to visualize sparsity pattern
        '''
        # import matplotlib.pyplot as plt
        # plt.spy(self.Abar)
        # plt.show()

    def x_update(
        self,
        z: np.array,
        q: np.array,
        A: np.array,
        P: np.array,
        mu: np.array,
        rho: float
    ) -> np.array:
        """
        Update x in ADMM by minimizing (1/2)x^TQx + q^Tx + (rho/2)||Ax - z||^2 + mu^T(Ax-z)

        Args:
            z: current value of z
            q: linear term in the objective
            A: constraint matrix
            P: precomputed matrix (Q + rho * A^T * A)
            mu: dual variable
            rho: consensus parameter
        Returns:
            x: updated variable
        """
        rhs = -q + A.T @ (-mu + rho * z)
        return np.linalg.solve(P, rhs)

    def z_update(
        self,
        x: np.array,
        A: np.array,
        mu: np.array,
        rho: float,
        projection: ADMMProjection = ADMMProjection.STANDARD
    ) -> np.array:
        """
        Update z in ADMM with the correct projection
        Args:
            A: constraint matrix
            x: current value of x
            mu: dual variable
            rho: consensus parameter
            mode: 'standard' or 'ADMMslack'
        Returns:
            z: updated variable
        """
        z_tilde = A @ x + mu / rho
        if projection == ADMMProjection.ADMMSLACK:
            z = np.where(
                z_tilde < self.l,
                (rho * z_tilde + self.alpha * self.l) / (rho + self.alpha),
                np.where(
                    z_tilde > self.u,
                    (rho * z_tilde + self.alpha * self.u) / (rho + self.alpha),
                    z_tilde
                )
            )
        elif projection in [ADMMProjection.STANDARD, ADMMProjection.STANDARDSLACK]:
            # Standard projection onto [l, u]
            z = np.minimum(np.maximum(z_tilde, self.l), self.u)
        else:
            raise ValueError(f"Unknown projection {projection}.")
        return z

    def mu_update_and_residuals(
        self,
        x: np.array,
        z: np.array,
        Q: np.array,
        q: np.array,
        A: np.array,
        mu: np.array,
        rho: float,
    ) -> Tuple[np.array, float, float]:
        """
        Update mu and compute primal/dual residuals
        Args:
            x: current value of x
            z: current value of z
            Q: quadratic cost matrix
            q: linear cost vector
            A: constraint matrix
            mu: dual variable
            rho: consensus parameter
        Returns:
            mu: updated dual variable
            primal_residual_norm: norm of the primal residual
            dual_residual_norm: norm of the dual residual
        """
        Axmz = A @ x - z
        mu += rho * Axmz

        primal_residual_norm = np.linalg.norm(Axmz, ord=np.inf)
        dual_residual_norm = np.linalg.norm(Q @ x + q + A.T @ mu, ord=np.inf)

        return mu, primal_residual_norm, dual_residual_norm

    def update_rho_osqp(
        self,
        x: np.array,
        z: np.array,
        A: np.array,
        Q: np.array,
        q: np.array,
        mu: np.array,
        rho: float,
        primal_residual_norm: float,
        dual_residual_norm: float,
    ):
        """
        Update rho using the OSQP-style update rule
        Args:
            x: current value of x
            z: current value of z
            A: constraint matrix
            Q: quadratic cost matrix
            q: linear cost vector
            mu: dual variable
            rho: consensus parameter
            primal_residual_norm: norm of the primal residual
            dual_residual_norm: norm of the dual residual
        Returns:
            rho: updated rho parameter
        """
        # Compute norms
        norm_Ax = np.linalg.norm(A @ x, ord=np.inf)
        norm_z = np.linalg.norm(z, ord=np.inf)
        norm_Qx = np.linalg.norm(Q @ x, ord=np.inf)
        norm_q = np.linalg.norm(q, ord=np.inf)
        norm_ATmu = np.linalg.norm(A.T @ mu, ord=np.inf)

        # Compute the ratio for rho update
        top = primal_residual_norm / max(norm_Ax, norm_z, 1e-12)
        bottom = 1e-12 + dual_residual_norm / max(norm_Qx, norm_q, norm_ATmu, 1e-12)
        ratio = np.sqrt(top / (bottom + 1e-12))  # add epsilon to avoid divide-by-zero
        return rho * ratio
    
    def update_rho(
        self,
        x: np.array,
        z: np.array,
        A: np.array,
        Q: np.array,
        q: np.array,
        mu: np.array,
        rho: float,
        primal_residual_norm: float,
        dual_residual_norm: float,
        rho_update: ADMMRhoUpdate = ADMMRhoUpdate.NONE
    ) -> float:
        """
        Update rho using different strategies
        Args:
            x: current value of x
            z: current value of z
            A: constraint matrix
            Q: quadratic cost matrix
            q: linear cost vector
            mu: dual variable
            rho: consensus parameter
            primal_residual_norm: norm of the primal residual
            dual_residual_norm: norm of the dual residual
            rho_update: strategy for updating rho
        Returns:
            rho: updated rho parameter
        """
        if (primal_residual_norm > self.rho_update_ratio * dual_residual_norm) or (dual_residual_norm > self.rho_update_ratio * primal_residual_norm):
            if rho_update == ADMMRhoUpdate.NONE:
                return rho
            elif rho_update == ADMMRhoUpdate.BASIC:
                rho_basic = rho
                rho_basic *= 2 if (primal_residual_norm > dual_residual_norm) else 0.5
                rho_basic = max(self.min_rho, min(self.max_rho, rho_basic))
                return rho_basic
            elif rho_update == ADMMRhoUpdate.SQRT:
                rho_sqrt_ratio = rho * np.sqrt(primal_residual_norm / (dual_residual_norm + 1e-12)) # add epsilon to avoid divide-by-zero
                rho_sqrt_ratio = max(self.min_rho, min(self.max_rho, rho_sqrt_ratio))
                return rho_sqrt_ratio
            elif rho_update == ADMMRhoUpdate.OSQP:
                rho_osqp = rho
                rho_osqp = self.update_rho_osqp(x, z, A, Q, q, mu, rho, primal_residual_norm, dual_residual_norm)
                rho_osqp = max(self.min_rho, min(self.max_rho, rho_osqp))
                return rho_osqp
            else:
                raise ValueError(f"Unknown rho update approach {rho_update}")
        else:
            return rho

    def solve(
        self,
        projection: ADMMProjection = ADMMProjection.STANDARD,
        rho_update: ADMMRhoUpdate = ADMMRhoUpdate.NONE,
        return_traces: bool = False
    ) -> Tuple[np.array, int, float, float]:
        if not self.ready:
            raise ValueError("Problem data not set. Call update_problem_info() first.")
        """
        Solve the standard problem.

        Args:
            projection: projection used in the z update, see ADMMProjection for options.
            rho_update: rho update rule, see ADMMRhoUpdate for options.
        Returns:
            x: optimal solution
            num_iter: number of iterations
            primal_residual_norm: norm of the primal residual
            dual_residual_norm: norm of the dual residual
        """
        n, m = self.n, self.m
        if projection == ADMMProjection.STANDARDSLACK:
            # stack the slacks and augment the problem
            Q = self.Qbar
            A = self.Abar
            q = self.qbar
            x = np.zeros(n+m)
        elif projection in [ADMMProjection.STANDARD, ADMMProjection.ADMMSLACK]:
            Q = self.Q
            A = self.A
            q = self.q
            x = np.zeros(n)
        else:
            raise ValueError(f"Unknown projection {projection}.")

        # initialization
        z = np.zeros(m)
        mu = np.zeros(m)
        rho = self.rho
        P = Q + rho * (A.T @ A) # precompute for later use
        primal_residual_norm = np.inf
        dual_residual_norm = np.inf
        traces = {}
        if return_traces:
            traces['X'] = [x[:n]]
            traces['Z'] = [z]
            traces['mu'] = [mu]
            traces['rho'] = [rho]
            traces['primal'] = [primal_residual_norm]
            traces['dual'] = [dual_residual_norm]

        # ADMM loop
        for num_iter in range(self.max_iter):
            x = self.x_update(z, q, A, P, mu, rho)
            z = self.z_update(x, A, mu, rho, projection)
            mu, primal_residual_norm, dual_residual_norm = \
                self.mu_update_and_residuals(x, z, Q, q, A, mu, rho)
               
            if return_traces:
                traces['X'].append(x[:n])
                traces['Z'].append(z)
                traces['mu'].append(mu)
                traces['rho'].append(rho)
                traces['primal'].append(primal_residual_norm)
                traces['dual'].append(dual_residual_norm)

            # check for exit
            if (
                self.check_for_exit and 
                primal_residual_norm < self.primal_tol and
                dual_residual_norm < self.dual_tol
            ):
                
                break

            # update rho
            if not(num_iter % self.rho_update_steps):
                rho = self.update_rho(x, z, A, Q, q, mu, rho, primal_residual_norm, dual_residual_norm, rho_update)
                P = Q + rho * (A.T @ A) # precompute for later use
        
        num_iter += 1
        x = x[:n] # make sure to return only the original x part (if augmented)
        if return_traces:
            return x, num_iter, primal_residual_norm, dual_residual_norm, traces
        else:
            return x, num_iter, primal_residual_norm, dual_residual_norm