# Copyright (c) 2024, Kerry He, James Saunderson, and Hamza Fawzi
# This Python package QICS is licensed under the MIT license; see LICENSE.md
# file in the root directory or at https://github.com/kerry-he/qics
import itertools
import numba as nb
import numpy as np
import scipy as sp
from qics._utils.linalg import congr_multi, scale_axis, x_dot_dense
from qics.cones.base import SymCone
from qics.vectorize import vec_to_mat
[docs]
class PosSemidefinite(SymCone):
r"""A class representing a positive semidefinite cone defined either
on real symmetric matrices
.. math::
\mathbb{S}^n_+ = \{ X \in \mathbb{S}^n : X \succeq 0 \},
or complex Hermitian matrices
.. math::
\mathbb{H}^n_+ = \{ X \in \mathbb{H}^n : X \succeq 0 \}.
Parameters
----------
n : :obj:`int`
Dimension of the matrix :math:`X`.
iscomplex : :obj:`bool`
Whether the matrix :math:`X` is defined over :math:`\mathbb{H}^n`
(``True``), or restricted to :math:`\mathbb{S}^n` (``False``). The
default is ``False``.
"""
def __init__(self, n, iscomplex=False):
self.n = n
self.iscomplex = iscomplex
self.nu = n # Barrier parameter
if iscomplex:
self.dim = [2 * n * n]
self.type = ["h"]
self.dtype = np.complex128
else:
self.dim = [n * n]
self.type = ["s"]
self.dtype = np.float64
# Get LAPACK operators
X = np.eye(self.n, dtype=self.dtype)
self.cho_fact = sp.linalg.lapack.get_lapack_funcs("potrf", (X,))
self.cho_inv = sp.linalg.lapack.get_lapack_funcs("trtri", (X,))
self.svd = sp.linalg.lapack.get_lapack_funcs("gesdd", (X,))
self.svd_lwork = sp.linalg.lapack.get_lapack_funcs("gesdd_lwork", (X,))(n, n)
if iscomplex:
self.eigvalsh = sp.linalg.lapack.get_lapack_funcs("heevr", (X,))
else:
self.eigvalsh = sp.linalg.lapack.get_lapack_funcs("syevr", (X,))
# Update flags
self.feas_updated = False
self.grad_updated = False
self.nt_aux_updated = False
self.congr_aux_updated = False
return
def get_iscomplex(self):
return self.iscomplex
def get_init_point(self, out):
point = [np.eye(self.n, dtype=self.dtype)]
self.set_point(point, point)
out[0][:] = self.X
return out
def set_point(self, primal, dual=None, a=True):
self.X = primal[0] * a
self.Z = dual[0] * a if (dual is not None) else None
self.feas_updated = False
self.grad_updated = False
self.nt_aux_updated = False
def set_dual(self, dual, a=True):
self.Z = dual[0] * a
def get_feas(self):
if self.feas_updated:
return self.feas
self.feas_updated = True
# Check that X is PSD by trying to Cholesky factor it
self.X_chol, info = self.cho_fact(self.X, lower=True)
if info != 0:
self.feas = False
return self.feas
# Check that Z is PSD by trying to Cholesky factor it
if self.Z is not None:
self.Z_chol, info = self.cho_fact(self.Z, lower=True)
if info != 0:
self.feas = False
return self.feas
self.feas = True
return self.feas
def get_dual_feas(self):
self.Z_chol, info = self.cho_fact(self.Z, lower=True)
return info == 0
def get_val(self):
(sign, logabsdet) = np.linalg.slogdet(self.X)
return -sign * logabsdet
def update_grad(self):
assert not self.grad_updated
self.X_chol_inv, _ = self.cho_inv(self.X_chol, lower=True)
self.X_inv = self.X_chol_inv.conj().T @ self.X_chol_inv
self.grad = [-self.X_inv]
self.grad_updated = True
def hess_prod_ip(self, out, H):
if not self.grad_updated:
self.update_grad()
XHX = self.X_inv @ H[0] @ self.X_inv
np.add(XHX, XHX.conj().T, out=out[0])
out[0] *= 0.5
return out
def hess_congr(self, A):
if not self.grad_updated:
self.update_grad()
return self.base_congr(A, self.X_inv, self.X_chol_inv.conj().T)
def invhess_prod_ip(self, out, H):
XHX = self.X @ H[0] @ self.X
np.add(XHX, XHX.conj().T, out=out[0])
out[0] *= 0.5
return out
def invhess_congr(self, A):
return self.base_congr(A, self.X, self.X_chol)
def base_congr(self, A, X, X_rt2):
# Generalized function to compute the matrix [<Ai, X Aj X>]_ij for a
# a given matrix X
if not self.congr_aux_updated:
self.congr_aux(A)
(n, p, p_ds) = (self.n, A.shape[0], len(self.A_ds_idxs))
out = np.zeros((p, p))
if len(self.A_sp_idxs) > 0:
# Compute sparse-sparse component using Numba compiled functions
if self.iscomplex:
_sparse_congr_complex(out, self.A_sp_rows, self.A_sp_cols,
self.A_sp_data, self.A_sp_nnzs, X, self.A_sp_idxs) # fmt: skip
else:
_sparse_congr(out, self.A_sp_rows, self.A_sp_cols,
self.A_sp_data, self.A_sp_nnzs, X, self.A_sp_idxs) # fmt: skip
# Compute sparse-dense and dense-dense components
if p_ds > 0:
work = self.work
# Compute X Aj X for all dense Aj
lhs = np.zeros((p_ds, self.dim[0]))
lhs_view = lhs.reshape((p_ds, n, -1)).view(self.dtype)
congr_multi(lhs_view, X, self.Ai_ds, work=work)
# Compute inner products <Ai, X Aj X> for all dense Aj
temp = x_dot_dense(self.A_triu, lhs[:, self.triu_idxs].T)
out[:, self.A_ds_idxs] = temp
out[self.A_ds_idxs, :] = temp.T
else:
# If there are no (nonzero) sparse Aj, then compute congruence using
# the symmetric multiplication [A (L' kr L)] [A (L' kr L)]'
work = self.work
# Compute L' Aj L
lhs = np.zeros((p_ds, self.dim[0]))
lhs_view = lhs.reshape((p_ds, n, -1)).view(self.dtype)
congr_multi(lhs_view, X_rt2.conj().T, self.Ai_ds, work=work)
# Compute inner products <L' Ai L, L' Aj L>
out[self.A_ds_ds_idxs] = lhs @ lhs.T
return out
def third_dir_deriv_axpy(self, out, H, a=True):
if not self.grad_updated:
self.update_grad()
XHX_2 = self.X_inv @ H[0] @ self.X_chol_inv.conj().T
out[0] -= 2 * a * XHX_2 @ XHX_2.conj().T
return out
def prox(self):
assert self.feas_updated
XZX_I = self.X_chol.conj().T @ self.Z @ self.X_chol
XZX_I.flat[:: self.n + 1] -= 1
return np.linalg.norm(XZX_I) ** 2
# ==========================================================================
# Functions specific to symmetric cones for NT scaling
# ==========================================================================
# Computes the NT scaling point W and scaled variable Lambda such that
# H(W)[S] = Z <==> Lambda := P^-T(S) = P(Z)
# where H(W) = V^T V. To obtain for for the PSD cone, first let compute the
# SVD
# U D V^T = Z_chol^T S_chol
# Then compute
# R := S_chol V D^-1/2 = Z_chol^-T U D^1/2
# R^-1 := D^1/2 V^T S_chol^-1 = D^-1/2 U^T Z_chol^T
# Then we can find the scaling point as
# W := R R^T
# = S^1/2 (S^1/2 Z S^1/2)^-1/2 S^1/2
# = Z^-1/2 (Z^1/2 S Z^1/2)^1/2 Z^1/2 (i.e., geo. mean of Z and S)
# W^-1 := R^-T R^-1
# and the scaled point as
# Lambda := D
# Also, we have the linear transformations given by
# H(W)[S] = W^-1 S W^-1
# P^-T(S) = R^-1 S R^-T
# P(Z) = R^T Z R
# See: [Section 4.3]https://www.seas.ucla.edu/~vandenbe/publications/coneprog.pdf
def nt_aux(self):
assert not self.nt_aux_updated
# Take the SVD of Z_chol^T S_chol to get scaled point Lambda := D
RL = self.Z_chol.conj().T @ self.X_chol
_, self.Lambda, Vt, _ = self.svd(RL, lwork=self.svd_lwork)
D_rt2 = np.sqrt(self.Lambda)
# Compute the scaling point as
# R := S_chol V D^-1/2, and
# W := R R^T
self.R = self.X_chol @ (Vt.conj().T / D_rt2)
self.W = self.R @ self.R.conj().T
# Compute the inverse scaling point as
# R^-1 := D^1/2 V^T S_chol^-1, and
# W^-1 := R^-T R^-1
self.X_chol_inv, _ = self.cho_inv(self.X_chol, lower=True)
self.R_inv = (self.X_chol_inv.conj().T @ (Vt.conj().T * D_rt2)).conj().T
self.W_inv = self.R_inv.conj().T @ self.R_inv
self.nt_aux_updated = True
def nt_prod_ip(self, out, H):
if not self.nt_aux_updated:
self.nt_aux()
WHW = self.W_inv @ H[0] @ self.W_inv
np.add(WHW, WHW.conj().T, out=out[0])
out[0] *= 0.5
def nt_congr(self, A):
if not self.nt_aux_updated:
self.nt_aux()
return self.base_congr(A, self.W_inv, self.R_inv.conj().T)
def invnt_prod_ip(self, out, H):
if not self.nt_aux_updated:
self.nt_aux()
WHW = self.W @ H[0] @ self.W
np.add(WHW, WHW.conj().T, out=out[0])
out[0] *= 0.5
def invnt_congr(self, A):
if not self.nt_aux_updated:
self.nt_aux()
return self.base_congr(A, self.W, self.R)
def comb_dir(self, out, dS, dZ, sigma_mu):
# Compute the residual for rs where rs is given as the lhs of
# Lambda o (W dz + W^-T ds) = -Lambda o Lambda - (W^-T ds_a) o (W dz_a)
# + sigma * mu * I
# which is rearranged into the form H ds + dz = rs, i.e.,
# rs := W^-1 [ Lambda \ (-Lambda o Lambda - (W^-T ds_a) o (W dz_a) + sigma*mu I) ]
# See: [Section 5.4]https://www.seas.ucla.edu/~vandenbe/publications/coneprog.pdf
if not self.nt_aux_updated:
self.nt_aux()
# Compute (W^-T ds_a) o (W dz_a) = (R^-1 dS R^-T) o (R^T dZ R)
# = 0.5 * ([R^-1 dS dZ R] + [R^T dZ dS R^-T])
temp1 = self.R_inv @ dS[0]
temp2 = dZ[0] @ self.R
temp3 = temp1 @ temp2
np.add(temp3, temp3.conj().T, out=temp1)
temp1 *= -0.5
# Compute -Lambda o Lambda - [ ... ] + sigma*mu I
# Note that Lambda is a diagonal matrix
temp1.flat[:: self.n + 1] -= np.square(self.Lambda)
temp1.flat[:: self.n + 1] += sigma_mu
# Compute Lambda \ [ ... ]
# Since Lambda is diagonal, the solution to the Sylvester equation
# find X s.t. 0.5 * (Lambda X + X Lambda) = B
# is given by
# X = B .* (2 / [Lambda_ii + Lambda_jj]_ij)
Gamma = np.add.outer(self.Lambda, self.Lambda)
temp1 /= Gamma
# Compute W^-1 [ ... ] = R^-T [... ] R^-1
temp = self.R_inv.conj().T @ temp1 @ self.R_inv
np.add(temp, temp.conj().T, out=out[0])
def step_to_boundary(self, dS, dZ):
# Compute the maximum step alpha in [0, 1] we can take such that
# S + alpha dS >= 0
# Z + alpha dZ >= 0
# See: [Section 8.3]https://www.seas.ucla.edu/~vandenbe/publications/coneprog.pdf
if not self.nt_aux_updated:
self.nt_aux()
Lambda_irt2 = np.reciprocal(np.sqrt(self.Lambda))
# Compute rho := H(lambda)^1/2 W^-T dS
# = Lambda^-1/2 R^-1 dS R^-T Lambda^-1/2
rho = self.R_inv @ dS @ self.R_inv.conj().T
rho *= Lambda_irt2.reshape((-1, 1))
rho *= Lambda_irt2.reshape((1, -1))
# Compute sig := H(lambda)^1/2 W dS
# = Lambda^-1/2 R^T dS R Lambda^-1/2
sig = self.R.conj().T @ dZ @ self.R
sig *= Lambda_irt2.reshape((-1, 1))
sig *= Lambda_irt2.reshape((1, -1))
# Compute minimum eigenvalues of rho and sig
min_eig_rho = self.eigvalsh(rho, compute_v=False, range="I", iu=1)[0][0]
min_eig_sig = self.eigvalsh(sig, compute_v=False, range="I", iu=1)[0][0]
# Maximum step is given by
# alpha := 1 / max(0, -min(eig(rho)), -min(eig(sig)))
# Clamp this step between 0 and 1
if min_eig_rho >= 0 and min_eig_sig >= 0:
return 1.0
else:
return 1.0 / max(-min_eig_rho, -min_eig_sig)
# ==========================================================================
# Auxilliary functions
# ==========================================================================
def congr_aux(self, A):
assert not self.congr_aux_updated
n = self.n
if sp.sparse.issparse(A):
A = A.tocsr()
# Split A into sparse and dense groups
A_nnz = A.getnnz(1)
self.A_sp_idxs = np.where((A_nnz > 0) & (A_nnz < n))[0]
self.A_ds_idxs = np.where((A_nnz >= n))[0]
self.A_sp_idxs = self.A_sp_idxs[np.argsort(A_nnz[self.A_sp_idxs])]
self.A_ds_ds_idxs = np.ix_(self.A_ds_idxs, self.A_ds_idxs)
# Prepare things we need for Strategy 1, i.e., sparse-sparse
# congruence using Numba functions
if len(self.A_sp_idxs) > 0:
A_sp = A[self.A_sp_idxs]
triu_idxs = _get_triu_idxs(n)
if self.iscomplex:
A_sp_real = A_sp[:, ::2][:, triu_idxs]
A_sp_imag = A_sp[:, 1::2][:, triu_idxs]
A_sp_lil = (A_sp_real + A_sp_imag * 1j).tolil()
else:
A_sp_lil = A_sp[:, triu_idxs].tolil()
# Get number of nonzeros for each sparse Ai (to account for
# ragged arrays)
self.A_sp_nnzs = A_sp_lil.getnnz(1)
# Get rows and columns of nonzeros of sparse Ai
rowcols = _lil_to_array(A_sp_lil.rows)
self.A_sp_rows, self.A_sp_cols = _triu_idx_to_ij(rowcols)
# Get values of nonzeros of sparse Ai, and scale off-diagonal
# elements to account for only using upper triangular nonzeros
self.A_sp_data = _lil_to_array(A_sp_lil.data)
self.A_sp_data[self.A_sp_cols != self.A_sp_rows] *= 2
# Prepare things we need for Strategy 2, i.e., sparse-dense and
# dense-dense congruence by computing X Aj X for all dense Aj, then
# <Ai, X Aj X> for all Ai and dense Aj
if len(self.A_ds_idxs) > 0:
A_ds = A[self.A_ds_idxs, :]
# Turn rows of A into matrices Ai
if self.iscomplex:
A_ds = A_ds[:, ::2] + A_ds[:, 1::2] * 1j
A_ds = A_ds.toarray()
self.Ai_ds = np.array([Ai.reshape((n, n)) for Ai in A_ds])
self.work = np.zeros_like(self.Ai_ds)
# Get upper triangular slices of A so we can more efficiently
# do the inner product <Ai, X Aj X> as a matrix multiplication
if len(self.A_sp_idxs) > 0:
# Scale upper triangular elements by 2 to account for only
# using upper triangular elements
scale = 2 * np.ones(self.dim[0])
if self.iscomplex:
scale[:: 2 * n + 2] = 1
else:
scale[:: n + 1] = 1
self.triu_idxs = _get_triu_idxs(n, self.iscomplex)
self.A_triu = scale_axis(A, scale_cols=scale).tocsr()
self.A_triu = self.A_triu[:, self.triu_idxs].tocoo()
else:
# A and all Ai are dense matrices
# Just need to convert the rows of A into dense matrices
A = np.ascontiguousarray(A)
self.A_sp_idxs = np.array([])
self.A_ds_idxs = np.arange(A.shape[0])
self.A_ds_ds_idxs = np.ix_(self.A_ds_idxs, self.A_ds_idxs)
self.Ai_ds = np.array([vec_to_mat(Ak, self.iscomplex) for Ak in A])
self.work = np.zeros_like(self.Ai_ds)
self.congr_aux_updated = True
def _lil_to_array(ragged):
# Converts a list of lists (with possibly different lengths) into a numpy
# array padded with zeros
padded_list = list(itertools.zip_longest(*ragged, fillvalue=0))
return np.array(padded_list).T
def _triu_idx_to_ij(idx):
# Converts upper triangular indices to (i,j) coordinates
# [ 0 1 3 ] [ (0,0) (0,1) (0,2) ]
# [ 2 4 ... ] --> [ (1,1) (1,2) ... ]
# [ 5 ] [ (2,2) ]
# See: https://stackoverflow.com/questions/40950460/
j = np.ceil(np.sqrt(2 * (idx + 1) + 0.25) - 0.5) - 1
i = idx - (j + 1) * j / 2
return i.astype("int32"), j.astype("int32")
def _get_triu_idxs(n, iscomplex=False):
# Gets indices of a vectorized matrix corresponding to the upper triangular
# elements of the matrix
if iscomplex:
diag = [2 * (i + i * n) for i in range(n)]
triu_real = [2 * (j + i * n) for j in range(n) for i in range(j)]
triu_imag = [2 * (j + i * n) + 1 for j in range(n) for i in range(j)]
return np.array(diag + triu_real + triu_imag)
else:
return np.array([j + i * n for j in range(n) for i in range(j + 1)])
# ============================================================================
# Numba functions for computing Schur complement matrix when A is very sparse
# ============================================================================
@nb.njit(cache=True, parallel=True, fastmath=True)
def _sparse_congr(out, A_rows, A_cols, A_vals, A_nnz, X, indices):
# Computes the congruence transform A (X kron X) A' when A is very sparse
# See https://link.springer.com/article/10.1007/BF02614319
# We can cut the amount of operations in half by exploiting symetry of A and
# X as follows
# (AHA)_ij = Σ_a,b (Ai)_ab (Σ_c,d (Aj)_cd X_ac X_db)
# = Σ_a,b (Ai)_ab ( [Σ_c=d (Aj)_cd X_ac X_db]
# + [Σ_c<d (Aj)_cd X_ac X_db]
# + [Σ_c>d (Aj)_cd X_ac X_db] )
# = Σ_a,b (Ai)_ab ( [Σ_c=d (Aj)_cd X_ac X_db]
# + [Σ_c<d (Aj)_cd (X_ac X_db + X_ad X_cb)] )
# = [Σ_a=b (Ai)_ab ( ... )] + [Σ_a<b (Ai)_ab ( ... )] + [Σ_a>b (Ai)_ab ( ... )]
# = [Σ_a=b (Ai)_ab ( ... )] + 2 [Σ_a<b (Ai)_ab ( ... )]
# Note that we assume off-diagonal entries of Ai have been scaled by 2.
# Also note that we assume only upper triangular elements are given to us so
# c < d.
n = A_rows.shape[0]
# Loop through upper triangular entries of the Schur complement matrix (AHA)_ij
for j in nb.prange(n):
for i in nb.prange(j + 1):
i_AHA = indices[i]
j_AHA = indices[j]
tmp1 = 0.0
# Loop over nonzero entries of Ai
for alpha in range(A_nnz[i]):
a = A_rows[i, alpha]
b = A_cols[i, alpha]
tmp2 = 0.0
tmp3 = 0.0
# Loop over nonzero entries of Aj
for beta in range(A_nnz[j]):
c = A_rows[j, beta]
d = A_cols[j, beta]
if c < d:
tmp2 += A_vals[j, beta] * (
X[a, c] * X[d, b] + X[a, d] * X[c, b]
)
else:
tmp3 += A_vals[j, beta] * X[a, c] * X[d, b]
tmp1 += A_vals[i, alpha] * (0.5 * tmp2 + tmp3)
if i_AHA <= j_AHA:
out[i_AHA, j_AHA] = tmp1
else:
out[j_AHA, i_AHA] = tmp1
@nb.njit(cache=True, parallel=True, fastmath=True)
def _sparse_congr_complex(out, A_rows, A_cols, A_vals, A_nnz, X, indices):
# Computes the congruence transform A (X kron X) A' when A is very sparse
# See https://link.springer.com/article/10.1007/BF02614319
# We can cut the amount of operations in half by exploiting symetry of A and
# X as follows
# (AHA)_ij = Σ_a,b (Ai)_ab* (Σ_c,d (Aj)_cd X_ac X_db)
# = Σ_a,b (Ai)_ab* ( [Σ_c=d (Aj)_cd X_ac X_db]
# + [Σ_c>d (Aj)_cd X_ac X_db]
# + [Σ_c<d (Aj)_cd X_ac X_db] )
# = Σ_a,b (Ai)_ab* ( [Σ_c=d (Aj)_cd X_ac X_db]
# + [Σ_c<d (Aj)_cd X_ac X_db + (Aj)_cd* X_ad X_cb] )
# = [Σ_a=b (Ai)_ab ( ... )] + [Σ_a<b (Ai)_ab ( ... )] + [Σ_a>b (Ai)_ab ( ... )]
# = [Σ_a=b (Ai)_ab ( ... )] + [Σ_a>b (Ai)_ab ( ... ) + (Ai)_ab* ( ... )]
# Also note that off-diagonal entries of Ai have been scaled by 2
n = A_rows.shape[0]
# Loop through each entry of the Schur complement matrix (AHA)_ij
for j in nb.prange(n):
for i in nb.prange(j + 1):
i_AHA = indices[i]
j_AHA = indices[j]
tmp1 = 0.0
for alpha in range(A_nnz[i]):
a = A_rows[i, alpha]
b = A_cols[i, alpha]
tmp2 = 0.0
tmp3 = 0.0
for beta in range(A_nnz[j]):
c = A_rows[j, beta]
d = A_cols[j, beta]
if c < d:
tmp2 += A_vals[j, beta] * X[a, c] * X[d, b]
tmp2 += np.conj(A_vals[j, beta]) * X[a, d] * X[c, b]
else:
tmp3 += A_vals[j, beta].real * X[a, c] * X[d, b]
# Do addition slightly differently to guarantee a real number
# i.e., just take the inner product between Ai and X Aj X by
# Σ_ab Re[(X Aj X)_ab] * Re[(Ai)_ab] + 2 Im[(X Aj X)_ab] * Im[(Ai)_ab]
tmp3 += 0.5 * tmp2
tmp1 += A_vals[i, alpha].real * tmp3.real
tmp1 += A_vals[i, alpha].imag * tmp3.imag
if i_AHA <= j_AHA:
out[i_AHA, j_AHA] = tmp1
else:
out[j_AHA, i_AHA] = tmp1
QICS