Augmented Lagrangian Method

About

The Extended Lagrangian Method is a technique for solving constrained optimization problems, especially those with equality and inequality constraints. This method can be considered as a combination of the Lagrangian Relaxation and the Penalty Method.

Augmented Lagrangian Method

Problem Setting

What this algorithm is going to solve is following minimization problem with equality and inequality constraints.

\min_{x} f(x)\\
\text{subject to} \quad h_i(x) = 0, \quad i = 1, ..., m \\
g_j(x) \leq 0, \quad j = 1, ..., p \\

In this case, the conventional Lagrangian function is defined as follows

\mathcal{L}(x, \lambda, \mu) = f(x) + \sum_{i=1}^{m} \lambda_i h_i(x) + \sum_{j=1}^{p} \mu_j g_j(x)

In this conventional LM, constraints are not always strictly satisfied, and it is difficult to control when constraints are not satisfied. Because of that, ALM is introduced.

\mathcal{L}_{\rho}(x, \lambda) = f(x) + \sum_{i=1}^{m} \lambda_i h_i(x) + \frac{\rho}{2} \sum_{i=1}^{m} h_i(x)^2\\
+ \sum_{j=1}^{p} \frac{1}{\rho} \left[ \max(0, \mu_j + \rho g_j(x))^2 - \mu_j^2 \right]

Algorithm

Step 1: Initialize values of variables, x^(0), \lambda^(0), \mu^(0), \rho^(0)

Step 2: minimize augmented lagrangian

x^{(k+1)} = \arg \min_x \mathcal{L}_{\rho}(x, \lambda^{(k)}, \mu^{(k)})

Step 3: update both lagrange multipliers, for equality and inequality.

\lambda_i^{(k+1)} = \lambda_i^{(k)} + \rho h_i(x^{(k+1)})\\
\mu_j^{(k+1)} = \max(0, \mu_j^{(k)} + \rho g_j(x^{(k+1)}))

Step4: If the convergence condition ||h(x)|| \leq \epsilon or ||g(x)|| \leq \epsilon are satisfied, the process will end

Psudo Code from ChatGPT

# Augmented Lagrangian Method (ALM) - Pseudocode

# Function: Augmented Lagrangian Optimization
# Inputs:
#   f(x)       - Objective function
#   h(x)       - Equality constraint functions (h_i(x) = 0)
#   g(x)       - Inequality constraint functions (g_j(x) <= 0)
#   x_init     - Initial guess for x
#   lambda_init - Initial Lagrange multipliers for equality constraints
#   mu_init    - Initial Lagrange multipliers for inequality constraints
#   rho_init   - Initial penalty parameter
#   tol        - Convergence tolerance
#   max_iter   - Maximum number of iterations

function Augmented_Lagrangian_Method(f, h, g, x_init, lambda_init, mu_init, rho_init, tol, max_iter):
    
    # Initialize variables
    x = x_init
    lambda = lambda_init  # Lagrange multipliers for equality constraints
    mu = mu_init          # Lagrange multipliers for inequality constraints
    rho = rho_init        # Initial penalty parameter
    k = 0                 # Iteration counter

    while k < max_iter:
        
        # Step 1: Minimize the Augmented Lagrangian function w.r.t. x
        x_new = argmin_x L_augmented(x, lambda, mu, rho)
        
        # Step 2: Update Lagrange multipliers
        lambda = lambda + rho * h(x_new)  # Equality constraints update
        
        for j in range(len(mu)):
            mu[j] = max(0, mu[j] + rho * g[j](x_new))  # Inequality constraints update
        
        # Step 3: Check convergence
        if norm(h(x_new)) < tol and norm([max(0, g_j(x_new)) for g_j in g]) < tol:
            break  # Converged
        
        # Step 4: Update x
        x = x_new
        
        # Optionally update penalty parameter rho
        rho = update_rho(rho, h(x), g(x))
        
        k += 1
    
    return x, lambda, mu

ALM in the case of Force Equilibrium

What to be minimized is below energy function with boundary constraint BU

\begin{align}
\mathcal{L}_{\text{normal}}(U, \lambda) = \frac{1}{2} U^T K U - F^T U + \lambda^T B U
\end{align}

If the equation is differentiated with respect to U, we get following equation

\begin{align}
     \begin{bmatrix}
     K & B^T \\
     B & 0
     \end{bmatrix}
     \begin{bmatrix}
     U \\
     \lambda
     \end{bmatrix}
     =
     \begin{bmatrix}
     F \\
     0
     \end{bmatrix}
\end{align}

But sinice

\begin{align}
\mathcal{L}_{\text{aug}}(U, \lambda) = \frac{1}{2} U^T K U - F^T U + \lambda^T B U + \frac{\rho}{2} \|B U\|^2
\end{align}

 (K + \rho B^T B) U = F - B^T \lambda

Implementation

import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as spla

# Stiffness Matrix
K = sp.csc_matrix([[4, -1, 0, 0],
                   [-1, 4, -1, 0],
                   [0, -1, 4, -1],
                   [0, 0, -1, 3]])

# Force Vector
f = np.array([1, 0, 0, 0])

# Boundary Condition
B = sp.csc_matrix([[1, 0, 0, 0]]) 

# Parameters for Augmented Lagrangian
rho = 10.0
lambda_ = np.zeros(B.shape[0])
max_iter = 100
tol = 1e-6

A = K + rho * (B.T @ B)
A_factorized = spla.factorized(A)  # LU
U = np.zeros(K.shape[0])

#M = spla.spilu(A)
#M_x = lambda x: M.spsolve(x)
#M_prec = spla.LinearOperator(A.shape, M_x)

U, info = spla.cg(A, b, tol=1e-6, maxiter=1000, M=M_prec)

for i in range(max_iter):
    b = f - B.T @ lambda_

    U = A_factorized(b)
    
    #U, info = spla.cg(A, b)
    #U, info = spla.cg(A, b, M=M_prec)
    #if info != 0:
    #    print(f"Warning: CG did not converge, info={info}")

    lambda_ += rho * (B @ U)

    constraint_violation = np.linalg.norm(B @ U)
    if constraint_violation < tol:
        print(f"Converged at iteration {i+1}")
        break

print("Optimal U:", U)