# 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 numpy as np
import scipy as sp
from qics._utils.gradient import D1_log, D2_log, scnd_frechet
from qics._utils.linalg import (
cho_fact,
cho_solve,
congr_multi,
dense_dot_x,
inp,
x_dot_dense,
)
from qics.cones.base import Cone
from qics.vectorize import get_full_to_compact_op, vec_to_mat
[docs]
class QuantKeyDist(Cone):
r"""A class representing a quantum key distribution cone
.. math::
\mathcal{QKD}_{\mathcal{G},\mathcal{Z}} =
\text{cl}\{ (t, X) \in \mathbb{R} \times \mathbb{H}^n_{++} :
t \geq -S(\mathcal{G}(X)) + S(\mathcal{Z}(\mathcal{G}(X))) \},
where
.. math::
S(X) = -\text{tr}[X \log(X)],
is the quantum (von Neumann) entropy function,
:math:`\mathcal{G}:\mathbb{H}^n\rightarrow\mathbb{H}^{mr}` is a
positive linear map, and
:math:`\mathcal{Z}:\mathbb{H}^{mr}\rightarrow\mathbb{H}^{mr}` is a
pinching map that maps off-diagonal blocks to zero.
Parameters
----------
G_info : :obj:`int` or :obj:`list` of :class:`~numpy.ndarray`
Defines the linear map :math:`\mathcal{G}`. There are two ways to
specify this linear map.
- If ``G_info`` is an :obj:`int`, then :math:`\mathcal{G}` is the
identity map, i.e., :math:`\mathcal{G}(X)=X`, and ``G_info``
specifies the dimension of :math:`X`.
- If ``G_info`` is a :obj:`list` of :class:`~numpy.ndarray`, then
``G_info`` specifies the Kraus operators
:math:`K_i \in \mathbb{C}^{mr \times n }` corresponding to
:math:`\mathcal{G}` such that
.. math::
\mathcal{G}(X) = \sum_{i} K_i X K_i^\dagger.
Z_info : :obj:`int` or :obj:`tuple` or :obj:`list` of :class:`~numpy.ndarray`
Defines the pinching map :math:`\mathcal{Z}`, which is of the form
.. math::
\mathcal{Z}(Y) = \sum_{i} Z_i Y Z_i^\dagger.
There are three ways to specify this linear map.
- If ``Z_info`` is an :obj:`int`, then
:math:`Z_i=|i \rangle\langle i| \otimes\mathbb{I}` for
:math:`i=1,\ldots,r`, where ``r=Z_info``.
- If ``Z_info`` is a :obj:`tuple` of the form ``(dims, sys)``,
where ``dims=(n0, n1)`` is a :obj:`tuple` of :obj:`int` and
``sys`` is an :obj:`int`, then
- :math:`Z_i=|i \rangle\langle i| \otimes\mathbb{I}_{n_1}`
for :math:`i=1,\ldots,n_0` if ``sys=0``, and
- :math:`Z_i=\mathbb{I}_{n_0}\otimes |i \rangle\langle i|` for
:math:`i=1,\ldots,n_1` if ``sys=1``.
We generalize this definition to the case where ``dims`` and ``sys``
are lists of any length.
- If ``Z_info`` is a :obj:`list` of :class:`~numpy.ndarray`, then
``Z_info`` directly specifies the Kraus operators
:math:`Z_i \in \mathbb{C}^{mr \times mr}`. Note that these Kraus
operators must have a similar structure to those defined using
the other options, i.e., must be diagonal matrices consisting of
either ones or zeros, and :math:`Z_iZ_j=0` for all :math:`i\neq j`.
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``.
See also
--------
QuantRelEntr : Quantum relative entropy cone
Notes
-----
The quantum key distribution cone can also be modelled by the quantum
relative entropy by noting the identity
.. math::
S(\mathcal{G}(X) \| \mathcal{Z}(\mathcal{G}(X)))
= -S(\mathcal{G}(X)) + S(\mathcal{Z}(\mathcal{G}(X))).
However, the cone oracles for the quantum key distribution cone are
more efficient than those for the quantum relative entropy cone
(especially when :math:`\mathcal{G}` is the idenity map), so it is
recommended to use the quantum key distribution cone where possible.
"""
def __init__(self, G_info, Z_info, iscomplex=False):
self._process_G_info(G_info)
self._process_Z_info(Z_info)
self.iscomplex = iscomplex
self.nu = 1 + self.n # Barrier parameter
if iscomplex:
self.vn = self.n * self.n
self.vm = self.m * self.m
self.dim = [1, 2 * self.n * self.n]
self.type = ["r", "h"]
self.dtype = np.complex128
else:
self.vn = self.n * (self.n + 1) // 2
self.vm = self.m * (self.m + 1) // 2
self.dim = [1, self.n * self.n]
self.type = ["r", "s"]
self.dtype = np.float64
# Facial reduction on G(X) and Z(G(X))
ZK_list_raw = [[K[Z, :] for K in self.K_list_raw] for Z in self.Z_idxs]
self.K_list_blk = [facial_reduction(self.K_list_raw)]
self.ZK_list_blk = [facial_reduction(ZK) for ZK in ZK_list_raw]
self.Nk = [K_list[0].shape[0] for K_list in self.K_list_blk]
self.Nzk = [ZK_list[0].shape[0] for ZK_list in self.ZK_list_blk]
# Update flags
self.feas_updated = False
self.grad_updated = False
self.hess_aux_updated = False
self.invhess_aux_updated = False
self.dder3_aux_updated = False
self.invhess_aux_aux_updated = False
self.congr_aux_updated = False
self.hess_congr_aux_updated = False
self.invhess_congr_aux_updated = False
if self.G_is_Id:
self.precompute_computational_basis(self.m)
self.F2C_op = get_full_to_compact_op(self.m, iscomplex)
else:
self.precompute_computational_basis(self.n)
self.F2C_op = get_full_to_compact_op(self.n, iscomplex)
return
def _process_G_info(self, G_info):
if isinstance(G_info, int):
# Define G(X) as the identity map
self.n = n = G_info # Input dimension of X
self.N = N = G_info # Output dimension of G(X)
self.K_list_raw = [np.eye(self.N)]
self.G_is_Id = True
else:
# Define G(X) using given Kraus operators
self.n = n = G_info[0].shape[1] # Input dimension of X
self.N = N = G_info[0].shape[0] # Output dimension of G(X)
assert all([Ki.shape == (N, n) for Ki in G_info]), "Kraus " \
"operators specified by G_info must have the same dimensions." # fmt: skip
self.K_list_raw = G_info
self.G_is_Id = len(G_info) == 1 and n == N \
and np.linalg.norm(G_info[0] - np.eye(n)) < 1e-10 # fmt: skip
def _process_Z_info(self, Z_info):
N = self.N
if isinstance(Z_info, int):
# Define block structure for Z(X) with r x r blocks of size m x m
self.r = r = Z_info
self.m = m = N // r
assert m * r == N, "Number of blocks r specified by Z_info " \
"must be an integer factor of the dimension of G(X)." # fmt: skip
self.Z_list_raw = [np.zeros((N, N)) for _ in range(r)]
for k in range(r):
range_k = np.arange(k * m, (k + 1) * m)
self.Z_list_raw[k][range_k, range_k] = 1.0
self.Z_idxs = [np.arange(i * m, (i + 1) * m) for i in range(r)]
elif isinstance(Z_info, tuple):
# Define custom block structure for a given subsystem structure
(dims, sys) = Z_info
if isinstance(sys, int):
sys = [sys]
if isinstance(sys, tuple) or isinstance(sys, set):
sys = list(sys)
assert np.prod(dims) == N, "Total dimension of subsystems must equal " \
"the dimension of G(X)." # fmt: skip
assert all([sys_k < len(dims) for sys_k in sys]), "Invalid subsystems " \
"specified, exceeds total number of dimensions provided." # fmt: skip
self.r = np.prod([dims[k] for k in sys])
self.m = N // self.r
idxs = np.meshgrid(*[range(dims[k]) for k in sys])
idxs = list(np.array(idxs).reshape(len(sys), -1).T)
self.Z_list_raw = []
for i in range(len(idxs)):
Z_i = np.array([1])
counter = 0
for k, dim_k in enumerate(dims):
if k in sys:
Z_ik = np.zeros(dim_k)
Z_ik[idxs[i][counter]] = 1
Z_i = np.kron(Z_i, Z_ik)
counter += 1
else:
Z_i = np.kron(Z_i, np.ones(dim_k))
self.Z_list_raw += [np.diag(Z_i)]
self.Z_idxs = [np.where(Z)[0] for Z in self.Z_list_raw]
else:
# Define Z(X) using given Kraus operators
self.r = len(Z_info)
self.m = self.N // self.r
assert all([Zi.shape == (N, N) for Zi in Z_info]), "Kraus operators " \
"specified by Z_info must have the same dimensions." # fmt: skip
assert not any([np.any(Zi - np.diag(np.diag(Zi))) for Zi in Z_info]), \
"Kraus operators specified by Z_info must be diagonal." # fmt: skip
assert all([all([x in {0, 1} for x in np.diag(Zi)]) for Zi in Z_info]), \
"Diagonal elements of Kraus operators specified by Z_info must only " \
"be either 0s or 1s." # fmt: skip
assert np.all(sum(Z_info) == np.eye(N)), "Kraus operators specified by " \
"Z_info must satisfy ZiZj = 0 for all i =/= j." # fmt: skip
self.Z_list_raw = Z_info
self.Z_idxs = [np.where(Z)[0] for Z in Z_info]
def get_iscomplex(self):
return self.iscomplex
def get_init_point(self, out):
from qics.quantum import entropy
eye = np.eye(self.n)
GI_blk = [apply_kraus(eye, K_list) for K_list in self.K_list_blk]
ZGI_blk = [apply_kraus(eye, ZK_list) for ZK_list in self.ZK_list_blk]
entr_GI = sum([entropy(KK) for KK in GI_blk])
entr_ZGI = sum([entropy(ZKKZ) for ZKKZ in ZGI_blk])
f0 = -entr_GI + entr_ZGI
t0 = f0 / 2 + np.sqrt(1 + f0 * f0 / 4)
point = [np.array([[t0]]), np.eye(self.n, dtype=self.dtype)]
self.set_point(point, point)
out[0][:] = point[0]
out[1][:] = point[1]
return out
def get_feas(self):
if self.feas_updated:
return self.feas
self.feas_updated = True
(self.t, self.X) = (t, X) = self.primal
# Check that X is positive definite
self.Dx, self.Ux = np.linalg.eigh(X)
if any(self.Dx <= 0):
self.feas = False
return self.feas
# Check that G(X) is positive definite
# Note that G(X) should be positive definite if X is, but check
# just to be sure that we can safely take logarithms of G(X).
GX_blk = [apply_kraus(X, K_list) for K_list in self.K_list_blk]
DUgx_blk = [np.linalg.eigh(GX) for GX in GX_blk]
self.Dgx_blk = Dgx_blk = [DUgx[0] for DUgx in DUgx_blk]
self.Ugx_blk = [DUgx[1] for DUgx in DUgx_blk]
if any([any(Dgx <= 0) for Dgx in self.Dgx_blk]):
self.feas = False
return self.feas
# Check that Z(G(X)) is positive definite
# Note that Z(G(X)) should be positive definite if X is, but check
# just to be sure that we can safely take logarithms of Z(G(X)).
ZGX_blk = [apply_kraus(X, ZG_list) for ZG_list in self.ZK_list_blk]
DUzgx_blk = [np.linalg.eigh(ZGX) for ZGX in ZGX_blk]
self.Dzgx_blk = Dzgx_blk = [DUzkx[0] for DUzkx in DUzgx_blk]
self.Uzgx_blk = [DUzkx[1] for DUzkx in DUzgx_blk]
if any([any(Dzkx <= 0) for Dzkx in self.Dzgx_blk]):
self.feas = False
return self.feas
# Check that t > -S(G(X)) + S(Z(G(X)))
self.log_Dgx_blk = log_Dgx_blk = [np.log(D) for D in self.Dgx_blk]
self.log_Dzgx_blk = log_Dzgx_blk = [np.log(D) for D in self.Dzgx_blk]
entr_GX = [inp(D, log_D) for (D, log_D) in zip(Dgx_blk, log_Dgx_blk)]
entr_ZGX = [inp(D, log_D) for (D, log_D) in zip(Dzgx_blk, log_Dzgx_blk)]
self.z = t[0, 0] - (sum(entr_GX) - sum(entr_ZGX))
self.feas = self.z > 0
return self.feas
def get_val(self):
assert self.feas_updated
return -np.log(self.z) - np.sum(np.log(self.Dx))
def update_grad(self):
assert self.feas_updated
assert not self.grad_updated
# Compute gradients of quantum relative entropy
# DPhi(X) = G'[log(G(X)) + I] - G'Z'[log(G(Z(X))) + I]
log_GX_blk = [
(U * logD) @ U.conj().T + np.eye(n)
for (U, logD, n) in zip(self.Ugx_blk, self.log_Dgx_blk, self.Nk)
]
log_ZGX_blk = [
(U * logD) @ U.conj().T + np.eye(n)
for (U, logD, n) in zip(self.Uzgx_blk, self.log_Dzgx_blk, self.Nzk)
]
G_log_GX = sum(
[
apply_kraus(logX, K_list, adjoint=True)
for (logX, K_list) in zip(log_GX_blk, self.K_list_blk)
]
)
ZG_log_ZGX = sum(
[
apply_kraus(logX, ZK_list, adjoint=True)
for (logX, ZK_list) in zip(log_ZGX_blk, self.ZK_list_blk)
]
)
self.DPhi = G_log_GX - ZG_log_ZGX
# Compute X^-1
inv_Dx = np.reciprocal(self.Dx)
inv_X_rt2 = self.Ux * np.sqrt(inv_Dx)
self.inv_X = inv_X_rt2 @ inv_X_rt2.conj().T
# Compute gradient of barrier function
self.zi = np.reciprocal(self.z)
self.grad = [-self.zi, self.zi * self.DPhi - self.inv_X]
self.grad_updated = True
def update_hessprod_aux(self):
assert not self.hess_aux_updated
assert self.grad_updated
self.zi2 = self.zi * self.zi
self.D1gx_log_blk = [
D1_log(D, log_D) for (D, log_D) in zip(self.Dgx_blk, self.log_Dgx_blk)
]
self.D1zgx_log_blk = [
D1_log(D, log_D) for (D, log_D) in zip(self.Dzgx_blk, self.log_Dzgx_blk)
]
if self.G_is_Id:
D1x_inv = np.reciprocal(np.outer(self.Dx, self.Dx))
self.D1x_comb = self.zi * self.D1gx_log_blk[0] + D1x_inv
self.hess_aux_updated = True
return
def hess_prod_ip(self, out, H):
assert self.grad_updated
if not self.hess_aux_updated:
self.update_hessprod_aux()
(Ht, Hx) = H
K_list_blk, ZK_list_blk = self.K_list_blk, self.ZK_list_blk
Ugx_blk, Uzgx_blk = self.Ugx_blk, self.Uzgx_blk
D1gx_log_blk, D1zgx_log_blk = self.D1gx_log_blk, self.D1zgx_log_blk
GH_blk = [apply_kraus(Hx, K_list) for K_list in K_list_blk]
ZGH_blk = [apply_kraus(Hx, ZK_list) for ZK_list in ZK_list_blk]
UGHU_blk = [U.conj().T @ H @ U for (H, U) in zip(GH_blk, Ugx_blk)]
UZGHU_blk = [U.conj().T @ H @ U for (H, U) in zip(ZGH_blk, Uzgx_blk)]
# Hessian products for quantum key distribution
# D2Phi(X)[H]
# = G'(Ugx [log^[1](Dgx) .* (Ugx' G(H) Ugx)] Ugx')
# - G'(Z'(Uzgx [log^[1](Dzgx) .* (Uzgx' Z(G(H)) Uzgx)] Uzgx'))
D2PhiH = np.zeros_like(self.X)
# First term, i.e., G'(Ugx [log^[1](Dgx) .* (Ugx' G(H) Ugx)] Ugx')
for k, (U_k, D1_k) in enumerate(zip(Ugx_blk, D1gx_log_blk)):
temp = U_k @ (D1_k * UGHU_blk[k]) @ U_k.conj().T
D2PhiH += apply_kraus(temp, K_list_blk[k], adjoint=True)
# Second term, i.e., -G'(Z'(Uzgx [ ... ] Uzgx'))
for k, (U_k, D1_k) in enumerate(zip(Uzgx_blk, D1zgx_log_blk)):
temp = U_k @ (D1_k * UZGHU_blk[k]) @ U_k.conj().T
D2PhiH -= apply_kraus(temp, ZK_list_blk[k], adjoint=True)
# ======================================================================
# Hessian products with respect to t
# ======================================================================
# D2_t F(t, X)[Ht, Hx] = (Ht - D_X Phi(X)[Hx]) / z^2
out[0][:] = (Ht - inp(self.DPhi, Hx)) * self.zi2
# ======================================================================
# Hessian products with respect to X
# ======================================================================
# D2_X F(t, X)[Ht, Hx] = -D2_t F(t, X)[Ht, Hx] * DPhi(X)
# + D2Phi(X)[Hx] / z + X^-1 Hx X^-1
out_X = -out[0] * self.DPhi
out_X += self.zi * D2PhiH
out_X += self.inv_X @ Hx @ self.inv_X
out[1][:] = (out_X + out_X.conj().T) * 0.5
return out
def hess_congr(self, A):
assert self.grad_updated
if not self.hess_aux_updated:
self.update_hessprod_aux()
if not self.congr_aux_updated:
self.congr_aux(A)
if not self.hess_congr_aux_updated:
self.update_hess_congr_aux(A)
p = A.shape[0]
lhs = np.zeros((p, sum(self.dim)))
Z_idxs = self.Z_idxs
K_list_blk, ZK_list_blk = self.K_list_blk, self.ZK_list_blk
Ugx_blk, Uzgx_blk = self.Ugx_blk, self.Uzgx_blk
D1gx_log_blk, D1zgx_log_blk = self.D1gx_log_blk, self.D1zgx_log_blk
Work0, Work1 = self.Work0, self.Work1
Work2, Work3 = self.Work2, self.Work3
# ======================================================================
# Hessian products with respect to t
# ======================================================================
# D2_t F(t, X)[Ht, Hx] = (Ht - D_X Phi(X)[Hx]) / z^2
Ax_vec = self.Ax.view(np.float64).reshape((p, 1, -1))
DPhi_vec = self.DPhi.view(np.float64).reshape((-1, 1))
outt = self.At - (Ax_vec @ DPhi_vec).ravel()
outt *= self.zi2
lhs[:, 0] = outt
# ======================================================================
# Hessian products with respect to X
# ======================================================================
# Hessian products for quantum key distribution
# D2Phi(X)[H]
# = G'(Ugx [log^[1](Dgx) .* (Ugx' G(H) Ugx)] Ugx')
# - G'(Z'(Uzgx [log^[1](Dzgx) .* (Uzgx' Z(G(H)) Uzgx)] Uzgx'))
if self.G_is_Id:
work1, work2, work3 = self.work1, self.work2, self.work2
# If G(X) = X, then we can do some more efficient slicing operations
# Compute second term, i.e., -Z'(Uzx [log^[1](Dzx) .* ... ] Uzx')
Work0 *= 0
for Z_idxs_k, U_k, D1_log_k in zip(Z_idxs, Uzgx_blk, D1zgx_log_blk):
# Apply Z(Ax), i.e., extract k-th submatrix from Ax
temp = self.Ax[np.ix_(np.arange(p), Z_idxs_k, Z_idxs_k)]
# Compute Uzx [log^[1](Dzx) .* (Uzgx' Z(Ax) Uzx)] Uzx'
congr_multi(work2, U_k.conj().T, temp, work3)
work2 *= D1_log_k * self.zi
congr_multi(work1, U_k, work2, work3)
# Apply Z'( ... ), i.e., place result in k-th submatrix
Work0[np.ix_(np.arange(p), Z_idxs_k, Z_idxs_k)] = work1
# Compute first term, i.e., Ux [log^[1](Dx) .* (Ux' G(H) Ux)] Ux'
# Also combine this with 1/z ( ... ) + X^-1 Ax X^-1
congr_multi(Work2, self.Ux.conj().T, self.Ax, Work3)
self.Work2 *= self.D1x_comb
congr_multi(Work1, self.Ux, Work2, Work3)
# Subtract two terms to get D2Phi(X)[H]
Work1 -= Work0
else:
Work1 *= 0
# Get first term, i.e., G'(Ugx [log^[1](Dgx) .* ... ] Ugx')
for k in range(len(K_list_blk)):
worka, workb = self.Work4[k], self.Work4b[k]
workc, workd = self.Work5[k], self.Work5b[k]
workc *= 0
KU_list = [K.conj().T @ Ugx_blk[k] for K in K_list_blk[k]]
# Apply Ugx' G(Ax) Ugx
for KU in KU_list:
congr_multi(workd, KU.conj().T, self.Ax, work=workb)
workc += workd
# Apply log^[1](Dgx) .* ( ... )
workc *= D1gx_log_blk[k] * self.zi
# Apply G'(Ugx [ ... ] Ugx')
for KU in KU_list:
congr_multi(Work0, KU, workc, work=worka)
Work1 += Work0
# Get second term, i.e., G'(Z'(Uzgx [log^[1](Dzgx) .* ... ] Uzgx'))
for k in range(len(ZK_list_blk)):
worka, workb = self.Work6[k], self.Work6b[k]
workc, workd = self.Work7[k], self.Work7b[k]
workc *= 0
KU_list = [K.conj().T @ Uzgx_blk[k] for K in ZK_list_blk[k]]
# Apply Uzgx' Z(G(H)) Uzgx
for KU in KU_list:
congr_multi(workd, KU.conj().T, self.Ax, work=workb)
workc += workd
# Apply log^[1](Dzgx) .* ( ... )
workc *= D1zgx_log_blk[k] * self.zi
# Apply G'(Z'(Uzgx [ ... ] Uzgx'))
for KU in KU_list:
congr_multi(Work0, KU, workc, work=worka)
Work1 -= Work0
# Compute X^-1 Ax X^-1
congr_multi(Work0, self.inv_X, self.Ax, Work2)
Work1 += Work0
# Hessian product of barrier function
# D2_X F(t, X)[Ht, Hx] = -D2_t F(t, X)[Ht, Hx] * DPhi(X)
# + D2Phi(X)[Hx] / z + X^-1 Hx X^-1
np.outer(outt, self.DPhi, out=Work0.reshape(p, -1))
Work1 -= Work0
lhs[:, 1:] = Work1.reshape((p, -1)).view(np.float64)
# Multiply A (H A')
return dense_dot_x(lhs, A.T)
def invhess_prod_ip(self, out, H):
assert self.grad_updated
if not self.hess_aux_updated:
self.update_hessprod_aux()
if not self.invhess_aux_updated:
self.update_invhessprod_aux()
(Ht, Hx) = H
Wx = Hx + Ht * self.DPhi
# ======================================================================
# Inverse Hessian products with respect to X
# ======================================================================
# Compute X = M \ Wx, where we solve this using two different strategies
if self.G_is_Id:
# If G(X) = X, then solve M using the matrix inversion lemma
work = np.zeros((self.r * self.vm))
# Apply D2S(X)^-1 = (Ux⊗Ux) (1/z log + inv)^[1](Dx)^-1 (Ux'⊗Ux')
temp = self.Ux.conj().T @ Wx @ self.Ux
out_X = self.Ux @ (self.D1x_comb_inv * temp) @ self.Ux.conj().T
# Apply Z( ... ), i.e., extract k-th submatrix then get compact
# vectorization of this submatrix
for k, Z_idxs_k in enumerate(self.Z_idxs):
temp_k = out_X[np.ix_(Z_idxs_k, Z_idxs_k)]
temp_vec = temp_k.view(np.float64).reshape(-1)
work[k * self.vm : (k + 1) * self.vm] = self.F2C_op @ temp_vec
work *= -1
# Solve linear system N \ ( ... )
work = cho_solve(self.schur_fact, work)
# Apply Z'( ... ), i.e., recover submatrix from compact
# vectorization then place in k-th submatrix
Work3 = np.zeros((self.n, self.n), dtype=self.dtype)
for k, Z_idxs_k in enumerate(self.Z_idxs):
work_k = work[k * self.vm : (k + 1) * self.vm]
work_k = self.F2C_op.T @ work_k
work_k = work_k.view(self.dtype).reshape((self.m, self.m))
Work3[np.ix_(Z_idxs_k, Z_idxs_k)] = work_k
# Apply D2S(X)^-1 = (Ux ⊗ Ux) log^[1](Dx) (Ux' ⊗ Ux') and subtract
# from (D2S(X)^-1 Wx) to get X
temp = self.Ux.conj().T @ Work3 @ self.Ux
out_X -= self.Ux @ (self.D1x_comb_inv * temp) @ self.Ux.conj().T
else:
# Otherwise, directly build M and Cholesky factor to solve M \ Wx
# Convert matrices to compact real vectors
temp_vec = Wx.view(np.float64).reshape((-1, 1))
temp_vec = self.F2C_op @ temp_vec
# Solve system
temp_vec = cho_solve(self.hess_fact, temp_vec)
# Expand compact real vectors back into full matrices
temp_vec = self.F2C_op.T @ temp_vec
out_X = temp_vec.T.view(self.dtype).reshape((self.n, self.n))
out[1][:] = (out_X + out_X.conj().T) * 0.5
# ======================================================================
# Inverse Hessian products with respect to t
# ======================================================================
# Compute t = z^2 Ht + <DPhi(X), X>
out[0][:] = self.z2 * Ht + inp(self.DPhi, out[1])
return out
def invhess_congr(self, A):
assert self.grad_updated
if not self.hess_aux_updated:
self.update_hessprod_aux()
if not self.invhess_aux_updated:
self.update_invhessprod_aux()
if not self.congr_aux_updated:
self.congr_aux(A)
if not self.invhess_congr_aux_updated:
self.update_invhess_congr_aux(A)
# The inverse Hessian product applied on (Ht, Hx, Hy) for the OPT
# barrier is
# X = M \ Wx
# t = z^2 Ht + <DPhi(X), X>
# where Wx = Hx + Ht DPhi(X) and M = 1/z D2Phi + X^1 ⊗ X^-1
if self.G_is_Id:
# When G(X) = X, we can write the expression for M as
# M = 1/z D2S(X) - 1/z Z' D2S(Z(X)) Z + X^-1 ⊗ X^-1
# = (Ux ⊗ Ux) (1/z log + inv)^[1](Dx) (Ux' ⊗ Ux')
# - 1/z Z' (Uzx ⊗ Uzx) log^[1](Dzx) (Uzx' ⊗ Uzx') Z
# Treating [Z' D2S(Z(X)) Z] as a low-rank perturbation of D2S(X), we
# can solve linear systems with M by using the matrix inversion
# lemma
# X = [D2S(X)^-1 - D2S(X)^-1 Z' N^-1 Z D2S(X)^-1] Wx
# where
# N = 1/z (Uzx ⊗ Uzx) [log^[1](Dzx)]^-1 (Uzx' ⊗ Uzx')
# - Z (Ux ⊗ Ux) [(1/z log + inv)^[1](Dx)]^-1 (Ux' ⊗ Ux') Z'
p = A.shape[0]
lhs = np.zeros((p, sum(self.dim)))
work = np.zeros((self.r * self.vm, p))
Work0, Work1 = self.Work0, self.Work1
Work2, Work3 = self.Work2, self.Work3
# ==================================================================
# Inverse Hessian products with respect to X
# ==================================================================
# Compute Wx = Hx + Ht DPhi(X)
np.outer(self.At, self.DPhi, out=Work2.reshape((p, -1)))
np.add(self.Ax, Work2, out=Work0)
# Apply D2S(X)^-1 = (Ux⊗Ux) (1/z log + inv)^[1](Dx)^-1 (Ux'⊗Ux')
congr_multi(Work2, self.Ux.conj().T, Work0, Work3)
Work2 *= self.D1x_comb_inv
congr_multi(Work0, self.Ux, Work2, Work3)
# Apply Z( ... ), i.e., extract k-th submatrix then get compact
# vectorization of this submatrix
for k, Z_idxs_k in enumerate(self.Z_idxs):
temp = Work0[np.ix_(np.arange(p), Z_idxs_k, Z_idxs_k)]
temp = temp.view(np.float64).reshape((p, -1)).T
temp = x_dot_dense(self.F2C_op, temp)
work[k * self.vm : (k + 1) * self.vm] = temp
work *= -1
# Solve linear system N \ ( ... )
work = cho_solve(self.schur_fact, work)
# Apply Z'( ... ), i.e., recover submatrix from compact
# vectorization then place in k-th submatrix
Work1.fill(0.0)
for k, Z_idxs_k in enumerate(self.Z_idxs):
work_k = work[k * self.vm : (k + 1) * self.vm]
work_k = x_dot_dense(self.F2C_op.T, work_k)
work_k = work_k.T.view(self.dtype).reshape(p, self.m, self.m)
Work1[np.ix_(np.arange(p), Z_idxs_k, Z_idxs_k)] = work_k
# Apply D2S(X)^-1 = (Ux⊗Ux) (1/z log + inv)^[1](Dx)^-1 (Ux'⊗Ux')
congr_multi(Work2, self.Ux.conj().T, Work1, Work3)
Work2 *= self.D1x_comb_inv
congr_multi(Work1, self.Ux, Work2, Work3)
# Subtract previous expression from (D2S(X)^-1 Wx) to get X
Work0 -= Work1
out_X = Work0.reshape((p, -1)).view(np.float64)
lhs[:, 1:] = out_X
# ==================================================================
# Inverse Hessian products with respect to t
# ==================================================================
# Compute t = z^2 Ht + <DPhi(X), X>
DPhi_vec = self.DPhi.view(np.float64).reshape((-1, 1))
out_t = self.z2 * self.At
out_t += (out_X @ DPhi_vec).ravel()
lhs[:, 0] = out_t
# Multiply A (H A')
return dense_dot_x(lhs, A.T)
else:
# Otherwise, we will directly build M and Cholesky factor it to
# solve X = M \ Wx
# ==================================================================
# Inverse Hessian products with respect to X
# ==================================================================
# Compute Wx = Hx + Ht DPhi(X)
np.outer(self.At, self.DPhi_cvec, out=self.work0)
self.work0 += self.Ax_cvec
# Solve system X = M \ Wx
out_X = cho_solve(self.hess_fact, self.work0.T)
# ==================================================================
# Inverse Hessian products with respect to t
# ==================================================================
# Compute t = z^2 Ht + <DPhi(X), X>
out_t = self.z2 * self.At
out_t += (out_X.T @ self.DPhi_cvec).ravel()
# Multiply A (H A')
out = np.outer(out_t, self.At)
out += x_dot_dense(self.Ax_cvec, out_X)
return out
def third_dir_deriv_axpy(self, out, H, a=True):
assert self.grad_updated
if not self.hess_aux_updated:
self.update_hessprod_aux()
if not self.dder3_aux_updated:
self.update_dder3_aux()
(Ht, Hx) = H
K_list_blk, ZK_list_blk = self.K_list_blk, self.ZK_list_blk
Ugx_blk, Uzgx_blk = self.Ugx_blk, self.Uzgx_blk
D1gx_log_blk, D1zgx_log_blk = self.D1gx_log_blk, self.D1zgx_log_blk
D2gx_log_blk, D2zgx_log_blk = self.D2gx_log_blk, self.D2zgx_log_blk
GH_blk = [apply_kraus(Hx, K_list) for K_list in K_list_blk]
ZGH_blk = [apply_kraus(Hx, ZK_list) for ZK_list in ZK_list_blk]
UGHU_blk = [U.conj().T @ H @ U for (H, U) in zip(GH_blk, Ugx_blk)]
UZGHU_blk = [U.conj().T @ H @ U for (H, U) in zip(ZGH_blk, Uzgx_blk)]
# Quantum key distribution Hessians
D2PhiH = np.zeros_like(self.X)
# Hessians of S(G(X))
for k, (U_k, D1_k) in enumerate(zip(Ugx_blk, D1gx_log_blk)):
temp = U_k @ (D1_k * UGHU_blk[k]) @ U_k.conj().T
D2PhiH += apply_kraus(temp, K_list_blk[k], adjoint=True)
# Hessians of S(Z(G(X)))
for k, (U_k, D1_k) in enumerate(zip(Uzgx_blk, D1zgx_log_blk)):
temp = U_k @ (D1_k * UZGHU_blk[k]) @ U_k.conj().T
D2PhiH -= apply_kraus(temp, ZK_list_blk[k], adjoint=True)
# Quantum key distribution third order derivatives
D3PhiHH = np.zeros_like(self.X)
# Third order derivatives of S(G(X))
for k, (U_k, D2_k) in enumerate(zip(Ugx_blk, D2gx_log_blk)):
temp = scnd_frechet(D2_k * UGHU_blk[k], UGHU_blk[k], U=U_k)
D3PhiHH += apply_kraus(temp, K_list_blk[k], adjoint=True)
# Third order derivatives of S(Z(G(X)))
for k, (U_k, D2_k) in enumerate(zip(Uzgx_blk, D2zgx_log_blk)):
temp = scnd_frechet(D2_k * UZGHU_blk[k], UZGHU_blk[k], U=U_k)
D3PhiHH -= apply_kraus(temp, ZK_list_blk[k], adjoint=True)
# Third derivative of barrier
DPhiH = inp(self.DPhi, Hx)
D2PhiHH = inp(D2PhiH, Hx)
chi = Ht - DPhiH
chi2 = chi * chi
dder3_t = -2 * self.zi3 * chi2 - self.zi2 * D2PhiHH
dder3_X = -dder3_t * self.DPhi
dder3_X -= 2 * self.zi2 * chi * D2PhiH
dder3_X += self.zi * D3PhiHH
dder3_X -= 2 * self.inv_X @ Hx @ self.inv_X @ Hx @ self.inv_X
dder3_X = (dder3_X + dder3_X.conj().T) * 0.5
out[0][:] += dder3_t * a
out[1][:] += dder3_X * a
return out
# ========================================================================
# Auxilliary functions
# ========================================================================
def congr_aux(self, A):
assert not self.congr_aux_updated
iscomplex = self.iscomplex
# Get slices and views of A matrix to be used in congruence computations
if sp.sparse.issparse(A):
A = A.tocsr()
self.Ax_vec = A[:, 1:]
if sp.sparse.issparse(A):
A = A.toarray()
Ax_dense = np.ascontiguousarray(A[:, 1:])
self.At = A[:, 0]
self.Ax = np.array([vec_to_mat(Ax_k, iscomplex) for Ax_k in Ax_dense])
self.congr_aux_updated = True
def update_hess_congr_aux(self, A):
assert not self.hess_congr_aux_updated
p = A.shape[0]
self.Work0 = np.zeros_like(self.Ax, dtype=self.dtype)
self.Work1 = np.zeros_like(self.Ax, dtype=self.dtype)
self.Work2 = np.zeros_like(self.Ax, dtype=self.dtype)
self.Work3 = np.zeros_like(self.Ax, dtype=self.dtype)
if self.G_is_Id:
self.work1 = np.empty((p, self.m, self.m), dtype=self.dtype)
self.work2 = np.empty((p, self.m, self.m), dtype=self.dtype)
self.work3 = np.empty((p, self.m, self.m), dtype=self.dtype)
else:
n, Nk, Nzk, dtype = self.n, self.Nk, self.Nzk, self.dtype
self.Work4 = [np.zeros((p, n, nk), dtype=dtype) for nk in Nk]
self.Work4b = [np.zeros((p, nk, n), dtype=dtype) for nk in Nk]
self.Work5 = [np.zeros((p, nk, nk), dtype=dtype) for nk in Nk]
self.Work5b = [np.zeros((p, nk, nk), dtype=dtype) for nk in Nk]
self.Work6 = [np.zeros((p, n, nzk), dtype=dtype) for nzk in Nzk]
self.Work6b = [np.zeros((p, nzk, n), dtype=dtype) for nzk in Nzk]
self.Work7 = [np.zeros((p, nzk, nzk), dtype=dtype) for nzk in Nzk]
self.Work7b = [np.zeros((p, nzk, nzk), dtype=dtype) for nzk in Nzk]
self.hess_congr_aux_updated = True
def update_invhess_congr_aux(self, A):
assert not self.invhess_congr_aux_updated
if self.G_is_Id:
self.Work0 = np.zeros_like(self.Ax, dtype=self.dtype)
self.Work1 = np.zeros_like(self.Ax, dtype=self.dtype)
self.Work2 = np.zeros_like(self.Ax, dtype=self.dtype)
self.Work3 = np.zeros_like(self.Ax, dtype=self.dtype)
else:
if sp.sparse.issparse(A):
A = A.tocsr()
self.Ax_cvec = (self.F2C_op @ A[:, 1:].T).T
if sp.sparse.issparse(A):
self.Ax_cvec = self.Ax_cvec.tocoo()
self.work0 = np.zeros(self.Ax_cvec.shape)
self.work1 = np.zeros(self.Ax_cvec.shape)
self.invhess_congr_aux_updated = True
def update_invhessprod_aux_aux(self):
assert not self.invhess_aux_aux_updated
if self.G_is_Id:
self.work6 = np.zeros((self.vm, self.m, self.m), dtype=self.dtype)
self.work7 = np.zeros((self.vm, self.m, self.m), dtype=self.dtype)
self.work8 = np.zeros((self.vm, self.m, self.m), dtype=self.dtype)
# Computational basis for symmetric/Hermitian matrices
rt2 = np.sqrt(0.5)
n, r, m, vm = self.n, self.r, self.m, self.vm
self.E_blk = np.zeros((r, vm, n, n), dtype=self.dtype)
for b, Z_idxs_k in enumerate(self.Z_idxs):
k = 0
for j_subblk in range(m):
j = Z_idxs_k[j_subblk]
for i_subblk in range(j_subblk):
i = Z_idxs_k[i_subblk]
self.E_blk[b, k, i, j] = rt2
self.E_blk[b, k, j, i] = rt2
k += 1
if self.iscomplex:
self.E_blk[b, k, i, j] = rt2 * 1j
self.E_blk[b, k, j, i] = rt2 * -1j
k += 1
self.E_blk[b, k, j, j] = 1.0
k += 1
self.E_blk = self.E_blk.reshape((-1, self.n, self.n))
self.Work6 = np.zeros((r * vm, n, n), dtype=self.dtype)
self.Work7 = np.zeros((r * vm, n, n), dtype=self.dtype)
self.Work8 = np.zeros((r * vm, n, n), dtype=self.dtype)
else:
dtype = self.dtype
n, vn, Nk, Nzk = self.n, self.vn, self.Nk, self.Nzk
self.work2 = [np.zeros((vn, n, nk), dtype=dtype) for nk in Nk]
self.work2b = [np.zeros((vn, nk, n), dtype=dtype) for nk in Nk]
self.work3 = [np.zeros((vn, nk, nk), dtype=dtype) for nk in Nk]
self.work3b = [np.zeros((vn, nk, nk), dtype=dtype) for nk in Nk]
self.work4 = [np.zeros((vn, n, nzk), dtype=dtype) for nzk in Nzk]
self.work4b = [np.zeros((vn, nzk, n), dtype=dtype) for nzk in Nzk]
self.work5 = [np.zeros((vn, nzk, nzk), dtype=dtype) for nzk in Nzk]
self.work5b = [np.zeros((vn, nzk, nzk), dtype=dtype) for nzk in Nzk]
self.work6 = np.zeros((vn, n, n), dtype=dtype)
self.work7 = np.zeros((vn, n, n), dtype=dtype)
self.work8 = np.zeros((vn, n, n), dtype=dtype)
self.invhess_aux_aux_updated = True
def update_invhessprod_aux(self):
assert not self.invhess_aux_updated
assert self.grad_updated
assert self.hess_aux_updated
if not self.invhess_aux_aux_updated:
self.update_invhessprod_aux_aux()
self.z2 = self.z * self.z
if self.G_is_Id:
# Precompute and factorize the matrix
# N = z (Uzx ⊗ Uzx) [log^[1](Dzx)]^-1 (Uzx' ⊗ Uzx')
# - Z (Ux ⊗ Ux) [(1/z log + inv)^[1](Dx)]^-1 (Ux' ⊗ Ux') Z'
# which we will need to solve linear systems with the Hessian of
# our barrier function
self.D1x_comb_inv = np.reciprocal(self.D1x_comb)
r, vm = self.r, self.vm
Uzgx_blk, D1zgx_log_blk = self.Uzgx_blk, self.D1zgx_log_blk
work6, work7, work8 = self.work6, self.work7, self.work8
Work6, Work7, Work8 = self.Work6, self.Work7, self.Work8
self.schur = np.zeros((r * vm, r * vm))
# ==================================================================
# Get first term, i.e., 1/z (Uzx ⊗ Uzx) ... (Uzx' ⊗ Uzx')
# ==================================================================
for k, (U, D1_log) in enumerate(zip(Uzgx_blk, D1zgx_log_blk)):
# Begin with (Uzx' ⊗ Uzx')
congr_multi(work8, U.conj().T, self.E, work=work7)
# Apply z [log^[1](Dzx)]^-1
work8 *= self.z * np.reciprocal(D1_log)
# Apply (Uzx ⊗ Uzx)
congr_multi(work6, U, work8, work=work7)
work = work6.view(np.float64).reshape((vm, -1))
work = x_dot_dense(self.F2C_op, work.T)
self.schur[k * vm : (k + 1) * vm, k * vm : (k + 1) * vm] = work
# ==================================================================
# Get second term, i.e., Z (Ux ⊗ Ux) ... (Ux' ⊗ Ux') Z'
# ==================================================================
# Begin with (Ux' ⊗ Ux') Z'
congr_multi(Work8, self.Ux.conj().T, self.E_blk, work=Work7)
# Apply [(1/z log + inv)^[1](Dx)]^-1
Work8 *= self.D1x_comb_inv
# Apply (Ux ⊗ Ux)
congr_multi(Work6, self.Ux, Work8, work=Work7)
# Apply Z, i.e., extract submatrices
for k, Z_idxs_k in enumerate(self.Z_idxs):
temp = Work6[np.ix_(np.arange(r * vm), Z_idxs_k, Z_idxs_k)]
temp = temp.view(np.float64).reshape((r * vm, -1))
temp = x_dot_dense(self.F2C_op, temp.T).T
self.schur[:, k * vm : (k + 1) * vm] -= temp
# Subtract terms to obtain N then Cholesky factor
self.schur_fact = cho_fact(self.schur)
else:
# Precompute and factorize the matrix
# M = 1/z G' (Ugx ⊗ Ugx) log^[1](Dgx) (Ugx' ⊗ Ugx') G
# - 1/z G'Z' (Uzgx ⊗ Uzgx) log^[1](Dzgx) (Uzgx' ⊗ Uzgx') ZG
# + X^-1 ⊗ X^-1
# which we will need to solve linear systems with the Hessian of
# our barrier function
K_list_blk, ZK_list_blk = self.K_list_blk, self.ZK_list_blk
Ugx_blk, Uzgx_blk = self.Ugx_blk, self.Uzgx_blk
D1gx_log_blk, D1zgx_log_blk = self.D1gx_log_blk, self.D1zgx_log_blk
work7, work8 = self.work7, self.work8
DPhi_vec = self.DPhi.view(np.float64).reshape(-1, 1)
self.DPhi_cvec = self.F2C_op @ DPhi_vec
# ==================================================================
# Get third term, i.e., X^-1 ⊗ X^-1
# ==================================================================
congr_multi(work8, self.inv_X, self.E, work=work7)
# ==================================================================
# Get first term, i.e., 1/z G' (Ugx ⊗ Ugx) ... (Ugx' ⊗ Ugx') G
# ==================================================================
for k in range(len(K_list_blk)):
worka, workb = self.work2[k], self.work2b[k]
workc, workd = self.work3[k], self.work3b[k]
workc *= 0
KU_list = [K.conj().T @ Ugx_blk[k] for K in K_list_blk[k]]
# Apply Ugx' G(Ax) Ugx
for KU in KU_list:
congr_multi(workd, KU.conj().T, self.E, work=workb)
workc += workd
# Apply log^[1](Dgx) .* ( ... )
workc *= D1gx_log_blk[k] * self.zi
# Apply G'(Ugx [ ... ] Ugx')
for KU in KU_list:
congr_multi(work7, KU, workc, work=worka)
work8 += work7
# ==================================================================
# Get second term, i.e., 1/z G'Z' [ ... ] ZG
# ==================================================================
for k in range(len(ZK_list_blk)):
worka, workb = self.work4[k], self.work4b[k]
workc, workd = self.work5[k], self.work5b[k]
workc *= 0
KU_list = [K.conj().T @ Uzgx_blk[k] for K in ZK_list_blk[k]]
# Apply Uzgx' Z(G(Ax)) Uzgx
for KU in KU_list:
congr_multi(workd, KU.conj().T, self.E, work=workb)
workc += workd
# Apply log^[1](Dzgx) .* ( ... )
workc *= D1zgx_log_blk[k] * self.zi
# Apply G'Z'((Uzgx [ ... ] Uzgx'))
for KU in KU_list:
congr_multi(work7, KU, workc, work=worka)
work8 -= work7
# Get Hessian and factorize
work = work8.view(np.float64).reshape((self.vn, -1))
self.hess = x_dot_dense(self.F2C_op, work.T)
self.hess = (self.hess + self.hess.T) * 0.5
self.hess_fact = cho_fact(self.hess)
self.invhess_aux_updated = True
def update_dder3_aux(self):
assert not self.dder3_aux_updated
assert self.hess_aux_updated
self.zi3 = self.zi2 * self.zi
self.D2gx_log_blk = [
D2_log(D, D1) for (D, D1) in zip(self.Dgx_blk, self.D1gx_log_blk)
]
self.D2zgx_log_blk = [
D2_log(D, D1) for (D, D1) in zip(self.Dzgx_blk, self.D1zgx_log_blk)
]
self.dder3_aux_updated = True
return
def facial_reduction(K_list):
# For a set of Kraus operators i.e., Σ_i K_i @ X @ K_i.T, returns a set of
# Kraus operators which preserves positive definiteness
nk = K_list[0].shape[0]
# Pass identity matrix (maximally mixed state) through the Kraus operators
KK = sum([K @ K.conj().T for K in K_list])
# Determine if output is low rank, in which case we need to perform facial
# reduction
Dkk, Ukk = np.linalg.eigh(KK)
KKnzidx = np.where(Dkk > 1e-12)[0]
nk_fr = np.size(KKnzidx)
if nk == nk_fr:
return K_list
# Perform facial reduction
Qkk = Ukk[:, KKnzidx]
K_list_fr = [Qkk.conj().T @ K for K in K_list]
return K_list_fr
def apply_kraus(x, Klist, adjoint=False):
# Compute congruence map
if adjoint:
return sum([K.conj().T @ x @ K for K in Klist])
else:
return sum([K @ x @ K.conj().T for K in Klist])
QICS